submission 74098

# !POPCORN leaderboard ranked
import torch

SF_VEC = 16  # scale granularity (per 16 K elements)

def ceil_div(a: int, b: int) -> int:
    return (a + b - 1) // b

@torch.jit.script_if_tracing
def _to_blocked(sf_2d: torch.Tensor) -> torch.Tensor:
    """
    Convert FP8 scaling tensor from (rows, K//16) to the flattened CuTe/Blackwell
    blocked layout expected by torch._scaled_mm. View/permute/reshape only.
    """
    rows = sf_2d.size(0)
    sf_k = sf_2d.size(1)
    n_row_blocks = (rows + 127) // 128
    n_col_blocks = (sf_k + 3) // 4

    # [nrb,128,ncb,4] -> permute -> reshape -> [*,32,16] -> flatten
    t = sf_2d.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
    t = t.reshape(-1, 128, 4).view(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
    return t.flatten()

def _scaled_gemv_N1(A_mk: torch.Tensor,
                    b1k: torch.Tensor,
                    sfa_mk16: torch.Tensor,
                    sfb_1k16: torch.Tensor) -> torch.Tensor:
    """
    Fast path: use torch._scaled_mm with N=1 (no padding). Returns [M, 1] fp16.
    """
    scale_a = _to_blocked(sfa_mk16)
    scale_b = _to_blocked(sfb_1k16)
    outM1 = torch._scaled_mm(
        A_mk,              # [M,K]  (nvfp4)
        b1k,               # [1,K]  (nvfp4)
        scale_a,           # flattened
        scale_b,           # flattened
        bias=None,
        out_dtype=torch.float16,
    )  # -> [M,1] fp16
    return outM1

def _scaled_gemv_N128(A_mk: torch.Tensor,
                      b1k: torch.Tensor,
                      sfa_mk16: torch.Tensor,
                      sfb_1k16: torch.Tensor,
                      scratch_B128K: torch.Tensor,
                      scratch_SFB128: torch.Tensor) -> torch.Tensor:
    """
    Fallback path: pad N to 128 using reusable scratch buffers.
    Returns [M,1] in fp16 (as a narrowed view of the GEMM output).
    """
    K = A_mk.size(1)
    sfk = sfa_mk16.size(1)  # K//16

    # zero/one reset in-place (cheap)
    scratch_B128K.zero_()
    scratch_SFB128.fill_(1)

    # write real row-0 only
    scratch_B128K[0, :].copy_(b1k[0, :])
    scratch_SFB128[0, :].copy_(sfb_1k16[0, :])

    scale_a = _to_blocked(sfa_mk16)
    scale_b = _to_blocked(scratch_SFB128)

    outMN = torch._scaled_mm(
        A_mk,
        scratch_B128K,
        scale_a,
        scale_b,
        bias=None,
        out_dtype=torch.float16,
    )  # -> [M,128]
    return outMN[:, :1]  # keep the true N=1 column

def custom_kernel(data):
    """
    Supports NVFP4 batched GEMV:
      inputs: (a[M,K,L], b[1,K,L], sfa[M,K//16,L], sfb[1,K//16,L], c[M,1,L])

    Also gracefully handles 2-tensor practice checks (RGB->Gray) if the runner probes.
    """
    # Handle practice/warmup probes that pass (x, out)
    if len(data) == 2:
        x, out = data
        if x.ndim == 3 and x.shape[-1] == 3:
            w = torch.tensor([0.2989, 0.5870, 0.1140], device=x.device, dtype=x.dtype)
            out.copy_(torch.einsum("hwc,c->hw", x, w))
        else:
            out.copy_(x)
        return out

    # Ranked path (5 tensors)
    a, b, sfa, sfb, c = data
    M, K, L = a.shape
    assert b.shape[0] == 1
    assert a.dtype == torch.float4_e2m1fn_x2 and b.dtype == torch.float4_e2m1fn_x2
    assert c.dtype == torch.float16
    assert sfa.dtype in (torch.float8_e4m3fn, getattr(torch, "float8_e4m3fnuz", torch.float8_e4m3fn))
    assert sfb.dtype in (torch.float8_e4m3fn, getattr(torch, "float8_e4m3fnuz", torch.float8_e4m3fn))

    # Preallocate scratch for fallback (N=128). Reused for all L.
    N_PAD = 128
    scratch_B128K = torch.empty((N_PAD, K), device=b.device, dtype=b.dtype)
    scratch_SFB128 = torch.empty((N_PAD, K // SF_VEC), device=sfb.device, dtype=sfb.dtype)

    # Try the super-fast N=1 path once; if it errors, use padded path thereafter.
    use_N1 = True
    for l in range(L):
        A_l = a[:, :, l].contiguous()       # [M,K] nvfp4
        b_l = b[:, :, l].contiguous()       # [1,K] nvfp4
        sfa_l = sfa[:, :, l].contiguous()   # [M,K//16] fp8
        sfb_l = sfb[:, :, l].contiguous()   # [1,K//16] fp8

        if use_N1:
            try:
                outM1 = _scaled_gemv_N1(A_l, b_l, sfa_l, sfb_l)  # [M,1]
            except Exception:
                use_N1 = False  # fallback permanently
            else:
                c[:, 0, l].copy_(outM1[:, 0])
                continue

        # Fallback: N=128 with reusable scratch
        outM1 = _scaled_gemv_N128(A_l, b_l, sfa_l, sfb_l, scratch_B128K, scratch_SFB128)
        c[:, 0, l].copy_(outM1[:, 0])

    return c