[Refactor] Worker/ModelRunner runner correctness fixes (G2/N) (#5452)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

tests/model_executor/models/test_gepard_wiring.py

@@ -10,9 +10,10 @@ survives the runner's sparse-audio routing.
 
 These drive the real ``preprocess`` -> ``forward`` -> ``make_omni_output``
 sequence through a stand-in for ``GPUARModelRunner.execute_model``, and use the
-runner's own ``_resolve_sparse_mm_routing`` to decide what actually reaches the
-caller. Only the LM backbone and the frame sampler are stubbed out; every
-method under test is the real one, so a test here goes red when the engine
+shared ``sparse_audio.resolve_sparse_mm_routing`` helper used by the runner to
+decide what actually reaches the caller. Only the LM backbone and the frame
+sampler are stubbed out; every method under test is the real one, so a test
+here goes red when the engine
 path breaks — which a test that calls the flush by hand cannot do.
 """
 
@@ -33,7 +34,7 @@ from vllm_omni.model_executor.models.gepard.gepard_talker import (
 )
 from vllm_omni.model_executor.models.gepard.nanocodec import NanoCodec
 from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
-from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+from vllm_omni.worker import sparse_audio
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -219,7 +220,7 @@ class _MiniEngine:
                 self.committed_frames[req_id] = state.frame_count
 
         # --- runner: sparse-audio routing decides what reaches the caller ---
-        downstream, sparse_index, is_sparse = GPUARModelRunner._resolve_sparse_mm_routing(
+        downstream, sparse_index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
             engine_output_type="audio",
             req_ids_output_copy=list(live),
             downstream_req_ids=list(live),

tests/model_executor/models/voxcpm2/test_talker_output_marker.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""VoxCPM2 dense-mode audio outputs must carry the sparse-alignment marker.
+
+Dense mode (`_uses_sparse_audio_outputs()` False) still yields a strict
+SUBSET of the batch whenever some requests emit no audio in a step (e.g.
+prefill phase). The runner routes per-request `model_outputs` lists by batch
+position unless `meta.sparse_audio` declares `meta.req_id` alignment — so a
+subset without the marker misroutes audio across requests (RFC #5450 C3
+review follow-up).
+"""
+
+from __future__ import annotations
+
+import functools
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+torch = pytest.importorskip("torch")
+
+
+@functools.lru_cache(maxsize=1)
+def _talker_cls():
+    """Defer talker import (pulls vLLM model_executor) until first use."""
+    from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
+        VoxCPM2TalkerForConditionalGeneration,
+    )
+
+    return VoxCPM2TalkerForConditionalGeneration
+
+
+def _make_dense_talker():
+    cls = _talker_cls()
+    talker = cls.__new__(cls)
+    # Dense mode: emit every step, no delayed copy, no batched VAE decode.
+    talker._audio_emit_every = 1
+    talker._vae_decode_every = 1
+    talker._enable_delayed_audio_copy = False
+    talker._coalesce_audio_d2h = False
+    talker._sample_rate = 16000
+    talker._audio_queue = []
+    return talker
+
+
+def test_dense_subset_audio_carries_sparse_alignment_marker():
+    talker = _make_dense_talker()
+    # Batch had [r0, r1, r2]; only r1/r2 produced audio this step.
+    talker._audio_queue = [("r1", torch.zeros(4)), ("r2", torch.ones(4))]
+
+    out = talker.make_omni_output(torch.zeros(1))
+
+    mm = out.multimodal_outputs
+    assert mm["meta"] == {"req_id": ["r1", "r2"], "sparse_audio": ["1"]}
+    assert len(mm["model_outputs"]) == 2
+    assert len(mm["sr"]) == 2
+    # (The coalesce-D2H sibling branch carries the same marker; it needs
+    # device-resident chunks, so its marker contract is pinned here via the
+    # shared dense-subset shape rather than a GPU-only test.)
+
+
+def test_dense_full_batch_audio_still_carries_marker_and_order():
+    # Full coverage is the degenerate subset: the marker stays truthful and
+    # `meta.req_id` order matches the emitted list order.
+    talker = _make_dense_talker()
+    talker._audio_queue = [("r0", torch.zeros(2)), ("r1", torch.ones(2))]
+
+    out = talker.make_omni_output(torch.zeros(1))
+
+    mm = out.multimodal_outputs
+    assert mm["meta"]["req_id"] == ["r0", "r1"]
+    assert mm["meta"]["sparse_audio"] == ["1"]

tests/platforms/npu/test_minicpmo_code2wav_npugraph.py

