Repository navigation
[SM90] Add W4A8 MegaMoE support - #6343
Merged
Merged
Conversation
Framework-agnostic kernel headers for the SM90 (Hopper) NVFP4 MegaMoE fused family -- a single persistent kernel that fuses the cross-rank token dispatch, both FFN GEMMs (W4A8: NVFP4 weights, FP8 activations), SwiGLU and the combine reduction over an 8-rank NVLink symmetric buffer. This is part 1 of 2. It covers the BN256 fused family, which the host policy selects for per-rank M <= family_threshold; the BN128 split family for larger M follows in part 2 and reuses every support header added here. Contents (include/flashinfer/megamoe/, no torch dependency per the framework separation rule): sm90_nvfp4_mega_moe_fused.cuh / _body.inl the fused kernel scheduler.cuh dev-m dynamic task scheduler layout.cuh workspace / token-pool layout nvfp4_dequant.cuh NVFP4 -> FP8 dequant barrier.cuh, sym_buffer.cuh NVLink barriers, symmetric buffer detail/ ptx, sm90 mma, tma, shared helpers Dispatch synchronisation uses counter-based push dispatch (the source rank writes rows straight into the destination pool and release-adds a DONE count) with a parity-rotated pool, gated by an M-range policy; outside those ranges it falls back to the NVLink barrier protocol. Requires -std=c++20 (the kernels take a float non-type template parameter for the activation clamp). Verified: 8 template instantiations build clean on sm_90a -- push+rotated, barrier-only, push-only, static scheduler, BLOCK_M=24, and the Flash / Pro / MiMo geometries, plus the phase-stamp diagnostic arm. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Wires the SM90 NVFP4 MegaMoE fused kernel into FlashInfer's JIT pipeline.
The kernel resolves its shared-memory layout, task schedule and dispatch
protocol at compile time, so the plan is selected in Python and rendered into
the config header; one module is built per (shape, plan) pair.
flashinfer/megamoe/heuristics.py plan selection (tuning tables, all-M arm,
counter-dispatch policy, pool sizing)
flashinfer/megamoe/__init__.py MegaMoESymmBuffer + nvfp4_mega_moe
flashinfer/jit/megamoe.py Jinja render + JitSpec
csrc/megamoe/*.cu, *.jinja launcher, TMA descriptors, TVM-FFI binding
tests/moe/test_megamoe_nvfp4_sm90.py
The symmetric-buffer layout is computed in C++ and returned as byte offsets so
Python slices exactly what the kernel addresses. The scale-factor region is
sized for the largest BLOCK_M candidate, since it is allocated once while the
plan varies per call -- sizing it from the current plan silently corrupts every
BLOCK_M != 8 schedule.
Needs -std=c++20 (float non-type template parameter for the activation clamp)
and builds against the pinned CUTLASS 4.4.
Verified on 8x H20-3e (EP8) against the DeepGEMM reference in the same process:
- correctness: 27/27 bit-identical over Flash / Pro / MiMo at
M = 1, 2, 4, 8, 16, 32, 64, 128, 256
- performance: same kineto kernel timer, max over 8 ranks -- 18 of 27 points
within +-2%, all within +-8%; the spread is measurement noise on a shared
host (at large M, where the signal dominates, every point is within 2%)
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The workspace sizing helpers are CUTLASS_HOST_DEVICE -- shared between the
launcher and the kernel -- so their assertion failure path is also compiled for
the host, where `asm("trap;")` is not a valid instruction. It only stayed hidden
because the fused launcher happens to reach those helpers through constant
evaluation.
Route both paths through FI_MEGAMOE_TRAP(): `trap;` under __CUDA_ARCH__,
__builtin_trap() otherwise. Device code is unchanged -- the coverage build's
cubin is byte-identical.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ast_math is set gen_jit_spec() adds -use_fast_math to every module unless told otherwise. For this kernel that overrides the fast_math argument: kFastMath=false picks expf() and a true division for SwiGLU in the source, but -use_fast_math makes nvcc lower expf() to the __expf() intrinsic and turns on --ftz anyway, so fast_math=False still compiled to the fast-math kernel. The DeepGEMM reference passes --use_fast_math to this family only when fast_math is set, and fast_math=False disagreed with it on 1-2 bf16 ULPs (M=64, 256 on Flash). Pass use_fast_math=fast_math so the compiler flag follows the template parameter. fast_math=True, the default, builds exactly as before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The pass/fail verdict was reduced across all ranks, but bit-identical, finite, max_abs and the mismatch count were rank 0's local values. When another rank diverged the test printed contradictory lines such as "FAIL ... bit-identical=True finite=True". Reduce each field as well (MIN for the flags, MAX for max_abs, SUM for the count) so the printed line describes all ranks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…files
Formatting from clang-format, ruff format and trailing-whitespace, plus the
fixes ruff and mypy asked for:
- heuristics.py: the tuning rows are typed tuples instead of mixed
int/bool lists, so the bool fields keep their type through FusedPlan
- jit/megamoe.py: read the copied sources under a context manager
- test: drop an unused import
No functional change.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The test spawns one worker per GPU itself (torch.multiprocessing.spawn), so launching it under torchrun as the docstring said would start 8 copies of the 8-process job. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Adds the BN128 split family, which the host policy selects for per-rank M > family_threshold, where a single fused persistent kernel no longer fits the working set. L1 runs dispatch-dequant with half-stream mode4 and the calibrated physical-warp remap; L2 is the warp-private contiguous scatter. Stacked on part 1 -- every support header (scheduler, layout, barriers, symmetric buffer, ptx/mma/tma helpers) is reused as-is. The only change to an existing file is restoring ptx::st_shared_cluster_async, which is used solely by the split L1 body. New files: sm90_nvfp4_mega_moe.cuh split entry (L1 + L2) sm90_nvfp4_mega_moe_split_l1_body.inl L1: dispatch-dequant + GEMM sm90_nvfp4_mega_moe_split_l2_body.inl L2: GEMM + scatter sm90_nvfp4_mega_moe_mode2_dequant.cuh Mode2 braided-sign dequant With both parts in, the host selector covers the full M range: fused below the threshold, split above it. Verified: the part-1 instantiations still build unchanged (identical cubin), and split L1 (with the two production defines the host always emits) and L2 build clean on sm_90a. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Adds the BN128 split path to the Python API. `nvfp4_mega_moe` now dispatches on
tokens-per-rank: the fused kernel at or below `family_threshold` (256), the
split pair above it, where a single persistent kernel no longer fits the
working set.
The split family runs two kernels per call -- L1 (dispatch-dequant + first GEMM,
2-CTA cluster) then L2 (second GEMM + scatter) -- each with its own pipeline
depth and shared-memory budget. Those are template parameters, so they have to
be known before the module is built. Rather than paraphrasing ~100 lines of
shared-memory arithmetic in Python, the reference sizing code is carried over as
`megamoe_sm90_nvfp4_split_plan.cuh` and compiled into a small template-free plan
module that Python queries first; the kernel module is then rendered from its
answer. Mis-sizing that pipeline would be silent, so it is worth not
re-deriving.
csrc/megamoe/megamoe_sm90_nvfp4_split_plan.{cuh,cu} plan query
csrc/megamoe/megamoe_sm90_nvfp4_split.cu two-phase launcher
csrc/megamoe/megamoe_sm90_nvfp4_split_binding.cu TVM-FFI binding
csrc/megamoe/megamoe_sm90_nvfp4_split_config.jinja
flashinfer/jit/megamoe.py, flashinfer/megamoe/__init__.py
Two descriptor details differ from the fused path and are easy to get wrong:
the L1 output box is 64x64 (each of the two M warpgroups stores its own
post-SwiGLU tile, so it is not (BLOCK_N/2, BLOCK_M)), and the L2 activation
scales are per-128 like L1's, not per-BLOCK_N/2.
Verified on 8x H20-3e (EP8) against the DeepGEMM reference in the same process:
- correctness: 6/6 bit-identical over Flash / Pro / MiMo at M = 512, 1024
- performance: same kineto timer summed over both phases, max over 8 ranks --
all six points within +-0.5%
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
gen_jit_spec() adds -use_fast_math to every module by default, which implies --ftz=true and lowers expf()/division to approximate intrinsics. The DeepGEMM reference builds the split L1/L2 kernels with no fast-math flag at all, for either value of fast_math: kFastMath alone picks __expf/fast_rcp or expf/division in the SwiGLU epilogue. With the default flag, our split kernels flushed denormal intermediates where the reference keeps them, and fast_math=False still ran the intrinsic path (1-2 bf16 ULP mismatches at M = 512, 1024, 2048 on Flash). A SASS diff showed the extra .FTZ suffixes. Pass use_fast_math=False. Appending --ftz=false --prec-div=true --prec-sqrt=true is not enough, because it does not undo the expf() lowering. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
kFastMath is a kernel template parameter, so it has to match on both the FlashInfer and DeepGEMM side for the bit-identical comparison to mean anything. The test hardcoded fast_math=True on the reference call and relied on the MegaMoESymmBuffer default for ours, so the fast_math=False path had no coverage at all -- which is how the -use_fast_math mismatch in the JIT specs went unnoticed. Thread a --no-fast-math flag through both sides and report it in the result line. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The prepacked NVFP4 scales are tile-major (E, N/block_n, K/128, block_n, 8), so the weight layout belongs to the kernel family: the fused kernel is BN256 and split L1/L2 is BN128 (kBlockN is hardcoded to 128 in the split config). The test transformed once with block_n=256 and then let nvfp4_mega_moe auto-dispatch, so every M above the threshold handed BN256 weights to the BN128 split kernel. Nothing above M=256 was ever covered correctly. Transform per family on first use, and pin kernel_family explicitly on both the FlashInfer and the DeepGEMM call instead of relying on the two "auto" rules agreeing -- a silent disagreement would compare two different kernels and still report a pass. The family is now printed in the result line. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
ruff (SIM115) flags the bare open().read() calls that copy the split plan and kernel sources into the generated directory. No functional change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Provide FP8-activation, FP4-weight expert-parallel execution with persistent fused kernels and the two-launch NVFP4 split path. Synchronize SS L2 math warps before decoded-weight reuse, validate packed layouts and distributed outputs, and document the collective APIs. Include H20/H200 workload tables in the kernel guide. Mark the H200 MiMo EP8 M4096 latency pending verification while retaining all observed batches. Validation: repository pre-commit hooks; CPU weight-layout tests; eight-rank H20 and H200 correctness and CUDA graph replay tests.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: 🟡 Moderate · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 26.56% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 64 functions across 15 files. (2 skipped:… | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Title check The title clearly and concisely identifies the main change: adding SM90 W4A8 MegaMoE support. Description check The description follows the repository template and provides the implementation scope, usage, performance data, validation results, related-issue status, checklist completion, experimental-track statu…
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 26.56% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 64 functions across 15 files. (2 skipped: 2 unsupported.)
- Fix all pre-merge checks with AI
✨ Finishing Touches 💡 1
-
🛠️ Fix failing CI checks 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Comment @coderabbitai help to get the list of available commands.
ouershua2299
marked this pull request as ready for review
October 12, 2026 06:17
ouershua2299
requested review from
Aneureka,
IwakuraRein,
Jingfan-Sun,
StudyingShao,
YangXu1990uiuc,
aleozlx,
bkryu,
feih-nv,
jiahanc,
nv-yunzheq,
samuellees,
yyihuang and
yzh119
as code owners
October 12, 2026 06:17
yzh119
approved these changes
Oct 13, 2026
Collaborator
|
@flashinfer-bot run |
Collaborator
|
/bot run tests/moe/test_megamoe_fp4_sm90.py tests/moe/test_megamoe_fp4_weights.py tests/moe/test_megamoe_nvfp4_sm90.py |
Collaborator
Parameterize fused rank count and host layouts, isolate EP-specific JIT modules, and preserve EP8 tuning. Validate EP4 fused/split and push barrier modes on H20/H200, and record standard cold-L2 timing in the kernel guide.
Remove the standalone SM90 Markdown guide and clean up its references. Retain the API documentation and original root README.
Collaborator
|
@flashinfer-bot run |
yyihuang
enabled auto-merge (squash)
October 13, 2026 20:58
Merged
5 tasks done
ouershua2299
added a commit
to ouershua2299/sglang
that referenced
this pull request
Oct 19, 2026
…4 FP4 checkpoints On SM90, --moe-a2a-backend megamoe with the original DeepSeek-V4 FP4 checkpoints converts the experts losslessly to NVFP4 and runs flashinfer.megamoe.nvfp4_mega_moe (W4A8; fused for M <= 256 tokens per rank, split above). Needs flashinfer-ai/flashinfer#6343 and sgl-project#6517; without them the load fails with an explicit error. Adds the SGLANG_FLASHINFER_MEGAMOE_SM90_{PREWARM,KERNEL_FAMILY} env vars, a DeepSeek-V4 cookbook section, a CPU conversion unit test, and a label-gated 8x H200 e2e test that skips when FlashInfer lacks flashinfer.megamoe. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
yyihuang
pushed a commit
that referenced
this pull request
Oct 21, 2026
…support other models (#6517) <!-- .github/pull_request_template.md --> ## 📌 Description Follow-up to #6343 (SM90 FP4 MegaMoE): 1. **Cleanup fix:** with `single_active_dispatch_warp`, the inactive dispatch warp cleaned the workspace with empty expert counts and left pool blocks ≥ 32 of hot experts stale, so the next launch failed (H20, 97–128 tokens per rank, e.g. CUDA-graph capture); it now fetches the counts first (outputs bitwise unchanged, timing within noise). 2. **H200 push dispatch:** add measured H200 EP8 buckets that enable push for M ≤ 32 on DeepSeek-V4-Flash and MiMo-V2.5-Pro, and M ≤ 64 on DeepSeek-V4-Pro (rotated pool at M 33–64), saving about 10 µs per call (1.03–1.16x kernel time, bitwise-identical outputs; H20 and EP4 plans unchanged). 3. **Other models:** accept any expert count that divides across ranks, keep the selected plans compilable for untuned geometries (H20 buckets match the full geometry; the planner checks the full fused shared-memory layout, which grows with the expert count, and falls back to the next legal smaller tile when the chosen one does not fit; unsupported shapes raise `ValueError`), and allow push only on validated plans, so Kimi-K2, DeepSeek-V3, GLM-4.x and Qwen3 shapes now run correctly while plans of the three tuned geometries are unchanged. 4. **Usage docs:** the weight-preparation docs now state what callers must convert before `transform_{nvfp4,mxfp4}_weights_for_mega_moe`: Marlin byte order (checkpoints are usually low-nibble-first), scales inside the decoder range (NVFP4 `6 * scale <= 448`; MXFP4 `2^-5`–`2^6` for groups with nonzero codes, otherwise outputs silently saturate or flush), and the excess moved into per-expert global scales passed at run time; the module example shows the full flow. 5. **Framework hooks:** add `MegaMoESymmBuffer.prewarm()`, which builds every NVFP4 module `nvfp4_mega_moe` can select up to capacity (sharded across ranks, before CUDA graph capture), and `commit_staged()`, which accepts rows a producer wrote directly into the now-documented staging views (`x`, `x_sf`, `topk_idx`, `topk_weights`), so integrations no longer reach into JIT internals or buffer fields. **How to use:** FlashInfer only provides the kernels and the layout transforms; frameworks convert checkpoint tensors to the format above at load time, call `prewarm()` once, and stage inputs with `stage()` or in place plus `commit_staged()` (the SGLang integration does this for `nvidia/DeepSeek-V4-Flash-NVFP4`, where about 51% of blocks would saturate without rebalancing). ## 🔍 Related Issues Follow-up to #6343. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). - `test_hot_expert_workspace_cleanup` reproduces fix 1: it fails before the fix on 8×H20 and passes after. - New CPU tests cover the H200 buckets, the counter-dispatch limits and plan legality for other geometries, including 2048 and 8192 experts (every plan fits shared memory); every fused and split plan selected for 12 geometries (the three tuned models, six common shapes, 1024 / 2048 / 8192 experts) compiles on SM90, and plans of the tuned and previously validated geometries are unchanged. - `test_nvfp4_prewarm_and_direct_staging` checks that after `prewarm()` no fused or split call compiles a new module, that in-place staging plus `commit_staged()` gives correct results under CUDA graphs, and that out-of-range counts are rejected. - `tests/moe/test_megamoe_fp4_{sm90,ep,weights}.py`: 73 passed on 8×H200. ## Reviewer Notes <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * Added optional JIT prewarming for NVFP4 workloads across supported token counts. * Added direct staging support for efficiently preparing dispatch inputs. * Added support for transformed NVFP4 weights with per-expert global scales. * Improved dispatch planning across supported hardware configurations and workload shapes. * **Bug Fixes** * Fixed workspace cleanup for inactive dispatch paths, improving repeated CUDA Graph execution reliability. * Improved expert-count handling during dispatch cleanup. * **Documentation** * Clarified supported weight formats, scale requirements, dispatch behavior, and configuration limits. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: ouershua2299 <180947746+ouershua2299@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
5 tasks done
ouershua2299
added a commit
to ouershua2299/sglang
that referenced
this pull request
Oct 21, 2026
…4 NVFP4 checkpoints On SM90, --moe-a2a-backend megamoe with the ModelOpt NVFP4 checkpoints (nvidia/DeepSeek-V4-Flash-NVFP4) repacks the expert tensors at load into the input of flashinfer.megamoe.transform_nvfp4_weights_for_mega_moe (Marlin byte order, local scales rebalanced into the decoder range, one global scale per expert) and runs flashinfer.megamoe.nvfp4_mega_moe (W4A8; fused for small batches, split above). Needs flashinfer-ai/flashinfer#6343 and sgl-project#6517; without them the load fails with an explicit error. EPLB is rejected on this path. Adds the SGLANG_FLASHINFER_MEGAMOE_SM90_{PREWARM,KERNEL_FAMILY} env vars, a DeepSeek-V4 cookbook section, a CPU unit test for the conversion, and a label-gated 8x H200 e2e test that skips when FlashInfer lacks flashinfer.megamoe. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
3 of 11 tasks
ouershua2299
added a commit
to ouershua2299/sglang
that referenced
this pull request
Oct 25, 2026
…4 NVFP4 checkpoints On SM90, --moe-a2a-backend megamoe with the ModelOpt NVFP4 checkpoints (nvidia/DeepSeek-V4-Flash-NVFP4) repacks the expert tensors at load into the input of flashinfer.megamoe.transform_nvfp4_weights_for_mega_moe (Marlin byte order, local scales rebalanced into the decoder range, one global scale per expert) and runs flashinfer.megamoe.nvfp4_mega_moe (W4A8; fused for small batches, split above). Needs flashinfer-ai/flashinfer#6343 and sgl-project#6517; without them the load fails with an explicit error. EPLB is rejected on this path. Adds the SGLANG_FLASHINFER_MEGAMOE_SM90_{PREWARM,KERNEL_FAMILY} env vars, a DeepSeek-V4 cookbook section, a CPU unit test for the conversion, and a label-gated 8x H200 e2e test that skips when FlashInfer lacks flashinfer.megamoe. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 task done
Closed
3 of 11 tasks
jiahanc
pushed a commit
that referenced
this pull request
Oct 27, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description This adds CuTe DSL NVFP4×FP8 to unified `moe_ep.MoEEpLayer`; [#6343](#6343) provides the separate `flashinfer.megamoe` CUDA API. NV inherits MXFP4 frontend/kernel classes, overriding weight-format, conversion and scheduling behavior with separate cache/workspace identities. The default is `nvfp4_kperm=64`. For H7168/I3072/E384/top-6/EP4, 13 capacities span 8–32768 tokens/rank. Other capacities use the next larger entry, capped at 32768; other shapes use heuristics. MX gains optional `compact_pull_buffer=True`, disabled by default; MX auto/candidates are unchanged. ### Performance and reproduction We compare NV defaults with MX main defaults with manual compact=True. Positive means NV is slower. | T/rank | NV µs | MX µs | Paired NV/MX Δ | |---:|---:|---:|---:| | 8 | 452.4 | 430.1 | +5.13% | | 16 | 642.8 | 643.2 | -0.03% | | 32 | 807.2 | 777.6 | +3.73% | | 64 | 908.8 | 878.3 | +3.51% | | 128 | 938.8 | 910.9 | +3.16% | | 256 | 1014.7 | 990.7 | +2.42% | | 512 | 1227.9 | 1191.2 | +3.08% | | 1024 | 1607.8 | 1608.8 | -0.13% | | 2048 | 2733.0 | 2868.6 | -4.92% | | 4096 | 5020.7 | 5229.2 | -4.18% | | 8192 | 9530.0 | 9845.9 | -3.26% | | 16384 | 18507.6 | 19150.0 | -3.34% | | 32768 | 36690.8 | 38506.6 | -4.80% | We measured MegaMoE plus top-k reduction on 4×H200. Each block takes the maximum of four rank medians. Paired ratios are summarized across three rounds and three fresh processes. SM clocks were set to 1830 MHz. Software power-cap events were recorded, so a constant clock in every timed block is not guaranteed. NV stores 5.9% more weight bytes, adding traffic. T512 observations were retained: rounds ranged +2.71–13.08%, while process medians ranged +2.89–3.67%. ```bash nv_bench_root="$(mktemp -d -p . sm90-pr1-current-XXXXXX)" for fresh in 1 2 3; do FLASHINFER_MOE_EP_KNOB_CACHE="$(pwd)/${nv_bench_root}/knob-${fresh}.json" \ torchrun --standalone --nproc_per_node=4 benchmarks/bench_moe_ep_sm90_nvfp4.py \ --tokens 8 16 32 64 128 256 512 1024 2048 4096 8192 16384 32768 \ --hidden 7168 --intermediate 3072 --experts 384 --top-k 6 --seed 1234 \ --nv-kperm 64 --nv-overlap auto --mx-compact \ --rounds 3 --warmup 10 --iters 50 --padding-us 300 --cooldown-seconds 5 \ --output "${nv_bench_root}/fresh-${fresh}.json" || exit $? done ``` ## 🔍 Related Issues Stacked follow-up: ruoqianguo#1. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). The existing NCCL-EP high-throughput test (tests/moe_ep/test_moe_ep_cudagraph_multirank.py, `high_throughput`) also fails on unmodified main (eed63cd) in our 4×H200 environment; it is unrelated to this PR. Default MXFP4/FP8/BF16 cubins match tested main byte-for-byte. The four-format EP4 end-to-end, host/unit and registration tests passed. NV requires relative L2<2.5% against FP32 with promoted operands. All 55 public auto candidates compiled and ran; a new layer reused the cache. This was functional smoke, not per-candidate FP32/performance qualification. Regression: ```bash for dtype in nvfp4 mxfp4 fp8 bf16; do torchrun --standalone --nproc_per_node=4 -m pytest \ "tests/moe_ep/test_moe_ep_sm90_pull_${dtype}_mega_multirank.py" \ -vv -ra -rs --tb=short --full -m arch_hopper || exit $? done ``` ## Reviewer Notes Please review weight preprocessing and the `nvfp4_kperm` layout tags, NV's inheritance from the MX classes and the hooks added to the MX/FP8 base classes, the default tactic table, and MX's optional `compact_pull_buffer`. The sf2 scale ring is covered by the stacked [PR2](ruoqianguo#1). ## 🔬 Experimental Track <!-- Only for PRs submitted under the experimental policy (CONTRIBUTING.md → "Experimental APIs and Backends"). Leave this section untouched for normal PRs. --> - [ ] This PR is **experimental**: it adds or changes code under `flashinfer/experimental/` and/or an `@flashinfer_experimental_api`. Tracking issue: # - [ ] The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release. - [ ] Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff). - [ ] Tests live in `tests/experimental/` and were validated on the intended hardware; a runnable example is included. - [ ] Nothing is registered in `flashinfer/aot.py`, and no experimental backend is reachable from `backend="auto"` without `FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. (Calling an `@flashinfer_experimental_api` or naming a backend explicitly is itself the opt-in and needs no environment variable.) - [ ] **Test scope declared below.** The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows. <!-- Required for experimental PRs. Replace the commented lines below with your targets. Do not delete the fence or change its `experimental-tests` tag — the experimental-track watcher reads it verbatim to decide which targets to ask CI for. --> ```experimental-tests # One target per line: a directory or a file. (A pytest ::selector is not # supported -- the sharding runner cannot consume one.) Must be under # tests/experimental/ and must exist. Delete these comment lines and add yours, e.g. # # tests/experimental/test_my_backend.py # tests/experimental/my_backend/ # # Declaring the whole tree (tests/experimental/) is allowed but means every # experimental PR pays for every other feature's tests, in every matrix cell. ``` <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * Added SM90 NVFP4 support for expert-parallel MegaMoE, including weight preprocessing, configurable layouts, and automatic or custom tuning. * Added options for K-permuted weights and FC1 scheduling, with validation for incompatible configurations. * Added a distributed benchmark for measuring NVFP4 performance. * **Documentation** * Documented the backend’s numerical behavior, weight requirements, configuration, tuning, and benchmark methodology. * **Tests** * Added multirank coverage for output accuracy, repeatability, layout validation, and CUDA graph scheduling. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
6 tasks done
5 of 12 tasks
2 of 7 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
Co-authors: @AichenF, @jychen21, @YijiaZhao.
Submitted by: xuechengw.
flashinfer.megamoeruns expert-parallel MoE on NVLink-connected Hopper GPUs,using NVFP4/MXFP4 weights, FP8 E4M3 activations and BF16 outputs. The kernels
perform token dispatch, two expert GEMMs, clamped SwiGLU and top-k reduction.
Usage
Allocate
MegaMoESymmBuffercollectively before graph capture. Prepare weightswith
transform_nvfp4_weights_for_mega_moeortransform_mxfp4_weights_for_mega_moe; all ranks stage inputs and call matchingkernels in matching order.
NVFP4 supports 256 or 384 total experts. The process group selects EP4 or EP8;
no separate execution API is needed. EP4 uses the generic fused schedule;
EP8 retains its measured tuning policies. Auto selection uses fused through
256 tokens per rank, then split.
Input scales are FP32 per 128 channels; routing uses INT64 expert indices
(
-1masks a route) and FP32 weights. Hidden width is a multiple of 256 andintermediate width a multiple of 128. Raw weight scales use UINT8 E4M3 per
16 values for NVFP4 or E8M0 per 32 values for MXFP4.
Performance
Cold L2 before each call; one MoE call per graph; CUDA events; maximum time
across ranks per call, then the median of 50 samples. Split timing includes
both launches. Packing, compilation and cache clearing are outside timing.
Mis tokens per rank; capacity is the allocated maximum tokens per rank.All 82 workloads: H20 and H200 latency
H200 EP8 regression against the original NVFP4 kernels covers 12 points
(Flash/MiMo/Pro, M=1/32/96/384, capacity 384): latency changes range from
-1.07% to +0.64%. Flash M96 is 227.808 us before and
227.440 us after.
Validation
CPU plan/cache and weight-layout checks pass. EP4 numerical, receive-count,
workspace-reuse and CUDA graph checks pass on H20 and H200, including Flash,
MiMo and Pro shapes and both push-dispatch barrier modes. H200 EP8 regression
outputs match the original NVFP4 fused/split kernels bitwise.
The benchmark defaults to warm-cache repeated calls; pass
--cold-l2and--repeats 1for one cold-L2 invocation per sample.🔍 Related Issues
None linked.
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).CPU targets: 32 passed. GPU targets: 8 passed on H200 (EP4/EP8),
6 passed on H20 (EP4). Tests cover dense-reference numerics, receive counts,
workspace reuse, graph replay, production shapes and push-dispatch barrier modes.
H200 NVFP4 EP8 source comparisons are bitwise exact.
🔬 Experimental Track
flashinfer/experimental/and/or an@flashinfer_experimental_api. Tracking issue: #tests/experimental/and were validated on the intended hardware; a runnable example is included.flashinfer/aot.py, and no experimental backend is reachable frombackend="auto"withoutFLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an@flashinfer_experimental_apior naming a backend explicitly is itself the opt-in and needs no environment variable.)Reviewer Notes
Reviewer Notes
The test checkbox covers PR-related targets; upstream multi-architecture CI
requires authorization. EP4 uses the generic fused schedule and has not been
tuned independently. This run remeasured EP8 performance on H200; H20 had four
free GPUs and was used for EP4 correctness verification. The previously accepted
H20 NVFP4 hot-routing M1 deviation matches the original EP8 output bitwise.
The H200 timing point marked Pending is not a settled performance result.
Root README is unchanged.
Summary by CodeRabbit
New Features
Tests