submission 698485

"""
FP4 quant + FP4 GEMM reference: bf16 A, MXFP4 B -> MXFP4 per-1x32 quant A -> gemm_a4w4 -> bf16 C.
Quant logic follows aiter op_tests/test_gemm_a4w4.py (get_triton_quant(QuantType.per_1x32)).
"""
from task import input_t, output_t
import torch
from torch import Tensor
import triton
import triton.language as tl

@triton.jit
def shuffle_8_32(
    a,
    b,
    n,
    m,
    tm,
    BLK_X: tl.constexpr,
    BLK_Y: tl.constexpr,
) :
    p_x = tl.program_id(axis=0)
    p_y = tl.program_id(axis=1)

    l_x = tl.arange(0, BLK_X)[:, None]
    l_y = tl.arange(0, BLK_Y)[None, :]
    x = p_x * BLK_X + l_x
    y = p_y * BLK_Y + l_y

    pos_a = y * n + x
    p_id = p_x * (tm // BLK_Y) + p_y

    l_id = l_y % (BLK_Y // 2) * (2 * BLK_X) + l_x % (BLK_X // 2) * 4 \
        + l_y // (BLK_Y // 2) * 2 + l_x // (BLK_X // 2)

    id = p_id * BLK_Y * BLK_X + l_id
    mask = (x < n) & (y < m)

    element = tl.load(a + pos_a, mask=mask)
    tl.store(b + id, element, mask=mask)

def e8m0_shuffle(scale):
    if scale is None:
        return scale
    if scale.dtype == torch.float32:
        return scale
    assert scale.ndim == 2, "scale must be a 2D tensor"
    n, m = scale.shape
    scale_padded = torch.empty(
        (n + 255) // 256 * 256,
        (m + 7) // 8 * 8,
        dtype=scale.dtype,
        device=scale.device,
    )
    _, tm = scale_padded.shape

    grid = (
        triton.cdiv(n, 32),
        triton.cdiv(m, 8)
    )
    shuffle_8_32[grid](scale, scale_padded, n, m, tm, BLK_X = 32, BLK_Y = 8)
    return scale_padded

def custom_kernel(data: input_t) -> output_t:
    """
    Reference: MXFP4 per-1x32 quant on A; B_shuffle, B_scale_sh from generate_input.
    gemm_a4w4 with bpreshuffle=True.
    """
    import aiter
    from aiter import QuantType, dtypes
    from aiter.ops.triton.quant import dynamic_mxfp4_quant 

    def _quant_mxfp4(x, shuffle=True):
        x_fp4, bs_e8m0 = dynamic_mxfp4_quant(x)
        if shuffle:
            bs_e8m0 = e8m0_shuffle(bs_e8m0)
        return x_fp4.view(dtypes.fp4x2), bs_e8m0.view(dtypes.fp8_e8m0)
    
    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    B = B.contiguous()
    m, k = A.shape
    n, _ = B.shape

    A_q, A_scale_sh = _quant_mxfp4(A, shuffle=True)
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    return out_gemm