submission 109475

import torch
import math
import operator
from task import input_t, output_t
from typing import Callable, Tuple

import cutlass
import cutlass.cute as cute
from cutlass.cute.tensor import TensorSSA
from cutlass.cute.runtime import make_ptr
import cutlass.utils.blockscaled_layout as blockscaled_utils

from cutlass._mlir.dialects import nvvm, llvm, builtin
from cutlass.cute.nvgpu import cpasync

from cutlass.cutlass_dsl import (
    dsl_user_op,
    T,
)

from cutlass._mlir import ir
from cutlass._mlir.dialects.cute import ReductionOp as ReductionOp
from cutlass._mlir.dialects import vector, arith

# 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
sf_vec_size = 16  # Scale factor block size (16 elements share one scale)

@cute.jit
def warp_reduce(
    val: cute.TensorSSA | cute.Numeric,
    op: Callable,
    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(res[i], op, width)
        return res.load()
    else:
        for i in cutlass.range_constexpr(int(math.log2(width))):
            val = op(val, cute.arch.shuffle_sync_bfly(val, offset=1 << i))
    return val

@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(cutlass.Int16.mlir_type, src_vec2, loc=loc, ip=ip)
    rst_i32 = llvm.inline_asm(
        cutlass.Int32.mlir_type,
        [src_i16],
        """{\n\t
            cvt.rn.f16x2.e4m3x2 $0, $1;\n\t
        }""",
        "=r,h",
    )
    vec_f16x2_type = ir.VectorType.get([2], cutlass.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(cutlass.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], cutlass.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], cutlass.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], cutlass.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(cutlass.Int32.mlir_type, 0), loc=loc, ip=ip)
    src_hi = llvm.extractelement(src_i32x2, arith.constant(cutlass.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], cutlass.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], cutlass.Float16.mlir_type, loc=loc)
    vec_f16x8 = llvm.bitcast(vec_f16x8_type, vec_i32x4, loc=loc, ip=ip)
    return vec_f16x8

@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], cutlass.Int8.mlir_type, loc=loc)],
        [vec_f8e4m3],
        loc=loc,
        ip=ip,
    )
    vec_i8x8_type = ir.VectorType.get([8], cutlass.Int8.mlir_type, loc=loc)
    vec_i8x4_type = ir.VectorType.get([4], cutlass.Int8.mlir_type, loc=loc)
    vec_i8x2_type = ir.VectorType.get([2], cutlass.Int8.mlir_type, loc=loc)
    vec_dst_type = ir.VectorType.get([length], cutlass.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(cutlass.Int32.mlir_type, src_pos),
                loc=loc,
                ip=ip,
            ),
            loc=loc,
            ip=ip,
        )
        vec_dst = vector.insertelement(
            val_f16,
            vec_dst,
            position=arith.constant(cutlass.Int32.mlir_type, src_pos),
            loc=loc,
            ip=ip,
        )

    return vec_dst


