submission 116627

import math
from typing import Tuple
from functools import partial
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass import Float16, Int8, Int16, Int32
from cutlass._mlir import ir
from cutlass._mlir.dialects import arith, builtin, llvm, vector
from cutlass._mlir.dialects.cute import ReductionOp as ReductionOp
from cutlass.cute.runtime import make_ptr
from cutlass.cutlass_dsl import T, dsl_user_op
from task import input_t, output_t

"""
M    K     L time[us]
7168 16384 1 8.622
4096 7168  8 17.275
7168 2048  4 4.317

Compared to v8
1. res += minor change
2. transpose: multiple tidx same tidy load 1 row
3. use warp reduce to replace shared mem for final compute: minor improvement
4. reduce sf load
5. auto tune
6. 32 ele
7. mul f16x2, no improvement, sass already does that

1:
num_row_threads = 4, num_col_threads = 32, k_tile 128: 27.424μs
k: 16384; l: 1; m: 7168; seed: 1111
 ⏱ 24.9 ± 0.05 µs
 ⚡ 24.4 µs 🐌 26.7 µs

k: 7168; l: 8; m: 4096; seed: 1111
 ⏱ 41.0 ± 0.04 µs
 ⚡ 40.9 µs 🐌 43.0 µs

k: 2048; l: 4; m: 7168; seed: 1111
 ⏱ 20.2 ± 0.05 µs
 ⚡ 18.4 µs 🐌 20.6 µs

1 + 2:
num_row_threads = 4, num_col_threads = 32, k_tile 128: 40.176μs
k: 16384; l: 1; m: 7168; seed: 1111
 ⏱ 39.1 ± 0.06 µs
 ⚡ 36.8 µs 🐌 41.1 µs

k: 7168; l: 8; m: 4096; seed: 1111
 ⏱ 63.6 ± 0.06 µs
 ⚡ 63.4 µs 🐌 64.5 µs

k: 2048; l: 4; m: 7168; seed: 1111
 ⏱ 26.1 ± 0.06 µs
 ⚡ 24.5 µs 🐌 26.9 µs

num_row_threads = 4, num_col_threads = 32, k_tile 32: 34.241μs
k: 16384; l: 1; m: 7168; seed: 1111
 ⏱ 34.8 ± 0.03 µs
 ⚡ 33.6 µs 🐌 36.9 µs

k: 7168; l: 8; m: 4096; seed: 1111
 ⏱ 55.2 ± 0.05 µs
 ⚡ 54.1 µs 🐌 56.4 µs

k: 2048; l: 4; m: 7168; seed: 1111
 ⏱ 20.9 ± 0.05 µs
 ⚡ 20.4 µs 🐌 22.8 µs

1 + 2 + 3:
num_row_threads = 4, num_col_threads = 32, k_tile 32: 33.785μs
k: 16384; l: 1; m: 7168; seed: 1111
 ⏱ 34.7 ± 0.03 µs
 ⚡ 32.7 µs 🐌 35.7 µs

k: 7168; l: 8; m: 4096; seed: 1111
 ⏱ 54.3 ± 0.07 µs
 ⚡ 53.1 µs 🐌 56.3 µs

k: 2048; l: 4; m: 7168; seed: 1111
 ⏱ 20.5 ± 0.01 µs
 ⚡ 20.4 µs 🐌 20.5 µs

1 + 2 + 3 + 4:
num_row_threads = 4, num_col_threads = 32, k_tile 32: 46.756μs
k: 16384; l: 1; m: 7168; seed: 1111
 ⏱ 46.0 ± 0.05 µs
 ⚡ 44.0 µs 🐌 46.3 µs

k: 7168; l: 8; m: 4096; seed: 1111
 ⏱ 74.7 ± 0.07 µs
 ⚡ 72.8 µs 🐌 75.8 µs

k: 2048; l: 4; m: 7168; seed: 1111
 ⏱ 29.7 ± 0.03 µs
 ⚡ 29.6 µs 🐌 30.5 µs

✅ 1 + 2 + 3 + 4 + 5: 27.098μs
"num_row_threads": [4, 8, 16, 32],
"k_tile_size": [32, 64],
best: 8, 32
k: 16384; l: 1; m: 7168; seed: 1111
 ⏱ 26.7 ± 0.03 µs
 ⚡ 26.5 µs 🐌 27.6 µs

k: 7168; l: 8; m: 4096; seed: 1111
 ⏱ 45.1 ± 0.04 µs
 ⚡ 43.9 µs 🐌 46.1 µs

k: 2048; l: 4; m: 7168; seed: 1111
 ⏱ 16.5 ± 0.04 µs
 ⚡ 16.3 µs 🐌 18.5 µs

1 + 2 + 3 + 5: 27.958μs
"num_row_threads": [4, 8, 16, 32],
"k_tile_size": [32, 64, 128],
k: 16384; l: 1; m: 7168; seed: 1111
 ⏱ 28.3 ± 0.06 µs
 ⚡ 26.6 µs 🐌 29.4 µs

k: 7168; l: 8; m: 4096; seed: 1111
 ⏱ 45.1 ± 0.04 µs
 ⚡ 44.9 µs 🐌 46.2 µs

k: 2048; l: 4; m: 7168; seed: 1111
 ⏱ 17.1 ± 0.07 µs
 ⚡ 16.2 µs 🐌 18.7 µs

✅ 1 + 2 + 3 + 4 + 5 + 6
very minor improvement
k: 16384; l: 1; m: 7168; seed: 1111
 ⏱ 26.7 ± 0.03 µs
 ⚡ 25.5 µs 🐌 27.7 µs

k: 7168; l: 8; m: 4096; seed: 1111
 ⏱ 44.9 ± 0.04 µs
 ⚡ 43.0 µs 🐌 46.2 µs

k: 2048; l: 4; m: 7168; seed: 1111
 ⏱ 16.5 ± 0.03 µs
 ⚡ 16.3 µs 🐌 18.5 µs

✅ 1 + 2 + 3 + 4 + 5 + 6 + 7
no change
k: 16384; l: 1; m: 7168; seed: 1111
 ⏱ 26.6 ± 0.01 µs
 ⚡ 26.5 µs 🐌 26.8 µs

k: 7168; l: 8; m: 4096; seed: 1111
 ⏱ 44.7 ± 0.06 µs
 ⚡ 42.9 µs 🐌 46.1 µs

k: 2048; l: 4; m: 7168; seed: 1111
 ⏱ 16.4 ± 0.01 µs
 ⚡ 16.2 µs 🐌 16.5 µs
"""

