Packed MXFP4 & NVFP4 GEMM + mega-MoE (SM100, true 2-CTA multicast)

[ipiszy]

Summary

Adds packed-FP4 × packed-FP4 GEMM and mega-MoE kernels for SM100 (Blackwell), supporting both MXFP4 and NVFP4 through a single unified kernel keyed on MmaKind:

• Standalone GEMM: mxfp4_gemm_nt, nvfp4_gemm_nt — [M,K] @ [N,K].T -> [M,N] (BF16 out).
• Mega-MoE: mxfp4_mxfp4_mega_moe, nvfp4_nvfp4_mega_moe (mma_type="mxfp4xmxfp4" / "nvfp4xnvfp4").

Both operands are packed E2M1 (2 elems/byte). MXFP4 uses UE8M0 scale factors at gran-32; NVFP4 uses E4M3/UE4M3 scale factors at gran-16 with per-tensor global scales passed as CPU-side scalar kernel params (dequant: acc * gs_act * gs_weight; the MoE L1 output is requantized with the L2-activation global scale).

Key technical details

• True 2-CTA cta_group::2 multicast TMA (SM100_TMA_2SM_LOAD) for activations/weights, with leader-routed barriers and per-CTA scale-factor loads.
• MMA atoms: kind::mxf4 (block32) for MXFP4, kind::mxf4nvf4 (block16 / scale_vec::4X) for NVFP4; UE8M0 vs UE4M3 instruction descriptors.
• gran-16 scale-factor handling (the subtle part): NVFP4 has 2 SF int32s per K-block vs 1 for MXFP4. The per-K-block SF stride, smem/tmem sizing, and UTCCP placement are sized accordingly. Crucially, crossing K-uint32s is done via the SF tensor-memory address — not sf_id, which is a 2-bit descriptor field (a_sf_id_/b_sf_id_) and only selects within a K-uint32.
• Unified on the refactored mega-MoE API (parse_mma_kind / ring-token buffering / sub-byte token byte math via get_element_bits).
• The shared swap-AB epilogue gains an optional out_scale (default 1.0, no-op for fp8/bf16/mxfp4).

Validation (B200)

Test	diff
standalone MXFP4 GEMM	0.00000
standalone NVFP4 GEMM	0.00000
MXFP4 mega-MoE	7.5e-4
NVFP4 mega-MoE	8.4e-4 (sweep: 6 configs, BLOCK_K 128 & 256)

Benchmark (device time via bench_kineto)

• mega-MoE: nvfp4 ≈ 1.04–1.06× the mxfp4 time
• standalone GEMM: nvfp4 ≈ 1.09–1.16× the mxfp4 time

The residual nvfp4 overhead is gran-16 doubling the scale factors (2× UTCCP copies + 2× SF traffic).

Tests

tests/test_mxfp4_gemm.py, tests/test_nvfp4_gemm.py, tests/test_mxfp4_mega_moe.py, tests/test_nvfp4_mega_moe.py, tests/bench_packed_fp4.py.

Known limitations / follow-ups

• Single-rank only in tests; multi-rank / expert-parallel dispatch+combine for packed-FP4 is untested.
• Coverage vs fp8_fp4: only gemm_nt (standalone) + mega-MoE; no nn/tn/tt, m_grouped, or k_grouped variants yet. The standalone GEMM is a fixed-config "de-risk" kernel (no autotuning; requires N%256, M%128, K%128).
• Weight global scales are single per-tensor CPU scalars, not per-expert device tensors.
• L2-input (intermediate) global scale must be pre-calibrated (CPU param).
• MXFP4 BLOCK_K=256 (small-token configs) is routed through the new address-based SF path but only BLOCK_K=128 is covered by mxfp4 tests; nvfp4 BLOCK_K=256 is tested.
• nvfp4 GEMM overhead (~9–16%) is optimizable (fewer UTCCP copies / better SF staging).

---

Note

High Risk
Large new SM100 kernel surface (GEMM + fused mega-MoE/EP) with distinct MXFP4 vs NVFP4 numerics and symmetric-buffer layout changes; regressions would affect core MoE performance and correctness on Blackwell.

Overview
Adds SM100-only packed FP4×FP4 compute: standalone mxfp4_gemm_nt / nvfp4_gemm_nt and fused mxfp4_mxfp4_mega_moe / nvfp4_nvfp4_mega_moe, wired through Python/C++ APIs and new JIT CUDA kernels (2-CTA mxf4 / NVF4 MMA, gran-32 vs gran-16 scale factors, NVFP4 global-scale dequant via optional out_scale on the swap-AB epilogue).

Mega-MoE plumbing is extended for sub-byte activations: new MmaKind::MXFP4 / NVFP4, get_element_bits / get_sf_gran_k, mma_type strings mxfp4xmxfp4 / nvfp4xnvfp4, and symmetric-buffer views that use int8 packed tokens and SF layouts sized by granularity (NVFP4 MoE takes per-expert gate_alpha / up_alpha / l2_input_global_scale / down_alpha).

Ships a benchmarks/ suite (internal MXFP4 vs NVFP4, DeepGEMM vs FlashInfer single-GPU, multi-GPU EP vs NCCL combine) plus a README with captured B200 numbers and run instructions.

^{Reviewed by Cursor Bugbot for commit 5de63a6. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

EP benchmark uses mean not max

Low Severity

The module header and several helpers claim worst-rank (max-over-ranks) latency, but rank_avg_ms all-reduces with SUM and divides by world size, so printed DeepGEMM and FlashInfer EP numbers are means. Straggler ranks are under-represented versus the documented metric.

Additional Locations (2)

• benchmarks/bench_ep_multi_gpu.py#L52-L60
• benchmarks/bench_ep_multi_gpu.py#L85-L104

 

^{Reviewed by Cursor Bugbot for commit 0d7c72a. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 5de63a6. Configure here.}

[cursor]

EP bench mean not worst-rank

Medium Severity

The multi-GPU EP benchmark text says it reports worst-rank latency, but rank_avg_ms all-reduces per-rank times with SUM and divides by world size, so printed DeepGEMM and FlashInfer numbers are cross-rank means. Straggler ranks are smoothed out, which misstates end-to-end EP latency when ranks diverge.

Additional Locations (2)

• benchmarks/bench_ep_multi_gpu.py#L17-L18
• benchmarks/README.md#L63-L64

 

^{Reviewed by Cursor Bugbot for commit 5de63a6. Configure here.}