[Model][Core] Avoid unused CosyVoice3 payloads and redundant AR output work (#8473)

Signed-off-by: Sorbett <sorbet111@163.com>

docs/design/module/ar_runtime.md

@@ -47,6 +47,29 @@ adapter, not injected through unrelated scheduler state.
 
 **Rule:** Workers and model runners MUST NOT implement cross-stage routing.
 
+## Token-only output payloads
+
+An AR stage whose downstream decoder consumes sampled tokens rather than
+hidden states can set `omni_pooler_payload_include_hidden = False`. Hidden
+states remain available for logits and sampling; prefill conditioning and
+other required multimodal outputs still cross the stage boundary.
+
+A chunk processor that needs IDs on steps without a tensor payload declares
+`requires_token_updates = True`. The scheduler then enqueues nonempty sampled
+IDs while preserving final-stage exclusions and the processor's EOF handling.
+Processors without this declaration retain their tensor-payload trigger.
+
+When hidden and multimodal payloads are absent, the shared runner skips empty
+per-request payload construction. Prefix-cache merges and pending model
+postprocess retain the normal path. This policy is independent of whether
+output materialization runs inline or in the background.
+
+Inline construction consumes query offsets, scheduled-token counts and
+scheduler metadata before the next execution can reuse them. Background
+construction snapshots those inputs. Both paths reuse the request IDs already
+copied by upstream bookkeeping; connector collection stays on the main thread
+because it can run a tensor-parallel collective.
+
 ## Safe-change guide
 
 Test request lifecycle, abort, cache state, and every affected worker execution

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -23,6 +23,7 @@ from vllm.v1.engine import FinishReason
 from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
+from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import talker2code2wav_async_chunk
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
@@ -287,7 +288,12 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
     )
 
 
-def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> None:
+@pytest.mark.parametrize("processor", [None, talker2code2wav_async_chunk], ids=["no_opt_in", "cosyvoice3"])
+@pytest.mark.parametrize("new_token_ids", [[], [42]], ids=["no_tokens", "sampled_token"])
+@pytest.mark.parametrize("omit_chunk_transfer", [False, True], ids=["downstream", "stage0_final"])
+def test_running_decode_step_without_inter_stage_payload_does_not_raise(
+    processor, new_token_ids, omit_chunk_transfer, mocker
+) -> None:
     """A decode step that neither stops nor carries an inter-stage payload.
 
     ``finished`` is only assigned when the request stops, yet the async-chunk
@@ -296,14 +302,16 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     """
     session = _make_request()
     session.status = RequestStatus.RUNNING
+    session.additional_information = {"omni_final_stage_id": 0} if omit_chunk_transfer else None
 
-    sched = MagicMock()
+    sched = mocker.MagicMock()
     sched.requests = {session.request_id: session}
     sched.perf_metrics = None
     sched.structured_output_manager.accept_tokens.return_value = True
-    sched._update_request_with_output.return_value = ([42], False)
+    sched._update_request_with_output.return_value = (new_token_ids, False)
     sched._process_kv_transfer_trigger.return_value = False
-    sched.chunk_transfer_adapter = MagicMock()
+    sched.chunk_transfer_adapter = mocker.MagicMock()
+    sched.chunk_transfer_adapter.custom_process_next_stage_input_func = processor
     sched.running = [session]
     sched.waiting_for_transfer_free = set()
     sched.transfer_triggered_requests = set()
@@ -316,13 +324,13 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.make_stats.return_value = None
     bind_omits_transfer_helpers(sched)
 
-    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output = mocker.MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
     scheduler_output.scheduled_spec_decode_tokens = {}
     scheduler_output.num_invalid_spec_tokens = 0
 
-    model_runner_output = MagicMock(spec=ModelRunnerOutput)
-    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output = mocker.MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = [new_token_ids]
     model_runner_output.logprobs = None
     model_runner_output.prompt_logprobs_dict = {}
     model_runner_output.pooler_output = None
@@ -335,8 +343,16 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
 
     OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
 
-    # Nothing to hand downstream: no payload, no segment boundary, not finished.
-    sched.chunk_transfer_adapter.save_async.assert_not_called()
+    if omit_chunk_transfer or processor is None or not new_token_ids:
+        sched.chunk_transfer_adapter.save_async.assert_not_called()
+    else:
+        sched.chunk_transfer_adapter.save_async.assert_called_once_with(
+            None,
+            session,
+            False,
+            new_token_ids=[42],
+            confirmed_num_computed_tokens=None,
+        )
 
 
 def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None:

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -3,32 +3,32 @@
 
 from __future__ import annotations
 
-import functools
+from collections import OrderedDict
 from threading import Lock
 from types import SimpleNamespace
