From 40fd9224774ba88862018d0de005f1a2c7ee6e49 Mon Sep 17 00:00:00 2001
From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Date: Wed, 18 Feb 2026 03:25:40 +0000
Subject: [PATCH 1/6] fix sm120 int4

Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
---
 cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
index 50794769b5ec..f5ef0fe82bf5 100644
--- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
+++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
@@ -587,7 +587,8 @@ std::vector<CutlassGemmConfig> get_candidate_configs(
     {
         return get_candidate_configs_sm100(config_type_param, sm);
     }
-    if (sm >= 120 && (config_type_param & CutlassGemmConfig::BLACKWELL))
+    if (sm >= 120 && (config_type_param & CutlassGemmConfig::BLACKWELL)
+        && !(config_type_param & CutlassGemmConfig::WEIGHT_ONLY))
     {
         return get_candidate_configs_sm120(config_type_param);
     }

From 2e877073feae70ae3a46457a2700930ea671fe9d Mon Sep 17 00:00:00 2001
From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Date: Wed, 29 Apr 2026 13:01:17 -0700
Subject: [PATCH 2/6] [None][fix] Enable W4A16 finegrained GEMM on SM120/121

The Python preprocess_weights_for_mixed_gemm had its
sm_ >= 90 -> sm_ = 80 adjustment trapped inside an elif tied to the 2-D
shape check, so vanilla 2-D Linear weights on SM120/121 kept sm_ = 120
and skipped both row permutation (gated sm_ < 100) and column interleave
(gated sm_ < 90). The CUTLASS Sm80 finegrained kernel then read a
mismatched layout, producing 80%+ mismatch with ground truth.

Move the adjustment out of the elif so it applies to both 2-D and 3-D
inputs, mirroring the C++ preprocessor.

Also lift the FinegrainedMixedDtypeGemm activation-aware version gate so
W4A16 is enabled up through SM121 while W4A8 stays gated at SM103
(no FP8-activation dispatch on the SM120 Sm80 fallback). Tests skip
correctly under the new gate.

Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/custom_ops/torch_custom_ops.py    | 11 ++++++++++-
 tensorrt_llm/quantization/functional.py               | 10 ++++++++--
 .../parallel/test_finegrained_mixed_dtype_gemm.py     |  4 +++-
 .../_torch/thop/parallel/test_w4a16_linear.py         |  6 +++---
 .../unittest/_torch/thop/parallel/test_w4a8_linear.py |  5 +++--
 5 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
index 42f41de8b911..d406b3f0b2ff 100644
--- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
@@ -1412,7 +1412,13 @@ def _(
 
 class FinegrainedMixedDtypeGemm(TunableRunner):
     _runner_dict = dict()
+    # Max SM for W4A8 (FP8 activation). The Sm80 fallback used on SM120/121
+    # has no FP8 mma, so W4A8 remains capped here.
     MAX_SUPPORTED_SM_VERSION = 103
+    # Max SM for W4A16 (FP16/BF16 activation). SM120/121 dispatch through
+    # cutlass::arch::Sm80 MMA kernels which are MMA-compatible with Blackwell
+    # consumer GPUs (RTX 5090/5080).
+    MAX_SUPPORTED_SM_VERSION_W4A16 = 121
 
     def __init__(self, activation_dtype: torch.dtype, output_dtype: torch.dtype,
                  quant_mode: int):
@@ -1445,7 +1451,10 @@ def forward(self,
                 do_preparation: bool = False,
                 **kwargs) -> torch.Tensor:
 
-        if get_sm_version() > self.MAX_SUPPORTED_SM_VERSION:
+        max_sm = (self.MAX_SUPPORTED_SM_VERSION
+                  if self.activation_dtype == torch.float8_e4m3fn else
+                  self.MAX_SUPPORTED_SM_VERSION_W4A16)
+        if get_sm_version() > max_sm:
             raise ValueError(
                 f"SM version {get_sm_version()} is not supported for W4A16/W4A8 finegrained mixed dtype GEMM"
             )
diff --git a/tensorrt_llm/quantization/functional.py b/tensorrt_llm/quantization/functional.py
index bdfa6a07964d..0364347c0f41 100644
--- a/tensorrt_llm/quantization/functional.py
+++ b/tensorrt_llm/quantization/functional.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -959,7 +959,13 @@ def preprocess_weights_for_mixed_gemm(
     sm_ = sm_ if sm_ > 0 else get_sm_version()
     if len(tensor.shape) == 2:
         tensor = tensor.unsqueeze(0)
-    elif sm_ >= 90:
+    # Hopper / Blackwell weight-only paths reuse the SM80 layout (TMA-WS or
+    # SM80-fallback kernels read the same interleaved/row-permuted format).
+    # This adjustment must apply for both 2-D and 3-D inputs; previously it
+    # was guarded by an elif tied to the 2-D shape check, which silently
+    # skipped row permutation + column interleave for vanilla 2-D linears on
+    # SM>=90 and produced wrong-layout weights on SM120/121.
+    if sm_ >= 90:
         sm_ = 80
     if sm_ == 100 or sm_ == 103:
         do_weight_interleave = False
diff --git a/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py b/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py
index 0041f11da6b7..736ff60b6772 100644
--- a/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py
+++ b/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py
@@ -51,7 +51,9 @@ def test_matmul_activation_int4_input(m, n, k, group_size, activation_dtype,
     torch.manual_seed(0)
     device = "cuda"
 
-    if get_sm_version() > FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION:
+    max_sm = (FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION if use_w4a8_awq
+              else FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION_W4A16)
+    if get_sm_version() > max_sm:
         pytest.skip(
             f"W4A16/W4A8 not supported for SM version {get_sm_version()}")
 
diff --git a/tests/unittest/_torch/thop/parallel/test_w4a16_linear.py b/tests/unittest/_torch/thop/parallel/test_w4a16_linear.py
index 8aac068211a8..48fdcb16b915 100644
--- a/tests/unittest/_torch/thop/parallel/test_w4a16_linear.py
+++ b/tests/unittest/_torch/thop/parallel/test_w4a16_linear.py
@@ -17,10 +17,10 @@
 )
 def test_w4a16_linear(dtype, weights_dtype, has_zero=False):
 
-    if get_sm_version() > FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION:
+    if get_sm_version(
+    ) > FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION_W4A16:
         pytest.skip(
-            f"W4A16/W4A8 is not supported in this SM version {get_sm_version()}"
-        )
+            f"W4A16 is not supported in this SM version {get_sm_version()}")
 
     SEQ_LEN = 10
     HIDDEN_SIZE = 128
diff --git a/tests/unittest/_torch/thop/parallel/test_w4a8_linear.py b/tests/unittest/_torch/thop/parallel/test_w4a8_linear.py
index 20187385a6d0..518a1e408200 100644
--- a/tests/unittest/_torch/thop/parallel/test_w4a8_linear.py
+++ b/tests/unittest/_torch/thop/parallel/test_w4a8_linear.py
@@ -18,10 +18,11 @@
 )
 def test_w4a8_linear(dtype, weights_dtype, has_zero=False):
 
+    # W4A8 requires an FP8-capable mixed-dtype dispatch; SM120/121 fall back
+    # to the Sm80 MMA path which has no FP8, so W4A8 stays gated on SM>103.
     if get_sm_version() > FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION:
         pytest.skip(
-            f"W4A16/W4A8 is not supported in this SM version {get_sm_version()}"
-        )
+            f"W4A8 is not supported in this SM version {get_sm_version()}")
 
     SEQ_LEN = 10
     HIDDEN_SIZE = 128

