[#15853][fix] Gate B12x MoE on CUDA13 CuTe DSL

tensorrt_llm/_torch/modules/fused_moe/create_moe.py

@@ -85,10 +85,19 @@ 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:
+                    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(

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_B12X = 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,51 @@ 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]:
+        """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}"
+
+        # 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 not cls._is_cutlass_dsl_runtime_available():
+            return (
+                "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 "
+                "'_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_B12X}.x or newer."
+            )
+        return None
+
+    @classmethod
+    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_B12X
+
+    @staticmethod
+    def _get_cutlass_dsl_cuda_major() -> Optional[int]:
+        try:
+            import cutlass
+        except ImportError:
+            return None
+        cuda_version = getattr(cutlass, "CUDA_VERSION", 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

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
@@ -48,13 +51,69 @@ def test_can_implement_rejects_unsupported_sm(sm_version):
 
 
 @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):
+def test_can_implement_rejects_supported_sm_with_cuda12_cute_dsl(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=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
+
+
+@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=sm_version),
+        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
+
+
+@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=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_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",
     [
@@ -63,26 +122,52 @@ def test_can_implement_accepts_supported_sm_with_nvfp4(sm_version):
         QuantAlgo.FP8_BLOCK_SCALES,
         QuantAlgo.W4A16_MXFP4,
         QuantAlgo.W4A8_MXFP4_FP8,
+        QuantAlgo.W4A8_MXFP4_MXFP8,
     ],
 )
 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_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):
+    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
 
 
 @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
@@ -131,18 +216,37 @@ 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_selects_b12x_on_supported_sm(sm_version):
-    """CUTEDSL + NVFP4 + SM120/121 + flashinfer importable → CuteDslB12xFusedMoE."""
+@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):
+    with (
+        patch("tensorrt_llm._utils.get_sm_version", return_value=sm_version),
+        patch.object(CuteDslB12xFusedMoE, "_get_cutlass_dsl_cuda_major", return_value=cuda_major),
+    ):
+        cls = get_moe_cls(cfg)
+    assert cls is CutlassFusedMoE
+
+
+@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=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_plain_cutedsl_when_flashinfer_missing(monkeypatch):
-    """CUTEDSL + NVFP4 + SM120/121 + flashinfer NOT importable → CuteDslFusedMoE."""
+    """CUTEDSL + NVFP4 + SM12x + flashinfer NOT importable -> CuteDslFusedMoE."""
     import builtins
 
     cfg = ModelConfig()
@@ -157,7 +261,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