From 307b17ce33165b41c490a099f54f0d0cd12e7f76 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Fri, 1 May 2026 00:57:27 +0800 Subject: [PATCH] [DSV4] Avoid redundant dtype conversion. (#41374) Signed-off-by: Jee Jee Li --- vllm/model_executor/models/deepseek_v4.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/models/deepseek_v4.py b/vllm/model_executor/models/deepseek_v4.py index baf28d04581..5521a9764a9 100644 --- a/vllm/model_executor/models/deepseek_v4.py +++ b/vllm/model_executor/models/deepseek_v4.py @@ -854,10 +854,9 @@ class DeepseekV4MoE(nn.Module): def forward( self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None ) -> torch.Tensor: - if self.gate.tid2eid is not None: - if input_ids is None: - raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") - input_ids = input_ids.to(dtype=self.hash_indices_dtype) + if self.gate.tid2eid is not None and input_ids is None: + raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + if not self.use_mega_moe: return self._forward_fused_moe(hidden_states, input_ids) @@ -1225,7 +1224,12 @@ class DeepseekV4Model(nn.Module): config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config self.config = config - + if vllm_config.parallel_config.enable_expert_parallel: + self.use_mega_moe = ( + vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe" + ) + else: + self.use_mega_moe = False self.vocab_size = config.vocab_size self.hc_eps = config.hc_eps self.hc_mult = config.hc_mult @@ -1309,7 +1313,8 @@ class DeepseekV4Model(nn.Module): ) -> torch.Tensor | IntermediateTensors: hidden_states = self.embed_input_ids(input_ids) hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1) - + if self.use_mega_moe: + input_ids = input_ids.to(torch.int64) for layer in islice(self.layers, self.start_layer, self.end_layer): hidden_states = layer( hidden_states,