[Core] Add LLM-stage queue/prefill/decode split to StageRequestStats

docs/design/metrics.md

@@ -272,8 +272,43 @@ After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency,
 
 ## Logging vs. Prometheus
 
-`OrchestratorAggregator` (in `vllm_omni/metrics/stats.py`) is the logging-oriented metrics path. It collects detailed per-request, per-stage, and per-transfer statistics and prints formatted tables to the `INFO` log. This is designed for development and debugging — individual request traces, transfer bandwidth, inter-stage timing.
+`OrchestratorAggregator` (in `vllm_omni/metrics/stats.py`) is the logging-oriented metrics path. It collects detailed per-request, per-stage, and per-transfer statistics and prints formatted tables to the log (per-request tables at `DEBUG`, the `[OmniTiming]` summary at `INFO`). This is designed for development and debugging — individual request traces, transfer bandwidth, inter-stage timing.
 
 `OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics` form the Prometheus-oriented path. They record aggregate counters, gauges, and histograms suitable for time-series monitoring and alerting. Both paths share the same source data (`StageRequestStats`, `TransferEdgeStats`) — `OrchestratorAggregator.record_transfer_tx/rx` in particular calls both the existing accumulator code and the Prometheus emit hook in the same method body. The two consumption models can run simultaneously without coupling.
 
 The separation follows upstream vLLM's pattern of `LoggingStatLogger` vs. `PrometheusStatLogger` — same underlying data, different consumption models.
+
+### LLM-stage phase split in `StageRequestStats`
+
+With `--log-stats`, a stage served by a vLLM engine core (AR and generation stages) adds the engine-core phase split of the finished request to its `StageRequestStats` row:
+
+| 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_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`.
+
+`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:
+
+- With `async_chunk`, a downstream stage that receives chunks is submitted together with stage 0, and the request stays in `WAITING_FOR_CHUNK` until the first upstream chunk arrives. Later waits for the next chunk fall inside `vllm_decode_ms`.
+- A stage with `requires_full_payload_input` stays in `WAITING_FOR_INPUT` until the connector delivers the payload.
+
+For these stages `vllm_queued_ms` includes input wait, and with `async_chunk` so does `vllm_decode_ms`; neither measures scheduler contention alone. The split does not separate input wait. `vllm_queued_ms` is also separate from the orchestration-layer wait in `vllm_omni:request_queue_wait_s` (`pipeline_timings["queue_wait_ms"]`) and from the diffusion scheduler wait in `vllm_omni:stage_in_queue_s`.
+
+These intervals can overlap with other timings of the same request: the upstream stage's run, the transfer times in `TransferEdgeStats` and `vllm_omni:transfer_tx_s` / `vllm_omni:transfer_rx_s`, and stage-to-stage handoff events. Do not add them together into an end-to-end breakdown of a request unless the intervals are known not to overlap and use the same clock.
+
+Two more cases where an interval covers more than its name suggests:
+
+- A KV-transfer sender (a stage with `kv_transfer_criteria`) keeps the request running until the KV extraction is acknowledged, and the token-less `kv_ready` output it emits still advances the last-token timestamp. `vllm_decode_ms` therefore includes the wait for that acknowledgement.
+- `vllm_num_preemptions` counts `PREEMPTED` events. In the default (non-windowed) path, when the chunk transfer adapter moves running requests over `max_num_seqs` back to the waiting queue, it sets `PREEMPTED` without recording an event, so that wait is not counted as a preemption and stays inside prefill or decode.
+
+The per-request `[StageRequestStats]` table is logged at `DEBUG` (for example `VLLM_LOGGING_LEVEL=DEBUG`); at `INFO` only the `[OmniTiming]` line is printed.
+
+Missing and zero are different. A field is `None` when its interval was not observed: a diffusion stage, `--log-stats` off, or an engine-core event that never arrived. A measured `0` stays `0`. In the `[StageRequestStats]` table a missing value prints as `None`, and, as for every other field, a row whose values are all zero or missing is hidden.
+
+For a streaming-input request, only the output that finishes the request carries the split. On an AR stage, vLLM's output processor finishes the request only after the input stream has ended; at each earlier segment stop it applies the next streaming update instead, which restarts the request stats, so the split covers the last input segment. The last segment's `QUEUED` event can arrive with the previous segment's final output, before that restart, so its `vllm_queued_ms` may be missing. A generation stage finishes on the first output that is not a segment stop and does not restart its stats between segments, so its split spans the whole request: `vllm_prefill_ms` is the first segment's, `vllm_decode_ms` runs from the first segment's first output to the last output, gaps between segments included, and `vllm_queued_ms` is missing when a later segment was queued after the first scheduling. The fields are not added to the `stage_metrics` snapshot returned to clients and benchmarks. The duplex per-response table leaves them out (`DUPLEX_STAGE_TABLE_EXCLUDE`): they describe one engine-core request, not a response turn.

