submission 96902

import torch
from task import input_t, output_t

import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import make_ptr
import cutlass.utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass._mlir.dialects import nvvm
from cutlass.cute.nvgpu import tcgen05

# Scale factor vector size: 16 elements share one scale factor in block-scaled format
# This matches the block size used in the scale factor layout (16 FP4 elements per scale)
sf_vec_size = 16
# Number of threads handling the M dimension per block
# Default path: 32 threads allows us to process 32 rows in parallel, matching warp size for efficient SIMD
threads_per_m = 32
# tcgen05 optimized - no specialized big kernel needed
# Number of threads handling the K dimension per block (base value, can be overridden)
# This is used as a default but get_threads_per_k() provides dynamic values based on K size
threads_per_k = 8
# Input matrices A and B use FP4 (4-bit float) format: Float4E2M1FN
# This is NVIDIA's FP4 format with 2 exponent bits and 1 mantissa bit
ab_dtype = cutlass.Float4E2M1FN
# Scale factors use FP8 (8-bit float) format: Float8E4M3FN
# FP8 provides sufficient precision for scale factors while keeping memory footprint small
sf_dtype = cutlass.Float8E4M3FN
# Output matrix C uses FP16 (16-bit float) format
# FP16 provides good precision for accumulation while being efficient on modern GPUs
c_dtype = cutlass.Float16

# tcgen05 MMA configuration for Blackwell
# Using standard tcgen05 instruction shape (K must be 16 for MmaF16BF16Op)
mma_inst_shape = (64, 64, 16)  # (M, N, K) - K=16 required for tcgen05
cta_group = tcgen05.CtaGroup.ONE
threads_per_cta = 128


def ceil_div(a, b):
    # Standard ceiling division: (a + b - 1) // b rounds up instead of down
    # This ensures we allocate enough tiles/blocks to cover all elements
    return (a + b - 1) // b


def get_k_tile_size(k):
    # Adaptive K tile size based on problem size
    # Smaller K (<=2048): use 128-element tiles for better cache locality
    # Larger K (>2048): use 256-element tiles to reduce loop overhead and improve throughput
    if k <= 2048:
        return 128
    else:
        return 256


def get_threads_per_k(k):
    # Adaptive thread count per K dimension based on problem size
    # Smaller K (<=2048): 4 threads - less parallelism needed, better register usage
    # Medium K (2048-8192): 6 threads - balanced parallelism
    # Large K (>8192): 8 threads - maximize parallelism to hide memory latency
    if k <= 2048:
        return 4
    elif k <= 8192:
        return 6
    else:
        return 8


@dsl_user_op
def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer:
    # Helper function to compute element pointer from tensor and coordinate
    # Converts coordinate to linear index and adds to tensor's base iterator
    return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)


