submission 101021

import torch
from task import input_t, output_t


@torch.no_grad()
def custom_kernel(data: input_t) -> output_t:
    """
    Optimized NVFP4 GEMV using torch._scaled_mm with hardware acceleration.
    
    Key optimizations:
    1. Uses pre-permuted scale factors (avoids CPU to_blocked computation)
    2. Makes scale factors contiguous for better memory access
    3. torch.no_grad() eliminates autograd overhead
    """
    a_ref, b_ref, _, _, sfa_permuted, sfb_permuted, c_ref = data
    
    # Get dimensions
    _, _, l = c_ref.shape
    
    # Pre-compute flattened scale factors for all batches
    # sfa_permuted: (32, 4, mn_blocks, 4, k_blocks, l)
    # Target: (l, flattened) matching to_blocked output
    scale_a_all = sfa_permuted.permute(5, 2, 4, 0, 1, 3).reshape(l, -1).contiguous()
    scale_b_all = sfb_permuted.permute(5, 2, 4, 0, 1, 3).reshape(l, -1).contiguous()
    
    # Process each batch
    for l_idx in range(l):
        res = torch._scaled_mm(
            a_ref[:, :, l_idx],
            b_ref[:, :, l_idx].t(),
            scale_a_all[l_idx],
            scale_b_all[l_idx],
            bias=None,
            out_dtype=torch.float16,
        )
        c_ref[:, 0, l_idx] = res[:, 0]
    
    return c_ref