From 84aaa69795f3588c0d384e034d060d4436d28672 Mon Sep 17 00:00:00 2001 From: inkcherry Date: Fri, 6 Mar 2026 16:57:12 +0800 Subject: [PATCH] [AMD] Use bfloat16 for correction_bias in AITER FP8 path to avoid runtime dtype conversion for dsv3 (#19843) --- python/sglang/srt/models/deepseek_v2.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index a63efb0fd..e9a4f511b 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -271,13 +271,18 @@ class MoEGate(nn.Module): torch.empty((config.n_routed_experts, config.hidden_size)) ) if config.topk_method == "noaux_tc": - correction_bias_dtype = ( - torch.bfloat16 - if quant_config is not None - and quant_config.get_name() == "modelopt_fp4" - and get_moe_runner_backend().is_flashinfer_trtllm() - else torch.float32 - ) + correction_bias_dtype = torch.float32 + if quant_config is not None: + if ( + quant_config.get_name() == "modelopt_fp4" + and get_moe_runner_backend().is_flashinfer_trtllm() + ): + correction_bias_dtype = torch.bfloat16 + elif _use_aiter and quant_config.get_name() in ( + "fp8", + "compressed_tensors", + ): + correction_bias_dtype = torch.bfloat16 self.e_score_correction_bias = nn.Parameter( torch.empty((config.n_routed_experts), dtype=correction_bias_dtype) )