Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dongfengy merged 4 commits into from
Oct 9, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
338e668
[None][fix] Restrict tinygemm use to sm90/sm100 only
dongfengy Oct 7, 2025
9ae47b7
Update tinygemm2.cpp
dongfengy Oct 8, 2025
7f753d9
Update modeling_gpt_oss.py
dongfengy Oct 8, 2025
a60ca92
Update tinygemm2.cpp
dongfengy Oct 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions cpp/tensorrt_llm/thop/tinygemm2.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Comment thread
dongfengy marked this conversation as resolved.
TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
Expand Down
5 changes: 4 additions & 1 deletion tensorrt_llm/_torch/models/modeling_gpt_oss.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Comment thread
dongfengy marked this conversation as resolved.
weight = self.gate.weight
bias = self.gate.bias
g = torch.ops.trtllm.tinygemm2(x, weight, bias)
Expand Down