[#11932][fix] Enable FP4 MoE dispatch for SM120/SM121 (DGX Spark)

cpp/include/tensorrt_llm/common/cudaUtils.h

@@ -304,10 +304,10 @@ inline int getSMVersion(bool queryRealSmArch = false)
     return sm;
 }
 
-inline bool isSM100Family()
+inline bool isBlackwellFamily()
 {
     int const sm = getSMVersion();
-    return sm == 100 || sm == 103; // To be continued...
+    return sm == 100 || sm == 103 || sm == 120 || sm == 121;
 }
 
 inline int getDevice()

cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp

@@ -52,7 +52,7 @@ FmhaDispatcher::FmhaDispatcher(MHARunnerFixedParams fixedParams)
     // TRTLLM-GEN only supports power of 2 head sizes.
     // The exception will fall back to fmha v2.
     // Please update fmha_v2/setup.py if you want to add more supported head sizes.
-    , mUseTllmGen(tensorrt_llm::common::isSM100Family() && fixedParams.headSize != 80 && fixedParams.headSize != 72)
+    , mUseTllmGen(tensorrt_llm::common::isBlackwellFamily() && fixedParams.headSize != 80 && fixedParams.headSize != 72)
 {
     if (mUseTllmGen)
     {

cpp/tensorrt_llm/kernels/xqaDispatcher.cpp

@@ -131,7 +131,7 @@ QKVPreprocessingParams<T, KVCacheBuffer> makeQKVPreprocessingParams(XQAParams co
 XqaDispatcher::XqaDispatcher(XqaFixedParams fixedParams)
     : mFixedParams(fixedParams)
     , mQDataType(mFixedParams.inputDataType)
-    , mUseTllmGen(tensorrt_llm::common::isSM100Family())
+    , mUseTllmGen(tensorrt_llm::common::isBlackwellFamily())
     , mMultiProcessorCount(getMultiProcessorCount())
 {
     if (mUseTllmGen)

cpp/tensorrt_llm/thop/attentionOp.cpp

@@ -493,7 +493,7 @@ class Runner : public RunnerBase
             }
             if (op.mIsSpecDecodingEnabled && op.mUseSpecDecoding)
             {
-                bool useTllmGen = tensorrt_llm::common::isSM100Family();
+                bool useTllmGen = tensorrt_llm::common::isBlackwellFamily();
                 if (useTllmGen)
                 {
                     TORCH_CHECK(spec_decoding_tensor_params.size() == 6,

cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp

@@ -48,7 +48,8 @@ std::vector<torch::Tensor> run_fp4_block_scale_moe_runner(torch::optional<torch:
     torch::optional<torch::Tensor> const& out_tensor = torch::nullopt)
 {
     TORCH_CHECK(dtype == btg::Dtype::E4m3 || dtype == btg::Dtype::E2m1, "dtype can only be e4m3 or e2m1.");
-    TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP4 block scale MOE");
+    TORCH_CHECK(tensorrt_llm::common::isBlackwellFamily(),
+        "Blackwell family (SM100/103/120/121) is required for FP4 block scale MOE");
     TORCH_CHECK(tile_tokens_dim == 8 || tile_tokens_dim == 16 || tile_tokens_dim == 32 || tile_tokens_dim == 64
             || tile_tokens_dim == 128 || tile_tokens_dim == 256,
         "tile_tokens_dim must be 8, 16, 32, 64, 128, 256");

cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp

@@ -46,7 +46,8 @@ at::Tensor run_fp8_block_scale_moe(at::optional<at::Tensor> const& routing_logit
     MoeRunnerType& moe_runner, int64_t moeConfigIndex, std::optional<at::Tensor> const& topk_weights,
     std::optional<at::Tensor> const& topk_ids, std::optional<at::Tensor> const& out_tensor = std::nullopt)
 {
-    TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP8 block scale MOE");
+    TORCH_CHECK(tensorrt_llm::common::isBlackwellFamily(),
+        "Blackwell family (SM100/103/120/121) is required for FP8 block scale MOE");
 
     if (topk_ids.has_value() && topk_weights.has_value())
     {

cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp

@@ -38,7 +38,8 @@ torch::Tensor fp8_per_tensor_scale_moe_runner(torch::optional<torch::Tensor> con
     int64_t const tile_tokens_dim, int64_t const routing_method_type,
     torch::optional<torch::Tensor> const& topk_weights, torch::optional<torch::Tensor> const& topk_ids)
 {
-    TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP8 block scale MOE");
+    TORCH_CHECK(tensorrt_llm::common::isBlackwellFamily(),
+        "Blackwell family (SM100/103/120/121) is required for FP8 per-tensor scale MOE");
     TORCH_CHECK(tile_tokens_dim == 8 || tile_tokens_dim == 16 || tile_tokens_dim == 32 || tile_tokens_dim == 64
             || tile_tokens_dim == 128 || tile_tokens_dim == 192 || tile_tokens_dim == 256,
         "tile_tokens_dim must be 8, 16, 32, 64, 128, 256");

cpp/tensorrt_llm/thop/fp8Quantize.cpp

@@ -69,7 +69,7 @@ std::tuple<at::Tensor, at::Tensor> fp8_quantize_1x128(at::Tensor const& self, bo
         act_buffer, act_scale_buffer, reinterpret_cast<__nv_bfloat16 const*>(self.data_ptr()), n, m, stream, use_ue8m0);
 
     // Post-process the scale tensor for sm100 gemm/moe kernel
-    if (tensorrt_llm::common::isSM100Family())
+    if (tensorrt_llm::common::isBlackwellFamily())
     {
         auto const num_n_blocks = (n + 127) / 128;
         auto const act_scal_elesize = num_n_blocks * m_padded;

cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp

@@ -52,7 +52,8 @@ torch::Tensor dtype_mxe2m1_block_scale_moe_runner(torch::optional<torch::Tensor>
     torch::optional<torch::Tensor> const& topk_weights, torch::optional<torch::Tensor> const& topk_ids,
     torch::optional<torch::Tensor> const& out_tensor)
 {
-    TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by MXFP4 block scale MOE");
+    TORCH_CHECK(tensorrt_llm::common::isBlackwellFamily(),
+        "Blackwell family (SM100/103/120/121) is required for MXFP4 block scale MOE");
     TORCH_CHECK(tile_tokens_dim == 8 || tile_tokens_dim == 16 || tile_tokens_dim == 32 || tile_tokens_dim == 64
             || tile_tokens_dim == 128 || tile_tokens_dim == 256,
         "tile_tokens_dim must be 8, 16, 32, 64, 128, 256");

cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp

@@ -170,7 +170,7 @@ class PrepareCustomMaskTest : public ::testing::Test
 protected:
     static bool shouldSkip()
     {
-        return !tensorrt_llm::common::isSM100Family();
+        return !tensorrt_llm::common::isBlackwellFamily();
     }
 
     void SetUp() override

tensorrt_llm/_torch/autotuner.py

@@ -1301,6 +1301,17 @@ def _optimization_profiles(
         for spec in tuning_config.dynamic_tensor_specs:
             assert callable(spec.gen_tuning_buckets) or isinstance(spec.gen_tuning_buckets, (list, tuple)), \
                 "The given dynamic dimension must provide a opt value generation function or a list of opt values"
+            if spec.input_idx >= len(base_profile.shapes):
+                logger.debug(
+                    f"[Autotuner] Skipping DynamicTensorSpec with input_idx={spec.input_idx}: "
+                    f"only {len(base_profile.shapes)} inputs available.")
+                continue
+            if spec.dim_idx >= len(base_profile.shapes[spec.input_idx]):
+                logger.debug(
+                    f"[Autotuner] Skipping DynamicTensorSpec with dim_idx={spec.dim_idx} for "
+                    f"input {spec.input_idx}: shape has only "
+                    f"{len(base_profile.shapes[spec.input_idx])} dims.")
+                continue
             if self.skip_dynamic_tuning_buckets:
                 if spec.map_to_tuning_buckets is not None:
                     # Still include the bucketed value of the actual shape so the
@@ -1358,10 +1369,21 @@ def _optimization_profiles(
 
             # Adjust the profile to satisfy the constraints
             for spec in tuning_config.constraint_specs:
-                min_value = opt_value = max_value = spec.infer_shape(
-                    p.get_opt_shapes())
+                if spec.input_idx >= len(p.shapes):
+                    logger.debug(
+                        f"[Autotuner] Skipping ConstraintSpec with input_idx={spec.input_idx}: "
+                        f"only {len(p.shapes)} inputs available.")
+                    continue
                 if p.shapes[spec.input_idx] == [StaticDim(0)]:
                     continue
+                if spec.dim_idx >= len(p.shapes[spec.input_idx]):
+                    logger.debug(
+                        f"[Autotuner] Skipping ConstraintSpec with dim_idx={spec.dim_idx} for "
+                        f"input {spec.input_idx}: shape has only "
+                        f"{len(p.shapes[spec.input_idx])} dims.")
+                    continue
+                min_value = opt_value = max_value = spec.infer_shape(
+                    p.get_opt_shapes())
                 p.shapes[spec.input_idx][spec.dim_idx] = DynamicDim(
                     min_value, opt_value, max_value)
             generated_profiles.append(p)
@@ -1395,6 +1417,21 @@ def _find_nearest_profile(
         base_profile = list(list(shape) for shape in shapes)
 
         for spec in dynamic_tensor_specs:
+            # Bounds check: skip specs that reference inputs or dimensions not present in the
+            # current shapes tuple. This can happen on hardware (e.g. SM121 / DGX Spark) where
+            # ops produce fewer or differently-shaped tensors than the specs were authored for.
+            if spec.input_idx >= len(base_profile):
+                logger.debug(
+                    f"[Autotuner] Skipping DynamicTensorSpec with input_idx={spec.input_idx}: "
+                    f"only {len(base_profile)} inputs available.")
+                continue
+            if spec.dim_idx >= len(base_profile[spec.input_idx]):
+                logger.debug(
+                    f"[Autotuner] Skipping DynamicTensorSpec with dim_idx={spec.dim_idx} for "
+                    f"input {spec.input_idx}: shape has only {len(base_profile[spec.input_idx])} dims."
+                )
+                continue
+
             # During runtime: apply map_to_tuning_buckets to map input to bucket
             # During tuning: no mapper, use raw bucket value
             if apply_map_to_tuning_buckets:
@@ -1409,8 +1446,20 @@ def _find_nearest_profile(
 
         # associated dimensions dependent on other free dynamic dimensions, so assign -1 in the profile
         for spec in constraint_specs:
+            # Bounds check: same defensive guard as above for constraint specs.
+            if spec.input_idx >= len(base_profile):
+                logger.debug(
+                    f"[Autotuner] Skipping ConstraintSpec with input_idx={spec.input_idx}: "
+                    f"only {len(base_profile)} inputs available.")
+                continue
             if base_profile[spec.input_idx] == [0]:
                 continue
+            if spec.dim_idx >= len(base_profile[spec.input_idx]):
+                logger.debug(
+                    f"[Autotuner] Skipping ConstraintSpec with dim_idx={spec.dim_idx} for "
+                    f"input {spec.input_idx}: shape has only {len(base_profile[spec.input_idx])} dims."
+                )
+                continue
             base_profile[spec.input_idx][spec.dim_idx] = -1
 
         return tuple(tuple(shape) for shape in base_profile)

tensorrt_llm/_torch/custom_ops/torch_custom_ops.py

@@ -809,13 +809,13 @@ def get_valid_tactics(self, inputs: List[torch.Tensor],
         # Add CuteDSL tactics if available
         if self._is_backend_allowed("cutedsl"):
             if IS_CUTLASS_DSL_AVAILABLE:
-                # Check SM version first - CuteDSL NVFP4 only supports SM 100 (B200)
+                # Check SM version first - CuteDSL NVFP4 only supports Blackwell
                 sm_version = get_sm_version()
-                if sm_version not in [100, 103]:
+                if sm_version not in [100, 103, 120, 121]:
                     if self._is_only_backend("cutedsl"):
                         # Explicitly forced CuteDSL but SM version not supported
                         raise ValueError(
-                            f"CuteDSL NVFP4 backend requires SM 100 (B200) or SM 103 (B300), but got SM {sm_version}. "
+                            f"CuteDSL NVFP4 backend requires Blackwell (SM100/103/120/121), but got SM {sm_version}. "
                             f"CuteDSL NVFP4 is not supported on this GPU architecture. "
                             "Please add other backends to allowed_backends.")
                 else:

tensorrt_llm/_torch/model_config.py

@@ -11,11 +11,11 @@
 import transformers
 from transformers.utils import HF_MODULES_CACHE
 
+import tensorrt_llm._utils as _utils
 from tensorrt_llm import logger
 from tensorrt_llm._torch.pyexecutor.config_utils import (
     get_qwen3_hybrid_num_attention_layers, is_nemotron_hybrid, is_qwen3_hybrid,
     load_pretrained_config)
-from tensorrt_llm._utils import get_sm_version, torch_dtype_to_binding
 from tensorrt_llm.bindings import LayerType as LayerTypeCpp
 from tensorrt_llm.functional import AllReduceStrategy
 from tensorrt_llm.llmapi.llm_args import (DeepSeekSparseAttentionConfig,
@@ -251,15 +251,14 @@ def resolve_moe_backend(moe_backend: str, architecture: str) -> str:
         if moe_backend.upper() != "AUTO":
             return moe_backend
 
+        # SM100 family: use TRTLLM (trtllm-gen kernels need tcgen05.mma, not available on SM120/SM121)
+        if _utils.is_sm_100f():
+            return "TRTLLM"
+
         if architecture == "GptOssForCausalLM":
-            sm_version = get_sm_version()
-            # Select the best performing backend based on SM version
-            if 100 <= sm_version < 120:  # Blackwell
-                return "TRTLLM"
-            elif 90 <= sm_version < 100:  # Hopper
+            sm_version = _utils.get_sm_version()
+            if 90 <= sm_version < 100:  # Hopper
                 return "TRITON"
-            else:
-                return "CUTLASS"  # Fallback to CUTLASS for other SM versions (e.g., SM120)
 
         return "CUTLASS"
 
@@ -335,7 +334,7 @@ def load_modelopt_quant_config(quant_config_file, checkpoint_dir,
     def get_mxfp4_quant_algo(moe_backend, is_dynamic_quant=False):
         quant_algo = ModelConfig.override_quant_algo()
         if quant_algo is None and not is_dynamic_quant:
-            if get_sm_version() >= 100:
+            if _utils.get_sm_version() >= 100:
                 if moe_backend == 'TRITON':
                     return QuantAlgo.W4A8_MXFP4_FP8
                 else:
@@ -677,7 +676,7 @@ def ceil_div(a, b):
             num_rnn_layers=0,
             num_heads=num_heads,
             hidden_size=hidden_size,
-            data_type=torch_dtype_to_binding(
+            data_type=_utils.torch_dtype_to_binding(
                 self.pretrained_config.torch_dtype))
 
         # For kv cache size calculation: set tokens_per_block

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py

@@ -19,7 +19,7 @@
 import torch
 import torch.nn.functional as F
 
-from tensorrt_llm._utils import get_sm_version, is_sm_100f
+import tensorrt_llm._utils as _utils
 from tensorrt_llm.models.modeling_utils import QuantAlgo
 
 from ...autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec,
@@ -62,8 +62,8 @@ def cute_dsl_fp8_group_blockwise_gemm_ref(
     a_tmp = a.as_strided((m, k, 1), (k, 1, m * k))
     b_tmp = b.permute(1, 2, 0)
 
-    # Note: we have different output scale shape for fp8_quantize_1x128, so we need to handle it differently for sm100 and other archs.
-    if is_sm_100f():
+    # Note: we have different output scale shape for fp8_quantize_1x128, so we need to handle it differently for Blackwell and other archs.
+    if _utils.is_blackwell():
         input_scale_tmp = a_sf.permute(1, 0).as_strided((m, w_k, 1),
                                                         (1, m, m * w_k))
     else:
@@ -339,7 +339,7 @@ def can_implement(
         Check if CuteDslFusedMoE can implement the given quantization algorithm.
 
         CuteDslFusedMoE supports:
-        - NVFP4: SM in {100, 103}
+        - NVFP4: SM in {100, 103, 120, 121}
 
         Does NOT support unquantized mode. Output dtype is hardcoded to bfloat16.
         Does NOT support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit).
@@ -356,7 +356,7 @@ def can_implement(
         """
         from .interface import _warn_and_return
 
-        sm_version = get_sm_version()
+        sm_version = _utils.get_sm_version()
 
         # CuteDslFusedMoE requires at least SM90
         if sm_version < 90:
@@ -381,11 +381,12 @@ def can_implement(
                 "CuteDslFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)"
             )
 
-        # NVFP4 - SM in {100, 103}
+        # NVFP4 - SM100 family only (SM120/SM121 has scale dtype mismatch)
         if quant_algo == QuantAlgo.NVFP4:
-            if sm_version not in {100, 103}:
+            if not _utils.is_sm_100f(sm_version):
                 return _warn_and_return(
-                    f"NVFP4 requires SM100 or SM103, got SM{sm_version}")
+                    f"CuteDSL NVFP4 requires SM100 family (SM120/SM121 excluded due to scale dtype mismatch), got SM{sm_version}"
+                )
             return True, None
 
         return _warn_and_return(

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

@@ -21,9 +21,9 @@
 import torch
 from torch import nn
 
+import tensorrt_llm._utils as _utils
 from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe
 from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll
-from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.logger import logger
 from tensorrt_llm.models.modeling_utils import QuantAlgo
 
@@ -129,7 +129,7 @@ def can_implement(
         """
         from .interface import _warn_and_return
 
-        sm_version = get_sm_version()
+        sm_version = _utils.get_sm_version()
 
         # TRTLLMGenFusedMoE requires SM in {100, 103}
         if sm_version not in {100, 103}:
@@ -203,7 +203,7 @@ def __init__(
             activation_type=activation_type,
         )
 
-        sm_version = get_sm_version()
+        sm_version = _utils.get_sm_version()
         if sm_version >= 120:
             raise NotImplementedError(
                 "TRTLLMGenFusedMoE does not support SM120 and above.")

tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py

@@ -223,11 +223,11 @@ def select_op(module: 'MoE') -> MoEOp:
         from .moe_op_deepgemm import DeepGemmMoEOp
 
         # Check if we should use DeepGemm op
-        # Blackwell has SM version 100
-        is_blackwell = is_sm_100f()
+        # DeepGemm supports SM100/SM103 (datacenter Blackwell) only
+        use_deepgemm_arch = is_sm_100f()
         has_block_fp8 = module.has_deepseek_fp8_block_scales
 
-        if is_blackwell and has_block_fp8:
+        if use_deepgemm_arch and has_block_fp8:
             # Use DeepGemm op for Blackwell with block FP8
             return DeepGemmMoEOp()
         else:

tensorrt_llm/_torch/modules/fused_moe/quantization.py

@@ -1099,7 +1099,7 @@ class DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm(
         DeepSeekFP8BlockScalesFusedMoEMethod):
 
     def _needs_e8m0_resmooth(self):
-        return is_sm_100f() or get_sm_version() == 120
+        return is_sm_100f() or get_sm_version() in (120, 121)
 
     def post_load_weights(self, module: torch.nn.Module):
         if self._needs_e8m0_resmooth():

tensorrt_llm/_utils.py

@@ -796,6 +796,18 @@ def is_sm_100f(sm_version=None):
     return sm_version == 100 or sm_version == 103
 
 
+def is_sm_120f(sm_version=None):
+    if sm_version is None:
+        sm_version = get_sm_version()
+    return sm_version == 120 or sm_version == 121
+
+
+def is_blackwell(sm_version=None):
+    if sm_version is None:
+        sm_version = get_sm_version()
+    return is_sm_100f(sm_version) or is_sm_120f(sm_version)
+
+
 def print_all_stacks():
     """Print stack traces for all threads"""
     for thread_id, frame in sys._current_frames().items():