[Core] Performance fix for orchestrator bottleneck: separates inter-stage from client outputs (#4527)

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Chenguang ZHENG <645327136@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

docs/contributing/profiling.md

@@ -286,3 +286,54 @@ Recommended viewers:
 - Nsight Systems GUI for CUDA kernel timelines
 
 For upstream background on the underlying vLLM profiling infrastructure, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).
+
+## 7. Orchestrator Monitor
+
+> **Warning:** Diagnostic only. Adds lightweight counters on the orchestrator poll loop and per-replica queue depth. Disable in production unless you are actively debugging orchestrator saturation.
+
+Multi-stage omni pipelines route client-facing outputs and inter-stage connector traffic through a single-process orchestrator. When that loop is saturated (for example by large multimodal payloads on the poll path), TTFT and stability can degrade even when GPU stages are healthy. The orchestrator monitor records time-series signals to confirm whether the orchestrator or downstream stage queues are the bottleneck.
+
+Enable on the stage-0 (orchestrator) process:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+  --stage-id 0 \
+  --enable-orch-monitor
+```
+
+Optional output path override:
+
+```bash
+export VLLM_OMNI_ORCH_MONITOR_PATH=/tmp/vllm_omni_orch_monitor.json
+```
+
+If unset, the monitor writes `./vllm_omni_orch_monitor_<MMDDHHMM>.json` at shutdown.
+
+### Recorded metrics
+
+Each 1-second window records:
+
+| Series | Description |
+|---|---|
+| `windows.duration_s` | Wall time covered by the window |
+| `windows.loop_idle` / `windows.loop_active` | Orchestrator poll-loop iterations with no work vs. active forwarding |
+| `replicas.<stage,replica>.outputs_queue_size` | MP client `outputs_queue` backlog for that replica |
+| `replicas.<stage,replica>.inflight` | Requests currently bound/routed to the replica |
+
+On shutdown the server also logs a short summary (`loop_active_pct`, per-replica queue averages/maxima).
+
+### Relationship to other diagnostics
+
+This monitor is intentionally separate from the existing profiling tools:
+
+| Tool | Scope | Output |
+|---|---|---|
+| `--enable-diffusion-pipeline-profiler` | Diffusion stage functions (`vae.decode`, `diffuse`, …) | Per-request `stage_durations` / logs |
+| `--enable-ar-profiler` | AR stage generation time | Per-request `stage_durations` |
+| `profiler_config` (`torch` / `cuda`) | GPU/CPU kernels inside a stage worker | `trace.json`, Nsight, operator tables |
+| Prometheus `/metrics` (`vllm:omni_*`) | Pipeline SLOs and cross-stage transfer | Continuous scrape endpoint |
+| `--enable-orch-monitor` | Orchestrator poll loop + replica queue backlog | Single JSON file at shutdown |
+
+The orchestrator monitor does not use `torch.profiler` because the bottleneck signal lives in the orchestrator process (poll-loop duty cycle and queue depth), not inside stage workers. It complements Prometheus metrics when you need a post-run time series for a benchmark or regression run.
+
+Process liveness monitors (`monitor_engine_liveness`, worker/engine-core monitors) are unrelated: they detect crashed workers, not orchestrator load.

tests/engine/test_orchestrator_monitor.py

@@ -0,0 +1,101 @@
+import json
+import time
+
+import pytest
+
+import vllm_omni.engine.orchestrator_monitor as monitor_mod
+from vllm_omni.engine.orchestrator_monitor import (
+    _NULL_ORCH_MONITOR,
+    OrchestratorMonitor,
+    create_orch_monitor,
+    replica_key,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_replica_key_format():
+    assert replica_key(0, 0) == "stage=0,replica=0"
+    assert replica_key(2, 1) == "stage=2,replica=1"
+
+
+def test_create_orch_monitor_disabled_returns_null_singleton():
+    monitor = create_orch_monitor(enabled=False, replica_sampler=lambda: {})
+    assert monitor is _NULL_ORCH_MONITOR
+    assert create_orch_monitor(enabled=False, replica_sampler=lambda: {}) is monitor
+    monitor.note_loop(idle=True)
+    monitor.flush()
+
+
+def test_note_loop_and_flush_json(tmp_path, monkeypatch):
+    out_path = tmp_path / "orch_monitor.json"
+    monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))
+
+    now = [0.0]
+
+    def fake_monotonic() -> float:
+        return now[0]
+
+    monkeypatch.setattr(time, "monotonic", fake_monotonic)
+
+    def sampler() -> dict[str, tuple[int, int]]:
+        return {replica_key(0, 0): (5, 10)}
+
+    monitor = create_orch_monitor(enabled=True, replica_sampler=sampler)
+    monitor.register_replica(0, 0)
+
+    monitor.note_loop(idle=False)
+    monitor.note_loop(idle=True)
+
+    now[0] = 1.5
+    monitor.flush()
+
+    payload = json.loads(out_path.read_text())
+    assert payload["configured_window_s"] == monitor_mod._WINDOW_S
+    assert payload["windows"]["loop_idle"] == [1]
+    assert payload["windows"]["loop_active"] == [1]
+    assert payload["replicas"]["stage=0,replica=0"]["outputs_queue_size"] == [5]
+    assert payload["replicas"]["stage=0,replica=0"]["inflight"] == [10]
+
+
+def test_register_replica_backfills_prior_windows(tmp_path, monkeypatch):
+    out_path = tmp_path / "orch_monitor.json"
+    monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))
+
+    now = [0.0]
+    monkeypatch.setattr(time, "monotonic", lambda: now[0])
+
+    monitor = OrchestratorMonitor(replica_sampler=lambda: {})
+    monitor.note_loop(idle=True)
+    now[0] = 1.5
+    monitor.register_replica(1, 0)
+    monitor.note_loop(idle=False)
+    monitor.flush()
+
+    series = json.loads(out_path.read_text())["replicas"]["stage=1,replica=0"]
+    assert series["outputs_queue_size"] == [0, 0]
+    assert series["inflight"] == [0, 0]
+
+
+def test_note_loop_survives_sampler_failure():
+    def bad_sampler() -> dict[str, tuple[int, int]]:
+        raise RuntimeError("sampler failed")
+
+    monitor = OrchestratorMonitor(replica_sampler=bad_sampler)
+    monitor.register_replica(0, 0)
+    monitor.note_loop(idle=True)
+    monitor.flush()
+
+
+def test_flush_is_idempotent(tmp_path, monkeypatch):
+    out_path = tmp_path / "orch_monitor.json"
+    monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path))
+
+    monitor = OrchestratorMonitor(replica_sampler=lambda: {})
+    monitor.register_replica(0, 0)
+    monitor.note_loop(idle=False)
+    monitor.flush()
+    monitor.flush()
+
+    payload = json.loads(out_path.read_text())
+    assert len(payload["windows"]["duration_s"]) == 1

tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py

@@ -156,7 +156,7 @@ class SyntheticModel(nn.Module):
         )
 
         audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0))
-        mm_tokens = {"audio": audio_list}
+        mm_tokens = {"codes": {"audio": audio_list}}
         return fake_eos, mm_tokens
 
 
@@ -192,9 +192,9 @@ def _cfg_alpha(batch_size, value=1.2, device=DEVICE):
 
 
 def _unpack_audio_codes(result):
-    """Unpack (fake_eos, {"audio": [list of tensors]}) into (fake_eos, audio_codes)."""
+    """Unpack (fake_eos, {"codes": {"audio": [list of tensors]}}) into (fake_eos, audio_codes)."""
     fake_eos, mm_tokens = result
-    audio_list = mm_tokens["audio"]
+    audio_list = mm_tokens["codes"]["audio"]
     # Each element is (1, 1, 1+n_acoustic_codebook), concat along batch dim
     audio_codes = torch.cat(audio_list, dim=0).squeeze(1)  # (B, 1+n_acoustic)
     return fake_eos, audio_codes

tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py

@@ -23,9 +23,9 @@ def _make_stage(engine_outputs):
 
 
 def _make_output(audio_tensors):
-    """Build a single generator output with multimodal_output["audio"] = audio_tensors."""
+    """Build a single generator output with multimodal_output codes.audio."""
     return SimpleNamespace(
-        outputs=[SimpleNamespace(multimodal_output={"audio": audio_tensors})],
+        outputs=[SimpleNamespace(multimodal_output={"codes": {"audio": audio_tensors}})],
     )
 
 
@@ -92,7 +92,7 @@ def test_empty_chunk_when_not_finished():
 
     payload = generator2tokenizer_async_chunk(
         transfer_manager=transfer_manager,
-        multimodal_output={"audio": torch.zeros((0,))},
+        multimodal_output={"codes": {"audio": torch.zeros((0,))}},
         request=request,
     )
 
@@ -153,7 +153,7 @@ def test_normal_chunk_emission():
 
     # Feed 25 frames one by one
     for i in range(25):
-        multimodal_output = {"audio": torch.tensor([float(i)] * 4)}
+        multimodal_output = {"codes": {"audio": torch.tensor([float(i)] * 4)}}
         payload = generator2tokenizer_async_chunk(
             transfer_manager=transfer_manager,
             multimodal_output=multimodal_output,
@@ -181,7 +181,7 @@ def test_small_initial_chunks():
 
     # Feed 5 frames (should trigger emission because codec_chunk_frames_at_begin=5)
     for i in range(5):
-        multimodal_output = {"audio": torch.tensor([float(i + 1)] * 3)}
+        multimodal_output = {"codes": {"audio": torch.tensor([float(i + 1)] * 3)}}
         payload = generator2tokenizer_async_chunk(
             transfer_manager=transfer_manager,
             multimodal_output=multimodal_output,
@@ -208,7 +208,7 @@ def test_no_emission_between_boundaries():
 
     # Feed 3 frames (3 % 5 != 0, should not emit)
     for i in range(3):
-        multimodal_output = {"audio": torch.tensor([float(i)] * 4)}
+        multimodal_output = {"codes": {"audio": torch.tensor([float(i)] * 4)}}
         payload = generator2tokenizer_async_chunk(
             transfer_manager=transfer_manager,
             multimodal_output=multimodal_output,
@@ -230,7 +230,7 @@ def test_context_handling_format():
 
     # Feed 5 frames to trigger a chunk (chunk_size_at_begin=5)
     for i in range(5):
-        multimodal_output = {"audio": torch.tensor([float(i + 10)] * 2)}  # codebook_dim=2
+        multimodal_output = {"codes": {"audio": torch.tensor([float(i + 10)] * 2)}}  # codebook_dim=2
         payload = generator2tokenizer_async_chunk(
             transfer_manager=transfer_manager,
             multimodal_output=multimodal_output,

tests/utils/test_mm_outputs_partition.py

@@ -0,0 +1,44 @@
+import torch
+
+from vllm_omni.utils.mm_outputs import partition_flat_payload, partition_payload_list
+
+
+def test_partition_thinker_latent_payload():
+    payload = {
+        "hidden_states.layer_0": torch.zeros(2, 4),
+        "hidden_states.layer_24": torch.zeros(2, 8),
+        "embed.tts_bos": [torch.zeros(1, 1, 4)],
+    }
+    inter, client = partition_flat_payload(payload)
+    assert inter == payload
+    assert client == {}
+
+
+def test_partition_talker_intermediate_codes():
+    payload = {
+        "codes.audio": torch.zeros(3, 2),
+        "hidden": torch.zeros(3, 16),
+    }
+    inter, client = partition_flat_payload(payload)
+    assert inter == payload
+    assert client == {}
+
+
+def test_partition_code2wav_client_audio():
+    payload = {
+        "model_outputs": torch.zeros(1, 2400),
+        "sr": torch.tensor(24000, dtype=torch.int32),
+    }
+    inter, client = partition_flat_payload(payload)
+    assert inter == {}
+    assert client == payload
+
+
+def test_partition_payload_list_preserves_request_alignment():
+    payloads = [
+        {"hidden_states.layer_0": torch.zeros(1, 2)},
+        {"model_outputs": torch.zeros(1, 10)},
+    ]
+    inter_list, client_list = partition_payload_list(payloads)
+    assert inter_list == [payloads[0], None]
+    assert client_list == [None, payloads[1]]

tests/worker/test_gpu_ar_model_runner.py

@@ -115,6 +115,7 @@ def _make_async_output_runner(engine_output_type: str = "audio"):
     )
     runner.vllm_config = SimpleNamespace(model_config=model_config)
     runner.model_config = model_config
+    runner._async_chunk = True
     runner.omni_prefix_cache = None
     runner.requests = {"r1": object(), "r2": object()}
     runner.supports_mm_inputs = False
@@ -174,8 +175,10 @@ def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monke
     )
 
     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.inter_stage_outputs is not None
+    assert torch.equal(output.inter_stage_outputs[0]["hidden"], torch.tensor([[1.0]]))
+    assert torch.equal(output.inter_stage_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
+    assert output.multimodal_outputs is None
     assert output.kv_extracted_req_ids == ["r2"]
     assert output.omni_connector_output == "connector-output"
     assert events == ["accumulate:r1", "accumulate:r2", "connector"]
@@ -217,10 +220,11 @@ def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatc
         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] == {}
+    assert output.inter_stage_outputs is not None
+    assert output.inter_stage_outputs[0] is None
+    assert torch.equal(output.inter_stage_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]]))
+    assert output.inter_stage_outputs[2] is None
+    assert output.multimodal_outputs is None
 
 
 def test_process_additional_information_uses_snapshot_request_order(monkeypatch):
@@ -331,10 +335,11 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
         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))
+    assert output.inter_stage_outputs is not None
+    assert len(output.inter_stage_outputs) == 1
+    assert "hidden" not in output.inter_stage_outputs[0]
+    assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], torch.tensor([[7, 8], [9, 10]], dtype=torch.long))
+    assert output.multimodal_outputs is None
 
 
 def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
@@ -395,6 +400,7 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
     runner.vllm_config = SimpleNamespace(
         model_config=SimpleNamespace(engine_output_type="audio"),
     )
+    runner._async_chunk = False
 
     monkeypatch.setattr(
         GPUARModelRunner, "_sample", lambda self, logits, spec_decode_metadata: SimpleNamespace(sampled_token_ids=[])
@@ -428,6 +434,8 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
 
     output = GPUARModelRunner.sample_tokens(runner, grammar_output=None)
 
+    assert output.inter_stage_outputs is None
+    assert output.multimodal_outputs is not None
     assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]]))
     assert torch.equal(
         output.multimodal_outputs[1]["hidden"],
@@ -482,5 +490,6 @@ def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch
         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])
+    assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1])
+    assert torch.equal(output.inter_stage_outputs[1]["codes.audio"], codes[1:2])
+    assert output.multimodal_outputs is None

tests/worker/test_gpu_generation_model_runner.py

@@ -40,6 +40,7 @@ def _make_runner(multimodal_outputs):
     runner.supports_mm_inputs = False
     runner.speculative_config = None
     runner.routed_experts_initialized = False
+    runner._async_chunk = False
     return runner
 
 

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -289,6 +289,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         num_scheduled_tokens = scheduler_output.num_scheduled_tokens
         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)
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
         cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
@@ -391,6 +392,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             new_token_ids = generated_token_ids
             pooler_output = pooler_outputs[req_index] if pooler_outputs else None
             mm_output = mm_outputs[req_index] if mm_outputs else None
+            inter_stage_output = inter_stage_outputs[req_index] if inter_stage_outputs else None
             kv_transfer_params = None
             status_before_stop = request.status
             finish_reason = None
@@ -487,12 +489,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                         new_prompt_len_snapshot=self._new_prompt_len_snapshot.get(req_id, None),
                     )
                 )
-                if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.save_async(mm_output, request, is_segment_finished)
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
 
+            if self.chunk_transfer_adapter is not None and (inter_stage_output is not None or is_segment_finished):
+                self.chunk_transfer_adapter.save_async(
+                    inter_stage_output,
+                    request,
+                    is_segment_finished,
+                )
+
         # Remove the stopped requests from the running and waiting queues.
         if stopped_running_reqs:
             self.running = remove_all(self.running, stopped_running_reqs)

vllm_omni/engine/arg_utils.py

@@ -438,6 +438,7 @@ class OrchestratorArgs:
 
     # === Observability ===
     log_stats: bool = False
+    enable_orch_monitor: bool = False
 
     # === Headless Mode (also forwarded to engine — see SHARED_FIELDS) ===
     stage_id: int | None = None

vllm_omni/engine/async_omni_engine.py

@@ -204,6 +204,7 @@ class AsyncOmniEngine:
     _log_stats: bool = False
     _coordinator_runtime: Any = None
     _transfer_emitter: Any = None
+    _enable_orch_monitor: bool = False
 
     def __init__(
         self,
@@ -233,6 +234,7 @@ class AsyncOmniEngine:
         # replica) vllm:* wrap stays registered but reads zero. Respects the
         # --log-stats CLI flag set by the user via OmniBase.
         self._log_stats = log_stats
+        self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False))
 
         logger.info(f"[AsyncOmniEngine] Initializing with model {model}")
 
@@ -425,6 +427,7 @@ class AsyncOmniEngine:
                 running_counter=self._running_counter,
                 transfer_emitter=self._transfer_emitter,
                 log_stats=self._log_stats,
+                enable_orch_monitor=self._enable_orch_monitor,
             )
             if not startup_future.done():
                 startup_future.set_result(asyncio.get_running_loop())

vllm_omni/engine/orchestrator.py

@@ -44,6 +44,7 @@ from vllm_omni.engine.messages import (
     StageSubmissionMessage,
     UnregisterRemoteReplicaMessage,
 )
+from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_key
 from vllm_omni.engine.serialization import serialize_additional_information
 from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.metrics.prometheus import OmniRequestCounter
@@ -221,6 +222,7 @@ class Orchestrator:
         running_counter: OmniRequestCounter | None = None,
         transfer_emitter: Any = None,
         log_stats: bool = False,
+        enable_orch_monitor: bool = False,
     ) -> None:
         self.request_async_queue = request_async_queue
         self.output_async_queue = output_async_queue
@@ -229,6 +231,13 @@ class Orchestrator:
         self.async_chunk = bool(async_chunk)
         self.num_stages = len(stage_pools)
         self.stage_pools: list[StagePool] = stage_pools
+        self._orch_monitor = create_orch_monitor(
+            enabled=enable_orch_monitor,
+            replica_sampler=self._sample_replica_metrics,
+        )
+        for stage_id, pool in enumerate(self.stage_pools):
+            for replica_id in pool.live_replica_ids():
+                self._orch_monitor.register_replica(stage_id, replica_id)
 
         # PD disaggregation state
         self._pd_pair: tuple[int, int] | None = None
@@ -372,6 +381,7 @@ class Orchestrator:
                 await self._membership.drain_tasks(timeout=10.0)
                 self._membership.shutdown()
 
+            self._orch_monitor.flush()
             self._shutdown_stages()
 
             loop = asyncio.get_running_loop()
@@ -402,6 +412,7 @@ class Orchestrator:
             elif isinstance(msg, RegisterRemoteReplicaMessage):
                 if self._membership is not None:
                     await self._membership.handle_register(msg.stage_id, msg.replica_id)
+                    self._orch_monitor.register_replica(msg.stage_id, msg.replica_id)
             elif isinstance(msg, UnregisterRemoteReplicaMessage):
                 if self._membership is not None:
                     await self._membership.handle_unregister(msg.stage_id, msg.input_addr)
@@ -629,6 +640,14 @@ class Orchestrator:
 
     # ---- Orchestration loop ----
 
+    def _sample_replica_metrics(self) -> dict[str, tuple[int, int]]:
+        samples: dict[str, tuple[int, int]] = {}
+        for stage_id, pool in enumerate(self.stage_pools):
+            for replica_id in pool.live_replica_ids():
+                key = replica_key(stage_id, replica_id)
+                samples[key] = pool.replica_monitor_sample(replica_id)
+        return samples
+
     async def _orchestration_output_handler(self) -> None:
         """Poll all stages, handle transfers, send final outputs to main."""
         try:
@@ -737,6 +756,7 @@ class Orchestrator:
                         await self._handle_processed_outputs(stage_id, replica_id, raw_output)
                         idle = False
 
+            self._orch_monitor.note_loop(idle=idle)
             if idle:
                 await asyncio.sleep(0.001)
             else:

vllm_omni/engine/orchestrator_monitor.py

@@ -0,0 +1,193 @@
+"""Optional orchestrator window monitor (diagnostic only).
+
+Enable with:
+  --enable-orch-monitor
+
+Optional output path override:
+  export VLLM_OMNI_ORCH_MONITOR_PATH=/path/to/monitor.json
+
+Each 1s window records:
+  - duration_s, loop_idle, loop_active (orchestrator poll-loop counts)
+  - per-replica outputs_queue_size (MP client outputs_queue backlog)
+  - per-replica inflight (requests currently bound/routed to the replica)
+
+See docs/contributing/profiling.md (Orchestrator Monitor).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from collections.abc import Callable
+from pathlib import Path
+from typing import Protocol, TypedDict
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_DEFAULT_OUTPUT_PREFIX = "vllm_omni_orch_monitor"
+_WINDOW_S = 1.0
+
+ReplicaSample = tuple[int, int]
+ReplicaSampler = Callable[[], dict[str, ReplicaSample]]
+
+
+class _ReplicaSeries(TypedDict):
+    outputs_queue_size: list[int]
+    inflight: list[int]
+
+
+def replica_key(stage_id: int, replica_id: int) -> str:
+    return f"stage={stage_id},replica={replica_id}"
+
+
+def _resolve_output_path() -> str:
+    env_path = os.environ.get("VLLM_OMNI_ORCH_MONITOR_PATH")
+    if env_path:
+        return env_path
+    timestamp = time.strftime("%m%d%H%M", time.localtime())
+    return str(Path.cwd() / f"{_DEFAULT_OUTPUT_PREFIX}_{timestamp}.json")
+
+
+class OrchestratorMonitorBase(Protocol):
+    def register_replica(self, stage_id: int, replica_id: int) -> None: ...
+
+    def note_loop(self, *, idle: bool) -> None: ...
+
+    def flush(self) -> None: ...
+
+
+class _NullOrchestratorMonitor:
+    def register_replica(self, stage_id: int, replica_id: int) -> None:
+        return
+
+    def note_loop(self, *, idle: bool) -> None:
+        return
+
+    def flush(self) -> None:
+        return
+
+
+_NULL_ORCH_MONITOR = _NullOrchestratorMonitor()
+
+
+def create_orch_monitor(
+    *,
+    enabled: bool,
+    replica_sampler: ReplicaSampler,
+) -> OrchestratorMonitorBase:
+    if not enabled:
+        return _NULL_ORCH_MONITOR
+    return OrchestratorMonitor(replica_sampler=replica_sampler)
+
+
+class OrchestratorMonitor:
+    def __init__(self, *, replica_sampler: ReplicaSampler) -> None:
+        self._output_path = _resolve_output_path()
+        self._flushed = False
+        self._replica_sampler = replica_sampler
+        self._window_start_mono = time.monotonic()
+        self._loop_idle = 0
+        self._loop_active = 0
+        self._duration_s: list[float] = []
+        self._loop_idle_buf: list[int] = []
+        self._loop_active_buf: list[int] = []
+        self._replicas: dict[str, _ReplicaSeries] = {}
+
+    def register_replica(self, stage_id: int, replica_id: int) -> None:
+        self._ensure_replica(replica_key(stage_id, replica_id))
+
+    def note_loop(self, *, idle: bool) -> None:
+        now_mono = time.monotonic()
+        self._roll_window_if_needed(now_mono)
+        if idle:
+            self._loop_idle += 1
+        else:
+            self._loop_active += 1
+
+    def flush(self) -> None:
+        if self._flushed:
+            return
+        self._roll_window(time.monotonic())
+        self._flushed = True
+        payload = {
+            "configured_window_s": _WINDOW_S,
+            "windows": {
+                "duration_s": self._duration_s,
+                "loop_idle": self._loop_idle_buf,
+                "loop_active": self._loop_active_buf,
+            },
+            "replicas": dict(sorted(self._replicas.items())),
+        }
+        try:
+            output_path = Path(self._output_path).expanduser()
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            with open(output_path, "w", encoding="utf-8") as f:
+                json.dump(payload, f, sort_keys=True)
+                f.write("\n")
+            self._log_summary(str(output_path))
+        except OSError:
+            logger.exception("[OrchestratorMonitor] Failed to write log to %s", self._output_path)
+
+    def _ensure_replica(self, key: str) -> None:
+        if key in self._replicas:
+            return
+        # Backfill zeros for windows recorded before this replica appeared.
+        n = len(self._duration_s)
+        self._replicas[key] = {
+            "outputs_queue_size": [0] * n,
+            "inflight": [0] * n,
+        }
+
+    def _roll_window_if_needed(self, now_mono: float) -> None:
+        if now_mono - self._window_start_mono >= _WINDOW_S:
+            self._roll_window(now_mono)
+
+    def _roll_window(self, now_mono: float) -> None:
+        try:
+            sampled = self._replica_sampler()
+        except Exception:
+            logger.exception("[OrchestratorMonitor] replica sampling failed")
+            sampled = {}
+        self._duration_s.append(max(now_mono - self._window_start_mono, 1e-9))
+        self._loop_idle_buf.append(self._loop_idle)
+        self._loop_active_buf.append(self._loop_active)
+        for key, series in self._replicas.items():
+            outputs_qsize, inflight = sampled.get(key, (0, 0))
+            series["outputs_queue_size"].append(outputs_qsize)
+            series["inflight"].append(inflight)
+
+        self._loop_idle = 0
+        self._loop_active = 0
+        self._window_start_mono = now_mono
+
+    def _log_summary(self, output_path: str) -> None:
+        loop_idle = sum(self._loop_idle_buf)
+        loop_active = sum(self._loop_active_buf)
+        loop_total = loop_idle + loop_active
+        loop_active_pct = (float(loop_active) / float(loop_total) * 100.0) if loop_total else 0.0
+        logger.info(
+            "[OrchestratorMonitor] wrote %s windows=%d duration=%.3fs loop_active=%.1f%%",
+            output_path,
+            len(self._duration_s),
+            sum(self._duration_s),
+            loop_active_pct,
+        )
+        for key, series in sorted(self._replicas.items()):
+            queue_values = series["outputs_queue_size"]
+            inflight_values = series["inflight"]
+            queue_avg = sum(queue_values) / float(len(queue_values)) if queue_values else 0.0
+            queue_max = max(queue_values) if queue_values else 0
+            inflight_avg = sum(inflight_values) / float(len(inflight_values)) if inflight_values else 0.0
+            inflight_max = max(inflight_values) if inflight_values else 0
+            logger.info(
+                "[OrchestratorMonitor] %s outputs_queue_avg=%.2f outputs_queue_max=%d "
+                "inflight_avg=%.2f inflight_max=%d",
+                key,
+                queue_avg,
+                queue_max,
+                inflight_avg,
+                inflight_max,
+            )

vllm_omni/engine/stage_pool.py

@@ -429,6 +429,37 @@ class StagePool:
             return None
         return cast(StagePoolLLMClient, client)
 
+    def replica_monitor_sample(self, replica_id: int) -> tuple[int, int]:
+        """Return (outputs_queue_size, inflight) for orchestrator load diagnostics."""
+        if replica_id < 0 or replica_id >= len(self.clients):
+            return (0, 0)
+        client = self.clients[replica_id]
+        if client is None:
+            return (0, 0)
+
+        outputs_qsize = 0
+        outputs_queue = getattr(client, "outputs_queue", None)
+        if outputs_queue is not None:
+            try:
+                outputs_qsize = max(int(outputs_queue.qsize()), 0)
+            except (AttributeError, NotImplementedError, RuntimeError, ValueError):
+                pass
+
+        if self._hub is not None:
+            input_addr = self._client_input_addr(client)
+            if input_addr is not None:
+                for replica in self._hub.get_replicas_for_stage(self.stage_id).replicas:
+                    if replica.input_addr == input_addr:
+                        return (outputs_qsize, max(int(replica.queue_length), 0))
+            # Distributed dispatch (pick): request_id -> input_addr in _affinity.
+            inflight = sum(
+                1 for input_addr in self._affinity.values() if self._addr_to_replica_id.get(input_addr) == replica_id
+            )
+        else:
+            # Legacy local dispatch (select_replica_id): request_id -> replica_id.
+            inflight = sum(1 for bound_replica_id in self._request_bindings.values() if bound_replica_id == replica_id)
+        return (outputs_qsize, inflight)
+
     def release_binding(self, request_id: str) -> None:
         """Drop the route binding for *request_id* in this stage."""
         self._request_bindings.pop(request_id, None)

vllm_omni/entrypoints/cli/serve.py

@@ -710,6 +710,11 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable AR stage profiler to include AR stage timing in stage_durations.",
         )
+        omni_config_group.add_argument(
+            "--enable-orch-monitor",
+            action="store_true",
+            help="Enable orchestrator window monitor and write a JSON log at shutdown.",
+        )
 
         # Supplementary auxiliary text encoder parameters
         # (e.g., the meta llama/meta llama-3.1-8b-instrument used by hidream)

vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py

@@ -262,8 +262,8 @@ class CUDAGraphAcousticTransformerWrapper:
         fake_eos = self.static_fake_eos[padded_size][:actual_size].clone()
         audio_codes = self.static_audio_codes[padded_size][:actual_size].clone()
 
-        # Package into expected format: (fake_eos, {"audio": [list of tensors]})
+        # Package into expected format: (fake_eos, {"codes": {"audio": [list of tensors]}})
         audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0))
-        mm_tokens = {"audio": audio_list}
+        mm_tokens = {"codes": {"audio": audio_list}}
 
         return fake_eos, mm_tokens

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py

@@ -197,6 +197,16 @@ class VoxtralTTSForConditionalGeneration(
             return self.model.sampler
         return Sampler()
 
+    @staticmethod
+    def _is_prefill_step(info_dict: Mapping[str, object], input_ids: torch.Tensor) -> bool:
+        is_prefill_raw = info_dict.get("_omni_is_prefill")
+        if isinstance(is_prefill_raw, bool):
+            return is_prefill_raw
+        try:
+            return int(info_dict["_omni_num_computed_tokens"]) < int(info_dict["_omni_prompt_len"])
+        except (KeyError, TypeError, ValueError):
+            return input_ids.shape[0] > 1
+
     def tts_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict | None):
         self.post_process_idx = 0
         audio_tokens = info_dict.pop("audio", None)
@@ -206,8 +216,9 @@ class VoxtralTTSForConditionalGeneration(
             if input_ids[0] == self._audio_token_id:
                 input_embeds = multimodal_embeddings[0]
             return input_ids, input_embeds, info_dict
-        voice = info_dict.pop("voice", None)
-        if voice is not None:
+        voice = info_dict.get("voice")
+        if voice is not None and self._is_prefill_step(info_dict, input_ids):
+            voice = info_dict.pop("voice")
             if isinstance(voice, list):
                 voice = voice[0]
             multimodal_embeddings = self.voice_to_embedding[voice].to(input_ids.device).clone().detach()
@@ -220,12 +231,15 @@ class VoxtralTTSForConditionalGeneration(
 
     def tts_postprocess(self, hidden_states: torch.Tensor, multimodal_outputs: object, **info_dict: object | None):
         update_dict = {}
-        if isinstance(multimodal_outputs, Mapping) and "audio" in multimodal_outputs:
-            assert self.post_process_idx < len(multimodal_outputs["audio"]), (
-                f"Expect {self.post_process_idx=} < {len(multimodal_outputs['audio'])=}"
-            )
-            update_dict["audio"] = multimodal_outputs["audio"][self.post_process_idx]
-            self.post_process_idx += 1
+        if isinstance(multimodal_outputs, Mapping):
+            codes = multimodal_outputs.get("codes")
+            if isinstance(codes, Mapping) and "audio" in codes:
+                audio_frames = codes["audio"]
+                assert self.post_process_idx < len(audio_frames), (
+                    f"Expect {self.post_process_idx=} < {len(audio_frames)=}"
+                )
+                update_dict["codes"] = {"audio": audio_frames[self.post_process_idx]}
+                self.post_process_idx += 1
         return update_dict
 
     def embed_input_ids(

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py

@@ -1060,7 +1060,7 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti
         # BxC -> Bx1xC since this is per-step
         # Make it a list for vllm-omni processing
         audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0))
-        mm_tokens = {"audio": audio_list}
+        mm_tokens = {"codes": {"audio": audio_list}}
 
         return fake_eos, mm_tokens
 

vllm_omni/model_executor/stage_input_processors/voxtral_tts.py

@@ -16,6 +16,21 @@ from vllm_omni.inputs.data import OmniTokensPrompt
 logger = init_logger(__name__)
 
 
+def _codec_audio_tensors(multimodal_output: OmniPayload | dict[str, Any]) -> list[torch.Tensor] | None:
+    """Inter-stage codec frames from ``codes.audio``."""
+    if not isinstance(multimodal_output, Mapping):
+        return None
+    codes = multimodal_output.get("codes")
+    if not isinstance(codes, Mapping) or "audio" not in codes:
+        return None
+    audio = codes["audio"]
+    if isinstance(audio, torch.Tensor):
+        return [audio]
+    if isinstance(audio, list) and audio:
+        return audio
+    return None
+
+
 def generator2tokenizer(
     source_outputs: list[Any],
     _prompt: OmniTokensPrompt | TextPrompt = None,
@@ -24,7 +39,10 @@ def generator2tokenizer(
     tokenizer_inputs = []
     for generator_output in source_outputs:
         output = generator_output.outputs[0]
-        audio_tokens = torch.cat(output.multimodal_output["audio"], dim=-1).flatten().detach().cpu().tolist()
+        audio_tensors = _codec_audio_tensors(output.multimodal_output)
+        if not audio_tensors:
+            raise ValueError("Missing codes.audio in generator multimodal_output")
+        audio_tokens = torch.cat(audio_tensors, dim=-1).flatten().detach().cpu().tolist()
         tokenizer_inputs.append(
             OmniTokensPrompt(
                 prompt_token_ids=audio_tokens,
@@ -37,10 +55,13 @@ def generator2tokenizer(
 
 
 def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None:
-    audio = multimodal_output.get("audio")
-    if not isinstance(audio, torch.Tensor) or audio.numel() == 0:
+    audio_tensors = _codec_audio_tensors(multimodal_output)
+    if not audio_tensors:
         return None
-    return audio.flatten()
+    frame = audio_tensors[-1]
+    if not isinstance(frame, torch.Tensor) or frame.numel() == 0:
+        return None
+    return frame.flatten()
 
 
 def generator2tokenizer_async_chunk(

vllm_omni/outputs.py

@@ -44,12 +44,15 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
     that may be produced by non-autoregressive stages.
 
     Attributes:
-        multimodal_outputs: Optional per-request list of multimodal output
-            dicts, indexed by req_index. Each element is a dict mapping
-            output keys to tensors/values (e.g., [{"audio": tensor}, ...]).
+        multimodal_outputs: Optional per-request list of client-facing multimodal
+            output dicts, indexed by req_index.
+        inter_stage_outputs: Optional per-request list of inter-stage payload dicts
+            for connector transport (``save_async`` / full_payload).  Not forwarded
+            to the orchestrator output processor.
     """
 
     multimodal_outputs: list[dict[str, object]] | None = None
+    inter_stage_outputs: list[dict[str, object]] | None = None
     # IDs of requests whose KV cache has been extracted from GPU/NPU to CPU.
     # The Scheduler can safely free the block tables for these requests.
     kv_extracted_req_ids: list[str] | None = None

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -45,7 +45,7 @@ from vllm_omni.data_entry_keys import flatten_payload
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
-from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element
+from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element
 
 
 def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]:
@@ -110,6 +110,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
         self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
         # Initialize KV cache manager (preserve vllm_config fallback behavior)
         self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
+        self._async_chunk = getattr(self.model_config, "async_chunk", False)
         self._downstream_payload_cache: dict[str, bool] = {}
 
     def _make_buffer(self, *size, dtype, numpy=True):
@@ -252,13 +253,21 @@ class NPUARModelRunner(OmniNPUModelRunner):
 
     def _build_multimodal_outputs(
         self,
-        per_req_payloads: list[dict[str, object]] | None,
-    ) -> list[dict[str, torch.Tensor]] | None:
+        per_req_payloads: list[dict[str, object] | None] | None,
+    ) -> list[dict[str, torch.Tensor] | None] | None:
         if self.vllm_config.model_config.engine_output_type == "text":
             return None
         if per_req_payloads is None:
             return None
-        return [_ensure_tensor_values(payload) if payload else {} for payload in per_req_payloads]
+        wire_payloads: list[dict[str, torch.Tensor] | None] = []
+        for payload in per_req_payloads:
+            if not payload:
+                wire_payloads.append(None)
+            else:
+                wire_payloads.append(_ensure_tensor_values(payload))
+        if all(item is None for item in wire_payloads):
+            return None
+        return wire_payloads
 
 
     def _request_final_stage_id(self, req_id: str) -> int | None:
@@ -1132,7 +1141,13 @@ class NPUARModelRunner(OmniNPUModelRunner):
                     payload.update(mm_payload)
                 pooler_output.append(flatten_payload(payload))
 
-        multimodal_outputs = self._build_multimodal_outputs(pooler_output)
+        pooler_output = pooler_output or []
+        if self._async_chunk:
+            pooler_inter, pooler_client = partition_payload_list(pooler_output)
+        else:
+            pooler_inter, pooler_client = None, pooler_output
+        inter_stage_outputs = self._build_multimodal_outputs(pooler_inter)
+        multimodal_outputs = self._build_multimodal_outputs(pooler_client)
         model_runner_output = OmniModelRunnerOutput(
             req_ids=req_ids_output_copy,
             req_id_to_index=req_id_to_index_output_copy,
@@ -1141,6 +1156,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
             prompt_logprobs_dict=prompt_logprobs_dict,
             pooler_output=None,
             multimodal_outputs=multimodal_outputs,
+            inter_stage_outputs=inter_stage_outputs,
             kv_connector_output=kv_connector_output,
             ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
             cudagraph_stats=cudagraph_stats,

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -35,11 +35,16 @@ from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState, _ensure_tensor_values
 from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
+from vllm_omni.utils.mm_outputs import partition_payload_list
 
 
 class NPUGenerationModelRunner(OmniNPUModelRunner):
     """Generation model runner for vLLM-omni on NPU (non-autoregressive)."""
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._async_chunk = getattr(self.model_config, "async_chunk", False)
+
     def _update_request_states(self, scheduler_output: SchedulerOutput):
         # remove requests
         for req_id in scheduler_output.finished_req_ids:
@@ -435,20 +440,20 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         #  -------------------------------------- Omni-new -------------------------------------------------
         # Build per-request multimodal_outputs list (dedicated channel).
         # pooler_output is no longer used for multimodal data.
-        multimodal_outputs: list[dict[str, object]] = []
+        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"
             )
             assert multimodal_outputs_raw.shape[0] == self.input_batch.num_reqs
             for i in range(self.input_batch.num_reqs):
-                multimodal_outputs.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()})
+                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"
             )
             for out in multimodal_outputs_raw:
-                multimodal_outputs.append(
+                per_req_payloads.append(
                     {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None}
                 )
         elif isinstance(multimodal_outputs_raw, Mapping):
@@ -467,9 +472,14 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                         mm_payload[key] = out.detach().to("cpu").contiguous()
                     else:
                         logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}")
-                multimodal_outputs.append(_ensure_tensor_values(mm_payload))
+                per_req_payloads.append(_ensure_tensor_values(mm_payload))
         else:
             raise RuntimeError("Unsupported diffusion output type")
+
+        if self._async_chunk:
+            inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads)
+        else:
+            inter_stage_outputs, multimodal_outputs = None, per_req_payloads
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
         req_ids_output_copy = self.input_batch.req_ids.copy()
         req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
@@ -486,6 +496,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             prompt_logprobs_dict={},
             pooler_output=None,
             multimodal_outputs=multimodal_outputs,
+            inter_stage_outputs=inter_stage_outputs,
             kv_connector_output=kv_connector_output,
             num_nans_in_logits={},
             cudagraph_stats=cudagraph_stats,

vllm_omni/utils/mm_outputs.py

@@ -9,6 +9,55 @@ from vllm.logger import init_logger
 
 logger = init_logger(__name__)
 
+# Flat payload keys partitioned at worker output into inter-stage connector
+# payloads vs client-facing multimodal outputs.  Only final output roots are
+# listed here; everything else remains available for stage-to-stage transport.
+_CLIENT_MM_ROOT_KEYS: frozenset[str] = frozenset(
+    {
+        "model_outputs",
+        "sr",
+        "audio",
+        "image",
+        "images",
+        "video",
+        "videos",
+        "trajectory_latents",
+        "latents",
+    }
+)
+
+
+def partition_flat_payload(
+    payload: Mapping[str, object],
+) -> tuple[dict[str, object], dict[str, object]]:
+    """Split a flattened per-request payload into inter-stage vs client mm dicts."""
+    if not payload:
+        return {}, {}
+    inter_stage: dict[str, object] = {}
+    client_mm: dict[str, object] = {}
+    for key, value in payload.items():
+        root = key.split(".", 1)[0]
+        if root in _CLIENT_MM_ROOT_KEYS:
+            client_mm[key] = value
+        else:
+            inter_stage[key] = value
+    return inter_stage, client_mm
+
+
+def partition_payload_list(
+    payloads: list[dict[str, object]],
+) -> tuple[list[dict[str, object] | None] | None, list[dict[str, object] | None] | None]:
+    inter_stage_list: list[dict[str, object] | None] = []
+    client_mm_list: list[dict[str, object] | None] = []
+    for payload in payloads:
+        inter_stage, client_mm = partition_flat_payload(payload)
+        inter_stage_list.append(inter_stage or None)
+        client_mm_list.append(client_mm or None)
+    return (
+        None if all(item is None for item in inter_stage_list) else inter_stage_list,
+        None if all(item is None for item in client_mm_list) else client_mm_list,
+    )
+
 
 def build_mm_cpu(multimodal_outputs: dict) -> dict[str, object]:
     """Pre-copies multimodal tensor to CPU once (not per-request) to avoid

vllm_omni/worker/gpu_ar_model_runner.py

@@ -42,7 +42,7 @@ from vllm.v1.worker.utils import is_residual_scattered_for_sp
 from vllm_omni.data_entry_keys import flatten_payload
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element
+from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
@@ -288,6 +288,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
         # Initialize KV cache manager (preserve vllm_config fallback behavior)
         self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
+        self._async_chunk = getattr(self.model_config, "async_chunk", False)
         # Worker-connector init is gated by a per-`model_arch` allowlist
         # (covers both producer-side and consumer-side runners for the
         # arches below).  Consumer-wait stages must be registered
@@ -1322,8 +1323,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
 
     def _build_multimodal_outputs(
         self,
-        per_req_payloads: list[dict[str, object]] | None,
-    ) -> list[dict[str, torch.Tensor]] | None:
+        per_req_payloads: list[dict[str, object] | None] | None,
+    ) -> list[dict[str, torch.Tensor] | None] | None:
         """Build per-request multimodal output payloads (dedicated channel).
 
         Reuses the per-request payloads assembled by the pooler-payload loop
@@ -1338,7 +1339,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             return None
         if per_req_payloads is None:
             return None
-        return [_ensure_tensor_values(payload) if payload else {} for payload in per_req_payloads]
+        wire_payloads: list[dict[str, torch.Tensor] | None] = []
+        for payload in per_req_payloads:
+            if not payload:
+                wire_payloads.append(None)
+            else:
+                wire_payloads.append(_ensure_tensor_values(payload))
+        if all(item is None for item in wire_payloads):
+            return None
+        return wire_payloads
 
     def _snapshot_query_start_loc_cpu(self) -> Any:
         query_start_loc_cpu = self.query_start_loc.cpu
@@ -1641,15 +1650,22 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     )
                     pooler_output.append(flatten_payload(payload))
 
-        if pooler_output and self._should_accumulate_full_payload_output():
+        pooler_output = pooler_output or []
+        if self._async_chunk:
+            pooler_inter, pooler_client = partition_payload_list(pooler_output)
+        else:
+            pooler_inter, pooler_client = None, pooler_output
+
+        if pooler_inter and self._should_accumulate_full_payload_output():
             with record_function_or_nullcontext("omni_output_builder:accumulate_full_payload_output"):
                 for i, rid in enumerate(req_ids_output_copy):
                     req_state = self.requests.get(rid)
-                    if req_state is not None and pooler_output[i]:
-                        self.accumulate_full_payload_output(rid, pooler_output[i], req_state)
+                    if req_state is not None and pooler_inter[i]:
+                        self.accumulate_full_payload_output(rid, pooler_inter[i], req_state)
 
         with record_function_or_nullcontext("omni_output_builder:build_multimodal_outputs"):
-            multimodal_outputs = self._build_multimodal_outputs(pooler_output)
+            inter_stage_outputs = self._build_multimodal_outputs(pooler_inter)
+            multimodal_outputs = self._build_multimodal_outputs(pooler_client)
 
         with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
             routed_experts_lists = None
@@ -1663,6 +1679,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 prompt_logprobs_dict=prompt_logprobs_dict,
                 pooler_output=None,
                 multimodal_outputs=multimodal_outputs,
+                inter_stage_outputs=inter_stage_outputs,
                 kv_connector_output=kv_connector_output,
                 ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
                 num_nans_in_logits=num_nans_in_logits,

vllm_omni/worker/gpu_generation_model_runner.py

@@ -37,6 +37,7 @@ from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 
 from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.utils.mm_outputs import partition_payload_list
 from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, _ensure_tensor_values
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
@@ -54,6 +55,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self._async_chunk = getattr(self.model_config, "async_chunk", False)
         # Mirrors the init allowlist in gpu_ar_model_runner.py.
         _OMNI_CONNECTOR_INIT_ARCHS = {
             "Qwen3OmniMoeForConditionalGeneration",
@@ -418,20 +420,20 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
 
         # Build per-request multimodal_outputs list (dedicated channel).
         # pooler_output is no longer used for multimodal data.
-        multimodal_outputs: list[dict[str, object]] = []
+        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"
             )
             assert multimodal_outputs_raw.shape[0] == self.input_batch.num_reqs
             for i in range(self.input_batch.num_reqs):
-                multimodal_outputs.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()})
+                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"
             )
             for out in multimodal_outputs_raw:
-                multimodal_outputs.append(
+                per_req_payloads.append(
                     {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None}
                 )
         elif isinstance(multimodal_outputs_raw, Mapping):
@@ -450,20 +452,26 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                         mm_payload[key] = out.detach().to("cpu").contiguous()
                     else:
                         logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}")
-                multimodal_outputs.append(_ensure_tensor_values(mm_payload))
+                per_req_payloads.append(_ensure_tensor_values(mm_payload))
         else:
             raise RuntimeError("Unsupported diffusion output type")
+
+        if self._async_chunk:
+            inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads)
+        else:
+            inter_stage_outputs, multimodal_outputs = None, per_req_payloads
+
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
         req_ids_output_copy = self.input_batch.req_ids.copy()
         req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
         routed_experts_lists = None
         if self.routed_experts_initialized:
             routed_experts_lists = self._omni_extract_routed_experts(scheduler_output)
-        if self._should_accumulate_full_payload_output():
+        if inter_stage_outputs and self._should_accumulate_full_payload_output():
             for i, rid in enumerate(req_ids_output_copy):
                 req_state = self.requests.get(rid)
-                if req_state is not None and multimodal_outputs[i]:
-                    self.accumulate_full_payload_output(rid, multimodal_outputs[i], req_state)
+                if req_state is not None and inter_stage_outputs[i]:
+                    self.accumulate_full_payload_output(rid, inter_stage_outputs[i], req_state)
 
         output = OmniModelRunnerOutput(
             req_ids=req_ids_output_copy,
@@ -473,6 +481,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             prompt_logprobs_dict={},
             pooler_output=None,
             multimodal_outputs=multimodal_outputs,
+            inter_stage_outputs=inter_stage_outputs,
             kv_connector_output=kv_connector_output,
             num_nans_in_logits={},
             cudagraph_stats=cudagraph_stats,