fix: avoid SM120 MXFP8 MXFP4 autotune crash - #4066
waynehacking8 wants to merge 1 commit into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughThe fused-MoE autotuning path now bypasses SM120/SM121 profiling for MXFP8 activation scaling with int64 expert weights. Tests can enable autotuning and add SM120/SM121-specific regression coverage. ChangesSM120 fused-MoE autotune guard
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to This localized change avoids the SM120/SM121 MXFP8×MXFP4 autotune crash by using fallback tactics and adds focused regression coverage; no actionable merge-blocking risk remains after normal checks and review. Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Description checkExplanation The description provides a clear summary, links issue Full details: Linked Issues checkExplanation The changes satisfy issue ✨ 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request addresses an issue where the SM120/SM121 profiler could poison the CUDA context when profiling MXFP8 x MXFP4 TMA inputs. It bypasses profiling and uses fallback tactics when the backend is SM120/SM121, use_mxfp8_act_scaling is enabled, and fc1_expert_weights are torch.int64. Additionally, a regression test test_moe_mxfp8_mxfp4_autotune_sm120 has been added to verify this behavior under autotuning. I have no further feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
Independent hardware validation on this exact PR head ( While reproducing #4049 on current main, I also isolated all 20 GEMM1 tactics in fresh processes. Every tactic failed for this quantization path; tactics I had opened #4121 after missing this earlier PR; its core guard is equivalent, so I am closing mine as a duplicate and consolidating the validation here. Test orchestration and code navigation were AI-assisted; the builds and GPU results above were executed against this PR head. |
0268e33 to
0923f51
Compare
|
@flashinfer-bot run |
|
Rebased onto current |
## 📌 Description `GemmProfilerBackend::init()` receives `use_mxfp8_act_scaling` but does not retain it, so `prepareQuantParams()` always constructs `QuantParams::FP8MXFP4` on the FP8-activation/MXFP4-weight path. That leaves the `mxfp8_mxfp4` weight block-scale pointers null, so the MXFPX scale-factor descriptor setup is skipped during profiling. VictoriaLogs captured the affected DeepSeek-V4-Flash-0731 TP2 startup reaching fused-MoE profiling and then exiting after this error surfaced: ```text [AutoTuner]: Tuning trtllm::fused_moe::gemm1 ... torch.AcceleratorError: CUDA error: an illegal instruction was encountered ... [2026-07-31 15:42:01] Received sigquit from a child process. It usually means the child failed. ``` This change: - stores `use_mxfp8_act_scaling` in the profiler; - sizes the fc1 weight scale-factor workspace for the doubled gated-activation output; - initializes the profiler's fc1/fc2 weight block scale-factor workspaces to UE8M0 unity on the MXFP8-activation path; - constructs `QuantParams::MXFP8MXFP4` for MXFP8-activation/MXFP4-weight profiling; and - runs the existing numerical test with autotuning disabled and enabled. When `use_mxfp8_act_scaling` is false, the existing `QuantParams::FP8MXFP4` branch remains in use. The shared fc1 scale-factor workspace sizing is corrected for the gated output in both branches. PR #4066 skips profiling and selects fallback tactics for MXFP8×MXFP4 on SM120/SM121; this change repairs the profiler inputs and retains tactic profiling. Both changes update the same regression test. If this change is accepted, PR #4066's SM120/SM121 fallback-tactic override should be removed, or #4066 closed. That override keeps SM120/SM121 on fallback tactics for this mode, which makes this profiler fix inert there. ## 🔍 Related Issues Fixes #4049. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I used `uvx pre-commit`. - [ ] I installed the hooks with `pre-commit install`. Not used; hooks were run directly with `uvx`. - [x] Focused pre-commit checks pass for all three changed paths: `csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh`, `csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h`, and `tests/moe/test_trtllm_cutlass_fused_moe.py`. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [ ] The complete repository test suite was not run. - `uv run --no-project python -m pytest -q tests/moe/test_trtllm_cutlass_fused_moe.py -k mxfp8_mxfp4` on SM120: 12 passed, 6 skipped, 132 deselected. - With this change applied, a `DeepSeek-V4-Flash-0731` TP2 startup completed all 10 gemm1 and 10 gemm2 profiler candidates on both ranks, completed FlashInfer autotuning, captured the target and draft CUDA graphs, and reached ready state with zero restarts. Test limitations: the autotuned numerical test exercises this path and checks numerical parity, but it does not directly inspect the fabricated workspace contents. The non-MXFP8 activation branch touched by the shared fc1 scale-factor sizing has no focused autotune regression case in this validation. ## Reviewer Notes The table below compares heuristic fallback tactics with profiler-selected tactics. It is not a direct patched-versus-unpatched throughput comparison; the patch is what allows the profiled-tactic side to start successfully. Both sides used the same SGLang revision, model, TP2 configuration, FP8 KV cache, DSpark width 5, hardware, and workload. The control skipped `trtllm::fused_moe::gemm1` and `gemm2`; the candidate profiled them. Methodology: each decode cell is one unreplicated 30-second run at C1, C2, C4, C8, C16, and C32. The coding row is the median of five sequential requests with a 2,000-token cap. The prefill rows are the median client-observed throughput from standalone cold-prefill runs at exact 8K, 64K, and 128K prompts with 7, 2, and 1 samples, respectively. All rows are single comparison runs with no variance estimate, and no effect size is claimed from any row. | Cell | Fallback tactics | Profiled tactics | |---|---:|---:| | C1 decode (tok/s) | 187.7 | 240.5 | | C2 decode (tok/s) | 241.4 | 349.5 | | C4 decode (tok/s) | 347.1 | 515.2 | | C8 decode (tok/s) | 456.5 | 662.9 | | C16 decode (tok/s) | 683.8 | 992.8 | | C32 decode (tok/s) | 1030.7 | 1413.6 | | Coding median (tok/s) | 218.1 | 264.8 | | 8K prefill (tok/s) | 7444 | 7638 | | 64K prefill (tok/s) | 8557 | 8351 | | 128K prefill (tok/s) | 7922 | 7744 | The pinned 1,319-row GSM8K gate was not counted as passing because each run had one response finish at the 1,024-token cap. Two profiled-tactic runs scored 1,241 and 1,239 correct; the fallback-tactic control scored 1,240 correct. No quality difference is claimed from these runs. AI was used to assist with implementation. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved profiling for fused mixture-of-experts workloads using MXFP8 activations and MXFP4 weights. * Corrected workspace sizing and scaling behavior for quantized weight formats. * Added validation for unsupported scaling configurations. * **Tests** * Expanded coverage to verify both autotuned and default execution paths for quantized MoE workloads. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
0923f51 to
9a99fb8
Compare
|
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. |
3 participants
Summary
Fixes #4049
Testing
ruff checkandruff format --checkon changed filesSummary by CodeRabbit
Bug Fixes
Tests
use_autotuneoption to the MXFP8/MXFP4 test to run with autotuning enabled or disabled.