[Bugfix][Metrics] Measure the codec stage TTFT and phase split after talker first audio

docs/design/metrics.md

@@ -285,13 +285,15 @@ With `--log-stats`, a stage served by a vLLM engine core (AR and generation stag
 | Field | Interval |
 | ----- | -------- |
 | `vllm_queued_ms` | `QUEUED` event to first `SCHEDULED` event |
-| `vllm_prefill_ms` | First `SCHEDULED` event to first engine-core output (the first token for an AR stage) |
+| `vllm_prefill_ms` | First `SCHEDULED` event to first engine-core output |
 | `vllm_decode_ms` | First engine-core output to last engine-core output |
 | `vllm_num_preemptions` | Number of `PREEMPTED` events |
 
 The intervals match upstream `FinishedRequestStats`, so time spent preempted stays inside prefill or decode. The wrapped `vllm:request_queue_time_seconds`, `vllm:request_prefill_time_seconds`, `vllm:request_decode_time_seconds`, and `vllm:request_num_preemptions` histograms, which already carry `{stage, replica}` labels, use the same intervals. Two differences: upstream subtracts an unobserved timestamp as `0.0`, and it reports a negative interval as is, while these fields leave either interval out. No Prometheus family is added.
 
-The first engine-core output ends prefill even when it carries no token. A generation stage that emits one output when it stops therefore reports its whole run as `vllm_prefill_ms`, with a measured `vllm_decode_ms` of `0`.
+For an AR stage that emits an output every step, the first engine-core output carries its first token. The first engine-core output ends prefill even when it carries no token. A generation stage that emits one output when it stops therefore reports its whole run as `vllm_prefill_ms`, with a measured `vllm_decode_ms` of `0`. The same holds for an AR stage that does not emit per-step outputs: with Model Runner V2 and `async_chunk`, a non-final AR stage (for example the Qwen3-TTS talker) emits one output when it stops, so its `vllm_prefill_ms` also covers its decode steps, and its `vllm_ttft_ms` is the time to that output.
+
+With talker first audio (`talker_first_audio`), the orchestrator feeds the upstream stage's first audio frame into the codec stage's output processor. That frame is not an engine-core output of the codec stage and does not end its prefill: the codec stage's split, its `vllm_ttft_ms`, and the samples it adds to the wrapped vLLM histograms start at the codec's own first engine-core output. Codec outputs that arrive before the frame are held until it has been processed, then processed in arrival order, each with the engine-core timestamp of the batch it arrived in, so the split and the inter-token latencies leave out the wait for the frame. `vllm_ttft_ms` and its histogram sample include that wait: as in vLLM, TTFT runs from the request's arrival to the time the output processor handles the first output, which for a held output is after the frame.
 
 `vllm_queued_ms` runs from the stage's engine core accepting the request to its first scheduling. For stage 0 this is scheduler wait. A downstream stage can also wait there for its input:
 

tests/engine/test_orchestrator_first_audio.py

@@ -8,7 +8,7 @@
 import pytest
 import torch
 from vllm.sampling_params import RequestOutputKind
-from vllm.v1.engine import FinishReason
+from vllm.v1.engine import EngineCoreEvent, EngineCoreEventType, FinishReason
 
 from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
 from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
@@ -60,9 +60,9 @@ def get_bound_replica_id(self, request_id):
     def schedule_release_request_resources(self, request_ids):
         pass
 
-    async def process_llm_raw_outputs(self, replica_id, raw, **kwargs):
+    async def process_llm_raw_outputs(self, replica_id, raw, iteration_stats=None):
         assert replica_id == self.replica_id
-        return self.output_processor.process_outputs(raw.outputs, raw.timestamp).request_outputs
+        return self.output_processor.process_outputs(raw.outputs, raw.timestamp, iteration_stats).request_outputs
 
 
 def _orchestrator(output_kind=RequestOutputKind.DELTA, *, registered=True, replica_id=2, final_stage_id=1):
@@ -274,6 +274,90 @@ async def test_first_audio_flush_uses_set_with_real_raw_terminal_handlers(termin
     assert messages[-1].finished is terminal
 
 
+@pytest.mark.asyncio
+async def test_codec_stats_start_at_its_own_first_output_after_first_audio():
+    # The first-audio frame reaches the codec's output processor without
+    # IterationStats, ahead of the codec's own engine-core outputs. The codec
+    # stage's TTFT and queue/prefill/decode split must still be measured.
+    obj = _orchestrator()
+    obj._prom_metrics = None
+    obj._stat_logger = Mock()
+    obj._stage_replica_to_engine_idx = {}
+    obj._handle_kv_ready_raw_outputs = AsyncMock()
+    obj._apply_raw_terminal_stage_finish = AsyncMock(return_value=False)
+    obj._report_duplex_session_request_error = AsyncMock()
+    obj._process_llm_stage_outputs = Orchestrator._process_llm_stage_outputs.__get__(obj)
+
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
+    queued_and_scheduled = [
+        EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0),
+        EngineCoreEvent(EngineCoreEventType.SCHEDULED, 1.25),
+    ]
+    for timestamp, events, terminal in [(1.75, queued_and_scheduled, False), (2.25, None, False), (3.0, None, True)]:
+        raw = _raw(terminal=terminal)
+        raw.timestamp = timestamp
+        raw.outputs[0].events = events
+        await obj._process_llm_stage_outputs(1, 2, raw, set())
+
+    record = obj.stage_pools[1].output_processor.pop_native_text_metrics("r")
+    assert {key: record.get(key) for key in ("vllm_queued_ms", "vllm_prefill_ms", "vllm_decode_ms")} == pytest.approx(
+        {"vllm_queued_ms": 250.0, "vllm_prefill_ms": 500.0, "vllm_decode_ms": 1250.0}
+    )
+    assert record["vllm_ttft_ms"] > 0
+    recorded = [call.args[1] for call in obj._stat_logger.record.call_args_list]
+    assert sum(len(stats.time_to_first_tokens_iter) for stats in recorded) == 1
+    assert [itl for stats in recorded for itl in stats.inter_token_latencies_iter] == pytest.approx([0.5, 0.75])
+
+
+@pytest.mark.asyncio
+async def test_codec_outputs_parked_for_first_audio_keep_engine_core_timestamps(mocker):
+    # Codec outputs that overtake the first-audio frame are parked and replayed
+    # after it. The replay keeps the engine-core timestamp of the batch each
+    # output arrived in, so the codec's split and ITLs leave out the wait for
+    # the frame. TTFT is taken when the output is processed, as in vLLM, so it
+    # is measured at replay and includes that wait.
+    obj = _orchestrator()
+    obj._prom_metrics = None
+    obj._stat_logger = Mock()
+    obj._stage_replica_to_engine_idx = {}
+    obj._handle_kv_ready_raw_outputs = AsyncMock()
+    obj._apply_raw_terminal_stage_finish = AsyncMock(return_value=False)
+    obj._report_duplex_session_request_error = AsyncMock()
+    obj._process_llm_stage_outputs = Orchestrator._process_llm_stage_outputs.__get__(obj)
+    # Wall clock read by IterationStats; the request arrived at 100.0.
+    clock = SimpleNamespace(now=101.0)
+    mocker.patch("vllm.v1.metrics.stats.time", SimpleNamespace(time=lambda: clock.now))
+    obj.stage_pools[1].output_processor.request_states["r"].native_text_stats.arrival_time = 100.0
+
+    queued_and_scheduled = [
+        EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0),
+        EngineCoreEvent(EngineCoreEventType.SCHEDULED, 1.25),
+    ]
+    for timestamp, events, required in [(1.75, queued_and_scheduled, True), (2.25, None, False)]:
+        raw = _raw(required=required)
+        raw.timestamp = timestamp
+        raw.outputs[0].events = events
+        await obj._process_llm_stage_outputs(1, 2, raw, set())
+    assert len(obj.request_states["r"].pending_first_audio_outputs) == 2
+    obj._stat_logger.record.assert_not_called()
+
+    clock.now = 104.0
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
+    raw = _raw(terminal=True)
+    raw.timestamp = 3.0
+    await obj._process_llm_stage_outputs(1, 2, raw, set())
+
+    record = obj.stage_pools[1].output_processor.pop_native_text_metrics("r")
+    assert {key: record.get(key) for key in ("vllm_queued_ms", "vllm_prefill_ms", "vllm_decode_ms")} == pytest.approx(
+        {"vllm_queued_ms": 250.0, "vllm_prefill_ms": 500.0, "vllm_decode_ms": 1250.0}
+    )
+    assert record["vllm_itls_ms"] == pytest.approx([500.0, 750.0])
+    assert record["vllm_ttft_ms"] == pytest.approx(4000.0)
+    recorded = [call.args[1] for call in obj._stat_logger.record.call_args_list]
+    assert [ttft for stats in recorded for ttft in stats.time_to_first_tokens_iter] == pytest.approx([4.0])
+    assert [itl for stats in recorded for itl in stats.inter_token_latencies_iter] == pytest.approx([0.5, 0.75])
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("processing_failed", [False, True])
 async def test_prewarm_flushes_first_audio_after_successful_submission(mocker, processing_failed):