# Kernel configuration parameters
ab_dtype = cutlass.Float4E2M1FN  # FP4 data type for A and B
sf_dtype = cutlass.Float8E4M3FN  # FP8 data type for scale factors
c_dtype = cutlass.Float16  # FP16 output type
acc_dtype = cutlass.Float32
sf_vec_size = 16  # Scale factor block size (16 elements share one scale)


# Helper function for ceiling division
def ceil_div(a, b):
    return (a + b - 1) // b


@cute.jit
def warp_reduce_sum(
    val: cute.TensorSSA | cute.Numeric,
    width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE,
) -> cute.TensorSSA | cute.Numeric:
    if cutlass.const_expr(isinstance(val, cute.TensorSSA)):
        res = cute.make_fragment(val.shape, val.dtype)
        res.store(val)
        for i in cutlass.range_constexpr(cute.size(val.shape)):
            res[i] = warp_reduce_sum(res[i], width)
        return res.load()
    else:
        for i in cutlass.range_constexpr(int(math.log2(width))):
            val += cute.arch.shuffle_sync_bfly(val, offset=1 << i)
    return val


@dsl_user_op
def fma_f16x2(
    a: Tuple[Float16, Float16],
    b: Tuple[Float16, Float16],
    c: Tuple[Float16, Float16],
    *,
    loc=None,
    ip=None,
) -> Tuple[Float16, Float16]:
    # Pack two Float16 values into vector<2xf16>
    vec_type = ir.VectorType.get([2], Float16.mlir_type, loc=loc)

    vec_a = vector.from_elements(
        vec_type,
        [a[0].ir_value(loc=loc, ip=ip), a[1].ir_value(loc=loc, ip=ip)],
        loc=loc,
        ip=ip,
    )
    vec_b = vector.from_elements(
        vec_type,
        [b[0].ir_value(loc=loc, ip=ip), b[1].ir_value(loc=loc, ip=ip)],
        loc=loc,
        ip=ip,
    )
    vec_c = vector.from_elements(
        vec_type,
        [c[0].ir_value(loc=loc, ip=ip), c[1].ir_value(loc=loc, ip=ip)],
        loc=loc,
        ip=ip,
    )

    # Bitcast to i32 for PTX (f16x2 is packed into 32 bits)
    a_i32 = llvm.bitcast(Int32.mlir_type, vec_a, loc=loc, ip=ip)
    b_i32 = llvm.bitcast(Int32.mlir_type, vec_b, loc=loc, ip=ip)
    c_i32 = llvm.bitcast(Int32.mlir_type, vec_c, loc=loc, ip=ip)

    # Simple single-line PTX like cvt_f16x2_f32
    result_i32 = llvm.inline_asm(
        Int32.mlir_type,
        [a_i32, b_i32, c_i32],
        "fma.rn.f16x2 $0, $1, $2, $3;",
        "=r,r,r,r",
        has_side_effects=False,
        is_align_stack=False,
        asm_dialect=llvm.AsmDialect.AD_ATT,
        loc=loc,
        ip=ip,
    )

    # Bitcast back to vector<2xf16>
    vec_result = llvm.bitcast(vec_type, result_i32, loc=loc, ip=ip)

    # Extract results
    result0 = Float16(
        vector.extract(
            vec_result, dynamic_position=[], static_position=[0], loc=loc, ip=ip
        )
    )
    result1 = Float16(
        vector.extract(
            vec_result, dynamic_position=[], static_position=[1], loc=loc, ip=ip
        )
    )

    return result0, result1


