From dd1a973bcbc18d0f04c65e27a41a08c50cb7f182 Mon Sep 17 00:00:00 2001 From: Ying Zhang Date: Sat, 27 Jun 2026 00:28:41 +0000 Subject: [PATCH 1/9] Add packed MXFP4 x MXFP4 GEMM and mega-MoE kernels (SM100) Adds a standalone packed-FP4 x packed-FP4 GEMM (`mxfp4_gemm_nt`) and a packed MXFP4 mega-MoE (`mxfp4_mxfp4_mega_moe`) for SM100, both using a true 2-CTA `cta_group::2` multicast TMA load path (`SM100_TMA_2SM_LOAD`) with leader-routed barriers and per-CTA scale-factor loads. Integrated onto the unified mega-MoE API (`mma_type`/`parse_mma_kind`, ring-token buffering): adds `MmaKind::MXFP4` with sub-byte (`get_element_bits`) token byte math, packed-FP4 symmetric-buffer slicing, and the packed-FP4 epilogue that writes E2M1 nibbles directly to their [token][inter] positions. Validated on B200: standalone diff=0.0; mega-MoE diff=7.5e-4. Co-authored-by: Cursor --- HANDOFF.md | 445 +++++ csrc/apis/gemm.hpp | 24 + csrc/apis/mega.hpp | 116 +- csrc/jit_kernels/heuristics/mega_moe.hpp | 6 +- csrc/jit_kernels/heuristics/sm100.hpp | 3 +- csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp | 148 ++ .../impls/sm100_mxfp4_mxfp4_mega_moe.hpp | 304 ++++ deep_gemm/__init__.py | 3 + deep_gemm/include/deep_gemm/common/types.cuh | 15 + .../deep_gemm/impls/sm100_mxfp4_gemm.cuh | 402 +++++ .../impls/sm100_mxfp4_mxfp4_mega_moe.cuh | 1474 +++++++++++++++++ deep_gemm/include/deep_gemm/ptx/tcgen05.cuh | 24 + deep_gemm/mega/__init__.py | 24 + tests/test_mxfp4_gemm.py | 48 + tests/test_mxfp4_mega_moe.py | 96 ++ 15 files changed, 3122 insertions(+), 10 deletions(-) create mode 100644 HANDOFF.md create mode 100644 csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp create mode 100644 csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp create mode 100644 deep_gemm/include/deep_gemm/impls/sm100_mxfp4_gemm.cuh create mode 100644 deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh create mode 100644 tests/test_mxfp4_gemm.py create mode 100644 tests/test_mxfp4_mega_moe.py diff --git a/HANDOFF.md b/HANDOFF.md new file mode 100644 index 0000000000..2c8283155f --- /dev/null +++ b/HANDOFF.md @@ -0,0 +1,445 @@ +# Handoff: Packed MXFP4 × MXFP4 for SM100 (Mega MoE) + +Status: **DONE. Packed MXFP4 × MXFP4 mega-MoE kernel PASSES end-to-end on B200 +(`tests/test_mxfp4_mega_moe.py`, `diff = 0.00075`).** Standalone GEMM also validated. + +## Update (final) — mega-MoE FIXED end-to-end + +Two fixes closed it out: +1. **Hang** (see below): per-CTA `SM90_TMA_LOAD_2D` + cross-CTA loads-done barrier. +2. **L1 packed-FP4 epilogue numerics**: the swap-AB epilogue does NOT need a + transposing `stmatrix`. Using the verified `SM100_TMEM_LOAD_16dp256b1x` fragment + map, each lane writes its e2m1 nibbles DIRECTLY to the `[token][inter]` packed + smem position (plain layout, swizzle-0 TMA store). The map (lane `a=lane%4`, + `b=lane/4`, warp `w`): `swiglu[i*2+0].{x,y} -> (tok 2a/2a+1, inter w*16+b)`, + `swiglu[i*2+1].{x,y} -> (tok 2a/2a+1, inter w*16+8+b)`. K-major packing of 2 + consecutive inter per byte is done by an even-`b` lane fetching its `b+1` + partner's nibble via `__shfl_sync(lane+4)`. Verified `warp_reduce<4,true>` + reduces over `b` (inter) via `shfl_xor 4/8/16`, and the SF write is guarded by + `lane_idx < 4` (b==0), so the existing UE8M0 SF code is correct as-is. + Isolation (`/tmp/iso2.py`, identity L2 weights): `diff 0.98 -> 0.00078`. + +Run: `python tests/test_mxfp4_mega_moe.py` -> `diff = 0.00075 MXFP4 mega MoE passed.` + +### Update — TRUE 2-CTA `cta_group::2` multicast TMA (standalone + mega) + +Both kernels now use genuine `cute::SM100_TMA_2SM_LOAD` (`cp.async.bulk.tensor.cta_group::2`) +instead of per-CTA `SM90_TMA_LOAD`. Key facts were obtained by **instrumenting the real +CUTLASS kernel** (example 72b forced to a 2x1 cluster, `-DDG_TMA_LOG` printf in +`cute/arch/copy_sm100_tma.hpp` + `cutlass/pipeline/sm100_pipeline.hpp`): + +- Each CTA issues its OWN per-CTA box at its OWN coord (acts: leader `m=0`, peer `m=128`; + i.e. keep `load_block_m` box + the per-CTA M offset). It is NOT full-box + shared-coord. +- Peer→leader tx routing is AUTOMATIC under cluster launch: the peer's smem/barrier address + has bit 24 set (cluster addressing), and the atom masks it (`Sm100MmaPeerBitMask=0xFEFFFFFF`) + → all tx lands on the LEADER's `full` barrier. No manual cluster-mapping needed. +- For a 2x1 cluster, B is NOT multicast (cluster N=1); both A and B use `SM100_TMA_2SM_LOAD` + (weights replicated by using the same n coord on both CTAs). +- SF must stay a PER-CTA SM90 load into each CTA's own smem (the `..._2cta` UTCCP reads both + CTAs' SF). So the leader's `full` expects `2x data + own SF` (both CTAs' data routed in), + while the non-leader's `full` expects only its own SF; both CTAs' transposer/sync warps wait + their own `full` and arrive on the leader's `with_sf`. + +Standalone `tests/test_mxfp4_gemm.py`: all sizes `diff=0.0`. Mega `diff=0.00075`. Perf +unchanged vs the per-CTA path (2x1 has no multicast bandwidth saving) — still ~1.12-1.23x +over FP8xFP4. Bug that caused the earlier hang: full-box+shared-coord + under-counted +`expect_tx`; fixed by per-CTA box/coord + correct `2x data + own SF` accounting. + +### Performance vs FP8×FP4 (single B200, kernel-only via kineto, BF16 out) + +MXFP4×MXFP4 is consistently faster than `fp8_fp4_mega_moe` (same shapes/routing): + +| tokens | topk/experts | hidden×inter | mxfp4 | fp8fp4 | speedup | +|---|---|---|---|---|---| +| 512 | 4/32 | 4096×4096 | 183µs/1126TF | 192µs/1076TF | 1.05x | +| 1024 | 2/8 | 4096×4096 | 109µs/1886TF | 122µs/1689TF | 1.12x | +| 4096 | 2/8 | 4096×4096 | 306µs/2697TF | 342µs/2409TF | 1.12x | +| 8192 | 2/8 | 4096×4096 | 490µs/3365TF | 600µs/2748TF | 1.22x | +| 4096 | 4/32 | 7168×2048 | 499µs/2893TF | 641µs/2253TF | 1.28x | + +Both share the FP4 tensor-core path + FP4 L2 input; MXFP4 wins because L1 activations +are FP4 (0.5 B/elem) vs FP8 (1 B/elem), halving L1 activation traffic. Gap widens with +size (more bandwidth/compute-bound). Trade-off: lower L1-input precision than FP8. +Bench: `/tmp/bench_compare.py [tokens experts topk hidden inter]`. + +--- +(earlier-session notes below) + +## Update (latest session) — mega-MoE hang FIXED, GEMM path working + +The long-standing 2-CTA TMA **deadlock is resolved**. Root cause + fix: + +- The mega used `cute::SM100_TMA_2SM_LOAD_2D` (`cta_group::2` multicast) for the + packed-FP4 A/B loads, but with DeepGEMM's hand-rolled half-box descriptor + + per-CTA arrive scheme. cute's true 2-SM multicast needs a different (full-box, + cluster-layout) descriptor, so the leader's `full_barrier` `expect_tx` never + matched what the atom delivered → MMA never fired → hang. +- Fix = adopt the **validated standalone's proven load mechanism** (per-CTA + `cute::SM90_TMA_LOAD_2D` + a cross-CTA "loads-done" barrier funneled to the + leader). The host descriptors were already correct for this (acts half-box + `load_block_m`, weights full `load_block_n`). Changes in + `deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh`: + - A/B load warps: `SM100_TMA_2SM_LOAD_2D` → `SM90_TMA_LOAD_2D`; SF `tma::copy` + multicast `2`→`1`; both CTAs `arrive_and_expect_tx(own bytes)` (dropped the + leader/non-leader split). + - New `with_sf_full_barriers[kNumStages]` in `SharedStorage`; init `full=2` + (A+B per CTA), `with_sf = 2*32`. + - Idle warp `kNumDispatchWarps+3` repurposed as the **cross-CTA sync warp**: + waits its CTA's `full_barriers[stage]`, then all 32 lanes `arrive(0u)` on the + leader's `with_sf_full_barriers[stage]`. The leader-only MMA now waits + `with_sf_full_barriers` instead of `full_barriers`. + - Host smem sizing updated for the 3rd per-stage barrier (`2*8`→`3*8`) in BOTH + `csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp` and + `csrc/jit_kernels/heuristics/mega_moe.hpp` (this was the illegal-access cause + after enlarging `SharedStorage`). + +Result: `python tests/test_mxfp4_mega_moe.py` runs the kernel cleanly; output +magnitude is sane (`y~0.23` vs `ref~0.19`) but **`diff≈0.99`** (values scrambled). + +### Remaining task: L1 packed-FP4 swap-AB epilogue transpose (numerics) + +The first-GEMM (L1) epilogue writes packed E2M1 + UE8M0 SF to the `l2_acts` pool +buffer that the L2 GEMM reads. The working FP8 sibling +(`sm100_fp8_fp4_mega_moe.cuh`) uses a **transposing** store-matrix +(`SM100_U8x4_STSM_T`, `.b8`) for the swap-AB layout; the MXFP4 port currently does +a **plain, non-transposed** `uint16` write (lines ~1091-1101 of the mega cuh, +marked `// VALIDATE` / stage-1 placeholder) → scrambled L1 output → `diff≈0.99`. +Isolation (`/tmp/iso_l1.py`, 1 expert / top-1) confirms L1 packed bytes are +scrambled (~11% match, confounded by gate/up gran-8 interleave + pool order). + +Direction: transpose at **b16 granularity** (each `e2m1x4` = 4 packed FP4 = the +K-major packing unit must stay together). A `SM90_U32x1_STSM_T` +(`stmatrix.x1.m8n8.b16.trans`) primitive was added to `ptx/ld_st.cuh`. Still need +to: gather `kNumAtomsPerStore` `e2m1x4` per lane into `uint32`s, pick x1/x2/x4 by +`kNumAtomsPerStore/2`, and get the transposed smem addressing + the L1-output TMA +descriptor swizzle (`swizzle_acts_mode/2`) consistent. Validate with an L1-output +isolation that accounts for the gate/up interleave and pool token order. + +--- +(Previous session below) + +Status: **Standalone packed `mxf4` 2-CTA GEMM is COMPILED + VALIDATED on B200 +(numerically exact, `calc_diff == 0.0`).** Remaining work is the port into the +mega-MoE kernel. + +## Update (this session) — standalone de-risk GEMM DONE + +Environment: this box has CUDA 13.0 + 8× B200 (unlike the original box), so the +standalone path was fully built and validated on hardware. + +What now exists and passes: +- `deep_gemm/include/deep_gemm/impls/sm100_mxfp4_gemm.cuh` — **rewritten** to mirror + the proven `sm100_fp8_fp4_gemm_1d1d` structure (its scheduler, warp roles, barrier + scheme incl. the dedicated **SF warp-transpose warp + `with_sf` barrier**, and the + swap-AB BF16 epilogue), specialized to packed `mxf4`: both operands packed E2M1 + (2/byte, byte-addressed smem), `UMMA_K = 64`, `SM100_MMA_MXF4_2x1SM_SS`, UE8M0 SF + gran-32 (`sf_id = k*2`), K-major packed swizzle = `BLOCK_K/2` bytes, descriptors via + `make_smem_desc` (mirrors `sm100_fp4_mqa_logits`). +- `csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp` — host/JIT `LaunchRuntime` with packed + A/B tensormaps (`make_tma_2d_desc(..., fp4_unpacked_smem=false)` → `16U4_ALIGN8B`), + UE8M0 SF descriptors, BF16 D, smem sizing, cluster-2 launch. Registered in + `csrc/apis/gemm.hpp` as `mxfp4_gemm_nt` (+ pybind), exported in `deep_gemm/__init__.py`. +- `tests/test_mxfp4_gemm.py` — quantizes A/B to packed E2M1 + UE8M0 SF (gran-32, via + `per_token_cast_to_fp4` + `get_mn_major_tma_aligned_packed_ue8m0_tensor`), runs the + kernel, compares to a dequant reference. All cases `diff == 0.0`. + +The three original `// VALIDATE` spots are now CONFIRMED correct on SM100. Two bugs were +fixed during bring-up: +1. The original draft did UTCCP **without the mandatory SF warp-transpose**; the proven + kernel runs a separate transposer warp on **all** CTAs + a `with_sf` barrier — now done. +2. Packed-FP4 data must be loaded with a **single raw `cute::SM90_TMA_LOAD_2D::copy`** + (the descriptor's smem box already spans the full packed `BLOCK_K`); the `tma::copy` + atom-splitter assumes byte-sized elems and overruns the smem stage for sub-byte FP4. + +Run it: `PYTHONPATH=$PWD python tests/test_mxfp4_gemm.py` (after `./develop.sh`). + +### Current standalone limitations (intentional for the de-risk) +- Hardcoded config: `BLOCK_M=BLOCK_N=BLOCK_K=128`, `kNumStages=4`, swap-AB, + **2-CTA (`cluster_n=2`)**, 1 epilogue warpgroup. So it requires **N % 256 == 0** + (even N-blocks), `M % 128 == 0`, `K % 128 == 0` (host-asserted; smaller N deadlocks + by construction, exactly like the production heuristics would reject it). +- No autotuning / `get_best_config` — explicit template params only. + +Original starting commit: `641d7a3` ("Add packed MXFP4 2-CTA foundations + standalone GEMM (WIP)"). + +## Goal + +Add a **pure, packed MXFP4 × MXFP4** path to DeepGEMM on SM100 (Blackwell / B200): +both operands E2M1 (FP4), UE8M0 scale factors at gran-K 32, **packed 2 elements per +byte**, using the dedicated `tcgen05.mma.kind::mxf4` tensor-core path in **2-CTA** +(`cta_group::2`) mode. + +The end target is the **Mega MoE** kernel (`sm100_fp8_fp4_mega_moe.cuh`), whose two +GEMMs currently run FP8×FP4 via `kind::mxf8f6f4`. The plan is to first de-risk the +packed `mxf4` 2-CTA path in a **standalone GEMM**, validate it on B200, then port the +proven pieces into the mega kernel. + +## Why packed `mxf4` (not `mxf8f6f4`) + +Key facts established during design (see also the chat that produced this work): + +- `kind::mxf8f6f4` is the flexible mixed kind (fp8/fp6/fp4, A and B may differ). It + stores every element in an **8-bit container** (FP4 = 1 byte, low nibble) and + contracts **K=32 per instruction**. So FP4 under `mxf8f6f4` gets **neither** + bandwidth/footprint savings **nor** compute speedup — it runs at FP8 storage *and* + FP8 TFLOPS. +- `kind::mxf4` is FP4-only, **packed 2 elements/byte** (`float_e2m1_t`, + `sizeof_bits == 4`), and contracts **K=64 per instruction** → ~2× MACs/instruction → + ~2× FP8 TFLOPS, plus half the smem/gmem/dispatch bandwidth. This is the real FP4 win + and the reason we must move to `mxf4`. +- Hardware constraint: a single block-scaled MMA has **one** `scale_format_` (E4M3 vs + E8M0) and **one** scale-vector size for *both* operands (see + `cute::UMMA::InstrDescriptorBlockScaled` in `cute/arch/mma_sm100_desc.hpp`). Element + data formats (`a_format_`/`b_format_`) can differ, but scale formats cannot. So you + cannot mix e.g. MXFP4 (E8M0/32) and NVFP4 (E4M3/16) **within one GEMM**; the two mega + GEMMs are independent instructions and may differ from each other. + +## Design decisions (carried from the chat) + +- Scope is **MXFP4 only** for now (UE8M0 SF, gran-32, no global scale). NVFP4 (E4M3 SF, + gran-16, per-tensor global scale) was discussed and deferred — it adds an on-device + global-scale problem that MXFP4 avoids entirely. +- The only existing **packed** FP4 precedent in the repo is **1-CTA** in + `sm100_mqa_logits.cuh` (`SM100_MMA_MXF4_SS`, `float_e2m1_t`, `make_smem_desc` with + `/2` addressing, `sf_id = k*2`). The production GEMM (`sm100_fp8_fp4_gemm_1d1d.cuh`) + and the mega kernel both use **unpacked** `mxf8f6f4`. So packed-2CTA is new territory + and needs hardware iteration. + +## What's done (committed in `641d7a3`) + +1. `deep_gemm/include/deep_gemm/ptx/tcgen05.cuh` + - Added `SM100_MMA_MXF4_2x1SM_SS`: `tcgen05.mma.cta_group::2.kind::mxf4.block_scale` + with `.block32` (CUDA ≥ 12.9) / `.scale_vec::2X` fallback. +2. `deep_gemm/include/deep_gemm/common/types.cuh` + - Added `MmaKind::MXFP4` and `get_element_bits()` (4-bit) so byte math can handle + sub-byte FP4. NOTE: `get_element_size()` still returns 1 for MXFP4; callers must + use bit math / divide element counts by 2 for packed storage. +3. `deep_gemm/include/deep_gemm/impls/sm100_mxfp4_gemm.cuh` (new) + - First-draft **standalone** packed MXFP4 2-CTA GEMM. Conventions mirror the mega L1 + GEMM so the validated pieces port back cleanly: + - swap-AB, K-major, `UMMA_M = 256` (2-CTA), `UMMA_N = BLOCK_M`, `UMMA_K = 64`. + - `smem_a` = activations (M axis, multicast, `LOAD_BLOCK_M = BLOCK_M/2`), + `smem_b` = weights (N axis). SFA on M, SFB on N. UTCCP-2cta SF → TMEM. + - packed-FP4 byte addressing: smem tiles are `uint8_t[... * BLOCK_K/2]`, swizzle = + `BLOCK_K/2` bytes, MMA descriptors via `make_smem_desc` with `/2` offsets, + `sf_id = (umma_k_block * (UMMA_BLOCK_K/UMMA_K) + k) * 2`. + - BF16 epilogue (TMEM → STSM → `SM90_TMA_STORE_2D`). + +## `// VALIDATE` spots (need SM100 compile/run to confirm) + +Grep `VALIDATE` in `sm100_mxfp4_gemm.cuh`: + +1. **Packed-FP4 TMA load + swizzle**: that `tma::copy` produces the correct K-major packed layout, + and the `arrive_and_expect_tx` byte counts match (A multicast counts ×2). +2. **`make_smem_desc` layout/stride for E2M1**: `to_umma_layout_type()` and `stride_byte_offset = 8 * BLOCK_K/2` are right for packed + FP4 (cross-check against the 1-CTA MQA usage). +3. **`sf_id` / K64 mapping**: that 2 gran-32 SFs per K64 instruction with base + `sf_id = k*2` and the UTCCP TMEM column placement is correct for 2-CTA `mxf4`. + +Other likely tuning points: epilogue STSM swizzle/`STORE_BLOCK_M`, register split +(`kNumEpilogueRegisters`/`kNumNonEpilogueRegisters`), `__grid_constant__` arg order. + +## Next steps + +1. ~~Toolchain~~ — DONE (CUDA 13.0 + B200; submodules initialized). +2. ~~Host runtime + JIT wiring~~ — DONE (`sm100_mxfp4_gemm.hpp`). +3. ~~Python entry + test~~ — DONE (`mxfp4_gemm_nt`, `tests/test_mxfp4_gemm.py`). +4. ~~Iterate the `// VALIDATE` spots on B200~~ — DONE (all confirmed, `diff == 0.0`). +## Mega-MoE port — STAGE 1 done (compiles + launches on B200), STAGE 2 = numerics + +New files / edits (all compile into `_C` and JIT-compile on B200): +- `deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh` (fork of the fp8 mega + kernel). Converted to **uniform packed `mxf4`**: both operands packed E2M1, token byte + layouts halved, dispatch pulls `hidden/2` bytes, both L1/L2 GEMMs use the validated + packed path (raw 2-CTA TMA loads, `make_smem_desc`, `SM100_MMA_MXF4_2x1SM_SS`, + `UMMA_K=64`, `sf_id=k*2`, swizzle `BLOCK_K/2`), L1 SwiGLU epilogue emits packed E2M1 + (`cvt.rn.satfinite.e2m1x2.f32`) + UE8M0 SF (`amax/6`). SF UTCCP path kept (SF is + pre-transposed in dispatch / weight transform). +- `csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp` — host runtime + `get_mxfp4_mega_moe_config` + (packed byte math: swizzle `block_k/2`, A/B/output smem halved, `num_bytes_per_pull = hidden/2`). +- `csrc/jit_kernels/heuristics/sm100.hpp` — `get_sf_uttcp_aligned_block_sizes` handles `MmaKind::MXFP4`. +- `csrc/apis/mega.hpp` — `get_symm_buffer_size_for_mxfp4_mega_moe` (packed `hidden/2` token + buffers, int8 views) + `mxfp4_mxfp4_mega_moe` entry (registered in pybind). +- `deep_gemm/mega/__init__.py` — `SymmBuffer(mma_type='mxfp4xmxfp4')` + `mxfp4_mxfp4_mega_moe`; + reuses `transform_weights_for_mega_moe` (interleave is on N, works on packed weights). +- Exported `mxfp4_mxfp4_mega_moe` from `deep_gemm/__init__.py`. + +STAGE 1 result: the kernel **compiles** (168 regs, 16 barriers, smem ~225KB < 232KB cap) and +**launches** on B200 (single-rank harness `/tmp/mega_compile.py`, shapes hidden=512, +inter=512, E8/top2). It currently **hangs / wrong numerics** — expected; the FP4 L1 epilogue +store layout (the only `// VALIDATE` piece) is a plain non-transposed store that almost +certainly doesn't match the L2 `mxf4` TMA read-back, and there may be dispatch/barrier +interplay to debug. + +STAGE 2 (in progress): +- DONE: single-rank **torch MoE reference** at `tests/test_mxfp4_mega_moe.py` + (FP4 quant → grouped L1×W1 → clamp+SwiGLU×weight → per-32 UE8M0 FP4 requant → + grouped L2×W2 → top-k combine). Single-rank harness: `/tmp/mega_compile.py`. +- DONE: confirmed the **FP8 mega kernel runs in the same single-rank harness** (so the + harness/dispatch setup is correct; the hang is in the mxfp4 changes). +- DIAGNOSED the hang (via `DG_JIT_WITH_LINEINFO=1` + `cuda-gdb --batch -ex run -ex 'info cuda + threads'`, interrupting the hung kernel): the **L2-GEMM acts-load warp spins at + `sm100_mxfp4_mxfp4_mega_moe.cuh:695`** — `while (ld_acq_gpu(l2_arrival_mask) != kExpectedMask)` — + waiting for the L2 arrival mask (`0xFF`, 8 bits) that the **L1 SwiGLU epilogue** must set via + `red_or_rel_gpu(l2_arrival_mask, 1 << n_block_idx)`. All other warps are downstream waiters + (dispatch clean-workspace `:605`, epilogue combine grid-sync `barrier.cuh:32/51`). GEMM + warps 4–7 have exited on idle SMs. So the culprit is the **packed-FP4 L1 epilogue** (the only + `// VALIDATE` piece): its placeholder non-transposed store is numerically wrong AND something + in that path prevents the L2 arrival mask from reaching `0xFF` (epilogues appear to complete, + so suspect: the SF/`red_or` interaction, or a per-`n_block` epilogue that never runs because + its L1 MMA output layout is wrong). The OOB that originally corrupted barriers was fixed by + over-sizing `smem_d.l1`. + +UPDATE (debugging done — root cause is NOT the epilogue): instrumented the kernel with +`printf`s (now removed) and traced across all SMs. Findings: +- `L2 stuck ... mask=0` — the L2 arrival mask is **0** (never set), not partial. +- `LA_WAIT/LA_PASS` fire (load warp clears the L1-arrival wait) and `LA_ISSUE` fires (leader + issues the TMA + `arrive_and_expect_tx`), **but `MMAFULL` never fires on any SM** — the + GEMM `full_barriers` never complete, so the MMA never runs, no `tmem_full`, L1 epilogue + never runs (`EPIFULL`=0), mask stays 0, L2 load spins at `:695`. +- => ROOT CAUSE: the GEMM **`full_barrier` `arrive_and_expect_tx` byte count does not match + what the raw `cute::SM100_TMA_2SM_LOAD_2D` delivers for packed FP4**, so the barrier waits + forever for tx that never arrives. The `expect_tx` formulas (`sizeof(smem_a[0])*2 + sizeof(smem_sfa[0])*2` + for the A/SFA warp; `sizeof(smem_b[0]) + sizeof(smem_sfb[0])*2` for the B/SFB warp) were + carried over from the FP8 kernel and are wrong for the packed-FP4 2-CTA load tx accounting. + +FURTHER PROBING (done — rules out the easy fix): tried `expect_tx` for the A/B data at +`*2`, `*1`, and even **SF-only** (data dropped). `MMAFULL` stays 0 in ALL cases — i.e. the +barrier never completes even when almost nothing is expected. Since the 4 arrivals do happen, +this means the packed A/B **2-CTA TMA loads deliver ~0 tx**: `cute::SM100_TMA_2SM_LOAD_2D` +does not work with the packed `16U4_ALIGN8B` descriptor (it's silently rejected, and the +failed load poisons the `full_barrier` so no tx is counted). This is consistent with the fact +that BOTH the validated standalone (`sm100_mxfp4_gemm.cuh`, diff=0) and `sm100_fp4_mqa_logits.cuh` +load packed FP4 with **per-CTA `cute::SM90_TMA_LOAD_2D` (1-CTA), never the 2-CTA multicast load**. + +ROOT CAUSE (from CUTLASS dig, 2026-06): CUTLASS's SM100 packed-FP4 (mxf4/nvf4) GEMMs use +**`MmaTileShape_MNK = Shape<_128,_128,_256>`** — i.e. **K-tile = 256**. For packed FP4 +(0.5 B/elem) that is a **128-byte** K extent → `cute::detail::sm100_smem_selector` picks +`Layout_K_SW128_Atom` (**128B swizzle**). The cute TMA atom is built with the packed +`float_e2m1_t` element type and `CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B` (identical to DeepGEMM's +`make_tma_2d_desc(..., fp4_unpacked_smem=false)`), and the 2-SM-ness comes purely from the copy +atom (`SM100_TMA_2SM_LOAD[_MULTICAST]` = `cp.async.bulk.tensor.2d.cta_group::2...`) + multicast +masks. So packed-FP4 + 2-CTA multicast IS supported. + +The DeepGEMM mega uses `BLOCK_K=128` packed → only **64B** (`SW64`) swizzle. Every working 2-SM +kernel (fp8) uses **128B** swizzle (fp8 gets it free: `128 elems × 1B = 128B`). Strong hypothesis: +the `cta_group::2` 2-SM TMA path requires (or is only exercised with) a 128B-swizzle smem layout, +so the 64B packed layout silently fails to deliver. **FIX: use `BLOCK_K = 256` for the mxfp4 +mega+standalone** → `kSwizzleAMode = BLOCK_K/2 = 128` → `SW128`, matching CUTLASS and fp8. This +keeps the desired **2-CTA multicast** (`SM100_TMA_2SM_LOAD_2D`). Mechanical follow-through: +`BLOCK_K/UMMA_BLOCK_K = 2` UMMA-K sub-blocks (already handled), and SF load/indexing for 2 packed +SF-ints per K-block (vs 1 today; `sf_smem_outer_dim = BLOCK_K/(32*4) = 2`). + +UPDATE (BLOCK_K=256 tested, swizzle hypothesis REFUTED): set `BLOCK_K=256` (confirmed in the +instantiation: template arg `256`, TMA desc `swizzle:128`, smem box `[256,32]`) — the 2-SM load +STILL delivers 0 (`MMAFULL`=0, hang). So swizzle (64B vs 128B) is NOT the cause. Empirical matrix: + - packed `16U4` + `SM90` 1-CTA + DeepGEMM `make_tma_2d_desc` => WORKS (standalone, diff=0) + - 8-bit + `cta_group::2` (2-SM) + DeepGEMM `make_tma_2d_desc` => WORKS (fp8 mega) + - packed `16U4` + `cta_group::2` (2-SM) + DeepGEMM `make_tma_2d_desc` => DELIVERS 0 (hang), any swizzle +So the blocker is specifically the **(`cta_group::2` 2-SM load) × (`16U4` packed descriptor built by +DeepGEMM's hand-rolled `make_tma_2d_desc`)** combination. CUTLASS makes the same combination work, +but it builds the descriptor via cute's `make_tma_atom_A_sm100(... ClusterLayout_VMNK ...)`, +which derives the 2-SM box partitioning / coords from the cluster layout — something DeepGEMM's +`make_tma_2d_desc` does not replicate for the 2-SM packed case. + +*** ROOT CAUSE FOUND (descriptor field-diff, option A) *** +Built a cute harness (`/tmp/tma_dump.cu`, compiles with `nvcc -arch=sm_100a -I third-party/cutlass/include +-I third-party/cutlass/tools/util/include`) that instantiates a CUTLASS 2-SM mxf4 mainloop +(`mx_float4_t`, MmaTileShape `<128,128,256>`, ClusterShape `<2,1,1>`), extracts its +A-operand `CUtensorMap` via `params.tma_load_a.get_tma_descriptor()`, and diffs it against +DeepGEMM's `make_tma_2d_desc` (replicated) for the same A. Only 2 of 16 words differ: + word[7]: cute `0x7f` (=127 → box-M **128**) vs deepgemm `0x3f` (=63 → box-M **64**) + word[1]: the dependent box-stride field (differs for the same reason) +=> **cute uses the FULL tile-M (128) as the TMA box and lets `cta_group::2` split it 64/64 across +the 2 CTAs; DeepGEMM uses the pre-split `load_block_m` (64).** For 8-bit (fp8) the half-box works, +but for packed `16U4` the `cta_group::2` load needs the FULL-tile box and silently delivers 0 with +the half-box — the exact hang observed. + +THE FIX: build the mxfp4 mega's A/B packed-FP4 TMA descriptors with the **full block-M / block-N** +box (not `load_block_m`/`load_block_n`), and issue the **base coord** (drop the per-CTA +`m_idx += get_valid_m/2` offset) so `cta_group::2` does the split — matching cute. Keep the smem +tile sized per-CTA (`load_block_m`). NOTE: the mega is swap-AB (acts=UMMA-N, weights=UMMA-M=256), so +apply the full-box to the operand actually split across the 2 SMs and re-derive the descriptor via a +packed-specific `make_tma_2d_desc` variant that takes the full tile extent for the box. + +UPDATE (full-box fix tested in mega — insufficient): applied cute's convention to the ACTS +(box = full `block_m` via host, dropped per-CTA `m_idx` offset; confirmed in DG_JIT_DEBUG: acts +smem box outer 32 -> 64). Still hangs (`MMAFULL`=0). Harness also dumped cute's smem: `SmemLayoutA` +per-CTA M = 64 = MmaTile_M/2 (half), box = 128 (full) — confirms the "full-box descriptor + +half-smem, cta_group::2 splits" model. But the mega is **swap-AB** (weights=UMMA-A spanning +UMMA_M=256 across the 2 CTAs; acts=UMMA-B/UMMA_N), so the 2-SM-split operand mapping differs from +cute's standard GEMM (where A=M is the split operand). Fixing only the acts box doesn't line up, +and A+B share one `full_barrier`. Cracking cta_group::2 here ≈ adopting cute's full non-swap 2-SM +mainloop structure (effectively a GEMM rewrite), not a descriptor tweak. + +DECISION POINT (both are real work; pick one): + (A) Keep 2-CTA multicast: rebuild the mega GEMM around cute's exact 2-SM model (full-tile boxes + for the correct split operand, base coords, matching smem). Largest; verify each descriptor + byte-matches cute via `/tmp/tma_dump.cu`. + (B) Solve now via the PROVEN path: use the validated standalone's load mechanism — per-CTA + `SM90_TMA_LOAD_2D` (each CTA loads its own split half) + a cross-CTA "loads-done" barrier + (repurpose the idle GEMM warp, like the standalone's SF-transpose warp arrives `with_sf`). + This is exactly what `sm100_fp4_mqa_logits` and the validated standalone (`diff=0`) do for + packed FP4. Reworks the mega's `full_barrier` arrive scheme (each CTA expects its own tx). + +(superseded options below) +REMAINING OPTIONS to get 2-CTA multicast + packed FP4: + (A) Match cute exactly: dump the `CUtensorMap` cute produces for a 2-SM packed-FP4 atom + (`make_tma_atom_A_sm100` with a 2-SM `ClusterLayout`) and diff its fields vs + `make_tma_2d_desc`; fix the box/coords in `make_tma_2d_desc` (or call cute's builder in the host). + This is the path that keeps true 2-CTA multicast. + (B) Fall back to the repo's own packed-FP4 precedent — per-CTA `SM90_TMA_LOAD_2D` (what + `sm100_fp4_mqa_logits` and the validated standalone use) + a cross-CTA "loads-done" barrier. + +(Superseded) earlier candidate — per-CTA SM90 loads: +- Replace the GEMM A/B loads with per-CTA `cute::SM90_TMA_LOAD_2D` (`num_tma_multicast=1`): + each CTA loads its own A half (split `m_idx`) and its own B into its own smem, and + `arrive_and_expect_tx`-es ITS OWN bytes on ITS OWN `full_barrier` (init per-CTA, not the + `2*2` leader-combined scheme). +- Add the cross-CTA "loads done" sync the 2-CTA TMA used to provide implicitly. The standalone + gets this for free via its SF-transpose warp + `with_sf` barrier (both CTAs arrive at the + leader after their own `full` completes, and the leader MMA waits on `with_sf`). The mega + removed the transpose warp (SF pre-transposed), so add an equivalent: after each CTA's + `full_barrier` is satisfied, have it `arrive` at a shared (leader) barrier that the MMA warp + waits on, so the leader doesn't issue the 2-CTA UMMA until BOTH CTAs' smem is loaded. + (SFA/SFB still load fine via the existing `tma::copy<...,2>`; only the packed A/B data loads + must move off the 2-CTA TMA.) + +After the GEMM loads complete, the next blocker is still the **packed-FP4 L1 epilogue** layout +(currently a placeholder, numerically wrong) — see below. + +REMAINING (focused): implement the **correct packed-FP4 L1 epilogue** — derive the TMEM→register +mapping (`SM100_TMEM_LOAD_16dp256b1x`) and a packed-E2M1 smem store (no 4-bit `stmatrix` +exists; needs a manual transpose/pack) whose layout matches the L2 `mxf4` TMA read +(`swizzle = BLOCK_K/2`), and verify the L2 arrival mask reaches `0xFF`. Then validate against +the torch reference and add a `mxfp4xmxfp4` case to `tests/test_mega_moe.py`. Debug loop: +`DG_JIT_WITH_LINEINFO=1 cuda-gdb --batch -ex run -ex 'info cuda threads' --args python -u +tests/test_mxfp4_mega_moe.py` (interrupt with `timeout -s INT`); inspect the workspace +`l2_arrival_mask` value for a stuck pool block. + +### (superseded) original step 5 + +### (Original) Port to mega notes +- create `sm100_mxfp4_mxfp4_mega_moe.cuh` (fork of the fp8 mega impl) + reusing the proven packed paths — dispatch pulls `hidden/2` bytes, the L1-output + epilogue emits packed E2M1 + UE8M0 SF (SF math is unchanged from today, only the + cast/STSM packing differs), and both GEMMs use the new `mxf4` 2-CTA MMA. Then add + `MmaKind::MXFP4` handling in `heuristics/mega_moe.hpp` + `heuristics/sm100.hpp` + (FP4 0.5-byte arithmetic, per-stage smem, SF aligned sizes), the host + `sm100_mxfp4_mxfp4_mega_moe.hpp`, the API in `csrc/apis/mega.hpp` + (`mma_type="mxfp4xmxfp4"` + buffer sizing), Python + weight transform, and a + `tests/test_mega_moe.py` case. + +## Useful references in-repo + +- `deep_gemm/include/deep_gemm/impls/sm100_mqa_logits.cuh` — the only packed `mxf4` + (1-CTA) usage; authoritative for descriptor `/2` addressing and `sf_id` mapping. +- `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` — the 2-CTA structure, + UTCCP-2cta SF, TMEM pipeline, and epilogue this draft mirrors. +- `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh` + + `csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp` — host/JIT/tensormap wiring to + copy for the standalone host runtime. +- `deep_gemm/include/deep_gemm/mma/sm100.cuh` — `make_smem_desc`, + `to_umma_layout_type`, `make_runtime_instr_desc_with_sf_id`. diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 991eabca11..1726007bbc 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -7,6 +7,7 @@ #include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp" #include "../jit_kernels/impls/sm90_bf16_gemm.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp" +#include "../jit_kernels/impls/sm100_mxfp4_gemm.hpp" #include "../jit_kernels/impls/sm100_bf16_gemm.hpp" #endif @@ -163,6 +164,27 @@ static void fp8_fp4_gemm_tt(const std::pair& a, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast); } +static void mxfp4_gemm_nt(const std::pair& a, + const std::pair& b, + const torch::Tensor& d) { + // Standalone packed MXFP4 x MXFP4 de-risk GEMM: `[M, K] @ [N, K].T -> [M, N]`. + // A/B packed E2M1 stored as int8 `[*, K/2]`; SFs are int32-packed UE8M0 (gran-32). + const auto& [a_val, sfa] = a; + const auto& [b_val, sfb] = b; + DG_HOST_ASSERT(a_val.dim() == 2 and b_val.dim() == 2 and d.dim() == 2); + DG_HOST_ASSERT(a_val.is_contiguous() and b_val.is_contiguous()); + + const int m = static_cast(a_val.size(0)); + const int k = static_cast(a_val.size(1)) * 2; // packed: 2 elements per byte + const int n = static_cast(b_val.size(0)); + DG_HOST_ASSERT(static_cast(b_val.size(1)) * 2 == k); + DG_HOST_ASSERT(static_cast(d.size(0)) == m and static_cast(d.size(1)) == n); + + const auto arch_major = device_runtime->get_arch_major(); + DG_HOST_ASSERT(arch_major == 10 and "MXFP4 GEMM requires SM100"); + sm100_mxfp4_gemm(a_val, sfa, b_val, sfb, d, m, n, k); +} + static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair& a, const std::pair& b, const torch::Tensor& d, @@ -670,6 +692,8 @@ static void register_apis(pybind11::module_& m) { py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, py::arg("compiled_dims") = "mn", py::arg("disable_ue8m0_cast") = false); + m.def("mxfp4_gemm_nt", &mxfp4_gemm_nt, + py::arg("a"), py::arg("b"), py::arg("d")); m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), py::arg("recipe") = std::nullopt, diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 435db6e62f..443e542053 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -12,6 +12,7 @@ #include "../jit/device_runtime.hpp" #include "../jit_kernels/impls/sm100_bf16_mega_moe.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp" +#include "../jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp" namespace deep_gemm::mega { @@ -48,15 +49,18 @@ get_symm_buffer_size_for_mega_moe( const auto mma_kind = parse_mma_kind(mma_type); const auto num_mma_elem_bytes = get_num_mma_elem_bytes(mma_kind); const auto with_sf = is_mma_with_sf(mma_kind); + // NOTES: packed FP4 is sub-byte (0.5 B/elem), so derive token byte sizes from element BITS + // (`get_element_bits/8`): BF16->2, FP8->1, packed FP4->0.5. `hidden` is a multiple of 128. + const auto token_bytes = [&](const int& num_elems) { return num_elems * get_element_bits(mma_kind) / 8; }; // Workspace const auto workspace = layout::Workspace( nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_ring_tokens); // Layouts - const auto input_token_layout = layout::Data(hidden * num_mma_elem_bytes); + const auto input_token_layout = layout::Data(token_bytes(hidden)); const auto bf16_token_layout = layout::Data(hidden * 2); - const auto intermediate_token_layout = layout::Data(intermediate_hidden * num_mma_elem_bytes); + const auto intermediate_token_layout = layout::Data(token_bytes(intermediate_hidden)); const auto input_sf_layout = layout::Data(with_sf ? hidden / 32 : 0); const auto intermediate_sf_layout = layout::Data(with_sf ? intermediate_hidden / 32 : 0); const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false); @@ -116,13 +120,20 @@ get_symm_buffer_size_for_mega_moe( DG_HOST_ASSERT(num_sf_ring_tokens % 4 == 0); } + // Token view shape/dtype: packed MXFP4 stores 2 elems/byte (int8, half the inner dim), + // FP8 uses 1 byte/elem (e4m3, full inner dim), BF16 uses bf16. + const bool is_mxfp4 = (mma_kind == MmaKind::MXFP4); + const auto token_dtype = with_sf ? (is_mxfp4 ? torch::kInt8 : torch::kFloat8_e4m3fn) : torch::kBFloat16; + const int x_token_inner = is_mxfp4 ? hidden / 2 : hidden; + const int l2_token_inner = is_mxfp4 ? intermediate_hidden / 2 : intermediate_hidden; + // Slice function: creates `(x, x_sf, topk_weights, topk_idx, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf)` tensor views from the raw buffer // NOTES: `x_sf` is K-major, while `l1_acts_sf` and `l2_acts_sf` are M-major auto slice_input_buffers = [=](const torch::Tensor& buffer) { auto x = torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_token_buffer.base)), - {num_max_tokens_per_rank, hidden}, - torch::TensorOptions().dtype(with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())); + {num_max_tokens_per_rank, x_token_inner}, + torch::TensorOptions().dtype(token_dtype).device(buffer.device())); auto x_sf = with_sf ? torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_sf_buffer.base)), {num_max_tokens_per_rank, hidden / 128}, @@ -137,8 +148,8 @@ get_symm_buffer_size_for_mega_moe( torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); auto l1_acts = torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_token_buffer.base)), - {num_ring_tokens, hidden}, - torch::TensorOptions().dtype(with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())); + {num_ring_tokens, x_token_inner}, + torch::TensorOptions().dtype(token_dtype).device(buffer.device())); auto l1_acts_sf = with_sf ? torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_sf_buffer.base)), {num_sf_ring_tokens, hidden / 128}, @@ -146,8 +157,8 @@ get_symm_buffer_size_for_mega_moe( torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); auto l2_acts = torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l2_token_buffer.base)), - {num_ring_tokens, intermediate_hidden}, - torch::TensorOptions().dtype(with_sf ? torch::kFloat8_e4m3fn : torch::kBFloat16).device(buffer.device())); + {num_ring_tokens, l2_token_inner}, + torch::TensorOptions().dtype(token_dtype).device(buffer.device())); auto l2_acts_sf = with_sf ? torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l2_sf_buffer.base)), {num_sf_ring_tokens, intermediate_hidden / 128}, @@ -253,6 +264,94 @@ static void fp8_fp4_mega_moe( sym_buffer.zero_(); } +static void mxfp4_mxfp4_mega_moe( + const torch::Tensor& y, + const std::tuple& l1_weights_tuple, + const std::tuple& l2_weights_tuple, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const std::vector& sym_buffer_ptrs, const int& rank_idx, + const int& num_max_tokens_per_rank, + const int& num_experts, const int& num_topk, + const std::tuple& recipe, + const std::string& activation, + const std::optional& activation_clamp_opt, + const bool& fast_math, + const int& num_ring_tokens +) { + // Packed MXFP4 x MXFP4 mega MoE (both activations and weights are packed E2M1 + UE8M0 SF). + const auto [l1_weights, l1_weights_sf] = l1_weights_tuple; + const auto [l2_weights, l2_weights_sf] = l2_weights_tuple; + + const auto num_tokens = static_cast(y.size(0)); + const auto [rm, rn, rk] = recipe; + DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 32); + DG_HOST_ASSERT(activation == "swiglu"); + + const auto activation_clamp = + activation_clamp_opt.value_or(std::numeric_limits::infinity()); + DG_HOST_ASSERT(activation_clamp >= 0); + + // Tensor checks (weights are packed FP4, same as the fp8xfp4 path) + DG_HOST_ASSERT(get_major_type_ab(l1_weights) == cute::UMMA::Major::K); + DG_HOST_ASSERT(get_major_type_ab(l2_weights) == cute::UMMA::Major::K); + const auto arch_major = device_runtime->get_arch_major(); + const auto [num_experts_per_rank, intermediate_hidden_2, hidden] = + check_grouped_ab_fp8_fp4(l1_weights, cute::UMMA::Major::K, arch_major); + const auto [num_experts_per_rank_, hidden_, intermediate_hidden] = + check_grouped_ab_fp8_fp4(l2_weights, cute::UMMA::Major::K, arch_major); + DG_HOST_ASSERT(num_tokens <= num_max_tokens_per_rank); + DG_HOST_ASSERT(num_experts_per_rank == num_experts_per_rank_); + DG_HOST_ASSERT(hidden == hidden_); + DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden); + DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous()); + + // Check weight SF layout (UE8M0, MN-major, TMA-aligned, gran-32) + constexpr int kGranMN = 1, kGranK = 32; + check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK, + num_experts_per_rank, true, false, torch::kInt); + check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK, + num_experts_per_rank, true, false, torch::kInt); + + if (cumulative_local_expert_recv_stats.has_value()) { + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt); + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->numel() == num_experts_per_rank); + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->is_contiguous()); + } + + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + const auto num_experts_ = num_experts_per_rank * num_ranks; + const auto [num_required_bytes, slice] = get_symm_buffer_size_for_mega_moe( + num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + "mxfp4xmxfp4", activation, num_ring_tokens); + DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); + DG_HOST_ASSERT(num_experts == num_experts_); + + const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + + if (arch_major == 10) { + sm100_mxfp4_mxfp4_mega_moe(y, + l1_acts, l1_acts_sf, + l2_acts, l2_acts_sf, + l1_weights, l2_weights, + l1_weights_sf, l2_weights_sf, + cumulative_local_expert_recv_stats, + sym_buffer_ptrs, + rank_idx, num_max_tokens_per_rank, + num_experts_per_rank, + num_tokens, num_topk, + hidden, intermediate_hidden, + activation_clamp, fast_math); + } else { + DG_HOST_UNREACHABLE("Unsupported architecture"); + } + + if (get_env("DG_COMM_KERNEL_DEBUG")) + sym_buffer.zero_(); +} + static void bf16_mega_moe( const torch::Tensor& y, const torch::Tensor& l1_weights, @@ -339,6 +438,7 @@ static void register_apis(pybind11::module_& m) { m.def("get_ring_limit_for_mega_moe", &get_ring_limit_for_mega_moe); m.def("get_symm_buffer_size_for_mega_moe", &get_symm_buffer_size_for_mega_moe); m.def("fp8_fp4_mega_moe", &fp8_fp4_mega_moe); + m.def("mxfp4_mxfp4_mega_moe", &mxfp4_mxfp4_mega_moe); m.def("bf16_mega_moe", &bf16_mega_moe); #endif } diff --git a/csrc/jit_kernels/heuristics/mega_moe.hpp b/csrc/jit_kernels/heuristics/mega_moe.hpp index ac23eb4442..8ca46d4cdf 100644 --- a/csrc/jit_kernels/heuristics/mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/mega_moe.hpp @@ -65,16 +65,20 @@ struct MegaMoEConfig { static MmaKind parse_mma_kind(const std::string& mma_type_str) { if (mma_type_str == "bf16xbf16") return MmaKind::BF16; + if (mma_type_str == "mxfp4xmxfp4") + return MmaKind::MXFP4; DG_HOST_ASSERT(mma_type_str == "fp8xfp4"); return MmaKind::MXFP8FP4; } static int get_num_mma_elem_bytes(const MmaKind& mma_kind) { + // NOTES: packed MXFP4 is sub-byte; callers that need its true byte size use + // `get_element_bits(mma_kind) / 8`. Here MXFP4 returns 1 so `block_k` stays in elements. return mma_kind == MmaKind::BF16 ? 2 : 1; } static bool is_mma_with_sf(const MmaKind& mma_kind) { - return mma_kind == MmaKind::MXFP8FP4; + return mma_kind == MmaKind::MXFP8FP4 or mma_kind == MmaKind::MXFP4; } static int get_num_wave_pool_tokens( diff --git a/csrc/jit_kernels/heuristics/sm100.hpp b/csrc/jit_kernels/heuristics/sm100.hpp index c8e9e2e07f..122839cfc9 100644 --- a/csrc/jit_kernels/heuristics/sm100.hpp +++ b/csrc/jit_kernels/heuristics/sm100.hpp @@ -19,7 +19,8 @@ struct SM100ArchSpec { constexpr int num_utccp_aligned_elems = 128; switch (mma_kind) { case MmaKind::BF16: return {0, 0}; - case MmaKind::MXFP8FP4: return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)}; + case MmaKind::MXFP8FP4: + case MmaKind::MXFP4: return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)}; default: DG_HOST_UNREACHABLE("Unknown dtype"); } } diff --git a/csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp b/csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp new file mode 100644 index 0000000000..709dbf4f72 --- /dev/null +++ b/csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp @@ -0,0 +1,148 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "../../utils/math.hpp" +#include "../heuristics/sm100.hpp" + +#include "runtime_utils.hpp" + +namespace deep_gemm { + +// Standalone packed MXFP4 x MXFP4 (E2M1 data, UE8M0 SF gran-32, 2-CTA) GEMM. +// De-risk vehicle: explicit (compile-time) template params, no `get_best_config`. +class SM100MXFP4GemmRuntime final: public LaunchRuntime { +public: + struct Args { + int m, n, k; + int block_m, block_n, block_k; + int num_stages; + int num_non_epilogue_threads, num_epilogue_threads; + int num_sms; + + CUtensorMap tensor_map_a; + CUtensorMap tensor_map_sfa; + CUtensorMap tensor_map_b; + CUtensorMap tensor_map_sfb; + CUtensorMap tensor_map_cd; + + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_mxfp4_gemm_impl< + {}, {}, {}, + {}, {}, {}, + {}, + {}, {}, + {} + >); +}}; +)", + args.m, args.n, args.k, + args.block_m, args.block_n, args.block_k, + args.num_stages, + args.num_non_epilogue_threads, args.num_epilogue_threads, + args.num_sms); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + static_cast(args.m), static_cast(args.n), static_cast(args.k), + args.tensor_map_a, args.tensor_map_sfa, + args.tensor_map_b, args.tensor_map_sfb, + args.tensor_map_cd)); + } +}; + +static void sm100_mxfp4_gemm(const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const torch::Tensor& d, + const int& m, const int& n, const int& k) { + // Fixed de-risk configuration + constexpr int block_m = 128, block_n = 128, block_k = 128; + constexpr int num_stages = 4; + constexpr int num_non_epilogue_threads = 128, num_epilogue_threads = 128; + constexpr int gran_k = 32; + + // Packed FP4 tensors are stored as int8 (`kPackedFP4`), 2 elements per byte + DG_HOST_ASSERT(a.scalar_type() == kPackedFP4 and b.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(sfa.scalar_type() == torch::kInt and sfb.scalar_type() == torch::kInt); + DG_HOST_ASSERT(m % block_m == 0 and n % block_n == 0 and k % block_k == 0); + // 2-CTA (cluster_n = 2) requires an even number of N blocks + DG_HOST_ASSERT((n / block_n) % 2 == 0 and "MXFP4 de-risk GEMM requires N divisible by 256 (2-CTA)"); + + // Even SM count for 2-CTA clusters + int num_sms = device_runtime->get_num_sms(); + num_sms -= num_sms % 2; + + const int load_block_m = block_m / 2; // acts split on M across the cluster (cluster_n = 2) + const int load_block_n = block_n; + constexpr int store_block_m = 16, store_block_n = 128, swizzle_cd = 128; + const int swizzle_ab = block_k / 2; // packed FP4 K-major swizzle in bytes (64) + + // Packed-FP4 A/B TMA descriptors (K-major, `fp4_unpacked_smem = false` -> 16U4_ALIGN8B) + // 2-CTA `cta_group::2`: per CUTLASS, each CTA loads its OWN per-CTA box (`load_block_m`) at + // its own M coord (offset by rank); the 2-SM atom routes all tx to the leader's barrier. + const auto tensor_map_a = make_tma_2d_desc(a, k, m, block_k, load_block_m, + static_cast(a.stride(0)), + swizzle_ab, 0, false, false); + const auto tensor_map_b = make_tma_2d_desc(b, k, n, block_k, load_block_n, + static_cast(b.stride(0)), + swizzle_ab, 0, false, false); + // UE8M0 SF descriptors (MN-major, no swizzle) + const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, block_m, gran_k, 1, 0); + const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, block_n, gran_k, 1, 0); + // BF16 output (N-major) + const auto tensor_map_cd = make_tma_cd_desc(d, m, n, store_block_m, store_block_n, + static_cast(d.stride(-2)), 1, swizzle_cd); + + // Shared memory size (must mirror the kernel's `SharedStorage` layout) + const int sf_block_m = align(block_m, 128), sf_block_n = align(block_n, 128); + constexpr int num_epilogue_stages = 2, num_tma_store_stages = 2; + const int smem_cd = store_block_m * store_block_n * static_cast(sizeof(int16_t)) * num_tma_store_stages; + const int smem_a = load_block_m * (block_k / 2); + const int smem_b = load_block_n * (block_k / 2); + const int smem_sfa = sf_block_m * static_cast(sizeof(int)); + const int smem_sfb = sf_block_n * static_cast(sizeof(int)); + const int smem_barriers = (num_stages * 3 + num_epilogue_stages * 2) * 8; + const int smem_size = smem_cd + + num_stages * (smem_a + smem_b) + + num_stages * (smem_sfa + smem_sfb) + + smem_barriers + 4; + DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity); + + const SM100MXFP4GemmRuntime::Args args = { + .m = m, .n = n, .k = k, + .block_m = block_m, .block_n = block_n, .block_k = block_k, + .num_stages = num_stages, + .num_non_epilogue_threads = num_non_epilogue_threads, + .num_epilogue_threads = num_epilogue_threads, + .num_sms = num_sms, + .tensor_map_a = tensor_map_a, + .tensor_map_sfa = tensor_map_sfa, + .tensor_map_b = tensor_map_b, + .tensor_map_sfb = tensor_map_sfb, + .tensor_map_cd = tensor_map_cd, + .launch_args = LaunchArgs(num_sms, + num_non_epilogue_threads + num_epilogue_threads, + smem_size, /*cluster_dim=*/2) + }; + const auto code = SM100MXFP4GemmRuntime::generate(args); + const auto runtime = compiler->build("sm100_mxfp4_gemm", code); + SM100MXFP4GemmRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp new file mode 100644 index 0000000000..17c064dbb2 --- /dev/null +++ b/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp @@ -0,0 +1,304 @@ +#pragma once + +#include +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "runtime_utils.hpp" + +#include +#include + +#include "../heuristics/mega_moe.hpp" + +namespace deep_gemm { + +// Packed MXFP4 x MXFP4 mega-MoE config: reuse the FP8 block heuristics, but with +// packed-FP4 (2 elems/byte) byte math: swizzle = block_k/2, A/B/output smem halved. +static MegaMoEConfig get_mxfp4_mega_moe_config( + const int& num_ranks, const int& num_experts, const int& num_experts_per_rank, + const int& num_max_tokens_per_rank, const int& num_tokens, const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const int& num_ring_tokens, const int& num_sf_ring_tokens) { + + const auto [cluster_size, block_m, store_block_m, block_k, num_epilogue_threads] = + get_block_config_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens, MmaKind::MXFP4); + const int block_n = 128; + const int load_block_m = block_m / 2; + const int load_block_n = block_n; + const auto [sf_block_m, sf_block_n] = + SM100ArchSpec::get_sf_uttcp_aligned_block_sizes(block_m, block_n, MmaKind::MXFP4); + + // Packed FP4 K-major swizzle == K extent in bytes (block_k / 2) + const int swizzle_acts_mode = block_k / 2; + const int swizzle_weights_mode = block_k / 2; + constexpr int gran_k = 32; + + const int num_sms = device_runtime->get_num_sms(); + const int num_experts_per_wave = get_num_experts_per_wave_for_mega_moe( + num_experts_per_rank, num_tokens, num_topk, intermediate_hidden, block_m, block_n, num_sms, + num_ring_tokens, num_max_tokens_per_rank, num_ranks); + + const int num_dispatch_threads = 128; + const int num_non_epilogue_threads = 128; + + // Pull: packed token bytes = hidden / 2 + constexpr int kPullThreshold = 4096; + int num_bytes_per_pull = hidden / 2; + while (num_bytes_per_pull > kPullThreshold) { + DG_HOST_ASSERT(num_bytes_per_pull % 2 == 0); + num_bytes_per_pull /= 2; + } + + // Shared-memory sizing (mirrors `get_pipeline_config_for_mega_moe`, with packed-FP4 byte math) + constexpr int kSmemAlignment = 1024; + constexpr int kNumEpilogueStages = 2, kNumTMAStoreStages = 2; + const int num_dispatch_warps = num_dispatch_threads / 32; + const int num_epilogue_warps = num_epilogue_threads / 32; + const int num_epilogue_warpgroups = num_epilogue_warps / 4; + + const int smem_expert_count_size = align(num_experts * static_cast(sizeof(uint32_t)), kSmemAlignment); + const int smem_send_buffers_size = align( + static_cast(layout::Buffer(layout::Data(num_bytes_per_pull), num_dispatch_warps, 1).get_num_bytes()), + kSmemAlignment); + const int smem_dispatch_size = smem_expert_count_size + smem_send_buffers_size; + + // L1 output is packed FP4 (block_n/2 elems -> block_n/4 bytes); L2 is BF16 + const int smem_cd_l1 = num_epilogue_warpgroups * store_block_m * (block_n / 4) * kNumTMAStoreStages; + const int smem_cd_l2 = num_epilogue_warpgroups * store_block_m * block_n * static_cast(sizeof(nv_bfloat16)); + const int smem_cd = align(std::max(smem_cd_l1, smem_cd_l2), kSmemAlignment); + + const int smem_barriers = (num_dispatch_warps + kNumEpilogueStages * 2 + num_epilogue_warps * 2) * 8; + const int smem_amax_reduction = store_block_m * num_epilogue_warps * static_cast(sizeof(float)); + const int smem_tmem_ptr = 4; + + const int smem_sfa_per_stage = sf_block_m * (block_k / gran_k); + const int smem_sfb_per_stage = sf_block_n * (block_k / gran_k); + // Packed FP4: A/B tiles are byte-addressed at block_k / 2 bytes per row + const int smem_a_size_per_stage = load_block_m * (block_k / 2); + const int smem_b_size_per_stage = block_n * (block_k / 2); + // 3 per-stage barriers: full + empty + with-sf (cross-CTA loads-done) + const int smem_size_per_stage = smem_a_size_per_stage + smem_b_size_per_stage + smem_sfa_per_stage + smem_sfb_per_stage + 3 * 8; + + const int smem_fixed = smem_dispatch_size + smem_cd + smem_amax_reduction + smem_barriers + smem_tmem_ptr; + const int num_stages = (SM100ArchSpec::smem_capacity - smem_fixed) / smem_size_per_stage; + DG_HOST_ASSERT(num_stages >= 2); + const int smem_size = smem_fixed + num_stages * smem_size_per_stage; + + return MegaMoEConfig { + block_m, block_n, block_k, + load_block_m, load_block_n, store_block_m, + sf_block_m, sf_block_n, + num_ring_tokens, num_sf_ring_tokens, + swizzle_acts_mode, swizzle_weights_mode, + num_experts_per_wave, + num_stages, smem_size, + num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads, + num_bytes_per_pull + }; +} + +class SM100MXFP4MegaMoERuntime final : public LaunchRuntime { +public: + struct Args { + int num_max_tokens_per_rank; + int hidden, intermediate_hidden; + int num_experts, num_topk; + int num_ranks; + float activation_clamp; + bool fast_math; + MegaMoEConfig config; + + void* y; + int* cumulative_local_expert_recv_stats; + int num_tokens; + layout::SymBuffer<> sym_buffer_ptrs; + + CUtensorMap tensor_map_l1_acts; + CUtensorMap tensor_map_l1_acts_sf; + CUtensorMap tensor_map_l1_weights; + CUtensorMap tensor_map_l1_weights_sf; + CUtensorMap tensor_map_l1_output; + CUtensorMap tensor_map_l2_acts; + CUtensorMap tensor_map_l2_acts_sf; + CUtensorMap tensor_map_l2_weights; + CUtensorMap tensor_map_l2_weights_sf; + + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_mxfp4_mxfp4_mega_moe_impl< + {}, + {}, {}, + {}, {}, + {}, + {}, {}, {}, + {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, {}, {}, + {}, {}, + {}, + {} + >); +}}; +)", args.num_max_tokens_per_rank, + args.hidden, args.intermediate_hidden, + args.num_experts, args.num_topk, + args.config.num_experts_per_wave, + args.config.block_m, args.config.block_n, args.config.block_k, + args.config.store_block_m, + args.config.sf_block_m, args.config.sf_block_n, + args.config.num_ring_tokens, + args.config.num_sf_ring_tokens, + args.config.num_stages, + args.config.num_bytes_per_pull, + args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads, + args.launch_args.grid_dim.first, args.num_ranks, + to_string(args.activation_clamp), + args.fast_math ? "true" : "false"); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.y, + args.cumulative_local_expert_recv_stats, + args.num_tokens, + args.sym_buffer_ptrs, + args.tensor_map_l1_acts, + args.tensor_map_l1_acts_sf, + args.tensor_map_l1_weights, + args.tensor_map_l1_weights_sf, + args.tensor_map_l1_output, + args.tensor_map_l2_acts, + args.tensor_map_l2_acts_sf, + args.tensor_map_l2_weights, + args.tensor_map_l2_weights_sf + )); + } +}; + +static void sm100_mxfp4_mxfp4_mega_moe( + const torch::Tensor& y, + const torch::Tensor& l1_acts, const torch::Tensor& l1_acts_sf, + const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf, + const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, + const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf, + const std::optional cumulative_local_expert_recv_stats, + const std::vector& sym_buffer_ptrs, + const int& rank_idx, const int& num_max_tokens_per_rank, + const int& num_experts_per_rank, + const int& num_tokens, const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const float& activation_clamp, + const bool& fast_math +) { + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + const auto num_experts = num_experts_per_rank * num_ranks; + const auto num_ring_tokens = static_cast(l1_acts.size(0)); + const auto num_sf_ring_tokens = static_cast(l1_acts_sf.size(0)); + + const auto config = get_mxfp4_mega_moe_config( + num_ranks, num_experts, num_experts_per_rank, + num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, + num_ring_tokens, num_sf_ring_tokens); + + constexpr int kGranK = 32; + const int sf_smem_outer_dim = config.block_k / (kGranK * 4); + + // Packed FP4 token/weight TMA descriptors (fp4_unpacked_smem = false -> 16U4_ALIGN8B) + const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts, + hidden, config.num_ring_tokens, + config.block_k, config.load_block_m, + static_cast(l1_acts.stride(-2)), + config.swizzle_acts_mode, 0, false, false); + const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf, + config.num_sf_ring_tokens, hidden, + config.sf_block_m, kGranK, + 1, 0, 0, false, sf_smem_outer_dim); + const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights, + hidden, num_experts_per_rank * intermediate_hidden * 2, + config.block_k, config.load_block_n, + static_cast(l1_weights.stride(-2)), + config.swizzle_weights_mode, 0, false, false); + const auto tensor_map_l1_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_weights_sf, + intermediate_hidden * 2, hidden, + config.block_n, kGranK, + num_experts_per_rank, 0, 0, false, sf_smem_outer_dim); + // L1 output / L2 activations are the same packed-FP4 tensor; post-SwiGLU N width = block_n/2 elems. + // NOTES: the L1 epilogue writes a PLAIN (non-swizzled) packed-FP4 [store_block_m, block_n/2] + // tile directly at each value's [token][inter] position, so the store uses swizzle 0. + const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts, + intermediate_hidden, config.num_ring_tokens, + config.block_n / 2, config.store_block_m, + static_cast(l2_acts.stride(-2)), + 0, 0, false, false); + const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts, + intermediate_hidden, config.num_ring_tokens, + config.block_k, config.load_block_m, + static_cast(l2_acts.stride(-2)), + config.swizzle_acts_mode, 0, false, false); + const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf, + config.num_sf_ring_tokens, intermediate_hidden, + config.sf_block_m, kGranK, + 1, 0, 0, false, sf_smem_outer_dim); + const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights, + intermediate_hidden, num_experts_per_rank * hidden, + config.block_k, config.load_block_n, + static_cast(l2_weights.stride(-2)), + config.swizzle_weights_mode, 0, false, false); + const auto tensor_map_l2_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_weights_sf, + hidden, intermediate_hidden, + config.block_n, kGranK, + num_experts_per_rank, 0, 0, false, sf_smem_outer_dim); + + int* cumulative_local_expert_recv_stats_ptr = nullptr; + if (cumulative_local_expert_recv_stats.has_value()) + cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr(); + + const auto num_sms = device_runtime->get_num_sms(); + const SM100MXFP4MegaMoERuntime::Args args = { + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .hidden = hidden, .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, .num_topk = num_topk, + .num_ranks = num_ranks, + .activation_clamp = activation_clamp, + .fast_math = fast_math, + .config = config, + .y = y.data_ptr(), + .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, + .num_tokens = num_tokens, + .sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), + .tensor_map_l1_acts = tensor_map_l1_acts, + .tensor_map_l1_acts_sf = tensor_map_l1_acts_sf, + .tensor_map_l1_weights = tensor_map_l1_weights, + .tensor_map_l1_weights_sf = tensor_map_l1_weights_sf, + .tensor_map_l1_output = tensor_map_l1_output, + .tensor_map_l2_acts = tensor_map_l2_acts, + .tensor_map_l2_acts_sf = tensor_map_l2_acts_sf, + .tensor_map_l2_weights = tensor_map_l2_weights, + .tensor_map_l2_weights_sf = tensor_map_l2_weights_sf, + .launch_args = LaunchArgs(num_sms, + config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, + config.smem_size, 2) + }; + + const auto code = SM100MXFP4MegaMoERuntime::generate(args); + const auto runtime = compiler->build("sm100_mxfp4_mxfp4_mega_moe", code); + SM100MXFP4MegaMoERuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index 4e9c924e66..e06aca4797 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -34,6 +34,8 @@ try: # DeepGEMM Kernels from ._C import ( + # Packed MXFP4 x MXFP4 GEMM (standalone de-risk) + mxfp4_gemm_nt, # FP8 FP4 GEMMs fp8_fp4_gemm_nt, fp8_fp4_gemm_nn, fp8_fp4_gemm_tn, fp8_fp4_gemm_tt, @@ -87,6 +89,7 @@ transform_weights_for_mega_moe, fp8_fp4_mega_moe, bf16_mega_moe, + mxfp4_mxfp4_mega_moe, ) # Some utils diff --git a/deep_gemm/include/deep_gemm/common/types.cuh b/deep_gemm/include/deep_gemm/common/types.cuh index ee4e14c031..216e9fdb07 100644 --- a/deep_gemm/include/deep_gemm/common/types.cuh +++ b/deep_gemm/include/deep_gemm/common/types.cuh @@ -7,12 +7,27 @@ namespace deep_gemm { enum class MmaKind { BF16 = 0, MXFP8FP4 = 1, + // Pure packed MXFP4 x MXFP4 (E2M1 data, UE8M0 SF, 2 elements per byte) + MXFP4 = 2, }; +// NOTES: element size in *bits*, since packed FP4 is sub-byte (4 bits) +constexpr CUTLASS_HOST_DEVICE int get_element_bits(const MmaKind& mma_kind) { + switch (mma_kind) { + case MmaKind::BF16: return 16; + case MmaKind::MXFP8FP4: return 8; + case MmaKind::MXFP4: return 4; + default: return 0; + } +} + constexpr CUTLASS_HOST_DEVICE int get_element_size(const MmaKind& mma_kind) { switch (mma_kind) { case MmaKind::BF16: return 2; case MmaKind::MXFP8FP4: return 1; + // NOTES: packed FP4 is 0.5 byte/elem; callers must use byte math that + // divides element counts by 2 (see `get_element_bits`) + case MmaKind::MXFP4: return 1; default: return 0; } } diff --git a/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_gemm.cuh b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_gemm.cuh new file mode 100644 index 0000000000..f069aa4ca9 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_gemm.cuh @@ -0,0 +1,402 @@ +#pragma once +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Standalone packed MXFP4 x MXFP4 (E2M1 data, UE8M0 SF, gran-K 32), 2-CTA UMMA. +// +// NOTES: de-risking vehicle for the packed `mxf4` 2-CTA path before porting it +// into the mega-MoE kernel. This deliberately mirrors the proven +// `sm100_fp8_fp4_gemm_1d1d` kernel structure (scheduler, warp roles, barriers, +// SF warp-transpose + with-SF barrier, and the swap-AB BF16 epilogue), changing +// ONLY what packed `mxf4` requires: +// - both operands packed E2M1 (2 elems/byte), byte-addressed smem (`/2`). +// - `UMMA_K = 64` and the `tcgen05.mma.kind::mxf4` 2-CTA wrapper. +// - UE8M0 SF fixed at gran-K 32 (one gran-32 SF per 32 K, `sf_id = k * 2`). +// - K-major packed swizzle == `BLOCK_K / 2` bytes; UMMA descriptors via +// `make_smem_desc` with `BLOCK_K/2` byte strides (mirrors `sm100_fp4_mqa_logits`). +// Spots needing on-SM100 confirmation are marked `// VALIDATE`. + +namespace deep_gemm { + +template +CUTLASS_GLOBAL void __launch_bounds__(kNumNonEpilogueThreads + kNumEpilogueThreads, 1) +sm100_mxfp4_gemm_impl(uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, + const __grid_constant__ cute::TmaDescriptor tensor_map_a, // acts, [M, K] K-major (packed E2M1) + const __grid_constant__ cute::TmaDescriptor tensor_map_sfa, // acts SF, UE8M0 (int32-packed, gran-32) + const __grid_constant__ cute::TmaDescriptor tensor_map_b, // weights, [N, K] K-major (packed E2M1) + const __grid_constant__ cute::TmaDescriptor tensor_map_sfb, // weights SF, UE8M0 + const __grid_constant__ cute::TmaDescriptor tensor_map_cd) { // out, [M, N] BF16 +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using Allocator = cute::TMEM::Allocator2Sm; + + // Packed FP4 data type (2 elements per byte) and UE8M0 scale factor + using ab_dtype_t = cutlass::float_e2m1_t; + using cd_dtype_t = cutlass::bfloat16_t; + + // Fixed configuration for this de-risk kernel + constexpr cute::UMMA::Major kMajorA = cute::UMMA::Major::K; + constexpr cute::UMMA::Major kMajorB = cute::UMMA::Major::K; + constexpr uint32_t kGemmType = static_cast(GemmType::Normal); + constexpr uint32_t kNumGroups = 1; + constexpr uint32_t kNumMulticast = 2; + + // MMA configs (swap-AB: weights take UMMA "A" / UMMA_M, acts take UMMA "B" / UMMA_N) + constexpr uint32_t LAYOUT_AD_M = 128; + constexpr uint32_t UMMA_M = LAYOUT_AD_M * kNumMulticast; + constexpr uint32_t UMMA_N = BLOCK_M; + constexpr uint32_t UMMA_K = 64; // packed FP4 contracts K=64 per instruction + constexpr uint32_t LOAD_BLOCK_M = BLOCK_M / kNumMulticast; // acts split on M across the cluster (cluster_n) + constexpr uint32_t LOAD_BLOCK_N = BLOCK_N; // weights full per CTA + DG_STATIC_ASSERT(BLOCK_K == 128, "Invalid block K"); + DG_STATIC_ASSERT(BLOCK_N == LAYOUT_AD_M, "Swap-AB requires BLOCK_N == 128"); + DG_STATIC_ASSERT(BLOCK_K % UMMA_K == 0, "Invalid K tiling"); + + // Packed-FP4 byte math: 2 elements per byte + constexpr uint32_t BLOCK_K_BYTES = BLOCK_K / 2; + constexpr uint32_t UMMA_K_BYTES = UMMA_K / 2; + // K-major packed swizzle == K extent in bytes (matches `sm100_fp4_mqa_logits`) + constexpr uint32_t kSwizzleABMode = BLOCK_K_BYTES; // 64 for BLOCK_K = 128 + constexpr uint32_t kSwizzleCDMode = 128; + + // SF configs (UE8M0, gran-K 32, UTCCP 128-aligned) + constexpr uint32_t kGranK = 32; + constexpr uint32_t kNumUTCCPAlignedElems = 128; + constexpr uint32_t SF_BLOCK_M = math::constexpr_align(BLOCK_M, kNumUTCCPAlignedElems); + constexpr uint32_t SF_BLOCK_N = math::constexpr_align(BLOCK_N, kNumUTCCPAlignedElems); + // One int32 packs 4 gran-32 SFs along K, i.e. covers 128 K + constexpr uint32_t kNumSFKPerLoad = BLOCK_K / (kGranK * 4); + DG_STATIC_ASSERT(kNumSFKPerLoad == 1, "BLOCK_K must be 128 for a single packed SF int per load"); + + // Epilogue configs (swap-AB) + constexpr uint32_t kNumEpilogueStages = 2; + constexpr uint32_t kNumTMAStoreStages = 2; + constexpr uint32_t STORE_BLOCK_M = 16; // swap-AB stores `umma_step_n = 16` rows at a time + constexpr uint32_t STORE_BLOCK_N = BLOCK_N; + constexpr uint32_t kNumUMMAStoreThreads = kNumEpilogueThreads; + DG_STATIC_ASSERT(kNumUMMAStoreThreads % 32 == 0, "Invalid store block M"); + + // Shared memory sizes (data in bytes via uint8, since packed FP4 is sub-byte) + constexpr uint32_t SMEM_CD_SIZE_PER_STAGE = STORE_BLOCK_M * STORE_BLOCK_N * sizeof(cd_dtype_t); + constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_SIZE_PER_STAGE * kNumTMAStoreStages; + constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K_BYTES; + constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K_BYTES; + constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * sizeof(uint32_t); + constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = SF_BLOCK_N * sizeof(uint32_t); + + // Tensor memory size and offsets + constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumEpilogueStages; + constexpr uint32_t kNumSFATmemCols = SF_BLOCK_M / 32; + constexpr uint32_t kNumSFBTmemCols = SF_BLOCK_N / 32; + constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); + constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols; + constexpr uint32_t kTmemStartColOfSFB = kNumAccumTmemCols + kNumSFATmemCols; + DG_STATIC_ASSERT(32 <= kNumTmemCols and kNumTmemCols <= 512, "Invalid tensor memory columns"); + + // Synchronize the cluster before 2-CTA TMEM allocation + comm::cluster_sync_with_relaxed_arrive(); + + // Utils + const bool is_leader_cta = cute::block_rank_in_cluster() == 0; + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto lane_idx = ptx::get_lane_idx(); + + // Prefetch TMA descriptors + if (warp_idx == 0) { + cute::prefetch_tma_descriptor(&tensor_map_a); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_sfa); + cute::prefetch_tma_descriptor(&tensor_map_sfb); + cute::prefetch_tma_descriptor(&tensor_map_cd); + } + + // Overwrite shapes if the compiler provides them + shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; + shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n; + shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k; + const auto shape_sf_k = math::ceil_div(shape_k, kGranK * 4); + + // Align to 1024 bytes for swizzle-128B (C/D) + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + + // D / A / B shared memory + auto smem_cd = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + i * SMEM_CD_SIZE_PER_STAGE); + }); + auto smem_a = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); + }); + auto smem_b = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); + }); + + // SFA / SFB shared memory + auto sf_start_ptr = reinterpret_cast(smem_b[kNumStages]); + auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); + }); + auto smem_sfb = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast(sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE + i * SMEM_SFB_SIZE_PER_STAGE); + }); + + // Barriers and tensor memory pointer + auto barrier_start_ptr = reinterpret_cast(smem_sfb[kNumStages]); + auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (i); }); + auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages + i); }); + auto with_sf_full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 2 + i); }); + auto tmem_full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 3 + i); }); + auto tmem_empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 3 + kNumEpilogueStages + i); }); + auto tmem_ptr_in_smem = reinterpret_cast(barrier_start_ptr + kNumStages * 3 + kNumEpilogueStages * 2); + + // Initialize barriers + if (warp_idx == 1 and cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(1); + // Both CTAs' transposers arrive on the leader's `with_sf` (cross-CTA loads-done sync) + with_sf_full_barriers[i]->init(kNumMulticast * 32); + } + #pragma unroll + for (uint32_t i = 0; i < kNumEpilogueStages; ++ i) { + tmem_full_barriers[i]->init(1); + tmem_empty_barriers[i]->init(kNumMulticast * kNumUMMAStoreThreads); + } + cutlass::arch::fence_barrier_init(); + } else if (warp_idx == 2) { + Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem); + } + comm::cluster_sync_with_relaxed_arrive(); + + // Wait for primary kernel completion (PDL) + cudaGridDependencySynchronize(); + + // Block scheduler (reuse the proven scheduler in Normal mode) + uint32_t m_block_idx, n_block_idx; + auto scheduler = sched::Scheduler( + shape_m, shape_n, shape_k, nullptr); + + // Pipeline and TMA phases + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + if (warp_idx == 0 and cute::elect_one_sync()) { + // TMA load warp — TRUE 2-CTA (cta_group::2), issued by BOTH CTAs. Matches CUTLASS for a + // 2x1 cluster (verified by instrumenting cute's SM100_TMA_2SM_LOAD on a real GEMM): + // - each CTA loads its OWN per-CTA box at its OWN coord (acts split on M via the rank + // offset; weights replicated by using the same n coord on both CTAs), + // - the 2-SM atom's peer bit routes ALL tx to the LEADER's `full` barrier automatically + // (cluster smem addressing sets bit 24 on the peer; the atom masks it to CTA0), + // - leader sets `expect_tx` for BOTH CTAs' contributions + loads SF (SM90, leader-resident). + const auto cache_hint = static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL); + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + uint32_t m_idx = scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx); + uint32_t n_idx = scheduler.template get_global_idx(shape_n, BLOCK_N, n_block_idx, m_block_idx); + uint32_t k_idx = k_block_idx * BLOCK_K; + auto* mbar = reinterpret_cast(full_barriers[stage_idx]); + + // acts: per-CTA M offset (each CTA loads its `LOAD_BLOCK_M`-row half) — 2-SM routes tx to leader + m_idx += cute::block_rank_in_cluster() * LOAD_BLOCK_M; + cute::SM100_TMA_2SM_LOAD_2D::copy(&tensor_map_a, mbar, cache_hint, smem_a[stage_idx], k_idx, m_idx); + // weights: same n coord on both CTAs (each loads the full BLOCK_N replica) — 2-SM routes tx to leader + cute::SM100_TMA_2SM_LOAD_2D::copy(&tensor_map_b, mbar, cache_hint, smem_b[stage_idx], k_idx, n_idx); + + // SF: per-CTA SM90 load into THIS CTA's smem (the 2-CTA UTCCP reads both CTAs' + // SF), signaling THIS CTA's own `full` barrier (NOT routed to the leader). So both + // CTAs' transposers can wait their own `full` and the cross-CTA `with_sf` sync holds. + uint32_t sfa_m_idx = m_block_idx * BLOCK_M; + uint32_t sfa_k_idx = scheduler.template get_global_idx(shape_sf_k, 1, k_block_idx); + tma::copy(&tensor_map_sfa, full_barriers[stage_idx], smem_sfa[stage_idx], sfa_m_idx, sfa_k_idx); + uint32_t sfb_n_idx = n_block_idx * BLOCK_N; + uint32_t sfb_k_idx = scheduler.template get_global_idx(shape_sf_k, 1, k_block_idx, m_block_idx); + tma::copy(&tensor_map_sfb, full_barriers[stage_idx], smem_sfb[stage_idx], sfb_n_idx, sfb_k_idx); + + // Expect: the leader collects BOTH CTAs' data (2-SM routed) + its own SF; the + // non-leader's `full` only sees its own SF (its data tx went to the leader). + const auto sf_bytes = BLOCK_M * sizeof(uint32_t) + BLOCK_N * sizeof(uint32_t); + if (is_leader_cta) + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE * kNumMulticast + SMEM_B_SIZE_PER_STAGE * kNumMulticast + sf_bytes); + else + full_barriers[stage_idx]->arrive_and_expect_tx(sf_bytes); + } + } + } else if (warp_idx == 1 and is_leader_cta) { + // MMA issue warp (leader CTA only) + // Swap-AB: weights -> UMMA "A" (UMMA_M), acts -> UMMA "B" (UMMA_N) + auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< + ab_dtype_t, ab_dtype_t, float, cutlass::float_ue8m0_t, + UMMA_M, UMMA_N, kMajorB, kMajorA>(); + + DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); + // VALIDATE: make_smem_desc layout/stride for packed E2M1 (mirrors mqa-logits) + constexpr auto kFP4Layout = mma::sm100::to_umma_layout_type(); + + DG_STATIC_ASSERT((UMMA_M == 256 and UMMA_N % 16 == 0 and 16 <= UMMA_N and UMMA_N <= 256), + "Invalid MMA instruction shape"); + + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + auto accum_stage_idx = scheduler.current_iter % kNumEpilogueStages; + auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; + tmem_empty_barriers[accum_stage_idx]->wait(accum_phase_idx ^ 1); + ptx::tcgen05_after_thread_sync(); + + auto empty_barrier_arrive = [&](const bool& do_tmem_full_arrive) { + auto umma_arrive = [](const uint64_t* barrier) { + constexpr uint16_t kCTAMask = (1 << kNumMulticast) - 1; + cutlass::arch::umma_arrive_multicast_2x1SM(barrier, kCTAMask); + }; + umma_arrive(reinterpret_cast(empty_barriers[stage_idx])); + if (do_tmem_full_arrive) + umma_arrive(reinterpret_cast(tmem_full_barriers[accum_stage_idx])); + __syncwarp(); + }; + + // Dynamic UMMA N based on effective M (swap-AB) + { + uint32_t umma_n = scheduler.get_aligned_effective_m_in_block(m_block_idx); + mma::sm100::update_instr_desc_with_umma_n(instr_desc, umma_n); + } + + const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + #pragma unroll 2 + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + with_sf_full_barriers[stage_idx]->wait(phase); + ptx::tcgen05_after_thread_sync(); + + if (cute::elect_one_sync()) { + // UTCCP copy SFA / SFB into TMEM (transposed in warp 3 already) + using cute_utccp_t = cute::SM100_UTCCP_4x32dp128bit_2cta; + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { + auto sf_desc = mma::sm100::make_sf_desc(smem_sfa[stage_idx] + i * kNumUTCCPAlignedElems); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + i * 4); + } + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { + auto sf_desc = mma::sm100::make_sf_desc(smem_sfb[stage_idx] + i * kNumUTCCPAlignedElems); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + i * 4); + } + + // Issue UMMA over UMMA_K (=64) sub-tiles + // VALIDATE: descriptor /2 byte addressing + sf_id = k*2 + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++ k) { + const uint32_t sf_id = k * 2; + const auto runtime_instr_desc = mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, sf_id, sf_id); + auto a_desc = mma::sm100::make_smem_desc( + kFP4Layout, reinterpret_cast(smem_b[stage_idx]) + k * UMMA_K_BYTES, + 8 * kSwizzleABMode, 0); + auto b_desc = mma::sm100::make_smem_desc( + kFP4Layout, reinterpret_cast(smem_a[stage_idx]) + k * UMMA_K_BYTES, + 8 * kSwizzleABMode, 0); + // Swap-AB: weights (b_desc) first, SFB id first + ptx::SM100_MMA_MXF4_2x1SM_SS::fma( + a_desc, b_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or k > 0, runtime_instr_desc, + kTmemStartColOfSFB, kTmemStartColOfSFA); + } + } + __syncwarp(); + empty_barrier_arrive(k_block_idx == num_total_k_blocks - 1); + } + } + + const auto iter_idx = scheduler.current_iter - 1; + if (iter_idx >= 0) { + const auto accum_phase_idx = (iter_idx / kNumEpilogueStages) & 1; + tmem_empty_barriers[iter_idx % kNumEpilogueStages]->wait(accum_phase_idx); + } + } else if (warp_idx == 2) { + // UTCCP transposer (BOTH CTAs): each waits its own `full` (its SF tx), transposes its SF, + // then arrives on the leader's `with_sf`. The leader's `full` also gates the 2-SM data. + auto utccp_required_smem_warp_transpose = [&](uint32_t* smem_ptr) { + DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, "Invalid aligned elements"); + uint32_t values[4]; + #pragma unroll + for (uint32_t i = 0; i < 4; ++ i) + values[i] = ptx::ld_shared(smem_ptr + i * 32 + lane_idx); + __syncwarp(); + ptx::st_shared(smem_ptr + lane_idx * 4, values[0], values[1], values[2], values[3]); + }; + + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) + utccp_required_smem_warp_transpose(smem_sfa[stage_idx] + i * kNumUTCCPAlignedElems); + cutlass::arch::fence_view_async_shared(); + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) + utccp_required_smem_warp_transpose(smem_sfb[stage_idx] + i * kNumUTCCPAlignedElems); + cutlass::arch::fence_view_async_shared(); + with_sf_full_barriers[stage_idx]->arrive(0u); + } + } + } else if (warp_idx >= kNumNonEpilogueThreads / 32 and warp_idx < (kNumNonEpilogueThreads + kNumUMMAStoreThreads) / 32) { + // Epilogue warp groups (swap-AB BF16 store) + const auto epilogue_warp_idx = warp_idx - (kNumNonEpilogueThreads / 32); + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(tmem_ptr_in_smem) == 0); + + uint32_t tma_stage_idx = 0; + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + auto accum_stage_idx = scheduler.current_iter % kNumEpilogueStages; + auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; + + tmem_full_barriers[accum_stage_idx]->wait(accum_phase_idx); + ptx::tcgen05_after_thread_sync(); + + const auto tmem_base_addr = accum_stage_idx * UMMA_N; + const auto base_m_idx = scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx); + const auto base_n_idx = n_block_idx * BLOCK_N; + const auto effective_m = scheduler.get_aligned_effective_m_in_block(m_block_idx); + epilogue::sm100_store_cd_swap_ab< + BLOCK_M, BLOCK_N, STORE_BLOCK_M, STORE_BLOCK_N, + kSwizzleCDMode, kNumTMAStoreStages, kNumUMMAStoreThreads, + GemmType::Normal, false, + cd_dtype_t, epilogue::transform::EpilogueIdentity> + (smem_cd, tma_stage_idx, tmem_base_addr, + base_m_idx, base_n_idx, scheduler.current_group_idx, + effective_m, epilogue_warp_idx, lane_idx, + tmem_empty_barriers[accum_stage_idx], tensor_map_cd); + } + } + + comm::cluster_sync_with_relaxed_arrive(); + if (warp_idx == 0) + Allocator().free(0, kNumTmemCols); +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports sm_100f"); +#endif +} + +} // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh new file mode 100644 index 0000000000..e13c3e8a96 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh @@ -0,0 +1,1474 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + +template < + uint32_t kNumMaxTokensPerRank, + uint32_t kHidden, uint32_t kIntermediateHidden, + uint32_t kNumExperts, uint32_t kNumTopk, + uint32_t kNumExpertsPerWave, + uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K, + uint32_t STORE_BLOCK_M, + uint32_t SF_BLOCK_M, uint32_t SF_BLOCK_N, + uint32_t kNumRingTokens, + uint32_t kNumSFRingTokens, + uint32_t kNumStages, + uint32_t kNumBytesPerPull, + uint32_t kNumDispatchThreads, uint32_t kNumNonEpilogueThreads, + uint32_t kNumEpilogueThreads, + uint32_t kNumSMs, uint32_t kNumRanks, + float kActivationClamp, + bool kFastMath, + uint32_t L1_SHAPE_N = kIntermediateHidden * 2, + uint32_t L1_SHAPE_K = kHidden, + uint32_t L2_SHAPE_N = kHidden, + uint32_t L2_SHAPE_K = kIntermediateHidden, + uint32_t kNumDispatchWarps = kNumDispatchThreads / 32, + uint32_t kNumMMANonEpilogueWarps = kNumNonEpilogueThreads / 32, + uint32_t kNumEpilogueWarps = kNumEpilogueThreads / 32, + uint32_t kNumEpilogueWarpgroups = kNumEpilogueWarps / 4, + uint32_t kNumThreads = kNumDispatchThreads + kNumNonEpilogueThreads + kNumEpilogueThreads, + uint32_t kNumTokensPerWarp = 32 / kNumTopk, + uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks, + uint32_t kNumRingBlocks = kNumRingTokens / BLOCK_M +> +CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void +sm100_mxfp4_mxfp4_mega_moe_impl(void* y, + int* cumulative_local_expert_recv_stats, + const uint32_t num_tokens, + const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights_sf) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using Allocator = cute::TMEM::Allocator2Sm; + + // Template checks + DG_STATIC_ASSERT(kNumDispatchThreads % 128 == 0, "Invalid number of dispatch threads"); + DG_STATIC_ASSERT(kNumNonEpilogueThreads == 128, "Invalid number of MMA non-epilogue threads"); + DG_STATIC_ASSERT(kNumEpilogueThreads % 128 == 0, "Invalid number of MMA epilogue and combine threads"); + DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); + + // Thread indices + const bool is_leader_cta = cute::block_rank_in_cluster() == 0; + const uint32_t sm_idx = blockIdx.x; + const uint32_t thread_idx = threadIdx.x; + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t lane_idx = ptx::get_lane_idx(); + + // Prefetch TMA descriptors at the very beginning + if (warp_idx == 0) { + cute::prefetch_tma_descriptor(&tensor_map_l1_acts); + cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l1_weights); + cute::prefetch_tma_descriptor(&tensor_map_l1_weights_sf); + cute::prefetch_tma_descriptor(&tensor_map_l1_output); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l2_weights); + cute::prefetch_tma_descriptor(&tensor_map_l2_weights_sf); + } + + // Workspaces + const auto workspace = layout::Workspace( + sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk, kNumRingTokens); + + // Token and buffer layouts + // NOTES: MXFP4 packs 2 elements per byte, so token byte sizes are halved vs FP8. + // UE8M0 scale-factor layouts (gran-32, int32-packed) are unchanged. + constexpr auto fp8_token_layout = layout::Data(kHidden / 2); + constexpr auto bf16_token_layout = layout::Data(kHidden * sizeof(nv_bfloat16)); + constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden / 2); + constexpr auto fp8_sf_layout = layout::Data(kHidden / 32); + constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 32); + constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * sizeof(int64_t), false); + constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * sizeof(float), false); + constexpr auto l1_topk_weights_layout = layout::Data(sizeof(float), false); + + // Registered inputs + const auto input_token_buffer = layout::Buffer( + fp8_token_layout, 1, kNumMaxTokensPerRank, + workspace.get_end_ptr()); + const auto input_sf_buffer = layout::Buffer( + fp8_sf_layout, 1, kNumMaxTokensPerRank, + input_token_buffer.get_end_ptr()); + const auto input_topk_idx_buffer = layout::Buffer( + input_topk_idx_layout, 1, kNumMaxTokensPerRank, + input_sf_buffer.get_end_ptr()); + const auto input_topk_weights_buffer = layout::Buffer( + input_topk_weights_layout, 1, kNumMaxTokensPerRank, + input_topk_idx_buffer.get_end_ptr()); + + // SF and its buffer configs + constexpr uint32_t kGranK = 32; + constexpr uint32_t kNumUTCCPAlignedElems = 128; + DG_STATIC_ASSERT(SF_BLOCK_M == math::constexpr_align(BLOCK_M, kNumUTCCPAlignedElems), "Invalid SF_BLOCK_M"); + DG_STATIC_ASSERT(SF_BLOCK_N == BLOCK_N, "No padding is needed for SFB"); + + // UTCCP 4x32 transpose index mapping within each 128-element group + const auto transform_sf_token_idx = [](const uint32_t& token_idx_in_expert) { + const uint32_t idx = token_idx_in_expert % BLOCK_M; + return token_idx_in_expert / BLOCK_M * SF_BLOCK_M + + (idx & ~127u) + (idx & 31u) * 4 + ((idx >> 5) & 3u); + }; + + // L1 inputs + const auto l1_token_buffer = layout::Buffer( + fp8_token_layout, 1, kNumRingTokens, + input_topk_weights_buffer.get_end_ptr()); + const auto l1_sf_buffer = layout::Buffer( + fp8_sf_layout, 1, kNumSFRingTokens, + l1_token_buffer.get_end_ptr()); + const auto l1_topk_weights_buffer = layout::Buffer( + l1_topk_weights_layout, 1, kNumRingTokens, + l1_sf_buffer.get_end_ptr()); + + // L2 inputs + const auto l2_token_buffer = layout::Buffer( + fp8_intermediate_token_layout, 1, kNumRingTokens, + l1_topk_weights_buffer.get_end_ptr() + ); + const auto l2_sf_buffer = layout::Buffer( + fp8_intermediate_sf_layout, 1, kNumSFRingTokens, + l2_token_buffer.get_end_ptr() + ); + + // Combine inputs + const auto combine_token_buffer = layout::Buffer( + bf16_token_layout, kNumTopk, kNumMaxTokensPerRank, + l2_sf_buffer.get_end_ptr() + ); + + // Data types + // NOTES: both activations and weights are packed MXFP4 (E2M1, 2 elems/byte) + using a_dtype_t = cutlass::float_e2m1_t; + using b_dtype_t = cutlass::float_e2m1_t; + + // MMA configs + // NOTES: always swap A/B, 2-CTA MMA, and matrices are K-major + constexpr uint32_t LAYOUT_AD_M = 128; + constexpr uint32_t UMMA_M = LAYOUT_AD_M * 2; + constexpr uint32_t UMMA_N = BLOCK_M; // Swap AB + constexpr uint32_t UMMA_BLOCK_K = 128; + constexpr uint32_t UMMA_K = 64; // packed FP4 contracts K=64 per instruction + constexpr uint32_t LOAD_BLOCK_M = BLOCK_M / 2; // Multicast on A + constexpr uint32_t LOAD_BLOCK_N = BLOCK_N; + DG_STATIC_ASSERT(BLOCK_M % 16 == 0, "Invalid block M"); + DG_STATIC_ASSERT(BLOCK_N == LAYOUT_AD_M, "Invalid block N"); + + // Packed-FP4 byte math (2 elements per byte) + constexpr uint32_t BLOCK_K_BYTES = BLOCK_K / 2; + constexpr uint32_t UMMA_BLOCK_K_BYTES = UMMA_BLOCK_K / 2; + constexpr uint32_t UMMA_K_BYTES = UMMA_K / 2; + + // Swizzle configs + // NOTES: K-major packed FP4 uses a swizzle equal to the K extent in bytes (BLOCK_K/2) + constexpr uint32_t kSwizzleAMode = BLOCK_K_BYTES; + constexpr uint32_t kSwizzleBMode = BLOCK_K_BYTES; + constexpr uint32_t kSwizzleCDMode = 128; + DG_STATIC_ASSERT(BLOCK_N % kSwizzleCDMode == 0, "Invalid block N"); + + // Epilogue configs + constexpr uint32_t kNumEpilogueStages = 2; + constexpr uint32_t kNumTMAStoreStages = 2; + + // Shared memory + constexpr uint32_t kSharedMemoryAlignment = 1024; + extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[]; + + // Shared memory sizes + // NOTES: FP8 CD output for L1 (2 TMA stages, BLOCK_N/2 post-SwiGLU), BF16 output for L2 (no TMA, a single stage) + constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; + constexpr uint32_t AMAX_REDUCTION_WARP_BUFFER_SIZE = STORE_BLOCK_M / 2; // float2 + + struct SharedStorage { + alignas(kSharedMemoryAlignment) uint32_t expert_token_count[kNumExperts]; + alignas(kSharedMemoryAlignment) uint8_t dispatch_send_buffer[kNumDispatchWarps][kNumBytesPerPull]; + union { + // L1 output is packed FP4 (E2M1, 2 elems/byte). NOTES: over-sized to the FP8 + // footprint (the BF16 `l2` member dominates the union anyway) so the stage-1 + // placeholder store stays in-bounds; the real packed layout is finalized in stage 2. + alignas(kSharedMemoryAlignment) uint8_t l1[kNumEpilogueWarpgroups][kNumTMAStoreStages][STORE_BLOCK_M * L1_OUT_BLOCK_N]; + alignas(kSharedMemoryAlignment) nv_bfloat16 l2[kNumEpilogueWarpgroups][STORE_BLOCK_M * BLOCK_N]; + } smem_d; + // Packed-FP4 A/B tiles: byte-addressed (`BLOCK_K / 2` bytes per row) + alignas(kSharedMemoryAlignment) uint8_t smem_a[kNumStages][LOAD_BLOCK_M * BLOCK_K_BYTES]; + alignas(kSharedMemoryAlignment) uint8_t smem_b[kNumStages][LOAD_BLOCK_N * BLOCK_K_BYTES]; + uint32_t smem_sfa[kNumStages][SF_BLOCK_M * (BLOCK_K / 128)]; + uint32_t smem_sfb[kNumStages][SF_BLOCK_N * (BLOCK_K / 128)]; + float2 amax_reduction[kNumEpilogueWarps][AMAX_REDUCTION_WARP_BUFFER_SIZE]; + Barrier dispatch_barriers[kNumDispatchWarps]; + Barrier full_barriers[kNumStages]; + Barrier empty_barriers[kNumStages]; + // Cross-CTA "loads-done" barrier: per-CTA SM90 loads each fill their own CTA's + // `full_barriers`; both CTAs' sync warp then funnel arrival to the leader so the + // (leader-only) 2-SM MMA waits until BOTH CTAs' operands are resident. + Barrier with_sf_full_barriers[kNumStages]; + Barrier tmem_full_barriers[kNumEpilogueStages]; + Barrier tmem_empty_barriers[kNumEpilogueStages]; + Barrier combine_barriers[kNumEpilogueWarps * 2]; + uint32_t tmem_ptr_in_smem; + }; + constexpr uint32_t kNumReusableSmemBytes = offsetof(SharedStorage, dispatch_barriers); + SharedStorage &shared_storage = *reinterpret_cast(smem_buffer); + + // Send buffers + constexpr auto pull_layout = layout::Data(kNumBytesPerPull); + const auto smem_send_buffers = layout::Buffer( + pull_layout, kNumDispatchWarps, 1, + static_cast(shared_storage.dispatch_send_buffer)); + + // Tensor memory size + constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumEpilogueStages; + constexpr uint32_t kNumSFATmemCols = SF_BLOCK_M / 32; + constexpr uint32_t kNumSFBTmemCols = SF_BLOCK_N / 32; + constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); + constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols; + constexpr uint32_t kTmemStartColOfSFB = kNumAccumTmemCols + kNumSFATmemCols; + DG_STATIC_ASSERT(32 <= kNumTmemCols and kNumTmemCols <= 512, "Invalid tensor memory columns"); + + // A cluster sync is essential for 2CTA tensor memory allocation + comm::cluster_sync_with_relaxed_arrive(); + + // Initialization + if (warp_idx == 0) { + // Clean shared memory + if (cute::elect_one_sync()) { + // The bytes must be 8 bytes aligned + ptx::st_shared_bulk( + shared_storage.expert_token_count, + math::constexpr_align(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment) + ); + } + } else if (warp_idx == 1) { + // Init m-barriers for dispatch + #pragma unroll + for (uint32_t i = lane_idx; i < kNumDispatchWarps; i += 32) + shared_storage.dispatch_barriers[i].init(1); + cutlass::arch::fence_barrier_init(); + } else if (warp_idx == 2) { + // Init GEMM barriers + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + // Per-CTA SM90 loads: A-load + B-load warps each arrive on their own CTA's full barrier + shared_storage.full_barriers[i].init(2); + shared_storage.empty_barriers[i].init(1); + // Both CTAs' sync warp (32 lanes each) arrive on the leader CTA + shared_storage.with_sf_full_barriers[i].init(2 * 32); + } + #pragma unroll + for (uint32_t i = 0; i < kNumEpilogueStages; ++ i) { + // Arrive at all CTAs + shared_storage.tmem_full_barriers[i].init(1); + // Arrive only at the leader CTA + shared_storage.tmem_empty_barriers[i].init(2 * kNumEpilogueThreads); + } + #pragma unroll + for (uint32_t i = 0; i < kNumEpilogueWarps * 2; ++ i) + shared_storage.combine_barriers[i].init(1); + } + cutlass::arch::fence_barrier_init(); + } else if (warp_idx == 3) { + // Allocate tensor memory + Allocator().allocate(kNumTmemCols, &shared_storage.tmem_ptr_in_smem); + } + // NOTES: Using `.relaxed` is allowed here since `fence_barrier_init` is `.release.cluster`, + // and `barrier.cluster.wait.aligned` is by default `.acquire` + comm::cluster_sync_with_relaxed_arrive(); + + // Task scheduler + auto scheduler = sched::MegaMoEScheduler< + BLOCK_M, BLOCK_N, BLOCK_K, + L1_SHAPE_N, L1_SHAPE_K, + L2_SHAPE_N, L2_SHAPE_K, + kNumExpertsPerRank, + kNumExpertsPerWave, + kNumSMs, kNumRanks>(workspace); + + // MMA pipeline and TMA phases + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + + // Flip phases only if reach the next first stage + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + // Intra-SM Barrier indices + constexpr uint32_t kDispatchBarrierIdx = 0; + constexpr uint32_t kDispatchWithEpilogueBarrierIdx = 1; + constexpr uint32_t kEpilogueFullBarrierIdx = 2; + constexpr uint32_t kEpilogueWGBarrierStartIdx = 3; + + // NVLink barrier tags + constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; + constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; + constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + + // Adjust registers + // NOTES: more experts per rank will cost more schedulers' registers + constexpr bool kUseMoreEpilogueRegisters = kNumExpertsPerRank <= 64; + constexpr uint32_t kNumDispatchRegisters = kUseMoreEpilogueRegisters ? 48 : 96; + constexpr uint32_t kNumNonEpilogueRegisters = kUseMoreEpilogueRegisters ? 40 : 88; + constexpr uint32_t kNumEpilogueRegisters = kUseMoreEpilogueRegisters ? 208 : 160; + DG_STATIC_ASSERT(kNumDispatchRegisters * kNumDispatchThreads + + kNumNonEpilogueRegisters * kNumNonEpilogueThreads + + kNumEpilogueRegisters * kNumEpilogueThreads <= 64512, + "Too many registers"); + + // Grid sync index assignments (dispatch and epilogue use separate counters to avoid conflicts) + constexpr uint32_t kDispatchGridSyncIndex = 0; + constexpr uint32_t kEpilogueGridSyncIndex = 1; + + // Different warp roles + if (warp_idx < kNumDispatchWarps) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // Dispatch warps + DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk; + const auto read_topk_idx = [&](const auto& process) { + // TODO: figure out better unrolling + // Now, `unroll` is better than `unroll 8` + #pragma unroll + for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp; + i < num_tokens; + i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) { + // Allocate slots for each token-topk + int expert_idx = -1; + if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { + expert_idx = static_cast( + __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + if (expert_idx >= 0) + process(i * kNumTopk + lane_idx, expert_idx); + } + __syncwarp(); + } + }; + + // Count experts' tokens + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + atomicAdd_block(shared_storage.expert_token_count + expert_idx, 1); + }); + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Get SM offset (~6.5 us) + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const uint64_t send_value = (1ull << 32) | static_cast(shared_storage.expert_token_count[i]); + shared_storage.expert_token_count[i] = static_cast( + ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value)); + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Write source indices (~2 us with 512 tokens) + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + const auto dst_rank_idx = expert_idx / kNumExpertsPerRank; + const auto dst_slot_idx = atomicAdd_block(shared_storage.expert_token_count + expert_idx, 1); + const auto dst_ptr = workspace.get_src_token_topk_idx_ptr( + expert_idx % kNumExpertsPerRank, sym_buffer.rank_idx, dst_slot_idx); + *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx; + }); + + // Grid sync + comm::grid_sync( + workspace, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); } + ); + + // Write expert count + if (sm_idx == 0) { + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const auto dst_rank_idx = i / kNumExpertsPerRank; + const auto dst_local_expert_idx = i % kNumExpertsPerRank; + const auto expert_status = *workspace.get_expert_send_count_ptr(i); + *sym_buffer.map( + workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), + dst_rank_idx) = expert_status & 0xffffffff; + ptx::atomic_add_sys( + sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx), + expert_status); + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Barrier before pulling + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* After the grid sync above, there is no more writes by other SMs (except 0) */ false, + /* After the NVLink barrier, there is a grid sync */ true + ); + + // Ensure the epilogue barrier cannot run with the pull barrier + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Pull token data and SF from remote ranks into local L1 buffer + uint32_t pull_mbarrier_phase = 0; + const auto pull_buffer = smem_send_buffers.get_rank_buffer(warp_idx).get_data_buffer(0); + const auto pull_mbarrier = &shared_storage.dispatch_barriers[warp_idx]; + + // Cache expert token counts in registers (same pattern as scheduler) + scheduler.fetch_expert_recv_count(); + + // Per-rank counts for current expert (re-loaded when expert changes) + constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); + int current_expert_idx = -1; + uint32_t stored_rank_count[kNumRanksPerLane] = {}; + uint32_t expert_start_idx = 0, expert_end_idx = 0; + uint32_t expert_pool_block_offset = 0; + + constexpr uint32_t kNumGlobalWarps = kNumSMs * kNumDispatchWarps; + for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalWarps) { + // Advance expert until within the range + int old_expert_idx = current_expert_idx; + while (token_idx >= expert_end_idx) { + if (++ current_expert_idx >= kNumExpertsPerRank) + break; + + // Update pool block offset for the new expert + expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M); + + // Move start and end to the next expert + expert_start_idx = expert_end_idx; + expert_end_idx += scheduler.get_num_tokens(current_expert_idx); + } + + // Finish all tokens + if (current_expert_idx >= kNumExpertsPerRank) + break; + + // Load per-rank counts when expert changes + if (old_expert_idx != current_expert_idx) { + old_expert_idx = current_expert_idx; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t j = i * 32 + lane_idx; + // TODO: this is not coalesced + stored_rank_count[i] = j < kNumRanks ? + static_cast(*workspace.get_expert_recv_count_ptr(j, current_expert_idx)) : 0; + } + } + + // Round-robin rank selection via iterative min-peeling + uint32_t current_rank_in_expert_idx; + uint32_t remaining[kNumRanksPerLane]; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] = stored_rank_count[i]; + uint32_t offset = 0; + uint32_t token_idx_in_expert = token_idx - expert_start_idx; + uint32_t slot_idx = token_idx_in_expert; + uint32_t token_idx_in_rank; + while (true) { + // Compute active count and min across all ranks + // NOTES: reduce within each lane first, then warp-reduce once + uint32_t num_actives_in_lane = 0; + uint32_t min_in_lane = 0xffffffff; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + num_actives_in_lane += remaining[i] > 0; + if (remaining[i] > 0) + min_in_lane = cute::min(min_in_lane, remaining[i]); + } + const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane); + const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane); + + // Hit in the current round + const uint32_t num_round_tokens = length * num_active_ranks; + if (slot_idx < num_round_tokens) { + const uint32_t slot_idx_in_round = slot_idx % num_active_ranks; + uint32_t num_seen_ranks = 0; + current_rank_in_expert_idx = 0; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0); + const uint32_t num_active_lanes = __popc(mask); + if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes) + current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1); + num_seen_ranks += num_active_lanes; + } + token_idx_in_rank = offset + (slot_idx / num_active_ranks); + break; + } + + // Move into the next round + slot_idx -= num_round_tokens; + offset += length; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] -= cute::min(remaining[i], length); + } + + // Read source token-topk index (written by remote dispatch via NVLink) + const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr( + current_expert_idx, current_rank_in_expert_idx, token_idx_in_rank); + const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; + const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk; + + // Packed-FP4 token bytes are divided into chunks (hidden / 2 bytes per token) + constexpr uint32_t kNumTokenBytes = kHidden / 2; + constexpr uint32_t kNumChunks = kNumTokenBytes / kNumBytesPerPull; + DG_STATIC_ASSERT(kNumChunks * kNumBytesPerPull == kNumTokenBytes, "kNumBytesPerPull must divide packed token bytes"); + + // TMA load token from remote rank and store into local + const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; + const uint32_t pool_block_idx = pool_token_idx / BLOCK_M; + + // Wait for ring buffer slot to be available (previous consumer must have finished all N blocks) + constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; + const auto l1_empty_count_target = (pool_block_idx / kNumRingBlocks) * kNumL1BlockNs; + if (l1_empty_count_target > 0) { + const auto empty_ptr = workspace.get_l1_empty_count_ptr(pool_block_idx % kNumRingBlocks); + while (ptx::ld_acq(empty_ptr) < l1_empty_count_target); + } + + const auto src_base_ptr = sym_buffer.map( + input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), current_rank_in_expert_idx); + const auto dst_base_ptr = l1_token_buffer.get_data_buffer(pool_token_idx % kNumRingTokens).get_base_ptr(); + const auto issue_and_wait_pull_store = [&](const uint32_t& i) { + ptx::mbarrier_wait_and_flip_phase(pull_mbarrier, pull_mbarrier_phase); + ptx::tma_store_1d( + math::advance_ptr(dst_base_ptr, i * kNumBytesPerPull), + pull_buffer.get_base_ptr(), kNumBytesPerPull + ); + cute::tma_store_arrive(); + ptx::tma_store_wait<0>(); + }; + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumChunks; ++ i) { + ptx::tma_load_1d( + pull_buffer.get_base_ptr(), + math::advance_ptr(src_base_ptr, i * kNumBytesPerPull), + pull_mbarrier, kNumBytesPerPull + ); + ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kNumBytesPerPull); + i != (kNumChunks - 1) ? issue_and_wait_pull_store(i) : void(); + } + } + __syncwarp(); + + // Load and store SF (overlaps with last chunk's TMA load from remote) + constexpr uint32_t kNumSFUint32 = kHidden / 128; + DG_STATIC_ASSERT(kNumSFUint32 > 0 and kHidden % 128 == 0, "Invalid SF"); + const auto remote_sf_ptr = sym_buffer.map( + input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx); + const auto local_sf_ptr = l1_sf_buffer.get_base_ptr(); + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + const uint32_t token_idx_in_block = token_idx_in_expert % BLOCK_M; + const auto sf_ring_token_idx = ring_block_idx * SF_BLOCK_M + + transform_sf_token_idx(token_idx_in_block); + #pragma unroll + for (uint32_t i = 0; i < math::constexpr_ceil_div(kNumSFUint32, 32u); ++ i) { + const uint32_t j = i * 32 + lane_idx; + if (j < kNumSFUint32) + local_sf_ptr[j * kNumSFRingTokens + sf_ring_token_idx] = remote_sf_ptr[j]; + } + __syncwarp(); + + // Store weights and metadata + if (cute::elect_one_sync()) { + // Load weights + const auto weight = *sym_buffer.map( + input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx % kNumRingTokens).template get_base_ptr() = weight; + + // Write source metadata for combine write-back (logical pool token) + *workspace.get_token_src_metadata_ptr(pool_token_idx) = + {current_rank_in_expert_idx, src_token_idx, src_topk_idx}; + + // Complete last chunk's store + issue_and_wait_pull_store(kNumChunks - 1); + const bool is_last_token = (token_idx == expert_end_idx - 1); + ptx::red_add_rel( + workspace.get_l1_full_count_ptr(pool_block_idx % kNumRingBlocks), + is_last_token ? BLOCK_M - (token_idx_in_expert % BLOCK_M) : 1u + ); + } + __syncwarp(); + } + + // Clean workspace for the next usage, and also do cumulative stats + // NOTES: it is overlapped with combine reduction epilogue + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + DG_STATIC_ASSERT(kNumSMs > 1, "Invalid SM count"); + if (sm_idx == 0) { + // SM 0: clear expert send count + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) + *workspace.get_expert_send_count_ptr(i) = 0; + } else { + // Other SMs: clean blocks + for (uint32_t i = sm_idx - 1; i < kNumExpertsPerRank; i += kNumSMs - 1) { + // Read expert token count before clearing + const auto num_recv_tokens = static_cast( + *workspace.get_expert_recv_count_sum_ptr(i)); + const auto num_recv_m_blocks = math::ceil_div(num_recv_tokens, BLOCK_M); + + // Compute expert pool block offset + expert_pool_block_offset = scheduler.get_pool_block_offset(i); + + // Wait read count ready + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Clean expert token count, and add cumulative results + DG_STATIC_ASSERT(kNumDispatchWarps >= 2, "Not enough dispatch warps"); + if (warp_idx == 0) { + *workspace.get_expert_recv_count_sum_ptr(i) = 0; + } else if (warp_idx == 1) { + if (cute::elect_one_sync() and cumulative_local_expert_recv_stats != nullptr) + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); + __syncwarp(); + } + + // Clean per-rank token count + for (uint32_t j = thread_idx; j < kNumRanks; j += kNumDispatchThreads) + *workspace.get_expert_recv_count_ptr(j, i) = 0; + __syncwarp(); + + // Clean L1 and L2 full stuffs and ring buffer counts + for (uint32_t j = thread_idx; j < num_recv_m_blocks; j += kNumDispatchThreads) { + *workspace.get_l1_full_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l1_empty_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l2_full_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l2_empty_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + } + __syncwarp(); + } + } + + // Wait for all ranks to finish cleaning + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + /* Before the NVLink barrier, there is a grid sync */ true, + /* At the end of kernel does not need to sync */ false + ); + } else if (warp_idx == kNumDispatchWarps) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // GEMM TMA load warp for tokens with SFA + scheduler.for_each_block([&](const sched::BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + const auto tensor_map_a_ptr = block_phase == sched::BlockPhase::Linear2 + ? &tensor_map_l2_acts : &tensor_map_l1_acts; + const auto tensor_map_sfa_ptr = block_phase == sched::BlockPhase::Linear2 + ? &tensor_map_l2_acts_sf : &tensor_map_l1_acts_sf; + + const auto shape_k = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_K : L1_SHAPE_K; + const auto shape_sfa_k = math::ceil_div(shape_k, kGranK * 4u); + + // Compute pool block offset for this expert + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + + // Wait the entire token arrival for linear 1 + if (block_phase == sched::BlockPhase::Linear1) { + const auto ptr = workspace.get_l1_full_count_ptr(ring_block_idx); + const auto num_expected_tokens = BLOCK_M * (pool_block_idx / kNumRingBlocks + 1); + while (ptx::ld_acq(ptr) != num_expected_tokens); + } else { + const auto ptr = workspace.get_l2_full_count_ptr(ring_block_idx); + const auto num_expected_blocks = (L2_SHAPE_K / BLOCK_N) * 2 * (pool_block_idx / kNumRingBlocks + 1); + while (ptx::ld_acq(ptr) != num_expected_blocks); + } + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + // Wait consumer release + shared_storage.empty_barriers[stage_idx].wait(phase ^ 1); + + // Compute token offsets from ring block index + uint32_t ring_m_idx = ring_block_idx * BLOCK_M; + uint32_t k_idx = k_block_idx * BLOCK_K; + uint32_t sfa_ring_m_idx = ring_block_idx * SF_BLOCK_M; + uint32_t sfa_k_idx = k_block_idx * (BLOCK_K / 128); + + // Add 2 CTA offsets for non-leader CTA + if (not is_leader_cta) + ring_m_idx += scheduler.template get_valid_m() / 2; + + // TMA copy tokens (TRUE 2-CTA cta_group::2) + SFA (per-CTA SM90), then arrive. + // NOTES (validated against CUTLASS via instrumentation, and in the standalone GEMM): + // - acts use `SM100_TMA_2SM_LOAD` with the per-CTA M offset above; the 2-SM atom + // routes ALL data tx to the LEADER's `full` barrier (cluster peer-bit masking). + // - SFA uses a per-CTA SM90 load into THIS CTA's smem (the 2-CTA UTCCP reads both + // CTAs' SF), signaling THIS CTA's own `full`. So leader expects 2x data + own SF, + // while the non-leader's `full` only sees its own SF (its data went to the leader). + // - Packed FP4 is a single raw 2D TMA (the `tma::copy` byte-splitter overruns FP4). + if (cute::elect_one_sync()) { + // acts: TRUE 2-CTA cta_group::2 load (routes tx to leader); use the ring offset + cute::SM100_TMA_2SM_LOAD_2D::copy( + tensor_map_a_ptr, reinterpret_cast(&shared_storage.full_barriers[stage_idx]), + static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + shared_storage.smem_a[stage_idx], k_idx, ring_m_idx); + // SFA: per-CTA SM90 load into THIS CTA's smem (the 2-CTA UTCCP reads both CTAs' SF) + tma::copy( + tensor_map_sfa_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_sfa[stage_idx], sfa_ring_m_idx, sfa_k_idx); + // Leader collects BOTH CTAs' acts (2-SM routed) + its own SFA; non-leader sees only its SFA + if (is_leader_cta) + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx( + sizeof(SharedStorage::smem_a[0]) * 2 + sizeof(SharedStorage::smem_sfa[0])); + else + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx( + sizeof(SharedStorage::smem_sfa[0])); + } + __syncwarp(); + } + }); + } else if (warp_idx == kNumDispatchWarps + 1) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // GEMM TMA load warp for weights with SF + scheduler.for_each_block([&](const sched::BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + const auto tensor_map_b_ptr = + block_phase == sched::BlockPhase::Linear2 ? &tensor_map_l2_weights : &tensor_map_l1_weights; + const auto tensor_map_sfb_ptr = + block_phase == sched::BlockPhase::Linear2 ? &tensor_map_l2_weights_sf : &tensor_map_l1_weights_sf; + + const auto shape_k = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_K : L1_SHAPE_K; + const auto shape_n = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_N : L1_SHAPE_N; + const auto shape_sfb_k = math::ceil_div(shape_k, kGranK * 4u); + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + // Wait consumer release + shared_storage.empty_barriers[stage_idx].wait(phase ^ 1); + + // Compute weight offset + uint32_t n_idx = local_expert_idx * shape_n + n_block_idx * BLOCK_N; + uint32_t k_idx = k_block_idx * BLOCK_K; + uint32_t sfb_n_idx = n_block_idx * BLOCK_N; + uint32_t sfb_k_idx = local_expert_idx * shape_sfb_k + k_block_idx * (BLOCK_K / 128); + + // TMA copy weights (TRUE 2-CTA cta_group::2; both CTAs use the same n coord so each + // loads the full BLOCK_N replica) + SFB (per-CTA SM90). Data tx routes to the leader. + if (cute::elect_one_sync()) { + cute::SM100_TMA_2SM_LOAD_2D::copy( + tensor_map_b_ptr, reinterpret_cast(&shared_storage.full_barriers[stage_idx]), + static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + shared_storage.smem_b[stage_idx], k_idx, n_idx); + tma::copy( + tensor_map_sfb_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_sfb[stage_idx], sfb_n_idx, sfb_k_idx); + if (is_leader_cta) + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx( + sizeof(SharedStorage::smem_b[0]) * 2 + sizeof(SharedStorage::smem_sfb[0])); + else + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx( + sizeof(SharedStorage::smem_sfb[0])); + } + __syncwarp(); + } + }); + } else if (warp_idx == kNumDispatchWarps + 2) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // GEMM MMA issue warp (only the leader CTA will run) + if (is_leader_cta) { + // Make instruction descriptor with block scaling + // NOTES: always swap A/B + auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< + b_dtype_t, a_dtype_t, float, cutlass::float_ue8m0_t, + UMMA_M, UMMA_N, + cute::UMMA::Major::K, cute::UMMA::Major::K + >(); + auto sf_desc = mma::sm100::make_sf_desc(nullptr); + + DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); + // Packed-FP4 K-major smem layout (byte-addressed; swizzle == BLOCK_K/2 bytes) + constexpr auto kFP4Layout = mma::sm100::to_umma_layout_type(); + + // Checks for MMA instructions + DG_STATIC_ASSERT((UMMA_M == 64 and UMMA_N % 8 == 0 and 8 <= UMMA_N and UMMA_N <= 256) or + (UMMA_M == 128 and UMMA_N % 16 == 0 and 16 <= UMMA_N and UMMA_N <= 256) or + (UMMA_M == 256 and UMMA_N % 16 == 0 and 16 <= UMMA_N and UMMA_N <= 256), + "Invalid MMA instruction shape"); + + // Persistently schedule over blocks + uint32_t current_iter_idx = 0; + scheduler.for_each_block([&](const sched::BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + // Dynamic update of UMMA N based on effective M + mma::sm100::update_instr_desc_with_umma_n(instr_desc, scheduler.template get_valid_m()); + + // Wait tensor memory empty barrier arrival + const auto accum_stage_idx = current_iter_idx % kNumEpilogueStages; + const auto accum_phase = (current_iter_idx ++ / kNumEpilogueStages) & 1; + shared_storage.tmem_empty_barriers[accum_stage_idx].wait(accum_phase ^ 1); + ptx::tcgen05_after_thread_sync(); + + // Empty barrier arrival + auto empty_barrier_arrive = [&](const bool& do_tmem_full_arrive) { + auto umma_arrive = [](const uint64_t* barrier) { + constexpr uint16_t kCTAMask = (1 << 2) - 1; + cutlass::arch::umma_arrive_multicast_2x1SM(barrier, kCTAMask); + }; + umma_arrive(reinterpret_cast(&shared_storage.empty_barriers[stage_idx])); + + // NOTES: the tensor memory accumulator pipeline has nothing to do with multicasting + if (do_tmem_full_arrive) + umma_arrive(reinterpret_cast(&shared_storage.tmem_full_barriers[accum_stage_idx])); + __syncwarp(); + }; + + // Launch MMAs + #pragma unroll 2 + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + // Wait both CTAs' TMA loads (funneled to the leader by the sync warp) + shared_storage.with_sf_full_barriers[stage_idx].wait(phase); + ptx::tcgen05_after_thread_sync(); + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t umma_k_block_idx = 0; umma_k_block_idx < BLOCK_K / UMMA_BLOCK_K; ++ umma_k_block_idx) { + // UTCCP copy SFA and SFB to TMEM (SF pre-transposed during dispatch / weight transform) + using cute_utccp_t = cute::SM100_UTCCP_4x32dp128bit_2cta; + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { + auto smem_ptr = shared_storage.smem_sfa[stage_idx] + umma_k_block_idx * SF_BLOCK_M + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + i * 4); + } + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { + auto smem_ptr = shared_storage.smem_sfb[stage_idx] + umma_k_block_idx * SF_BLOCK_N + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + i * 4); + } + + // Issue packed-mxf4 UMMA over UMMA_K (=64) sub-tiles (byte-addressed smem) + const auto a_stage_ptr = reinterpret_cast(shared_storage.smem_a[stage_idx]) + umma_k_block_idx * UMMA_BLOCK_K_BYTES; + const auto b_stage_ptr = reinterpret_cast(shared_storage.smem_b[stage_idx]) + umma_k_block_idx * UMMA_BLOCK_K_BYTES; + #pragma unroll + for (uint32_t k = 0; k < UMMA_BLOCK_K / UMMA_K; ++ k) { + const uint32_t sf_id = (umma_k_block_idx * (UMMA_BLOCK_K / UMMA_K) + k) * 2; + const auto runtime_instr_desc = + mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, sf_id, sf_id); + auto a_desc = mma::sm100::make_smem_desc(kFP4Layout, a_stage_ptr + k * UMMA_K_BYTES, 8 * kSwizzleAMode, 0); + auto b_desc = mma::sm100::make_smem_desc(kFP4Layout, b_stage_ptr + k * UMMA_K_BYTES, 8 * kSwizzleBMode, 0); + // Swap-AB: weights (b_desc) take the UMMA "A" slot + SFB + ptx::SM100_MMA_MXF4_2x1SM_SS::fma( + b_desc, a_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or umma_k_block_idx > 0 or k > 0, runtime_instr_desc, + kTmemStartColOfSFB, kTmemStartColOfSFA); + } + } + } + __syncwarp(); + + // Commit to the mbarrier object + // No explicit `tcgen05.fence::before_thread_sync` is needed, as this is implicitly performed by `tcgen05.commit` + empty_barrier_arrive(k_block_idx == num_k_blocks - 1); + } + }); + + // To safely deconstruct barriers, we need another round of waits + if (current_iter_idx > 0) { + const auto accum_phase_idx = ((current_iter_idx - 1) / kNumEpilogueStages) & 1; + shared_storage.tmem_empty_barriers[(current_iter_idx - 1) % kNumEpilogueStages].wait(accum_phase_idx); + } + } + } else if (warp_idx == kNumDispatchWarps + 3) { + // Adjust registers + cutlass::arch::warpgroup_reg_dealloc(); + + // Cross-CTA loads-done sync warp (runs on BOTH CTAs). + // Each CTA's per-CTA SM90 loads fill its own `full_barriers`; this warp waits that and + // funnels arrival to the leader's `with_sf_full_barriers`, so the leader-only 2-SM MMA + // only fires once BOTH CTAs' operands are resident. Mirrors the validated standalone. + scheduler.for_each_block([&](const sched::BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + shared_storage.full_barriers[stage_idx].wait(phase); + // All 32 lanes arrive on the leader CTA (init == 2 * 32) + shared_storage.with_sf_full_barriers[stage_idx].arrive(0u); + } + }); + } else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) { + // Adjust registers + cutlass::arch::warpgroup_reg_alloc(); + + // NOTES: tensor memory addresses are simplified, as the hardware will ignore the warp index bits, + // i.e., no need for `tmem_ptr |= (epilogue_warp_idx * 32) << 16`. + // NOTES: we also forbid two CTAs to share the same SM and its tensor memory + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&shared_storage.tmem_ptr_in_smem) == 0); + + // GEMM epilogue warps + const auto epilogue_warp_idx = warp_idx - (kNumDispatchWarps + kNumMMANonEpilogueWarps); + const auto epilogue_wg_idx = epilogue_warp_idx / 4; + const auto epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx; + const auto warp_idx_in_wg = epilogue_warp_idx % 4; + DG_STATIC_ASSERT((kNumDispatchWarps + kNumMMANonEpilogueWarps) % 4 == 0 and + kNumEpilogueWarps % 4 == 0, "Invalid epilogue warps"); + + // TODO: support effective block M + // NOTES: + // - 2 warpgroups divide the whole BM into BM / 2 + // - 4 warps divide the whole BN into BN / 4 + // - BM / 2 is further divided into stored blocks, i.e. with `STORE_BLOCK_M` size + // - `STORE_BLOCK_M` in further divided into `ATOM_M` + constexpr uint32_t WG_BLOCK_M = BLOCK_M / kNumEpilogueWarpgroups; + constexpr uint32_t ATOM_M = 8; + constexpr uint32_t kNumBankGroupBytes = 16u; + constexpr uint32_t kNumAtomsPerStore = STORE_BLOCK_M / ATOM_M; + DG_STATIC_ASSERT(BLOCK_M % kNumEpilogueWarpgroups == 0, "Invalid block M"); + DG_STATIC_ASSERT(WG_BLOCK_M % STORE_BLOCK_M == 0, "Invalid warpgroup block M"); + DG_STATIC_ASSERT(STORE_BLOCK_M % ATOM_M == 0, "Invalid store block M"); + DG_STATIC_ASSERT(BLOCK_N == 128, "Invalid block N"); + + // Ensure the epilogue barrier cannot run with the pull barrier + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Persistently schedule over blocks + uint32_t current_iter_idx = 0; + scheduler.for_each_block([&](const sched::BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + // Wait UMMA arrival + const auto accum_stage_idx = current_iter_idx % kNumEpilogueStages; + const auto accum_phase = (current_iter_idx ++ / kNumEpilogueStages) & 1; + shared_storage.tmem_full_barriers[accum_stage_idx].wait(accum_phase); + ptx::tcgen05_after_thread_sync(); + + // Compute offsets + // NOTES: use shuffle here to let NVCC know warp divergence won't happen + const uint32_t valid_m = ptx::exchange(scheduler.template get_valid_m(), 0); + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t ring_block_idx = pool_block_idx % kNumRingBlocks; + const uint32_t ring_m_idx = ring_block_idx * BLOCK_M; // Ring-buffer offset for reusable data buffers + const uint32_t pool_m_idx = pool_block_idx * BLOCK_M; // Full-pool offset for non-ring metadata + uint32_t n_idx = n_block_idx * BLOCK_N; + + if (block_phase == sched::BlockPhase::Linear1) { + // Wait L2 block empty + const auto l2_empty_ptr = workspace.get_l2_empty_count_ptr(ring_block_idx); + const auto num_expected_blocks = (L2_SHAPE_N / BLOCK_N) * (pool_block_idx / kNumRingBlocks); + while (ptx::ld_acq(l2_empty_ptr) != num_expected_blocks); + + // Unified L1 epilogue: SwiGLU in-place using granularity 8 interleaved weights + // With `SM100_TMEM_LOAD_16dp256b1x`, gate/up pairs are: + float stored_cached_weight = 0; + + #pragma unroll + for (uint32_t s = 0; s < WG_BLOCK_M / STORE_BLOCK_M; ++ s) { + // Early break if the entire store block is beyond the valid token range + if (epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M >= valid_m) { + ptx::tcgen05_before_thread_sync(); + shared_storage.tmem_empty_barriers[accum_stage_idx].arrive(0u); + break; + } + + // Iterate all atoms in the store block + float2 activation_values[kNumAtomsPerStore][2]; + float2 amax_values[kNumAtomsPerStore]; + #pragma unroll + for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) { + const uint32_t j = s * kNumAtomsPerStore + i; + + // Load weights from global into register cache per 32 tokens + DG_STATIC_ASSERT(32 % ATOM_M == 0, "Invalid block size"); + if ((j * ATOM_M) % 32 == 0 and (WG_BLOCK_M % 32 == 0 or j * ATOM_M + lane_idx < WG_BLOCK_M)) { + stored_cached_weight = *l1_topk_weights_buffer + .get_data_buffer(ring_m_idx + epilogue_wg_idx * WG_BLOCK_M + j * ATOM_M + lane_idx) + .template get_base_ptr(); + } + + // Load weights from register cache + const float2 weights = { + ptx::exchange(stored_cached_weight, (j * ATOM_M) % 32 + (lane_idx % 4) * 2 + 0), + ptx::exchange(stored_cached_weight, (j * ATOM_M) % 32 + (lane_idx % 4) * 2 + 1) + }; + + // Load from TMEM + uint2 raw_values[4]; + uint32_t tmem_addr = accum_stage_idx * UMMA_N + epilogue_wg_idx * WG_BLOCK_M + j * ATOM_M; + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, + raw_values[0].x, raw_values[0].y, raw_values[1].x, raw_values[1].y); + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr | 0x00100000, + raw_values[2].x, raw_values[2].y, raw_values[3].x, raw_values[3].y); + cutlass::arch::fence_view_async_tmem_load(); + + // Signal tensor memory consumed on the last atom + if (j == WG_BLOCK_M / ATOM_M - 1) { + ptx::tcgen05_before_thread_sync(); + shared_storage.tmem_empty_barriers[accum_stage_idx].arrive(0u); + } + + // Apply SwiGLU: silu(gate) * up + auto fp32_values = reinterpret_cast(raw_values); + #pragma unroll + for (uint32_t k = 0; k < 2; ++ k) { + auto bf16_gate = __float22bfloat162_rn(fp32_values[k * 2 + 0]); + auto bf16_up = __float22bfloat162_rn(fp32_values[k * 2 + 1]); + + // Clamp + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) { + bf16_gate = __hmin2(bf16_gate, {kActivationClamp, kActivationClamp}); + bf16_up = __hmax2(bf16_up, {-kActivationClamp, -kActivationClamp}); + bf16_up = __hmin2(bf16_up, {kActivationClamp, kActivationClamp}); + } + + // SwiGLU + auto gate = __bfloat1622float2(bf16_gate); + auto neg_gate_exp = make_float2( + kFastMath ? __expf(-gate.x) : expf(-gate.x), + kFastMath ? __expf(-gate.y) : expf(-gate.y)); + const auto denom = __fadd2_rn({1.0f, 1.0f}, neg_gate_exp); + if constexpr (kFastMath) { + gate = __fmul2_rn(gate, {math::fast_rcp(denom.x), math::fast_rcp(denom.y)}); + } else { + gate = {gate.x / denom.x, gate.y / denom.y}; + } + const auto up = __bfloat1622float2(bf16_up); + activation_values[i][k] = __fmul2_rn(__fmul2_rn(gate, up), weights); + } + + // Amax reduction (thread-level) + float2 thread_local_amax = {0.f, 0.f}; + #pragma unroll + for (uint32_t k = 0; k < 2; ++ k) { + thread_local_amax.x = cute::max(thread_local_amax.x, cute::abs(activation_values[i][k].x)); + thread_local_amax.y = cute::max(thread_local_amax.y, cute::abs(activation_values[i][k].y)); + } + + // Amax reduction (warp-level) + amax_values[i].x = math::warp_reduce<4, true>( + thread_local_amax.x, math::ReduceMax()); + amax_values[i].y = math::warp_reduce<4, true>( + thread_local_amax.y, math::ReduceMax()); + + // Reduce amax (warp-pair-level) + if (lane_idx < 4) + shared_storage.amax_reduction[epilogue_warp_idx][i * (ATOM_M / 2) + lane_idx] = amax_values[i]; + __syncwarp(); + } + + // Wait shared memory release from previous TMA store + // And fence `shared_storage.amax_reduction` + const uint32_t tma_stage_idx = s % kNumTMAStoreStages; + ptx::tma_store_wait(); + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Cast to FP8 E4M3 and store into shared memory + #pragma unroll + for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) { + // Reduce amax (warp-pair-level) + const float2 wp_amax = + shared_storage.amax_reduction[epilogue_warp_idx ^ 1][i * (ATOM_M / 2) + lane_idx % 4]; + amax_values[i].x = cute::max(amax_values[i].x, wp_amax.x); + amax_values[i].y = cute::max(amax_values[i].y, wp_amax.y); + + // Calculate UE8M0 SF (MXFP4: E2M1 max is 6.0, so scale by amax/6) + float2 sf, sf_inv; + { + const float2 scaled = __fmul2_rn(amax_values[i], {1.0f / 6.0f, 1.0f / 6.0f}); + const int ex = math::fast_log2_ceil(scaled.x); + const int ey = math::fast_log2_ceil(scaled.y); + sf.x = math::fast_pow2(ex); sf_inv.x = math::fast_pow2(-ex); + sf.y = math::fast_pow2(ey); sf_inv.y = math::fast_pow2(-ey); + } + + // Cast 2 activation outputs -> e2m1 nibbles, then write them DIRECTLY to their + // [token][inter] packed positions (no stmatrix / no smem swizzle). This + // relies on the verified `SM100_TMEM_LOAD_16dp256b1x` fragment map for the + // (swap-AB) accumulator: with lane = (a = lane%4, b = lane/4) and warp w, + // activation_values[i][0].{x,y} -> (token 2a / 2a+1, inter w*16 + b) + // activation_values[i][1].{x,y} -> (token 2a / 2a+1, inter w*16 + 8 + b) + // Packed K-major requires 2 consecutive inter per byte, so the even-`b` lane + // fetches the odd inter nibble from its `b+1` partner (lane + 4) via shuffle. + const float2 upper = __fmul2_rn(activation_values[i][0], sf_inv); + const float2 lower = __fmul2_rn(activation_values[i][1], sf_inv); + uint16_t e2m1x4; + asm volatile( + "{\n\t" + ".reg .b8 b0, b1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 b0, %2, %1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 b1, %4, %3;\n\t" + "mov.b16 %0, {b0, b1};\n\t" + "}\n" + : "=h"(e2m1x4) + : "f"(upper.x), "f"(upper.y), "f"(lower.x), "f"(lower.y)); + + // My 4 nibbles (low bits): [upper.x | upper.y | lower.x | lower.y] + const uint32_t my_nibs = + ((uint32_t)(e2m1x4 & 0xF)) | + ((uint32_t)((e2m1x4 >> 4) & 0xF) << 8) | + ((uint32_t)((e2m1x4 >> 8) & 0xF) << 16)| + ((uint32_t)((e2m1x4 >> 12) & 0xF) << 24); + // Partner = lane + 4 (the `b+1` inter); its nibbles become the high half. + const uint32_t pr_nibs = __shfl_sync(0xffffffffu, my_nibs, lane_idx + 4); + + // Only even-`b` lanes write (they own the even/low inter of each packed byte) + if (((lane_idx / 4) & 1) == 0) { + constexpr uint32_t L1_OUT_BLOCK_N_BYTES = L1_OUT_BLOCK_N / 2; + const uint32_t a = lane_idx % 4, bb = lane_idx / 4; + const uint32_t tok0 = i * ATOM_M + 2 * a; // local token in store tile + const uint32_t j0 = warp_idx_in_wg * 8 + bb / 2; // byte for inter (w*16+bb, +1) + const uint32_t j1 = warp_idx_in_wg * 8 + 4 + bb / 2; // byte for inter (w*16+8+bb, +1) + auto* p = reinterpret_cast(shared_storage.smem_d.l1[epilogue_wg_idx][tma_stage_idx]); + p[(tok0 ) * L1_OUT_BLOCK_N_BYTES + j0] = (my_nibs & 0xF) | (((pr_nibs) & 0xF) << 4); + p[(tok0 + 1) * L1_OUT_BLOCK_N_BYTES + j0] = ((my_nibs >> 8) & 0xF) | (((pr_nibs >> 8) & 0xF) << 4); + p[(tok0 ) * L1_OUT_BLOCK_N_BYTES + j1] = ((my_nibs >> 16) & 0xF) | (((pr_nibs >> 16) & 0xF) << 4); + p[(tok0 + 1) * L1_OUT_BLOCK_N_BYTES + j1] = ((my_nibs >> 24) & 0xF) | (((pr_nibs >> 24) & 0xF) << 4); + } + + // Store SF to `l2_sf_buffer` as UE8M0 (MN-major layout) + // Only one warp per pair writes (both hold the same SF after cross-warp reduce) + // Each lane < 4 holds SF for 2 rows (sf.x and sf.y) + if (warp_idx_in_wg % 2 == 0 and lane_idx < 4) { + const uint32_t k_idx = n_block_idx * 2 + warp_idx_in_wg / 2; + const uint32_t k_uint_idx = k_idx / 4, byte_idx = k_idx % 4; + const uint32_t mn_stride = kNumSFRingTokens * sizeof(uint32_t); + const auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + // NOTES: consecutive tokens (t, t + 1) are in the same 32-group, so `sf_idx` differs by 4 + // NOTES: originally there was: + // - `const uint32_t token_idx_in_expert = m_block_idx * BLOCK_M + epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + i * ATOM_M + lane_idx * 2 + // - `scheduler.get_current_pool_block_offset() * SF_BLOCK_M + transform_sf_token_idx(token_idx_in_expert)` + // We find out that + // 1. `m_block_idx * BLOCK_M` mod `BLOCK_M` is 0, and `epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + i * ATOM_M + lane_idx * 2` is always < `BLOCK_M`, so we can put `m_block_idx * BLOCK_M` outside + // 2. `lane_idx * 2` controls the lowest 3 bit of `token_idx_in_expert`, and `transform_sf_token_idx` is a bitwise-independent transformation if the input is less than `BLOCK_M`, so we can put `lane_idx * 2` outside + // This reduce the number of computation instructions. + const uint32_t token_base_idx = epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + i * ATOM_M; + __builtin_assume(token_base_idx < BLOCK_M); + const auto sf_ring_token_idx = ring_block_idx * SF_BLOCK_M + + transform_sf_token_idx(token_base_idx) + (lane_idx * 2) * 4; + const auto sf_addr = k_uint_idx * mn_stride + sf_ring_token_idx * static_cast(sizeof(uint32_t)) + byte_idx; + sf_base_ptr[sf_addr] = + (*reinterpret_cast(&sf.x) >> 23); + sf_base_ptr[sf_addr + 4 * static_cast(sizeof(uint32_t))] = + (*reinterpret_cast(&sf.y) >> 23); + } + __syncwarp(); + } + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Issue TMA store after all atoms in this store block + if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { + uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N; + cute::tma_store_fence(); + cute::SM90_TMA_STORE_2D::copy( + &tensor_map_l1_output, + shared_storage.smem_d.l1[epilogue_wg_idx][tma_stage_idx], + out_n_idx, + ring_m_idx + epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M); + cute::tma_store_arrive(); + } + __syncwarp(); + } + + // Notify L2 and increment L1 empty count + // TODO: less epilogue sync scope + ptx::tma_store_wait<0>(); + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_add_rel( + workspace.get_l2_full_count_ptr(ring_block_idx), 1u); + + // Increment L1 empty count for this physical slot (one per N block) + ptx::red_add( + workspace.get_l1_empty_count_ptr(ring_block_idx), 1u); + } + __syncwarp(); + } else { + // Increment L2 empty count for this physical slot (one per N block) + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_add( + workspace.get_l2_empty_count_ptr(ring_block_idx), 1u); + } + __syncwarp(); + + DG_STATIC_ASSERT(STORE_BLOCK_M % 8 == 0, "Invalid store M"); + constexpr uint32_t kNumRowsPerWarp = STORE_BLOCK_M / 8; + + // L2 BF16 epilogue: write GEMM output to remote combine buffer via NVLink + #pragma unroll + for (uint32_t s = 0; s < WG_BLOCK_M / STORE_BLOCK_M; ++ s) { + // Early break if the entire store block is beyond the valid token range + // TODO: check performance + if (epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M >= valid_m) { + ptx::tcgen05_before_thread_sync(); + shared_storage.tmem_empty_barriers[accum_stage_idx].arrive(0u); + break; + } + + #pragma unroll + for (uint32_t i = 0; i < STORE_BLOCK_M / ATOM_M; ++ i) { + // Load from TMEM using .16x256b shape to satisfy STSM layout requirements + // Start from lane index 0 and 16 + uint32_t tmem_addr = accum_stage_idx * UMMA_N + epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + i * ATOM_M; + uint32_t values[ATOM_M]; + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, + values[0], values[1], values[2], values[3]); + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr | 0x00100000, + values[4], values[5], values[6], values[7]); + cutlass::arch::fence_view_async_tmem_load(); + + // Wait shared memory release from previous NVLink store + // NOTES: skip for the first store block since the prior full barrier already ensures completion + if (i == 0 and s > 0) + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Signal tensor memory consumed + if (s == WG_BLOCK_M / STORE_BLOCK_M - 1 and i == STORE_BLOCK_M / ATOM_M - 1) { + ptx::tcgen05_before_thread_sync(); + shared_storage.tmem_empty_barriers[accum_stage_idx].arrive(0u); + } + + // Store into shared memory + // NOTES: each lane provides its own address for stmatrix; 2 warps share a BF16 swizzle atom + uint32_t row = lane_idx % 8; + uint32_t col = (epilogue_warp_idx % 2) * 4 + lane_idx / 8; + const auto smem_ptr = reinterpret_cast(shared_storage.smem_d.l2[epilogue_wg_idx]) + + (warp_idx_in_wg / 2) * STORE_BLOCK_M * kSwizzleCDMode + + i * ATOM_M * kSwizzleCDMode + + row * (kNumBankGroupBytes * 8) + + (col ^ row) * kNumBankGroupBytes; + ptx::SM90_U32x4_STSM_T::copy( + math::cast_into_bf16_and_pack(values[0], values[1]), + math::cast_into_bf16_and_pack(values[2], values[3]), + math::cast_into_bf16_and_pack(values[4], values[5]), + math::cast_into_bf16_and_pack(values[6], values[7]), + smem_ptr + ); + } + + // Wait shared memory ready + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Write into remote buffers + // Each warp writes 2 rows (lane_idx/16 splits the warp into two halves, one per row) + const uint32_t row_in_atom = (warp_idx_in_wg * 2 + lane_idx / 16) % ATOM_M; + const uint32_t bank_group_idx = lane_idx % 8; + + #pragma unroll + for (uint32_t j = 0; j < kNumRowsPerWarp; ++ j) { + const uint32_t row_in_store = j * 8 + warp_idx_in_wg * 2 + lane_idx / 16; + const uint32_t m_idx_in_block = epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + row_in_store; + + // Skip padding rows beyond the actual token count for this expert + if (m_idx_in_block >= valid_m) + break; + + const auto src_metadata = *workspace.get_token_src_metadata_ptr(pool_m_idx + m_idx_in_block); + const uint32_t dst_rank_idx = src_metadata.rank_idx; + const uint32_t dst_token_idx = src_metadata.token_idx; + const uint32_t dst_topk_idx = src_metadata.topk_idx; + + // Read from shared memory + const auto smem_ptr = reinterpret_cast(shared_storage.smem_d.l2[epilogue_wg_idx]) + + (lane_idx % 16 / 8) * STORE_BLOCK_M * kSwizzleCDMode + + row_in_store * kSwizzleCDMode + + (bank_group_idx ^ row_in_atom) * kNumBankGroupBytes; + const auto packed = ptx::ld_shared(reinterpret_cast(smem_ptr)); + + // Write into remote + const auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_idx) + .get_data_buffer(dst_token_idx); + const auto dst_ptr = math::advance_ptr( + dst_token.get_base_ptr(), + n_idx * static_cast(sizeof(nv_bfloat16)) + (lane_idx % 16) * static_cast(sizeof(float4))); + *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; + } + } + + // Ensure the next epilogue safe to use shared memory + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + }); + + // Deallocate tensor memory + // NOTES: must be called by the same logical warp ID on both CTAs + if (epilogue_warp_idx == 0) + Allocator().free(0, kNumTmemCols); + + // NVLink barrier (grid sync + cross-rank signal + grid sync): ~4 us + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + + // Barrier with dispatch warps, so that they can do clean workspace + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Combine: reduce top-k results and write back + // NOTES: reuse shared memory from start up to the barriers + // 1 token, 1 topk latency: ~3 us + constexpr uint32_t kNumHiddenBytes = kHidden * sizeof(nv_bfloat16); + constexpr uint32_t kNumElemsPerUint4 = sizeof(uint4) / sizeof(nv_bfloat162); + + // 3 slots of chunk is needed: 2 load stages and 1 store + constexpr uint32_t kNumChunkSlots = 3; + constexpr uint32_t kNumMaxRegistersForBuffer = 128; + + // NOTES: either 1 or 2 chunks for simplicity + // NOTES: Restrict on both smem and register + constexpr uint32_t kNumChunks = + kNumChunkSlots * kNumEpilogueWarps * kNumHiddenBytes <= kNumReusableSmemBytes and kHidden <= 32 * kNumMaxRegistersForBuffer ? 1 : 2; + constexpr uint32_t kNumChunkBytes = kNumHiddenBytes / kNumChunks; + constexpr uint32_t kNumChunkUint4 = kNumChunkBytes / sizeof(uint4); + constexpr uint32_t kNumUint4PerLane = kNumChunkUint4 / 32; + DG_STATIC_ASSERT(kHidden % kNumChunks == 0, "Hidden must be divisible by number of chunks"); + DG_STATIC_ASSERT(kNumChunkSlots * kNumEpilogueWarps * kNumHiddenBytes / kNumChunks <= kNumReusableSmemBytes, "Hidden is too large"); + DG_STATIC_ASSERT(kNumChunkBytes % 16 == 0, "Combine chunk must be TMA-aligned (16 bytes)"); + DG_STATIC_ASSERT(kNumChunkBytes % sizeof(uint4) == 0, "Combine chunk must be divisible by 16 bytes"); + DG_STATIC_ASSERT(kNumChunkUint4 % 32 == 0, "Combine chunk must be a multiple of 32 16-byte elements (one per lane)"); + DG_STATIC_ASSERT(kNumTopk <= 32, "Top-k must fit in a single warp"); + + // Verify combined shared memory budget at runtime + DG_DEVICE_ASSERT(kNumChunkSlots * kNumEpilogueWarps * kNumChunkBytes <= kNumReusableSmemBytes); + + // Per-warp buffer: 2 stage load buffers + 1 store buffer + const auto combine_load_buffer = utils::PatternVisitor([&](const uint32_t& i) { + return math::advance_ptr(smem_buffer, (epilogue_warp_idx + i * kNumEpilogueWarps) * kNumChunkBytes); + }); + const auto combine_store_buffer = math::advance_ptr(smem_buffer, (epilogue_warp_idx + kNumEpilogueWarps * 2) * kNumChunkBytes); + + // Per-warp barriers + auto combine_load_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return &shared_storage.combine_barriers[i + epilogue_warp_idx * 2]; + }); + + // Iterate over all tokens + uint32_t combine_phase = 0; + uint32_t load_stage_idx = 0; + for (uint32_t token_idx = sm_idx * kNumEpilogueWarps + epilogue_warp_idx; + token_idx < num_tokens; + token_idx += kNumSMs * kNumEpilogueWarps) { + // Read top-k slot indices: each lane reads one slot, then broadcast via exchange + DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + const int stored_topk_slot_idx = lane_idx < kNumTopk ? + static_cast(__ldg(input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : -1; + const uint32_t total_mask = __ballot_sync(0xffffffff, stored_topk_slot_idx >= 0); + + // Iterate all chunks + for (uint32_t chunk = 0; chunk < kNumChunks; ++ chunk) { + const uint32_t chunk_byte_offset = chunk * kNumChunkBytes; + + // Move mask and load + uint32_t mask = total_mask; + const auto move_mask_and_load = [&](const uint32_t& i) { + if (mask) { + // Move + const uint32_t slot_idx = __ffs(mask) - 1; + mask ^= 1 << slot_idx; + + // Load + if (cute::elect_one_sync()) { + const auto src_ptr = math::advance_ptr( + combine_token_buffer.get_rank_buffer(slot_idx) + .get_data_buffer(token_idx).get_base_ptr(), + chunk_byte_offset); + ptx::tma_load_1d(combine_load_buffer[i], src_ptr, combine_load_barriers[i], kNumChunkBytes); + ptx::mbarrier_arrive_and_set_tx(combine_load_barriers[i], kNumChunkBytes); + } + __syncwarp(); + return true; + } + return false; + }; + + // Load the first selection + bool do_reduce = move_mask_and_load(load_stage_idx); + + // Accumulate all top-k contributions for this chunk in float registers + float2 reduced[kNumUint4PerLane * kNumElemsPerUint4] = {}; + while (do_reduce) { + // Prefetch next top-k into the buffer while current is being accumulated + do_reduce = move_mask_and_load(load_stage_idx ^ 1); + + // Accumulate + combine_load_barriers[load_stage_idx]->wait(combine_phase); + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + const auto uint4_values = combine_load_buffer[load_stage_idx][j * 32 + lane_idx]; + const auto bf16_values = reinterpret_cast(&uint4_values); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + ptx::accumulate(reduced[j * kNumElemsPerUint4 + l], bf16_values[l]); + } + combine_phase ^= load_stage_idx; + load_stage_idx ^= 1; + } + + // Cast + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + uint4 casted; + auto casted_bf16 = reinterpret_cast(&casted); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + casted_bf16[l] = __float22bfloat162_rn(reduced[j * kNumElemsPerUint4 + l]); + + // Wait share memory release and write + if (j == 0) { + ptx::tma_store_wait<0>(); + __syncwarp(); + } + ptx::st_shared(combine_store_buffer + j * 32 + lane_idx, + casted.x, casted.y, casted.z, casted.w); + } + __syncwarp(); + + // TMA store the token chunk + if (cute::elect_one_sync()) { + cute::tma_store_fence(); + ptx::tma_store_1d( + math::advance_ptr(y, static_cast(token_idx) * kNumHiddenBytes + chunk_byte_offset), + combine_store_buffer, kNumChunkBytes); + cute::tma_store_arrive(); + } + __syncwarp(); + } + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only support sm_100f"); +#endif +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh b/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh index 528b3dd103..0df6f8c971 100644 --- a/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh +++ b/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh @@ -139,6 +139,30 @@ struct SM100_MMA_MXF4_SS { } }; +struct SM100_MMA_MXF4_2x1SM_SS { + CUTLASS_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t const& tmem_c, + uint32_t const& scale_c, + uint64_t const& desc, + uint32_t const& tmem_sfa, + uint32_t const& tmem_sfb) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" +#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9) + "tcgen05.mma.cta_group::2.kind::mxf4.block_scale.block32 [%0], %1, %2, %3, [%5], [%6], p; \n\t" +#else + "tcgen05.mma.cta_group::2.kind::mxf4.block_scale.scale_vec::2X [%0], %1, %2, %3, [%5], [%6], p; \n\t" +#endif + "}\n" + :: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast(desc >> 32)), "r"(scale_c), + "r"(tmem_sfa), "r"(tmem_sfb)); + } +}; + struct SM100_MMA_F16BF16_WS_SS { CUTLASS_DEVICE static void fma(uint64_t const& desc_a, diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index bb5de820cf..6970d52e9d 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -175,6 +175,30 @@ def fp8_fp4_mega_moe(y: torch.Tensor, sym_buffer.num_ring_tokens ) +def mxfp4_mxfp4_mega_moe(y: torch.Tensor, + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor], + sym_buffer: SymmBuffer, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + recipe: Tuple[int, int, int] = (1, 1, 32), + activation: str = 'swiglu', + activation_clamp: Optional[float] = None, + fast_math: bool = True): + # Packed MXFP4 x MXFP4 mega MoE (both activations and weights are E2M1 + UE8M0). + _C.mxfp4_mxfp4_mega_moe( + y, + l1_weights, l2_weights, + cumulative_local_expert_recv_stats, + sym_buffer.buffer, + sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), + sym_buffer.num_max_tokens_per_rank, + sym_buffer.num_experts, sym_buffer.num_topk, + recipe, + activation, activation_clamp, + fast_math, + sym_buffer.num_ring_tokens + ) + def bf16_mega_moe(y: torch.Tensor, l1_weights: torch.Tensor, l2_weights: torch.Tensor, diff --git a/tests/test_mxfp4_gemm.py b/tests/test_mxfp4_gemm.py new file mode 100644 index 0000000000..2c37185aa6 --- /dev/null +++ b/tests/test_mxfp4_gemm.py @@ -0,0 +1,48 @@ +import random +import torch + +import deep_gemm +from deep_gemm.testing import calc_diff +from deep_gemm.utils.layout import get_mn_major_tma_aligned_packed_ue8m0_tensor +from deep_gemm.utils.math import per_token_cast_to_fp4, cast_back_from_fp4 + + +def _prepare(x: torch.Tensor, gran_k: int = 32): + # Packed E2M1 (int8, K/2) + float UE8M0 SF (gran-K 32) + packed, sf_f = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=gran_k) + # Dequantized reference using the *same* packed values + SF + deq = cast_back_from_fp4(packed, sf_f, gran_k=gran_k) + # SF in the MN-major TMA-aligned packed UE8M0 layout the kernel expects + sf_packed = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf_f) + return packed, sf_packed, deq + + +def test_mxfp4_gemm() -> None: + print('Testing packed MXFP4 x MXFP4 GEMM:') + # NOTES: this de-risk kernel is hardcoded to a 2-CTA (cluster_n = 2) config, so it + # requires N divisible by 256, M divisible by 128, and K divisible by 128. + for m, n, k in ((256, 256, 256), (256, 256, 512), (512, 256, 1024), (128, 512, 256), (1024, 768, 512)): + a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16) + b = torch.randn((n, k), device='cuda', dtype=torch.bfloat16) + + a_packed, sfa, a_deq = _prepare(a) + b_packed, sfb, b_deq = _prepare(b) + + ref_d = (a_deq.float() @ b_deq.float().T).to(torch.bfloat16) + + d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16) + deep_gemm.mxfp4_gemm_nt((a_packed, sfa), (b_packed, sfb), d) + + diff = calc_diff(d, ref_d) + status = 'OK' if diff < 0.05 else 'FAIL' + print(f' > m={m:5}, n={n:5}, k={k:5}: diff={diff:.5f} [{status}]') + assert diff < 0.05, f'{m=}, {n=}, {k=}, diff={diff:.5f}' + print('All MXFP4 GEMM cases passed.\n') + + +if __name__ == '__main__': + torch.manual_seed(0) + random.seed(0) + print('Library path:') + print(f' > {deep_gemm.__path__}\n') + test_mxfp4_gemm() diff --git a/tests/test_mxfp4_mega_moe.py b/tests/test_mxfp4_mega_moe.py new file mode 100644 index 0000000000..b5028e8442 --- /dev/null +++ b/tests/test_mxfp4_mega_moe.py @@ -0,0 +1,96 @@ +import os +import torch +import torch.distributed as dist + +import deep_gemm +from deep_gemm.testing import calc_diff +from deep_gemm.utils.math import per_token_cast_to_fp4, cast_back_from_fp4 + + +def _fp4_roundtrip(x: torch.Tensor, gran_k: int = 32) -> torch.Tensor: + """Quantize bf16 -> packed MXFP4 (UE8M0, gran-32) and dequantize back to the + exact values the kernel operates on.""" + packed, sf = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=gran_k) + return cast_back_from_fp4(packed, sf, gran_k=gran_k).to(x.dtype) + + +def reference_mxfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, + activation_clamp: float = 10.0): + """Single-rank MXFP4 MoE reference matching the fused kernel's math: + FP4(x) @ FP4(W1).T -> SwiGLU*weight -> FP4 requant -> @ FP4(W2).T -> combine (sum over top-k).""" + num_tokens = x.shape[0] + x_deq = _fp4_roundtrip(x) + w1_deq = torch.stack([_fp4_roundtrip(l1w[e]) for e in range(l1w.shape[0])]) + w2_deq = torch.stack([_fp4_roundtrip(l2w[e]) for e in range(l2w.shape[0])]) + + y = torch.zeros((num_tokens, hidden), dtype=torch.float, device=x.device) + for t in range(num_tokens): + for k in range(topk_idx.shape[1]): + e = int(topk_idx[t, k].item()) + if e < 0: + continue + l1 = x_deq[t].float() @ w1_deq[e].float().T # [inter*2] + gate, up = l1[:inter], l1[inter:] + gate = gate.clamp(max=activation_clamp) + up = up.clamp(min=-activation_clamp, max=activation_clamp) + act = (gate * torch.sigmoid(gate)) * up * float(topk_weights[t, k].item()) + act_deq = _fp4_roundtrip(act.to(torch.bfloat16).unsqueeze(0)).squeeze(0) + l2 = act_deq.float() @ w2_deq[e].float().T # [hidden] + y[t] += l2 + return y.to(torch.bfloat16) + + +def test_mxfp4_mega_moe(): + os.environ.setdefault('MASTER_ADDR', '127.0.0.1') + os.environ.setdefault('MASTER_PORT', '12399') + dist.init_process_group('nccl', rank=0, world_size=1) + group = dist.group.WORLD + + num_max_tokens, num_tokens = 128, 128 + hidden, inter = 512, 512 + num_experts, num_topk = 8, 2 + ne_per_rank = num_experts + + buf = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, + hidden, inter, mma_type='mxfp4xmxfp4') + + def cast_w(w): + g, n, k = w.shape + wp = torch.empty((g, n, k // 2), device='cuda', dtype=torch.int8) + wsf = torch.empty((g, n, k // 32), device='cuda', dtype=torch.float) + for i in range(g): + wp[i], wsf[i] = per_token_cast_to_fp4(w[i], use_ue8m0=True, gran_k=32) + return wp, deep_gemm.transform_sf_into_required_layout(wsf, n, k, (1, 32), g) + + torch.manual_seed(0) + x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + l1w = torch.randn((ne_per_rank, inter * 2, hidden), dtype=torch.bfloat16, device='cuda') / (hidden ** 0.5) + l2w = torch.randn((ne_per_rank, hidden, inter), dtype=torch.bfloat16, device='cuda') / (inter ** 0.5) + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') + topk_weights, topk_idx = torch.topk(scores.softmax(-1), num_topk, dim=-1) + + xp, xsf = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + l1 = cast_w(l1w) + l2 = cast_w(l2w) + tl1, tl2 = deep_gemm.transform_weights_for_mega_moe(l1, l2) + + buf.x[:num_tokens].copy_(xp) + buf.x_sf[:num_tokens].copy_(xsf) + buf.topk_idx[:num_tokens].copy_(topk_idx) + buf.topk_weights[:num_tokens].copy_(topk_weights) + + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.mxfp4_mxfp4_mega_moe(y=y, l1_weights=tl1, l2_weights=tl2, sym_buffer=buf, + activation_clamp=10.0, fast_math=True) + torch.cuda.synchronize() + + ref = reference_mxfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, 10.0) + diff = calc_diff(y, ref) + print(f'diff = {diff:.5f} (y~{y.float().abs().mean():.3f}, ref~{ref.float().abs().mean():.3f})') + assert diff < 0.05, f'{diff=}' + print('MXFP4 mega MoE passed.') + dist.destroy_process_group() + + +if __name__ == '__main__': + test_mxfp4_mega_moe() From c6ade0b825e34d7ad048540fbb135e9574e16bf1 Mon Sep 17 00:00:00 2001 From: Ying Zhang Date: Sat, 27 Jun 2026 02:40:35 +0000 Subject: [PATCH 2/9] Add NVFP4 x NVFP4 support to packed-FP4 GEMM and mega-MoE (SM100) Extends the unified packed-FP4 kernels (standalone GEMM and mega-MoE) to also support NVFP4 (E2M1 data, E4M3/UE4M3 scale factors at gran-16, with per-tensor global scales), keyed on `MmaKind`. Adds `nvfp4_gemm_nt` and `nvfp4_nvfp4_mega_moe` entries (mma_type="nvfp4xnvfp4"), the NVF4 2-CTA `kind::mxf4nvf4` MMA atom, NVFP4 quant utils, and tests. Global scales are CPU-side scalar kernel params: the GEMM/MoE accumulators are dequantized by gs_act * gs_weight, and the MoE L1 output is requantized with the L2 activation global scale. Key gran-16 scale-factor handling (vs gran-32 MXFP4): 2 scale-factor int32s per K-block (stride/smem/tmem sized accordingly), and crossing K-uint32s via the SF tensor-memory ADDRESS since `a_sf_id_`/`b_sf_id_` are 2-bit descriptor fields (sf_id only selects within a K-uint32). Validated on B200: standalone GEMM diff=0.0; mega-MoE diff=8.4e-4; MXFP4 paths unchanged (regression intact). Co-authored-by: Cursor --- csrc/apis/gemm.hpp | 27 +++ csrc/apis/mega.hpp | 118 ++++++++++-- csrc/jit_kernels/heuristics/mega_moe.hpp | 4 +- csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp | 31 ++- .../impls/sm100_mxfp4_mxfp4_mega_moe.hpp | 38 +++- deep_gemm/__init__.py | 4 +- deep_gemm/include/deep_gemm/common/types.cuh | 11 +- .../epilogue/sm100_store_cd_swap_ab.cuh | 16 +- .../deep_gemm/impls/sm100_mxfp4_gemm.cuh | 105 +++++++---- .../impls/sm100_mxfp4_mxfp4_mega_moe.cuh | 178 +++++++++++++----- deep_gemm/include/deep_gemm/ptx/tcgen05.cuh | 24 +++ deep_gemm/mega/__init__.py | 31 +++ deep_gemm/utils/math.py | 40 ++++ tests/test_nvfp4_gemm.py | 64 +++++++ tests/test_nvfp4_mega_moe.py | 143 ++++++++++++++ 15 files changed, 713 insertions(+), 121 deletions(-) create mode 100644 tests/test_nvfp4_gemm.py create mode 100644 tests/test_nvfp4_mega_moe.py diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 1726007bbc..e59dfc75b8 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -185,6 +185,30 @@ static void mxfp4_gemm_nt(const std::pair& a, sm100_mxfp4_gemm(a_val, sfa, b_val, sfb, d, m, n, k); } +static void nvfp4_gemm_nt(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const float& a_global_scale, + const float& b_global_scale) { + // Standalone packed NVFP4 x NVFP4 GEMM: `[M, K] @ [N, K].T -> [M, N]`. + // A/B packed E2M1 stored as int8 `[*, K/2]`; SFs are int32-packed E4M3 (gran-16); + // per-tensor global scales (CPU scalars) dequant the output: D = (A@B.T) * gs_a * gs_b. + const auto& [a_val, sfa] = a; + const auto& [b_val, sfb] = b; + DG_HOST_ASSERT(a_val.dim() == 2 and b_val.dim() == 2 and d.dim() == 2); + DG_HOST_ASSERT(a_val.is_contiguous() and b_val.is_contiguous()); + + const int m = static_cast(a_val.size(0)); + const int k = static_cast(a_val.size(1)) * 2; // packed: 2 elements per byte + const int n = static_cast(b_val.size(0)); + DG_HOST_ASSERT(static_cast(b_val.size(1)) * 2 == k); + DG_HOST_ASSERT(static_cast(d.size(0)) == m and static_cast(d.size(1)) == n); + + const auto arch_major = device_runtime->get_arch_major(); + DG_HOST_ASSERT(arch_major == 10 and "NVFP4 GEMM requires SM100"); + sm100_mxfp4_gemm(a_val, sfa, b_val, sfb, d, m, n, k, MmaKind::NVFP4, a_global_scale, b_global_scale); +} + static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair& a, const std::pair& b, const torch::Tensor& d, @@ -694,6 +718,9 @@ static void register_apis(pybind11::module_& m) { py::arg("disable_ue8m0_cast") = false); m.def("mxfp4_gemm_nt", &mxfp4_gemm_nt, py::arg("a"), py::arg("b"), py::arg("d")); + m.def("nvfp4_gemm_nt", &nvfp4_gemm_nt, + py::arg("a"), py::arg("b"), py::arg("d"), + py::arg("a_global_scale"), py::arg("b_global_scale")); m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), py::arg("recipe") = std::nullopt, diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 443e542053..98de0884d6 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -47,8 +47,9 @@ get_symm_buffer_size_for_mega_moe( // Parse MMA type const auto mma_kind = parse_mma_kind(mma_type); - const auto num_mma_elem_bytes = get_num_mma_elem_bytes(mma_kind); const auto with_sf = is_mma_with_sf(mma_kind); + // SF granularity: MXFP8FP4/MXFP4 -> 32, NVFP4 -> 16 (2x as many scale factors) + const int gran_k = get_sf_gran_k(mma_kind); // NOTES: packed FP4 is sub-byte (0.5 B/elem), so derive token byte sizes from element BITS // (`get_element_bits/8`): BF16->2, FP8->1, packed FP4->0.5. `hidden` is a multiple of 128. const auto token_bytes = [&](const int& num_elems) { return num_elems * get_element_bits(mma_kind) / 8; }; @@ -61,8 +62,8 @@ get_symm_buffer_size_for_mega_moe( const auto input_token_layout = layout::Data(token_bytes(hidden)); const auto bf16_token_layout = layout::Data(hidden * 2); const auto intermediate_token_layout = layout::Data(token_bytes(intermediate_hidden)); - const auto input_sf_layout = layout::Data(with_sf ? hidden / 32 : 0); - const auto intermediate_sf_layout = layout::Data(with_sf ? intermediate_hidden / 32 : 0); + const auto input_sf_layout = layout::Data(with_sf ? hidden / gran_k : 0); + const auto intermediate_sf_layout = layout::Data(with_sf ? intermediate_hidden / gran_k : 0); const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false); const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false); const auto l1_topk_weights_layout = layout::Data(sizeof(float), false); @@ -120,12 +121,12 @@ get_symm_buffer_size_for_mega_moe( DG_HOST_ASSERT(num_sf_ring_tokens % 4 == 0); } - // Token view shape/dtype: packed MXFP4 stores 2 elems/byte (int8, half the inner dim), + // Token view shape/dtype: packed FP4 (MXFP4/NVFP4) stores 2 elems/byte (int8, half the inner dim), // FP8 uses 1 byte/elem (e4m3, full inner dim), BF16 uses bf16. - const bool is_mxfp4 = (mma_kind == MmaKind::MXFP4); - const auto token_dtype = with_sf ? (is_mxfp4 ? torch::kInt8 : torch::kFloat8_e4m3fn) : torch::kBFloat16; - const int x_token_inner = is_mxfp4 ? hidden / 2 : hidden; - const int l2_token_inner = is_mxfp4 ? intermediate_hidden / 2 : intermediate_hidden; + const bool is_packed_fp4 = (mma_kind == MmaKind::MXFP4 or mma_kind == MmaKind::NVFP4); + const auto token_dtype = with_sf ? (is_packed_fp4 ? torch::kInt8 : torch::kFloat8_e4m3fn) : torch::kBFloat16; + const int x_token_inner = is_packed_fp4 ? hidden / 2 : hidden; + const int l2_token_inner = is_packed_fp4 ? intermediate_hidden / 2 : intermediate_hidden; // Slice function: creates `(x, x_sf, topk_weights, topk_idx, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf)` tensor views from the raw buffer // NOTES: `x_sf` is K-major, while `l1_acts_sf` and `l2_acts_sf` are M-major @@ -136,7 +137,7 @@ get_symm_buffer_size_for_mega_moe( torch::TensorOptions().dtype(token_dtype).device(buffer.device())); auto x_sf = with_sf ? torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_sf_buffer.base)), - {num_max_tokens_per_rank, hidden / 128}, + {num_max_tokens_per_rank, hidden / (gran_k * 4)}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); auto topk_idx = torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_topk_idx_buffer.base)), @@ -152,7 +153,7 @@ get_symm_buffer_size_for_mega_moe( torch::TensorOptions().dtype(token_dtype).device(buffer.device())); auto l1_acts_sf = with_sf ? torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_sf_buffer.base)), - {num_sf_ring_tokens, hidden / 128}, + {num_sf_ring_tokens, hidden / (gran_k * 4)}, {1, num_sf_ring_tokens}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); auto l2_acts = torch::from_blob( @@ -161,7 +162,7 @@ get_symm_buffer_size_for_mega_moe( torch::TensorOptions().dtype(token_dtype).device(buffer.device())); auto l2_acts_sf = with_sf ? torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l2_sf_buffer.base)), - {num_sf_ring_tokens, intermediate_hidden / 128}, + {num_sf_ring_tokens, intermediate_hidden / (gran_k * 4)}, {1, num_sf_ring_tokens}, torch::TensorOptions().dtype(torch::kInt).device(buffer.device())) : torch::Tensor(); return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf); @@ -352,6 +353,100 @@ static void mxfp4_mxfp4_mega_moe( sym_buffer.zero_(); } +static void nvfp4_nvfp4_mega_moe( + const torch::Tensor& y, + const std::tuple& l1_weights_tuple, + const std::tuple& l2_weights_tuple, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const std::vector& sym_buffer_ptrs, const int& rank_idx, + const int& num_max_tokens_per_rank, + const int& num_experts, const int& num_topk, + // Per-tensor global scales (CPU scalars): activations (L1 input, L2 input) and weights (L1, L2). + const float& l1_act_global_scale, const float& l2_act_global_scale, + const float& l1_weight_global_scale, const float& l2_weight_global_scale, + const std::tuple& recipe, + const std::string& activation, + const std::optional& activation_clamp_opt, + const bool& fast_math, + const int& num_ring_tokens +) { + // Packed NVFP4 x NVFP4 mega MoE (E2M1 data, E4M3 SF gran-16, per-tensor global scales). + const auto [l1_weights, l1_weights_sf] = l1_weights_tuple; + const auto [l2_weights, l2_weights_sf] = l2_weights_tuple; + + const auto num_tokens = static_cast(y.size(0)); + const auto [rm, rn, rk] = recipe; + DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 16); + DG_HOST_ASSERT(activation == "swiglu"); + + const auto activation_clamp = + activation_clamp_opt.value_or(std::numeric_limits::infinity()); + DG_HOST_ASSERT(activation_clamp >= 0); + + // Tensor checks (weights are packed FP4, same layout as the fp8xfp4 path) + DG_HOST_ASSERT(get_major_type_ab(l1_weights) == cute::UMMA::Major::K); + DG_HOST_ASSERT(get_major_type_ab(l2_weights) == cute::UMMA::Major::K); + const auto arch_major = device_runtime->get_arch_major(); + const auto [num_experts_per_rank, intermediate_hidden_2, hidden] = + check_grouped_ab_fp8_fp4(l1_weights, cute::UMMA::Major::K, arch_major); + const auto [num_experts_per_rank_, hidden_, intermediate_hidden] = + check_grouped_ab_fp8_fp4(l2_weights, cute::UMMA::Major::K, arch_major); + DG_HOST_ASSERT(num_tokens <= num_max_tokens_per_rank); + DG_HOST_ASSERT(num_experts_per_rank == num_experts_per_rank_); + DG_HOST_ASSERT(hidden == hidden_); + DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden); + DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous()); + + // Weight SF: E4M3 bytes packed 4-per-int32 (same container as MXFP4's UE8M0), MN-major, TMA-aligned, gran-16 + constexpr int kGranMN = 1, kGranK = 16; + check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK, + num_experts_per_rank, true, false, torch::kInt); + check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK, + num_experts_per_rank, true, false, torch::kInt); + + if (cumulative_local_expert_recv_stats.has_value()) { + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt); + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->numel() == num_experts_per_rank); + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->is_contiguous()); + } + + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + const auto num_experts_ = num_experts_per_rank * num_ranks; + const auto [num_required_bytes, slice] = get_symm_buffer_size_for_mega_moe( + num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + "nvfp4xnvfp4", activation, num_ring_tokens); + DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); + DG_HOST_ASSERT(num_experts == num_experts_); + + const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + + if (arch_major == 10) { + sm100_mxfp4_mxfp4_mega_moe(y, + l1_acts, l1_acts_sf, + l2_acts, l2_acts_sf, + l1_weights, l2_weights, + l1_weights_sf, l2_weights_sf, + cumulative_local_expert_recv_stats, + sym_buffer_ptrs, + rank_idx, num_max_tokens_per_rank, + num_experts_per_rank, + num_tokens, num_topk, + hidden, intermediate_hidden, + activation_clamp, fast_math, + MmaKind::NVFP4, + l1_act_global_scale, l1_weight_global_scale, + l2_act_global_scale, l2_weight_global_scale); + } else { + DG_HOST_UNREACHABLE("Unsupported architecture"); + } + + if (get_env("DG_COMM_KERNEL_DEBUG")) + sym_buffer.zero_(); +} + static void bf16_mega_moe( const torch::Tensor& y, const torch::Tensor& l1_weights, @@ -439,6 +534,7 @@ static void register_apis(pybind11::module_& m) { m.def("get_symm_buffer_size_for_mega_moe", &get_symm_buffer_size_for_mega_moe); m.def("fp8_fp4_mega_moe", &fp8_fp4_mega_moe); m.def("mxfp4_mxfp4_mega_moe", &mxfp4_mxfp4_mega_moe); + m.def("nvfp4_nvfp4_mega_moe", &nvfp4_nvfp4_mega_moe); m.def("bf16_mega_moe", &bf16_mega_moe); #endif } diff --git a/csrc/jit_kernels/heuristics/mega_moe.hpp b/csrc/jit_kernels/heuristics/mega_moe.hpp index 8ca46d4cdf..d8994168a5 100644 --- a/csrc/jit_kernels/heuristics/mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/mega_moe.hpp @@ -67,6 +67,8 @@ static MmaKind parse_mma_kind(const std::string& mma_type_str) { return MmaKind::BF16; if (mma_type_str == "mxfp4xmxfp4") return MmaKind::MXFP4; + if (mma_type_str == "nvfp4xnvfp4") + return MmaKind::NVFP4; DG_HOST_ASSERT(mma_type_str == "fp8xfp4"); return MmaKind::MXFP8FP4; } @@ -78,7 +80,7 @@ static int get_num_mma_elem_bytes(const MmaKind& mma_kind) { } static bool is_mma_with_sf(const MmaKind& mma_kind) { - return mma_kind == MmaKind::MXFP8FP4 or mma_kind == MmaKind::MXFP4; + return mma_kind == MmaKind::MXFP8FP4 or mma_kind == MmaKind::MXFP4 or mma_kind == MmaKind::NVFP4; } static int get_num_wave_pool_tokens( diff --git a/csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp b/csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp index 709dbf4f72..2cbfeb7585 100644 --- a/csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_mxfp4_gemm.hpp @@ -24,6 +24,8 @@ class SM100MXFP4GemmRuntime final: public LaunchRuntime { int num_stages; int num_non_epilogue_threads, num_epilogue_threads; int num_sms; + bool is_nvfp4; + float ab_global_scale; // NVFP4 output dequant scale (gs_a * gs_b); 1.0 for MXFP4 CUtensorMap tensor_map_a; CUtensorMap tensor_map_sfa; @@ -46,6 +48,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, {}, + {}, {} >); }}; @@ -54,7 +57,8 @@ static void __instantiate_kernel() {{ args.block_m, args.block_n, args.block_k, args.num_stages, args.num_non_epilogue_threads, args.num_epilogue_threads, - args.num_sms); + args.num_sms, + args.is_nvfp4 ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -62,19 +66,25 @@ static void __instantiate_kernel() {{ static_cast(args.m), static_cast(args.n), static_cast(args.k), args.tensor_map_a, args.tensor_map_sfa, args.tensor_map_b, args.tensor_map_sfb, - args.tensor_map_cd)); + args.tensor_map_cd, + args.ab_global_scale)); } }; static void sm100_mxfp4_gemm(const torch::Tensor& a, const torch::Tensor& sfa, const torch::Tensor& b, const torch::Tensor& sfb, const torch::Tensor& d, - const int& m, const int& n, const int& k) { + const int& m, const int& n, const int& k, + const MmaKind& mma_kind = MmaKind::MXFP4, + const float& a_global_scale = 1.0f, + const float& b_global_scale = 1.0f) { // Fixed de-risk configuration constexpr int block_m = 128, block_n = 128, block_k = 128; constexpr int num_stages = 4; constexpr int num_non_epilogue_threads = 128, num_epilogue_threads = 128; - constexpr int gran_k = 32; + const bool is_nvfp4 = (mma_kind == MmaKind::NVFP4); + const int gran_k = get_sf_gran_k(mma_kind); // 32 (mxfp4) / 16 (nvfp4) + const int num_sf_k_per_load = block_k / (gran_k * 4); // int32s per token per K-block (1 / 2) // Packed FP4 tensors are stored as int8 (`kPackedFP4`), 2 elements per byte DG_HOST_ASSERT(a.scalar_type() == kPackedFP4 and b.scalar_type() == kPackedFP4); @@ -102,9 +112,10 @@ static void sm100_mxfp4_gemm(const torch::Tensor& a, const torch::Tensor& sfa, const auto tensor_map_b = make_tma_2d_desc(b, k, n, block_k, load_block_n, static_cast(b.stride(0)), swizzle_ab, 0, false, false); - // UE8M0 SF descriptors (MN-major, no swizzle) - const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, block_m, gran_k, 1, 0); - const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, block_n, gran_k, 1, 0); + // SF descriptors (MN-major, no swizzle): UE8M0 gran-32 (mxfp4) / E4M3 gran-16 (nvfp4). + // `smem_outer_dim = num_sf_k_per_load` loads all K-uint32s of a K-block (2 for gran-16). + const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, block_m, gran_k, 1, 0, 0, false, num_sf_k_per_load); + const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, block_n, gran_k, 1, 0, 0, false, num_sf_k_per_load); // BF16 output (N-major) const auto tensor_map_cd = make_tma_cd_desc(d, m, n, store_block_m, store_block_n, static_cast(d.stride(-2)), 1, swizzle_cd); @@ -115,8 +126,8 @@ static void sm100_mxfp4_gemm(const torch::Tensor& a, const torch::Tensor& sfa, const int smem_cd = store_block_m * store_block_n * static_cast(sizeof(int16_t)) * num_tma_store_stages; const int smem_a = load_block_m * (block_k / 2); const int smem_b = load_block_n * (block_k / 2); - const int smem_sfa = sf_block_m * static_cast(sizeof(int)); - const int smem_sfb = sf_block_n * static_cast(sizeof(int)); + const int smem_sfa = sf_block_m * num_sf_k_per_load * static_cast(sizeof(int)); + const int smem_sfb = sf_block_n * num_sf_k_per_load * static_cast(sizeof(int)); const int smem_barriers = (num_stages * 3 + num_epilogue_stages * 2) * 8; const int smem_size = smem_cd + num_stages * (smem_a + smem_b) @@ -131,6 +142,8 @@ static void sm100_mxfp4_gemm(const torch::Tensor& a, const torch::Tensor& sfa, .num_non_epilogue_threads = num_non_epilogue_threads, .num_epilogue_threads = num_epilogue_threads, .num_sms = num_sms, + .is_nvfp4 = is_nvfp4, + .ab_global_scale = a_global_scale * b_global_scale, .tensor_map_a = tensor_map_a, .tensor_map_sfa = tensor_map_sfa, .tensor_map_b = tensor_map_b, diff --git a/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp index 17c064dbb2..298e058703 100644 --- a/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp @@ -22,10 +22,11 @@ static MegaMoEConfig get_mxfp4_mega_moe_config( const int& num_ranks, const int& num_experts, const int& num_experts_per_rank, const int& num_max_tokens_per_rank, const int& num_tokens, const int& num_topk, const int& hidden, const int& intermediate_hidden, - const int& num_ring_tokens, const int& num_sf_ring_tokens) { + const int& num_ring_tokens, const int& num_sf_ring_tokens, + const MmaKind& mma_kind) { const auto [cluster_size, block_m, store_block_m, block_k, num_epilogue_threads] = - get_block_config_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens, MmaKind::MXFP4); + get_block_config_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens, mma_kind); const int block_n = 128; const int load_block_m = block_m / 2; const int load_block_n = block_n; @@ -35,7 +36,7 @@ static MegaMoEConfig get_mxfp4_mega_moe_config( // Packed FP4 K-major swizzle == K extent in bytes (block_k / 2) const int swizzle_acts_mode = block_k / 2; const int swizzle_weights_mode = block_k / 2; - constexpr int gran_k = 32; + const int gran_k = get_sf_gran_k(mma_kind); const int num_sms = device_runtime->get_num_sms(); const int num_experts_per_wave = get_num_experts_per_wave_for_mega_moe( @@ -110,6 +111,10 @@ class SM100MXFP4MegaMoERuntime final : public LaunchRuntime); }}; @@ -169,7 +175,8 @@ static void __instantiate_kernel() {{ args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads, args.launch_args.grid_dim.first, args.num_ranks, to_string(args.activation_clamp), - args.fast_math ? "true" : "false"); + args.fast_math ? "true" : "false", + args.is_nvfp4 ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -186,7 +193,10 @@ static void __instantiate_kernel() {{ args.tensor_map_l2_acts, args.tensor_map_l2_acts_sf, args.tensor_map_l2_weights, - args.tensor_map_l2_weights_sf + args.tensor_map_l2_weights_sf, + args.l1_acc_scale, + args.l2_act_global_scale, + args.l2_acc_scale )); } }; @@ -204,19 +214,25 @@ static void sm100_mxfp4_mxfp4_mega_moe( const int& num_tokens, const int& num_topk, const int& hidden, const int& intermediate_hidden, const float& activation_clamp, - const bool& fast_math + const bool& fast_math, + // MmaKind::MXFP4 or MmaKind::NVFP4. NVFP4 also needs the per-tensor global scales below. + const MmaKind& mma_kind = MmaKind::MXFP4, + // NVFP4 global scales (CPU scalars). Derived dequant: L1 acc -> real, L1-output requant, L2 acc -> real. + const float& l1_act_gs = 1.0f, const float& l1_weight_gs = 1.0f, + const float& l2_act_gs = 1.0f, const float& l2_weight_gs = 1.0f ) { const auto num_ranks = static_cast(sym_buffer_ptrs.size()); const auto num_experts = num_experts_per_rank * num_ranks; const auto num_ring_tokens = static_cast(l1_acts.size(0)); const auto num_sf_ring_tokens = static_cast(l1_acts_sf.size(0)); + const bool is_nvfp4 = (mma_kind == MmaKind::NVFP4); const auto config = get_mxfp4_mega_moe_config( num_ranks, num_experts, num_experts_per_rank, num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, - num_ring_tokens, num_sf_ring_tokens); + num_ring_tokens, num_sf_ring_tokens, mma_kind); - constexpr int kGranK = 32; + const int kGranK = get_sf_gran_k(mma_kind); const int sf_smem_outer_dim = config.block_k / (kGranK * 4); // Packed FP4 token/weight TMA descriptors (fp4_unpacked_smem = false -> 16U4_ALIGN8B) @@ -277,6 +293,12 @@ static void sm100_mxfp4_mxfp4_mega_moe( .num_ranks = num_ranks, .activation_clamp = activation_clamp, .fast_math = fast_math, + .is_nvfp4 = is_nvfp4, + // L1 acc -> real = gs_l1_act * gs_l1_weight; L1-output requant uses gs_l2_act; + // L2 acc -> real = gs_l2_act * gs_l2_weight. (All 1.0 for MXFP4.) + .l1_acc_scale = l1_act_gs * l1_weight_gs, + .l2_act_global_scale = l2_act_gs, + .l2_acc_scale = l2_act_gs * l2_weight_gs, .config = config, .y = y.data_ptr(), .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index e06aca4797..adede057ea 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -34,8 +34,9 @@ try: # DeepGEMM Kernels from ._C import ( - # Packed MXFP4 x MXFP4 GEMM (standalone de-risk) + # Packed MXFP4 / NVFP4 GEMM (standalone de-risk) mxfp4_gemm_nt, + nvfp4_gemm_nt, # FP8 FP4 GEMMs fp8_fp4_gemm_nt, fp8_fp4_gemm_nn, fp8_fp4_gemm_tn, fp8_fp4_gemm_tt, @@ -90,6 +91,7 @@ fp8_fp4_mega_moe, bf16_mega_moe, mxfp4_mxfp4_mega_moe, + nvfp4_nvfp4_mega_moe, ) # Some utils diff --git a/deep_gemm/include/deep_gemm/common/types.cuh b/deep_gemm/include/deep_gemm/common/types.cuh index 216e9fdb07..0d810a25cc 100644 --- a/deep_gemm/include/deep_gemm/common/types.cuh +++ b/deep_gemm/include/deep_gemm/common/types.cuh @@ -7,8 +7,10 @@ namespace deep_gemm { enum class MmaKind { BF16 = 0, MXFP8FP4 = 1, - // Pure packed MXFP4 x MXFP4 (E2M1 data, UE8M0 SF, 2 elements per byte) + // Pure packed MXFP4 x MXFP4 (E2M1 data, UE8M0 SF gran-32, 2 elements per byte) MXFP4 = 2, + // Pure packed NVFP4 x NVFP4 (E2M1 data, E4M3 SF gran-16, per-tensor global scale) + NVFP4 = 3, }; // NOTES: element size in *bits*, since packed FP4 is sub-byte (4 bits) @@ -17,10 +19,16 @@ constexpr CUTLASS_HOST_DEVICE int get_element_bits(const MmaKind& mma_kind) { case MmaKind::BF16: return 16; case MmaKind::MXFP8FP4: return 8; case MmaKind::MXFP4: return 4; + case MmaKind::NVFP4: return 4; default: return 0; } } +// SF granularity (elements per scale factor): MXFP4 -> 32, NVFP4 -> 16 +constexpr CUTLASS_HOST_DEVICE int get_sf_gran_k(const MmaKind& mma_kind) { + return mma_kind == MmaKind::NVFP4 ? 16 : 32; +} + constexpr CUTLASS_HOST_DEVICE int get_element_size(const MmaKind& mma_kind) { switch (mma_kind) { case MmaKind::BF16: return 2; @@ -28,6 +36,7 @@ constexpr CUTLASS_HOST_DEVICE int get_element_size(const MmaKind& mma_kind) { // NOTES: packed FP4 is 0.5 byte/elem; callers must use byte math that // divides element counts by 2 (see `get_element_bits`) case MmaKind::MXFP4: return 1; + case MmaKind::NVFP4: return 1; default: return 0; } } diff --git a/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh b/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh index f3f5351e6a..6e8fcfebcb 100644 --- a/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh +++ b/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh @@ -26,7 +26,9 @@ sm100_store_cd_swap_ab(const utils::PatternVisitor& smem_cd, uint3 const uint32_t& effective_m, const uint32_t& epilogue_warp_idx, const uint32_t& lane_idx, const cutlass::arch::ClusterTransactionBarrier* tmem_empty_barrier, - const cute::TmaDescriptor& tensor_map_cd) { + const cute::TmaDescriptor& tensor_map_cd, + // Optional per-tensor output scale (NVFP4 global-scale dequant); 1.0 = no-op. + const float& out_scale = 1.0f) { // NOTES: The epilogue requires a full warpgroup to read all 128 TMEM rows, // implying STORE_BLOCK_N must be 128. DG_STATIC_ASSERT(STORE_BLOCK_N == 128, "STORE_BLOCK_N must be 128 to match TMEM rows"); @@ -73,6 +75,11 @@ sm100_store_cd_swap_ab(const utils::PatternVisitor& smem_cd, uint3 // NOTES: Swizzling is not required in this case, but used here for consistency with other cases cute::SM100_TMEM_LOAD_32dp32b8x::copy(tmem_addr, values[0], values[1], values[2], values[3], values[4], values[5], values[6], values[7]); + if (out_scale != 1.0f) { + #pragma unroll + for (uint32_t vi = 0; vi < kNumSwizzleAtomRows; ++ vi) + values[vi] = __float_as_uint(__uint_as_float(values[vi]) * out_scale); + } uint32_t col = lane_idx / 4; #pragma unroll @@ -92,6 +99,13 @@ sm100_store_cd_swap_ab(const utils::PatternVisitor& smem_cd, uint3 values[4], values[5], values[6], values[7]); cutlass::arch::fence_view_async_tmem_load(); + // NVFP4: dequant the MMA accumulator by the per-tensor global scale before BF16 cast + if (out_scale != 1.0f) { + #pragma unroll + for (uint32_t vi = 0; vi < 8; ++ vi) + values[vi] = __float_as_uint(__uint_as_float(values[vi]) * out_scale); + } + // Destination shared memory address uint32_t row = lane_idx % 8; uint32_t col = (epilogue_warp_idx % 2) * 4 + lane_idx / 8; diff --git a/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_gemm.cuh b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_gemm.cuh index f069aa4ca9..8dd5fa1585 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_gemm.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_gemm.cuh @@ -34,14 +34,18 @@ template + uint32_t kNumSMs, + // NVFP4 (E4M3/UE4M3 SF gran-16 + per-tensor global scale) vs MXFP4 (UE8M0 SF gran-32) + bool kIsNVFP4 = false> CUTLASS_GLOBAL void __launch_bounds__(kNumNonEpilogueThreads + kNumEpilogueThreads, 1) sm100_mxfp4_gemm_impl(uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, const __grid_constant__ cute::TmaDescriptor tensor_map_a, // acts, [M, K] K-major (packed E2M1) - const __grid_constant__ cute::TmaDescriptor tensor_map_sfa, // acts SF, UE8M0 (int32-packed, gran-32) + const __grid_constant__ cute::TmaDescriptor tensor_map_sfa, // acts SF (int32-packed): UE8M0 gran-32 (mxfp4) / E4M3 gran-16 (nvfp4) const __grid_constant__ cute::TmaDescriptor tensor_map_b, // weights, [N, K] K-major (packed E2M1) - const __grid_constant__ cute::TmaDescriptor tensor_map_sfb, // weights SF, UE8M0 - const __grid_constant__ cute::TmaDescriptor tensor_map_cd) { // out, [M, N] BF16 + const __grid_constant__ cute::TmaDescriptor tensor_map_sfb, // weights SF + const __grid_constant__ cute::TmaDescriptor tensor_map_cd, // out, [M, N] BF16 + // NVFP4 output dequant scale = gs_a * gs_b (CPU scalar; 1.0 for MXFP4) + const float ab_global_scale = 1.0f) { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; using Allocator = cute::TMEM::Allocator2Sm; @@ -75,14 +79,14 @@ sm100_mxfp4_gemm_impl(uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, constexpr uint32_t kSwizzleABMode = BLOCK_K_BYTES; // 64 for BLOCK_K = 128 constexpr uint32_t kSwizzleCDMode = 128; - // SF configs (UE8M0, gran-K 32, UTCCP 128-aligned) - constexpr uint32_t kGranK = 32; + // SF configs: MXFP4 -> UE8M0 gran-32, NVFP4 -> E4M3/UE4M3 gran-16. UTCCP 128-aligned. + constexpr uint32_t kGranK = kIsNVFP4 ? 16 : 32; constexpr uint32_t kNumUTCCPAlignedElems = 128; constexpr uint32_t SF_BLOCK_M = math::constexpr_align(BLOCK_M, kNumUTCCPAlignedElems); constexpr uint32_t SF_BLOCK_N = math::constexpr_align(BLOCK_N, kNumUTCCPAlignedElems); - // One int32 packs 4 gran-32 SFs along K, i.e. covers 128 K + // One int32 packs 4 SFs along K: gran-32 -> covers 128 K (1/load), gran-16 -> covers 64 K (2/load) constexpr uint32_t kNumSFKPerLoad = BLOCK_K / (kGranK * 4); - DG_STATIC_ASSERT(kNumSFKPerLoad == 1, "BLOCK_K must be 128 for a single packed SF int per load"); + DG_STATIC_ASSERT(kNumSFKPerLoad == 1 or kNumSFKPerLoad == 2, "Invalid packed SF int count per load"); // Epilogue configs (swap-AB) constexpr uint32_t kNumEpilogueStages = 2; @@ -97,13 +101,17 @@ sm100_mxfp4_gemm_impl(uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_SIZE_PER_STAGE * kNumTMAStoreStages; constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K_BYTES; constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K_BYTES; - constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * sizeof(uint32_t); - constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = SF_BLOCK_N * sizeof(uint32_t); + constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * kNumSFKPerLoad * sizeof(uint32_t); + constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = SF_BLOCK_N * kNumSFKPerLoad * sizeof(uint32_t); - // Tensor memory size and offsets + // Tensor memory size and offsets. Each K-uint32 occupies (SF_BLOCK/32) cols (the MN rows); + // gran-16 has 2 K-uint32s per K-block. The 2-bit `sf_id` selects within a K-uint32; crossing + // K-uint32s is done via the SF tmem ADDRESS. constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumEpilogueStages; - constexpr uint32_t kNumSFATmemCols = SF_BLOCK_M / 32; - constexpr uint32_t kNumSFBTmemCols = SF_BLOCK_N / 32; + constexpr uint32_t kSFAColsPerKUint = SF_BLOCK_M / 32; + constexpr uint32_t kSFBColsPerKUint = SF_BLOCK_N / 32; + constexpr uint32_t kNumSFATmemCols = kSFAColsPerKUint * kNumSFKPerLoad; + constexpr uint32_t kNumSFBTmemCols = kSFBColsPerKUint * kNumSFKPerLoad; constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols; constexpr uint32_t kTmemStartColOfSFB = kNumAccumTmemCols + kNumSFATmemCols; @@ -228,16 +236,17 @@ sm100_mxfp4_gemm_impl(uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, // SF: per-CTA SM90 load into THIS CTA's smem (the 2-CTA UTCCP reads both CTAs' // SF), signaling THIS CTA's own `full` barrier (NOT routed to the leader). So both // CTAs' transposers can wait their own `full` and the cross-CTA `with_sf` sync holds. + // K-uint32 stride per K-block: gran-32 -> 1, gran-16 -> 2 uint32_t sfa_m_idx = m_block_idx * BLOCK_M; - uint32_t sfa_k_idx = scheduler.template get_global_idx(shape_sf_k, 1, k_block_idx); + uint32_t sfa_k_idx = scheduler.template get_global_idx(shape_sf_k, kNumSFKPerLoad, k_block_idx); tma::copy(&tensor_map_sfa, full_barriers[stage_idx], smem_sfa[stage_idx], sfa_m_idx, sfa_k_idx); uint32_t sfb_n_idx = n_block_idx * BLOCK_N; - uint32_t sfb_k_idx = scheduler.template get_global_idx(shape_sf_k, 1, k_block_idx, m_block_idx); + uint32_t sfb_k_idx = scheduler.template get_global_idx(shape_sf_k, kNumSFKPerLoad, k_block_idx, m_block_idx); tma::copy(&tensor_map_sfb, full_barriers[stage_idx], smem_sfb[stage_idx], sfb_n_idx, sfb_k_idx); // Expect: the leader collects BOTH CTAs' data (2-SM routed) + its own SF; the // non-leader's `full` only sees its own SF (its data tx went to the leader). - const auto sf_bytes = BLOCK_M * sizeof(uint32_t) + BLOCK_N * sizeof(uint32_t); + const auto sf_bytes = (BLOCK_M + BLOCK_N) * kNumSFKPerLoad * sizeof(uint32_t); if (is_leader_cta) full_barriers[stage_idx]->arrive_and_expect_tx( SMEM_A_SIZE_PER_STAGE * kNumMulticast + SMEM_B_SIZE_PER_STAGE * kNumMulticast + sf_bytes); @@ -248,8 +257,9 @@ sm100_mxfp4_gemm_impl(uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, } else if (warp_idx == 1 and is_leader_cta) { // MMA issue warp (leader CTA only) // Swap-AB: weights -> UMMA "A" (UMMA_M), acts -> UMMA "B" (UMMA_N) + using sf_dtype_t = cute::conditional_t; auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< - ab_dtype_t, ab_dtype_t, float, cutlass::float_ue8m0_t, + ab_dtype_t, ab_dtype_t, float, sf_dtype_t, UMMA_M, UMMA_N, kMajorB, kMajorA>(); DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); @@ -289,24 +299,32 @@ sm100_mxfp4_gemm_impl(uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, ptx::tcgen05_after_thread_sync(); if (cute::elect_one_sync()) { - // UTCCP copy SFA / SFB into TMEM (transposed in warp 3 already) + // UTCCP copy SFA / SFB into TMEM (transposed in warp 2 already). + // gran-16 has `kNumSFKPerLoad` K-uint32s per K-block, each at kSF*ColsPerKUint apart. using cute_utccp_t = cute::SM100_UTCCP_4x32dp128bit_2cta; #pragma unroll - for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { - auto sf_desc = mma::sm100::make_sf_desc(smem_sfa[stage_idx] + i * kNumUTCCPAlignedElems); - cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + i * 4); - } - #pragma unroll - for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { - auto sf_desc = mma::sm100::make_sf_desc(smem_sfb[stage_idx] + i * kNumUTCCPAlignedElems); - cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + i * 4); + for (uint32_t ku = 0; ku < kNumSFKPerLoad; ++ ku) { + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { + auto sf_desc = mma::sm100::make_sf_desc(smem_sfa[stage_idx] + ku * SF_BLOCK_M + i * kNumUTCCPAlignedElems); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + ku * kSFAColsPerKUint + i * 4); + } + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { + auto sf_desc = mma::sm100::make_sf_desc(smem_sfb[stage_idx] + ku * SF_BLOCK_N + i * kNumUTCCPAlignedElems); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + ku * kSFBColsPerKUint + i * 4); + } } - // Issue UMMA over UMMA_K (=64) sub-tiles - // VALIDATE: descriptor /2 byte addressing + sf_id = k*2 + // Issue UMMA over UMMA_K (=64) sub-tiles. The 2-bit `sf_id` selects WITHIN a + // K-uint32; crossing K-uint32s (gran-16) uses the SF tmem ADDRESS. #pragma unroll for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++ k) { - const uint32_t sf_id = k * 2; + const uint32_t global_sf_idx = k * (UMMA_K / kGranK); + const uint32_t sf_kuint = global_sf_idx / 4; + const uint32_t sf_id = global_sf_idx % 4; + const uint32_t tmem_sfa = kTmemStartColOfSFA + sf_kuint * kSFAColsPerKUint; + const uint32_t tmem_sfb = kTmemStartColOfSFB + sf_kuint * kSFBColsPerKUint; const auto runtime_instr_desc = mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, sf_id, sf_id); auto a_desc = mma::sm100::make_smem_desc( kFP4Layout, reinterpret_cast(smem_b[stage_idx]) + k * UMMA_K_BYTES, @@ -315,10 +333,16 @@ sm100_mxfp4_gemm_impl(uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, kFP4Layout, reinterpret_cast(smem_a[stage_idx]) + k * UMMA_K_BYTES, 8 * kSwizzleABMode, 0); // Swap-AB: weights (b_desc) first, SFB id first - ptx::SM100_MMA_MXF4_2x1SM_SS::fma( - a_desc, b_desc, accum_stage_idx * UMMA_N, - k_block_idx > 0 or k > 0, runtime_instr_desc, - kTmemStartColOfSFB, kTmemStartColOfSFA); + if constexpr (kIsNVFP4) + ptx::SM100_MMA_NVF4_2x1SM_SS::fma( + a_desc, b_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or k > 0, runtime_instr_desc, + tmem_sfb, tmem_sfa); + else + ptx::SM100_MMA_MXF4_2x1SM_SS::fma( + a_desc, b_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or k > 0, runtime_instr_desc, + tmem_sfb, tmem_sfa); } } __syncwarp(); @@ -349,12 +373,16 @@ sm100_mxfp4_gemm_impl(uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { full_barriers[stage_idx]->wait(phase); #pragma unroll - for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) - utccp_required_smem_warp_transpose(smem_sfa[stage_idx] + i * kNumUTCCPAlignedElems); + for (uint32_t ku = 0; ku < kNumSFKPerLoad; ++ ku) + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) + utccp_required_smem_warp_transpose(smem_sfa[stage_idx] + ku * SF_BLOCK_M + i * kNumUTCCPAlignedElems); cutlass::arch::fence_view_async_shared(); #pragma unroll - for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) - utccp_required_smem_warp_transpose(smem_sfb[stage_idx] + i * kNumUTCCPAlignedElems); + for (uint32_t ku = 0; ku < kNumSFKPerLoad; ++ ku) + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) + utccp_required_smem_warp_transpose(smem_sfb[stage_idx] + ku * SF_BLOCK_N + i * kNumUTCCPAlignedElems); cutlass::arch::fence_view_async_shared(); with_sf_full_barriers[stage_idx]->arrive(0u); } @@ -384,7 +412,8 @@ sm100_mxfp4_gemm_impl(uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, (smem_cd, tma_stage_idx, tmem_base_addr, base_m_idx, base_n_idx, scheduler.current_group_idx, effective_m, epilogue_warp_idx, lane_idx, - tmem_empty_barriers[accum_stage_idx], tensor_map_cd); + tmem_empty_barriers[accum_stage_idx], tensor_map_cd, + ab_global_scale); } } diff --git a/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh index e13c3e8a96..60c2a1f802 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh @@ -35,6 +35,8 @@ template < uint32_t kNumSMs, uint32_t kNumRanks, float kActivationClamp, bool kFastMath, + // NVFP4 (E4M3 SF gran-16 + per-tensor global scale) vs MXFP4 (UE8M0 SF gran-32) + bool kIsNVFP4 = false, uint32_t L1_SHAPE_N = kIntermediateHidden * 2, uint32_t L1_SHAPE_K = kHidden, uint32_t L2_SHAPE_N = kHidden, @@ -61,7 +63,14 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, - const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights_sf) { + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights_sf, + // NVFP4 global-scale dequant params (CPU-side scalars; ignored for MXFP4). + // l1_acc_scale = gs_l1_act * gs_l1_weight (L1 MMA acc -> real) + // l2_act_global_scale = gs_l2_act (L1-output requant) + // l2_acc_scale = gs_l2_act * gs_l2_weight (L2 MMA acc -> real) + const float l1_acc_scale = 1.0f, + const float l2_act_global_scale = 1.0f, + const float l2_acc_scale = 1.0f) { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; using Allocator = cute::TMEM::Allocator2Sm; @@ -102,8 +111,9 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, constexpr auto fp8_token_layout = layout::Data(kHidden / 2); constexpr auto bf16_token_layout = layout::Data(kHidden * sizeof(nv_bfloat16)); constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden / 2); - constexpr auto fp8_sf_layout = layout::Data(kHidden / 32); - constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 32); + constexpr uint32_t kGranK = kIsNVFP4 ? 16 : 32; + constexpr auto fp8_sf_layout = layout::Data(kHidden / kGranK); + constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / kGranK); constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * sizeof(int64_t), false); constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * sizeof(float), false); constexpr auto l1_topk_weights_layout = layout::Data(sizeof(float), false); @@ -123,7 +133,6 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, input_topk_idx_buffer.get_end_ptr()); // SF and its buffer configs - constexpr uint32_t kGranK = 32; constexpr uint32_t kNumUTCCPAlignedElems = 128; DG_STATIC_ASSERT(SF_BLOCK_M == math::constexpr_align(BLOCK_M, kNumUTCCPAlignedElems), "Invalid SF_BLOCK_M"); DG_STATIC_ASSERT(SF_BLOCK_N == BLOCK_N, "No padding is needed for SFB"); @@ -183,6 +192,10 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, constexpr uint32_t BLOCK_K_BYTES = BLOCK_K / 2; constexpr uint32_t UMMA_BLOCK_K_BYTES = UMMA_BLOCK_K / 2; constexpr uint32_t UMMA_K_BYTES = UMMA_K / 2; + // SF count along K per token, in int32 units (4 SF bytes each): gran-32 -> BLOCK_K/128, gran-16 -> BLOCK_K/64. + constexpr uint32_t kNumSFKUint32 = BLOCK_K / (kGranK * 4); + // SF int32s per token per UMMA_BLOCK_K (=128): gran-32 -> 1, gran-16 -> 2. + constexpr uint32_t kNumSFKUint32PerUmmaBlock = UMMA_BLOCK_K / (kGranK * 4); // Swizzle configs // NOTES: K-major packed FP4 uses a swizzle equal to the K extent in bytes (BLOCK_K/2) @@ -217,8 +230,8 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, // Packed-FP4 A/B tiles: byte-addressed (`BLOCK_K / 2` bytes per row) alignas(kSharedMemoryAlignment) uint8_t smem_a[kNumStages][LOAD_BLOCK_M * BLOCK_K_BYTES]; alignas(kSharedMemoryAlignment) uint8_t smem_b[kNumStages][LOAD_BLOCK_N * BLOCK_K_BYTES]; - uint32_t smem_sfa[kNumStages][SF_BLOCK_M * (BLOCK_K / 128)]; - uint32_t smem_sfb[kNumStages][SF_BLOCK_N * (BLOCK_K / 128)]; + uint32_t smem_sfa[kNumStages][SF_BLOCK_M * kNumSFKUint32]; + uint32_t smem_sfb[kNumStages][SF_BLOCK_N * kNumSFKUint32]; float2 amax_reduction[kNumEpilogueWarps][AMAX_REDUCTION_WARP_BUFFER_SIZE]; Barrier dispatch_barriers[kNumDispatchWarps]; Barrier full_barriers[kNumStages]; @@ -243,8 +256,13 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, // Tensor memory size constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumEpilogueStages; - constexpr uint32_t kNumSFATmemCols = SF_BLOCK_M / 32; - constexpr uint32_t kNumSFBTmemCols = SF_BLOCK_N / 32; + // Each K-uint32 of SF occupies (SF_BLOCK/32) tmem cols (the MN rows; cell = 4 SF bytes). + // The 2-bit `sf_id` field selects the SF WITHIN a K-uint32 (<=3); crossing K-uint32s + // (gran-16 needs 2/K-block) is done via the SF tmem ADDRESS, not sf_id. + constexpr uint32_t kSFAColsPerKUint = SF_BLOCK_M / 32; + constexpr uint32_t kSFBColsPerKUint = SF_BLOCK_N / 32; + constexpr uint32_t kNumSFATmemCols = kSFAColsPerKUint * kNumSFKUint32; + constexpr uint32_t kNumSFBTmemCols = kSFBColsPerKUint * kNumSFKUint32; constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols; constexpr uint32_t kTmemStartColOfSFB = kNumAccumTmemCols + kNumSFATmemCols; @@ -578,7 +596,7 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, __syncwarp(); // Load and store SF (overlaps with last chunk's TMA load from remote) - constexpr uint32_t kNumSFUint32 = kHidden / 128; + constexpr uint32_t kNumSFUint32 = kHidden / (kGranK * 4); DG_STATIC_ASSERT(kNumSFUint32 > 0 and kHidden % 128 == 0, "Invalid SF"); const auto remote_sf_ptr = sym_buffer.map( input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), @@ -717,7 +735,8 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, uint32_t ring_m_idx = ring_block_idx * BLOCK_M; uint32_t k_idx = k_block_idx * BLOCK_K; uint32_t sfa_ring_m_idx = ring_block_idx * SF_BLOCK_M; - uint32_t sfa_k_idx = k_block_idx * (BLOCK_K / 128); + // K-uint32 stride per K-block: gran-32 -> BLOCK_K/128, gran-16 -> BLOCK_K/64 + uint32_t sfa_k_idx = k_block_idx * kNumSFKUint32; // Add 2 CTA offsets for non-leader CTA if (not is_leader_cta) @@ -777,7 +796,7 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, uint32_t n_idx = local_expert_idx * shape_n + n_block_idx * BLOCK_N; uint32_t k_idx = k_block_idx * BLOCK_K; uint32_t sfb_n_idx = n_block_idx * BLOCK_N; - uint32_t sfb_k_idx = local_expert_idx * shape_sfb_k + k_block_idx * (BLOCK_K / 128); + uint32_t sfb_k_idx = local_expert_idx * shape_sfb_k + k_block_idx * kNumSFKUint32; // TMA copy weights (TRUE 2-CTA cta_group::2; both CTAs use the same n coord so each // loads the full BLOCK_N replica) + SFB (per-CTA SM90). Data tx routes to the leader. @@ -806,8 +825,10 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, if (is_leader_cta) { // Make instruction descriptor with block scaling // NOTES: always swap A/B + // NVFP4 uses UE4M3 block scale factors; MXFP4 uses UE8M0 + using sf_dtype_t = cute::conditional_t; auto instr_desc = cute::UMMA::make_instr_desc_block_scaled< - b_dtype_t, a_dtype_t, float, cutlass::float_ue8m0_t, + b_dtype_t, a_dtype_t, float, sf_dtype_t, UMMA_M, UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K >(); @@ -861,19 +882,26 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, if (cute::elect_one_sync()) { #pragma unroll for (uint32_t umma_k_block_idx = 0; umma_k_block_idx < BLOCK_K / UMMA_BLOCK_K; ++ umma_k_block_idx) { - // UTCCP copy SFA and SFB to TMEM (SF pre-transposed during dispatch / weight transform) + // UTCCP copy SFA and SFB to TMEM (SF pre-transposed during dispatch / weight transform). + // NVFP4 (gran-16) packs 2x the K-scale-factors per token (kNumSFKUint32PerUmmaBlock), + // laid out [K-uint][M] in smem and [M-group][K-uint][4 cols] in tmem. using cute_utccp_t = cute::SM100_UTCCP_4x32dp128bit_2cta; #pragma unroll - for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { - auto smem_ptr = shared_storage.smem_sfa[stage_idx] + umma_k_block_idx * SF_BLOCK_M + i * kNumUTCCPAlignedElems; - mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); - cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + i * 4); - } - #pragma unroll - for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { - auto smem_ptr = shared_storage.smem_sfb[stage_idx] + umma_k_block_idx * SF_BLOCK_N + i * kNumUTCCPAlignedElems; - mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); - cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + i * 4); + for (uint32_t ku = 0; ku < kNumSFKUint32PerUmmaBlock; ++ ku) { + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { + auto smem_ptr = shared_storage.smem_sfa[stage_idx] + + (umma_k_block_idx * kNumSFKUint32PerUmmaBlock + ku) * SF_BLOCK_M + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + (umma_k_block_idx * kNumSFKUint32PerUmmaBlock + ku) * kSFAColsPerKUint + i * 4); + } + #pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { + auto smem_ptr = shared_storage.smem_sfb[stage_idx] + + (umma_k_block_idx * kNumSFKUint32PerUmmaBlock + ku) * SF_BLOCK_N + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + (umma_k_block_idx * kNumSFKUint32PerUmmaBlock + ku) * kSFBColsPerKUint + i * 4); + } } // Issue packed-mxf4 UMMA over UMMA_K (=64) sub-tiles (byte-addressed smem) @@ -881,16 +909,29 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, const auto b_stage_ptr = reinterpret_cast(shared_storage.smem_b[stage_idx]) + umma_k_block_idx * UMMA_BLOCK_K_BYTES; #pragma unroll for (uint32_t k = 0; k < UMMA_BLOCK_K / UMMA_K; ++ k) { - const uint32_t sf_id = (umma_k_block_idx * (UMMA_BLOCK_K / UMMA_K) + k) * 2; + // 4 SFs pack into one TMEM word (one K-uint32). The 2-bit `sf_id` + // selects WITHIN that K-uint32 (<=3); crossing K-uint32s (gran-16 needs + // 2 per K-block) is done via the SF tmem ADDRESS, not sf_id. + const uint32_t global_sf_idx = (umma_k_block_idx * (UMMA_BLOCK_K / UMMA_K) + k) * (UMMA_K / kGranK); + const uint32_t sf_kuint = global_sf_idx / 4; + const uint32_t sf_id = global_sf_idx % 4; + const uint32_t tmem_sfa = kTmemStartColOfSFA + sf_kuint * kSFAColsPerKUint; + const uint32_t tmem_sfb = kTmemStartColOfSFB + sf_kuint * kSFBColsPerKUint; const auto runtime_instr_desc = mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, sf_id, sf_id); auto a_desc = mma::sm100::make_smem_desc(kFP4Layout, a_stage_ptr + k * UMMA_K_BYTES, 8 * kSwizzleAMode, 0); auto b_desc = mma::sm100::make_smem_desc(kFP4Layout, b_stage_ptr + k * UMMA_K_BYTES, 8 * kSwizzleBMode, 0); // Swap-AB: weights (b_desc) take the UMMA "A" slot + SFB - ptx::SM100_MMA_MXF4_2x1SM_SS::fma( - b_desc, a_desc, accum_stage_idx * UMMA_N, - k_block_idx > 0 or umma_k_block_idx > 0 or k > 0, runtime_instr_desc, - kTmemStartColOfSFB, kTmemStartColOfSFA); + if constexpr (kIsNVFP4) + ptx::SM100_MMA_NVF4_2x1SM_SS::fma( + b_desc, a_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or umma_k_block_idx > 0 or k > 0, runtime_instr_desc, + tmem_sfb, tmem_sfa); + else + ptx::SM100_MMA_MXF4_2x1SM_SS::fma( + b_desc, a_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or umma_k_block_idx > 0 or k > 0, runtime_instr_desc, + tmem_sfb, tmem_sfa); } } } @@ -1039,6 +1080,13 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, // Apply SwiGLU: silu(gate) * up auto fp32_values = reinterpret_cast(raw_values); + // NVFP4: dequant the raw MMA accumulator by the L1 global scale BEFORE the + // (nonlinear) SwiGLU. MXFP4's UE8M0 block SFs already fully dequant the accumulator. + if constexpr (kIsNVFP4) { + #pragma unroll + for (uint32_t kk = 0; kk < 4; ++ kk) + fp32_values[kk] = __fmul2_rn(fp32_values[kk], {l1_acc_scale, l1_acc_scale}); + } #pragma unroll for (uint32_t k = 0; k < 2; ++ k) { auto bf16_gate = __float22bfloat162_rn(fp32_values[k * 2 + 0]); @@ -1080,9 +1128,12 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, amax_values[i].y = math::warp_reduce<4, true>( thread_local_amax.y, math::ReduceMax()); - // Reduce amax (warp-pair-level) - if (lane_idx < 4) - shared_storage.amax_reduction[epilogue_warp_idx][i * (ATOM_M / 2) + lane_idx] = amax_values[i]; + // Reduce amax (warp-pair-level) -- only MXFP4 (gran-32 SF spans 2 warps). + // NVFP4 (gran-16) has one SF per warp, so each warp keeps its own amax. + if constexpr (not kIsNVFP4) { + if (lane_idx < 4) + shared_storage.amax_reduction[epilogue_warp_idx][i * (ATOM_M / 2) + lane_idx] = amax_values[i]; + } __syncwarp(); } @@ -1095,20 +1146,37 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, // Cast to FP8 E4M3 and store into shared memory #pragma unroll for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) { - // Reduce amax (warp-pair-level) - const float2 wp_amax = - shared_storage.amax_reduction[epilogue_warp_idx ^ 1][i * (ATOM_M / 2) + lane_idx % 4]; - amax_values[i].x = cute::max(amax_values[i].x, wp_amax.x); - amax_values[i].y = cute::max(amax_values[i].y, wp_amax.y); - - // Calculate UE8M0 SF (MXFP4: E2M1 max is 6.0, so scale by amax/6) - float2 sf, sf_inv; - { + // Reduce amax (warp-pair-level), MXFP4 only + if constexpr (not kIsNVFP4) { + const float2 wp_amax = + shared_storage.amax_reduction[epilogue_warp_idx ^ 1][i * (ATOM_M / 2) + lane_idx % 4]; + amax_values[i].x = cute::max(amax_values[i].x, wp_amax.x); + amax_values[i].y = cute::max(amax_values[i].y, wp_amax.y); + } + + // Output block SF + the inverse used to scale activations before the E2M1 cast. + // MXFP4: UE8M0 (pow2 of amax/6). NVFP4: E4M3 of amax/(6*gs), code = act/(sf_e4m3*gs). + float2 sf_inv; + uint8_t sf_byte_x, sf_byte_y; + if constexpr (kIsNVFP4) { + const float inv_gs6 = 1.0f / (6.0f * l2_act_global_scale); + const __nv_fp8_e4m3 e4x(amax_values[i].x * inv_gs6); + const __nv_fp8_e4m3 e4y(amax_values[i].y * inv_gs6); + sf_byte_x = e4x.__x; sf_byte_y = e4y.__x; + const float sx = static_cast(e4x) * l2_act_global_scale; + const float sy = static_cast(e4y) * l2_act_global_scale; + sf_inv.x = sx > 0.f ? 1.0f / sx : 0.f; + sf_inv.y = sy > 0.f ? 1.0f / sy : 0.f; + } else { const float2 scaled = __fmul2_rn(amax_values[i], {1.0f / 6.0f, 1.0f / 6.0f}); const int ex = math::fast_log2_ceil(scaled.x); const int ey = math::fast_log2_ceil(scaled.y); - sf.x = math::fast_pow2(ex); sf_inv.x = math::fast_pow2(-ex); - sf.y = math::fast_pow2(ey); sf_inv.y = math::fast_pow2(-ey); + sf_inv.x = math::fast_pow2(-ex); + sf_inv.y = math::fast_pow2(-ey); + float2 sf; + sf.x = math::fast_pow2(ex); sf.y = math::fast_pow2(ey); + sf_byte_x = static_cast(*reinterpret_cast(&sf.x) >> 23); + sf_byte_y = static_cast(*reinterpret_cast(&sf.y) >> 23); } // Cast 2 activation outputs -> e2m1 nibbles, then write them DIRECTLY to their @@ -1155,11 +1223,14 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, p[(tok0 + 1) * L1_OUT_BLOCK_N_BYTES + j1] = ((my_nibs >> 24) & 0xF) | (((pr_nibs >> 24) & 0xF) << 4); } - // Store SF to `l2_sf_buffer` as UE8M0 (MN-major layout) - // Only one warp per pair writes (both hold the same SF after cross-warp reduce) - // Each lane < 4 holds SF for 2 rows (sf.x and sf.y) - if (warp_idx_in_wg % 2 == 0 and lane_idx < 4) { - const uint32_t k_idx = n_block_idx * 2 + warp_idx_in_wg / 2; + // Store SF to `l2_sf_buffer` (MN-major): MXFP4 UE8M0 byte, NVFP4 E4M3 byte. + // MXFP4 (gran-32): one SF per warp-pair (2 SFs per n-block, even warp writes). + // NVFP4 (gran-16): one SF per warp (4 SFs per n-block, every warp writes). + const bool sf_writer = kIsNVFP4 ? (lane_idx < 4) + : (warp_idx_in_wg % 2 == 0 and lane_idx < 4); + if (sf_writer) { + const uint32_t k_idx = kIsNVFP4 ? (n_block_idx * 4 + warp_idx_in_wg) + : (n_block_idx * 2 + warp_idx_in_wg / 2); const uint32_t k_uint_idx = k_idx / 4, byte_idx = k_idx % 4; const uint32_t mn_stride = kNumSFRingTokens * sizeof(uint32_t); const auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); @@ -1176,10 +1247,8 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, const auto sf_ring_token_idx = ring_block_idx * SF_BLOCK_M + transform_sf_token_idx(token_base_idx) + (lane_idx * 2) * 4; const auto sf_addr = k_uint_idx * mn_stride + sf_ring_token_idx * static_cast(sizeof(uint32_t)) + byte_idx; - sf_base_ptr[sf_addr] = - (*reinterpret_cast(&sf.x) >> 23); - sf_base_ptr[sf_addr + 4 * static_cast(sizeof(uint32_t))] = - (*reinterpret_cast(&sf.y) >> 23); + sf_base_ptr[sf_addr] = sf_byte_x; + sf_base_ptr[sf_addr + 4 * static_cast(sizeof(uint32_t))] = sf_byte_y; } __syncwarp(); } @@ -1246,6 +1315,13 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, values[4], values[5], values[6], values[7]); cutlass::arch::fence_view_async_tmem_load(); + // NVFP4: dequant the L2 MMA accumulator by the L2 global scale before BF16 cast + if constexpr (kIsNVFP4) { + #pragma unroll + for (uint32_t v = 0; v < ATOM_M; ++ v) + values[v] = __float_as_uint(__uint_as_float(values[v]) * l2_acc_scale); + } + // Wait shared memory release from previous NVLink store // NOTES: skip for the first store block since the prior full barrier already ensures completion if (i == 0 and s > 0) diff --git a/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh b/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh index 0df6f8c971..799d801a56 100644 --- a/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh +++ b/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh @@ -163,6 +163,30 @@ struct SM100_MMA_MXF4_2x1SM_SS { } }; +struct SM100_MMA_NVF4_2x1SM_SS { + CUTLASS_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t const& tmem_c, + uint32_t const& scale_c, + uint64_t const& desc, + uint32_t const& tmem_sfa, + uint32_t const& tmem_sfb) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" +#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9) + "tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%5], [%6], p; \n\t" +#else + "tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X [%0], %1, %2, %3, [%5], [%6], p; \n\t" +#endif + "}\n" + :: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast(desc >> 32)), "r"(scale_c), + "r"(tmem_sfa), "r"(tmem_sfb)); + } +}; + struct SM100_MMA_F16BF16_WS_SS { CUTLASS_DEVICE static void fma(uint64_t const& desc_a, diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 6970d52e9d..37e1976a7a 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -199,6 +199,37 @@ def mxfp4_mxfp4_mega_moe(y: torch.Tensor, sym_buffer.num_ring_tokens ) +def nvfp4_nvfp4_mega_moe(y: torch.Tensor, + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor], + sym_buffer: SymmBuffer, + l1_act_global_scale: float, + l2_act_global_scale: float, + l1_weight_global_scale: float, + l2_weight_global_scale: float, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + recipe: Tuple[int, int, int] = (1, 1, 16), + activation: str = 'swiglu', + activation_clamp: Optional[float] = None, + fast_math: bool = True): + # Packed NVFP4 x NVFP4 mega MoE (E2M1 data, E4M3 SF gran-16, per-tensor global scales). + # Global scales are CPU-side scalar kernel params (dequant: acc * gs_act * gs_weight). + _C.nvfp4_nvfp4_mega_moe( + y, + l1_weights, l2_weights, + cumulative_local_expert_recv_stats, + sym_buffer.buffer, + sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), + sym_buffer.num_max_tokens_per_rank, + sym_buffer.num_experts, sym_buffer.num_topk, + float(l1_act_global_scale), float(l2_act_global_scale), + float(l1_weight_global_scale), float(l2_weight_global_scale), + recipe, + activation, activation_clamp, + fast_math, + sym_buffer.num_ring_tokens + ) + def bf16_mega_moe(y: torch.Tensor, l1_weights: torch.Tensor, l2_weights: torch.Tensor, diff --git a/deep_gemm/utils/math.py b/deep_gemm/utils/math.py index fa2f31a0de..97d623b478 100644 --- a/deep_gemm/utils/math.py +++ b/deep_gemm/utils/math.py @@ -111,6 +111,46 @@ def per_token_cast_to_fp4(x: torch.Tensor, use_ue8m0: bool, gran_k: int = 128, return packed[:, :n // 2].contiguous(), sf +def nvfp4_global_scale(x: torch.Tensor) -> float: + # Per-tensor global scale s.t. per-block E4M3 SFs (amax_block / (6*gs)) stay <= E4M3 max (448). + amax = x.abs().float().amax().clamp_min(1e-6) + return float(amax / (6.0 * 448.0)) + + +def per_token_cast_to_nvfp4(x: torch.Tensor, global_scale: float, + gran_k: int = 16) -> Tuple[torch.Tensor, torch.Tensor]: + # NVFP4: packed E2M1 codes + per-(gran_k)-block E4M3 scale factors, with a per-tensor global scale. + # Dequant: x ~= code_e2m1 * sf_e4m3 * global_scale. + m, n = x.shape + assert n % 2 == 0 + padded_n = align(n, gran_k) + x_padded = torch.zeros((m, padded_n), dtype=x.dtype, device=x.device) + x_padded[:, :n] = x + x_view = x_padded.view(m, -1, gran_k) + amax = x_view.abs().float().amax(dim=2).clamp_min(1e-6) + block_sf_e4m3 = (amax / (6.0 * global_scale)).to(torch.float8_e4m3fn) + block_sf = block_sf_e4m3.float().clamp_min(1e-9) + x_scaled = x_view / (block_sf.unsqueeze(2) * global_scale) + codes = _quantize_to_fp4_e2m1(x_scaled).view(m, padded_n) + codes2 = codes.view(m, padded_n // 2, 2) + packed = (codes2[:, :, 0] & 0x0F) | ((codes2[:, :, 1] & 0x0F) << 4) + sf_bytes = block_sf_e4m3.view(torch.uint8) # [m, n // gran_k] + return packed[:, :n // 2].contiguous(), sf_bytes + + +def cast_back_from_nvfp4(packed: torch.Tensor, sf_bytes: torch.Tensor, global_scale: float, + gran_k: int = 16) -> torch.Tensor: + m, n2 = packed.shape + n = n2 * 2 + unpacked = torch.zeros((m, n), dtype=torch.int8, device=packed.device) + unpacked[:, ::2] = packed & 0x0F + unpacked[:, 1::2] = (packed >> 4) & 0x0F + x_deq = _dequantize_from_fp4_e2m1(unpacked) + sf = sf_bytes.view(torch.float8_e4m3fn).float() * global_scale + group_idx = torch.arange(n, device=packed.device) // gran_k + return x_deq * sf[:, group_idx] + + def transpose_packed_fp4(a: torch.Tensor) -> torch.Tensor: assert a.dtype == torch.int8 assert a.dim() == 2 diff --git a/tests/test_nvfp4_gemm.py b/tests/test_nvfp4_gemm.py new file mode 100644 index 0000000000..20c0766a55 --- /dev/null +++ b/tests/test_nvfp4_gemm.py @@ -0,0 +1,64 @@ +import random +import torch + +import deep_gemm +from deep_gemm.testing import calc_diff +from deep_gemm.utils.math import ( + align, per_token_cast_to_nvfp4, cast_back_from_nvfp4, nvfp4_global_scale, +) +from deep_gemm.utils.layout import get_tma_aligned_size + +GRAN_K = 16 + + +def _mn_major_packed_e4m3(sf_bytes: torch.Tensor) -> torch.Tensor: + """[mn, k_sf] uint8 (E4M3 bytes) -> MN-major TMA-aligned int32 (4 SF bytes/int32), + matching `get_mn_major_tma_aligned_packed_ue8m0_tensor` but for raw E4M3 bytes.""" + mn, k = sf_bytes.shape + aligned_mn = get_tma_aligned_size(mn, 4) + aligned_k = align(k, 4) + padded = torch.zeros((aligned_mn, aligned_k), dtype=torch.uint8, device=sf_bytes.device) + padded[:mn, :k] = sf_bytes + padded = padded.reshape(-1).view(torch.int32).view(aligned_mn, aligned_k // 4) + out = torch.empty_strided((aligned_mn, aligned_k // 4), (1, aligned_mn), + dtype=torch.int32, device=sf_bytes.device) + return out.copy_(padded)[:mn] + + +def _prepare(x: torch.Tensor, gs: float): + packed, sf_bytes = per_token_cast_to_nvfp4(x, gs, gran_k=GRAN_K) + deq = cast_back_from_nvfp4(packed, sf_bytes, gs, gran_k=GRAN_K) + sf_packed = _mn_major_packed_e4m3(sf_bytes) + return packed, sf_packed, deq + + +def test_nvfp4_gemm() -> None: + print('Testing packed NVFP4 x NVFP4 GEMM:') + for m, n, k in ((256, 256, 256), (256, 256, 512), (512, 256, 1024), (128, 512, 256), (1024, 768, 512)): + a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16) + b = torch.randn((n, k), device='cuda', dtype=torch.bfloat16) + gs_a = nvfp4_global_scale(a) + gs_b = nvfp4_global_scale(b) + + a_packed, sfa, a_deq = _prepare(a, gs_a) + b_packed, sfb, b_deq = _prepare(b, gs_b) + + ref_d = (a_deq.float() @ b_deq.float().T).to(torch.bfloat16) + + d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16) + deep_gemm.nvfp4_gemm_nt((a_packed, sfa), (b_packed, sfb), d, + a_global_scale=gs_a, b_global_scale=gs_b) + + diff = calc_diff(d, ref_d) + status = 'OK' if diff < 0.05 else 'FAIL' + print(f' > m={m:5}, n={n:5}, k={k:5}: diff={diff:.5f} [{status}]') + assert diff < 0.05, f'{m=}, {n=}, {k=}, diff={diff:.5f}' + print('All NVFP4 GEMM cases passed.\n') + + +if __name__ == '__main__': + torch.manual_seed(0) + random.seed(0) + print('Library path:') + print(f' > {deep_gemm.__path__}\n') + test_nvfp4_gemm() diff --git a/tests/test_nvfp4_mega_moe.py b/tests/test_nvfp4_mega_moe.py new file mode 100644 index 0000000000..07c96826d7 --- /dev/null +++ b/tests/test_nvfp4_mega_moe.py @@ -0,0 +1,143 @@ +import os +import torch +import torch.distributed as dist + +import deep_gemm +from deep_gemm.testing import calc_diff +from deep_gemm.utils.math import ( + align, per_token_cast_to_nvfp4, cast_back_from_nvfp4, nvfp4_global_scale, +) +from deep_gemm.utils.layout import get_tma_aligned_size + +GRAN_K = 16 + + +def _nvfp4_roundtrip(x: torch.Tensor, gs: float) -> torch.Tensor: + """Quantize bf16 -> NVFP4 (E4M3 SF gran-16 + global scale) and dequantize back to the + exact values the kernel operates on.""" + packed, sf = per_token_cast_to_nvfp4(x, gs, gran_k=GRAN_K) + return cast_back_from_nvfp4(packed, sf, gs, gran_k=GRAN_K).to(x.dtype) + + +def _mn_major_packed_e4m3(sf_bytes: torch.Tensor) -> torch.Tensor: + """Replicate `get_mn_major_tma_aligned_packed_ue8m0_tensor_torch` for raw E4M3 bytes: + pad to TMA-aligned MN and K%4, pack 4 SF bytes -> int32, transpose to MN-major.""" + g, mn, k = sf_bytes.shape + aligned_mn = get_tma_aligned_size(mn, 4) + aligned_k = align(k, 4) + padded = torch.zeros((g, aligned_mn, aligned_k), dtype=torch.uint8, device=sf_bytes.device) + padded[:, :mn, :k] = sf_bytes + padded = padded.reshape(-1).view(torch.int32).view(g, aligned_mn, aligned_k // 4) + out = torch.empty_strided((g, aligned_mn, aligned_k // 4), + (aligned_mn * (aligned_k // 4), 1, aligned_mn), + dtype=torch.int32, device=sf_bytes.device) + out = out.copy_(padded)[:, :mn] + return out + + +def _cast_w_nvfp4(w: torch.Tensor, gs: float): + g, n, k = w.shape + wp = torch.empty((g, n, k // 2), device='cuda', dtype=torch.int8) + wsf = torch.empty((g, n, k // GRAN_K), device='cuda', dtype=torch.uint8) + for i in range(g): + wp[i], wsf[i] = per_token_cast_to_nvfp4(w[i], gs, gran_k=GRAN_K) + return wp, _mn_major_packed_e4m3(wsf) + + +def reference_nvfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, + gs_x, gs_w1, gs_w2, gs_l2, activation_clamp: float = 10.0): + num_tokens = x.shape[0] + x_deq = _nvfp4_roundtrip(x, gs_x) + w1_deq = torch.stack([_nvfp4_roundtrip(l1w[e], gs_w1) for e in range(l1w.shape[0])]) + w2_deq = torch.stack([_nvfp4_roundtrip(l2w[e], gs_w2) for e in range(l2w.shape[0])]) + + y = torch.zeros((num_tokens, hidden), dtype=torch.float, device=x.device) + for t in range(num_tokens): + for k in range(topk_idx.shape[1]): + e = int(topk_idx[t, k].item()) + if e < 0: + continue + l1 = x_deq[t].float() @ w1_deq[e].float().T + gate, up = l1[:inter], l1[inter:] + gate = gate.clamp(max=activation_clamp) + up = up.clamp(min=-activation_clamp, max=activation_clamp) + act = (gate * torch.sigmoid(gate)) * up * float(topk_weights[t, k].item()) + act_deq = _nvfp4_roundtrip(act.to(torch.bfloat16).unsqueeze(0), gs_l2).squeeze(0) + l2 = act_deq.float() @ w2_deq[e].float().T + y[t] += l2 + return y.to(torch.bfloat16) + + +def _estimate_l2_global_scale(x, l1w, topk_idx, topk_weights, inter, gs_x, gs_w1, clamp): + # The L1-output (L2-input) global scale is a CPU param; estimate it from the activation amax. + x_deq = _nvfp4_roundtrip(x, gs_x) + w1_deq = torch.stack([_nvfp4_roundtrip(l1w[e], gs_w1) for e in range(l1w.shape[0])]) + amax = torch.tensor(1e-6, device=x.device) + for t in range(x.shape[0]): + for k in range(topk_idx.shape[1]): + e = int(topk_idx[t, k].item()) + l1 = x_deq[t].float() @ w1_deq[e].float().T + gate, up = l1[:inter].clamp(max=clamp), l1[inter:].clamp(min=-clamp, max=clamp) + act = (gate * torch.sigmoid(gate)) * up * float(topk_weights[t, k].item()) + amax = torch.maximum(amax, act.abs().amax()) + return float(amax / (6.0 * 448.0)) + + +def test_nvfp4_mega_moe(): + os.environ.setdefault('MASTER_ADDR', '127.0.0.1') + os.environ.setdefault('MASTER_PORT', '12400') + dist.init_process_group('nccl', rank=0, world_size=1) + group = dist.group.WORLD + + num_max_tokens, num_tokens = 128, 128 + hidden, inter = 512, 512 + num_experts, num_topk = 8, 2 + ne_per_rank = num_experts + clamp = 10.0 + + buf = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, + hidden, inter, mma_type='nvfp4xnvfp4') + + torch.manual_seed(0) + x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + l1w = torch.randn((ne_per_rank, inter * 2, hidden), dtype=torch.bfloat16, device='cuda') / (hidden ** 0.5) + l2w = torch.randn((ne_per_rank, hidden, inter), dtype=torch.bfloat16, device='cuda') / (inter ** 0.5) + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') + topk_weights, topk_idx = torch.topk(scores.softmax(-1), num_topk, dim=-1) + + # Per-tensor global scales (CPU scalars) + gs_x = nvfp4_global_scale(x) + gs_w1 = nvfp4_global_scale(l1w) + gs_w2 = nvfp4_global_scale(l2w) + gs_l2 = _estimate_l2_global_scale(x, l1w, topk_idx, topk_weights, inter, gs_x, gs_w1, clamp) + + # Quantize inputs + xp, xsf = per_token_cast_to_nvfp4(x, gs_x, gran_k=GRAN_K) + xsf_packed = xsf.contiguous().view(torch.int32) # K-major int32 (4 E4M3 bytes / int32) + l1 = _cast_w_nvfp4(l1w, gs_w1) + l2 = _cast_w_nvfp4(l2w, gs_w2) + tl1, tl2 = deep_gemm.transform_weights_for_mega_moe(l1, l2) + + buf.x[:num_tokens].copy_(xp) + buf.x_sf[:num_tokens].copy_(xsf_packed) + buf.topk_idx[:num_tokens].copy_(topk_idx) + buf.topk_weights[:num_tokens].copy_(topk_weights) + + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.nvfp4_nvfp4_mega_moe(y=y, l1_weights=tl1, l2_weights=tl2, sym_buffer=buf, + l1_act_global_scale=gs_x, l2_act_global_scale=gs_l2, + l1_weight_global_scale=gs_w1, l2_weight_global_scale=gs_w2, + activation_clamp=clamp, fast_math=True) + torch.cuda.synchronize() + + ref = reference_nvfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, + gs_x, gs_w1, gs_w2, gs_l2, clamp) + diff = calc_diff(y, ref) + print(f'diff = {diff:.5f} (y~{y.float().abs().mean():.3f}, ref~{ref.float().abs().mean():.3f})') + assert diff < 0.05, f'{diff=}' + print('NVFP4 mega MoE passed.') + dist.destroy_process_group() + + +if __name__ == '__main__': + test_nvfp4_mega_moe() From d4f3c6682c63ccc15307b9caaf91a00d433e247f Mon Sep 17 00:00:00 2001 From: Ying Zhang Date: Sat, 27 Jun 2026 02:45:28 +0000 Subject: [PATCH 3/9] Add MXFP4-vs-NVFP4 perf benchmark for packed-FP4 GEMM and mega-MoE Benchmarks the standalone GEMM and mega-MoE for both MXFP4 and NVFP4 using `bench_kineto` (pure device time, kernel-name filtered + L2-flushed, the repo-standard method), reporting per-shape latency, GEMM TFLOPS, and the nvfp4/mxfp4 ratio. Device-time results (B200): mega-MoE nvfp4 is ~1.04-1.06x the mxfp4 time, standalone GEMM ~1.09-1.16x. The residual is gran-16 doubling the scale factors (2x UTCCP copies + 2x SF traffic) vs gran-32 MXFP4. Co-authored-by: Cursor --- tests/bench_packed_fp4.py | 117 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 117 insertions(+) create mode 100644 tests/bench_packed_fp4.py diff --git a/tests/bench_packed_fp4.py b/tests/bench_packed_fp4.py new file mode 100644 index 0000000000..75cb4ff5f5 --- /dev/null +++ b/tests/bench_packed_fp4.py @@ -0,0 +1,117 @@ +"""Performance benchmarks for the packed-FP4 paths (MXFP4 vs NVFP4). + +Covers the standalone GEMM (`mxfp4_gemm_nt` / `nvfp4_gemm_nt`) and the +mega-MoE (`mxfp4_mxfp4_mega_moe` / `nvfp4_nvfp4_mega_moe`). +""" +import os +import sys +import torch +import torch.distributed as dist + +import deep_gemm +from deep_gemm.testing import bench_kineto +from deep_gemm.utils.math import ( + per_token_cast_to_fp4, per_token_cast_to_nvfp4, nvfp4_global_scale, +) + +sys.path.insert(0, os.path.dirname(__file__)) +from test_mxfp4_gemm import _prepare as _prep_mxfp4_gemm +from test_nvfp4_gemm import _prepare as _prep_nvfp4_gemm +from test_nvfp4_mega_moe import _cast_w_nvfp4, _estimate_l2_global_scale, GRAN_K + + +def bench_gemm(): + print('=== Standalone packed-FP4 GEMM (2-CTA de-risk kernel) ===') + print(f'{"M":>6} {"N":>6} {"K":>6} | {"mxfp4 us":>9} {"TFLOPS":>8} | {"nvfp4 us":>9} {"TFLOPS":>8} | nv/mx') + for m, n, k in ((4096, 4096, 4096), (4096, 4096, 8192), (8192, 8192, 8192), (2048, 4096, 16384)): + a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16) + b = torch.randn((n, k), device='cuda', dtype=torch.bfloat16) + d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16) + flops = 2.0 * m * n * k + + ap, asf, _ = _prep_mxfp4_gemm(a) + bp, bsf, _ = _prep_mxfp4_gemm(b) + t_mx = bench_kineto(lambda: deep_gemm.mxfp4_gemm_nt((ap, asf), (bp, bsf), d), + 'mxfp4_gemm', suppress_kineto_output=True) + + gsa, gsb = nvfp4_global_scale(a), nvfp4_global_scale(b) + ap2, asf2, _ = _prep_nvfp4_gemm(a, gsa) + bp2, bsf2, _ = _prep_nvfp4_gemm(b, gsb) + t_nv = bench_kineto(lambda: deep_gemm.nvfp4_gemm_nt((ap2, asf2), (bp2, bsf2), d, a_global_scale=gsa, b_global_scale=gsb), + 'mxfp4_gemm', suppress_kineto_output=True) + + print(f'{m:>6} {n:>6} {k:>6} | {t_mx*1e6:>9.1f} {flops/t_mx/1e12:>8.1f} | ' + f'{t_nv*1e6:>9.1f} {flops/t_nv/1e12:>8.1f} | {t_nv/t_mx:>4.2f}x') + print() + + +def _cast_w_mxfp4(w): + g, n, k = w.shape + wp = torch.empty((g, n, k // 2), device='cuda', dtype=torch.int8) + wsf = torch.empty((g, n, k // 32), device='cuda', dtype=torch.float) + for i in range(g): + wp[i], wsf[i] = per_token_cast_to_fp4(w[i], use_ue8m0=True, gran_k=32) + return wp, deep_gemm.transform_sf_into_required_layout(wsf, n, k, (1, 32), g) + + +def bench_mega(): + os.environ.setdefault('MASTER_ADDR', '127.0.0.1') + os.environ.setdefault('MASTER_PORT', '12533') + dist.init_process_group('nccl', rank=0, world_size=1) + group = dist.group.WORLD + clamp = 10.0 + + print('=== Packed-FP4 mega-MoE (single rank) ===') + print(f'{"tok":>5} {"exp":>4} {"topk":>4} {"hid":>5} {"int":>5} | {"mxfp4 us":>9} | {"nvfp4 us":>9} | nv/mx') + torch.manual_seed(0) + for num_tokens, num_experts, num_topk, hidden, inter in ( + (128, 8, 2, 2048, 2048), + (512, 8, 2, 2048, 2048), + (1024, 32, 4, 4096, 1536), + ): + num_max_tokens = max(128, num_tokens) + x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + l1w = torch.randn((num_experts, inter * 2, hidden), dtype=torch.bfloat16, device='cuda') / (hidden ** 0.5) + l2w = torch.randn((num_experts, hidden, inter), dtype=torch.bfloat16, device='cuda') / (inter ** 0.5) + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') + topk_weights, topk_idx = torch.topk(scores.softmax(-1), num_topk, dim=-1) + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + + # MXFP4 + buf_mx = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, + hidden, inter, mma_type='mxfp4xmxfp4') + xp, xsf = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + tl1, tl2 = deep_gemm.transform_weights_for_mega_moe(_cast_w_mxfp4(l1w), _cast_w_mxfp4(l2w)) + buf_mx.x[:num_tokens].copy_(xp); buf_mx.x_sf[:num_tokens].copy_(xsf) + buf_mx.topk_idx[:num_tokens].copy_(topk_idx); buf_mx.topk_weights[:num_tokens].copy_(topk_weights) + t_mx = bench_kineto(lambda: deep_gemm.mxfp4_mxfp4_mega_moe(y=y, l1_weights=tl1, l2_weights=tl2, sym_buffer=buf_mx, + activation_clamp=clamp, fast_math=True), + 'mega_moe', suppress_kineto_output=True) + buf_mx.destroy() + + # NVFP4 + buf_nv = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, + hidden, inter, mma_type='nvfp4xnvfp4') + gs_x, gs_w1, gs_w2 = nvfp4_global_scale(x), nvfp4_global_scale(l1w), nvfp4_global_scale(l2w) + gs_l2 = _estimate_l2_global_scale(x, l1w, topk_idx, topk_weights, inter, gs_x, gs_w1, clamp) + xpn, xsfn = per_token_cast_to_nvfp4(x, gs_x, gran_k=GRAN_K) + nl1, nl2 = deep_gemm.transform_weights_for_mega_moe(_cast_w_nvfp4(l1w, gs_w1), _cast_w_nvfp4(l2w, gs_w2)) + buf_nv.x[:num_tokens].copy_(xpn); buf_nv.x_sf[:num_tokens].copy_(xsfn.contiguous().view(torch.int32)) + buf_nv.topk_idx[:num_tokens].copy_(topk_idx); buf_nv.topk_weights[:num_tokens].copy_(topk_weights) + t_nv = bench_kineto(lambda: deep_gemm.nvfp4_nvfp4_mega_moe(y=y, l1_weights=nl1, l2_weights=nl2, sym_buffer=buf_nv, + l1_act_global_scale=gs_x, l2_act_global_scale=gs_l2, + l1_weight_global_scale=gs_w1, l2_weight_global_scale=gs_w2, + activation_clamp=clamp, fast_math=True), + 'mega_moe', suppress_kineto_output=True) + buf_nv.destroy() + + print(f'{num_tokens:>5} {num_experts:>4} {num_topk:>4} {hidden:>5} {inter:>5} | ' + f'{t_mx*1e6:>9.1f} | {t_nv*1e6:>9.1f} | {t_nv/t_mx:>4.2f}x') + dist.destroy_process_group() + print() + + +if __name__ == '__main__': + torch.manual_seed(0) + bench_gemm() + bench_mega() From 5bd36ac53779caa3d73f57c78d2d12c1c7057ef9 Mon Sep 17 00:00:00 2001 From: Ying Zhang Date: Sat, 27 Jun 2026 05:33:17 +0000 Subject: [PATCH 4/9] NVFP4 mega-MoE: per-expert global scales (TRT-LLM convention) Switch the NVFP4 mega-MoE from per-tensor scalar global scales to per-expert (num_experts_per_rank,) float32 device tensors, matching the TRT-LLM format: - gate_alpha / up_alpha = 1/(l1_input_gs * gate|up_weight_gs) -> L1 acc dequant (applied per gate/up column before SwiGLU) - l2_input_global_scale (= 448*6/amax) -> L1-output requant - down_alpha = 1/(l2_input_gs * down_weight_gs) -> L2 acc dequant All indexed per local expert in the fused epilogue. Updates the C++/Python entries, the test (per-expert gate/up/down weight scales + per-expert intermediate global scale), and the benchmark. Validated on B200: nvfp4 mega-MoE diff=5.8e-4; mxfp4 + standalone GEMMs unchanged. Co-authored-by: Cursor --- csrc/apis/mega.hpp | 12 +- .../impls/sm100_mxfp4_mxfp4_mega_moe.hpp | 53 +++++-- .../impls/sm100_mxfp4_mxfp4_mega_moe.cuh | 51 ++++--- deep_gemm/mega/__init__.py | 18 ++- tests/bench_packed_fp4.py | 18 ++- tests/test_nvfp4_mega_moe.py | 138 ++++++++++-------- 6 files changed, 175 insertions(+), 115 deletions(-) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 98de0884d6..84555c2373 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -362,9 +362,12 @@ static void nvfp4_nvfp4_mega_moe( const std::vector& sym_buffer_ptrs, const int& rank_idx, const int& num_max_tokens_per_rank, const int& num_experts, const int& num_topk, - // Per-tensor global scales (CPU scalars): activations (L1 input, L2 input) and weights (L1, L2). - const float& l1_act_global_scale, const float& l2_act_global_scale, - const float& l1_weight_global_scale, const float& l2_weight_global_scale, + // NVFP4 per-expert global scales (TRT-LLM convention), each (num_experts_per_rank,) float32: + // gate_alpha / up_alpha = 1/(l1_input_gs * gate|up_weight_gs) (L1 dequant combiner) + // l2_input_global_scale = L2-input per-expert global scale (L1-output requant) + // down_alpha = 1/(l2_input_gs * down_weight_gs) (L2 dequant combiner) + const torch::Tensor& gate_alpha, const torch::Tensor& up_alpha, + const torch::Tensor& l2_input_global_scale, const torch::Tensor& down_alpha, const std::tuple& recipe, const std::string& activation, const std::optional& activation_clamp_opt, @@ -437,8 +440,7 @@ static void nvfp4_nvfp4_mega_moe( hidden, intermediate_hidden, activation_clamp, fast_math, MmaKind::NVFP4, - l1_act_global_scale, l1_weight_global_scale, - l2_act_global_scale, l2_weight_global_scale); + gate_alpha, up_alpha, l2_input_global_scale, down_alpha); } else { DG_HOST_UNREACHABLE("Unsupported architecture"); } diff --git a/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp index 298e058703..f44461b67b 100644 --- a/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp @@ -113,8 +113,11 @@ class SM100MXFP4MegaMoERuntime final : public LaunchRuntime real, L1-output requant, L2 acc -> real. - const float& l1_act_gs = 1.0f, const float& l1_weight_gs = 1.0f, - const float& l2_act_gs = 1.0f, const float& l2_weight_gs = 1.0f + // NVFP4 per-expert global scales (TRT-LLM convention), each (num_experts_per_rank,) float32: + // gate/up_alpha = 1/(l1_input_gs * gate|up_weight_gs); down_alpha = 1/(l2_input_gs * down_weight_gs); + // l2_input_global_scale = the L2-input per-expert global scale (= 448*6/amax). + const std::optional& gate_alpha = std::nullopt, + const std::optional& up_alpha = std::nullopt, + const std::optional& l2_input_global_scale = std::nullopt, + const std::optional& down_alpha = std::nullopt ) { const auto num_ranks = static_cast(sym_buffer_ptrs.size()); const auto num_experts = num_experts_per_rank * num_ranks; @@ -227,6 +235,23 @@ static void sm100_mxfp4_mxfp4_mega_moe( const auto num_sf_ring_tokens = static_cast(l1_acts_sf.size(0)); const bool is_nvfp4 = (mma_kind == MmaKind::NVFP4); + // Extract per-expert global-scale device pointers (NVFP4 only) + const float* gate_alpha_ptr = nullptr; + const float* up_alpha_ptr = nullptr; + const float* l2_input_gs_ptr = nullptr; + const float* down_alpha_ptr = nullptr; + if (is_nvfp4) { + DG_HOST_ASSERT(gate_alpha and up_alpha and l2_input_global_scale and down_alpha); + for (const auto& t : {gate_alpha, up_alpha, l2_input_global_scale, down_alpha}) { + DG_HOST_ASSERT(t->scalar_type() == torch::kFloat and t->is_contiguous()); + DG_HOST_ASSERT(static_cast(t->numel()) == num_experts_per_rank); + } + gate_alpha_ptr = gate_alpha->data_ptr(); + up_alpha_ptr = up_alpha->data_ptr(); + l2_input_gs_ptr = l2_input_global_scale->data_ptr(); + down_alpha_ptr = down_alpha->data_ptr(); + } + const auto config = get_mxfp4_mega_moe_config( num_ranks, num_experts, num_experts_per_rank, num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, @@ -294,11 +319,11 @@ static void sm100_mxfp4_mxfp4_mega_moe( .activation_clamp = activation_clamp, .fast_math = fast_math, .is_nvfp4 = is_nvfp4, - // L1 acc -> real = gs_l1_act * gs_l1_weight; L1-output requant uses gs_l2_act; - // L2 acc -> real = gs_l2_act * gs_l2_weight. (All 1.0 for MXFP4.) - .l1_acc_scale = l1_act_gs * l1_weight_gs, - .l2_act_global_scale = l2_act_gs, - .l2_acc_scale = l2_act_gs * l2_weight_gs, + // Per-expert NVFP4 combiners/global-scales (device ptrs); nullptr for MXFP4. + .gate_alpha = gate_alpha_ptr, + .up_alpha = up_alpha_ptr, + .l2_input_global_scale = l2_input_gs_ptr, + .down_alpha = down_alpha_ptr, .config = config, .y = y.data_ptr(), .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, diff --git a/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh index 60c2a1f802..b9420c917a 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh @@ -64,13 +64,15 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights_sf, - // NVFP4 global-scale dequant params (CPU-side scalars; ignored for MXFP4). - // l1_acc_scale = gs_l1_act * gs_l1_weight (L1 MMA acc -> real) - // l2_act_global_scale = gs_l2_act (L1-output requant) - // l2_acc_scale = gs_l2_act * gs_l2_weight (L2 MMA acc -> real) - const float l1_acc_scale = 1.0f, - const float l2_act_global_scale = 1.0f, - const float l2_acc_scale = 1.0f) { + // NVFP4 per-expert global-scale tensors (device ptrs, (num_experts_per_rank,) f32; + // ignored for MXFP4). TRT-LLM convention (global_scale = 448*6/amax): + // gate_alpha[e]/up_alpha[e] = 1/(l1_input_gs * gate|up_weight_gs) (L1 acc -> real) + // l2_input_global_scale[e] (L1-output requant) + // down_alpha[e] = 1/(l2_input_gs * down_weight_gs) (L2 acc -> real) + const float* gate_alpha = nullptr, + const float* up_alpha = nullptr, + const float* l2_input_global_scale = nullptr, + const float* down_alpha = nullptr) { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; using Allocator = cute::TMEM::Allocator2Sm; @@ -1080,12 +1082,17 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, // Apply SwiGLU: silu(gate) * up auto fp32_values = reinterpret_cast(raw_values); - // NVFP4: dequant the raw MMA accumulator by the L1 global scale BEFORE the - // (nonlinear) SwiGLU. MXFP4's UE8M0 block SFs already fully dequant the accumulator. + // NVFP4: dequant the raw MMA accumulator BEFORE the (nonlinear) SwiGLU, using the + // per-expert combiners. fp32_values are [gate0, up0, gate1, up1] -> even=gate, odd=up. + // MXFP4's UE8M0 block SFs already fully dequant the accumulator. if constexpr (kIsNVFP4) { + const float ga = gate_alpha[local_expert_idx]; + const float ua = up_alpha[local_expert_idx]; #pragma unroll - for (uint32_t kk = 0; kk < 4; ++ kk) - fp32_values[kk] = __fmul2_rn(fp32_values[kk], {l1_acc_scale, l1_acc_scale}); + for (uint32_t kk = 0; kk < 4; ++ kk) { + const float s = (kk & 1) ? ua : ga; + fp32_values[kk] = __fmul2_rn(fp32_values[kk], {s, s}); + } } #pragma unroll for (uint32_t k = 0; k < 2; ++ k) { @@ -1159,14 +1166,17 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, float2 sf_inv; uint8_t sf_byte_x, sf_byte_y; if constexpr (kIsNVFP4) { - const float inv_gs6 = 1.0f / (6.0f * l2_act_global_scale); - const __nv_fp8_e4m3 e4x(amax_values[i].x * inv_gs6); - const __nv_fp8_e4m3 e4y(amax_values[i].y * inv_gs6); + // TRT-LLM: per-expert L2-input global scale (= 448*6/amax). Stored block SF is + // E4M3(amax_block * gs / 6); the E2M1 code uses sf_inv = gs / sf_e4m3. + const float gs = l2_input_global_scale[local_expert_idx]; + const float gs6 = gs * (1.0f / 6.0f); + const __nv_fp8_e4m3 e4x(amax_values[i].x * gs6); + const __nv_fp8_e4m3 e4y(amax_values[i].y * gs6); sf_byte_x = e4x.__x; sf_byte_y = e4y.__x; - const float sx = static_cast(e4x) * l2_act_global_scale; - const float sy = static_cast(e4y) * l2_act_global_scale; - sf_inv.x = sx > 0.f ? 1.0f / sx : 0.f; - sf_inv.y = sy > 0.f ? 1.0f / sy : 0.f; + const float fx = static_cast(e4x); + const float fy = static_cast(e4y); + sf_inv.x = fx > 0.f ? gs / fx : 0.f; + sf_inv.y = fy > 0.f ? gs / fy : 0.f; } else { const float2 scaled = __fmul2_rn(amax_values[i], {1.0f / 6.0f, 1.0f / 6.0f}); const int ex = math::fast_log2_ceil(scaled.x); @@ -1315,11 +1325,12 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, values[4], values[5], values[6], values[7]); cutlass::arch::fence_view_async_tmem_load(); - // NVFP4: dequant the L2 MMA accumulator by the L2 global scale before BF16 cast + // NVFP4: dequant the L2 MMA accumulator by the per-expert combiner before BF16 cast if constexpr (kIsNVFP4) { + const float da = down_alpha[local_expert_idx]; #pragma unroll for (uint32_t v = 0; v < ATOM_M; ++ v) - values[v] = __float_as_uint(__uint_as_float(values[v]) * l2_acc_scale); + values[v] = __float_as_uint(__uint_as_float(values[v]) * da); } // Wait shared memory release from previous NVLink store diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 37e1976a7a..03c31eaf98 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -203,17 +203,20 @@ def nvfp4_nvfp4_mega_moe(y: torch.Tensor, l1_weights: Tuple[torch.Tensor, torch.Tensor], l2_weights: Tuple[torch.Tensor, torch.Tensor], sym_buffer: SymmBuffer, - l1_act_global_scale: float, - l2_act_global_scale: float, - l1_weight_global_scale: float, - l2_weight_global_scale: float, + gate_alpha: torch.Tensor, + up_alpha: torch.Tensor, + l2_input_global_scale: torch.Tensor, + down_alpha: torch.Tensor, cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, recipe: Tuple[int, int, int] = (1, 1, 16), activation: str = 'swiglu', activation_clamp: Optional[float] = None, fast_math: bool = True): - # Packed NVFP4 x NVFP4 mega MoE (E2M1 data, E4M3 SF gran-16, per-tensor global scales). - # Global scales are CPU-side scalar kernel params (dequant: acc * gs_act * gs_weight). + # Packed NVFP4 x NVFP4 mega MoE (E2M1 data, E4M3 SF gran-16). Per-expert global scales + # (TRT-LLM convention), each (num_experts_per_rank,) float32 device tensor: + # gate_alpha / up_alpha = 1/(l1_input_gs * gate|up_weight_gs) (L1 acc -> real) + # l2_input_global_scale = L2-input per-expert global scale (L1-output requant) + # down_alpha = 1/(l2_input_gs * down_weight_gs) (L2 acc -> real) _C.nvfp4_nvfp4_mega_moe( y, l1_weights, l2_weights, @@ -222,8 +225,7 @@ def nvfp4_nvfp4_mega_moe(y: torch.Tensor, sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), sym_buffer.num_max_tokens_per_rank, sym_buffer.num_experts, sym_buffer.num_topk, - float(l1_act_global_scale), float(l2_act_global_scale), - float(l1_weight_global_scale), float(l2_weight_global_scale), + gate_alpha, up_alpha, l2_input_global_scale, down_alpha, recipe, activation, activation_clamp, fast_math, diff --git a/tests/bench_packed_fp4.py b/tests/bench_packed_fp4.py index 75cb4ff5f5..14a3c3a78b 100644 --- a/tests/bench_packed_fp4.py +++ b/tests/bench_packed_fp4.py @@ -17,7 +17,7 @@ sys.path.insert(0, os.path.dirname(__file__)) from test_mxfp4_gemm import _prepare as _prep_mxfp4_gemm from test_nvfp4_gemm import _prepare as _prep_nvfp4_gemm -from test_nvfp4_mega_moe import _cast_w_nvfp4, _estimate_l2_global_scale, GRAN_K +from test_nvfp4_mega_moe import _cast_l1_w, _cast_l2_w, _estimate_l2act_gs, GRAN_K def bench_gemm(): @@ -89,18 +89,22 @@ def bench_mega(): 'mega_moe', suppress_kineto_output=True) buf_mx.destroy() - # NVFP4 + # NVFP4 (per-expert global scales, TRT-LLM convention) buf_nv = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, hidden, inter, mma_type='nvfp4xnvfp4') - gs_x, gs_w1, gs_w2 = nvfp4_global_scale(x), nvfp4_global_scale(l1w), nvfp4_global_scale(l2w) - gs_l2 = _estimate_l2_global_scale(x, l1w, topk_idx, topk_weights, inter, gs_x, gs_w1, clamp) + gs_x = nvfp4_global_scale(x) + l1, gate_gs, up_gs = _cast_l1_w(l1w) + l2, down_gs = _cast_l2_w(l2w) + l2act_gs = _estimate_l2act_gs(x, l1w, topk_idx, topk_weights, inter, gs_x, gate_gs, up_gs, clamp, num_experts) + gate_alpha = (gs_x * gate_gs).contiguous(); up_alpha = (gs_x * up_gs).contiguous() + down_alpha = (l2act_gs * down_gs).contiguous(); l2_input_gs = (1.0 / l2act_gs).contiguous() xpn, xsfn = per_token_cast_to_nvfp4(x, gs_x, gran_k=GRAN_K) - nl1, nl2 = deep_gemm.transform_weights_for_mega_moe(_cast_w_nvfp4(l1w, gs_w1), _cast_w_nvfp4(l2w, gs_w2)) + nl1, nl2 = deep_gemm.transform_weights_for_mega_moe(l1, l2) buf_nv.x[:num_tokens].copy_(xpn); buf_nv.x_sf[:num_tokens].copy_(xsfn.contiguous().view(torch.int32)) buf_nv.topk_idx[:num_tokens].copy_(topk_idx); buf_nv.topk_weights[:num_tokens].copy_(topk_weights) t_nv = bench_kineto(lambda: deep_gemm.nvfp4_nvfp4_mega_moe(y=y, l1_weights=nl1, l2_weights=nl2, sym_buffer=buf_nv, - l1_act_global_scale=gs_x, l2_act_global_scale=gs_l2, - l1_weight_global_scale=gs_w1, l2_weight_global_scale=gs_w2, + gate_alpha=gate_alpha, up_alpha=up_alpha, + l2_input_global_scale=l2_input_gs, down_alpha=down_alpha, activation_clamp=clamp, fast_math=True), 'mega_moe', suppress_kineto_output=True) buf_nv.destroy() diff --git a/tests/test_nvfp4_mega_moe.py b/tests/test_nvfp4_mega_moe.py index 07c96826d7..844e96c1f4 100644 --- a/tests/test_nvfp4_mega_moe.py +++ b/tests/test_nvfp4_mega_moe.py @@ -12,16 +12,12 @@ GRAN_K = 16 -def _nvfp4_roundtrip(x: torch.Tensor, gs: float) -> torch.Tensor: - """Quantize bf16 -> NVFP4 (E4M3 SF gran-16 + global scale) and dequantize back to the - exact values the kernel operates on.""" - packed, sf = per_token_cast_to_nvfp4(x, gs, gran_k=GRAN_K) - return cast_back_from_nvfp4(packed, sf, gs, gran_k=GRAN_K).to(x.dtype) +def _rt(x, gs): + p, s = per_token_cast_to_nvfp4(x, gs, gran_k=GRAN_K) + return cast_back_from_nvfp4(p, s, gs, gran_k=GRAN_K).to(x.dtype) -def _mn_major_packed_e4m3(sf_bytes: torch.Tensor) -> torch.Tensor: - """Replicate `get_mn_major_tma_aligned_packed_ue8m0_tensor_torch` for raw E4M3 bytes: - pad to TMA-aligned MN and K%4, pack 4 SF bytes -> int32, transpose to MN-major.""" +def _mn_major_packed_e4m3(sf_bytes): g, mn, k = sf_bytes.shape aligned_mn = get_tma_aligned_size(mn, 4) aligned_k = align(k, 4) @@ -31,56 +27,73 @@ def _mn_major_packed_e4m3(sf_bytes: torch.Tensor) -> torch.Tensor: out = torch.empty_strided((g, aligned_mn, aligned_k // 4), (aligned_mn * (aligned_k // 4), 1, aligned_mn), dtype=torch.int32, device=sf_bytes.device) - out = out.copy_(padded)[:, :mn] - return out - - -def _cast_w_nvfp4(w: torch.Tensor, gs: float): - g, n, k = w.shape + return out.copy_(padded)[:, :mn] + + +def _cast_l1_w(l1w): + # l1w: [E, inter*2, hidden]. gate = rows [:inter], up = rows [inter:]; per-expert global scales. + g, n2, k = l1w.shape + inter = n2 // 2 + wp = torch.empty((g, n2, k // 2), device='cuda', dtype=torch.int8) + wsf = torch.empty((g, n2, k // GRAN_K), device='cuda', dtype=torch.uint8) + gate_gs = torch.empty(g, device='cuda', dtype=torch.float32) + up_gs = torch.empty(g, device='cuda', dtype=torch.float32) + for e in range(g): + ggs = nvfp4_global_scale(l1w[e][:inter]); ugs = nvfp4_global_scale(l1w[e][inter:]) + gate_gs[e], up_gs[e] = ggs, ugs + wp[e][:inter], wsf[e][:inter] = per_token_cast_to_nvfp4(l1w[e][:inter], ggs, gran_k=GRAN_K) + wp[e][inter:], wsf[e][inter:] = per_token_cast_to_nvfp4(l1w[e][inter:], ugs, gran_k=GRAN_K) + return (wp, _mn_major_packed_e4m3(wsf)), gate_gs, up_gs + + +def _cast_l2_w(l2w): + g, n, k = l2w.shape wp = torch.empty((g, n, k // 2), device='cuda', dtype=torch.int8) wsf = torch.empty((g, n, k // GRAN_K), device='cuda', dtype=torch.uint8) - for i in range(g): - wp[i], wsf[i] = per_token_cast_to_nvfp4(w[i], gs, gran_k=GRAN_K) - return wp, _mn_major_packed_e4m3(wsf) + down_gs = torch.empty(g, device='cuda', dtype=torch.float32) + for e in range(g): + dgs = nvfp4_global_scale(l2w[e]); down_gs[e] = dgs + wp[e], wsf[e] = per_token_cast_to_nvfp4(l2w[e], dgs, gran_k=GRAN_K) + return (wp, _mn_major_packed_e4m3(wsf)), down_gs -def reference_nvfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, - gs_x, gs_w1, gs_w2, gs_l2, activation_clamp: float = 10.0): - num_tokens = x.shape[0] - x_deq = _nvfp4_roundtrip(x, gs_x) - w1_deq = torch.stack([_nvfp4_roundtrip(l1w[e], gs_w1) for e in range(l1w.shape[0])]) - w2_deq = torch.stack([_nvfp4_roundtrip(l2w[e], gs_w2) for e in range(l2w.shape[0])]) +def _swiglu(l1, inter, weight, clamp): + gate = l1[:inter].clamp(max=clamp) + up = l1[inter:].clamp(min=-clamp, max=clamp) + return (gate * torch.sigmoid(gate)) * up * weight - y = torch.zeros((num_tokens, hidden), dtype=torch.float, device=x.device) - for t in range(num_tokens): + +def _estimate_l2act_gs(x, l1w, topk_idx, topk_weights, inter, gs_x, gate_gs, up_gs, clamp, num_experts): + # Per-expert L2-input (intermediate) global scale (my convention: amax/(6*448)). + x_deq = _rt(x, gs_x) + w1_deq = [torch.cat([_rt(l1w[e][:inter], gate_gs[e].item()), _rt(l1w[e][inter:], up_gs[e].item())]) + for e in range(num_experts)] + amax = torch.full((num_experts,), 1e-6, device=x.device) + for t in range(x.shape[0]): for k in range(topk_idx.shape[1]): e = int(topk_idx[t, k].item()) - if e < 0: - continue l1 = x_deq[t].float() @ w1_deq[e].float().T - gate, up = l1[:inter], l1[inter:] - gate = gate.clamp(max=activation_clamp) - up = up.clamp(min=-activation_clamp, max=activation_clamp) - act = (gate * torch.sigmoid(gate)) * up * float(topk_weights[t, k].item()) - act_deq = _nvfp4_roundtrip(act.to(torch.bfloat16).unsqueeze(0), gs_l2).squeeze(0) - l2 = act_deq.float() @ w2_deq[e].float().T - y[t] += l2 - return y.to(torch.bfloat16) + act = _swiglu(l1, inter, float(topk_weights[t, k].item()), clamp) + amax[e] = torch.maximum(amax[e], act.abs().amax()) + return amax / (6.0 * 448.0) -def _estimate_l2_global_scale(x, l1w, topk_idx, topk_weights, inter, gs_x, gs_w1, clamp): - # The L1-output (L2-input) global scale is a CPU param; estimate it from the activation amax. - x_deq = _nvfp4_roundtrip(x, gs_x) - w1_deq = torch.stack([_nvfp4_roundtrip(l1w[e], gs_w1) for e in range(l1w.shape[0])]) - amax = torch.tensor(1e-6, device=x.device) - for t in range(x.shape[0]): +def reference_nvfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, + gs_x, gate_gs, up_gs, down_gs, l2act_gs, clamp): + num_tokens, num_experts = x.shape[0], l1w.shape[0] + x_deq = _rt(x, gs_x) + w1_deq = [torch.cat([_rt(l1w[e][:inter], gate_gs[e].item()), _rt(l1w[e][inter:], up_gs[e].item())]) + for e in range(num_experts)] + w2_deq = [_rt(l2w[e], down_gs[e].item()) for e in range(num_experts)] + y = torch.zeros((num_tokens, hidden), dtype=torch.float, device=x.device) + for t in range(num_tokens): for k in range(topk_idx.shape[1]): e = int(topk_idx[t, k].item()) l1 = x_deq[t].float() @ w1_deq[e].float().T - gate, up = l1[:inter].clamp(max=clamp), l1[inter:].clamp(min=-clamp, max=clamp) - act = (gate * torch.sigmoid(gate)) * up * float(topk_weights[t, k].item()) - amax = torch.maximum(amax, act.abs().amax()) - return float(amax / (6.0 * 448.0)) + act = _swiglu(l1, inter, float(topk_weights[t, k].item()), clamp).to(torch.bfloat16) + act_deq = _rt(act.unsqueeze(0), l2act_gs[e].item()).squeeze(0) + y[t] += act_deq.float() @ w2_deq[e].float().T + return y.to(torch.bfloat16) def test_nvfp4_mega_moe(): @@ -92,7 +105,6 @@ def test_nvfp4_mega_moe(): num_max_tokens, num_tokens = 128, 128 hidden, inter = 512, 512 num_experts, num_topk = 8, 2 - ne_per_rank = num_experts clamp = 10.0 buf = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, @@ -100,38 +112,42 @@ def test_nvfp4_mega_moe(): torch.manual_seed(0) x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') - l1w = torch.randn((ne_per_rank, inter * 2, hidden), dtype=torch.bfloat16, device='cuda') / (hidden ** 0.5) - l2w = torch.randn((ne_per_rank, hidden, inter), dtype=torch.bfloat16, device='cuda') / (inter ** 0.5) + l1w = torch.randn((num_experts, inter * 2, hidden), dtype=torch.bfloat16, device='cuda') / (hidden ** 0.5) + l2w = torch.randn((num_experts, hidden, inter), dtype=torch.bfloat16, device='cuda') / (inter ** 0.5) scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') topk_weights, topk_idx = torch.topk(scores.softmax(-1), num_topk, dim=-1) - # Per-tensor global scales (CPU scalars) + # Per-tensor input global scale (my convention amax/(6*448)); per-expert weight global scales gs_x = nvfp4_global_scale(x) - gs_w1 = nvfp4_global_scale(l1w) - gs_w2 = nvfp4_global_scale(l2w) - gs_l2 = _estimate_l2_global_scale(x, l1w, topk_idx, topk_weights, inter, gs_x, gs_w1, clamp) + l1, gate_gs, up_gs = _cast_l1_w(l1w) + l2, down_gs = _cast_l2_w(l2w) + l2act_gs = _estimate_l2act_gs(x, l1w, topk_idx, topk_weights, inter, gs_x, gate_gs, up_gs, clamp, num_experts) + + # Per-expert kernel params (TRT-LLM convention): + # alpha = 1/(input_gs_trt * weight_gs_trt) = my_gs_input * my_gs_weight + # l2_input_global_scale (TRT) = 1 / my_gs + gate_alpha = (gs_x * gate_gs).contiguous() + up_alpha = (gs_x * up_gs).contiguous() + down_alpha = (l2act_gs * down_gs).contiguous() + l2_input_global_scale = (1.0 / l2act_gs).contiguous() - # Quantize inputs - xp, xsf = per_token_cast_to_nvfp4(x, gs_x, gran_k=GRAN_K) - xsf_packed = xsf.contiguous().view(torch.int32) # K-major int32 (4 E4M3 bytes / int32) - l1 = _cast_w_nvfp4(l1w, gs_w1) - l2 = _cast_w_nvfp4(l2w, gs_w2) tl1, tl2 = deep_gemm.transform_weights_for_mega_moe(l1, l2) + xp, xsf = per_token_cast_to_nvfp4(x, gs_x, gran_k=GRAN_K) buf.x[:num_tokens].copy_(xp) - buf.x_sf[:num_tokens].copy_(xsf_packed) + buf.x_sf[:num_tokens].copy_(xsf.contiguous().view(torch.int32)) buf.topk_idx[:num_tokens].copy_(topk_idx) buf.topk_weights[:num_tokens].copy_(topk_weights) y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') deep_gemm.nvfp4_nvfp4_mega_moe(y=y, l1_weights=tl1, l2_weights=tl2, sym_buffer=buf, - l1_act_global_scale=gs_x, l2_act_global_scale=gs_l2, - l1_weight_global_scale=gs_w1, l2_weight_global_scale=gs_w2, + gate_alpha=gate_alpha, up_alpha=up_alpha, + l2_input_global_scale=l2_input_global_scale, down_alpha=down_alpha, activation_clamp=clamp, fast_math=True) torch.cuda.synchronize() ref = reference_nvfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, - gs_x, gs_w1, gs_w2, gs_l2, clamp) + gs_x, gate_gs, up_gs, down_gs, l2act_gs, clamp) diff = calc_diff(y, ref) print(f'diff = {diff:.5f} (y~{y.float().abs().mean():.3f}, ref~{ref.float().abs().mean():.3f})') assert diff < 0.05, f'{diff=}' From c6842aa4f490d11f522eed123c5af7015d489dd1 Mon Sep 17 00:00:00 2001 From: Ying Zhang Date: Sat, 27 Jun 2026 05:39:05 +0000 Subject: [PATCH 5/9] NVFP4 mega-MoE: hoist per-expert scale loads out of the per-atom loops The per-expert global scales (gate/up/down alpha, l2_input_global_scale) were re-read from global memory inside the per-atom epilogue loops. Hoist them to once-per-block registers. NVFP4 mega-MoE overhead drops from ~1.10x to ~1.05-1.08x vs MXFP4 (the residual is gran-16 SF traffic, not scale loading). Correctness unchanged (diff=5.8e-4). Co-authored-by: Cursor --- .../impls/sm100_mxfp4_mxfp4_mega_moe.cuh | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh index b9420c917a..65fe7420a4 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_mxfp4_mxfp4_mega_moe.cuh @@ -1025,6 +1025,15 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, const uint32_t pool_m_idx = pool_block_idx * BLOCK_M; // Full-pool offset for non-ring metadata uint32_t n_idx = n_block_idx * BLOCK_N; + // NVFP4: load per-expert global scales ONCE per block (hoisted out of the per-atom loops). + float nv_gate_alpha = 1.0f, nv_up_alpha = 1.0f, nv_l2_in_gs = 1.0f, nv_down_alpha = 1.0f; + if constexpr (kIsNVFP4) { + nv_gate_alpha = gate_alpha[local_expert_idx]; + nv_up_alpha = up_alpha[local_expert_idx]; + nv_l2_in_gs = l2_input_global_scale[local_expert_idx]; + nv_down_alpha = down_alpha[local_expert_idx]; + } + if (block_phase == sched::BlockPhase::Linear1) { // Wait L2 block empty const auto l2_empty_ptr = workspace.get_l2_empty_count_ptr(ring_block_idx); @@ -1086,11 +1095,9 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, // per-expert combiners. fp32_values are [gate0, up0, gate1, up1] -> even=gate, odd=up. // MXFP4's UE8M0 block SFs already fully dequant the accumulator. if constexpr (kIsNVFP4) { - const float ga = gate_alpha[local_expert_idx]; - const float ua = up_alpha[local_expert_idx]; #pragma unroll for (uint32_t kk = 0; kk < 4; ++ kk) { - const float s = (kk & 1) ? ua : ga; + const float s = (kk & 1) ? nv_up_alpha : nv_gate_alpha; fp32_values[kk] = __fmul2_rn(fp32_values[kk], {s, s}); } } @@ -1168,7 +1175,7 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, if constexpr (kIsNVFP4) { // TRT-LLM: per-expert L2-input global scale (= 448*6/amax). Stored block SF is // E4M3(amax_block * gs / 6); the E2M1 code uses sf_inv = gs / sf_e4m3. - const float gs = l2_input_global_scale[local_expert_idx]; + const float gs = nv_l2_in_gs; const float gs6 = gs * (1.0f / 6.0f); const __nv_fp8_e4m3 e4x(amax_values[i].x * gs6); const __nv_fp8_e4m3 e4y(amax_values[i].y * gs6); @@ -1327,10 +1334,9 @@ sm100_mxfp4_mxfp4_mega_moe_impl(void* y, // NVFP4: dequant the L2 MMA accumulator by the per-expert combiner before BF16 cast if constexpr (kIsNVFP4) { - const float da = down_alpha[local_expert_idx]; #pragma unroll for (uint32_t v = 0; v < ATOM_M; ++ v) - values[v] = __float_as_uint(__uint_as_float(values[v]) * da); + values[v] = __float_as_uint(__uint_as_float(values[v]) * nv_down_alpha); } // Wait shared memory release from previous NVLink store From 883dca527e93bc38498a282fcfe08e23e5c8102d Mon Sep 17 00:00:00 2001 From: Ying Zhang Date: Sun, 28 Jun 2026 08:20:59 +0000 Subject: [PATCH 6/9] Add DeepGEMM-vs-FlashInfer NVFP4 MoE benchmarks (single-device + EP) Adds device-time benchmarks comparing DeepGEMM mega-MoE (fp8xfp4 / nvfp4) against FlashInfer's NVFP4 MoE backends (cute_dsl, cutlass, trtllm-gen): - tests/bench_flashinfer_vs_deepgemm.py: single-device comparison (CUDA-graph device time) across all backends; EP-aware builders (local_expert_offset / ep_rank) for cute_dsl and cutlass. - tests/bench_ep_multi_gpu.py: multi-GPU expert-parallel benchmark. DeepGEMM uses its native fused dispatch+combine; FlashInfer uses replicated input + expert-sharded compute + all_reduce combine. Per-kernel device time summed from traces; moe/all_reduce breakdown; NVLS combine via symm-registered buffer; sleep-aligned collective timing; avg-over-ranks aggregation. - tests/bench_kernel_breakdown.py: per-kernel device-time breakdown per backend. - tests/fi_trtllm.py + tests/_fi_vendor/: thin driver over FlashInfer's trtllm_fp4_block_scale_moe, reusing the vendored (v0.6.11) weight-shuffle and routing harness. - tests/bench_packed_fp4.py: add fp8xfp4 mega-MoE column to the mega benchmark. Co-authored-by: Cursor --- tests/_fi_vendor/__init__.py | 0 tests/_fi_vendor/test_trtllm_gen_fused_moe.py | 4496 +++++++++++++++++ tests/_fi_vendor/utils.py | 173 + tests/bench_ep_multi_gpu.py | 338 ++ tests/bench_flashinfer_vs_deepgemm.py | 280 + tests/bench_kernel_breakdown.py | 94 + tests/bench_packed_fp4.py | 20 +- tests/fi_trtllm.py | 96 + 8 files changed, 5493 insertions(+), 4 deletions(-) create mode 100644 tests/_fi_vendor/__init__.py create mode 100644 tests/_fi_vendor/test_trtllm_gen_fused_moe.py create mode 100644 tests/_fi_vendor/utils.py create mode 100644 tests/bench_ep_multi_gpu.py create mode 100644 tests/bench_flashinfer_vs_deepgemm.py create mode 100644 tests/bench_kernel_breakdown.py create mode 100644 tests/fi_trtllm.py diff --git a/tests/_fi_vendor/__init__.py b/tests/_fi_vendor/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/_fi_vendor/test_trtllm_gen_fused_moe.py b/tests/_fi_vendor/test_trtllm_gen_fused_moe.py new file mode 100644 index 0000000000..231db811d8 --- /dev/null +++ b/tests/_fi_vendor/test_trtllm_gen_fused_moe.py @@ -0,0 +1,4496 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import pytest +from abc import ABC, abstractmethod +from typing import Dict +import torch +from cuda.bindings import runtime +from torch.nn import functional as F + +from flashinfer import ( + ActivationType, + RoutingMethodType, + e2m1_and_ufp8sf_scale_to_float, + fp4_quantize, + mxfp8_dequantize_host, + mxfp8_quantize, + reorder_rows_for_gated_act_gemm, + shuffle_matrix_a, + shuffle_matrix_sf_a, +) +from flashinfer.autotuner import autotune +from flashinfer.fp4_quantization import block_scale_interleave +from flashinfer.fused_moe import ( + WeightLayout, + convert_to_block_layout, + trtllm_fp4_block_scale_moe, + trtllm_fp8_block_scale_moe, + trtllm_fp8_block_scale_routed_moe, + trtllm_fp8_per_tensor_scale_moe, + trtllm_bf16_moe, + trtllm_mxint4_block_scale_moe, +) +from flashinfer.fused_moe.core import ( + get_w2_permute_indices_with_cache, + _maybe_get_cached_w3_w1_permute_indices, + Fp8QuantizationType, +) +from flashinfer.utils import get_compute_capability +from .utils import is_gated_activation, skip_checks, QuantMode + + +# Max num tokens to tune for trtllm-gen fused moe +TUNE_MAX_NUM_TOKENS = 4096 + + +def check_cuda(err): + """Unified CUDA error checking function used throughout the file.""" + if err != runtime.cudaError_t.cudaSuccess: + error_name = runtime.cudaGetErrorName(err) + error_string = runtime.cudaGetErrorString(err) + raise RuntimeError(f"CUDA error: {error_name[1]}: {error_string[1]}") + + +class CUDAGraphMoE: + """ + Simple CUDA Graph wrapper for MoE operations. + + The graph captures tensor references and automatically updates them during execution. + + Three core methods: capture(), launch(), cleanup() + + Usage: + cuda_graph = CUDAGraphMoE(moe_impl, static_data, **config) + cuda_graph.capture(hidden_states_sample, expert_logits=logits, routing_bias=bias) + output = cuda_graph.launch(new_hidden_states) # Repeat as needed + cuda_graph.cleanup() + """ + + def __init__(self, moe_impl, static_data, **config): + self.moe_impl = moe_impl + self.static_data = static_data + self.config = config + self.enable_autotune = config.get("enable_autotune", True) + self.graph = None + self.graph_exec = None + self.stream = None + self.input_tensor = None + self.output_tensor = None + self.is_captured = False + + def capture(self, hidden_states_sample, **runtime_args): + """Capture CUDA graph with the given sample input.""" + if self.is_captured: + raise RuntimeError( + "Graph already captured. Call cleanup() first to re-capture." + ) + if not isinstance(self.moe_impl, FP4Moe): + raise NotImplementedError( + f"CUDA graph capture not yet implemented for {type(self.moe_impl)}" + ) + + # Create stream + err, self.stream = runtime.cudaStreamCreate() + check_cuda(err) + + # Get the raw stream pointer for PyTorch + stream_ptr = int(self.stream) + torch_stream = torch.cuda.ExternalStream(stream_ptr) + + # Store input tensor reference (will be updated in place during launch) + self.input_tensor = hidden_states_sample.clone() + + # Warmup + with torch.cuda.stream(torch_stream), autotune(self.enable_autotune): + for _ in range(1): + self._run_moe_computation(runtime_args) + + # Synchronize our stream after warmup + err = runtime.cudaStreamSynchronize(self.stream)[0] + check_cuda(err) + + # Begin capture + err, self.graph = runtime.cudaGraphCreate(0) + check_cuda(err) + err = runtime.cudaStreamBeginCapture( + self.stream, runtime.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal + )[0] + check_cuda(err) + + try: + # Capture computation on our stream + with torch.cuda.stream(torch_stream): + self.output_tensor = self._run_moe_computation(runtime_args) + err, self.graph = runtime.cudaStreamEndCapture(self.stream) + check_cuda(err) + err, self.graph_exec = runtime.cudaGraphInstantiate(self.graph, 0) + check_cuda(err) + self.is_captured = True + except Exception as e: + self.cleanup() + raise RuntimeError(f"CUDA graph capture failed: {e}") from e + + def launch(self, hidden_states_new): + """Launch captured CUDA graph with new input.""" + if not self.is_captured: + raise RuntimeError("Graph not captured. Call capture() first.") + + # Update input tensor in place + self.input_tensor.copy_(hidden_states_new) + + # Launch graph + err = runtime.cudaGraphLaunch(self.graph_exec, self.stream)[0] + check_cuda(err) + err = runtime.cudaStreamSynchronize(self.stream)[0] + check_cuda(err) + + # Return output tensor (automatically updated by graph execution) + return self.output_tensor + + def cleanup(self): + """Clean up all CUDA graph resources.""" + if self.graph_exec is not None: + err = runtime.cudaGraphExecDestroy(self.graph_exec)[0] + check_cuda(err) + self.graph_exec = None + if self.graph is not None: + err = runtime.cudaGraphDestroy(self.graph)[0] + check_cuda(err) + self.graph = None + if self.stream is not None: + err = runtime.cudaStreamDestroy(self.stream)[0] + check_cuda(err) + self.stream = None + self.input_tensor = None + self.output_tensor = None + self.is_captured = False + + def _run_moe_computation(self, runtime_args): + """Run the MoE computation.""" + input_quantized = self.moe_impl.quantize_inputs( + self.input_tensor, + self.config["hidden_states_scale_global"], + is_swizzling=False, + ) + + output = trtllm_fp4_block_scale_moe( + routing_logits=runtime_args["expert_logits"], + routing_bias=runtime_args["routing_bias"], + hidden_states=input_quantized["hidden_states"], + hidden_states_scale=input_quantized["hidden_states_scale"], + gemm1_weights=self.static_data["gemm1_weights_fp4_shuffled"], + gemm1_weights_scale=self.static_data["gemm1_scales_fp4_shuffled"], + gemm1_bias=self.config["gemm1_bias"], + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=self.static_data["gemm2_weights_fp4_shuffled"], + gemm2_weights_scale=self.static_data["gemm2_scales_fp4_shuffled"], + gemm2_bias=self.config["gemm2_bias"], + output1_scale_scalar=self.static_data["scale_c_fc1"], + output1_scale_gate_scalar=self.static_data["scale_gate_fc1"], + output2_scale_scalar=self.static_data["scale_c_fc2"], + num_experts=self.config["num_experts"], + top_k=self.config["top_k"], + n_group=self.config["n_groups"], + topk_group=self.config["top_k_groups"], + intermediate_size=self.config["intermediate_size"], + local_expert_offset=0, + local_num_experts=self.config["num_experts"], + routed_scaling_factor=self.config["routed_scaling"], + routing_method_type=self.config["routing_method_type"], + activation_type=self.config["activation_type"], + do_finalize=True, + tune_max_num_tokens=TUNE_MAX_NUM_TOKENS, + norm_topk_prob=self.config.get("norm_topk_prob", True), + ) + return output # Extract tensor from tuple + + +# ==================================================================================== +# Abstract Base Class for MoE Implementations +# ==================================================================================== + + +class Moe(ABC): + """Abstract base class for MoE implementations.""" + + def __init__(self): + self.name = self.__class__.__name__ + + @property + @abstractmethod + def quant_mode(self) -> QuantMode: + """Get the quantization mode of this MoE implementation.""" + pass + + @abstractmethod + def quantize_weights(self, gemm1_weights, gemm2_weights, hidden_states_sample): + """Quantize static weights and compute global scale factors (done offline).""" + pass + + @abstractmethod + def quantize_inputs(self, hidden_states, hidden_states_scale_global): + """Quantize dynamic inputs/hidden states using pre-computed global scale (done at runtime).""" + pass + + @abstractmethod + def prepare_static_weights_for_kernel( + self, + args_dequant, + args, + gemm1_weights_orig, + gemm2_weights_orig, + hidden_size, + intermediate_size, + num_experts, + weight_processing, + ): + """ + Prepare quantized weights for kernel (done offline with weights). + + Args: + args_dequant: Contains c_global_sf and other dequantization parameters + args: Contains already quantized weights (gemm1_weights, gemm2_weights) and scales + gemm1_weights_orig: Original unquantized FC1 weights (used by FP4 for re-quantization) + gemm2_weights_orig: Original unquantized FC2 weights (used by FP4 for re-quantization) + + Note: + - FP4 implementations use both original weights (for linear layout quantization) + and args.gemm*_weights (for swizzled layout) + - FP8 implementations typically only use args.gemm*_weights (already quantized) + """ + pass + + @abstractmethod + def call_moe( + self, static_data, hidden_states_orig, hidden_states_scale_global, **kwargs + ): + """Call MoE with runtime input quantization + kernel execution (done at runtime).""" + pass + + @abstractmethod + def compute_reference(self, args): + """Compute reference output using dequantized operations.""" + pass + + def compute_production(self, args_dequant, args, **kwargs): + """Unified actual computation that delegates to implementation-specific methods.""" + return _compute_moe_actual_unified(self, args_dequant, args, **kwargs) + + @abstractmethod + def get_tolerances(self): + """Get accuracy tolerances for this quantization mode.""" + pass + + def __str__(self): + return self.name + + +# ==================================================================================== +# FP4 Quantization Implementation +# ==================================================================================== + + +class FP4Moe(Moe): + """ + FP4 NvFP4 / MxFP4 MoE implementation with block scaling. + Args: + is_mxfp4: Whether to use MxFP4 or NvFP4 weight quantization + If True, the activation is quantized to MxFP8, else the activation is quantized to NvFP4 + """ + + def __init__(self, quant_mode: QuantMode): + super().__init__() + self._quant_mode = quant_mode + self.is_mxfp4 = ( + quant_mode == QuantMode.FP4_MXFP4_MXFP8 + or quant_mode == QuantMode.FP4_MXFP4_Bf16 + ) + self.sf_vec_size = 32 if self.is_mxfp4 else 16 + + @property + def quant_mode(self) -> QuantMode: + return self._quant_mode + + def quantize_weights(self, gemm1_weights, gemm2_weights, hidden_states_sample): + """Quantize weights to FP4 format and compute global scale factors.""" + num_experts = gemm1_weights.shape[0] + # Compute global scale factor for hidden states (offline calibration) + if self.quant_mode == QuantMode.FP4_NVFP4_NVFP4: + # nvfp4 hidden states + hidden_states_scale_global = calculate_fp4_global_scale_factor( + hidden_states_sample, + False, + ) + else: + # mxfp8 / bf16 hidden states + hidden_states_scale_global = 1.0 + + # Quantize the weights for FC1 + gemm1_weights_fp4_bytes, gemm1_scales_fp4_bytes, gemm1_scales_global = ( + quant_fp4_batches(gemm1_weights, num_experts, self.is_mxfp4, True) + ) + + # Quantize the weights for FC2 + gemm2_weights_fp4_bytes, gemm2_scales_fp4_bytes, gemm2_scales_global = ( + quant_fp4_batches(gemm2_weights, num_experts, self.is_mxfp4, True) + ) + + return { + "hidden_states_scale_global": hidden_states_scale_global, + "gemm1_weights": gemm1_weights_fp4_bytes, + "gemm1_scales": gemm1_scales_fp4_bytes, + "gemm1_scales_global": gemm1_scales_global, + "gemm2_weights": gemm2_weights_fp4_bytes, + "gemm2_scales": gemm2_scales_fp4_bytes, + "gemm2_scales_global": gemm2_scales_global, + } + + def quantize_inputs( + self, hidden_states, hidden_states_scale_global, is_swizzling=True + ): + if self.quant_mode == QuantMode.FP4_MXFP4_MXFP8: + """Quantize hidden states to MxFP8 format.""" + hidden_states_quant, hidden_states_scale = mxfp8_quantize( + hidden_states, is_swizzling + ) + hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape( + *hidden_states.shape[:-1], -1 + ) + return { + "hidden_states": hidden_states_quant, + "hidden_states_scale": hidden_states_scale, + } + elif self.quant_mode == QuantMode.FP4_NVFP4_NVFP4: + """Quantize hidden states to NvFP4 format using pre-computed global scale.""" + ( + hidden_states_fp4_bytes, + hidden_states_scale_fp4_bytes, + _, + ) = quant_fp4( + hidden_states, hidden_states_scale_global, False, is_swizzling + ) + hidden_states_scale_fp4_bytes = hidden_states_scale_fp4_bytes.view( + torch.float8_e4m3fn + ).reshape(*hidden_states.shape[:-1], -1) + + return { + "hidden_states": hidden_states_fp4_bytes, + "hidden_states_scale": hidden_states_scale_fp4_bytes, + } + else: # bf16 + return { + "hidden_states": hidden_states.to(torch.bfloat16), + "hidden_states_scale": None, + } + + def prepare_static_weights_for_kernel( + self, + args_dequant, + args, + gemm1_weights_orig, + gemm2_weights_orig, + hidden_size, + intermediate_size, + num_experts, + weight_processing, + ): + """Prepare quantized weights for kernel (done offline with weights).""" + use_ue8m0 = self.is_mxfp4 + epilogue_tile_m = 128 # FIXME: this depends on the kernel internals + + # Quantize weights with linear layout for kernels + _, gemm1_scales_linear_fp4_bytes, _ = quant_fp4_batches( + gemm1_weights_orig, num_experts, use_ue8m0, False + ) + _, gemm2_scales_linear_fp4_bytes, _ = quant_fp4_batches( + gemm2_weights_orig, num_experts, use_ue8m0, False + ) + + # Convert quantized weights to proper formats + intermediate_size_factor = 2 if is_gated_activation(args.activation_type) else 1 + gemm1_weights_fp4 = args.gemm1_weights.view(torch.float8_e4m3fn).reshape( + num_experts, intermediate_size_factor * intermediate_size, hidden_size // 2 + ) # packed fp4 + gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view( + torch.float8_e4m3fn + ).reshape( + num_experts, + intermediate_size_factor * intermediate_size, + hidden_size // self.sf_vec_size, + ) # fp8 scaling factors + + gemm2_weights_fp4 = args.gemm2_weights.view(torch.float8_e4m3fn).reshape( + num_experts, hidden_size, intermediate_size // 2 + ) # packed fp4 + gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view( + torch.float8_e4m3fn + ).reshape( + num_experts, hidden_size, intermediate_size // self.sf_vec_size + ) # fp8 scaling factors + + # Using cached permute index calculation can speed up weights preprocessing + gemm1_weights_fp4_shuffled = [] + gemm1_scales_fp4_shuffled = [] + gemm2_weights_fp4_shuffled = [] + gemm2_scales_fp4_shuffled = [] + for i in range(num_experts): + # Calculate the permute indices for the following: + # 1. Reorder rows of W1 and scales for fused gated activation + # 2. Shuffle weights and scaling factors for transposed mma output + # for both w3_w1 and w2 weights and scale factors + permute_indices = _maybe_get_cached_w3_w1_permute_indices( + self._cache_permute_indices, + gemm1_weights_fp4[i].view(torch.uint8), + epilogue_tile_m, + is_gated_act_gemm=is_gated_activation(args.activation_type), + ) + gemm1_weights_fp4_shuffled.append( + gemm1_weights_fp4[i] + .view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)] + .contiguous() + ) + + permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices( + self._cache_permute_indices, + gemm1_scales_linear_fp4[i].view(torch.uint8), + epilogue_tile_m, + num_elts_per_sf=16, + is_gated_act_gemm=is_gated_activation(args.activation_type), + ) + gemm1_scales_fp4_shuffled.append( + block_scale_interleave( + gemm1_scales_linear_fp4[i] + .view(torch.uint8)[ + permute_sf_indices.to(gemm1_scales_linear_fp4.device) + ] + .contiguous() + ) + ) + + permute_indices = get_w2_permute_indices_with_cache( + self._cache_permute_indices, + gemm2_weights_fp4[i].view(torch.uint8), + epilogue_tile_m, + ) + gemm2_weights_fp4_shuffled.append( + gemm2_weights_fp4[i] + .view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)] + .contiguous() + ) + + permute_sf_indices = get_w2_permute_indices_with_cache( + self._cache_permute_indices, + gemm2_scales_linear_fp4[i].view(torch.uint8), + epilogue_tile_m, + num_elts_per_sf=16, + ) + gemm2_scales_fp4_shuffled.append( + block_scale_interleave( + gemm2_scales_linear_fp4[i] + .view(torch.uint8)[ + permute_sf_indices.to(gemm2_scales_linear_fp4.device) + ] + .contiguous() + ) + ) + + # Stack weights for all experts + gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled) + gemm1_scales_fp4_shuffled = ( + torch.stack(gemm1_scales_fp4_shuffled) + .view(torch.float8_e4m3fn) + .reshape( + num_experts, + intermediate_size_factor * intermediate_size, + hidden_size // self.sf_vec_size, + ) + ) + + gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled) + gemm2_scales_fp4_shuffled = ( + torch.stack(gemm2_scales_fp4_shuffled) + .view(torch.float8_e4m3fn) + .reshape(num_experts, hidden_size, intermediate_size // self.sf_vec_size) + ) + + # Calculate scaling factors that depend on weights + if is_gated_activation(args.activation_type): + scale_c_fc1 = ( + args_dequant.c_global_sf + * (1.0 / args.gemm1_scales_global) + * (1.0 / args.hidden_states_scale_global) + ) + else: + scale_c_fc1 = torch.full_like( + args.gemm1_scales_global, args_dequant.c_global_sf + ) + scale_gate_fc1 = (1.0 / args.gemm1_scales_global) * ( + 1.0 / args.hidden_states_scale_global + ) + scale_c_fc2 = (1.0 / args_dequant.c_global_sf) * ( + 1.0 / args.gemm2_scales_global + ) + + return { + "gemm1_weights_fp4_shuffled": gemm1_weights_fp4_shuffled, + "gemm1_scales_fp4_shuffled": gemm1_scales_fp4_shuffled, + "gemm2_weights_fp4_shuffled": gemm2_weights_fp4_shuffled, + "gemm2_scales_fp4_shuffled": gemm2_scales_fp4_shuffled, + "scale_c_fc1": scale_c_fc1, + "scale_gate_fc1": scale_gate_fc1, + "scale_c_fc2": scale_c_fc2, + } + + def call_moe( + self, static_data, hidden_states_orig, hidden_states_scale_global, **kwargs + ): + """Call MoE using CUDA graph for maximum performance (create, capture, launch).""" + # Extract runtime arguments + expert_logits = kwargs["expert_logits"] + routing_bias = kwargs["routing_bias"] + num_experts = kwargs["num_experts"] + top_k = kwargs["top_k"] + n_groups = kwargs["n_groups"] + top_k_groups = kwargs["top_k_groups"] + intermediate_size = kwargs["intermediate_size"] + routed_scaling = kwargs["routed_scaling"] + activation_type = kwargs["activation_type"] + routing_method_type = kwargs["routing_method_type"] + enable_autotune = kwargs.get("enable_autotune", True) + gemm1_bias = kwargs["gemm1_bias"] + gemm2_bias = kwargs["gemm2_bias"] + norm_topk_prob = kwargs.get("norm_topk_prob", True) + + # Create CUDA graph configuration + config = { + "hidden_states_scale_global": hidden_states_scale_global, + "num_experts": num_experts, + "top_k": top_k, + "n_groups": n_groups, + "top_k_groups": top_k_groups, + "intermediate_size": intermediate_size, + "routed_scaling": routed_scaling, + "activation_type": activation_type, + "routing_method_type": routing_method_type, + "enable_autotune": enable_autotune, + "gemm1_bias": gemm1_bias, + "gemm2_bias": gemm2_bias, + "norm_topk_prob": norm_topk_prob, + } + + runtime_args = { + "expert_logits": expert_logits, + "routing_bias": routing_bias, + } + + # Create, capture and launch CUDA graph in one shot + cuda_graph = CUDAGraphMoE(self, static_data, **config) + try: + cuda_graph.capture(hidden_states_orig, **runtime_args) + output = cuda_graph.launch(hidden_states_orig) + return output[0].to(torch.float) + finally: + cuda_graph.cleanup() + + def compute_reference(self, args): + return run_moe_reference_fp4(args, self.quant_mode) + + def get_tolerances(self): + """Get FP4-specific accuracy tolerances.""" + return {"atol": 0.1, "rtol": 0.85, "percent": 0.92} + + +# ==================================================================================== +# MxInt4 Block Scale Quantization Implementation +# ==================================================================================== + + +def mxint4_quantize( + x: torch.Tensor, sf_vec_size: int = 32 +) -> tuple[torch.Tensor, torch.Tensor]: + x_reshaped = x.reshape(-1, sf_vec_size) + x_max = x_reshaped.max(dim=-1, keepdim=True)[0].to(torch.float32) + x_min = x_reshaped.min(dim=-1, keepdim=True)[0].to(torch.float32) + x_max = x_max * 8.0 / 7.0 + amax = torch.where(x_max > -x_min, x_max, -x_min) + scales = amax / 8.0 + x_scaled = x_reshaped * scales.reciprocal() + x_int8 = ( + x_scaled.round().clamp(-8, 7).to(torch.int8).reshape(-1, sf_vec_size // 2, 2) + ) + x_int4 = (x_int8[..., 0] & 0x0F) | ((x_int8[..., 1] & 0x0F) << 4) + return x_int4.reshape(*x.shape[:-1], x.shape[-1] // 2), scales.reshape( + -1, sf_vec_size + ) + + +class MxInt4BlockScaleMoe(Moe): + """MxInt4 MoE implementation with block scaling (DeepSeek style).""" + + @property + def quant_mode(self) -> QuantMode: + return QuantMode.MXINT4_BF16_BF16 + + def quantize_weights(self, gemm1_weights, gemm2_weights, hidden_states_sample): + """Quantize weights to MxInt4 with block scaling.""" + num_experts = gemm1_weights.shape[0] + intermediate_size = gemm1_weights.shape[1] // 2 + hidden_size = gemm1_weights.shape[ + 2 + ] # [num_experts, 2*intermediate_size, hidden_size] + + # Quantize weights to MxInt4 + sf_vec_size = 32 + gemm1_weights_int4, gemm1_scales = mxint4_quantize(gemm1_weights, sf_vec_size) + gemm2_weights_int4, gemm2_scales = mxint4_quantize(gemm2_weights, sf_vec_size) + gemm1_scales = gemm1_scales.to(torch.bfloat16).reshape( + num_experts, + 2 * intermediate_size, + hidden_size // sf_vec_size, + ) + gemm2_scales = gemm2_scales.to(torch.bfloat16).reshape( + num_experts, hidden_size, intermediate_size // sf_vec_size + ) + return { + "hidden_states_scale_global": None, + "gemm1_weights": gemm1_weights_int4, + "gemm2_weights": gemm2_weights_int4, + "gemm1_scales": gemm1_scales, + "gemm2_scales": gemm2_scales, + "gemm1_scales_global": None, + "gemm2_scales_global": None, + } + + def quantize_inputs(self, hidden_states, *unused_args): + """No scaling for hidden states.""" + return { + "hidden_states": hidden_states.to(torch.bfloat16), + "hidden_states_scale": None, + } + + def prepare_static_weights_for_kernel( + self, + args_dequant, + args, + gemm1_weights_orig, + gemm2_weights_orig, + hidden_size, + intermediate_size, + num_experts, + weight_processing, + ): + """Prepare quantized weights for kernel (done offline with weights).""" + + epilogue_tile_m = 128 + gemm1_weights_mxint4_shuffled = [] + gemm1_scales_shuffled = [] + gemm2_weights_mxint4_shuffled = [] + gemm2_scales_shuffled = [] + + for i in range(num_experts): + # Calculate the permute indices for the following: + # 1. Reorder rows of W1 and scales for fused gated activation + # 2. Shuffle weights and scaling factors for transposed mma output + # for both w3_w1 and w2 weights and scale factors + permute_indices = _maybe_get_cached_w3_w1_permute_indices( + self._cache_permute_indices, + args.gemm1_weights[i].view(torch.uint8), + epilogue_tile_m, + ) + gemm1_weights_shuffled = ( + args.gemm1_weights[i] + .view(torch.uint8)[permute_indices.to(args.gemm1_weights.device)] + .contiguous() + ) + permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices( + self._cache_permute_indices, + args.gemm1_scales[i].view(torch.bfloat16), + epilogue_tile_m, + num_elts_per_sf=32, + ) + gemm1_scales_shuffled.append( + block_scale_interleave( + args.gemm1_scales[i] + .view(torch.bfloat16)[ + permute_sf_indices.to(args.gemm1_scales.device) + ] + .contiguous() + ) + ) + + permute_indices = get_w2_permute_indices_with_cache( + self._cache_permute_indices, + args.gemm2_weights[i].view(torch.uint8), + epilogue_tile_m, + ) + gemm2_weights_shuffled = ( + args.gemm2_weights[i] + .view(torch.uint8)[permute_indices.to(args.gemm2_weights.device)] + .contiguous() + ) + + permute_sf_indices = get_w2_permute_indices_with_cache( + self._cache_permute_indices, + args.gemm2_scales[i].view(torch.bfloat16), + epilogue_tile_m, + num_elts_per_sf=16, + ) + gemm2_scales_shuffled.append( + block_scale_interleave( + args.gemm2_scales[i] + .view(torch.bfloat16)[ + permute_sf_indices.to(args.gemm2_scales.device) + ] + .contiguous() + ) + ) + + block_k = 128 + gemm1_weights_shuffled = convert_to_block_layout( + gemm1_weights_shuffled, block_k + ) + gemm2_weights_shuffled = convert_to_block_layout( + gemm2_weights_shuffled.view(torch.uint8), block_k + ) + + gemm1_weights_mxint4_shuffled.append(gemm1_weights_shuffled) + gemm2_weights_mxint4_shuffled.append(gemm2_weights_shuffled) + + gemm1_weights_mxint4_shuffled = torch.stack(gemm1_weights_mxint4_shuffled) + gemm2_weights_mxint4_shuffled = torch.stack(gemm2_weights_mxint4_shuffled) + gemm1_scales_shuffled = torch.stack(gemm1_scales_shuffled).view(torch.bfloat16) + gemm2_scales_shuffled = torch.stack(gemm2_scales_shuffled).view(torch.bfloat16) + + return { + "gemm1_weights": gemm1_weights_mxint4_shuffled, + "gemm1_scales": gemm1_scales_shuffled, + "gemm2_weights": gemm2_weights_mxint4_shuffled, + "gemm2_scales": gemm2_scales_shuffled, + } + + def call_moe( + self, static_data, hidden_states_orig, hidden_states_scale_global, **kwargs + ): + """Call MoE with runtime input quantization + kernel execution (done at runtime).""" + expert_logits = kwargs["expert_logits"] + routing_bias = kwargs["routing_bias"] + num_experts = kwargs["num_experts"] + top_k = kwargs["top_k"] + n_groups = kwargs["n_groups"] + top_k_groups = kwargs["top_k_groups"] + intermediate_size = kwargs["intermediate_size"] + routing_method_type = kwargs["routing_method_type"] + enable_autotune = kwargs.get("enable_autotune", True) + routed_scaling = kwargs.get("routed_scaling", 1.0) + norm_topk_prob = kwargs.get("norm_topk_prob", True) + + # Use autotuner for optimal kernel selection + with autotune(enable_autotune): + output = trtllm_mxint4_block_scale_moe( + expert_logits, # float + routing_bias, + hidden_states_orig, + static_data["gemm1_weights"], + static_data["gemm1_scales"], + None, + None, + None, + static_data["gemm2_weights"], + static_data["gemm2_scales"], + num_experts, + top_k, + n_groups, + top_k_groups, + intermediate_size, + 0, + num_experts, + routed_scaling, + routing_method_type=routing_method_type, + tune_max_num_tokens=TUNE_MAX_NUM_TOKENS, + norm_topk_prob=norm_topk_prob, + ) + return output[0].to(torch.float) + + def compute_reference(self, args): + return run_moe_reference_mxint4(args) + + def get_tolerances(self): + """Get MXINT4-specific accuracy tolerances.""" + return {"atol": 0.1, "rtol": 0.85, "percent": 0.925} + + +# ==================================================================================== +# FP8 Block Scale Quantization Implementation +# ==================================================================================== + + +class FP8BlockScaleMoe(Moe): + """FP8 MoE implementation with block scaling (DeepSeek style or MxFp8 x MxFp8).""" + + def __init__( + self, fp8_quantization_type: QuantMode = QuantMode.FP8_BLOCK_SCALE_DEEPSEEK + ): + super().__init__() + self.fp8_quantization_type = fp8_quantization_type + + @property + def quant_mode(self) -> QuantMode: + return self.fp8_quantization_type + + def quantize_weights(self, gemm1_weights, gemm2_weights, hidden_states_sample): + """Quantize weights to FP8 with block scaling.""" + num_experts = gemm1_weights.shape[0] + # Non-gated activations (e.g. Relu2) use [E, I, H], gated use [E, 2I, H]. + intermediate_size = gemm2_weights.shape[2] + intermediate_size_factor = gemm1_weights.shape[1] // intermediate_size + assert intermediate_size_factor in (1, 2) + hidden_size = gemm1_weights.shape[ + 2 + ] # [num_experts, 2*intermediate_size, hidden_size] + + if self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_DEEPSEEK: + # Quantize weights to FP8 + gemm1_weights_fp8 = gemm1_weights.to(torch.float8_e4m3fn) + gemm1_scales = 2 * torch.rand( + ( + num_experts, + intermediate_size_factor * intermediate_size // 128, + hidden_size // 128, + ), + device="cuda", + ).to(torch.float) + + gemm2_weights_fp8 = gemm2_weights.to(torch.float8_e4m3fn) + gemm2_scales = 2 * torch.rand( + (num_experts, hidden_size // 128, intermediate_size // 128), + device="cuda", + ).to(torch.float) + elif self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_MXFP8: + gemm1_weights_fp8, gemm1_scales = mxfp8_quantize_batches( + gemm1_weights, False + ) + gemm2_weights_fp8, gemm2_scales = mxfp8_quantize_batches( + gemm2_weights, False + ) + else: + raise ValueError( + f"Unsupported FP8 quantization type: {self.fp8_quantization_type}" + ) + + return { + "hidden_states_scale_global": None, # Block scales computed at runtime + "gemm1_weights": gemm1_weights_fp8, + "gemm1_scales": gemm1_scales, + "gemm1_scales_global": None, + "gemm2_weights": gemm2_weights_fp8, + "gemm2_scales": gemm2_scales, + "gemm2_scales_global": None, + } + + def quantize_inputs( + self, + hidden_states: torch.Tensor, + hidden_states_scale_global: torch.Tensor = None, + is_swizzling: bool = False, + ): + """For FP8 block scaling, no pre-quantization - everything happens at runtime.""" + + def to_float8_blockwise( + x, + block_size_m=128, + block_size_n=128, + dtype=torch.float8_e4m3fn, + transpose_scale=True, + is_blockm=False, + is_blockn=True, + ): + assert x.dtype == torch.bfloat16 + x = x.contiguous() + assert x.dim() == 2 + m, n = x.shape + + m_tile = block_size_m if is_blockm else 1 + n_tile = block_size_n if is_blockn else 1 + num_blocks_m = m // m_tile + num_blocks_n = n // n_tile + + # Initialize output tensors + quantized_x = torch.empty_like(x, dtype=dtype, device=x.device) + scales = torch.empty( + (num_blocks_m, num_blocks_n), dtype=torch.float32, device=x.device + ) + + # Quantize tensor in blocks + finfo = torch.finfo(dtype) + for i in range(num_blocks_m): + for j in range(num_blocks_n): + # Determine block slices + start_m, end_m = i * m_tile, min((i + 1) * m_tile, m) + start_n, end_n = j * n_tile, min((j + 1) * n_tile, n) + + # Extract the block + block = x[start_m:end_m, start_n:end_n] + + # Per-block quantization logic + min_val, max_val = block.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) + scale = finfo.max / amax + + # Quantize the block and store the scale + quantized_block = (block * scale).clamp( + min=finfo.min, max=finfo.max + ) + quantized_x[start_m:end_m, start_n:end_n] = quantized_block.to( + dtype + ) + scales[i, j] = scale.float().reciprocal() + + if transpose_scale: + scales = scales.t() + + return quantized_x, scales + + # todo(Yingyi):quantize bf16 to fp8 + if self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_DEEPSEEK: + hidden_states_quant, hidden_states_scale = to_float8_blockwise( + hidden_states + ) + elif self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_MXFP8: + hidden_states_quant, hidden_states_scale = mxfp8_quantize( + hidden_states, is_swizzling + ) + hidden_states_scale = hidden_states_scale.view(torch.uint8).reshape( + *hidden_states.shape[:-1], -1 + ) + else: + raise ValueError( + f"Unsupported FP8 quantization type: {self.fp8_quantization_type}" + ) + return { + "hidden_states": hidden_states_quant, + "hidden_states_scale": hidden_states_scale, + } + + def prepare_static_weights_for_kernel( + self, + args_dequant, + args, + gemm1_weights_orig, + gemm2_weights_orig, + hidden_size, + intermediate_size, + num_experts, + weight_processing, + ): + """Prepare quantized weights for kernel (done offline with weights).""" + + # Use shuffled weights with BlockMajorK layout for better performance + use_shuffled_weight = weight_processing["use_shuffled_weight"] + weight_layout = weight_processing["layout"] + + if use_shuffled_weight: + # FIXME: this depends on the kernel internals + epilogue_tile_m = ( + 64 + if self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_DEEPSEEK + else 128 + ) + + intermediate_size_factor = ( + 2 if is_gated_activation(args.activation_type) else 1 + ) + + gemm1_weights_fp8_interleaved = args.gemm1_weights.clone() + gemm1_scales_fp8_interleaved = args.gemm1_scales.clone() + if self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_MXFP8: + # Reorder rows of W1 only for fused gated activation. + gemm1_weights_fp8_interleaved = [] + gemm1_scales_fp8_interleaved = [] + for i in range(num_experts): + gemm1_w = ( + args.gemm1_weights[i] + .clone() + .reshape(intermediate_size_factor * intermediate_size, -1) + ) + gemm1_s = ( + args.gemm1_scales[i] + .clone() + .reshape(intermediate_size_factor * intermediate_size, -1) + ) + if is_gated_activation(args.activation_type): + gemm1_w = reorder_rows_for_gated_act_gemm(gemm1_w) + gemm1_s = reorder_rows_for_gated_act_gemm(gemm1_s) + gemm1_weights_fp8_interleaved.append(gemm1_w) + gemm1_scales_fp8_interleaved.append(gemm1_s) + + # Stack weights and scales for all experts + gemm1_weights_fp8_interleaved = torch.stack( + gemm1_weights_fp8_interleaved + ).reshape(args.gemm1_weights.shape) + gemm1_scales_fp8_interleaved = torch.stack( + gemm1_scales_fp8_interleaved + ).reshape(args.gemm1_scales.shape) + + gemm1_weights_fp8_shuffled = [] + gemm2_weights_fp8_shuffled = [] + gemm1_scales_fp8_shuffled = [] + gemm2_scales_fp8_shuffled = [] + for i in range(num_experts): + tmp_weights1 = shuffle_matrix_a( + gemm1_weights_fp8_interleaved[i].view(torch.uint8), epilogue_tile_m + ) + tmp_weights2 = shuffle_matrix_a( + args.gemm2_weights[i].view(torch.uint8), epilogue_tile_m + ) + if self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_MXFP8: + tmp_scales1 = shuffle_matrix_sf_a( + gemm1_scales_fp8_interleaved[i] + .view(torch.uint8) + .reshape(intermediate_size_factor * intermediate_size, -1), + epilogue_tile_m, + ) + tmp_scales2 = shuffle_matrix_sf_a( + args.gemm2_scales[i].view(torch.uint8).reshape(hidden_size, -1), + epilogue_tile_m, + ) + gemm1_scales_fp8_shuffled.append(tmp_scales1) + gemm2_scales_fp8_shuffled.append(tmp_scales2) + + if weight_layout == WeightLayout.BlockMajorK: + block_k = 128 + tmp_weights1 = convert_to_block_layout(tmp_weights1, block_k) + tmp_weights2 = convert_to_block_layout(tmp_weights2, block_k) + + gemm1_weights_fp8_shuffled.append(tmp_weights1) + gemm2_weights_fp8_shuffled.append(tmp_weights2) + + kernel_gemm1_weights = torch.stack(gemm1_weights_fp8_shuffled).view( + torch.float8_e4m3fn + ) + kernel_gemm2_weights = torch.stack(gemm2_weights_fp8_shuffled).view( + torch.float8_e4m3fn + ) + if self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_MXFP8: + kernel_gemm1_scales = torch.stack(gemm1_scales_fp8_shuffled).reshape( + args.gemm1_scales.shape + ) + kernel_gemm2_scales = torch.stack(gemm2_scales_fp8_shuffled).reshape( + args.gemm2_scales.shape + ) + else: + kernel_gemm1_scales = args.gemm1_scales + kernel_gemm2_scales = args.gemm2_scales + else: + kernel_gemm1_weights = args.gemm1_weights + kernel_gemm2_weights = args.gemm2_weights + kernel_gemm1_scales = args.gemm1_scales + kernel_gemm2_scales = args.gemm2_scales + + return { + "gemm1_weights": kernel_gemm1_weights, + "gemm1_scales": kernel_gemm1_scales, + "gemm2_weights": kernel_gemm2_weights, + "gemm2_scales": kernel_gemm2_scales, + "use_shuffled_weight": use_shuffled_weight, + "weight_layout": weight_layout, + } + + def call_moe( + self, static_data, hidden_states_orig, hidden_states_scale_global, **kwargs + ): + """Call MoE with runtime block scale generation + kernel execution.""" + expert_logits = kwargs["expert_logits"] + routing_bias = kwargs["routing_bias"] + num_experts = kwargs["num_experts"] + top_k = kwargs["top_k"] + n_groups = kwargs["n_groups"] + top_k_groups = kwargs["top_k_groups"] + intermediate_size = kwargs["intermediate_size"] + routed_scaling = kwargs["routed_scaling"] + routing_method_type = kwargs["routing_method_type"] + activation_type = kwargs["activation_type"] + enable_autotune = kwargs.get("enable_autotune", True) + enable_pdl = kwargs.get("enable_pdl") + hidden_states_scale = kwargs["hidden_states_scale"] + hidden_states_quant = kwargs["hidden_states_quant"] + norm_topk_prob = kwargs.get("norm_topk_prob", True) + + # Generate block scales and quantize hidden states at runtime + hidden_states_fp8 = hidden_states_quant.to(torch.float8_e4m3fn) + assert not torch.isnan(hidden_states_fp8.float()).any(), ( + "NaN detected in hidden_states_fp8" + ) + + if self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_MXFP8: + quantization_mode = Fp8QuantizationType.MxFp8 + elif self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_DEEPSEEK: + quantization_mode = Fp8QuantizationType.DeepSeekFp8 + else: + raise ValueError( + f"Unsupported FP8 quantization type: {self.fp8_quantization_type}" + ) + + # Use autotuner for optimal kernel selection + with autotune(enable_autotune): + output = trtllm_fp8_block_scale_moe( + expert_logits, + routing_bias, + hidden_states_fp8, + hidden_states_scale, + static_data["gemm1_weights"], + static_data["gemm1_scales"], + static_data["gemm2_weights"], + static_data["gemm2_scales"], + num_experts, + top_k, + n_groups, + top_k_groups, + intermediate_size, + 0, + num_experts, + routed_scaling, + routing_method_type, + use_shuffled_weight=static_data["use_shuffled_weight"], + weight_layout=static_data["weight_layout"], + enable_pdl=enable_pdl, + tune_max_num_tokens=TUNE_MAX_NUM_TOKENS, + fp8_quantization_type=quantization_mode, + activation_type=activation_type, + norm_topk_prob=norm_topk_prob, + ) + return output.to(torch.float) + + def compute_reference(self, args): + """FP8 block-scale reference implementation.""" + if self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_DEEPSEEK: + return run_moe_reference_dsfp8(args) + elif self.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_MXFP8: + return run_moe_reference_mxfp8(args) + else: + raise ValueError( + f"Unsupported FP8 quantization type: {self.fp8_quantization_type}" + ) + + def get_tolerances(self): + """Get FP8 block-scale accuracy tolerances.""" + return {"atol": 0.1, "rtol": 0.85, "percent": 0.79} + + +# ==================================================================================== +# FP8 Per-Tensor Quantization Implementation +# ==================================================================================== + + +class FP8PerTensorMoe(Moe): + """FP8 MoE implementation with per-tensor scaling (Llama4 style).""" + + @property + def quant_mode(self) -> QuantMode: + return QuantMode.FP8_PER_TENSOR + + def quantize_weights(self, gemm1_weights, gemm2_weights, hidden_states_sample): + """Quantize weights to FP8 per-tensor and compute global scale factors.""" + # Compute global scale factor for hidden states (offline calibration) + hidden_states_global_scale = calculate_fp8_global_scale_factor( + hidden_states_sample + ) + + # Quantize to FP8 per-tensor + gemm1_weights_quant, gemm1_global_scales = quant_fp8_per_tensor_batches( + gemm1_weights + ) + gemm2_weights_quant, gemm2_global_scales = quant_fp8_per_tensor_batches( + gemm2_weights + ) + + return { + "hidden_states_scale_global": hidden_states_global_scale, + "gemm1_weights": gemm1_weights_quant, + "gemm1_scales": None, + "gemm1_scales_global": gemm1_global_scales, + "gemm2_weights": gemm2_weights_quant, + "gemm2_scales": None, + "gemm2_scales_global": gemm2_global_scales, + } + + def quantize_inputs(self, hidden_states, hidden_states_scale_global): + """Quantize hidden states to FP8 per-tensor using pre-computed global scale.""" + # Quantize to FP8 per-tensor using pre-computed global scale factor + hidden_states_quant, _ = quant_fp8_per_tensor( + hidden_states, hidden_states_scale_global + ) + + return { + "hidden_states": hidden_states_quant, + "hidden_states_scale": None, + } + + def prepare_static_weights_for_kernel( + self, + args_dequant, + args, + gemm1_weights_orig, + gemm2_weights_orig, + hidden_size, + intermediate_size, + num_experts, + weight_processing, + ): + """Prepare quantized weights for kernel (done offline with weights).""" + # FIXME: this depends on the kernel internals + epilogue_tile_m = 128 + + # Reorder rows of W1 for fused gated activation + gemm1_weights_fp8_interleaved = [] + for i in range(num_experts): + if is_gated_activation(args.activation_type): + weights = reorder_rows_for_gated_act_gemm(args.gemm1_weights[i].clone()) + else: + weights = args.gemm1_weights[i].clone() + gemm1_weights_fp8_interleaved.append(weights) + + # Stack weights and scales for all experts + gemm1_weights_fp8_interleaved = torch.stack( + gemm1_weights_fp8_interleaved + ).reshape( + num_experts, + (2 if is_gated_activation(args.activation_type) else 1) * intermediate_size, + hidden_size, + ) + + # Shuffle weights and scaling factors for transposed mma output + gemm1_weights_fp8_shuffled = [] + gemm2_weights_fp8_shuffled = [] + for i in range(num_experts): + gemm1_weights_fp8_shuffled.append( + shuffle_matrix_a( + gemm1_weights_fp8_interleaved[i].view(torch.uint8), epilogue_tile_m + ) + ) + + gemm2_weights_fp8_shuffled.append( + shuffle_matrix_a( + args.gemm2_weights[i].view(torch.uint8), epilogue_tile_m + ) + ) + + # Stack weights for all experts + gemm1_weights_fp8_shuffled = torch.stack(gemm1_weights_fp8_shuffled).view( + torch.float8_e4m3fn + ) + gemm2_weights_fp8_shuffled = torch.stack(gemm2_weights_fp8_shuffled).view( + torch.float8_e4m3fn + ) + + # Calculate scaling factors that depend on weights + if is_gated_activation(args.activation_type): + scale_c_fc1 = ( + args_dequant.c_global_sf + * (1.0 / args.gemm1_scales_global) + * (1.0 / args.hidden_states_scale_global) + ) + else: + scale_c_fc1 = torch.full_like( + args.gemm1_scales_global, args_dequant.c_global_sf + ) + scale_gate_fc1 = (1.0 / args.gemm1_scales_global) * ( + 1.0 / args.hidden_states_scale_global + ) + scale_c_fc2 = (1.0 / args_dequant.c_global_sf) * ( + 1.0 / args.gemm2_scales_global + ) + + return { + "gemm1_weights": gemm1_weights_fp8_shuffled, + "gemm2_weights": gemm2_weights_fp8_shuffled, + "scale_c_fc1": scale_c_fc1, + "scale_gate_fc1": scale_gate_fc1, + "scale_c_fc2": scale_c_fc2, + } + + def call_moe( + self, static_data, hidden_states_orig, hidden_states_scale_global, **kwargs + ): + """Call MoE with runtime input quantization + kernel execution (done at runtime).""" + expert_logits = kwargs["expert_logits"] + routing_bias = kwargs["routing_bias"] + num_experts = kwargs["num_experts"] + top_k = kwargs["top_k"] + n_groups = kwargs["n_groups"] + top_k_groups = kwargs["top_k_groups"] + intermediate_size = kwargs["intermediate_size"] + routed_scaling = kwargs["routed_scaling"] + routing_method_type = kwargs["routing_method_type"] + enable_autotune = kwargs.get("enable_autotune", True) + activation_type = kwargs["activation_type"] + norm_topk_prob = kwargs.get("norm_topk_prob", True) + + # Quantize to FP8 per-tensor using pre-computed global scale factor + hidden_states_fp8, _ = quant_fp8_per_tensor( + hidden_states_orig, hidden_states_scale_global + ) + + # Use autotuner for optimal kernel selection + with autotune(enable_autotune): + output = trtllm_fp8_per_tensor_scale_moe( + expert_logits, + routing_bias, + hidden_states_fp8, + static_data["gemm1_weights"], + static_data["scale_c_fc1"], + static_data["scale_gate_fc1"], + static_data["gemm2_weights"], + static_data["scale_c_fc2"], + num_experts, + top_k, + n_groups, + top_k_groups, + intermediate_size, + 0, + num_experts, + routed_scaling, + routing_method_type + == RoutingMethodType.Llama4, # Use_routing_scales_on_input + routing_method_type, + tune_max_num_tokens=TUNE_MAX_NUM_TOKENS, + activation_type=activation_type, + norm_topk_prob=norm_topk_prob, + ) + + return output.to(torch.float) + + def compute_reference(self, args): + """FP8 per-tensor reference implementation.""" + return run_moe_reference_per_tensor_scale_fp8(args) + + def get_tolerances(self): + """Get FP8 per-tensor accuracy tolerances.""" + return {"atol": 0.1, "rtol": 0.85, "percent": 0.92} + + +# ==================================================================================== +# BF16 Implementation +# ==================================================================================== + + +class BF16Moe(Moe): + """BF16 MoE implementation.""" + + @property + def quant_mode(self) -> QuantMode: + return QuantMode.BF16 + + def quantize_weights(self, gemm1_weights, gemm2_weights, hidden_states_sample): + """No scaling for weights.""" + return { + "hidden_states_scale_global": None, + "gemm1_weights": gemm1_weights.to(torch.bfloat16), + "gemm1_scales": None, + "gemm1_scales_global": None, + "gemm2_weights": gemm2_weights.to(torch.bfloat16), + "gemm2_scales": None, + "gemm2_scales_global": None, + } + + def quantize_inputs(self, hidden_states, *unused_args): + """No scaling for hidden states.""" + return { + "hidden_states": hidden_states.to(torch.bfloat16), + "hidden_states_scale": None, + } + + def prepare_static_weights_for_kernel( + self, + args_dequant, + args, + gemm1_weights_orig, + gemm2_weights_orig, + hidden_size, + intermediate_size, + num_experts, + weight_processing, + ): + """Prepare quantized weights for kernel (done offline with weights).""" + + # Use shuffled weights with BlockMajorK layout for better performance + use_shuffled_weight = weight_processing["use_shuffled_weight"] + weight_layout = weight_processing["layout"] + + if use_shuffled_weight: + # FIXME: this depends on the kernel internals + epilogue_tile_m = 128 + + # Reorder rows of W1 for fused gated activation and shuffle for both W1 and W2 + # Using cached permute index calculation can speed up weights preprocessing + gemm1_weights_bf16_shuffled = [] + gemm2_weights_bf16_shuffled = [] + for i in range(num_experts): + permute_indices = _maybe_get_cached_w3_w1_permute_indices( + self._cache_permute_indices, + args.gemm1_weights[i].view(torch.uint8), + epilogue_tile_m, + is_gated_act_gemm=is_gated_activation(args.activation_type), + ) + tmp_weights1 = ( + args.gemm1_weights[i] + .view(torch.uint8)[permute_indices.to(args.gemm1_weights.device)] + .contiguous() + ) + + permute_indices = get_w2_permute_indices_with_cache( + self._cache_permute_indices, + args.gemm2_weights[i].view(torch.uint8), + epilogue_tile_m, + ) + tmp_weights2 = ( + args.gemm2_weights[i] + .view(torch.uint8)[permute_indices.to(args.gemm2_weights.device)] + .contiguous() + ) + + if weight_layout == WeightLayout.BlockMajorK: + block_k = 128 + tmp_weights1 = convert_to_block_layout( + tmp_weights1.view(torch.uint8), block_k + ) + tmp_weights2 = convert_to_block_layout( + tmp_weights2.view(torch.uint8), block_k + ) + + gemm1_weights_bf16_shuffled.append(tmp_weights1.view(torch.bfloat16)) + gemm2_weights_bf16_shuffled.append(tmp_weights2.view(torch.bfloat16)) + + # Stack weights for all experts + gemm1_weights_bf16_shuffled = ( + torch.stack(gemm1_weights_bf16_shuffled) + .view(torch.bfloat16) + .contiguous() + ) + gemm2_weights_bf16_shuffled = ( + torch.stack(gemm2_weights_bf16_shuffled) + .view(torch.bfloat16) + .contiguous() + ) + + return { + "gemm1_weights": gemm1_weights_bf16_shuffled, + "gemm2_weights": gemm2_weights_bf16_shuffled, + "use_shuffled_weight": use_shuffled_weight, + "weight_layout": weight_layout, + } + + def call_moe( + self, static_data, hidden_states_orig, hidden_states_scale_global, **kwargs + ): + """Call MoE with runtime input quantization + kernel execution (done at runtime).""" + expert_logits = kwargs["expert_logits"] + routing_bias = kwargs["routing_bias"] + num_experts = kwargs["num_experts"] + top_k = kwargs["top_k"] + n_groups = kwargs["n_groups"] + top_k_groups = kwargs["top_k_groups"] + intermediate_size = kwargs["intermediate_size"] + routed_scaling = kwargs["routed_scaling"] + routing_method_type = kwargs["routing_method_type"] + enable_autotune = kwargs.get("enable_autotune", True) + activation_type = kwargs["activation_type"] + norm_topk_prob = kwargs.get("norm_topk_prob", True) + + # Use autotuner for optimal kernel selection + with autotune(enable_autotune): + output = trtllm_bf16_moe( + expert_logits, # float + routing_bias, + hidden_states_orig, + static_data["gemm1_weights"], + static_data["gemm2_weights"], + num_experts, + top_k, + n_groups, + top_k_groups, + intermediate_size, + 0, + num_experts, + routed_scaling, + use_shuffled_weight=static_data["use_shuffled_weight"], + weight_layout=static_data["weight_layout"], + routing_method_type=routing_method_type, + tune_max_num_tokens=TUNE_MAX_NUM_TOKENS, + activation_type=activation_type, + norm_topk_prob=norm_topk_prob, + ) + return output.to(torch.float) + + def compute_reference(self, args): + """BF16 reference implementation.""" + return run_moe_reference_bf16(args) + + def get_tolerances(self): + """Get BF16 accuracy tolerances.""" + return {"atol": 0.1, "rtol": 0.85, "percent": 0.925} + + +# ==================================================================================== +# Quantizer Factory +# ==================================================================================== +def get_moe_impl(quant_mode: QuantMode): + """Factory function to get the appropriate MoE implementation.""" + if quant_mode == QuantMode.FP8_BLOCK_SCALE_DEEPSEEK: + return FP8BlockScaleMoe( + fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_DEEPSEEK + ) + elif quant_mode == QuantMode.FP8_BLOCK_SCALE_MXFP8: + return FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_MXFP8) + elif quant_mode == QuantMode.FP8_PER_TENSOR: + return FP8PerTensorMoe() + else: + return FP4Moe(quant_mode) + + +class moe_args: + """Arguments container for MoE operations.""" + + def __init__( + self, + num_tokens, + num_experts, + hidden_size, + intermediate_size, + top_k, + padding, + hidden_states, + hidden_states_scale, + hidden_states_scale_global, + expert_logits, + gemm1_weights, + gemm1_scales, + gemm1_scales_global, + gemm2_weights, + gemm2_scales, + gemm2_scales_global, + permute_info, + use_routing_scales_on_input, + activation_type, + gemm1_bias=None, + gemm2_bias=None, + ): + self.num_tokens = num_tokens + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.top_k = top_k + self.padding = padding + self.hidden_states = hidden_states + self.hidden_states_scale = hidden_states_scale + self.hidden_states_scale_global = hidden_states_scale_global + self.expert_logits = expert_logits + self.gemm1_weights = gemm1_weights + self.gemm1_scales = gemm1_scales + self.gemm1_scales_global = gemm1_scales_global + self.gemm2_weights = gemm2_weights + self.gemm2_scales = gemm2_scales + self.gemm2_scales_global = gemm2_scales_global + self.permute_info = permute_info + self.use_routing_scales_on_input = use_routing_scales_on_input + self.activation_type = activation_type + self.gemm1_bias = gemm1_bias + self.gemm2_bias = gemm2_bias + + +class moe_args_dequant: + """Arguments container for dequantized MoE operations.""" + + def __init__( + self, + num_tokens, + num_experts, + hidden_size, + intermediate_size, + top_k, + padding, + hidden_states, + expert_logits, + gemm1_weights, + gemm2_weights, + permute_info, + use_routing_scales_on_input, + activation_type, + hidden_states_scale=None, + gemm1_bias=None, + gemm2_bias=None, + ): + self.num_tokens = num_tokens + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.top_k = top_k + self.padding = padding + self.hidden_states = hidden_states + self.expert_logits = expert_logits + self.gemm1_weights = gemm1_weights + self.gemm2_weights = gemm2_weights + self.permute_info = permute_info + self.use_routing_scales_on_input = use_routing_scales_on_input + self.activation_type = activation_type + self.hidden_states_scale = hidden_states_scale + self.gemm1_bias = gemm1_bias + self.gemm2_bias = gemm2_bias + + +def routing_reference(expertLogits, topK, padding): + """Reference routing implementation for permutation calculation.""" + originalDevice = expertLogits.device + expertLogits = expertLogits.cpu() + numTokens, numExperts = expertLogits.shape + assert topK <= numExperts + + numTokensPerExpert = torch.zeros(numExperts, dtype=torch.int64) + expandedTokenIdxToExpert = -torch.ones(numTokens * topK, dtype=torch.int64) + expandedTokenIdxToIdxInExpert = -torch.ones(numTokens * topK, dtype=torch.int64) + + topKLogits, topKIndices = torch.topk(expertLogits, topK, dim=1) + for tokenIdx in range(numTokens): + for k in range(topK): + expandedIdx = tokenIdx * topK + k + expertIndex = topKIndices[tokenIdx, k] + expandedTokenIdxToExpert[expandedIdx] = expertIndex + expandedTokenIdxToIdxInExpert[expandedIdx] = numTokensPerExpert[expertIndex] + numTokensPerExpert[expertIndex] += 1 + + paddedTokensPerExpertPrefixSum = torch.zeros(numExperts + 1, dtype=torch.int64) + for ii in range(numExperts): + + def divUpMul(a, b): + return (a + b - 1) // b * b + + paddedTokensPerExpertPrefixSum[ii + 1] = paddedTokensPerExpertPrefixSum[ + ii + ] + divUpMul(numTokensPerExpert[ii], padding) + permutedBufferSize = paddedTokensPerExpertPrefixSum[numExperts] + + expandedTokenIdxToPermutedIdx = -torch.ones(numTokens * topK, dtype=torch.int64) + permutedIdxToExpandedIdx = -torch.ones(permutedBufferSize, dtype=torch.int64) + permutedIdxToTokenIdx = -torch.ones(permutedBufferSize, dtype=torch.int64) + for tokenIdx in range(numTokens): + for k in range(topK): + expandedIdx = tokenIdx * topK + k + expert = expandedTokenIdxToExpert[expandedIdx] + offsetWithinExpert = expandedTokenIdxToIdxInExpert[expandedIdx] + offsetForExpert = paddedTokensPerExpertPrefixSum[expert] + permutedIdx = offsetForExpert + offsetWithinExpert + + expandedTokenIdxToPermutedIdx[expandedIdx] = permutedIdx + permutedIdxToExpandedIdx[permutedIdx] = expandedIdx + permutedIdxToTokenIdx[permutedIdx] = tokenIdx + return { + "paddedTokensPerExpertPrefixSum": paddedTokensPerExpertPrefixSum.to( + originalDevice + ), + "permutedBufferSize": permutedBufferSize.item(), + "expandedTokenIdxToPermutedIdx": expandedTokenIdxToPermutedIdx.to( + originalDevice + ), + "permutedIdxToExpandedIdx": permutedIdxToExpandedIdx.to(originalDevice), + "numTokensPerExpert": numTokensPerExpert.to(originalDevice), + "expandedTokenIdxToExpert": expandedTokenIdxToExpert.to(originalDevice), + "topKLogits": topKLogits.to(originalDevice), + "permutedIdxToTokenIdx": permutedIdxToTokenIdx.to(originalDevice), + "topKIndices": topKIndices.to(originalDevice), + } + + +def noaux_tc_ref(logits, bias, n_group, topk_group, top_k, routed_scaling_factor): + """DeepSeek-style no-aux routing reference implementation.""" + scores = F.sigmoid(logits) + scores_with_bias = scores + bias + if n_group > 1: + scores_shape = list(scores_with_bias.shape) + group_scores = torch.sum( + torch.topk( + scores_with_bias.view( + scores_shape[:-1] + [n_group, scores_shape[-1] // n_group] + ), + k=2, + dim=-1, + largest=True, + sorted=True, + )[0], + dim=-1, + ) + _, group_idx = torch.topk( + group_scores, k=topk_group, dim=-1, largest=True, sorted=True + ) + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(-1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(scores_shape[:-1] + [n_group, scores_shape[-1] // n_group]) + .reshape(scores_shape) + ) + scores_with_bias = scores_with_bias * score_mask + + _, topk_idx = torch.topk( + scores_with_bias, k=top_k, dim=-1, largest=True, sorted=True + ) + new_mask = torch.zeros_like(scores) + new_mask.scatter_(-1, topk_idx, 1) + scores = scores * new_mask + score_sum = torch.sum(scores, dim=-1, keepdim=True) + 1e-20 + scores = scores / score_sum * routed_scaling_factor + return scores + + +def routing_reference_no_aux( + expert_logits, + routing_bias, + top_k, + n_groups, + top_k_groups, + routed_scaling, + padding, + use_routing_scales_on_input=False, +): + """Tiered TopK routing used by DeepSeek.""" + routing_logits = expert_logits.to(dtype=torch.float, device="cuda") + if use_routing_scales_on_input: + # if using routing scales on input, topK == 1 and the score is a plain sigmoid + scores = F.sigmoid(routing_logits) + else: + scores = noaux_tc_ref( + routing_logits, routing_bias, n_groups, top_k_groups, top_k, routed_scaling + ) + permute_info = routing_reference(scores, top_k, padding) + return permute_info, scores + + +def routing_reference_default(expert_logits, top_k, num_experts, padding): + """Softmax -> TopK routing reference (Default method).""" + scores = torch.nn.functional.softmax(expert_logits.float(), dim=-1) + topk_values, topk_idx = torch.topk(scores, k=top_k, dim=-1) + topk_values = topk_values.to(expert_logits.dtype) + + scores = torch.zeros_like(scores, dtype=expert_logits.dtype) + for i in range(topk_idx.shape[0]): + for j in range(topk_idx.shape[1]): + scores[i, topk_idx[i, j]] = topk_values[i, j] + permute_info = routing_reference(scores, top_k, padding) + return permute_info, scores + + +def routing_reference_renormalize(expert_logits, top_k, num_experts, padding): + """TopK -> Softmax routing reference.""" + topk_values, topk_idx = torch.topk(expert_logits, k=top_k, dim=-1) + topk_values = torch.nn.functional.softmax(topk_values.float(), dim=-1) + + new_mask = torch.zeros_like(expert_logits) + new_mask.scatter_(-1, topk_idx, 1) + scores = expert_logits * new_mask + + for i in range(topk_idx.shape[0]): + for j in range(topk_idx.shape[1]): + scores[i, topk_idx[i, j]] = topk_values[i, j] + permute_info = routing_reference(scores, top_k, padding) + return permute_info, scores + + +def routing_reference_renormalize_naive(expert_logits, top_k, num_experts, padding): + """Softmax->TopK -> Normalize routing reference.""" + norm_topk_prob = True + scores = torch.nn.functional.softmax(expert_logits.float(), dim=-1) + topk_values, topk_idx = torch.topk(scores, k=top_k, dim=-1) + + if norm_topk_prob: # only diff with mixtral sparse moe block! + topk_values /= topk_values.sum(dim=-1, keepdim=True) + topk_values = topk_values.to(expert_logits.dtype) + scores = scores.to(expert_logits.dtype) + + new_mask = torch.zeros_like(expert_logits) + new_mask.scatter_(-1, topk_idx, 1) + scores = expert_logits * new_mask + + for i in range(topk_idx.shape[0]): + for j in range(topk_idx.shape[1]): + scores[i, topk_idx[i, j]] = topk_values[i, j] + permute_info = routing_reference(scores, top_k, padding) + return permute_info, scores + + +def routing_reference_topk(expert_logits, top_k, num_experts, padding): + """TopK only (no softmax) routing reference.""" + topk_values, topk_idx = torch.topk(expert_logits, k=top_k, dim=-1) + + new_mask = torch.zeros_like(expert_logits) + new_mask.scatter_(-1, topk_idx, 1) + scores = expert_logits * new_mask + + for i in range(topk_idx.shape[0]): + for j in range(topk_idx.shape[1]): + scores[i, topk_idx[i, j]] = topk_values[i, j] + permute_info = routing_reference(scores, top_k, padding) + return permute_info, scores + + +def routing_reference_sigmoid_renorm( + expert_logits, top_k, num_experts, padding, norm_topk_prob=True +): + """Sigmoid -> TopK -> Renormalize routing reference.""" + sigmoid_scores = torch.sigmoid(expert_logits.float()) + topk_values, topk_idx = torch.topk(sigmoid_scores, k=top_k, dim=-1) + + if norm_topk_prob: + topk_values = topk_values / (topk_values.sum(dim=-1, keepdim=True) + 1e-20) + topk_values = topk_values.to(expert_logits.dtype) + + scores = torch.zeros_like(sigmoid_scores, dtype=expert_logits.dtype) + for i in range(topk_idx.shape[0]): + for j in range(topk_idx.shape[1]): + scores[i, topk_idx[i, j]] = topk_values[i, j] + permute_info = routing_reference(scores, top_k, padding) + return permute_info, scores + + +def routing_reference_minimax2( + expert_logits, routing_bias, top_k, num_experts, padding +): + """Sigmoid + Bias -> TopK -> ScaledSumNormalize routing reference (MiniMax2). + Bias affects expert selection but NOT the final weights. + Weights = sigmoid(logit) / (sum_of_selected_sigmoid + 1e-20). + """ + sigmoid_scores = torch.sigmoid(expert_logits.float()) + selection_scores = sigmoid_scores.clone() + if routing_bias is not None: + selection_scores = selection_scores + routing_bias.float() + _, topk_idx = torch.topk(selection_scores, k=top_k, dim=-1) + + # Weights use un-biased sigmoid scores + raw_weights = torch.gather(sigmoid_scores, -1, topk_idx) + raw_weights = raw_weights / (raw_weights.sum(dim=-1, keepdim=True) + 1e-20) + raw_weights = raw_weights.to(expert_logits.dtype) + + scores = torch.zeros_like(sigmoid_scores, dtype=expert_logits.dtype) + for i in range(topk_idx.shape[0]): + for j in range(topk_idx.shape[1]): + scores[i, topk_idx[i, j]] = raw_weights[i, j] + permute_info = routing_reference(scores, top_k, padding) + return permute_info, scores + + +def check_accuracy(a, b, atol, rtol, percent): + """Unified accuracy checking function with detailed error reporting.""" + if not torch.isfinite(a).all(): + raise Exception("Non-finite values in reference output") + if not torch.isfinite(b).all(): + raise Exception("Non-finite values in actual output") + assert a.shape == b.shape, f"Shape mismatch: {a.shape} vs {b.shape}" + + close = torch.isclose(a, b, atol=atol, rtol=rtol) + match_ratio = close.float().mean() + if match_ratio >= percent: + return + + mismatch_percent = 1.0 - match_ratio.item() + if mismatch_percent > 1 - percent: + raise Exception( + f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} " + f"(threshold: {1 - percent:.4f})" + ) + + +# ==================================================================================== +# FP4 Quantization Functions +# ==================================================================================== + + +def calculate_fp4_global_scale_factor(tensor, use_ue8m0=False): + """ + Calculate FP4 global scale factor for a tensor. + + NOTE: In production, global scale factors are typically obtained offline during: + - Post-Training Quantization (PTQ) calibration process + - Quantization-Aware Training (QAT) process + + This function is used here for testing/reference purposes. + Formula: (448 * 6) represents max representable value in FP4 format. + """ + if use_ue8m0: + return torch.tensor(1.0, dtype=torch.float32) + else: + return (448 * 6) / tensor.float().abs().nan_to_num().max() + + +def e2m1_and_ufp8_scale_batches( + mat_fp4: torch.Tensor, + scale_tensor: torch.Tensor, + global_scale_tensor: torch.Tensor, + sf_vec_size: int, + ufp8_type: int = 1, +): + """Batch FP4 dequantization helper.""" + num_batches = mat_fp4.size(0) + scale_tensor = scale_tensor.view(num_batches, -1) + + tensors = [ + e2m1_and_ufp8sf_scale_to_float( + mat_fp4[b, :, :].cpu(), + scale_tensor[b, :].cpu().reshape(-1), + global_scale_tensor[b].cpu(), + sf_vec_size, + ufp8_type, + True, # is_sf_swizzled_layout + ) + for b in range(num_batches) + ] + + result = torch.stack(tensors) + return result + + +def quant_fp4(a, a_global_sf, use_ue8m0=False, is_sf_swizzled_layout=True): + """ + Quantize FP4 with pre-computed global scale factor. + + This function expects global scale factors that have been pre-computed offline + during PTQ/QAT calibration process. The global scale factor should NOT be + computed at runtime to avoid performance overhead. + + Pure function - same inputs always produce same outputs. + """ + sf_vec_size = 32 if use_ue8m0 else 16 + + a_fp4, a_sf = fp4_quantize( + a.cuda(), a_global_sf.cuda(), sf_vec_size, use_ue8m0, is_sf_swizzled_layout + ) + + return a_fp4, a_sf, a_global_sf + + +def quant_fp4_batches(a, num_experts, use_ue8m0=False, is_sf_swizzled_layout=True): + """FP4 batch quantization function with centralized global scale factor calculation.""" + quant_a = [] + sfs = [] + global_sfs = [] + for i in range(num_experts): + # Use centralized global scale factor calculation + a_global_sf = calculate_fp4_global_scale_factor(a[i], use_ue8m0) + a_fp4, a_sf, _ = quant_fp4(a[i], a_global_sf, use_ue8m0, is_sf_swizzled_layout) + quant_a.append(a_fp4) + sfs.append(a_sf) + global_sfs.append(a_global_sf) + + result_quant_a = torch.stack(quant_a) + result_sfs = torch.stack(sfs) + result_global_sfs = torch.stack(global_sfs) + + return result_quant_a, result_sfs, result_global_sfs + + +def quant_dequant_fp4(a, use_ue8m0=False, is_sf_swizzled_layout=True): + """FP4 quantize-dequantize roundtrip function with centralized global scale factor calculation.""" + # Use centralized global scale factor calculation + a_global_sf = calculate_fp4_global_scale_factor(a, use_ue8m0) + sf_vec_size = 32 if use_ue8m0 else 16 + + a_fp4, a_sf = fp4_quantize( + a.cuda(), a_global_sf.cuda(), sf_vec_size, use_ue8m0, is_sf_swizzled_layout + ) + + a_pt = e2m1_and_ufp8sf_scale_to_float( + a_fp4.cpu(), + a_sf.cpu().reshape(-1), + (1 / a_global_sf).cpu(), + sf_vec_size, + 1 if not use_ue8m0 else 0, # ufp8_type + is_sf_swizzled_layout, + ) + + return a_pt.cuda(), a_global_sf + + +# ==================================================================================== +# FP8 Quantization Functions +# ==================================================================================== + + +def calculate_fp8_global_scale_factor(tensor): + """ + Calculate FP8 global scale factor for a tensor. + + NOTE: In production, global scale factors are typically obtained offline during: + - Post-Training Quantization (PTQ) calibration process + - Quantization-Aware Training (QAT) process + + This function is used here for testing/reference purposes. + Formula: 448 represents max representable value in FP8 E4M3 format. + """ + return 448 / tensor.float().abs().nan_to_num().max() + + +def quant_fp8_per_tensor(a, a_global_sf): + """ + Quantize FP8 per-tensor with pre-computed global scale factor. + + This function expects global scale factors that have been pre-computed offline + during PTQ/QAT calibration process. The global scale factor should NOT be + computed at runtime to avoid performance overhead. + + Pure function - same inputs always produce same outputs. + """ + a_fp8 = (a * a_global_sf).to(torch.float8_e4m3fn) + return a_fp8, a_global_sf + + +def quant_fp8_per_tensor_batches(a): + """FP8 per-tensor batch quantization function with centralized global scale factor calculation.""" + num_batches = a.size(0) + a_quant = [] + a_scales = [] + + for i in range(num_batches): + # Use centralized global scale factor calculation + a_global_sf = calculate_fp8_global_scale_factor(a[i]) + a_fp8, _ = quant_fp8_per_tensor(a[i], a_global_sf) + a_quant.append(a_fp8) + a_scales.append(a_global_sf) + + result_a_quant = torch.stack(a_quant) + result_a_scales = torch.stack(a_scales) + + return result_a_quant, result_a_scales + + +def quant_dequant_per_tensor_fp8(a): + """FP8 per-tensor quantize-dequantize roundtrip function with centralized global scale factor calculation.""" + # Use centralized global scale factor calculation + a_global_sf = calculate_fp8_global_scale_factor(a) + a_fp8, _ = quant_fp8_per_tensor(a, a_global_sf) + a_pt = a_fp8.to(torch.float) / a_global_sf + return a_pt.cuda(), a_global_sf + + +def dequant_reference_dsfp8(input, scale, transpose_scale, block_m, block_n): + """Reference FP8 block-scale dequantization.""" + input = input.to(torch.float) + scale = scale.to(torch.float) + if transpose_scale: + scale = scale.t() + + m, n = input.shape + m_tile = 128 if block_m else 1 + n_tile = 128 if block_n else 1 + + assert m % m_tile == 0 + assert n % n_tile == 0 + assert scale.shape == (m // m_tile, n // n_tile) + + # Expand scale to match input dimensions using tensor operations + if m_tile > 1: + scale = torch.repeat_interleave(scale, m_tile, dim=0) + if n_tile > 1: + scale = torch.repeat_interleave(scale, n_tile, dim=1) + + # Element-wise multiplication (equivalent to the nested loop logic) + output = input * scale + return output + + +def mxfp8_quantize_batches(a, is_swizzling=True): + """MxFp8 batch quantization function with centralized global scale factor calculation.""" + num_batches = a.size(0) + a_quant = [] + a_scales = [] + for i in range(num_batches): + mx_fp8_quant, mx_fp8_scale = mxfp8_quantize(a[i], is_swizzling) + a_quant.append(mx_fp8_quant) + a_scales.append(mx_fp8_scale.view(torch.uint8)) + + result_a_quant = torch.stack(a_quant) + result_a_scales = torch.stack(a_scales) + + return result_a_quant, result_a_scales + + +def mxfp8_dequantize_batches(a, a_scales, is_swizzling=True): + """MxFp8 batch dequantization function.""" + num_batches = a.size(0) + a_dequant = [] + for i in range(num_batches): + mx_fp8_dequant = mxfp8_dequantize_host( + a[i].cpu().view(torch.uint8), + a_scales[i].cpu().view(torch.uint8).reshape(-1), + is_swizzling, + ) + a_dequant.append(mx_fp8_dequant.cuda()) + + result_a_dequant = torch.stack(a_dequant) + + return result_a_dequant + + +# ==================================================================================== +# Common MoE Reference Implementation +# ==================================================================================== + + +def run_moe_dequant(args, quant_mode: QuantMode): + """Common dequantized MoE reference implementation.""" + # Permute + total_num_padded_tokens = args.permute_info["permutedBufferSize"] + expanded_idx_to_permuted_idx = args.permute_info[ + "expandedTokenIdxToPermutedIdx" + ].cpu() + num_tokens_per_expert = args.permute_info["numTokensPerExpert"].cpu() + permute_output = torch.full( + (total_num_padded_tokens, args.hidden_size), float("nan"), device="cuda" + ).to(torch.float) + for i in range(args.num_tokens): + for j in range(args.top_k): + permuted_idx = expanded_idx_to_permuted_idx[i * args.top_k + j] + permute_output[permuted_idx] = args.hidden_states[i] + + # Gemm1 + gemm1_output = torch.full( + ( + total_num_padded_tokens, + (2 if is_gated_activation(args.activation_type) else 1) + * args.intermediate_size, + ), + float("nan"), + device="cuda", + ).to(torch.float) + i = 0 + for expert_idx in range(args.num_experts): + my_num_tokens = num_tokens_per_expert[expert_idx] + if my_num_tokens == 0: + continue + my_a = permute_output[i : i + my_num_tokens] + my_b = args.gemm1_weights[expert_idx] + my_c = my_a @ my_b.t() + if args.gemm1_bias is not None: + my_c = my_c + args.gemm1_bias[expert_idx].to(torch.float) + gemm1_output[i : i + my_num_tokens] = my_c + i += my_num_tokens + i = (i + args.padding - 1) // args.padding * args.padding + + if args.use_routing_scales_on_input: + assert args.top_k == 1 + # For each token and its top_k experts + for token_idx in range(args.num_tokens): + for k in range(args.top_k): + # Get the permuted index for this token's k-th expert + expanded_idx = token_idx * args.top_k + k + permuted_idx = expanded_idx_to_permuted_idx[expanded_idx] + expert_weight = args.permute_info["topKLogits"].to(torch.float) + # Get the expert weight for this token and expert + weight = expert_weight[token_idx, k] + # Scale the corresponding row in gemm1_output + gemm1_output[permuted_idx] *= weight + + # Activation + activation_output = torch.full( + (total_num_padded_tokens, args.intermediate_size), float("nan"), device="cuda" + ).to(torch.float) + + activation_type = args.activation_type + activation_type_to_func = { + ActivationType.Swiglu: F.silu, + ActivationType.Geglu: F.gelu, + ActivationType.Relu2: lambda x: F.relu(x) ** 2, + } + activation_func = activation_type_to_func[activation_type] + + i = 0 + for expert_idx in range(args.num_experts): + my_num_tokens = num_tokens_per_expert[expert_idx] + if my_num_tokens == 0: + continue + my_a = gemm1_output[i : i + my_num_tokens] + if is_gated_activation(args.activation_type): + my_x1 = my_a[:, : args.intermediate_size] + my_x2 = my_a[:, args.intermediate_size :] + activation_output[i : i + my_num_tokens] = activation_func(my_x2) * my_x1 + else: + my_x1 = my_a[:, : args.intermediate_size] + activation_output[i : i + my_num_tokens] = activation_func(my_x1) + i += my_num_tokens + i = (i + args.padding - 1) // args.padding * args.padding + + if quant_mode == QuantMode.FP4_NVFP4_NVFP4: + # Use centralized function for activation quantization + activation_output, c_global_sf = quant_dequant_fp4( + activation_output.to(torch.bfloat16), False, True + ) + activation_output = activation_output.to(torch.float) + args.c_global_sf = c_global_sf + elif quant_mode == QuantMode.FP8_PER_TENSOR: + activation_output, c_global_sf = quant_dequant_per_tensor_fp8( + activation_output.to(torch.bfloat16) + ) + activation_output = activation_output.to(torch.float) + args.c_global_sf = c_global_sf + elif ( + quant_mode == QuantMode.FP4_MXFP4_MXFP8 + or quant_mode == QuantMode.FP8_BLOCK_SCALE_MXFP8 + ): + activation_output, scale_bytes = mxfp8_quantize( + activation_output.to(torch.bfloat16), True + ) + scale_bytes = scale_bytes.view(torch.uint8).reshape(-1).cpu() + activation_output = ( + mxfp8_dequantize_host( + activation_output.cpu().view(torch.uint8), scale_bytes + ) + .cuda() + .to(torch.float) + ) + args.c_global_sf = 1.0 + else: # Bf16, MxFp4xBf16, MxInt4xBf16 + activation_output = activation_output.to(torch.bfloat16).to(torch.float) + args.c_global_sf = 1.0 + + # Gemm2 + gemm2_output = torch.full( + (total_num_padded_tokens, args.hidden_size), float("nan"), device="cuda" + ).to(torch.float) + i = 0 + for expert_idx in range(args.num_experts): + my_num_tokens = num_tokens_per_expert[expert_idx] + if my_num_tokens == 0: + continue + my_a = activation_output[i : i + my_num_tokens] + my_b = args.gemm2_weights[expert_idx] + my_c = my_a @ my_b.t() + if args.gemm2_bias is not None: + my_c = my_c + args.gemm2_bias[expert_idx].to(torch.float) + gemm2_output[i : i + my_num_tokens] = my_c + i += my_num_tokens + i = (i + args.padding - 1) // args.padding * args.padding + + # Finalize + expert_weight = args.permute_info["topKLogits"].to(torch.float) + finalize_output = torch.full( + (args.num_tokens, args.hidden_size), float("nan"), device="cuda" + ).to(torch.float) + for i in range(args.num_tokens): + acc = torch.zeros(args.hidden_size, dtype=torch.float, device="cuda") + for top_k_idx in range(args.top_k): + expanded_idx = i * args.top_k + top_k_idx + permuted_idx = expanded_idx_to_permuted_idx[expanded_idx] + original_vector = gemm2_output[permuted_idx] + weight = ( + expert_weight[i, top_k_idx] + if not args.use_routing_scales_on_input + else 1.0 + ) + acc += original_vector * weight + finalize_output[i] = acc + return finalize_output + + +# ==================================================================================== +# Quantization-Specific Reference Implementations +# ==================================================================================== + + +def run_moe_reference_fp4(args, quant_mode: QuantMode): + sf_vec_size = 16 if quant_mode == QuantMode.FP4_NVFP4_NVFP4 else 32 + ufp8_type_weights = 1 if quant_mode == QuantMode.FP4_NVFP4_NVFP4 else 0 + + if quant_mode == QuantMode.FP4_NVFP4_NVFP4: + hidden_states_dequant = e2m1_and_ufp8sf_scale_to_float( + args.hidden_states.cpu(), + args.hidden_states_scale.cpu().view(torch.uint8).reshape(-1), + (1 / args.hidden_states_scale_global).cpu(), + sf_vec_size, + ufp8_type_weights, + True, # is_sf_swizzled_layout + ).cuda() + elif quant_mode == QuantMode.FP4_MXFP4_MXFP8: + hidden_states_dequant = mxfp8_dequantize_host( + args.hidden_states.cpu().view(torch.uint8), + args.hidden_states_scale.cpu().view(torch.uint8).reshape(-1), + True, # is_sf_swizzled_layout + ).cuda() + else: + hidden_states_dequant = args.hidden_states.to(torch.bfloat16).to(torch.float) + + gemm1_weights_dequant = e2m1_and_ufp8_scale_batches( + args.gemm1_weights, + args.gemm1_scales, + 1 / args.gemm1_scales_global, + sf_vec_size, + ufp8_type_weights, + ).cuda() + + gemm2_weights_dequant = e2m1_and_ufp8_scale_batches( + args.gemm2_weights, + args.gemm2_scales, + 1 / args.gemm2_scales_global, + sf_vec_size, + ufp8_type_weights, + ).cuda() + + args_dequant = moe_args_dequant( + args.num_tokens, + args.num_experts, + args.hidden_size, + args.intermediate_size, + args.top_k, + args.padding, + hidden_states_dequant, + args.expert_logits, + gemm1_weights_dequant, + gemm2_weights_dequant, + args.permute_info, + args.use_routing_scales_on_input, + args.activation_type, + gemm1_bias=args.gemm1_bias, + gemm2_bias=args.gemm2_bias, + ) + + return run_moe_dequant(args_dequant, quant_mode), args_dequant + + +def run_moe_reference_mxfp8(args): + hidden_states_dequant = mxfp8_dequantize_host( + args.hidden_states.cpu().view(torch.uint8), + args.hidden_states_scale.cpu().view(torch.uint8).reshape(-1), + False, # is_sf_swizzled_layout + ).cuda() + + gemm1_weights_dequant = mxfp8_dequantize_batches( + args.gemm1_weights, + args.gemm1_scales, + False, + ).cuda() + + gemm2_weights_dequant = mxfp8_dequantize_batches( + args.gemm2_weights, + args.gemm2_scales, + False, + ).cuda() + + args_dequant = moe_args_dequant( + args.num_tokens, + args.num_experts, + args.hidden_size, + args.intermediate_size, + args.top_k, + args.padding, + hidden_states_dequant, + args.expert_logits, + gemm1_weights_dequant, + gemm2_weights_dequant, + args.permute_info, + args.use_routing_scales_on_input, + args.activation_type, + ) + + return run_moe_dequant(args_dequant, QuantMode.FP8_BLOCK_SCALE_MXFP8), args_dequant + + +def run_moe_reference_dsfp8(args): + """FP8 block-scale reference implementation (DeepSeek style).""" + # Generate block scales at runtime for FP8 block scaling + + def dequant_reference_dsfp8(input, scale, transpose_scale, block_m, block_n): + """Reference FP8 block-scale dequantization.""" + input = input.to(torch.float) + scale = scale.to(torch.float) + if transpose_scale: + scale = scale.t() + + m, n = input.shape + m_tile = 128 if block_m else 1 + n_tile = 128 if block_n else 1 + + assert m % m_tile == 0 + assert n % n_tile == 0 + assert scale.shape == (m // m_tile, n // n_tile) + + # Expand scale to match input dimensions using tensor operations + if m_tile > 1: + scale = torch.repeat_interleave(scale, m_tile, dim=0) + if n_tile > 1: + scale = torch.repeat_interleave(scale, n_tile, dim=1) + + # Element-wise multiplication (equivalent to the nested loop logic) + output = input * scale + return output + + # todo(Yingyi): use original hidden_states?? + hidden_states_dequant = dequant_reference_dsfp8( + args.hidden_states, args.hidden_states_scale, True, False, True + ) + + gemm1_weights_dequant = {} + for i in range(args.num_experts): + gemm1_weights_dequant[i] = dequant_reference_dsfp8( + args.gemm1_weights[i], args.gemm1_scales[i], False, True, True + ) + + gemm2_weights_dequant = {} + for i in range(args.num_experts): + gemm2_weights_dequant[i] = dequant_reference_dsfp8( + args.gemm2_weights[i], args.gemm2_scales[i], False, True, True + ) + + args_dequant = moe_args_dequant( + args.num_tokens, + args.num_experts, + args.hidden_size, + args.intermediate_size, + args.top_k, + args.padding, + hidden_states_dequant, + args.expert_logits, + gemm1_weights_dequant, + gemm2_weights_dequant, + args.permute_info, + args.use_routing_scales_on_input, + args.activation_type, + gemm1_bias=args.gemm1_bias, + gemm2_bias=args.gemm2_bias, + ) + + return run_moe_dequant( + args_dequant, QuantMode.FP8_BLOCK_SCALE_DEEPSEEK + ), args_dequant + + +def run_moe_reference_per_tensor_scale_fp8(args): + """FP8 per-tensor reference implementation.""" + hidden_states_dequant = ( + args.hidden_states.to(torch.float) / args.hidden_states_scale_global + ) + + gemm1_weights_dequant = {} + for i in range(args.num_experts): + gemm1_weights_dequant[i] = ( + args.gemm1_weights[i].to(torch.float) / args.gemm1_scales_global[i] + ) + + gemm2_weights_dequant = {} + for i in range(args.num_experts): + gemm2_weights_dequant[i] = ( + args.gemm2_weights[i].to(torch.float) / args.gemm2_scales_global[i] + ) + + args_dequant = moe_args_dequant( + args.num_tokens, + args.num_experts, + args.hidden_size, + args.intermediate_size, + args.top_k, + args.padding, + hidden_states_dequant, + args.expert_logits, + gemm1_weights_dequant, + gemm2_weights_dequant, + args.permute_info, + args.use_routing_scales_on_input, + args.activation_type, + gemm1_bias=args.gemm1_bias, + gemm2_bias=args.gemm2_bias, + ) + + return run_moe_dequant(args_dequant, QuantMode.FP8_PER_TENSOR), args_dequant + + +def run_moe_reference_bf16(args): + """BF16 reference implementation.""" + + # no scaling for hidden states and weights + hidden_states_dequant = args.hidden_states.to(torch.float) + gemm1_weights_dequant = {} + for i in range(args.num_experts): + gemm1_weights_dequant[i] = args.gemm1_weights[i].to(torch.float) + gemm2_weights_dequant = {} + for i in range(args.num_experts): + gemm2_weights_dequant[i] = args.gemm2_weights[i].to(torch.float) + + args_dequant = moe_args_dequant( + args.num_tokens, + args.num_experts, + args.hidden_size, + args.intermediate_size, + args.top_k, + args.padding, + hidden_states_dequant, + args.expert_logits, + gemm1_weights_dequant, + gemm2_weights_dequant, + args.permute_info, + args.use_routing_scales_on_input, + args.activation_type, + gemm1_bias=args.gemm1_bias, + gemm2_bias=args.gemm2_bias, + ) + + return run_moe_dequant(args_dequant, QuantMode.BF16), args_dequant + + +def run_moe_reference_mxint4(args): + sf_vec_size = 32 + + hidden_states_dequant = args.hidden_states.to(torch.bfloat16).to(torch.float) + + num_experts = args.gemm1_weights.shape[0] + + def dequantize(weights, scales): + k = weights.shape[-1] * 2 + n = weights.shape[-2] + # Unpack two 4-bit values (stored in two's-complement) from each byte + weights_int8 = ( + torch.stack([weights & 0x0F, (weights >> 4) & 0x0F], dim=-1) + .reshape(num_experts, n, k) + .to(torch.int8) + ) + + # Interpret nibbles as signed 4-bit two's-complement values in [-8, 7] + weights_int8 = torch.where(weights_int8 < 8, weights_int8, weights_int8 - 16) + + weights_float = weights_int8.to(torch.float) + scales_expanded = ( + scales.to(torch.bfloat16) + .to(torch.float) + .repeat_interleave(sf_vec_size, dim=-1) + .reshape(weights_float.shape) + ) + return weights_float * scales_expanded + + gemm1_weights_dequant = dequantize(args.gemm1_weights, args.gemm1_scales) + gemm2_weights_dequant = dequantize(args.gemm2_weights, args.gemm2_scales) + + args_dequant = moe_args_dequant( + args.num_tokens, + args.num_experts, + args.hidden_size, + args.intermediate_size, + args.top_k, + args.padding, + hidden_states_dequant, + args.expert_logits, + gemm1_weights_dequant, + gemm2_weights_dequant, + args.permute_info, + args.use_routing_scales_on_input, + args.activation_type, + gemm1_bias=args.gemm1_bias, + gemm2_bias=args.gemm2_bias, + ) + + return run_moe_dequant(args_dequant, QuantMode.MXINT4_BF16_BF16), args_dequant + + +def _compute_moe_actual_unified(moe_impl, args_dequant, args, **kwargs): + """Unified actual computation that delegates to implementation-specific methods.""" + # 1. Prepare static weights for the kernel (offline processing) + static_data = moe_impl.prepare_static_weights_for_kernel( + args_dequant, + args, + kwargs["gemm1_weights_orig"], + kwargs["gemm2_weights_orig"], + args.hidden_size, + args.intermediate_size, + args.num_experts, + kwargs["weight_processing"], + ) + + # 2. Call MoE with runtime input quantization + kernel execution + kernel_kwargs = { + "expert_logits": kwargs["expert_logits"], + "routing_bias": kwargs["routing_bias"], + "num_experts": args.num_experts, + "num_tokens": args.num_tokens, + "hidden_size": args.hidden_size, + "top_k": args.top_k, + "n_groups": kwargs["n_groups"], + "top_k_groups": kwargs["top_k_groups"], + "intermediate_size": args.intermediate_size, + "routed_scaling": kwargs["routed_scaling"], + "routing_method_type": kwargs["routing_method_type"], + "do_finalize": True, + "activation_type": args.activation_type, + "hidden_states_scale": args.hidden_states_scale, + "hidden_states_quant": kwargs["hidden_states_quant"], + "enable_autotune": kwargs.get("enable_autotune", True), + "gemm1_bias": args.gemm1_bias, + "gemm2_bias": args.gemm2_bias, + "norm_topk_prob": kwargs.get("norm_topk_prob", True), + } + + return moe_impl.call_moe( + static_data, + kwargs["hidden_states_orig"], + args.hidden_states_scale_global, + **kernel_kwargs, + ) + + +@pytest.fixture(scope="module") +def cache_permute_indices(): + # The cache key is now a tuple of (weight_type, shape) + _cache_permute_indices: Dict[tuple, torch.Tensor] = {} + return _cache_permute_indices + + +def run_moe_test( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices, + routing_logits_dtype=torch.bfloat16, + zero_hidden_states=False, + gemm1_bias=None, + gemm2_bias=None, + routing_bias_dtype=None, + norm_topk_prob=True, +): + """Common test logic for all routing methods.""" + skip_checks( + moe_impl, + routing_config, + weight_processing, + activation_type, + num_tokens, + hidden_size, + intermediate_size, + routing_logits_dtype, + zero_hidden_states=zero_hidden_states, + ) + + torch.cuda.synchronize() + + moe_impl._cache_permute_indices = cache_permute_indices + + seed = 0 + torch.random.manual_seed(seed) + + # Extract routing configuration + top_k = routing_config["top_k"] + padding = routing_config["padding"] + n_groups = routing_config["n_groups"] + top_k_groups = routing_config["top_k_groups"] + routed_scaling = routing_config["routed_scaling"] + num_experts = routing_config["num_experts"] + routing_method_type = routing_config["routing_method_type"] + + # Validation checks + assert top_k <= num_experts + if (top_k_groups is not None) and (n_groups is not None) and (n_groups > 0): + assert top_k_groups <= 4 + assert num_experts > n_groups + assert num_experts % n_groups == 0 + assert num_experts % 4 == 0 + assert top_k < (top_k_groups * num_experts / n_groups) + + # Create test data based on routing method + expert_logits = torch.randn((num_tokens, num_experts), device="cuda").to( + routing_logits_dtype + ) + + if routing_config["has_routing_bias"]: + bias_dtype = ( + routing_bias_dtype if routing_bias_dtype is not None else torch.bfloat16 + ) + routing_bias = torch.randn(num_experts, device="cuda", dtype=bias_dtype) + else: + routing_bias = None + + hidden_states_fn = torch.zeros if zero_hidden_states else torch.randn + hidden_states = 2 * hidden_states_fn( + (num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16 + ) + gemm1_weights = torch.randn( + ( + num_experts, + (2 if is_gated_activation(activation_type) else 1) * intermediate_size, + hidden_size, + ), + device="cuda", + dtype=torch.bfloat16, + ) + gemm2_weights = torch.randn( + (num_experts, hidden_size, intermediate_size), + device="cuda", + dtype=torch.bfloat16, + ) + + # Generate routing info + use_routing_scales_on_input = routing_method_type == RoutingMethodType.Llama4 + + if routing_method_type == RoutingMethodType.Default: + permute_info, scores = routing_reference_default( + expert_logits, top_k, num_experts, padding + ) + elif routing_method_type == RoutingMethodType.DeepSeekV3: + permute_info, scores = routing_reference_no_aux( + expert_logits, + routing_bias, + top_k, + n_groups, + top_k_groups, + routed_scaling, + padding, + use_routing_scales_on_input, + ) + elif routing_method_type == RoutingMethodType.Renormalize: + permute_info, scores = routing_reference_renormalize( + expert_logits, top_k, num_experts, padding + ) + elif routing_method_type == RoutingMethodType.RenormalizeNaive: + # RenormalizeNaive (Softmax -> TopK -> SumNormalize) is mathematically equivalent + # to Renormalize (TopK -> Softmax), so we use the same reference implementation. + permute_info, scores = routing_reference_renormalize( + expert_logits, top_k, num_experts, padding + ) + elif routing_method_type == RoutingMethodType.TopK: + permute_info, scores = routing_reference_topk( + expert_logits, top_k, num_experts, padding + ) + elif routing_method_type == RoutingMethodType.SigmoidRenorm: + permute_info, scores = routing_reference_sigmoid_renorm( + expert_logits, top_k, num_experts, padding, norm_topk_prob=norm_topk_prob + ) + elif routing_method_type == RoutingMethodType.MiniMax2: + permute_info, scores = routing_reference_minimax2( + expert_logits, routing_bias, top_k, num_experts, padding + ) + elif routing_method_type == RoutingMethodType.Sigmoid: + permute_info, scores = routing_reference_sigmoid_renorm( + expert_logits, top_k, num_experts, padding, norm_topk_prob=False + ) + elif routing_method_type == RoutingMethodType.Llama4: + permute_info, scores = routing_reference_no_aux( + expert_logits, + routing_bias, + top_k, + n_groups, + top_k_groups, + routed_scaling, + padding, + use_routing_scales_on_input=True, + ) + else: + raise NotImplementedError( + f"Routing method {routing_method_type} not implemented" + ) + + # 1. Quantize weights offline + weights_data = moe_impl.quantize_weights( + gemm1_weights, gemm2_weights, hidden_states + ) + + # 2. Quantize inputs at runtime + inputs_data = moe_impl.quantize_inputs( + hidden_states, weights_data["hidden_states_scale_global"] + ) + + # 3. Combine quantized data + quant_data = {**weights_data, **inputs_data} + + # Create arguments for reference computation + args = moe_args( + num_tokens, + num_experts, + hidden_size, + intermediate_size, + top_k, + padding, + quant_data["hidden_states"], + quant_data["hidden_states_scale"], + quant_data["hidden_states_scale_global"], + scores, + quant_data["gemm1_weights"], + quant_data["gemm1_scales"], + quant_data["gemm1_scales_global"], + quant_data["gemm2_weights"], + quant_data["gemm2_scales"], + quant_data["gemm2_scales_global"], + permute_info, + use_routing_scales_on_input, + activation_type, + gemm1_bias=gemm1_bias, + gemm2_bias=gemm2_bias, + ) + + # Compute reference output + output_dequant_reference, args_dequant = moe_impl.compute_reference(args) + + if output_dequant_reference is None: + pytest.fail("Reference computation failed to produce output") + + # Compute actual output + enable_autotune = routing_config.get("enable_autotune", True) + + output_dequant_actual = moe_impl.compute_production( + args_dequant, + args, + expert_logits=expert_logits, + routing_bias=routing_bias, + hidden_states_orig=hidden_states, + gemm1_weights_orig=gemm1_weights, + gemm2_weights_orig=gemm2_weights, + n_groups=n_groups, + top_k_groups=top_k_groups, + routed_scaling=routed_scaling, + routing_method_type=routing_method_type, + weight_processing=weight_processing, + enable_pdl=True, + hidden_states_quant=inputs_data["hidden_states"], + enable_autotune=enable_autotune, + norm_topk_prob=norm_topk_prob, + ) + + # Compare outputs + tolerances = moe_impl.get_tolerances() + check_accuracy( + output_dequant_reference, + output_dequant_actual, + atol=tolerances["atol"], + rtol=tolerances["rtol"], + percent=tolerances["percent"], + ) + + +# Test: Renormalize routing +@pytest.mark.parametrize( + "zero_hidden_states", + [ + pytest.param(True, id="ZeroHiddenStates"), + pytest.param(False, id="RandomHiddenStates"), + ], +) +@pytest.mark.parametrize("num_tokens", [8, 768, 3072]) +@pytest.mark.parametrize("hidden_size", [1024]) +@pytest.mark.parametrize("intermediate_size", [1024, 768, 512, 384]) +@pytest.mark.parametrize( + "moe_impl", + [ + pytest.param(BF16Moe(), id="BF16xBF16"), + pytest.param( + FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_DEEPSEEK), + id="FP8_Block_DeepSeek", + ), + pytest.param( + FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_MXFP8), + id="FP8_Block_MxFp8", + ), + pytest.param(FP8PerTensorMoe(), id="FP8_Tensor"), + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_NVFP4_NVFP4), id="NvFP4xNvFP4"), + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_MXFP4_MXFP8), id="MxFP4xMxFP8"), + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_MXFP4_Bf16), id="MxFP4xBf16"), + pytest.param(MxInt4BlockScaleMoe(), id="MxInt4xBf16"), + ], +) +@pytest.mark.parametrize( + "routing_config", + [ + pytest.param( + { + "num_experts": 128, + "top_k": 8, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [ + FP8PerTensorMoe, + FP8BlockScaleMoe, + FP4Moe, + BF16Moe, + MxInt4BlockScaleMoe, + ], + "compatible_intermediate_size": [384, 768, 1024], + "enable_autotune": True, + }, + id="Qwen3_MOE", + ), + pytest.param( + { + "num_experts": 256, + "top_k": 8, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [ + FP8PerTensorMoe, + FP8BlockScaleMoe, + FP4Moe, + BF16Moe, + MxInt4BlockScaleMoe, + ], + "compatible_intermediate_size": [384, 1024], + "enable_autotune": False, + }, + id="Renorm", + ), + pytest.param( + { + "num_experts": 512, + "top_k": 10, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [ + FP8PerTensorMoe, + FP8BlockScaleMoe, + FP4Moe, + BF16Moe, + MxInt4BlockScaleMoe, + ], + "compatible_intermediate_size": [512], + "enable_autotune": True, + }, + id="Qwen3_next", + ), + pytest.param( + { + "num_experts": 2048, + "top_k": 32, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [ + FP8BlockScaleMoe, + FP4Moe, + BF16Moe, + MxInt4BlockScaleMoe, + ], + "compatible_intermediate_size": [384], + "enable_autotune": True, + }, + id="RoutingRenormalize_large_experts", + ), + pytest.param( + { + "num_experts": 128, + "top_k": 8, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Default, + "compatible_moe_impls": [ + FP8PerTensorMoe, + FP8BlockScaleMoe, + FP4Moe, + BF16Moe, + MxInt4BlockScaleMoe, + ], + "compatible_intermediate_size": [384, 768, 1024], + "enable_autotune": False, + }, + id="Default_128e_top8", + ), + pytest.param( + { + "num_experts": 128, + "top_k": 8, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.SigmoidRenorm, + "compatible_moe_impls": [ + FP8PerTensorMoe, + FP8BlockScaleMoe, + FP4Moe, + BF16Moe, + MxInt4BlockScaleMoe, + ], + "compatible_intermediate_size": [384, 768, 1024], + "enable_autotune": False, + }, + id="SigmoidRenorm_128e_top8", + ), + pytest.param( + { + "num_experts": 256, + "top_k": 6, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.MiniMax2, + "compatible_moe_impls": [ + FP8PerTensorMoe, + FP8BlockScaleMoe, + FP4Moe, + BF16Moe, + MxInt4BlockScaleMoe, + ], + "compatible_intermediate_size": [384, 768, 1024], + "enable_autotune": False, + }, + id="MiniMax2_256e_top6", + ), + ], +) +@pytest.mark.parametrize( + "weight_processing", + [ + pytest.param( + { + "use_shuffled_weight": False, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8BlockScaleMoe], + }, + id="NoShuffle_MajorK", + ), + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP4Moe, FP8PerTensorMoe, FP8BlockScaleMoe], + }, + id="Shuffled_MajorK", + ), + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.BlockMajorK, + "compatible_moe_impls": [ + FP8BlockScaleMoe, + BF16Moe, + MxInt4BlockScaleMoe, + ], + }, + id="Shuffled_BlockMajorK", + ), + ], +) +@pytest.mark.parametrize( + "activation_type", + [ + pytest.param(ActivationType.Swiglu.value, id="Swiglu"), + pytest.param(ActivationType.Geglu.value, id="Geglu"), + ], +) +@pytest.mark.parametrize( + "routing_logits_dtype", + [ + pytest.param(torch.float32, id="FP32_logits"), + pytest.param(torch.bfloat16, id="BF16_logits"), + ], +) +def test_renormalize_routing( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices, + routing_logits_dtype, + zero_hidden_states, +): + """Test Renormalize routing configurations.""" + run_moe_test( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices, + routing_logits_dtype, + zero_hidden_states=zero_hidden_states, + ) + + +# Test: Sigmoid routing (Sigmoid -> TopK, no renormalization) +@pytest.mark.parametrize("num_tokens", [8, 768, 3072]) +@pytest.mark.parametrize("hidden_size", [1024]) +@pytest.mark.parametrize("intermediate_size", [1024, 768, 512, 384]) +@pytest.mark.parametrize( + "moe_impl", + [ + pytest.param(BF16Moe(), id="BF16xBF16"), + pytest.param( + FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_DEEPSEEK), + id="FP8_Block_DeepSeek", + ), + pytest.param( + FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_MXFP8), + id="FP8_Block_MxFp8", + ), + pytest.param(FP8PerTensorMoe(), id="FP8_Tensor"), + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_NVFP4_NVFP4), id="NvFP4xNvFP4"), + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_MXFP4_MXFP8), id="MxFP4xMxFP8"), + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_MXFP4_Bf16), id="MxFP4xBf16"), + pytest.param(MxInt4BlockScaleMoe(), id="MxInt4xBf16"), + ], +) +@pytest.mark.parametrize( + "routing_config", + [ + pytest.param( + { + "num_experts": 128, + "top_k": 8, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Sigmoid, + "compatible_moe_impls": [ + FP8PerTensorMoe, + FP8BlockScaleMoe, + FP4Moe, + BF16Moe, + MxInt4BlockScaleMoe, + ], + "compatible_intermediate_size": [384, 768, 1024], + "enable_autotune": True, + }, + id="Sigmoid_128e_top8", + ), + ], +) +@pytest.mark.parametrize( + "weight_processing", + [ + pytest.param( + { + "use_shuffled_weight": False, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8BlockScaleMoe], + }, + id="NoShuffle_MajorK", + ), + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP4Moe, FP8PerTensorMoe, FP8BlockScaleMoe], + }, + id="Shuffled_MajorK", + ), + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.BlockMajorK, + "compatible_moe_impls": [ + FP8BlockScaleMoe, + BF16Moe, + MxInt4BlockScaleMoe, + ], + }, + id="Shuffled_BlockMajorK", + ), + ], +) +@pytest.mark.parametrize( + "activation_type", + [ + pytest.param(ActivationType.Swiglu.value, id="Swiglu"), + pytest.param(ActivationType.Geglu.value, id="Geglu"), + ], +) +def test_sigmoid_routing( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices, +): + """Test Sigmoid routing configurations (Sigmoid -> TopK, no renormalization).""" + run_moe_test( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices, + ) + + +# Test: DeepSeekV3 routing +@pytest.mark.parametrize("num_tokens", [8, 768, 3072]) +@pytest.mark.parametrize("hidden_size", [1024]) +@pytest.mark.parametrize("intermediate_size", [2688, 2048, 1024, 768, 512, 384]) +@pytest.mark.parametrize( + "moe_impl", + [ + pytest.param(FP8PerTensorMoe(), id="FP8_PerTensor"), + pytest.param( + FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_DEEPSEEK), + id="FP8_Block_DeepSeek", + ), + pytest.param( + FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_MXFP8), + id="FP8_Block_MxFp8", + ), + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_NVFP4_NVFP4), id="NvFP4xNvFP4"), + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_MXFP4_MXFP8), id="MxFP4xMxFP8"), + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_MXFP4_Bf16), id="MxFP4xBf16"), + pytest.param(MxInt4BlockScaleMoe(), id="MxInt4xBf16"), + pytest.param(BF16Moe(), id="Bf16xBf16"), + ], +) +@pytest.mark.parametrize( + "routing_config", + [ + pytest.param( + { + "num_experts": 512, + "top_k": 22, + "padding": 8, + "n_groups": 1, + "top_k_groups": 1, + "routed_scaling": 2.5, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.DeepSeekV3, + "compatible_moe_impls": [BF16Moe, FP8PerTensorMoe, FP4Moe], + "compatible_intermediate_size": [2688], + "compatible_activation_types": [ActivationType.Relu2], + "enable_autotune": True, + }, + id="nemotron_3_super", + ), + pytest.param( + { + "num_experts": 384, + "top_k": 8, + "padding": 8, + "n_groups": 1, + "top_k_groups": 1, + "routed_scaling": 2.5, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.DeepSeekV3, + "compatible_moe_impls": [FP4Moe, FP8BlockScaleMoe], + "compatible_intermediate_size": [1024, 2048], + "compatible_activation_types": [ + ActivationType.Swiglu, + ActivationType.Geglu, + ], + "enable_autotune": True, + }, + id="kimi_k2", + ), + pytest.param( + { + "num_experts": 256, + "top_k": 8, + "padding": 8, + "n_groups": 8, + "top_k_groups": 4, + "routed_scaling": 2.5, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.DeepSeekV3, + "compatible_moe_impls": [ + FP4Moe, + FP8BlockScaleMoe, + MxInt4BlockScaleMoe, + BF16Moe, + ], + "compatible_intermediate_size": [512, 1024, 2048], + "compatible_activation_types": [ + ActivationType.Swiglu, + ActivationType.Geglu, + ], + "enable_autotune": True, + }, + id="DSv3", + ), + pytest.param( + { + "num_experts": 72, + "top_k": 6, + "padding": 8, + "n_groups": 1, + "top_k_groups": 1, + "routed_scaling": 2.5, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.DeepSeekV3, + "compatible_moe_impls": [FP4Moe, FP8BlockScaleMoe], + "compatible_intermediate_size": [384, 768], + "compatible_activation_types": [ + ActivationType.Swiglu, + ActivationType.Geglu, + ], + "enable_autotune": False, + }, + id="DSLite", + ), + pytest.param( + { + "num_experts": 160, + "top_k": 8, + "padding": 8, + "n_groups": 1, + "top_k_groups": 1, + "routed_scaling": 2.5, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.DeepSeekV3, + "compatible_moe_impls": [FP4Moe, FP8BlockScaleMoe, BF16Moe], + "compatible_intermediate_size": [512, 1024, 1536], + "compatible_activation_types": [ + ActivationType.Swiglu, + ActivationType.Geglu, + ], + "enable_autotune": False, + }, + id="GLM4_MoE", + ), + ], +) +@pytest.mark.parametrize( + "weight_processing", + [ + pytest.param( + { + "use_shuffled_weight": False, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8BlockScaleMoe], + }, + id="NoShuffle_MajorK", + ), + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP4Moe, FP8PerTensorMoe, FP8BlockScaleMoe], + }, + id="Shuffled_MajorK", + ), + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.BlockMajorK, + "compatible_moe_impls": [ + FP8BlockScaleMoe, + MxInt4BlockScaleMoe, + BF16Moe, + ], + }, + id="Shuffled_BlockMajorK", + ), + ], +) +@pytest.mark.parametrize( + "activation_type", + [ + pytest.param(ActivationType.Swiglu.value, id="Swiglu"), + pytest.param(ActivationType.Geglu.value, id="Geglu"), + pytest.param(ActivationType.Relu2.value, id="Relu2"), + ], +) +@pytest.mark.parametrize( + "routing_logits_dtype", + [ + pytest.param(torch.float32, id="FP32_logits"), + ], +) +def test_deepseekv3_routing( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + routing_logits_dtype, + cache_permute_indices, +): + """Test DeepSeekV3 routing configurations.""" + run_moe_test( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices, + routing_logits_dtype, + ) + + +# Test: TopK routing +@pytest.mark.parametrize("num_tokens", [8, 128]) # Limited for GeGlu +@pytest.mark.parametrize("hidden_size", [1024]) +@pytest.mark.parametrize("intermediate_size", [384, 512, 768, 1024]) +@pytest.mark.parametrize( + "moe_impl", + [ + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_NVFP4_NVFP4), id="NvFP4xNvFP4"), + pytest.param(FP4Moe(quant_mode=QuantMode.FP4_MXFP4_MXFP8), id="MxFP4xMxFP8"), + ], +) +@pytest.mark.parametrize( + "routing_config", + [ + pytest.param( + { + "num_experts": 16, + "top_k": 2, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.TopK, + "compatible_moe_impls": [FP4Moe], + "compatible_intermediate_size": [512, 768, 1024], + "enable_autotune": True, + }, + id="TopK", + ), + ], +) +@pytest.mark.parametrize( + "weight_processing", + [ + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP4Moe, FP8PerTensorMoe, FP8BlockScaleMoe], + }, + id="Shuffled_MajorK", + ), + ], +) +@pytest.mark.parametrize( + "activation_type", + [ + pytest.param(ActivationType.Swiglu.value, id="Swiglu"), + pytest.param(ActivationType.Geglu.value, id="Geglu"), + ], +) +@pytest.mark.parametrize( + "routing_logits_dtype", + [ + pytest.param(torch.float32, id="FP32_logits"), + pytest.param(torch.bfloat16, id="BF16_logits"), + ], +) +def test_topk_routing( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + routing_logits_dtype, + cache_permute_indices, +): + """Test TopK routing configuration.""" + run_moe_test( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices, + routing_logits_dtype, + ) + + +# Test: Llama4 routing +@pytest.mark.parametrize("num_tokens", [8, 768, 3072]) +@pytest.mark.parametrize("hidden_size", [1024]) +@pytest.mark.parametrize("intermediate_size", [1024, 2048]) +@pytest.mark.parametrize( + "moe_impl", + [ + pytest.param(FP8PerTensorMoe(), id="FP8_Tensor"), + ], +) +@pytest.mark.parametrize( + "routing_config", + [ + pytest.param( + { + "num_experts": 128, + "top_k": 1, + "padding": 8, + "n_groups": 0, + "top_k_groups": 0, + "routed_scaling": 2.5, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.Llama4, + "compatible_moe_impls": [FP8PerTensorMoe], + "compatible_intermediate_size": [1024, 2048], + "enable_autotune": True, + }, + id="Llama4", + ), + ], +) +@pytest.mark.parametrize( + "weight_processing", + [ + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP4Moe, FP8PerTensorMoe, FP8BlockScaleMoe], + }, + id="Shuffled_MajorK", + ), + ], +) +@pytest.mark.parametrize( + "activation_type", + [ + pytest.param(ActivationType.Swiglu.value, id="Swiglu"), + ], +) +@pytest.mark.parametrize( + "routing_logits_dtype", + [ + pytest.param(torch.bfloat16, id="BF16_logits"), + ], +) +def test_llama4_routing( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + routing_logits_dtype, + cache_permute_indices, +): + """Test Llama4 routing configuration with FP8 per-tensor.""" + run_moe_test( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices, + routing_logits_dtype, + ) + + +@pytest.mark.parametrize("num_tokens", [32, 768, 3072]) +@pytest.mark.parametrize("hidden_size", [1024]) +@pytest.mark.parametrize("intermediate_size", [2048, 1024, 768, 512]) +@pytest.mark.parametrize("bias", ["gemm2", "gemm1", "gemm1_and_gemm2"]) +def test_nvfp4_moe_gemm_bias( + num_tokens, hidden_size, intermediate_size, bias, cache_permute_indices +): + """Test NvFP4 MoE with GEMM bias support.""" + num_experts = 8 + top_k = 2 + device = "cuda" + + gemm1_bias = None + gemm2_bias = None + if "gemm1" in bias: + gemm1_bias = torch.randn( + (num_experts, 2 * intermediate_size), device=device, dtype=torch.float32 + ) + if "gemm2" in bias: + gemm2_bias = torch.randn( + (num_experts, hidden_size), device=device, dtype=torch.float32 + ) + + run_moe_test( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + moe_impl=FP4Moe(quant_mode=QuantMode.FP4_NVFP4_NVFP4), + routing_config={ + "num_experts": num_experts, + "top_k": top_k, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [FP4Moe], + "compatible_intermediate_size": [512, 768, 1024, 2048], + "enable_autotune": True, + }, + weight_processing={ + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP4Moe, FP8PerTensorMoe, FP8BlockScaleMoe], + }, + activation_type=ActivationType.Swiglu, + cache_permute_indices=cache_permute_indices, + routing_logits_dtype=torch.bfloat16, + gemm1_bias=gemm1_bias, + gemm2_bias=gemm2_bias, + ) + + +@pytest.mark.parametrize("num_tokens", [1, 16, 64, 256, 1000, 4000]) +@pytest.mark.parametrize("hidden_size", [512, 1024]) +@pytest.mark.parametrize("intermediate_size", [512, 1024]) +@pytest.mark.parametrize( + "zero_hidden_states", + [ + pytest.param(True, id="ZeroHiddenStates"), + pytest.param(False, id="RandomHiddenStates"), + ], +) +@pytest.mark.parametrize( + "routing_config", + [ + pytest.param( + { + "num_experts": 32, + "top_k": 4, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [512, 1024], + "compatible_activation_types": [ActivationType.Relu2], + "enable_autotune": False, + }, + id="E32_K4", + ), + pytest.param( + { + "num_experts": 64, + "top_k": 8, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [512, 1024], + "compatible_activation_types": [ActivationType.Relu2], + "enable_autotune": False, + }, + id="E64_K8", + ), + ], +) +@pytest.mark.parametrize( + "weight_processing", + [ + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8BlockScaleMoe], + }, + id="Shuffled_MajorK", + ), + ], +) +def test_mxfp8_block_scale_moe_relu2_non_gated( + num_tokens, + hidden_size, + intermediate_size, + zero_hidden_states, + routing_config, + weight_processing, + cache_permute_indices, +): + """Test MXFP8 block-scale TRTLLM MoE with non-gated RELU2.""" + run_moe_test( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + moe_impl=FP8BlockScaleMoe( + fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_MXFP8 + ), + routing_config=routing_config, + weight_processing=weight_processing, + activation_type=ActivationType.Relu2, + cache_permute_indices=cache_permute_indices, + routing_logits_dtype=torch.bfloat16, + zero_hidden_states=zero_hidden_states, + ) + + +def test_mxfp8_block_scale_moe_relu2_deepseekv3_topk22(cache_permute_indices): + """Targeted coverage for MXFP8 non-gated Relu2 with DeepSeekV3 routing top_k=22.""" + run_moe_test( + num_tokens=128, + hidden_size=1024, + intermediate_size=512, + moe_impl=FP8BlockScaleMoe( + fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_MXFP8 + ), + routing_config={ + # top_k=22 is only supported when num_experts > NumKimiK2Experts (384) + "num_experts": 512, + "top_k": 22, + "padding": 8, + "n_groups": 1, + "top_k_groups": 1, + "routed_scaling": 2.5, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.DeepSeekV3, + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [512], + "compatible_activation_types": [ActivationType.Relu2], + "enable_autotune": False, + }, + weight_processing={ + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8BlockScaleMoe], + }, + activation_type=ActivationType.Relu2, + cache_permute_indices=cache_permute_indices, + routing_logits_dtype=torch.float32, + ) + + +@pytest.mark.parametrize( + "autotune_case", + [ + pytest.param( + { + "num_tokens": 1, + "hidden_size": 1024, + "intermediate_size": 1024, + "quant_mode": QuantMode.FP8_BLOCK_SCALE_MXFP8, + "activation_type": ActivationType.Relu2, + "num_experts": 64, + "top_k": 8, + "routing_method_type": RoutingMethodType.Renormalize, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + }, + id="MxFp8_Relu2_T1_H1024_I1024_K8", + ), + pytest.param( + { + "num_tokens": 64, + "hidden_size": 512, + "intermediate_size": 512, + "quant_mode": QuantMode.FP8_BLOCK_SCALE_MXFP8, + "activation_type": ActivationType.Relu2, + "num_experts": 512, + "top_k": 22, + # top_k=22 is only valid on DeepSeekV3 routing path for large expert counts. + "routing_method_type": RoutingMethodType.DeepSeekV3, + "n_groups": 1, + "top_k_groups": 1, + "routed_scaling": 2.5, + "has_routing_bias": True, + }, + id="MxFp8_Relu2_T64_H512_I512_K22", + ), + pytest.param( + { + "num_tokens": 256, + "hidden_size": 1024, + "intermediate_size": 512, + "quant_mode": QuantMode.FP8_BLOCK_SCALE_DEEPSEEK, + "activation_type": ActivationType.Swiglu, + "num_experts": 256, + "top_k": 8, + "routing_method_type": RoutingMethodType.Renormalize, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + }, + id="DeepSeek_Swiglu_T256_H1024_I512_K8", + ), + ], +) +def test_fp8_block_scale_autotune_valid_configs(autotune_case, cache_permute_indices): + """Autotune smoke matrix to exercise C++ getValidConfigs across FP8 modes/shapes.""" + run_moe_test( + num_tokens=autotune_case["num_tokens"], + hidden_size=autotune_case["hidden_size"], + intermediate_size=autotune_case["intermediate_size"], + moe_impl=FP8BlockScaleMoe(fp8_quantization_type=autotune_case["quant_mode"]), + routing_config={ + "num_experts": autotune_case["num_experts"], + "top_k": autotune_case["top_k"], + "padding": 8, + "n_groups": autotune_case["n_groups"], + "top_k_groups": autotune_case["top_k_groups"], + "routed_scaling": autotune_case["routed_scaling"], + "has_routing_bias": autotune_case["has_routing_bias"], + "routing_method_type": autotune_case["routing_method_type"], + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [autotune_case["intermediate_size"]], + "compatible_activation_types": [autotune_case["activation_type"]], + "enable_autotune": True, + }, + weight_processing={ + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8BlockScaleMoe], + }, + activation_type=autotune_case["activation_type"], + cache_permute_indices=cache_permute_indices, + routing_logits_dtype=torch.float32, + zero_hidden_states=False, + ) + + +@pytest.mark.parametrize( + "autotune_case", + [ + pytest.param( + { + "num_tokens": 64, + "hidden_size": 1024, + "intermediate_size": 1024, + "activation_type": ActivationType.Swiglu, + "num_experts": 64, + "top_k": 8, + }, + id="PerTensor_Swiglu_T64_H1024_I1024_K8", + ), + pytest.param( + { + "num_tokens": 32, + "hidden_size": 512, + "intermediate_size": 512, + "activation_type": ActivationType.Relu2, + "num_experts": 64, + "top_k": 8, + }, + id="PerTensor_Relu2_T32_H512_I512_K8", + ), + ], +) +def test_fp8_per_tensor_autotune_valid_configs_nonefp8( + autotune_case, cache_permute_indices +): + """Exercise per-tensor autotune path that uses NoneFp8 in valid-config dispatch.""" + run_moe_test( + num_tokens=autotune_case["num_tokens"], + hidden_size=autotune_case["hidden_size"], + intermediate_size=autotune_case["intermediate_size"], + moe_impl=FP8PerTensorMoe(), + routing_config={ + "num_experts": autotune_case["num_experts"], + "top_k": autotune_case["top_k"], + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [FP8PerTensorMoe], + "compatible_intermediate_size": [autotune_case["intermediate_size"]], + "compatible_activation_types": [autotune_case["activation_type"]], + "enable_autotune": True, + }, + weight_processing={ + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8PerTensorMoe], + }, + activation_type=autotune_case["activation_type"], + cache_permute_indices=cache_permute_indices, + routing_logits_dtype=torch.bfloat16, + zero_hidden_states=False, + ) + + +@pytest.mark.parametrize( + "num_tokens", + [5, 8, 12, 16], + ids=lambda t: f"T{t}", +) +@pytest.mark.parametrize("hidden_size", [512]) +@pytest.mark.parametrize("intermediate_size", [512]) +@pytest.mark.parametrize( + "moe_impl", + [ + pytest.param( + FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_DEEPSEEK), + id="FP8_Block_DeepSeek", + ), + ], +) +@pytest.mark.parametrize( + "routing_config", + [ + pytest.param( + { + "num_experts": 64, + "top_k": 4, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [512], + "enable_autotune": False, + }, + id="Renormalize_64e_top4", + ), + ], +) +@pytest.mark.parametrize( + "weight_processing", + [ + pytest.param( + { + "use_shuffled_weight": False, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8BlockScaleMoe], + }, + id="NoShuffle_MajorK", + ), + ], +) +@pytest.mark.parametrize("activation_type", [ActivationType.Swiglu]) +def test_dyn_block_kernel_routing( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, +): + """Test token counts 5-16 that exercise the dynamic block kernel path (BlockKernelMaxNumTokens < tokens <= DynBlockKernelMaxNumTokens).""" + run_moe_test( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices=False, + ) + + +@pytest.mark.parametrize("num_tokens", [8]) +@pytest.mark.parametrize("hidden_size", [512]) +@pytest.mark.parametrize("intermediate_size", [512]) +@pytest.mark.parametrize( + "moe_impl", + [ + pytest.param( + FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_DEEPSEEK), + id="FP8_Block_DeepSeek", + ), + ], +) +@pytest.mark.parametrize( + "routing_config", + [ + pytest.param( + { + "num_experts": 1024, + "top_k": 8, + "padding": 8, + "n_groups": 1, + "top_k_groups": 1, + "routed_scaling": 1.0, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.DeepSeekV3, + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [512], + "enable_autotune": False, + }, + id="DeepSeekV3_1024e_top8", + ), + ], +) +@pytest.mark.parametrize( + "weight_processing", + [ + pytest.param( + { + "use_shuffled_weight": False, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8BlockScaleMoe], + }, + id="NoShuffle_MajorK", + ), + ], +) +@pytest.mark.parametrize("activation_type", [ActivationType.Swiglu]) +def test_tier_1024_experts_routing( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, +): + """Test 1024-expert routing to exercise Tier<1024, 32> in SigmoidBias+ScaledSumNormalize policy.""" + run_moe_test( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices=False, + ) + + +# num_tokens is chosen to straddle the dispatch thresholds in routingCustom::run +# (see trtllm_fused_moe_routing_custom.cu): +# - tokens == 8 : dyn-block kernel path (tokens <= DynBlockKernelMaxNumTokens=16, +# numExperts <= DynBlockKernelMaxNumExperts=512) +# - tokens == 32 : block-per-token "split" path on the single-cluster kernel +# (17 <= tokens <= 256, numExperts >= 160, policy pair opts into +# PolicyPairSupportsBlockPerToken) — exercises the +# routingIndicesBlockScoresKernel path added for this feature. +@pytest.mark.parametrize("num_tokens", [8, 32]) +@pytest.mark.parametrize("hidden_size", [512]) +@pytest.mark.parametrize("intermediate_size", [512]) +@pytest.mark.parametrize( + "moe_impl", + [ + pytest.param( + FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_DEEPSEEK), + id="FP8_Block_DeepSeek", + ), + pytest.param(BF16Moe(), id="BF16xBF16"), + pytest.param(FP8PerTensorMoe(), id="FP8_PerTensor"), + pytest.param( + FP4Moe(quant_mode=QuantMode.FP4_NVFP4_NVFP4), + id="NvFP4xNvFP4", + ), + ], +) +@pytest.mark.parametrize( + "routing_config", + [ + # DeepSeekV3 + nGroup == 1 routes through routingCustom with + # (SigmoidBiasPreprocess, ScaledSumNormalizePostprocess) — the policy pair + # that opts into the block-per-token BlockScores kernel. + pytest.param( + { + "num_experts": 384, + "top_k": 6, + "padding": 8, + "n_groups": 1, + "top_k_groups": 1, + "routed_scaling": 2.5, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.DeepSeekV3, + "compatible_moe_impls": [ + FP8BlockScaleMoe, + BF16Moe, + FP8PerTensorMoe, + FP4Moe, + ], + "compatible_intermediate_size": [512], + "enable_autotune": False, + }, + id="DeepSeekV3_ngroup1_384e_top6", + ), + pytest.param( + { + # top_k=22 requires num_experts > NumKimiK2Experts (384). + "num_experts": 512, + "top_k": 22, + "padding": 8, + "n_groups": 1, + "top_k_groups": 1, + "routed_scaling": 2.5, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.DeepSeekV3, + "compatible_moe_impls": [ + FP8BlockScaleMoe, + BF16Moe, + FP8PerTensorMoe, + FP4Moe, + ], + "compatible_intermediate_size": [512], + "enable_autotune": False, + }, + id="DeepSeekV3_ngroup1_512e_top22", + ), + ], +) +@pytest.mark.parametrize( + "weight_processing", + [ + pytest.param( + { + "use_shuffled_weight": False, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8BlockScaleMoe], + }, + id="NoShuffle_MajorK", + ), + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP4Moe, FP8PerTensorMoe, FP8BlockScaleMoe], + }, + id="Shuffled_MajorK", + ), + pytest.param( + { + "use_shuffled_weight": True, + "layout": WeightLayout.BlockMajorK, + "compatible_moe_impls": [ + FP8BlockScaleMoe, + MxInt4BlockScaleMoe, + BF16Moe, + ], + }, + id="Shuffled_BlockMajorK", + ), + ], +) +@pytest.mark.parametrize( + "activation_type", + [ + pytest.param(ActivationType.Swiglu, id="Swiglu"), + pytest.param(ActivationType.Relu2, id="Relu2"), + ], +) +def test_deepseek_ngroup1_block_per_token_routing( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices, +): + """Exercise the block-per-token BlockScores kernel in routingCustom. + + DeepSeekV3 with n_group == 1 dispatches to routingCustom with the + (SigmoidBiasPreprocess, ScaledSumNormalizePostprocess) policy pair, which opts + into PolicyPairSupportsBlockPerToken. For num_experts >= 160 and + 17 <= num_tokens <= 256, routingIndicesBlockScoresKernel replaces the + fused single-cluster kernel. We intentionally use independent + parametrization here; incompatible combinations are filtered by skip_checks. + + Covered tiers: + - Tier<384, 8> via num_experts=384, top_k=6 + - Tier<512, 22> via num_experts=512, top_k=22 + """ + run_moe_test( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices=cache_permute_indices, + ) + + +@pytest.mark.parametrize("num_tokens", [8]) +@pytest.mark.parametrize("hidden_size", [512]) +@pytest.mark.parametrize("intermediate_size", [512]) +@pytest.mark.parametrize( + "moe_impl", + [ + pytest.param( + FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_DEEPSEEK), + id="FP8_Block_DeepSeek", + ), + ], +) +@pytest.mark.parametrize( + "routing_config", + [ + pytest.param( + { + "num_experts": 256, + "top_k": 6, + "padding": 8, + "n_groups": 8, + "top_k_groups": 4, + "routed_scaling": 2.5, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.DeepSeekV3, + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [512], + "enable_autotune": False, + }, + id="DeepSeekV3_256e", + ), + pytest.param( + { + "num_experts": 128, + "top_k": 8, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Renormalize, + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [512], + "enable_autotune": False, + }, + id="Renormalize_128e", + ), + pytest.param( + { + "num_experts": 128, + "top_k": 8, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.Default, + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [512], + "enable_autotune": False, + }, + id="Default_128e", + ), + pytest.param( + { + "num_experts": 128, + "top_k": 8, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": False, + "routing_method_type": RoutingMethodType.SigmoidRenorm, + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [512], + "enable_autotune": False, + }, + id="SigmoidRenorm_128e", + ), + pytest.param( + { + "num_experts": 256, + "top_k": 6, + "padding": 8, + "n_groups": None, + "top_k_groups": None, + "routed_scaling": None, + "has_routing_bias": True, + "routing_method_type": RoutingMethodType.MiniMax2, + "compatible_moe_impls": [FP8BlockScaleMoe], + "compatible_intermediate_size": [512], + "enable_autotune": False, + }, + id="MiniMax2_256e", + ), + ], +) +@pytest.mark.parametrize( + "weight_processing", + [ + pytest.param( + { + "use_shuffled_weight": False, + "layout": WeightLayout.MajorK, + "compatible_moe_impls": [FP8BlockScaleMoe], + }, + id="NoShuffle_MajorK", + ), + ], +) +@pytest.mark.parametrize("activation_type", [ActivationType.Swiglu]) +@pytest.mark.parametrize( + "routing_logits_dtype", + [ + pytest.param(torch.bfloat16, id="BF16_logits"), + pytest.param(torch.float32, id="FP32_logits"), + ], +) +@pytest.mark.parametrize( + "routing_bias_dtype", + [ + pytest.param(None, id="default_bias"), + pytest.param(torch.float32, id="FP32_bias"), + ], +) +def test_routing_dtype_flexibility( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + routing_logits_dtype, + routing_bias_dtype, +): + """Test that routing works with both bfloat16 and float32 logits/bias across all routing methods.""" + run_moe_test( + num_tokens, + hidden_size, + intermediate_size, + moe_impl, + routing_config, + weight_processing, + activation_type, + cache_permute_indices=False, + routing_logits_dtype=routing_logits_dtype, + routing_bias_dtype=routing_bias_dtype, + ) + + +def test_fp8_block_scale_routed_activation_type_relu2_smoke(): + """Smoke test routed FP8 block-scale call path with explicit non-gated activation_type.""" + compute_capability = get_compute_capability(torch.device(device="cuda")) + if compute_capability[0] not in [10]: + pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.") + + torch.manual_seed(0) + device = torch.device("cuda:0") + + num_tokens = 32 + hidden_size = 512 + intermediate_size = 512 + num_experts = 64 + top_k = 8 + routing_method_type = RoutingMethodType.Renormalize + activation_type = ActivationType.Relu2.value + fp8_quantization_type = Fp8QuantizationType.MxFp8 + + routing_logits = torch.randn((num_tokens, num_experts), device=device).to( + torch.bfloat16 + ) + hidden_states = torch.randn((num_tokens, hidden_size), device=device).to( + torch.bfloat16 + ) + gemm1_weights = torch.randn( + (num_experts, intermediate_size, hidden_size), + device=device, + dtype=torch.bfloat16, + ) + gemm2_weights = torch.randn( + (num_experts, hidden_size, intermediate_size), + device=device, + dtype=torch.bfloat16, + ) + + quant_impl = FP8BlockScaleMoe(fp8_quantization_type=QuantMode.FP8_BLOCK_SCALE_MXFP8) + quant_weights = quant_impl.quantize_weights( + gemm1_weights, gemm2_weights, hidden_states + ) + quant_inputs = quant_impl.quantize_inputs(hidden_states) + + output_ref = trtllm_fp8_block_scale_moe( + routing_logits=routing_logits, + routing_bias=None, + hidden_states=quant_inputs["hidden_states"], + hidden_states_scale=quant_inputs["hidden_states_scale"], + gemm1_weights=quant_weights["gemm1_weights"], + gemm1_weights_scale=quant_weights["gemm1_scales"], + gemm2_weights=quant_weights["gemm2_weights"], + gemm2_weights_scale=quant_weights["gemm2_scales"], + num_experts=num_experts, + top_k=top_k, + n_group=None, + topk_group=None, + intermediate_size=intermediate_size, + local_expert_offset=0, + local_num_experts=num_experts, + routed_scaling_factor=None, + routing_method_type=routing_method_type.value, + use_shuffled_weight=True, + weight_layout=WeightLayout.MajorK.value, + enable_pdl=True, + fp8_quantization_type=fp8_quantization_type, + activation_type=activation_type, + ).to(torch.float) + + permute_info, expert_weights_full = routing_reference_renormalize( + routing_logits, top_k, num_experts, 8 + ) + topk_ids = permute_info["topKIndices"].to(torch.int32) + expert_weights = expert_weights_full.view(num_tokens, num_experts)[ + torch.arange(num_tokens, device=device).unsqueeze(1), topk_ids + ].to(torch.bfloat16) + packed_topk_ids = (topk_ids << 16) | expert_weights.view(torch.int16).to( + torch.int32 + ) + + output_routed = trtllm_fp8_block_scale_routed_moe( + topk_ids=packed_topk_ids, + routing_bias=None, + hidden_states=quant_inputs["hidden_states"], + hidden_states_scale=quant_inputs["hidden_states_scale"], + gemm1_weights=quant_weights["gemm1_weights"], + gemm1_weights_scale=quant_weights["gemm1_scales"], + gemm2_weights=quant_weights["gemm2_weights"], + gemm2_weights_scale=quant_weights["gemm2_scales"], + num_experts=num_experts, + top_k=top_k, + n_group=None, + topk_group=None, + intermediate_size=intermediate_size, + local_expert_offset=0, + local_num_experts=num_experts, + routed_scaling_factor=None, + routing_method_type=routing_method_type.value, + use_shuffled_weight=True, + weight_layout=WeightLayout.MajorK.value, + enable_pdl=True, + fp8_quantization_type=fp8_quantization_type, + activation_type=activation_type, + ).to(torch.float) + + close = torch.isclose(output_ref, output_routed, atol=1e-2, rtol=1e-2) + mismatch_pct = (~close).float().mean().item() * 100 + assert mismatch_pct < 10, f"Mismatch percentage is {mismatch_pct:.2f}%" diff --git a/tests/_fi_vendor/utils.py b/tests/_fi_vendor/utils.py new file mode 100644 index 0000000000..39abf18717 --- /dev/null +++ b/tests/_fi_vendor/utils.py @@ -0,0 +1,173 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import pytest +import torch +from enum import IntEnum +from flashinfer import ActivationType, RoutingMethodType +from flashinfer.utils import get_compute_capability +from flashinfer.fused_moe import WeightLayout + + +class QuantMode(IntEnum): + """Supported quantization modes for MoE testing.""" + + FP4_NVFP4_NVFP4 = 1 + FP4_MXFP4_MXFP8 = 2 + FP4_MXFP4_Bf16 = 3 + FP8_BLOCK_SCALE_DEEPSEEK = 4 + FP8_BLOCK_SCALE_MXFP8 = 5 + FP8_PER_TENSOR = 6 + BF16 = 7 + MXINT4_BF16_BF16 = 8 + + +NON_GATED_ACTIVATION_SUPPORTED_QUANT_MODES = [ + QuantMode.FP4_NVFP4_NVFP4, + QuantMode.FP8_BLOCK_SCALE_MXFP8, + QuantMode.FP8_PER_TENSOR, + QuantMode.BF16, +] + + +def is_gated_activation(activation_type: ActivationType) -> bool: + return activation_type in [ + ActivationType.Swiglu, + ActivationType.Geglu, + ActivationType.SwigluBias, + ] + + +def skip_checks( + moe_impl, + routing_config, + weight_processing, + activation_type, + num_tokens, + hidden_size, + intermediate_size, + logits_dtype, + zero_hidden_states=False, +): + """Common skip logic for all tests.""" + compute_capability = get_compute_capability(torch.device(device="cuda")) + if compute_capability[0] not in [10]: + pytest.skip("These tests are only guaranteed to work on SM100 and SM103 GPUs.") + + # Check moe_impl class by name to avoid circular imports + is_fp4_moe = type(moe_impl).__name__ == "FP4Moe" + is_fp8_block_scale_moe = type(moe_impl).__name__ == "FP8BlockScaleMoe" + + # Skip zero hidden states tests for non-FP8 Block Scale MoE implementations + if zero_hidden_states and not is_fp8_block_scale_moe: + pytest.skip("Skipping zero hidden states tests for non-FP8 Block Scale MoE.") + + # Skip incompatible combinations + if activation_type == ActivationType.Geglu and ( + not is_fp4_moe + or moe_impl.quant_mode != QuantMode.FP4_NVFP4_NVFP4 + or routing_config["routing_method_type"] != RoutingMethodType.TopK + or num_tokens > 128 + ): + pytest.skip( + f"Incompatible: {moe_impl.name} + {activation_type} + {routing_config['routing_method_type']} + {num_tokens}" + ) + elif activation_type == ActivationType.Swiglu and ( + hidden_size > 1024 or intermediate_size > 1024 + ): + pytest.skip( + f"Skip for testing speed: {activation_type} + {hidden_size} + {intermediate_size}" + ) + + compatible_activation_types = routing_config.get( + "compatible_activation_types", None + ) + if ( + compatible_activation_types is not None + and activation_type not in compatible_activation_types + ): + pytest.skip( + f"Incompatible: activation_type={activation_type} not in compatible_activation_types ({compatible_activation_types})" + ) + + if ( + not is_gated_activation(activation_type) + and moe_impl.quant_mode not in NON_GATED_ACTIVATION_SUPPORTED_QUANT_MODES + ): + pytest.skip( + f"Incompatible: {moe_impl.name} + {activation_type=} + quant_mode={moe_impl.quant_mode}: non-gated activations only supported with these quant modes: {NON_GATED_ACTIVATION_SUPPORTED_QUANT_MODES}" + ) + + # Skip large intermediate sizes for configurations with many experts + if routing_config["num_experts"] > 512 and intermediate_size > 512: + pytest.skip( + f"Skipping for testing speed: intermediate_size={intermediate_size} with {routing_config['num_experts']} experts" + ) + + if type(moe_impl) not in routing_config["compatible_moe_impls"]: + pytest.skip( + f"Incompatible: {moe_impl.name} + {routing_config['routing_method_type'].name}" + ) + if type(moe_impl) not in weight_processing["compatible_moe_impls"]: + pytest.skip( + f"Incompatible: {moe_impl.name} + {weight_processing['use_shuffled_weight']} + {weight_processing['layout']}" + ) + if ( + is_fp8_block_scale_moe + and moe_impl.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_MXFP8 + and not weight_processing["use_shuffled_weight"] + ): + pytest.skip("use_shuffled_weight must be true for MxFp8.") + if ( + is_fp8_block_scale_moe + and moe_impl.fp8_quantization_type == QuantMode.FP8_BLOCK_SCALE_MXFP8 + and weight_processing["layout"] != WeightLayout.MajorK + ): + pytest.skip("weight_layout must be MajorK for MxFp8.") + + if intermediate_size not in routing_config["compatible_intermediate_size"]: + pytest.skip( + f"Incompatible: intermediate_size={intermediate_size} with {routing_config['routing_method_type'].name} routing ({routing_config['num_experts']} experts)" + ) + + if moe_impl.quant_mode == QuantMode.MXINT4_BF16_BF16 and ( + intermediate_size % 256 != 0 or hidden_size % 256 != 0 + ): + pytest.skip( + f"Incompatible: intermediate_size={intermediate_size} or hidden_size={hidden_size} with MXINT4_BF16_BF16 quantization" + ) + + # TODO(jimmzhou): enable MxFP4xBf16 on SM103 + if ( + is_fp4_moe + and moe_impl.quant_mode == QuantMode.FP4_MXFP4_Bf16 + and compute_capability[0] == 10 + and compute_capability[1] == 3 + ): + pytest.xfail( + "Note(jimmzhou): Make MxFP4xBf16 nonfunctional on SM103 to avoid B200 regression" + ) + + if logits_dtype == torch.float32 and moe_impl.quant_mode not in [ + QuantMode.FP4_NVFP4_NVFP4, + QuantMode.FP8_PER_TENSOR, + QuantMode.FP8_BLOCK_SCALE_DEEPSEEK, + QuantMode.FP8_BLOCK_SCALE_MXFP8, + QuantMode.BF16, + ]: + pytest.skip( + f"Incompatible: logits_dtype={logits_dtype} with {type(moe_impl).__name__} + {moe_impl.quant_mode}" + ) diff --git a/tests/bench_ep_multi_gpu.py b/tests/bench_ep_multi_gpu.py new file mode 100644 index 0000000000..bbc533cec8 --- /dev/null +++ b/tests/bench_ep_multi_gpu.py @@ -0,0 +1,338 @@ +"""Multi-GPU expert-parallel (EP) MoE benchmark: DeepGEMM mega-MoE vs FlashInfer NVFP4. + +Layout: total_experts=512 sharded across `num_processes` GPUs (512/N experts per +shard); each shard holds `num_tokens` local tokens; top_k routing. + + * DeepGEMM `nvfp4_nvfp4_mega_moe`: native EP kernel (dispatch + grouped GEMM + + combine fused via symmetric memory). Timed end-to-end (comm is internal). + + * FlashInfer `cute_dsl_fused_moe_nvfp4`: single-device kernel, benchmarked with + the user-specified recipe per iteration: + 1) barrier + cuda sync across all devices + 2) start timer + 3) moe (local shard: 512/N experts, num_tokens tokens, top_k) + 4) all_reduce the output across all devices (combine proxy) + 5) stop timer + so its number includes the cross-device communication cost. + +We report the worst-rank average latency (max over ranks). + +Run: + python tests/bench_ep_multi_gpu.py --num-processes 4 + python tests/bench_ep_multi_gpu.py --num-processes 8 +""" +import argparse +import os +import sys + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +import deep_gemm +from deep_gemm.testing import bench_kineto +from deep_gemm.utils.dist import init_dist, dist_print +from deep_gemm.utils.math import per_token_cast_to_nvfp4, nvfp4_global_scale, per_token_cast_to_fp8 + +sys.path.insert(0, os.path.dirname(__file__)) +from test_nvfp4_mega_moe import _cast_l1_w, _cast_l2_w, GRAN_K +from bench_flashinfer_vs_deepgemm import build_flashinfer, build_flashinfer_cutlass +from fi_trtllm import build_flashinfer_trtllm +from bench_packed_fp4 import _cast_w_mxfp4 + +HIDDEN, INTER = 4608, 2560 +TOTAL_EXPERTS, TOP_K = 512, 16 +CLAMP = 10.0 + + +def rank_avg_ms(per_iter_ms, group): + """Aggregate a per-rank per-iter device time (ms) across GPUs by AVERAGING + (mean over ranks). Each rank's value is its own pure device time from the + trace; we just collect and average them.""" + world = dist.get_world_size(group) + t = torch.tensor([per_iter_ms], device='cuda', dtype=torch.float64) + dist.all_reduce(t, op=dist.ReduceOp.SUM, group=group) + return t.item() / world + + +def _self_dev_us(evt): + return getattr(evt, 'self_device_time_total', None) or getattr(evt, 'self_cuda_time_total', 0.0) + + +def kernel_rows(fn, group, warmup=10, iters=50): + """Return per-kernel (name, launches/iter, us/iter) device-time rows for `fn`.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + dist.barrier(group) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + fn() + torch.cuda.synchronize() + rows = [] + for evt in prof.key_averages(): + dev = getattr(evt, 'self_device_time_total', None) or getattr(evt, 'self_cuda_time_total', 0.0) + if dev and dev > 0: + rows.append((evt.key, evt.count / iters, dev / iters)) + rows.sort(key=lambda r: -r[2]) + return rows + + +def allreduce_time(buf, group, warmup=10, iters=50): + """Device time of the all_reduce NCCL kernel, with ~10ms sleep alignment so the + collective isn't inflated by cross-rank launch skew. Returns worst-rank avg ms.""" + fn = lambda: dist.all_reduce(buf, group=group) + for _ in range(warmup): + fn() + torch.cuda.synchronize() + dist.barrier(group) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + torch.cuda._sleep(int(2e7)) # ~10ms: align launches across ranks + fn() + torch.cuda.synchronize() + us, name = 0.0, '' + for evt in prof.key_averages(): + if 'allreduce' in evt.key.lower() or 'nccl' in evt.key.lower(): + us += (getattr(evt, 'self_device_time_total', None) or getattr(evt, 'self_cuda_time_total', 0.0)) + name = evt.key + return rank_avg_ms(us / iters / 1e3, group), name + + +def ep_total_device_time(step_fn, group, warmup=10, iters=50): + """Profile the full EP path in ONE window (steps 2->6: dispatch + slice + moe + + combine + all_reduce, all inside step_fn) and SUM the device time of every GPU + kernel. A dist sync is done once at the beginning to align ranks. Returns the + avg-over-ranks per-iter device time (ms).""" + for _ in range(warmup): + step_fn() + torch.cuda.synchronize() + dist.barrier(group) # sync at the beginning + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + step_fn() + torch.cuda.synchronize() + total_us = sum((getattr(e, 'self_device_time_total', None) or getattr(e, 'self_cuda_time_total', 0.0)) + for e in prof.key_averages()) + return rank_avg_ms(total_us / iters / 1e3, group) + + +def device_time_sum(fn, group, warmup=10, iters=50): + """Sum the self device time of every GPU kernel in the trace (per iter). + Used for the local FlashInfer MoE + its short combine all_reduce (no hard + cross-rank data dependency, so no launch alignment needed). NCCL/all_reduce + kernels (the combine) ARE counted. Returns worst-rank avg in ms.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + dist.barrier(group) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + fn() + torch.cuda.synchronize() + total_us = sum(_self_dev_us(evt) for evt in prof.key_averages()) + return rank_avg_ms(total_us / iters / 1e3, group) + + +def build_deepgemm_ep(group, x_bf16, sel, rw, local_experts): + num_tokens = x_bf16.shape[0] + num_max = 128 + w1 = torch.randn(local_experts, INTER * 2, HIDDEN, dtype=torch.bfloat16, device='cuda') / 10 + w2 = torch.randn(local_experts, HIDDEN, INTER, dtype=torch.bfloat16, device='cuda') / 10 + buf = deep_gemm.get_symm_buffer_for_mega_moe(group, TOTAL_EXPERTS, num_max, TOP_K, + HIDDEN, INTER, mma_type='nvfp4xnvfp4') + gs_x = nvfp4_global_scale(x_bf16) + l1, gate_gs, up_gs = _cast_l1_w(w1) + l2, down_gs = _cast_l2_w(w2) + l2act_gs = torch.ones(local_experts, device='cuda', dtype=torch.float32) + gate_alpha = (gs_x * gate_gs).contiguous() + up_alpha = (gs_x * up_gs).contiguous() + down_alpha = (l2act_gs * down_gs).contiguous() + l2_input_gs = (1.0 / l2act_gs).contiguous() + nl1, nl2 = deep_gemm.transform_weights_for_mega_moe(l1, l2) + xpn, xsfn = per_token_cast_to_nvfp4(x_bf16, gs_x, gran_k=GRAN_K) + buf.x[:num_tokens].copy_(xpn) + buf.x_sf[:num_tokens].copy_(xsfn.contiguous().view(torch.int32)) + buf.topk_idx[:num_tokens].copy_(sel) + buf.topk_weights[:num_tokens].copy_(rw) + y = torch.empty((num_tokens, HIDDEN), dtype=torch.bfloat16, device='cuda') + + def run(): + deep_gemm.nvfp4_nvfp4_mega_moe(y=y, l1_weights=nl1, l2_weights=nl2, sym_buffer=buf, + gate_alpha=gate_alpha, up_alpha=up_alpha, + l2_input_global_scale=l2_input_gs, down_alpha=down_alpha, + activation_clamp=CLAMP, fast_math=True) + return run, buf + + +def build_deepgemm_fp8_ep(group, x_bf16, sel, rw, local_experts): + num_tokens = x_bf16.shape[0] + num_max = 128 + w1 = torch.randn(local_experts, INTER * 2, HIDDEN, dtype=torch.bfloat16, device='cuda') / 10 + w2 = torch.randn(local_experts, HIDDEN, INTER, dtype=torch.bfloat16, device='cuda') / 10 + buf = deep_gemm.get_symm_buffer_for_mega_moe(group, TOTAL_EXPERTS, num_max, TOP_K, + HIDDEN, INTER, mma_type='fp8xfp4') + xp8, xsf8 = per_token_cast_to_fp8(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + tl1, tl2 = deep_gemm.transform_weights_for_mega_moe(_cast_w_mxfp4(w1), _cast_w_mxfp4(w2)) + buf.x[:num_tokens].copy_(xp8) + buf.x_sf[:num_tokens].copy_(xsf8) + buf.topk_idx[:num_tokens].copy_(sel) + buf.topk_weights[:num_tokens].copy_(rw) + y = torch.empty((num_tokens, HIDDEN), dtype=torch.bfloat16, device='cuda') + + def run(): + deep_gemm.fp8_fp4_mega_moe(y=y, l1_weights=tl1, l2_weights=tl2, sym_buffer=buf, + activation_clamp=CLAMP, fast_math=True) + return run, buf + + +def worker(local_rank, num_local_ranks, args): + global INTER + INTER = args.inter + rank, world, group = init_dist(local_rank, num_local_ranks) + local_experts = TOTAL_EXPERTS // world + local_offset = rank * local_experts + # Global batch = num_tokens (per-rank for DeepGEMM dispatch) x world. + dg_tokens = args.num_tokens # DeepGEMM: tokens per rank (dispatch model) + global_tokens = dg_tokens * world # FlashInfer: full batch replicated on every GPU + tk = min(TOP_K, local_experts) + + results = {} + + def mega_kineto(run): + dg_s = bench_kineto(run, 'mega_moe', suppress_kineto_output=True, + barrier=lambda: dist.barrier(group)) + return rank_avg_ms(dg_s * 1e3, group) + + # ---- DeepGEMM native EP mega kernel (dispatch+GEMM+combine fused), dg_tokens/rank ---- + torch.manual_seed(rank) + x_dg = torch.randn(dg_tokens, HIDDEN, dtype=torch.bfloat16, device='cuda') / 10 + rw_dg, sel_dg = torch.topk(F.softmax(torch.randn(dg_tokens, TOTAL_EXPERTS, device='cuda'), -1), TOP_K, dim=-1) + rw_dg = (rw_dg / rw_dg.sum(-1, keepdim=True)).float() + sel_dg = sel_dg.to(torch.int64) + for key, builder in (('dg_nvfp4', build_deepgemm_ep), ('dg_fp8', build_deepgemm_fp8_ep)): + try: + run, buf = builder(group, x_dg, sel_dg, rw_dg, local_experts) + results[key] = mega_kineto(run) + buf.destroy() + except Exception as ex: + results[key] = None + if rank == 0: + print(f' {key} err: {str(ex).splitlines()[-1][:120]}') + + # ---- FlashInfer: replicated input + expert-shard + all_reduce combine ---- + # Every GPU holds the SAME full batch (replicated) and the SAME global routing; + # each GPU owns experts [local_offset : local_offset+local_experts] and computes + # ONLY their contributions (via local_expert_offset / ep_rank). all_reduce then + # sums the partial [global_tokens, hidden] outputs into the full result. + torch.manual_seed(12345) # identical inputs/routing on all ranks (replicated) + x_rep = torch.randn(global_tokens, HIDDEN, dtype=torch.bfloat16, device='cuda') / 10 + rw_g, sel_g = torch.topk(F.softmax(torch.randn(global_tokens, TOTAL_EXPERTS, device='cuda'), -1), TOP_K, dim=-1) + rw_g = (rw_g / rw_g.sum(-1, keepdim=True)).float() + sel_g = sel_g.to(torch.int64) + # Local expert weight slice (random values; only shapes/positions matter for timing). + torch.manual_seed(1000 + rank) + w1 = torch.randn(local_experts, INTER * 2, HIDDEN, dtype=torch.bfloat16, device='cuda') / 10 + w2 = torch.randn(local_experts, HIDDEN, INTER, dtype=torch.bfloat16, device='cuda') / 10 + + # all_reduce-only device time (combine of the full [global_tokens, hidden] output) + # Allocate the combine buffer from NCCL's registered (multicast) allocator when + # available, so NCCL can dispatch the NVLS all_reduce kernel (plain torch tensors + # fall back to RING+LL even when NVLS is forced). + ar_buf, ar_pool, ar_be = None, None, None + dev_t = torch.device('cuda', torch.cuda.current_device()) + try: + ar_be = group._get_backend(dev_t) + ar_pool = torch.cuda.MemPool(ar_be.mem_allocator) + with torch.cuda.use_mem_pool(ar_pool): + ar_buf = torch.empty(global_tokens, HIDDEN, dtype=torch.bfloat16, device=dev_t) + ar_be.register_mem_pool(ar_pool, symm=True) # symm=True -> NVLS multicast registration + except Exception as ex: + if rank == 0: + print(f' (NVLS-registered buffer failed, using plain tensor: {str(ex).splitlines()[-1][:80]})') + ar_buf, ar_pool = torch.empty(global_tokens, HIDDEN, dtype=torch.bfloat16, device='cuda'), None + results['allreduce'], ar_kname = allreduce_time(ar_buf, group) + + # GeGLU (gelu-gated) where supported; cute_dsl and DeepGEMM are SwiGLU-only. + import flashinfer + geglu = flashinfer.ActivationType.Geglu + + def run_fi(name): + if name == 'fi_cutedsl': # SwiGLU-only (no activation param) + return build_flashinfer(x_rep, w1, w2, sel_g, rw_g, TOTAL_EXPERTS, tk, HIDDEN, INTER, + local_expert_offset=local_offset) + if name == 'fi_cutlass': + return build_flashinfer_cutlass(x_rep, w1, w2, sel_g, rw_g, local_experts, tk, HIDDEN, INTER, + ep_size=world, ep_rank=rank, activation_type=geglu) + return build_flashinfer_trtllm(x_rep, w1, w2, TOTAL_EXPERTS, tk, HIDDEN, INTER, + local_expert_offset=local_offset, activation_type=geglu) + + for name in ('fi_cutedsl', 'fi_cutlass', 'fi_trtllm'): + try: + fi_run = run_fi(name) + results[name + '_moe'] = device_time_sum(fi_run, group) # moe only (breakdown) + # Full path (steps 2->6) in one profiling window: moe + NVLS all_reduce combine. + def step(fi_run=fi_run): + fi_run() + dist.all_reduce(ar_buf, group=group) + results[name + '_total'] = ep_total_device_time(step, group) + except Exception as ex: + results[name + '_moe'] = results[name + '_total'] = None + if rank == 0: + print(f' {name} err: {str(ex).splitlines()[-1][:120]}') + + if rank == 0: + def us(v): + return f'{v*1e3:7.1f}' if v is not None else f'{"n/a":>7}' + ar = results.get('allreduce') + m_e = global_tokens * TOP_K / TOTAL_EXPERTS + ar_kib = global_tokens * HIDDEN * 2 / 1024 + print(f'\n=== EP MoE (device time from trace): {world} GPUs | total_experts={TOTAL_EXPERTS} | ' + f'{local_experts} experts/gpu | global batch={global_tokens} tokens ' + f'(DeepGEMM {dg_tokens}/rank dispatch; FlashInfer replicated+expert-sharded) | ' + f'~{m_e:g} tokens/expert | hidden={HIDDEN} inter={INTER} top_k={TOP_K} ===') + print(f' {"kernel":<34} {"moe":>8} {"allreduce":>10} {"sum":>8} {"total(2-6)":>11} (us)') + print(f' {"DeepGEMM fp8xfp4 mega (fused EP)":<34} {us(results.get("dg_fp8")):>8} ' + f'{"fused":>10} {us(results.get("dg_fp8")):>8} {us(results.get("dg_fp8")):>11}') + print(f' {"DeepGEMM nvfp4 mega (fused EP)":<34} {us(results.get("dg_nvfp4")):>8} ' + f'{"fused":>10} {us(results.get("dg_nvfp4")):>8} {us(results.get("dg_nvfp4")):>11}') + for name, lbl in (('fi_cutedsl', 'FlashInfer nvfp4 cute_dsl (SwiGLU)'), + ('fi_cutlass', 'FlashInfer nvfp4 cutlass (GeGLU)'), + ('fi_trtllm', 'FlashInfer nvfp4 trtllm-gen (GeGLU)')): + moe = results.get(name + '_moe') + tot_sum = (moe + ar) if (moe is not None and ar is not None) else None # separate-measure sum + tot_win = results.get(name + '_total') # one-window sum (2-6) + print(f' {lbl:<34} {us(moe):>8} {us(ar):>10} {us(tot_sum):>8} {us(tot_win):>11}') + print(f'\n all_reduce combine ([{global_tokens}, {HIDDEN}] bf16, {ar_kib:.0f} KiB, sleep-aligned, ' + f'NCCL_ALGO={os.environ.get("NCCL_ALGO", "auto")}): {us(ar)} us {ar_kname[:72]}') + print(' total(2-6) = one profiling window: dispatch+slice+moe+combine+all_reduce, summed device time') + + # Clean up the NVLS-registered combine pool before teardown (all ranks). + if ar_pool is not None: + try: + ar_be.deregister_mem_pool(ar_pool) + except Exception: + pass + dist.destroy_process_group() + + +def main(): + p = argparse.ArgumentParser() + p.add_argument('--num-processes', type=int, default=4) + p.add_argument('--num-tokens', type=int, default=32) + p.add_argument('--inter', type=int, default=2560, + help='intermediate size; fp8xfp4 needs a multiple of 512 (2304 is nvfp4-only)') + p.add_argument('--nccl-algo', type=str, default='', + help='force NCCL_ALGO for the combine all_reduce, e.g. NVLS, Ring, Tree') + args = p.parse_args() + if args.nccl_algo: + os.environ['NCCL_ALGO'] = args.nccl_algo + if 'NVLS' in args.nccl_algo.upper(): + os.environ.setdefault('NCCL_NVLS_ENABLE', '1') + os.environ.setdefault('MASTER_ADDR', '127.0.0.1') + os.environ.setdefault('MASTER_PORT', '12577') + torch.multiprocessing.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes) + + +if __name__ == '__main__': + main() diff --git a/tests/bench_flashinfer_vs_deepgemm.py b/tests/bench_flashinfer_vs_deepgemm.py new file mode 100644 index 0000000000..daeaa19d5a --- /dev/null +++ b/tests/bench_flashinfer_vs_deepgemm.py @@ -0,0 +1,280 @@ +"""Single-device perf comparison: + + DeepGEMM `fp8_fp4_mega_moe` (FP8 E4M3 activations x FP4 weights) + vs + FlashInfer `cute_dsl_fused_moe_nvfp4` (NVFP4 activations x NVFP4 weights) + +Both run the *same* MoE problem (identical shapes + identical routing) so the +comparison is apples-to-apples at the op level. We report end-to-end device +time (CUDA events), which includes all sub-kernels (FlashInfer dispatches +moe_sort + gemm1 + gemm2/finalize; DeepGEMM is a single mega kernel). +""" +import os +import sys +import torch +import torch.distributed as dist +import torch.nn.functional as F + +import deep_gemm +from deep_gemm.utils.math import ( + per_token_cast_to_fp8, per_token_cast_to_nvfp4, nvfp4_global_scale, +) + +sys.path.insert(0, os.path.dirname(__file__)) +from bench_packed_fp4 import _cast_w_mxfp4 +from test_nvfp4_mega_moe import _cast_l1_w, _cast_l2_w, GRAN_K + +import flashinfer +from flashinfer import fused_moe +from flashinfer.fused_moe import cute_dsl_fused_moe_nvfp4 +from flashinfer.fp4_quantization import fp4_quantize +from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout + +from fi_trtllm import build_flashinfer_trtllm + +SF_VEC = 16 +FLOAT8_E4M3_MAX = 448.0 +FLOAT4_E2M1_MAX = 6.0 +_round_up = lambda x, y: (x + y - 1) // y * y + + +def build_flashinfer_cutlass(x_bf16, w1_bf16, w2_bf16, sel, rw, num_experts, top_k, hidden, inter, + ep_size=1, ep_rank=0, activation_type=None): + """FlashInfer CUTLASS NVFP4 fused MoE (`cutlass_fused_moe`), prep per the + reference test. `num_experts` here is the LOCAL expert count (w1 groups); + EP slicing is handled via ep_size/ep_rank with global `sel` in [0, local*ep_size).""" + dev = x_bf16.device + e, n, k = num_experts, inter, hidden + w1_n = 2 * n + w1_q = torch.empty((e, w1_n, k // 2), device=dev, dtype=torch.uint8) + w2_q = torch.empty((e, k, n // 2), device=dev, dtype=torch.uint8) + w1_bs = torch.empty((e, _round_up(w1_n, 128), _round_up(k // 16, 4)), device=dev, dtype=torch.float8_e4m3fn) + w2_bs = torch.empty((e, _round_up(k, 128), _round_up(n // 16, 4)), device=dev, dtype=torch.float8_e4m3fn) + w1_gs = torch.empty(e, device=dev, dtype=torch.float32) + w2_gs = torch.empty(e, device=dev, dtype=torch.float32) + for ex in range(e): + w1_gs[ex] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_bf16[ex].abs().max().float() + w2_gs[ex] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_bf16[ex].abs().max().float() + w1_q[ex], w1_bs[ex] = fp4_quantize(w1_bf16[ex], w1_gs[ex]) + w2_q[ex], w2_bs[ex] = fp4_quantize(w2_bf16[ex], w2_gs[ex]) + a1_gs = torch.tensor(1.0, device=dev, dtype=torch.float32) + a2_gs = torch.tensor(1.0, device=dev, dtype=torch.float32) + quant_scales = [a1_gs, w1_bs.view(torch.int32), 1.0 / (a1_gs * w1_gs), + a2_gs, w2_bs.view(torch.int32), 1.0 / (a2_gs * w2_gs)] + hidden_states, input_sf = fp4_quantize(x_bf16, a1_gs) + out = torch.zeros_like(x_bf16) + w1l = w1_q.contiguous().view(torch.long) + w2l = w2_q.contiguous().view(torch.long) + seli = sel.to(torch.int32) + act = activation_type if activation_type is not None else flashinfer.ActivationType.Swiglu + + def run(): + return fused_moe.cutlass_fused_moe(hidden_states, seli, rw, w1l, w2l, + torch.bfloat16, quant_scales=quant_scales, + input_sf=input_sf, output=out, + ep_size=ep_size, ep_rank=ep_rank, + activation_type=act) + return run + + +def _time_events(launch, iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start.record() + for _ in range(iters): + launch() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters / 1e3 # seconds + + +def cuda_time(fn, warmup=10, iters=50): + """Pure device time via CUDA-graph replay (eliminates CPU dispatch/autotuner + overhead). Falls back to eager event timing if capture is unsupported.""" + # Warm up first so any JIT/autotuning happens before capture. + for _ in range(5): + fn() + torch.cuda.synchronize() + try: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + fn() + torch.cuda.synchronize() + for _ in range(warmup): + g.replay() + return _time_events(g.replay, iters), 'graph' + except Exception: + for _ in range(warmup): + fn() + return _time_events(fn, iters), 'eager' + + +def _interleave_gate(x, group_size=64, dim=1): + sizes = x.size() + dim = dim % x.dim() + x = x.view(*sizes[:dim], 2, sizes[dim] // (group_size * 2), group_size, *sizes[dim + 1:]) + return x.transpose(dim, dim + 1).contiguous().view(*sizes) + + +def build_flashinfer(x_bf16, w1_bf16, w2_bf16, sel, rw, num_experts, top_k, hidden, inter, + local_expert_offset=0): + """cute_dsl NVFP4 MoE. `num_experts` is the GLOBAL expert count (routing space); + `w1_bf16`/`w2_bf16` hold only this rank's local expert slice. For single device, + num_experts == w1.shape[0] and local_expert_offset == 0.""" + dev = x_bf16.device + g = torch.tensor([1.0], device=dev, dtype=torch.float32) + e = w1_bf16.shape[0] # local experts + + xq, xsf = fp4_quantize(x_bf16, global_scale=g, sf_vec_size=SF_VEC, is_sf_swizzled_layout=False) + xsf = xsf.unsqueeze(-1) + + fc1_rows = 2 * inter + w1_il = _interleave_gate(w1_bf16, group_size=64, dim=1) + w1q, w1sf = fp4_quantize(w1_il.reshape(e * fc1_rows, hidden), global_scale=g, + sf_vec_size=SF_VEC, is_sf_swizzled_layout=True) + w1q = w1q.view(e, fc1_rows, hidden // 2) + w1sf = convert_sf_to_mma_layout(w1sf, m=fc1_rows, k=hidden, num_groups=e, sf_vec_size=SF_VEC) + + w2q, w2sf = fp4_quantize(w2_bf16.view(e * hidden, inter), global_scale=g, + sf_vec_size=SF_VEC, is_sf_swizzled_layout=True) + w2q = w2q.view(e, hidden, inter // 2) + w2sf = convert_sf_to_mma_layout(w2sf, m=hidden, k=inter, num_groups=e, sf_vec_size=SF_VEC) + + ones = torch.ones(e, device=dev, dtype=torch.float32) + fc2_in = torch.tensor([1.0], device=dev, dtype=torch.float32) + + def run(): + return cute_dsl_fused_moe_nvfp4( + x=xq, x_sf=xsf, token_selected_experts=sel.to(torch.int32), + token_final_scales=rw, w1_weight=w1q, w1_weight_sf=w1sf, w1_alpha=ones, + fc2_input_scale=fc2_in, w2_weight=w2q, w2_weight_sf=w2sf, w2_alpha=ones, + num_experts=num_experts, top_k=top_k, + num_local_experts=e, local_expert_offset=local_expert_offset) + return run + + +def build_deepgemm(group, x_bf16, w1_bf16, w2_bf16, sel, rw, num_experts, top_k, hidden, inter): + num_tokens = x_bf16.shape[0] + num_max = max(128, num_tokens) + clamp = 10.0 + buf = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max, top_k, + hidden, inter, mma_type='fp8xfp4') + xp8, xsf8 = per_token_cast_to_fp8(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + tl1, tl2 = deep_gemm.transform_weights_for_mega_moe(_cast_w_mxfp4(w1_bf16), _cast_w_mxfp4(w2_bf16)) + buf.x[:num_tokens].copy_(xp8) + buf.x_sf[:num_tokens].copy_(xsf8) + buf.topk_idx[:num_tokens].copy_(sel) + buf.topk_weights[:num_tokens].copy_(rw) + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + + def run(): + deep_gemm.fp8_fp4_mega_moe(y=y, l1_weights=tl1, l2_weights=tl2, sym_buffer=buf, + activation_clamp=clamp, fast_math=True) + return run, buf + + +def build_deepgemm_nvfp4(group, x_bf16, w1_bf16, w2_bf16, sel, rw, num_experts, top_k, hidden, inter): + num_tokens = x_bf16.shape[0] + num_max = max(128, num_tokens) + clamp = 10.0 + buf = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max, top_k, + hidden, inter, mma_type='nvfp4xnvfp4') + gs_x = nvfp4_global_scale(x_bf16) + l1, gate_gs, up_gs = _cast_l1_w(w1_bf16) + l2, down_gs = _cast_l2_w(w2_bf16) + # Per-expert global scales: values don't affect timing, so skip the costly + # per-token L2-activation amax estimation and use unit L2-input scale. + l2act_gs = torch.ones(num_experts, device='cuda', dtype=torch.float32) + gate_alpha = (gs_x * gate_gs).contiguous() + up_alpha = (gs_x * up_gs).contiguous() + down_alpha = (l2act_gs * down_gs).contiguous() + l2_input_gs = (1.0 / l2act_gs).contiguous() + nl1, nl2 = deep_gemm.transform_weights_for_mega_moe(l1, l2) + xpn, xsfn = per_token_cast_to_nvfp4(x_bf16, gs_x, gran_k=GRAN_K) + buf.x[:num_tokens].copy_(xpn) + buf.x_sf[:num_tokens].copy_(xsfn.contiguous().view(torch.int32)) + buf.topk_idx[:num_tokens].copy_(sel) + buf.topk_weights[:num_tokens].copy_(rw) + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + + def run(): + deep_gemm.nvfp4_nvfp4_mega_moe(y=y, l1_weights=nl1, l2_weights=nl2, sym_buffer=buf, + gate_alpha=gate_alpha, up_alpha=up_alpha, + l2_input_global_scale=l2_input_gs, down_alpha=down_alpha, + activation_clamp=clamp, fast_math=True) + return run, buf + + +def bench_one(group, num_tokens, hidden, inter, num_experts, top_k): + torch.manual_seed(0) + dev = 'cuda' + x_bf16 = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=dev) / 10 + w1_bf16 = torch.randn(num_experts, inter * 2, hidden, dtype=torch.bfloat16, device=dev) / 10 + w2_bf16 = torch.randn(num_experts, hidden, inter, dtype=torch.bfloat16, device=dev) / 10 + + logits = torch.randn(num_tokens, num_experts, device=dev) + probs = F.softmax(logits, dim=1, dtype=torch.float) + rw, sel = torch.topk(probs, top_k, dim=-1) + rw = (rw / rw.sum(-1, keepdim=True)).float() + sel = sel.to(torch.int64) + + res = {} + + def try_run(name, builder): + try: + out = builder() + run = out[0] if isinstance(out, tuple) else out + res[name], res[name + '_m'] = cuda_time(run) + if isinstance(out, tuple): + out[1].destroy() + except Exception as ex: + res[name] = None + res[name + '_e'] = str(ex).splitlines()[-1][:90] + + try_run('dg8', lambda: build_deepgemm(group, x_bf16, w1_bf16, w2_bf16, sel, rw, + num_experts, top_k, hidden, inter)) + try_run('dg4', lambda: build_deepgemm_nvfp4(group, x_bf16, w1_bf16, w2_bf16, sel, rw, + num_experts, top_k, hidden, inter)) + try_run('fi_cutedsl', lambda: build_flashinfer(x_bf16, w1_bf16, w2_bf16, sel, rw, + num_experts, top_k, hidden, inter)) + try_run('fi_cutlass', lambda: build_flashinfer_cutlass(x_bf16, w1_bf16, w2_bf16, sel, rw, + num_experts, top_k, hidden, inter)) + try_run('fi_trtllm', lambda: build_flashinfer_trtllm(x_bf16, w1_bf16, w2_bf16, + num_experts, top_k, hidden, inter)) + + def f(k): + if res.get(k) is None: + return f'{"n/a":>9}' + return f'{res[k]*1e6:>8.1f}{res.get(k + "_m", "")[0]}' # suffix g/e + + print(f'{num_tokens:>5} {num_experts:>4} {top_k:>3} {hidden:>5} {inter:>5} | ' + f'{f("dg8")} | {f("dg4")} | {f("fi_cutedsl")} | {f("fi_cutlass")} | {f("fi_trtllm")}') + for k in ('dg8', 'dg4', 'fi_cutedsl', 'fi_cutlass', 'fi_trtllm'): + if res.get(k + '_e'): + print(f' {k} err: {res[k + "_e"]}') + + +def main(): + os.environ.setdefault('MASTER_ADDR', '127.0.0.1') + os.environ.setdefault('MASTER_PORT', '12566') + dist.init_process_group('nccl', rank=0, world_size=1) + group = dist.group.WORLD + + print(f'torch {torch.__version__}, flashinfer {flashinfer.__version__}, {torch.cuda.get_device_name(0)}') + print('DeepGEMM mega-MoE (fp8xfp4 / nvfp4) vs FlashInfer NVFP4 MoE backends (device us, CUDA graph; g=graph e=eager)') + print(f'{"tok":>5} {"exp":>4} {"tk":>3} {"hid":>5} {"int":>5} | {"dg_fp8":>9} | {"dg_nvfp4":>9} | ' + f'{"fi_cutedsl":>9} | {"fi_cutlass":>9} | {"fi_trtllm":>9}') + for cfg in ( + (128, 2048, 2048, 32, 4), + (512, 2048, 2048, 32, 4), + (1024, 7168, 2560, 256, 8), + # inter=2304 is nvfp4-only (fp8xfp4 needs inter % 512 == 0) + (32, 4608, 2304, 256, 16), + ): + bench_one(group, *cfg) + dist.destroy_process_group() + + +if __name__ == '__main__': + main() diff --git a/tests/bench_kernel_breakdown.py b/tests/bench_kernel_breakdown.py new file mode 100644 index 0000000000..3479c77dfb --- /dev/null +++ b/tests/bench_kernel_breakdown.py @@ -0,0 +1,94 @@ +"""Per-kernel device-time breakdown for each MoE backend (single device). + +Profiles each backend's run closure and lists every GPU kernel with its +per-iteration count and device time, so you can see how many kernels each +backend launches and where the time goes. +""" +import os +import sys + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +import deep_gemm + +sys.path.insert(0, os.path.dirname(__file__)) +from bench_flashinfer_vs_deepgemm import ( + build_deepgemm, build_deepgemm_nvfp4, build_flashinfer, build_flashinfer_cutlass, +) +from fi_trtllm import build_flashinfer_trtllm + + +def kernel_breakdown(run, warmup=10, iters=50): + for _ in range(warmup): + run() + torch.cuda.synchronize() + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + run() + torch.cuda.synchronize() + rows = [] + for evt in prof.key_averages(): + dev = getattr(evt, 'self_device_time_total', None) or getattr(evt, 'self_cuda_time_total', 0.0) + if dev and dev > 0: + rows.append((evt.key, evt.count / iters, dev / iters)) # name, launches/iter, us/iter + rows.sort(key=lambda r: -r[2]) + return rows + + +def show(title, build, *, is_tuple=False): + print(f'\n===== {title} =====') + try: + out = build() + run = out[0] if is_tuple else out + rows = kernel_breakdown(run) + if is_tuple: + out[1].destroy() + except Exception as ex: + print(f' ERROR: {str(ex).splitlines()[-1][:120]}') + return + total = sum(r[2] for r in rows) + n_launches = sum(r[1] for r in rows) + print(f' {len(rows)} distinct kernels, {n_launches:.0f} launches/iter, total device {total:.1f} us/iter') + print(f' {"#/it":>5} {"us/it":>8} kernel') + for name, cnt, us in rows: + print(f' {cnt:>5.0f} {us:>8.1f} {name[:88]}') + + +def main(): + os.environ.setdefault('MASTER_ADDR', '127.0.0.1') + os.environ.setdefault('MASTER_PORT', '12599') + dist.init_process_group('nccl', rank=0, world_size=1) + group = dist.group.WORLD + + # 4-GPU-shard view of the previous shape: 128 experts, 32 tokens, top-16, m_e=4. + num_tokens, hidden, inter, num_experts, top_k = 32, 4608, 2304, 128, 16 + print(f'torch {torch.__version__}, flashinfer {__import__("flashinfer").__version__}, ' + f'{torch.cuda.get_device_name(0)}') + print(f'shape: tokens={num_tokens} experts={num_experts} top_k={top_k} hidden={hidden} inter={inter}') + + torch.manual_seed(0) + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device='cuda') / 10 + w1 = torch.randn(num_experts, inter * 2, hidden, dtype=torch.bfloat16, device='cuda') / 10 + w2 = torch.randn(num_experts, hidden, inter, dtype=torch.bfloat16, device='cuda') / 10 + rw, sel = torch.topk(F.softmax(torch.randn(num_tokens, num_experts, device='cuda'), -1), top_k, dim=-1) + rw = (rw / rw.sum(-1, keepdim=True)).float() + sel = sel.to(torch.int64) + + show('DeepGEMM fp8xfp4 mega (fused EP)', + lambda: build_deepgemm(group, x, w1, w2, sel, rw, num_experts, top_k, hidden, inter), is_tuple=True) + show('DeepGEMM nvfp4 mega (fused EP)', + lambda: build_deepgemm_nvfp4(group, x, w1, w2, sel, rw, num_experts, top_k, hidden, inter), is_tuple=True) + show('FlashInfer nvfp4 (cute_dsl)', + lambda: build_flashinfer(x, w1, w2, sel, rw, num_experts, top_k, hidden, inter)) + show('FlashInfer nvfp4 (cutlass)', + lambda: build_flashinfer_cutlass(x, w1, w2, sel, rw, num_experts, top_k, hidden, inter)) + show('FlashInfer nvfp4 (trtllm-gen)', + lambda: build_flashinfer_trtllm(x, w1, w2, num_experts, top_k, hidden, inter)) + + dist.destroy_process_group() + + +if __name__ == '__main__': + main() diff --git a/tests/bench_packed_fp4.py b/tests/bench_packed_fp4.py index 14a3c3a78b..43025b952f 100644 --- a/tests/bench_packed_fp4.py +++ b/tests/bench_packed_fp4.py @@ -11,7 +11,7 @@ import deep_gemm from deep_gemm.testing import bench_kineto from deep_gemm.utils.math import ( - per_token_cast_to_fp4, per_token_cast_to_nvfp4, nvfp4_global_scale, + per_token_cast_to_fp4, per_token_cast_to_fp8, per_token_cast_to_nvfp4, nvfp4_global_scale, ) sys.path.insert(0, os.path.dirname(__file__)) @@ -61,8 +61,8 @@ def bench_mega(): group = dist.group.WORLD clamp = 10.0 - print('=== Packed-FP4 mega-MoE (single rank) ===') - print(f'{"tok":>5} {"exp":>4} {"topk":>4} {"hid":>5} {"int":>5} | {"mxfp4 us":>9} | {"nvfp4 us":>9} | nv/mx') + print('=== mega-MoE (single rank): fp8xfp4 vs mxfp4 vs nvfp4 ===') + print(f'{"tok":>5} {"exp":>4} {"topk":>4} {"hid":>5} {"int":>5} | {"fp8fp4 us":>9} | {"mxfp4 us":>9} | {"nvfp4 us":>9} | nv/fp8 | nv/mx') torch.manual_seed(0) for num_tokens, num_experts, num_topk, hidden, inter in ( (128, 8, 2, 2048, 2048), @@ -77,6 +77,18 @@ def bench_mega(): topk_weights, topk_idx = torch.topk(scores.softmax(-1), num_topk, dim=-1) y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + # FP8xFP4 (FP8 E4M3 activations, FP4 weights) + buf_f8 = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, + hidden, inter, mma_type='fp8xfp4') + xp8, xsf8 = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + f8l1, f8l2 = deep_gemm.transform_weights_for_mega_moe(_cast_w_mxfp4(l1w), _cast_w_mxfp4(l2w)) + buf_f8.x[:num_tokens].copy_(xp8); buf_f8.x_sf[:num_tokens].copy_(xsf8) + buf_f8.topk_idx[:num_tokens].copy_(topk_idx); buf_f8.topk_weights[:num_tokens].copy_(topk_weights) + t_f8 = bench_kineto(lambda: deep_gemm.fp8_fp4_mega_moe(y=y, l1_weights=f8l1, l2_weights=f8l2, sym_buffer=buf_f8, + activation_clamp=clamp, fast_math=True), + 'mega_moe', suppress_kineto_output=True) + buf_f8.destroy() + # MXFP4 buf_mx = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, hidden, inter, mma_type='mxfp4xmxfp4') @@ -110,7 +122,7 @@ def bench_mega(): buf_nv.destroy() print(f'{num_tokens:>5} {num_experts:>4} {num_topk:>4} {hidden:>5} {inter:>5} | ' - f'{t_mx*1e6:>9.1f} | {t_nv*1e6:>9.1f} | {t_nv/t_mx:>4.2f}x') + f'{t_f8*1e6:>9.1f} | {t_mx*1e6:>9.1f} | {t_nv*1e6:>9.1f} | {t_nv/t_f8:>4.2f}x | {t_nv/t_mx:>4.2f}x') dist.destroy_process_group() print() diff --git a/tests/fi_trtllm.py b/tests/fi_trtllm.py new file mode 100644 index 0000000000..d8407751e6 --- /dev/null +++ b/tests/fi_trtllm.py @@ -0,0 +1,96 @@ +"""FlashInfer TRT-LLM-gen NVFP4 fused MoE driver (`trtllm_fp4_block_scale_moe`). + +Reuses flashinfer's own (version-matched, vendored) offline weight-shuffle + +routing-reference harness from `_fi_vendor/test_trtllm_gen_fused_moe.py`, then +exposes a thin timeable closure that calls the kernel directly (mirroring the +harness's `_run_moe_computation`, but without the per-call CUDA-graph wrapper). +""" +import os +import sys + +import torch + +sys.path.insert(0, os.path.dirname(__file__)) +from _fi_vendor.test_trtllm_gen_fused_moe import ( + FP4Moe, moe_args, routing_reference_renormalize, +) +from _fi_vendor.utils import QuantMode +from flashinfer import RoutingMethodType, ActivationType +from flashinfer.fused_moe import trtllm_fp4_block_scale_moe + + +class _Dequant: + """Minimal stand-in for moe_args_dequant: prepare_static_weights only reads + `c_global_sf`. Scale values don't affect kernel timing, so we stub it (avoids + the slow global reference MoE and lets each rank build only its expert slice).""" + def __init__(self, c_global_sf): + self.c_global_sf = c_global_sf + + +def build_flashinfer_trtllm(x_bf16, w1_bf16, w2_bf16, num_experts, top_k, hidden, inter, + local_expert_offset=0, activation_type=None): + """TRT-LLM-gen NVFP4 MoE. `num_experts` is the GLOBAL expert count (routing + space); `w1_bf16`/`w2_bf16` hold only this rank's local expert slice. The + kernel computes only the local experts via local_expert_offset/local_num_experts.""" + dev = x_bf16.device + act = activation_type if activation_type is not None else ActivationType.Swiglu + moe = FP4Moe(QuantMode.FP4_NVFP4_NVFP4) + moe._cache_permute_indices = {} + num_tokens = x_bf16.shape[0] + num_local = w1_bf16.shape[0] + padding = 8 + + weights_data = moe.quantize_weights(w1_bf16, w2_bf16, x_bf16) + inputs_data = moe.quantize_inputs(x_bf16, weights_data['hidden_states_scale_global']) + q = {**weights_data, **inputs_data} + + args = moe_args( + num_tokens, num_local, hidden, inter, top_k, padding, + q['hidden_states'], q['hidden_states_scale'], q['hidden_states_scale_global'], + None, q['gemm1_weights'], q['gemm1_scales'], q['gemm1_scales_global'], + q['gemm2_weights'], q['gemm2_scales'], q['gemm2_scales_global'], + None, False, act) + + args_dequant = _Dequant(torch.tensor(1.0, device=dev, dtype=torch.float32)) + static = moe.prepare_static_weights_for_kernel( + args_dequant, args, w1_bf16, w2_bf16, hidden, inter, num_local, None) + + # Global routing space (top-k over all `num_experts`); each rank computes only its slice. + expert_logits = torch.randn(num_tokens, num_experts, device=dev).to(torch.bfloat16) + inp = moe.quantize_inputs(x_bf16, weights_data['hidden_states_scale_global'], is_swizzling=False) + hs, hs_sf = inp['hidden_states'], inp['hidden_states_scale'] + + def run(): + return trtllm_fp4_block_scale_moe( + routing_logits=expert_logits, routing_bias=None, + hidden_states=hs, hidden_states_scale=hs_sf, + gemm1_weights=static['gemm1_weights_fp4_shuffled'], + gemm1_weights_scale=static['gemm1_scales_fp4_shuffled'], + gemm1_bias=None, gemm1_alpha=None, gemm1_beta=None, gemm1_clamp_limit=None, + gemm2_weights=static['gemm2_weights_fp4_shuffled'], + gemm2_weights_scale=static['gemm2_scales_fp4_shuffled'], + gemm2_bias=None, + output1_scale_scalar=static['scale_c_fc1'], + output1_scale_gate_scalar=static['scale_gate_fc1'], + output2_scale_scalar=static['scale_c_fc2'], + num_experts=num_experts, top_k=top_k, n_group=None, topk_group=None, + intermediate_size=inter, local_expert_offset=local_expert_offset, + local_num_experts=num_local, + routed_scaling_factor=None, routing_method_type=RoutingMethodType.Renormalize, + activation_type=act, do_finalize=True, + tune_max_num_tokens=8192, norm_topk_prob=True) + return run + + +if __name__ == '__main__': + torch.manual_seed(0) + he, it, ne, tk, nt = 4608, 2560, 128, 16, 32 + x = torch.randn(nt, he, dtype=torch.bfloat16, device='cuda') / 10 + w1 = torch.randn(ne, it * 2, he, dtype=torch.bfloat16, device='cuda') / 10 + w2 = torch.randn(ne, he, it, dtype=torch.bfloat16, device='cuda') / 10 + run = build_flashinfer_trtllm(x, w1, w2, ne, tk, he, it) + out = run() + torch.cuda.synchronize() + o = out[0] if isinstance(out, (list, tuple)) else out + print('trtllm output:', type(out), o.shape, o.dtype, 'mean', o.float().abs().mean().item()) + print('OK') From 87458760616b6889636af890cd0edce5c2fb812c Mon Sep 17 00:00:00 2001 From: Ying Zhang Date: Mon, 6 Jul 2026 00:20:16 +0000 Subject: [PATCH 7/9] docs: move HANDOFF.md to doc/fp4_moe.md Reorganize the FP4 MoE handoff note under a new doc/ directory and rename it to reflect its subject. Co-Authored-By: Claude --- HANDOFF.md => doc/fp4_moe.md | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename HANDOFF.md => doc/fp4_moe.md (100%) diff --git a/HANDOFF.md b/doc/fp4_moe.md similarity index 100% rename from HANDOFF.md rename to doc/fp4_moe.md From 0d7c72a05c4c398c1d7e75b3dfe04bb0ca1b0a02 Mon Sep 17 00:00:00 2001 From: Ying Zhang Date: Mon, 6 Jul 2026 01:08:00 +0000 Subject: [PATCH 8/9] test: reorganize FP4 tests, move benchmarks, expand mega-MoE coverage * Move the four bench_*.py scripts from tests/ into a new benchmarks/ directory; fix their sys.path (add tests/) and repoint helper imports at the consolidated test modules (signatures preserved via aliases). * Consolidate the mxfp4/nvfp4 test pairs into two parameterized scripts: tests/test_fp4_gemm.py and tests/test_fp4_mega_moe.py, each looping over ('mxfp4', 'nvfp4') via a per-format strategy. Delete the four old test_mxfp4_*/test_nvfp4_* files. Update doc/fp4_moe.md test paths. * Expand mega-MoE coverage: 1-rank shape matrix (8 shapes: small, odd, single-token, large-asymmetric, masked) and a multi-rank EP matrix (world 2/4/8) via torch.multiprocessing.spawn + init_dist. Multi-rank uses replicated x/weights/routing so the NVFP4 per-expert l2act_gs reference is computable without a cross-rank all-gather; each rank asserts against the full reference. Helper signatures kept stable for the benchmark imports. Verified on B200: 32 mega-MoE cases pass (baseline mxfp4 0.00075 / nvfp4 0.00058 unchanged; 1-rank 16 + EP2 6 + EP4 4 + EP8 4, all diff < 0.05, incl. masked routing); GEMM 10/10 diff=0.0. Co-Authored-By: Claude --- {tests => benchmarks}/bench_ep_multi_gpu.py | 11 +- .../bench_flashinfer_vs_deepgemm.py | 7 +- .../bench_kernel_breakdown.py | 0 {tests => benchmarks}/bench_packed_fp4.py | 19 +- doc/fp4_moe.md | 18 +- tests/test_fp4_gemm.py | 112 +++++ tests/test_fp4_mega_moe.py | 433 ++++++++++++++++++ tests/test_mxfp4_gemm.py | 48 -- tests/test_mxfp4_mega_moe.py | 96 ---- tests/test_nvfp4_gemm.py | 64 --- tests/test_nvfp4_mega_moe.py | 159 ------- 11 files changed, 580 insertions(+), 387 deletions(-) rename {tests => benchmarks}/bench_ep_multi_gpu.py (98%) rename {tests => benchmarks}/bench_flashinfer_vs_deepgemm.py (98%) rename {tests => benchmarks}/bench_kernel_breakdown.py (100%) rename {tests => benchmarks}/bench_packed_fp4.py (93%) create mode 100644 tests/test_fp4_gemm.py create mode 100644 tests/test_fp4_mega_moe.py delete mode 100644 tests/test_mxfp4_gemm.py delete mode 100644 tests/test_mxfp4_mega_moe.py delete mode 100644 tests/test_nvfp4_gemm.py delete mode 100644 tests/test_nvfp4_mega_moe.py diff --git a/tests/bench_ep_multi_gpu.py b/benchmarks/bench_ep_multi_gpu.py similarity index 98% rename from tests/bench_ep_multi_gpu.py rename to benchmarks/bench_ep_multi_gpu.py index bbc533cec8..46b0650166 100644 --- a/tests/bench_ep_multi_gpu.py +++ b/benchmarks/bench_ep_multi_gpu.py @@ -18,8 +18,8 @@ We report the worst-rank average latency (max over ranks). Run: - python tests/bench_ep_multi_gpu.py --num-processes 4 - python tests/bench_ep_multi_gpu.py --num-processes 8 + python benchmarks/bench_ep_multi_gpu.py --num-processes 4 + python benchmarks/bench_ep_multi_gpu.py --num-processes 8 """ import argparse import os @@ -35,7 +35,12 @@ from deep_gemm.utils.math import per_token_cast_to_nvfp4, nvfp4_global_scale, per_token_cast_to_fp8 sys.path.insert(0, os.path.dirname(__file__)) -from test_nvfp4_mega_moe import _cast_l1_w, _cast_l2_w, GRAN_K +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'tests')) +from test_fp4_mega_moe import ( + _cast_l1_w_nvfp4 as _cast_l1_w, + _cast_l2_w_nvfp4 as _cast_l2_w, + NVFP4_GRAN_K as GRAN_K, +) from bench_flashinfer_vs_deepgemm import build_flashinfer, build_flashinfer_cutlass from fi_trtllm import build_flashinfer_trtllm from bench_packed_fp4 import _cast_w_mxfp4 diff --git a/tests/bench_flashinfer_vs_deepgemm.py b/benchmarks/bench_flashinfer_vs_deepgemm.py similarity index 98% rename from tests/bench_flashinfer_vs_deepgemm.py rename to benchmarks/bench_flashinfer_vs_deepgemm.py index daeaa19d5a..f455aca92a 100644 --- a/tests/bench_flashinfer_vs_deepgemm.py +++ b/benchmarks/bench_flashinfer_vs_deepgemm.py @@ -21,8 +21,13 @@ ) sys.path.insert(0, os.path.dirname(__file__)) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'tests')) from bench_packed_fp4 import _cast_w_mxfp4 -from test_nvfp4_mega_moe import _cast_l1_w, _cast_l2_w, GRAN_K +from test_fp4_mega_moe import ( + _cast_l1_w_nvfp4 as _cast_l1_w, + _cast_l2_w_nvfp4 as _cast_l2_w, + NVFP4_GRAN_K as GRAN_K, +) import flashinfer from flashinfer import fused_moe diff --git a/tests/bench_kernel_breakdown.py b/benchmarks/bench_kernel_breakdown.py similarity index 100% rename from tests/bench_kernel_breakdown.py rename to benchmarks/bench_kernel_breakdown.py diff --git a/tests/bench_packed_fp4.py b/benchmarks/bench_packed_fp4.py similarity index 93% rename from tests/bench_packed_fp4.py rename to benchmarks/bench_packed_fp4.py index 43025b952f..578a3fc37c 100644 --- a/tests/bench_packed_fp4.py +++ b/benchmarks/bench_packed_fp4.py @@ -15,9 +15,14 @@ ) sys.path.insert(0, os.path.dirname(__file__)) -from test_mxfp4_gemm import _prepare as _prep_mxfp4_gemm -from test_nvfp4_gemm import _prepare as _prep_nvfp4_gemm -from test_nvfp4_mega_moe import _cast_l1_w, _cast_l2_w, _estimate_l2act_gs, GRAN_K +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'tests')) +from test_fp4_gemm import _prepare_mxfp4 as _prep_mxfp4_gemm, _prepare_nvfp4 as _prep_nvfp4_gemm +from test_fp4_mega_moe import ( + _cast_l1_w_nvfp4 as _cast_l1_w, + _cast_l2_w_nvfp4 as _cast_l2_w, + _estimate_l2act_gs, + NVFP4_GRAN_K as GRAN_K, +) def bench_gemm(): @@ -29,14 +34,14 @@ def bench_gemm(): d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16) flops = 2.0 * m * n * k - ap, asf, _ = _prep_mxfp4_gemm(a) - bp, bsf, _ = _prep_mxfp4_gemm(b) + ap, asf, _, _ = _prep_mxfp4_gemm(a) + bp, bsf, _, _ = _prep_mxfp4_gemm(b) t_mx = bench_kineto(lambda: deep_gemm.mxfp4_gemm_nt((ap, asf), (bp, bsf), d), 'mxfp4_gemm', suppress_kineto_output=True) gsa, gsb = nvfp4_global_scale(a), nvfp4_global_scale(b) - ap2, asf2, _ = _prep_nvfp4_gemm(a, gsa) - bp2, bsf2, _ = _prep_nvfp4_gemm(b, gsb) + ap2, asf2, _, _ = _prep_nvfp4_gemm(a) + bp2, bsf2, _, _ = _prep_nvfp4_gemm(b) t_nv = bench_kineto(lambda: deep_gemm.nvfp4_gemm_nt((ap2, asf2), (bp2, bsf2), d, a_global_scale=gsa, b_global_scale=gsb), 'mxfp4_gemm', suppress_kineto_output=True) diff --git a/doc/fp4_moe.md b/doc/fp4_moe.md index 2c8283155f..4f972dcb75 100644 --- a/doc/fp4_moe.md +++ b/doc/fp4_moe.md @@ -1,7 +1,7 @@ # Handoff: Packed MXFP4 × MXFP4 for SM100 (Mega MoE) Status: **DONE. Packed MXFP4 × MXFP4 mega-MoE kernel PASSES end-to-end on B200 -(`tests/test_mxfp4_mega_moe.py`, `diff = 0.00075`).** Standalone GEMM also validated. +(`tests/test_fp4_mega_moe.py`, `diff = 0.00075`).** Standalone GEMM also validated. ## Update (final) — mega-MoE FIXED end-to-end @@ -19,7 +19,7 @@ Two fixes closed it out: `lane_idx < 4` (b==0), so the existing UE8M0 SF code is correct as-is. Isolation (`/tmp/iso2.py`, identity L2 weights): `diff 0.98 -> 0.00078`. -Run: `python tests/test_mxfp4_mega_moe.py` -> `diff = 0.00075 MXFP4 mega MoE passed.` +Run: `python tests/test_fp4_mega_moe.py` -> `diff = 0.00075 MXFP4 mega MoE passed.` ### Update — TRUE 2-CTA `cta_group::2` multicast TMA (standalone + mega) @@ -40,7 +40,7 @@ CUTLASS kernel** (example 72b forced to a 2x1 cluster, `-DDG_TMA_LOG` printf in while the non-leader's `full` expects only its own SF; both CTAs' transposer/sync warps wait their own `full` and arrive on the leader's `with_sf`. -Standalone `tests/test_mxfp4_gemm.py`: all sizes `diff=0.0`. Mega `diff=0.00075`. Perf +Standalone `tests/test_fp4_gemm.py`: all sizes `diff=0.0`. Mega `diff=0.00075`. Perf unchanged vs the per-CTA path (2x1 has no multicast bandwidth saving) — still ~1.12-1.23x over FP8xFP4. Bug that caused the earlier hang: full-box+shared-coord + under-counted `expect_tx`; fixed by per-CTA box/coord + correct `2x data + own SF` accounting. @@ -93,7 +93,7 @@ The long-standing 2-CTA TMA **deadlock is resolved**. Root cause + fix: `csrc/jit_kernels/heuristics/mega_moe.hpp` (this was the illegal-access cause after enlarging `SharedStorage`). -Result: `python tests/test_mxfp4_mega_moe.py` runs the kernel cleanly; output +Result: `python tests/test_fp4_mega_moe.py` runs the kernel cleanly; output magnitude is sane (`y~0.23` vs `ref~0.19`) but **`diff≈0.99`** (values scrambled). ### Remaining task: L1 packed-FP4 swap-AB epilogue transpose (numerics) @@ -139,7 +139,7 @@ What now exists and passes: A/B tensormaps (`make_tma_2d_desc(..., fp4_unpacked_smem=false)` → `16U4_ALIGN8B`), UE8M0 SF descriptors, BF16 D, smem sizing, cluster-2 launch. Registered in `csrc/apis/gemm.hpp` as `mxfp4_gemm_nt` (+ pybind), exported in `deep_gemm/__init__.py`. -- `tests/test_mxfp4_gemm.py` — quantizes A/B to packed E2M1 + UE8M0 SF (gran-32, via +- `tests/test_fp4_gemm.py` — quantizes A/B to packed E2M1 + UE8M0 SF (gran-32, via `per_token_cast_to_fp4` + `get_mn_major_tma_aligned_packed_ue8m0_tensor`), runs the kernel, compares to a dequant reference. All cases `diff == 0.0`. @@ -151,7 +151,7 @@ fixed during bring-up: (the descriptor's smem box already spans the full packed `BLOCK_K`); the `tma::copy` atom-splitter assumes byte-sized elems and overruns the smem stage for sub-byte FP4. -Run it: `PYTHONPATH=$PWD python tests/test_mxfp4_gemm.py` (after `./develop.sh`). +Run it: `PYTHONPATH=$PWD python tests/test_fp4_gemm.py` (after `./develop.sh`). ### Current standalone limitations (intentional for the de-risk) - Hardcoded config: `BLOCK_M=BLOCK_N=BLOCK_K=128`, `kNumStages=4`, swap-AB, @@ -245,7 +245,7 @@ Other likely tuning points: epilogue STSM swizzle/`STORE_BLOCK_M`, register spli 1. ~~Toolchain~~ — DONE (CUDA 13.0 + B200; submodules initialized). 2. ~~Host runtime + JIT wiring~~ — DONE (`sm100_mxfp4_gemm.hpp`). -3. ~~Python entry + test~~ — DONE (`mxfp4_gemm_nt`, `tests/test_mxfp4_gemm.py`). +3. ~~Python entry + test~~ — DONE (`mxfp4_gemm_nt`, `tests/test_fp4_gemm.py`). 4. ~~Iterate the `// VALIDATE` spots on B200~~ — DONE (all confirmed, `diff == 0.0`). ## Mega-MoE port — STAGE 1 done (compiles + launches on B200), STAGE 2 = numerics @@ -274,7 +274,7 @@ certainly doesn't match the L2 `mxf4` TMA read-back, and there may be dispatch/b interplay to debug. STAGE 2 (in progress): -- DONE: single-rank **torch MoE reference** at `tests/test_mxfp4_mega_moe.py` +- DONE: single-rank **torch MoE reference** at `tests/test_fp4_mega_moe.py` (FP4 quant → grouped L1×W1 → clamp+SwiGLU×weight → per-32 UE8M0 FP4 requant → grouped L2×W2 → top-k combine). Single-rank harness: `/tmp/mega_compile.py`. - DONE: confirmed the **FP8 mega kernel runs in the same single-rank harness** (so the @@ -416,7 +416,7 @@ exists; needs a manual transpose/pack) whose layout matches the L2 `mxf4` TMA re (`swizzle = BLOCK_K/2`), and verify the L2 arrival mask reaches `0xFF`. Then validate against the torch reference and add a `mxfp4xmxfp4` case to `tests/test_mega_moe.py`. Debug loop: `DG_JIT_WITH_LINEINFO=1 cuda-gdb --batch -ex run -ex 'info cuda threads' --args python -u -tests/test_mxfp4_mega_moe.py` (interrupt with `timeout -s INT`); inspect the workspace +tests/test_fp4_mega_moe.py` (interrupt with `timeout -s INT`); inspect the workspace `l2_arrival_mask` value for a stuck pool block. ### (superseded) original step 5 diff --git a/tests/test_fp4_gemm.py b/tests/test_fp4_gemm.py new file mode 100644 index 0000000000..1c3b8cd971 --- /dev/null +++ b/tests/test_fp4_gemm.py @@ -0,0 +1,112 @@ +import random +import torch + +import deep_gemm +from deep_gemm.testing import calc_diff +from deep_gemm.utils.layout import ( + get_mn_major_tma_aligned_packed_ue8m0_tensor, + get_tma_aligned_size, +) +from deep_gemm.utils.math import ( + align, + per_token_cast_to_fp4, + cast_back_from_fp4, + per_token_cast_to_nvfp4, + cast_back_from_nvfp4, + nvfp4_global_scale, +) + +# Packed FP4 GEMM is a de-risk kernel hardcoded to a 2-CTA (cluster_n = 2) config, +# so it requires N divisible by 256, M divisible by 128, and K divisible by 128. +GEMM_SHAPES = ( + (256, 256, 256), + (256, 256, 512), + (512, 256, 1024), + (128, 512, 256), + (1024, 768, 512), +) + +DIFF_TOL = 0.05 + +MXFP4_GRAN_K = 32 +NVFP4_GRAN_K = 16 + + +def _mn_major_packed_e4m3(sf_bytes: torch.Tensor) -> torch.Tensor: + """[mn, k_sf] uint8 (E4M3 bytes) -> MN-major TMA-aligned int32 (4 SF bytes/int32), + matching ``get_mn_major_tma_aligned_packed_ue8m0_tensor`` but for raw E4M3 bytes.""" + mn, k = sf_bytes.shape + aligned_mn = get_tma_aligned_size(mn, 4) + aligned_k = align(k, 4) + padded = torch.zeros((aligned_mn, aligned_k), dtype=torch.uint8, device=sf_bytes.device) + padded[:mn, :k] = sf_bytes + padded = padded.reshape(-1).view(torch.int32).view(aligned_mn, aligned_k // 4) + out = torch.empty_strided((aligned_mn, aligned_k // 4), (1, aligned_mn), + dtype=torch.int32, device=sf_bytes.device) + return out.copy_(padded)[:mn] + + +def _prepare_mxfp4(x: torch.Tensor): + # Packed E2M1 (int8, K/2) + float UE8M0 SF (gran-K 32); no global scale. + packed, sf_f = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=MXFP4_GRAN_K) + deq = cast_back_from_fp4(packed, sf_f, gran_k=MXFP4_GRAN_K) + sf_packed = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf_f) + return packed, sf_packed, deq, None + + +def _prepare_nvfp4(x: torch.Tensor): + # Packed E2M1 (int8, K/2) + E4M3 SF bytes (gran-K 16) + per-tensor global scale. + gs = nvfp4_global_scale(x) + packed, sf_bytes = per_token_cast_to_nvfp4(x, gs, gran_k=NVFP4_GRAN_K) + deq = cast_back_from_nvfp4(packed, sf_bytes, gs, gran_k=NVFP4_GRAN_K) + sf_packed = _mn_major_packed_e4m3(sf_bytes) + return packed, sf_packed, deq, gs + + +def _run_mxfp4_gemm(a, b, d, gs_a, gs_b) -> None: + deep_gemm.mxfp4_gemm_nt(a, b, d) + + +def _run_nvfp4_gemm(a, b, d, gs_a, gs_b) -> None: + deep_gemm.nvfp4_gemm_nt(a, b, d, a_global_scale=gs_a, b_global_scale=gs_b) + + +# Per-format strategy: (prepare, kernel runner). ``prepare`` returns +# (packed, sf_packed, deq, global_scale); ``run`` takes the two (packed, sf) pairs, +# the output, and the two global scales (None for MXFP4). +FP4_GEMM = { + 'mxfp4': (_prepare_mxfp4, _run_mxfp4_gemm), + 'nvfp4': (_prepare_nvfp4, _run_nvfp4_gemm), +} + + +def test_fp4_gemm(fmt: str = 'mxfp4') -> None: + assert fmt in FP4_GEMM, f'unknown FP4 format {fmt!r}; expected one of {list(FP4_GEMM)}' + prepare, run_gemm = FP4_GEMM[fmt] + print(f'Testing packed {fmt.upper()} x {fmt.upper()} GEMM:') + for m, n, k in GEMM_SHAPES: + a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16) + b = torch.randn((n, k), device='cuda', dtype=torch.bfloat16) + + a_packed, sfa, a_deq, gs_a = prepare(a) + b_packed, sfb, b_deq, gs_b = prepare(b) + + ref_d = (a_deq.float() @ b_deq.float().T).to(torch.bfloat16) + + d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16) + run_gemm((a_packed, sfa), (b_packed, sfb), d, gs_a, gs_b) + + diff = calc_diff(d, ref_d) + status = 'OK' if diff < DIFF_TOL else 'FAIL' + print(f' > [{fmt}] m={m:5}, n={n:5}, k={k:5}: diff={diff:.5f} [{status}]') + assert diff < DIFF_TOL, f'{fmt=}, {m=}, {n=}, {k=}, diff={diff:.5f}' + print(f'All {fmt.upper()} GEMM cases passed.\n') + + +if __name__ == '__main__': + torch.manual_seed(0) + random.seed(0) + print('Library path:') + print(f' > {deep_gemm.__path__}\n') + for fmt in FP4_GEMM: + test_fp4_gemm(fmt) diff --git a/tests/test_fp4_mega_moe.py b/tests/test_fp4_mega_moe.py new file mode 100644 index 0000000000..07de92e8a8 --- /dev/null +++ b/tests/test_fp4_mega_moe.py @@ -0,0 +1,433 @@ +import argparse +import os +from types import SimpleNamespace + +import torch +import torch.distributed as dist + +import deep_gemm +from deep_gemm.testing import calc_diff +from deep_gemm.utils.dist import init_dist, dist_print +from deep_gemm.utils.layout import get_tma_aligned_size +from deep_gemm.utils.math import ( + align, + per_token_cast_to_fp4, + cast_back_from_fp4, + per_token_cast_to_nvfp4, + cast_back_from_nvfp4, + nvfp4_global_scale, +) + +DIFF_TOL = 0.05 + +MXFP4_GRAN_K = 32 +NVFP4_GRAN_K = 16 +CLAMP = 10.0 + +# Kernel constraints (verified in csrc/jit_kernels/impls/sm100_mxfp4_mxfp4_mega_moe.hpp +# and deep_gemm/include/deep_gemm/layout/mega_moe.cuh): +# - block_n = 128 -> hidden % 128 == 0 and inter % 128 == 0 +# - num_max_tokens_per_rank is aligned up to kLCMCandidateBlockM = 384 internally; +# num_tokens may be any value <= num_max_tokens_per_rank (partial M-blocks ok) +# - num_experts % num_ranks == 0 (expert-parallel sharding) +# - topk_idx entries may be -1 (masked / dropped selections) +assert_msg = 'hidden/inter must be multiples of 128 (block_n=128)' + + +def _swiglu(l1: torch.Tensor, inter: int, weight: float, clamp: float) -> torch.Tensor: + gate = l1[:inter].clamp(max=clamp) + up = l1[inter:].clamp(min=-clamp, max=clamp) + return (gate * torch.sigmoid(gate)) * up * weight + + +def _mn_major_packed_e4m3_3d(sf_bytes: torch.Tensor) -> torch.Tensor: + """[g, mn, k_sf] uint8 (E4M3 bytes) -> MN-major TMA-aligned int32 (4 SF bytes/int32).""" + g, mn, k = sf_bytes.shape + aligned_mn = get_tma_aligned_size(mn, 4) + aligned_k = align(k, 4) + padded = torch.zeros((g, aligned_mn, aligned_k), dtype=torch.uint8, device=sf_bytes.device) + padded[:, :mn, :k] = sf_bytes + padded = padded.reshape(-1).view(torch.int32).view(g, aligned_mn, aligned_k // 4) + out = torch.empty_strided((g, aligned_mn, aligned_k // 4), + (aligned_mn * (aligned_k // 4), 1, aligned_mn), + dtype=torch.int32, device=sf_bytes.device) + return out.copy_(padded)[:, :mn] + + +# --------------------------------------------------------------------------- # +# MXFP4 strategy +# --------------------------------------------------------------------------- # + +def _fp4_roundtrip_mxfp4(x: torch.Tensor, gran_k: int = MXFP4_GRAN_K) -> torch.Tensor: + packed, sf = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=gran_k) + return cast_back_from_fp4(packed, sf, gran_k=gran_k).to(x.dtype) + + +def _cast_w_mxfp4(w: torch.Tensor): + """w: [g, n, k] (a local expert shard) -> (packed_e2m1, UE8M0 SF in kernel layout).""" + g, n, k = w.shape + wp = torch.empty((g, n, k // 2), device='cuda', dtype=torch.int8) + wsf = torch.empty((g, n, k // MXFP4_GRAN_K), device='cuda', dtype=torch.float) + for i in range(g): + wp[i], wsf[i] = per_token_cast_to_fp4(w[i], use_ue8m0=True, gran_k=MXFP4_GRAN_K) + return wp, deep_gemm.transform_sf_into_required_layout(wsf, n, k, (1, MXFP4_GRAN_K), g) + + +def _run_mxfp4(s) -> torch.Tensor: + lo, le = s.local_offset, s.local_experts + xp, xsf = per_token_cast_to_fp4(s.x, use_ue8m0=True, gran_k=MXFP4_GRAN_K, use_packed_ue8m0=True) + l1 = _cast_w_mxfp4(s.l1w[lo:lo + le]) + l2 = _cast_w_mxfp4(s.l2w[lo:lo + le]) + tl1, tl2 = deep_gemm.transform_weights_for_mega_moe(l1, l2) + + s.buf.x[:s.num_tokens].copy_(xp) + s.buf.x_sf[:s.num_tokens].copy_(xsf) + s.buf.topk_idx[:s.num_tokens].copy_(s.topk_idx) + s.buf.topk_weights[:s.num_tokens].copy_(s.topk_weights) + + y = torch.empty((s.num_tokens, s.hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.mxfp4_mxfp4_mega_moe(y=y, l1_weights=tl1, l2_weights=tl2, sym_buffer=s.buf, + activation_clamp=s.clamp, fast_math=True) + torch.cuda.synchronize() + return y + + +def _reference_mxfp4(s) -> torch.Tensor: + """Full MoE reference over ALL experts (the kernel's combine produces the full + per-rank output, so the reference ignores EP sharding).""" + x_deq = _fp4_roundtrip_mxfp4(s.x) + w1_deq = torch.stack([_fp4_roundtrip_mxfp4(s.l1w[e]) for e in range(s.num_experts)]) + w2_deq = torch.stack([_fp4_roundtrip_mxfp4(s.l2w[e]) for e in range(s.num_experts)]) + + y = torch.zeros((s.num_tokens, s.hidden), dtype=torch.float, device=s.x.device) + for t in range(s.num_tokens): + for k in range(s.topk_idx.shape[1]): + e = int(s.topk_idx[t, k].item()) + if e < 0: + continue + l1 = x_deq[t].float() @ w1_deq[e].float().T + act = _swiglu(l1, s.inter, float(s.topk_weights[t, k].item()), s.clamp) + act_deq = _fp4_roundtrip_mxfp4(act.to(torch.bfloat16).unsqueeze(0)).squeeze(0) + y[t] += act_deq.float() @ w2_deq[e].float().T + return y.to(torch.bfloat16) + + +# --------------------------------------------------------------------------- # +# NVFP4 strategy +# --------------------------------------------------------------------------- # + +def _rt_nvfp4(x: torch.Tensor, gs) -> torch.Tensor: + p, sb = per_token_cast_to_nvfp4(x, gs, gran_k=NVFP4_GRAN_K) + return cast_back_from_nvfp4(p, sb, gs, gran_k=NVFP4_GRAN_K).to(x.dtype) + + +def _cast_l1_w_nvfp4(l1w: torch.Tensor): + """l1w: [g, inter*2, hidden] (a local expert shard). gate = rows [:inter], up = [inter:].""" + g, n2, k = l1w.shape + inter = n2 // 2 + wp = torch.empty((g, n2, k // 2), device='cuda', dtype=torch.int8) + wsf = torch.empty((g, n2, k // NVFP4_GRAN_K), device='cuda', dtype=torch.uint8) + gate_gs = torch.empty(g, device='cuda', dtype=torch.float32) + up_gs = torch.empty(g, device='cuda', dtype=torch.float32) + for e in range(g): + ggs = nvfp4_global_scale(l1w[e][:inter]) + ugs = nvfp4_global_scale(l1w[e][inter:]) + gate_gs[e], up_gs[e] = ggs, ugs + wp[e][:inter], wsf[e][:inter] = per_token_cast_to_nvfp4(l1w[e][:inter], ggs, gran_k=NVFP4_GRAN_K) + wp[e][inter:], wsf[e][inter:] = per_token_cast_to_nvfp4(l1w[e][inter:], ugs, gran_k=NVFP4_GRAN_K) + return (wp, _mn_major_packed_e4m3_3d(wsf)), gate_gs, up_gs + + +def _cast_l2_w_nvfp4(l2w: torch.Tensor): + g, n, k = l2w.shape + wp = torch.empty((g, n, k // 2), device='cuda', dtype=torch.int8) + wsf = torch.empty((g, n, k // NVFP4_GRAN_K), device='cuda', dtype=torch.uint8) + down_gs = torch.empty(g, device='cuda', dtype=torch.float32) + for e in range(g): + dgs = nvfp4_global_scale(l2w[e]) + down_gs[e] = dgs + wp[e], wsf[e] = per_token_cast_to_nvfp4(l2w[e], dgs, gran_k=NVFP4_GRAN_K) + return (wp, _mn_major_packed_e4m3_3d(wsf)), down_gs + + +def _estimate_l2act_gs(x, l1w, topk_idx, topk_weights, inter, gs_x, gate_gs, up_gs, + clamp, num_experts) -> torch.Tensor: + # Per-expert L2-input (intermediate) global scale (convention: amax/(6*448)). + x_deq = _rt_nvfp4(x, gs_x) + w1_deq = [torch.cat([_rt_nvfp4(l1w[e][:inter], gate_gs[e].item()), + _rt_nvfp4(l1w[e][inter:], up_gs[e].item())]) + for e in range(num_experts)] + amax = torch.full((num_experts,), 1e-6, device=x.device) + for t in range(x.shape[0]): + for k in range(topk_idx.shape[1]): + e = int(topk_idx[t, k].item()) + if e < 0: + continue + l1 = x_deq[t].float() @ w1_deq[e].float().T + act = _swiglu(l1, inter, float(topk_weights[t, k].item()), clamp) + amax[e] = torch.maximum(amax[e], act.abs().amax()) + return amax / (6.0 * 448.0) + + +def _run_nvfp4(s) -> torch.Tensor: + s.gs_x = nvfp4_global_scale(s.x) + lo, le = s.local_offset, s.local_experts + + # Full per-expert global scales (reference uses ALL experts; kernel uses the local slice). + # nvfp4_global_scale returns a python float, so build the tensors explicitly. + s.gate_gs = torch.tensor([nvfp4_global_scale(s.l1w[e][:s.inter]) for e in range(s.num_experts)], + dtype=torch.float32, device='cuda') + s.up_gs = torch.tensor([nvfp4_global_scale(s.l1w[e][s.inter:]) for e in range(s.num_experts)], + dtype=torch.float32, device='cuda') + s.down_gs = torch.tensor([nvfp4_global_scale(s.l2w[e]) for e in range(s.num_experts)], + dtype=torch.float32, device='cuda') + s.l2act_gs = _estimate_l2act_gs(s.x, s.l1w, s.topk_idx, s.topk_weights, s.inter, + s.gs_x, s.gate_gs, s.up_gs, s.clamp, s.num_experts) + + # Local-shard weight cast + per-local-expert alphas (TRT-LLM convention). + l1_local, _gate_gs_local, _up_gs_local = _cast_l1_w_nvfp4(s.l1w[lo:lo + le]) + l2_local, _down_gs_local = _cast_l2_w_nvfp4(s.l2w[lo:lo + le]) + gate_gs_local = s.gate_gs[lo:lo + le].contiguous() + up_gs_local = s.up_gs[lo:lo + le].contiguous() + down_gs_local = s.down_gs[lo:lo + le].contiguous() + l2act_local = s.l2act_gs[lo:lo + le].contiguous() + gate_alpha = (s.gs_x * gate_gs_local).contiguous() + up_alpha = (s.gs_x * up_gs_local).contiguous() + down_alpha = (l2act_local * down_gs_local).contiguous() + l2_input_global_scale = (1.0 / l2act_local).contiguous() + + tl1, tl2 = deep_gemm.transform_weights_for_mega_moe(l1_local, l2_local) + xp, xsf = per_token_cast_to_nvfp4(s.x, s.gs_x, gran_k=NVFP4_GRAN_K) + + s.buf.x[:s.num_tokens].copy_(xp) + s.buf.x_sf[:s.num_tokens].copy_(xsf.contiguous().view(torch.int32)) + s.buf.topk_idx[:s.num_tokens].copy_(s.topk_idx) + s.buf.topk_weights[:s.num_tokens].copy_(s.topk_weights) + + y = torch.empty((s.num_tokens, s.hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.nvfp4_nvfp4_mega_moe(y=y, l1_weights=tl1, l2_weights=tl2, sym_buffer=s.buf, + gate_alpha=gate_alpha, up_alpha=up_alpha, + l2_input_global_scale=l2_input_global_scale, down_alpha=down_alpha, + activation_clamp=s.clamp, fast_math=True) + torch.cuda.synchronize() + return y + + +def _reference_nvfp4(s) -> torch.Tensor: + x_deq = _rt_nvfp4(s.x, s.gs_x) + w1_deq = [torch.cat([_rt_nvfp4(s.l1w[e][:s.inter], s.gate_gs[e].item()), + _rt_nvfp4(s.l1w[e][s.inter:], s.up_gs[e].item())]) + for e in range(s.num_experts)] + w2_deq = [_rt_nvfp4(s.l2w[e], s.down_gs[e].item()) for e in range(s.num_experts)] + y = torch.zeros((s.num_tokens, s.hidden), dtype=torch.float, device=s.x.device) + for t in range(s.num_tokens): + for k in range(s.topk_idx.shape[1]): + e = int(s.topk_idx[t, k].item()) + if e < 0: + continue + l1 = x_deq[t].float() @ w1_deq[e].float().T + act = _swiglu(l1, s.inter, float(s.topk_weights[t, k].item()), s.clamp).to(torch.bfloat16) + act_deq = _rt_nvfp4(act.unsqueeze(0), s.l2act_gs[e].item()).squeeze(0) + y[t] += act_deq.float() @ w2_deq[e].float().T + return y.to(torch.bfloat16) + + +# --------------------------------------------------------------------------- # +# Per-format strategy table +# --------------------------------------------------------------------------- # + +FP4_MOE = { + 'mxfp4': SimpleNamespace(mma_type='mxfp4xmxfp4', run=_run_mxfp4, reference=_reference_mxfp4), + 'nvfp4': SimpleNamespace(mma_type='nvfp4xnvfp4', run=_run_nvfp4, reference=_reference_nvfp4), +} + +# Baseline single-shape topology (matches the original pre-consolidation test). +BASELINE_SHAPE = (128, 8, 2, 512, 512, 0.0) + +# 1-rank coverage: (num_tokens, num_experts, num_topk, hidden, inter, masked_ratio). +# Covers small / odd-token / single-token / large-asymmetric / masked routing. +SHAPES_1RANK = [ + (32, 8, 2, 512, 512, 0.0), # small + (128, 8, 2, 512, 512, 0.0), # baseline + (96, 8, 2, 512, 512, 0.0), # odd token count (partial M-block) + (1, 8, 2, 512, 512, 0.0), # single token (smallest block_m path) + (128, 8, 2, 2048, 2048, 0.0), # larger hidden/inter + (512, 8, 2, 2048, 2048, 0.0), # more tokens + (1024, 32, 4, 4096, 1536, 0.0), # large + asymmetric + more experts + (128, 8, 2, 512, 512, 0.3), # masked routing +] + +# Multi-rank EP coverage: (world, total_experts, num_topk, num_tokens, hidden, inter, masked_ratio). +# Tokens/weights/routing are REPLICATED across ranks (seed=0) so the NVFP4 per-expert +# l2act_gs reference is computable without a cross-rank all-gather; the cross-rank +# dispatch+combine path is still exercised (tokens leave to the expert-owning rank). +SHAPES_MULTIRANK = [ + (2, 8, 2, 128, 512, 512, 0.0), + (4, 8, 2, 128, 512, 512, 0.0), + (8, 8, 2, 128, 512, 512, 0.0), + (2, 32, 4, 512, 2048, 2048, 0.0), + (4, 32, 4, 512, 2048, 2048, 0.0), + (8, 32, 4, 1024, 4096, 1536, 0.0), + (2, 8, 2, 128, 512, 512, 0.3), # masked EP +] + + +def _build_s(spec, shape, rank, world, group) -> SimpleNamespace: + num_tokens, num_experts, num_topk, hidden, inter, masked_ratio = shape + assert num_experts % world == 0, f'num_experts={num_experts} not divisible by world={world}' + assert hidden % 128 == 0 and inter % 128 == 0, assert_msg + local_experts = num_experts // world + local_offset = rank * local_experts + + # num_max_tokens is aligned up to 384 internally; num_tokens may be any value <= it. + num_max_tokens = max(num_tokens, 1) + buf = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, + hidden, inter, mma_type=spec.mma_type) + + # Replicated inputs across ranks (same seed) -> every rank computes the same full + # reference; each rank passes only its local expert shard to the kernel. + torch.manual_seed(0) + x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + l1w = torch.randn((num_experts, inter * 2, hidden), dtype=torch.bfloat16, device='cuda') / (hidden ** 0.5) + l2w = torch.randn((num_experts, hidden, inter), dtype=torch.bfloat16, device='cuda') / (inter ** 0.5) + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') + topk_weights, topk_idx = torch.topk(scores.softmax(-1), num_topk, dim=-1) + if masked_ratio > 0: + rand_mask = torch.rand_like(topk_idx, dtype=torch.float) + topk_idx = topk_idx.masked_fill(rand_mask < masked_ratio, -1) + topk_weights = topk_weights.masked_fill(topk_idx < 0, 0) + + return SimpleNamespace( + num_tokens=num_tokens, num_experts=num_experts, num_topk=num_topk, + hidden=hidden, inter=inter, clamp=CLAMP, + local_experts=local_experts, local_offset=local_offset, + x=x, l1w=l1w, l2w=l2w, topk_idx=topk_idx, topk_weights=topk_weights, + buf=buf, + ) + + +def _run_one(spec, shape, rank, world, group): + """Build the shape, run the kernel + reference, return (y, ref, diff).""" + s = _build_s(spec, shape, rank, world, group) + try: + y = spec.run(s) + ref = spec.reference(s) + return y, ref, calc_diff(y, ref), s + except Exception: + s.buf.destroy() + raise + + +def _shape_tag(shape, world=None): + if world is None: + tok, exp, topk, h, i, m = shape + return f'1rank tok={tok} E={exp} topk={topk} h={h} inter={i} mask={m}' + _w, exp, topk, tok, h, i, m = shape + return f'EP{world} tok={tok} E={exp} topk={topk} h={h} inter={i} mask={m}' + + +def _check_and_print(label, fmt, diff, y, ref): + """Print (only on local_rank 0) and assert on ALL ranks. ``dist_print`` barriers + on every rank, so all ranks must call this — never gate it behind a rank check.""" + ok = diff < DIFF_TOL + dist_print(f'[{label} {fmt}] diff={diff:.5f} ' + f'(y~{y.float().abs().mean():.3f}, ref~{ref.float().abs().mean():.3f}) ' + f'{"OK" if ok else "FAIL"}', once_in_node=True) + assert ok, f'{fmt=}, {label}, diff={diff:.5f}' + + +def _ensure_master(port: int): + os.environ.setdefault('MASTER_ADDR', '127.0.0.1') + os.environ['MASTER_PORT'] = str(port) + + +# --------------------------------------------------------------------------- # +# Entry points +# --------------------------------------------------------------------------- # + +def test_fp4_mega_moe(fmt: str = 'mxfp4') -> None: + """Single-shape 1-rank smoke test (baseline topology). Kept for backward compat.""" + assert fmt in FP4_MOE, f'unknown FP4 format {fmt!r}; expected one of {list(FP4_MOE)}' + _ensure_master(13799) + rank, world, group = init_dist(0, 1) + spec = FP4_MOE[fmt] + try: + y, ref, diff, s = _run_one(spec, BASELINE_SHAPE, rank, world, group) + try: + print(f'[{fmt}] diff = {diff:.5f} ' + f'(y~{y.float().abs().mean():.3f}, ref~{ref.float().abs().mean():.3f})') + assert diff < DIFF_TOL, f'{fmt=}, {diff=}' + print(f'{fmt.upper()} mega MoE passed.') + finally: + s.buf.destroy() + finally: + dist.destroy_process_group() + + +def test_fp4_mega_moe_1rank(fmt: str = 'mxfp4') -> None: + """Run the full 1-rank shape matrix for one format.""" + assert fmt in FP4_MOE, f'unknown FP4 format {fmt!r}; expected one of {list(FP4_MOE)}' + _ensure_master(13799) + rank, world, group = init_dist(0, 1) + spec = FP4_MOE[fmt] + print(f'=== 1-rank {fmt.upper()} x {fmt.upper()} mega-MoE ({len(SHAPES_1RANK)} shapes) ===') + try: + for shape in SHAPES_1RANK: + y, ref, diff, s = _run_one(spec, shape, rank, world, group) + try: + _check_and_print(_shape_tag(shape), fmt, diff, y, ref) + finally: + s.buf.destroy() + print(f'All 1-rank {fmt.upper()} cases passed.\n') + finally: + dist.destroy_process_group() + + +def _multirank_worker(local_rank, num_local_ranks, cases): + """Spawned per-rank worker: init dist once, loop (fmt, shape) cases, assert per rank.""" + rank, world, group = init_dist(local_rank, num_local_ranks) + try: + for fmt, shape in cases: + _w, exp, topk, tok, h, i, m = shape + inner = (tok, exp, topk, h, i, m) # strip the leading world field + spec = FP4_MOE[fmt] + y, ref, diff, s = _run_one(spec, inner, rank, world, group) + try: + # All ranks compute the same replicated reference; rank 0 reports. + _check_and_print(_shape_tag(shape, world), fmt, diff, y, ref) + finally: + s.buf.destroy() + dist_print(f'All EP{world} cases passed.', once_in_node=True) + finally: + dist.destroy_process_group() + + +def test_fp4_mega_moe_multirank(world: int = 2) -> None: + """Run the multi-rank EP shape matrix for a given world size (both formats).""" + cases = [(fmt, shape) for shape in SHAPES_MULTIRANK if shape[0] == world for fmt in FP4_MOE] + if not cases: + print(f'No multi-rank shapes for world={world}; skipping.') + return + _ensure_master(13800 + world) + print(f'=== multi-rank EP (world={world}, {len(cases)} cases) ===', flush=True) + torch.multiprocessing.spawn(_multirank_worker, args=(world, cases), nprocs=world) + print() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Packed FP4 mega-MoE correctness (1-rank + multi-rank EP)') + parser.add_argument('--scope', choices=['all', '1rank', 'multirank', 'baseline'], + default='all', help='Test scope (default: all)') + parser.add_argument('--worlds', type=str, default='2,4,8', + help='Comma-separated world sizes for the multi-rank scope') + args = parser.parse_args() + + if args.scope in ('all', 'baseline'): + for fmt in FP4_MOE: + test_fp4_mega_moe(fmt) + if args.scope in ('all', '1rank'): + for fmt in FP4_MOE: + test_fp4_mega_moe_1rank(fmt) + if args.scope in ('all', 'multirank'): + for w in [int(x) for x in args.worlds.split(',') if x.strip()]: + test_fp4_mega_moe_multirank(w) diff --git a/tests/test_mxfp4_gemm.py b/tests/test_mxfp4_gemm.py deleted file mode 100644 index 2c37185aa6..0000000000 --- a/tests/test_mxfp4_gemm.py +++ /dev/null @@ -1,48 +0,0 @@ -import random -import torch - -import deep_gemm -from deep_gemm.testing import calc_diff -from deep_gemm.utils.layout import get_mn_major_tma_aligned_packed_ue8m0_tensor -from deep_gemm.utils.math import per_token_cast_to_fp4, cast_back_from_fp4 - - -def _prepare(x: torch.Tensor, gran_k: int = 32): - # Packed E2M1 (int8, K/2) + float UE8M0 SF (gran-K 32) - packed, sf_f = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=gran_k) - # Dequantized reference using the *same* packed values + SF - deq = cast_back_from_fp4(packed, sf_f, gran_k=gran_k) - # SF in the MN-major TMA-aligned packed UE8M0 layout the kernel expects - sf_packed = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf_f) - return packed, sf_packed, deq - - -def test_mxfp4_gemm() -> None: - print('Testing packed MXFP4 x MXFP4 GEMM:') - # NOTES: this de-risk kernel is hardcoded to a 2-CTA (cluster_n = 2) config, so it - # requires N divisible by 256, M divisible by 128, and K divisible by 128. - for m, n, k in ((256, 256, 256), (256, 256, 512), (512, 256, 1024), (128, 512, 256), (1024, 768, 512)): - a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16) - b = torch.randn((n, k), device='cuda', dtype=torch.bfloat16) - - a_packed, sfa, a_deq = _prepare(a) - b_packed, sfb, b_deq = _prepare(b) - - ref_d = (a_deq.float() @ b_deq.float().T).to(torch.bfloat16) - - d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16) - deep_gemm.mxfp4_gemm_nt((a_packed, sfa), (b_packed, sfb), d) - - diff = calc_diff(d, ref_d) - status = 'OK' if diff < 0.05 else 'FAIL' - print(f' > m={m:5}, n={n:5}, k={k:5}: diff={diff:.5f} [{status}]') - assert diff < 0.05, f'{m=}, {n=}, {k=}, diff={diff:.5f}' - print('All MXFP4 GEMM cases passed.\n') - - -if __name__ == '__main__': - torch.manual_seed(0) - random.seed(0) - print('Library path:') - print(f' > {deep_gemm.__path__}\n') - test_mxfp4_gemm() diff --git a/tests/test_mxfp4_mega_moe.py b/tests/test_mxfp4_mega_moe.py deleted file mode 100644 index b5028e8442..0000000000 --- a/tests/test_mxfp4_mega_moe.py +++ /dev/null @@ -1,96 +0,0 @@ -import os -import torch -import torch.distributed as dist - -import deep_gemm -from deep_gemm.testing import calc_diff -from deep_gemm.utils.math import per_token_cast_to_fp4, cast_back_from_fp4 - - -def _fp4_roundtrip(x: torch.Tensor, gran_k: int = 32) -> torch.Tensor: - """Quantize bf16 -> packed MXFP4 (UE8M0, gran-32) and dequantize back to the - exact values the kernel operates on.""" - packed, sf = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=gran_k) - return cast_back_from_fp4(packed, sf, gran_k=gran_k).to(x.dtype) - - -def reference_mxfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, - activation_clamp: float = 10.0): - """Single-rank MXFP4 MoE reference matching the fused kernel's math: - FP4(x) @ FP4(W1).T -> SwiGLU*weight -> FP4 requant -> @ FP4(W2).T -> combine (sum over top-k).""" - num_tokens = x.shape[0] - x_deq = _fp4_roundtrip(x) - w1_deq = torch.stack([_fp4_roundtrip(l1w[e]) for e in range(l1w.shape[0])]) - w2_deq = torch.stack([_fp4_roundtrip(l2w[e]) for e in range(l2w.shape[0])]) - - y = torch.zeros((num_tokens, hidden), dtype=torch.float, device=x.device) - for t in range(num_tokens): - for k in range(topk_idx.shape[1]): - e = int(topk_idx[t, k].item()) - if e < 0: - continue - l1 = x_deq[t].float() @ w1_deq[e].float().T # [inter*2] - gate, up = l1[:inter], l1[inter:] - gate = gate.clamp(max=activation_clamp) - up = up.clamp(min=-activation_clamp, max=activation_clamp) - act = (gate * torch.sigmoid(gate)) * up * float(topk_weights[t, k].item()) - act_deq = _fp4_roundtrip(act.to(torch.bfloat16).unsqueeze(0)).squeeze(0) - l2 = act_deq.float() @ w2_deq[e].float().T # [hidden] - y[t] += l2 - return y.to(torch.bfloat16) - - -def test_mxfp4_mega_moe(): - os.environ.setdefault('MASTER_ADDR', '127.0.0.1') - os.environ.setdefault('MASTER_PORT', '12399') - dist.init_process_group('nccl', rank=0, world_size=1) - group = dist.group.WORLD - - num_max_tokens, num_tokens = 128, 128 - hidden, inter = 512, 512 - num_experts, num_topk = 8, 2 - ne_per_rank = num_experts - - buf = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, - hidden, inter, mma_type='mxfp4xmxfp4') - - def cast_w(w): - g, n, k = w.shape - wp = torch.empty((g, n, k // 2), device='cuda', dtype=torch.int8) - wsf = torch.empty((g, n, k // 32), device='cuda', dtype=torch.float) - for i in range(g): - wp[i], wsf[i] = per_token_cast_to_fp4(w[i], use_ue8m0=True, gran_k=32) - return wp, deep_gemm.transform_sf_into_required_layout(wsf, n, k, (1, 32), g) - - torch.manual_seed(0) - x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') - l1w = torch.randn((ne_per_rank, inter * 2, hidden), dtype=torch.bfloat16, device='cuda') / (hidden ** 0.5) - l2w = torch.randn((ne_per_rank, hidden, inter), dtype=torch.bfloat16, device='cuda') / (inter ** 0.5) - scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') - topk_weights, topk_idx = torch.topk(scores.softmax(-1), num_topk, dim=-1) - - xp, xsf = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) - l1 = cast_w(l1w) - l2 = cast_w(l2w) - tl1, tl2 = deep_gemm.transform_weights_for_mega_moe(l1, l2) - - buf.x[:num_tokens].copy_(xp) - buf.x_sf[:num_tokens].copy_(xsf) - buf.topk_idx[:num_tokens].copy_(topk_idx) - buf.topk_weights[:num_tokens].copy_(topk_weights) - - y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') - deep_gemm.mxfp4_mxfp4_mega_moe(y=y, l1_weights=tl1, l2_weights=tl2, sym_buffer=buf, - activation_clamp=10.0, fast_math=True) - torch.cuda.synchronize() - - ref = reference_mxfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, 10.0) - diff = calc_diff(y, ref) - print(f'diff = {diff:.5f} (y~{y.float().abs().mean():.3f}, ref~{ref.float().abs().mean():.3f})') - assert diff < 0.05, f'{diff=}' - print('MXFP4 mega MoE passed.') - dist.destroy_process_group() - - -if __name__ == '__main__': - test_mxfp4_mega_moe() diff --git a/tests/test_nvfp4_gemm.py b/tests/test_nvfp4_gemm.py deleted file mode 100644 index 20c0766a55..0000000000 --- a/tests/test_nvfp4_gemm.py +++ /dev/null @@ -1,64 +0,0 @@ -import random -import torch - -import deep_gemm -from deep_gemm.testing import calc_diff -from deep_gemm.utils.math import ( - align, per_token_cast_to_nvfp4, cast_back_from_nvfp4, nvfp4_global_scale, -) -from deep_gemm.utils.layout import get_tma_aligned_size - -GRAN_K = 16 - - -def _mn_major_packed_e4m3(sf_bytes: torch.Tensor) -> torch.Tensor: - """[mn, k_sf] uint8 (E4M3 bytes) -> MN-major TMA-aligned int32 (4 SF bytes/int32), - matching `get_mn_major_tma_aligned_packed_ue8m0_tensor` but for raw E4M3 bytes.""" - mn, k = sf_bytes.shape - aligned_mn = get_tma_aligned_size(mn, 4) - aligned_k = align(k, 4) - padded = torch.zeros((aligned_mn, aligned_k), dtype=torch.uint8, device=sf_bytes.device) - padded[:mn, :k] = sf_bytes - padded = padded.reshape(-1).view(torch.int32).view(aligned_mn, aligned_k // 4) - out = torch.empty_strided((aligned_mn, aligned_k // 4), (1, aligned_mn), - dtype=torch.int32, device=sf_bytes.device) - return out.copy_(padded)[:mn] - - -def _prepare(x: torch.Tensor, gs: float): - packed, sf_bytes = per_token_cast_to_nvfp4(x, gs, gran_k=GRAN_K) - deq = cast_back_from_nvfp4(packed, sf_bytes, gs, gran_k=GRAN_K) - sf_packed = _mn_major_packed_e4m3(sf_bytes) - return packed, sf_packed, deq - - -def test_nvfp4_gemm() -> None: - print('Testing packed NVFP4 x NVFP4 GEMM:') - for m, n, k in ((256, 256, 256), (256, 256, 512), (512, 256, 1024), (128, 512, 256), (1024, 768, 512)): - a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16) - b = torch.randn((n, k), device='cuda', dtype=torch.bfloat16) - gs_a = nvfp4_global_scale(a) - gs_b = nvfp4_global_scale(b) - - a_packed, sfa, a_deq = _prepare(a, gs_a) - b_packed, sfb, b_deq = _prepare(b, gs_b) - - ref_d = (a_deq.float() @ b_deq.float().T).to(torch.bfloat16) - - d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16) - deep_gemm.nvfp4_gemm_nt((a_packed, sfa), (b_packed, sfb), d, - a_global_scale=gs_a, b_global_scale=gs_b) - - diff = calc_diff(d, ref_d) - status = 'OK' if diff < 0.05 else 'FAIL' - print(f' > m={m:5}, n={n:5}, k={k:5}: diff={diff:.5f} [{status}]') - assert diff < 0.05, f'{m=}, {n=}, {k=}, diff={diff:.5f}' - print('All NVFP4 GEMM cases passed.\n') - - -if __name__ == '__main__': - torch.manual_seed(0) - random.seed(0) - print('Library path:') - print(f' > {deep_gemm.__path__}\n') - test_nvfp4_gemm() diff --git a/tests/test_nvfp4_mega_moe.py b/tests/test_nvfp4_mega_moe.py deleted file mode 100644 index 844e96c1f4..0000000000 --- a/tests/test_nvfp4_mega_moe.py +++ /dev/null @@ -1,159 +0,0 @@ -import os -import torch -import torch.distributed as dist - -import deep_gemm -from deep_gemm.testing import calc_diff -from deep_gemm.utils.math import ( - align, per_token_cast_to_nvfp4, cast_back_from_nvfp4, nvfp4_global_scale, -) -from deep_gemm.utils.layout import get_tma_aligned_size - -GRAN_K = 16 - - -def _rt(x, gs): - p, s = per_token_cast_to_nvfp4(x, gs, gran_k=GRAN_K) - return cast_back_from_nvfp4(p, s, gs, gran_k=GRAN_K).to(x.dtype) - - -def _mn_major_packed_e4m3(sf_bytes): - g, mn, k = sf_bytes.shape - aligned_mn = get_tma_aligned_size(mn, 4) - aligned_k = align(k, 4) - padded = torch.zeros((g, aligned_mn, aligned_k), dtype=torch.uint8, device=sf_bytes.device) - padded[:, :mn, :k] = sf_bytes - padded = padded.reshape(-1).view(torch.int32).view(g, aligned_mn, aligned_k // 4) - out = torch.empty_strided((g, aligned_mn, aligned_k // 4), - (aligned_mn * (aligned_k // 4), 1, aligned_mn), - dtype=torch.int32, device=sf_bytes.device) - return out.copy_(padded)[:, :mn] - - -def _cast_l1_w(l1w): - # l1w: [E, inter*2, hidden]. gate = rows [:inter], up = rows [inter:]; per-expert global scales. - g, n2, k = l1w.shape - inter = n2 // 2 - wp = torch.empty((g, n2, k // 2), device='cuda', dtype=torch.int8) - wsf = torch.empty((g, n2, k // GRAN_K), device='cuda', dtype=torch.uint8) - gate_gs = torch.empty(g, device='cuda', dtype=torch.float32) - up_gs = torch.empty(g, device='cuda', dtype=torch.float32) - for e in range(g): - ggs = nvfp4_global_scale(l1w[e][:inter]); ugs = nvfp4_global_scale(l1w[e][inter:]) - gate_gs[e], up_gs[e] = ggs, ugs - wp[e][:inter], wsf[e][:inter] = per_token_cast_to_nvfp4(l1w[e][:inter], ggs, gran_k=GRAN_K) - wp[e][inter:], wsf[e][inter:] = per_token_cast_to_nvfp4(l1w[e][inter:], ugs, gran_k=GRAN_K) - return (wp, _mn_major_packed_e4m3(wsf)), gate_gs, up_gs - - -def _cast_l2_w(l2w): - g, n, k = l2w.shape - wp = torch.empty((g, n, k // 2), device='cuda', dtype=torch.int8) - wsf = torch.empty((g, n, k // GRAN_K), device='cuda', dtype=torch.uint8) - down_gs = torch.empty(g, device='cuda', dtype=torch.float32) - for e in range(g): - dgs = nvfp4_global_scale(l2w[e]); down_gs[e] = dgs - wp[e], wsf[e] = per_token_cast_to_nvfp4(l2w[e], dgs, gran_k=GRAN_K) - return (wp, _mn_major_packed_e4m3(wsf)), down_gs - - -def _swiglu(l1, inter, weight, clamp): - gate = l1[:inter].clamp(max=clamp) - up = l1[inter:].clamp(min=-clamp, max=clamp) - return (gate * torch.sigmoid(gate)) * up * weight - - -def _estimate_l2act_gs(x, l1w, topk_idx, topk_weights, inter, gs_x, gate_gs, up_gs, clamp, num_experts): - # Per-expert L2-input (intermediate) global scale (my convention: amax/(6*448)). - x_deq = _rt(x, gs_x) - w1_deq = [torch.cat([_rt(l1w[e][:inter], gate_gs[e].item()), _rt(l1w[e][inter:], up_gs[e].item())]) - for e in range(num_experts)] - amax = torch.full((num_experts,), 1e-6, device=x.device) - for t in range(x.shape[0]): - for k in range(topk_idx.shape[1]): - e = int(topk_idx[t, k].item()) - l1 = x_deq[t].float() @ w1_deq[e].float().T - act = _swiglu(l1, inter, float(topk_weights[t, k].item()), clamp) - amax[e] = torch.maximum(amax[e], act.abs().amax()) - return amax / (6.0 * 448.0) - - -def reference_nvfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, - gs_x, gate_gs, up_gs, down_gs, l2act_gs, clamp): - num_tokens, num_experts = x.shape[0], l1w.shape[0] - x_deq = _rt(x, gs_x) - w1_deq = [torch.cat([_rt(l1w[e][:inter], gate_gs[e].item()), _rt(l1w[e][inter:], up_gs[e].item())]) - for e in range(num_experts)] - w2_deq = [_rt(l2w[e], down_gs[e].item()) for e in range(num_experts)] - y = torch.zeros((num_tokens, hidden), dtype=torch.float, device=x.device) - for t in range(num_tokens): - for k in range(topk_idx.shape[1]): - e = int(topk_idx[t, k].item()) - l1 = x_deq[t].float() @ w1_deq[e].float().T - act = _swiglu(l1, inter, float(topk_weights[t, k].item()), clamp).to(torch.bfloat16) - act_deq = _rt(act.unsqueeze(0), l2act_gs[e].item()).squeeze(0) - y[t] += act_deq.float() @ w2_deq[e].float().T - return y.to(torch.bfloat16) - - -def test_nvfp4_mega_moe(): - os.environ.setdefault('MASTER_ADDR', '127.0.0.1') - os.environ.setdefault('MASTER_PORT', '12400') - dist.init_process_group('nccl', rank=0, world_size=1) - group = dist.group.WORLD - - num_max_tokens, num_tokens = 128, 128 - hidden, inter = 512, 512 - num_experts, num_topk = 8, 2 - clamp = 10.0 - - buf = deep_gemm.get_symm_buffer_for_mega_moe(group, num_experts, num_max_tokens, num_topk, - hidden, inter, mma_type='nvfp4xnvfp4') - - torch.manual_seed(0) - x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') - l1w = torch.randn((num_experts, inter * 2, hidden), dtype=torch.bfloat16, device='cuda') / (hidden ** 0.5) - l2w = torch.randn((num_experts, hidden, inter), dtype=torch.bfloat16, device='cuda') / (inter ** 0.5) - scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') - topk_weights, topk_idx = torch.topk(scores.softmax(-1), num_topk, dim=-1) - - # Per-tensor input global scale (my convention amax/(6*448)); per-expert weight global scales - gs_x = nvfp4_global_scale(x) - l1, gate_gs, up_gs = _cast_l1_w(l1w) - l2, down_gs = _cast_l2_w(l2w) - l2act_gs = _estimate_l2act_gs(x, l1w, topk_idx, topk_weights, inter, gs_x, gate_gs, up_gs, clamp, num_experts) - - # Per-expert kernel params (TRT-LLM convention): - # alpha = 1/(input_gs_trt * weight_gs_trt) = my_gs_input * my_gs_weight - # l2_input_global_scale (TRT) = 1 / my_gs - gate_alpha = (gs_x * gate_gs).contiguous() - up_alpha = (gs_x * up_gs).contiguous() - down_alpha = (l2act_gs * down_gs).contiguous() - l2_input_global_scale = (1.0 / l2act_gs).contiguous() - - tl1, tl2 = deep_gemm.transform_weights_for_mega_moe(l1, l2) - xp, xsf = per_token_cast_to_nvfp4(x, gs_x, gran_k=GRAN_K) - - buf.x[:num_tokens].copy_(xp) - buf.x_sf[:num_tokens].copy_(xsf.contiguous().view(torch.int32)) - buf.topk_idx[:num_tokens].copy_(topk_idx) - buf.topk_weights[:num_tokens].copy_(topk_weights) - - y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') - deep_gemm.nvfp4_nvfp4_mega_moe(y=y, l1_weights=tl1, l2_weights=tl2, sym_buffer=buf, - gate_alpha=gate_alpha, up_alpha=up_alpha, - l2_input_global_scale=l2_input_global_scale, down_alpha=down_alpha, - activation_clamp=clamp, fast_math=True) - torch.cuda.synchronize() - - ref = reference_nvfp4_moe(x, l1w, l2w, topk_idx, topk_weights, hidden, inter, - gs_x, gate_gs, up_gs, down_gs, l2act_gs, clamp) - diff = calc_diff(y, ref) - print(f'diff = {diff:.5f} (y~{y.float().abs().mean():.3f}, ref~{ref.float().abs().mean():.3f})') - assert diff < 0.05, f'{diff=}' - print('NVFP4 mega MoE passed.') - dist.destroy_process_group() - - -if __name__ == '__main__': - test_nvfp4_mega_moe() From 5de63a64de7a583b5596b5d927f83ff9e795899d Mon Sep 17 00:00:00 2001 From: Ying Zhang Date: Mon, 6 Jul 2026 01:25:17 +0000 Subject: [PATCH 9/9] bench: merge bench_kernel_breakdown into bench_flashinfer_vs_deepgemm; add README * Fold the per-kernel diagnostic into bench_flashinfer_vs_deepgemm.py as a --breakdown mode (kernel_breakdown + show_breakdown helpers, breakdown branch in bench_one, --breakdown CLI flag). Delete the standalone bench_kernel_breakdown.py (nothing imported it). * Add benchmarks/README.md documenting the three benchmarks (FP4-format, single-device DeepGEMM-vs-FlashInfer, multi-GPU EP), prerequisites, how to run, and captured B200 results (DeepGEMM ~1.6-1.9x faster than FlashInfer cutlass single-device; ~1.7x faster fused-EP at world=2), plus a --breakdown example and compatibility notes. Co-Authored-By: Claude --- benchmarks/README.md | 166 +++++++++++++++++++++ benchmarks/bench_flashinfer_vs_deepgemm.py | 85 ++++++++++- benchmarks/bench_kernel_breakdown.py | 94 ------------ 3 files changed, 246 insertions(+), 99 deletions(-) create mode 100644 benchmarks/README.md delete mode 100644 benchmarks/bench_kernel_breakdown.py diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000000..90852eb32d --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,166 @@ +# DeepGEMM packed-FP4 benchmarks + +Performance benchmarks for the packed-FP4 (MXFP4 / NVFP4) GEMM and mega-MoE +kernels on SM100 (Blackwell). Three scripts, each measuring a different axis: + +| Script | Axis | What it compares | GPUs | +|---|---|---|---| +| `bench_packed_fp4.py` | **FP4 format** | MXFP4 vs NVFP4 (vs `fp8xfp4`) within DeepGEMM | 1 | +| `bench_flashinfer_vs_deepgemm.py` | **framework (single box)** | DeepGEMM mega-MoE vs FlashInfer NVFP4 MoE backends | 1 | +| `bench_ep_multi_gpu.py` | **scale (multi-GPU EP)** | DeepGEMM native fused EP vs FlashInfer + NCCL `all_reduce` combine | N (2/4/8) | + +`bench_flashinfer_vs_deepgemm.py` also has a `--breakdown` mode that prints a +per-kernel device-time breakdown (launches/iter + us/iter per kernel) for each +backend, so you can see how many kernels each launches and where the time goes. + +## Prerequisites + +- **SM100 (Blackwell)** GPU — the packed-FP4 kernels are SM100-only. +- **DeepGEMM built in-place** from the repo root: + ```bash + git submodule update --init --recursive # cutlass + fmt + ./develop.sh # builds deep_gemm._C and symlinks the .so + ``` +- **FlashInfer** (`pip install flashinfer`) for the two vs-benchmarks. Some + FlashInfer MoE backends are sensitive to the installed version (see + *Compatibility notes* below). +- Tests are imported by the benchmarks for shared quantization helpers, so the + repo root must be on the path (the scripts add `tests/` to `sys.path` themselves). + +## How to run + +### 1. `bench_packed_fp4.py` — FP4 format comparison (single device) + +DeepGEMM's own packed-FP4 paths: standalone GEMM (`mxfp4_gemm_nt` / +`nvfp4_gemm_nt`) and single-rank mega-MoE (`mxfp4_mxfp4_mega_moe` / +`nvfp4_nvfp4_mega_moe`), with `fp8_fp4_mega_moe` as the FP8×FP4 baseline. + +```bash +python benchmarks/bench_packed_fp4.py +``` + +Output: a GEMM table (mxfp4 vs nvfp4 us + TFLOPS) and a MoE table +(fp8xfp4 vs mxfp4 vs nvfp4 us, with nv/fp8 and nv/mx ratios). + +### 2. `bench_flashinfer_vs_deepgemm.py` — DeepGEMM vs FlashInfer (single device) + +Same MoE problem (identical shapes + routing) for every backend, apples-to-apples. +Backends: DeepGEMM `fp8_fp4_mega_moe` (dg_fp8), DeepGEMM `nvfp4_nvfp4_mega_moe` +(dg_nvfp4), FlashInfer `cute_dsl_fused_moe_nvfp4` (fi_cutedsl), FlashInfer +`cutlass_fused_moe` (fi_cutlass), FlashInfer trtllm-gen (fi_trtllm). Timing is +end-to-end device time via CUDA-graph replay (suffix `g` = graph, `e` = eager). + +```bash +python benchmarks/bench_flashinfer_vs_deepgemm.py # timing table +python benchmarks/bench_flashinfer_vs_deepgemm.py --breakdown # per-kernel breakdown +``` + +### 3. `bench_ep_multi_gpu.py` — multi-GPU expert-parallel EP + +`total_experts=512` sharded across N GPUs; each shard holds `num_tokens` local +tokens; top-k routing. DeepGEMM's native fused EP (dispatch + grouped GEMM + +combine via symmetric memory, timed end-to-end with comm internal) vs FlashInfer +single-device kernel + manual NCCL `all_reduce` combine (timed as +barrier → moe → all_reduce). Reports worst-rank average latency. + +```bash +python benchmarks/bench_ep_multi_gpu.py --num-processes 2 --num-tokens 32 +python benchmarks/bench_ep_multi_gpu.py --num-processes 4 --num-tokens 32 +python benchmarks/bench_ep_multi_gpu.py --num-processes 8 --num-tokens 32 +``` + +Flags: `--num-processes` (world size), `--num-tokens` (per-rank tokens for the +DeepGEMM dispatch model), `--inter` (intermediate size; `fp8xfp4` needs a +multiple of 512, e.g. 2304 is nvfp4-only), `--nccl-algo` (force `NCCL_ALGO`, +e.g. `NVLS`, for the combine `all_reduce`). + +## Results + +Environment: **NVIDIA B200 (SM100)**, torch `2.9.0a0+145a3a7bda.nv25.10`, +flashinfer `0.6.11`, single B200 for the single-device benches, 2× B200 for EP. +Captured 2026-07-06 on branch `mxfp4-mxfp4-mega-moe`. Times are kernel-only +device time (CUDA graph / profiler); lower is better. + +### bench_packed_fp4 — standalone packed-FP4 GEMM (2-CTA de-risk kernel) + +| M | N | K | mxfp4 us | mxfp4 TFLOPS | nvfp4 us | nvfp4 TFLOPS | nv/mx | +|---:|---:|---:|---:|---:|---:|---:|---:| +| 4096 | 4096 | 4096 | 59.9 | 2293 | 69.8 | 1968 | 1.16× | +| 4096 | 4096 | 8192 | 119.0 | 2309 | 134.3 | 2047 | 1.13× | +| 8192 | 8192 | 8192 | 456.6 | 2408 | 507.4 | 2167 | 1.11× | +| 2048 | 4096 | 16384 | 140.5 | 1957 | 153.0 | 1797 | 1.09× | + +MXFP4 is ~1.1–1.16× faster than NVFP4 on the standalone GEMM (same FP4 tensor-core +path; MXFP4's UE8M0 gran-32 SFs are cheaper than NVFP4's E4M3 gran-16 + global scale). + +### bench_packed_fp4 — single-rank mega-MoE (fp8xfp4 vs mxfp4 vs nvfp4) + +| tokens | experts | topk | hidden | inter | fp8xfp4 us | mxfp4 us | nvfp4 us | nv/fp8 | nv/mx | +|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 128 | 8 | 2 | 2048 | 2048 | 39.7 | 40.9 | 45.3 | 1.14× | 1.11× | +| 512 | 8 | 2 | 2048 | 2048 | 47.5 | 51.4 | 54.5 | 1.15× | 1.06× | +| 1024 | 32 | 4 | 4096 | 1536 | 124.6 | 117.2 | 122.4 | 0.98× | 1.04× | + +fp8xfp4 (FP8 E4M3 activations × FP4 weights) and the FP4×FP4 variants are within +~15%; which wins depends on shape (the 1024/32/4 shape is compute-bound enough +that the FP4-activation paths pull ahead of fp8xfp4). + +### bench_flashinfer_vs_deepgemm — single-device DeepGEMM vs FlashInfer (device us, CUDA graph) + +| tokens | experts | topk | hidden | inter | dg_fp8 | dg_nvfp4 | fi_cutedsl | fi_cutlass | fi_trtllm | +|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 128 | 32 | 4 | 2048 | 2048 | 63.7 | 69.1 | n/a | 110.1 | n/a | +| 512 | 32 | 4 | 2048 | 2048 | 77.7 | 84.4 | n/a | 107.5 | n/a | +| 1024 | 256 | 8 | 7168 | 2560 | 1347.7 | 1322.7 | n/a | 2242.7 | n/a | +| 32 | 256 | 16 | 4608 | 2304 | n/a | 604.2 | n/a | 1127.5 | n/a | + +DeepGEMM's fused mega kernel is ~1.6–1.9× faster than FlashInfer's +`cutlass_fused_moe` on these shapes. `dg_fp8` is `n/a` for `inter=2304` +(fp8xfp4 needs `inter % 512 == 0`); that shape is nvfp4-only. + +### bench_ep_multi_gpu — multi-GPU EP (world=2, total_experts=512, 32 tokens/rank, hidden=4608 inter=2560 top_k=16) + +| kernel | moe us | allreduce us | total us | +|---|---:|---:|---:| +| DeepGEMM fp8xfp4 mega (fused EP) | 721.9 | fused | 721.9 | +| DeepGEMM nvfp4 mega (fused EP) | 742.6 | fused | 742.6 | +| FlashInfer nvfp4 cutlass (GeGLU) | 1252.8 | 14.4 | 1271.0 | + +DeepGEMM's fused EP (dispatch + GEMM + combine via symmetric memory, comm +internal) is ~1.7× faster than FlashInfer `cutlass_fused_moe` + NCCL +`all_reduce` combine at world=2. The combine `all_reduce` itself is only +~14 µs here — the gap is the MoE kernel, not the combine. + +### `--breakdown` example (`bench_flashinfer_vs_deepgemm.py --breakdown`) + +For shape `tokens=32 experts=256 top_k=16 hidden=4608 inter=2304`: + +``` +===== DeepGEMM nvfp4 mega (fused EP) ===== + 1 distinct kernels, 1 launches/iter, total device 604.2 us/iter + #/it us/it kernel + 1 604.2 void deep_gemm::sm100_mxfp4_mxfp4_mega_moe_impl<...> + +===== FlashInfer nvfp4 (cutlass) ===== + 8 distinct kernels, 9 launches/iter, total device 1127.5 us/iter + #/it us/it kernel + 2 1058.4 cutlass::device_kernel<...GemmUniversal...> + 1 36.8 tensorrt_llm::finalizeMoeRoutingKernel + 1 15.1 tensorrt_llm::blockExpertPrefixSumKernel<256> + ... +``` + +DeepGEMM launches **1** fused kernel; FlashInfer cutlass launches **9** across +8 distinct kernels (two GEMMs + finalize + prefix-sum/expand/activation helpers). + +## Compatibility notes + +- The FlashInfer `cute_dsl` and `trtllm-gen` MoE backends are sensitive to the + installed FlashInfer version. On flashinfer `0.6.11` they fail with + `Module has no function 'flashinfer_moe_output_memset_inplace_bf16'` / + trtllm argument-mismatch errors; the `cutlass` backend works. The DeepGEMM + paths are unaffected. (The benches catch and report these per-backend errors + rather than aborting the sweep.) +- `fp8xfp4` requires `inter % 512 == 0`; `inter = 2304` is nvfp4-only. +- All single-device benches use CUDA-graph replay for timing; the suffix `g` + (graph) or `e` (eager) is printed next to each number. diff --git a/benchmarks/bench_flashinfer_vs_deepgemm.py b/benchmarks/bench_flashinfer_vs_deepgemm.py index f455aca92a..0a1cb57e24 100644 --- a/benchmarks/bench_flashinfer_vs_deepgemm.py +++ b/benchmarks/bench_flashinfer_vs_deepgemm.py @@ -8,7 +8,12 @@ comparison is apples-to-apples at the op level. We report end-to-end device time (CUDA events), which includes all sub-kernels (FlashInfer dispatches moe_sort + gemm1 + gemm2/finalize; DeepGEMM is a single mega kernel). + +With ``--breakdown``, instead of the timing table we print a per-kernel +device-time breakdown (launches/iter + us/iter per kernel) for each backend, +so you can see how many kernels each launches and where the time goes. """ +import argparse import os import sys import torch @@ -115,6 +120,45 @@ def cuda_time(fn, warmup=10, iters=50): return _time_events(fn, iters), 'eager' +def kernel_breakdown(run, warmup=10, iters=50): + """Profile `run` and return per-kernel (name, launches/iter, us/iter), sorted by us/iter.""" + for _ in range(warmup): + run() + torch.cuda.synchronize() + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + run() + torch.cuda.synchronize() + rows = [] + for evt in prof.key_averages(): + dev = getattr(evt, 'self_device_time_total', None) or getattr(evt, 'self_cuda_time_total', 0.0) + if dev and dev > 0: + rows.append((evt.key, evt.count / iters, dev / iters)) # name, launches/iter, us/iter + rows.sort(key=lambda r: -r[2]) + return rows + + +def show_breakdown(title, builder, *, is_tuple=False): + """Build a backend, profile its run closure, and print the per-kernel device-time + breakdown (launches/iter + us/iter per kernel).""" + print(f'\n===== {title} =====') + try: + out = builder() + run = out[0] if is_tuple else out + rows = kernel_breakdown(run) + if is_tuple: + out[1].destroy() + except Exception as ex: + print(f' ERROR: {str(ex).splitlines()[-1][:120]}') + return + total = sum(r[2] for r in rows) + n_launches = sum(r[1] for r in rows) + print(f' {len(rows)} distinct kernels, {n_launches:.0f} launches/iter, total device {total:.1f} us/iter') + print(f' {"#/it":>5} {"us/it":>8} kernel') + for name, cnt, us in rows: + print(f' {cnt:>5.0f} {us:>8.1f} {name[:88]}') + + def _interleave_gate(x, group_size=64, dim=1): sizes = x.size() dim = dim % x.dim() @@ -211,7 +255,7 @@ def run(): return run, buf -def bench_one(group, num_tokens, hidden, inter, num_experts, top_k): +def bench_one(group, num_tokens, hidden, inter, num_experts, top_k, breakdown=False): torch.manual_seed(0) dev = 'cuda' x_bf16 = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=dev) / 10 @@ -224,6 +268,26 @@ def bench_one(group, num_tokens, hidden, inter, num_experts, top_k): rw = (rw / rw.sum(-1, keepdim=True)).float() sel = sel.to(torch.int64) + if breakdown: + print(f'\n### shape: tokens={num_tokens} experts={num_experts} top_k={top_k} ' + f'hidden={hidden} inter={inter} ###') + show_breakdown('DeepGEMM fp8xfp4 mega', + lambda: build_deepgemm(group, x_bf16, w1_bf16, w2_bf16, sel, rw, + num_experts, top_k, hidden, inter), is_tuple=True) + show_breakdown('DeepGEMM nvfp4 mega', + lambda: build_deepgemm_nvfp4(group, x_bf16, w1_bf16, w2_bf16, sel, rw, + num_experts, top_k, hidden, inter), is_tuple=True) + show_breakdown('FlashInfer nvfp4 (cute_dsl)', + lambda: build_flashinfer(x_bf16, w1_bf16, w2_bf16, sel, rw, + num_experts, top_k, hidden, inter)) + show_breakdown('FlashInfer nvfp4 (cutlass)', + lambda: build_flashinfer_cutlass(x_bf16, w1_bf16, w2_bf16, sel, rw, + num_experts, top_k, hidden, inter)) + show_breakdown('FlashInfer nvfp4 (trtllm-gen)', + lambda: build_flashinfer_trtllm(x_bf16, w1_bf16, w2_bf16, + num_experts, top_k, hidden, inter)) + return + res = {} def try_run(name, builder): @@ -261,15 +325,26 @@ def f(k): def main(): + parser = argparse.ArgumentParser( + description='Single-device DeepGEMM vs FlashInfer NVFP4 MoE benchmark.') + parser.add_argument('--breakdown', action='store_true', + help='Print the per-kernel device-time breakdown for each backend ' + 'instead of the end-to-end timing table') + args = parser.parse_args() + os.environ.setdefault('MASTER_ADDR', '127.0.0.1') os.environ.setdefault('MASTER_PORT', '12566') dist.init_process_group('nccl', rank=0, world_size=1) group = dist.group.WORLD print(f'torch {torch.__version__}, flashinfer {flashinfer.__version__}, {torch.cuda.get_device_name(0)}') - print('DeepGEMM mega-MoE (fp8xfp4 / nvfp4) vs FlashInfer NVFP4 MoE backends (device us, CUDA graph; g=graph e=eager)') - print(f'{"tok":>5} {"exp":>4} {"tk":>3} {"hid":>5} {"int":>5} | {"dg_fp8":>9} | {"dg_nvfp4":>9} | ' - f'{"fi_cutedsl":>9} | {"fi_cutlass":>9} | {"fi_trtllm":>9}') + if args.breakdown: + print('Per-kernel device-time breakdown per backend (torch.profiler, single device)') + else: + print('DeepGEMM mega-MoE (fp8xfp4 / nvfp4) vs FlashInfer NVFP4 MoE backends ' + '(device us, CUDA graph; g=graph e=eager)') + print(f'{"tok":>5} {"exp":>4} {"tk":>3} {"hid":>5} {"int":>5} | {"dg_fp8":>9} | {"dg_nvfp4":>9} | ' + f'{"fi_cutedsl":>9} | {"fi_cutlass":>9} | {"fi_trtllm":>9}') for cfg in ( (128, 2048, 2048, 32, 4), (512, 2048, 2048, 32, 4), @@ -277,7 +352,7 @@ def main(): # inter=2304 is nvfp4-only (fp8xfp4 needs inter % 512 == 0) (32, 4608, 2304, 256, 16), ): - bench_one(group, *cfg) + bench_one(group, *cfg, breakdown=args.breakdown) dist.destroy_process_group() diff --git a/benchmarks/bench_kernel_breakdown.py b/benchmarks/bench_kernel_breakdown.py deleted file mode 100644 index 3479c77dfb..0000000000 --- a/benchmarks/bench_kernel_breakdown.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Per-kernel device-time breakdown for each MoE backend (single device). - -Profiles each backend's run closure and lists every GPU kernel with its -per-iteration count and device time, so you can see how many kernels each -backend launches and where the time goes. -""" -import os -import sys - -import torch -import torch.distributed as dist -import torch.nn.functional as F - -import deep_gemm - -sys.path.insert(0, os.path.dirname(__file__)) -from bench_flashinfer_vs_deepgemm import ( - build_deepgemm, build_deepgemm_nvfp4, build_flashinfer, build_flashinfer_cutlass, -) -from fi_trtllm import build_flashinfer_trtllm - - -def kernel_breakdown(run, warmup=10, iters=50): - for _ in range(warmup): - run() - torch.cuda.synchronize() - with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as prof: - for _ in range(iters): - run() - torch.cuda.synchronize() - rows = [] - for evt in prof.key_averages(): - dev = getattr(evt, 'self_device_time_total', None) or getattr(evt, 'self_cuda_time_total', 0.0) - if dev and dev > 0: - rows.append((evt.key, evt.count / iters, dev / iters)) # name, launches/iter, us/iter - rows.sort(key=lambda r: -r[2]) - return rows - - -def show(title, build, *, is_tuple=False): - print(f'\n===== {title} =====') - try: - out = build() - run = out[0] if is_tuple else out - rows = kernel_breakdown(run) - if is_tuple: - out[1].destroy() - except Exception as ex: - print(f' ERROR: {str(ex).splitlines()[-1][:120]}') - return - total = sum(r[2] for r in rows) - n_launches = sum(r[1] for r in rows) - print(f' {len(rows)} distinct kernels, {n_launches:.0f} launches/iter, total device {total:.1f} us/iter') - print(f' {"#/it":>5} {"us/it":>8} kernel') - for name, cnt, us in rows: - print(f' {cnt:>5.0f} {us:>8.1f} {name[:88]}') - - -def main(): - os.environ.setdefault('MASTER_ADDR', '127.0.0.1') - os.environ.setdefault('MASTER_PORT', '12599') - dist.init_process_group('nccl', rank=0, world_size=1) - group = dist.group.WORLD - - # 4-GPU-shard view of the previous shape: 128 experts, 32 tokens, top-16, m_e=4. - num_tokens, hidden, inter, num_experts, top_k = 32, 4608, 2304, 128, 16 - print(f'torch {torch.__version__}, flashinfer {__import__("flashinfer").__version__}, ' - f'{torch.cuda.get_device_name(0)}') - print(f'shape: tokens={num_tokens} experts={num_experts} top_k={top_k} hidden={hidden} inter={inter}') - - torch.manual_seed(0) - x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device='cuda') / 10 - w1 = torch.randn(num_experts, inter * 2, hidden, dtype=torch.bfloat16, device='cuda') / 10 - w2 = torch.randn(num_experts, hidden, inter, dtype=torch.bfloat16, device='cuda') / 10 - rw, sel = torch.topk(F.softmax(torch.randn(num_tokens, num_experts, device='cuda'), -1), top_k, dim=-1) - rw = (rw / rw.sum(-1, keepdim=True)).float() - sel = sel.to(torch.int64) - - show('DeepGEMM fp8xfp4 mega (fused EP)', - lambda: build_deepgemm(group, x, w1, w2, sel, rw, num_experts, top_k, hidden, inter), is_tuple=True) - show('DeepGEMM nvfp4 mega (fused EP)', - lambda: build_deepgemm_nvfp4(group, x, w1, w2, sel, rw, num_experts, top_k, hidden, inter), is_tuple=True) - show('FlashInfer nvfp4 (cute_dsl)', - lambda: build_flashinfer(x, w1, w2, sel, rw, num_experts, top_k, hidden, inter)) - show('FlashInfer nvfp4 (cutlass)', - lambda: build_flashinfer_cutlass(x, w1, w2, sel, rw, num_experts, top_k, hidden, inter)) - show('FlashInfer nvfp4 (trtllm-gen)', - lambda: build_flashinfer_trtllm(x, w1, w2, num_experts, top_k, hidden, inter)) - - dist.destroy_process_group() - - -if __name__ == '__main__': - main()