minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[ROCm] fix: remove extra bf16 -> fp32 cast in jit grouped topk kernel path (#39775)

Browse Source
This commit is contained in:
RuibinCheung
2026-09-21 15:43:21 -07:00
committed by GitHub
parent d47b8c454c
commit 0c53fec476
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/topk.py
+4
View File
@@ -2471,10 +2471,14 @@ def select_experts(
info=expert_location_dispatch_info, 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 ( if (
_use_aiter _use_aiter
and use_grouped_topk and use_grouped_topk
and correction_bias is not None and correction_bias is not None
and not envs.SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK.get()
and dynamic_expert_bias is None and dynamic_expert_bias is None
): ):
correction_bias = topk_config.correction_bias_for_dtype(router_logits.dtype) correction_bias = topk_config.correction_bias_for_dtype(router_logits.dtype)