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

[Perf] Restore the 16-token router GEMM threshold on SM10X (#34953)

Browse Source 
Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>
This commit is contained in:
Mohammad Miadh Angkad
2026-08-18 19:53:52 -07:00
committed by GitHub
co-authored by Mohammad Angkad
parent ee1f2e8dfd
commit 1beb805356
1 changed files with 1 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+1 -3
View File
@@ -525,11 +525,9 @@ class MoEGate(nn.Module):
return linear_bf16_fp32(hidden_states, self.weight) return linear_bf16_fp32(hidden_states, self.weight)
return F.linear(hidden_states, self.weight, None) return 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] <= max_router_gemm_tokens and hidden_states.shape[0] <= 16
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