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
    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)
    per_expert_alpha=None,  # (E_per_rank,) float32 — per-expert alpha overrides scalar alpha
):
    """Per-expert grouped GEMM for MoE dispatch using CUTLASS NVFP4.

    Takes 1D per-slot expert IDs and token indices (pre-built by caller).
    SFB weight scales are remapped per-expert inside CUTLASS on each call.
    NO prepack cache — see nvfp4_mega_moe.py for rationale.

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

    If per_expert_alpha is provided, each expert uses its own alpha value
    (activation_global_scale * weight_global_scale[expert]) instead of the
    scalar alpha. This preserves full float32 precision — no lossy float8
    folding of weight global scales.

    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
    provided_slot_token = slot_token

    if provided_slot_token is None:
        slot_token_out = torch.arange(num_slots, device=x_fp4.device)
        slot_x = x_fp4
        slot_x_sf = x_sf
    else:
        slot_token_out = provided_slot_token
        slot_x = x_fp4[provided_slot_token].contiguous()
        slot_x_sf = x_sf[provided_slot_token].contiguous()

    if MEGA_MOE_DEBUG:
        print(f"[cutlass_grouped_gemm] slots={num_slots} K={K} N={N} "
              f"experts={num_experts} per_expert_alpha={'yes' if per_expert_alpha is not None else 'no'}")

    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]

        # Per-expert alpha: activation_gs * weight_gs (float32, no precision loss)
        expert_alpha = float(per_expert_alpha[e]) if per_expert_alpha is not None else alpha

        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} alpha={expert_alpha:.4e}")

        # Shape/dtype contract asserts — SFB bugs hide in silent shape mismatches
        assert expert_x.shape == (M_expert, K // 2), f"expert_x shape {expert_x.shape} != ({M_expert}, {K // 2})"
        assert expert_x_sf.shape == (M_expert, K // 16), f"SFA shape {expert_x_sf.shape} != ({M_expert}, {K // 16})"
        assert expert_w.shape == (K // 2, N), f"expert_w shape {expert_w.shape} != ({K // 2}, {N})"
        assert expert_w_sf.shape == (K // 16, N), f"SFB shape {expert_w_sf.shape} != ({K // 16}, {N})"
        assert expert_x_sf.dtype == torch.float8_e4m3fn, f"SFA dtype {expert_x_sf.dtype}"
        assert expert_w_sf.dtype == torch.float8_e4m3fn, f"SFB dtype {expert_w_sf.dtype}"

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

        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} alpha={expert_alpha:.4e}")

        slot_out[e_idx] = expert_out

    return slot_out, slot_token_out