submission 487055

#!POPCORN leaderboard nvfp4_group_gemm
#!POPCORN gpu B200

"""
NVFP4 Group GEMM - Production submission (v9)

Performance: 246μs (8-group), 181μs (8-group k=2048), 49μs (2-group), 41μs (2-group)
vs Speed of Light: 18.8μs, 10.7μs, 2.4μs, 1.5μs

Optimizations applied:
1. Custom CUDA kernel for strided scale factor conversion (no .contiguous() copy)
2. Direct read from non-contiguous 6D permuted tensor using stride arithmetic
3. Use pre-reordered scale factors from input data (3rd tuple element)
4. Minimal Python overhead in hot loop
5. Lean CUDA kernel with bit-shift arithmetic
"""

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void convert_sf(
    const uint8_t* __restrict__ src,
    uint8_t* __restrict__ dst,
    const int rest_m,
    const int rest_k,
    const long s0, const long s1, const long s2,
    const long s3, const long s4,
    const int total
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total) return;
    
    int l16 = i & 15;
    int t = i >> 4;
    int l32 = t & 31;
    int fb = t >> 5;
    
    int rb = fb / rest_k;
    int cb = fb - rb * rest_k;
    
    if (rb >= rest_m) { dst[i] = 0; return; }
    
    dst[i] = src[(long)l32*s0 + (long)(l16>>2)*s1 + (long)rb*s2 + (long)(l16&3)*s3 + (long)cb*s4];
}

torch::Tensor convert_one(torch::Tensor src) {
    auto sz = src.sizes();
    auto st = src.strides();
    int rm = sz[2], rk = sz[4];
    int total = rm * rk * 512;
    
    auto dst = torch::empty({total}, torch::TensorOptions().dtype(torch::kUInt8).device(src.device()));
    
    convert_sf<<<(total+511)/512, 512>>>(
        reinterpret_cast<const uint8_t*>(src.data_ptr()),
        dst.data_ptr<uint8_t>(),
        rm, rk, st[0], st[1], st[2], st[3], st[4], total
    );
    return dst;
}
"""

cpp_source = r"""
#include <torch/extension.h>
torch::Tensor convert_one(torch::Tensor src);
"""

_m = load_inline(
    name='nvfp4_prod',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['convert_one'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False,
)
_convert = _m.convert_one


def custom_kernel(data: input_t) -> output_t:
    abc, _, sf_reord, _ = data
    n = len(abc)
    r = [None] * n
    
    for i in range(n):
        a, b, c = abc[i]
        sfa_r, sfb_r = sf_reord[i]
        
        # Strided convert - no .contiguous()
        sa = _convert(sfa_r.view(torch.uint8)).view(torch.float8_e4m3fn)
        sb = _convert(sfb_r.view(torch.uint8)).view(torch.float8_e4m3fn)
        
        # GEMM
        c[:, :, 0] = torch._scaled_mm(
            a[:, :, 0].view(torch.float4_e2m1fn_x2),
            b[:, :, 0].transpose(0, 1).view(torch.float4_e2m1fn_x2),
            sa, sb, bias=None, out_dtype=torch.float16,
        )
        r[i] = c
    
    return r