tests/engine/test_output_processor.py

@@ -155,7 +155,7 @@ def test_native_text_metrics_include_segment_generation_token_count(monkeypatch)
     processor._native_text_metrics_by_request = {}
     processor.lora_states = {}
     state = _make_state(RequestOutputKind.DELTA)
-    state.is_prefilling = False
+    state.native_is_prefilling = False
     iteration_stats = MagicMock()
 
     def update_native_stats(_output, _timestamp, _was_prefilling, native_stats, *_args):
@@ -208,14 +208,12 @@ def update_segment(_output, timestamp, _was_prefilling, native_stats, *_args):
             arrival_time=10.01,
         )
     )
-    state.is_prefilling = True
     processor._update_stats_from_output(
         state,
         MagicMock(),
         10.02,
         iteration_stats,
     )
-    state.is_prefilling = False
     processor._update_stats_from_output(
         state,
         MagicMock(),
@@ -224,6 +222,8 @@ def update_segment(_output, timestamp, _was_prefilling, native_stats, *_args):
     )
     second_segment = processor.pop_native_text_metrics("r")
 
+    # The streaming update starts a new prefill for the second segment.
+    assert [call.args[2] for call in iteration_stats.update_from_output.call_args_list] == [True, True, False]
     assert second_segment["num_generation_tokens"] == 2
     assert second_segment["vllm_itls_ms"] == pytest.approx([10.0])
     assert second_segment["vllm_tpot_ms"] == pytest.approx(10.0)