tests/engine/duplex/test_engine_session.py

@@ -949,6 +949,104 @@ def _capture(self: OrchestratorAggregator) -> dict[str, object]:
     assert [event.serving_time_to_first_output_ms for event in events] == pytest.approx([80.631, 418.829, 487.109])
 
 
+def test_end_response_omits_engine_core_phase_split_from_stage_table(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    phase_fields = ("vllm_queued_ms", "vllm_prefill_ms", "vllm_decode_ms", "vllm_num_preemptions")
+    original = OrchestratorAggregator.build_and_log_summary
+    captured: list[tuple[OrchestratorAggregator, dict[str, object]]] = []
+
+    def _capture(self: OrchestratorAggregator) -> dict[str, object]:
+        summary = original(self)
+        captured.append((self, summary))
+        return summary
+
+    monkeypatch.setattr(OrchestratorAggregator, "build_and_log_summary", _capture)
+    session = _session(num_stages=2, log_stats=True)
+    response_id = session.begin_response()
+    for stage_id in (0, 1):
+        stats = _stage_stats(stage_id=stage_id)
+        stats.vllm_queued_ms = 1.5
+        stats.vllm_prefill_ms = 20.0
+        stats.vllm_decode_ms = 300.0
+        stats.vllm_num_preemptions = 0
+        session.observe_stage_request_stats(stage_id, stats)
+    session.end_response()
+
+    assert len(captured) == 1
+    aggregator, summary = captured[0]
+    stage_table = summary.get("stage_table")
+    assert isinstance(stage_table, list) and stage_table
+    rows = stage_table[0]["stages"]
+    assert isinstance(rows, list)
+    assert [row["stage_id"] for row in rows] == [0, 1]
+    assert all(field not in row for row in rows for field in phase_fields)
+    assert all("vllm_ttft_ms" in row for row in rows)
+    events = aggregator.stage_events[response_id]
+    assert [event.vllm_prefill_ms for event in events] == pytest.approx([20.0, 20.0])
+
+
+def test_end_response_fold_keeps_the_chunk_that_carries_the_phase_split(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    phase_fields = ("vllm_queued_ms", "vllm_prefill_ms", "vllm_decode_ms", "vllm_num_preemptions")
+    logged: list[OrchestratorAggregator] = []
+
+    def _record(self: OrchestratorAggregator) -> dict[str, object]:
+        logged.append(self)
+        return {}
+
+    monkeypatch.setattr(OrchestratorAggregator, "build_and_log_summary", _record)
+    session = _session(num_stages=3, log_stats=True)
+    response_id = session.begin_response()
+
+    def _chunk(stage_id: int, split: tuple[float | None, float | None, float | None, int | None]):
+        stats = _stage_stats(stage_id=stage_id)
+        for field, value in zip(phase_fields, split):
+            setattr(stats, field, value)
+        return stats
+
+    no_split = (None, None, None, None)
+    # Stage 0: the split arrives on the last chunk, after the first (template) chunk.
+    session.observe_stage_request_stats(0, _chunk(0, no_split))
+    session.observe_stage_request_stats(0, _chunk(0, (1.5, 20.0, 300.0, 0)))
+    # Stage 1: a later chunk without a split does not erase it, and the split
+    # is taken from one chunk, not merged field by field.
+    session.observe_stage_request_stats(1, _chunk(1, (2.0, 30.0, 400.0, 1)))
+    session.observe_stage_request_stats(1, _chunk(1, (None, 35.0, 450.0, 0)))
+    session.observe_stage_request_stats(1, _chunk(1, no_split))
+    # Stage 2: no chunk carries a split.
+    session.observe_stage_request_stats(2, _chunk(2, no_split))
+    session.observe_stage_request_stats(2, _chunk(2, no_split))
+    session.end_response()
+
+    rows = logged[0].stage_events[response_id]
+    assert [row.stage_id for row in rows] == [0, 1, 2]
+    assert [tuple(getattr(row, field) for field in phase_fields) for row in rows] == [
+        (1.5, 20.0, 300.0, 0),
+        (None, 35.0, 450.0, 0),
+        no_split,
+    ]
+
+
+def test_http_stage_table_keeps_engine_core_phase_split() -> None:
+    agg = OrchestratorAggregator(num_stages=1, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=0)
+    stats = _stage_stats(stage_id=0)
+    stats.vllm_queued_ms = 1.5
+    stats.vllm_prefill_ms = 20.0
+    stats.vllm_decode_ms = 300.0
+    stats.vllm_num_preemptions = 0
+    agg.on_stage_metrics(0, "r1", stats)
+    agg.on_finalize_request(0, "r1", req_start_ts=0.0)
+
+    summary = agg.build_and_log_summary()
+    stage_table = summary.get("stage_table")
+    assert isinstance(stage_table, list) and stage_table
+    row = stage_table[0]["stages"][0]
+    assert row["vllm_prefill_ms"] == pytest.approx(20.0)
+    assert row["vllm_num_preemptions"] == 0
+
+
 def test_http_stage_table_keeps_serving_time_to_first_output() -> None:
     agg = OrchestratorAggregator(num_stages=2, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=1)
     agg.on_stage_metrics(0, "r1", _stage_stats(stage_id=0, serving_time_to_first_output_ms=80.0))

tests/engine/test_orchestrator.py

@@ -1810,6 +1810,71 @@ def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]:
     assert metrics.finish_reason == "length"
 
 
+@pytest.mark.parametrize(
+    ("native_record", "expected"),
+    [
+        pytest.param(
+            {"vllm_queued_ms": 2.5, "vllm_prefill_ms": 30.0, "vllm_decode_ms": 0.0, "vllm_num_preemptions": 0},
+            {"vllm_queued_ms": 2.5, "vllm_prefill_ms": 30.0, "vllm_decode_ms": 0.0, "vllm_num_preemptions": 0},
+            id="observed",
+        ),
+        pytest.param(
+            # An absent key is an unobserved phase and must not become 0.
+            {"vllm_decode_ms": 12.0},
+            {"vllm_queued_ms": None, "vllm_prefill_ms": None, "vllm_decode_ms": 12.0, "vllm_num_preemptions": None},
+            id="partly_unobserved",
+        ),
+    ],
+)
+def test_stage_pool_metrics_copy_native_phase_split(native_record, expected) -> None:
+    class PhaseMetricsOutputProcessor(FakeOutputProcessor):
+        def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]:
+            assert request_id == "req-ar"
+            return {"num_generation_tokens": 2, **native_record}
+
+    pool = StagePool(0, [FakeStageClient(stage_type="llm")], output_processor=PhaseMetricsOutputProcessor())
+
+    metrics = pool.build_stage_metrics(
+        [_build_request_output("req-ar")],
+        submit_ts=time.time(),
+        request_timestamp=time.time(),
+        replica_id=0,
+    )
+
+    assert {key: getattr(metrics, key) for key in expected} == expected
+
+
+def test_stage_pool_metrics_skip_native_phase_split_when_metrics_disabled() -> None:
+    popped: list[str] = []
+
+    class PhaseMetricsOutputProcessor(FakeOutputProcessor):
+        def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]:
+            popped.append(request_id)
+            return {
+                "num_generation_tokens": 2,
+                "vllm_queued_ms": 2.5,
+                "vllm_prefill_ms": 30.0,
+                "vllm_decode_ms": 12.0,
+                "vllm_num_preemptions": 1,
+            }
+
+    pool = StagePool(0, [FakeStageClient(stage_type="llm")], output_processor=PhaseMetricsOutputProcessor())
+    pool.collect_stage_metrics = False
+
+    metrics = pool.build_stage_metrics(
+        [_build_request_output("req-ar")],
+        submit_ts=time.time(),
+        request_timestamp=time.time(),
+        replica_id=0,
+    )
+
+    # The record is still popped so it does not leak, but the optional split stays unset.
+    assert popped == ["req-ar"]
+    assert metrics.num_tokens_out == 2
+    assert metrics.vllm_queued_ms is None and metrics.vllm_prefill_ms is None
+    assert metrics.vllm_decode_ms is None and metrics.vllm_num_preemptions is None
+
+
 def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None:
     stage = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
     pool = StagePool(
@@ -1830,6 +1895,9 @@ def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None:
 
     assert metrics.serving_time_to_first_output_ms == pytest.approx(12000.0)
     assert metrics.image_time_to_first_output_ms == pytest.approx(2000.0)
+    # A diffusion stage has no engine-core phase split.
+    assert metrics.vllm_queued_ms is None and metrics.vllm_prefill_ms is None
+    assert metrics.vllm_decode_ms is None and metrics.vllm_num_preemptions is None
 
 
 @pytest.mark.asyncio