submission 418649

import torch
from task import input_t, output_t

# Scaling factor vector size
sf_vec_size = 16


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


# Helper function to convert scale factor tensor to blocked format
def to_blocked(input_matrix):
    rows, cols = input_matrix.shape

    # Please 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)
    padded_rows = n_row_blocks * 128
    padded_cols = n_col_blocks * 4

    # Pad the input matrix if necessary
    if padded_rows != rows or padded_cols != cols:
        padded = torch.nn.functional.pad(
            input_matrix,
            (0, padded_cols - cols, 0, padded_rows - rows),
            mode="constant",
            value=0,
        )
    else:
        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(data: input_t) -> output_t:
    """
    PyTorch implementation of NVFP4 block-scaled group GEMM.
    Using torch._scaled_mm for each group as baseline.

    Args:
        data: list of tuples (abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes)
    Returns:
        list of output tensors [c_0, c_1, ..., c_G-1]
    """
    abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes = data
    result_tensors = []

    for i, ((a, b, c), (sfa, sfb), (m, n, k, l)) in enumerate(
        zip(abc_tensors, sfasfb_tensors, problem_sizes)
    ):
        # Process each group's matrices
        for l_idx in range(l):
            # Convert the scale factor tensor to blocked format
            scale_a = to_blocked(sfa[:, :, l_idx])
            scale_b = to_blocked(sfb[:, :, l_idx])

            # Compute GEMM using torch._scaled_mm
            # (m, k) @ (n, k).T -> (m, n)
            res = torch._scaled_mm(
                a[:, :, l_idx].view(torch.float4_e2m1fn_x2),
                b[:, :, l_idx].transpose(0, 1).view(torch.float4_e2m1fn_x2),
                scale_a.cuda(),
                scale_b.cuda(),
                bias=None,
                out_dtype=torch.float16,
            )
            c[:, :, l_idx] = res

        result_tensors.append(c)

    return result_tensors