@dsl_user_op
def cvt_f8e4m3_f16_intrinsic(vec_f8e4m3, length, *, loc=None, ip=None):
    """
    Convert a vector of float8e4m3 to a vector of float16.
    :param vec_f8e4m3: The input vector of float8e4m3.
    :type vec_f8e4m3: 1D vector of float8e4m3
    :param length: The length of the input vector.
    :type length: int
    :return: The output 1D vector of float16 with the same length as the input vector.
    :rtype: 1D vector of float16
    """
    src_pos = 0
    vec_src_i8 = builtin.unrealized_conversion_cast(
        [ir.VectorType.get([length], Int8.mlir_type, loc=loc)],
        [vec_f8e4m3],
        loc=loc,
        ip=ip,
    )
    vec_i8x8_type = ir.VectorType.get([8], Int8.mlir_type, loc=loc)
    vec_i8x4_type = ir.VectorType.get([4], Int8.mlir_type, loc=loc)
    vec_i8x2_type = ir.VectorType.get([2], Int8.mlir_type, loc=loc)
    vec_dst_type = ir.VectorType.get([length], Float16.mlir_type, loc=loc)
    vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip)

    # try to use vectorized version
    if length >= 8:
        num_vec8 = length // 8
        for _ in range(num_vec8):
            vec_f8e4m3x8 = vector.extract_strided_slice(
                vec_i8x8_type, vec_src_i8, [src_pos], [8], [1], loc=loc, ip=ip
            )
            vec_f16x8 = cvt_f8e4m3x8_to_f16x8(vec_f8e4m3x8, loc=loc, ip=ip)
            vec_dst = vector.insert_strided_slice(
                vec_f16x8, vec_dst, [src_pos], [1], loc=loc, ip=ip
            )
            src_pos += 8
            length -= 8

    if length >= 4:
        vec_f8e4m3x4 = vector.extract_strided_slice(
            vec_i8x4_type, vec_src_i8, [src_pos], [4], [1], loc=loc, ip=ip
        )
        vec_f16x4 = cvt_f8e4m3x4_to_f16x4(vec_f8e4m3x4, loc=loc, ip=ip)
        vec_dst = vector.insert_strided_slice(
            vec_f16x4, vec_dst, [src_pos], [1], loc=loc, ip=ip
        )
        src_pos += 4
        length -= 4

    if length >= 2:
        vec_f8e4m3x2 = vector.extract_strided_slice(
            vec_i8x2_type, vec_src_i8, [src_pos], [2], [1], loc=loc, ip=ip
        )
        vec_f16x2 = cvt_f8e4m3x2_to_f16x2(vec_f8e4m3x2, loc=loc, ip=ip)
        vec_dst = vector.insert_strided_slice(
            vec_f16x2, vec_dst, [src_pos], [1], loc=loc, ip=ip
        )
        src_pos += 2
        length -= 2

    if length >= 1:
        val_f16 = cvt_f8e4m3_f16(
            vector.extractelement(
                vec_src_i8,
                position=arith.constant(Int32.mlir_type, src_pos),
                loc=loc,
                ip=ip,
            ),
            loc=loc,
            ip=ip,
        )
        vec_dst = vector.insertelement(
            val_f16,
            vec_dst,
            position=arith.constant(Int32.mlir_type, src_pos),
            loc=loc,
            ip=ip,
        )

    return vec_dst


