fix(moe): handle CuTe DSL finalize output tails - #4186
Merged
Merged
Conversation
…ut (gh flashinfer-ai#3957) The SM100 CuteDSL MoE epilogues store full CTA-tile rows with no column predicate: the finalize kernel's raw-pointer bulk scatter (cp.reduce.async.bulk ... add) and gemm1's SFC autovec_copy (the epilogue TODO, dating to flashinfer-ai#2398) both write past the real output columns whenever the N-tiling leaves a partial CTA tile or a cluster-padding CTA (the persistent scheduler pads the grid to a cluster multiple with an M-only validity guard). The stray read-modify-writes land in neighboring caching-allocator memory: off the end of the allocation for the last token (IMA), a silent add-of-zero otherwise -- which canonicalizes NaN bit patterns and thereby corrupts integer data reinterpreted as bf16 (the gh flashinfer-ai#3957 gather-assert on trtllm's long-lived permute-index cache, surfacing ~46 configs after the writes in accumulated runs). Fix (stop-the-bleeding layer): - finalize: can_implement requires n % (mma_tiler_n * cluster_n) == 0 (rejects both partial N-tiles and cluster padding along N). - gemm1: can_implement requires n % mma_tiler_n == 0 (cluster_n == 1 is already enforced there). - tuner: the DEFAULT_MOE_TACTIC fallback is gated on the same can_implement -- never fall back to an unvalidated tactic; refuse the shape (empty tactic list) if even the default cannot run safely. - directed host-side unit tests pin the accept/reject matrix. Validated on a B200-class SM100: the previously deterministic accumulated sweep (aborts at item ~54 on unpatched main with the flashinfer-ai#3957 cascade) runs all 98 configs to completion with the guards; residual mxfp8/trtllm_fp8 failures reproduce on unpatched main (different backend, untouched by this change) and are tracked separately. Defense-in-depth follow-ups for the kernel owner (out of scope here): padding CTAs should keep cluster synchronization but skip the global scatter; predicated tail handling if partial N is ever to be supported; a real coordinate predicate on gemm1's SFC store. The same kernels ship in TensorRT-LLM with weaker filtering -- forward this guard upstream. AI-assisted (root-caused and validated on live SM100 hardware). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DEFAULT_MOE_TACTIC is a member of ALL_MOE_TACTICS, so re-checking it after the filtered list comes back empty was dead code. Keep the early refusal (clear warning + empty list) but state its real role honestly: the kernel wrappers re-validate can_implement at launch and raise, so this is diagnostics/defense-in-depth, not the OOB barrier. Making MoELayer skip a runner with no valid tactics (instead of surfacing the wrapper's error) is a separate multi-backend dispatch improvement, out of scope here. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
📝 WalkthroughWalkthroughThe changes add N-tile validation, make fused finalize operations handle partial and padded output tiles, and prevent tactic selection from falling back to an unsupported default. New tests cover host-side validation and functional accuracy. ChangesFused MoE N-tile support
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🧪 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 |
S1ro1
force-pushed
the
fix/cutedsl-moe-tail-store
branch
2 times, most recently
from
July 29, 2026 23:46
b21e926 to
39d0715
Compare
S1ro1
force-pushed
the
fix/cutedsl-moe-tail-store
branch
from
July 29, 2026 23:51
39d0715 to
6d361ab
Compare
S1ro1
force-pushed
the
fix/cutedsl-moe-tail-store
branch
from
July 29, 2026 23:57
6d361ab to
16d9e82
Compare
S1ro1
marked this pull request as ready for review
July 30, 2026 00:30
S1ro1
requested review from
IwakuraRein,
aleozlx,
jiahanc,
nv-yunzheq,
samuellees and
yzh119
as code owners
July 30, 2026 00:30
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Member
|
/bot run tests/moe |
Collaborator
Collaborator
|
Hi @S1ro1, the current PR tries to merge into |
Contributor
Author
Hi, fixed the base to main, however this depends on the PR, if these 2 don't get merged in order I think it'd leave main in a "broken" state. I see the previous base is approved but just noting it down here |
![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 `@flashinfer/fused_moe/cute_dsl/tuner.py`:
- Around line 534-541: Update the no-valid-tactics comment in get_valid_tactics
so it accurately states that an empty result means no tactic, including
DEFAULT_MOE_TACTIC, is selected or profiled; remove any wording implying
fallback to the default tactic while preserving the existing early refusal
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: 815eb161-f979-457f-860c-d5a3b224db20
📒 Files selected for processing (5)
flashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_gather_grouped_gemm_act_fusion.pyflashinfer/fused_moe/cute_dsl/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.pyflashinfer/fused_moe/cute_dsl/tuner.pytests/moe/test_cute_dsl_fused_moe.pytests/moe/test_cute_dsl_moe_can_implement.py
Comment on lines
534
to
+541
| if not valid_tactics: | ||
| # DEFAULT_MOE_TACTIC is a member of ALL_MOE_TACTICS, so an empty | ||
| # list means even the default fails can_implement -- do not fall | ||
| # back to it unvalidated (gh #3957). This early refusal is | ||
| # diagnostics/defense-in-depth: the kernel wrappers re-validate | ||
| # can_implement at launch and raise, so an unvalidated tactic | ||
| # cannot reach the device -- but refusing here avoids pointless | ||
| # profiling of a tactic that can only throw, and says why. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the no-valid-tactics warning.
get_valid_tactics returns an empty list. It does not fall back to DEFAULT_MOE_TACTIC. The warning at Line 545 reports the opposite behavior. This can mislead users during autotuning failures.
Proposed fix
logger.warning(
"No valid tactics found for problem dims "
"(tokens=%d, hidden=%d, intermediate=%d, experts=%d, top_k=%d). "
- "Falling back to default tactic.",
+ "Returning no tactics.",As per coding guidelines, keep documentation synchronized with code changes.
🤖 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/fused_moe/cute_dsl/tuner.py` around lines 534 - 541, Update the
no-valid-tactics comment in get_valid_tactics so it accurately states that an
empty result means no tactic, including DEFAULT_MOE_TACTIC, is selected or
profiled; remove any wording implying fallback to the default tactic while
preserving the existing early refusal behavior.
Source: Coding guidelines
aleozlx
approved these changes
Aug 3, 2026
2 tasks
Member
|
/bot run tests/moe |
Member
|
added 0.6.18 label due to cherry picking to 0.6.17rc2 before merging to main |
Collaborator
Collaborator
|
[FAILED] Pipeline #60947895 — 17/18 executed test jobs passed Compared with nightly #60831563. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
5 tasks
This was referenced Aug 7, 2026
10 tasks
aleozlx
added a commit
that referenced
this pull request
Aug 19, 2026
…4475) ## 📌 Description Takes over and supersedes #3958: rebase onto tot and make the accumulated fuzzer the default regression for the #3957 CUDA-context corruption (fixed by #4186). - Enable `tests/moe/test_unified_moe_fuzz.py` by default. `FLASHINFER_UMOE_FUZZ=0` remains an emergency waiver. Randomized sweep default is 160 configs (was 80). - Keep that accumulated sequence in one pytest process via `shard_group("unified-moe-accumulated")`, so node-level CI sharding cannot split the #3957 regression. - Add a shared finding/quarantine ledger (`tests/test_helpers/fuzz_ledger.py`): - Wrong-answer findings still run, then report XFAIL. - Crash findings are quarantined before kernel launch. - All-backend quarantines report XFAIL rather than SKIP. - Unexpected passes fail strictly. - Every curated fuzzer seed is unique; duplicates are rejected at import. - Add CUTLASS backends to the unified MoE fuzzer (`CutlassBf16Config`, `CutlassW4A16Config`) with a shared BF16-grid routing-weight contract and an SM90-safe Torch MXFP4 reference. - Replace large GEMM/MoE Cartesian grids with curated smoke/regression cases. Randomized shape breadth moves to the default-on unified fuzzers; backend × quant × routing × layout matrices and error-path anchors stay in the original files. - Fix the MxFP8 B-layout used by the cuDNN override-shape path (column-major `[b, k, n]` view). - Document that #3547 and #3957 are fixed. The live ledger is empty; those cases remain as regression coverage, not active waivers. ## 🔍 Related Issues - Supersedes #3958 - #3957 — cumulative CUDA-context corruption; fixed by #4186 - #3547 — expert-offset all-zeros; fixed - #3605 — release-quality / fuzzing plan ## 🚀 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. - [ ] All tests are passing (`unittest`, etc.). Focused local checks (SM100): - [x] `FuzzLedger` unit tests - [x] Targeted unified-MoE fuzzer cases (CUTLASS BF16, W4A16 reference, seed 99, historical MXFP4 config) - [x] #4186 output-tail / tactic guards - [x] Full post-rebase accumulated sweep: 191 passed, 2 skipped - [ ] CUTLASS W4A16 fuzzer path on SM90/H100 (needs GPU CI) A previous `tests/gemm` + `tests/moe` GitLab run passed 18/18 jobs, but that pipeline started before the duplicate-seed fix. Re-run after this description lands. ## Reviewer Notes Legacy-test reductions are intentional: keep kernel-selection and error-path anchors in the original files, and put randomized shape breadth in the default-enabled unified fuzzer. Model-relevant 1024/768 routing sizes remain where the fuzzer does not reproduce the full implementation × weight-layout × activation matrix. The sigmoid grid dropping `intermediate_size=512` matches that test’s compatible sizes (`384/768/1024`). Linear MxFP8 scale layout is not represented by the public 3D BMM API; that coverage stays in `tests/gemm/test_unified_gemm_fuzz.py`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added public access to the BF16 routed MoE runner. - Enabled unified MoE fuzz testing by default in CI, with failure tracking, quarantine handling, and unexpected-pass detection. - **Bug Fixes** - Retained regression coverage for expert-offset handling and improved reference validation for quantized MoE cases. - **Tests** - Streamlined GEMM and MoE tests into focused smoke suites. - Expanded randomized coverage through unified fuzz testing across backends, layouts, dtypes, routing, and autotuning scenarios. - **Documentation** - Added contributor guidance explaining smoke-test scope and randomized coverage responsibilities. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yang Xu <yanxu@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Alex Yang <aleyang@nvidia.com>
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
…lashinfer-ai#4475) ## 📌 Description Takes over and supersedes flashinfer-ai#3958: rebase onto tot and make the accumulated fuzzer the default regression for the flashinfer-ai#3957 CUDA-context corruption (fixed by flashinfer-ai#4186). - Enable `tests/moe/test_unified_moe_fuzz.py` by default. `FLASHINFER_UMOE_FUZZ=0` remains an emergency waiver. Randomized sweep default is 160 configs (was 80). - Keep that accumulated sequence in one pytest process via `shard_group("unified-moe-accumulated")`, so node-level CI sharding cannot split the flashinfer-ai#3957 regression. - Add a shared finding/quarantine ledger (`tests/test_helpers/fuzz_ledger.py`): - Wrong-answer findings still run, then report XFAIL. - Crash findings are quarantined before kernel launch. - All-backend quarantines report XFAIL rather than SKIP. - Unexpected passes fail strictly. - Every curated fuzzer seed is unique; duplicates are rejected at import. - Add CUTLASS backends to the unified MoE fuzzer (`CutlassBf16Config`, `CutlassW4A16Config`) with a shared BF16-grid routing-weight contract and an SM90-safe Torch MXFP4 reference. - Replace large GEMM/MoE Cartesian grids with curated smoke/regression cases. Randomized shape breadth moves to the default-on unified fuzzers; backend × quant × routing × layout matrices and error-path anchors stay in the original files. - Fix the MxFP8 B-layout used by the cuDNN override-shape path (column-major `[b, k, n]` view). - Document that flashinfer-ai#3547 and flashinfer-ai#3957 are fixed. The live ledger is empty; those cases remain as regression coverage, not active waivers. ## 🔍 Related Issues - Supersedes flashinfer-ai#3958 - flashinfer-ai#3957 — cumulative CUDA-context corruption; fixed by flashinfer-ai#4186 - flashinfer-ai#3547 — expert-offset all-zeros; fixed - flashinfer-ai#3605 — release-quality / fuzzing plan ## 🚀 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. - [ ] All tests are passing (`unittest`, etc.). Focused local checks (SM100): - [x] `FuzzLedger` unit tests - [x] Targeted unified-MoE fuzzer cases (CUTLASS BF16, W4A16 reference, seed 99, historical MXFP4 config) - [x] flashinfer-ai#4186 output-tail / tactic guards - [x] Full post-rebase accumulated sweep: 191 passed, 2 skipped - [ ] CUTLASS W4A16 fuzzer path on SM90/H100 (needs GPU CI) A previous `tests/gemm` + `tests/moe` GitLab run passed 18/18 jobs, but that pipeline started before the duplicate-seed fix. Re-run after this description lands. ## Reviewer Notes Legacy-test reductions are intentional: keep kernel-selection and error-path anchors in the original files, and put randomized shape breadth in the default-enabled unified fuzzer. Model-relevant 1024/768 routing sizes remain where the fuzzer does not reproduce the full implementation × weight-layout × activation matrix. The sigmoid grid dropping `intermediate_size=512` matches that test’s compatible sizes (`384/768/1024`). Linear MxFP8 scale layout is not represented by the public 3D BMM API; that coverage stays in `tests/gemm/test_unified_gemm_fuzz.py`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added public access to the BF16 routed MoE runner. - Enabled unified MoE fuzz testing by default in CI, with failure tracking, quarantine handling, and unexpected-pass detection. - **Bug Fixes** - Retained regression coverage for expert-offset handling and improved reference validation for quantized MoE cases. - **Tests** - Streamlined GEMM and MoE tests into focused smoke suites. - Expanded randomized coverage through unified fuzz testing across backends, layouts, dtypes, routing, and autotuning scenarios. - **Documentation** - Added contributor guidance explaining smoke-test scope and randomized coverage responsibilities. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yang Xu <yanxu@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Alex Yang <aleyang@nvidia.com>
kahyunnam
pushed a commit
that referenced
this pull request
Aug 20, 2026
…4475) Takes over and supersedes #3958: rebase onto tot and make the accumulated fuzzer the default regression for the #3957 CUDA-context corruption (fixed by #4186). - Enable `tests/moe/test_unified_moe_fuzz.py` by default. `FLASHINFER_UMOE_FUZZ=0` remains an emergency waiver. Randomized sweep default is 160 configs (was 80). - Keep that accumulated sequence in one pytest process via `shard_group("unified-moe-accumulated")`, so node-level CI sharding cannot split the #3957 regression. - Add a shared finding/quarantine ledger (`tests/test_helpers/fuzz_ledger.py`): - Wrong-answer findings still run, then report XFAIL. - Crash findings are quarantined before kernel launch. - All-backend quarantines report XFAIL rather than SKIP. - Unexpected passes fail strictly. - Every curated fuzzer seed is unique; duplicates are rejected at import. - Add CUTLASS backends to the unified MoE fuzzer (`CutlassBf16Config`, `CutlassW4A16Config`) with a shared BF16-grid routing-weight contract and an SM90-safe Torch MXFP4 reference. - Replace large GEMM/MoE Cartesian grids with curated smoke/regression cases. Randomized shape breadth moves to the default-on unified fuzzers; backend × quant × routing × layout matrices and error-path anchors stay in the original files. - Fix the MxFP8 B-layout used by the cuDNN override-shape path (column-major `[b, k, n]` view). - Document that #3547 and #3957 are fixed. The live ledger is empty; those cases remain as regression coverage, not active waivers. - Supersedes #3958 - #3957 — cumulative CUDA-context corruption; fixed by #4186 - #3547 — expert-offset all-zeros; fixed - #3605 — release-quality / fuzzing plan - [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. - [x] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). Focused local checks (SM100): - [x] `FuzzLedger` unit tests - [x] Targeted unified-MoE fuzzer cases (CUTLASS BF16, W4A16 reference, seed 99, historical MXFP4 config) - [x] #4186 output-tail / tactic guards - [x] Full post-rebase accumulated sweep: 191 passed, 2 skipped - [ ] CUTLASS W4A16 fuzzer path on SM90/H100 (needs GPU CI) A previous `tests/gemm` + `tests/moe` GitLab run passed 18/18 jobs, but that pipeline started before the duplicate-seed fix. Re-run after this description lands. Legacy-test reductions are intentional: keep kernel-selection and error-path anchors in the original files, and put randomized shape breadth in the default-enabled unified fuzzer. Model-relevant 1024/768 routing sizes remain where the fuzzer does not reproduce the full implementation × weight-layout × activation matrix. The sigmoid grid dropping `intermediate_size=512` matches that test’s compatible sizes (`384/768/1024`). Linear MxFP8 scale layout is not represented by the public 3D BMM API; that coverage stays in `tests/gemm/test_unified_gemm_fuzz.py`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> - **New Features** - Added public access to the BF16 routed MoE runner. - Enabled unified MoE fuzz testing by default in CI, with failure tracking, quarantine handling, and unexpected-pass detection. - **Bug Fixes** - Retained regression coverage for expert-offset handling and improved reference validation for quantized MoE cases. - **Tests** - Streamlined GEMM and MoE tests into focused smoke suites. - Expanded randomized coverage through unified fuzz testing across backends, layouts, dtypes, routing, and autotuning scenarios. - **Documentation** - Added contributor guidance explaining smoke-test scope and randomized coverage responsibilities. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yang Xu <yanxu@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Alex Yang <aleyang@nvidia.com> (cherry picked from commit 693fed4)
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
5 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
in the current N tile
uses a separate unpredicated store path
Why
The fused-finalize epilogue previously transferred one full compile-time CTA
tile for every valid row. This is unsafe when:
#4086 prevents those cases by rejecting any finalize configuration where
N % (mma_n * cluster_n) != 0. That is safe, but it removes otherwise usefulkernel configurations and can leave some widths with no eligible configuration.
This PR fixes the transfer itself. It computes the remaining columns for the
current tile, transfers only that many bytes, and performs no transfer when the
tile starts beyond the output width. All CTAs still execute the existing
commit/wait/barrier sequence.
Validation
Correctness
Validated on NVIDIA GB200 with the newly enabled
tile_m=256, mma_n=256, cluster_n=2finalize configuration forced:N=256: exercises an empty cluster-padding CTAN=384: exercises a partially populated final N tileleft a sentinel row immediately after the output unchanged
The focused host-side
can_implementchecks, Ruff checks, and formatting checksalso pass.
Configuration coverage
The table counts complete GEMM1 + GEMM2 configurations exposed by the MoE
runner:
For
N=2880, all eight newly eligible finalize configurations were alsolaunched directly and produced finite output.
Existing-case performance
The finalize kernel was benchmarked in isolation so the epilogue change was not
hidden by routing or GEMM1. The #4086 base and this PR were loaded in the same
process and measured in alternating pairs on one GB200:
N=512, K=512andN=4096, K=1024Lower is better. Across the 16 exact-tile comparisons:
No previously supported configuration showed a hot-cache slowdown. Cold-cache
single-launch measurements were noisier, but showed no systematic regression.
Scope and dependencies
This PR changes only the finalize output transfer. GEMM1 keeps the conservative
guard from #4086.
The PR is stacked on #4086 and targets its upstream
fix-3957-cluster-paddingbranch so this diff contains only the durabletail-handling change. It should be retargeted to
mainafter #4086 merges.The vLLM consumer is vllm-project/vllm#50030.
Summary by CodeRabbit
Bug Fixes
Tests