submission 90140

import torch
from task import input_t, output_t

@torch.no_grad()
def custom_kernel(data: input_t) -> output_t:
    """
    Optimized NVFP4 GEMV Kernel:
    1. Uses pre-uploaded 'sfa_perm'/'sfb_perm' to skip PCIe transfer and CPU blocking.
    2. Fixes scale layout on GPU with a single vectorized permute.
    3. Reorders 'b' to batch-contiguous layout for optimal read performance.
    4. Executes torch._scaled_mm in a tight loop.
    """
    a, b, sfa, sfb, sfa_perm, sfb_perm, c = data
    
    # Dimensions
    # c: [M, 1, L]
    # b: [N_padded=128, K, L] (physically)
    M, _, L = c.shape

    # ---------------------------------------------------------------------------
    # 1. Optimize Scale Factors
    # ---------------------------------------------------------------------------
    # sfa_perm is provided on GPU with shape: [32, 4, M_blk, 4, K_blk, L]
    # The torch._scaled_mm kernel expects a specific flattened block layout.
    # We permute sfa_perm to: [L, M_blk, K_blk, 32, 4, 4] and flatten.
    # Permutation indices derived from layout analysis: 
    # 5(L), 2(M_blk), 4(K_blk), 0(mm32), 1(mm4), 3(kk4)
    
    sfa_gpu = sfa_perm.permute(5, 2, 4, 0, 1, 3).contiguous().flatten(1)
    sfb_gpu = sfb_perm.permute(5, 2, 4, 0, 1, 3).contiguous().flatten(1)

    # ---------------------------------------------------------------------------
    # 2. Optimize Input Vector B
    # ---------------------------------------------------------------------------
    # b is [N, K, L] in float4. Accessing b[:, :, l] is strided and slow.
    # We permute it to [L, N, K] so that b[l] is a contiguous chunk.
    # Since b is float4, we use a uint8 view to manipulate the memory.
    
    b_u8 = b.view(torch.uint8)                  # View as [N, K//2, L]
    b_u8_opt = b_u8.permute(2, 0, 1).contiguous() # Permute to [L, N, K//2]
    b_opt = b_u8_opt.view(torch.float4_e2m1fn_x2) # View back to [L, N, K]

    # ---------------------------------------------------------------------------
    # 3. Compute Loop
    # ---------------------------------------------------------------------------
    # Iterate over batch dimension L. 
    # All inputs are now GPU-resident and layout-optimized.
    
    for l_idx in range(L):
        # b_opt[l_idx] is [N, K]. Transpose to [K, N] for the kernel.
        # The transpose here is a metadata-only operation on a contiguous block.
        
        res = torch._scaled_mm(
            a[:, :, l_idx],               # [M, K]
            b_opt[l_idx].transpose(0, 1), # [K, N]
            sfa_gpu[l_idx],               # [BlockScale]
            sfb_gpu[l_idx],               # [BlockScale]
            bias=None,
            out_dtype=torch.float16
        )
        
        # Store result. N=0 is the valid vector output.
        c[:, 0, l_idx] = res[:, 0]

    return c