@@ -244,9 +244,7 @@ def update_segment(_output, timestamp, _was_prefilling, native_stats, *_args):
         native_stats.first_token_latency = 0.01
 
     iteration_stats.update_from_output.side_effect = update_segment
-    state.is_prefilling = True
     processor._update_stats_from_output(state, MagicMock(), 10.0, iteration_stats)
-    state.is_prefilling = False
     processor._update_stats_from_output(state, MagicMock(), 10.03, iteration_stats)
     record = processor.pop_native_text_metrics("r")
 
@@ -274,7 +272,6 @@ def test_coalesced_outputs_skip_itl_but_preserve_token_and_finish_stats(is_coale
         output = OmniEngineCoreOutput(request_id="r", new_token_ids=tokens, is_coalesced=coalesced)
         output = MsgpackDecoder(OmniEngineCoreOutput).decode(MsgpackEncoder().encode(output))
         processor._update_stats_from_output(state, output, timestamp, stats)
-        state.is_prefilling = False
     processor._update_stats_from_finished(state, FinishReason.STOP, stats)
     record = processor.pop_native_text_metrics("r")
     expected_itls = [10.0] if is_coalesced else [160.0, 10.0]
@@ -1246,30 +1243,52 @@ def test_one_shot_generation_phase_split_reports_whole_run_as_prefill():
     )
 
 
