From 0c53fec4768a46a003a7afe27c938469db956368 Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Tue, 22 Sep 2026 06:43:21 +0800 Subject: [PATCH] [ROCm] fix: remove extra bf16 -> fp32 cast in jit grouped topk kernel path (#39775) --- python/sglang/srt/layers/moe/topk.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 97c04a694..3175d230d 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -2471,10 +2471,14 @@ def select_experts( info=expert_location_dispatch_info, ) + # Only aiter_biased_grouped_topk wants the bias in the gating dtype. The JIT + # router returns before it in biased_grouped_topk_gpu and upcasts the bias + # in-register, so downcasting here would only buy it a per-call cast back. if ( _use_aiter and use_grouped_topk and correction_bias is not None + and not envs.SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK.get() and dynamic_expert_bias is None ): correction_bias = topk_config.correction_bias_for_dtype(router_logits.dtype)