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

[AMD] Use bfloat16 for correction_bias in AITER FP8 path to avoid runtime dtype conversion for dsv3 (#19843)

Browse Source
This commit is contained in:
inkcherry
2026-03-06 00:57:12 -08:00
committed by GitHub
parent 27053aa5ed
commit 84aaa69795
1 changed files with 12 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+12 -7
View File
@@ -271,13 +271,18 @@ class MoEGate(nn.Module):
torch.empty((config.n_routed_experts, config.hidden_size)) torch.empty((config.n_routed_experts, config.hidden_size))
) )
if config.topk_method == "noaux_tc": if config.topk_method == "noaux_tc":
correction_bias_dtype = ( correction_bias_dtype = torch.float32
torch.bfloat16 if quant_config is not None:
if quant_config is not None if (
and quant_config.get_name() == "modelopt_fp4" quant_config.get_name() == "modelopt_fp4"
and get_moe_runner_backend().is_flashinfer_trtllm() and get_moe_runner_backend().is_flashinfer_trtllm()
else torch.float32 ):
) 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( self.e_score_correction_bias = nn.Parameter(
torch.empty((config.n_routed_experts), dtype=correction_bias_dtype) torch.empty((config.n_routed_experts), dtype=correction_bias_dtype)
) )