[Bugfix] Forward input IDs across DeepSeek V4 pipeline stages

Signed-off-by: ZacharyZcR <payasonorahc@protonmail.com>

Author: ZacharyZcR

tests/models/test_deepseek_v4_mega_moe.py

@@ -12,6 +12,7 @@
 from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM
 from vllm.models.deepseek_v4.nvidia.model import (
     DeepseekV4ForCausalLM,
+    DeepseekV4Model,
     DeepseekV4MegaMoEExperts,
     DeepseekV4MoE,
     make_deepseek_v4_expert_params_mapping,
@@ -26,6 +27,22 @@
 )
 
 
+def test_deepseek_v4_pp_intermediate_tensors_include_input_ids():
+    model = DeepseekV4Model.__new__(DeepseekV4Model)
+    model.hc_mult = 4
+    model.config = SimpleNamespace(hidden_size=16)
+
+    tensors = model.make_empty_intermediate_tensors(
+        batch_size=8,
+        dtype=torch.bfloat16,
+        device=torch.device("cpu"),
+    )
+
+    assert tensors["hidden_states"].shape == (8, 4, 16)
+    assert tensors["input_ids"].shape == (8,)
+    assert tensors["input_ids"].dtype == torch.int64
+
+
 def test_deepseek_v4_mega_moe_expert_mapping():
     mapping = make_deepseek_v4_expert_params_mapping(2)
 

vllm/models/deepseek_v4/nvidia/model.py

@@ -1386,24 +1386,31 @@ def make_empty_intermediate_tensors(
                     dtype=dtype,
                     device=device,
                 ),
+                "input_ids": torch.zeros(
+                    batch_size,
+                    dtype=torch.int64,
+                    device=device,
+                ),
             }
         )
 
     def forward(
         self,
-        input_ids: torch.Tensor,
+        input_ids: torch.Tensor | None,
         positions: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None,
         inputs_embeds: torch.Tensor | None = None,
     ) -> torch.Tensor | IntermediateTensors:
         if get_pp_group().is_first_rank:
+            assert input_ids is not None
             if inputs_embeds is not None:
                 hidden_states = inputs_embeds
             else:
                 hidden_states = self.embed_input_ids(input_ids)
         else:
             assert intermediate_tensors is not None
             hidden_states = intermediate_tensors["hidden_states"]
+            input_ids = intermediate_tensors["input_ids"]
 
         if self.use_mega_moe:
             input_ids = input_ids.to(torch.int64)
@@ -1457,6 +1464,7 @@ def forward(
             return IntermediateTensors(
                 {
                     "hidden_states": hidden_states,
+                    "input_ids": input_ids,
                     **self.pack_local_aux_hidden_states(aux_hidden_states),
                 }
             )