submission 183047

#!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 lifirght 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!


'''
ARITHMETIC INTENSITY

CI = FLOPS / NUM BYTES ACCESSED

ELEMENTS COMPUTED FOR GEMM: 2 * ( M * N * K) --> 2: mult & sum  
Num Byte Acces: 2 * ( M*N + M*K + K*N) --> 2: assuming f16 


'''
from task import input_t, output_t
import torch


def _scale_vec(sf_perm, l_idx: int):
    # [32, 4, rest_m, 4, rest_k, L] -> 1D esperado por _scaled_mm
    return (sf_perm[..., l_idx]
            .permute(2, 4, 0, 1, 3)   # (rest_m, rest_k, 32, 4, 4)
            .contiguous()
            .view(-1))


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]
    """

    a, b, sfa, sfb, sfa_perm, sfb_perm, c = data

    scale_a = _scale_vec(sfa_perm, 0)
    scale_b = _scale_vec(sfb_perm, 0)

    out0 = c[:, :, 0]
    torch._scaled_mm(
        a[:, :, 0],
        b[:, :, 0].t(),
        scale_a,
        scale_b,
        bias=None,
        out_dtype=torch.float16,
        out=out0,                # <-- clave
    )
    return c