[Bugfix][MRV2] Keep input_ids on later PP stages for raw-token models (#61016)

Signed-off-by: booth-algo <kevinlauofficial01@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>

tests/distributed/test_distributed_oot.py

@@ -1,6 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
+
+import pytest
+
 from tests.plugins_tests.test_oot_registration_online import (
     run_and_test_dummy_opt_api_server,
 )
@@ -8,3 +12,17 @@ from tests.plugins_tests.test_oot_registration_online import (
 
 def test_distributed_oot(dummy_opt_path: str):
     run_and_test_dummy_opt_api_server(dummy_opt_path, tp=2)
+
+
+@pytest.mark.parametrize("enforce_eager", [True, False])
+def test_distributed_oot_pp_raw_input_tokens(dummy_opt_path: str, enforce_eager: bool):
+    # The model asserts that input_ids reaches every pipeline stage.
+    overrides = {"architectures": ["MyOPTRawInputTokensForCausalLM"]}
+    extra_args = ["-pp", "2", "--hf-overrides", json.dumps(overrides)]
+    if enforce_eager:
+        extra_args.append("--enforce-eager")
+    run_and_test_dummy_opt_api_server(
+        dummy_opt_path,
+        extra_args=extra_args,
+        env_dict={"VLLM_USE_V2_MODEL_RUNNER": "1"},
+    )

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py

@@ -6,11 +6,16 @@ from vllm import ModelRegistry
 
 def register():
     # Test directly passing the model
-    from .my_opt import MyOPTForCausalLM
+    from .my_opt import MyOPTForCausalLM, MyOPTRawInputTokensForCausalLM
 
     if "MyOPTForCausalLM" not in ModelRegistry.get_supported_archs():
         ModelRegistry.register_model("MyOPTForCausalLM", MyOPTForCausalLM)
 
+    if "MyOPTRawInputTokensForCausalLM" not in ModelRegistry.get_supported_archs():
+        ModelRegistry.register_model(
+            "MyOPTRawInputTokensForCausalLM", MyOPTRawInputTokensForCausalLM
+        )
+
     # Test passing lazy model
     if "MyGemma2Embedding" not in ModelRegistry.get_supported_archs():
         ModelRegistry.register_model(

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py

@@ -5,6 +5,7 @@
 import torch
 
 from vllm.model_executor.models.opt import OPTForCausalLM
+from vllm.sequence import IntermediateTensors
 
 
 class MyOPTForCausalLM(OPTForCausalLM):
@@ -15,3 +16,21 @@ class MyOPTForCausalLM(OPTForCausalLM):
             logits.zero_()
             logits[:, 0] += 1.0
         return logits
+
+
+class MyOPTRawInputTokensForCausalLM(MyOPTForCausalLM):
+    # Like models whose decoder layers read token ids, this model needs
+    # input_ids on every pipeline-parallel stage, not only the first.
+    requires_raw_input_tokens = True
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        assert input_ids is not None, "requires_raw_input_tokens model got no input_ids"
+        return super().forward(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )

tests/plugins_tests/test_oot_registration_online.py

@@ -7,7 +7,7 @@ chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja"
 assert chatml_jinja_path.exists()
 
 
-def run_and_test_dummy_opt_api_server(model, tp=1):
+def run_and_test_dummy_opt_api_server(model, tp=1, extra_args=(), env_dict=None):
     # the model is registered through the plugin
     server_args = [
         "--gpu-memory-utilization",
@@ -20,8 +20,9 @@ def run_and_test_dummy_opt_api_server(model, tp=1):
         "dummy",
         "-tp",
         f"{tp}",
+        *extra_args,
     ]
-    with RemoteOpenAIServer(model, server_args) as server:
+    with RemoteOpenAIServer(model, server_args, env_dict=env_dict) as server:
         client = server.get_client()
         completion = client.chat.completions.create(
             model=model,

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -33,6 +33,7 @@ from vllm.distributed.parallel_state import (
 )
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces import requires_raw_input_tokens
 from vllm.model_executor.offloader.base import get_offloader
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
@@ -690,7 +691,8 @@ class ModelCudaGraphManager(CudaGraphManager):
             }
             if not self.is_first_pp_rank:
                 # Update for non-first PP ranks.
-                model_inputs["input_ids"] = None
+                if not requires_raw_input_tokens(model):
+                    model_inputs["input_ids"] = None
                 model_inputs["inputs_embeds"] = None
                 assert intermediate_tensors is not None
                 model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens]

vllm/v1/worker/gpu/model_runner.py

@@ -1970,7 +1970,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         }
         if not self.is_first_pp_rank:
             # Update for non-first PP ranks.
-            model_inputs["input_ids"] = None
+            if not requires_raw_input_tokens(self.model):
+                model_inputs["input_ids"] = None
             model_inputs["inputs_embeds"] = None
 
             # Prepare the intermediate tensors.