[AMD] Use aiter_dsv3_router_gemm kernel if number of experts <= 256. (#18451)

2026-03-19 07:40:48 +02:00
parent 126cd5cfae
commit 85fe8c6793
1 changed files with 5 additions and 1 deletions
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -326,7 +326,11 @@ class MoEGate(nn.Module):
                logits = dsv3_router_gemm(
                    hidden_states, self.weight, out_dtype=torch.float32
                )
-            elif _use_aiter_gfx95 and hidden_states.shape[0] <= 256:
+            elif (
+                _use_aiter_gfx95
+                and hidden_states.shape[0] <= 256
+                and self.weight.shape[0] <= 256
+            ):
                logits = aiter_dsv3_router_gemm(
                    hidden_states, self.weight, gemm_output_zero_allocator
                )