From 338e66800a9e5aaeea703a16b09d3f7f0e504def Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Tue, 7 Oct 2025 22:08:34 +0000 Subject: [PATCH 1/4] [None][fix] Restrict tinygemm use to sm90/sm100 only Signed-off-by: Dongfeng Yu --- cpp/tensorrt_llm/thop/tinygemm2.cpp | 2 ++ tensorrt_llm/_torch/models/modeling_gpt_oss.py | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/thop/tinygemm2.cpp b/cpp/tensorrt_llm/thop/tinygemm2.cpp index 2c78151357a2..2bd6e72f3379 100644 --- a/cpp/tensorrt_llm/thop/tinygemm2.cpp +++ b/cpp/tensorrt_llm/thop/tinygemm2.cpp @@ -30,6 +30,8 @@ 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, "tinygemm2 only supports SM90 and SM100."); 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"); diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py index 80456c331952..83ee19de39d7 100644 --- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py +++ b/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 + ] 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) From 9ae47b779391d4029b69704f8ae19a54cf6262c5 Mon Sep 17 00:00:00 2001 From: dongfengy <99041270+dongfengy@users.noreply.github.com> Date: Wed, 8 Oct 2025 09:24:47 -0700 Subject: [PATCH 2/4] Update tinygemm2.cpp Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> --- cpp/tensorrt_llm/thop/tinygemm2.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/thop/tinygemm2.cpp b/cpp/tensorrt_llm/thop/tinygemm2.cpp index 2bd6e72f3379..2bb7ba3e34b2 100644 --- a/cpp/tensorrt_llm/thop/tinygemm2.cpp +++ b/cpp/tensorrt_llm/thop/tinygemm2.cpp @@ -31,7 +31,7 @@ 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, "tinygemm2 only supports SM90 and SM100."); + 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"); From 7f753d94d4bc1850286ae03df5a3f64800278d26 Mon Sep 17 00:00:00 2001 From: dongfengy <99041270+dongfengy@users.noreply.github.com> Date: Wed, 8 Oct 2025 09:28:31 -0700 Subject: [PATCH 3/4] Update modeling_gpt_oss.py Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_gpt_oss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py index 83ee19de39d7..677711419848 100644 --- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py +++ b/tensorrt_llm/_torch/models/modeling_gpt_oss.py @@ -227,7 +227,7 @@ def _create_ideal_expert_load_balanced_logits( def compute_gate_output(self, x: torch.Tensor) -> torch.Tensor: if get_sm_version() in [ - 90, 100 + 90, 100, 103 ] and x.shape[0] <= MIN_LATENCY_TINYGEMM_NUM_TOKENS: weight = self.gate.weight bias = self.gate.bias From a60ca92aaf514822ca438d5a3004b0af66662caf Mon Sep 17 00:00:00 2001 From: dongfengy <99041270+dongfengy@users.noreply.github.com> Date: Wed, 8 Oct 2025 09:40:30 -0700 Subject: [PATCH 4/4] Update tinygemm2.cpp Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> --- cpp/tensorrt_llm/thop/tinygemm2.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/thop/tinygemm2.cpp b/cpp/tensorrt_llm/thop/tinygemm2.cpp index 2bb7ba3e34b2..3be0bea04b1f 100644 --- a/cpp/tensorrt_llm/thop/tinygemm2.cpp +++ b/cpp/tensorrt_llm/thop/tinygemm2.cpp @@ -31,7 +31,8 @@ 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( + 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");