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

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 @@ def compute_mm_logits(
         )
 
         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