@cute.kernel
def cutedsl_kernel_tcgen(
    mA_mkl: cute.Tensor,
    mB_nkl: cute.Tensor,
    mSFA_mkl: cute.Tensor,
    mSFB_nkl: cute.Tensor,
    mC_mnl: cute.Tensor,
    k_tile_size: cutlass.Constexpr,
    threads_per_k_param: cutlass.Constexpr,
):
    # Get CUDA block and thread indices
    # bidx: block index in M dimension
    # bidy: block index in N dimension (unused here)
    # bidz: block index in L dimension
    bidx, bidy, bidz = cute.arch.block_idx()
    # tidx: thread index in M dimension (0-31)
    # tidy: thread index in K dimension (0 to threads_per_k_param-1)
    tidx, tidy, _ = cute.arch.thread_idx()

    # Define tile size: (M threads, 1 element in N dimension, k_tile_size elements in K dimension)
    # N=1 because this is matrix-vector multiplication (B has shape 1xKxL)
    mma_tiler_mnk = (threads_per_m, 1, k_tile_size)

    # Extract tiles from global memory tensors for this block
    # slice_(mma_tiler_mnk, (None, 0, None)) means:
    #   take full M tile, element 0 in N, full K tile
    # This gives each block a tile of A: [threads_per_m, 1, k_tile_size] elements
    gA_mkl = cute.local_tile(
        mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
    )
    # Scale factors for A: same shape as A but with K dimension
    # divided by sf_vec_size (16)
    gSFA_mkl = cute.local_tile(
        mSFA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
    )
    # B is broadcast across M dimension, so slice_(mma_tiler_mnk, (0, None, None))
    # takes element 0 in M
    # This gives us a single row of B: [1, 1, k_tile_size] elements
    gB_nkl = cute.local_tile(
        mB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
    )
    # Scale factors for B: same shape as B
    gSFB_nkl = cute.local_tile(
        mSFB_nkl, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
    )
    # Output C: slice_(mma_tiler_mnk, (None, None, 0)) takes full M, full N, element 0 in K
    # Since N=1 and K reduction happens in kernel, we get [threads_per_m, 1, 1] per block
    gC_mnl = cute.local_tile(
        mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
    )

    # Get the output element this thread is responsible for: thread tidx in block bidx, batch bidz
    # Each thread handles one element in the M dimension
    tCgC = gC_mnl[tidx, None, bidx, 0, bidz]
    tCgC = cute.make_tensor(tCgC.iterator, 1)
    # Accumulate result in FP32 for precision (FP32 has more precision than FP16 for accumulation)
    # This prevents precision loss during the dot product accumulation
    res = cute.zeros_like(tCgC, cutlass.Float32)

    # Allocate shared memory for partial results reduction
    # Shared memory allows threads to collaborate on reducing partial dot products
    allocator = cutlass.utils.SmemAllocator()
    # Layout: [threads_per_m, threads_per_k_param] with column-major stride
    # Each thread (tidx, tidy) stores its partial result at shared_res[(tidx, tidy)]
    # Column-major stride (stride=(threads_per_k_param, 1)) ensures coalesced memory access
    smem_layout = cute.make_layout(
        (threads_per_m, threads_per_k_param), stride=(threads_per_k_param, 1)
    )
    shared_res = allocator.allocate_tensor(element_type=cutlass.Float32, layout=smem_layout)

    # Number of K tiles to process (K dimension is divided into tiles of size k_tile_size)
    k_tile_cnt = gA_mkl.layout[3].shape

    # Using manual FP4 computation - tcgen05 setup removed for now

    # Parallelize K dimension across threads: each thread handles a subset of K tiles
    # tidy-thread processes tiles: tidy, tidy+threads_per_k_param, tidy+2*threads_per_k_param, ...
    # unroll_full=True: fully unroll the loop for better performance (no loop overhead)
    for k_tile in range(tidy, k_tile_cnt, threads_per_k_param, unroll_full=True):
        # Get tiles for this thread: A tile for thread tidx, B tile (same for all threads since N=1)
        # Each thread loads its portion of A: [1, 1, k_tile_size] elements
        tAgA = gA_mkl[tidx, None, bidx, k_tile, bidz]
        tAgSFA = gSFA_mkl[tidx, None, bidx, k_tile, bidz]

        # B is broadcast, so all threads load the same B tile: [1, 1, k_tile_size] elements
        tBgB = gB_nkl[0, None, 0, k_tile, bidz]
        tBgSFB = gSFB_nkl[0, None, 0, k_tile, bidz]

        # Create register memory tensors (fast on-chip memory) for computation
        # Convert FP4 to FP16 for A and B: FP16 provides better precision
        # for intermediate computation
        # Keep scale factors in FP32 for maximum precision during scaling
        tArA = cute.make_rmem_tensor_like(tAgA, cutlass.Float16)
        tBrB = cute.make_rmem_tensor_like(tBgB, cutlass.Float16)
        tArSFA = cute.make_rmem_tensor_like(tAgSFA, cutlass.Float32)
        tBrSFB = cute.make_rmem_tensor_like(tBgSFB, cutlass.Float32)

        # Load data from global memory and convert to computation types
        # FP4 -> FP16 conversion happens here (dequantization)
        tArA.store(tAgA.load().to(cutlass.Float16))
        tBrB.store(tBgB.load().to(cutlass.Float16))
        # Scale factors stay in FP32 for precision
        tArSFA.store(tAgSFA.load().to(cutlass.Float32))
        tBrSFB.store(tBgSFB.load().to(cutlass.Float32))

        # Load data from global memory and convert to computation types
        # FP4 -> FP16 conversion happens here (dequantization)
        tArA.store(tAgA.load().to(cutlass.Float16))
        tBrB.store(tBgB.load().to(cutlass.Float16))
        # Scale factors stay in FP32 for precision
        tArSFA.store(tAgSFA.load().to(cutlass.Float32))
        tBrSFB.store(tBgSFB.load().to(cutlass.Float32))

        # Manual FP4 computation (original working approach)
        # Process K dimension in blocks of 8 elements (unrolled for performance)
        for k_block in cutlass.range_constexpr(k_tile_size // 8):
            # Process 8 elements per iteration (fully unrolled)
            idx0 = k_block * 8 + 0
            res += (tArA[idx0] * tBrB[idx0]).to(cutlass.Float32) * (
                tArSFA[idx0] * tBrSFB[idx0]
            )
            idx1 = k_block * 8 + 1
            res += (tArA[idx1] * tBrB[idx1]).to(cutlass.Float32) * (
                tArSFA[idx1] * tBrSFB[idx1]
            )
            idx2 = k_block * 8 + 2
            res += (tArA[idx2] * tBrB[idx2]).to(cutlass.Float32) * (
                tArSFA[idx2] * tBrSFB[idx2]
            )
            idx3 = k_block * 8 + 3
            res += (tArA[idx3] * tBrB[idx3]).to(cutlass.Float32) * (
                tArSFA[idx3] * tBrSFB[idx3]
            )
            idx4 = k_block * 8 + 4
            res += (tArA[idx4] * tBrB[idx4]).to(cutlass.Float32) * (
                tArSFA[idx4] * tBrSFB[idx4]
            )
            idx5 = k_block * 8 + 5
            res += (tArA[idx5] * tBrB[idx5]).to(cutlass.Float32) * (
                tArSFA[idx5] * tBrSFB[idx5]
            )
            idx6 = k_block * 8 + 6
            res += (tArA[idx6] * tBrB[idx6]).to(cutlass.Float32) * (
                tArSFA[idx6] * tBrSFB[idx6]
            )
            idx7 = k_block * 8 + 7
            res += (tArA[idx7] * tBrB[idx7]).to(cutlass.Float32) * (
                tArSFA[idx7] * tBrSFB[idx7]
            )

        # Store partial result to shared memory: each thread (tidx, tidy) contributes one partial sum
    # After this, shared_res[tidx, :] contains all partial sums for output element tidx
    shared_res[(tidx, tidy)] = res[0]
    # Synchronize threads to ensure all partial results are written before reduction
    cute.arch.sync_threads()

    # Only thread 0 in K dimension (tidy==0) performs the final reduction
    # This avoids race conditions and ensures only one thread writes the final result
    if tidy == 0:
        out = cute.zeros_like(tCgC, cutlass.Float32)

        # Reduce across K dimension: sum all partial results from threads_per_k_param threads
        # Each thread handled different K tiles, so we sum their contributions
        for i in cutlass.range_constexpr(threads_per_k_param):
            out += shared_res[(tidx, i)]

        # Convert final result from FP32 to FP16 and store to global memory
        # FP32->FP16 conversion happens here (quantization for output)
        tCgC.store(out.to(cutlass.Float16))

    return


@cute.jit
def cutedsl_path_tcgen(
    a_ptr: cute.Pointer,
    b_ptr: cute.Pointer,
    sfa_ptr: cute.Pointer,
    sfb_ptr: cute.Pointer,
    c_ptr: cute.Pointer,
    problem_size: tuple,
    k_tile_size: cutlass.Constexpr,
    threads_per_k_param: cutlass.Constexpr,
):
    # Extract problem dimensions: m (rows), k (cols), l (batch)
    # Note: second dimension is ignored (it's 1 for matrix-vector multiplication)
    m, _, k, l = problem_size

    # Create tensor layout for A: [m, k, l] with K-major order
    # cute.assume(k, 32) tells compiler that k is divisible by 32 (enables optimizations)
    # Stride: (k, 1, m*k) means row-major within each batch, batches are contiguous
    a_tensor = 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)),
        ),
    )
    # B tensor is padded to 128 in N dimension (even though actual N=1)
    # Padding to 128 ensures alignment and enables vectorized memory access patterns
    # This is a common optimization for memory-bound kernels
    n_padded_128 = 128
    b_tensor = cute.make_tensor(
        b_ptr,
        cute.make_layout(
            (n_padded_128, cute.assume(k, 32), l),
            stride=(cute.assume(k, 32), 1, cute.assume(n_padded_128 * k, 32)),
        ),
    )
    # C tensor: [m, 1, l] with column-major stride (stride=(1, 1, m))
    # Column-major allows efficient writes when each thread writes
    # one element
    c_tensor = cute.make_tensor(
        c_ptr, cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m))
    )
    # Scale factor layouts: automatically computed based on tensor shape
    # and sf_vec_size
    # The layout matches the block-scaled format where every sf_vec_size (16)
    # elements share one scale
    sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a_tensor.shape, sf_vec_size)
    sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
    sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b_tensor.shape, sf_vec_size)
    sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)

    # Compute grid dimensions: (blocks in M dimension, blocks in N dimension, batches)
    # Each block handles threads_per_m rows, so we need ceil(m / threads_per_m) blocks
    # N dimension is 1 (matrix-vector multiplication), L dimension is batch size
    grid = (
        cute.ceil_div(c_tensor.shape[0], threads_per_m),
        1,
        c_tensor.shape[2],
    )

    # Launch kernel with computed grid and block dimensions
    # block: [threads_per_m, threads_per_k_param, 1] threads per block
    # cluster: (1, 1, 1) means no multi-cluster launch (standard CUDA launch)
    cutedsl_kernel_tcgen(
        a_tensor,
        b_tensor,
        sfa_tensor,
        sfb_tensor,
        c_tensor,
        k_tile_size,
        threads_per_k_param,
    ).launch(
        grid=grid,
        block=[threads_per_m, threads_per_k_param, 1],
        cluster=(1, 1, 1),
    )
    return


