submission 112289

import torch
from task import input_t, output_t

# Kernel configuration parameters
sf_vec_size = 16
ROW_BLOCK_SIZE = 128
COL_BLOCK_SIZE = 4


# Helper function for ceiling division
def ceil_div(a, b):
    """Compute ceiling division: ceil(a / b)"""
    return (a + b - 1) // b


# Helper function to convert scale factor tensor to blocked format
def to_blocked(input_matrix, device=None):
    """
    Convert scale factor tensor to blocked format.
    
    Args:
        input_matrix: 2D tensor with shape (rows, cols)
        device: Target device for output tensor
    
    Returns:
        Flattened blocked format tensor
    """
    rows, cols = input_matrix.shape
    
    # Ensure rows and cols are multiples of 128 and 4 respectively
    n_row_blocks = ceil_div(rows, ROW_BLOCK_SIZE)
    n_col_blocks = ceil_div(cols, COL_BLOCK_SIZE)
    
    # Get device from input if not specified
    if device is None:
        device = input_matrix.device
    
    # Ensure tensor is on correct device and contiguous
    padded = input_matrix.contiguous().to(device)
    
    # Reshape and permute to blocked format
    blocks = padded.view(n_row_blocks, ROW_BLOCK_SIZE, n_col_blocks, COL_BLOCK_SIZE).permute(0, 2, 1, 3)
    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
    
    return rearranged.flatten()


def custom_kernel(
    data: input_t,
) -> output_t:
    """
    PyTorch reference implementation of NVFP4 block-scaled GEMV.
    
    Args:
        data: Tuple containing (a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref)
    
    Returns:
        Output tensor c_ref with computed results
    """
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data
    
    # Get device from input tensors (prefer CUDA if available)
    device = None
    if isinstance(a_ref, torch.Tensor):
        device = a_ref.device
    elif isinstance(b_ref, torch.Tensor):
        device = b_ref.device
    elif isinstance(c_ref, torch.Tensor):
        device = c_ref.device
    
    # Default to CUDA if device not found and CUDA is available
    if device is None or (device.type == 'cpu' and torch.cuda.is_available()):
        device = torch.device('cuda')
    
    # Ensure input tensors are on correct device
    a_ref = a_ref.to(device=device)
    b_ref = b_ref.to(device=device)
    c_ref = c_ref.to(device=device)
    
    # Get dimensions from MxNxL layout
    _, _, l = c_ref.shape
    
    # Process each slice in the L dimension
    for l_idx in range(l):
        # Extract current slice
        a_slice = a_ref[:, :, l_idx]
        b_slice = b_ref[:, :, l_idx]
        sfa_slice = sfa_ref_cpu[:, :, l_idx]
        sfb_slice = sfb_ref_cpu[:, :, l_idx]
        
        # Convert the scale factor tensors to blocked format
        scale_a = to_blocked(sfa_slice, device=device)
        scale_b = to_blocked(sfb_slice, device=device)
        
        # Ensure scale tensors are on correct device
        scale_a = scale_a.to(device=device)
        scale_b = scale_b.to(device=device)
        
        # Compute scaled matrix multiplication: (m, k) @ (n, k).T -> (m, n)
        res = torch._scaled_mm(
            a_slice,
            b_slice.transpose(0, 1),
            scale_a,
            scale_b,
            bias=None,
            out_dtype=torch.float16,
        )
        
        # Store result in output tensor
        c_ref[:, 0, l_idx] = res[:, 0]
    
    return c_ref