submission 81672

"""
CuTe/CUTLASS kernel for NVFP4 block-scaled GEMV - Version 6
Conservative optimizations with guaranteed correctness

Strategy: Keep V3 tile sizes but optimize the computation pattern
- Same tile configuration as V3 (128, 1, 64) - proven correct
- Optimize only the inner computation loop
- Focus on reducing memory stalls through better scheduling
"""
import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import make_ptr
import cutlass.utils.blockscaled_layout as blockscaled_utils

# ============================================================================
# KERNEL CONFIGURATION - SAME AS V3 FOR SAFETY
# ============================================================================

TILE_SIZE_M = 128
TILE_SIZE_N = 1
TILE_SIZE_K = 64
mma_tiler_mnk = (TILE_SIZE_M, TILE_SIZE_N, TILE_SIZE_K)

INPUT_DTYPE = cutlass.Float4E2M1FN
SCALE_FACTOR_DTYPE = cutlass.Float8E4M3FN
OUTPUT_DTYPE = cutlass.Float16
ACCUMULATOR_DTYPE = cutlass.Float32

SCALE_FACTOR_BLOCK_SIZE = 16

THREADS_PER_M_DIM = 128
THREADS_PER_K_DIM = 4
TOTAL_THREADS = THREADS_PER_M_DIM * THREADS_PER_K_DIM


# ============================================================================
# DEVICE KERNEL - COMPUTATION OPTIMIZED
# ============================================================================