# tcgen05 kernels only - no specialized variants needed

# Cache for compiled kernels: avoids recompiling for same (k_tile_size, threads_per_k) parameters
# Compilation is expensive, so caching significantly speeds up repeated calls with same parameters
_compiled_cutedsl_cache_tcgen = {}


def compile_cutedsl_tcgen(k_tile_size, threads_per_k_val):
    global _compiled_cutedsl_cache_tcgen
    # Use (k_tile_size, threads_per_k_val) as cache key since these affect kernel code generation
    cache_key = (k_tile_size, threads_per_k_val)
    if cache_key in _compiled_cutedsl_cache_tcgen:
        return _compiled_cutedsl_cache_tcgen[cache_key]

    # Create dummy pointers for compilation (address 0, will be replaced at runtime)
    # assumed_align=16 for A, B, C: 16-byte alignment enables vectorized loads (128-bit loads)
    # assumed_align=32 for scale factors: 32-byte alignment for optimal memory access
    a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
    b_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
    c_ptr = make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
    sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
    sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)

    # Compile kernel with maximum optimization level
    # --opt-level 3: maximum LLVM optimization
    # --ptxas-options '--opt-level=3': maximum PTX assembler optimization
    # These optimizations are critical for performance on B200
    _compiled_cutedsl_cache_tcgen[cache_key] = cute.compile(
        cutedsl_path_tcgen,
        a_ptr,
        b_ptr,
        sfa_ptr,
        sfb_ptr,
        c_ptr,
        (0, 0, 0, 0),
        k_tile_size,
        threads_per_k_val,
        options="--opt-level 3 --ptxas-options '--opt-level=3'",
    )

    return _compiled_cutedsl_cache_tcgen[cache_key]


