From 93442c6fbc07c48a73959851e5274210d0f49f2f Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Wed, 12 Aug 2026 11:41:44 -0700 Subject: [PATCH 1/2] [https://nvbugs/5986434][fix] Fall back from TRTLLM MoE backend on pre-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 #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 --- .../_torch/modules/fused_moe/create_moe.py | 13 ++++++++++ .../modules/fused_moe/fused_moe_trtllm_gen.py | 6 ++++- .../_torch/modules/moe/test_moe_backend.py | 26 +++++++++++++++++++ 3 files changed, 44 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 0bcb92f07b0a..7f094cf8ebfe 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -158,6 +158,19 @@ def get_moe_cls( return CutlassFusedMoE return DenseGEMMFusedMoE elif moe_backend.upper() == "TRTLLM": + # TRTLLM-Gen MoE kernels only exist for Blackwell-family GPUs; on other + # architectures the backend fails at engine init (autotuner profile + # lookup or "No kernel found" during kernel selection). Fall back like + # the other arch-gated backends. CutlassFusedMoE matches what AUTO + # 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: + logger.warning( + f"{layer_prefix}TRTLLMGenFusedMoE only supports SM " + f"{list(TRTLLMGenFusedMoE._SUPPORTED_SM_VERSIONS)} " + f"(got SM {sm_version}). Using CutlassFusedMoE instead.") + return CutlassFusedMoE has_quant = quant_config is not None and quant_config.quant_mode.has_any_quant( exclude_kv_cache=True) if has_quant and (quant_config.quant_mode.has_fp8_block_scales() diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 79c2f064e1ac..f864d084d188 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -98,6 +98,10 @@ class TRTLLMGenFusedMoE(MoE): onesided_workspace_dtype=torch.bfloat16, ) + # 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) + # Supported quantization algorithms for TRTLLMGenFusedMoE _SUPPORTED_QUANT_ALGOS = { QuantAlgo.NVFP4, @@ -161,7 +165,7 @@ def can_implement( sm_version = get_sm_version() # TRTLLMGenFusedMoE requires SM in {100, 103} - if sm_version not in {100, 103}: + if sm_version not in cls._SUPPORTED_SM_VERSIONS: return _warn_and_return( f"TRTLLMGenFusedMoE requires SM100 or SM103, got SM{sm_version}" ) diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 4b8764c548a6..7715e33a4c11 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -61,6 +61,7 @@ from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe_backend, get_moe_cls from tensorrt_llm._torch.modules.fused_moe.fused_moe_cutlass import CutlassFusedMoE from tensorrt_llm._torch.modules.fused_moe.fused_moe_marlin import MarlinFusedMoE +from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE from tensorrt_llm._torch.modules.fused_moe.impl_contract import MoECommPlan, MoERunContext from tensorrt_llm._torch.modules.fused_moe.interface import ( MoE, @@ -407,6 +408,31 @@ def test_get_moe_cls_marlin_override_quant_config_per_layer(): ) +def _trtllm_model_config(quant_algo=QuantAlgo.FP8_BLOCK_SCALES): + cfg = ModelConfig() + cfg.moe_backend = "TRTLLM" + cfg.quant_config = QuantConfig(quant_algo=quant_algo) if quant_algo else None + return cfg + + +@pytest.mark.parametrize("sm", [90, 89, 120]) +def test_get_moe_cls_trtllm_falls_back_to_cutlass_on_unsupported_sm(sm, monkeypatch): + """TRTLLM-Gen MoE kernels only exist for SM100/SM103. Requesting the TRTLLM + backend elsewhere (e.g. DeepSeek FP8_BLOCK_SCALES on Hopper) must fall back + to Cutlass instead of failing at engine init: on 1.3.0rc7 this crashed with + 'list assignment index out of range' in AutoTuner._find_nearest_profile, + later with 'No kernel found' during trtllm-gen kernel selection.""" + monkeypatch.setattr("tensorrt_llm._utils.get_sm_version", lambda: sm) + assert get_moe_cls(_trtllm_model_config()) is CutlassFusedMoE + + +@pytest.mark.parametrize("sm", [100, 103]) +def test_get_moe_cls_trtllm_selects_trtllm_gen_on_blackwell(sm, monkeypatch): + """The SM gate must not change selection on Blackwell-family GPUs.""" + monkeypatch.setattr("tensorrt_llm._utils.get_sm_version", lambda: sm) + assert get_moe_cls(_trtllm_model_config()) is TRTLLMGenFusedMoE + + def test_megamoe_cutedsl_post_load_weights_uses_staged_hooks(): moe = MegaMoECuteDsl.__new__(MegaMoECuteDsl) torch.nn.Module.__init__(moe) From fa9c5fc23cfb1ebf5b9bbf0b1c11eeea477f021b Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Wed, 12 Aug 2026 12:52:24 -0700 Subject: [PATCH 2/2] [https://nvbugs/5986434][fix] Address review: runnable fallback SM params, register tests in CI Per CodeRabbit review on #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 --- tests/integration/test_lists/test-db/l0_b200.yml | 2 ++ tests/integration/test_lists/test-db/l0_b300.yml | 2 ++ tests/unittest/_torch/modules/moe/test_moe_backend.py | 6 ++++-- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index bc719ed06229..2ab8a6b3f8bd 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -132,6 +132,8 @@ l0_b200: - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout + - unittest/_torch/modules/moe/test_moe_backend.py::test_get_moe_cls_trtllm_falls_back_to_cutlass_on_unsupported_sm + - unittest/_torch/modules/moe/test_moe_backend.py::test_get_moe_cls_trtllm_selects_trtllm_gen_on_blackwell # ------------- MoE: test_single_gpu (by backend) --------------- - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "TRTLLM" diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 95938a31cd98..b79406916490 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -41,6 +41,8 @@ l0_b300: - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout + - unittest/_torch/modules/moe/test_moe_backend.py::test_get_moe_cls_trtllm_falls_back_to_cutlass_on_unsupported_sm + - unittest/_torch/modules/moe/test_moe_backend.py::test_get_moe_cls_trtllm_selects_trtllm_gen_on_blackwell # ------------- MoE: test_single_gpu (specific quant per backend) --------------- # CUTLASS backend: FP8, NVFP4, W4A8_MXFP4_MXFP8, W8A16 - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu[e60_k4_h2048_i1408-seq=1-dtype=torch.bfloat16-backend=CUTLASS-quant=FP8-routing=Renormalize] diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 7715e33a4c11..3d958e51de34 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -415,13 +415,15 @@ def _trtllm_model_config(quant_algo=QuantAlgo.FP8_BLOCK_SCALES): return cfg -@pytest.mark.parametrize("sm", [90, 89, 120]) +@pytest.mark.parametrize("sm", [90, 120]) def test_get_moe_cls_trtllm_falls_back_to_cutlass_on_unsupported_sm(sm, monkeypatch): """TRTLLM-Gen MoE kernels only exist for SM100/SM103. Requesting the TRTLLM backend elsewhere (e.g. DeepSeek FP8_BLOCK_SCALES on Hopper) must fall back to Cutlass instead of failing at engine init: on 1.3.0rc7 this crashed with 'list assignment index out of range' in AutoTuner._find_nearest_profile, - later with 'No kernel found' during trtllm-gen kernel selection.""" + later with 'No kernel found' during trtllm-gen kernel selection. + The parametrized SMs are ones CutlassFusedMoE supports for FP8_BLOCK_SCALES, + so the fallback selection is also runnable.""" monkeypatch.setattr("tensorrt_llm._utils.get_sm_version", lambda: sm) assert get_moe_cls(_trtllm_model_config()) is CutlassFusedMoE