From 9b2730945805145abe797dae68bc23ec39a84382 Mon Sep 17 00:00:00 2001
From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Date: Sun, 3 May 2026 02:32:07 +0000
Subject: [PATCH 3/6] Fix unreachable SM100/103 branch after SM remap.

Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
---
 tensorrt_llm/quantization/functional.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/tensorrt_llm/quantization/functional.py b/tensorrt_llm/quantization/functional.py
index 0364347c0f41..2685523fc9c5 100644
--- a/tensorrt_llm/quantization/functional.py
+++ b/tensorrt_llm/quantization/functional.py
@@ -959,13 +959,17 @@ def preprocess_weights_for_mixed_gemm(
     sm_ = sm_ if sm_ > 0 else get_sm_version()
     if len(tensor.shape) == 2:
         tensor = tensor.unsqueeze(0)
-    # Hopper / Blackwell weight-only paths reuse the SM80 layout (TMA-WS or
-    # SM80-fallback kernels read the same interleaved/row-permuted format).
-    # This adjustment must apply for both 2-D and 3-D inputs; previously it
-    # was guarded by an elif tied to the 2-D shape check, which silently
-    # skipped row permutation + column interleave for vanilla 2-D linears on
-    # SM>=90 and produced wrong-layout weights on SM120/121.
-    if sm_ >= 90:
+    elif sm_ >= 90:
+        # 3-D inputs (MoE expert weights) have no specialised Hopper/Blackwell
+        # kernels and fall back to the SM80 interleaved layout, mirroring the
+        # `force_interleave && arch >= 90` clause in the C++ preprocessor.
+        sm_ = 80
+    if sm_ >= 120:
+        # SM120/SM121 (RTX 5090 / 5080) reuse the SM80 layout for both 2-D and
+        # 3-D inputs because the SM120 dispatch routes through cutlass::arch::Sm80.
+        # Without this, vanilla 2-D Linear weights on SM120 skipped both row
+        # permutation (gated sm_ < 100) and column interleave (gated sm_ < 90)
+        # and produced wrong-layout weights for the Sm80 finegrained kernel.
         sm_ = 80
     if sm_ == 100 or sm_ == 103:
         do_weight_interleave = False

From bc390bb9f8539061fc3c4667e7aaae0b35b0ed7a Mon Sep 17 00:00:00 2001
From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Date: Tue, 5 May 2026 03:46:48 +0000
Subject: [PATCH 4/6] fix mxfp4_mxfp8

Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
---
 .../kernels/cutlass_kernels/cutlass_heuristic.cpp     | 10 +++++++++-
 .../_torch/modules/fused_moe/fused_moe_cutlass.py     | 11 ++++++++---
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
index f5ef0fe82bf5..06fa9b5e9090 100644
--- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
+++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
@@ -587,8 +587,16 @@ std::vector<CutlassGemmConfig> get_candidate_configs(
     {
         return get_candidate_configs_sm100(config_type_param, sm);
     }
+    // SM120/121 path: take it for any BLACKWELL request unless this is a
+    // pure integer weight-only GEMM (W4A16 / W8A16). Those need the
+    // Ampere-style fallback because the SM120 candidate list only covers
+    // the FP4-mixed shapes (FP4xFP4, FP8xFP4 / W4A8_MXFP4_*) added by the
+    // GPT-OSS SM120/121 enablement. WEIGHT_ONLY here means T != WeightType,
+    // which is also true for FP8xFP4, so we must additionally check the
+    // FP4 flags before falling through.
     if (sm >= 120 && (config_type_param & CutlassGemmConfig::BLACKWELL)
-        && !(config_type_param & CutlassGemmConfig::WEIGHT_ONLY))
+        && (!(config_type_param & CutlassGemmConfig::WEIGHT_ONLY)
+            || (config_type_param & (CutlassGemmConfig::FP4_ONLY | CutlassGemmConfig::FP8FP4_MIXED))))
     {
         return get_candidate_configs_sm120(config_type_param);
     }
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
index fcd31ada5596..4d5ce0b386ce 100755
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
@@ -111,9 +111,14 @@ class CutlassFusedMoE(MoE):
             "sm_constraint": ("in", {100, 103}),
             "dtypes": {torch.float16, torch.bfloat16, torch.float32},
         },
