submission 119628

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
from utils import make_match_reference

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

# --------- C++/CUDA inline extension ---------

_cpp_src = r"""
#include <torch/extension.h>

void to_blocked_batched_cuda_launcher(
    torch::Tensor input,
    torch::Tensor output
);

torch::Tensor to_blocked_batched(torch::Tensor input) {
    TORCH_CHECK(input.dim() == 3, "input must be 3D (l, rows, cols)");
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.is_contiguous(), "input must be contiguous");

    auto l = input.size(0);
    auto rows = input.size(1);
    auto cols = input.size(2);
    
    TORCH_CHECK(rows % 128 == 0, "Rows must be multiple of 128");
    TORCH_CHECK(cols % 4 == 0, "Cols must be multiple of 4");

    int64_t n_row_blocks = rows / 128;
    int64_t n_col_blocks = cols / 4;
    
    auto output = torch::empty({l, n_row_blocks * n_col_blocks * 512}, 
                               input.options());

    to_blocked_batched_cuda_launcher(input, output);

    return output;
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cstdint>

#ifndef TILES_PER_CTA
#define TILES_PER_CTA 8  // how many 128x4 tiles each block handles along N
#endif

template <typename T>
__global__ void to_blocked_batched_kernel(
    const T* __restrict__ input_ptr,
    T* __restrict__ output_ptr,
    int64_t rows,
    int64_t cols,
    int n_row_blocks,
    int n_col_blocks
) { 
    int batch_idx = blockIdx.x;
    int block_row_idx = blockIdx.y;  // Which 128-row block (0 to n_row_blocks-1)
    int base_col_blk  = blockIdx.z * TILES_PER_CTA; // Which 4-col block (0 to n_col_blocks-1)
    
    int tid = threadIdx.x;  
    
    int local_row = tid / 4;   
    int local_col = tid % 4;  

    for (int t = 0; t < TILES_PER_CTA; ++t) {
        int block_col_idx = base_col_blk + t;
        if (block_col_idx >= n_col_blocks) break;
        
        int global_row = block_row_idx * 128 + local_row;
        int global_col = block_col_idx * 4 + local_col;

        if (global_row >= rows || global_col >= cols) continue;
        
        int input_idx = batch_idx * (rows * cols) + global_row * cols + global_col;

        T value = input_ptr[input_idx];
    
        int row_in_32 = local_row % 32;
        int which_32_block = local_row / 32;
        
        int block_id = block_row_idx * n_col_blocks + block_col_idx;

        int output_base = batch_idx * (n_row_blocks * n_col_blocks * 512) + block_id * 512;

        int output_offset = row_in_32 * 16 + which_32_block * 4 + local_col;
    
        output_ptr[output_base + output_offset] = value;
    }
}

void to_blocked_batched_cuda_launcher(
    torch::Tensor input,
    torch::Tensor output
) {
    auto l = input.size(0);
    auto rows = input.size(1);
    auto cols = input.size(2);

    int n_row_blocks = rows / 128;
    int n_col_blocks = cols / 4;

    dim3 block(512);

    int grid_z = (n_col_blocks + TILES_PER_CTA - 1) / TILES_PER_CTA;

    dim3 grid(
        static_cast<unsigned int>(l),
        static_cast<unsigned int>(n_row_blocks),
        static_cast<unsigned int>(grid_z)
    );

    to_blocked_batched_kernel<uint8_t><<<grid, block>>>(
        reinterpret_cast<const uint8_t*>(input.data_ptr()),
        reinterpret_cast<uint8_t*>(output.data_ptr()),
        rows,
        cols,
        n_row_blocks,
        n_col_blocks
    );

    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "CUDA kernel failed: ", cudaGetErrorString(err));
}
"""

# Load the extension
to_blocked_ext = load_inline(
    name="to_blocked_batched_ext",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["to_blocked_batched"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=True,
)

# Wrapper function
def to_blocked_batched(input_matrix):
    """
    CUDA-based version using inline extension.
    Input: (rows, cols, l) on GPU
    Output: (l, flattened_size) on GPU
    """
    return to_blocked_ext.to_blocked_batched(input_matrix)

def custom_kernel(
    data: input_t,
) -> output_t:
    """
    Batched PyTorch implementation of NVFP4 block-scaled GEMM.
    Optimized for batched transfer and minimal memory overhead.
    """
    a_ref, b_ref, sfa_ref_cpu, sfb_ref_cpu, _, _, c_ref = data
    
    M, K, l = a_ref.shape
    K2, N, _ = b_ref.shape

    a_batched = a_ref.permute(2, 0, 1).contiguous()  # [l, M, K]
    b_batched = b_ref.permute(2, 0, 1).contiguous()  # [l, K, N]
    c_batched = c_ref.permute(2, 0, 1).contiguous()  # [l, M, N]

    # Transfer scale factors to GPU and transform
    sfa_gpu = sfa_ref_cpu.cuda().permute(2, 0, 1).contiguous()  # [l, rows, cols]
    sfb_gpu = sfb_ref_cpu.cuda().permute(2, 0, 1).contiguous()  # [l, rows, cols]

    # Process with CUDA kernel
    scale_a_batched = to_blocked_batched(sfa_gpu)
    scale_b_batched = to_blocked_batched(sfb_gpu)

    for l_idx in range(l):
        torch._scaled_mm(
            a_batched[l_idx], 
            b_batched[l_idx].transpose(0, 1),
            scale_a_batched[l_idx],
            scale_b_batched[l_idx],
            bias=None,
            out_dtype=torch.float16,
            out=c_batched[l_idx]
        )
    
    # Copy back to original c_ref layout
    c_ref.copy_(c_batched.permute(1, 2, 0))
    
    return c_ref

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