From 373fb880a8499f1d94eaa116cae3c5075244c19d Mon Sep 17 00:00:00 2001
From: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Date: Tue, 17 Mar 2026 23:03:25 -0700
Subject: [PATCH 1/6] [#11932][fix] Enable FP4 MoE dispatch for SM120/SM121
 (DGX Spark)

Remove the NotImplementedError gate in TRTLLMGenFusedMoE.__init__ that
blocked ALL MoE models on SM120+ (DGX Spark / RTX 5090).  The underlying
CUTLASS kernels already have SM120 templates and PR #12141 fixed the FP4
GEMM shared-memory overflow on SM121, so the Python-side SM version
checks were the only remaining barrier.

Changes:
- tensorrt_llm/_utils.py: add is_sm_120f() and is_blackwell() helpers
- fused_moe_trtllm_gen.py: remove __init__ SM>=120 gate; extend
  can_implement() SM set {100,103} -> {100,103,120,121}
- fused_moe_cute_dsl.py: extend NVFP4 SM check to include 120/121;
  use is_blackwell() for FP8 scale layout (shared across Blackwell)
- model_config.py: route SM120/121 to TRTLLM backend in
  resolve_moe_backend() (was falling back to CUTLASS)
- torch_custom_ops.py: extend CuTE DSL NVFP4 dense GEMM SM check
- tests/integration/defs/conftest.py: add matching test helpers

Signed-off-by: Mihai <mihai@dgx-spark>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
---
 .../_torch/custom_ops/torch_custom_ops.py         |  6 +++---
 tensorrt_llm/_torch/model_config.py               |  4 ++--
 .../modules/fused_moe/fused_moe_cute_dsl.py       | 14 +++++++-------
 .../modules/fused_moe/fused_moe_trtllm_gen.py     | 15 +++++----------
 tensorrt_llm/_utils.py                            | 12 ++++++++++++
 tests/integration/defs/conftest.py                | 12 ++++++++++++
 6 files changed, 41 insertions(+), 22 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
index ee150d1be98a..2d5d37acfb06 100644
--- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
+++ b/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:
diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py
index b1a932953b06..0490d1fb1e0b 100644
--- a/tensorrt_llm/_torch/model_config.py
+++ b/tensorrt_llm/_torch/model_config.py
@@ -254,12 +254,12 @@ def resolve_moe_backend(moe_backend: str, architecture: str) -> str:
         if architecture == "GptOssForCausalLM":
             sm_version = get_sm_version()
             # Select the best performing backend based on SM version
-            if 100 <= sm_version < 120:  # Blackwell
+            if 100 <= sm_version < 120 or sm_version in (120, 121):  # Blackwell
                 return "TRTLLM"
             elif 90 <= sm_version < 100:  # Hopper
                 return "TRITON"
             else:
-                return "CUTLASS"  # Fallback to CUTLASS for other SM versions (e.g., SM120)
+                return "CUTLASS"  # Fallback for other SM versions
 
         return "CUTLASS"
 
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
index 1273262f5f42..4ec9c778a0f7 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
+++ b/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
+from tensorrt_llm._utils import get_sm_version, is_blackwell, is_sm_100f
 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 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).
@@ -381,11 +381,11 @@ def can_implement(
                 "CuteDslFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)"
             )
 
-        # NVFP4 - SM in {100, 103}
+        # NVFP4 - SM in {100, 103, 120, 121} (Blackwell family)
         if quant_algo == QuantAlgo.NVFP4:
