submission 183399

import torch
import triton
import triton.language as tl


@triton.jit
def block_scales_kernel(
    input_ptr, output_ptr,
    M, K_scales, L,
    stride_m, stride_k, stride_l,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_k = tl.program_id(1)
    pid_l = tl.program_id(2)

    rm = tl.arange(0, BLOCK_SIZE_M)
    rk = tl.arange(0, BLOCK_SIZE_K)

    m_offsets = pid_m * BLOCK_SIZE_M + rm
    k_offsets = pid_k * BLOCK_SIZE_K + rk

    mask = (m_offsets[:, None] < M) & (k_offsets[None, :] < K_scales)

    src_ptr = (
        input_ptr
        + pid_l * stride_l
        + m_offsets[:, None] * stride_m
        + k_offsets[None, :] * stride_k
    )
    scales = tl.load(src_ptr, mask=mask, other=0.0)

    # Blackwell 6D row mapping
    mm32 = rm % 32
    mm4 = rm // 32

    # 128x4 tile laid out contiguously
    dst_internal_offset = mm32[:, None] * 16 + mm4[:, None] * 4 + rk[None, :]

    block_offset = (
        pid_l * M * K_scales
        + pid_m * BLOCK_SIZE_M * K_scales
        + pid_k * BLOCK_SIZE_M * BLOCK_SIZE_K
    )

    dst_ptr = output_ptr + block_offset + dst_internal_offset
    tl.store(dst_ptr, scales, mask=mask)


def triton_to_blocked(scales_tensor: torch.Tensor) -> torch.Tensor:
    """
    scales_tensor: [M_or_N, K_scales, L]
    returns: [M_or_N * K_scales * L] in the NVFP4 blocked layout
    expected by torch._scaled_mm.
    """
    M, K_s, L = scales_tensor.shape
    output = torch.empty(
        M * K_s * L,
        device=scales_tensor.device,
        dtype=scales_tensor.dtype,
    )

    grid = (triton.cdiv(M, 128), triton.cdiv(K_s, 4), L)

    block_scales_kernel[grid](
        scales_tensor, output,
        M, K_s, L,
        scales_tensor.stride(0),
        scales_tensor.stride(1),
        scales_tensor.stride(2),
        BLOCK_SIZE_M=128,
        BLOCK_SIZE_K=4,
    )
    return output


def custom_kernel(data):
    """
    Leaderboard kernel.

    Accepts:
      (a, b, sfa, sfb, c)
      or
      (a, b, sfa, sfb, sfa_perm, sfb_perm, c)

    Shapes:
      a   : [M, K, L], float4_e2m1fn_x2
      b   : [N, K, L], float4_e2m1fn_x2
      sfa : [M, K_s, L], float8_e4m3fn
      sfb : [N, K_s, L], float8_e4m3fn
      c   : [M, N, L], float16
    """
    if len(data) >= 7:
        a, b, sfa_raw, sfb_raw, _, _, c = data
    else:
        a, b, sfa_raw, sfb_raw, c = data

    device = a.device
    _, _, L = c.shape

    sfa_gpu = sfa_raw.to(device)
    sfb_gpu = sfb_raw.to(device)

    # Convert scale factors to blocked layout once on GPU
    sfa_blocked = triton_to_blocked(sfa_gpu)  # [M * K_s * L]
    sfb_blocked = triton_to_blocked(sfb_gpu)  # [N * K_s * L]

    s_size_a = sfa_gpu.shape[0] * sfa_gpu.shape[1]  # M * K_s
    s_size_b = sfb_gpu.shape[0] * sfb_gpu.shape[1]  # N * K_s

    for i in range(L):
        sfa_i = sfa_blocked[i * s_size_a:(i + 1) * s_size_a]
        sfb_i = sfb_blocked[i * s_size_b:(i + 1) * s_size_b]

        c[:, :, i] = torch._scaled_mm(
            a[:, :, i],           # [M, K]
            b[:, :, i].t(),       # [K, N]
            sfa_i,                # blocked scales for A
            sfb_i,                # blocked scales for B
            out_dtype=torch.float16,
        )

    return c