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/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 4b8764c548a6..3d958e51de34 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,33 @@ 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, 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. + 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 + + +@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)