[#15673][fix] Enable CUDA core fast path for SM89/SM120/SM121

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(

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()