Fix GLM/DeepSeek NVFP4 + flashinfer_trtllm long-context "!!!!" collapse (NaN routing) (#31001)
This commit is contained in:
@@ -457,6 +457,14 @@ class MoEGate(nn.Module):
|
|||||||
"quark",
|
"quark",
|
||||||
):
|
):
|
||||||
correction_bias_dtype = torch.bfloat16
|
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(
|
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)
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user