-def test_first_output_without_stats_leaves_prefill_and_decode_out():
-    # The orchestrator flushes a codec stage's first audio frame without
-    # iteration stats, so the native first-token timestamp is never set.
-    processor = _stats_processor(_make_no_detok_state(RequestOutputKind.CUMULATIVE))
-    processor.process_outputs([_stats_output()], engine_core_timestamp=1.5, iteration_stats=None)
-    processor.process_outputs(
-        [
+@pytest.mark.parametrize(
+    "make_state",
+    [_make_text_stage_state, lambda: _make_no_detok_state(RequestOutputKind.CUMULATIVE)],
+    ids=["ar_text_path", "mm_only_path"],
+)
+def test_output_processed_without_stats_does_not_end_prefill(make_state):
+    # With talker first audio (async_chunk), the orchestrator feeds the
+    # upstream first-audio frame into the codec stage without IterationStats,
+    # before the codec's own first engine-core output. The frame clears
+    # ``is_prefilling`` but must not end the codec stage's prefill.
+    processor = _stats_processor(make_state())
+    processor.process_outputs([_stats_output(new_token_ids=())], engine_core_timestamp=0.0, iteration_stats=None)
+    steps = [
+        (
+            1.75,
             _stats_output(
                 events=[
                     EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0),
                     EngineCoreEvent(EngineCoreEventType.SCHEDULED, 1.25),
-                ],
-                finish_reason=FinishReason.STOP,
-            )
-        ],
-        engine_core_timestamp=2.0,
-        iteration_stats=IterationStats(),
-    )
+                ]
+            ),
+        ),
+        (2.25, _stats_output()),
+        (3.0, _stats_output(finish_reason=FinishReason.STOP)),
+    ]
+    ttfts, itls = [], []
+    for timestamp, output in steps:
+        iteration_stats = IterationStats()
+        iteration_stats.iteration_timestamp = timestamp  # one clock, so TTFT is exact
+        processor.process_outputs([output], engine_core_timestamp=timestamp, iteration_stats=iteration_stats)
+        ttfts += iteration_stats.time_to_first_tokens_iter
+        itls += iteration_stats.inter_token_latencies_iter
 
     record = processor.pop_native_text_metrics("r")
 
