diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 6eb96aadf..32c4c990d 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -457,6 +457,14 @@ class MoEGate(nn.Module): "quark", ): correction_bias_dtype = torch.bfloat16 + # NOTE(kpham-sgl): flashinfer trtllm routing requires a bf16 + # routing_bias; an fp32 bias yields NaN routing on exact ties. + # Mirror the fp8 path's cast. + if ( + quant_config.get_name() == "modelopt_fp4" + and get_moe_runner_backend().is_flashinfer_trtllm() + ): + correction_bias_dtype = torch.bfloat16 self.e_score_correction_bias = nn.Parameter( torch.empty((config.n_routed_experts), dtype=correction_bias_dtype) )