From 504698dec824b99ea42765a7c21823cb97eb1713 Mon Sep 17 00:00:00 2001 From: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com> Date: Mon, 29 Jun 2026 23:43:36 +0800 Subject: [PATCH] [Core] Performance fix for orchestrator bottleneck: separates inter-stage from client outputs (#4527) Signed-off-by: yinpeiqi Co-authored-by: Chenguang ZHENG <645327136@qq.com> Co-authored-by: Cursor Co-authored-by: Hongsheng Liu Co-authored-by: Gao Han --- docs/contributing/profiling.md | 51 +++++ tests/engine/test_orchestrator_monitor.py | 101 +++++++++ .../test_cuda_graph_acoustic_transformer.py | 6 +- .../test_voxtral_tts_async_chunk.py | 14 +- tests/utils/test_mm_outputs_partition.py | 44 ++++ tests/worker/test_gpu_ar_model_runner.py | 33 +-- .../test_gpu_generation_model_runner.py | 1 + vllm_omni/core/sched/omni_ar_scheduler.py | 11 +- vllm_omni/engine/arg_utils.py | 1 + vllm_omni/engine/async_omni_engine.py | 3 + vllm_omni/engine/orchestrator.py | 20 ++ vllm_omni/engine/orchestrator_monitor.py | 193 ++++++++++++++++++ vllm_omni/engine/stage_pool.py | 31 +++ vllm_omni/entrypoints/cli/serve.py | 5 + ...cuda_graph_acoustic_transformer_wrapper.py | 4 +- .../models/voxtral_tts/voxtral_tts.py | 30 ++- .../voxtral_tts_audio_generation.py | 2 +- .../stage_input_processors/voxtral_tts.py | 29 ++- vllm_omni/outputs.py | 9 +- .../npu/worker/npu_ar_model_runner.py | 26 ++- .../npu/worker/npu_generation_model_runner.py | 19 +- vllm_omni/utils/mm_outputs.py | 49 +++++ vllm_omni/worker/gpu_ar_model_runner.py | 33 ++- .../worker/gpu_generation_model_runner.py | 23 ++- 24 files changed, 672 insertions(+), 66 deletions(-) create mode 100644 tests/engine/test_orchestrator_monitor.py create mode 100644 tests/utils/test_mm_outputs_partition.py create mode 100644 vllm_omni/engine/orchestrator_monitor.py diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 400d2a3c0..b57de441e 100644 --- a/docs/contributing/profiling.md +++ b/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_.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..outputs_queue_size` | MP client `outputs_queue` backlog for that replica | +| `replicas..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. diff --git a/tests/engine/test_orchestrator_monitor.py b/tests/engine/test_orchestrator_monitor.py new file mode 100644 index 000000000..ee9126cf4 --- /dev/null +++ b/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 diff --git a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py index 3b645632e..cf40f9ad4 100644 --- a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py +++ b/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 diff --git a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py index 5235d7c68..3fd5675ce 100644 --- a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py +++ b/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, diff --git a/tests/utils/test_mm_outputs_partition.py b/tests/utils/test_mm_outputs_partition.py new file mode 100644 index 000000000..f9affa471 --- /dev/null +++ b/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]] diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 60872636d..3088a3dff 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/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 diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index be7fadcb6..0cd9aeb26 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/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 diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 984844bdb..b18d090de 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index acfc0a83f..0be99bc79 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 5b8d2430f..45bc7c025 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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()) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 41207102b..7139f7e06 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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: diff --git a/vllm_omni/engine/orchestrator_monitor.py b/vllm_omni/engine/orchestrator_monitor.py new file mode 100644 index 000000000..8be903a88 --- /dev/null +++ b/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, + ) diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 4e1f3d5a5..d93dbf33e 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/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) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index d4c0b462d..5ff58700e 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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) diff --git a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py index 491f185a4..5898136c9 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py +++ b/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 diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py index 8b6872cfe..f262ede61 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py +++ b/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( diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index ce005936c..370639fb0 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py index 2928a7cf3..d3c1219b7 100644 --- a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py +++ b/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( diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index b3a597365..2f2a6f2ee 100644 --- a/vllm_omni/outputs.py +++ b/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 diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index b0c120ca7..2ed47ee15 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/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, diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index 8cb6b86a7..5b5308f14 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/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, diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py index 656ced316..0ca0c2c3a 100644 --- a/vllm_omni/utils/mm_outputs.py +++ b/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 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 23e6a4aac..0bf5cfc01 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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, diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 07f394842..a8e96bf83 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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,