@dsl_user_op
def cvt_f8e4m3_f16(src, *, loc=None, ip=None):
    # 0 padding for upper 8 bits
    zero = arith.constant(src.type, 0, loc=loc, ip=ip)
    vec2 = vector.from_elements(
        ir.VectorType.get([2], src.type, loc=loc), [src, zero], loc=loc, ip=ip
    )
    rst_vec2 = cvt_f8e4m3x2_to_f16x2(vec2, loc=loc, ip=ip)
    # only the 1st element is valid
    rst = vector.extract(
        rst_vec2, dynamic_position=[], static_position=[0], loc=loc, ip=ip
    )
    return rst


# Convert 2 float8e4m3 values to 2 float16 values
@dsl_user_op
def cvt_f8e4m3x2_to_f16x2(src_vec2, *, loc=None, ip=None):
    # pack 2 float8e4m3 into 1 int16 value
    src_i16 = llvm.bitcast(Int16.mlir_type, src_vec2, loc=loc, ip=ip)
    rst_i32 = llvm.inline_asm(
        Int32.mlir_type,
        [src_i16],
        """{\n\t
            cvt.rn.f16x2.e4m3x2 $0, $1;\n\t
        }""",
        "=r,h",
    )
    vec_f16x2_type = ir.VectorType.get([2], Float16.mlir_type, loc=loc)
    vec_f16x2 = llvm.bitcast(vec_f16x2_type, rst_i32, loc=loc, ip=ip)
    return vec_f16x2


# Convert 4 float8e4m3 values to 4 float16 values
@dsl_user_op
def cvt_f8e4m3x4_to_f16x4(src_vec4, *, loc=None, ip=None):
    # pack 4 float8e4m3 into 1 int32 value
    src_i32 = llvm.bitcast(Int32.mlir_type, src_vec4, loc=loc, ip=ip)
    rst_i32x2 = llvm.inline_asm(
        llvm.StructType.get_literal([T.i32(), T.i32()]),
        [src_i32],
        """{\n\t
            .reg .b16 h0, h1;\n\t
            mov.b32 {h0, h1}, $2;\n\t
            cvt.rn.f16x2.e4m3x2 $0, h0;\n\t
            cvt.rn.f16x2.e4m3x2 $1, h1;\n\t
        }""",
        "=r,=r,r",
    )
    res0 = llvm.extractvalue(T.i32(), rst_i32x2, [0])
    res1 = llvm.extractvalue(T.i32(), rst_i32x2, [1])
    vec_i32x2_type = ir.VectorType.get([2], Int32.mlir_type, loc=loc)
    vec_i32x2 = vector.from_elements(vec_i32x2_type, [res0, res1], loc=loc, ip=ip)
    vec_f16x4_type = ir.VectorType.get([4], Float16.mlir_type, loc=loc)
    vec_f16x4 = llvm.bitcast(vec_f16x4_type, vec_i32x2, loc=loc, ip=ip)
    return vec_f16x4


