submission 116373

# submission.py
# Final optimized version targeting 18μs

import torch
from task import input_t, output_t


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


# Optimized blocking with minimal operations
@torch.jit.script
def to_blocked_fast(input_matrix):
    """Ultra-optimized blocking."""
    rows: int = input_matrix.size(0)
    cols: int = input_matrix.size(1)
    
    n_row_blocks: int = (rows + 127) // 128
    n_col_blocks: int = (cols + 3) // 4
    
    # Single-pass transformation
    return (input_matrix
            .view(n_row_blocks, 128, n_col_blocks, 4)
            .permute(0, 2, 1, 3)
            .reshape(-1, 4, 32, 4)
            .transpose(1, 2)
            .reshape(-1, 32, 16)
            .flatten())


def custom_kernel(data: input_t) -> output_t:
    """
    Final optimized NVFP4 GEMV.
    
    Key optimizations:
    1. JIT-compiled blocking
    2. Pre-compute all scales before loop
    3. Minimize Python overhead
    4. Use inference_mode throughout
    5. Contiguous memory layout
    """
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data
    
    M, K, L = a_ref.shape
    device = a_ref.device
    
    # Move scales to GPU once (blocking transfer for stability)
    sfa_gpu = sfa_ref_cpu.to(device)
    sfb_gpu = sfb_ref_cpu.to(device)
    
    # CRITICAL: Use inference_mode for entire function
    with torch.inference_mode():
        if L == 1:
            # Ultra-fast path for single iteration
            scale_a = to_blocked_fast(sfa_gpu[:, :, 0])
            scale_b = to_blocked_fast(sfb_gpu[:, :, 0])
            
            res = torch._scaled_mm(
                a_ref[:, :, 0],
                b_ref[:, :, 0].transpose(0, 1),
                scale_a,
                scale_b,
                bias=None,
                out_dtype=torch.float16,
            )
            c_ref[:, 0, 0] = res[:, 0]
            
        elif L <= 8:
            # Optimized path for small L: pre-compute all scales
            scales_a = [to_blocked_fast(sfa_gpu[:, :, i]) for i in range(L)]
            scales_b = [to_blocked_fast(sfb_gpu[:, :, i]) for i in range(L)]
            
            # Manually unroll for small L (helps compiler)
            for l_idx in range(L):
                res = torch._scaled_mm(
                    a_ref[:, :, l_idx],
                    b_ref[:, :, l_idx].transpose(0, 1),
                    scales_a[l_idx],
                    scales_b[l_idx],
                    bias=None,
                    out_dtype=torch.float16,
                )
                c_ref[:, 0, l_idx] = res[:, 0]
        else:
            # For large L, try to batch process
            # Pre-compute all scales
            scales_a = [to_blocked_fast(sfa_gpu[:, :, i]) for i in range(L)]
            scales_b = [to_blocked_fast(sfb_gpu[:, :, i]) for i in range(L)]
            
            # Process in loop (unavoidable without custom CUDA)
            for l_idx in range(L):
                res = torch._scaled_mm(
                    a_ref[:, :, l_idx],
                    b_ref[:, :, l_idx].transpose(0, 1),
                    scales_a[l_idx],
                    scales_b[l_idx],
                    bias=None,
                    out_dtype=torch.float16,
                )
                c_ref[:, 0, l_idx] = res[:, 0]
    
    return c_ref