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

Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>
2026-08-18 19:53:52 -07:00
parent ee1f2e8dfd
commit 1beb805356
1 changed files with 1 additions and 3 deletions
@@ -525,11 +525,9 @@ class MoEGate(nn.Module):
                return linear_bf16_fp32(hidden_states, self.weight)
            return 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] <= max_router_gemm_tokens
+                and hidden_states.shape[0] <= 16
                and hidden_states.shape[1] % 1024 == 0
                and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
                and _device_sm >= 90