# Convert 8 float8e4m3 values to 8 float16 values
@dsl_user_op
def cvt_f8e4m3x8_to_f16x8(src_vec8, *, loc=None, ip=None):
    # Split into two i32 values instead of using i64
    vec_i32x2_type = ir.VectorType.get([2], Int32.mlir_type, loc=loc)
    src_i32x2 = llvm.bitcast(vec_i32x2_type, src_vec8, loc=loc, ip=ip)
    src_lo = llvm.extractelement(
        src_i32x2, arith.constant(Int32.mlir_type, 0), loc=loc, ip=ip
    )
    src_hi = llvm.extractelement(
        src_i32x2, arith.constant(Int32.mlir_type, 1), loc=loc, ip=ip
    )

    # Process lower 4 bytes (4 fp8 values)
    rst_lo_i32x2 = llvm.inline_asm(
        llvm.StructType.get_literal([T.i32(), T.i32()]),
        [src_lo],
        """{\n\t
            .reg .b16 h0, h1;\n\t
            mov.b32 {h0, h1}, $2;\n\t
            cvt.rn.f16x2.e4m3x2 $0, h0;\n\t
            cvt.rn.f16x2.e4m3x2 $1, h1;\n\t
        }""",
        "=r,=r,r",
    )

    # Process upper 4 bytes (4 fp8 values)
    rst_hi_i32x2 = llvm.inline_asm(
        llvm.StructType.get_literal([T.i32(), T.i32()]),
        [src_hi],
        """{\n\t
            .reg .b16 h0, h1;\n\t
            mov.b32 {h0, h1}, $2;\n\t
            cvt.rn.f16x2.e4m3x2 $0, h0;\n\t
            cvt.rn.f16x2.e4m3x2 $1, h1;\n\t
        }""",
        "=r,=r,r",
    )

    res0 = llvm.extractvalue(T.i32(), rst_lo_i32x2, [0])
    res1 = llvm.extractvalue(T.i32(), rst_lo_i32x2, [1])
    res2 = llvm.extractvalue(T.i32(), rst_hi_i32x2, [0])
    res3 = llvm.extractvalue(T.i32(), rst_hi_i32x2, [1])

    vec_i32x4_type = ir.VectorType.get([4], Int32.mlir_type, loc=loc)
    vec_i32x4 = vector.from_elements(
        vec_i32x4_type, [res0, res1, res2, res3], loc=loc, ip=ip
    )
    vec_f16x8_type = ir.VectorType.get([8], Float16.mlir_type, loc=loc)
    vec_f16x8 = llvm.bitcast(vec_f16x8_type, vec_i32x4, loc=loc, ip=ip)
    return vec_f16x8