-            if sm_version not in {100, 103}:
+            if sm_version not in {100, 103, 120, 121}:
                 return _warn_and_return(
-                    f"NVFP4 requires SM100 or SM103, got SM{sm_version}")
+                    f"NVFP4 requires Blackwell (SM100/103/120/121), got SM{sm_version}")
             return True, None
 
         return _warn_and_return(
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 23354f5a5b34..bb79ae91d3fc 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
@@ -62,7 +62,7 @@ class TRTLLMGenFusedMoE(MoE):
         aux_stream_dict (Optional[Dict[AuxStreamType, torch.cuda.Stream]]): Auxiliary CUDA streams for overlapping.
 
     MoE torch custom op:
-        Only support min-latency mode now (SM100 Blackwell only).
+        Only support min-latency mode now (Blackwell: SM100/103/120/121).
         Quant: fp8 block scales quant and nvfp4 quant and w4a16_mxfp4 quant
             FusedMoE Op: routing(topK, etc.) + scatter + gemm1 + swiglu + gemm2 + finalize MoeRoute
 
@@ -107,7 +107,7 @@ def can_implement(
         """
         Check if TRTLLMGenFusedMoE can implement the given quantization algorithm.
 
-        TRTLLMGenFusedMoE only supports SM in {100, 103} and the following quantizations:
+        TRTLLMGenFusedMoE only supports SM in {100, 103, 120, 121} and the following quantizations:
         - NVFP4
         - FP8_BLOCK_SCALES
         - W4A8_NVFP4_FP8
@@ -131,10 +131,10 @@ def can_implement(
 
         sm_version = get_sm_version()
 
-        # TRTLLMGenFusedMoE requires SM in {100, 103}
-        if sm_version not in {100, 103}:
+        # TRTLLMGenFusedMoE requires SM in {100, 103, 120, 121} (Blackwell family)
+        if sm_version not in {100, 103, 120, 121}:
             return _warn_and_return(
-                f"TRTLLMGenFusedMoE requires SM100 or SM103, got SM{sm_version}"
+                f"TRTLLMGenFusedMoE requires Blackwell (SM100/103/120/121), got SM{sm_version}"
             )
 
         # Check dtype_activation: only bfloat16 is supported
@@ -203,11 +203,6 @@ def __init__(
             activation_type=activation_type,
         )
 
-        sm_version = get_sm_version()
-        if sm_version >= 120:
-            raise NotImplementedError(
-                "TRTLLMGenFusedMoE does not support SM120 and above.")
-
         assert not self.smart_router, "Smart router is not supported in TRTLLMGenFusedMoE."
 
         self.use_flashinfer = self._check_op_backend_support()
diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py
index 47a6a88499ea..580928962507 100644
--- a/tensorrt_llm/_utils.py
+++ b/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():
diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py
index cd399acb4155..82948213b539 100644
--- a/tests/integration/defs/conftest.py
+++ b/tests/integration/defs/conftest.py
@@ -1919,6 +1919,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 get_gpu_device_list():
     "get device list"
     with tempfile.TemporaryDirectory() as temp_dirname:

From 92379a98bcd1d1f63143088cadbeae4a610b218d Mon Sep 17 00:00:00 2001
From: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Date: Tue, 17 Mar 2026 23:12:57 -0700
Subject: [PATCH 2/6] fix: address code review findings for MoE SM120 gate

- model_config.py: use is_blackwell() helper instead of redundant conditional
- quantization.py: add SM121 to e8m0 resmooth check (was SM120 only)
- fused_moe_cute_dsl.py: remove dead is_sm_100f import

Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
---
 tensorrt_llm/_torch/model_config.py                         | 2 +-
 tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py | 2 +-
 tensorrt_llm/_torch/modules/fused_moe/quantization.py       | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py
index 0490d1fb1e0b..28b48d4a6ede 100644
--- a/tensorrt_llm/_torch/model_config.py
+++ b/tensorrt_llm/_torch/model_config.py
@@ -254,7 +254,7 @@ def resolve_moe_backend(moe_backend: str, architecture: str) -> str:
         if architecture == "GptOssForCausalLM":
             sm_version = get_sm_version()
             # Select the best performing backend based on SM version
-            if 100 <= sm_version < 120 or sm_version in (120, 121):  # Blackwell
+            if is_blackwell(sm_version):  # Blackwell (SM100/103/120/121)
                 return "TRTLLM"
             elif 90 <= sm_version < 100:  # Hopper
                 return "TRITON"
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
index 4ec9c778a0f7..6633054b4104 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
+++ b/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_blackwell, is_sm_100f
+from tensorrt_llm._utils import get_sm_version, is_blackwell
 from tensorrt_llm.models.modeling_utils import QuantAlgo
 
 from ...autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec,
diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py
index c30b7d771aa9..91fe5b66ae97 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py
+++ b/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():

From 813a3ac57667525a8713727e2f5fb3afc78fe8ca Mon Sep 17 00:00:00 2001
From: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Date: Tue, 17 Mar 2026 23:19:33 -0700
Subject: [PATCH 3/6] fix: rename shadowed is_blackwell local in moe_op.py

The local variable `is_blackwell = is_sm_100f()` shadowed the new
module-level `is_blackwell()` utility from _utils.py, which covers
SM100/103/120/121. Renamed to `use_deepgemm_arch` to clarify that
DeepGemm only supports SM100/SM103, avoiding confusion with the
broader is_blackwell() predicate.

Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
---
 tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py b/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py
index c9771996d55b..96358c82d6ac 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op.py
+++ b/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:

From 9e87af11af7e2f75848434525ca6df4d9ea9abab Mon Sep 17 00:00:00 2001
From: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Date: Tue, 17 Mar 2026 23:34:21 -0700
Subject: [PATCH 4/6] fix: apply pre-commit formatting

Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
---
 tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
index 6633054b4104..997601196a27 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
@@ -385,7 +385,8 @@ def can_implement(
         if quant_algo == QuantAlgo.NVFP4:
             if sm_version not in {100, 103, 120, 121}:
                 return _warn_and_return(
-                    f"NVFP4 requires Blackwell (SM100/103/120/121), got SM{sm_version}")
+                    f"NVFP4 requires Blackwell (SM100/103/120/121), got SM{sm_version}"
+                )
             return True, None
 
         return _warn_and_return(

From d67a9b25cff95099c6099366ce79dd2a345fcbdb Mon Sep 17 00:00:00 2001
From: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Date: Mon, 23 Mar 2026 14:32:11 -0700
Subject: [PATCH 5/6] [#11932][fix] Add C++ isBlackwellFamily() for SM120/SM121
 MoE dispatch

Companion C++ change required for the Python gates to work end-to-end.
Without this, TORCH_CHECK in the thop layer blocks SM120/121 after
Python dispatch succeeds.

- Rename isSM100Family() to isBlackwellFamily() in cudaUtils.h
- Add SM120/SM121 to the Blackwell family check
- Update all callers across thop and kernel dispatchers
- Make AUTO MoE backend select TRTLLM for all architectures on Blackwell
- Fix import style per coding guidelines (CodeRabbit feedback)

Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
---
 cpp/include/tensorrt_llm/common/cudaUtils.h   |   4 +-
 cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp   |   2 +-
 cpp/tensorrt_llm/kernels/xqaDispatcher.cpp    |   2 +-
 cpp/tensorrt_llm/thop/attentionOp.cpp         |   2 +-
 cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp    |   3 +-
 cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp    |   3 +-
 .../thop/fp8PerTensorScaleMoe.cpp             |   3 +-
 cpp/tensorrt_llm/thop/fp8Quantize.cpp         |   2 +-
 cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp  |   3 +-
 .../kernels/prepareCustomMaskTest.cpp         |   2 +-
 run_e2e_test.sh                               | 115 ++++++++++++++++++
 run_precommit.sh                              |  10 ++
 run_test.sh                                   |  76 ++++++++++++
 tensorrt_llm/_torch/autotuner.py              |  53 +++++++-
 tensorrt_llm/_torch/model_config.py           |  19 ++-
 .../modules/fused_moe/fused_moe_trtllm_gen.py |  19 +--
 .../_torch/thop/parallel/test_cute_dsl_moe.py |   2 +-
 17 files changed, 289 insertions(+), 31 deletions(-)
 create mode 100755 run_e2e_test.sh
 create mode 100755 run_precommit.sh
 create mode 100755 run_test.sh

diff --git a/cpp/include/tensorrt_llm/common/cudaUtils.h b/cpp/include/tensorrt_llm/common/cudaUtils.h
index cd58a7abb5d9..ef1986e5c34f 100644
--- a/cpp/include/tensorrt_llm/common/cudaUtils.h
+++ b/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()
diff --git a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp
index 68e3e4d60040..28c81138d9c9 100644
--- a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp
+++ b/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)
     {
diff --git a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp
index 41aab02d3043..8d06164b9963 100644
--- a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp
+++ b/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)
diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp
index 9a7af4da49f6..dbdf12464496 100644
--- a/cpp/tensorrt_llm/thop/attentionOp.cpp
+++ b/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,
diff --git a/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp
index 3440a59737fc..95e399084bd2 100644
--- a/cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp
+++ b/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");
diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp
index 3c13c695991c..79f9e24e7602 100644
--- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp
+++ b/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())
     {
diff --git a/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp
index 092f8f013620..77285c37492d 100644
--- a/cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp
+++ b/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");
diff --git a/cpp/tensorrt_llm/thop/fp8Quantize.cpp b/cpp/tensorrt_llm/thop/fp8Quantize.cpp
index 9e94e02950f7..5bdab0143020 100644
--- a/cpp/tensorrt_llm/thop/fp8Quantize.cpp
+++ b/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;
diff --git a/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp
index 5e8331b77c3f..e3df2ff19e94 100644
--- a/cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp
+++ b/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");
diff --git a/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp b/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp
index 116013de2ffe..387eeae53f89 100644
--- a/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp
+++ b/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
diff --git a/run_e2e_test.sh b/run_e2e_test.sh
new file mode 100755
index 000000000000..b98da8174287
--- /dev/null
+++ b/run_e2e_test.sh
@@ -0,0 +1,115 @@
+#!/bin/bash
+# E2E test: serve a small MoE model on SM121 with NVFP4, verify it generates text
+# Tests the full path: model load → autotuner → MoE dispatch → CUTLASS GEMM → generation
+
+set -euo pipefail
+
+WORKTREE="/home/mihai/workspace/trtllm-pr12309"
+LOGDIR="${WORKTREE}/test_logs"
+mkdir -p "${LOGDIR}"
+TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+LOGFILE="${LOGDIR}/e2e_${TIMESTAMP}.log"
+MODEL="${1:-Qwen/Qwen1.5-MoE-A2.7B-Chat}"
+PORT=9123
+
+echo "=== E2E Test: $(date) ===" | tee "${LOGFILE}"
+echo "Model: ${MODEL}" | tee -a "${LOGFILE}"
+echo "Log: ${LOGFILE}" | tee -a "${LOGFILE}"
+free -h | tee -a "${LOGFILE}"
+
+echo "" | tee -a "${LOGFILE}"
+echo "=== Starting trtllm-serve in docker ===" | tee -a "${LOGFILE}"
+
+# Run trtllm-serve inside the container
+docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \
+  -v "${WORKTREE}:/workspace/TensorRT-LLM" \
+  -v /home/mihai/.cache/huggingface:/root/.cache/huggingface \
+  --network host \
+  -w /workspace/TensorRT-LLM \
+  -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+  -e HF_HOME=/root/.cache/huggingface \
+  --name trtllm-e2e-test \
+  tensorrt_llm/devel:latest \
+  bash -c "
+    set -euo pipefail
+    echo '=== Container started: '\$(date)' ==='
+
+    echo '=== Installing TRT-LLM ==='
+    pip install --no-build-isolation -e '.[devel]' 2>&1 | tail -3
+    echo '=== Install done: '\$(date)' ==='
+
+    echo ''
+    echo '=== Starting trtllm-serve on port ${PORT} ==='
+    echo '=== Model: ${MODEL} ==='
+    echo '=== Backend: pytorch (default AUTO MoE) ==='
+    echo ''
+
+    # Start the server in background, log output
+    trtllm-serve ${MODEL} \
+      --backend pytorch \
+      --port ${PORT} \
+      --max_batch_size 1 \
+      --max_seq_len 256 \
+      2>&1 &
+    SERVER_PID=\$!
+
+    # Wait for server to be ready (poll health endpoint)
+    echo 'Waiting for server to start...'
+    for i in \$(seq 1 1200); do
+      if curl -s http://localhost:${PORT}/health > /dev/null 2>&1; then
+        echo \"Server ready after \${i}s\"
+        break
+      fi
+      if ! kill -0 \$SERVER_PID 2>/dev/null; then
+        echo 'Server process died!'
+        wait \$SERVER_PID || true
+        exit 1
+      fi
+      sleep 1
+    done
+
+    if ! curl -s http://localhost:${PORT}/health > /dev/null 2>&1; then
+      echo 'Server failed to start within 1200s'
+      kill \$SERVER_PID 2>/dev/null || true
+      exit 1
+    fi
+
+    echo ''
+    echo '=== Running inference test ==='
+    RESPONSE=\$(curl -s http://localhost:${PORT}/v1/completions \
+      -H 'Content-Type: application/json' \
+      -d '{
+        \"model\": \"${MODEL}\",
+        \"prompt\": \"The capital of France is\",
+        \"max_tokens\": 32,
+        \"temperature\": 0.0
+      }')
+
+    echo \"Response: \${RESPONSE}\"
+    echo ''
+
+    # Check if response contains generated text
+    if echo \"\${RESPONSE}\" | python3 -c 'import json,sys; d=json.load(sys.stdin); text=d[\"choices\"][0][\"text\"]; print(f\"Generated: {text}\"); assert len(text.strip()) > 0, \"Empty response\"' 2>&1; then
+      echo '=== E2E TEST PASSED ==='
+    else
+      echo '=== E2E TEST FAILED: no valid response ==='
+      echo \"Raw response: \${RESPONSE}\"
+      kill \$SERVER_PID 2>/dev/null || true
+      exit 1
+    fi
+
+    # Cleanup
+    kill \$SERVER_PID 2>/dev/null || true
+    wait \$SERVER_PID 2>/dev/null || true
+    echo '=== Done: '\$(date)' ==='
+  " 2>&1 | tee -a "${LOGFILE}"
+
+EXIT_CODE=${PIPESTATUS[0]}
+echo "" | tee -a "${LOGFILE}"
+echo "=== Exit code: ${EXIT_CODE} ===" | tee -a "${LOGFILE}"
+echo "=== Log saved to: ${LOGFILE} ===" | tee -a "${LOGFILE}"
+
+# Cleanup container if still running
+docker rm -f trtllm-e2e-test 2>/dev/null || true
+
+exit ${EXIT_CODE}
diff --git a/run_precommit.sh b/run_precommit.sh
new file mode 100755
index 000000000000..a5aeac31f272
--- /dev/null
+++ b/run_precommit.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+set -e
+cd /code/tensorrt_llm
+git config --global --add safe.directory /code/tensorrt_llm
+git config --global --add safe.directory /home/mihai/workspace/TensorRT-LLM
+git config --global --add safe.directory /home/mihai/workspace/trtllm-pr12309
+pip install pre-commit 2>&1 | tail -1
+CHANGED=$(git diff --name-only main...HEAD)
+echo "Changed files: $CHANGED"
+pre-commit run --show-diff-on-failure --files $CHANGED 2>&1
diff --git a/run_test.sh b/run_test.sh
new file mode 100755
index 000000000000..78d5f38315ba
--- /dev/null
+++ b/run_test.sh
@@ -0,0 +1,76 @@
+#!/bin/bash
+# Run CuTEDSL MoE tests for PR #12309 with proper logging and memory management
+# Usage: ./run_test.sh [test_filter]
+# Example: ./run_test.sh "test_nvfp4_grouped_gemm_blackwell"
+# Example: ./run_test.sh  (runs all CuTEDSL MoE tests)
+
+set -euo pipefail
+
+WORKTREE="/home/mihai/workspace/trtllm-pr12309"
+LOGDIR="${WORKTREE}/test_logs"
+mkdir -p "${LOGDIR}"
+TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+LOGFILE="${LOGDIR}/test_${TIMESTAMP}.log"
+
+TEST_FILTER="${1:-}"
+FILTER_ARGS=""
+if [[ -n "${TEST_FILTER}" ]]; then
+    FILTER_ARGS="-k '${TEST_FILTER}'"
+fi
+
+echo "=== Test Run: $(date) ===" | tee "${LOGFILE}"
+echo "Filter: ${TEST_FILTER:-<all>}" | tee -a "${LOGFILE}"
+echo "Log: ${LOGFILE}" | tee -a "${LOGFILE}"
+
+# Show system state before test
+echo "" | tee -a "${LOGFILE}"
+echo "=== PRE-TEST SYSTEM STATE ===" | tee -a "${LOGFILE}"
+free -h | tee -a "${LOGFILE}"
+swapon --show | tee -a "${LOGFILE}"
+nvidia-smi 2>/dev/null | tee -a "${LOGFILE}"
+echo "" | tee -a "${LOGFILE}"
+
+docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \
+  -v "${WORKTREE}:/workspace/TensorRT-LLM" \
+  -v /home/mihai/.cache/huggingface:/root/.cache/huggingface \
+  --network host \
+  -w /workspace/TensorRT-LLM \
+  -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+  -e CUDA_MODULE_LOADING=LAZY \
+  tensorrt_llm/devel:latest \
+  bash -c "
+    set -euo pipefail
+    echo '=== Container started: '\$(date)' ==='
+    echo '=== Installing TRT-LLM ==='
+    pip install --no-build-isolation -e '.[devel]' 2>&1 | tail -5
+    pip install -r requirements-dev.txt 2>&1 | tail -5
+    echo '=== Upgrading nvidia-cutlass-dsl to 4.3.5 ==='
+    pip install nvidia-cutlass-dsl==4.3.5 2>&1 | tail -3
+    echo '=== Install done: '\$(date)' ==='
+    echo ''
+    echo '=== PYTORCH_CUDA_ALLOC_CONF='\${PYTORCH_CUDA_ALLOC_CONF}' ==='
+    echo '=== Running pytest ==='
+    echo ''
+    cd tests/unittest
+    python3 -m pytest \
+      _torch/thop/parallel/test_cute_dsl_moe.py \
+      ${FILTER_ARGS} \
+      -v -s \
+      --tb=short \
+      --timeout=7200 \
+      -p no:cacheprovider \
+      2>&1
+    echo ''
+    echo '=== Tests complete: '\$(date)' ==='
+  " 2>&1 | tee -a "${LOGFILE}"
+
+EXIT_CODE=${PIPESTATUS[0]}
+
+# Show system state after test
+echo "" | tee -a "${LOGFILE}"
+echo "=== POST-TEST SYSTEM STATE ===" | tee -a "${LOGFILE}"
+free -h | tee -a "${LOGFILE}"
+echo "" | tee -a "${LOGFILE}"
+echo "=== Exit code: ${EXIT_CODE} ===" | tee -a "${LOGFILE}"
+echo "=== Log saved to: ${LOGFILE} ===" | tee -a "${LOGFILE}"
+exit ${EXIT_CODE}
diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py
index eaee69e54ebb..9c324129a2dd 100644
--- a/tensorrt_llm/_torch/autotuner.py
+++ b/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)
diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py
index 28b48d4a6ede..8169ae29afcf 100644
--- a/tensorrt_llm/_torch/model_config.py
+++ b/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
 
+        # Blackwell family: use TRTLLM for all MoE architectures
+        if _utils.is_blackwell():
+            return "TRTLLM"
+
         if architecture == "GptOssForCausalLM":
-            sm_version = get_sm_version()
-            # Select the best performing backend based on SM version
-            if is_blackwell(sm_version):  # Blackwell (SM100/103/120/121)
-                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 for other SM versions
 
         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
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 bb79ae91d3fc..8122aab97668 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
@@ -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
 
@@ -62,7 +62,7 @@ class TRTLLMGenFusedMoE(MoE):
         aux_stream_dict (Optional[Dict[AuxStreamType, torch.cuda.Stream]]): Auxiliary CUDA streams for overlapping.
 
     MoE torch custom op:
-        Only support min-latency mode now (Blackwell: SM100/103/120/121).
+        Only support min-latency mode now (SM100 Blackwell only).
         Quant: fp8 block scales quant and nvfp4 quant and w4a16_mxfp4 quant
             FusedMoE Op: routing(topK, etc.) + scatter + gemm1 + swiglu + gemm2 + finalize MoeRoute
 
@@ -107,7 +107,7 @@ def can_implement(
         """
         Check if TRTLLMGenFusedMoE can implement the given quantization algorithm.
 
-        TRTLLMGenFusedMoE only supports SM in {100, 103, 120, 121} and the following quantizations:
+        TRTLLMGenFusedMoE only supports SM in {100, 103} and the following quantizations:
         - NVFP4
         - FP8_BLOCK_SCALES
         - W4A8_NVFP4_FP8
@@ -129,12 +129,12 @@ 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, 120, 121} (Blackwell family)
-        if sm_version not in {100, 103, 120, 121}:
+        # TRTLLMGenFusedMoE requires SM in {100, 103}
+        if sm_version not in {100, 103}:
             return _warn_and_return(
-                f"TRTLLMGenFusedMoE requires Blackwell (SM100/103/120/121), got SM{sm_version}"
+                f"TRTLLMGenFusedMoE requires SM100 or SM103, got SM{sm_version}"
             )
 
         # Check dtype_activation: only bfloat16 is supported
@@ -203,6 +203,11 @@ def __init__(
             activation_type=activation_type,
         )
 
+        sm_version = _utils.get_sm_version()
+        if sm_version >= 120:
+            raise NotImplementedError(
+                "TRTLLMGenFusedMoE does not support SM120 and above.")
+
         assert not self.smart_router, "Smart router is not supported in TRTLLMGenFusedMoE."
 
         self.use_flashinfer = self._check_op_backend_support()
diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py
index 91024f5e4b77..c463e7ff5b1b 100644
--- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py
+++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py
@@ -313,7 +313,7 @@ def test_moe_swiglu(dtype: str, num_tokens: int, top_k: int, tile_size: int):
 
 
 @pytest.mark.skipif(
-    get_sm_version() not in (100, 103),
+    get_sm_version() not in (100, 103, 120, 121),
     reason="This test is only supported on SM 100 and SM 103 GPUs",
 )
 @pytest.mark.parametrize("tile_size", [128, 256])

From a8b3c230599195ef78056906d5e8630591f22d3b Mon Sep 17 00:00:00 2001
From: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Date: Mon, 23 Mar 2026 21:41:12 -0700
Subject: [PATCH 6/6] [#11932][fix] Scope SM120/SM121 MoE to CUTLASS backend
 only

- Narrow TRTLLM can_implement() to SM100 family (trtllm-gen kernels use
  tcgen05.mma instructions not available on SM120/SM121)
- Narrow CuteDSL can_implement() NVFP4 to SM100 family (scale dtype
  mismatch on SM121)
- Fix skip_no_sm120 test mark to include SM121 via is_sm_120f()
- Add test for resolve_moe_backend returning CUTLASS on SM121

Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
---
 run_e2e_test.sh                               | 115 ------------------
 run_precommit.sh                              |  10 --
 run_test.sh                                   |  76 ------------
 tensorrt_llm/_torch/model_config.py           |   4 +-
 .../modules/fused_moe/fused_moe_cute_dsl.py   |  12 +-
 tests/integration/defs/conftest.py            |   2 +-
 .../_torch/modules/moe/test_moe_backend.py    |  18 +++
 7 files changed, 27 insertions(+), 210 deletions(-)
 delete mode 100755 run_e2e_test.sh
 delete mode 100755 run_precommit.sh
 delete mode 100755 run_test.sh

diff --git a/run_e2e_test.sh b/run_e2e_test.sh
deleted file mode 100755
index b98da8174287..000000000000
--- a/run_e2e_test.sh
+++ /dev/null
@@ -1,115 +0,0 @@
-#!/bin/bash
-# E2E test: serve a small MoE model on SM121 with NVFP4, verify it generates text
-# Tests the full path: model load → autotuner → MoE dispatch → CUTLASS GEMM → generation
-
-set -euo pipefail
-
-WORKTREE="/home/mihai/workspace/trtllm-pr12309"
-LOGDIR="${WORKTREE}/test_logs"
-mkdir -p "${LOGDIR}"
-TIMESTAMP=$(date +%Y%m%d_%H%M%S)
-LOGFILE="${LOGDIR}/e2e_${TIMESTAMP}.log"
-MODEL="${1:-Qwen/Qwen1.5-MoE-A2.7B-Chat}"
-PORT=9123
-
-echo "=== E2E Test: $(date) ===" | tee "${LOGFILE}"
-echo "Model: ${MODEL}" | tee -a "${LOGFILE}"
-echo "Log: ${LOGFILE}" | tee -a "${LOGFILE}"
-free -h | tee -a "${LOGFILE}"
-
-echo "" | tee -a "${LOGFILE}"
-echo "=== Starting trtllm-serve in docker ===" | tee -a "${LOGFILE}"
-
-# Run trtllm-serve inside the container
-docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \
-  -v "${WORKTREE}:/workspace/TensorRT-LLM" \
-  -v /home/mihai/.cache/huggingface:/root/.cache/huggingface \
-  --network host \
-  -w /workspace/TensorRT-LLM \
-  -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
-  -e HF_HOME=/root/.cache/huggingface \
-  --name trtllm-e2e-test \
-  tensorrt_llm/devel:latest \
-  bash -c "
-    set -euo pipefail
-    echo '=== Container started: '\$(date)' ==='
-
-    echo '=== Installing TRT-LLM ==='
-    pip install --no-build-isolation -e '.[devel]' 2>&1 | tail -3
-    echo '=== Install done: '\$(date)' ==='
-
-    echo ''
-    echo '=== Starting trtllm-serve on port ${PORT} ==='
-    echo '=== Model: ${MODEL} ==='
-    echo '=== Backend: pytorch (default AUTO MoE) ==='
-    echo ''
-
-    # Start the server in background, log output
-    trtllm-serve ${MODEL} \
-      --backend pytorch \
-      --port ${PORT} \
-      --max_batch_size 1 \
-      --max_seq_len 256 \
-      2>&1 &
-    SERVER_PID=\$!
-
-    # Wait for server to be ready (poll health endpoint)
-    echo 'Waiting for server to start...'
-    for i in \$(seq 1 1200); do
-      if curl -s http://localhost:${PORT}/health > /dev/null 2>&1; then
-        echo \"Server ready after \${i}s\"
-        break
-      fi
-      if ! kill -0 \$SERVER_PID 2>/dev/null; then
-        echo 'Server process died!'
-        wait \$SERVER_PID || true
-        exit 1
-      fi
-      sleep 1
-    done
-
-    if ! curl -s http://localhost:${PORT}/health > /dev/null 2>&1; then
-      echo 'Server failed to start within 1200s'
-      kill \$SERVER_PID 2>/dev/null || true
-      exit 1
-    fi
-
-    echo ''
-    echo '=== Running inference test ==='
-    RESPONSE=\$(curl -s http://localhost:${PORT}/v1/completions \
-      -H 'Content-Type: application/json' \
-      -d '{
-        \"model\": \"${MODEL}\",
-        \"prompt\": \"The capital of France is\",
-        \"max_tokens\": 32,
-        \"temperature\": 0.0
-      }')
-
-    echo \"Response: \${RESPONSE}\"
-    echo ''
-
-    # Check if response contains generated text
-    if echo \"\${RESPONSE}\" | python3 -c 'import json,sys; d=json.load(sys.stdin); text=d[\"choices\"][0][\"text\"]; print(f\"Generated: {text}\"); assert len(text.strip()) > 0, \"Empty response\"' 2>&1; then
-      echo '=== E2E TEST PASSED ==='
-    else
-      echo '=== E2E TEST FAILED: no valid response ==='
-      echo \"Raw response: \${RESPONSE}\"
-      kill \$SERVER_PID 2>/dev/null || true
-      exit 1
-    fi
-
-    # Cleanup
-    kill \$SERVER_PID 2>/dev/null || true
-    wait \$SERVER_PID 2>/dev/null || true
-    echo '=== Done: '\$(date)' ==='
-  " 2>&1 | tee -a "${LOGFILE}"
-
-EXIT_CODE=${PIPESTATUS[0]}
-echo "" | tee -a "${LOGFILE}"
-echo "=== Exit code: ${EXIT_CODE} ===" | tee -a "${LOGFILE}"
-echo "=== Log saved to: ${LOGFILE} ===" | tee -a "${LOGFILE}"
-
-# Cleanup container if still running
-docker rm -f trtllm-e2e-test 2>/dev/null || true
-
-exit ${EXIT_CODE}
diff --git a/run_precommit.sh b/run_precommit.sh
deleted file mode 100755
index a5aeac31f272..000000000000
--- a/run_precommit.sh
+++ /dev/null
@@ -1,10 +0,0 @@
-#!/bin/bash
-set -e
-cd /code/tensorrt_llm
-git config --global --add safe.directory /code/tensorrt_llm
-git config --global --add safe.directory /home/mihai/workspace/TensorRT-LLM
-git config --global --add safe.directory /home/mihai/workspace/trtllm-pr12309
-pip install pre-commit 2>&1 | tail -1
-CHANGED=$(git diff --name-only main...HEAD)
-echo "Changed files: $CHANGED"
-pre-commit run --show-diff-on-failure --files $CHANGED 2>&1
diff --git a/run_test.sh b/run_test.sh
deleted file mode 100755
index 78d5f38315ba..000000000000
--- a/run_test.sh
+++ /dev/null
@@ -1,76 +0,0 @@
-#!/bin/bash
-# Run CuTEDSL MoE tests for PR #12309 with proper logging and memory management
-# Usage: ./run_test.sh [test_filter]
-# Example: ./run_test.sh "test_nvfp4_grouped_gemm_blackwell"
-# Example: ./run_test.sh  (runs all CuTEDSL MoE tests)
-
-set -euo pipefail
-
-WORKTREE="/home/mihai/workspace/trtllm-pr12309"
-LOGDIR="${WORKTREE}/test_logs"
-mkdir -p "${LOGDIR}"
-TIMESTAMP=$(date +%Y%m%d_%H%M%S)
-LOGFILE="${LOGDIR}/test_${TIMESTAMP}.log"
-
-TEST_FILTER="${1:-}"
-FILTER_ARGS=""
-if [[ -n "${TEST_FILTER}" ]]; then
-    FILTER_ARGS="-k '${TEST_FILTER}'"
-fi
-
-echo "=== Test Run: $(date) ===" | tee "${LOGFILE}"
-echo "Filter: ${TEST_FILTER:-<all>}" | tee -a "${LOGFILE}"
-echo "Log: ${LOGFILE}" | tee -a "${LOGFILE}"
-
-# Show system state before test
-echo "" | tee -a "${LOGFILE}"
-echo "=== PRE-TEST SYSTEM STATE ===" | tee -a "${LOGFILE}"
-free -h | tee -a "${LOGFILE}"
-swapon --show | tee -a "${LOGFILE}"
-nvidia-smi 2>/dev/null | tee -a "${LOGFILE}"
-echo "" | tee -a "${LOGFILE}"
-
-docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \
-  -v "${WORKTREE}:/workspace/TensorRT-LLM" \
-  -v /home/mihai/.cache/huggingface:/root/.cache/huggingface \
-  --network host \
-  -w /workspace/TensorRT-LLM \
-  -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
-  -e CUDA_MODULE_LOADING=LAZY \
-  tensorrt_llm/devel:latest \
-  bash -c "
-    set -euo pipefail
-    echo '=== Container started: '\$(date)' ==='
-    echo '=== Installing TRT-LLM ==='
-    pip install --no-build-isolation -e '.[devel]' 2>&1 | tail -5
-    pip install -r requirements-dev.txt 2>&1 | tail -5
-    echo '=== Upgrading nvidia-cutlass-dsl to 4.3.5 ==='
-    pip install nvidia-cutlass-dsl==4.3.5 2>&1 | tail -3
-    echo '=== Install done: '\$(date)' ==='
-    echo ''
-    echo '=== PYTORCH_CUDA_ALLOC_CONF='\${PYTORCH_CUDA_ALLOC_CONF}' ==='
-    echo '=== Running pytest ==='
-    echo ''
-    cd tests/unittest
-    python3 -m pytest \
-      _torch/thop/parallel/test_cute_dsl_moe.py \
-      ${FILTER_ARGS} \
-      -v -s \
-      --tb=short \
-      --timeout=7200 \
-      -p no:cacheprovider \
-      2>&1
-    echo ''
-    echo '=== Tests complete: '\$(date)' ==='
-  " 2>&1 | tee -a "${LOGFILE}"
-
-EXIT_CODE=${PIPESTATUS[0]}
-
-# Show system state after test
-echo "" | tee -a "${LOGFILE}"
-echo "=== POST-TEST SYSTEM STATE ===" | tee -a "${LOGFILE}"
-free -h | tee -a "${LOGFILE}"
-echo "" | tee -a "${LOGFILE}"
-echo "=== Exit code: ${EXIT_CODE} ===" | tee -a "${LOGFILE}"
-echo "=== Log saved to: ${LOGFILE} ===" | tee -a "${LOGFILE}"
-exit ${EXIT_CODE}
diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py
index 8169ae29afcf..2d6d825a7b44 100644
--- a/tensorrt_llm/_torch/model_config.py
+++ b/tensorrt_llm/_torch/model_config.py
@@ -251,8 +251,8 @@ def resolve_moe_backend(moe_backend: str, architecture: str) -> str:
         if moe_backend.upper() != "AUTO":
             return moe_backend
 
-        # Blackwell family: use TRTLLM for all MoE architectures
-        if _utils.is_blackwell():
+        # 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":
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
index 997601196a27..07f8af12b0d4 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py
+++ b/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_blackwell
+import tensorrt_llm._utils as _utils
 from tensorrt_llm.models.modeling_utils import QuantAlgo
 
 from ...autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec,
@@ -63,7 +63,7 @@ def cute_dsl_fp8_group_blockwise_gemm_ref(
     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 Blackwell and other archs.
-    if is_blackwell():
+    if _utils.is_blackwell():
         input_scale_tmp = a_sf.permute(1, 0).as_strided((m, w_k, 1),
                                                         (1, m, m * w_k))
     else:
@@ -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,11 @@ def can_implement(
                 "CuteDslFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)"
             )
 
-        # NVFP4 - SM in {100, 103, 120, 121} (Blackwell family)
+        # NVFP4 - SM100 family only (SM120/SM121 has scale dtype mismatch)
         if quant_algo == QuantAlgo.NVFP4:
-            if sm_version not in {100, 103, 120, 121}:
+            if not _utils.is_sm_100f(sm_version):
                 return _warn_and_return(
-                    f"NVFP4 requires Blackwell (SM100/103/120/121), 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
 
diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py
index 82948213b539..c1c2c74da172 100644
--- a/tests/integration/defs/conftest.py
+++ b/tests/integration/defs/conftest.py
@@ -1982,7 +1982,7 @@ def check_device_contain(keyword_list):
     get_sm_version() != 90,
     reason="This test is only  supported in Hopper architecture")
 
-skip_no_sm120 = pytest.mark.skipif(get_sm_version() != 120,
+skip_no_sm120 = pytest.mark.skipif(not is_sm_120f(),
                                    reason="This test is for SM120")
 
 skip_arm = pytest.mark.skipif(
diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py
index 2e1d97326d4a..7fd21d2be46c 100644
--- a/tests/unittest/_torch/modules/moe/test_moe_backend.py
+++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py
@@ -645,3 +645,21 @@ def run_moe():
             with torch.inference_mode():
                 output = run_moe()
                 ref_fused_moe.check_accuracy(output, ref_output)
+
+
+class TestResolveMoeBackendSM121:
+    """Test that resolve_moe_backend returns CUTLASS (not TRTLLM) on SM121.
+
+    SM120/SM121 lack tcgen05.mma instructions required by trtllm-gen kernels,
+    so the AUTO backend must fall back to CUTLASS.
+    """
+
+    def test_resolve_moe_backend_returns_cutlass_on_sm121(self):
+        from unittest.mock import patch
+
+        with patch("tensorrt_llm._utils.get_sm_version", return_value=121):
+            result = ModelConfig.resolve_moe_backend("AUTO", "SomeArchitecture")
+        assert result == "CUTLASS", (
+            f"Expected CUTLASS on SM121 but got {result}; "
+            "trtllm-gen kernels use tcgen05.mma which is unavailable on SM120/SM121"
+        )