[https://nvbugs/5986434][fix] Fall back from TRTLLM MoE backend on pre-Blackwell GPUs - #17570
[https://nvbugs/5986434][fix] Fall back from TRTLLM MoE backend on pre-Blackwell GPUs#17570moraxu wants to merge 2 commits into
Conversation
…e-Blackwell GPUs moe_config.backend: TRTLLM selects TRTLLMGenFusedMoE purely on quantization, but the TRTLLM-Gen MoE kernels only exist for SM100/SM103. On Hopper this fails at engine init: on 1.3.0rc7 with "IndexError: list assignment index out of range" in AutoTuner._find_nearest_profile (the op's hidden_states_scale ConstraintSpec assumes Blackwell's 2-D scale layout while SM90 1x128 quantization produces a rank-1 scale), and since the NVIDIA#12310 bounds guard with "No kernel found" during trtllm-gen kernel selection. Gate the TRTLLM branch of get_moe_cls on TRTLLMGenFusedMoE._SUPPORTED_SM_VERSIONS and fall back to CutlassFusedMoE with a warning, mirroring the DENSEGEMM branch. Cutlass matches what AUTO resolves to for these quantizations on non-SM100 GPUs. Verified on 8xH200 with DeepSeek-R1 (TP8, MTP-3, fp8 KV cache, CUDA graphs, autotuner on): 1.3.0rc7 reproduces the reported crash on all ranks; with this change the same config serves successfully with MTP active. Signed-off-by: Michal Guzek <mguzek@nvidia.com>
|
/bot run --disable-fail-fast |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughTRTLLM-Gen MoE backend selection now checks GPU architecture support. Unsupported SM versions fall back to ChangesMoE backend selection
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tensorrt_llm/_torch/modules/fused_moe/create_moe.py`:
- Around line 161-173: Update the unsupported-SM fallback in the MoE backend
factory to validate `CutlassFusedMoE` with `can_implement()` for the requested
quantization and configuration before returning it. If CUTLASS cannot implement
the configuration, select another compatible backend through its capability
checks; otherwise raise a clear unsupported-configuration error instead of
returning `CutlassFusedMoE` unconditionally.
In `@tests/unittest/_torch/modules/moe/test_moe_backend.py`:
- Around line 411-435: Register both
test_get_moe_cls_trtllm_falls_back_to_cutlass_on_unsupported_sm and
test_get_moe_cls_trtllm_selects_trtllm_gen_on_blackwell explicitly in the test
entries for l0_b200.yml and l0_b300.yml. Keep the existing TRTLLM test entries
intact and ensure the new entries invoke these exact test functions so CI
collects them.
🪄 Autofix
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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 08a7ed58-202b-4bbc-b3c6-b3762ce83fc3
📒 Files selected for processing (3)
tensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.pytests/unittest/_torch/modules/moe/test_moe_backend.py
|
PR_Github #65650 [ run ] triggered by Bot. Commit: |
|
/bot kill |
…rams, register tests in CI Per CodeRabbit review on NVIDIA#17570: - Parametrize the fallback test with SM 90/120 only (SMs CutlassFusedMoE supports for FP8_BLOCK_SCALES), so the asserted fallback selection is also runnable. - Register both new get_moe_cls tests explicitly in l0_b200.yml and l0_b300.yml; the existing -k "TRTLLM" selectors only match the uppercase backend-parametrized ids and would not collect them. Signed-off-by: Michal Guzek <mguzek@nvidia.com>
|
PR_Github #65663 [ kill ] triggered by Bot. Commit: |
|
PR_Github #65650 [ run ] completed with state |
|
PR_Github #65663 [ kill ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #65670 [ run ] triggered by Bot. Commit: |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The gate is in the right place — ConfigurableMoE routes through get_moe_cls, so every construction path picks it up — but the fallback only changes the class, not model_config.moe_backend, and a fair number of sites branch on that string:
modeling_gpt_oss.py:174— routingoutput_dtype=bfloat16when the string isTRTLLM.CutlassFusedMoE.input_requirement.routing_scales_dtypeisfloat32andmoe_scheduler.py:407asserts rather than casts, so GptOss + explicitTRTLLMon Hopper/SM120 still fails, just with a different message. The fallback doesn't rescue that model.modeling_deepseekv3.py:872— gatee_score_correction_biasstays bf16, so the fallback run doesn't numerically match a realCUTLASSrun (works via promotion, which is why the H200 run looked fine).model_config.py:492— FP8-block-scalesexclude_modulesstill gets the TRTLLM-specific list.- Also
modeling_qwen3_moe.py:333,modeling_glm.py:722,modeling_deepseekv3.py:1546/deepseekv4.py:2159(do_finalizefusion, reachable on SM120+NVFP4).
Gating in ModelConfig.resolve_moe_backend (or rewriting moe_backend to CUTLASS on fallback) keeps the string and the class in agreement. The new tests assert class identity only; a case covering one string-driven consumer would catch this.
| # resolves to for these quantizations on non-SM100 GPUs. | ||
| from tensorrt_llm._utils import get_sm_version | ||
| sm_version = get_sm_version() | ||
| if sm_version not in TRTLLMGenFusedMoE._SUPPORTED_SM_VERSIONS: |
There was a problem hiding this comment.
Returning CutlassFusedMoE here without updating model_config.moe_backend leaves the rest of the stack configured for TRTLLM-Gen. The concrete failure: modeling_gpt_oss.py:174 picks output_dtype=torch.bfloat16 for the routing method whenever the string is "TRTLLM", and moe_scheduler.py:407 asserts float32 for Cutlass instead of casting — so GptOss with an explicit TRTLLM backend on SM90/SM120 still dies at init after this change.
Suggest doing the SM gate in ModelConfig.resolve_moe_backend (or normalizing model_config.moe_backend = "CUTLASS" when falling back here) so the resolved name and the instantiated class agree for every string-based consumer.
There was a problem hiding this comment.
Please do not add additional fallbacks here.
Fallbacks will soon be deprecated. Users should explicitly set CUTLASS rather than relying on fallbacks provided by TRTLLM.
The new MOE design should honor user‑provided configurations and throw an error when unsupported settings are encountered.
Please update the test case and close this PR.
|
|
||
| # SM versions the TRTLLM-Gen MoE kernels are built for (Blackwell family). | ||
| # Read by create_moe.get_moe_cls to fall back on other architectures. | ||
| _SUPPORTED_SM_VERSIONS = (100, 103) |
There was a problem hiding this comment.
Two more hardcoded copies of this set remain in the same file: _validate_backend_local_activation (get_sm_version() not in {100, 103}) and the sm_version >= 120 raise in __init__. Worth pointing both at _SUPPORTED_SM_VERSIONS while you're consolidating, otherwise the next SM addition has to find three sites.
|
PR_Github #65670 [ run ] completed with state
|
| # resolves to for these quantizations on non-SM100 GPUs. | ||
| from tensorrt_llm._utils import get_sm_version | ||
| sm_version = get_sm_version() | ||
| if sm_version not in TRTLLMGenFusedMoE._SUPPORTED_SM_VERSIONS: |
There was a problem hiding this comment.
Please do not add additional fallbacks here.
Fallbacks will soon be deprecated. Users should explicitly set CUTLASS rather than relying on fallbacks provided by TRTLLM.
The new MOE design should honor user‑provided configurations and throw an error when unsupported settings are encountered.
Please update the test case and close this PR.
|
@xxi-nv , thanks for clarifying. Closing |
4 participants
Dev Engineer Review
get_moe_clsnow gatesTRTLLMGenFusedMoEon SM100 and SM103.CutlassFusedMoEwith a warning._SUPPORTED_SM_VERSIONScentralizes the supported architectures.QA Engineer Review
test_get_moe_cls_trtllm_falls_back_to_cutlass_on_unsupported_sm.test_get_moe_cls_trtllm_selects_trtllm_gen_on_blackwell.tests/integration/test_lists/test-db/l0_b200.yml.tests/integration/test_lists/test-db/l0_b300.yml.Verdict: sufficient
Description
Problem. Setting
moe_config.backend: TRTLLMon a pre-Blackwell GPU (e.g. DeepSeek-R1 FP8 on 8xH200) fails at engine init, becausecreate_moe.get_moe_clsselectsTRTLLMGenFusedMoEbased on quantization alone while the TRTLLM-Gen MoE kernels only exist for SM100/SM103:IndexError: list assignment index out of rangeinAutoTuner._find_nearest_profile: thetrtllm::fp8_block_scale_moe_runnerop declaresConstraintSpec(input_idx=3, dim_idx=1)for Blackwell's 2-Dhidden_states_scalelayout, but Hopper's 1x128 activation quantization produces a rank-1 scale, so the profile-normalization write goes out of bounds. The[Autotuner] Autotuning process endsline in the reported log is printed during exception unwind, which is why the crash appears in the same second.No kernel found for the given options ... mUseDeepSeekFp8: 1fromtrtllmGenKernels/batchedGemm/KernelRunner.cpp— the cubins genuinely don't exist on SM90.Fix. Gate the
TRTLLMbranch ofget_moe_clson the SM versions the backend supports and fall back toCutlassFusedMoEwith a warning, mirroring the existingDENSEGEMM/MEGAMOE_*fallback pattern.CutlassFusedMoEis also whatAUTOresolves to for these quantizations on non-SM100 GPUs (ModelConfig.resolve_moe_backend). The supported-SM set is hoisted intoTRTLLMGenFusedMoE._SUPPORTED_SM_VERSIONS(previously hardcoded inline incan_implement), following theDenseGEMMFusedMoE._SUPPORTED_SM_VERSIONSprecedent already read by the factory.Verification on 8xH200 NVL (SM90), DeepSeek-R1 FP8, exact config from the bug (TP8/EP1, MTP
num_nextn_predict_layers=3, FP8 KV cache, CUDA graphs bs 1–128 with padding, autotuner enabled):release:1.3.0rc7container: reproduces the reported crash byte-for-byte on all 8 ranks.No kernel foundduring the first warmup forward.trtllm-servereachesApplication startup complete, and completions return correct output withavg_decoded_tokens_per_iter = 3.0(MTP-3 active). Also verified with DeepSeek-V3-Lite FP8 on 1 GPU.Test Coverage
tests/unittest/_torch/modules/moe/test_moe_backend.py::test_get_moe_cls_trtllm_falls_back_to_cutlass_on_unsupported_sm(SM 89/90/120, mockedget_sm_version).tests/unittest/_torch/modules/moe/test_moe_backend.py::test_get_moe_cls_trtllm_selects_trtllm_gen_on_blackwell(SM 100/103 selection unchanged).-k "get_moe_cls"selection tests pass (9 passed locally).PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.🤖 Generated with Claude Code