def custom_kernel(data: input_t) -> output_t:
    # Unpack input data: a (FP4), b (FP4), scale factors (both CPU and permuted versions), output c
    # sfa_permuted and sfb_permuted are pre-permuted for optimal memory access pattern
    a, b, sfa_cpu, sfb_cpu, sfa_permuted, sfb_permuted, c = data

    m, k, l = a.shape
    # FP4 packs 2 elements per byte, so actual K dimension is k*2
    # This accounts for the packed storage format of FP4 tensors
    k = k * 2

    # tcgen05 optimized - single kernel for all problem sizes
    # Select optimal tile size and thread count based on problem size
    # Adaptive selection ensures best performance across different problem sizes
    k_tile_size = get_k_tile_size(k)
    threads_per_k_val = get_threads_per_k(k)
    # Compile (or retrieve from cache) the tcgen05 kernel with these parameters
    compiled_func = compile_cutedsl_tcgen(k_tile_size, threads_per_k_val)
    block_threads_m = threads_per_m

    # Create pointers to GPU memory with alignment assumptions
    # These pointers will be used by the compiled kernel to access data
    # assumed_align tells compiler about memory alignment for optimization
    a_ptr = make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    b_ptr = make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    c_ptr = make_ptr(c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    # Use permuted scale factors: they're pre-arranged for optimal memory access
    sfa_ptr = make_ptr(sf_dtype, sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)
    sfb_ptr = make_ptr(sf_dtype, sfb_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32)

    # Launch kernel: (m, 1, k, l) = (rows, columns=1 for vector, cols, batches)
    # The kernel computes: C = A @ B where A is [m, k, l] and B is [1, k, l]
    compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, 1, k, l))

    # Return output tensor (modified in-place)
    return c