-from typing import TYPE_CHECKING
 
 import pytest
 import torch
 import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.logits_processor.state import LogitsProcessors
 from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.worker.gpu_input_batch import InputBatch
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+from vllm_omni.config.model import OmniModelConfig
+from vllm_omni.data_entry_keys import to_struct
+from vllm_omni.model_executor.models.cosyvoice3 import cosyvoice3, cosyvoice3_talker
+from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import CosyVoice3Model, _normalize_request_conditioning
+from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
+from vllm_omni.worker import gpu_ar_model_runner
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-if TYPE_CHECKING:
-    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import CosyVoice3Model
-
-
-@functools.lru_cache(maxsize=1)
-def _cosyvoice3_model_and_runner():
-    """Defer heavy Omni/vLLM imports until a test runs (avoids duplicate CustomOp init)."""
-    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import CosyVoice3Model
-    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
-
-    return CosyVoice3Model, GPUARModelRunner
-
 
 class _DummyCode2Wav:
     def __init__(
@@ -89,7 +89,6 @@ def _make_code2wav_model(
     num_samples: int = 32,
     outputs: list[tuple[torch.Tensor, dict[str, object] | None]] | None = None,
 ) -> CosyVoice3Model:
-    CosyVoice3Model, _ = _cosyvoice3_model_and_runner()
     model = object.__new__(CosyVoice3Model)
     nn.Module.__init__(model)
     model.model_stage = "cosyvoice3_code2wav"
@@ -115,7 +114,6 @@ def _make_code2wav_model(
 
 
 def _make_talker_model() -> CosyVoice3Model:
-    CosyVoice3Model, _ = _cosyvoice3_model_and_runner()
     model = object.__new__(CosyVoice3Model)
     nn.Module.__init__(model)
     model.model_stage = "cosyvoice3_talker"
@@ -160,6 +158,162 @@ def _make_sampling_metadata(
     )
 
 
+@pytest.fixture
+def output_payload_vllm_config(tmp_path, request):
+    hf_config = CosyVoice3Config()
+    hf_config.llm.update(llm_input_size=16, llm_output_size=16, speech_token_size=32)
+    # Exercise the real config fields without model downloads or engine setup.
+    model_config = object.__new__(OmniModelConfig)
+    model_config.hf_config = hf_config
+    model_config.model_stage = "cosyvoice3_talker"
+    model_config.model = str(tmp_path)
+    model_config.async_chunk = getattr(request, "param", True)
+    model_config.enable_return_routed_experts = False
+    model_config.engine_output_type = "latent"
+    model_config.stage_connector_config = {"name": "SharedMemoryConnector", "extra": {"role": "sender"}}
+    vllm_config = object.__new__(VllmConfig)
+    vllm_config.model_config = model_config
+    return vllm_config
+
+
+@pytest.fixture
+def output_payload_talker(monkeypatch, output_payload_vllm_config):
+    class DummyEncoder(nn.Module):
+        def __init__(self, **kwargs):
+            super().__init__()
+
+        def forward(self, inputs_embeds, positions):
+            return inputs_embeds
+
+    class DummyTalker(nn.Module):
+        def __init__(self, *, llm: nn.Module, **kwargs):
+            super().__init__()
+            self.llm = llm
+
+    monkeypatch.setattr(cosyvoice3_talker, "VLLMQwen2Encoder", DummyEncoder)
+    monkeypatch.setattr(cosyvoice3_talker, "CosyVoice3LM", DummyTalker)
+    monkeypatch.setattr(CosyVoice3Model, "_create_llm_vllm_config", lambda self, config: config)
+
+    return CosyVoice3Model(vllm_config=output_payload_vllm_config)
+
+
+@pytest.mark.parametrize(
+    "output_payload_vllm_config", [True, False], indirect=True, ids=["async_chunks", "sync_chunks"]
+)
+@pytest.mark.parametrize("async_scheduling", [True, False], ids=["async_scheduling", "sync_scheduling"])
+def test_talker_skips_unused_hidden_payload_without_async_materialization(
+    output_payload_talker, output_payload_vllm_config, async_scheduling
+):
+    runner = object.__new__(GPUARModelRunner)
+    runner.use_async_scheduling = async_scheduling
+    runner.omni_prefix_cache = None
+    runner.speculative_config = None
+    runner.model_config = output_payload_vllm_config.model_config
+    runner.model = output_payload_talker
+    # load_model caches this policy before the output path runs.
+    runner._pooler_payload_include_hidden_flag = bool(getattr(runner.model, "omni_pooler_payload_include_hidden", True))
+
+    assert runner._model_omni_pooler_payload_include_hidden() is (not runner.model_config.async_chunk)
+    assert runner._should_use_async_omni_output() is False
+
+
+@pytest.mark.parametrize("is_prefill", [True, False], ids=["prefill", "decode"])
+@pytest.mark.parametrize(
+    "output_payload_vllm_config", [True, False], indirect=True, ids=["async_chunks", "sync_chunks"]
+)
+def test_talker_inline_output_preserves_conditioning_and_tokens(
+    output_payload_talker, output_payload_vllm_config, monkeypatch, mocker, is_prefill
+):
+    """Token-only output retains prefill conditioning and sampled codec IDs."""
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = output_payload_talker
+    runner._pooler_payload_include_hidden_flag = bool(getattr(runner.model, "omni_pooler_payload_include_hidden", True))
+    runner.vllm_config = output_payload_vllm_config
+    runner.model_config = output_payload_vllm_config.model_config
+    runner._async_chunk = runner.model_config.async_chunk
+    runner.omni_prefix_cache = None
+    runner.supports_mm_inputs = False
+    runner.routed_experts_initialized = False
+    runner.model_intermediate_buffer = {}
+    runner.input_batch = object.__new__(InputBatch)
+    runner.input_batch._req_ids = ["r1", "r2"]
+    runner.input_batch.req_id_to_index = {"r1": 0, "r2": 1}
+    monkeypatch.setattr(GPUARModelRunner, "_resolve_pooler_payload_req_ids", lambda self, req_ids: ("latent", req_ids))
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+
+    conditioning = {
+        "speech_token": torch.tensor([[11, 12], [21, 0]], dtype=torch.long),
+        "speech_token_len": torch.tensor([2, 1], dtype=torch.long),
+        "speech_feat": torch.arange(16, dtype=torch.float32).reshape(2, 4, 2),
+        "embedding": torch.tensor([[0.1, 0.2], [0.3, 0.4]]),
+    }
+    expected = {key: value.clone() for key, value in conditioning.items()}
+    hidden_states = output_payload_talker.forward(
+        input_ids=torch.tensor([1, 2, 3]),
+        positions=torch.arange(3),
+        inputs_embeds=torch.ones(3, 16),
+        **(conditioning if is_prefill else {}),
+    )
+    model_output = output_payload_talker.make_omni_output(hidden_states, **(conditioning if is_prefill else {}))
+
+    cpu_materialization = mocker.spy(gpu_ar_model_runner, "_to_cpu_contiguous")
+    scheduler_output = object.__new__(SchedulerOutput)
+    scheduler_output.total_num_scheduled_tokens = 3
+    scheduler_output.num_scheduled_tokens = {"r1": 2, "r2": 1}
+    output = runner._build_omni_model_runner_output_from_snapshot(
+        scheduler_output=scheduler_output,
+        hidden_states=model_output.text_hidden_states,
+        staged_hidden_states_cpu=None,
+        multimodal_outputs=model_output.multimodal_outputs,
+        req_ids_output_copy=["r1", "r2"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1},
+        valid_sampled_token_ids=[[101], [102]],
+        logprobs_lists=None,
+        prompt_logprobs_dict={},
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        ec_connector_output=None,
+        cudagraph_stats=None,
+        kv_extracted_req_ids=None,
+        num_scheduled_tokens_np=torch.tensor([2, 1], dtype=torch.int32).numpy(),
+        query_start_loc_cpu=torch.tensor([0, 2], dtype=torch.long),
+    )
+    assert output.req_ids == ["r1", "r2"]
+    assert output.sampled_token_ids == [[101], [102]]
+    hidden_storage = model_output.text_hidden_states.untyped_storage().data_ptr()
+    hidden_copied = any(
+        call.args[0].untyped_storage().data_ptr() == hidden_storage for call in cpu_materialization.call_args_list
+    )
+    assert hidden_copied is (not runner._async_chunk)
+    if runner._async_chunk:
+        assert output.multimodal_outputs is None
+    if not is_prefill and runner._async_chunk:
+        assert not output.inter_stage_outputs or all(not item for item in output.inter_stage_outputs)
+        return
+
+    assert len(output.inter_stage_outputs) == 2
+    for idx, prompt_len in enumerate([2, 1]):
+        item = output.inter_stage_outputs[idx]
+        assert ("hidden" not in item) is runner._async_chunk
+        if not is_prefill:
+            continue
+        # The downstream processor removes padding using speech_token_len.
+        torch.testing.assert_close(item["embed.speech_token"], expected["speech_token"][idx : idx + 1])
+        torch.testing.assert_close(item["embed.speech_feat"], expected["speech_feat"][idx : idx + 1])
+        assert item["embed.speech_token_len"].item() == prompt_len
+        torch.testing.assert_close(item["embed.embedding"], expected["embedding"][idx : idx + 1])
+
+
+def test_talker_preserves_packed_payload_policy(output_payload_talker, output_payload_vllm_config, monkeypatch):
+    output_payload_vllm_config.model_config.async_chunk = False
+    monkeypatch.setattr(cosyvoice3, "cosyvoice3_packed_inference_enabled", lambda: True)
+    model = CosyVoice3Model(vllm_config=output_payload_vllm_config)
+
+    assert model.omni_pooler_payload_include_hidden is False
+    assert not getattr(model, "use_async_omni_output", False)
+
+
 def test_forward_prefers_token_offset_when_present():
     model = _make_code2wav_model()
 
@@ -561,7 +715,6 @@ def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
             # runs for all sampling paths
             self.updated = True
 
-    _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
     runner.input_batch = _DummyInputBatch()
 
@@ -602,7 +755,6 @@ def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
             # runs for all sampling paths
             self.update_async_called = True
 
-    _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
     runner.input_batch = _DummyInputBatch()
 
@@ -649,7 +801,6 @@ def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
             # runs for all sampling paths (model sampler + fallback sampler)
             self.update_async_called = True
 
-    _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
     runner.input_batch = _DummyInputBatch()
 
@@ -722,8 +873,6 @@ def test_embed_input_ids_preserves_interleaved_request_boundaries(order):
 
 
 def test_mrv2_conditioning_follows_encoded_prompts_to_batch_rows():
-    from collections import OrderedDict
-
     model = _make_talker_model()
     speech = nn.Embedding.from_pretrained(torch.arange(100 * 4, dtype=torch.float32).reshape(100, 4))
     text = nn.Embedding.from_pretrained(torch.arange(100 * 4, dtype=torch.float32).reshape(100, 4) + 1000)
@@ -791,10 +940,8 @@ def test_embed_input_ids_rejects_cross_request_placeholder_block():
 
 
 def test_full_response_singleton_uses_the_optimized_batch_path(monkeypatch):
-    import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 as module
-
     monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1")
-    monkeypatch.setattr(module, "cosyvoice3_packed_inference_enabled", lambda: True)
+    monkeypatch.setattr(cosyvoice3, "cosyvoice3_packed_inference_enabled", lambda: True)
     model = _make_code2wav_model()
     model.forward(
         input_ids=torch.tensor([0, 1, 2]),
@@ -816,10 +963,6 @@ def test_full_response_singleton_uses_the_optimized_batch_path(monkeypatch):
 
 
 def test_full_response_adapter_uses_live_prefill_then_empty_decode(monkeypatch):
-    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-
-    from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
-
     model = _make_talker_model()
     hidden = torch.ones(2, 4)
     runner = object.__new__(OmniGPUModelRunner)
@@ -840,13 +983,11 @@ def test_full_response_adapter_uses_live_prefill_then_empty_decode(monkeypatch):
 
 
 def test_packed_flow_preserves_torch_estimator_when_speaker_trt_is_enabled(monkeypatch):
-    import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 as module
-
     model = _make_code2wav_model()
     model._code2wav_trt_done = False
-    monkeypatch.setattr(module.torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(module, "cosyvoice3_packed_inference_enabled", lambda: True)
-    monkeypatch.setattr(module, "_cosyvoice3_trt_enabled", lambda: True)
+    monkeypatch.setattr(cosyvoice3.torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(cosyvoice3, "cosyvoice3_packed_inference_enabled", lambda: True)
+    monkeypatch.setattr(cosyvoice3, "_cosyvoice3_trt_enabled", lambda: True)
 
     def unexpected_estimator_resolution():
         raise AssertionError("Packed Flow must not resolve or replace its torch estimator")
@@ -904,21 +1045,18 @@ def test_sampling_mode_preserves_or_merges_control_logits(mode):
 
 
 def test_cancelled_stream_releases_only_its_vocoder_state():
-    model_cls, _ = _cosyvoice3_model_and_runner()
     model = SimpleNamespace(
         _stream_audio_cache_lock=Lock(),
         _stream_vocoder_cache_by_req={"cancelled": {"mel": torch.ones(2)}, "live": {"mel": torch.zeros(2)}},
     )
-    model_cls.on_requests_finished(model, {"cancelled", "unknown"})
-    model_cls.on_requests_finished(model, {"cancelled"})
+    CosyVoice3Model.on_requests_finished(model, {"cancelled", "unknown"})
+    CosyVoice3Model.on_requests_finished(model, {"cancelled"})
     assert set(model._stream_vocoder_cache_by_req) == {"live"}
 
 
 @pytest.mark.parametrize("mrv2,packed", [(False, False), (False, True), (True, False)])
 def test_talker_output_contract(monkeypatch, mrv2, packed):
-    import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 as mod
-
-    monkeypatch.setattr(mod, "cosyvoice3_packed_inference_enabled", lambda: packed)
+    monkeypatch.setattr(cosyvoice3, "cosyvoice3_packed_inference_enabled", lambda: packed)
     model = _make_talker_model()
     hidden = torch.zeros(2, 4)
     model.model = SimpleNamespace(llm=lambda embeddings, positions: embeddings)
@@ -932,9 +1070,6 @@ def test_talker_output_contract(monkeypatch, mrv2, packed):
 
 
 def test_request_conditioning_normalizes_singletons_without_changing_schema():
-    from vllm_omni.data_entry_keys import to_struct
-    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
-
     token = torch.tensor([[7, 8]], dtype=torch.int32)
     feat = torch.randn(1, 4, 80)
     speaker = torch.randn(1, 192)
@@ -949,8 +1084,6 @@ def test_request_conditioning_normalizes_singletons_without_changing_schema():
 
 
 def test_request_conditioning_rejects_unsplit_voices():
-    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
-
     with pytest.raises(ValueError, match="unsplit batch"):
         _normalize_request_conditioning({"embed": {"embedding": [torch.ones(1, 192), torch.zeros(1, 192)]}})
 

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -6,6 +6,7 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.sampling_params import SamplingParams
 
 from vllm_omni.data_entry_keys import serialize_payload
 from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
@@ -14,6 +15,7 @@ from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     text2flow_token_only,
 )
 from vllm_omni.outputs.mm_outputs import MultimodalPayload
+from vllm_omni.request import OmniRequest
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -376,6 +378,61 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
     assert payload_final.codes.audio.tolist() == []
 
 
+def test_interleaved_token_only_requests_keep_conditioning_and_eof_separate():
+    transfer_manager = _transfer_manager()
+    requests = [
+        OmniRequest(
+            request_id=f"token-only-{index}",
+            external_req_id=f"token-only-{index}",
+            prompt_token_ids=[1],
+            sampling_params=SamplingParams(max_tokens=16),
+            pooling_params=None,
+            arrival_time=0.0,
+            block_hasher=None,
+        )
+        for index in range(2)
+    ]
+    conditioning = [
+        {
+            "embed": {
+                "speech_token": torch.tensor([[101 + index, 201 + index]]),
+                "speech_feat": torch.arange(8, dtype=torch.float32).reshape(1, 4, 2) + index,
+                "embedding": torch.tensor([[0.1 + index, 0.2 + index]]),
+            }
+        }
+        for index in range(2)
+    ]
+
+    requests[0].append_output_token_ids([8, 9])
+    requests[1].append_output_token_ids([18, 19])
+    first_chunks = [
+        talker2code2wav_async_chunk(transfer_manager, payload, request)
+        for request, payload in zip(requests, conditioning)
+    ]
+    for index, chunk in enumerate(first_chunks):
+        assert chunk is not None
+        assert chunk.codes.audio.tolist() == ([8, 9] if index == 0 else [18, 19])
+        torch.testing.assert_close(chunk.embed.embedding, conditioning[index]["embed"]["embedding"])
+
+    requests[0].append_output_token_ids([10, 11])
+    final_a = talker2code2wav_async_chunk(transfer_manager, None, requests[0], is_finished=True)
+    requests[1].append_output_token_ids([20])
+    pending_b = talker2code2wav_async_chunk(transfer_manager, None, requests[1])
+    final_b = talker2code2wav_async_chunk(transfer_manager, None, requests[1], is_finished=True)
+
+    assert final_a.codes.audio.tolist() == [8, 9, 10, 11]
+    assert final_a.meta.left_context_size == 2
+    assert final_a.meta.finished.item() is True
+    assert final_a.embed is None
+    assert pending_b is None
+    assert final_b.codes.audio.tolist() == [18, 19, 20]
+    assert final_b.meta.left_context_size == 2
+    assert final_b.meta.finished.item() is True
+    assert final_b.embed is None
+    for request in requests:
+        assert talker2code2wav_async_chunk(transfer_manager, None, request, is_finished=True) is None
+
+
 def test_full_payload_retains_true_length_for_padded_conditioning():
     from vllm_omni.data_entry_keys import to_struct
 

tests/worker/test_gpu_ar_model_runner.py

@@ -15,6 +15,8 @@ import numpy as np
 import pytest
 import torch
 from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.utils import CpuGpuBuffer
 from vllm.v1.worker import gpu_input_batch
@@ -25,7 +27,7 @@ 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 import gpu_ar_model_runner, sparse_audio
 from vllm_omni.worker.gpu_ar_model_runner import (
     ExecuteModelState,
     GPUARModelRunner,
@@ -614,8 +616,10 @@ def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatc
     assert output.multimodal_outputs is None
 
 
-def test_process_additional_information_uses_snapshot_request_order(monkeypatch):
+@pytest.mark.parametrize("include_hidden", [True, False])
+def test_process_additional_information_uses_snapshot_request_order(monkeypatch, include_hidden):
     runner = _make_async_output_runner()
+    runner._pooler_payload_include_hidden_flag = include_hidden
     seen = []
 
     class PostprocessModel:
@@ -958,6 +962,11 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
     monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
     monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
 
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_snapshot_scheduler_output_for_async_omni_output",
+        lambda *args: pytest.fail("inline output must not snapshot scheduler metadata"),
+    )
     output = GPUARModelRunner.sample_tokens(runner, grammar_output=None)
 
     # Non-async-chunk now ships the full payload to the next stage, so
@@ -1946,3 +1955,176 @@ def test_prepare_hook_keeps_legacy_signature_and_orders_opt_in_metadata(accepts_
         assert seen["discard"] == [True, False]
     else:
         assert set(seen) == {"req_ids"}
+
+
+@pytest.fixture
+def token_only_builder(monkeypatch):
+    runner = _make_async_output_runner()
+    runner._pooler_payload_include_hidden_flag = False
+    runner.supports_mm_inputs = True
+    monkeypatch.setattr(GPUARModelRunner, "_resolve_pooler_payload_req_ids", lambda self, req_ids: ("latent", req_ids))
+    scheduler = SchedulerOutput.make_empty()
+    scheduler.total_num_scheduled_tokens = 2
+    scheduler.num_scheduled_tokens = {"r1": 1, "r2": 1}
+    kwargs = dict(
+        scheduler_output=scheduler,
+        hidden_states=torch.tensor([[1.0], [2.0]]),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs=None,
+        req_ids_output_copy=["r1", "r2"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1},
+        valid_sampled_token_ids=[[11], [22]],
+        logprobs_lists=object(),
+        prompt_logprobs_dict={"r1": object()},
+        num_nans_in_logits={"r1": 0},
+        kv_connector_output=object(),
+        ec_connector_output=object(),
+        cudagraph_stats=object(),
+        kv_extracted_req_ids=["r2"],
+        num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32),
+    )
+    return runner, kwargs
+
+
+@pytest.mark.parametrize("empty_mm", [None, {}])
+def test_token_only_output_preserves_sampling_and_connector_metadata(token_only_builder, mocker, empty_mm):
+    runner, kwargs = token_only_builder
+    kwargs["multimodal_outputs"] = empty_mm
+    payload_builder = mocker.spy(runner, "_build_omni_pooler_payload")
+    output = runner._build_omni_model_runner_output_from_snapshot(**kwargs)
+
+    assert output.req_ids == ["r1", "r2"]
+    assert output.req_id_to_index == {"r1": 0, "r2": 1}
+    assert output.sampled_token_ids == [[11], [22]]
+    assert output.inter_stage_outputs is None
+    assert output.multimodal_outputs is None
+    assert output.pooler_output is None
+    assert output.logprobs is kwargs["logprobs_lists"]
+    for name in (
+        "prompt_logprobs_dict",
+        "num_nans_in_logits",
+        "kv_connector_output",
+        "ec_connector_output",
+        "cudagraph_stats",
+        "kv_extracted_req_ids",
+    ):
+        assert getattr(output, name) is kwargs[name]
+    payload_builder.assert_not_called()
+
+
+def test_token_only_output_keeps_prefix_cache_merge(token_only_builder, mocker):
+    runner, kwargs = token_only_builder
+    runner.omni_prefix_cache = object()
+    cached = {"codes.audio": {"r1": torch.tensor([31]), "r2": torch.tensor([32])}}
+    merge = mocker.patch.object(
+        runner, "_prepare_prefix_cache_pooler_payload_sources", return_value=(None, None, cached)
+    )
+    output = runner._build_omni_model_runner_output_from_snapshot(**kwargs)
+
+    merge.assert_called_once()
+    assert output.inter_stage_outputs[0]["codes.audio"].tolist() == [31]
+    assert output.inter_stage_outputs[1]["codes.audio"].tolist() == [32]
+
+
+def test_token_only_output_keeps_routed_experts(token_only_builder, mocker):
+    runner, kwargs = token_only_builder
+    runner.model_config.enable_return_routed_experts = True
+    runner.routed_experts_initialized = True
+    routes = np.array([[[1, 2]], [[3, 4]]], dtype=np.int32)
+    extract = mocker.patch.object(runner, "_omni_extract_routed_experts", return_value=routes)
+    output = runner._build_omni_model_runner_output_from_snapshot(**kwargs)
+
+    extract.assert_called_once_with(kwargs["scheduler_output"])
+    assert output.routed_experts is routes
+
+
+@pytest.mark.parametrize("snapshot", [True, False], ids=["background", "inline"])
+def test_query_metadata_lifetime_matches_materialization_mode(snapshot):
+    runner = object.__new__(GPUARModelRunner)
+    runner.query_start_loc = object.__new__(CpuGpuBuffer)
+    runner.query_start_loc.cpu = torch.tensor([0, 2, 3], dtype=torch.int32)
+    offsets = runner._get_query_start_loc_cpu(snapshot=snapshot)
+    runner.query_start_loc.cpu[1] = 9
+
+    if snapshot:
+        assert offsets.tolist() == [0, 2, 3]
+    else:
+        assert offsets is runner.query_start_loc.cpu
+        assert offsets.tolist() == [0, 9, 3]
+
+
+def test_background_output_survives_step_metadata_reuse(monkeypatch, mocker):
+    """Delay the real builder until offsets, batch IDs and counts are reused."""
+
+    runner = _make_async_output_runner()
+    runner.use_async_scheduling = True
+    runner.speculative_config = None
+    runner.model.use_async_omni_output = True
+    runner.broadcast_pp_output = False
+    runner.kv_connector_output = None
+    runner.device = torch.device("cpu")
+    runner.async_output_copy_stream = None
+    runner.input_batch = mocker.Mock(
+        spec=InputBatch, req_ids=["r1", "r2"], req_id_to_index={"r1": 0, "r2": 1}, vocab_size=10
+    )
+    runner.query_start_loc = object.__new__(CpuGpuBuffer)
+    runner.query_start_loc.cpu = torch.tensor([0, 1, 3], dtype=torch.int32)
+    runner._omni_num_scheduled_tokens_np = np.array([1, 2], dtype=np.int32)
+    scheduler = SchedulerOutput.make_empty()
+    scheduler.num_scheduled_tokens = {"r1": 1, "r2": 2}
+    scheduler.total_num_scheduled_tokens = 3
+    hidden = torch.tensor([[1.0], [2.0], [3.0]])
+    runner.execute_model_state = ExecuteModelState(
+        scheduler, None, None, None, hidden, None, None, None, None, None, {}, None
+    )
+    sampler = SamplerOutput(sampled_token_ids=torch.tensor([[7], [8]]), logprobs_tensors=None)
+    mocker.patch.object(runner, "_sample", return_value=sampler)
+    mocker.patch.object(runner, "_update_states_after_model_execute")
+    # Match upstream bookkeeping's ownership contract: these are already copies.
+    mocker.patch.object(
+        runner,
+        "_bookkeeping_sync",
+        return_value=(
+            {},
+            None,
+            None,
+            [],
+            {},
+            runner.input_batch.req_ids.copy(),
+            runner.input_batch.req_id_to_index.copy(),
+            [],
+        ),
+    )
+    mocker.patch.object(runner, "_run_post_sample_talker_mtp", return_value={})
+    mocker.patch.object(runner, "eplb_step")
+    mocker.patch.object(runner, "get_omni_connector_output", return_value=None)
+    monkeypatch.setattr(gpu_ar_model_runner, "get_pp_group", lambda: mocker.Mock(world_size=1))
+    mocker.patch.object(runner, "_resolve_pooler_payload_req_ids", side_effect=lambda ids: ("audio", ids))
+    # Device copying is covered separately; this test isolates metadata lifetime.
+    mocker.patch.object(
+        runner,
+        "_snapshot_omni_output_tensors_for_async_output",
+        return_value=gpu_ar_model_runner._OmniOutputTensorSnapshot(
+            hidden_states=hidden, staged_hidden_states_cpu=None, multimodal_outputs={}
+        ),
+    )
+    async_output = mocker.Mock(
+        spec=OmniAsyncGPUModelRunnerOutput,
+        sampled_token_ids_cpu=sampler.sampled_token_ids,
+        async_copy_ready_event=object(),
+    )
+    constructor = mocker.patch.object(gpu_ar_model_runner, "OmniAsyncGPUModelRunnerOutput", return_value=async_output)
+
+    assert runner.sample_tokens(None) is async_output
+    runner.query_start_loc.cpu.fill_(99)
+    runner._omni_num_scheduled_tokens_np.fill(99)
+    scheduler.num_scheduled_tokens.clear()
+    runner.input_batch.req_ids[:] = ["next-step"]
+    runner.input_batch.req_id_to_index.clear()
+    output = constructor.call_args.kwargs["model_runner_output_builder"]()
+
+    assert output.req_ids == ["r1", "r2"]
+    assert output.req_id_to_index == {"r1": 0, "r2": 1}
+    torch.testing.assert_close(output.inter_stage_outputs[0]["hidden"], hidden[:1])
+    torch.testing.assert_close(output.inter_stage_outputs[1]["hidden"], hidden[1:])

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -432,6 +432,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         pooler_outputs = model_runner_output.pooler_output
         mm_outputs = getattr(model_runner_output, "multimodal_outputs", None)
         inter_stage_outputs = getattr(model_runner_output, "inter_stage_outputs", None)
+        # Token-only processors need sampled IDs even without a tensor payload.
+        processor = getattr(getattr(self, "chunk_transfer_adapter", None), "custom_process_next_stage_input_func", None)
+        requires_token_updates = getattr(processor, "requires_token_updates", False) is True
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
         ec_connector_output = getattr(model_runner_output, "ec_connector_output", None)
@@ -689,7 +692,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             # Only when this step might save. Read additional_information
             # before _free_request rewrites it.
             omits_chunk_transfer = False
-            if self.chunk_transfer_adapter is not None and (inter_stage_output is not None or stopped):
+            if self.chunk_transfer_adapter is not None and (
+                inter_stage_output is not None or (new_token_ids and requires_token_updates) or stopped
+            ):
                 omits_chunk_transfer = self._request_omits_chunk_transfer_to_next_stage(request)
             # Capture before resumable stop handling can clear token history.
             output_token_ids: Any = getattr(request, "output_token_ids", None)
@@ -818,7 +823,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if (
                 self.chunk_transfer_adapter is not None
                 and not omits_chunk_transfer
-                and (inter_stage_output is not None or is_segment_finished or finished)
+                and (
+                    inter_stage_output is not None
+                    or (new_token_ids and requires_token_updates)
+                    or is_segment_finished
+                    or finished
+                )
             ):
                 save_kwargs = {
                     "new_token_ids": new_token_ids,

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -773,6 +773,10 @@ class CosyVoice3Model(
         self.model_dir = model_dir
         self.model = None
         if self.model_stage == "cosyvoice3_talker":
+            # Code2Wav consumes sampled tokens and prompt conditioning, not
+            # hidden states. The processor opts into token-only chunk updates.
+            if getattr(vllm_config.model_config, "async_chunk", False):
+                self.omni_pooler_payload_include_hidden = False
             # Initialize talker stage (text to speech tokens)
             from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_talker import CosyVoice3LM, VLLMQwen2Encoder
 

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -297,6 +297,10 @@ _FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset(
 )
 
 
+# This processor consumes sampled codec IDs even when no tensor payload is emitted.
+talker2code2wav_async_chunk.requires_token_updates = True  # type: ignore[attr-defined]
+
+
 def text2flow_token_only(
     source_outputs: list,
     prompt: OmniTokensPrompt | TextPrompt = None,

vllm_omni/worker/gpu_ar_model_runner.py

@@ -11,7 +11,7 @@ from __future__ import annotations
 
 import gc
 import threading
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from contextlib import nullcontext
 from copy import copy
 from dataclasses import replace
@@ -1596,10 +1596,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
             return None
         return wire_payloads
 
-    def _snapshot_query_start_loc_cpu(self) -> Any:
+    def _get_query_start_loc_cpu(self, *, snapshot: bool) -> Any:
         query_start_loc_cpu = self.query_start_loc.cpu
         if callable(query_start_loc_cpu):
             query_start_loc_cpu = query_start_loc_cpu()
+        if not snapshot:
+            return query_start_loc_cpu
         if isinstance(query_start_loc_cpu, torch.Tensor):
             return query_start_loc_cpu.detach().cpu().clone()
         if isinstance(query_start_loc_cpu, np.ndarray):
@@ -1888,7 +1890,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         # OmniModelRunnerOutput.pooler_output (which is set to None below).
         # The actual multimodal wire transport uses multimodal_outputs instead.
         pooler_output: list[dict[str, object]] | None = None
-        if needs_pooler_payload:
+        # Token-only steps have no pooler payload to construct. Keep the full
+        # path for cached payloads and model postprocess, even without MM input.
+        has_mm_payload = multimodal_outputs is not None and (
+            not isinstance(multimodal_outputs, Mapping) or len(multimodal_outputs) > 0
+        )
+        needs_payload_work = (
+            include_hidden_payload
+            or has_mm_payload
+            or self.omni_prefix_cache is not None
+            or (not postprocess_already_applied and self._runner_model_omni_flag("has_postprocess"))
+        )
+        if needs_pooler_payload and needs_payload_work:
             hidden_seq_len = int(hidden_states.shape[0])
             scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens)
             mm_cpu = None
@@ -2186,6 +2199,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         kv_connector_output = self.kv_connector_output
         self.kv_connector_output = None
 
+        use_async_omni_output = self._should_use_async_omni_output()
         num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
         if num_scheduled_tokens_np is None:
             num_scheduled_tokens_np = np.array(
@@ -2193,20 +2207,29 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
                 dtype=np.int32,
             )
         else:
-            num_scheduled_tokens_np = np.asarray(num_scheduled_tokens_np, dtype=np.int32).copy()
+            num_scheduled_tokens_np = np.asarray(num_scheduled_tokens_np, dtype=np.int32)
+            if use_async_omni_output:
+                num_scheduled_tokens_np = num_scheduled_tokens_np.copy()
 
-        query_start_loc_cpu = self._snapshot_query_start_loc_cpu()
-        scheduler_output_snapshot = self._snapshot_scheduler_output_for_async_omni_output(scheduler_output)
-        req_ids_output_snapshot = list(req_ids_output_copy)
-        req_id_to_index_output_snapshot = dict(req_id_to_index_output_copy)
-        valid_sampled_token_ids_snapshot = [list(token_ids) for token_ids in valid_sampled_token_ids]
-        logprobs_lists_snapshot = copy(logprobs_lists) if logprobs_lists is not None else None
-        prompt_logprobs_dict_snapshot = dict(prompt_logprobs_dict) if prompt_logprobs_dict is not None else {}
-        num_nans_in_logits_snapshot = (
-            dict(num_nans_in_logits) if isinstance(num_nans_in_logits, dict) else num_nans_in_logits
-        )
-
-        use_async_omni_output = self._should_use_async_omni_output()
+        # Inline materialization finishes before the next execute_model can
+        # mutate step metadata. Only the background builder needs snapshots.
+        query_start_loc_cpu = self._get_query_start_loc_cpu(snapshot=use_async_omni_output)
+        scheduler_output_snapshot = scheduler_output
+        valid_sampled_token_ids_snapshot = valid_sampled_token_ids
+        logprobs_lists_snapshot = logprobs_lists
+        prompt_logprobs_dict_snapshot = prompt_logprobs_dict
+        num_nans_in_logits_snapshot = num_nans_in_logits
+        if use_async_omni_output:
+            scheduler_output_snapshot = self._snapshot_scheduler_output_for_async_omni_output(scheduler_output)
+            valid_sampled_token_ids_snapshot = [list(token_ids) for token_ids in valid_sampled_token_ids]
+            logprobs_lists_snapshot = copy(logprobs_lists) if logprobs_lists is not None else None
+            prompt_logprobs_dict_snapshot = dict(prompt_logprobs_dict) if prompt_logprobs_dict is not None else {}
+            num_nans_in_logits_snapshot = (
+                dict(num_nans_in_logits) if isinstance(num_nans_in_logits, dict) else num_nans_in_logits
+            )
+        # _bookkeeping_sync already detached these from the mutable input batch.
+        req_ids_output_snapshot = req_ids_output_copy
+        req_id_to_index_output_snapshot = req_id_to_index_output_copy
         omni_postprocess_already_applied = False
         if use_async_omni_output:
             omni_postprocess_already_applied = self._maybe_run_eager_omni_postprocess_before_async_output(