submission 149959

"""
NVFP4 Block-Scaled GEMM Optimized Submission

Key optimizations over baseline:
1. All operations stay on GPU (no CPU round-trips)
2. Uses contiguous memory for scale factors
3. Proper memory layout for torch._scaled_mm

For truly competitive performance, a CuTeDSL kernel targeting
Blackwell's TCGen05 tensor cores would be needed.
"""

import torch
from task import input_t, output_t


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


@torch.jit.script
def to_blocked_gpu(input_matrix: torch.Tensor) -> torch.Tensor:
    """
    Convert scale factor tensor to blocked format expected by torch._scaled_mm.
    
    This is the GPU-only version - no CPU transfers!
    The blocked format arranges scale factors for optimal tensor core access.
    
    Layout transformation:
    (rows, cols) -> (n_row_blocks, n_col_blocks, 128, 4) -> (-1, 32, 16) -> flat
    """
    rows, cols = input_matrix.shape
    n_row_blocks = (rows + 127) // 128
    n_col_blocks = (cols + 3) // 4
    
    # Reshape and permute on GPU (no CPU transfer!)
    blocks = input_matrix.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
    
    return rearranged.flatten().contiguous()


def custom_kernel(data: input_t) -> output_t:
    """
    Optimized NVFP4 block-scaled GEMM: C = A @ B.T with block scaling
    
    This kernel uses torch._scaled_mm which calls optimized CUTLASS kernels
    for NVFP4 on Blackwell GPUs.
    
    Args:
        data: Tuple containing:
            a: [m, k//2, l] in float4_e2m1fn_x2 (2 FP4 values packed per byte)
            b: [n, k//2, l] in float4_e2m1fn_x2
            sfa: [m, k//16, l] scale factors in float8_e4m3fnuz
            sfb: [n, k//16, l] scale factors in float8_e4m3fnuz  
            sfa_permuted: pre-permuted scale factors (unused)
            sfb_permuted: pre-permuted scale factors (unused)
            c: [m, n, l] output buffer in float16
    
    Returns:
        c: [m, n, l] result in float16
    """
    a, b, sfa, sfb, sfa_permuted, sfb_permuted, c = data
    
    # Get dimensions
    m, n, l = c.shape
    
    # Process each batch element
    for l_idx in range(l):
        # Get current batch slices (already contiguous on GPU)
        a_slice = a[:, :, l_idx]
        b_slice = b[:, :, l_idx]
        
        # Convert scale factors to blocked format ON GPU
        # This is the key optimization - no CPU transfer!
        scale_a = to_blocked_gpu(sfa[:, :, l_idx])
        scale_b = to_blocked_gpu(sfb[:, :, l_idx])
        
        # Perform scaled GEMM: C = (A * scale_a) @ (B * scale_b).T
        # torch._scaled_mm handles NVFP4 block-scaled operations
        result = torch._scaled_mm(
            a_slice,
            b_slice.transpose(0, 1),  # Transpose B for (m,k) @ (k,n)
            scale_a,
            scale_b,
            bias=None,
            out_dtype=torch.float16,
        )
        
        c[:, :, l_idx] = result
    
    return c