fix(moe_ep): preserve singleton expert TMA modes - #4296
Merged
Merged
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
📝 WalkthroughWalkthroughThe NVFP4 MoE path now supports singleton-local-expert weights with dynamically compact tensor shapes. FC1 and FC2 TMA layouts preserve runtime expert extents. The kernel/reference test adds configurable expert and top-k cases. ChangesNVFP4 runtime expert shapes
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
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 |
cherichy
marked this pull request as ready for review
July 31, 2026 11:20
cherichy
requested review from
Anerudhan,
Aneureka,
IwakuraRein,
aleozlx,
feih-nv,
jiahanc,
mhoqueanik,
nv-yunzheq,
samuellees and
yzh119
as code owners
July 31, 2026 11:20
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.py`:
- Around line 329-341: The test test_nvfp4_kernel_matches_torch_reference
currently uses only the arch_blackwell guard, which permits unsupported SM100
environments with CUDA older than 12.8. Replace or augment that marker with the
repository’s flashinfer.utils backend_requirement capability marker, specifying
the required Blackwell architecture and CUDA version, while preserving the
existing parameterization and test behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: c7054a26-d125-4f77-87fe-4263291cf5c7
📥 Commits
Reviewing files that changed from the base of the PR and between 668a1ba and 4d5f6a5c880d42c21c1c1983161d473f2ac01e81.
📒 Files selected for processing (3)
flashinfer/moe_ep/kernel_src/cutedsl_megamoe/shim/nvfp4.pyflashinfer/moe_ep/kernel_src/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.pytests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.py
Comment on lines
+329
to
+341
| "hidden,intermediate,num_experts,topk", | ||
| [ | ||
| (2048, 1024), | ||
| pytest.param(2048, 1024, 4, 4, id="regular-e4"), | ||
| # 128-misaligned (hidden % 128 == 64): exercises the ceil-div K-tail | ||
| # and predicated epilogue paths the %64 validation relaxation opened | ||
| # up (gpt-oss-120b geometry class). | ||
| (2880, 2880), | ||
| pytest.param(2880, 2880, 4, 4, id="tail-e4"), | ||
| pytest.param(2048, 1024, 1, 1, id="singleton-e1"), | ||
| ], | ||
| ) | ||
| def test_nvfp4_kernel_matches_torch_reference(monkeypatch, hidden, intermediate): | ||
| def test_nvfp4_kernel_matches_torch_reference( | ||
| monkeypatch, hidden, intermediate, num_experts, topk | ||
| ): |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
ast-grep outline flashinfer/utils.py --type function --match is_sm100a_supported --view expanded
rg -n -C 3 'def is_sm100a_supported|is_sm100a_supported\(' \
flashinfer/utils.py \
tests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.pyRepository: flashinfer-ai/flashinfer
Length of output: 657
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
sed -n '320,350p' tests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.py
rg -n -C 5 'torch\.cuda\.get_device_capability|backend_requirement|is_sm100a_supported|nvfp4_mega_moe' tests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.py flashinfer/utils.py docs docs/code_review_guidance.md 2>/dev/null || true
python3 - <<'PY'
from pathlib import Path
import re
path = Path('tests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.py')
text = path.read_text()
imports_end = text.index('\n\n', text.index('import torch'))
top = text.splitlines()[:120]
for i, line in enumerate(top, 1):
print(f'{i}: {line}')
PYRepository: flashinfer-ai/flashinfer
Length of output: 18930
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
sed -n '595,622p' flashinfer/utils.py
rg -n -C 3 'arch_blackwell|`@pytest`\.mark\.arch_blackwell|backend_requirement|supported_compute_capability' tests moe_ep flashinfer 2>/dev/null || true
sed -n '342,370p' tests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.pyRepository: flashinfer-ai/flashinfer
Length of output: 50381
Use the backend_requirement architecture guard for this test.
@pytest.mark.arch_blackwell only checks the major compute capability and the conftest only skips non-Blackwell devices, so this still allows SM100 runs on CUDA < 12.8. Use a flashinfer.utils support check/capability marker that includes the CUDA version requirement instead.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.py` around lines 329 -
341, The test test_nvfp4_kernel_matches_torch_reference currently uses only the
arch_blackwell guard, which permits unsupported SM100 environments with CUDA
older than 12.8. Replace or augment that marker with the repository’s
flashinfer.utils backend_requirement capability marker, specifying the required
Blackwell architecture and CUDA version, while preserving the existing
parameterization and test behavior.
Sources: Coding guidelines, Learnings
Collaborator
|
/bot run tests/moe_op |
Collaborator
Collaborator
|
[FAILED] Pipeline #60465367 — 5/18 executed test jobs passed Compared with nightly #60101730. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
Collaborator
|
/bot run tests/moe_ep |
Collaborator
Collaborator
|
[SUCCESS] Pipeline #61378845: 18/18 executed test jobs passed |
cherichy
force-pushed
the
yuxic/fix_megamoe_e_local_1
branch
from
August 8, 2026 06:13
4d5f6a5 to
c60aa11
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/nvfp4.py (1)
777-795: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the singleton-only dynamic-layout choice.
The new compact-shape path changes the JIT layout contract. Add a comment that identifies mode
0as the expert/TMA mode, explains whytensor.dim_order()is used for the transposed weight views, and states why other modes remain static.Suggested comment
cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) for mode in dynamic_compact_shape_modes: + # Mode 0 is the expert/TMA basis for singleton weights. Preserve + # the physical order of the transposed view. Do not make hidden + # or intermediate modes dynamic; that would reduce specialization. cute_tensor = cute_tensor.mark_compact_shape_dynamic(As per coding guidelines, performance-critical hot paths must document the rationale for special algorithmic choices and potential alternatives.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/nvfp4.py` around lines 777 - 795, Add a concise comment in the conversion logic around dynamic_compact_shape_modes explaining that mode 0 is the expert/TMA mode, tensor.dim_order() preserves the correct layout for transposed weight views, and all other modes remain static because only this singleton dynamic-layout path is supported. Keep the implementation unchanged.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/nvfp4.py`:
- Around line 777-795: Add a concise comment in the conversion logic around
dynamic_compact_shape_modes explaining that mode 0 is the expert/TMA mode,
tensor.dim_order() preserves the correct layout for transposed weight views, and
all other modes remain static because only this singleton dynamic-layout path is
supported. Keep the implementation unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 8f456c99-f136-4cd4-83ea-738dfae8e5c1
📒 Files selected for processing (3)
flashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/shim/nvfp4.pyflashinfer/moe_ep/kernel_src/sm100/cutedsl_megamoe/src/moe_nvfp4_swapab/kernel_fc12.pytests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/moe_ep/test_nvfp4_cutedsl_kernel_vs_reference.py
jiahanc
enabled auto-merge (squash)
August 8, 2026 06:48
mhoqueanik
added a commit
to mhoqueanik/flashinfer-moe_ep
that referenced
this pull request
Aug 13, 2026
Aligns the restructure branch with flashinfer TOT (95 commits, incl. the v0.6.17 release line). moe_ep conflict resolution: - All taxonomy-vs-flat naming/import conflicts (backends, tests, docs, configs, run_tests.sh, pre-commit/pyproject excludes) resolve to the restructure spellings; upstream's non-conflicting deltas auto-merge. - kernel_src/sm100/cutedsl_megamoe (old path, re-added by upstream's rename detection) is dissolved into kernel_src/cutedsl_megamoe; the one real upstream advance there — the 4fbac49 singleton-expert TMA-modes fix (PR flashinfer-ai#4296: shim/nvfp4.py + src/moe_nvfp4_swapab/kernel_fc12.py) — is ported onto the renamed paths and recorded in VENDOR.md. - kernel_src/sm90/pull_style_cutedsl_megakernel add/add conflicts take our tree verbatim (upstream b1d9585 is an older squash of the same drop; no upstream advances since). - tune.py keeps the restructure layout (upstream's _restage_skewed_routing hunk is our own pre-restructure code, now in backends/mega/kernel/tuning.py). - De-duplicated the mxfp8 multirank torch-oracle block the auto-merge re-added from upstream's copy (ours carries the try/finally symm-buffer cleanup from 640b75f). Notable upstream picks now in-tree: BootstrapConfig.device (3c57ef1, PR flashinfer-ai#4348 — resolves the vLLM smoke blocker recorded in TODO_sm100_next_resync.md) and the E_local=1 nvfp4 oracle regression test. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
aleozlx
pushed a commit
that referenced
this pull request
Aug 14, 2026
…ush-style FP8 backend; sync CuTe-DSL 4.7 quant-staging fix (#4449) ## Summary Three things: a layout/naming refactor of `flashinfer.moe_ep`'s mega-kernel layer, the incorporation of the SM90 push-style FP8 backend (#4069, since merged upstream) as the first new backend added in the restructured shape, and one vendored-kernel sync that fixes the fused activation-quant staging crash on CuTe-DSL 4.7 — un-blocking 4.7.x and lifting the temporary `==4.6.1` pin. The branch is merged up to upstream/main (2febce5, past the v0.6.17 line and the #4069 squash). The refactor organizes the layer around two orthogonal views: 1. **Taxonomy (user view)** — backends move to `backends/mega/kernel/sm<arch>/<act_dtype>_<weight_dtype>_<out_dtype>_<kernel_style>/`, and registry `kernel_name` strings plus config classes carry the same fully-qualified names. One glance at a name now tells you the architecture, the activation/weight/output dtypes, and the kernel style: | old kernel_name | new kernel_name | new config class | |---|---|---| | `deep_gemm_mega` | `sm100_fp8_fp4_bf16_deepgemm` | `Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig` | | `nvfp4_cutedsl` | `sm100_nvfp4_nvfp4_bf16_cutedsl` | `Sm100_Nvfp4_Nvfp4_Bf16_Cutedsl_MegaMoeConfig` | | `mxfp8_cutedsl` | `sm100_mxfp8_mxfp8_bf16_cutedsl` | `Sm100_Mxfp8_Mxfp8_Bf16_Cutedsl_MegaMoeConfig` | | `sm90_pull_fp8` | `sm90_fp8_fp8_bf16_pull_cutedsl` | `Sm90_Fp8_Fp8_Bf16_PullCutedsl_MegaMoeConfig` | | `sm90_push_fp8` (new, from #4069) | `sm90_fp8_fp8_bf16_push_cuda` | `Sm90_Fp8_Fp8_Bf16_PushCuda_MegaMoeConfig` | Naming conventions: deep_gemm dtypes are plain `fp8`/`fp4`, matching upstream `deep_gemm.fp8_fp4_mega_moe`; the mx/nv prefixes are reserved for the cutedsl kernels' block-scaled formats. Output dtype is always bf16 — nvfp4's `combine_dtype` is comm-wire compression, not an output format. 2. **Provenance (kernel-dev view)** — vendored kernel sources in `kernel_src/` are keyed by upstream repo snapshot, not by architecture: `kernel_src/sm100/cutedsl_megamoe` moves to `kernel_src/cutedsl_megamoe` (the mother repo ships kernels for multiple arches, so an smXX level misrepresents it). Each drop mirrors the vendor repo layout — `src/` byte-for-byte upstream, all adaptation in `shim/` — and gains a `VENDOR.md` recording upstream repo/commit/sync state and pending local diffs. A new `kernel_src/README.md` states the contract explicitly: **no edits to `src/` of any kind — including docstrings, comments, and lint fixes**; tool warnings against vendored files (docstring-coverage gates, review bots) are handled by excluding the path, never by editing the file. The sm90 fork trees (`kernel_src/sm90/pull_style_cutedsl_megakernel` from #4113, `kernel_src/sm90/push_style_megamoe` from #4069) intentionally stay separate snapshots — one kernel_src dir = one upstream commit — and fold into the mother drop if/when upstream merges them. **Why:** verbatim snapshots must stay diffable against one upstream commit, and splitting vendored trees per-dtype or per-arch breaks re-sync; meanwhile users navigate by architecture and dtype, not by which vendor repo a kernel came from. Putting each concern where its audience looks resolves the tension. The layout rule is documented in `docs/design_docs/moe_ep_architecture.md`, and it is what makes new backend families routine — demonstrated in this very PR by the SM90 push-style incorporation below, and next by the follow-up backend-family PRs (SM100 BF16 #4386, SM120 MXFP8). ## Directories affected All changes live under `flashinfer/moe_ep/` plus its tests and docs: - `backends/mega/kernel/sm100/{fp8_fp4_bf16_deepgemm,nvfp4_nvfp4_bf16_cutedsl,mxfp8_mxfp8_bf16_cutedsl}/` and `backends/mega/kernel/sm90/{fp8_fp8_bf16_pull_cutedsl,fp8_fp8_bf16_push_cuda}/` — taxonomy backend wrappers (moved/renamed; push_cuda is new). - `kernel_src/cutedsl_megamoe/` (moved from `kernel_src/sm100/cutedsl_megamoe/`), `kernel_src/sm90/pull_style_cutedsl_megakernel/`, `kernel_src/sm90/push_style_megamoe/` (new) — provenance-keyed vendored drops, each with `VENDOR.md`; new `kernel_src/README.md` states the no-edits contract. - `backends/mega/kernel/tuning.py` + per-backend `tuner.py` files — tuning machinery moved out of `tune.py` (now a CLI shim). - `core/kernel/registry.py`, `moe_ep/__init__.py` — deprecated-alias resolution and re-exports. - `tests/moe_ep/`, `docs/design_docs/moe_ep_{architecture,runbook}.md`, `pyproject.toml`/`.pre-commit-config.yaml` excludes, `run_tests.sh` (new 2-GPU `sm90_push` target). ## Test results - **Full `run_tests.sh` matrix — all 12 targets green** on 4xH100 (job 2389821, 2026-08-13), including the new `sm90_push` Hopper target and the fault-tolerance suites after the deadlock fixes. - **B200** (jobs 2388315/2388326): registry/alias smoke, deprecated aliases, unit x3 green — 396 passed / 72 skipped (push cpu/packaging/contract tests run; Hopper-marked kernel tests skip). - **Unit target re-validated green** after the second upstream merge (job 2389880) and again after the round-2 CodeRabbit fixes (job 2389916), same 396/72 counts, B200. - **8x B200** (jobs 2384640/2384641/2384650): quant-staging sync matrix fully green on both dsl 4.6.1 and 4.7.0 (details in the vendored-sync section below). - **GB200 + B200**: mxfp8/nvfp4 multirank oracle suites with the per-cell tolerance band. - Microbenchmark re-run: no regressions vs pre-restructure reference numbers (deep_gemm parity; cutedsl kernels at or above their previous points). - `pre-commit run -a` fully green at the branch head (e9f791a). ## SM90 push-style FP8 backend (incorporates #4069) Ports #4069 (head 301f8ce; since merged to main as f9b13ef — re-diffed, byte-identical, no post-review deltas) onto the taxonomy/provenance layout, serving as the first proof of the "one taxonomy backend dir + one provenance-keyed kernel drop" recipe: - **`kernel_src/sm90/push_style_megamoe/`** — verbatim byte-for-byte drop from the PR head (`src/{a2a,fp8_gemm}` CUDA sources, `shim/`, ACKNOWLEDGEMENT.md) plus a `VENDOR.md` provenance record. - **`backends/mega/kernel/sm90/fp8_fp8_bf16_push_cuda/`** — the five wrapper files relocated from upstream's flat `kernel/sm90_push_fp8/`, config renamed to `Sm90_Fp8_Fp8_Bf16_PushCuda_MegaMoeConfig`, registered with `deprecated_aliases=("sm90_push_fp8",)`. - **Core deltas carried from the PR:** `mega_layer.py` allocates the output before `stage_inputs`; pyproject package-data ships the drop's `.cu`/`.cuh` for non-editable installs; the `isolated_deep_gemm_cache` conftest fixture; the mega-layer allocation-order regression test. - **Tests:** the nine sm90_push_fp8 test files (names kept to minimize re-sync friction) rewritten to the taxonomy. Deviation from upstream: `run_tests.sh` exposes `sm90_push` as its own 2-GPU Hopper target instead of folding it into multirank — on non-Hopper nodes the arch-marked files collect 0 tests and torchrun turns pytest exit 5 into a failure. ## CuTe-DSL 4.7 quant-staging fix (vendored sync) The `CUDA_ERROR_MISALIGNED_ADDRESS` crash on cutlass-dsl 4.7.0 — which presented as a deep_gemm mega multirank failure — was root-caused to the **fused bf16→quantized activation staging** (`DataPreprocess` in the vendored cutedsl_megamoe tree), which every mega staging path shares, deep_gemm's included. The kernel team's fix is synced in as a single-file partial re-sync per the vendoring policy: - `kernel_src/cutedsl_megamoe/src/src/inputs_process.py` + `src/common/host_utils.py` taken **verbatim** from upstream `bangyus/cutedsl_megamoe @ 50117315d`, recorded in `VENDOR.md` under pending-diffs (resolves at the next full re-sync). The mxfp8 quant kernel is reworked so each lane owns one contiguous 16-byte fp8 store (adjacent lanes reduce the 32-element block amax via `shuffle_sync_bfly`, even lane writes the E8M0 scale), and `__init__` gains a hidden-size row-alignment guard. - Also fixes a stale pre-commit exclude left by the directory move (`kernel_src/sm100/cutedsl_megamoe` → `kernel_src/cutedsl_megamoe`) so hooks stop reformatting the verbatim `src/` tree. Validated on 8x B200 (jobs 2384640/2384641/2384650), full matrix green on **both** DSL versions: | section | dsl 4.6.1 | dsl 4.7.0 | |---|---|---| | drop's own harness (`python -m src.inputs_process`: bit-exact scales + SNR vs reference, nvfp4 offline/online + mxfp8) | 3/3 | 3/3 | | `test_fused_quant_stage.py` | 11/11 | 11/11 | | mega multirank x4 ranks (deep_gemm + nvfp4 + mxfp8) | 20/rank | 20/rank | | single-rank kernel-vs-reference oracles | 6/6 | 6/6 | The deep_gemm multirank suite previously crashed deterministically on 4.7.0; it now passes there. On the strength of this, the runbook's temporary `==4.6.1` pin is lifted (see the DSL guidance bullet below). ## Also in this PR - **Per-backend tuners.** `flashinfer/moe_ep/tune.py` becomes a pure CLI shim (surface unchanged: `python -m flashinfer.moe_ep.tune`); dtype-specific tuning moves into the backends (`sm100/{nvfp4,mxfp8}.../tuner.py`), shared sweep machinery (dist lifecycle, skewed restage, schedule grid, timed sweep tail) into `backends/mega/kernel/tuning.py`. - **CUTLASS DSL guidance updated (pin lifted).** The test-container recipe briefly carried a hard `nvidia-cutlass-dsl==4.6.1` pin because 4.7.0 crashed the mega multirank path; with the crash root-caused and fixed above, the runbook now allows `-U` installs again. 4.6.1 remains the perf-validated reference (pin it when producing numbers meant to compare against the TUNING.md tables); 4.7.0 is correctness-validated. The library's supported floor remains 4.5.2 (the MR!27 WAR already in main). - **Per-cell bf16 term-magnitude tolerance band** for the mxfp8 multirank oracle compares — a principled per-cell bound derived from the bf16 accumulation term magnitudes, replacing the global rtol that produced rare single-cell false failures. Validated on GB200 and B200. - **One-direction import layering rules** codified in the architecture doc, with all `cutedsl_megamoe` access routed through the drop's `__init__` rather than deep-path imports. ## Merge with upstream/main and follow-up fixes The branch is merged up to upstream/main in two steps. First to aaf97df (95 commits, incl. the v0.6.17 release line): conflict resolution keeps the restructure spellings everywhere; upstream's one real kernel advance in the moved tree — the 4fbac49 singleton-expert TMA-modes fix (#4296) — is ported onto the renamed paths and recorded in `VENDOR.md`. Notable upstream picks now in-tree: `BootstrapConfig.device` (#4348) and the E_local=1 nvfp4 oracle regression test. Second merge to 2febce5 (13 commits), resolving the conflicts created when #4069 itself squash-merged upstream (f9b13ef) with the same moe_ep files in the pre-restructure flat layout. Every conflict resolves to the taxonomy spellings (upstream's side is the flat spelling of content this branch already carries); upstream's flat `backends/mega/kernel/sm90_push_fp8/` wrapper and its re-folding of `sm90_push` into the multirank target are dropped in favor of this branch's layout. The vendored push drop was re-diffed against the merged SHA: byte-for-byte identical, no post-review deltas (recorded in `VENDOR.md`). Post-merge hardening found and fixed by full-suite runs: - **Merge fallout:** auto-merged regions had re-introduced pre-restructure `kernel_src.sm100.cutedsl_megamoe` spellings in 12 files, silently skipping entire GPU test files via `importorskip`; restored, and upstream's re-added flat `sm90_pull_fp8/` wrapper removed. - **FT test deadlocks (4xH100):** the fault-tolerance multirank test's evicted victim ran a collective `destroy()` against the survivors' barrier sequence, deadlocking until the NCCL watchdog — the victim tail now mirrors the survivors' barrier→destroy→barrier shape. The FT smoke's survivors now keep forwarding past the kill window so they actually observe the fault, and `run_tests.sh` judges the smoke by counting `SMOKE_RESULT` markers (torchrun interleaves lines). - **Unit-suite crasher isolation:** the long-known in-suite-only interpreter abort (heap corruption accumulating over the ~200-test single-process run, firing during a plain module import or in CPython teardown) is worked around by running the trigger test in its own pytest process and exiting the unit invocations via `os._exit(pytest_rc)`; rationale in the runbook, root cause tracked (needs ASAN). All tests pass — this is process-teardown hygiene, not a kernel bug. **CodeRabbit review responses.** Two rounds of actionable findings are fixed in-branch (640b75f, 57926a9) — highlights from round 2: the push packaging test's import-boundary gate was building the pre-taxonomy flat backend path and passing vacuously (fixed, now validates all 5 wrapper files); the test baseline's weight cache gains weakref eviction; `cutedsl_megamoe/shim/__main__.py` added so the documented `python -m ...shim` commands resolve; the cutedsl_megamoe `VENDOR.md` provenance TODOs are filled. Findings inside verbatim-vendored `kernel_src/**/src/` trees are deliberately not patched locally — they route upstream per the vendoring policy in `kernel_src/README.md`. **Lint.** `pre-commit run -a` is fully green (clang-format, mypy, ruff check/format, whitespace hooks). The final e9f791a is a pure ruff-format pass over 13 moe_ep files — line wraps where the longer taxonomy class names pushed calls past the limit. The vendored `src/` trees are untouched by hooks (the exclude set holds). ## Backward compatibility External callers keep working unchanged — both the old config-class names and the old kernel_name strings remain as deprecated aliases: - **Config classes**: `DeepGemmMegaMoeConfig`, `Nvfp4CutedslMegaMoeConfig`, `Mxfp8CutedslMegaMoeConfig`, `Sm90PullFp8MegaMoeConfig`, and `Sm90PushFp8MegaMoeConfig` are plain aliases of the new `Sm<arch>_..._MegaMoeConfig` classes, defined (with a removal note) in `flashinfer/moe_ep/__init__.py` right below the taxonomy imports, and still exported via `__all__`. - **Registry kernel_name strings**: `deep_gemm_mega`, `nvfp4_cutedsl`, `mxfp8_cutedsl`, `sm90_pull_fp8`, and `sm90_push_fp8` resolve to the taxonomy backends through the `deprecated_aliases=` parameter of each backend's `@register_mega_kernel(...)` decoration; the resolution machinery lives in `flashinfer/moe_ep/core/kernel/registry.py`. Using one emits a `DeprecationWarning`, and aliases are excluded from the available-kernels listing. - Both alias families WILL BE REMOVED in a future release (noted at both locations above). ## Testing - Directory moves and renames are behavior-preserving by construction; registry tests exercise both the taxonomy names and the deprecated aliases (alias use warns; the kernel listing shows taxonomy names only). - Full `run_tests.sh` matrix (all 12 targets) green on 4xH100 (job 2389821); B200 unit/registry/alias validation (jobs 2388315/2388326) — see Test results above. - The quant-staging sync validated on both dsl 4.6.1 and 4.7.0 (matrix above); mxfp8/nvfp4 multirank oracle suites validated on GB200 and B200. - The standalone MoE-EP microbenchmark was re-run against this branch with no regressions vs the pre-restructure reference numbers (deep_gemm parity; cutedsl kernels at or above their previous points). ## Relation to other PRs Re-layering on top of #4113 (SM90 pull-style FP8 backend, merged) and incorporating #4069 (SM90 push-style FP8 backend, merged upstream 2026-08-12; the vendored drop was re-diffed against the merged SHA f9b13ef and is byte-identical). This is the base branch for the upcoming backend-family PRs — SM100 BF16 (#4386) and SM120 MXFP8 — each of which adds one taxonomy backend directory plus one provenance-keyed kernel drop in the shape this restructure establishes. Both follow-up branches are already rebased onto this branch's head (unit target green on each), so they apply as exactly their backend-specific commits once this merges. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Md Anik <mhoqueanik@cw-dfw-cs-001-login-01.cm.cluster> Co-authored-by: Md Saidul Hoque Anik <mhoqueanik@login-preos01.a51.clusters.nvidia.com>
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
…ush-style FP8 backend; sync CuTe-DSL 4.7 quant-staging fix (flashinfer-ai#4449) ## Summary Three things: a layout/naming refactor of `flashinfer.moe_ep`'s mega-kernel layer, the incorporation of the SM90 push-style FP8 backend (flashinfer-ai#4069, since merged upstream) as the first new backend added in the restructured shape, and one vendored-kernel sync that fixes the fused activation-quant staging crash on CuTe-DSL 4.7 — un-blocking 4.7.x and lifting the temporary `==4.6.1` pin. The branch is merged up to upstream/main (2febce5, past the v0.6.17 line and the flashinfer-ai#4069 squash). The refactor organizes the layer around two orthogonal views: 1. **Taxonomy (user view)** — backends move to `backends/mega/kernel/sm<arch>/<act_dtype>_<weight_dtype>_<out_dtype>_<kernel_style>/`, and registry `kernel_name` strings plus config classes carry the same fully-qualified names. One glance at a name now tells you the architecture, the activation/weight/output dtypes, and the kernel style: | old kernel_name | new kernel_name | new config class | |---|---|---| | `deep_gemm_mega` | `sm100_fp8_fp4_bf16_deepgemm` | `Sm100_Fp8_Fp4_Bf16_Deepgemm_MegaMoeConfig` | | `nvfp4_cutedsl` | `sm100_nvfp4_nvfp4_bf16_cutedsl` | `Sm100_Nvfp4_Nvfp4_Bf16_Cutedsl_MegaMoeConfig` | | `mxfp8_cutedsl` | `sm100_mxfp8_mxfp8_bf16_cutedsl` | `Sm100_Mxfp8_Mxfp8_Bf16_Cutedsl_MegaMoeConfig` | | `sm90_pull_fp8` | `sm90_fp8_fp8_bf16_pull_cutedsl` | `Sm90_Fp8_Fp8_Bf16_PullCutedsl_MegaMoeConfig` | | `sm90_push_fp8` (new, from flashinfer-ai#4069) | `sm90_fp8_fp8_bf16_push_cuda` | `Sm90_Fp8_Fp8_Bf16_PushCuda_MegaMoeConfig` | Naming conventions: deep_gemm dtypes are plain `fp8`/`fp4`, matching upstream `deep_gemm.fp8_fp4_mega_moe`; the mx/nv prefixes are reserved for the cutedsl kernels' block-scaled formats. Output dtype is always bf16 — nvfp4's `combine_dtype` is comm-wire compression, not an output format. 2. **Provenance (kernel-dev view)** — vendored kernel sources in `kernel_src/` are keyed by upstream repo snapshot, not by architecture: `kernel_src/sm100/cutedsl_megamoe` moves to `kernel_src/cutedsl_megamoe` (the mother repo ships kernels for multiple arches, so an smXX level misrepresents it). Each drop mirrors the vendor repo layout — `src/` byte-for-byte upstream, all adaptation in `shim/` — and gains a `VENDOR.md` recording upstream repo/commit/sync state and pending local diffs. A new `kernel_src/README.md` states the contract explicitly: **no edits to `src/` of any kind — including docstrings, comments, and lint fixes**; tool warnings against vendored files (docstring-coverage gates, review bots) are handled by excluding the path, never by editing the file. The sm90 fork trees (`kernel_src/sm90/pull_style_cutedsl_megakernel` from flashinfer-ai#4113, `kernel_src/sm90/push_style_megamoe` from flashinfer-ai#4069) intentionally stay separate snapshots — one kernel_src dir = one upstream commit — and fold into the mother drop if/when upstream merges them. **Why:** verbatim snapshots must stay diffable against one upstream commit, and splitting vendored trees per-dtype or per-arch breaks re-sync; meanwhile users navigate by architecture and dtype, not by which vendor repo a kernel came from. Putting each concern where its audience looks resolves the tension. The layout rule is documented in `docs/design_docs/moe_ep_architecture.md`, and it is what makes new backend families routine — demonstrated in this very PR by the SM90 push-style incorporation below, and next by the follow-up backend-family PRs (SM100 BF16 flashinfer-ai#4386, SM120 MXFP8). ## Directories affected All changes live under `flashinfer/moe_ep/` plus its tests and docs: - `backends/mega/kernel/sm100/{fp8_fp4_bf16_deepgemm,nvfp4_nvfp4_bf16_cutedsl,mxfp8_mxfp8_bf16_cutedsl}/` and `backends/mega/kernel/sm90/{fp8_fp8_bf16_pull_cutedsl,fp8_fp8_bf16_push_cuda}/` — taxonomy backend wrappers (moved/renamed; push_cuda is new). - `kernel_src/cutedsl_megamoe/` (moved from `kernel_src/sm100/cutedsl_megamoe/`), `kernel_src/sm90/pull_style_cutedsl_megakernel/`, `kernel_src/sm90/push_style_megamoe/` (new) — provenance-keyed vendored drops, each with `VENDOR.md`; new `kernel_src/README.md` states the no-edits contract. - `backends/mega/kernel/tuning.py` + per-backend `tuner.py` files — tuning machinery moved out of `tune.py` (now a CLI shim). - `core/kernel/registry.py`, `moe_ep/__init__.py` — deprecated-alias resolution and re-exports. - `tests/moe_ep/`, `docs/design_docs/moe_ep_{architecture,runbook}.md`, `pyproject.toml`/`.pre-commit-config.yaml` excludes, `run_tests.sh` (new 2-GPU `sm90_push` target). ## Test results - **Full `run_tests.sh` matrix — all 12 targets green** on 4xH100 (job 2389821, 2026-08-13), including the new `sm90_push` Hopper target and the fault-tolerance suites after the deadlock fixes. - **B200** (jobs 2388315/2388326): registry/alias smoke, deprecated aliases, unit x3 green — 396 passed / 72 skipped (push cpu/packaging/contract tests run; Hopper-marked kernel tests skip). - **Unit target re-validated green** after the second upstream merge (job 2389880) and again after the round-2 CodeRabbit fixes (job 2389916), same 396/72 counts, B200. - **8x B200** (jobs 2384640/2384641/2384650): quant-staging sync matrix fully green on both dsl 4.6.1 and 4.7.0 (details in the vendored-sync section below). - **GB200 + B200**: mxfp8/nvfp4 multirank oracle suites with the per-cell tolerance band. - Microbenchmark re-run: no regressions vs pre-restructure reference numbers (deep_gemm parity; cutedsl kernels at or above their previous points). - `pre-commit run -a` fully green at the branch head (e9f791a). ## SM90 push-style FP8 backend (incorporates flashinfer-ai#4069) Ports flashinfer-ai#4069 (head 301f8ce; since merged to main as f9b13ef — re-diffed, byte-identical, no post-review deltas) onto the taxonomy/provenance layout, serving as the first proof of the "one taxonomy backend dir + one provenance-keyed kernel drop" recipe: - **`kernel_src/sm90/push_style_megamoe/`** — verbatim byte-for-byte drop from the PR head (`src/{a2a,fp8_gemm}` CUDA sources, `shim/`, ACKNOWLEDGEMENT.md) plus a `VENDOR.md` provenance record. - **`backends/mega/kernel/sm90/fp8_fp8_bf16_push_cuda/`** — the five wrapper files relocated from upstream's flat `kernel/sm90_push_fp8/`, config renamed to `Sm90_Fp8_Fp8_Bf16_PushCuda_MegaMoeConfig`, registered with `deprecated_aliases=("sm90_push_fp8",)`. - **Core deltas carried from the PR:** `mega_layer.py` allocates the output before `stage_inputs`; pyproject package-data ships the drop's `.cu`/`.cuh` for non-editable installs; the `isolated_deep_gemm_cache` conftest fixture; the mega-layer allocation-order regression test. - **Tests:** the nine sm90_push_fp8 test files (names kept to minimize re-sync friction) rewritten to the taxonomy. Deviation from upstream: `run_tests.sh` exposes `sm90_push` as its own 2-GPU Hopper target instead of folding it into multirank — on non-Hopper nodes the arch-marked files collect 0 tests and torchrun turns pytest exit 5 into a failure. ## CuTe-DSL 4.7 quant-staging fix (vendored sync) The `CUDA_ERROR_MISALIGNED_ADDRESS` crash on cutlass-dsl 4.7.0 — which presented as a deep_gemm mega multirank failure — was root-caused to the **fused bf16→quantized activation staging** (`DataPreprocess` in the vendored cutedsl_megamoe tree), which every mega staging path shares, deep_gemm's included. The kernel team's fix is synced in as a single-file partial re-sync per the vendoring policy: - `kernel_src/cutedsl_megamoe/src/src/inputs_process.py` + `src/common/host_utils.py` taken **verbatim** from upstream `bangyus/cutedsl_megamoe @ 50117315d`, recorded in `VENDOR.md` under pending-diffs (resolves at the next full re-sync). The mxfp8 quant kernel is reworked so each lane owns one contiguous 16-byte fp8 store (adjacent lanes reduce the 32-element block amax via `shuffle_sync_bfly`, even lane writes the E8M0 scale), and `__init__` gains a hidden-size row-alignment guard. - Also fixes a stale pre-commit exclude left by the directory move (`kernel_src/sm100/cutedsl_megamoe` → `kernel_src/cutedsl_megamoe`) so hooks stop reformatting the verbatim `src/` tree. Validated on 8x B200 (jobs 2384640/2384641/2384650), full matrix green on **both** DSL versions: | section | dsl 4.6.1 | dsl 4.7.0 | |---|---|---| | drop's own harness (`python -m src.inputs_process`: bit-exact scales + SNR vs reference, nvfp4 offline/online + mxfp8) | 3/3 | 3/3 | | `test_fused_quant_stage.py` | 11/11 | 11/11 | | mega multirank x4 ranks (deep_gemm + nvfp4 + mxfp8) | 20/rank | 20/rank | | single-rank kernel-vs-reference oracles | 6/6 | 6/6 | The deep_gemm multirank suite previously crashed deterministically on 4.7.0; it now passes there. On the strength of this, the runbook's temporary `==4.6.1` pin is lifted (see the DSL guidance bullet below). ## Also in this PR - **Per-backend tuners.** `flashinfer/moe_ep/tune.py` becomes a pure CLI shim (surface unchanged: `python -m flashinfer.moe_ep.tune`); dtype-specific tuning moves into the backends (`sm100/{nvfp4,mxfp8}.../tuner.py`), shared sweep machinery (dist lifecycle, skewed restage, schedule grid, timed sweep tail) into `backends/mega/kernel/tuning.py`. - **CUTLASS DSL guidance updated (pin lifted).** The test-container recipe briefly carried a hard `nvidia-cutlass-dsl==4.6.1` pin because 4.7.0 crashed the mega multirank path; with the crash root-caused and fixed above, the runbook now allows `-U` installs again. 4.6.1 remains the perf-validated reference (pin it when producing numbers meant to compare against the TUNING.md tables); 4.7.0 is correctness-validated. The library's supported floor remains 4.5.2 (the MR!27 WAR already in main). - **Per-cell bf16 term-magnitude tolerance band** for the mxfp8 multirank oracle compares — a principled per-cell bound derived from the bf16 accumulation term magnitudes, replacing the global rtol that produced rare single-cell false failures. Validated on GB200 and B200. - **One-direction import layering rules** codified in the architecture doc, with all `cutedsl_megamoe` access routed through the drop's `__init__` rather than deep-path imports. ## Merge with upstream/main and follow-up fixes The branch is merged up to upstream/main in two steps. First to aaf97df (95 commits, incl. the v0.6.17 release line): conflict resolution keeps the restructure spellings everywhere; upstream's one real kernel advance in the moved tree — the 4fbac49 singleton-expert TMA-modes fix (flashinfer-ai#4296) — is ported onto the renamed paths and recorded in `VENDOR.md`. Notable upstream picks now in-tree: `BootstrapConfig.device` (flashinfer-ai#4348) and the E_local=1 nvfp4 oracle regression test. Second merge to 2febce5 (13 commits), resolving the conflicts created when flashinfer-ai#4069 itself squash-merged upstream (f9b13ef) with the same moe_ep files in the pre-restructure flat layout. Every conflict resolves to the taxonomy spellings (upstream's side is the flat spelling of content this branch already carries); upstream's flat `backends/mega/kernel/sm90_push_fp8/` wrapper and its re-folding of `sm90_push` into the multirank target are dropped in favor of this branch's layout. The vendored push drop was re-diffed against the merged SHA: byte-for-byte identical, no post-review deltas (recorded in `VENDOR.md`). Post-merge hardening found and fixed by full-suite runs: - **Merge fallout:** auto-merged regions had re-introduced pre-restructure `kernel_src.sm100.cutedsl_megamoe` spellings in 12 files, silently skipping entire GPU test files via `importorskip`; restored, and upstream's re-added flat `sm90_pull_fp8/` wrapper removed. - **FT test deadlocks (4xH100):** the fault-tolerance multirank test's evicted victim ran a collective `destroy()` against the survivors' barrier sequence, deadlocking until the NCCL watchdog — the victim tail now mirrors the survivors' barrier→destroy→barrier shape. The FT smoke's survivors now keep forwarding past the kill window so they actually observe the fault, and `run_tests.sh` judges the smoke by counting `SMOKE_RESULT` markers (torchrun interleaves lines). - **Unit-suite crasher isolation:** the long-known in-suite-only interpreter abort (heap corruption accumulating over the ~200-test single-process run, firing during a plain module import or in CPython teardown) is worked around by running the trigger test in its own pytest process and exiting the unit invocations via `os._exit(pytest_rc)`; rationale in the runbook, root cause tracked (needs ASAN). All tests pass — this is process-teardown hygiene, not a kernel bug. **CodeRabbit review responses.** Two rounds of actionable findings are fixed in-branch (640b75f, 57926a9) — highlights from round 2: the push packaging test's import-boundary gate was building the pre-taxonomy flat backend path and passing vacuously (fixed, now validates all 5 wrapper files); the test baseline's weight cache gains weakref eviction; `cutedsl_megamoe/shim/__main__.py` added so the documented `python -m ...shim` commands resolve; the cutedsl_megamoe `VENDOR.md` provenance TODOs are filled. Findings inside verbatim-vendored `kernel_src/**/src/` trees are deliberately not patched locally — they route upstream per the vendoring policy in `kernel_src/README.md`. **Lint.** `pre-commit run -a` is fully green (clang-format, mypy, ruff check/format, whitespace hooks). The final e9f791a is a pure ruff-format pass over 13 moe_ep files — line wraps where the longer taxonomy class names pushed calls past the limit. The vendored `src/` trees are untouched by hooks (the exclude set holds). ## Backward compatibility External callers keep working unchanged — both the old config-class names and the old kernel_name strings remain as deprecated aliases: - **Config classes**: `DeepGemmMegaMoeConfig`, `Nvfp4CutedslMegaMoeConfig`, `Mxfp8CutedslMegaMoeConfig`, `Sm90PullFp8MegaMoeConfig`, and `Sm90PushFp8MegaMoeConfig` are plain aliases of the new `Sm<arch>_..._MegaMoeConfig` classes, defined (with a removal note) in `flashinfer/moe_ep/__init__.py` right below the taxonomy imports, and still exported via `__all__`. - **Registry kernel_name strings**: `deep_gemm_mega`, `nvfp4_cutedsl`, `mxfp8_cutedsl`, `sm90_pull_fp8`, and `sm90_push_fp8` resolve to the taxonomy backends through the `deprecated_aliases=` parameter of each backend's `@register_mega_kernel(...)` decoration; the resolution machinery lives in `flashinfer/moe_ep/core/kernel/registry.py`. Using one emits a `DeprecationWarning`, and aliases are excluded from the available-kernels listing. - Both alias families WILL BE REMOVED in a future release (noted at both locations above). ## Testing - Directory moves and renames are behavior-preserving by construction; registry tests exercise both the taxonomy names and the deprecated aliases (alias use warns; the kernel listing shows taxonomy names only). - Full `run_tests.sh` matrix (all 12 targets) green on 4xH100 (job 2389821); B200 unit/registry/alias validation (jobs 2388315/2388326) — see Test results above. - The quant-staging sync validated on both dsl 4.6.1 and 4.7.0 (matrix above); mxfp8/nvfp4 multirank oracle suites validated on GB200 and B200. - The standalone MoE-EP microbenchmark was re-run against this branch with no regressions vs the pre-restructure reference numbers (deep_gemm parity; cutedsl kernels at or above their previous points). ## Relation to other PRs Re-layering on top of flashinfer-ai#4113 (SM90 pull-style FP8 backend, merged) and incorporating flashinfer-ai#4069 (SM90 push-style FP8 backend, merged upstream 2026-08-12; the vendored drop was re-diffed against the merged SHA f9b13ef and is byte-identical). This is the base branch for the upcoming backend-family PRs — SM100 BF16 (flashinfer-ai#4386) and SM120 MXFP8 — each of which adds one taxonomy backend directory plus one provenance-keyed kernel drop in the shape this restructure establishes. Both follow-up branches are already rebased onto this branch's head (unit target green on each), so they apply as exactly their backend-specific commits once this merges. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Md Anik <mhoqueanik@cw-dfw-cs-001-login-01.cm.cluster> Co-authored-by: Md Saidul Hoque Anik <mhoqueanik@login-preos01.a51.clusters.nvidia.com>
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
3 participants
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.
Summary
E_local=1,topk=1NVFP4 MegaMoE regression to the existing independent torch-oracle test.Root cause
NVFP4 MegaMoE always supplies
static_expert_shape. When the local expert count is one, static shape refinement rewrites the weight expert mode to a Python1. CuTeDSL 4.6.0 canonicalizes that singleton TMA batch basis before the descriptor is derefined to the kernel ABI, causing a segmentation fault during the first kernel compilation/launch.The frontend now keeps the singleton weight expert mode runtime-dynamic, and the kernel saves that runtime extent before static refinement and reuses it for the FC1 and FC2 weight TMA descriptors. Multi-expert launches follow the existing path unchanged.
Validation
Tested on NVIDIA B200 with CUDA 13 and
nvidia-cutlass-dsl==4.6.0.Before the fix:
E_local=4,topk=4: passes the torch oracle.E_local=1,topk=1: reproducibly segfaults at the first kernel compile/launch after preprocessing, staging, and oracle computation complete.After the fix:
Summary by CodeRabbit