submission 202638

#!POPCORN leaderboard nvfp4_dual_gemm

import torch

@torch.compile(mode="reduce-overhead")
def _sota_logic(a_p, b1_t, b2_t, sa, sb1, sb2, c):
    """
    Compiled core logic to eliminate Python overhead and fuse operations.
    """
    l = a_p.shape[0]
    for i in range(l):
        # Hardware-accelerated GEMMs using Blackwell Tensor Cores
        # We use float32 for intermediate results to maintain reference precision
        r1 = torch._scaled_mm(a_p[i], b1_t[i], sa[i], sb1[i], out_dtype=torch.float32)
        r2 = torch._scaled_mm(a_p[i], b2_t[i], sa[i], sb2[i], out_dtype=torch.float32)
        
        # Fused SiLU and multiplication
        c[:, :, i] = (torch.nn.functional.silu(r1) * r2).to(torch.float16)

@torch.inference_mode()
def custom_kernel(data):
    """
    Hyper-optimized NVFP4 Dual GEMM kernel.
    Targets 14us performance on NVIDIA B200.
    """
    a, b1, b2, _, _, _, sfa_p, sfb1_p, sfb2_p, c = data
    l = c.shape[2]
    
    # 1. Zero-copy preparation of scale factors from pre-permuted inputs
    # sfa_p is [32, 4, rest_m, 4, rest_k, l]
    # We permute to [l, rest_m, rest_k, 32, 4, 4] which matches the hardware layout
    sa = sfa_p.permute(5, 2, 4, 0, 1, 3).reshape(l, -1)
    sb1 = sfb1_p.permute(5, 2, 4, 0, 1, 3).reshape(l, -1)
    sb2 = sfb2_p.permute(5, 2, 4, 0, 1, 3).reshape(l, -1)
    
    # 2. Pre-transpose/permute inputs to batch-first format for efficient slicing
    a_p = a.permute(2, 0, 1)    # [L, M, K]
    b1_t = b1.permute(2, 1, 0)  # [L, K, N]
    b2_t = b2.permute(2, 1, 0)  # [L, K, N]
    
    # 3. Call the compiled logic
    _sota_logic(a_p, b1_t, b2_t, sa, sb1, sb2, c)
    
    return c