diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index b97168161..0cd0a1ab3 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -1,9 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import contextlib from types import SimpleNamespace import pytest import torch -from vllm.config import CacheConfig +from vllm.config import CacheConfig, CUDAGraphMode from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT import vllm_omni.worker.gpu_generation_model_runner as gen_runner_module @@ -15,6 +18,10 @@ from vllm_omni.worker.gpu_generation_model_runner import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +class _StopExecutionError(Exception): + """Halts execute_model at the model boundary.""" + + class _DummyInputBatch: def __init__(self, num_reqs: int = 1): self.req_ids = [f"req-{i + 1}" for i in range(num_reqs)] @@ -98,6 +105,9 @@ class _StubSchedulerOutput: self.num_scheduled_tokens = {"req-1": total_num_scheduled_tokens} self.finished_req_ids = set() self.kv_connector_metadata = None + self.scheduled_encoder_inputs = {} + self.scheduled_spec_decode_tokens = {} + self.num_common_prefix_blocks = 0 def _make_guard_runner(): @@ -109,6 +119,7 @@ def _make_guard_runner(): runner.speculative_config = None runner.model_config = SimpleNamespace(async_chunk=False) runner.cache_config = CacheConfig() + runner.kv_cache_config = SimpleNamespace(kv_cache_groups=[]) runner.input_batch = _DummyInputBatch() runner.model = object() runner.parallel_config = SimpleNamespace( @@ -236,3 +247,74 @@ class TestSampleTokensBatched: with pytest.raises(ValueError, match="length 1 .* 2 requests"): GPUGenerationModelRunner.sample_tokens(runner) + + +@pytest.mark.parametrize("exact_shape, expected_len", [(True, 3), (False, 4)]) +def test_exact_shape_models_receive_unpadded_input_ids(monkeypatch, exact_shape, expected_len): + """#6712: the cudagraph dispatcher rounds a 3-token batch up to a capture + size of 4, so ``_preprocess`` returns a 4-element ``input_ids`` while + ``seq_token_counts`` stays ``[3]``. Models opting into + ``requires_exact_input_shape`` must receive the trimmed view; everything + else must keep the padded buffer. + """ + monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False) + monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False) + + runner = _make_guard_runner() + runner.model = SimpleNamespace( + has_preprocess=True, + requires_request_ids=False, + requires_exact_input_shape=exact_shape, + ) + runner.cascade_attn_enabled = False + runner.num_prompt_logprobs = None + runner.vllm_config = None + runner.parallel_config = SimpleNamespace( + distributed_executor_backend=None, + data_parallel_size=1, + use_ubatching=False, + num_ubatches=1, + ) + + monkeypatch.setattr(GPUGenerationModelRunner, "_prepare_inputs", lambda self, so, nst: (None, None, 1)) + monkeypatch.setattr( + GPUGenerationModelRunner, + "_determine_batch_execution_and_padding", + lambda self, **kw: ( + CUDAGraphMode.PIECEWISE, + SimpleNamespace(num_tokens=4, num_reqs=1), + False, + None, + None, + ), + ) + monkeypatch.setattr(GPUGenerationModelRunner, "_get_slot_mappings", lambda self, **kw: (None, None)) + monkeypatch.setattr(GPUGenerationModelRunner, "_build_attention_metadata", lambda self, **kw: (None, None)) + monkeypatch.setattr( + GPUGenerationModelRunner, + "_maybe_attach_attention_metadata_extensions", + lambda self, **kw: None, + ) + monkeypatch.setattr(gen_runner_module, "set_forward_context", lambda *a, **kw: contextlib.nullcontext()) + + def _fake_preprocess(self, so, num_input_tokens, itensors=None): + return (torch.zeros(num_input_tokens, dtype=torch.int32), None, None, None, {}, None) + + monkeypatch.setattr(GPUGenerationModelRunner, "_preprocess", _fake_preprocess) + + seen = {} + + def _capture(self, **kwargs): + seen["input_ids"] = kwargs["input_ids"] + seen["counts"] = kwargs["model_kwargs"]["seq_token_counts"] + raise _StopExecutionError + + monkeypatch.setattr(GPUGenerationModelRunner, "_run_generation_model", _capture) + + with pytest.raises(_StopExecutionError): + GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(3)) + + assert seen["input_ids"].numel() == expected_len + assert seen["counts"] == [3] + if exact_shape: + assert sum(seen["counts"]) == seen["input_ids"].numel() diff --git a/tests/worker/test_mixins.py b/tests/worker/test_mixins.py new file mode 100644 index 000000000..c6e2ff83c --- /dev/null +++ b/tests/worker/test_mixins.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the shared runner helpers in vllm_omni.worker.mixins.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.worker.mixins import maybe_unpad_input_ids + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + "flag, padded_len, unpadded, expected", + [ + (True, 4, 3, 3), + (False, 4, 3, 4), + (True, 3, 3, 3), + ], +) +def test_maybe_unpad_input_ids(flag, padded_len, unpadded, expected): + """#6712: opted-in models get input_ids at the real token count; everything + else keeps the graph-bucket padded buffer.""" + model = SimpleNamespace(requires_exact_input_shape=flag) + ids = torch.zeros(padded_len, dtype=torch.int32) + assert maybe_unpad_input_ids(model, ids, unpadded).numel() == expected + + +def test_maybe_unpad_input_ids_without_capability_attribute(): + ids = torch.zeros(4, dtype=torch.int32) + assert maybe_unpad_input_ids(SimpleNamespace(), ids, 3).numel() == 4 + + +def test_maybe_unpad_input_ids_handles_none(): + model = SimpleNamespace(requires_exact_input_shape=True) + assert maybe_unpad_input_ids(model, None, 3) is None diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 7017f703f..8b3926cc0 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -143,6 +143,7 @@ class MiniCPMO45Code2Wav(nn.Module): replace_runtime_additional_information = True requires_raw_input_tokens = True requires_request_ids = True + requires_exact_input_shape = True has_preprocess = False has_postprocess = False diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index 2a1e9a468..231e33857 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). @@ -226,6 +229,7 @@ class MossTTSCodecDecoder(nn.Module): has_postprocess: bool = False enable_update_additional_information: bool = True requires_raw_input_tokens: bool = True + requires_exact_input_shape: bool = True _OUTPUT_SAMPLE_RATE: int = 24_000 diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 71e2408d7..a5856718d 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The Qwen team. """Inference-only Qwen3-Omni-Moe unified model (thinker + talker + code2wav).""" @@ -263,6 +263,8 @@ class Qwen3OmniMoeForConditionalGeneration( ) self.model = self.code2wav self.requires_raw_input_tokens = True + # torch.split below requires sum(counts) == input_ids.numel() (#6712) + self.requires_exact_input_shape = True else: raise ValueError( f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'talker', 'code2wav'" diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index eeac7d059..fc26882cd 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -38,6 +38,7 @@ from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState, _ensure_tensor_values from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner from vllm_omni.utils.mm_outputs import partition_payload_list +from vllm_omni.worker.mixins import maybe_unpad_input_ids from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, needs_omni_connector, @@ -336,6 +337,8 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin ) # -------------------------------------- Omni-new ------------------------------------------------- + # [Omni] Exact-shape models need input_ids at its real token count (#6712). + input_ids = maybe_unpad_input_ids(self.model, input_ids, scheduler_output.total_num_scheduled_tokens) # [Omni] Pass token counts per request for code2wav output slicing model_kwargs["seq_token_counts"] = tokens if getattr(self.model, "requires_request_ids", False): diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 307a5a428..e43c6f608 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -42,6 +42,7 @@ from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import partition_payload_list from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, _ensure_tensor_values from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner +from vllm_omni.worker.mixins import maybe_unpad_input_ids from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, needs_omni_connector, @@ -307,6 +308,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin num_tokens_padded, intermediate_tensors, ) + # [Omni] Exact-shape models need input_ids at its real token count (#6712). + input_ids = maybe_unpad_input_ids(self.model, input_ids, num_tokens_unpadded) + # [Omni] Pass token counts per request for code2wav output slicing model_kwargs["seq_token_counts"] = tokens if getattr(self.model, "requires_request_ids", False): diff --git a/vllm_omni/worker/mixins.py b/vllm_omni/worker/mixins.py index 5b25b8362..5f0fd8300 100644 --- a/vllm_omni/worker/mixins.py +++ b/vllm_omni/worker/mixins.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from typing import Any @@ -11,3 +14,11 @@ class OmniWorkerMixin: from vllm_omni.plugins import load_omni_general_plugins load_omni_general_plugins() + + +def maybe_unpad_input_ids(model, input_ids, num_tokens_unpadded): + """Trim cudagraph-bucket padding for models that split a flat ``input_ids`` + by per-request ``seq_token_counts``. See #6712.""" + if input_ids is not None and getattr(model, "requires_exact_input_shape", False): + return input_ids[:num_tokens_unpadded] + return input_ids