[Feature] Support Async Omni output materialization (#4476)

Signed-off-by: Chenguang ZHENG <645327136@qq.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/worker/test_gpu_ar_model_runner.py

@@ -1,9 +1,15 @@
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 import torch
 
-from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, GPUARModelRunner
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.gpu_ar_model_runner import (
+    ExecuteModelState,
+    GPUARModelRunner,
+    OmniAsyncGPUModelRunnerOutput,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -52,6 +58,300 @@ def test_sparse_mm_req_ids_requires_sparse_audio_marker():
     assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
 
 
+def test_omni_async_gpu_model_runner_output_builds_lazily_once():
+    async_output = object.__new__(OmniAsyncGPUModelRunnerOutput)
+    calls = []
+    sync_calls = []
+
+    def builder():
+        calls.append("build")
+        return OmniModelRunnerOutput(req_ids=["r1"], req_id_to_index={"r1": 0})
+
+    async_output._model_runner_output = None
+    async_output._model_runner_output_builder = builder
+    async_output._invalid_req_indices = []
+    async_output.sampled_token_ids_cpu = torch.tensor([[7]], dtype=torch.long)
+    async_output.async_copy_ready_event = SimpleNamespace(synchronize=lambda: sync_calls.append("sync"))
+    async_output._sampled_token_ids = torch.tensor([[7]], dtype=torch.long)
+    async_output._logprobs_tensors = None
+    async_output._logprobs_tensors_cpu = None
+    async_output._routed_experts = None
+    async_output._routed_experts_cpu = None
+    async_output.vocab_size = 10
+
+    output = async_output.get_output()
+
+    assert calls == ["build"]
+    assert sync_calls == ["sync"]
+    assert async_output._model_runner_output_builder is None
+    assert output.req_ids == ["r1"]
+    assert output.sampled_token_ids == [[7]]
+
+
+def test_omni_async_gpu_model_runner_output_reraises_background_exception():
+    async_output = object.__new__(OmniAsyncGPUModelRunnerOutput)
+    joined = []
+
+    class FakeThread:
+        def join(self):
+            joined.append("join")
+
+    async_output._background_thread = FakeThread()
+    async_output._background_exception = RuntimeError("background failed")
+
+    with pytest.raises(RuntimeError, match="background failed"):
+        async_output.get_output()
+
+    assert joined == ["join"]
+    assert async_output._background_thread is None
+
+
+def _make_async_output_runner(engine_output_type: str = "audio"):
+    runner = object.__new__(GPUARModelRunner)
+    model_config = SimpleNamespace(
+        engine_output_type=engine_output_type,
+        async_chunk=True,
+        enable_return_routed_experts=False,
+    )
+    runner.vllm_config = SimpleNamespace(model_config=model_config)
+    runner.model_config = model_config
+    runner.omni_prefix_cache = None
+    runner.requests = {"r1": object(), "r2": object()}
+    runner.supports_mm_inputs = False
+    runner.routed_experts_initialized = False
+    runner.model = SimpleNamespace(has_postprocess=False)
+    runner.model_intermediate_buffer = {}
+    runner.input_batch = SimpleNamespace(
+        req_ids=["mutated"],
+        req_id_to_index={"mutated": 0},
+    )
+    return runner
+
+
+def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monkeypatch):
+    runner = _make_async_output_runner()
+    events = []
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("audio", req_ids),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: True)
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "accumulate_full_payload_output",
+        lambda self, rid, payload, request: events.append(f"accumulate:{rid}"),
+    )
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "get_omni_connector_output",
+        lambda self: events.append("connector") or "connector-output",
+    )
+
+    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=3,
+            num_scheduled_tokens={"r1": 1, "r2": 2},
+        ),
+        hidden_states=torch.tensor([[1.0], [2.0], [3.0]]),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={"foo": torch.tensor([10.0, 20.0, 30.0])},
+        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=["r2"],
+        seq_len=3,
+        num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
+        query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
+    )
+
+    assert output.req_ids == ["r1", "r2"]
+    assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0]]))
+    assert torch.equal(output.multimodal_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
+    assert output.kv_extracted_req_ids == ["r2"]
+    assert output.omni_connector_output == "connector-output"
+    assert events == ["accumulate:r1", "accumulate:r2", "connector"]
+
+
+def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatch):
+    runner = _make_async_output_runner(engine_output_type="latent")
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("latent", ["r2"]),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
+
+    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=6,
+            num_scheduled_tokens={"r1": 1, "r2": 2, "r3": 3},
+        ),
+        hidden_states=torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]]),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={},
+        req_ids_output_copy=["r1", "r2", "r3"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
+        valid_sampled_token_ids=[[], [], []],
+        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,
+        seq_len=6,
+        num_scheduled_tokens_np=np.array([1, 2, 3], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.long),
+    )
+
+    assert output.multimodal_outputs is not None
+    assert output.multimodal_outputs[0] == {}
+    assert torch.equal(output.multimodal_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
+    assert output.multimodal_outputs[2] == {}
+
+
+def test_process_additional_information_uses_snapshot_request_order(monkeypatch):
+    runner = _make_async_output_runner()
+    seen = []
+
+    class PostprocessModel:
+        has_postprocess = True
+
+        def postprocess(self, hidden_states, **kwargs):
+            seen.append(hidden_states.clone())
+            return {}
+
+    runner.model = PostprocessModel()
+    runner.model_intermediate_buffer = {"r1": {}, "r2": {}}
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("audio", req_ids),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+    monkeypatch.setattr(GPUARModelRunner, "_update_intermediate_buffer", lambda *args, **kwargs: None)
+
+    GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=3,
+            num_scheduled_tokens={"r1": 1, "r2": 2},
+        ),
+        hidden_states=torch.tensor([[1.0], [2.0], [3.0]]),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={},
+        req_ids_output_copy=["r1", "r2"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1},
+        valid_sampled_token_ids=[[], []],
+        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,
+        seq_len=3,
+        num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
+        query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
+    )
+
+    assert len(seen) == 2
+    assert torch.equal(seen[0], torch.tensor([[1.0]]))
+    assert torch.equal(seen[1], torch.tensor([[2.0], [3.0]]))
+
+
+def test_async_omni_output_guard_requires_safe_conditions():
+    runner = _make_async_output_runner()
+    runner.use_async_scheduling = True
+    runner.speculative_config = None
+    runner.model.use_async_omni_output = True
+
+    assert GPUARModelRunner._should_use_async_omni_output(runner)
+
+    runner.omni_prefix_cache = object()
+    assert not GPUARModelRunner._should_use_async_omni_output(runner)
+
+    runner.omni_prefix_cache = None
+    runner.model.has_postprocess = True
+    assert not GPUARModelRunner._should_use_async_omni_output(runner)
+
+    runner.model.eager_omni_postprocess_before_async_output = True
+    assert GPUARModelRunner._should_use_async_omni_output(runner)
+
+
+def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
+    runner = _make_async_output_runner(engine_output_type="latent")
+    runner.model.omni_pooler_payload_include_hidden = False
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("latent", req_ids),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
+
+    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=2,
+            num_scheduled_tokens={"r1": 2},
+        ),
+        hidden_states=torch.tensor([[1.0], [2.0]]),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={"codes": {"audio": torch.tensor([[7, 8], [9, 10]], dtype=torch.long)}},
+        req_ids_output_copy=["r1"],
+        req_id_to_index_output_copy={"r1": 0},
+        valid_sampled_token_ids=[[101]],
+        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,
+        seq_len=2,
+        num_scheduled_tokens_np=np.array([2], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0], dtype=torch.long),
+    )
+
+    assert output.multimodal_outputs is not None
+    assert len(output.multimodal_outputs) == 1
+    assert "hidden" not in output.multimodal_outputs[0]
+    assert torch.equal(output.multimodal_outputs[0]["codes.audio"], torch.tensor([[7, 8], [9, 10]], dtype=torch.long))
+
+
+def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
+    runner = _make_async_output_runner()
+    runner.model.omni_pooler_payload_include_hidden = False
+
+    payload = GPUARModelRunner._build_omni_async_snapshot_payload(
+        runner,
+        hidden_states=torch.tensor([[1.0], [2.0]]),
+        staged_hidden_states_cpu=torch.tensor([[3.0]]),
+        multimodal_outputs={"codes": {"audio": torch.tensor([[1]], dtype=torch.long)}},
+    )
+
+    assert set(payload.keys()) == {"multimodal_outputs"}
+    assert payload["multimodal_outputs"]["codes"]["audio"].tolist() == [[1]]
+
+
 @pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"])
 def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr):
     runner = object.__new__(GPUARModelRunner)
