submission 161208

#!POPCORN leaderboard nvfp4_gemm

# This is a submission template for popcorn leaderboard 'nvfp4_gemm'.
# Your task is as follows:
# > 
# > You will implement a block scaled matrix-matrix multiplication kernel optimized for NVIDIA B200.
# > To be explicit, you will be given a tuple of tensors:
# > ```
# > (a, b, sfa, sfb, c)
# > ```
# > where:
# > * `a` is M x K x L in K-major order in nvfp4(e2m1)
# > * `b` is N x K x L in K-major order in nvfp4(e2m1)
# > * `sfa` is M x (K // 16) x L in K-major order in fp8(e4m3fnuz)
# > * `sfb` is N x (K // 16) x L in K-major order in fp8(e4m3fnuz)
# > * `c` is M x N x L in fp16
# > 
# > Matrix sizes `M` is divisible by mma_tiler_mn[0], `N` is divisible by mma_tiler_mn[1], `K` is divisible by 256.
# > The ranking criteria is the geometric mean of the benchmark results.
# > For the grand price, your kernel will be evaluated against the speed of light analysis
# > and the solution closest to the speed of light will be awarded the grand price.
# > ```
# > The speed of light analysis based on the max(FP4 Tensor Core math throughput, DRAM memory throughput) of B200 and tested under 1.5Ghz clock:
# > M   N    K   L time[us]
# > 128 7168 16384 1 8.994
# > 128 4096 7168  1 2.354
# > 128 7168 2048  1 1.333
# > ```
# The deadline for this leaderboard is 2025-12-20 06:59:00+00:00

# You can automatically route this file to specific GPUs by adding a line
# `#!POPCORN gpus <GPUs>` to the header of this file.
# Happy hacking!

from task import input_t, output_t


sf_vec_size = 16


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


def _to_cublas_blocked_scale(scale_2d):
    """
    Convert scale factors from shape (rows, cols) where cols == K//sf_vec_size
    into the 1D "blocked" layout expected by torch._scaled_mm / cuBLAS.

    This matches the reference implementation's `to_blocked`.
    """
    rows, cols = scale_2d.shape
    # cuBLAS block scaling layout assumes rows multiple of 128 and cols multiple of 4
    n_row_blocks = rows // 128
    n_col_blocks = cols // 4
    # (n_row_blocks, 128, n_col_blocks, 4) -> (n_row_blocks, n_col_blocks, 128, 4)
    blocks = scale_2d.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
    # -> (-1, 4, 32, 4) -> (-1, 32, 4, 4) -> (-1, 32, 16)
    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
    return rearranged.flatten()


def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of block-scale fp4 gemm
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float4e2m1fn] of shape [m, k, l],
            b: torch.Tensor[float4e2m1fn] of shape [n, k, l],
            sfa: torch.Tensor[float8_e4m3fnuz] of shape [m, k // 16, l],
            sfb: torch.Tensor[float8_e4m3fnuz] of shape [n, k // 16, l],
            sfa_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_m, 4, rest_k, l],
            sfb_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_n, 4, rest_k, l],
            c: torch.Tensor[float16] of shape [m, n, l]
    Returns:
        Tensor containing output in float16
        c: torch.Tensor[float16] of shape [m, n, l]
    """
    # c: [m, n, l] is pre-allocated memory to avoid timing allocation overhead.
    a, b, sfa, sfb, sfa_permuted, sfb_permuted, c = data

    import torch

    m, k_packed, l = a.shape
    n = b.shape[0]

    # The benchmark uses NVFP4 packed dtypes (commonly torch.float4_e2m1fn_x2) where
    # the tensor K dimension is *packed* (K/2). Derive the *logical* K for scale shapes.
    pack_factor = 2 if a.dtype == torch.float4_e2m1fn_x2 else 1
    k = k_packed * pack_factor

    # Scale factors are per 16 elements along logical K.
    sf_k = k // sf_vec_size
    assert k % sf_vec_size == 0, "Logical K must be divisible by 16"

    assert sfa.shape == (m, sf_k, l), f"sfa shape {sfa.shape} != {(m, sf_k, l)}"
    assert sfb.shape == (n, sf_k, l), f"sfb shape {sfb.shape} != {(n, sf_k, l)}"
    assert c.shape == (m, n, l)

    # Use torch's built-in scaled GEMM path (cuBLAS) to avoid any JIT compilation.
    # aten::_scaled_mm schema:
    #   (Tensor self, Tensor mat2, Tensor scale_a, Tensor scale_b, Tensor? bias=None,
    #    Tensor? scale_result=None, ScalarType? out_dtype=None, bool use_fast_accum=False) -> Tensor
    #
    # It expects scale_a/scale_b in cuBLAS "blocked" 1D layout.
    assert sf_k % 4 == 0, "K//sf_vec_size must be divisible by 4 for blocked scale layout"
    assert m % 128 == 0 and n % 128 == 0, "cuBLAS blocked scale layout requires M and N divisible by 128"

    # Some environments provide fp8(e4m3fnuz); cuBLAS path generally expects e4m3fn.
    if sfa.dtype == torch.float8_e4m3fnuz:
        sfa_use = sfa.to(torch.float8_e4m3fn)
        sfb_use = sfb.to(torch.float8_e4m3fn)
    else:
        sfa_use, sfb_use = sfa, sfb

    for l_idx in range(l):
        scale_a = _to_cublas_blocked_scale(sfa_use[:, :, l_idx])
        scale_b = _to_cublas_blocked_scale(sfb_use[:, :, l_idx])
        res = torch._scaled_mm(
            a[:, :, l_idx],
            b[:, :, l_idx].transpose(0, 1),
            scale_a,
            scale_b,
            bias=None,
            out_dtype=torch.float16,
        )
        c[:, :, l_idx].copy_(res)

    return c