From d12ea8e309d71a79d2001255f500e9642cfe87ef Mon Sep 17 00:00:00 2001 From: Kevin <107279223+booth-algo@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:54:21 +0100 Subject: [PATCH] [Bugfix][MRV2] Keep input_ids on later PP stages for raw-token models (#61016) Signed-off-by: booth-algo Co-authored-by: Claude Opus 5.5 Co-authored-by: Kevin H. Luu --- tests/distributed/test_distributed_oot.py | 18 ++++++++++++++++++ .../vllm_add_dummy_model/__init__.py | 7 ++++++- .../vllm_add_dummy_model/my_opt.py | 19 +++++++++++++++++++ .../test_oot_registration_online.py | 5 +++-- vllm/v1/worker/gpu/cudagraph_utils.py | 4 +++- vllm/v1/worker/gpu/model_runner.py | 3 ++- 6 files changed, 51 insertions(+), 5 deletions(-) diff --git a/tests/distributed/test_distributed_oot.py b/tests/distributed/test_distributed_oot.py index 5f7f3ffa8a..51ae8b35f8 100644 --- a/tests/distributed/test_distributed_oot.py +++ b/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"}, + ) diff --git a/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py b/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py index 457187e4b4..74ffac1f24 100644 --- a/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py +++ b/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( diff --git a/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py b/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py index f1e6e7b10f..da3ad64253 100644 --- a/tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py +++ b/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 + ) diff --git a/tests/plugins_tests/test_oot_registration_online.py b/tests/plugins_tests/test_oot_registration_online.py index 151373d82f..1c0880bf1d 100644 --- a/tests/plugins_tests/test_oot_registration_online.py +++ b/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, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 0a74b685b4..a1955b09c7 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/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] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 5ba90b792e..cb6ab75f63 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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.