[None][fix] Restrict tinygemm use to certain SMs (#8182)

Signed-off-by: Dongfeng Yu <[email protected]>
Signed-off-by: dongfengy <[email protected]>

Author: dongfengy

cpp/tensorrt_llm/thop/tinygemm2.cpp

@@ -30,6 +30,9 @@ namespace torch_ext
 {
 torch::Tensor tinygemm2_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias)
 {
+    auto const smVersion = tensorrt_llm::common::getSMVersion();
+    TORCH_CHECK(
+        smVersion == 90 || smVersion == 100 || smVersion == 103, "tinygemm2 only supports SM90, SM100, and SM103.");
     TORCH_CHECK(input.dim() == 2, "input must be 2D");
     TORCH_CHECK(weight.dim() == 2, "weight must be 2D");
     TORCH_CHECK(bias.dim() == 1, "bias must be 1D");

tensorrt_llm/_torch/models/modeling_gpt_oss.py

@@ -7,6 +7,7 @@
 from tqdm import tqdm
 from transformers import GptOssConfig
 
+from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType
 
 from ..attention_backend import AttentionMetadata
@@ -225,7 +226,9 @@ def _create_ideal_expert_load_balanced_logits(
             dtype=pretrained_config.torch_dtype)
 
     def compute_gate_output(self, x: torch.Tensor) -> torch.Tensor:
-        if x.shape[0] <= MIN_LATENCY_TINYGEMM_NUM_TOKENS:
+        if get_sm_version() in [
+                90, 100, 103
+        ] and x.shape[0] <= MIN_LATENCY_TINYGEMM_NUM_TOKENS:
             weight = self.gate.weight
             bias = self.gate.bias
             g = torch.ops.trtllm.tinygemm2(x, weight, bias)