@@ -32,6 +32,7 @@ pytestmark = [
 class _Flow(nn.Module):
     def __init__(self, estimator: nn.Module):
         super().__init__()
+        self.encoder = nn.Identity()
         self.decoder = SimpleNamespace(estimator=estimator)
 
 

tests/worker/test_capture_replay_contract.py

@@ -14,9 +14,9 @@ cannot be masked by an adjacent ``else``/``elif`` that happens to read it. It
 matches by attribute chain (``self.input_ids.gpu``), so local variable renames
 do not break it.
 
-Scope: the BASE runner (``OmniGPUModelRunner``) only. The generation runner is
-KNOWN to lack the ``has_preprocess`` branch; extend this
-check to the generation file when it adds the branch.
+Scope: the BASE runner (``OmniGPUModelRunner``) and the generation runner's
+``_dummy_run`` override (which gained the branch in RFC #5450 C6; its
+``_preprocess`` is inherited from the base).
 """
 
 import ast
@@ -25,6 +25,7 @@ import textwrap
 
 import pytest
 
+from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -69,11 +70,24 @@ def _has_preprocess_branch_reads_both_buffers(method) -> bool:
     return any(_REQUIRED_BUFFERS <= _attr_chains_in(node.body) for node in guarded)
 
 
-@pytest.mark.parametrize("method_name", ["_dummy_run", "_preprocess"])
-def test_has_preprocess_branch_reads_shared_buffers(method_name):
-    method = getattr(OmniGPUModelRunner, method_name)
+@pytest.mark.parametrize(
+    ("runner_cls", "method_name"),
+    [
+        (OmniGPUModelRunner, "_dummy_run"),
+        (OmniGPUModelRunner, "_preprocess"),
+        (GPUGenerationModelRunner, "_dummy_run"),
+    ],
+)
+def test_has_preprocess_branch_reads_shared_buffers(runner_cls, method_name):
+    method = getattr(runner_cls, method_name)
     assert _has_preprocess_branch_reads_both_buffers(method), (
-        f"{OmniGPUModelRunner.__name__}.{method_name}: the has_preprocess branch no "
+        f"{runner_cls.__name__}.{method_name}: the has_preprocess branch no "
         "longer reads BOTH self.input_ids.gpu and self.inputs_embeds.gpu — capture "
         "and replay would diverge."
     )
+
+
+def test_generation_dummy_run_overrides_with_branch():
+    # Guard against the override being deleted (which would silently fall back
+    # to the base implementation and vacuously pass the parametrized check).
+    assert "_dummy_run" in GPUGenerationModelRunner.__dict__

tests/worker/test_gpu_ar_model_runner.py

@@ -1,8 +1,12 @@
 import asyncio
 import copy
+import re
 from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
+from threading import Barrier
 from types import SimpleNamespace
 from typing import Any
+from unittest.mock import patch
 
 import numpy as np
 import pytest
@@ -15,12 +19,15 @@ from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechReques
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
 from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker import sparse_audio
 from vllm_omni.worker.gpu_ar_model_runner import (
     ExecuteModelState,
     GPUARModelRunner,
     OmniAsyncGPUModelRunnerOutput,
 )
+from vllm_omni.worker.output import payload_build
 from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
+from vllm_omni.worker.sampling_utils import clamp_prompt_ids_to_penalty_padding
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -301,11 +308,13 @@ def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests(
 
 
 def test_sparse_mm_req_ids_requires_sparse_audio_marker():
-    assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None
-    assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None
+    assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None
+    assert sparse_audio.resolve_sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None
 
-    assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"]
-    assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
+    assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"]
+    assert sparse_audio.resolve_sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
+    assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"]}, "meta.sparse_audio": ["1"]}) == ["r1"]
+    assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"sparse_audio": ["1"]}, "meta.req_id": ["r1"]}) == ["r1"]
 
 
 def test_runner_assisted_full_attention_metadata_request_is_opt_in():
@@ -1198,3 +1207,642 @@ def test_build_omni_output_uses_combined_prefix_cache_mm_payload_for_partial_dow
 
     assert torch.equal(output.inter_stage_outputs[0]["codes.ref"], ref_codes[0])
     assert torch.equal(output.inter_stage_outputs[2]["codes.ref"], ref_codes[2])
+
+
+class TestPromptIdsClampPaddingConvention:
+    """Pin `sampling_utils.clamp_prompt_ids_to_penalty_padding` (the clamp
+    `sample_tokens` applies before penalties).
+
+    The audit (RFC #5450 C2) originally flagged the clamp as an off-by-one
+    (OOB index). It is not: upstream penalties allocate `vocab_size + 1` bins
+    and drop the last column, so `vocab_size` is the designed padding value
+    (vllm/model_executor/layers/utils.py::get_token_bin_counts_and_mask).
+    These tests call the runner's own helper, so a future "fix" to the clamp
+    fails here.
+    """
+
+    LOGITS_VOCAB = 8
+    BATCH_VOCAB = 10  # wider input_batch vocab; its pad value is BATCH_VOCAB
+
+    def _padded_prompt(self):
+        # Two real tokens + two padding slots padded with the batch vocab size.
+        return torch.tensor([[1, 2, self.BATCH_VOCAB, self.BATCH_VOCAB]])
+
+    def test_clamp_to_logits_vocab_keeps_padding_inert(self):
+        from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask
+
+        # The runner's own helper, not a re-implementation.
+        clamped = clamp_prompt_ids_to_penalty_padding(self._padded_prompt(), self.LOGITS_VOCAB)
+
+        _, mask = get_token_bin_counts_and_mask(clamped, self.LOGITS_VOCAB, 1)
+
+        expected = torch.zeros(1, self.LOGITS_VOCAB, dtype=torch.bool)
+        expected[0, 1] = True
+        expected[0, 2] = True
+        # Pad ids landed in the dropped (vocab_size-th) bin: no penalty effect.
+        assert torch.equal(mask, expected)
+
+    def test_clamp_one_lower_would_corrupt_last_token_penalties(self):
+        from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask
+
+        # The "fix" the audit first proposed - shown here to be the actual bug:
+        # padding would count as occurrences of the last real vocab token.
+        clamped = self._padded_prompt().clamp(max=self.LOGITS_VOCAB - 1)
+
+        _, mask = get_token_bin_counts_and_mask(clamped, self.LOGITS_VOCAB, 1)
+
+        assert bool(mask[0, self.LOGITS_VOCAB - 1])
+
+    def test_unclamped_padding_overflows_penalty_bins(self):
+        from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask
+
+        # Why the clamp exists: the batch-level pad value (BATCH_VOCAB) is out
+        # of range even for the vocab_size+1 penalty bins of a narrower head.
+        with pytest.raises(RuntimeError):
+            get_token_bin_counts_and_mask(self._padded_prompt(), self.LOGITS_VOCAB, 1)
+
+
+class TestSparseAudioMarkerRobustness:
+    """RFC #5450 C3: `sparse_audio.is_sparse_audio_marker` must not crash on tensor/array
+    markers, and the combined prefix-cache payload builder must never ship one
+    request's data as another's."""
+
+    def test_marker_list_and_str_forms(self):
+        marker = sparse_audio.is_sparse_audio_marker
+        assert marker(["1"]) is True
+        assert marker(["0"]) is False
+        assert marker("true") is True
+        assert marker("off") is False
+
+    def test_marker_enforced_canonical_forms_no_error(self):
+        # The designed carrier (list-of-str; bare str accepted) and marker
+        # absence (None) are the ONLY silent forms.
+        marker = sparse_audio.is_sparse_audio_marker
+        with patch.object(sparse_audio.logger, "error") as err:
+            assert marker(["1"]) is True
+            assert marker(["0"]) is False
+            assert marker(["0", "1"]) is True
+            assert marker("true") is True
+            assert marker("off") is False
+            assert marker(None) is False
+        err.assert_not_called()
+
+    def test_marker_illegal_carriers_raise_one_error_per_type(self):
+        # Enforcement, not normalization: any carrier outside the design
+        # contract is a producer bug - it raises InvalidSparseDeclarationError
+        # (caught inside the module by resolve_sparse_mm_routing, which
+        # fails closed)
+        # and logs at error level once per offending type. Interpreting
+        # deviants would need a second truth predicate, and predicate
+        # divergence is exactly how the np.array(["0"])-routes-as-sparse bug
+        # happened.
+        marker = sparse_audio.is_sparse_audio_marker
+        for value in (torch.tensor(1), torch.tensor([1, 1]), np.array(["1"]), 1, 1.0, True, ["1", 1]):
+            with (
+                patch.object(sparse_audio, "_logged_offenders", set()),
+                patch.object(sparse_audio.logger, "error") as err,
+            ):
+                with pytest.raises(sparse_audio.InvalidSparseDeclarationError):
+                    marker(value)
+                assert err.call_count == 1, f"expected one error for {value!r}"
+                with pytest.raises(sparse_audio.InvalidSparseDeclarationError):
+                    marker(value)
+                assert err.call_count == 1  # deduped per type, not per step
+
+    def test_marker_device_tensor_rejected_without_read(self):
+        # Deviant carriers are never read element-wise, so a device tensor
+        # can never introduce a D2H sync here. Meta tensors hold no data -
+        # any read (`.item()`/`bool()`) raises a RuntimeError, so seeing OUR
+        # exception type (not RuntimeError) proves the data is never touched.
+        marker = sparse_audio.is_sparse_audio_marker
+        with patch.object(sparse_audio, "_logged_offenders", set()):
+            with pytest.raises(sparse_audio.InvalidSparseDeclarationError):
+                marker(torch.ones((), device="meta"))
+            with pytest.raises(sparse_audio.InvalidSparseDeclarationError):
+                marker(torch.ones(3, device="meta"))
+
+    def test_invalid_marker_fails_closed_at_routing(self):
+        # A present-but-invalid sparse declaration most likely belongs to a
+        # producer that INTENDED sparse output; falling back to dense would
+        # batch-index-assign its payloads to the WRONG requests. Routing
+        # must fail closed: no payload routed for the step (empty sparse
+        # set), never dense misrouting.
+        with patch.object(sparse_audio, "_logged_offenders", set()):
+            downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+                engine_output_type="audio",
+                req_ids_output_copy=["r0", "r1"],
+                downstream_req_ids=["r0", "r1"],
+                multimodal_outputs={
+                    "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])},
+                    "model_outputs": [torch.zeros(1)],
+                },
+            )
+        assert downstream == []
+        assert index == {}
+        assert is_sparse is True
+
+    def test_fail_closed_routing_error_remains_observable(self):
+        def route_invalid_marker():
+            return sparse_audio.resolve_sparse_mm_routing(
+                engine_output_type="audio",
+                req_ids_output_copy=["r0"],
+                downstream_req_ids=["r0"],
+                multimodal_outputs={
+                    "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])},
+                    "model_outputs": [torch.zeros(1)],
+                },
+            )
+
+        with (
+            patch.object(sparse_audio, "_logged_offenders", set()),
+            patch.object(sparse_audio, "_fail_closed_last_log_at", None),
+            patch.object(sparse_audio, "_fail_closed_suppressed", 0),
+            patch.object(sparse_audio.time, "monotonic", side_effect=[100.0, 120.0, 161.0]),
+            patch.object(sparse_audio.logger, "error") as err,
+        ):
+            for _ in range(3):
+                assert route_invalid_marker() == ([], {}, True)
+
+        routing_errors = [call for call in err.call_args_list if "no multimodal payload" in call.args[0]]
+        assert len(routing_errors) == 2
+        assert [call.args[1] for call in routing_errors] == [0, 1]
+
+    def test_fail_closed_logging_is_thread_safe(self):
+        workers = 8
+        ready = Barrier(workers)
+
+        def route_invalid_marker():
+            return sparse_audio.resolve_sparse_mm_routing(
+                engine_output_type="audio",
+                req_ids_output_copy=["r0"],
+                downstream_req_ids=["r0"],
+                multimodal_outputs={
+                    "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])},
+                },
+            )
+
+        def route_concurrently(_):
+            ready.wait(timeout=10)
+            return route_invalid_marker()
+
+        with (
+            patch.object(sparse_audio, "_logged_offenders", set()),
+            patch.object(sparse_audio, "_fail_closed_last_log_at", None),
+            patch.object(sparse_audio, "_fail_closed_suppressed", 0),
+            patch.object(sparse_audio.time, "monotonic", return_value=100.0) as clock,
+            patch.object(sparse_audio.logger, "error") as err,
+        ):
+            with ThreadPoolExecutor(max_workers=workers) as executor:
+                results = list(executor.map(route_concurrently, range(workers)))
+            assert results == [([], {}, True)] * workers
+
+            clock.return_value = 161.0
+            assert route_invalid_marker() == ([], {}, True)
+
+        routing_errors = [call for call in err.call_args_list if "no multimodal payload" in call.args[0]]
+        assert len(routing_errors) == 2
+        assert [call.args[1] for call in routing_errors] == [0, workers - 1]
+        classifier_errors = [call for call in err.call_args_list if "Illegal sparse-audio marker" in call.args[0]]
+        assert len(classifier_errors) == 1
+
+    def test_duplicate_sparse_metadata_must_be_consistent(self):
+        equal = {
+            "meta": {"req_id": ["r1"], "sparse_audio": ["1"]},
+            "meta.req_id": ["r1"],
+            "meta.sparse_audio": ["true"],
+        }
+        assert sparse_audio.resolve_sparse_mm_req_ids(equal) == ["r1"]
+
+        invalid_declarations = (
+            {
+                "meta": {"req_id": ["r1"], "sparse_audio": ["1"]},
+                "meta.req_id": ["r2"],
+            },
+            {
+                "meta": {"req_id": ["r1"], "sparse_audio": ["1"]},
+                "meta.req_id": "r1",
+            },
+            {
+                "meta": {"req_id": ["r1"], "sparse_audio": ["0"]},
+                "meta.sparse_audio": ["1"],
+            },
+            {"meta": {"req_id": ["r1", "r1"], "sparse_audio": ["1"]}},
+        )
+        with patch.object(sparse_audio, "_logged_offenders", set()):
+            for declaration in invalid_declarations:
+                with pytest.raises(sparse_audio.InvalidSparseDeclarationError):
+                    sparse_audio.resolve_sparse_mm_req_ids(declaration)
+
+    def test_conflicting_sparse_metadata_fails_closed_at_public_boundary(self):
+        downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+            engine_output_type="audio",
+            req_ids_output_copy=["r1", "r2"],
+            downstream_req_ids=["r1", "r2"],
+            multimodal_outputs={
+                "meta": {"req_id": ["r1"], "sparse_audio": ["1"]},
+                "meta.req_id": ["r2"],
+            },
+        )
+        assert downstream == []
+        assert index == {}
+        assert is_sparse is True
+
+    def test_flat_marker_with_nested_req_ids_is_not_erased(self):
+        downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+            engine_output_type="audio",
+            req_ids_output_copy=["r0", "r1"],
+            downstream_req_ids=["r0", "r1"],
+            multimodal_outputs={"meta": {"req_id": ["r1"]}, "meta.sparse_audio": ["1"]},
+        )
+        assert downstream == ["r1"]
+        assert index == {"r1": 0}
+        assert is_sparse is True
+
+    def test_sparse_marker_with_unusable_req_ids_fails_closed(self):
+        # The marker says "sparse" but the alignment list is unusable - the
+        # declaration as a whole is out of contract; same fail-closed shape.
+        with patch.object(sparse_audio, "_logged_offenders", set()):
+            downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+                engine_output_type="audio",
+                req_ids_output_copy=["r0"],
+                downstream_req_ids=["r0"],
+                multimodal_outputs={"meta": {"req_id": "r0", "sparse_audio": ["1"]}},
+            )
+        assert downstream == []
+        assert is_sparse is True
+
+    def test_nested_marker_without_req_id_fails_closed(self):
+        # {"meta": {"sparse_audio": ["1"]}} - marker set, req_id missing.
+        # The flattened-encoding fallback must not ERASE the nested marker
+        # (which used to silently demote the declaration to dense); this is
+        # an unusable-req_id invalid declaration -> fail closed.
+        with patch.object(sparse_audio, "_logged_offenders", set()):
+            downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+                engine_output_type="audio",
+                req_ids_output_copy=["r0"],
+                downstream_req_ids=["r0"],
+                multimodal_outputs={"meta": {"sparse_audio": ["1"]}},
+            )
+        assert downstream == []
+        assert is_sparse is True
+
+    def test_invalid_marker_on_non_audio_pipeline_keeps_dense_routing(self):
+        # Non-audio pipelines never consult sparse routing; the carrier
+        # error is logged by the classifier but routing is unaffected.
+        with patch.object(sparse_audio, "_logged_offenders", set()):
+            downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+                engine_output_type="latent",
+                req_ids_output_copy=["r0", "r1"],
+                downstream_req_ids=["r0", "r1"],
+                multimodal_outputs={"meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])}},
+            )
+        assert downstream == ["r0", "r1"]
+        assert is_sparse is False
+
+    def test_legal_marker_still_routes_sparse(self):
+        downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+            engine_output_type="audio",
+            req_ids_output_copy=["r0", "r1", "r2"],
+            downstream_req_ids=["r0", "r1", "r2"],
+            multimodal_outputs={"meta": {"req_id": ["r1", "r2"], "sparse_audio": ["1"]}},
+        )
+        assert downstream == ["r1", "r2"]
+        assert index == {"r1": 0, "r2": 1}
+        assert is_sparse is True
+
+    def test_combined_payload_unwraps_aligned_list(self):
+        combined = {"codes.audio": {"req-1": [torch.zeros(1), torch.ones(1)]}}
+        payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-1", list_idx=1)
+        assert torch.equal(payload["codes.audio"], torch.ones(1))
+
+    def test_combined_payload_broadcasts_qwen_request_invariant_singletons(self):
+        req_ids = ("r0", "r1")
+        tts_embeddings = {
+            "embed.tts_bos": torch.full((1, 1, 2), 1.0),
+            "embed.tts_eos": torch.full((1, 1, 2), 2.0),
+            "embed.tts_pad": torch.full((1, 1, 2), 3.0),
+        }
+        combined = {key: {rid: [value.clone()] for rid in req_ids} for key, value in tts_embeddings.items()}
+
+        for idx, rid in enumerate(req_ids):
+            payload = payload_build.build_omni_mm_payload(
+                combined_multimodal_outputs=combined,
+                mm_cpu=None,
+                rid=rid,
+                idx=idx,
+                start=idx,
+                end=idx + 1,
+                audio_sparse_output=False,
+                sparse_mm_index={},
+                hidden_seq_len=len(req_ids),
+                scheduled_seq_len=len(req_ids),
+            )
+            assert payload.keys() == tts_embeddings.keys()
+            for key, expected in tts_embeddings.items():
+                assert torch.equal(payload[key], expected)
+
+    def test_combined_payload_drops_sparse_protocol_metadata(self):
+        combined = {
+            "meta.sparse_audio": {"req-1": ["1"]},
+            "meta.req_id": {"req-1": ["req-1"]},
+            "sr": {"req-1": 22050},
+        }
+        payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-1", list_idx=2)
+        assert payload == {"sr": 22050}
+
+    def test_combined_payload_misalignment_drops_key_not_request_zero(self):
+        # A short multi-element list used to silently return `v[0]` - request
+        # 0's payload shipped under request `rid`. Never substitute another
+        # request's data; the key is dropped with an error log (same policy
+        # as the non-combined sparse mismatch) rather than raised - v1
+        # runners don't catch per-request exceptions, so a raise would turn a
+        # per-request protocol break into a whole-stage outage.
+        combined = {
+            "codes.audio": {"req-3": [torch.zeros(1), torch.ones(1)]},
+            "sr": {"req-3": 22050},
+        }
+        payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-3", list_idx=2)
+        assert "codes.audio" not in payload
+        assert payload["sr"] == 22050
+
+    def test_direct_path_sparse_mismatch_drops_key_and_keeps_payload_shape(self):
+        # Mirror of the combined-path policy on the non-prefix-cache path:
+        # a sparse index beyond the per-request list drops that key (error
+        # log) and the rest of the request's payload survives.
+        payload = payload_build.build_omni_mm_payload(
+            combined_multimodal_outputs=None,
+            mm_cpu={"codes.audio": [torch.zeros(1)]},
+            rid="r2",
+            idx=1,
+            start=0,
+            end=1,
+            audio_sparse_output=True,
+            sparse_mm_index={"r2": 1},
+            hidden_seq_len=1,
+            scheduled_seq_len=1,
+        )
+        assert payload == {}
+
+    def test_combined_payload_uses_sparse_index_under_sparse_routing(self):
+        # Batch [r0, r1, r2], sparse outputs only for [r1, r2]: producers
+        # align per-request lists to the SPARSE order, so r1 (batch idx 1,
+        # sparse idx 0) must get element 0 and r2 (batch idx 2, sparse idx 1)
+        # element 1. Indexing by batch idx shipped r2's payload to r1 and
+        # went out of range for r2. The sparse->list_idx resolution lives in
+        # `_build_omni_mm_payload` (the combined helper itself is
+        # sparse-agnostic), so this exercises the caller.
+        sparse_mm_index = {"r1": 0, "r2": 1}
+        combined = {
+            "codes.audio": {
+                "r1": [torch.zeros(1), torch.ones(1)],
+                "r2": [torch.zeros(1), torch.ones(1)],
+            }
+        }
+        payloads = {}
+        for rid, idx in (("r1", 1), ("r2", 2)):
+            payloads[rid] = payload_build.build_omni_mm_payload(
+                combined_multimodal_outputs=combined,
+                mm_cpu=None,
+                rid=rid,
+                idx=idx,
+                start=0,
+                end=1,
+                audio_sparse_output=True,
+                sparse_mm_index=sparse_mm_index,
+                hidden_seq_len=1,
+                scheduled_seq_len=1,
+            )
+        assert torch.equal(payloads["r1"]["codes.audio"], torch.zeros(1))
+        assert torch.equal(payloads["r2"]["codes.audio"], torch.ones(1))
+
+
+class TestMergeModelKvTransferMetadata:
+    """RFC #5450 C4: model KV-transfer metadata must merge into a COPY of the
+    scheduler's finished-request data - `scheduler_output` can be shared with
+    the engine-core process and must never be mutated by the runner."""
+
+    def _runner_with_model(self, model):
+        runner = object.__new__(GPUARModelRunner)
+        runner.model = model
+        return runner
+
+    def test_merge_does_not_mutate_scheduler_data(self):
+        class _Model:
+            def get_kv_transfer_metadata(self, req_id, num_computed_tokens):
+                return {"talker_codes": [7], "seen_tokens": num_computed_tokens}
+
+        original_meta = {"a": 1}
+        original = {"r1": {"seq_len": 3, "block_ids": [4], "custom_metadata": original_meta}}
+        runner = self._runner_with_model(_Model())
+
+        merged = runner._merge_model_kv_transfer_metadata(original)
+
+        # The manager sees the merged view...
+        assert merged["r1"]["custom_metadata"] == {"a": 1, "talker_codes": [7], "seen_tokens": 3}
+        assert merged["r1"]["seq_len"] == 3
+        # ...while the engine-shared originals are untouched.
+        assert original["r1"]["custom_metadata"] is original_meta
+        assert original_meta == {"a": 1}
+        assert "talker_codes" not in original["r1"]["custom_metadata"]
+
+    def test_merge_without_model_meta_passes_entry_through(self):
+        class _Model:
+            def get_kv_transfer_metadata(self, req_id, num_computed_tokens):
+                return None
+
+        data = {"seq_len": 2, "block_ids": [1]}
+        runner = self._runner_with_model(_Model())
+
+        merged = runner._merge_model_kv_transfer_metadata({"r1": data})
+
+        assert merged["r1"] is data
+
+    def test_merge_failure_raises_instead_of_shipping_partial_transfer(self):
+        # RFC #5450 C7: the bare `except Exception -> logger.warning` used to
+        # drop the model's custom metadata and let the KV transfer proceed
+        # incomplete; the failure then surfaced much later on the consumer
+        # stage. It must fail at the defect instead.
+        class _Model:
+            def get_kv_transfer_metadata(self, req_id, num_computed_tokens):
+                raise RuntimeError("boom")
+
+        data = {"seq_len": 2, "block_ids": [1], "custom_metadata": {"a": 1}}
+        runner = self._runner_with_model(_Model())
+
+        with pytest.raises(RuntimeError, match="boom"):
+            runner._merge_model_kv_transfer_metadata({"r1": data})
+        assert data["custom_metadata"] == {"a": 1}
+
+
+class TestDownstreamPayloadMemoization:
+    """RFC #5450 C7: `_request_needs_downstream_stage_payload` must not
+    memoize before the deciding marker exists - the final-stage id arrives via
+    `model_intermediate_buffer`, which can be unpopulated on the first call."""
+
+    def _runner(self, stages):
+        runner = object.__new__(GPUARModelRunner)
+        runner._downstream_payload_cache = {}
+        runner._request_final_stage_id = lambda rid: stages[rid]
+        return runner
+
+    def test_missing_marker_defaults_true_without_memoizing(self):
+        stages = {"r1": None}
+        runner = self._runner(stages)
+
+        assert runner._request_needs_downstream_stage_payload("r1") is True
+        # Not cached: the answer must refresh once the marker arrives.
+        assert "r1" not in runner._downstream_payload_cache
+
+        stages["r1"] = 0  # marker lands: this stage IS the final stage
+        assert runner._request_needs_downstream_stage_payload("r1") is False
+        assert runner._downstream_payload_cache["r1"] is False
+
+    def test_memoizes_once_marker_known(self):
+        stages = {"r1": 2}
+        runner = self._runner(stages)
+
+        assert runner._request_needs_downstream_stage_payload("r1") is True
+        assert runner._downstream_payload_cache["r1"] is True
+        # Cached value now authoritative for the request's lifetime.
+        stages["r1"] = 0
+        assert runner._request_needs_downstream_stage_payload("r1") is True
+
+
+# Matches a real ASSIGNMENT of the capability (optionally `self.`-qualified
+# and/or annotated), not a bare mention: a comment, a docstring, or a read
+# such as `getattr(self.model, "prefer_model_sampler", False)` must not count
+# as declaring it. `=(?!=)` keeps `==` comparisons out.
+_PREFER_MODEL_SAMPLER_ASSIGNMENT = re.compile(
+    r"^[ \t]*(?:self\.)?prefer_model_sampler[ \t]*(?::[^=\n]*)?=(?!=)[ \t]*(?P<rhs>[^\n]*)",
+    re.MULTILINE,
+)
+
+
+def _declares_prefer_model_sampler(source: str) -> bool:
+    """Whether `source` opts INTO the model-sampler contract.
+
+    An explicit `prefer_model_sampler = False` is an opt-OUT and does not
+    count; any other right-hand side does, including runtime-conditional
+    forms (minicpmo_4_5 assigns `self.model_stage in {...}`), which can
+    evaluate True.
+    """
+    for match in _PREFER_MODEL_SAMPLER_ASSIGNMENT.finditer(source):
+        if match.group("rhs").split("#")[0].strip() != "False":
+            return True
+    return False
+
+
+class TestPreferModelSamplerNoneFallback:
+    """RFC #5450 C7: a `prefer_model_sampler` model returning None from
+    `sample()` DECLINES the step and the runner falls back to the default
+    sampler. The fallback is load-bearing (cosyvoice3/glm_tts decline for
+    empty-logits / inapplicable-stage steps), so it must keep working - now
+    as a documented contract with a one-time visibility log."""
+
+    def _runner(self, model_sample_result, default_result):
+        runner = object.__new__(GPUARModelRunner)
+        runner.input_batch = SimpleNamespace(
+            sampling_metadata="smd",
+            update_async_output_token_ids=lambda: None,
+        )
+        runner.model = SimpleNamespace(
+            prefer_model_sampler=True,
+            sample=lambda logits, metadata: model_sample_result,
+        )
+        runner._sampling_metadata_for_model_sampler = lambda smd: smd
+        runner._resolve_duplex_sampling_hook = lambda: None
+        runner.sampler = lambda logits, sampling_metadata: default_result
+        return runner
+
+    def test_none_falls_back_to_default_sampler_and_warns_once(self):
+        import logging
+
+        from vllm_omni.worker import gpu_ar_model_runner as ar_module
+
+        class _UniqueDecliningSampler:
+            prefer_model_sampler = True
+
+            @staticmethod
+            def sample(logits, metadata):
+                return None
+
+        sentinel = object()
+        runner = self._runner(model_sample_result=None, default_result=sentinel)
+        runner.model = _UniqueDecliningSampler()
+
+        records: list[logging.LogRecord] = []
+
+        class _Capture(logging.Handler):
+            def emit(self, record):
+                records.append(record)
+
+        handler = _Capture(level=logging.WARNING)
+        ar_module.logger.addHandler(handler)
+        try:
+            for _ in range(3):
+                out = GPUARModelRunner._sample(runner, torch.zeros(1, 4), None)
+                assert out is sentinel
+        finally:
+            ar_module.logger.removeHandler(handler)
+
+        fallback_records = [r for r in records if "falling back to the default sampler" in r.getMessage()]
+        assert len(fallback_records) == 1
+        assert "_UniqueDecliningSampler" in fallback_records[0].getMessage()
+
+    def test_model_sampler_output_short_circuits_default(self):
+        model_out = object()
+        runner = self._runner(model_sample_result=model_out, default_result=object())
+
+        out = GPUARModelRunner._sample(runner, torch.zeros(1, 4), None)
+
+        assert out is model_out
+
+    def test_declaration_matcher_ignores_mentions_and_opt_outs(self):
+        # Guards the guard: the inventory below is only meaningful if
+        # "declares" means an actual opt-in assignment.
+        assert _declares_prefer_model_sampler("    prefer_model_sampler = True")
+        assert _declares_prefer_model_sampler("    prefer_model_sampler: bool = True")
+        assert _declares_prefer_model_sampler('        self.prefer_model_sampler = stage in {"llm"}')
+        assert not _declares_prefer_model_sampler("    prefer_model_sampler = False")
+        assert not _declares_prefer_model_sampler("    prefer_model_sampler = False  # opted out")
+        assert not _declares_prefer_model_sampler("    # prefer_model_sampler = True (not anymore)")
+        assert not _declares_prefer_model_sampler('    """Does not use prefer_model_sampler."""')
+        assert not _declares_prefer_model_sampler('    x = getattr(model, "prefer_model_sampler", False)')
+        assert not _declares_prefer_model_sampler("    if prefer_model_sampler == True:")
+
+    def test_in_tree_declarer_inventory_is_exact(self):
+        # The contract's audience: an exact inventory, so ADDING a declarer
+        # fails this test and the new model owner consciously signs up for the
+        # documented None-fallback semantics (update the set when they do).
+        import pathlib
+
+        import vllm_omni
+
+        models_dir = pathlib.Path(vllm_omni.__file__).parent / "model_executor" / "models"
+        declarers = {
+            p.parent.name
+            for p in models_dir.rglob("*.py")
+            if _declares_prefer_model_sampler(p.read_text(encoding="utf-8", errors="ignore"))
+        }
+        expected = {
+            "cosyvoice3",
+            "glm_tts",
+            "higgs_audio_v2",
+            "higgs_audio_v3",
+            "hunyuan_image3",
+            "minicpmo_4_5",
+            "minimax_music3",
+        }
+        assert declarers == expected, (
+            "The set of models declaring `prefer_model_sampler` changed:\n"
+            f"  added:   {sorted(declarers - expected) or 'none'}\n"
+            f"  removed: {sorted(expected - declarers) or 'none'}\n\n"
+            "This is a deliberate consent gate, not a broken test. Declaring "
+            "`prefer_model_sampler` opts a model into the runner's documented "
+            "None-fallback contract: `model.sample()` may return None to decline "
+            "a step, and the runner then falls back to the DEFAULT sampler "
+            "(see `GPUARModelRunner._sample`; RFC #5450 C7.3).\n\n"
+            "If you ADDED a declarer: confirm your model relies on -- or at "
+            "least tolerates -- that fallback, then add its directory name to "
+            "`expected` above. If you REMOVED one, drop its name."
+        )

tests/worker/test_gpu_generation_model_runner.py

@@ -16,14 +16,14 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _DummyInputBatch:
-    def __init__(self):
-        self.req_ids = ["req-1"]
-        self.req_id_to_index = {"req-1": 0}
-        self.num_reqs = 1
+    def __init__(self, num_reqs: int = 1):
+        self.req_ids = [f"req-{i + 1}" for i in range(num_reqs)]
+        self.req_id_to_index = {rid: i for i, rid in enumerate(self.req_ids)}
+        self.num_reqs = num_reqs
         self.vocab_size = 10
 
 
-def _make_runner(multimodal_outputs):
+def _make_runner(multimodal_outputs, num_reqs: int = 1):
     runner = object.__new__(GPUGenerationModelRunner)
     runner.execute_model_state = ExecuteModelState(
         None,
@@ -40,7 +40,7 @@ def _make_runner(multimodal_outputs):
         None,
     )
     runner.kv_connector_output = None
-    runner.input_batch = _DummyInputBatch()
+    runner.input_batch = _DummyInputBatch(num_reqs)
     runner.use_async_scheduling = False
     runner.device = torch.device("cpu")
     runner.supports_mm_inputs = False
@@ -111,6 +111,12 @@ def _make_guard_runner():
     runner.cache_config = CacheConfig()
     runner.input_batch = _DummyInputBatch()
     runner.model = object()
+    runner.parallel_config = SimpleNamespace(
+        distributed_executor_backend=None,
+        data_parallel_size=1,
+    )
+    runner._dummy_run_calls = []
+    runner._dummy_run = lambda n: runner._dummy_run_calls.append(n)
     runner._update_states = lambda scheduler_output: None
     runner.synchronize_input_prep = contextlib.nullcontext
     runner.attach_omni_connector_output = lambda result: result
@@ -121,9 +127,112 @@ def _make_guard_runner():
 def test_execute_model_skips_non_positive_scheduled_span(monkeypatch, total):
     """#5196: a negative span is truthy, so it used to reach `_prepare_inputs`,
     whose `assert total_num_scheduled_tokens > 0` killed the stage EngineCore."""
+    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()
 
     output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(total))
 
     assert output is EMPTY_MODEL_RUNNER_OUTPUT
+    assert runner._dummy_run_calls == []
+
+
+@pytest.mark.parametrize("total", [-1, 0])
+def test_execute_model_zero_tokens_runs_dp_dummy_run(monkeypatch, total):
+    """A 0-token step under external_launcher + DP>1 must `_dummy_run(1)` so
+    `coordinate_batch_across_dp` stays in sync (it used to be shadowed by an
+    unconditional early return and never ran -> DP hang risk)."""
+    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.parallel_config = SimpleNamespace(
+        distributed_executor_backend="external_launcher",
+        data_parallel_size=2,
+    )
+
+    output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(total))
+
+    assert output is EMPTY_MODEL_RUNNER_OUTPUT
+    assert runner._dummy_run_calls == [1]
+
+
+def test_execute_model_zero_tokens_kv_connector_no_forward(monkeypatch):
+    """With a KV-transfer group, a 0-token step must go through
+    `kv_connector_no_forward` (also shadowed by the old early return)."""
+    monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False)
+    monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: True)
+    runner = _make_guard_runner()
+    sentinel = object()
+    calls = []
+    runner.kv_connector_no_forward = lambda scheduler_output, vllm_config: (
+        calls.append(scheduler_output),
+        sentinel,
+    )[1]
+    runner.vllm_config = object()
+
+    output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(0))
+
+    assert output is sentinel
+    assert len(calls) == 1
+
+
+def test_execute_model_zero_tokens_ec_producer_takes_encoder_path(monkeypatch):
+    """AR lock-step: the EC-producer block sits ABOVE the zero-token block
+    (matching GPUARModelRunner), so a 0-token step on an EC producer returns
+    the empty ENCODER output. Before the C1 reorder, the unconditional early
+    return won and produced EMPTY_MODEL_RUNNER_OUTPUT instead — this is the
+    one cell whose behaviour the reorder changed."""
+    monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: True)
+    monkeypatch.setattr(gen_runner_module, "get_ec_transfer", lambda: SimpleNamespace(is_consumer=False))
+    monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False)
+    sentinel = object()
+    monkeypatch.setattr(gen_runner_module, "make_empty_encoder_model_runner_output", lambda so: sentinel)
+    runner = _make_guard_runner()
+    runner.encoder_cache = {}
+    runner.maybe_get_ec_connector_output = lambda scheduler_output, encoder_cache: contextlib.nullcontext()
+    encoder_calls = []
+    runner._execute_mm_encoder = lambda so: encoder_calls.append(so)
+
+    output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(0))
+
+    assert output is sentinel
+    assert len(encoder_calls) == 1
+    assert runner._dummy_run_calls == []
+
+
+class TestSampleTokensBatched:
+    """RFC #5450 C5: `sample_tokens` must slice per request for num_reqs > 1.
+    The tensor branch's asserts used to force num_reqs == 1 and the list
+    branch built a length-1 payload misaligned with `req_ids`."""
+
+    def test_tensor_output_sliced_per_request(self):
+        multimodal_outputs = torch.arange(12, dtype=torch.float32).reshape(2, 2, 3)
+        runner = _make_runner(multimodal_outputs, num_reqs=2)
+
+        output = GPUGenerationModelRunner.sample_tokens(runner)
+
+        assert len(output.multimodal_outputs) == 2
+        assert torch.equal(output.multimodal_outputs[0]["model_outputs"], multimodal_outputs[0])
+        assert torch.equal(output.multimodal_outputs[1]["model_outputs"], multimodal_outputs[1])
+
+    def test_list_output_one_entry_per_request(self):
+        entries = [torch.zeros(2, 1), torch.ones(2, 1)]
+        runner = _make_runner(list(entries), num_reqs=2)
+
+        output = GPUGenerationModelRunner.sample_tokens(runner)
+
+        assert len(output.multimodal_outputs) == 2
+        assert torch.equal(output.multimodal_outputs[0]["model_outputs"], entries[0])
+        assert torch.equal(output.multimodal_outputs[1]["model_outputs"], entries[1])
+
+    def test_tensor_output_batch_mismatch_raises(self):
+        runner = _make_runner(torch.randn(1, 2, 3), num_reqs=2)
+
+        with pytest.raises(ValueError, match="leading dim 1 .* 2 requests"):
+            GPUGenerationModelRunner.sample_tokens(runner)
+
+    def test_list_output_batch_mismatch_raises(self):
+        runner = _make_runner([torch.randn(2, 1)], num_reqs=2)
+
+        with pytest.raises(ValueError, match="length 1 .* 2 requests"):
+            GPUGenerationModelRunner.sample_tokens(runner)

vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py

@@ -3013,9 +3013,18 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                     merged_cpu = merged.detach().cpu().float()
                     mm["model_outputs"] = list(merged_cpu.split(sizes))
                     mm["sr"] = [sr for _ in ready_req_ids]
+                    # Dense mode still yields a strict SUBSET of the batch when
+                    # some requests emit no audio this step (e.g. prefill
+                    # phase). Without the marker the runner indexes these
+                    # per-request lists by batch position and misroutes audio
+                    # across requests; the marker declares meta.req_id
+                    # alignment.
+                    mm["meta"] = {"req_id": ready_req_ids, "sparse_audio": ["1"]}
                 else:
                     mm["model_outputs"] = list(audio_by_req.values())
                     mm["sr"] = [sr for _ in audio_by_req]
+                    # Same subset hazard as the coalesce branch above.
+                    mm["meta"] = {"req_id": list(audio_by_req), "sparse_audio": ["1"]}
             elif self._uses_sparse_audio_outputs():
                 mm["model_outputs"] = []
                 mm["sr"] = []

vllm_omni/worker/gpu_ar_model_runner.py

@@ -51,12 +51,13 @@ from vllm_omni.utils.mm_outputs import (
     build_mm_cpu,
     partition_payload_list,
     snapshot_mm_payload,
-    to_payload_element,
 )
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+from vllm_omni.worker.output.payload_build import build_omni_mm_payload
 from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
-from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
+from vllm_omni.worker.sampling_utils import clamp_prompt_ids_to_penalty_padding, sanitize_min_tokens_stop_ids
+from vllm_omni.worker.sparse_audio import resolve_sparse_mm_routing
 
 logger = init_logger(__name__)
 
@@ -457,9 +458,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         cached = self._downstream_payload_cache.get(req_id)
         if cached is not None:
             return cached
-        # Conservative default: keep payload if marker is missing.
         final_stage_id = self._request_final_stage_id(req_id)
-        needs_payload = final_stage_id is None or final_stage_id > 0
+        if final_stage_id is None:
+            # Conservative default while the marker is missing: keep the
+            # payload, but do NOT memoize - the marker arrives via
+            # `model_intermediate_buffer`, which may be unpopulated on the
+            # first call (memoizing here pinned the request to True forever,
+            # never refreshing once the marker landed).
+            return True
+        needs_payload = final_stage_id > 0
         self._downstream_payload_cache[req_id] = needs_payload
         return needs_payload
 
@@ -472,50 +479,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
             downstream_req_ids = req_ids_output_copy
         return engine_output_type, downstream_req_ids
 
-    @staticmethod
-    def _sparse_mm_req_ids(multimodal_outputs: Any) -> list[str] | None:
-        if not isinstance(multimodal_outputs, dict):
-            return None
-        meta = multimodal_outputs.get("meta")
-        req_ids = None
-        sparse_audio = False
-        if isinstance(meta, dict):
-            req_ids = meta.get("req_id")
-            sparse_audio = GPUARModelRunner._is_sparse_audio_marker(meta.get("sparse_audio"))
-        if req_ids is None:
-            req_ids = multimodal_outputs.get("meta.req_id")
-            sparse_audio = GPUARModelRunner._is_sparse_audio_marker(multimodal_outputs.get("meta.sparse_audio"))
-        if not sparse_audio:
-            return None
-        if not isinstance(req_ids, list):
-            return None
-        return [rid for rid in req_ids if isinstance(rid, str)]
-
-    @staticmethod
-    def _resolve_sparse_mm_routing(
-        *,
-        engine_output_type: str,
-        req_ids_output_copy: list[str],
-        downstream_req_ids: list[str],
-        multimodal_outputs: Any,
-    ) -> tuple[list[str], dict[str, int], bool]:
-        sparse_mm_req_ids = GPUARModelRunner._sparse_mm_req_ids(multimodal_outputs)
-        sparse_mm_index = {rid: i for i, rid in enumerate(sparse_mm_req_ids or [])}
-        if engine_output_type != "audio" or sparse_mm_req_ids is None:
-            return downstream_req_ids, sparse_mm_index, False
-
-        sparse_req_id_set = set(sparse_mm_req_ids)
-        sparse_downstream_req_ids = [rid for rid in req_ids_output_copy if rid in sparse_req_id_set]
-        return sparse_downstream_req_ids, sparse_mm_index, True
-
-    @staticmethod
-    def _is_sparse_audio_marker(value: Any) -> bool:
-        if isinstance(value, list):
-            return any(str(item).lower() in ("1", "true", "yes", "on") for item in value)
-        if isinstance(value, str):
-            return value.lower() in ("1", "true", "yes", "on")
-        return bool(value)
-
     def capture_model(self) -> int:
         result = super().capture_model()
         self._capture_talker_mtp_graphs()
@@ -862,79 +825,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         )
         return hidden_states_cpu, combined_hidden_states, combined_multimodal_outputs
 
-    @staticmethod
-    def _build_combined_prefix_cache_mm_payload(
-        combined_multimodal_outputs: dict,
-        *,
-        rid: str,
-        idx: int,
-    ) -> dict[str, object]:
-        def _unwrap_lists(v):
-            if isinstance(v, list):
-                return v[idx] if idx < len(v) else v[0]
-            if isinstance(v, dict):
-                return {k: _unwrap_lists(sv) for k, sv in v.items()}
-            return v
-
-        return {
-            mm_key: _unwrap_lists(combined_multimodal_outputs[mm_key][rid])
-            for mm_key in combined_multimodal_outputs.keys()
-        }
-
-    def _build_omni_mm_payload(
-        self,
-        *,
-        combined_multimodal_outputs: dict | None,
-        mm_cpu: dict[str, object] | None,
-        rid: str,
-        idx: int,
-        start: int,
-        end: int,
-        audio_sparse_output: bool,
-        sparse_mm_index: dict[str, int],
-        hidden_seq_len: int,
-        scheduled_seq_len: int,
-    ) -> dict[str, object]:
-        if combined_multimodal_outputs:
-            return self._build_combined_prefix_cache_mm_payload(
-                combined_multimodal_outputs,
-                rid=rid,
-                idx=idx,
-            )
-
-        mm_payload: dict[str, object] = {}
-        if not mm_cpu:
-            return mm_payload
-
-        for mm_key, mm_val in mm_cpu.items():
-            if mm_key in {"meta.req_id", "meta.sparse_audio"}:
-                continue
-            if audio_sparse_output and isinstance(mm_val, list):
-                sparse_idx = sparse_mm_index.get(rid)
-                if sparse_idx is None:
-                    continue
-                if sparse_idx >= len(mm_val):
-                    logger.warning(
-                        "Sparse multimodal payload mismatch for request %s: index %d >= %d.",
-                        rid,
-                        sparse_idx,
-                        len(mm_val),
-                    )
-                    continue
-                sparse_val = mm_val[sparse_idx]
-                mm_payload[mm_key] = sparse_val.clone() if isinstance(sparse_val, torch.Tensor) else sparse_val
-                continue
-            mm_payload[mm_key] = to_payload_element(
-                element=mm_val,
-                idx=idx,
-                start=start,
-                end=end,
-                pass_lists_through=False,
-                seq_len=hidden_seq_len,
-                scheduled_seq_len=scheduled_seq_len,
-            )
-        return mm_payload
-
     def _build_omni_pooler_payload(
         self,
         *,
@@ -967,7 +857,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
             if req_hidden_states is not None:
                 payload["hidden"] = req_hidden_states
 
-        mm_payload = self._build_omni_mm_payload(
+        mm_payload = build_omni_mm_payload(
             combined_multimodal_outputs=combined_multimodal_outputs,
             mm_cpu=mm_cpu,
             rid=rid,
@@ -982,6 +872,47 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         payload.update(mm_payload)
         return payload
 
+    def _merge_model_kv_transfer_metadata(
+        self,
+        finished_reqs: dict[str, dict[str, Any]],
+    ) -> dict[str, dict[str, Any]]:
+        """Merge model-provided KV-transfer metadata copy-on-write.
+
+        The outer mapping is always new; a per-request dict is shallow-copied
+        (with a fresh merged ``custom_metadata``) only when the model supplies
+        metadata for it, and passes through unchanged otherwise.
+        `scheduler_output` (and its per-request dicts) can be shared with the
+        engine-core process, so nothing here mutates them — the ngram block
+        in `execute_model` `replace()`-copies for exactly the same reason. The
+        merged mapping is only consumed by
+        `kv_transfer_manager.handle_finished_requests_kv_transfer`.
+        """
+        merged: dict[str, dict[str, Any]] = {}
+        for req_id, data in finished_reqs.items():
+            try:
+                # NOTE: seq_len is the same as num_computed_tokens_cpu in current
+                # async scheduling, since both exclude async placeholders. We use
+                # seq_len since we control it, just in case upstream async scheduler
+                # semantics change in the future.
+                model_meta = self.model.get_kv_transfer_metadata(
+                    req_id,
+                    num_computed_tokens=data.get("seq_len"),
+                )
+            except Exception:
+                # A swallowed failure here used to ship the KV transfer WITHOUT
+                # the model's custom metadata and surface much later as a
+                # corrupt/incomplete transfer on the consumer stage. Fail at
+                # the point of the defect instead.
+                logger.exception("Failed to get custom KV-transfer metadata from model for %s", req_id)
+                raise
+            if model_meta:
+                data = {
+                    **data,
+                    "custom_metadata": {**(data.get("custom_metadata") or {}), **model_meta},
+                }
+            merged[req_id] = data
+        return merged
+
     @torch.inference_mode()
     def execute_model(
         self,
@@ -1012,24 +943,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
         finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
         if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"):
-            for req_id, data in finished_reqs.items():
-                try:
-                    # NOTE: seq_len is the same as num_computed_tokens_cpu in current
-                    # async scheduling, since both exclude async placeholders. We use
-                    # seq_len since we control it, just in case upstream async scheduler
-                    # semantics change in the future.
-                    num_computed = data.get("seq_len")
-
-                    model_meta = self.model.get_kv_transfer_metadata(
-                        req_id,
-                        num_computed_tokens=num_computed,
-                    )
-                    if model_meta:
-                        existing = data.get("custom_metadata") or {}
-                        existing.update(model_meta)
-                        data["custom_metadata"] = existing
-                except Exception as e:
-                    logger.warning(f"Failed to get custom metadata from model for {req_id}: {e}")
+            finished_reqs = self._merge_model_kv_transfer_metadata(finished_reqs)
         self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
             finished_reqs=finished_reqs,
             kv_caches=self.kv_caches,
@@ -1467,6 +1381,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         logits: torch.Tensor | None,
         spec_decode_metadata: Any,
     ):
+        """Sample tokens, preferring the model's own sampler when declared.
+
+        Model-interface contract for `prefer_model_sampler` models: a custom
+        `model.sample(logits, sampling_metadata)` may return ``None`` to
+        DECLINE the step, and the runner then falls back to the default
+        sampler. Declarers rely on this deliberately for empty-logits or
+        inapplicable-stage steps (e.g. cosyvoice3, glm_tts), so the
+        fallthrough is load-bearing, not an error path; it is logged once per
+        model for visibility.
+        """
         sampling_metadata = self.input_batch.sampling_metadata
         if spec_decode_metadata is None:
             model_sample = getattr(self.model, "sample", None)
@@ -1487,6 +1411,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                 sampler_output = model_sample(logits, prepared_sampling_metadata)
                 if sampler_output is not None:
                     return sampler_output
+                # Contract: None => fall back to the default sampler (see
+                # docstring). If a custom sampler returns None by accident,
+                # its tokens silently come from the default sampler - the
+                # one-time log is the visibility hook for that case.
+                logger.warning_once(
+                    "prefer_model_sampler model %s returned None from sample(); falling back to the default sampler.",
+                    type(self.model).__name__,
+                )
             return self.sampler(
                 logits=logits,
                 sampling_metadata=sampling_metadata,
@@ -1889,7 +1821,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         combined_multimodal_outputs = None
 
         engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy)
-        downstream_req_ids, sparse_mm_index, audio_sparse_output = self._resolve_sparse_mm_routing(
+        downstream_req_ids, sparse_mm_index, audio_sparse_output = resolve_sparse_mm_routing(
             engine_output_type=engine_output_type,
             req_ids_output_copy=req_ids_output_copy,
             downstream_req_ids=downstream_req_ids,
@@ -2089,13 +2021,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         if grammar_output is not None:
             apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits)
 
-        # Correct padding values of prompt_token_ids to match the logits vocabulary size
+        # Correct padding values of prompt_token_ids to match the logits vocabulary size.
         if logits is not None and not self.input_batch.sampling_metadata.no_penalties:
             smd = self.input_batch.sampling_metadata
             if smd.prompt_token_ids is not None:
                 logits_vocab = logits.shape[-1]
                 if self.input_batch.vocab_size > logits_vocab:
-                    smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab)
+                    smd.prompt_token_ids = clamp_prompt_ids_to_penalty_padding(smd.prompt_token_ids, logits_vocab)
 
         # Drop min-tokens stop ids the head cannot emit (e.g. the text
         # tokenizer EOS folded into all_stop_token_ids on a narrow codec

vllm_omni/worker/gpu_generation_model_runner.py

@@ -167,10 +167,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
                 self.model.on_requests_finished(scheduler_output.finished_req_ids)
 
-            # `<= 0`: upstream can schedule a negative span, which is truthy (#5196).
-            if scheduler_output.total_num_scheduled_tokens <= 0:
-                return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
-
             if has_ec_transfer() and not get_ec_transfer().is_consumer:
                 with self.maybe_get_ec_connector_output(
                     scheduler_output,
@@ -179,7 +175,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                     self._execute_mm_encoder(scheduler_output)
                     return self.attach_omni_connector_output(make_empty_encoder_model_runner_output(scheduler_output))
 
-            if not num_scheduled_tokens:
+            # OMNI: keep this block in lock-step with the same block in
+            # GPUARModelRunner.execute_model.
+            # `<= 0`: upstream can schedule a negative span, which is truthy (#5196).
+            if num_scheduled_tokens <= 0:
                 if (
                     self.parallel_config.distributed_executor_backend == "external_launcher"
                     and self.parallel_config.data_parallel_size > 1
@@ -439,15 +438,28 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         # pooler_output is no longer used for multimodal data.
         per_req_payloads: list[dict[str, object]] = []
         if isinstance(multimodal_outputs_raw, torch.Tensor):
-            assert multimodal_outputs_raw.shape[0] == 1, (
-                "model should return a single tensor, to return multiple tensors, use a dict"
+            # One row per request. The old asserts (`shape[0] == 1` AND
+            # `shape[0] == num_reqs`) jointly forced num_reqs == 1, silently
+            # rejecting batched steps the generation scheduler does admit.
+            num_reqs = self.input_batch.num_reqs
+            if multimodal_outputs_raw.shape[0] != num_reqs:
+                raise ValueError(
+                    f"Multimodal output tensor has leading dim {multimodal_outputs_raw.shape[0]} "
+                    f"but the batch has {num_reqs} requests (one row per request; "
+                    "to return multiple tensors per request, use a dict)."
                 )
-            assert multimodal_outputs_raw.shape[0] == self.input_batch.num_reqs
-            for i in range(self.input_batch.num_reqs):
+            for i in range(num_reqs):
                 per_req_payloads.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()})
         elif isinstance(multimodal_outputs_raw, list):
-            assert len(multimodal_outputs_raw) == 1, (
-                "model should return a single list, to return multiple lists, use a dict"
+            # One entry per request. The old `len == 1` assert did not check
+            # num_reqs, so a batched step built a length-1 payload list that
+            # misaligned with `req_ids` downstream.
+            num_reqs = self.input_batch.num_reqs
+            if len(multimodal_outputs_raw) != num_reqs:
+                raise ValueError(
+                    f"Multimodal output list has length {len(multimodal_outputs_raw)} "
+                    f"but the batch has {num_reqs} requests (one entry per request; "
+                    "to return multiple lists per request, use a dict)."
                 )
             for out in multimodal_outputs_raw:
                 per_req_payloads.append(
@@ -794,6 +806,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                 input_ids = None
                 inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
                 model_kwargs = self._init_model_kwargs()
+            elif getattr(getattr(self, "model", None), "has_preprocess", False):
+                # Capture CUDA graph with inputs_embeds path so replay reads
+                # from the same buffer that _preprocess writes into.
+                input_ids = self.input_ids.gpu[:num_tokens_padded]
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
             else:
                 input_ids = self.input_ids.gpu[:num_tokens_padded]
                 inputs_embeds = None

vllm_omni/worker/output/__init__.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Worker output helpers.
+
+Package seeded ahead of the async-output refactor (RFC #5450 series, G5-C2),
+which relocates the remaining payload-copy helpers, snapshot types, and
+`ExecuteModelState` here. Currently hosts the per-request multimodal payload
+builders.
+"""
+
+from vllm_omni.worker.output.payload_build import (
+    build_combined_prefix_cache_mm_payload,
+    build_omni_mm_payload,
+    unwrap_combined_payload_value,
+)
+
+__all__ = [
+    "build_combined_prefix_cache_mm_payload",
+    "build_omni_mm_payload",
+    "unwrap_combined_payload_value",
+]

vllm_omni/worker/output/payload_build.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-request multimodal payload builders.
+
+State-free extraction helpers that turn a step's `multimodal_outputs`
+(combined prefix-cache-merged form, or the direct `mm_cpu` form) into one
+request's payload dict. Sparse-routing context (`audio_sparse_output`,
+`sparse_mm_index`) is resolved upstream by
+`vllm_omni.worker.sparse_audio.resolve_sparse_mm_routing` and arrives here
+as plain arguments.
+"""
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.utils.mm_outputs import to_payload_element
+
+logger = init_logger(__name__)
+
+_SPARSE_AUDIO_PROTOCOL_KEYS = frozenset({"meta.req_id", "meta.sparse_audio"})
+
+
+def unwrap_combined_payload_value(
+    value: Any,
+    *,
+    rid: str,
+    list_idx: int,
+    mm_key: str,
+) -> tuple[bool, Any]:
+    """Unwrap one value from the prefix-cache-merged payload.
+
+    Returns ``(keep, unwrapped)``: ``keep`` is False when a non-singleton
+    per-request list is misaligned and must be dropped. A singleton list is
+    request-invariant passthrough data and is shared across the batch.
+    """
+    if isinstance(value, list):
+        if len(value) == 1:
+            # Prefix-cache passthrough intentionally preserves list
+            # containers and copies the whole list under every request id.
+            # Qwen3-Omni uses this shape for its request-invariant TTS
+            # BOS/EOS/PAD embeddings, so every batch index unwraps the same
+            # singleton. Sparse one-request lists are safe here too: routing
+            # limits them to the declared request at sparse index zero.
+            return True, value[0]
+        if list_idx < len(value):
+            return True, value[list_idx]
+        # A non-singleton list shorter than the request's index is per-request
+        # data that lost alignment. Never substitute another request's
+        # element (`v[0]` used to ship request 0's payload as `rid`'s); drop
+        # the key loudly instead — the SAME protocol-break policy as the
+        # sparse mismatch in
+        # `build_omni_mm_payload`. Not a raise: v1 runners do not catch
+        # per-request exceptions, so raising would turn a per-request
+        # protocol break into a whole-stage outage.
+        logger.error(
+            "Combined prefix-cache mm payload misaligned for request %s: index %d out of "
+            "range for per-request list of length %d; dropping key %s.",
+            rid,
+            list_idx,
+            len(value),
+            mm_key,
+        )
+        return False, None
+    if isinstance(value, dict):
+        nested = {}
+        for k, sv in value.items():
+            keep, unwrapped = unwrap_combined_payload_value(sv, rid=rid, list_idx=list_idx, mm_key=mm_key)
+            if keep:
+                nested[k] = unwrapped
+        return True, nested
+    return True, value
+
+
+def build_combined_prefix_cache_mm_payload(
+    combined_multimodal_outputs: dict,
+    *,
+    rid: str,
+    list_idx: int,
+) -> dict[str, object]:
+    """Unwrap one request's entry from the prefix-cache-merged payload.
+
+    Sparse-agnostic by design: ``list_idx`` arrives pre-resolved by the
+    caller (`build_omni_mm_payload`, which owns the sparse-routing context)
+    — batch index normally, sparse index under sparse routing.
+    """
+    payload: dict[str, object] = {}
+    for mm_key in combined_multimodal_outputs.keys():
+        if mm_key in _SPARSE_AUDIO_PROTOCOL_KEYS:
+            continue
+        keep, unwrapped = unwrap_combined_payload_value(
+            combined_multimodal_outputs[mm_key][rid],
+            rid=rid,
+            list_idx=list_idx,
+            mm_key=mm_key,
+        )
+        if keep:
+            payload[mm_key] = unwrapped
+    return payload
+
+
+def build_omni_mm_payload(
+    *,
+    combined_multimodal_outputs: dict | None,
+    mm_cpu: dict[str, object] | None,
+    rid: str,
+    idx: int,
+    start: int,
+    end: int,
+    audio_sparse_output: bool,
+    sparse_mm_index: dict[str, int],
+    hidden_seq_len: int,
+    scheduled_seq_len: int,
+) -> dict[str, object]:
+    """Build one request's multimodal payload for the step."""
+    if combined_multimodal_outputs:
+        # Sparse-audio producers emit per-request lists aligned to the
+        # sparse request order (`meta.req_id`), not the batch order; under
+        # sparse routing the request's sparse index selects its element.
+        list_idx = sparse_mm_index.get(rid, idx) if audio_sparse_output else idx
+        return build_combined_prefix_cache_mm_payload(
+            combined_multimodal_outputs,
+            rid=rid,
+            list_idx=list_idx,
+        )
+
+    mm_payload: dict[str, object] = {}
+    if not mm_cpu:
+        return mm_payload
+
+    for mm_key, mm_val in mm_cpu.items():
+        if mm_key in _SPARSE_AUDIO_PROTOCOL_KEYS:
+            continue
+        if audio_sparse_output and isinstance(mm_val, list):
+            sparse_idx = sparse_mm_index.get(rid)
+            if sparse_idx is None:
+                continue
+            if sparse_idx >= len(mm_val):
+                # Misalignment means the sparse-marker protocol broke; the
+                # request's output for this key is dropped, so be loud.
+                logger.error(
+                    "Sparse multimodal payload mismatch for request %s: index %d >= %d; dropping key %s.",
+                    rid,
+                    sparse_idx,
+                    len(mm_val),
+                    mm_key,
+                )
+                continue
+            sparse_val = mm_val[sparse_idx]
+            mm_payload[mm_key] = sparse_val.clone() if isinstance(sparse_val, torch.Tensor) else sparse_val
+            continue
+        mm_payload[mm_key] = to_payload_element(
+            element=mm_val,
+            idx=idx,
+            start=start,
+            end=end,
+            pass_lists_through=False,
+            seq_len=hidden_seq_len,
+            scheduled_seq_len=scheduled_seq_len,
+        )
+    return mm_payload

vllm_omni/worker/sampling_utils.py

@@ -8,7 +8,21 @@ from vllm.v1.sample.logits_processor import LogitsProcessors, MinTokensLogitsPro
 
 logger = init_logger(__name__)
 
-__all__ = ["sanitize_min_tokens_stop_ids"]
+__all__ = ["clamp_prompt_ids_to_penalty_padding", "sanitize_min_tokens_stop_ids"]
+
+
+def clamp_prompt_ids_to_penalty_padding(prompt_token_ids: torch.Tensor, logits_vocab: int) -> torch.Tensor:
+    """Clamp batch-level pad ids down to ``logits_vocab`` — upstream's
+    designed penalty padding value.
+
+    ``max=logits_vocab`` (NOT ``logits_vocab - 1``) is deliberate: upstream
+    penalty computation allocates ``vocab_size + 1`` bins and drops the last
+    column, so ``vocab_size`` is the padding value that never affects
+    penalties (vllm/model_executor/layers/utils.py::
+    get_token_bin_counts_and_mask). Clamping one lower would count padding
+    as real occurrences of the last vocab token.
+    """
+    return prompt_token_ids.clamp(max=logits_vocab)
 
 
 def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: int) -> None:

