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

2026-06-29 14:19:24 -07:00
parent 11b7ed7c9e
commit 98d0e702c3
1 changed files with 3 additions and 1 deletions
@@ -496,9 +496,11 @@ class MoEGate(nn.Module):
        ):
            logits = F.linear(hidden_states, self.weight, None)
        else:
+            # NOTE(b8zhong): this threshold has been empirically verified
+            max_router_gemm_tokens = 4 if _device_sm in (100, 103) else 16
            if (
                _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 (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
                and _device_sm >= 90