@dsl_user_op
def fma_f16x2(
    a: Tuple[cutlass.Float16, cutlass.Float16],
    b: Tuple[cutlass.Float16, cutlass.Float16],
    c: Tuple[cutlass.Float16, cutlass.Float16],
    *,
    loc=None,
    ip=None,
) -> Tuple[cutlass.Float16, cutlass.Float16]:
    # Pack two Float16 values into vector<2xf16>
    vec_type = ir.VectorType.get([2], cutlass.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(cutlass.Int32.mlir_type, vec_a, loc=loc, ip=ip)
    b_i32 = llvm.bitcast(cutlass.Int32.mlir_type, vec_b, loc=loc, ip=ip)
    c_i32 = llvm.bitcast(cutlass.Int32.mlir_type, vec_c, loc=loc, ip=ip)

    result_i32 = llvm.inline_asm(
        cutlass.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 = cutlass.Float16(
        vector.extract(vec_result, dynamic_position=[], static_position=[0], loc=loc, ip=ip)
    )
    result1 = cutlass.Float16(
        vector.extract(vec_result, dynamic_position=[], static_position=[1], loc=loc, ip=ip)
    )

    return result0, result1

@dsl_user_op
def dot_f16xN_with_fma_f16x2(
    v1,
    v2,
    length: int,
    acc: cutlass.Float16,
    *,
    loc=None,
    ip=None,
) -> cutlass.Float16:
    """
    Compute acc + sum_{i=0}^{length-1} v1[i] * v2[i],
    using the fma.rn.f16x2 intrinsic via fma_f16x2.

    Assumes:
      - length is even
      - v1, v2 are 1D vectors of f16 with at least `length` elements
    """

    # Treat v1, v2 as already-typed MLIR vectors.
    # If a cast is still needed, just cast to their existing type instead of constructing a new one.
    vec_type = v1.type

    vec_v1 = builtin.unrealized_conversion_cast(
        [vec_type],
        [v1],
        loc=loc,
        ip=ip,
    )
    vec_v2 = builtin.unrealized_conversion_cast(
        [vec_type],
        [v2],
        loc=loc,
        ip=ip,
    )

    # Two-lane accumulator: (acc, 0.0)
    acc0 = acc
    acc1 = cutlass.Float16(
        llvm.mlir_zero(cutlass.Float16.mlir_type, loc=loc, ip=ip)
    )

    # Walk the first `length` elements in pairs: (0,1), (2,3), ...
    for i in range(0, length, 2):
        # v1[i], v1[i+1]
        a0 = cutlass.Float16(
            vector.extract(
                vec_v1,
                dynamic_position=[],
                static_position=[i],
                loc=loc,
                ip=ip,
            )
        )
        a1 = cutlass.Float16(
            vector.extract(
                vec_v1,
                dynamic_position=[],
                static_position=[i + 1],
                loc=loc,
                ip=ip,
            )
        )

        # v2[i], v2[i+1]
        b0 = cutlass.Float16(
            vector.extract(
                vec_v2,
                dynamic_position=[],
                static_position=[i],
                loc=loc,
                ip=ip,
            )
        )
        b1 = cutlass.Float16(
            vector.extract(
                vec_v2,
                dynamic_position=[],
                static_position=[i + 1],
                loc=loc,
                ip=ip,
            )
        )

        # FMA in packed f16x2: (acc0, acc1) = (a0, a1)*(b0, b1) + (acc0, acc1)
        acc0, acc1 = fma_f16x2(
            (a0, a1),
            (b0, b1),
            (acc0, acc1),
            loc=loc,
            ip=ip,
        )

    # Horizontal add: final = acc0 + acc1
    res_val = arith.addf(
        acc0.ir_value(loc=loc, ip=ip),
        acc1.ir_value(loc=loc, ip=ip),
        loc=loc,
        ip=ip,
    )

    return cutlass.Float16(res_val)

@dsl_user_op
def set_block_rank(
    smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: cute.Int32, *, loc=None, ip=None
) -> cutlass.Int32:
    """Map the given smem pointer to the address at another CTA rank in the cluster."""
    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
    return cutlass.Int32(
        llvm.inline_asm(
            T.i32(),
            [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()],
            "mapa.shared::cluster.u32 $0, $1, $2;",
            "=r,r,r",
            has_side_effects=False,
            is_align_stack=False,
            asm_dialect=llvm.AsmDialect.AD_ATT,
        )
    )

@dsl_user_op
def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer:
    return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)

@cute.jit
def cluster_all_gather_128(
    buffer: cute.Tensor,         # shape (128, cluster_n)
    mbar_ptr: cute.Pointer,      # cluster mbarrier object in smem
    cluster_n: int,
    *,
    loc=None,
    ip=None,
) -> None:
    # Thread / block identifiers
    tidx, _, _ = cute.arch.thread_idx()
    cta_rank  = cute.arch.block_idx_in_cluster()

    # Load the element that this CTA owns in its column
    my_val = buffer[tidx % 128, cta_rank]

    # Pointer to (row = lane_idx, col = cta_rank) within *this* CTA's SMEM
    local_elem_ptr = elem_pointer(
        buffer,
        (tidx % 128, cta_rank)
    )

    rank_mod8 = tidx // 128

    # Send this element to every CTA in the cluster
    for peer_rank_div8 in cutlass.range_constexpr(cute.ceil_div(cluster_n, 8)):
        peer_rank = peer_rank_div8 * 8 + rank_mod8
        if peer_rank < cluster_n:
            wrapped_peer_rank = cute.Int32(peer_rank)
            # Compute remote (peer CTA) pointer for this (row, col) location
            remote_smem_ptr_i32 = set_block_rank(
                local_elem_ptr,
                wrapped_peer_rank,
            ).ir_value()

            remote_mbar_ptr_i32 = set_block_rank(
                mbar_ptr,
                wrapped_peer_rank,
            ).ir_value()

            # Perform DSMEM async store into the peer CTA’s shared memory
            llvm.inline_asm(
                None,
                [
                    remote_smem_ptr_i32,       # $0 = remote smem address
                    cutlass.Float32(my_val).ir_value(),  # $1 = value
                    remote_mbar_ptr_i32        # $2 = peer mbar pointer
                ],
                "st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 "
                "[$0], $1, [$2];",
                "r,f,r",
                has_side_effects=True,
                is_align_stack=False,
                asm_dialect=llvm.AsmDialect.AD_ATT,
            )

    # ======== Synchronize all async DSMEM stores =========
    # mbarrier_wait() ensures that *this CTA's* SMEM has received all bytes
    cute.arch.mbarrier_wait(mbar_ptr, phase=0)

def _get_shared_storage_cls(
    N_ROWS_PER_CTA,
    N_COLS_PER_CTA,
    cluster_size,
):
    # TODO: Check, they probably don't need 1024 byte alignment?

    # (N_ROWS_PER_CTA, N_COLS_PER_CTA)
    sA_struct = cute.struct.Align[cute.struct.MemRange[ab_dtype, N_ROWS_PER_CTA * N_COLS_PER_CTA], 1024]
    sB_struct = cute.struct.Align[cute.struct.MemRange[ab_dtype,  N_COLS_PER_CTA], 1024]
    sAScales_struct = cute.struct.Align[cute.struct.MemRange[sf_dtype, N_ROWS_PER_CTA * N_COLS_PER_CTA // sf_vec_size], 1024]
    sBScales_struct = cute.struct.Align[cute.struct.MemRange[sf_dtype, N_COLS_PER_CTA // sf_vec_size], 1024]
    # For cluster reduce
    sC_struct = cute.struct.Align[cute.struct.MemRange[cutlass.Float32, N_ROWS_PER_CTA * cluster_size], 1024]
    # Maybe needs to be 16?
    mbar_struct = cute.struct.Align[cute.struct.MemRange[cutlass.Int64, 8], 1024]

    @cute.struct
    class SharedStorage:
        sA: sA_struct
        sB: sB_struct
        sAScales: sAScales_struct
        sBScales: sBScales_struct
        sC: sC_struct
        mbar: mbar_struct

    return SharedStorage

# Define layouts, atoms, etc.
# n_padded --> 128 for b
@cute.jit
def custom_launcher(
    a_ptr: cute.Pointer,
    b_ptr: cute.Pointer,
    sfa_ptr: cute.Pointer,
    sfb_ptr: cute.Pointer,
    c_ptr: cute.Pointer,
    problem_size: cutlass.Constexpr[tuple], # constexpr
):
    # CONSTANTS
    m, n, k, l = problem_size
    # Each block computes (128, a subset of k) --> use GMEM atomics for the full add
    # TODO: Change so that N_COLS_PER_CTA is max(1024, smallest power of 2 dividing k that splits k into <= 8 pieces)
    N_ROWS_PER_CTA, N_COLS_PER_CTA = 128, 1024# k // 8
    THREADS_PER_CTA = 1024

    unrounded_cluster_size = cute.ceil_div(k, N_COLS_PER_CTA)
    cluster_size = 0
    if cutlass.const_expr(unrounded_cluster_size > 8):
        cluster_size = 16
    elif cutlass.const_expr(unrounded_cluster_size > 4):
        cluster_size = 8
    elif cutlass.const_expr(unrounded_cluster_size > 2):
        cluster_size = 4
    elif cutlass.const_expr(unrounded_cluster_size > 1):
        cluster_size = 2
    else:
        cluster_size = 1
    # import pdb; pdb.set_trace()

    # COPY ATOMS
    copy_bits = 128
    universal_copy_atom_fp4 = cute.make_copy_atom(
        # op=cute.nvgpu.CopyUniversalOp(),
        op=cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
        copy_internal_type=ab_dtype,
        num_bits_per_copy=copy_bits,
    )
    universal_copy_atom_fp8 = cute.make_copy_atom(
        # op=cute.nvgpu.CopyUniversalOp(),
        op=cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
        copy_internal_type=sf_dtype,
        num_bits_per_copy=copy_bits,
    )
    half_width_copy_atom_fp8 = cute.make_copy_atom(
        op=cute.nvgpu.CopyUniversalOp(),
        # op=cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
        copy_internal_type=sf_dtype,
        num_bits_per_copy=copy_bits // 2,
    )


    sfb_copy_atom = cute.make_copy_atom(
        op=cute.nvgpu.CopyUniversalOp(),
        # op=cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
        copy_internal_type=sf_dtype,
        num_bits_per_copy=32,
    )

    # sfA "contiguous blocks" are 128 x 4 regions, we will take first n_rows_per_cta of those 128
    # they expand to 128 x 64 sized chunks --> we need to copy from k // 64 chunks per CTA
    # ^ not entirely true (4x4 regions interleaved, but a 128x4 region can be loaded contiguously)


    # MATRIX TENSORS
    a_tensor = cute.make_tensor(
        a_ptr,
        cute.make_layout(
            (m, k, l),
            stride=(k, 1, m * k),
        ),
    )
    n_padded_128 = 128
    b_tensor = cute.make_tensor(
        b_ptr,
        cute.make_layout(
            (n_padded_128, k, l),
            stride=(k, 1, (n_padded_128 * k)),
        ),
    )

    c_tensor = cute.make_tensor(
        c_ptr, cute.make_layout((m, 1, l, 16), stride=(1, 1, m, m * l))
    )

    # SCALE TENSORS
    sfa_layout = cute.make_layout(
        shape=((128, cute.ceil_div(m, 128)), (4, cute.ceil_div(k, 4 * sf_vec_size)), l),
        stride=((4, 128 * k // sf_vec_size), (1, 4 * 128), k * m // sf_vec_size),
    )

    # Have to use entire 128...
    sfb_layout = cute.make_layout(
        shape=(128, (4, cute.ceil_div(k, 4 * sf_vec_size)), l),
        stride=(4, (1, 4 * 128), k * 128 // sf_vec_size),
    )

    sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
    sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)

    SharedStorage = _get_shared_storage_cls(N_ROWS_PER_CTA, N_COLS_PER_CTA, cluster_size)

    # (32, 1024) blocks
    elems_per_copy_fp4 = copy_bits // 4
    gA_tiled_copy = cute.make_tiled_copy_tv(
        universal_copy_atom_fp4,
        thr_layout = cute.make_ordered_layout((32, 32), order=((1, 0))),
        val_layout = cute.make_layout((1, elems_per_copy_fp4)),
    )

    # (1024 * 32) blocks --> need to predicate
    gB_tiled_copy = cute.make_tiled_copy_tv(
        universal_copy_atom_fp4,
        thr_layout=cute.make_layout((1024,)),
        val_layout=cute.make_layout((elems_per_copy_fp4,)),
    )

    DIV_M = DIV_K = 1024

    # Want 2x4 squares for each thread (for now)
    # tiler_mn becomes (128, 64) --> no pred since 64 * sf_vec_size | 1024
    # AScales_threads_per_col = DIV_K // (4 * sf_vec_size) # 16
    # AScales_threads_per_row = 1024 // AScales_threads_per_col # 64
    # AScales_layout_tv = cute.make_layout(
    #     shape=((AScales_threads_per_col, AScales_threads_per_row), (2, 4)),
    #     stride=((2, 4 * N_ROWS_PER_CTA), (1, N_ROWS_PER_CTA)),
    # )

    # Not great...
    AScales_threads_per_col = DIV_K // (2 * 4 * sf_vec_size) # 16
    AScales_threads_per_row = 1024 // AScales_threads_per_col # 64
    AScales_layout_tv = cute.make_layout(
        shape=(AScales_threads_per_col * AScales_threads_per_row, 16),
        stride=(16, 1),
    )

    gAScales_tiled_copy = cute.make_tiled_copy(
        # universal_copy_atom_fp8,
        half_width_copy_atom_fp8,
        layout_tv=AScales_layout_tv,
        tiler_mn=(AScales_threads_per_row * AScales_threads_per_col * 8,),
    )

    # Leave gaps in this layout, store contiguously
    # gSFB_layout = cute.make_layout(shape=(1, DIV_K // (4 * sf_vec_size)), stride=(1, 4))
    gSFB_layout = cute.make_layout(shape=(4, N_COLS_PER_CTA // (4 * sf_vec_size)), stride=(1, 4))

    # CTA with (bidx, bidy, bidz) where bidy == 0 (for now)
    # should index sfa_tensor at (((bidx % cta_per_chunk_m, all), bidx // cta_per_chunk_m), all, bidz)
    grid = (
        cluster_size,
        cute.ceil_div(m, N_ROWS_PER_CTA),
        l, # batch size
    )

    # print(f"{grid=} {SharedStorage.size_in_bytes()=} {THREADS_PER_CTA=}")

    # Launch the CUDA kernel
    new_kernel(
        a_tensor,
        b_tensor,
        sfa_tensor,
        sfb_tensor,
        c_tensor,
        gA_tiled_copy,
        gAScales_tiled_copy,
        gB_tiled_copy,
        gSFB_layout,
        sfb_copy_atom,
        N_ROWS_PER_CTA,
        N_COLS_PER_CTA,
        unrounded_cluster_size,
        cluster_size,
        problem_size,
        SharedStorage,
    ).launch(
        grid=grid,
        block=[THREADS_PER_CTA, 1, 1],
        smem=SharedStorage.size_in_bytes(),
        cluster=[cluster_size, 1, 1],
    )

@cute.kernel
def new_kernel(
    mA: cute.Tensor,
    mB: cute.Tensor,
    mSFA: cute.Tensor,
    mSFB: cute.Tensor,
    mC: cute.Tensor,
    gA_tiled_copy: cute.TiledCopy,
    gAScales_tiled_copy: cute.TiledCopy,
    gB_tiled_copy: cute.TiledCopy,
    gSFB_layout: cute.Layout,
    sfb_copy_atom: cute.CopyAtom,
    N_ROWS_PER_CTA: cutlass.Constexpr[int],
    N_COLS_PER_CTA: cutlass.Constexpr[int],
    UNROUNDED_CLUSTER_SIZE: cutlass.Constexpr[int],
    CLUSTER_SIZE: cutlass.Constexpr[int],
    problem_size: cutlass.Constexpr[tuple],
    SharedStorage: cutlass.Constexpr,
):
    m, n, k, l = problem_size
    # (row chunk, col chunk, batch idx)
    bidy, bidx, bidz = cute.arch.block_idx()
    tidx, _, _ = cute.arch.thread_idx()
    cta_rank = cute.arch.block_idx_in_cluster()

    out = cutlass.Float16(0.0)

    # Initial SMEM layouts (for copy purposes)
    smem = cutlass.utils.SmemAllocator()
    storage = smem.allocate(SharedStorage)
    sC = cute.make_tensor(
        iterator=storage.sC.data_ptr(),
        layout=cute.make_ordered_layout((N_ROWS_PER_CTA, CLUSTER_SIZE), (0, 1)),
    )

    mbar_ptr = storage.mbar.data_ptr()

    # CTA-level matrices
    blkA_shape = (N_ROWS_PER_CTA, N_COLS_PER_CTA)
    gA = cute.local_tile(
        mA, blkA_shape, (bidx, bidy, bidz)
    )
    blkB_shape = (1, N_COLS_PER_CTA)
    gB = cute.local_tile(
        mB, blkB_shape, (0, bidy, bidz)
    )
    gB = gB[0, None]
    cB = cute.make_identity_tensor(shape=(N_COLS_PER_CTA,))

    blkC_shape = (N_ROWS_PER_CTA, 1)
    gC = cute.local_tile(
        mC, blkC_shape, (bidx, 0, bidz, bidy),
    )
    gC = gC[None, 0]

    # CTA-level scales
    blkAScale_shape = (N_ROWS_PER_CTA, N_COLS_PER_CTA // sf_vec_size)
    gSFA = cute.local_tile(
        mSFA, blkAScale_shape,(bidx, bidy, bidz)
    )

    # No predication needed anyway
    gSFA = cute.make_tensor(
        gSFA.iterator,
        cute.make_layout(
            shape=(blkAScale_shape[0] * blkAScale_shape[1],),
            stride=(1,)
        )
    )

    # cSFA = cute.make_identity_tensor(shape=blkAScale_shape)
    # gSFA = (128, (4, 32)) : (4, (1, 512)) for problem shapes 1
    blkBScale_shape = (1, N_COLS_PER_CTA // sf_vec_size)
    gSFB = cute.local_tile(
        mSFB, blkBScale_shape, (0, bidy, bidz)
    )
    gSFB = gSFB[0, None]
    # gSFB = ((4, 32)) : (1, 512) for problem shape 1

    sA = cute.make_tensor(
            iterator=storage.sA.data_ptr(),
            layout=cute.make_ordered_layout(gA.shape, (1, 0)),
    )
    sB = cute.make_tensor(
            iterator=storage.sB.data_ptr(),
            layout=gB.layout,
    )
    sAScales = cute.make_tensor(
            iterator=storage.sAScales.data_ptr(),
            layout=gSFA.layout,
    )
    sBScales = cute.make_tensor(
            iterator=storage.sBScales.data_ptr(),
            layout=gSFB_layout,
    )

    # A: Trivial copy, no predication needed since tiles are nicely divisible
    gA_thr_copy = gA_tiled_copy.get_slice(tidx)
    tAgA = gA_thr_copy.partition_S(gA)
    tAsA = gA_thr_copy.partition_D(sA)

    if bidy < UNROUNDED_CLUSTER_SIZE:
        cute.copy(gA_thr_copy, tAgA, tAsA)

    # B: Need to predicate
    gB_thr_copy = gB_tiled_copy.get_slice(tidx)
    tBgB = gB_thr_copy.partition_S(gB)
    tBsB = gB_thr_copy.partition_D(sB)
    tBcB = gB_thr_copy.partition_S(cB)

    # Assumes at most 1 copy issue per thread
    if (
        # Only want to use a sub-tile of the tiled copy if k is too small
        (tBcB[0][0] < N_COLS_PER_CTA)
        # Tiled copy k-dim might not evenly divide k
        and (tBcB[0][0] + bidy * N_COLS_PER_CTA < k)
        and bidy < UNROUNDED_CLUSTER_SIZE
    ):
        cute.copy(gB_thr_copy, tBgB, tBsB)

    # SFA: Copying (128, k // (8 * 16)) elems
    gAScales_thr_copy = gAScales_tiled_copy.get_slice(tidx)
    tSFAgA = gAScales_thr_copy.partition_S(gSFA)
    tSFAsA = gAScales_thr_copy.partition_D(sAScales)
    # tSFAcA = gAScales_thr_copy.partition_S(cSFA)

    # Now doing 64 bit copies
    # No predication neded, since tiler_mn[1] | 1024
    if tidx < 512 and bidy < UNROUNDED_CLUSTER_SIZE:
        cute.copy(gAScales_thr_copy, tSFAgA, tSFAsA)

    # maybe cycle from back of thread list? lol...
    if tidx < cute.size(gSFB.shape[0][1]) and bidy < UNROUNDED_CLUSTER_SIZE:
        # needs to copy 4 fp8s
        cute.copy(sfb_copy_atom, gSFB[(None, tidx),], sBScales[(None, tidx)])

    cute.arch.cp_async_commit_group()
    cute.arch.cp_async_wait_group(0)
    cute.arch.barrier() # Need after cp.async

    sfa_blockscaled_layout = blockscaled_utils.tile_atom_to_shape_SF(blkA_shape + (1,), sf_vec_size)((None, None, 0))
    sfb_blockscaled_layout = blockscaled_utils.tile_atom_to_shape_SF(blkB_shape + (1,), sf_vec_size)((0, None, 0))[0]

    # Need to do surgery on sfb blockscaled layout to get stride matching
    sfb_blockscaled_layout = cute.make_layout(
        shape=sfb_blockscaled_layout.shape,
        stride=(
            sfb_blockscaled_layout.stride[0],
            4,
        )
    )

    # Blockscaled layouts for SFA, SFB
    sSFA_BS = cute.make_tensor(
            iterator=storage.sAScales.data_ptr(),
            layout=sfa_blockscaled_layout,
    )
    sSFB_BS = cute.make_tensor(
            iterator=storage.sBScales.data_ptr(),
            layout=sfb_blockscaled_layout,
    )

    blkAtile_shape = (128, 16)
    sA_tiled = cute.local_tile(
        sA, blkAtile_shape, (0, None)
    )
    sSFA_tiled = cute.local_tile(
        sSFA_BS, blkAtile_shape, (0, None)
    )
    blkBtile_shape = (16,)
    sB_tiled = cute.local_tile(
        sB, blkBtile_shape, (None,)
    )
    sSFB_tiled = cute.local_tile(
        cute.group_modes(sSFB_BS, 0, 2), blkBtile_shape, (None,)
    )

    tArA = cute.make_rmem_tensor_like(sA_tiled[0, None, 0])
    tSFArA = cute.make_rmem_tensor_like(sSFA_tiled[0, None, 0])
    tBrB = cute.make_rmem_tensor_like(sB_tiled[None, 0])
    tSFBrB = cute.make_rmem_tensor_like(sSFB_tiled[None, 0])

    num_k_tiles = N_COLS_PER_CTA // 128

    # Process a (128, 128) tile at a time?
    # 8 threads per row, 16 elems per thread per iter
    res = cutlass.Float16(0.0)

    if bidy < UNROUNDED_CLUSTER_SIZE:
        for k_tile in cutlass.range_constexpr(num_k_tiles):
            tcAsA = sA_tiled[tidx // 8, None, (tidx % 8) + 8 * k_tile]
            tcSFAsA = sSFA_tiled[tidx // 8, None, (tidx % 8) + 8 * k_tile]
            tcBsB = sB_tiled[None, (tidx % 8) + 8 * k_tile]
            tcSFBsB = sSFB_tiled[None, (tidx % 8) + 8 * k_tile]

            cute.autovec_copy(tcAsA, tArA)
            cute.autovec_copy(tcSFAsA, tSFArA)
            cute.autovec_copy(tcBsB, tBrB)
            cute.autovec_copy(tcSFBsB, tSFBrB)

            ld_tSFArA = tSFArA.load()

            scaleA = TensorSSA(
                cvt_f8e4m3_f16_intrinsic(ld_tSFArA, cute.size(tSFArA)),
                ld_tSFArA._shape,
                cutlass.Float16
            )

            ld_tSFBrB = tSFBrB.load()

            scaleB = TensorSSA(
                cvt_f8e4m3_f16_intrinsic(ld_tSFBrB, cute.size(tSFBrB)),
                ld_tSFBrB._shape,
                cutlass.Float16
            )

            tempA = tArA.load().to(cutlass.Float16) * scaleA
            tempB = tBrB.load().to(cutlass.Float16) * scaleB

            res = dot_f16xN_with_fma_f16x2(
                tempA,
                tempB,
                cute.size(tArA),
                res,
            )

        out = warp_reduce(res, operator.add, 8)
    # Now each thread that is 0 % 8 has the reduced value for that row.
    if tidx % 8 == 0:
        sC[tidx // 8, cta_rank] = out

    # might not need this?
    cute.arch.sync_threads()

    if tidx == 0:
        cute.arch.mbarrier_init(mbar_ptr, 1)

    cute.arch.mbarrier_init_fence()
    if tidx == 0:
        # initialize memory barrier transaction counter.
        cute.arch.mbarrier_arrive_and_expect_tx(
            mbar_ptr,
            N_ROWS_PER_CTA * CLUSTER_SIZE * 4,
        )
        # send an “arrive” signal after barrier init
    cute.arch.cluster_arrive_relaxed()
    # wait until all warps in the cluster have initialized their local reduction buffer in their SMEM.
    cute.arch.cluster_wait()

    cluster_all_gather_128(sC, mbar_ptr, CLUSTER_SIZE)

    # Split row by row, need CLUSTER_SIZE threads per row
    if cta_rank == 0:
        for i in cutlass.range_constexpr(cute.ceil_div(N_ROWS_PER_CTA, 1024 // CLUSTER_SIZE)):
            row = i * (1024 // CLUSTER_SIZE) + (tidx // CLUSTER_SIZE)
            col = tidx % CLUSTER_SIZE
            if row < N_ROWS_PER_CTA:
                result = sC[row, col]

                result = warp_reduce(result, operator.add, CLUSTER_SIZE)
                if col == 0:
                    gC[row] = cutlass.Float16(result)

# REF STUFF BELOW

# Only for ref kernel
ref_mma_tiler_mnk = (128, 1, 64)  # Tile sizes for M, N, K dimensions
ref_threads_per_cta = 128  # Number of threads per CUDA thread block

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

# The CuTe reference implementation for NVFP4 block-scaled GEMV
@cute.kernel
def ref_kernel(
    mA_mkl: cute.Tensor,
    mB_nkl: cute.Tensor,
    mSFA_mkl: cute.Tensor,
    mSFB_nkl: cute.Tensor,
    mC_mnl: cute.Tensor,
):
    # Get CUDA block and thread indices
    bidx, bidy, bidz = cute.arch.block_idx()
    tidx, _, _ = cute.arch.thread_idx()

    # Extract the local tile for input matrix A (shape: [block_M, block_K, rest_M, rest_K, rest_L])
    gA_mkl = cute.local_tile(
        mA_mkl, cute.slice_(ref_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(
        mSFA_mkl, cute.slice_(ref_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(
        mB_nkl, cute.slice_(ref_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(
        mSFB_nkl, cute.slice_(ref_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(
        mC_mnl, cute.slice_(ref_mma_tiler_mnk, (None, None, 0)), (None, None, None)
    )

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

    # 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(k_tile_cnt):
        tAgA = gA_mkl[tidx, None, bidx, k_tile, bidz]
        tBgB = gB_nkl[0, None, bidy, k_tile, bidz]
        tAgSFA = gSFA_mkl[tidx, None, bidx, k_tile, bidz]
        tBgSFB = gSFB_nkl[0, None, bidy, k_tile, bidz]

        tArA = cute.make_rmem_tensor_like(tAgA, cutlass.Float32)
        tBrB = cute.make_rmem_tensor_like(tBgB, cutlass.Float32)
        tArSFA = cute.make_rmem_tensor_like(tAgSFA, cutlass.Float32)
        tBrSFB = cute.make_rmem_tensor_like(tBgSFB, cutlass.Float32)

        # 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.Float32)
        b_val = b_val_nvfp4.to(cutlass.Float32)
        sfa_val = sfa_val_fp8.to(cutlass.Float32)
        sfb_val = sfb_val_fp8.to(cutlass.Float32)


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

        # Iterate over SF vector tiles and compute the scale&matmul accumulation
        for i in cutlass.range_constexpr(ref_mma_tiler_mnk[2]):
            res += tArA[i] * tArSFA[i] * tBrB[i] * tBrSFB[i]

    # Store the final float16 result back to global memory
    tCgC.store(res.to(cutlass.Float16))
    return


@cute.jit
def ref_launcher(
    a_ptr: cute.Pointer,
    b_ptr: cute.Pointer,
    sfa_ptr: cute.Pointer,
    sfb_ptr: cute.Pointer,
    c_ptr: cute.Pointer,
    problem_size: cutlass.Constexpr[tuple], # constexpr
):
    """
    Host-side JIT function to prepare tensors and launch GPU kernel.
    """
    m, _, k, l = problem_size

    # Create CuTe Tensor via pointer and problem size.
    a_tensor = cute.make_tensor(
        a_ptr,
        cute.make_layout(
            (m, k, l),
            stride=(k, 1, m * k),
        ),
    )
    # 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, k, l),
            stride=(k, 1, (n_padded_128 * k)),
        ),
    )
    c_tensor = cute.make_tensor(
        c_ptr, cute.make_layout((m, 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)

    # 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], 128),
        1,
        c_tensor.shape[2],
    )

    # Launch the CUDA kernel
    ref_kernel(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor).launch(
        grid=grid,
        block=[ref_threads_per_cta, 1, 1],
        cluster=(1, 1, 1),
    )
    return


# Global cache for compiled kernel
_compiled_kernel_cache = {}


# 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(m, n, k, l):
    """
    Compile the kernel once and cache it.
    This should be called before any timing measurements.

    Returns:
        The compiled kernel function
    """
    global _compiled_kernel_cache

    key = (m, n, k, l)

    if key in _compiled_kernel_cache:
        return _compiled_kernel_cache[key]

    # Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
    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 the kernel
    # if key == (7168, 1, 16384, 1):
    # if False:
    if key[0] % 1024 == 0 and key[2] % 1024 == 0:
        _compiled_kernel_cache[key] = cute.compile(
            custom_launcher, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l),
            # options="--generate-line-info"
        )
    else:
        _compiled_kernel_cache[key] = cute.compile(
            ref_launcher, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, (m, n, k, l)
        )

    return _compiled_kernel_cache[key]


def custom_kernel(data: input_t) -> output_t:
    a, b, _, _, sfa_permuted, sfb_permuted, c = data

    m, k, l = a.shape
    # Torch use e2m1_x2 data type, thus k is halved
    k = k * 2
    n = 1

    compiled_func = compile_kernel(m, n, k, l)

    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
    )

    # if (m, n, k, l) == (7168, 1, 16384, 1):
    if m % 1024 == 0 and k % 1024 == 0:
        compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr)
    else:
        compiled_func(a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr)

    return c

if __name__ == "__main__":
    # m = 7168
    n = 1
    # k = 16384
    # l = 1
    m = 7168
    k = 2048
    l = 4
    # input = generate_input(m, k, l, 0)
    # out = custom_kernel(input)
    # print(f"{out=}")
    # print(f"{out.shape=}")
    compile_kernel(m, n, k, l)