class Nvfp4BatchedGemvKernel:
    def __init__(
        self,
    ):
        self.threads_per_cta = 0

    @cute.kernel
    def kernel(
        self,
        gA_mkl: cute.Tensor,
        gB_nkl: cute.Tensor,
        gSFA_mkl: cute.Tensor,
        gSFB_nkl: cute.Tensor,
        gC_mnl: cute.Tensor,
        num_col_threads: cutlass.Constexpr,
        k_tile_size: cutlass.Constexpr,
    ):
        # Get CUDA block and thread indices
        bidx, bidy, bidz = cute.arch.block_idx()
        tidx, tidy, _ = cute.arch.thread_idx()

        # Select output element corresponding to this thread and block indices
        tCgC = gC_mnl[tidy, None, bidx, bidy, bidz]
        tCgC = cute.make_tensor(tCgC.iterator, 1)
        res = cute.zeros_like(tCgC, acc_dtype)

        # Get the number of k tiles (depth dimension) for the reduction loop
        k_tile_cnt = gA_mkl.layout[3].shape
        for k_tile in range(tidx, k_tile_cnt, num_col_threads, unroll_full=True):
            tAgA = gA_mkl[tidy, None, bidx, k_tile, bidz]
            tBgB = gB_nkl[0, None, bidy, k_tile, bidz]
            if cutlass.const_expr(k_tile_size % 128 == 0):
                tAgSFA = gSFA_mkl[tidy, (0, None, None), bidx, k_tile, bidz]
                tBgSFB = gSFB_nkl[0, (0, None, None), bidy, k_tile, bidz]
            else:
                tAgSFA = gSFA_mkl[tidy, (0, None), bidx, k_tile, bidz]
                tBgSFB = gSFB_nkl[0, (0, None), bidy, k_tile, bidz]

            tArA = cute.make_rmem_tensor_like(tAgA, cutlass.Float16)
            tBrB = cute.make_rmem_tensor_like(tBgB, cutlass.Float16)
            tCrSFA = cute.make_rmem_tensor_like(tAgSFA, cutlass.Float16)
            tCrSFB = cute.make_rmem_tensor_like(tAgSFA, cutlass.Float16)

            # Load NVFP4 or FP8 values from global memory
            a_val_nvfp4 = tAgA.load()
            b_val_nvfp4 = tBgB.load()
            sfa_val_fp8 = tAgSFA.load()
            sfb_val_fp8 = tBgSFB.load()

            # Convert loaded values to float32 for computation (FFMA)
            a_val = a_val_nvfp4.to(cutlass.Float16)
            b_val = b_val_nvfp4.to(cutlass.Float16)
            sfa_val = cute.TensorSSA(
                cvt_f8e4m3_f16_intrinsic(sfa_val_fp8, cute.size(sfa_val_fp8.shape)),
                sfa_val_fp8.shape,
                cutlass.Float16,
            )
            sfb_val = cute.TensorSSA(
                cvt_f8e4m3_f16_intrinsic(sfb_val_fp8, cute.size(sfb_val_fp8.shape)),
                sfb_val_fp8.shape,
                cutlass.Float16,
            )

            # Store the converted values to RMEM CuTe tensors
            tArA.store(a_val)
            tBrB.store(b_val)
            tCrSFA.store(sfa_val)
            tCrSFB.store(sfb_val)

            # Iterate over SF vector tiles and compute the scale&matmul accumulation
            for i in cutlass.range_constexpr(k_tile_size // sf_vec_size):
                sfBlock = cute.make_rmem_tensor((2,), c_dtype)
                sfBlock.fill(0.0)
                offset = i * sf_vec_size
                for j in cutlass.range_constexpr(sf_vec_size // 2):
                    offset_ele = offset + j * 2
                    sfBlock[0], sfBlock[1] = fma_f16x2(
                        (tArA[offset_ele], tArA[offset_ele + 1]),
                        (tBrB[offset_ele], tBrB[offset_ele + 1]),
                        (sfBlock[0], sfBlock[1]),
                    )
                sf = tCrSFA[i] * tCrSFB[i]
                res += sfBlock[0] * sf + sfBlock[1] * sf

        out = warp_reduce_sum(res, width=num_col_threads)
        if tidx == 0:
            tCgC.store(out.to(c_dtype))
        return

    @cute.jit
    def __call__(
        self,
        a_ptr: cute.Pointer,
        b_ptr: cute.Pointer,
        sfa_ptr: cute.Pointer,
        sfb_ptr: cute.Pointer,
        c_ptr: cute.Pointer,
        problem_size: tuple,
        num_row_threads: cutlass.Constexpr = 4,
        num_col_threads: cutlass.Constexpr = 32,
        k_tile_size: cutlass.Constexpr = 128,
    ):
        """
        Host-side JIT function to prepare tensors and launch GPU kernel.
        """
        mma_tiler_mnk = (num_row_threads, 1, k_tile_size)
        m, _, k, l = problem_size
        # Create CuTe Tensor via pointer and problem size.
        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)),
            ),
        )
        # We use n=128 to create the torch tensor to do fp4 computation via torch._scaled_mm
        # then copy torch tensor to cute tensor for cute customize kernel computation
        # therefore we need to ensure b_tensor has the right stride with this 128 padded size on n.
        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 = cute.make_tensor(
            c_ptr, cute.make_layout((cute.assume(m, 32), 1, l), stride=(1, 1, m))
        )
        # Convert scale factor tensors to MMA layout
        # The layout matches Tensor Core requirements: (((32, 4), REST_M), ((SF_K, 4), REST_K), (1, REST_L))
        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)

        # Extract the local tile for input matrix A (shape: [block_M, block_K, rest_M, rest_K, rest_L])
        gA_mkl = cute.local_tile(
            a_tensor,
            cute.slice_(mma_tiler_mnk, (None, 0, None)),
            (None, None, None),
        )
        # Extract the local tile for scale factor tensor for A (same shape as gA_mkl)
        # Here, block_M = (32, 4); block_K = (16, 4)
        gSFA_mkl = cute.local_tile(
            sfa_tensor,
            cute.slice_(mma_tiler_mnk, (None, 0, None)),
            (None, None, None),
        )
        # Extract the local tile for input matrix B (shape: [block_N, block_K, rest_N, rest_K, rest_L])
        gB_nkl = cute.local_tile(
            b_tensor,
            cute.slice_(mma_tiler_mnk, (0, None, None)),
            (None, None, None),
        )
        # Extract the local tile for scale factor tensor for B (same shape as gB_nkl)
        gSFB_nkl = cute.local_tile(
            sfb_tensor,
            cute.slice_(mma_tiler_mnk, (0, None, None)),
            (None, None, None),
        )
        # Extract the local tile for output matrix C (shape: [block_M, block_N, rest_M, rest_N, rest_L])
        gC_mnl = cute.local_tile(
            c_tensor,
            cute.slice_(mma_tiler_mnk, (None, None, 0)),
            (None, None, None),
        )

        # Compute grid dimensions
        # Grid is (M_blocks, 1, L) where:
        # - M_blocks = ceil(M / 128) to cover all output rows
        # - L = batch size
        grid = (
            cute.ceil_div(c_tensor.shape[0], num_row_threads),
            1,
            c_tensor.shape[2],
        )

        # Launch the CUDA kernel
        self.kernel(
            gA_mkl,
            gB_nkl,
            gSFA_mkl,
            gSFB_nkl,
            gC_mnl,
            num_col_threads,
            k_tile_size,
        ).launch(
            grid=grid,
            block=[num_col_threads, num_row_threads, 1],
            cluster=(1, 1, 1),
        )
        return


