submission 80641

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

# Scaling factor vector size
sf_vec_size = 16

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


# --------- C++/CUDA inline extension (single-file) ---------

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

torch::Tensor to_blocked_batched_cuda(torch::Tensor input);

torch::Tensor to_blocked_batched_cuda(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);

    auto padded_rows = (rows + 127) / 128 * 128;
    auto padded_cols = (cols + 3)   / 4   * 4;

    auto options = input.options();
    at::Tensor padded;

    if (rows != padded_rows || cols != padded_cols) {
        padded = at::zeros({l, padded_rows, padded_cols}, options);

        // Copy the valid region as a flat slice to avoid complex indexing
        auto flat_padded = padded.view({l, padded_rows * padded_cols});
        auto flat_input  = input.view({l, rows * cols});
        flat_padded.slice(1, 0, rows * cols).copy_(flat_input);
    } else {
        padded = input;
    }

    auto n_row_blocks = padded_rows / 128;
    auto n_col_blocks = padded_cols / 4;

    int64_t per_batch_elems = n_row_blocks * n_col_blocks * 32 * 16;
    at::Tensor output = at::empty({l, per_batch_elems}, options);

    // Launch CUDA kernel (defined in the .cu source)
    extern void to_blocked_batched_cuda_launcher(
        const at::Tensor& padded,
        at::Tensor& output,
        int64_t n_row_blocks,
        int64_t n_col_blocks
    );

    to_blocked_batched_cuda_launcher(padded, output, n_row_blocks, n_col_blocks);

    return output;
}
"""

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

template <typename scalar_t>
__global__ void to_blocked_batched_kernel(
    const scalar_t* __restrict__ padded,
    scalar_t* __restrict__ out,
    int64_t l,
    int64_t padded_rows,
    int64_t padded_cols,
    int64_t n_row_blocks,
    int64_t n_col_blocks
) {
    int64_t batch = blockIdx.y;
    if (batch >= l) return;

    int64_t per_batch_elems = n_row_blocks * n_col_blocks * 32 * 16;

    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= per_batch_elems) return;

    // Decode idx -> (blk, i32, j16)
    int64_t blk  = idx / (32 * 16);
    int64_t rem  = idx % (32 * 16);
    int64_t i32  = rem / 16;   // 0..31
    int64_t j16  = rem % 16;   // 0..15

    int64_t i4   = j16 / 4;    // 0..3
    int64_t ic   = j16 % 4;    // 0..3

    // Correct mapping across row/col blocks
    int64_t br   = blk / n_col_blocks;    // row block
    int64_t bc   = blk % n_col_blocks;    // col block

    int64_t ir   = i4 * 32 + i32;         // 0..127

    int64_t row  = br * 128 + ir;
    int64_t col  = bc * 4   + ic;

    int64_t in_offset  = (batch * padded_rows + row) * padded_cols + col;
    int64_t out_offset = batch * per_batch_elems + idx;

    out[out_offset] = padded[in_offset];
}

void to_blocked_batched_cuda_launcher(
    const at::Tensor& padded,
    at::Tensor& output,
    int64_t n_row_blocks,
    int64_t n_col_blocks
) {
    auto l           = padded.size(0);
    auto padded_rows = padded.size(1);
    auto padded_cols = padded.size(2);

    const int threads = 256;
    int64_t per_batch_elems = n_row_blocks * n_col_blocks * 32 * 16;
    int64_t blocks_x = (per_batch_elems + threads - 1) / threads;

    dim3 grid(blocks_x, l);
    dim3 block(threads);

    auto scalar_type = padded.scalar_type();

    if (scalar_type == at::ScalarType::Float8_e4m3fn) {
        // Treat Float8 as uint8_t (we're just copying bytes, not doing math)
        to_blocked_batched_kernel<uint8_t><<<grid, block>>>(
            reinterpret_cast<const uint8_t*>(padded.data_ptr()),
            reinterpret_cast<uint8_t*>(output.data_ptr()),
            l,
            padded_rows,
            padded_cols,
            n_row_blocks,
            n_col_blocks
        );
    } else {
        // Fallback for regular floating point types (float16, float32, etc.)
        AT_DISPATCH_FLOATING_TYPES_AND_HALF(
            scalar_type, "to_blocked_batched_kernel", [&] {
                to_blocked_batched_kernel<scalar_t><<<grid, block>>>(
                    padded.data_ptr<scalar_t>(),
                    output.data_ptr<scalar_t>(),
                    l,
                    padded_rows,
                    padded_cols,
                    n_row_blocks,
                    n_col_blocks
                );
            }
        );
    }

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

_to_blocked_ext = load_inline(
    name="to_blocked_batched_ext",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    functions=["to_blocked_batched_cuda"],
    verbose=False,
)


# --------- Python wrapper using the CUDA kernel (with CPU fallback) ---------

def to_blocked_batched(input_matrix):
    """
    Batched version of to_blocked that processes all l batches in parallel.
    Input: (l, rows, cols)
    Output: (l, flattened_size)
    """
    # Fast path: CUDA extension
    if input_matrix.is_cuda and input_matrix.is_contiguous():
        return _to_blocked_ext.to_blocked_batched_cuda(input_matrix)

    # Fallback: original PyTorch implementation (CPU or non-contig)
    l, rows, cols = input_matrix.shape

    # Target multiples required by the layout
    padded_rows = ceil_div(rows, 128) * 128
    padded_cols = ceil_div(cols, 4) * 4

    # Pad if needed (broadcast padding across batch dimension)
    if rows != padded_rows or cols != padded_cols:
        padded = torch.zeros(
            (l, padded_rows, padded_cols),
            dtype=input_matrix.dtype,
            device=input_matrix.device,
        )
        padded[:, :rows, :cols] = input_matrix
    else:
        padded = input_matrix

    n_row_blocks = padded_rows // 128
    n_col_blocks = padded_cols // 4

    # Apply same transformations but keep batch dimension
    blocks = padded.view(l, n_row_blocks, 128, n_col_blocks, 4).permute(0, 1, 3, 2, 4)
    rearranged = blocks.reshape(l, -1, 4, 32, 4).transpose(2, 3).reshape(l, -1, 32, 16)

    return rearranged.flatten(start_dim=1)  # Flatten only spatial dims, keep batch


def custom_kernel(data):
    """
    PyTorch implementation of NVFP4 block-scaled GEMV using blocked scale factors.
    """
    a_ref, b_ref, sfa_ref, sfb_ref, _, _, c_ref = data

    # Get dimensions from MxNxL layout
    _, _, l = c_ref.shape

    # Permute to get (l, m, sf_k) and (l, n, sf_k)
    sfa_batched = sfa_ref.permute(2, 0, 1)  # (l, m, sf_k)
    sfb_batched = sfb_ref.permute(2, 0, 1)  # (l, n, sf_k)

    # Process all batches in parallel on GPU - ONCE
    scale_a_batched = to_blocked_batched(sfa_batched)  # (l, flattened_size)
    scale_b_batched = to_blocked_batched(sfb_batched)  # (l, flattened_size)

    # Directly write results to c_ref without intermediate storage
    for l_idx in range(l):
        res = torch._scaled_mm(
            a_ref[:, :, l_idx],
            b_ref[:, :, l_idx].transpose(0, 1),
            scale_a_batched[l_idx],
            scale_b_batched[l_idx],
            bias=None,
            out_dtype=torch.float16,
        )
        c_ref[:, 0, l_idx] = res[:, 0]

    return c_ref


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























































def generate_input(
    m: int,
    k: int,
    l: int,
    seed: int,
):
    """
    Generate input tensors for NVFP4 block-scaled GEMV.
    
    Args:
        m: Number of rows in matrix A
        k: Number of columns in A (and length of vector b)
        l: Batch size
        seed: Random seed for reproducibility
    
    Returns:
        Tuple of (a, b, scale_a, scale_b, c) where:
            a: [m, k, l] - Input matrix in torch.float4e2m1fn_x2 data type
            b: [1, k, l] - Input vector in torch.float4e2m1fn_x2 data type
            scale_a: [m, k, l] - Input scale factors in torch.float8e4m3fn data type
            scale_b: [1, k, l] - Input scale factors in torch.float8e4m3fn data type
            scale_a_permuted: [32, 4, rest_m, 4, rest_k, l] - Input scale factors in torch.float8e4m3fn data type
            scale_b_permuted: [32, 4, rest_n, 4, rest_k, l] - Input scale factors in torch.float8e4m3fn data type
            c: [m, 1, l] - Output vector in torch.float16 data type
    """
    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)
    a_ref = a_ref.view(torch.float4_e2m1fn_x2)
    b_ref = b_ref.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 for both reference
    # kernel and customize kernel. The customized data layout can be found in:
    # https://docs.nvidia.com/cuda/cublas/index.html?highlight=fp4#d-block-scaling-factors-layout
    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 uint8 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
        ref_f8_torch_tensor_permuted = ref_f8_torch_tensor.permute(*ref_permute_order)
        
        atom_m = (32, 4)
        atom_k = 4
        mma_shape = (
            l,  # batch size
            ceil_div(mn, atom_m[0] * atom_m[1]),
            ceil_div(sf_k, atom_k),
            atom_m[0],
            atom_m[1],
            atom_k,
        )

        # Reorder scale factor tensor to (32, 4, rest_m, 4, rest_k, l) layout
        # Which is needed by the CuTe customized kernel
        mma_permute_order = (3, 4, 1, 5, 2, 0)
        # Generate a random int8 tensor, then convert to float8_e4m3fn
        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)
        # Permute according to mma_permute_order
        reordered_f8_torch_tensor = reordered_f8_torch_tensor.permute(*mma_permute_order)

        # GPU-side vectorized reordering (replaces slow CPU nested loops)
        # Create index grids for all dimensions
        i_idx = torch.arange(mn, device='cuda')
        j_idx = torch.arange(sf_k, device='cuda')
        b_idx = torch.arange(l, device='cuda')
        
        # Create meshgrid for all combinations of (i, j, b)
        i_grid, j_grid, b_grid = torch.meshgrid(i_idx, j_idx, b_idx, indexing='ij')
        
        # Calculate target indices in vectorized manner
        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
        
        # Perform the reordering with advanced indexing (all on GPU)
        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)

    sfa_ref = sfa_ref_cpu.to("cuda")
    sfb_ref = sfb_ref_cpu.to("cuda")
    
    return (a_ref, b_ref, sfa_ref, sfb_ref, sfa_permuted, sfb_permuted, c_ref)

#custom_kernel(generate_input(4096, 7168, 8, 1111))