-    assert {key: record[key] for key in _PHASE_KEYS if key in record} == pytest.approx(
-        {"vllm_queued_ms": 250.0, "vllm_num_preemptions": 0}
+    assert _phase_split(record) == pytest.approx(
+        {"vllm_queued_ms": 250.0, "vllm_prefill_ms": 500.0, "vllm_decode_ms": 1250.0, "vllm_num_preemptions": 0}
     )
+    assert record["vllm_ttft_ms"] == pytest.approx(1750.0)
+    # These IterationStats also feed the wrapped vLLM histograms: one TTFT
+    # sample, and no ITL measured from a last-token timestamp of 0.
+    assert ttfts == pytest.approx([1.75])
+    assert itls == pytest.approx([0.5, 0.75])
+    (finished,) = iteration_stats.finished_requests
+    assert finished.prefill_time == pytest.approx(0.5)
+    assert finished.decode_time == pytest.approx(1.25)
 
 
 @pytest.mark.parametrize(

vllm_omni/engine/orchestrator.py

@@ -276,7 +276,8 @@ class OrchestratorRequestState:
     # First-frame delivery and final-stage outputs use different channels.
     upstream_first_audio: bool = False
     pending_upstream_first_audio: tuple[int, OmniEngineCoreOutput] | None = None
-    pending_first_audio_outputs: list[tuple[int, int, Any]] = field(default_factory=list)
+    # (stage_id, replica_id, output, engine-core timestamp of its batch)
+    pending_first_audio_outputs: list[tuple[int, int, Any, float]] = field(default_factory=list)
 
 
 @dataclass
@@ -1012,7 +1013,7 @@ async def _route_upstream_first_audio(self, stage_id: int, replica_id: int, raw_
                 and (requires_first or req_state.pending_first_audio_outputs)
                 and getattr(eco, "finish_reason", None) != FinishReason.ERROR
             ):
-                req_state.pending_first_audio_outputs.append((stage_id, replica_id, eco))
+                req_state.pending_first_audio_outputs.append((stage_id, replica_id, eco, raw_outputs.timestamp))
                 continue
             kept.append(eco)
         raw_outputs.outputs = kept
@@ -1054,12 +1055,17 @@ async def _flush_upstream_first_audio(self, req_state: OrchestratorRequestState)
         if self.request_states.get(req_state.request_id) is not req_state:
             return
         ready, req_state.pending_first_audio_outputs = req_state.pending_first_audio_outputs, []
-        for pending_stage, pending_replica, output in ready:
+        for pending_stage, pending_replica, output, timestamp in ready:
             if self.request_states.get(req_state.request_id) is not req_state:
                 break
             terminal: set[str] = set()
+            # Pass the timestamp of the batch the output arrived in; without
+            # one, the wrapper is stamped at replay time.
             processed = await self._process_llm_stage_outputs(
-                pending_stage, pending_replica, OmniEngineCoreOutputs(outputs=[output]), terminal
+                pending_stage,
+                pending_replica,
+                OmniEngineCoreOutputs(outputs=[output], timestamp=timestamp),
+                terminal,
             )
             await self._handle_processed_outputs(pending_stage, pending_replica, processed)
             await self._finish_raw_terminal_requests(pending_stage, pending_replica, terminal)

vllm_omni/outputs/output_processor.py

@@ -134,6 +134,11 @@ def __init__(
         arrival_time = kwargs.get("arrival_time")
         super().__init__(*args, **kwargs)
         self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0))
+        # Prefill flag of native_text_stats. vLLM clears ``is_prefilling`` on
+        # every processed output, also on one processed without IterationStats
+        # (the upstream first-audio frame fed into a codec stage), which would
+        # record the stage's first engine-core output as a decode step.
+        self.native_is_prefilling = True
         # Omni-specific: multimodal output accumulation
         # TODO: mm_type is per-request, not per-key. If a model ever produces
         # both audio and latent outputs, the modality type would flip on each
@@ -145,6 +150,7 @@ def __init__(
     def apply_streaming_update(self, update) -> None:
         super().apply_streaming_update(update)
         self.native_text_stats = RequestStateStats(arrival_time=float(update.arrival_time or 0.0))
+        self.native_is_prefilling = True
 
     def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None:
         """Accumulate a multimodal tensor payload into the request state.
@@ -834,7 +840,6 @@ def _update_stats_from_output(
         engine_core_timestamp: float | None,
         iteration_stats: IterationStats | None,
     ):
-        was_prefilling = req_state.is_prefilling
         native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None
         previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0
         previous_num_generation_tokens = int(native_stats.num_generation_tokens) if native_stats is not None else 0
@@ -861,6 +866,10 @@ def _update_stats_from_output(
         if iteration_stats is None or engine_core_timestamp is None or native_stats is None:
             return
 
+        assert isinstance(req_state, OmniRequestState)
+        # Not ``req_state.is_prefilling``: see ``OmniRequestState.native_is_prefilling``.
+        was_prefilling = req_state.native_is_prefilling
+        req_state.native_is_prefilling = False
         is_coalesced = getattr(engine_core_output, "is_coalesced", False) is True
         num_itls_before = len(iteration_stats.inter_token_latencies_iter) if is_coalesced else 0
         iteration_stats.update_from_output(