# Global cache for compiled kernel
_compiled_kernel_cache = {}

benchmark_problem_sizes = [
    (7168, 1, 16384, 1),
    (4096, 1, 7168, 8),
    (7168, 1, 2048, 4),
]


# This function is used to compile the kernel once and cache it and then allow users to
# run the kernel multiple times to get more accurate timing results.
def compile_kernel(a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, problem_size):
    """
    Compile the kernel once and cache it.
    This should be called before any timing measurements.

    Returns:
        The compiled kernel function
    """
    global _compiled_kernel_cache

    if problem_size not in benchmark_problem_sizes:
        problem_size = (0, 0, 0, 0)

    if problem_size in _compiled_kernel_cache:
        return _compiled_kernel_cache[problem_size]

    gemv = Nvfp4BatchedGemvKernel()

    def tune_func(
        a_ptr,
        b_ptr,
        sfa_ptr,
        sfb_ptr,
        c_ptr,
        problem_size,
        num_row_threads=4,
        num_col_threads=32,
        k_tile_size=128,
    ):
        compiled_func = cute.compile(
            gemv,
            a_ptr,
            b_ptr,
            sfa_ptr,
            sfb_ptr,
            c_ptr,
            problem_size,
            num_row_threads,
            num_col_threads,
            k_tile_size,
        )
        return lambda: compiled_func(
            a_ptr,
            b_ptr,
            sfa_ptr,
            sfb_ptr,
            c_ptr,
            problem_size,
        )

    auto_tune = False
    if problem_size == (7168, 1, 16384, 1):
        if auto_tune:
            params = testing.tune(
                tune_func,
                params_dict={
                    "num_row_threads": [8, 16, 32],
                    "num_col_threads": [4, 8, 16],
                    "k_tile_size": [32, 64, 128],
                },
                kernel_arguments=testing.JitArguments(
                    a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size
                ),
            )
        else:
            params = {
                "num_row_threads": 8,
                "num_col_threads": 32,
                "k_tile_size": 32,
            }
    elif problem_size == (4096, 1, 7168, 8):
        if auto_tune:
            params = testing.tune(
                tune_func,
                params_dict={
                    "num_row_threads": [8, 16, 32],
                    "num_col_threads": [4, 8, 16],
                    "k_tile_size": [32, 64, 128],
                },
                kernel_arguments=testing.JitArguments(
                    a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size
                ),
            )
        else:
            params = {
                "num_row_threads": 32,
                "num_col_threads": 8,
                "k_tile_size": 64,
            }
    elif problem_size == (7168, 1, 2048, 4):
        if auto_tune:
            params = testing.tune(
                tune_func,
                params_dict={
                    "num_row_threads": [8, 16, 32],
                    "num_col_threads": [4, 8, 16],
                    "k_tile_size": [32, 64, 128],
                },
                kernel_arguments=testing.JitArguments(
                    a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size
                ),
            )
        else:
            params = {
                "num_row_threads": 32,
                "num_col_threads": 8,
                "k_tile_size": 64,
            }
    else:
        params = {
            "num_row_threads": 4,
            "num_col_threads": 16,
            "k_tile_size": 32,
        }
    print(f"The best kernel configs found: {params}")
    # Compile the kernel
    _compiled_kernel_cache[problem_size] = cute.compile(
        gemv, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size, **params
    )

    return _compiled_kernel_cache[problem_size]