@@ -133,3 +433,54 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
         output.multimodal_outputs[1]["hidden"],
         torch.tensor([[2.0, 20.0], [3.0, 30.0]]),
     )
+
+
+def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch):
+    """Tail-only prefix-cache models still need per-step mm passthrough (e.g. codes.audio)."""
+    runner = _make_async_output_runner(engine_output_type="latent")
+    runner.omni_prefix_cache = object()
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("latent", req_ids),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"})
+    monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None)
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_prepare_prefix_cache_pooler_payload_sources",
+        lambda *args, **kwargs: (None, None, None),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+
+    codes = torch.tensor([[11.0, 12.0], [21.0, 22.0]], dtype=torch.float32)
+    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=2,
+            num_scheduled_tokens={"r1": 1, "r2": 1},
+        ),
+        hidden_states=torch.tensor([[1.0], [2.0]]),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={"codes.audio": codes},
+        req_ids_output_copy=["r1", "r2"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1},
+        valid_sampled_token_ids=[[], []],
+        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,
+        seq_len=2,
+        num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
+    )
+
+    assert torch.equal(output.multimodal_outputs[0]["codes.audio"], codes[0:1])
+    assert torch.equal(output.multimodal_outputs[1]["codes.audio"], codes[1:2])

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -28,7 +28,7 @@ stages:
     max_num_seqs: 64
     gpu_memory_utilization: 0.9
     trust_remote_code: true
