diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/quant.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/quant.py
index f2ebb177f2a2..197741281124 100644
--- a/tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/quant.py
+++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/quantization/quant.py
@@ -109,7 +109,8 @@ def _trtllm_fp8_prequant_linear_core(
     enable_cuda_core = False
     if torch.cuda.is_available():
         capability = torch.cuda.get_device_capability(0)
-        enable_cuda_core = capability == (8, 9) or capability == (12, 0)
+        # enable cuda core for sm89, sm120, and sm121
+        enable_cuda_core = capability in ((8, 9), (12, 0), (12, 1))
 
     if x.shape[0] <= 8 and enable_cuda_core:
         output = torch.ops.trtllm.cuda_scaled_mm(
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index 240c217a044e..2bbb53776abf 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -3344,9 +3344,8 @@ def __init__(
         if torch.cuda.is_available():
             capability = torch.cuda.get_device_capability(
                 torch.device('cuda:0'))
-            # enable cuda core for sm89 and sm120
-            self.enable_cuda_core = (capability[0] == 8 and capability[1] == 9) \
-                or (capability[0] == 12 and capability[1] == 0)
+            # enable cuda core for sm89, sm120, and sm121
+            self.enable_cuda_core = capability in ((8, 9), (12, 0), (12, 1))
 
         if not skip_create_weights_in_init:
             self.create_weights()