nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py

"""
CUTLASS NVFP4 Block-Scaled GEMM — Native Blackwell SM100 kernel.

Uses the pre-compiled PyTorch CUDA extension (cutlass_nvfp4_gemm._C)
which invokes native mxf8f6f4.block_scale tensor core instructions.
"""

import os
import torch

MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0"))

try:
    from cutlass_nvfp4_gemm import _C
    _CUTLASS_AVAILABLE = True
except ImportError:
    _CUTLASS_AVAILABLE = False


def cutlass_nvfp4_blockscaled_gemm(
    A_packed,   # (M, K_half) int8 packed E2M1
    SFA,        # scale factors for A (float8_e4m3fn)
    B_packed,   # (K_half, N) int8 packed E2M1, column-major for CUTLASS
    SFB,        # scale factors for B — either (sf_k, N) float8_e4m3fn row-major, or prepacked CUTLASS layout
    M, N, K,    # Problem dimensions (K in FP4 elements)
    alpha=1.0,  # fp32 scalar applied in epilogue: D = alpha * A @ B + beta * C
    sfb_prepacked=False,  # True if SFB is already in CUTLASS layout
):
    """Single NVFP4 block-scaled GEMM using CUTLASS.

    If sfb_prepacked=True, SFB is assumed to be in CUTLASS interleaved layout
    (from prepack_sfb) and the SFB remap is skipped.
    """
    if not _CUTLASS_AVAILABLE:
        raise RuntimeError("CUTLASS NVFP4 GEMM extension not available")
    if sfb_prepacked:
        return _C.forward_prepacked_sfb(A_packed, SFA, B_packed, SFB, M, N, K, alpha)
    else:
        return _C.forward(A_packed, SFA, B_packed, SFB, M, N, K, alpha)


def prepack_sfb(SFB, M, N, K):
    """Pre-remap SFB weight scales into CUTLASS interleaved layout.

    Call once after weight transform. Returns a tensor that can be passed
    to cutlass_nvfp4_blockscaled_gemm with sfb_prepacked=True.

    M is used for layout sizing. Test with different M values to confirm
    SFB layout is M-independent; if so, any valid M works (e.g. 128).
    """
    if not _CUTLASS_AVAILABLE:
        raise RuntimeError("CUTLASS NVFP4 GEMM extension not available")
    return _C.prepack_sfb(SFB, M, N, K)


def cutlass_grouped_nvfp4_gemm(
    x_fp4,        # (num_slots_or_tokens, K_half) int8 packed E2M1
    x_sf,         # (num_slots_or_tokens, sf_k) float8_e4m3fn block scales
    weights,      # (E_per_rank, K_half, N) int8 packed E2M1, column-major for CUTLASS
    weight_sf,    # (E_per_rank, sf_k, N) float8_e4m3fn, column-major — or prepacked (E_per_rank, sfb_size) if sfb_prepacked=True
    slot_expert_ids,  # (num_slots,) int32 — per-slot local expert IDs
    slot_token=None,  # (num_slots,) int64 — per-slot token indices (default: arange)
    alpha=1.0,    # fp32 scalar: D = alpha * A @ B (from stage_activation global scale)
    sfb_prepacked=False,  # True if weight_sf is already prepacked into CUTLASS layout
):
    """Per-expert grouped GEMM for MoE dispatch using CUTLASS NVFP4.

    Takes 1D per-slot expert IDs and token indices (pre-built by caller).
    Returns slot-based output: one row per (token, topk) slot.

    For L1: x_fp4 has num_tokens rows, slot_token maps slots→rows.
    For L2: x_fp4 has num_slots rows, slot_token is just arange(num_slots).

    Returns:
        slot_out:    (num_slots, N) bfloat16 — per-slot GEMM results
        slot_token:  (num_slots,) int64 — token index for each slot
    """
    num_slots = slot_expert_ids.shape[0]
    K_half = x_fp4.shape[1]
    K = K_half * 2
    N = weights.shape[2]
    num_experts = weights.shape[0]

    if num_slots == 0:
        slot_out = torch.empty(0, N, dtype=torch.bfloat16, device=x_fp4.device)
        slot_token_out = torch.empty(0, dtype=torch.int64, device=x_fp4.device)
        return slot_out, slot_token_out

    # Use provided slot_token or default to identity mapping
    if slot_token is None:
        slot_token = torch.arange(num_slots, device=x_fp4.device)

    if MEGA_MOE_DEBUG:
        print(f"[cutlass_grouped_gemm] slots={num_slots} K={K} N={N} "
              f"experts={num_experts} sfb_prepacked={sfb_prepacked}")

    # Gather input rows by slot_token when x_fp4 has more tokens than slots
    # (L1: x_fp4=num_tokens, L2: x_fp4=num_slots)
    if x_fp4.shape[0] != num_slots:
        slot_x = x_fp4[slot_token]
        slot_x_sf = x_sf[slot_token]
    else:
        slot_x = x_fp4
        slot_x_sf = x_sf

    slot_out = torch.empty(num_slots, N, dtype=torch.bfloat16, device=x_fp4.device)

    for e in range(num_experts):
        expert_slots = (slot_expert_ids == e)
        if not expert_slots.any():
            continue

        e_idx = expert_slots.nonzero(as_tuple=True)[0]
        expert_x = slot_x[e_idx]
        expert_x_sf = slot_x_sf[e_idx]
        expert_w = weights[e]
        expert_w_sf = weight_sf[e]
        M_expert = e_idx.shape[0]

        if MEGA_MOE_DEBUG and e < 3 and M_expert > 0:
            print(f"[GEMM-IN] expert={e} M={M_expert} N={N} K={K} "
                  f"w shape={expert_w.shape} sfb_prepacked={sfb_prepacked}")

        expert_out = cutlass_nvfp4_blockscaled_gemm(
            expert_x, expert_x_sf,
            expert_w, expert_w_sf,
            M_expert, N, K,
            alpha=alpha,
            sfb_prepacked=sfb_prepacked,
        )

        if MEGA_MOE_DEBUG:
            if torch.isnan(expert_out).any() or torch.isinf(expert_out).any():
                raise RuntimeError(
                    f"expert {e} of {num_experts}: GEMM emitted NaN/Inf. "
                    f"M={M_expert} N={N} K={K}")

        slot_out[e_idx] = expert_out

    return slot_out, slot_token