def custom_kernel(data: input_t) -> output_t:
    """
    Execute the block-scaled GEMV kernel.

    This is the main entry point called by the evaluation framework.
    It converts PyTorch tensors to CuTe tensors, launches the kernel,
    and returns the result.

    Args:
        data: Tuple of (a, b, sfa_cpu, sfb_cpu, c) PyTorch tensors
            a: [m, k, l] - Input matrix in float4e2m1fn
            b: [1, k, l] - Input vector in float4e2m1fn
            sfa_cpu: [m, k, l] - Scale factors in float8_e4m3fn
            sfb_cpu: [1, k, l] - Scale factors in float8_e4m3fn
            sfa_permuted: [32, 4, rest_m, 4, rest_k, l] - Scale factors in float8_e4m3fn
            sfb_permuted: [32, 4, rest_n, 4, rest_k, l] - Scale factors in float8_e4m3fn
            c: [m, 1, l] - Output vector in float16

    Returns:
        Output tensor c with computed GEMV results
    """
    a, b, _, _, sfa_permuted, sfb_permuted, c = data

    # Get dimensions from MxKxL layout
    m, k, l = a.shape
    # Torch use e2m1_x2 data type, thus k is halved
    k = k * 2
    # GEMV N dimension is always 1
    n = 1

    problem_size = (m, n, k, l)

    # Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
    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)
    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
    )

    compiled_func = compile_kernel(a_ptr, b_ptr, c_ptr, sfa_ptr, sfb_ptr, problem_size)

    # Execute the compiled kernel
    compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, problem_size)

    return c