submission 180443

#!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


# Cache internal op to avoid repeated global lookups in hot path
SCALED_MM = torch._scaled_mm

# Optional: disable Torch Dynamo graph capture to keep this kernel purely eager.
# This avoids one-time compilation overhead that can skew short benchmarks.
if hasattr(torch, "_dynamo"):
    dynamo_disable = torch._dynamo.disable
else:
    def dynamo_disable():
        def wrapper(fn):
            return fn
        return wrapper


# 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:
        # Empty placeholder on CUDA device; not expected in normal runs.
        device = s.device if (s is not None and s.is_cuda) else "cuda"
        dtype = s.dtype if s is not None else torch.float16
        return torch.empty((0,), device=device, dtype=dtype)

    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()


# Specialized helpers for the common leaderboard case L == 1
def blocked_from_permuted_single(s_perm: torch.Tensor) -> torch.Tensor | None:
    """
    Fast path for L == 1 using the permuted layout.
    s_perm: [32, 4, rest_rows, 4, rest_cols, 1]
    Returns: 1D tensor of length rest_rows * rest_cols * 32 * 16.
    """
    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 or L != 1:
        return None

    # Drop the singleton L dimension and reorder to [rest_rows, rest_cols, 32, 4a, 4b]
    x = s_perm[..., 0].permute(2, 4, 0, 1, 3)
    # Merge to [-1, 32, 16] then flatten
    x = x.reshape(rest_rows * rest_cols, 32, 16)
    return x.reshape(-1).contiguous()


def blocked_from_regular_single(s: torch.Tensor) -> torch.Tensor:
    """
    Fast path for L == 1 from regular [rows, cols, 1] layout.
    Returns a 1D blocked scale vector.
    """
    if s is None or s.numel() == 0 or s.ndim != 3:
        # Should not happen in normal runs; return empty placeholder.
        device = s.device if (s is not None and s.is_cuda) else "cuda"
        dtype = s.dtype if s is not None else torch.float16
        return torch.empty((0,), device=device, dtype=dtype)

    rows, cols, L = s.shape
    if L != 1:
        # Defer to the general routine if L != 1 unexpectedly.
        all_l = blocked_from_regular_all_l(s)
        return all_l[0] if all_l.numel() > 0 else all_l

    # Aligned fast path
    if rows % 128 == 0 and cols % 4 == 0:
        n_row_blocks = rows // 128
        n_col_blocks = cols // 4
        # s is [rows, cols, 1], slice out the single L
        x = s[..., 0]  # [rows, cols], still contiguous
        x = x.view(n_row_blocks, 128, n_col_blocks, 4)
        x = x.permute(0, 2, 1, 3).reshape(-1, 4, 32, 4)
        x = x.transpose(1, 2).reshape(-1, 32, 16)
        return x.reshape(-1).contiguous()

    # Fallback: use the 2D helper on the single slice
    return to_blocked_2d(s[:, :, 0])


@dynamo_disable()
@torch.inference_mode()
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.

    Includes a specialized fast path for the common leaderboard case L == 1
    to minimize Python-side overhead.
    """
    a, b, sfa, sfb, sfa_permuted, sfb_permuted, c = data

    # Batch dimension along the third axis
    L = c.shape[2]

    # Bring hot helpers into local scope to reduce attribute/global lookups
    scaled_mm = SCALED_MM
    perm_single = blocked_from_permuted_single
    reg_single = blocked_from_regular_single
    perm_all = blocked_from_permuted_all_l
    reg_all = blocked_from_regular_all_l

    # Fast path for the common leaderboard case L == 1:
    # avoids per-l Python loop and reduces scale preprocessing overhead.
    if L == 1:
        # Use simple slicing views to avoid extra shape computations
        a_l = a[..., 0]  # [M, K]
        b_l_t = b[..., 0].transpose(0, 1)  # [K, N], metadata-only view
        c_l = c[..., 0]  # [M, N]

        # Prefer permuted layout when available
        sfa_vec = perm_single(sfa_permuted)
        if sfa_vec is None:
            sfa_vec = reg_single(sfa)

        sfb_vec = perm_single(sfb_permuted)
        if sfb_vec is None:
            sfb_vec = reg_single(sfb)

        scaled_mm(
            a_l,
            b_l_t,
            sfa_vec,
            sfb_vec,
            bias=None,
            out=c_l,
            out_dtype=torch.float16,
        )
        return c

    # General path for arbitrary L (rare in leaderboard but needed for correctness)
    # Precompute blocked scale vectors for all l using the permuted layout when available
    scale_a_all = perm_all(sfa_permuted)
    scale_b_all = perm_all(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 = reg_all(sfa)
    if scale_b_all is None:
        scale_b_all = reg_all(sfb)

    # Reduce per-iteration transpose overhead by transposing B once (metadata-only)
    b_t = b.transpose(0, 1)  # [K, N, L]

    for l_idx in range(L):
        a_l = a.select(2, l_idx)        # [M, K]
        b_l_t = b_t.select(2, l_idx)    # [K, N]
        sfa_vec = scale_a_all[l_idx]    # [scale_len]
        sfb_vec = scale_b_all[l_idx]    # [scale_len]
        c_l = c.select(2, l_idx)        # [M, N]

        scaled_mm(
            a_l,
            b_l_t,
            sfa_vec,
            sfb_vec,
            bias=None,
            out=c_l,
            out_dtype=torch.float16,
        )

    return c