minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[GLM-5] Tune the threshold of router GEMM (#29470)

Browse Source
This commit is contained in:
Brayden Zhong
2026-06-29 14:19:24 -07:00
committed by GitHub
parent 11b7ed7c9e
commit 98d0e702c3
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+3 -1
View File
@@ -496,9 +496,11 @@ class MoEGate(nn.Module):
): ):
logits = F.linear(hidden_states, self.weight, None) logits = F.linear(hidden_states, self.weight, None)
else: else:
# NOTE(b8zhong): this threshold has been empirically verified
max_router_gemm_tokens = 4 if _device_sm in (100, 103) else 16
if ( if (
_is_cuda _is_cuda
and hidden_states.shape[0] <= 16 and hidden_states.shape[0] <= max_router_gemm_tokens
and hidden_states.shape[1] % 1024 == 0 and hidden_states.shape[1] % 1024 == 0
and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384) and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
and _device_sm >= 90 and _device_sm >= 90