@cute.kernel
def blockscaled_gemv_kernel_v6(
    matrix_a_mkl: cute.Tensor,
    matrix_b_nkl: cute.Tensor,
    scale_factors_a_mkl: cute.Tensor,
    scale_factors_b_nkl: cute.Tensor,
    output_c_mnl: cute.Tensor,
):
    """
    V6 GPU kernel - optimized computation with V3's proven tile configuration.

    Changes from V3:
    - Load all 4 tensors before any conversion (better memory scheduling)
    - Minimize load/store operations in inner loop
    - Explicit accumulator to avoid repeated indexing
    """

    # Get thread/block indices
    block_idx_m, block_idx_n, block_idx_l = cute.arch.block_idx()
    thread_idx_m, thread_idx_k, _ = cute.arch.thread_idx()

    # Extract tiles from global tensors (same as V3)
    global_tile_a = cute.local_tile(
        matrix_a_mkl,
        cute.slice_(mma_tiler_mnk, (None, 0, None)),
        (None, None, None)
    )
    global_tile_scale_a = cute.local_tile(
        scale_factors_a_mkl,
        cute.slice_(mma_tiler_mnk, (None, 0, None)),
        (None, None, None)
    )
    global_tile_b = cute.local_tile(
        matrix_b_nkl,
        cute.slice_(mma_tiler_mnk, (0, None, None)),
        (None, None, None)
    )
    global_tile_scale_b = cute.local_tile(
        scale_factors_b_nkl,
        cute.slice_(mma_tiler_mnk, (0, None, None)),
        (None, None, None)
    )
    global_tile_output = cute.local_tile(
        output_c_mnl,
        cute.slice_(mma_tiler_mnk, (None, None, 0)),
        (None, None, None)
    )

    # Thread-local output preparation
    thread_output_ref = global_tile_output[thread_idx_m, None, block_idx_m, block_idx_n, block_idx_l]
    thread_output = cute.make_tensor(thread_output_ref.iterator, 1)
    partial_result = cute.zeros_like(thread_output, ACCUMULATOR_DTYPE)

    # Shared memory for reduction
    smem_allocator = cutlass.utils.SmemAllocator()
    smem_layout = cute.make_layout(
        (THREADS_PER_M_DIM, THREADS_PER_K_DIM),
        stride=(THREADS_PER_K_DIM, 1)
    )
    shared_memory_buffer = smem_allocator.allocate_tensor(
        element_type=cutlass.Float32,
        layout=smem_layout
    )

    # K-dimension loop
    num_k_tiles = global_tile_a.layout[3].shape

    for k_tile_idx in range(thread_idx_k, num_k_tiles, THREADS_PER_K_DIM, unroll_full=True):
        # Get pointers to thread-local data
        thread_a_global = global_tile_a[thread_idx_m, None, block_idx_m, k_tile_idx, block_idx_l]
        thread_b_global = global_tile_b[0, None, block_idx_n, k_tile_idx, block_idx_l]
        thread_scale_a_global = global_tile_scale_a[thread_idx_m, None, block_idx_m, k_tile_idx, block_idx_l]
        thread_scale_b_global = global_tile_scale_b[0, None, block_idx_n, k_tile_idx, block_idx_l]

        # OPTIMIZATION 1: Load all values first (batch memory operations)
        # This allows the memory controller to better schedule and overlap requests
        a_value_fp4 = thread_a_global.load()
        b_value_fp4 = thread_b_global.load()
        scale_a_value_fp8 = thread_scale_a_global.load()
        scale_b_value_fp8 = thread_scale_b_global.load()

        # Allocate registers
        reg_a_fp16 = cute.make_rmem_tensor_like(thread_a_global, OUTPUT_DTYPE)
        reg_b_fp16 = cute.make_rmem_tensor_like(thread_b_global, OUTPUT_DTYPE)
        reg_ab_product = cute.make_rmem_tensor_like(thread_a_global, OUTPUT_DTYPE)
        reg_scale_a_fp32 = cute.make_rmem_tensor_like(thread_scale_a_global, ACCUMULATOR_DTYPE)
        reg_scale_b_fp32 = cute.make_rmem_tensor_like(thread_scale_b_global, ACCUMULATOR_DTYPE)
        reg_scale_product = cute.make_rmem_tensor_like(thread_scale_a_global, ACCUMULATOR_DTYPE)

        # OPTIMIZATION 2: Convert and store in one step
        reg_a_fp16.store(a_value_fp4.to(OUTPUT_DTYPE))
        reg_b_fp16.store(b_value_fp4.to(OUTPUT_DTYPE))
        reg_scale_a_fp32.store(scale_a_value_fp8.to(ACCUMULATOR_DTYPE))
        reg_scale_b_fp32.store(scale_b_value_fp8.to(ACCUMULATOR_DTYPE))

        # Pre-compute products
        reg_ab_product.store(reg_a_fp16.load() * reg_b_fp16.load())
        reg_scale_product.store(reg_scale_a_fp32.load() * reg_scale_b_fp32.load())

        # Accumulation loop - same as V3
        for element_idx in cutlass.range_constexpr(TILE_SIZE_K):
            partial_result += reg_ab_product[element_idx] * reg_scale_product[element_idx]

    # Write partial result to shared memory
    shared_memory_buffer[(thread_idx_m, thread_idx_k)] = partial_result[0]
    cute.arch.sync_threads()

    # Final reduction (same as V3)
    if thread_idx_k == 0:
        final_output = cute.zeros_like(thread_output, ACCUMULATOR_DTYPE)
        for i in cutlass.range_constexpr(THREADS_PER_K_DIM):
            final_output += shared_memory_buffer[(thread_idx_m, i)]
        thread_output.store(final_output.to(cutlass.Float16))

    return


# ============================================================================
# HOST-SIDE JIT COMPILATION WRAPPER
# ============================================================================

@cute.jit
def gemv_kernel_launcher(
    a_ptr: cute.Pointer,
    b_ptr: cute.Pointer,
    scale_a_ptr: cute.Pointer,
    scale_b_ptr: cute.Pointer,
    output_ptr: cute.Pointer,
    problem_size: tuple,
):
    """Host-side JIT function to prepare tensors and launch GPU kernel."""
    m, _, k, l = problem_size

    tensor_a = cute.make_tensor(
        a_ptr,
        cute.make_layout(
            (m, cute.assume(k, 32), l),
            stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
        ),
    )

    N_PADDED = 128
    tensor_b = cute.make_tensor(
        b_ptr,
        cute.make_layout(
            (N_PADDED, cute.assume(k, 32), l),
            stride=(cute.assume(k, 32), 1, cute.assume(N_PADDED * k, 32)),
        ),
    )

    tensor_output = cute.make_tensor(
        output_ptr,
        cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m))
    )

    scale_a_layout = blockscaled_utils.tile_atom_to_shape_SF(
        tensor_a.shape,
        SCALE_FACTOR_BLOCK_SIZE
    )
    tensor_scale_a = cute.make_tensor(scale_a_ptr, scale_a_layout)

    scale_b_layout = blockscaled_utils.tile_atom_to_shape_SF(
        tensor_b.shape,
        SCALE_FACTOR_BLOCK_SIZE
    )
    tensor_scale_b = cute.make_tensor(scale_b_ptr, scale_b_layout)

    grid_config = (
        cute.ceil_div(tensor_output.shape[0], TILE_SIZE_M),
        1,
        tensor_output.shape[2],
    )

    blockscaled_gemv_kernel_v6(
        tensor_a,
        tensor_b,
        tensor_scale_a,
        tensor_scale_b,
        tensor_output
    ).launch(
        grid=grid_config,
        block=[THREADS_PER_M_DIM, THREADS_PER_K_DIM, 1],
        cluster=(1, 1, 1),
    )
    return


