[ROCm][Bugfix][Kimi-K3] Preserve MoE correction bias in FP32 (#50761)
Signed-off-by: Fangzhou Ai <Fangzhou.Ai@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Fangzhou Ai committed
Aug 2, 2026 at 16:30 UTC
c6668106760e1f8222abd19a6ee9e93fc060cb19
1 file changed
+4
-1
vllm/models/kimi_k3/amd/linear.py
+4
-1
@@ -210,7 +210,10 @@ class KimiMoE(nn.Module):
210
prefix=f"{prefix}.gate",
211
)
212
213
- self.gate.e_score_correction_bias = nn.Parameter(torch.empty(num_experts))
213
+ # Preserve FP32 checkpoint values and match FP32 router logits.
214
+ self.gate.e_score_correction_bias = nn.Parameter(
215
+ torch.empty(num_experts, dtype=torch.float32)
216
+ )
217
218
if self.num_shared_experts is not None:
219
shared_intermediate_size = moe_intermediate_size * self.num_shared_experts