-        # W4A8_MXFP4_MXFP8: SM in {100, 103}
+        # W4A8_MXFP4_MXFP8: SM in {100, 103, 120, 121}
+        # SM120/121 added by commit 27a5091f ([None][feat] GPT-OSS Sm120/Sm121
+        # Support); the C++ isValidSM120MOESpecialisation<T=FP8, WeightType=FP4>
+        # path covers MXFP4_MXFP8 since use_wfp4afp8 is type-defined as
+        # (T==FP8 && WeightType==FP4), and the runtime block-scaling type for
+        # use_wfp4afp8 is hard-wired to MXFPX in moe_gemm_template_dispatch.h.
         QuantAlgo.W4A8_MXFP4_MXFP8: {
-            "sm_constraint": ("in", {100, 103}),
+            "sm_constraint": ("in", {100, 103, 120, 121}),
             "dtypes": {torch.float16, torch.bfloat16},
         },
     }
@@ -140,7 +145,7 @@ def can_implement(
         - W8A16: SM >= 80
         - W4A16_MXFP4: SM == 90 only
         - W4A8_MXFP4_FP8: SM in {100, 103}
-        - W4A8_MXFP4_MXFP8: SM in {100, 103}
+        - W4A8_MXFP4_MXFP8: SM in {100, 103, 120, 121}
 
         Args:
             quant_algo: The quantization algorithm to check (None for unquantized)

From 7346d5456b2d06b27930a5294b3cfd6fe1a9e86c Mon Sep 17 00:00:00 2001
From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Date: Tue, 5 May 2026 21:07:26 +0000
Subject: [PATCH 5/6] update comments

Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
---
 .../kernels/cutlass_kernels/cutlass_heuristic.cpp     | 11 ++++-------
 tensorrt_llm/_torch/custom_ops/torch_custom_ops.py    | 10 +++-------
 .../_torch/modules/fused_moe/fused_moe_cutlass.py     |  8 --------
 tensorrt_llm/quantization/functional.py               | 10 +++-------
 .../test_lists/test-db/l0_rtx_pro_6000.yml            |  7 +++++++
 .../parallel/test_finegrained_mixed_dtype_gemm.py     |  5 +++--
 .../unittest/_torch/thop/parallel/test_w4a8_linear.py |  3 ++-
 7 files changed, 22 insertions(+), 32 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
index 06fa9b5e9090..0e1586ca2ebf 100644
--- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
+++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
@@ -587,13 +587,10 @@ std::vector<CutlassGemmConfig> get_candidate_configs(
     {
         return get_candidate_configs_sm100(config_type_param, sm);
     }
-    // SM120/121 path: take it for any BLACKWELL request unless this is a
-    // pure integer weight-only GEMM (W4A16 / W8A16). Those need the
-    // Ampere-style fallback because the SM120 candidate list only covers
-    // the FP4-mixed shapes (FP4xFP4, FP8xFP4 / W4A8_MXFP4_*) added by the
-    // GPT-OSS SM120/121 enablement. WEIGHT_ONLY here means T != WeightType,
-    // which is also true for FP8xFP4, so we must additionally check the
-    // FP4 flags before falling through.
+    // Take the SM120/121 path for FP4-mixed GEMMs (FP4xFP4, FP8xFP4 /
+    // W4A8_MXFP4_*). Pure integer weight-only (W4A16 / W8A16) falls through
+    // to the Ampere candidate list. The WEIGHT_ONLY flag is set whenever
+    // T != WeightType (true for FP8xFP4 too), so guard with the FP4 flags.
     if (sm >= 120 && (config_type_param & CutlassGemmConfig::BLACKWELL)
         && (!(config_type_param & CutlassGemmConfig::WEIGHT_ONLY)
             || (config_type_param & (CutlassGemmConfig::FP4_ONLY | CutlassGemmConfig::FP8FP4_MIXED))))
diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
index d406b3f0b2ff..0711c2ded7e5 100644
--- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
@@ -1412,12 +1412,8 @@ def _(
 
 class FinegrainedMixedDtypeGemm(TunableRunner):
     _runner_dict = dict()
-    # Max SM for W4A8 (FP8 activation). The Sm80 fallback used on SM120/121
-    # has no FP8 mma, so W4A8 remains capped here.
-    MAX_SUPPORTED_SM_VERSION = 103
-    # Max SM for W4A16 (FP16/BF16 activation). SM120/121 dispatch through
-    # cutlass::arch::Sm80 MMA kernels which are MMA-compatible with Blackwell
-    # consumer GPUs (RTX 5090/5080).
+    MAX_SUPPORTED_SM_VERSION_W4A8 = 103
+    # W4A16 (FP16/BF16 activation): SM120/121 dispatch via cutlass::arch::Sm80.
     MAX_SUPPORTED_SM_VERSION_W4A16 = 121
 
     def __init__(self, activation_dtype: torch.dtype, output_dtype: torch.dtype,
@@ -1451,7 +1447,7 @@ def forward(self,
                 do_preparation: bool = False,
                 **kwargs) -> torch.Tensor:
 
-        max_sm = (self.MAX_SUPPORTED_SM_VERSION
+        max_sm = (self.MAX_SUPPORTED_SM_VERSION_W4A8
                   if self.activation_dtype == torch.float8_e4m3fn else
                   self.MAX_SUPPORTED_SM_VERSION_W4A16)
         if get_sm_version() > max_sm:
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
index 4d5ce0b386ce..ba4ce303fe23 100755
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
@@ -84,9 +84,6 @@ class CutlassFusedMoE(MoE):
             "dtypes": {torch.bfloat16},
         },
         # NVFP4: SM in {100, 103, 120, 121}
-        # SM 120 = desktop Blackwell (e.g. RTX 5090 / GB202)
-        # SM 121 = GB10 / DGX Spark
-        # C++ kernel: isValidSM120MOESpecialisation() supports FP4xFP4 and FP8xFP4
         QuantAlgo.NVFP4: {
             "sm_constraint": ("in", {100, 103, 120, 121}),
             "dtypes": {torch.float16, torch.bfloat16, torch.float8_e4m3fn},
@@ -112,11 +109,6 @@ class CutlassFusedMoE(MoE):
             "dtypes": {torch.float16, torch.bfloat16, torch.float32},
         },
         # W4A8_MXFP4_MXFP8: SM in {100, 103, 120, 121}
-        # SM120/121 added by commit 27a5091f ([None][feat] GPT-OSS Sm120/Sm121
-        # Support); the C++ isValidSM120MOESpecialisation<T=FP8, WeightType=FP4>
-        # path covers MXFP4_MXFP8 since use_wfp4afp8 is type-defined as
-        # (T==FP8 && WeightType==FP4), and the runtime block-scaling type for
-        # use_wfp4afp8 is hard-wired to MXFPX in moe_gemm_template_dispatch.h.
         QuantAlgo.W4A8_MXFP4_MXFP8: {
             "sm_constraint": ("in", {100, 103, 120, 121}),
             "dtypes": {torch.float16, torch.bfloat16},
diff --git a/tensorrt_llm/quantization/functional.py b/tensorrt_llm/quantization/functional.py
index 2685523fc9c5..9c330931151d 100644
--- a/tensorrt_llm/quantization/functional.py
+++ b/tensorrt_llm/quantization/functional.py
@@ -960,16 +960,12 @@ def preprocess_weights_for_mixed_gemm(
     if len(tensor.shape) == 2:
         tensor = tensor.unsqueeze(0)
     elif sm_ >= 90:
-        # 3-D inputs (MoE expert weights) have no specialised Hopper/Blackwell
-        # kernels and fall back to the SM80 interleaved layout, mirroring the
+        # 3-D inputs (MoE) fall back to SM80 interleaved layout, mirroring the
         # `force_interleave && arch >= 90` clause in the C++ preprocessor.
         sm_ = 80
     if sm_ >= 120:
-        # SM120/SM121 (RTX 5090 / 5080) reuse the SM80 layout for both 2-D and
-        # 3-D inputs because the SM120 dispatch routes through cutlass::arch::Sm80.
-        # Without this, vanilla 2-D Linear weights on SM120 skipped both row
-        # permutation (gated sm_ < 100) and column interleave (gated sm_ < 90)
-        # and produced wrong-layout weights for the Sm80 finegrained kernel.
+        # SM120/SM121 dispatch through cutlass::arch::Sm80; reuse SM80 layout
+        # for both 2-D and 3-D inputs.
         sm_ = 80
     if sm_ == 100 or sm_ == 103:
         do_weight_interleave = False
diff --git a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
index b4187e6a08cc..81afed83e2ca 100644
--- a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
+++ b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml
@@ -17,6 +17,13 @@ l0_rtx_pro_6000:
   - unittest/_torch/modeling -k "modeling_out_of_tree"
   # - unittest/_torch/modeling -k "modeling_qwen" # https://nvbugs/5234573
   - unittest/_torch/attention/test_attention_mla.py
+  # SM120 W4A16 / W4A8 mixed-dtype GEMM coverage (paired with FinegrainedMixedDtypeGemm
+  # MAX_SUPPORTED_SM_VERSION_W4A* gates and cutlass_heuristic.cpp SM120 routing).
+  - unittest/_torch/thop/parallel/test_w4a16_linear.py
+  - unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py
+  - unittest/_torch/thop/parallel/test_weight_only_quant_linear.py
+  - unittest/_torch/thop/parallel/test_weight_only_quant_gemm.py
+  - unittest/_torch/thop/parallel/test_w4a8_linear.py
   - test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity]
   - test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity]
   - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-BF16-llama-3.1-model/Meta-Llama-3.1-8B]
diff --git a/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py b/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py
index 736ff60b6772..b6df8f035b99 100644
--- a/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py
+++ b/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py
@@ -51,8 +51,9 @@ def test_matmul_activation_int4_input(m, n, k, group_size, activation_dtype,
     torch.manual_seed(0)
     device = "cuda"
 
-    max_sm = (FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION if use_w4a8_awq
-              else FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION_W4A16)
+    max_sm = (FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION_W4A8
+              if use_w4a8_awq else
+              FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION_W4A16)
     if get_sm_version() > max_sm:
         pytest.skip(
             f"W4A16/W4A8 not supported for SM version {get_sm_version()}")
diff --git a/tests/unittest/_torch/thop/parallel/test_w4a8_linear.py b/tests/unittest/_torch/thop/parallel/test_w4a8_linear.py
index 518a1e408200..458146a24c0c 100644
--- a/tests/unittest/_torch/thop/parallel/test_w4a8_linear.py
+++ b/tests/unittest/_torch/thop/parallel/test_w4a8_linear.py
@@ -20,7 +20,8 @@ def test_w4a8_linear(dtype, weights_dtype, has_zero=False):
 
     # W4A8 requires an FP8-capable mixed-dtype dispatch; SM120/121 fall back
     # to the Sm80 MMA path which has no FP8, so W4A8 stays gated on SM>103.
-    if get_sm_version() > FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION:
+    if get_sm_version(
+    ) > FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION_W4A8:
         pytest.skip(
             f"W4A8 is not supported in this SM version {get_sm_version()}")
 

From 282324741f9c0b30dc05d4e6afbf1f21ddf7e760 Mon Sep 17 00:00:00 2001
From: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
Date: Tue, 12 May 2026 21:07:54 +0000
Subject: [PATCH 6/6] address comments

Signed-off-by: Pamela <179191831+pamelap-nvidia@users.noreply.github.com>
---
 .../kernels/cutlass_kernels/cutlass_heuristic.cpp  | 11 +++++------
 tensorrt_llm/_torch/custom_ops/torch_custom_ops.py | 14 +++++++++-----
 tensorrt_llm/quantization/functional.py            | 12 ++++--------
 .../parallel/test_finegrained_mixed_dtype_gemm.py  | 13 +++++++------
 4 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
index 0e1586ca2ebf..b0ea6333a88c 100644
--- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
+++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cpp
@@ -587,13 +587,12 @@ std::vector<CutlassGemmConfig> get_candidate_configs(
     {
         return get_candidate_configs_sm100(config_type_param, sm);
     }
-    // Take the SM120/121 path for FP4-mixed GEMMs (FP4xFP4, FP8xFP4 /
-    // W4A8_MXFP4_*). Pure integer weight-only (W4A16 / W8A16) falls through
-    // to the Ampere candidate list. The WEIGHT_ONLY flag is set whenever
-    // T != WeightType (true for FP8xFP4 too), so guard with the FP4 flags.
+    // Use the Blackwell (SM120+) specialized candidate list only for these cases:
+    //   - FP4 dense GEMM (FP4_ONLY)
+    //   - Mixed FP8/FP4 (FP8FP4_MIXED, e.g., W4A8_MXFP4_*)
+    // If neither applies, fall through (e.g. pure int W4A16/W8A16 handled elsewhere).
     if (sm >= 120 && (config_type_param & CutlassGemmConfig::BLACKWELL)
-        && (!(config_type_param & CutlassGemmConfig::WEIGHT_ONLY)
-            || (config_type_param & (CutlassGemmConfig::FP4_ONLY | CutlassGemmConfig::FP8FP4_MIXED))))
+        && ((config_type_param & CutlassGemmConfig::FP4_ONLY) || (config_type_param & CutlassGemmConfig::FP8FP4_MIXED)))
     {
         return get_candidate_configs_sm120(config_type_param);
     }
diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
index 0711c2ded7e5..285faf5b762a 100644
--- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
@@ -1447,12 +1447,16 @@ def forward(self,
                 do_preparation: bool = False,
                 **kwargs) -> torch.Tensor:
 
-        max_sm = (self.MAX_SUPPORTED_SM_VERSION_W4A8
-                  if self.activation_dtype == torch.float8_e4m3fn else
-                  self.MAX_SUPPORTED_SM_VERSION_W4A16)
-        if get_sm_version() > max_sm:
+        sm = get_sm_version()
+        if (self.activation_dtype == torch.float8_e4m3fn
+                and sm > self.MAX_SUPPORTED_SM_VERSION_W4A8):
             raise ValueError(
-                f"SM version {get_sm_version()} is not supported for W4A16/W4A8 finegrained mixed dtype GEMM"
+                f"SM version {sm} is not supported for W4A8 finegrained mixed dtype GEMM"
+            )
+        if (self.activation_dtype in (torch.float16, torch.bfloat16)
+                and sm > self.MAX_SUPPORTED_SM_VERSION_W4A16):
+            raise ValueError(
+                f"SM version {sm} is not supported for W4A16 finegrained mixed dtype GEMM"
             )
 
         activation, weights_packed, scales = inputs
diff --git a/tensorrt_llm/quantization/functional.py b/tensorrt_llm/quantization/functional.py
index 9c330931151d..ddaa7394eb5b 100644
--- a/tensorrt_llm/quantization/functional.py
+++ b/tensorrt_llm/quantization/functional.py
@@ -957,16 +957,12 @@ def preprocess_weights_for_mixed_gemm(
         sm_: int = -1,
         do_weight_interleave: bool = True) -> torch.Tensor:
     sm_ = sm_ if sm_ > 0 else get_sm_version()
+    # 3-D inputs (MoE) on Hopper+ and any input on SM120/SM121 reuse the SM80
+    # interleaved layout. Check the original rank before unsqueeze.
+    if (len(tensor.shape) == 3 and sm_ >= 90) or sm_ >= 120:
+        sm_ = 80
     if len(tensor.shape) == 2:
         tensor = tensor.unsqueeze(0)
-    elif sm_ >= 90:
-        # 3-D inputs (MoE) fall back to SM80 interleaved layout, mirroring the
-        # `force_interleave && arch >= 90` clause in the C++ preprocessor.
-        sm_ = 80
-    if sm_ >= 120:
-        # SM120/SM121 dispatch through cutlass::arch::Sm80; reuse SM80 layout
-        # for both 2-D and 3-D inputs.
-        sm_ = 80
     if sm_ == 100 or sm_ == 103:
         do_weight_interleave = False
 
diff --git a/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py b/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py
index b6df8f035b99..211398c6028a 100644
--- a/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py
+++ b/tests/unittest/_torch/thop/parallel/test_finegrained_mixed_dtype_gemm.py
@@ -51,12 +51,13 @@ def test_matmul_activation_int4_input(m, n, k, group_size, activation_dtype,
     torch.manual_seed(0)
     device = "cuda"
 
-    max_sm = (FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION_W4A8
-              if use_w4a8_awq else
-              FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION_W4A16)
-    if get_sm_version() > max_sm:
-        pytest.skip(
-            f"W4A16/W4A8 not supported for SM version {get_sm_version()}")
+    sm = get_sm_version()
+    if (use_w4a8_awq
+            and sm > FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION_W4A8):
+        pytest.skip(f"W4A8 not supported for SM version {sm}")
+    if (not use_w4a8_awq
+            and sm > FinegrainedMixedDtypeGemm.MAX_SUPPORTED_SM_VERSION_W4A16):
+        pytest.skip(f"W4A16 not supported for SM version {sm}")
 
     total_groups = (k + group_size - 1) // group_size
     scale_zero_dtype = torch.float16 if use_w4a8_awq else activation_dtype