From 85fe8c6793a0b2bf8d5b2e98c88a8630515b0ac6 Mon Sep 17 00:00:00 2001 From: Matti Varjokallio Date: Thu, 19 Mar 2026 07:40:48 +0200 Subject: [PATCH] [AMD] Use aiter_dsv3_router_gemm kernel if number of experts <= 256. (#18451) --- python/sglang/srt/models/deepseek_v2.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 9715d147f..d6df9b6f9 100644 --- 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 )