vllm_omni/worker/sparse_audio.py

@@ -0,0 +1,232 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Sparse-audio output protocol: marker classification and payload routing.
+
+A step's per-request `multimodal_outputs` lists may cover only a SUBSET of
+the batch (requests that produced audio this step). The producer declares
+this with two metadata keys (nested ``meta`` dict or flattened ``meta.*``):
+
+* ``meta.req_id`` — the request ids that DO have output, in list order;
+* ``meta.sparse_audio`` — the marker saying "per-request lists are aligned
+  to ``meta.req_id`` order, not batch order".
+
+The marker's carrier form is ENFORCED, not normalized. The design contract
+is a list of strings (e.g. ``["1"]`` — what every in-tree producer emits); a
+bare string is also accepted. Interpreting deviant carriers would need a
+second truth predicate, and predicate divergence is exactly how
+``np.array(["0"])`` came to route as sparse while ``["0"]`` routed as dense;
+per the refactor series' Policy A the runner<->model interface owes no
+tolerance shim.
+
+A present-but-invalid declaration (illegal marker carrier, or marker set
+with an unusable ``meta.req_id``) FAILS CLOSED: it most likely belongs to a
+producer that INTENDED sparse output, and falling back to dense would
+batch-index-assign its payloads to the WRONG requests. `InvalidSparseDeclarationError`
+is raised by the classifiers and caught INSIDE this module by
+`resolve_sparse_mm_routing`, which routes no payload for the affected step —
+the exception never crosses the module boundary, so it can never escape a
+v1 runner (which does not catch per-request exceptions).
+
+This module is the single home of the contract until the model-state
+refactor series (G9-C2.6) turns the marker into a typed bool resolved at
+one adapter site and unifies the nested/flattened encodings. The NPU runner
+still carries a tolerant pre-enforcement copy — NPU-track handoff item.
+"""
+
+import threading
+import time
+from typing import Any
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+# The DESIGNED marker carrier is a list of strings (bare str also accepted);
+# truthiness of a string item:
+_SPARSE_MARKER_TRUTHY = ("1", "true", "yes", "on")
+
+# Per-process dedup so a broken producer logs once per offending shape, not
+# once per step:
+_logged_offenders: set[str] = set()
+_error_state_lock = threading.Lock()
+
+# A persistent invalid producer must remain observable without flooding the
+# per-step output path. Classifier errors above stay deduplicated per carrier
+# type; the routing consequence is reported at most once per interval, with a
+# count of repeats suppressed since the previous report.
+_FAIL_CLOSED_LOG_INTERVAL_S = 60.0
+_fail_closed_last_log_at: float | None = None
+_fail_closed_suppressed = 0
+
+
+class InvalidSparseDeclarationError(ValueError):
+    """Sparse-audio metadata is PRESENT but out of contract.
+
+    Internal control signal: raised by `is_sparse_audio_marker` /
+    `resolve_sparse_mm_req_ids` and caught by `resolve_sparse_mm_routing`
+    in this same module, which fails closed. Callers outside this module
+    only ever see the routing result, never the exception.
+    """
+
+
+def _error_once(key: str, msg: str, *args: Any) -> None:
+    with _error_state_lock:
+        if key in _logged_offenders:
+            return
+        _logged_offenders.add(key)
+    logger.error(msg, *args)
+
+
+def _error_fail_closed_rate_limited(cause: str) -> None:
+    global _fail_closed_last_log_at, _fail_closed_suppressed
+
+    now = time.monotonic()
+    with _error_state_lock:
+        if _fail_closed_last_log_at is not None and now - _fail_closed_last_log_at < _FAIL_CLOSED_LOG_INTERVAL_S:
+            _fail_closed_suppressed += 1
+            return
+        suppressed = _fail_closed_suppressed
+        _fail_closed_last_log_at = now
+        _fail_closed_suppressed = 0
+
+    logger.error(
+        "Invalid sparse-audio declaration: failing closed - no multimodal payload is routed "
+        "for the affected step(s). Fix the emitting model. Suppressed repeats before this "
+        "report: %d. Cause: %s",
+        suppressed,
+        cause,
+    )
+
+
+def is_sparse_audio_marker(value: Any) -> bool:
+    """Classify a ``meta.sparse_audio`` marker.
+
+    Legal carriers (list of strings, bare string) evaluate against the
+    literal truthy set; ``None`` means "marker absent" and returns ``False``.
+    An ILLEGAL carrier (tensors, arrays, numbers, lists with non-string
+    items) is a producer bug: logged at error level once per offending type,
+    then raised as `InvalidSparseDeclarationError`. Deviant carriers are never
+    read element-wise, so a device tensor can never introduce a D2H sync
+    here.
+    """
+    if value is None:
+        return False
+    if isinstance(value, str):
+        return value.lower() in _SPARSE_MARKER_TRUTHY
+    if isinstance(value, list) and all(isinstance(item, str) for item in value):
+        return any(item.lower() in _SPARSE_MARKER_TRUTHY for item in value)
+    vtype = type(value).__name__
+    _error_once(
+        vtype,
+        'Illegal sparse-audio marker of type %s: the marker MUST be a list of strings (e.g. ["1"]). '
+        "The emitting model is out of contract; routing fails closed for the affected step(s).",
+        vtype,
+    )
+    raise InvalidSparseDeclarationError(f"illegal sparse-audio marker carrier: {vtype}")
+
+
+def resolve_sparse_mm_req_ids(multimodal_outputs: Any) -> list[str] | None:
+    """Resolve the sparse request-id list.
+
+    Returns the request-id list when a legal sparse declaration is present,
+    and ``None`` when there is no sparse declaration (dense). Raises
+    `InvalidSparseDeclarationError` when the declaration is present but out of
+    contract (illegal marker carrier, or marker set with an unusable
+    ``meta.req_id``).
+    """
+    if not isinstance(multimodal_outputs, dict):
+        return None
+    meta = multimodal_outputs.get("meta")
+    nested_req_ids = meta.get("req_id") if isinstance(meta, dict) else None
+    flat_req_ids = multimodal_outputs.get("meta.req_id")
+
+    # Resolve both marker encodings independently, even when the request ids
+    # came from the other encoding. Previously a nested `req_id` prevented the
+    # flattened marker from being read, while the mirror combination worked.
+    # Evaluating both also preserves the strict contract: an illegal carrier
+    # on either side wins over an otherwise-valid declaration and fails closed.
+    nested_marker_value = meta.get("sparse_audio") if isinstance(meta, dict) else None
+    flat_marker_value = multimodal_outputs.get("meta.sparse_audio")
+    nested_marker = is_sparse_audio_marker(nested_marker_value)
+    flat_marker = is_sparse_audio_marker(flat_marker_value)
+    if nested_marker_value is not None and flat_marker_value is not None and nested_marker != flat_marker:
+        _error_once(
+            "marker:conflict",
+            "Nested and flattened sparse-audio markers disagree. Failing closed.",
+        )
+        raise InvalidSparseDeclarationError("nested and flattened sparse markers disagree")
+    sparse_audio = nested_marker or flat_marker
+    if not sparse_audio:
+        return None
+
+    req_id_carriers = [
+        (name, value) for name, value in (("nested", nested_req_ids), ("flattened", flat_req_ids)) if value is not None
+    ]
+    if not req_id_carriers:
+        req_id_carriers = [("missing", None)]
+
+    validated_req_ids: list[list[str]] = []
+    for carrier_name, req_ids in req_id_carriers:
+        if isinstance(req_ids, list) and all(isinstance(rid, str) for rid in req_ids):
+            validated_req_ids.append(req_ids)
+            continue
+        rtype = type(req_ids).__name__
+        _error_once(
+            f"req_id:{carrier_name}:{rtype}",
+            "Sparse-audio marker is set but the %s meta.req_id carrier is unusable (%s; "
+            "must be a list of request-id strings). Failing closed.",
+            carrier_name,
+            rtype,
+        )
+        raise InvalidSparseDeclarationError(f"sparse marker set but {carrier_name} meta.req_id is unusable: {rtype}")
+
+    req_ids = validated_req_ids[0]
+    if len(validated_req_ids) == 2 and validated_req_ids[1] != req_ids:
+        _error_once(
+            "req_id:conflict",
+            "Nested and flattened sparse-audio request-id lists disagree. Failing closed.",
+        )
+        raise InvalidSparseDeclarationError("nested and flattened sparse request-id lists disagree")
+    if len(req_ids) != len(set(req_ids)):
+        _error_once(
+            "req_id:duplicate",
+            "Sparse-audio request-id list contains duplicates. Failing closed.",
+        )
+        raise InvalidSparseDeclarationError("sparse request-id list contains duplicates")
+    return list(req_ids)
+
+
+def resolve_sparse_mm_routing(
+    *,
+    engine_output_type: str,
+    req_ids_output_copy: list[str],
+    downstream_req_ids: list[str],
+    multimodal_outputs: Any,
+) -> tuple[list[str], dict[str, int], bool]:
+    """Resolve `(downstream_req_ids, sparse_mm_index, audio_sparse_output)`.
+
+    The only entry point runners use; all protocol errors are absorbed here.
+    """
+    try:
+        sparse_mm_req_ids = resolve_sparse_mm_req_ids(multimodal_outputs)
+    except InvalidSparseDeclarationError as exc:
+        if engine_output_type != "audio":
+            # Non-audio pipelines never consult sparse routing; the carrier
+            # error is already logged by the classifier.
+            return downstream_req_ids, {}, False
+        # FAIL CLOSED: a present-but-invalid sparse declaration most likely
+        # belongs to a producer that INTENDED sparse output; routing it as
+        # dense would batch-index-assign its payloads to the WRONG requests.
+        # Route no payload for this step instead - a visible, bounded gap
+        # rather than silent cross-request corruption. (Empty-sparse is an
+        # already-exercised path: the zero-audio sparse step ships
+        # req_id=[] through the same shape.)
+        _error_fail_closed_rate_limited(str(exc))
+        return [], {}, True
+    sparse_mm_index = {rid: i for i, rid in enumerate(sparse_mm_req_ids or [])}
+    if engine_output_type != "audio" or sparse_mm_req_ids is None:
+        return downstream_req_ids, sparse_mm_index, False
+
+    sparse_req_id_set = set(sparse_mm_req_ids)
+    sparse_downstream_req_ids = [rid for rid in req_ids_output_copy if rid in sparse_req_id_set]
+    return sparse_downstream_req_ids, sparse_mm_index, True