A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 16:30 UTC
[Bugfix] Fix GLM-4 MoE router logits dtype for data parallel chunking (#31055)
Signed-off-by: ReinforcedKnowledge <reinforced.knowledge@gmail.com>
Yakine Tahtah committed
Jan 6, 2026 at 18:57 UTC
4e67a8f616f4f202fd7a549978914f398a738d49
3 files changed
+13
-1
vllm/model_executor/layers/fused_moe/config.py
+6
@@ -1006,6 +1006,9 @@ class FusedMoEConfig:
1006
# The activation type.
1007
in_dtype: torch.dtype
1008
1009
+ # Defaults to in_dtype if not specified.
1010
+ router_logits_dtype: torch.dtype | None = None
1011
+
1012
max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE
1013
1014
has_bias: bool = False
@@ -1022,6 +1025,9 @@ class FusedMoEConfig:
1025
1026
assert self.max_num_tokens > 0
1027
1028
+ if self.router_logits_dtype is None:
1029
+ self.router_logits_dtype = self.in_dtype
1030
+
1031
@property
1032
def tp_size(self):
1033
return self.moe_parallel_config.tp_size
vllm/model_executor/layers/fused_moe/layer.py
+6
-1
@@ -314,6 +314,7 @@ class FusedMoE(CustomOp):
314
renormalize: Whether to renormalize the logits in the fused_moe kernel
315
quant_config: Quantization configure.
316
enable_eplb: Whether to enable expert parallelism load balancer.
317
+ router_logits_dtype: Data type for router logits buffers.
318
"""
319
320
def __init__(
@@ -348,6 +349,7 @@ class FusedMoE(CustomOp):
349
expert_mapping: list[tuple[str, str, int, str]] | None = None,
350
n_shared_experts: int | None = None,
351
routing_method_type: int | None = None,
352
+ router_logits_dtype: torch.dtype | None = None,
353
):
354
super().__init__()
355
@@ -559,6 +561,7 @@ class FusedMoE(CustomOp):
561
num_local_experts=self.local_num_experts,
562
moe_parallel_config=self.moe_parallel_config,
563
in_dtype=moe_in_dtype,
564
+ router_logits_dtype=router_logits_dtype,
565
max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
566
has_bias=has_bias,
567
is_act_and_mul=is_act_and_mul,
@@ -1509,7 +1512,9 @@ class FusedMoE(CustomOp):
1512
)
1513
1514
self.batched_router_logits = torch.zeros(
1512
- logits_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()
1515
+ logits_shape,
1516
+ dtype=moe.router_logits_dtype,
1517
+ device=torch.cuda.current_device(),
1518
)
1519
1520
def select_experts(
vllm/model_executor/models/glm4_moe.py
+1
@@ -197,6 +197,7 @@ class Glm4MoE(nn.Module):
197
e_score_correction_bias=self.gate.e_score_correction_bias,
198
enable_eplb=self.enable_eplb,
199
num_redundant_experts=self.n_redundant_experts,
200
+ router_logits_dtype=torch.float32,
201
)
202
203
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: