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, 11:15 UTC
[DSV4] Avoid redundant dtype conversion. (#41374)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
May 1, 2026 at 00:57 UTC
307b17ce33165b41c490a099f54f0d0cd12e7f76
1 file changed
+11
-6
vllm/model_executor/models/deepseek_v4.py
+11
-6
@@ -854,10 +854,9 @@ class DeepseekV4MoE(nn.Module):
854
def forward(
855
self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
856
) -> torch.Tensor:
857
- if self.gate.tid2eid is not None:
858
- if input_ids is None:
859
- raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
860
- input_ids = input_ids.to(dtype=self.hash_indices_dtype)
857
+ if self.gate.tid2eid is not None and input_ids is None:
858
+ raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
859
+
860
if not self.use_mega_moe:
861
return self._forward_fused_moe(hidden_states, input_ids)
862
@@ -1225,7 +1224,12 @@ class DeepseekV4Model(nn.Module):
1224
config = vllm_config.model_config.hf_config
1225
quant_config = vllm_config.quant_config
1226
self.config = config
1228
-
1227
+ if vllm_config.parallel_config.enable_expert_parallel:
1228
+ self.use_mega_moe = (
1229
+ vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
1230
+ )
1231
+ else:
1232
+ self.use_mega_moe = False
1233
self.vocab_size = config.vocab_size
1234
self.hc_eps = config.hc_eps
1235
self.hc_mult = config.hc_mult
@@ -1309,7 +1313,8 @@ class DeepseekV4Model(nn.Module):
1313
) -> torch.Tensor | IntermediateTensors:
1314
hidden_states = self.embed_input_ids(input_ids)
1315
hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
1312
-
1316
+ if self.use_mega_moe:
1317
+ input_ids = input_ids.to(torch.int64)
1318
for layer in islice(self.layers, self.start_layer, self.end_layer):
1319
hidden_states = layer(
1320
hidden_states,