submission 179507

#!POPCORN leaderboard nvfp4_gemm
import torch
from task import input_t, output_t


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


# Fallback: Convert [rows, cols] scale matrix to the blocked 1D layout expected by torch._scaled_mm
def to_blocked_2d(input_matrix: torch.Tensor) -> torch.Tensor:
    rows, cols = input_matrix.shape
    n_row_blocks = ceil_div(rows, 128)
    n_col_blocks = ceil_div(cols, 4)
    blocks = (
        input_matrix.view(n_row_blocks, 128, n_col_blocks, 4)
        .permute(0, 2, 1, 3)
        .reshape(-1, 4, 32, 4)
        .transpose(1, 2)
        .reshape(-1, 32, 16)
    )
    return blocks.flatten()


# Preferred: Use evaluator-provided permuted layout to build blocked scales for all batches at once
def blocked_from_permuted_all_l(s_perm: torch.Tensor) -> torch.Tensor | None:
    """
    s_perm: Tensor with shape [32, 4, rest_rows, 4, rest_cols, L]
    Returns: Tensor with shape [L, rest_rows * rest_cols * 32 * 16], contiguous
    """
    if s_perm is None or s_perm.numel() == 0 or s_perm.ndim != 6:
        return None
    d0, d1, rest_rows, d3, rest_cols, L = s_perm.shape
    if d0 != 32 or d1 != 4 or d3 != 4:
        return None

    # Reorder to [rest_rows, rest_cols, 32, 4a, 4b, L]
    x = s_perm.permute(2, 4, 0, 1, 3, 5)
    # Merge block factors to [-1, 32, 16, L]
    x = x.reshape(rest_rows * rest_cols, 32, 16, L)
    # Move L to front and flatten the rest to produce per-batch 1D scale vectors
    x = x.permute(3, 0, 1, 2).reshape(L, -1).contiguous()
    return x


# Vectorized fallback when permuted layout is not available:
# build blocked scale vectors for all L in one shot from [rows, cols, L] layout.
def blocked_from_regular_all_l(s: torch.Tensor) -> torch.Tensor:
    """
    s: Tensor with shape [rows, cols, L]
    Returns: Tensor with shape [L, rows * cols] in blocked layout expected by torch._scaled_mm.
    Equivalent to applying to_blocked_2d on each slice s[:, :, l] and stacking.
    """
    if s is None or s.numel() == 0 or s.ndim != 3:
        return torch.empty(0, device=s.device if s is not None else "cuda")

    rows, cols, L = s.shape

    # Fast path: rows divisible by 128 and cols divisible by 4 (expected in benchmark)
    if rows % 128 == 0 and cols % 4 == 0:
        n_row_blocks = rows // 128
        n_col_blocks = cols // 4
        # [rows, cols, L] -> [n_row_blocks, 128, n_col_blocks, 4, L]
        x = s.view(n_row_blocks, 128, n_col_blocks, 4, L)
        # -> [n_row_blocks, n_col_blocks, 128, 4, L]
        x = x.permute(0, 2, 1, 3, 4)
        # -> [n_row_blocks * n_col_blocks, 4, 32, 4, L]
        x = x.reshape(-1, 4, 32, 4, L)
        # -> [n_row_blocks * n_col_blocks, 32, 4, 4, L]
        x = x.transpose(1, 2)
        # -> [n_row_blocks * n_col_blocks, 32, 16, L]
        x = x.reshape(-1, 32, 16, L)
        # Move L to front and flatten the rest: [L, -1]
        x = x.permute(3, 0, 1, 2).reshape(L, -1).contiguous()
        return x

    # Conservative slow path: build per-L with the 2D helper (should rarely trigger)
    scales = []
    for l in range(L):
        scales.append(to_blocked_2d(s[:, :, l]))
    return torch.stack(scales, dim=0).contiguous()


def custom_kernel(data: input_t) -> output_t:
    """
    Optimized implementation of block-scale NVFP4 GEMM using torch._scaled_mm.
    Uses evaluator-provided permuted scale layout when available, otherwise a
    vectorized transform from the regular [rows, cols, L] layout.
    """
    a, b, sfa, sfb, sfa_permuted, sfb_permuted, c = data

    L = c.shape[2]

    with torch.inference_mode():
        # Precompute blocked scale vectors for all l using the permuted layout when available
        scale_a_all = blocked_from_permuted_all_l(sfa_permuted)
        scale_b_all = blocked_from_permuted_all_l(sfb_permuted)

        # Fallback path: build from regular [rows, cols, L] layout in a single vectorized pass
        if scale_a_all is None:
            scale_a_all = blocked_from_regular_all_l(sfa)
        if scale_b_all is None:
            scale_b_all = blocked_from_regular_all_l(sfb)

        # Main loop over batch dimension L
        for l_idx in range(L):
            a_l = a.select(2, l_idx)  # [M, K]
            b_l_t = b.select(2, l_idx).transpose(0, 1)  # [K, N]
            sfa_vec = scale_a_all[l_idx]
            sfb_vec = scale_b_all[l_idx]
            c_l = c.select(2, l_idx)  # [M, N]

            # Perform scaled mm directly into the preallocated output slice
            torch._scaled_mm(
                a_l,
                b_l_t,
                sfa_vec,
                sfb_vec,
                bias=None,
                out=c_l,
                out_dtype=torch.float16,
            )

    return c