submission 94867

"""
Optimized submission using torch._scaled_mm for maximum performance.
Based on reference implementation approach.
"""

import torch
from typing import Tuple

# Type aliases
input_t = Tuple[
    torch.Tensor,  # a: M x K x L (nvfp4)
    torch.Tensor,  # b: N x K x L (nvfp4)
    torch.Tensor,  # sfa: M x (K/16) x L (fp8)
    torch.Tensor,  # sfb: N x (K/16) x L (fp8)
    torch.Tensor,  # sfa_permuted
    torch.Tensor,  # sfb_permuted
    torch.Tensor,  # c_ref: M x 1 x L (fp16)
]
output_t = torch.Tensor  # M x 1 x L (fp16)


def ceil_div(a, b):
    """Integer ceiling division"""
    return (a + b - 1) // b


def to_blocked(input_matrix):
    """Convert scale factors to blocked format for torch._scaled_mm"""
    rows, cols = input_matrix.shape
    
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)
    
    padded = input_matrix
    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 _custom_kernel_impl(data: input_t) -> output_t:
    """
    Optimized implementation using torch._scaled_mm.
    
    Uses PyTorch's highly optimized scaled matrix multiplication kernel.
    """
    a, b, sfa, sfb, sfa_permuted, sfb_permuted, c_ref = data
    
    M, K_packed, L = a.shape
    
    # Initialize output
    c = torch.zeros_like(c_ref)
    
    # Process each batch
    for l_idx in range(L):
        # Convert scale factors to blocked format
        scale_a = to_blocked(sfa[:, :, l_idx])
        scale_b = to_blocked(sfb[:, :, l_idx])
        
        # Use torch._scaled_mm for optimized computation
        res = torch._scaled_mm(
            a[:, :, l_idx],
            b[:, :, l_idx].transpose(0, 1),
            scale_a,
            scale_b,
            bias=None,
            out_dtype=torch.float16,
        )
        
        # Extract first column (GEMV result)
        c[:, 0, l_idx] = res[:, 0]
    
    return c


# JIT compile for additional speedup
try:
    custom_kernel = torch.compile(_custom_kernel_impl, mode="reduce-overhead", fullgraph=True)
except Exception:
    custom_kernel = _custom_kernel_impl