submission 70433

import torch
import ctypes
import os
from task import input_t, output_t
from utils import make_match_reference

# ============================================================================
# CONFIGURATION & HELPERS
# ============================================================================

sf_vec_size = 16

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

def to_blocked(input_matrix):
    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()

# ============================================================================
# REFERENCE KERNEL (PyTorch - for validation)
# ============================================================================

def ref_kernel(data: input_t) -> output_t:
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data
    _, _, l = c_ref.shape
    for l_idx in range(l):
        scale_a = to_blocked(sfa_ref_cpu[:, :, l_idx])
        scale_b = to_blocked(sfb_ref_cpu[:, :, l_idx])
        res = torch._scaled_mm(
            a_ref[:, :, l_idx],
            b_ref[:, :, l_idx].transpose(0, 1),
            scale_a.cuda(),
            scale_b.cuda(),
            bias=None,
            out_dtype=torch.float16,
        )
        c_ref[:, 0, l_idx] = res[:, 0]
    return c_ref

# ============================================================================
# HIP KERNEL LOADER & WRAPPER
# ============================================================================

class HIPKernel:
    """Wrapper to load and call HIP/CUDA kernel from shared library."""
    
    def __init__(self, lib_path="./libgemv.so"):
        """
        Load pre-compiled HIP kernel library.
        Assumes kernel compiled with: hipcc -fPIC -shared nvfp4_gemv_kernel.hip -o libgemv.so
        """
        if not os.path.exists(lib_path):
            raise FileNotFoundError(
                f"HIP kernel library not found: {lib_path}\n"
                "Compile with: hipcc -fPIC -shared nvfp4_gemv_kernel.hip -o libgemv.so"
            )
        
        self.lib = ctypes.CDLL(lib_path)
        self.launch_kernel = self.lib.launch_nvfp4_gemv
        self.launch_kernel.argtypes = [
            ctypes.c_void_p,  # d_A
            ctypes.c_void_p,  # d_B
            ctypes.c_void_p,  # d_scale_A
            ctypes.c_void_p,  # d_scale_B
            ctypes.c_void_p,  # d_C
            ctypes.c_int,     # M
            ctypes.c_int,     # K
            ctypes.c_int,     # L
        ]
        self.launch_kernel.restype = None
    
    def __call__(self, A, B, scale_A, scale_B, C, M, K, L):
        """Launch kernel with given tensors."""
        self.launch_kernel(
            A.data_ptr(),
            B.data_ptr(),
            scale_A.data_ptr(),
            scale_B.data_ptr(),
            C.data_ptr(),
            M, K, L
        )

# Try to load HIP kernel
try:
    hip_kernel = HIPKernel("./libgemv.so")
    HIP_AVAILABLE = True
except FileNotFoundError:
    HIP_AVAILABLE = False
    print("[WARNING] HIP kernel not found. Falling back to reference implementation.")

# ============================================================================
# CUSTOM KERNEL - MAIN ENTRY POINT
# ============================================================================

def custom_kernel(data: input_t) -> output_t:
    """
    Main entry point for kernel evaluation.
    
    Uses HIP kernel if available, otherwise falls back to PyTorch reference.
    
    Args:
        data: tuple of (a, b, sfa_ref_cpu, sfb_ref_cpu, sfa_permuted, sfb_permuted, c_ref)
    
    Returns:
        c_result: [M, 1, L] output in FP16
    """
    
    a, b, sfa_ref_cpu, sfb_ref_cpu, sfa_permuted, sfb_permuted, c_ref = data
    m, k_half, l = a.shape
    k = k_half * 2
    
    c_result = torch.zeros_like(c_ref)
    
    if HIP_AVAILABLE:
        # Use HIP kernel
        # Convert scales to GPU if not already
        sfa_gpu = sfa_ref_cpu.cuda() if sfa_ref_cpu.is_cpu else sfa_ref_cpu
        sfb_gpu = sfb_ref_cpu.cuda() if sfb_ref_cpu.is_cpu else sfb_ref_cpu
        
        # Launch HIP kernel for all batches at once
        hip_kernel(a, b, sfa_gpu, sfb_gpu, c_result, m, k, l)
        torch.cuda.synchronize()
    else:
        # Fallback to reference (slow, for testing/validation only)
        for batch_idx in range(l):
            result_slice = ref_kernel((
                a[:, :, batch_idx:batch_idx+1],
                b[:, :, batch_idx:batch_idx+1],
                sfa_ref_cpu[:, :, batch_idx:batch_idx+1],
                sfb_ref_cpu[:, :, batch_idx:batch_idx+1],
                None, None,
                c_ref[:, :, batch_idx:batch_idx+1],
            ))
            c_result[:, :, batch_idx] = result_slice[:, :, 0]
    
    return c_result

# ============================================================================
# INPUT GENERATION (for testing)
# ============================================================================

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

    n = 1
    n_padded_128 = 128
    
    a_ref = torch.randint(
        0, 2, (l, m, k // 2), dtype=torch.uint8, device="cuda"
    ).permute(1, 2, 0)
    b_ref = torch.randint(
        0, 2, (l, n_padded_128, k // 2), dtype=torch.uint8, device="cuda"
    ).permute(1, 2, 0)
    a_ref = a_ref.view(torch.float4_e2m1fn_x2)
    b_ref = b_ref.view(torch.float4_e2m1fn_x2)

    c_ref = torch.randn((l, m, n), dtype=torch.float16, device="cuda").permute(
        1, 2, 0
    )
    
    def create_scale_factor_tensors(l, mn, sf_k):
        ref_shape = (l, mn, sf_k)
        ref_permute_order = (1, 2, 0)
        ref_f8_random_int = torch.randint(1, 3, ref_shape, dtype=torch.int8, device='cuda')
        ref_f8_torch_tensor = ref_f8_random_int.to(dtype=torch.float8_e4m3fn)
        ref_f8_torch_tensor_permuted = ref_f8_torch_tensor.permute(*ref_permute_order)
        
        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, 2, 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)

        i_idx = torch.arange(mn, device='cuda')
        j_idx = torch.arange(sf_k, device='cuda')
        b_idx = torch.arange(l, device='cuda')
        
        i_grid, j_grid, b_grid = torch.meshgrid(i_idx, j_idx, b_idx, indexing='ij')
        
        mm = i_grid // (atom_m[0] * atom_m[1])
        mm32 = i_grid % atom_m[0]
        mm4 = (i_grid % 128) // atom_m[0]
        kk = j_grid // atom_k
        kk4 = j_grid % atom_k
        
        reordered_f8_torch_tensor[mm32, mm4, mm, kk4, kk, b_grid] = ref_f8_torch_tensor_permuted[i_grid, j_grid, b_grid]
        
        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)
    
    return (a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, sfa_permuted, sfb_permuted, c_ref)

# ============================================================================
# VALIDATION
# ============================================================================

check_implementation = make_match_reference(ref_kernel, rtol=1e-03, atol=1e-03)