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 4, 08:45 UTC
[ROCm][Bugfix] Fix init-time bias dtype cast when gate.out_dtype is None (#41405)
Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
Rita Brugarolas committed
May 1, 2026 at 21:13 UTC
c293ccc58ef6e1a0976a62f79f57bc045108073d
1 file changed
+2
-1
vllm/model_executor/models/deepseek_v2.py
+2
-1
@@ -351,8 +351,9 @@ class DeepseekV2MoE(nn.Module):
351
self.is_rocm_aiter_moe_enabled
352
and self.gate.e_score_correction_bias is not None
353
):
354
+ gate_out_dtype = self.gate.out_dtype or self.gate.weight.dtype
355
self.gate.e_score_correction_bias.data = (
355
- self.gate.e_score_correction_bias.data.to(self.gate.out_dtype)
356
+ self.gate.e_score_correction_bias.data.to(gate_out_dtype)
357
)
358
359
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: