[Bugfix] Give exact-shape codec models an unpadded input_ids view (#6775)

Signed-off-by: Inesh Reddy <ineshreddy249@gmail.com>
Signed-off-by: IneshReddy249 <ineshreddy249@gmail.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

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()

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

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
 

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
 

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'"

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):

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):

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