submission 242045

#!POPCORN leaderboard nvfp4_dual_gemm

import torch

from task import input_t, output_t


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


def to_blocked(input_matrix):
    """
    Convert scale factors to the blocked layout expected by torch._scaled_mm.

    input_matrix: [mn, k//16] in K-major order.
    """
    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()


def custom_kernel(data: input_t) -> output_t:
    """
    Reference-style implementation of NVFP4 block-scaled dual GEMM with silu activation:
    C = silu(A @ B1^T) * (A @ B2^T).

    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float4e2m1fn] of shape [m, k, l],
            b1: torch.Tensor[float4e2m1fn] of shape [n, k, l],
            b2: torch.Tensor[float4e2m1fn] of shape [n, k, l],
            sfa: torch.Tensor[float8_e4m3fnuz] of shape [m, k // 16, l],
            sfb1: torch.Tensor[float8_e4m3fnuz] of shape [n, k // 16, l],
            sfb2: torch.Tensor[float8_e4m3fnuz] of shape [n, k // 16, l],
            sfa_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_m, 4, rest_k, l],
            sfb1_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_n, 4, rest_k, l],
            sfb2_permuted: torch.Tensor[float8_e4m3fnuz] of shape [32, 4, rest_n, 4, rest_k, l],
            c: torch.Tensor[float16] of shape [m, n, l]
    Returns:
        Tensor containing output in float16: c of shape [m, n, l]
    """
    a, b1, b2, sfa, sfb1, sfb2, _sfa_p, _sfb1_p, _sfb2_p, c = data

    batch = c.shape[2]
    for l_idx in range(batch):
        scale_a = to_blocked(sfa[:, :, l_idx])
        scale_b1 = to_blocked(sfb1[:, :, l_idx])
        scale_b2 = to_blocked(sfb2[:, :, l_idx])

        x1 = torch._scaled_mm(
            a[:, :, l_idx],
            b1[:, :, l_idx].transpose(0, 1),
            scale_a,
            scale_b1,
            bias=None,
            out_dtype=torch.float32,
        )
        x2 = torch._scaled_mm(
            a[:, :, l_idx],
            b2[:, :, l_idx].transpose(0, 1),
            scale_a,
            scale_b2,
            bias=None,
            out_dtype=torch.float32,
        )

        torch.nn.functional.silu(x1, inplace=True)
        x1.mul_(x2)
        c[:, :, l_idx] = x1.to(torch.float16)

    return c