submission 105294

"""
Simple optimized NVFP4 block-scaled GEMV kernel.
No CUDA graphs - just pure vectorized operations.
"""
import torch
from task import input_t, output_t


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


@torch.no_grad()
def to_blocked_all(sf: torch.Tensor) -> list:
    """
    Convert (rows, cols, L) scale factors to blocked format.
    Returns list of L blocked tensors.
    """
    rows, cols, L = sf.shape
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)
    
    # Single vectorized operation for all L
    blocked = (sf.view(n_row_blocks, 128, n_col_blocks, 4, L)
                 .permute(0, 2, 1, 3, 4)
                 .reshape(-1, 4, 32, 4, L)
                 .transpose(1, 2)
                 .reshape(-1, 32, 16, L)
                 .flatten(0, 2))  # (blocked_size, L)
    
    # Return list of contiguous slices
    return [blocked[:, i].contiguous() for i in range(L)]


@torch.no_grad()
def to_blocked_single(sf: torch.Tensor) -> torch.Tensor:
    """Convert (rows, cols) scale factors to blocked format."""
    rows, cols = sf.shape
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)
    return (sf.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())


@torch.no_grad()
def custom_kernel(data: input_t) -> output_t:
    """Simple optimized NVFP4 block-scaled GEMV."""
    a, b, sfa, sfb, _, _, c = data
    _, _, l = c.shape
    
    if l == 1:
        # Single batch path
        scale_a = to_blocked_single(sfa[:, :, 0])
        scale_b = to_blocked_single(sfb[:, :, 0])
        result = torch._scaled_mm(
            a[:, :, 0], b[:, :, 0].t(),
            scale_a, scale_b,
            bias=None, out_dtype=torch.float16
        )
        c[:, 0, 0] = result[:, 0]
    else:
        # Multi-batch path with vectorized scale conversion
        blocked_sfa = to_blocked_all(sfa)
        blocked_sfb = to_blocked_all(sfb)
        
        for i in range(l):
            result = torch._scaled_mm(
                a[:, :, i], b[:, :, i].t(),
                blocked_sfa[i], blocked_sfb[i],
                bias=None, out_dtype=torch.float16
            )
            c[:, 0, i] = result[:, 0]
    
    return c