submission 85010

import torch
import math
from typing import Tuple

# --- Type Definitions (Assuming availability in competition environment) ---
# Note: These types must be available in the PyTorch build used for the competition.
input_t = Tuple[torch.Tensor, ...]
output_t = torch.Tensor

# Scaling factor vector size
sf_vec_size = 16

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

# --- 1. GPU-Optimized Scale Factor Pre-processing ---

def optimized_to_blocked_batched(input_tensor: torch.Tensor) -> torch.Tensor:
    """
    GPU-native, batched implementation of the scale factor permutation.
    
    This function replaces the slow CPU-based 'to_blocked' inside the loop 
    and is run only ONCE before the main computation, eliminating CPU-GPU overhead.
    
    Args:
        input_tensor: A 3D tensor of shape (L, Rows, Cols) on the GPU.
                      (L is the batch size)
                      
    Returns:
        A 2D tensor of shape (L, BlockedSize) containing the permuted scales 
        on the GPU, ready for use by torch._scaled_mm.
    """
    l, rows, cols = input_tensor.shape
    
    # Calculate padding dimensions based on required tile sizes (128 and 4)
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)
    
    # Pad the tensor on the GPU
    padded = torch.nn.functional.pad(
        input_tensor, 
        (0, n_col_blocks * 4 - cols, 0, n_row_blocks * 128 - rows)
    )
    
    # Reshape and Permute operations (Vectorized on GPU)
    
    # 1. Block the tensor: (L, n_row_blocks, 128, n_col_blocks, 4)
    blocks = padded.view(l, n_row_blocks, 128, n_col_blocks, 4)
    
    # 2. Permute: (L, n_row_blocks, n_col_blocks, 128, 4)
    blocks = blocks.permute(0, 1, 3, 2, 4)
    
    # 3. Reshape/View for final blocking structure: (L, -1, 4, 32, 4)
    # The -1 dimension combines the n_row_blocks and n_col_blocks * 128/4
    rearranged = blocks.reshape(l, -1, 4, 32, 4)
    
    # 4. Transpose/View final 4x32x4 structure to 32x4x4 (-> 32x16)
    rearranged = rearranged.transpose(2, 3)
    
    # 5. Final reshape: (L, -1, 32, 16)
    rearranged = rearranged.reshape(l, -1, 32, 16)
    
    # 6. Flatten the blocked dimensions (excluding the batch size L)
    return rearranged.flatten(start_dim=1)

# --- 2. Optimized Kernel (Renamed to custom_kernel) ---

def custom_kernel(
    data: input_t,
) -> output_t:
    """
    Optimized kernel for NVFP4 block-scaled GEMV.
    
    This implementation avoids the Host-to-Device latency by pre-processing 
    scale factors in a batched, GPU-native operation.
    """
    a_ref, b_ref, sfa_ref, sfb_ref, _, _, c_ref = data
    
    # Get batch size (L)
    l = c_ref.shape[2]

    # --- OPTIMIZATION STEP ---
    # 1. Pre-process the entire batch of scale factors on the GPU.
    # The input scales are (M, sf_k, L). Permute to (L, M, sf_k) for batching.
    scale_a_batched = optimized_to_blocked_batched(sfa_ref.permute(2, 0, 1))
    scale_b_batched = optimized_to_blocked_batched(sfb_ref.permute(2, 0, 1))
    # -------------------------

    # The loop is now lightweight and performs only GPU-side slicing and computation.
    for l_idx in range(l):
        # Slice the pre-computed scales, which are already on the GPU
        scale_a = scale_a_batched[l_idx]
        scale_b = scale_b_batched[l_idx]

        # Call the PyTorch internal function for NVFP4 GEMM
        res = torch._scaled_mm(
            a_ref[:, :, l_idx],
            b_ref[:, :, l_idx].transpose(0, 1),
            scale_a, # GPU-resident
            scale_b, # GPU-resident
            bias=None,
            out_dtype=torch.float16,
        )
        # Store the result back into the output tensor
        c_ref[:, 0, l_idx] = res[:, 0]
        
    return c_ref

# --- 3. Reference Kernel (Provided by the competition for comparison) ---
# Note: This kernel is provided here for the matching utility, but it is 
# inherently slow due to CPU work and HtoD transfers inside the loop.

def to_blocked(input_matrix):
    """Reference implementation of scale factor conversion (CPU-based)."""
    rows, cols = input_matrix.shape

    # Ensure rows and cols are multiples of 128 and 4 respectively
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)

    # Pad (Note: padding is required for correct view/reshape logic)
    padded = torch.nn.functional.pad(
        input_matrix,
        (0, n_col_blocks * 4 - cols, 0, n_row_blocks * 128 - rows)
    )
    
    # View and Permute sequence
    blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)

    return rearranged.flatten()


