From c932ccacc8b033d4b2d33a897d531ddf70b8b832 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Tue, 16 Jun 2026 19:42:57 +0800
Subject: [PATCH 01/21] separate mm output and inter output

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 tests/utils/test_mm_outputs_partition.py      |  48 ++++++
 vllm_omni/core/sched/omni_ar_scheduler.py     |  19 ++-
 vllm_omni/engine/orchestrator_profiler.py     | 155 ++++++++++++++++++
 vllm_omni/outputs.py                          |   9 +-
 vllm_omni/utils/mm_outputs.py                 |  69 ++++++++
 vllm_omni/worker/gpu_ar_model_runner.py       |  56 ++++---
 .../worker/gpu_generation_model_runner.py     |  46 ++++--
 .../omni_connector_model_runner_mixin.py      |  28 +++-
 8 files changed, 377 insertions(+), 53 deletions(-)
 create mode 100644 tests/utils/test_mm_outputs_partition.py
 create mode 100644 vllm_omni/engine/orchestrator_profiler.py

diff --git a/tests/utils/test_mm_outputs_partition.py b/tests/utils/test_mm_outputs_partition.py
new file mode 100644
index 00000000000..627d1d905e0
--- /dev/null
+++ b/tests/utils/test_mm_outputs_partition.py
@@ -0,0 +1,48 @@
+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, engine_output_type="latent")
+    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, engine_output_type="latent")
+    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, engine_output_type="audio")
+    assert inter == {}
+    assert client == payload
+
+
+def test_partition_payload_list_preserves_length():
+    payloads = [
+        {"hidden_states.layer_0": torch.zeros(1, 2)},
+        {"model_outputs": torch.zeros(1, 10)},
+    ]
+    inter_list, client_list = partition_payload_list(payloads, engine_output_type="latent")
+    assert len(inter_list) == 2
+    assert len(client_list) == 2
+    assert "hidden_states.layer_0" in inter_list[0]
+    assert client_list[0] == {}
+    assert inter_list[1] == {}
+    assert "model_outputs" in client_list[1]
diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py
index 984844bdbae..9dbac0f8b9a 100644
--- a/vllm_omni/core/sched/omni_ar_scheduler.py
+++ b/vllm_omni/core/sched/omni_ar_scheduler.py
@@ -289,6 +289,7 @@ def update_from_output(
         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 @@ def update_from_output(
             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
@@ -465,7 +467,14 @@ def update_from_output(
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
+            if (
+                new_token_ids
+                or mm_output is not None
+                or inter_stage_output is not None
+                or pooler_output is not None
+                or kv_transfer_params
+                or stopped
+            ):
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
                     OmniEngineCoreOutput(
@@ -488,7 +497,13 @@ def update_from_output(
                     )
                 )
                 if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.save_async(mm_output, request, is_segment_finished)
+                    connector_mm_output = inter_stage_output if inter_stage_output is not None else mm_output
+                    if connector_mm_output is not None:
+                        self.chunk_transfer_adapter.save_async(
+                            connector_mm_output,
+                            request,
+                            is_segment_finished,
+                        )
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
diff --git a/vllm_omni/engine/orchestrator_profiler.py b/vllm_omni/engine/orchestrator_profiler.py
new file mode 100644
index 00000000000..8daaa3a4a95
--- /dev/null
+++ b/vllm_omni/engine/orchestrator_profiler.py
@@ -0,0 +1,155 @@
+"""Optional orchestrator phase profiler (diagnostic only).
+
+Enable with:
+  export VLLM_OMNI_ORCH_PROFILE=1
+  export VLLM_OMNI_ORCH_PROFILE_PATH=/path/to/orch_profile.jsonl
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import threading
+import time
+from collections import defaultdict
+from contextlib import contextmanager
+from typing import Any, Iterator
+
+_ENABLED = os.environ.get("VLLM_OMNI_ORCH_PROFILE", "").lower() in ("1", "true", "yes")
+_OUTPUT_PATH = os.environ.get(
+    "VLLM_OMNI_ORCH_PROFILE_PATH",
+    "/tmp/vllm_omni_orch_profile.jsonl",
+)
+_FLUSH_INTERVAL_S = float(os.environ.get("VLLM_OMNI_ORCH_PROFILE_FLUSH_S", "1.0"))
+
+
+class OrchestratorProfiler:
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._totals_s: dict[str, float] = defaultdict(float)
+        self._counts: dict[str, int] = defaultdict(int)
+        self._value_sums: dict[str, float] = defaultdict(float)
+        self._value_counts: dict[str, int] = defaultdict(int)
+        self._value_max: dict[str, float] = defaultdict(float)
+        self._value_last: dict[str, float] = {}
+        self._window_start = time.monotonic()
+        self._last_flush = self._window_start
+        self._loop_iterations = 0
+        self._loop_idle_iterations = 0
+
+    @property
+    def enabled(self) -> bool:
+        return _ENABLED
+
+    def record(self, phase: str, elapsed_s: float, **tags: Any) -> None:
+        if not _ENABLED or elapsed_s < 0:
+            return
+        key = phase
+        if tags:
+            tag_str = ",".join(f"{k}={v}" for k, v in sorted(tags.items()))
+            key = f"{phase}|{tag_str}"
+        with self._lock:
+            self._totals_s[key] += elapsed_s
+            self._counts[key] += 1
+            now = time.monotonic()
+            if now - self._last_flush >= _FLUSH_INTERVAL_S:
+                self._flush_locked(now)
+
+    def record_value(self, name: str, value: float, **tags: Any) -> None:
+        if not _ENABLED:
+            return
+        key = name
+        if tags:
+            tag_str = ",".join(f"{k}={v}" for k, v in sorted(tags.items()))
+            key = f"{name}|{tag_str}"
+        with self._lock:
+            self._value_sums[key] += float(value)
+            self._value_counts[key] += 1
+            self._value_max[key] = max(self._value_max[key], float(value))
+            self._value_last[key] = float(value)
+            now = time.monotonic()
+            if now - self._last_flush >= _FLUSH_INTERVAL_S:
+                self._flush_locked(now)
+
+    def note_loop(self, *, idle: bool) -> None:
+        if not _ENABLED:
+            return
+        with self._lock:
+            self._loop_iterations += 1
+            if idle:
+                self._loop_idle_iterations += 1
+
+    def _flush_locked(self, now: float) -> None:
+        window_s = max(now - self._window_start, 1e-9)
+        profiled_s = sum(self._totals_s.values()) or 1e-9
+        payload = {
+            "ts": time.time(),
+            "window_s": window_s,
+            "loop_iterations": self._loop_iterations,
+            "loop_idle_iterations": self._loop_idle_iterations,
+            "phases": {
+                key: {
+                    "total_ms": totals * 1000.0,
+                    "count": self._counts[key],
+                    "avg_ms": (totals / self._counts[key]) * 1000.0 if self._counts[key] else 0.0,
+                    "pct": (totals / profiled_s) * 100.0,
+                }
+                for key, totals in self._totals_s.items()
+            },
+            "values": {
+                key: {
+                    "count": self._value_counts[key],
+                    "avg": self._value_sums[key] / self._value_counts[key] if self._value_counts[key] else 0.0,
+                    "max": self._value_max[key],
+                    "last": self._value_last.get(key, 0.0),
+                }
+                for key in self._value_counts
+            },
+        }
+        try:
+            with open(_OUTPUT_PATH, "a", encoding="utf-8") as f:
+                f.write(json.dumps(payload, sort_keys=True) + "\n")
+        except OSError:
+            pass
+        self._totals_s.clear()
+        self._counts.clear()
+        self._value_sums.clear()
+        self._value_counts.clear()
+        self._value_max.clear()
+        self._value_last.clear()
+        self._loop_iterations = 0
+        self._loop_idle_iterations = 0
+        self._window_start = now
+        self._last_flush = now
+
+    def flush(self) -> None:
+        if not _ENABLED:
+            return
+        with self._lock:
+            self._flush_locked(time.monotonic())
+
+
+_profiler: OrchestratorProfiler | None = None
+_profiler_lock = threading.Lock()
+
+
+def get_orchestrator_profiler() -> OrchestratorProfiler:
+    global _profiler
+    if _profiler is None:
+        with _profiler_lock:
+            if _profiler is None:
+                _profiler = OrchestratorProfiler()
+    return _profiler
+
+
+@contextmanager
+def profile_phase(phase: str, **tags: Any) -> Iterator[None]:
+    prof = get_orchestrator_profiler()
+    if not prof.enabled:
+        yield
+        return
+    t0 = time.perf_counter()
+    try:
+        yield
+    finally:
+        prof.record(phase, time.perf_counter() - t0, **tags)
diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py
index 2dfa1699fc0..b9426addb84 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/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py
index 3009e21d4d8..3b09aef3aca 100644
--- a/vllm_omni/utils/mm_outputs.py
+++ b/vllm_omni/utils/mm_outputs.py
@@ -9,6 +9,75 @@
 
 logger = init_logger(__name__)
 
+# Flat payload keys (``hidden_states.layer_0``, ``codes.audio``, …) partitioned
+# at worker ``sample_tokens`` into inter-stage vs client-facing wire channels.
+_INTER_STAGE_ROOT_KEYS: frozenset[str] = frozenset(
+    {
+        "hidden_states",
+        "hidden",
+        "embed",
+        "codes",
+        "latent",
+        "ids",
+        "meta",
+    }
+)
+_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],
+    *,
+    engine_output_type: str | None = None,
+) -> 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 _INTER_STAGE_ROOT_KEYS:
+            inter_stage[key] = value
+        elif root in _CLIENT_MM_ROOT_KEYS:
+            client_mm[key] = value
+        elif engine_output_type == "latent":
+            inter_stage[key] = value
+        elif engine_output_type in ("audio", "image", "video", "images", "videos"):
+            client_mm[key] = value
+        else:
+            client_mm[key] = value
+    return inter_stage, client_mm
+
+
+def partition_payload_list(
+    payloads: list[dict[str, object]],
+    *,
+    engine_output_type: str | None = None,
+) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
+    inter_stage_list: list[dict[str, object]] = []
+    client_mm_list: list[dict[str, object]] = []
+    for payload in payloads:
+        inter_stage, client_mm = partition_flat_payload(
+            payload,
+            engine_output_type=engine_output_type,
+        )
+        inter_stage_list.append(inter_stage)
+        client_mm_list.append(client_mm)
+    return inter_stage_list, 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 d29af7314fc..8531227fd5f 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -39,11 +39,17 @@
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
 
 from vllm_omni.data_entry_keys import flatten_payload
+from vllm_omni.core.sched.omni_scheduling_coordinator import (
+    uses_full_payload_input_coordinator,
+)
 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
+from vllm_omni.worker.omni_connector_model_runner_mixin import (
+    OmniConnectorModelRunnerMixin,
+    should_init_worker_omni_connector,
+)
 
 logger = init_logger(__name__)
 
@@ -115,23 +121,9 @@ def __init__(self, *args, **kwargs):
         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)
-        # 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
-        # separately as `(model_arch, model_stage)` tuples in
-        # `omni_scheduling_coordinator._FULL_PAYLOAD_INPUT_STAGES`;
-        # forgetting that produces a Stage-1 hang on the consumer.
-        _OMNI_CONNECTOR_INIT_ARCHS = {
-            "Qwen3OmniMoeForConditionalGeneration",
-            "Qwen2_5OmniForConditionalGeneration",
-            "CovoAudioForConditionalGeneration",
-            "MiMoAudioModel",
-            "Qwen3TTSTalkerForConditionalGeneration",
-            "Qwen3TTSCode2Wav",
-            "CosyVoice3Model",
-            "DyninOmniForConditionalGeneration",
-        }
-        if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
+        if should_init_worker_omni_connector(self.model_config) or uses_full_payload_input_coordinator(
+            self.model_config
+        ):
             self.init_omni_connectors(
                 vllm_config=self.vllm_config,
                 model_config=self.model_config,
@@ -964,7 +956,15 @@ def _build_multimodal_outputs(
             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: list[dict[str, torch.Tensor] | None] = []
+        for payload in per_req_payloads:
+            if not payload:
+                wire.append(None)
+            else:
+                wire.append(_ensure_tensor_values(payload))
+        if all(item is None for item in wire):
+            return None
+        return wire
 
     @torch.inference_mode()
     def sample_tokens(
@@ -1280,13 +1280,20 @@ def _unwrap_lists(v):
                     payload.update(mm_payload)
                 pooler_output.append(flatten_payload(payload))
 
-        if pooler_output and self._should_accumulate_full_payload_output():
+        engine_output_type = getattr(self.vllm_config.model_config, "engine_output_type", None)
+        pooler_inter, pooler_client = partition_payload_list(
+            pooler_output or [],
+            engine_output_type=engine_output_type,
+        )
+
+        if pooler_inter 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 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)
 
-        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
             if self.routed_experts_initialized:
@@ -1299,6 +1306,7 @@ def _unwrap_lists(v):
                 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 bc3d6ea2d82..1e9cf287a4b 100644
--- a/vllm_omni/worker/gpu_generation_model_runner.py
+++ b/vllm_omni/worker/gpu_generation_model_runner.py
@@ -37,9 +37,16 @@
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 
 from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.core.sched.omni_scheduling_coordinator import (
+    uses_full_payload_input_coordinator,
+)
+from vllm_omni.utils.mm_outputs import partition_flat_payload
 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
+from vllm_omni.worker.omni_connector_model_runner_mixin import (
+    OmniConnectorModelRunnerMixin,
+    should_init_worker_omni_connector,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -54,18 +61,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        # Mirrors the init allowlist in gpu_ar_model_runner.py.
-        _OMNI_CONNECTOR_INIT_ARCHS = {
-            "Qwen3OmniMoeForConditionalGeneration",
-            "Qwen2_5OmniForConditionalGeneration",
-            "CovoAudioForConditionalGeneration",
-            "MiMoAudioModel",
-            "Qwen3TTSTalkerForConditionalGeneration",
-            "Qwen3TTSCode2Wav",
-            "CosyVoice3Model",
-            "DyninOmniForConditionalGeneration",
-        }
-        if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
+        if should_init_worker_omni_connector(self.model_config) or uses_full_payload_input_coordinator(
+            self.model_config
+        ):
             self.init_omni_connectors(
                 vllm_config=self.vllm_config,
                 model_config=self.model_config,
@@ -452,6 +450,18 @@ def sample_tokens(
                 multimodal_outputs.append(_ensure_tensor_values(mm_payload))
         else:
             raise RuntimeError("Unsupported diffusion output type")
+
+        engine_output_type = getattr(self.vllm_config.model_config, "engine_output_type", None)
+        inter_stage_outputs: list[dict[str, object] | None] = []
+        client_mm_outputs: list[dict[str, object] | None] = []
+        for mm_payload in multimodal_outputs:
+            inter_payload, client_payload = partition_flat_payload(
+                mm_payload,
+                engine_output_type=engine_output_type,
+            )
+            inter_stage_outputs.append(inter_payload or None)
+            client_mm_outputs.append(client_payload or None)
+
         # [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()
@@ -461,8 +471,11 @@ def sample_tokens(
         if 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)
+
+        wire_inter_stage = None if all(x is None for x in inter_stage_outputs) else inter_stage_outputs
+        wire_client_mm = None if all(x is None for x in client_mm_outputs) else client_mm_outputs
 
         output = OmniModelRunnerOutput(
             req_ids=req_ids_output_copy,
@@ -471,7 +484,8 @@ def sample_tokens(
             logprobs=None,
             prompt_logprobs_dict={},
             pooler_output=None,
-            multimodal_outputs=multimodal_outputs,
+            multimodal_outputs=wire_client_mm,
+            inter_stage_outputs=wire_inter_stage,
             kv_connector_output=kv_connector_output,
             num_nans_in_logits={},
             cudagraph_stats=cudagraph_stats,
diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py
index 205a28c6d02..06a2ccfe2da 100644
--- a/vllm_omni/worker/omni_connector_model_runner_mixin.py
+++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py
@@ -46,6 +46,25 @@
 logger = init_logger(__name__)
 
 
+def is_declared_full_payload_producer_stage(model_config: Any) -> bool:
+    """Stage declares a non-chunked full-payload producer hook in the pipeline."""
+    if getattr(model_config, "async_chunk", False):
+        return False
+    if getattr(model_config, "final_output", False):
+        return False
+    next_stage_proc = getattr(model_config, "custom_process_next_stage_input_func", None)
+    if not isinstance(next_stage_proc, str) or not next_stage_proc:
+        return False
+    return getattr(model_config, "model_stage", None) is not None
+
+
+def should_init_worker_omni_connector(model_config: Any) -> bool:
+    """Whether the worker should initialize ``init_omni_connectors`` for this stage."""
+    if getattr(model_config, "async_chunk", False):
+        return True
+    return is_declared_full_payload_producer_stage(model_config)
+
+
 def should_accumulate_full_payload_output(model_config, custom_process_func) -> bool:
     """Producer-side structural gate.
 
@@ -57,14 +76,7 @@ def should_accumulate_full_payload_output(model_config, custom_process_func) ->
     """
     if custom_process_func is None:
         return False
-    if getattr(model_config, "async_chunk", False):
-        return False
-    if getattr(model_config, "final_output", False):
-        return False
-    next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None)
-    if not isinstance(next_stage_func, str) or not next_stage_func:
-        return False
-    return getattr(model_config, "model_stage", None) is not None
+    return is_declared_full_payload_producer_stage(model_config)
 
 
 class OmniConnectorModelRunnerMixin:

From dacc1cc226cec9590a0b21a0b0a58158e01337fa Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Wed, 17 Jun 2026 17:18:28 +0800
Subject: [PATCH 02/21] revert redundant modify

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 tests/utils/test_mm_outputs_partition.py      | 15 +------
 vllm_omni/core/sched/omni_ar_scheduler.py     | 24 ++++-------
 vllm_omni/utils/mm_outputs.py                 | 18 --------
 vllm_omni/worker/gpu_ar_model_runner.py       | 43 ++++++++++++-------
 .../worker/gpu_generation_model_runner.py     | 23 +++++-----
 .../omni_connector_model_runner_mixin.py      | 28 ++++--------
 6 files changed, 58 insertions(+), 93 deletions(-)

diff --git a/tests/utils/test_mm_outputs_partition.py b/tests/utils/test_mm_outputs_partition.py
index 627d1d905e0..9548b177ca2 100644
--- a/tests/utils/test_mm_outputs_partition.py
+++ b/tests/utils/test_mm_outputs_partition.py
@@ -1,6 +1,6 @@
 import torch
 
-from vllm_omni.utils.mm_outputs import partition_flat_payload, partition_payload_list
+from vllm_omni.utils.mm_outputs import partition_flat_payload
 
 
 def test_partition_thinker_latent_payload():
@@ -33,16 +33,3 @@ def test_partition_code2wav_client_audio():
     assert inter == {}
     assert client == payload
 
-
-def test_partition_payload_list_preserves_length():
-    payloads = [
-        {"hidden_states.layer_0": torch.zeros(1, 2)},
-        {"model_outputs": torch.zeros(1, 10)},
-    ]
-    inter_list, client_list = partition_payload_list(payloads, engine_output_type="latent")
-    assert len(inter_list) == 2
-    assert len(client_list) == 2
-    assert "hidden_states.layer_0" in inter_list[0]
-    assert client_list[0] == {}
-    assert inter_list[1] == {}
-    assert "model_outputs" in client_list[1]
diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py
index 9dbac0f8b9a..6645e08ed47 100644
--- a/vllm_omni/core/sched/omni_ar_scheduler.py
+++ b/vllm_omni/core/sched/omni_ar_scheduler.py
@@ -467,14 +467,7 @@ def update_from_output(
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if (
-                new_token_ids
-                or mm_output is not None
-                or inter_stage_output is not None
-                or pooler_output is not None
-                or kv_transfer_params
-                or stopped
-            ):
+            if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
                     OmniEngineCoreOutput(
@@ -496,18 +489,17 @@ def update_from_output(
                         new_prompt_len_snapshot=self._new_prompt_len_snapshot.get(req_id, None),
                     )
                 )
-                if self.chunk_transfer_adapter is not None:
-                    connector_mm_output = inter_stage_output if inter_stage_output is not None else mm_output
-                    if connector_mm_output is not None:
-                        self.chunk_transfer_adapter.save_async(
-                            connector_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:
+                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/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py
index 3b09aef3aca..7f4f229a12a 100644
--- a/vllm_omni/utils/mm_outputs.py
+++ b/vllm_omni/utils/mm_outputs.py
@@ -61,24 +61,6 @@ def partition_flat_payload(
             client_mm[key] = value
     return inter_stage, client_mm
 
-
-def partition_payload_list(
-    payloads: list[dict[str, object]],
-    *,
-    engine_output_type: str | None = None,
-) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
-    inter_stage_list: list[dict[str, object]] = []
-    client_mm_list: list[dict[str, object]] = []
-    for payload in payloads:
-        inter_stage, client_mm = partition_flat_payload(
-            payload,
-            engine_output_type=engine_output_type,
-        )
-        inter_stage_list.append(inter_stage)
-        client_mm_list.append(client_mm)
-    return inter_stage_list, 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
     redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU.
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index 8531227fd5f..35748477f83 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -39,17 +39,11 @@
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
 
 from vllm_omni.data_entry_keys import flatten_payload
-from vllm_omni.core.sched.omni_scheduling_coordinator import (
-    uses_full_payload_input_coordinator,
-)
 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, partition_payload_list, to_payload_element
+from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_flat_payload, to_payload_element
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
-from vllm_omni.worker.omni_connector_model_runner_mixin import (
-    OmniConnectorModelRunnerMixin,
-    should_init_worker_omni_connector,
-)
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 logger = init_logger(__name__)
 
@@ -121,9 +115,23 @@ def __init__(self, *args, **kwargs):
         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)
-        if should_init_worker_omni_connector(self.model_config) or uses_full_payload_input_coordinator(
-            self.model_config
-        ):
+        # 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
+        # separately as `(model_arch, model_stage)` tuples in
+        # `omni_scheduling_coordinator._FULL_PAYLOAD_INPUT_STAGES`;
+        # forgetting that produces a Stage-1 hang on the consumer.
+        _OMNI_CONNECTOR_INIT_ARCHS = {
+            "Qwen3OmniMoeForConditionalGeneration",
+            "Qwen2_5OmniForConditionalGeneration",
+            "CovoAudioForConditionalGeneration",
+            "MiMoAudioModel",
+            "Qwen3TTSTalkerForConditionalGeneration",
+            "Qwen3TTSCode2Wav",
+            "CosyVoice3Model",
+            "DyninOmniForConditionalGeneration",
+        }
+        if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
             self.init_omni_connectors(
                 vllm_config=self.vllm_config,
                 model_config=self.model_config,
@@ -1281,10 +1289,15 @@ def _unwrap_lists(v):
                 pooler_output.append(flatten_payload(payload))
 
         engine_output_type = getattr(self.vllm_config.model_config, "engine_output_type", None)
-        pooler_inter, pooler_client = partition_payload_list(
-            pooler_output or [],
-            engine_output_type=engine_output_type,
-        )
+        pooler_inter: list[dict[str, object]] = []
+        pooler_client: list[dict[str, object]] = []
+        for payload in pooler_output or []:
+            inter_payload, client_payload = partition_flat_payload(
+                payload,
+                engine_output_type=engine_output_type,
+            )
+            pooler_inter.append(inter_payload)
+            pooler_client.append(client_payload)
 
         if pooler_inter and self._should_accumulate_full_payload_output():
             for i, rid in enumerate(req_ids_output_copy):
diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py
index 1e9cf287a4b..1b307d798a3 100644
--- a/vllm_omni/worker/gpu_generation_model_runner.py
+++ b/vllm_omni/worker/gpu_generation_model_runner.py
@@ -37,16 +37,10 @@
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.core.sched.omni_scheduling_coordinator import (
-    uses_full_payload_input_coordinator,
-)
 from vllm_omni.utils.mm_outputs import partition_flat_payload
 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,
-    should_init_worker_omni_connector,
-)
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 logger = logging.getLogger(__name__)
 
@@ -61,9 +55,18 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        if should_init_worker_omni_connector(self.model_config) or uses_full_payload_input_coordinator(
-            self.model_config
-        ):
+        # Mirrors the init allowlist in gpu_ar_model_runner.py.
+        _OMNI_CONNECTOR_INIT_ARCHS = {
+            "Qwen3OmniMoeForConditionalGeneration",
+            "Qwen2_5OmniForConditionalGeneration",
+            "CovoAudioForConditionalGeneration",
+            "MiMoAudioModel",
+            "Qwen3TTSTalkerForConditionalGeneration",
+            "Qwen3TTSCode2Wav",
+            "CosyVoice3Model",
+            "DyninOmniForConditionalGeneration",
+        }
+        if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS:
             self.init_omni_connectors(
                 vllm_config=self.vllm_config,
                 model_config=self.model_config,
diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py
index 06a2ccfe2da..205a28c6d02 100644
--- a/vllm_omni/worker/omni_connector_model_runner_mixin.py
+++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py
@@ -46,25 +46,6 @@
 logger = init_logger(__name__)
 
 
-def is_declared_full_payload_producer_stage(model_config: Any) -> bool:
-    """Stage declares a non-chunked full-payload producer hook in the pipeline."""
-    if getattr(model_config, "async_chunk", False):
-        return False
-    if getattr(model_config, "final_output", False):
-        return False
-    next_stage_proc = getattr(model_config, "custom_process_next_stage_input_func", None)
-    if not isinstance(next_stage_proc, str) or not next_stage_proc:
-        return False
-    return getattr(model_config, "model_stage", None) is not None
-
-
-def should_init_worker_omni_connector(model_config: Any) -> bool:
-    """Whether the worker should initialize ``init_omni_connectors`` for this stage."""
-    if getattr(model_config, "async_chunk", False):
-        return True
-    return is_declared_full_payload_producer_stage(model_config)
-
-
 def should_accumulate_full_payload_output(model_config, custom_process_func) -> bool:
     """Producer-side structural gate.
 
@@ -76,7 +57,14 @@ def should_accumulate_full_payload_output(model_config, custom_process_func) ->
     """
     if custom_process_func is None:
         return False
-    return is_declared_full_payload_producer_stage(model_config)
+    if getattr(model_config, "async_chunk", False):
+        return False
+    if getattr(model_config, "final_output", False):
+        return False
+    next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None)
+    if not isinstance(next_stage_func, str) or not next_stage_func:
+        return False
+    return getattr(model_config, "model_stage", None) is not None
 
 
 class OmniConnectorModelRunnerMixin:

From 1a8553ee3ed44edfaba1ce388ba687055e0883fe Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Wed, 17 Jun 2026 17:26:52 +0800
Subject: [PATCH 03/21] fix branch

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 tests/utils/test_mm_outputs_partition.py      |  6 ++--
 vllm_omni/utils/mm_outputs.py                 | 28 ++++---------------
 vllm_omni/worker/gpu_ar_model_runner.py       |  6 +---
 .../worker/gpu_generation_model_runner.py     |  6 +---
 4 files changed, 10 insertions(+), 36 deletions(-)

diff --git a/tests/utils/test_mm_outputs_partition.py b/tests/utils/test_mm_outputs_partition.py
index 9548b177ca2..ca6011a3ad3 100644
--- a/tests/utils/test_mm_outputs_partition.py
+++ b/tests/utils/test_mm_outputs_partition.py
@@ -9,7 +9,7 @@ def test_partition_thinker_latent_payload():
         "hidden_states.layer_24": torch.zeros(2, 8),
         "embed.tts_bos": [torch.zeros(1, 1, 4)],
     }
-    inter, client = partition_flat_payload(payload, engine_output_type="latent")
+    inter, client = partition_flat_payload(payload)
     assert inter == payload
     assert client == {}
 
@@ -19,7 +19,7 @@ def test_partition_talker_intermediate_codes():
         "codes.audio": torch.zeros(3, 2),
         "hidden": torch.zeros(3, 16),
     }
-    inter, client = partition_flat_payload(payload, engine_output_type="latent")
+    inter, client = partition_flat_payload(payload)
     assert inter == payload
     assert client == {}
 
@@ -29,7 +29,7 @@ def test_partition_code2wav_client_audio():
         "model_outputs": torch.zeros(1, 2400),
         "sr": torch.tensor(24000, dtype=torch.int32),
     }
-    inter, client = partition_flat_payload(payload, engine_output_type="audio")
+    inter, client = partition_flat_payload(payload)
     assert inter == {}
     assert client == payload
 
diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py
index 7f4f229a12a..877796d2405 100644
--- a/vllm_omni/utils/mm_outputs.py
+++ b/vllm_omni/utils/mm_outputs.py
@@ -9,19 +9,9 @@
 
 logger = init_logger(__name__)
 
-# Flat payload keys (``hidden_states.layer_0``, ``codes.audio``, …) partitioned
-# at worker ``sample_tokens`` into inter-stage vs client-facing wire channels.
-_INTER_STAGE_ROOT_KEYS: frozenset[str] = frozenset(
-    {
-        "hidden_states",
-        "hidden",
-        "embed",
-        "codes",
-        "latent",
-        "ids",
-        "meta",
-    }
-)
+# 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",
@@ -39,8 +29,6 @@
 
 def partition_flat_payload(
     payload: Mapping[str, object],
-    *,
-    engine_output_type: str | None = None,
 ) -> tuple[dict[str, object], dict[str, object]]:
     """Split a flattened per-request payload into inter-stage vs client mm dicts."""
     if not payload:
@@ -49,16 +37,10 @@ def partition_flat_payload(
     client_mm: dict[str, object] = {}
     for key, value in payload.items():
         root = key.split(".", 1)[0]
-        if root in _INTER_STAGE_ROOT_KEYS:
-            inter_stage[key] = value
-        elif root in _CLIENT_MM_ROOT_KEYS:
-            client_mm[key] = value
-        elif engine_output_type == "latent":
-            inter_stage[key] = value
-        elif engine_output_type in ("audio", "image", "video", "images", "videos"):
+        if root in _CLIENT_MM_ROOT_KEYS:
             client_mm[key] = value
         else:
-            client_mm[key] = value
+            inter_stage[key] = value
     return inter_stage, client_mm
 
 def build_mm_cpu(multimodal_outputs: dict) -> dict[str, object]:
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index 35748477f83..906c7c0b14c 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -1288,14 +1288,10 @@ def _unwrap_lists(v):
                     payload.update(mm_payload)
                 pooler_output.append(flatten_payload(payload))
 
-        engine_output_type = getattr(self.vllm_config.model_config, "engine_output_type", None)
         pooler_inter: list[dict[str, object]] = []
         pooler_client: list[dict[str, object]] = []
         for payload in pooler_output or []:
-            inter_payload, client_payload = partition_flat_payload(
-                payload,
-                engine_output_type=engine_output_type,
-            )
+            inter_payload, client_payload = partition_flat_payload(payload)
             pooler_inter.append(inter_payload)
             pooler_client.append(client_payload)
 
diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py
index 1b307d798a3..e9a01d476df 100644
--- a/vllm_omni/worker/gpu_generation_model_runner.py
+++ b/vllm_omni/worker/gpu_generation_model_runner.py
@@ -454,14 +454,10 @@ def sample_tokens(
         else:
             raise RuntimeError("Unsupported diffusion output type")
 
-        engine_output_type = getattr(self.vllm_config.model_config, "engine_output_type", None)
         inter_stage_outputs: list[dict[str, object] | None] = []
         client_mm_outputs: list[dict[str, object] | None] = []
         for mm_payload in multimodal_outputs:
-            inter_payload, client_payload = partition_flat_payload(
-                mm_payload,
-                engine_output_type=engine_output_type,
-            )
+            inter_payload, client_payload = partition_flat_payload(mm_payload)
             inter_stage_outputs.append(inter_payload or None)
             client_mm_outputs.append(client_payload or None)
 

From 4ef63196f2fb8909b21cfeee391c8209e9d1f510 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Wed, 17 Jun 2026 17:37:02 +0800
Subject: [PATCH 04/21] update

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 tests/utils/test_mm_outputs_partition.py      | 12 ++++++++-
 vllm_omni/utils/mm_outputs.py                 | 16 ++++++++++++
 vllm_omni/worker/gpu_ar_model_runner.py       | 13 +++-------
 .../worker/gpu_generation_model_runner.py     | 26 +++++++------------
 4 files changed, 40 insertions(+), 27 deletions(-)

diff --git a/tests/utils/test_mm_outputs_partition.py b/tests/utils/test_mm_outputs_partition.py
index ca6011a3ad3..aa45eae2024 100644
--- a/tests/utils/test_mm_outputs_partition.py
+++ b/tests/utils/test_mm_outputs_partition.py
@@ -1,6 +1,6 @@
 import torch
 
-from vllm_omni.utils.mm_outputs import partition_flat_payload
+from vllm_omni.utils.mm_outputs import partition_flat_payload, partition_payload_list
 
 
 def test_partition_thinker_latent_payload():
@@ -33,3 +33,13 @@ def test_partition_code2wav_client_audio():
     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/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py
index 877796d2405..20cda344a06 100644
--- a/vllm_omni/utils/mm_outputs.py
+++ b/vllm_omni/utils/mm_outputs.py
@@ -43,6 +43,22 @@ def partition_flat_payload(
             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
     redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU.
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index 906c7c0b14c..a71032d3b77 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -41,7 +41,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.utils.mm_outputs import build_mm_cpu, partition_flat_payload, 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
 
@@ -948,8 +948,8 @@ def _resolve_req_hidden_states(
 
     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
@@ -1288,12 +1288,7 @@ def _unwrap_lists(v):
                     payload.update(mm_payload)
                 pooler_output.append(flatten_payload(payload))
 
-        pooler_inter: list[dict[str, object]] = []
-        pooler_client: list[dict[str, object]] = []
-        for payload in pooler_output or []:
-            inter_payload, client_payload = partition_flat_payload(payload)
-            pooler_inter.append(inter_payload)
-            pooler_client.append(client_payload)
+        pooler_inter, pooler_client = partition_payload_list(pooler_output or [])
 
         if pooler_inter and self._should_accumulate_full_payload_output():
             for i, rid in enumerate(req_ids_output_copy):
diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py
index e9a01d476df..68ecd0c5e40 100644
--- a/vllm_omni/worker/gpu_generation_model_runner.py
+++ b/vllm_omni/worker/gpu_generation_model_runner.py
@@ -37,7 +37,7 @@
 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_flat_payload
+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
@@ -418,20 +418,20 @@ def sample_tokens(
 
         # 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,16 +450,11 @@ def sample_tokens(
                         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")
 
-        inter_stage_outputs: list[dict[str, object] | None] = []
-        client_mm_outputs: list[dict[str, object] | None] = []
-        for mm_payload in multimodal_outputs:
-            inter_payload, client_payload = partition_flat_payload(mm_payload)
-            inter_stage_outputs.append(inter_payload or None)
-            client_mm_outputs.append(client_payload or None)
+        inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads)
 
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
         req_ids_output_copy = self.input_batch.req_ids.copy()
@@ -467,15 +462,12 @@ def sample_tokens(
         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 inter_stage_outputs[i]:
                     self.accumulate_full_payload_output(rid, inter_stage_outputs[i], req_state)
 
-        wire_inter_stage = None if all(x is None for x in inter_stage_outputs) else inter_stage_outputs
-        wire_client_mm = None if all(x is None for x in client_mm_outputs) else client_mm_outputs
-
         output = OmniModelRunnerOutput(
             req_ids=req_ids_output_copy,
             req_id_to_index=req_id_to_index_output_copy,
@@ -483,8 +475,8 @@ def sample_tokens(
             logprobs=None,
             prompt_logprobs_dict={},
             pooler_output=None,
-            multimodal_outputs=wire_client_mm,
-            inter_stage_outputs=wire_inter_stage,
+            multimodal_outputs=multimodal_outputs,
+            inter_stage_outputs=inter_stage_outputs,
             kv_connector_output=kv_connector_output,
             num_nans_in_logits={},
             cudagraph_stats=cudagraph_stats,

From 48528c61328cd4d189470655e9ff7f742c4bd858 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Wed, 17 Jun 2026 17:42:42 +0800
Subject: [PATCH 05/21] update precommit

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 tests/utils/test_mm_outputs_partition.py  | 1 -
 vllm_omni/engine/orchestrator_profiler.py | 3 ++-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/utils/test_mm_outputs_partition.py b/tests/utils/test_mm_outputs_partition.py
index aa45eae2024..f9affa47131 100644
--- a/tests/utils/test_mm_outputs_partition.py
+++ b/tests/utils/test_mm_outputs_partition.py
@@ -42,4 +42,3 @@ def test_partition_payload_list_preserves_request_alignment():
     inter_list, client_list = partition_payload_list(payloads)
     assert inter_list == [payloads[0], None]
     assert client_list == [None, payloads[1]]
-
diff --git a/vllm_omni/engine/orchestrator_profiler.py b/vllm_omni/engine/orchestrator_profiler.py
index 8daaa3a4a95..cf5cc06b1a2 100644
--- a/vllm_omni/engine/orchestrator_profiler.py
+++ b/vllm_omni/engine/orchestrator_profiler.py
@@ -12,8 +12,9 @@
 import threading
 import time
 from collections import defaultdict
+from collections.abc import Iterator
 from contextlib import contextmanager
-from typing import Any, Iterator
+from typing import Any
 
 _ENABLED = os.environ.get("VLLM_OMNI_ORCH_PROFILE", "").lower() in ("1", "true", "yes")
 _OUTPUT_PATH = os.environ.get(

From 6e2371dd8dacd7237ca039ebe6be3dc9141ae7b9 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Thu, 18 Jun 2026 11:18:12 +0800
Subject: [PATCH 06/21] update

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 vllm_omni/engine/arg_utils.py             |   1 +
 vllm_omni/engine/async_omni_engine.py     |   3 +
 vllm_omni/engine/orchestrator.py          |  12 +
 vllm_omni/engine/orchestrator_profiler.py | 330 +++++++++++++---------
 vllm_omni/engine/stage_pool.py            |  23 ++
 vllm_omni/entrypoints/cli/serve.py        |   5 +
 6 files changed, 245 insertions(+), 129 deletions(-)

diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 148412012cf..7caff49eb23 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -422,6 +422,7 @@ class OrchestratorArgs:
 
     # === Observability ===
     log_stats: bool = False
+    enable_orch_profiler: 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 4a298feeb6f..bdff542cf7f 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -263,6 +263,7 @@ class AsyncOmniEngine:
     _log_stats: bool = False
     _coordinator_runtime: Any = None
     _transfer_emitter: Any = None
+    _enable_orch_profiler: bool = False
 
     def __init__(
         self,
@@ -290,6 +291,7 @@ def __init__(
         # 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_profiler = bool(kwargs.get("enable_orch_profiler", False))
 
         logger.info(f"[AsyncOmniEngine] Initializing with model {model}")
 
@@ -1398,6 +1400,7 @@ async def _run_orchestrator() -> None:
                 remote_replica_factory=remote_replica_factory,
                 transfer_emitter=self._transfer_emitter,
                 log_stats=self._log_stats,
+                enable_orch_profiler=self._enable_orch_profiler,
             )
             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 f8ff84de8fd..7c7ac4e2974 100644
--- a/vllm_omni/engine/orchestrator.py
+++ b/vllm_omni/engine/orchestrator.py
@@ -53,6 +53,7 @@
     StageSubmissionMessage,
     UnregisterRemoteReplicaMessage,
 )
+from vllm_omni.engine.orchestrator_profiler import get_orchestrator_profiler
 from vllm_omni.engine.serialization import serialize_additional_information
 from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.metrics.prometheus import OmniRequestCounter
@@ -184,6 +185,7 @@ def __init__(
         remote_replica_factory: RemoteReplicaFactory | None = None,
         transfer_emitter: Any = None,
         log_stats: bool = False,
+        enable_orch_profiler: bool = False,
     ) -> None:
         self.request_async_queue = request_async_queue
         self.output_async_queue = output_async_queue
@@ -192,6 +194,10 @@ def __init__(
         self.async_chunk = bool(async_chunk)
         self.num_stages = len(stage_pools)
         self.stage_pools: list[StagePool] = stage_pools
+        self._profiler = get_orchestrator_profiler(enabled=enable_orch_profiler)
+        for stage_id, pool in enumerate(self.stage_pools):
+            for replica_id in pool.live_replica_ids():
+                self._profiler.register_replica(stage_id, replica_id)
 
         # PD disaggregation state
         self._pd_pair: tuple[int, int] | None = None
@@ -363,6 +369,7 @@ async def run(self) -> None:
                             t.cancel()
 
             self._shutdown_stages()
+            self._profiler.flush()
 
             # Close the hub last so any in-flight dispatch still has access.
             if self._hub is not None:
@@ -671,6 +678,8 @@ async def _orchestration_loop(self) -> None:
                 for replica_id in pool.live_replica_ids():
                     if self._shutdown_event.is_set():
                         return
+                    self._profiler.register_replica(stage_id, replica_id)
+                    self._profiler.record_queue(stage_id, replica_id, pool.replica_queue_size(replica_id))
 
                     if pool.stage_type == "diffusion":
                         output = pool.poll_diffusion_output(replica_id)
@@ -679,6 +688,7 @@ async def _orchestration_loop(self) -> None:
 
                         pool.record_output_timestamps([output])
                         await self._handle_processed_outputs(stage_id, replica_id, [output])
+                        self._profiler.record_outputs(stage_id, replica_id, [output])
                         idle = False
                     else:
                         try:
@@ -760,8 +770,10 @@ async def _orchestration_loop(self) -> None:
                             raise
 
                         await self._handle_processed_outputs(stage_id, replica_id, raw_output)
+                        self._profiler.record_outputs(stage_id, replica_id, raw_output)
                         idle = False
 
+            self._profiler.note_loop(idle=idle, active_requests=len(self.request_states))
             if idle:
                 await asyncio.sleep(0.001)
             else:
diff --git a/vllm_omni/engine/orchestrator_profiler.py b/vllm_omni/engine/orchestrator_profiler.py
index cf5cc06b1a2..86902cfe667 100644
--- a/vllm_omni/engine/orchestrator_profiler.py
+++ b/vllm_omni/engine/orchestrator_profiler.py
@@ -1,156 +1,228 @@
-"""Optional orchestrator phase profiler (diagnostic only).
+"""Optional orchestrator window profiler (diagnostic only).
 
 Enable with:
-  export VLLM_OMNI_ORCH_PROFILE=1
-  export VLLM_OMNI_ORCH_PROFILE_PATH=/path/to/orch_profile.jsonl
+  --enable-orch-profiler
 """
 
 from __future__ import annotations
 
 import json
-import os
-import threading
 import time
-from collections import defaultdict
-from collections.abc import Iterator
-from contextlib import contextmanager
-from typing import Any
+from dataclasses import dataclass
+from pathlib import Path
 
-_ENABLED = os.environ.get("VLLM_OMNI_ORCH_PROFILE", "").lower() in ("1", "true", "yes")
-_OUTPUT_PATH = os.environ.get(
-    "VLLM_OMNI_ORCH_PROFILE_PATH",
-    "/tmp/vllm_omni_orch_profile.jsonl",
-)
-_FLUSH_INTERVAL_S = float(os.environ.get("VLLM_OMNI_ORCH_PROFILE_FLUSH_S", "1.0"))
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_DEFAULT_OUTPUT_PREFIX = "vllm_omni_orch_profile"
+_WINDOW_S = 1.0
+
+
+def _default_output_path() -> str:
+    timestamp = time.strftime("%m%d%H%M", time.localtime())
+    return str(Path.cwd() / f"{_DEFAULT_OUTPUT_PREFIX}_{timestamp}.json")
+
+
+@dataclass
+class _ReplicaWindow:
+    queue_size_sum: float = 0.0
+    queue_size_samples: int = 0
+    queue_size_max: int = 0
+    outputs: int = 0
+    finished_outputs: int = 0
 
 
 class OrchestratorProfiler:
-    def __init__(self) -> None:
-        self._lock = threading.Lock()
-        self._totals_s: dict[str, float] = defaultdict(float)
-        self._counts: dict[str, int] = defaultdict(int)
-        self._value_sums: dict[str, float] = defaultdict(float)
-        self._value_counts: dict[str, int] = defaultdict(int)
-        self._value_max: dict[str, float] = defaultdict(float)
-        self._value_last: dict[str, float] = {}
-        self._window_start = time.monotonic()
-        self._last_flush = self._window_start
+    def __init__(
+        self,
+        *,
+        enabled: bool = False,
+    ) -> None:
+        self._enabled = bool(enabled)
+        self._output_path = _default_output_path()
+        self._dumped = False
+        self._started_at_wall = time.time()
+        self._window_start_wall = self._started_at_wall
+        self._window_start_mono = time.monotonic()
         self._loop_iterations = 0
         self._loop_idle_iterations = 0
+        self._active_requests = 0
+        self._replica_windows: dict[str, _ReplicaWindow] = {}
+        self._replica_keys: set[str] = set()
+        self._windows: dict[str, list[float | int]] = {
+            "start_ts": [],
+            "end_ts": [],
+            "duration_s": [],
+            "active_requests": [],
+            "loop_iterations": [],
+            "loop_idle": [],
+            "loop_running": [],
+        }
+        self._replicas: dict[str, dict[str, list[float | int]]] = {}
 
     @property
     def enabled(self) -> bool:
-        return _ENABLED
+        return self._enabled
+
+    @staticmethod
+    def _replica_key(stage_id: int, replica_id: int) -> str:
+        return f"stage={stage_id},replica={replica_id}"
+
+    def register_replica(self, stage_id: int, replica_id: int) -> None:
+        if not self._enabled:
+            return
+        self._ensure_replica(self._replica_key(stage_id, replica_id))
 
-    def record(self, phase: str, elapsed_s: float, **tags: Any) -> None:
-        if not _ENABLED or elapsed_s < 0:
+    def note_loop(self, *, idle: bool, active_requests: int) -> None:
+        if not self._enabled:
+            return
+        self._roll_window_if_needed(time.monotonic())
+        self._loop_iterations += 1
+        if idle:
+            self._loop_idle_iterations += 1
+        self._active_requests = int(active_requests)
+
+    def record_queue(self, stage_id: int, replica_id: int, queue_size: int) -> None:
+        if not self._enabled:
             return
-        key = phase
-        if tags:
-            tag_str = ",".join(f"{k}={v}" for k, v in sorted(tags.items()))
-            key = f"{phase}|{tag_str}"
-        with self._lock:
-            self._totals_s[key] += elapsed_s
-            self._counts[key] += 1
-            now = time.monotonic()
-            if now - self._last_flush >= _FLUSH_INTERVAL_S:
-                self._flush_locked(now)
-
-    def record_value(self, name: str, value: float, **tags: Any) -> None:
-        if not _ENABLED:
+        self._roll_window_if_needed(time.monotonic())
+        key = self._replica_key(stage_id, replica_id)
+        self._ensure_replica(key)
+        window = self._replica_windows.setdefault(key, _ReplicaWindow())
+        queue_size = max(int(queue_size), 0)
+        window.queue_size_sum += float(queue_size)
+        window.queue_size_samples += 1
+        window.queue_size_max = max(window.queue_size_max, queue_size)
+
+    def record_outputs(self, stage_id: int, replica_id: int, outputs: list[object]) -> None:
+        if not self._enabled:
             return
-        key = name
-        if tags:
-            tag_str = ",".join(f"{k}={v}" for k, v in sorted(tags.items()))
-            key = f"{name}|{tag_str}"
-        with self._lock:
-            self._value_sums[key] += float(value)
-            self._value_counts[key] += 1
-            self._value_max[key] = max(self._value_max[key], float(value))
-            self._value_last[key] = float(value)
-            now = time.monotonic()
-            if now - self._last_flush >= _FLUSH_INTERVAL_S:
-                self._flush_locked(now)
-
-    def note_loop(self, *, idle: bool) -> None:
-        if not _ENABLED:
+        self._roll_window_if_needed(time.monotonic())
+        key = self._replica_key(stage_id, replica_id)
+        self._ensure_replica(key)
+        window = self._replica_windows.setdefault(key, _ReplicaWindow())
+        window.outputs += len(outputs)
+        window.finished_outputs += sum(1 for output in outputs if bool(getattr(output, "finished", False)))
+
+    def _ensure_replica(self, key: str) -> None:
+        if key in self._replica_keys:
             return
-        with self._lock:
-            self._loop_iterations += 1
-            if idle:
-                self._loop_idle_iterations += 1
-
-    def _flush_locked(self, now: float) -> None:
-        window_s = max(now - self._window_start, 1e-9)
-        profiled_s = sum(self._totals_s.values()) or 1e-9
+        self._replica_keys.add(key)
+        window_count = len(self._windows["start_ts"])
+        self._replicas[key] = {
+            "queue_size_avg": [0.0] * window_count,
+            "queue_size_max": [0] * window_count,
+            "outputs": [0] * window_count,
+            "finished_outputs": [0] * window_count,
+        }
+
+    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:
+        end_wall = time.time()
+        duration_s = max(now_mono - self._window_start_mono, 1e-9)
+        self._windows["start_ts"].append(self._window_start_wall)
+        self._windows["end_ts"].append(end_wall)
+        self._windows["duration_s"].append(duration_s)
+        self._windows["active_requests"].append(self._active_requests)
+        self._windows["loop_iterations"].append(self._loop_iterations)
+        self._windows["loop_idle"].append(self._loop_idle_iterations)
+        self._windows["loop_running"].append(self._loop_iterations - self._loop_idle_iterations)
+
+        for key in sorted(self._replica_keys):
+            metrics = self._replicas[key]
+            window = self._replica_windows.get(key, _ReplicaWindow())
+            queue_avg = window.queue_size_sum / float(window.queue_size_samples) if window.queue_size_samples else 0.0
+            metrics["queue_size_avg"].append(queue_avg)
+            metrics["queue_size_max"].append(window.queue_size_max)
+            metrics["outputs"].append(window.outputs)
+            metrics["finished_outputs"].append(window.finished_outputs)
+
+        self._replica_windows.clear()
+        self._loop_iterations = 0
+        self._loop_idle_iterations = 0
+        self._window_start_wall = end_wall
+        self._window_start_mono = now_mono
+
+    def dump(self) -> None:
+        if not self._enabled:
+            return
+        output_path = self._output_path
+        if self._dumped:
+            return
+        self._roll_window(time.monotonic())
+        self._dumped = True
         payload = {
-            "ts": time.time(),
-            "window_s": window_s,
-            "loop_iterations": self._loop_iterations,
-            "loop_idle_iterations": self._loop_idle_iterations,
-            "phases": {
-                key: {
-                    "total_ms": totals * 1000.0,
-                    "count": self._counts[key],
-                    "avg_ms": (totals / self._counts[key]) * 1000.0 if self._counts[key] else 0.0,
-                    "pct": (totals / profiled_s) * 100.0,
-                }
-                for key, totals in self._totals_s.items()
-            },
-            "values": {
-                key: {
-                    "count": self._value_counts[key],
-                    "avg": self._value_sums[key] / self._value_counts[key] if self._value_counts[key] else 0.0,
-                    "max": self._value_max[key],
-                    "last": self._value_last.get(key, 0.0),
-                }
-                for key in self._value_counts
-            },
+            "started_at": self._started_at_wall,
+            "ended_at": time.time(),
+            "configured_window_s": _WINDOW_S,
+            "windows": self._windows,
+            "replicas": self._replicas,
         }
+        summary = self._build_summary()
         try:
-            with open(_OUTPUT_PATH, "a", encoding="utf-8") as f:
-                f.write(json.dumps(payload, sort_keys=True) + "\n")
+            parent = Path(output_path).expanduser().parent
+            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(output_path, summary)
         except OSError:
-            pass
-        self._totals_s.clear()
-        self._counts.clear()
-        self._value_sums.clear()
-        self._value_counts.clear()
-        self._value_max.clear()
-        self._value_last.clear()
-        self._loop_iterations = 0
-        self._loop_idle_iterations = 0
-        self._window_start = now
-        self._last_flush = now
+            logger.exception("[OrchestratorProfiler] Failed to write profile to %s", output_path)
 
     def flush(self) -> None:
-        if not _ENABLED:
-            return
-        with self._lock:
-            self._flush_locked(time.monotonic())
-
-
-_profiler: OrchestratorProfiler | None = None
-_profiler_lock = threading.Lock()
-
-
-def get_orchestrator_profiler() -> OrchestratorProfiler:
-    global _profiler
-    if _profiler is None:
-        with _profiler_lock:
-            if _profiler is None:
-                _profiler = OrchestratorProfiler()
-    return _profiler
-
-
-@contextmanager
-def profile_phase(phase: str, **tags: Any) -> Iterator[None]:
-    prof = get_orchestrator_profiler()
-    if not prof.enabled:
-        yield
-        return
-    t0 = time.perf_counter()
-    try:
-        yield
-    finally:
-        prof.record(phase, time.perf_counter() - t0, **tags)
+        self.dump()
+
+    def _build_summary(self) -> dict[str, float | int]:
+        duration_s = sum(float(x) for x in self._windows["duration_s"])
+        loop_iterations = sum(int(x) for x in self._windows["loop_iterations"])
+        loop_idle = sum(int(x) for x in self._windows["loop_idle"])
+        loop_running = sum(int(x) for x in self._windows["loop_running"])
+        active_values = [int(x) for x in self._windows["active_requests"]]
+        return {
+            "windows": len(self._windows["duration_s"]),
+            "duration_s": duration_s,
+            "loop_iterations": loop_iterations,
+            "loop_idle": loop_idle,
+            "loop_running": loop_running,
+            "loop_running_pct": (float(loop_running) / float(loop_iterations) * 100.0) if loop_iterations else 0.0,
+            "active_requests_avg": (sum(active_values) / float(len(active_values))) if active_values else 0.0,
+            "active_requests_max": max(active_values) if active_values else 0,
+        }
+
+    def _log_summary(self, output_path: str, summary: dict[str, float | int]) -> None:
+        logger.info(
+            "[OrchestratorProfiler] wrote %s windows=%d duration=%.3fs loop_running=%.1f%% "
+            "active_avg=%.2f active_max=%d",
+            output_path,
+            int(summary["windows"]),
+            float(summary["duration_s"]),
+            float(summary["loop_running_pct"]),
+            float(summary["active_requests_avg"]),
+            int(summary["active_requests_max"]),
+        )
+        for key, metrics in sorted(self._replicas.items()):
+            queue_avg_values = [float(x) for x in metrics["queue_size_avg"]]
+            queue_max_values = [int(x) for x in metrics["queue_size_max"]]
+            outputs = sum(int(x) for x in metrics["outputs"])
+            finished_outputs = sum(int(x) for x in metrics["finished_outputs"])
+            queue_avg = sum(queue_avg_values) / float(len(queue_avg_values)) if queue_avg_values else 0.0
+            queue_max = max(queue_max_values) if queue_max_values else 0
+            logger.info(
+                "[OrchestratorProfiler] %s queue_avg=%.2f queue_max=%d outputs=%d finished_outputs=%d",
+                key,
+                queue_avg,
+                queue_max,
+                outputs,
+                finished_outputs,
+            )
+
+
+def get_orchestrator_profiler(
+    *,
+    enabled: bool = False,
+) -> OrchestratorProfiler:
+    return OrchestratorProfiler(enabled=enabled)
diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py
index 4e1f3d5a55b..63683a0184c 100644
--- a/vllm_omni/engine/stage_pool.py
+++ b/vllm_omni/engine/stage_pool.py
@@ -429,6 +429,29 @@ def get_bound_llm_client(self, request_id: str) -> StagePoolLLMClient | None:
             return None
         return cast(StagePoolLLMClient, client)
 
+    def replica_queue_size(self, replica_id: int) -> int:
+        """Best-effort queue/backlog size for profiler sampling.
+
+        Distributed replicas publish scheduler queue length via coordinator
+        heartbeats. Local/in-process stages do not expose that value cheaply,
+        so fall back to the number of logical requests currently bound to the
+        replica in this stage.
+        """
+        if replica_id < 0 or replica_id >= len(self.clients) or self.clients[replica_id] is None:
+            return 0
+        if self._hub is not None:
+            client = self.clients[replica_id]
+            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 max(int(replica.queue_length), 0)
+        bound_legacy = sum(1 for bound_replica_id in self._request_bindings.values() if bound_replica_id == replica_id)
+        bound_affinity = sum(
+            1 for input_addr in self._affinity.values() if self._addr_to_replica_id.get(input_addr) == replica_id
+        )
+        return max(bound_legacy, bound_affinity)
+
     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 5b6c7d9b34f..86208ad9ea6 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -670,6 +670,11 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu
             action="store_true",
             help="Enable AR stage profiler to include AR stage timing in stage_durations.",
         )
+        omni_config_group.add_argument(
+            "--enable-orch-profiler",
+            action="store_true",
+            help="Enable orchestrator window profiler and write a JSON summary at shutdown.",
+        )
 
         # Supplementary auxiliary text encoder parameters
         # (e.g., the meta llama/meta llama-3.1-8b-instrument used by hidream)

From 314f3f269a57c14f987d52a3703f984e4bcc2d2d Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Thu, 18 Jun 2026 13:56:05 +0800
Subject: [PATCH 07/21] update to monitor

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 vllm_omni/engine/arg_utils.py             |   2 +-
 vllm_omni/engine/async_omni_engine.py     |   6 +-
 vllm_omni/engine/orchestrator.py          |  31 ++-
 vllm_omni/engine/orchestrator_monitor.py  | 207 ++++++++++++++++++++
 vllm_omni/engine/orchestrator_profiler.py | 228 ----------------------
 vllm_omni/engine/stage_pool.py            |  20 +-
 vllm_omni/entrypoints/cli/serve.py        |   4 +-
 7 files changed, 247 insertions(+), 251 deletions(-)
 create mode 100644 vllm_omni/engine/orchestrator_monitor.py
 delete mode 100644 vllm_omni/engine/orchestrator_profiler.py

diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 7caff49eb23..4daace5a909 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -422,7 +422,7 @@ class OrchestratorArgs:
 
     # === Observability ===
     log_stats: bool = False
-    enable_orch_profiler: 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 bdff542cf7f..8f107092a0c 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -263,7 +263,7 @@ class AsyncOmniEngine:
     _log_stats: bool = False
     _coordinator_runtime: Any = None
     _transfer_emitter: Any = None
-    _enable_orch_profiler: bool = False
+    _enable_orch_monitor: bool = False
 
     def __init__(
         self,
@@ -291,7 +291,7 @@ def __init__(
         # 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_profiler = bool(kwargs.get("enable_orch_profiler", False))
+        self._enable_orch_monitor = bool(kwargs.get("enable_orch_monitor", False))
 
         logger.info(f"[AsyncOmniEngine] Initializing with model {model}")
 
@@ -1400,7 +1400,7 @@ async def _run_orchestrator() -> None:
                 remote_replica_factory=remote_replica_factory,
                 transfer_emitter=self._transfer_emitter,
                 log_stats=self._log_stats,
-                enable_orch_profiler=self._enable_orch_profiler,
+                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 7c7ac4e2974..ea58611f2e5 100644
--- a/vllm_omni/engine/orchestrator.py
+++ b/vllm_omni/engine/orchestrator.py
@@ -53,7 +53,7 @@
     StageSubmissionMessage,
     UnregisterRemoteReplicaMessage,
 )
-from vllm_omni.engine.orchestrator_profiler import get_orchestrator_profiler
+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
@@ -185,7 +185,7 @@ def __init__(
         remote_replica_factory: RemoteReplicaFactory | None = None,
         transfer_emitter: Any = None,
         log_stats: bool = False,
-        enable_orch_profiler: bool = False,
+        enable_orch_monitor: bool = False,
     ) -> None:
         self.request_async_queue = request_async_queue
         self.output_async_queue = output_async_queue
@@ -194,10 +194,13 @@ def __init__(
         self.async_chunk = bool(async_chunk)
         self.num_stages = len(stage_pools)
         self.stage_pools: list[StagePool] = stage_pools
-        self._profiler = get_orchestrator_profiler(enabled=enable_orch_profiler)
+        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._profiler.register_replica(stage_id, replica_id)
+                self._orch_monitor.register_replica(stage_id, replica_id)
 
         # PD disaggregation state
         self._pd_pair: tuple[int, int] | None = None
@@ -368,8 +371,8 @@ async def run(self) -> None:
                         if not t.done():
                             t.cancel()
 
+            self._orch_monitor.flush()
             self._shutdown_stages()
-            self._profiler.flush()
 
             # Close the hub last so any in-flight dispatch still has access.
             if self._hub is not None:
@@ -448,6 +451,7 @@ async def _request_handler(self) -> None:
                     for client in pool.clients:
                         if hasattr(client, "_shutting_down"):
                             client._shutting_down = True
+                self._orch_monitor.flush()
                 self._shutdown_stages()
                 break
             else:
@@ -661,6 +665,17 @@ async def _handle_collective_rpc(self, msg: CollectiveRPCRequestMessage) -> None
 
     # ---- 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_outputs_queue_size(replica_id),
+                    pool.replica_inflight_count(replica_id),
+                )
+        return samples
+
     async def _orchestration_output_handler(self) -> None:
         """Poll all stages, handle transfers, send final outputs to main."""
         try:
@@ -678,8 +693,6 @@ async def _orchestration_loop(self) -> None:
                 for replica_id in pool.live_replica_ids():
                     if self._shutdown_event.is_set():
                         return
-                    self._profiler.register_replica(stage_id, replica_id)
-                    self._profiler.record_queue(stage_id, replica_id, pool.replica_queue_size(replica_id))
 
                     if pool.stage_type == "diffusion":
                         output = pool.poll_diffusion_output(replica_id)
@@ -688,7 +701,6 @@ async def _orchestration_loop(self) -> None:
 
                         pool.record_output_timestamps([output])
                         await self._handle_processed_outputs(stage_id, replica_id, [output])
-                        self._profiler.record_outputs(stage_id, replica_id, [output])
                         idle = False
                     else:
                         try:
@@ -770,10 +782,9 @@ async def _orchestration_loop(self) -> None:
                             raise
 
                         await self._handle_processed_outputs(stage_id, replica_id, raw_output)
-                        self._profiler.record_outputs(stage_id, replica_id, raw_output)
                         idle = False
 
-            self._profiler.note_loop(idle=idle, active_requests=len(self.request_states))
+            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 00000000000..e0ac7c522f5
--- /dev/null
+++ b/vllm_omni/engine/orchestrator_monitor.py
@@ -0,0 +1,207 @@
+"""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)
+"""
+
+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 _default_output_path() -> str:
+    timestamp = time.strftime("%m%d%H%M", time.localtime())
+    return str(Path.cwd() / f"{_DEFAULT_OUTPUT_PREFIX}_{timestamp}.json")
+
+
+def _resolve_output_path() -> str:
+    for env_name in (
+        "VLLM_OMNI_ORCH_MONITOR_PATH",
+        "VLLM_OMNI_ORCH_OBSERVER_PATH",
+        "VLLM_OMNI_ORCH_PROFILE_PATH",
+    ):
+        env_path = os.environ.get(env_name)
+        if env_path:
+            return env_path
+    return _default_output_path()
+
+
+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 | None = None,
+) -> OrchestratorMonitorBase:
+    if not enabled:
+        return _NULL_ORCH_MONITOR
+    return OrchestratorMonitor(replica_sampler=replica_sampler)
+
+
+class OrchestratorMonitor:
+    def __init__(self, *, replica_sampler: ReplicaSampler | None = None) -> None:
+        self._output_path = _resolve_output_path()
+        self._dumped = 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:
+        self.dump()
+
+    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:
+        sampled = self._replica_sampler() if self._replica_sampler is not None else {}
+        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(max(int(outputs_qsize), 0))
+            series["inflight"].append(max(int(inflight), 0))
+
+        self._loop_idle = 0
+        self._loop_active = 0
+        self._window_start_mono = now_mono
+
+    def dump(self) -> None:
+        if self._dumped:
+            return
+        self._roll_window(time.monotonic())
+        self._dumped = 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())),
+        }
+        summary = self._build_summary()
+        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), summary)
+        except OSError:
+            logger.exception("[OrchestratorMonitor] Failed to write log to %s", self._output_path)
+
+    def _build_summary(self) -> dict[str, float | int]:
+        loop_idle = sum(self._loop_idle_buf)
+        loop_active = sum(self._loop_active_buf)
+        loop_total = loop_idle + loop_active
+        return {
+            "windows": len(self._duration_s),
+            "duration_s": sum(self._duration_s),
+            "loop_idle": loop_idle,
+            "loop_active": loop_active,
+            "loop_active_pct": (float(loop_active) / float(loop_total) * 100.0) if loop_total else 0.0,
+        }
+
+    def _log_summary(self, output_path: str, summary: dict[str, float | int]) -> None:
+        logger.info(
+            "[OrchestratorMonitor] wrote %s windows=%d duration=%.3fs loop_active=%.1f%%",
+            output_path,
+            int(summary["windows"]),
+            float(summary["duration_s"]),
+            float(summary["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/orchestrator_profiler.py b/vllm_omni/engine/orchestrator_profiler.py
deleted file mode 100644
index 86902cfe667..00000000000
--- a/vllm_omni/engine/orchestrator_profiler.py
+++ /dev/null
@@ -1,228 +0,0 @@
-"""Optional orchestrator window profiler (diagnostic only).
-
-Enable with:
-  --enable-orch-profiler
-"""
-
-from __future__ import annotations
-
-import json
-import time
-from dataclasses import dataclass
-from pathlib import Path
-
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-
-_DEFAULT_OUTPUT_PREFIX = "vllm_omni_orch_profile"
-_WINDOW_S = 1.0
-
-
-def _default_output_path() -> str:
-    timestamp = time.strftime("%m%d%H%M", time.localtime())
-    return str(Path.cwd() / f"{_DEFAULT_OUTPUT_PREFIX}_{timestamp}.json")
-
-
-@dataclass
-class _ReplicaWindow:
-    queue_size_sum: float = 0.0
-    queue_size_samples: int = 0
-    queue_size_max: int = 0
-    outputs: int = 0
-    finished_outputs: int = 0
-
-
-class OrchestratorProfiler:
-    def __init__(
-        self,
-        *,
-        enabled: bool = False,
-    ) -> None:
-        self._enabled = bool(enabled)
-        self._output_path = _default_output_path()
-        self._dumped = False
-        self._started_at_wall = time.time()
-        self._window_start_wall = self._started_at_wall
-        self._window_start_mono = time.monotonic()
-        self._loop_iterations = 0
-        self._loop_idle_iterations = 0
-        self._active_requests = 0
-        self._replica_windows: dict[str, _ReplicaWindow] = {}
-        self._replica_keys: set[str] = set()
-        self._windows: dict[str, list[float | int]] = {
-            "start_ts": [],
-            "end_ts": [],
-            "duration_s": [],
-            "active_requests": [],
-            "loop_iterations": [],
-            "loop_idle": [],
-            "loop_running": [],
-        }
-        self._replicas: dict[str, dict[str, list[float | int]]] = {}
-
-    @property
-    def enabled(self) -> bool:
-        return self._enabled
-
-    @staticmethod
-    def _replica_key(stage_id: int, replica_id: int) -> str:
-        return f"stage={stage_id},replica={replica_id}"
-
-    def register_replica(self, stage_id: int, replica_id: int) -> None:
-        if not self._enabled:
-            return
-        self._ensure_replica(self._replica_key(stage_id, replica_id))
-
-    def note_loop(self, *, idle: bool, active_requests: int) -> None:
-        if not self._enabled:
-            return
-        self._roll_window_if_needed(time.monotonic())
-        self._loop_iterations += 1
-        if idle:
-            self._loop_idle_iterations += 1
-        self._active_requests = int(active_requests)
-
-    def record_queue(self, stage_id: int, replica_id: int, queue_size: int) -> None:
-        if not self._enabled:
-            return
-        self._roll_window_if_needed(time.monotonic())
-        key = self._replica_key(stage_id, replica_id)
-        self._ensure_replica(key)
-        window = self._replica_windows.setdefault(key, _ReplicaWindow())
-        queue_size = max(int(queue_size), 0)
-        window.queue_size_sum += float(queue_size)
-        window.queue_size_samples += 1
-        window.queue_size_max = max(window.queue_size_max, queue_size)
-
-    def record_outputs(self, stage_id: int, replica_id: int, outputs: list[object]) -> None:
-        if not self._enabled:
-            return
-        self._roll_window_if_needed(time.monotonic())
-        key = self._replica_key(stage_id, replica_id)
-        self._ensure_replica(key)
-        window = self._replica_windows.setdefault(key, _ReplicaWindow())
-        window.outputs += len(outputs)
-        window.finished_outputs += sum(1 for output in outputs if bool(getattr(output, "finished", False)))
-
-    def _ensure_replica(self, key: str) -> None:
-        if key in self._replica_keys:
-            return
-        self._replica_keys.add(key)
-        window_count = len(self._windows["start_ts"])
-        self._replicas[key] = {
-            "queue_size_avg": [0.0] * window_count,
-            "queue_size_max": [0] * window_count,
-            "outputs": [0] * window_count,
-            "finished_outputs": [0] * window_count,
-        }
-
-    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:
-        end_wall = time.time()
-        duration_s = max(now_mono - self._window_start_mono, 1e-9)
-        self._windows["start_ts"].append(self._window_start_wall)
-        self._windows["end_ts"].append(end_wall)
-        self._windows["duration_s"].append(duration_s)
-        self._windows["active_requests"].append(self._active_requests)
-        self._windows["loop_iterations"].append(self._loop_iterations)
-        self._windows["loop_idle"].append(self._loop_idle_iterations)
-        self._windows["loop_running"].append(self._loop_iterations - self._loop_idle_iterations)
-
-        for key in sorted(self._replica_keys):
-            metrics = self._replicas[key]
-            window = self._replica_windows.get(key, _ReplicaWindow())
-            queue_avg = window.queue_size_sum / float(window.queue_size_samples) if window.queue_size_samples else 0.0
-            metrics["queue_size_avg"].append(queue_avg)
-            metrics["queue_size_max"].append(window.queue_size_max)
-            metrics["outputs"].append(window.outputs)
-            metrics["finished_outputs"].append(window.finished_outputs)
-
-        self._replica_windows.clear()
-        self._loop_iterations = 0
-        self._loop_idle_iterations = 0
-        self._window_start_wall = end_wall
-        self._window_start_mono = now_mono
-
-    def dump(self) -> None:
-        if not self._enabled:
-            return
-        output_path = self._output_path
-        if self._dumped:
-            return
-        self._roll_window(time.monotonic())
-        self._dumped = True
-        payload = {
-            "started_at": self._started_at_wall,
-            "ended_at": time.time(),
-            "configured_window_s": _WINDOW_S,
-            "windows": self._windows,
-            "replicas": self._replicas,
-        }
-        summary = self._build_summary()
-        try:
-            parent = Path(output_path).expanduser().parent
-            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(output_path, summary)
-        except OSError:
-            logger.exception("[OrchestratorProfiler] Failed to write profile to %s", output_path)
-
-    def flush(self) -> None:
-        self.dump()
-
-    def _build_summary(self) -> dict[str, float | int]:
-        duration_s = sum(float(x) for x in self._windows["duration_s"])
-        loop_iterations = sum(int(x) for x in self._windows["loop_iterations"])
-        loop_idle = sum(int(x) for x in self._windows["loop_idle"])
-        loop_running = sum(int(x) for x in self._windows["loop_running"])
-        active_values = [int(x) for x in self._windows["active_requests"]]
-        return {
-            "windows": len(self._windows["duration_s"]),
-            "duration_s": duration_s,
-            "loop_iterations": loop_iterations,
-            "loop_idle": loop_idle,
-            "loop_running": loop_running,
-            "loop_running_pct": (float(loop_running) / float(loop_iterations) * 100.0) if loop_iterations else 0.0,
-            "active_requests_avg": (sum(active_values) / float(len(active_values))) if active_values else 0.0,
-            "active_requests_max": max(active_values) if active_values else 0,
-        }
-
-    def _log_summary(self, output_path: str, summary: dict[str, float | int]) -> None:
-        logger.info(
-            "[OrchestratorProfiler] wrote %s windows=%d duration=%.3fs loop_running=%.1f%% "
-            "active_avg=%.2f active_max=%d",
-            output_path,
-            int(summary["windows"]),
-            float(summary["duration_s"]),
-            float(summary["loop_running_pct"]),
-            float(summary["active_requests_avg"]),
-            int(summary["active_requests_max"]),
-        )
-        for key, metrics in sorted(self._replicas.items()):
-            queue_avg_values = [float(x) for x in metrics["queue_size_avg"]]
-            queue_max_values = [int(x) for x in metrics["queue_size_max"]]
-            outputs = sum(int(x) for x in metrics["outputs"])
-            finished_outputs = sum(int(x) for x in metrics["finished_outputs"])
-            queue_avg = sum(queue_avg_values) / float(len(queue_avg_values)) if queue_avg_values else 0.0
-            queue_max = max(queue_max_values) if queue_max_values else 0
-            logger.info(
-                "[OrchestratorProfiler] %s queue_avg=%.2f queue_max=%d outputs=%d finished_outputs=%d",
-                key,
-                queue_avg,
-                queue_max,
-                outputs,
-                finished_outputs,
-            )
-
-
-def get_orchestrator_profiler(
-    *,
-    enabled: bool = False,
-) -> OrchestratorProfiler:
-    return OrchestratorProfiler(enabled=enabled)
diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py
index 63683a0184c..b8e368b99d5 100644
--- a/vllm_omni/engine/stage_pool.py
+++ b/vllm_omni/engine/stage_pool.py
@@ -429,14 +429,20 @@ def get_bound_llm_client(self, request_id: str) -> StagePoolLLMClient | None:
             return None
         return cast(StagePoolLLMClient, client)
 
-    def replica_queue_size(self, replica_id: int) -> int:
-        """Best-effort queue/backlog size for profiler sampling.
+    def replica_outputs_queue_size(self, replica_id: int) -> int:
+        """Depth of the EngineCore async client's outputs_queue (MP receive backlog)."""
+        if replica_id < 0 or replica_id >= len(self.clients) or self.clients[replica_id] is None:
+            return 0
+        outputs_queue = getattr(self.clients[replica_id], "outputs_queue", None)
+        if outputs_queue is None:
+            return 0
+        try:
+            return max(int(outputs_queue.qsize()), 0)
+        except (AttributeError, NotImplementedError, RuntimeError, ValueError):
+            return 0
 
-        Distributed replicas publish scheduler queue length via coordinator
-        heartbeats. Local/in-process stages do not expose that value cheaply,
-        so fall back to the number of logical requests currently bound to the
-        replica in this stage.
-        """
+    def replica_inflight_count(self, replica_id: int) -> int:
+        """Number of logical requests currently bound/routed to this replica."""
         if replica_id < 0 or replica_id >= len(self.clients) or self.clients[replica_id] is None:
             return 0
         if self._hub is not None:
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index 86208ad9ea6..faaa5eec6b4 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -671,9 +671,9 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu
             help="Enable AR stage profiler to include AR stage timing in stage_durations.",
         )
         omni_config_group.add_argument(
-            "--enable-orch-profiler",
+            "--enable-orch-monitor",
             action="store_true",
-            help="Enable orchestrator window profiler and write a JSON summary at shutdown.",
+            help="Enable orchestrator window monitor and write a JSON log at shutdown.",
         )
 
         # Supplementary auxiliary text encoder parameters

From 02af9a73d011bc8e51d155bdf9aac18b8ff9c376 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Thu, 18 Jun 2026 15:41:46 +0800
Subject: [PATCH 08/21] fix precommit

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 vllm_omni/engine/orchestrator_monitor.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/vllm_omni/engine/orchestrator_monitor.py b/vllm_omni/engine/orchestrator_monitor.py
index e0ac7c522f5..79a7ca3542e 100644
--- a/vllm_omni/engine/orchestrator_monitor.py
+++ b/vllm_omni/engine/orchestrator_monitor.py
@@ -198,7 +198,8 @@ def _log_summary(self, output_path: str, summary: dict[str, float | int]) -> Non
             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",
+                "[OrchestratorMonitor] %s outputs_queue_avg=%.2f outputs_queue_max=%d "
+                "inflight_avg=%.2f inflight_max=%d",
                 key,
                 queue_avg,
                 queue_max,

From 1ed8b8e44ec3718e1e06653048450a23bf236b95 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Thu, 18 Jun 2026 16:07:31 +0800
Subject: [PATCH 09/21] update voxtral tts

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 tests/engine/test_orchestrator_monitor.py     | 91 +++++++++++++++++++
 .../test_cuda_graph_acoustic_transformer.py   |  6 +-
 .../test_voxtral_tts_async_chunk.py           | 14 +--
 ...cuda_graph_acoustic_transformer_wrapper.py |  4 +-
 .../models/voxtral_tts/voxtral_tts.py         | 15 +--
 .../voxtral_tts_audio_generation.py           |  2 +-
 .../stage_input_processors/voxtral_tts.py     | 29 +++++-
 7 files changed, 138 insertions(+), 23 deletions(-)
 create mode 100644 tests/engine/test_orchestrator_monitor.py

diff --git a/tests/engine/test_orchestrator_monitor.py b/tests/engine/test_orchestrator_monitor.py
new file mode 100644
index 00000000000..70f4ab61804
--- /dev/null
+++ b/tests/engine/test_orchestrator_monitor.py
@@ -0,0 +1,91 @@
+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)
+    assert monitor is _NULL_ORCH_MONITOR
+    assert create_orch_monitor(enabled=False) 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_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 3b645632e9c..cf40f9ad417 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 @@ 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
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 5235d7c68f7..3fd5675ce4b 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/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 491f185a4b8..5898136c997 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 @@ def __call__(
         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 d12034c243b..2502e27b886 100644
--- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
+++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
@@ -220,12 +220,15 @@ def tts_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **
 
     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 2e50e34f32d..1fbcc509f81 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
@@ -1083,7 +1083,7 @@ def compute_mm_logits(
         # 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 2928a7cf31b..d3c1219b7d5 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 @@
 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
+    frame = audio_tensors[-1]
+    if not isinstance(frame, torch.Tensor) or frame.numel() == 0:
         return None
-    return audio.flatten()
+    return frame.flatten()
 
 
 def generator2tokenizer_async_chunk(

From a94ae7caf594036c01d1773e10a00e128d487c0e Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Mon, 22 Jun 2026 20:06:11 +0800
Subject: [PATCH 10/21] fix

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 vllm_omni/core/sched/omni_ar_scheduler.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py
index 6645e08ed47..4761af8508d 100644
--- a/vllm_omni/core/sched/omni_ar_scheduler.py
+++ b/vllm_omni/core/sched/omni_ar_scheduler.py
@@ -493,7 +493,9 @@ def update_from_output(
                 # 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:
+            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,

From 121a2e219a5bc61a8392912763dcec4c5cb84db4 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Mon, 22 Jun 2026 20:07:20 +0800
Subject: [PATCH 11/21] fix

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 vllm_omni/engine/orchestrator.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py
index 92923da3933..b1f73f5db68 100644
--- a/vllm_omni/engine/orchestrator.py
+++ b/vllm_omni/engine/orchestrator.py
@@ -348,6 +348,7 @@ async def _request_handler(self) -> None:
             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)

From 5b03230f46789c0facb5895582c2afe99131ce43 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Mon, 22 Jun 2026 20:14:25 +0800
Subject: [PATCH 12/21] fix precommit

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 vllm_omni/core/sched/omni_ar_scheduler.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py
index 4761af8508d..b18d090de4c 100644
--- a/vllm_omni/core/sched/omni_ar_scheduler.py
+++ b/vllm_omni/core/sched/omni_ar_scheduler.py
@@ -493,9 +493,7 @@ def update_from_output(
                 # 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
-            ):
+            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,

From 4e912b0fcc61f02b7fe75c731f774e0d48d983cd Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Mon, 22 Jun 2026 20:15:48 +0800
Subject: [PATCH 13/21] fix tts

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 .../models/voxtral_tts/voxtral_tts.py             | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

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 2502e27b886..d446ed6923f 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 @@ def sampler(self):
             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 @@ def tts_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **
             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()

From 79d0dc30218a85fc5054793697a0a5f2fa322ee5 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Wed, 24 Jun 2026 12:50:30 +0800
Subject: [PATCH 14/21] update npu

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 .../npu/worker/npu_ar_model_runner.py         | 21 ++++++++++++++-----
 .../npu/worker/npu_generation_model_runner.py | 12 +++++++----
 2 files changed, 24 insertions(+), 9 deletions(-)

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 b0c120ca704..8b1bb3422c8 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.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]:
@@ -252,13 +252,21 @@ def _resolve_req_hidden_states(
 
     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: list[dict[str, torch.Tensor] | None] = []
+        for payload in per_req_payloads:
+            if not payload:
+                wire.append(None)
+            else:
+                wire.append(_ensure_tensor_values(payload))
+        if all(item is None for item in wire):
+            return None
+        return wire
 
 
     def _request_final_stage_id(self, req_id: str) -> int | None:
@@ -1132,7 +1140,9 @@ def _unwrap_lists(v):
                     payload.update(mm_payload)
                 pooler_output.append(flatten_payload(payload))
 
-        multimodal_outputs = self._build_multimodal_outputs(pooler_output)
+        pooler_inter, pooler_client = partition_payload_list(pooler_output or [])
+        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 +1151,7 @@ def _unwrap_lists(v):
             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 d2eeeb5789b..170589e97e7 100644
--- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
@@ -35,6 +35,7 @@
 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):
@@ -435,20 +436,20 @@ def sample_tokens(
         #  -------------------------------------- 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 +468,11 @@ def sample_tokens(
                         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")
+
+        inter_stage_outputs, multimodal_outputs = partition_payload_list(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 +489,7 @@ def sample_tokens(
             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,

From f208660a2eb6c17cd13c689a211a89902329654e Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Wed, 24 Jun 2026 16:30:13 +0800
Subject: [PATCH 15/21] update

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 tests/worker/test_gpu_ar_model_runner.py | 35 ++++++++++++++----------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py
index 60872636d7f..dcbe92b8b16 100644
--- a/tests/worker/test_gpu_ar_model_runner.py
+++ b/tests/worker/test_gpu_ar_model_runner.py
@@ -174,8 +174,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 +219,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 +334,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():
@@ -428,11 +432,13 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
 
     output = GPUARModelRunner.sample_tokens(runner, grammar_output=None)
 
-    assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]]))
+    assert output.inter_stage_outputs is not None
+    assert torch.equal(output.inter_stage_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]]))
     assert torch.equal(
-        output.multimodal_outputs[1]["hidden"],
+        output.inter_stage_outputs[1]["hidden"],
         torch.tensor([[2.0, 20.0], [3.0, 30.0]]),
     )
+    assert output.multimodal_outputs is None
 
 
 def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch):
@@ -482,5 +488,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

From 3216e1ae344a6fb9684697626a568807035933b9 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Thu, 25 Jun 2026 13:11:37 +0800
Subject: [PATCH 16/21] update code

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 docs/contributing/profiling.md            | 51 ++++++++++++
 tests/engine/test_orchestrator_monitor.py |  4 +-
 vllm_omni/engine/async_omni_engine.py     |  2 +-
 vllm_omni/engine/orchestrator.py          |  5 +-
 vllm_omni/engine/orchestrator_monitor.py  | 97 +++++++++--------------
 vllm_omni/engine/stage_pool.py            | 44 +++++-----
 6 files changed, 116 insertions(+), 87 deletions(-)

diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md
index 400d2a3c02c..b57de441ed1 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_<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.
diff --git a/tests/engine/test_orchestrator_monitor.py b/tests/engine/test_orchestrator_monitor.py
index 70f4ab61804..6a1f11d47f8 100644
--- a/tests/engine/test_orchestrator_monitor.py
+++ b/tests/engine/test_orchestrator_monitor.py
@@ -20,9 +20,9 @@ def test_replica_key_format():
 
 
 def test_create_orch_monitor_disabled_returns_null_singleton():
-    monitor = create_orch_monitor(enabled=False)
+    monitor = create_orch_monitor(enabled=False, replica_sampler=lambda: {})
     assert monitor is _NULL_ORCH_MONITOR
-    assert create_orch_monitor(enabled=False) is monitor
+    assert create_orch_monitor(enabled=False, replica_sampler=lambda: {}) is monitor
     monitor.note_loop(idle=True)
     monitor.flush()
 
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index 144f1770cbb..c8c40baf14d 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -234,7 +234,7 @@ def __init__(
         # 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.get("enable_orch_monitor", False))
+        self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False))
 
         logger.info(f"[AsyncOmniEngine] Initializing with model {model}")
 
diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py
index ea6d8e4b683..7139f7e06d1 100644
--- a/vllm_omni/engine/orchestrator.py
+++ b/vllm_omni/engine/orchestrator.py
@@ -645,10 +645,7 @@ def _sample_replica_metrics(self) -> 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_outputs_queue_size(replica_id),
-                    pool.replica_inflight_count(replica_id),
-                )
+                samples[key] = pool.replica_monitor_sample(replica_id)
         return samples
 
     async def _orchestration_output_handler(self) -> None:
diff --git a/vllm_omni/engine/orchestrator_monitor.py b/vllm_omni/engine/orchestrator_monitor.py
index 79a7ca3542e..debd689dd8c 100644
--- a/vllm_omni/engine/orchestrator_monitor.py
+++ b/vllm_omni/engine/orchestrator_monitor.py
@@ -10,6 +10,8 @@
   - 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
@@ -41,23 +43,14 @@ def replica_key(stage_id: int, replica_id: int) -> str:
     return f"stage={stage_id},replica={replica_id}"
 
 
-def _default_output_path() -> str:
+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")
 
 
-def _resolve_output_path() -> str:
-    for env_name in (
-        "VLLM_OMNI_ORCH_MONITOR_PATH",
-        "VLLM_OMNI_ORCH_OBSERVER_PATH",
-        "VLLM_OMNI_ORCH_PROFILE_PATH",
-    ):
-        env_path = os.environ.get(env_name)
-        if env_path:
-            return env_path
-    return _default_output_path()
-
-
 class OrchestratorMonitorBase(Protocol):
     def register_replica(self, stage_id: int, replica_id: int) -> None: ...
 
@@ -83,7 +76,7 @@ def flush(self) -> None:
 def create_orch_monitor(
     *,
     enabled: bool,
-    replica_sampler: ReplicaSampler | None = None,
+    replica_sampler: ReplicaSampler,
 ) -> OrchestratorMonitorBase:
     if not enabled:
         return _NULL_ORCH_MONITOR
@@ -91,9 +84,9 @@ def create_orch_monitor(
 
 
 class OrchestratorMonitor:
-    def __init__(self, *, replica_sampler: ReplicaSampler | None = None) -> None:
+    def __init__(self, *, replica_sampler: ReplicaSampler) -> None:
         self._output_path = _resolve_output_path()
-        self._dumped = False
+        self._flushed = False
         self._replica_sampler = replica_sampler
         self._window_start_mono = time.monotonic()
         self._loop_idle = 0
@@ -115,7 +108,28 @@ def note_loop(self, *, idle: bool) -> None:
             self._loop_active += 1
 
     def flush(self) -> None:
-        self.dump()
+        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:
@@ -132,63 +146,30 @@ def _roll_window_if_needed(self, now_mono: float) -> None:
             self._roll_window(now_mono)
 
     def _roll_window(self, now_mono: float) -> None:
-        sampled = self._replica_sampler() if self._replica_sampler is not None else {}
+        sampled = self._replica_sampler()
         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(max(int(outputs_qsize), 0))
-            series["inflight"].append(max(int(inflight), 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 dump(self) -> None:
-        if self._dumped:
-            return
-        self._roll_window(time.monotonic())
-        self._dumped = 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())),
-        }
-        summary = self._build_summary()
-        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), summary)
-        except OSError:
-            logger.exception("[OrchestratorMonitor] Failed to write log to %s", self._output_path)
-
-    def _build_summary(self) -> dict[str, float | int]:
+    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
-        return {
-            "windows": len(self._duration_s),
-            "duration_s": sum(self._duration_s),
-            "loop_idle": loop_idle,
-            "loop_active": loop_active,
-            "loop_active_pct": (float(loop_active) / float(loop_total) * 100.0) if loop_total else 0.0,
-        }
-
-    def _log_summary(self, output_path: str, summary: dict[str, float | int]) -> None:
+        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,
-            int(summary["windows"]),
-            float(summary["duration_s"]),
-            float(summary["loop_active_pct"]),
+            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"]
diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py
index b8e368b99d5..6539309abff 100644
--- a/vllm_omni/engine/stage_pool.py
+++ b/vllm_omni/engine/stage_pool.py
@@ -429,34 +429,34 @@ def get_bound_llm_client(self, request_id: str) -> StagePoolLLMClient | None:
             return None
         return cast(StagePoolLLMClient, client)
 
-    def replica_outputs_queue_size(self, replica_id: int) -> int:
-        """Depth of the EngineCore async client's outputs_queue (MP receive backlog)."""
-        if replica_id < 0 or replica_id >= len(self.clients) or self.clients[replica_id] is None:
-            return 0
-        outputs_queue = getattr(self.clients[replica_id], "outputs_queue", None)
-        if outputs_queue is None:
-            return 0
-        try:
-            return max(int(outputs_queue.qsize()), 0)
-        except (AttributeError, NotImplementedError, RuntimeError, ValueError):
-            return 0
+    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
 
-    def replica_inflight_count(self, replica_id: int) -> int:
-        """Number of logical requests currently bound/routed to this replica."""
-        if replica_id < 0 or replica_id >= len(self.clients) or self.clients[replica_id] is None:
-            return 0
         if self._hub is not None:
-            client = self.clients[replica_id]
             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 max(int(replica.queue_length), 0)
-        bound_legacy = sum(1 for bound_replica_id in self._request_bindings.values() if bound_replica_id == replica_id)
-        bound_affinity = sum(
-            1 for input_addr in self._affinity.values() if self._addr_to_replica_id.get(input_addr) == replica_id
-        )
-        return max(bound_legacy, bound_affinity)
+                        return (outputs_qsize, max(int(replica.queue_length), 0))
+            inflight = sum(
+                1 for input_addr in self._affinity.values() if self._addr_to_replica_id.get(input_addr) == replica_id
+            )
+        else:
+            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."""

From 8ba09ff65940cb2ee602fffc186fb932e8b19bdf Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Thu, 25 Jun 2026 13:30:05 +0800
Subject: [PATCH 17/21] update name

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 vllm_omni/engine/stage_pool.py                        |  2 ++
 vllm_omni/platforms/npu/worker/npu_ar_model_runner.py | 10 +++++-----
 vllm_omni/worker/gpu_ar_model_runner.py               | 10 +++++-----
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py
index 6539309abff..d93dbf33e84 100644
--- a/vllm_omni/engine/stage_pool.py
+++ b/vllm_omni/engine/stage_pool.py
@@ -451,10 +451,12 @@ def replica_monitor_sample(self, replica_id: int) -> tuple[int, int]:
                 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)
 
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 8b1bb3422c8..e98dcfb181f 100644
--- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
@@ -258,15 +258,15 @@ def _build_multimodal_outputs(
             return None
         if per_req_payloads is None:
             return None
-        wire: list[dict[str, torch.Tensor] | None] = []
+        wire_payloads: list[dict[str, torch.Tensor] | None] = []
         for payload in per_req_payloads:
             if not payload:
-                wire.append(None)
+                wire_payloads.append(None)
             else:
-                wire.append(_ensure_tensor_values(payload))
-        if all(item is None for item in wire):
+                wire_payloads.append(_ensure_tensor_values(payload))
+        if all(item is None for item in wire_payloads):
             return None
-        return wire
+        return wire_payloads
 
 
     def _request_final_stage_id(self, req_id: str) -> int | None:
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index 79b17dcbae2..68a9d7d9ec9 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -1338,15 +1338,15 @@ def _build_multimodal_outputs(
             return None
         if per_req_payloads is None:
             return None
-        wire: list[dict[str, torch.Tensor] | None] = []
+        wire_payloads: list[dict[str, torch.Tensor] | None] = []
         for payload in per_req_payloads:
             if not payload:
-                wire.append(None)
+                wire_payloads.append(None)
             else:
-                wire.append(_ensure_tensor_values(payload))
-        if all(item is None for item in wire):
+                wire_payloads.append(_ensure_tensor_values(payload))
+        if all(item is None for item in wire_payloads):
             return None
-        return wire
+        return wire_payloads
 
     def _snapshot_query_start_loc_cpu(self) -> Any:
         query_start_loc_cpu = self.query_start_loc.cpu

From 70625109ffb53ff3ca099433006fb6c91085fb76 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Thu, 25 Jun 2026 13:35:23 +0800
Subject: [PATCH 18/21] add try catch

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 tests/engine/test_orchestrator_monitor.py | 10 ++++++++++
 vllm_omni/engine/orchestrator_monitor.py  |  6 +++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/tests/engine/test_orchestrator_monitor.py b/tests/engine/test_orchestrator_monitor.py
index 6a1f11d47f8..ee9126cf470 100644
--- a/tests/engine/test_orchestrator_monitor.py
+++ b/tests/engine/test_orchestrator_monitor.py
@@ -77,6 +77,16 @@ def test_register_replica_backfills_prior_windows(tmp_path, monkeypatch):
     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))
diff --git a/vllm_omni/engine/orchestrator_monitor.py b/vllm_omni/engine/orchestrator_monitor.py
index debd689dd8c..8be903a88dd 100644
--- a/vllm_omni/engine/orchestrator_monitor.py
+++ b/vllm_omni/engine/orchestrator_monitor.py
@@ -146,7 +146,11 @@ def _roll_window_if_needed(self, now_mono: float) -> None:
             self._roll_window(now_mono)
 
     def _roll_window(self, now_mono: float) -> None:
-        sampled = self._replica_sampler()
+        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)

From 63426df57b9ff98257360d18143cb7c2e0cdbff7 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Thu, 25 Jun 2026 14:10:15 +0800
Subject: [PATCH 19/21] fix

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 vllm_omni/platforms/npu/worker/npu_ar_model_runner.py | 6 +++++-
 vllm_omni/worker/gpu_ar_model_runner.py               | 6 +++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

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 e98dcfb181f..dff639cc5d2 100644
--- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
@@ -1140,7 +1140,11 @@ def _unwrap_lists(v):
                     payload.update(mm_payload)
                 pooler_output.append(flatten_payload(payload))
 
-        pooler_inter, pooler_client = partition_payload_list(pooler_output or [])
+        pooler_output = pooler_output or []
+        if self.model_config.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(
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index 68a9d7d9ec9..7b61cb33d57 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -1649,7 +1649,11 @@ def _build_omni_model_runner_output_from_snapshot(
                     )
                     pooler_output.append(flatten_payload(payload))
 
-        pooler_inter, pooler_client = partition_payload_list(pooler_output or [])
+        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"):

From 4601afbc32b641ec06d9c75b8ceb01081836ff2e Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Sun, 28 Jun 2026 16:35:18 +0800
Subject: [PATCH 20/21] fix ci

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 vllm_omni/platforms/npu/worker/npu_ar_model_runner.py | 5 +----
 vllm_omni/worker/gpu_ar_model_runner.py               | 5 +----
 2 files changed, 2 insertions(+), 8 deletions(-)

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 dff639cc5d2..41777e32b70 100644
--- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
@@ -1141,10 +1141,7 @@ def _unwrap_lists(v):
                 pooler_output.append(flatten_payload(payload))
 
         pooler_output = pooler_output or []
-        if self.model_config.async_chunk:
-            pooler_inter, pooler_client = partition_payload_list(pooler_output)
-        else:
-            pooler_inter, pooler_client = None, pooler_output
+        pooler_inter, pooler_client = partition_payload_list(pooler_output)
         inter_stage_outputs = self._build_multimodal_outputs(pooler_inter)
         multimodal_outputs = self._build_multimodal_outputs(pooler_client)
         model_runner_output = OmniModelRunnerOutput(
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index 7b61cb33d57..0ac2b6a3bf2 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -1650,10 +1650,7 @@ def _build_omni_model_runner_output_from_snapshot(
                     pooler_output.append(flatten_payload(payload))
 
         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
+        pooler_inter, pooler_client = partition_payload_list(pooler_output)
 
         if pooler_inter and self._should_accumulate_full_payload_output():
             with record_function_or_nullcontext("omni_output_builder:accumulate_full_payload_output"):

From 5cec9c318da79057bc0c7efbf13a51e3dacf3336 Mon Sep 17 00:00:00 2001
From: yinpeiqi <yinpeiqi809@gmail.com>
Date: Mon, 29 Jun 2026 01:43:08 +0800
Subject: [PATCH 21/21] fix

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
---
 tests/worker/test_gpu_ar_model_runner.py               | 10 ++++++----
 tests/worker/test_gpu_generation_model_runner.py       |  1 +
 vllm_omni/platforms/npu/worker/npu_ar_model_runner.py  |  6 +++++-
 .../npu/worker/npu_generation_model_runner.py          |  9 ++++++++-
 vllm_omni/worker/gpu_ar_model_runner.py                |  6 +++++-
 vllm_omni/worker/gpu_generation_model_runner.py        |  6 +++++-
 6 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py
index dcbe92b8b16..3088a3dffa2 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
@@ -399,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=[])
@@ -432,13 +434,13 @@ 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 not None
-    assert torch.equal(output.inter_stage_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]]))
+    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.inter_stage_outputs[1]["hidden"],
+        output.multimodal_outputs[1]["hidden"],
         torch.tensor([[2.0, 20.0], [3.0, 30.0]]),
     )
-    assert output.multimodal_outputs is None
 
 
 def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch):
diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py
index be7fadcb643..0cd9aeb2683 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/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
index 41777e32b70..2ed47ee1598 100644
--- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
@@ -110,6 +110,7 @@ def __init__(self, *args, **kwargs):
         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):
@@ -1141,7 +1142,10 @@ def _unwrap_lists(v):
                 pooler_output.append(flatten_payload(payload))
 
         pooler_output = pooler_output or []
-        pooler_inter, pooler_client = partition_payload_list(pooler_output)
+        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(
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 170589e97e7..54a846c432b 100644
--- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
@@ -41,6 +41,10 @@
 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:
@@ -472,7 +476,10 @@ def sample_tokens(
         else:
             raise RuntimeError("Unsupported diffusion output type")
 
-        inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads)
+        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()
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index 0ac2b6a3bf2..0bf5cfc01f0 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -288,6 +288,7 @@ def __init__(self, *args, **kwargs):
         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
@@ -1650,7 +1651,10 @@ def _build_omni_model_runner_output_from_snapshot(
                     pooler_output.append(flatten_payload(payload))
 
         pooler_output = pooler_output or []
-        pooler_inter, pooler_client = partition_payload_list(pooler_output)
+        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"):
diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py
index f75d54bd80f..a8e96bf83f9 100644
--- a/vllm_omni/worker/gpu_generation_model_runner.py
+++ b/vllm_omni/worker/gpu_generation_model_runner.py
@@ -55,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",
@@ -455,7 +456,10 @@ def sample_tokens(
         else:
             raise RuntimeError("Unsupported diffusion output type")
 
-        inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads)
+        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()