support GLM-5.2 gate use FP32 (#47410)
Signed-off-by: zRzRzRzRzRzRzR <Yuxuan.Zhang2@liverpool.ac.uk>
Yuxuan Zhang committed
Jul 2, 2026 at 22:45 UTC
178fd5609427d4da1ca901ec0e78364d4d1c6d15
2 files changed
+19
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+1
@@ -74,6 +74,7 @@ class GateLinear(ReplicatedLinear):
74
self.allow_specialized_router_gemm = can_use_specialized_kernels
75
self.allow_dsv3_router_gemm = (
76
self.allow_specialized_router_gemm
77
+ and self.weight.dtype == torch.bfloat16
78
and output_size in self.DSV3_SUPPORTED_NUM_EXPERTS
79
and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
80
and (input_size, output_size) not in self.DSV3_UNSUPPORTED_SHAPES
vllm/model_executor/models/deepseek_v2.py
+18
@@ -117,6 +117,19 @@ from .utils import (
117
logger = init_logger(__name__)
118
119
120
+def _get_moe_router_dtype(
121
+ config: DeepseekV2Config | DeepseekV3Config,
122
+) -> torch.dtype | None:
123
+ router_dtype = getattr(config, "moe_router_dtype", None)
124
+ if getattr(config, "model_type", None) == "glm_moe_dsa":
125
+ # Older GLM-5/5.2 configs require fp32 routing but do not expose
126
+ # moe_router_dtype yet.
127
+ return torch.float32
128
+ if router_dtype == "float32":
129
+ return torch.float32
130
+ return None
131
+
132
+
133
class DeepseekAttention(nn.Module):
134
"""Normal MHA implementation used by Deepseek v1."""
135
@@ -288,9 +301,13 @@ class DeepseekV2MoE(nn.Module):
301
"Only silu is supported for now."
302
)
303
304
+ self.router_dtype = _get_moe_router_dtype(config)
305
self.gate = GateLinear(
306
config.hidden_size,
307
config.n_routed_experts,
308
+ params_dtype=self.router_dtype,
309
+ out_dtype=self.router_dtype,
310
+ force_fp32_compute=self.router_dtype == torch.float32,
311
prefix=f"{prefix}.gate",
312
)
313
if getattr(config, "topk_method", None) == "noaux_tc":
@@ -321,6 +338,7 @@ class DeepseekV2MoE(nn.Module):
338
if (
339
self.is_rocm_aiter_moe_enabled
340
and self.gate.e_score_correction_bias is not None
341
+ and self.gate.out_dtype is None
342
):
343
# Accumulates in fp32; avoids bf16->fp32 cast.
344
self.gate.set_out_dtype(self.gate.weight.dtype)