submission 192656

#!POPCORN leaderboard nvfp4_dual_gemm

# This is a submission template for popcorn leaderboard 'nvfp4_dual_gemm'.
# Your task is as follows:
# > 
# > You will implement a block scaled dual matrix-matrix multiplication kernel with silu activation optimized for NVIDIA B200.
# > To be explicit, you will be given a tuple of tensors:
# > ```
# > (a, b1, b2, sfa, sfb1, sfb2, c)
# > ```
# > where:
# > * `a` is M x K x L in K-major order in nvfp4(e2m1)
# > * `b1` is N x K x L in K-major order in nvfp4(e2m1)
# > * `b2` 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)
# > * `sfb1` is N x (K // 16) x L in K-major order in fp8(e4m3fnuz)
# > * `sfb2` 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] 
# > 256 4096 7168 1 4.708
# > 512 4096 7168 1 8.714
# > 256 3072 4096 1 2.125
# > 512 3072 7168 1 6.535
# > ```
# The deadline for this leaderboard is 2026-01-17 07: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!
'''
k: 7168; l: 1; m: 256; n: 4096; seed: 1111
 ⏱ 48.6 ± 0.05 µs
 ⚡ 48.5 µs 🐌 48.8 µs

k: 7168; l: 1; m: 512; n: 4096; seed: 1111
 ⏱ 51.3 ± 0.04 µs
 ⚡ 51.2 µs 🐌 51.3 µs

k: 4096; l: 1; m: 256; n: 3072; seed: 1111
 ⏱ 41.1 ± 0.01 µs
 ⚡ 41.1 µs 🐌 41.1 µs

k: 7168; l: 1; m: 512; n: 3072; seed: 1111
 ⏱ 48.8 ± 0.04 µs
 ⚡ 48.6 µs 🐌 48.8 µs
'''

from task import input_t, output_t
import torch

def _scale_vec(sf_perm, lidx: int):
    # [32, 4, rest_m, 4, rest_k, L] -> 1D esperado por _scaled_mm
    return (sf_perm[..., lidx]
            .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 dual gemm with silu activation
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float4e2m1fn] of shape [m, k, l],
            b1: torch.Tensor[float4e2m1fn] of shape [n, k, l],
            b2: torch.Tensor[float4e2m1fn] of shape [n, k, l],
            sfa: torch.Tensor[float8_e4m3fnuz] of shape [m, k // 16, l], used by reference implementation
            sfb1: torch.Tensor[float8_e4m3fnuz] of shape [n, k // 16, l], used by reference implementation
            sfb2: torch.Tensor[float8_e4m3fnuz] of shape [n, k // 16, l], used by reference implementation
            sfa_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_m, 4, rest_k, l],
            sfb1_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_n, 4, rest_k, l],
            sfb2_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_ref, b1_ref, b2_ref, sfa_ref_cpu, sfb1_ref_cpu, sfb2_ref_cpu, sfa_permuted, sfb1_permuted, sfb2_permuted, c_ref = data
    
    # Get dimensions from MxNxL layout
    m, n, l = c_ref.shape
    
        # Convert the scale factor tensor to blocked format
    scale_a = _scale_vec(sfa_permuted,0)
    scale_b1 =_scale_vec(sfb1_permuted,0)
    scale_b2 = _scale_vec(sfb2_permuted,0)
    

    ref1 = torch.empty((m, n, l), dtype=torch.float32, device="cuda")
    
    ref2 = torch.empty((m, n, l), dtype=torch.float32, device="cuda")
    
    # (m, k) @ (n, k).T -> (m, n)
    res1 =  torch._scaled_mm(
        a_ref[:, :, 0],
        b1_ref[:, :, 0].t(),
        scale_a,
        scale_b1,
        bias=None,
        out_dtype=torch.float32,
        
    )
    res2 =  torch._scaled_mm(
        a_ref[:, :, 0],
        b2_ref[:, :, 0].t(),
        scale_a,
        scale_b2,
        bias=None,
        out_dtype=torch.float32,
    )
    ref1[:, :, 0].copy_(res1)
    ref2[:, :, 0].copy_(res2)
    
    # Do silu on the first GEMM result and multiply with the second GEMM result
    c_ref = (torch.nn.functional.silu(ref1) * ref2).to(torch.float16)
    return c_ref