-    enable_prefix_caching: true
+    enable_prefix_caching: false
     devices: "0"
     default_sampling_params:
       temperature: 0.0

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -105,6 +105,7 @@ class Qwen3OmniMoeForConditionalGeneration(
         self.have_multimodal_outputs = True
         self.has_preprocess = False
         self.has_postprocess = False
+        self.use_async_omni_output = False
         config: Qwen3OmniMoeConfig = vllm_config.model_config.hf_config
         multimodal_config = vllm_config.model_config.multimodal_config
 
@@ -129,6 +130,7 @@ class Qwen3OmniMoeForConditionalGeneration(
         self.model_stage = vllm_config.model_config.model_stage
 
         if self.model_stage == "thinker":
+            self.use_async_omni_output = True
             # Initialize thinker model (multimodal processing + text generation)
             # Create a new vllm_config with thinker_config as the hf_config
             thinker_vllm_config = vllm_config.with_hf_config(
@@ -152,6 +154,13 @@ class Qwen3OmniMoeForConditionalGeneration(
             multimodal_config.skip_mm_profiling = True
             self.has_preprocess = True
             self.has_postprocess = True
+            # Talker only ships codec codes to code2wav; skip latent hidden D2H.
+            # Build Omni output asynchronously like thinker, but run lightweight
+            # postprocess eagerly so hidden_states.last stays on GPU before the
+            # next decode step.
+            self.use_async_omni_output = True
+            self.eager_omni_postprocess_before_async_output = True
+            self.omni_pooler_payload_include_hidden = False
             self.set_custom_preprocess(self.talker_preprocess)
             self.set_custom_postprocess(self.talker_postprocess)
             self.thinker = None

vllm_omni/worker/gpu_model_runner.py

@@ -1378,8 +1378,15 @@ class OmniGPUModelRunner(GPUModelRunner):
         combined_hidden_states: dict[str, torch.Tensor] | None = None,
         combined_multimodal_outputs: dict[str, object] | None = None,
         req_ids_filter: set[str] | None = None,
+        req_ids: list[str] | None = None,
+        query_start_loc_cpu: object | None = None,
     ) -> None:
         """Process model-provided per-request updates and merge into model_intermediate_buffer."""
+        req_ids = req_ids if req_ids is not None else self.input_batch.req_ids
+        if query_start_loc_cpu is None:
+            query_start_loc_cpu = self.query_start_loc.cpu
+            if callable(query_start_loc_cpu):
+                query_start_loc_cpu = query_start_loc_cpu()
         try:
             # execute the custom postprocess function
             # TODO(Peiqi): do we have a more elegant way to do this?
@@ -1387,7 +1394,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 postprocess_uses_hidden_states = getattr(self.model, "postprocess_uses_hidden_states", True)
                 postprocess_uses_multimodal_outputs = getattr(self.model, "postprocess_uses_multimodal_outputs", True)
                 postprocess_uses_req_infos = getattr(self.model, "postprocess_uses_req_infos", True)
-                for req_index, req_id in enumerate(self.input_batch.req_ids):
+                for req_index, req_id in enumerate(req_ids):
                     if req_ids_filter is not None and req_id not in req_ids_filter:
                         continue
                     req_infos = self.model_intermediate_buffer.get(req_id, {}) if postprocess_uses_req_infos else {}
@@ -1396,7 +1403,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                             # Combined hidden states contains all hidden states for every request
                             hidden_states_slice = combined_hidden_states[req_id]
                         else:
-                            start_offset = int(self.query_start_loc.cpu[req_index])
+                            start_offset = int(query_start_loc_cpu[req_index])
                             sched_tokens = int(num_scheduled_tokens_np[req_index])
                             s, e = start_offset, start_offset + sched_tokens
                             # only consider to store data into update dict.
@@ -1426,7 +1433,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                     )
                     self._update_intermediate_buffer(req_id, update_dict)
         except Exception as e:
-            logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}")
+            logger.error(f"Error merging for requests:{req_ids} additional information update: {e}")
             import traceback
 
             traceback.print_exc()