# ============================================================================
# KERNEL COMPILATION AND CACHING
# ============================================================================

_compiled_kernel_cache = None


def compile_kernel():
    """Compile the CuTe kernel and cache it for reuse."""
    global _compiled_kernel_cache

    if _compiled_kernel_cache is not None:
        return _compiled_kernel_cache

    dummy_a_ptr = make_ptr(INPUT_DTYPE, 0, cute.AddressSpace.gmem, assumed_align=16)
    dummy_b_ptr = make_ptr(INPUT_DTYPE, 0, cute.AddressSpace.gmem, assumed_align=16)
    dummy_output_ptr = make_ptr(OUTPUT_DTYPE, 0, cute.AddressSpace.gmem, assumed_align=16)
    dummy_scale_a_ptr = make_ptr(SCALE_FACTOR_DTYPE, 0, cute.AddressSpace.gmem, assumed_align=32)
    dummy_scale_b_ptr = make_ptr(SCALE_FACTOR_DTYPE, 0, cute.AddressSpace.gmem, assumed_align=32)
    dummy_problem_size = (0, 0, 0, 0)

    _compiled_kernel_cache = cute.compile(
        gemv_kernel_launcher,
        dummy_a_ptr,
        dummy_b_ptr,
        dummy_scale_a_ptr,
        dummy_scale_b_ptr,
        dummy_output_ptr,
        dummy_problem_size
    )

    return _compiled_kernel_cache


def reset_cache():
    """Reset the compiled kernel cache (useful for benchmarking)."""
    global _compiled_kernel_cache
    _compiled_kernel_cache = None


# ============================================================================
# PUBLIC API
# ============================================================================

def custom_kernel(data):
    """
    Execute the optimized block-scaled GEMV kernel.

    Args:
        data: Tuple of (a, b, sfa, sfb, sfa_permuted, sfb_permuted, c)

    Returns:
        Output tensor c with computed GEMV results
    """
    matrix_a, vector_b, _, _, scale_a_permuted, scale_b_permuted, output_c = data

    compiled_kernel = compile_kernel()

    m, k_packed, l = matrix_a.shape
    k_actual = k_packed * 2
    n = 1

    a_ptr = make_ptr(
        INPUT_DTYPE,
        matrix_a.data_ptr(),
        cute.AddressSpace.gmem,
        assumed_align=16
    )
    b_ptr = make_ptr(
        INPUT_DTYPE,
        vector_b.data_ptr(),
        cute.AddressSpace.gmem,
        assumed_align=16
    )
    output_ptr = make_ptr(
        OUTPUT_DTYPE,
        output_c.data_ptr(),
        cute.AddressSpace.gmem,
        assumed_align=16
    )
    scale_a_ptr = make_ptr(
        SCALE_FACTOR_DTYPE,
        scale_a_permuted.data_ptr(),
        cute.AddressSpace.gmem,
        assumed_align=32
    )
    scale_b_ptr = make_ptr(
        SCALE_FACTOR_DTYPE,
        scale_b_permuted.data_ptr(),
        cute.AddressSpace.gmem,
        assumed_align=32
    )

    compiled_kernel(
        a_ptr,
        b_ptr,
        scale_a_ptr,
        scale_b_ptr,
        output_ptr,
        (m, n, k_actual, l)
    )

    return output_c