Fix GLM/DeepSeek NVFP4 + flashinfer_trtllm long-context "!!!!" collapse (NaN routing) (#31001)

2026-07-13 12:54:38 -07:00
parent 86c59ac1aa
commit f49cbbd67d
1 changed files with 8 additions and 0 deletions
@@ -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)
            )