def reference_kernel(
    data: input_t,
) -> output_t:
    """PyTorch reference implementation (SLOW version)."""
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data
    
    # Get dimensions from MxNxL layout
    _, _, l = c_ref.shape

    # Call torch._scaled_mm to compute the GEMV result
    for l_idx in range(l):
        # This is the slow part: running CPU logic inside a loop
        scale_a = to_blocked(sfa_ref_cpu[:, :, l_idx].cpu())
        scale_b = to_blocked(sfb_ref_cpu[:, :, l_idx].cpu())
        
        # (m, k) @ (n, k).T -> (m, n)
        res = torch._scaled_mm(
            a_ref[:, :, l_idx],
            b_ref[:, :, l_idx].transpose(0, 1),
            scale_a.cuda(), # HtoD transfer inside loop
            scale_b.cuda(), # HtoD transfer inside loop
            bias=None,
            out_dtype=torch.float16,
        )
        c_ref[:, 0, l_idx] = res[:, 0]
    return c_ref

# --- 4. Input Generation (Required for matching the competition test harness) ---

def generate_input(
    m: int,
    k: int,
    l: int,
    seed: int,
):
    """
    Generate input tensors for NVFP4 block-scaled GEMV.
    """
    torch.manual_seed(seed)

    # GEMV N dimension is always 1
    n = 1
    # Scaling factor needs to pad the N size to 128
    n_padded_128 = 128
    
    # Generate uint8 tensor, then convert to float4e2m1fn_x2 data type
    a_ref = torch.randint(
        0, 4, (l, m, k // 2), dtype=torch.uint8, device="cuda"
    ).permute(1, 2, 0)
    # Pad b tensor's N dimension to 128 to call torch._scaled_mm for nvfp4 dot product computation
    b_ref = torch.randint(
        0, 4, (l, n_padded_128, k // 2), dtype=torch.uint8, device="cuda"
    ).permute(1, 2, 0)
    
    # Correctly cast the packed uint8 representation to the custom float4 type
    a_ref = a_ref.view(torch.uint8).view(torch.float4_e2m1fn_x2)
    b_ref = b_ref.view(torch.uint8).view(torch.float4_e2m1fn_x2)

    # Create float16 output tensor
    c_ref = torch.randn((l, m, n), dtype=torch.float16, device="cuda").permute(
        1, 2, 0
    )
    
    # Helper function to prepare the scale factor tensors
    def create_scale_factor_tensors(l, mn, sf_k):
        # Create the reference scale factor tensor (mn, sf_k, l) on CPU.
        ref_shape = (l, mn, sf_k)
        ref_permute_order = (1, 2, 0)
        # Init with int8 tensor, then convert to float8_e4m3fn
        ref_f8_random_int = torch.randint(0, 3, ref_shape, dtype=torch.int8, device='cuda')
        ref_f8_torch_tensor = ref_f8_random_int.to(dtype=torch.float8_e4m3fn)
        # permute to match ref_permute_order (M, K, L format)
        ref_f8_torch_tensor_permuted = ref_f8_torch_tensor.permute(*ref_permute_order)
        
        # --- Dummy tensors for the permuted output (required by the input_t structure) ---
        atom_m = (32, 4)
        atom_k = 4
        mma_shape = (
            l, ceil_div(mn, atom_m[0] * atom_m[1]), ceil_div(sf_k, atom_k),
            atom_m[0], atom_m[1], atom_k,
        )
        mma_permute_order = (3, 4, 1, 5, 2, 0)
        rand_int_tensor = torch.randint(0, 3, mma_shape, dtype=torch.int8, device='cuda')
        reordered_f8_torch_tensor = rand_int_tensor.to(dtype=torch.float8_e4m3fn)
        reordered_f8_torch_tensor = reordered_f8_torch_tensor.permute(*mma_permute_order)
        # End of dummy tensor creation
        
        # The key outputs for this problem are the CPU and GPU scale tensors
        return ref_f8_torch_tensor_permuted.cpu(), reordered_f8_torch_tensor

    sf_k = ceil_div(k, sf_vec_size)
    sfa_ref_cpu, sfa_permuted = create_scale_factor_tensors(l, m, sf_k)
    sfb_ref_cpu, sfb_permuted = create_scale_factor_tensors(l, n_padded_128, sf_k)

    sfa_ref = sfa_ref_cpu.to("cuda")
    sfb_ref = sfb_ref_cpu.to("cuda")
    
    # The output format for input_t is:
    # (a, b, sfa_ref_gpu, sfb_ref_gpu, sfa_permuted_gpu, sfb_permuted_gpu, c)
    return (a_ref, b_ref, sfa_ref, sfb_ref, sfa_permuted, sfb_permuted, c_ref)