From e9be4f8004278d8903d52c1e41741bc7bb532f02 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Wed, 1 Jul 2026 14:34:31 -0700 Subject: [PATCH 1/5] [None][fix] Gate SM121 B12x MoE on CUDA13 CuTe DSL Signed-off-by: Mihai Chiorean --- .../_torch/modules/fused_moe/create_moe.py | 9 +++ .../fused_moe/fused_moe_cute_dsl_b12x.py | 48 ++++++++++++-- .../moe/test_cute_dsl_b12x_moe_backend.py | 65 +++++++++++++++++-- 3 files changed, 114 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 1d22fc1ac47c..865047aebe94 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -89,6 +89,15 @@ def get_moe_cls( from tensorrt_llm._utils import get_sm_version sm_version = get_sm_version() if sm_version in CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS: + runtime_disable_reason = CuteDslB12xFusedMoE.get_runtime_disable_reason( + sm_version) + if runtime_disable_reason is not None: + logger.warning_once( + f"{layer_prefix}{runtime_disable_reason} " + "Using CutlassFusedMoE instead.", + key="cute_dsl_b12x_runtime_disabled", + ) + return CutlassFusedMoE try: import flashinfer # noqa: F401 logger.info( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py index f9cc59724ec1..3cd4abf4d09c 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py @@ -83,9 +83,10 @@ class CuteDslB12xFusedMoE(CuteDslFusedMoE): flashinfer-importable gates pass (see ``create_moe.get_moe_cls``). """ - # SM versions on which the FlashInfer b12x NVFP4 MoE kernel is available. + # SM versions on which FlashInfer exposes the b12x NVFP4 MoE kernel. # SM120 = desktop Blackwell (RTX 5090 / GB202); SM121 = GB10 / DGX Spark. _SUPPORTED_SM_VERSIONS = frozenset({120, 121}) + _MIN_CUTLASS_DSL_CUDA_MAJOR_FOR_SM121 = 13 # Prefill chunks (``x.shape[0] >= threshold``) route via CUTLASS NVFP4 # GroupGEMM; decode (``x.shape[0] < threshold``) uses b12x. 64 cleanly @@ -102,9 +103,9 @@ def can_implement( swiglu_gptoss_style: bool = False, ) -> Tuple[bool, Optional[str]]: sm_version = get_sm_version() - if sm_version not in cls._SUPPORTED_SM_VERSIONS: - sm_list = "/".join(f"SM{v}" for v in sorted(cls._SUPPORTED_SM_VERSIONS)) - return _warn_and_return(f"CuteDslB12xFusedMoE requires {sm_list}, got SM{sm_version}") + runtime_disable_reason = cls.get_runtime_disable_reason(sm_version) + if runtime_disable_reason is not None: + return _warn_and_return(runtime_disable_reason) if quant_algo != QuantAlgo.NVFP4: return _warn_and_return( f"CuteDslB12xFusedMoE only supports NVFP4 quantization " @@ -119,6 +120,45 @@ def can_implement( return _warn_and_return("CuteDslB12xFusedMoE does not support swiglu_gptoss_style") return True, None + @classmethod + def get_runtime_disable_reason(cls, sm_version: int) -> Optional[str]: + if sm_version not in cls._SUPPORTED_SM_VERSIONS: + sm_list = "/".join(f"SM{v}" for v in sorted(cls._SUPPORTED_SM_VERSIONS)) + return f"CuteDslB12xFusedMoE requires {sm_list}, got SM{sm_version}" + + # CUDA 12.x CuTe DSL lowers the SM121 NVFP4 MMA atom to the internal + # ``_mma.block_scale...`` spelling, which ptxas rejects. CUDA 13.x + # emits the public ``mma.sync.aligned...kind::mxf4nvf4`` opcode. + if sm_version == 121 and not cls._is_sm121_cutlass_dsl_runtime_available(): + return ( + "CuteDslB12xFusedMoE on SM121 requires the active " + "nvidia-cutlass-dsl CUDA 13 native payload. CUDA 12.x CuTe " + "DSL lowers FlashInfer's B12x NVFP4 MMA to PTX that ptxas " + "rejects with Unexpected instruction types specified for " + "'_mma'. Install nvidia-cutlass-dsl-libs-cu13 after " + "nvidia-cutlass-dsl-libs-base (or install " + "nvidia-cutlass-dsl[cu13]) so " + "cutlass.base_dsl.version_info.CUDA_VERSION reports CUDA " + f"{cls._MIN_CUTLASS_DSL_CUDA_MAJOR_FOR_SM121}.x or newer." + ) + return None + + @classmethod + def _is_sm121_cutlass_dsl_runtime_available(cls) -> bool: + cuda_major = cls._get_cutlass_dsl_cuda_major() + return cuda_major is not None and cuda_major >= cls._MIN_CUTLASS_DSL_CUDA_MAJOR_FOR_SM121 + + @staticmethod + def _get_cutlass_dsl_cuda_major() -> Optional[int]: + try: + from cutlass.base_dsl.version_info import CUDA_VERSION + except ImportError: + return None + major = getattr(CUDA_VERSION, "major", None) + if isinstance(major, int): + return major + return None + def __init__(self, *args, **kwargs): # ``ModelConfig`` is consumed by the inherited ``__init__`` for cache # / mapping setup but isn't kept on ``self``. b12x's wrapper needs the diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 1feac564b176..fb200bdfa83c 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -47,7 +47,7 @@ def test_can_implement_rejects_unsupported_sm(sm_version): assert reason is not None and f"SM{sm_version}" in reason -@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) +@pytest.mark.parametrize("sm_version", [120]) def test_can_implement_accepts_supported_sm_with_nvfp4(sm_version): with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version): ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) @@ -55,6 +55,38 @@ def test_can_implement_accepts_supported_sm_with_nvfp4(sm_version): assert reason is None +def test_can_implement_rejects_sm121_with_cuda12_cute_dsl(): + with ( + patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=121), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=12), + ): + ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) + assert not ok + assert reason is not None + assert "CUDA 13 native payload" in reason + assert "Unexpected instruction types" in reason + + +def test_can_implement_rejects_sm121_when_cute_dsl_version_unavailable(): + with ( + patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=121), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=None), + ): + ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) + assert not ok + assert reason is not None and "CUDA 13 native payload" in reason + + +def test_can_implement_accepts_sm121_with_cuda13_cute_dsl(): + with ( + patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=121), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), + ): + ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) + assert ok + assert reason is None + + @pytest.mark.parametrize( "quant_algo", [ @@ -130,9 +162,9 @@ def test_get_moe_cls_cutedsl_returns_plain_cutedsl_on_unsupported_sm(): assert cls is CuteDslFusedMoE -@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) +@pytest.mark.parametrize("sm_version", [120]) def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version): - """CUTEDSL + NVFP4 + SM120/121 + flashinfer importable → CuteDslB12xFusedMoE.""" + """CUTEDSL + NVFP4 + SM120 + flashinfer importable → CuteDslB12xFusedMoE.""" cfg = ModelConfig() cfg.moe_backend = "CUTEDSL" cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) @@ -141,8 +173,33 @@ def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version): assert cls is CuteDslB12xFusedMoE +def test_get_moe_cls_cutedsl_falls_back_to_cutlass_on_sm121_with_cuda12(): + """CUTEDSL + NVFP4 + SM121 avoids the CUDA 12 CuTe DSL JIT crash.""" + cfg = ModelConfig() + cfg.moe_backend = "CUTEDSL" + cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) + with ( + patch("tensorrt_llm._utils.get_sm_version", return_value=121), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=12), + ): + cls = get_moe_cls(cfg) + assert cls is CutlassFusedMoE + + +def test_get_moe_cls_cutedsl_selects_b12x_on_sm121_with_cuda13(): + cfg = ModelConfig() + cfg.moe_backend = "CUTEDSL" + cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) + with ( + patch("tensorrt_llm._utils.get_sm_version", return_value=121), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), + ): + cls = get_moe_cls(cfg) + assert cls is CuteDslB12xFusedMoE + + def test_get_moe_cls_cutedsl_falls_back_to_plain_cutedsl_when_flashinfer_missing(monkeypatch): - """CUTEDSL + NVFP4 + SM120/121 + flashinfer NOT importable → CuteDslFusedMoE.""" + """CUTEDSL + NVFP4 + SM120 + flashinfer NOT importable → CuteDslFusedMoE.""" import builtins cfg = ModelConfig() From d03e0f0af5c7d606ecb67dc687edeee7ea0221ec Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Wed, 1 Jul 2026 15:17:10 -0700 Subject: [PATCH 2/5] [#15853][fix] Gate B12x MoE on CUDA13 CuTe DSL for SM120 Signed-off-by: Mihai Chiorean --- .../fused_moe/fused_moe_cute_dsl_b12x.py | 14 ++-- .../moe/test_cute_dsl_b12x_moe_backend.py | 65 +++++++++++++------ 2 files changed, 51 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py index 3cd4abf4d09c..1e26bc8d3e1f 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py @@ -86,7 +86,7 @@ class CuteDslB12xFusedMoE(CuteDslFusedMoE): # SM versions on which FlashInfer exposes the b12x NVFP4 MoE kernel. # SM120 = desktop Blackwell (RTX 5090 / GB202); SM121 = GB10 / DGX Spark. _SUPPORTED_SM_VERSIONS = frozenset({120, 121}) - _MIN_CUTLASS_DSL_CUDA_MAJOR_FOR_SM121 = 13 + _MIN_CUTLASS_DSL_CUDA_MAJOR_FOR_B12X = 13 # Prefill chunks (``x.shape[0] >= threshold``) route via CUTLASS NVFP4 # GroupGEMM; decode (``x.shape[0] < threshold``) uses b12x. 64 cleanly @@ -126,12 +126,12 @@ def get_runtime_disable_reason(cls, sm_version: int) -> Optional[str]: sm_list = "/".join(f"SM{v}" for v in sorted(cls._SUPPORTED_SM_VERSIONS)) return f"CuteDslB12xFusedMoE requires {sm_list}, got SM{sm_version}" - # CUDA 12.x CuTe DSL lowers the SM121 NVFP4 MMA atom to the internal + # CUDA 12.x CuTe DSL lowers the SM12x NVFP4 MMA atom to the internal # ``_mma.block_scale...`` spelling, which ptxas rejects. CUDA 13.x # emits the public ``mma.sync.aligned...kind::mxf4nvf4`` opcode. - if sm_version == 121 and not cls._is_sm121_cutlass_dsl_runtime_available(): + if not cls._is_cutlass_dsl_runtime_available(): return ( - "CuteDslB12xFusedMoE on SM121 requires the active " + "CuteDslB12xFusedMoE requires the active " "nvidia-cutlass-dsl CUDA 13 native payload. CUDA 12.x CuTe " "DSL lowers FlashInfer's B12x NVFP4 MMA to PTX that ptxas " "rejects with Unexpected instruction types specified for " @@ -139,14 +139,14 @@ def get_runtime_disable_reason(cls, sm_version: int) -> Optional[str]: "nvidia-cutlass-dsl-libs-base (or install " "nvidia-cutlass-dsl[cu13]) so " "cutlass.base_dsl.version_info.CUDA_VERSION reports CUDA " - f"{cls._MIN_CUTLASS_DSL_CUDA_MAJOR_FOR_SM121}.x or newer." + f"{cls._MIN_CUTLASS_DSL_CUDA_MAJOR_FOR_B12X}.x or newer." ) return None @classmethod - def _is_sm121_cutlass_dsl_runtime_available(cls) -> bool: + def _is_cutlass_dsl_runtime_available(cls) -> bool: cuda_major = cls._get_cutlass_dsl_cuda_major() - return cuda_major is not None and cuda_major >= cls._MIN_CUTLASS_DSL_CUDA_MAJOR_FOR_SM121 + return cuda_major is not None and cuda_major >= cls._MIN_CUTLASS_DSL_CUDA_MAJOR_FOR_B12X @staticmethod def _get_cutlass_dsl_cuda_major() -> Optional[int]: diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index fb200bdfa83c..973f47739e1f 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -47,17 +47,21 @@ def test_can_implement_rejects_unsupported_sm(sm_version): assert reason is not None and f"SM{sm_version}" in reason -@pytest.mark.parametrize("sm_version", [120]) +@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) def test_can_implement_accepts_supported_sm_with_nvfp4(sm_version): - with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version): + with ( + patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), + ): ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) assert ok assert reason is None -def test_can_implement_rejects_sm121_with_cuda12_cute_dsl(): +@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) +def test_can_implement_rejects_supported_sm_with_cuda12_cute_dsl(sm_version): with ( - patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=121), + patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version), patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=12), ): ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) @@ -67,9 +71,10 @@ def test_can_implement_rejects_sm121_with_cuda12_cute_dsl(): assert "Unexpected instruction types" in reason -def test_can_implement_rejects_sm121_when_cute_dsl_version_unavailable(): +@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) +def test_can_implement_rejects_supported_sm_when_cute_dsl_version_unavailable(sm_version): with ( - patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=121), + patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version), patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=None), ): ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) @@ -77,9 +82,10 @@ def test_can_implement_rejects_sm121_when_cute_dsl_version_unavailable(): assert reason is not None and "CUDA 13 native payload" in reason -def test_can_implement_accepts_sm121_with_cuda13_cute_dsl(): +@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) +def test_can_implement_accepts_supported_sm_with_cuda13_cute_dsl(sm_version): with ( - patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=121), + patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version), patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), ): ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) @@ -99,14 +105,20 @@ def test_can_implement_accepts_sm121_with_cuda13_cute_dsl(): ) def test_can_implement_rejects_non_nvfp4(quant_algo): """Only NVFP4 is supported; everything else must be turned away.""" - with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=120): + with ( + patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=120), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), + ): ok, reason = CuteDslB12xFusedMoE.can_implement(quant_algo) assert not ok assert reason is not None and "NVFP4" in reason def test_can_implement_rejects_swiglu_gptoss_style(): - with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=120): + with ( + patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=120), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), + ): ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4, swiglu_gptoss_style=True) assert not ok assert reason is not None and "swiglu_gptoss_style" in reason @@ -114,7 +126,10 @@ def test_can_implement_rejects_swiglu_gptoss_style(): @pytest.mark.parametrize("dtype", [torch.float32, torch.float8_e4m3fn]) def test_can_implement_rejects_unsupported_activation_dtype(dtype): - with patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=120): + with ( + patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=120), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), + ): ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4, dtype_activation=dtype) assert not ok assert reason is not None @@ -162,36 +177,41 @@ def test_get_moe_cls_cutedsl_returns_plain_cutedsl_on_unsupported_sm(): assert cls is CuteDslFusedMoE -@pytest.mark.parametrize("sm_version", [120]) +@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version): - """CUTEDSL + NVFP4 + SM120 + flashinfer importable → CuteDslB12xFusedMoE.""" + """CUTEDSL + NVFP4 + SM12x + CUDA13 CuTe DSL -> CuteDslB12xFusedMoE.""" cfg = ModelConfig() cfg.moe_backend = "CUTEDSL" cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - with patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version): + with ( + patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), + ): cls = get_moe_cls(cfg) assert cls is CuteDslB12xFusedMoE -def test_get_moe_cls_cutedsl_falls_back_to_cutlass_on_sm121_with_cuda12(): - """CUTEDSL + NVFP4 + SM121 avoids the CUDA 12 CuTe DSL JIT crash.""" +@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) +def test_get_moe_cls_cutedsl_falls_back_to_cutlass_on_supported_sm_with_cuda12(sm_version): + """CUTEDSL + NVFP4 + SM12x avoids the CUDA 12 CuTe DSL JIT crash.""" cfg = ModelConfig() cfg.moe_backend = "CUTEDSL" cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) with ( - patch("tensorrt_llm._utils.get_sm_version", return_value=121), + patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version), patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=12), ): cls = get_moe_cls(cfg) assert cls is CutlassFusedMoE -def test_get_moe_cls_cutedsl_selects_b12x_on_sm121_with_cuda13(): +@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) +def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm_with_cuda13(sm_version): cfg = ModelConfig() cfg.moe_backend = "CUTEDSL" cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) with ( - patch("tensorrt_llm._utils.get_sm_version", return_value=121), + patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version), patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), ): cls = get_moe_cls(cfg) @@ -199,7 +219,7 @@ def test_get_moe_cls_cutedsl_selects_b12x_on_sm121_with_cuda13(): def test_get_moe_cls_cutedsl_falls_back_to_plain_cutedsl_when_flashinfer_missing(monkeypatch): - """CUTEDSL + NVFP4 + SM120 + flashinfer NOT importable → CuteDslFusedMoE.""" + """CUTEDSL + NVFP4 + SM12x + flashinfer NOT importable -> CuteDslFusedMoE.""" import builtins cfg = ModelConfig() @@ -214,7 +234,10 @@ def _raise_on_flashinfer(name, *args, **kwargs): return real_import(name, *args, **kwargs) monkeypatch.setattr(builtins, "__import__", _raise_on_flashinfer) - with patch("tensorrt_llm._utils.get_sm_version", return_value=120): + with ( + patch("tensorrt_llm._utils.get_sm_version", return_value=120), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), + ): cls = get_moe_cls(cfg) assert cls is CuteDslFusedMoE From ff022fc2de6b7bd8ea2c28a8b904c1141f47ea9d Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Wed, 1 Jul 2026 16:37:23 -0700 Subject: [PATCH 3/5] [#15853][fix] Use public CUTLASS DSL version API Signed-off-by: Mihai Chiorean --- .../fused_moe/fused_moe_cute_dsl_b12x.py | 10 +++- .../moe/test_cute_dsl_b12x_moe_backend.py | 57 +++++++++++-------- 2 files changed, 40 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py index 1e26bc8d3e1f..5ab771a796b9 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py @@ -122,6 +122,11 @@ def can_implement( @classmethod def get_runtime_disable_reason(cls, sm_version: int) -> Optional[str]: + """Return why this backend cannot run, or ``None`` when it can. + + Args: + sm_version: SM version as returned by ``get_sm_version()``. + """ if sm_version not in cls._SUPPORTED_SM_VERSIONS: sm_list = "/".join(f"SM{v}" for v in sorted(cls._SUPPORTED_SM_VERSIONS)) return f"CuteDslB12xFusedMoE requires {sm_list}, got SM{sm_version}" @@ -151,10 +156,11 @@ def _is_cutlass_dsl_runtime_available(cls) -> bool: @staticmethod def _get_cutlass_dsl_cuda_major() -> Optional[int]: try: - from cutlass.base_dsl.version_info import CUDA_VERSION + import cutlass except ImportError: return None - major = getattr(CUDA_VERSION, "major", None) + cuda_version = getattr(cutlass, "CUDA_VERSION", None) + major = getattr(cuda_version, "major", None) if isinstance(major, int): return major return None diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 973f47739e1f..f7c6a0167cda 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -23,6 +23,9 @@ SM120/SM121 hardware. """ +import builtins +import sys +import types from unittest.mock import patch import pytest @@ -47,17 +50,6 @@ def test_can_implement_rejects_unsupported_sm(sm_version): assert reason is not None and f"SM{sm_version}" in reason -@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -def test_can_implement_accepts_supported_sm_with_nvfp4(sm_version): - with ( - patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=sm_version), - patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), - ): - ok, reason = CuteDslB12xFusedMoE.can_implement(QuantAlgo.NVFP4) - assert ok - assert reason is None - - @pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) def test_can_implement_rejects_supported_sm_with_cuda12_cute_dsl(sm_version): with ( @@ -93,6 +85,35 @@ def test_can_implement_accepts_supported_sm_with_cuda13_cute_dsl(sm_version): assert reason is None +def test_get_cutlass_dsl_cuda_major_returns_none_when_cutlass_missing(monkeypatch): + real_import = builtins.__import__ + + def _raise_on_cutlass(name, *args, **kwargs): + if name == "cutlass": + raise ImportError("cutlass not installed (simulated)") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _raise_on_cutlass) + assert CuteDslB12xFusedMoE._get_cutlass_dsl_cuda_major() is None + + +@pytest.mark.parametrize( + ("cuda_version", "expected"), + [ + (None, None), + (types.SimpleNamespace(), None), + (types.SimpleNamespace(major="13"), None), + (types.SimpleNamespace(major=13), 13), + ], +) +def test_get_cutlass_dsl_cuda_major_reads_public_cutlass_api(monkeypatch, cuda_version, expected): + cutlass_module = types.ModuleType("cutlass") + if cuda_version is not None: + cutlass_module.CUDA_VERSION = cuda_version + monkeypatch.setitem(sys.modules, "cutlass", cutlass_module) + assert CuteDslB12xFusedMoE._get_cutlass_dsl_cuda_major() == expected + + @pytest.mark.parametrize( "quant_algo", [ @@ -177,20 +198,6 @@ def test_get_moe_cls_cutedsl_returns_plain_cutedsl_on_unsupported_sm(): assert cls is CuteDslFusedMoE -@pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -def test_get_moe_cls_cutedsl_selects_b12x_on_supported_sm(sm_version): - """CUTEDSL + NVFP4 + SM12x + CUDA13 CuTe DSL -> CuteDslB12xFusedMoE.""" - cfg = ModelConfig() - cfg.moe_backend = "CUTEDSL" - cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) - with ( - patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version), - patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=13), - ): - cls = get_moe_cls(cfg) - assert cls is CuteDslB12xFusedMoE - - @pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) def test_get_moe_cls_cutedsl_falls_back_to_cutlass_on_supported_sm_with_cuda12(sm_version): """CUTEDSL + NVFP4 + SM12x avoids the CUDA 12 CuTe DSL JIT crash.""" From 28b2baec44d24576e052093172aa06a210a8a7a3 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Sat, 4 Jul 2026 14:54:16 -0700 Subject: [PATCH 4/5] [#15853][fix] Limit B12x gate to pure NVFP4 Signed-off-by: Mihai Chiorean --- .../_torch/modules/fused_moe/create_moe.py | 2 +- .../moe/test_cute_dsl_b12x_moe_backend.py | 17 +++++++++++++++++ 2 files changed, 18 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 865047aebe94..695926fbc0c0 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -85,7 +85,7 @@ def get_moe_cls( # hybrid CUTLASS-prefill / FlashInfer NVFP4 MoE decode backend # (CuteDslB12xFusedMoE). Prefer it when flashinfer is importable; # otherwise fall through to CuteDslFusedMoE for SM100 / SM103. - if quant_config.quant_mode.has_nvfp4(): + if quant_config.quant_algo == QuantAlgo.NVFP4: from tensorrt_llm._utils import get_sm_version sm_version = get_sm_version() if sm_version in CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS: diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index f7c6a0167cda..820f49394727 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -122,6 +122,7 @@ def test_get_cutlass_dsl_cuda_major_reads_public_cutlass_api(monkeypatch, cuda_v QuantAlgo.FP8_BLOCK_SCALES, QuantAlgo.W4A16_MXFP4, QuantAlgo.W4A8_MXFP4_FP8, + QuantAlgo.W4A8_MXFP4_MXFP8, ], ) def test_can_implement_rejects_non_nvfp4(quant_algo): @@ -135,6 +136,22 @@ def test_can_implement_rejects_non_nvfp4(quant_algo): assert reason is not None and "NVFP4" in reason +def test_get_moe_cls_cutedsl_does_not_apply_b12x_gate_to_mixed_fp8_fp4(): + cfg = ModelConfig() + cfg.moe_backend = "CUTEDSL" + cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8) + with ( + patch("tensorrt_llm._utils.get_sm_version", return_value=120), + patch.object( + CuteDslB12xFusedMoE, + "_get_cutlass_dsl_cuda_major", + side_effect=AssertionError("B12x gate should only run for pure NVFP4"), + ), + ): + cls = get_moe_cls(cfg) + assert cls is CutlassFusedMoE + + def test_can_implement_rejects_swiglu_gptoss_style(): with ( patch(f"{_FUSED_MOE_MODULE}.get_sm_version", return_value=120), From 324d004783b953816bf524462778253325180913 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Thu, 6 Aug 2026 11:38:56 -0700 Subject: [PATCH 5/5] [#15853][test] Cover unavailable CUDA13 B12x runtime Signed-off-by: Mihai Chiorean --- .../_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py index 820f49394727..336e2742b9e5 100644 --- a/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_cute_dsl_b12x_moe_backend.py @@ -216,14 +216,17 @@ def test_get_moe_cls_cutedsl_returns_plain_cutedsl_on_unsupported_sm(): @pytest.mark.parametrize("sm_version", sorted(CuteDslB12xFusedMoE._SUPPORTED_SM_VERSIONS)) -def test_get_moe_cls_cutedsl_falls_back_to_cutlass_on_supported_sm_with_cuda12(sm_version): - """CUTEDSL + NVFP4 + SM12x avoids the CUDA 12 CuTe DSL JIT crash.""" +@pytest.mark.parametrize("cuda_major", [12, None]) +def test_get_moe_cls_cutedsl_falls_back_to_cutlass_on_unavailable_cuda13_runtime( + sm_version, cuda_major +): + """CUTEDSL + NVFP4 + SM12x avoids the unsupported CuTe DSL JIT path.""" cfg = ModelConfig() cfg.moe_backend = "CUTEDSL" cfg.quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) with ( patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version), - patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=12), + patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=cuda_major), ): cls = get_moe_cls(cfg) assert cls is CutlassFusedMoE