[Metrics] Add audio SLOs + cross-stage transfer families + per-(stage, replica) wrap for upstream vllm:* (#3576)

Signed-off-by: vraiti <vraiti@redhat.com>
Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>
Co-authored-by: vraiti <vraiti@redhat.com>

docs/design/index.md

@@ -13,6 +13,10 @@ This section contains design documents and architecture specifications for vLLM-
 - [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
 - [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
 
+## Infrastructure Design Documents
+
+- [Prometheus Metrics](metrics.md)
+
 ## Module Design Documents
 
 - [AR Module](module/ar_module.md)

docs/design/metrics.md

@@ -0,0 +1,208 @@
+# Prometheus Metrics Design
+
+This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
+
+## Objectives
+
+- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
+- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
+- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
+- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
+- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
+
+## Background
+
+### Upstream vLLM Metrics
+
+Upstream vLLM defines ~37 Prometheus metric families under the `vllm:` prefix. These are registered by `PrometheusStatLogger` and cover engine-level state: KV cache usage, running/waiting request counts, token throughput, TTFT, inter-token latency, e2e latency, and so on. They are served via the `/metrics` HTTP endpoint provided by `prometheus_fastapi_instrumentator` and the default `prometheus_client` WSGI handler.
+
+vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` collector whose `_name` attribute contains the substring `"vllm"`. This runs during engine initialization to clean up stale collectors from prior instantiations within the same process.
+
+### The Problem
+
+vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+
+Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
+
+## Architecture
+
+### Component Overview
+
+```text
+                       +------------------------+
+                       |  API Server (FastAPI)  |
+                       |   GET /metrics         |
+                       +-----------+------------+
+                                   |
+                  prometheus_client default registry
+                                   |
+        +--------+--------+--------+--------+--------+
+        |                                            |
+   vllm:omni_*                                    vllm:*
+   collectors                                  collectors
+        |                                            |
+   +----+--------+   +-----------+   +----------+   +-----------+
+   | OmniPromet- |   | OmniMod-  |   | OmniTra- |   | OmniProm- |
+   | heusMetrics |   | alityMet- |   | nsferMe- |   | etheusSt- |
+   |             |   | rics      |   | trics    |   | atLogger  |
+   +----+--------+   +-----+-----+   +----+-----+   +----+------+
+        |                  |              |              |
+     OmniBase           OmniBase     Orchestrator    Orchestrator
+   (request life-     (finalize +   (record_trans-  (per-(stage,
+    cycle, success/    streaming     fer_tx/rx        replica)
+    fail counter)      hooks via     hooks via        scheduler/
+                       observe_*     emit hook in     iteration
+                       APIs)         OrchestratorAg-  stats)
+                                     gregator)
+```
+
+### Data Flow
+
+There are four independent paths for metric collection.
+
+**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
+
+`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
+
+- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
+- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
+- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
+
+**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
+
+`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
+
+- `observe_modality_at_finalize(...)` — called from `omni_base._process_single_result` inside the existing `e2e_done` finalize guard. For `output_type == "audio"` it emits `audio_frames_total`, `audio_duration_s`, `audio_rtf` (or `audio_skipped_requests_total{reason="no_audio_data"}` when no audio was produced). Sample rate is resolved from `engine_outputs.multimodal_output` via `definitions.resolve_audio_sample_rate(...)` (fallback chain mirrors `serving_chat.py`'s audio response path).
+- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
+- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
+
+**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
+
+`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
+
+The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transfer_rx`. After the existing `TransferEdgeStats` accumulation, the aggregator calls `_emit_transfer_tx` / `_emit_transfer_rx`. Those:
+
+1. Resolve `from_replica` / `to_replica` via a `replica_resolver` callback supplied by `async_omni.py`. The resolver delegates to `stage_pool.get_bound_replica_id(request_id)` — i.e. the orchestrator's existing sticky-routing binding is the source of truth.
+2. Convert the underlying `_ms` accumulators to seconds and call the `_s`-suffixed observe methods on `OmniTransferMetrics`.
+
+Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
+
+**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
+
+The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
+
+### Shared State Between Threads
+
+The Orchestrator runs in a background thread. The API server (OmniBase) runs in the asyncio event loop thread. `OmniRequestCounter` bridges them — a plain Python object with an `int` field. The Orchestrator increments/decrements it; the entrypoint reads it for gauge updates. No lock is needed because the counter is advisory (a stale read by one Prometheus scrape interval is acceptable). It is created by `AsyncOmniEngine.__init__()` and passed to the Orchestrator at construction time.
+
+### Metric Registration and Lifecycle
+
+All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
+
+The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
+
+## OmniPrometheusStatLogger Wrap
+
+Upstream `PrometheusStatLogger.__init__` hard-codes `labelnames = ["model_name", "engine"]` as a local variable, references it across ~37 metric-family construction sites, and uses the `engine` label value in five different `.labels()` call shapes (kwarg with int engine, kwarg with str engine, positional with str engine in the middle, plus a `metrics_info["engine"] = str(...)` dict pattern). To reshape `engine` into `stage` + `replica` without forking the entire upstream `__init__`, the wrap uses three coordinated mechanisms:
+
+1. **Class-level metric class slot overrides.** `OmniPrometheusStatLogger` overrides `_gauge_cls`, `_counter_cls`, `_histogram_cls` (which upstream calls via `self._gauge_cls(...)` etc.) with `_RelabelGauge` / `_RelabelCounter` / `_RelabelHistogram` wrapper classes. These intercept the `labelnames` kwarg at metric family creation time and replace `engine` with `("stage", "replica")`.
+2. **Property descriptor for `per_engine_labelvalues`.** Upstream builds `self.per_engine_labelvalues = {idx: [model_name, str(idx)]}` inside `__init__` and then captures it into a local variable for `create_metric_per_engine` calls. By making `per_engine_labelvalues` a Python property on the subclass, the setter intercepts upstream's assignment and rewrites each 2-tuple into a 3-tuple `[model_name, stage, replica]` using the `stage_replica_map` supplied at construction time. The captured local then sees the rewritten dict.
+3. **Override of `.labels()` on the wrapper classes.** For the five call sites that pass `engine` directly (kwarg or positional, int or str), `_RelabelMixin.labels()` translates the engine value back to `(stage, replica)` via a process-level `_ENGINE_INDEX_MAP` populated by `OmniPrometheusStatLogger.__init__`. This handles `gauge_engine_sleep_state.labels(engine=idx, ...)`, `counter_request_success_base.labels(model_name, str(idx), str(reason))`, `info_gauge.labels(**metrics_info)`, etc.
+
+The three sub-helpers that upstream `PrometheusStatLogger.__init__` constructs (`spec_decoding_prom` / `kv_connector_prom` / `perf_metrics_prom`) use their own `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots and would otherwise build families with the raw 2-element labelnames. `_OmniPerfMetricsProm` / `_OmniSpecDecodingProm` / `_OmniKVConnectorProm` subclass each helper to route the same relabel mixin through their internal family construction.
+
+The `Orchestrator` constructs `stage_replica_map` from the static `stage_pools` configuration at startup:
+
+```python
+stage_replica_map = {
+    flat_idx: (str(stage_id), str(replica_id))
+    for flat_idx, (stage_id, replica_id) in enumerate(
+        (s, r)
+        for s, pool in enumerate(stage_pools)
+        for r in range(pool.num_replicas)
+    )
+}
+```
+
+A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestrator so the per-replica `record(engine_idx=...)` call site can look up the right flat index.
+
+> Dynamic add/remove of replicas at runtime is intentionally out of scope — the upstream `PrometheusStatLogger` materializes per-engine_idx child metrics at init time, and supporting hot-add would require non-trivial intervention into upstream's per-family child dictionaries.
+
+## Gating: `--log-stats`
+
+All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
+
+Behavior with `--log-stats=off` (default):
+
+- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
+- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
+- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
+
+Behavior with `--log-stats=on`: all metric paths fire normally; the orchestrator's per-replica recording is bounded only by `OmniSchedulerMixin.make_stats()`'s per-scheduler 1 Hz throttle (see next section).
+
+The overhead with the flag on is small enough that an A/B benchmark on Qwen3-Omni-30B single replica (30 sequential audio requests) showed a mean latency delta of +0.6% (Welch's t = 0.318, n=30, not statistically significant at α=0.05); the `/metrics` line count drops from 1358 to 124 lines when the flag is off.
+
+## Throttling: `make_stats()` Override
+
+Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, returning a SchedulerStats object for the orchestrator. Under vLLM's architecture, this is fine. But since vLLM-Omni requires that the object be serialized and transferred over ZMQ, receiving a SchedulerStats object on every step can introduce unacceptable overhead to the system.
+
+`OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead.
+
+The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second.
+
+## Metric Definitions
+
+### Pipeline (4)
+
+| Metric | Type | Labels | Description |
+|--------|------|--------|-------------|
+| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
+| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
+| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
+| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
+
+### Audio (7)
+
+Labels: `{model_name, stage, replica}` plus the listed extra label.
+
+| Metric | Type | Extra label | Description |
+|--------|------|-------------|-------------|
+| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
+| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
+| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
+| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
+| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
+| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
+| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
+
+### Cross-stage transfer (4)
+
+Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
+
+| Metric | Type | Description |
+|--------|------|-------------|
+| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
+| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
+| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
+| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
+
+### LLM stage-level (wrapped `vllm:*`)
+
+After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — carries `{model_name, stage, replica}` labels. For the full upstream catalog see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md); note that metrics depending on features unsupported in vLLM-Omni (e.g. speculative decoding, LoRA) will not be available.
+
+## Naming Convention
+
+- All time-bearing metrics use the `_s` suffix (values in seconds). Two bucket families are used:
+  - `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values.
+  - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
+- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
+- Sizes use the `_bytes` suffix.
+- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+
+## 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.
+
+`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.

docs/usage/metrics.md

@@ -0,0 +1,98 @@
+# Production Metrics
+
+vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
+
+```bash
+vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
+curl http://localhost:8000/metrics
+```
+
+**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
+
+## Metric Namespaces
+
+| Prefix | Source | Present when |
+|--------|--------|--------------|
+| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
+| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
+| `http_` / `process_` | Uvicorn / Python runtime | Always |
+
+## Pipeline-Level Metrics (`vllm:omni_`)
+
+Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
+
+### Request counts
+
+| Metric | Type | Labels | Description |
+|--------|------|--------|-------------|
+| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
+| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
+| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
+
+### Latency
+
+| Metric | Type | Labels | Description |
+|--------|------|--------|-------------|
+| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
+
+## Audio Modality Metrics (`vllm:omni_`)
+
+Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
+
+| Metric | Type | Extra label | Description |
+|--------|------|-------------|-------------|
+| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
+| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
+| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
+| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
+| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
+| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
+| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
+
+The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
+
+## Cross-Stage Transfer Metrics (`vllm:omni_`)
+
+Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
+
+| Metric | Type | Description |
+|--------|------|-------------|
+| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
+| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
+| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
+| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
+
+## vLLM Engine Metrics (`vllm:`)
+
+When the pipeline includes an LLM stage, the upstream vLLM engine exposes its full set of ~37 metric families under the `vllm:` prefix.
+
+vLLM-Omni wraps the upstream `vllm.v1.metrics.loggers.PrometheusStatLogger` with `OmniPrometheusStatLogger` so that the original `engine` single label is reshaped into `stage` + `replica`. Every `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — therefore gains per-`(stage, replica)` visibility automatically. No omni-side duplicate is needed for the text path.
+
+```text
+# Before wrap:
+vllm:num_requests_running{model_name="...", engine="1"}              3.0
+
+# After wrap:
+vllm:num_requests_running{model_name="...", stage="1", replica="0"}  2.0
+vllm:num_requests_running{model_name="...", stage="1", replica="1"}  1.0
+```
+
+For the full list of upstream metrics, see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md).
+
+## Metric Availability by Pipeline Type
+
+| Metric group | Multi-stage LLM (Qwen3-Omni) |
+|---|---|
+| `vllm:omni_` request tracking + latency | With `--log-stats` |
+| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
+| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
+| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
+| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
+
+## Naming Convention
+
+- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
+- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
+- Sizes use the `_bytes` suffix.
+- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.

tests/engine/test_async_omni_engine_stage_init.py

@@ -312,7 +312,7 @@ def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_inpu
     monkeypatch.setattr(
         engine_mod,
         "build_llm_stage_output_processor",
-        lambda plan, _cfg: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor,
+        lambda plan, _cfg, **_kw: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor,
     )
     monkeypatch.setattr(engine_mod, "build_stage0_input_processor", lambda _cfg: top_level_input_processor)
 

tests/engine/test_orchestrator.py

@@ -1157,3 +1157,37 @@ async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator
         assert stage0_r1.add_request_calls[1][0].request_id == "parent-neg"
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
+
+
+def test_orchestrator_does_not_re_introduce_global_stats_throttle() -> None:
+    """Regression: each (stage, replica) must independently publish its wrapped
+    vllm:* stats when its scheduler emits non-None scheduler_stats.
+
+    A previous version of Orchestrator carried a global self._last_stats_ts /
+    _stats_interval_s gate around _stat_logger.record(). Because
+    OmniSchedulerMixin.make_stats() already throttles at 1 Hz per scheduler
+    (one per (stage, replica)), the extra global gate starved every replica
+    other than the first to emit within each second — their {stage, replica}
+    gauges/counters went stale.
+
+    The fix removed the global gate entirely; the only signal needed is
+    'this replica's scheduler emitted non-None scheduler_stats'. This test
+    fails loudly if someone reintroduces the global throttle.
+    """
+    import inspect
+
+    from vllm_omni.engine.orchestrator import Orchestrator
+
+    source = inspect.getsource(Orchestrator)
+    assert "_last_stats_ts" not in source, (
+        "Orchestrator must not gate stat recording on a global timestamp. "
+        "OmniSchedulerMixin.make_stats() already throttles per scheduler "
+        "(per (stage, replica)); an outer global gate starves all but the "
+        "first replica to emit within each 1s window."
+    )
+    assert "_stats_interval_s" not in source
+    assert "raw_outputs.scheduler_stats is not None" in source, (
+        "Orchestrator must gate stat recording solely on "
+        "raw_outputs.scheduler_stats being non-None — the per-scheduler 1Hz "
+        "throttle in OmniSchedulerMixin.make_stats() is the only gate needed."
+    )

tests/metrics/test_modality.py

@@ -0,0 +1,399 @@
+from __future__ import annotations
+
+import pytest
+from prometheus_client import REGISTRY, generate_latest
+
+from vllm_omni.metrics import definitions as defs
+from vllm_omni.metrics.modality import (
+    OmniModalityMetrics,
+    observe_audio_first_packet,
+    observe_audio_streaming_finalize,
+    observe_modality_at_finalize,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+_MODEL = "test-modality-model"
+
+
+@pytest.fixture(scope="module")
+def mod() -> OmniModalityMetrics:
+    return OmniModalityMetrics(model_name=_MODEL)
+
+
+def _sample_value(output: str, line_prefix: str) -> float | None:
+    for line in output.splitlines():
+        if line.startswith(line_prefix):
+            return float(line.split()[-1])
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Family registration — 7 audio families served from modality.py
+# ---------------------------------------------------------------------------
+
+
+_EXPECTED_FAMILIES = [
+    defs.AUDIO_TTFP_S,
+    defs.AUDIO_DURATION_S,
+    defs.AUDIO_RTF_METRIC,
+    defs.AUDIO_FRAMES_METRIC,
+    defs.AUDIO_UNDERRUN_S,
+    defs.AUDIO_CONTINUITY_OK_METRIC,
+    defs.AUDIO_SKIPPED_REQUESTS_METRIC,
+]
+
+
+class TestRegistration:
+    def test_all_locked_families_present(self, mod: OmniModalityMetrics) -> None:
+        # Trigger at least one observation per family so the registry exposes them.
+        mod.observe_audio_ttfp("s", "r", 0.1)
+        mod.observe_audio_duration("s", "r", 1.0)
+        mod.observe_audio_rtf("s", "r", 0.5)
+        mod.inc_audio_frames("s", "r", 1)
+        mod.observe_audio_underrun("s", "r", 0.01)
+        mod.inc_audio_continuity_ok("s", "r", 100)
+        mod.inc_audio_skipped("s", "r", "malformed_codec")
+
+        out = generate_latest(REGISTRY).decode()
+        for name in _EXPECTED_FAMILIES:
+            assert f"# HELP {name}" in out, f"missing family: {name}"
+
+
+# ---------------------------------------------------------------------------
+# Audio observe API
+# ---------------------------------------------------------------------------
+
+
+class TestAudio:
+    def test_audio_ttfp_observed(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_ttfp", "0"
+        mod.observe_audio_ttfp(stage, replica, 0.42)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_TTFP_S}_count{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 1.0
+
+    def test_audio_duration_observed(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_dur", "0"
+        mod.observe_audio_duration(stage, replica, 3.5)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_DURATION_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 3.5
+
+    def test_audio_rtf_observed(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_rtf", "0"
+        mod.observe_audio_rtf(stage, replica, 0.45)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_RTF_METRIC}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 0.45
+
+    def test_audio_frames_inc(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_frames", "0"
+        mod.inc_audio_frames(stage, replica, 240)
+        mod.inc_audio_frames(stage, replica, 60)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_FRAMES_METRIC}_total{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 300.0
+
+    def test_audio_frames_zero_or_negative_skipped(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_zero", "0"
+        mod.inc_audio_frames(stage, replica, 0)
+        mod.inc_audio_frames(stage, replica, -5)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_FRAMES_METRIC}_total{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) is None
+
+    def test_audio_underrun_observed(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_under", "0"
+        mod.observe_audio_underrun(stage, replica, 0.05)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_UNDERRUN_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 0.05
+
+    def test_audio_underrun_negative_clamped_to_zero(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_under_neg", "0"
+        mod.observe_audio_underrun(stage, replica, -0.1)
+        out = generate_latest(REGISTRY).decode()
+        prefix = f'{defs.AUDIO_UNDERRUN_S}_sum{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+        assert _sample_value(out, prefix) == 0.0
+
+    def test_audio_continuity_ok_carries_threshold_label(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_cont", "0"
+        mod.inc_audio_continuity_ok(stage, replica, threshold_ms=100)
+        mod.inc_audio_continuity_ok(stage, replica, threshold_ms=200)
+        out = generate_latest(REGISTRY).decode()
+        # Counter family auto-suffixes _total at exposition.
+        v100 = _sample_value(
+            out,
+            f'{defs.AUDIO_CONTINUITY_OK_METRIC}_total{{model_name="{_MODEL}",'
+            f'replica="{replica}",stage="{stage}",threshold_ms="100"}}',
+        )
+        v200 = _sample_value(
+            out,
+            f'{defs.AUDIO_CONTINUITY_OK_METRIC}_total{{model_name="{_MODEL}",'
+            f'replica="{replica}",stage="{stage}",threshold_ms="200"}}',
+        )
+        assert v100 == 1.0
+        assert v200 == 1.0
+
+    def test_audio_skipped_carries_reason_label(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_skip", "0"
+        mod.inc_audio_skipped(stage, replica, "malformed_codec")
+        mod.inc_audio_skipped(stage, replica, "")  # → "unknown"
+        out = generate_latest(REGISTRY).decode()
+        good = _sample_value(
+            out,
+            f'{defs.AUDIO_SKIPPED_REQUESTS_METRIC}_total{{model_name="{_MODEL}",'
+            f'reason="malformed_codec",replica="{replica}",stage="{stage}"}}',
+        )
+        unknown = _sample_value(
+            out,
+            f'{defs.AUDIO_SKIPPED_REQUESTS_METRIC}_total{{model_name="{_MODEL}",'
+            f'reason="unknown",replica="{replica}",stage="{stage}"}}',
+        )
+        assert good == 1.0
+        assert unknown == 1.0
+
+
+# ---------------------------------------------------------------------------
+# Stub-driven routing tests for the finalize / streaming helpers.
+# ---------------------------------------------------------------------------
+
+
+class _StubModMetrics:
+    def __init__(self):
+        self.calls: list[tuple] = []
+
+    def inc_audio_frames(self, s, r, n):
+        self.calls.append(("inc_audio_frames", s, r, n))
+
+    def observe_audio_duration(self, s, r, d):
+        self.calls.append(("observe_audio_duration", s, r, d))
+
+    def observe_audio_rtf(self, s, r, rtf):
+        self.calls.append(("observe_audio_rtf", s, r, rtf))
+
+    def observe_audio_underrun(self, s, r, u):
+        self.calls.append(("observe_audio_underrun", s, r, u))
+
+    def inc_audio_continuity_ok(self, s, r, threshold_ms):
+        self.calls.append(("inc_audio_continuity_ok", s, r, threshold_ms))
+
+    def inc_audio_skipped(self, s, r, reason):
+        self.calls.append(("inc_audio_skipped", s, r, reason))
+
+
+class _Bag:
+    def __init__(self, **kw):
+        self.__dict__.update(kw)
+
+
+class TestObserveModalityAtFinalize:
+    def test_audio_path_full(self):
+        stub = _StubModMetrics()
+        stage_metrics = _Bag(stage_gen_time_ms=500.0, audio_generated_frames=24000)
+        engine_outputs = _Bag(multimodal_output={"audio_sample_rate": 24000})
+
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=0,
+            stage_metrics=stage_metrics,
+            engine_outputs=engine_outputs,
+        )
+        # 24000 frames / 24000 Hz = 1.0s duration; gen 0.5s → rtf 0.5
+        assert ("inc_audio_frames", "1", "0", 24000) in stub.calls
+        assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
+        assert ("observe_audio_rtf", "1", "0", 0.5) in stub.calls
+        # Continuity/underrun NOT emitted from finalize — they come from the
+        # streaming hook because they need the per-chunk arrival timeline.
+        assert not any(c[0] == "observe_audio_underrun" for c in stub.calls)
+        assert not any(c[0] == "inc_audio_continuity_ok" for c in stub.calls)
+
+    def test_audio_path_zero_frames_emits_skipped(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=0,
+            stage_metrics=_Bag(stage_gen_time_ms=300.0, audio_generated_frames=0),
+            engine_outputs=_Bag(multimodal_output={}),
+        )
+        assert ("inc_audio_frames", "1", "0", 0) in stub.calls
+        assert not any(c[0] == "observe_audio_duration" for c in stub.calls)
+        assert not any(c[0] == "observe_audio_rtf" for c in stub.calls)
+        assert ("inc_audio_skipped", "1", "0", "no_audio_data") in stub.calls
+
+    def test_audio_uses_resolved_sample_rate_from_multimodal_output(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=0,
+            stage_metrics=_Bag(stage_gen_time_ms=1000.0, audio_generated_frames=16000),
+            engine_outputs=_Bag(multimodal_output={"sample_rate": 16000}),
+        )
+        assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
+
+    def test_non_audio_output_type_skipped(self):
+        # Image / video / text output types are out of scope for this module.
+        stub = _StubModMetrics()
+        for output_type in ("text", "image", "video"):
+            observe_modality_at_finalize(
+                stub,
+                output_type=output_type,
+                stage_id=0,
+                replica_id=0,
+                stage_metrics=_Bag(stage_gen_time_ms=100.0),
+                engine_outputs=_Bag(),
+            )
+        assert stub.calls == []
+
+    def test_replica_id_none_skipped(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=None,
+            stage_metrics=_Bag(stage_gen_time_ms=500.0, audio_generated_frames=240),
+            engine_outputs=_Bag(multimodal_output={}),
+        )
+        assert stub.calls == []
+
+    def test_stage_metrics_none_skipped(self):
+        stub = _StubModMetrics()
+        observe_modality_at_finalize(
+            stub,
+            output_type="audio",
+            stage_id=1,
+            replica_id=0,
+            stage_metrics=None,
+            engine_outputs=_Bag(multimodal_output={}),
+        )
+        assert stub.calls == []
+
+
+class TestObserveAudioFirstPacket:
+    def test_observes_with_valid_inputs(self):
+        stub = _StubModMetrics()
+        stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t))
+
+        observe_audio_first_packet(
+            stub,
+            stage_id=1,
+            replica_id=0,
+            arrival_ts=100.0,
+            now_ts=100.42,
+        )
+        assert stub.calls == [("observe_audio_ttfp", "1", "0", pytest.approx(0.42))]
+
+    def test_replica_none_skipped(self):
+        stub = _StubModMetrics()
+        stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t))
+        observe_audio_first_packet(stub, stage_id=1, replica_id=None, arrival_ts=100.0, now_ts=100.5)
+        assert stub.calls == []
+
+    def test_arrival_ts_zero_skipped(self):
+        stub = _StubModMetrics()
+        stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t))
+        observe_audio_first_packet(stub, stage_id=1, replica_id=0, arrival_ts=0.0, now_ts=100.5)
+        assert stub.calls == []
+
+    def test_clock_skew_clamped_to_zero(self):
+        stub = _StubModMetrics()
+        stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t))
+        observe_audio_first_packet(stub, stage_id=1, replica_id=0, arrival_ts=100.5, now_ts=100.0)
+        assert stub.calls == [("observe_audio_ttfp", "1", "0", 0.0)]
+
+
+class TestObserveAudioStreamingFinalize:
+    def test_continuous_run_emits_underrun_and_continuity_ok(self):
+        stub = _StubModMetrics()
+        # Three back-to-back chunks of 0.5s each at 24 kHz s16le mono:
+        # 24000 Hz * 2 bytes * 0.5s = 24000 bytes per chunk.
+        arrivals = [0.5, 1.0, 1.5]
+        bytes_ = [24000, 24000, 24000]
+        observe_audio_streaming_finalize(
+            stub,
+            stage_id=1,
+            replica_id=0,
+            chunk_arrival_times_s=arrivals,
+            chunk_bytes=bytes_,
+            sample_rate=24000,
+            threshold_s=0.1,
+        )
+        # 0 underrun -> observe with 0.0, continuity_ok incremented at threshold 100.
+        assert ("observe_audio_underrun", "1", "0", 0.0) in stub.calls
+        assert ("inc_audio_continuity_ok", "1", "0", 100) in stub.calls
+
+    def test_late_chunk_emits_nonzero_underrun_and_no_continuity_inc(self):
+        stub = _StubModMetrics()
+        # Two chunks but second arrives 1s late; underrun > 0.1s threshold.
+        arrivals = [0.0, 2.0]
+        bytes_ = [24000, 24000]
+        observe_audio_streaming_finalize(
+            stub,
+            stage_id=1,
+            replica_id=0,
+            chunk_arrival_times_s=arrivals,
+            chunk_bytes=bytes_,
+            sample_rate=24000,
+            threshold_s=0.1,
+        )
+        underrun_calls = [c for c in stub.calls if c[0] == "observe_audio_underrun"]
+        assert underrun_calls
+        assert underrun_calls[0][-1] > 0.1
+        assert not any(c[0] == "inc_audio_continuity_ok" for c in stub.calls)
+
+    def test_empty_arrivals_skipped(self):
+        stub = _StubModMetrics()
+        observe_audio_streaming_finalize(
+            stub,
+            stage_id=1,
+            replica_id=0,
+            chunk_arrival_times_s=[],
+            chunk_bytes=[],
+            sample_rate=24000,
+        )
+        assert stub.calls == []
+
+    def test_replica_none_skipped(self):
+        stub = _StubModMetrics()
+        observe_audio_streaming_finalize(
+            stub,
+            stage_id=1,
+            replica_id=None,
+            chunk_arrival_times_s=[0.5],
+            chunk_bytes=[24000],
+            sample_rate=24000,
+        )
+        assert stub.calls == []
+
+
+class TestBucketSelection:
+    def test_audio_rtf_uses_rtf_buckets(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_buckets", "0"
+        mod.observe_audio_rtf(stage, replica, 0.5)
+        out = generate_latest(REGISTRY).decode()
+        rtf_marker = f'{defs.AUDIO_RTF_METRIC}_bucket{{le="0.9"'
+        assert rtf_marker in out, "audio_rtf should use RTF_BUCKETS containing le=0.9"
+
+    def test_audio_ttfp_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_seconds", "0"
+        mod.observe_audio_ttfp(stage, replica, 0.1)
+        out = generate_latest(REGISTRY).decode()
+        sec_marker = f'{defs.AUDIO_TTFP_S}_bucket{{le="0.05"'
+        assert sec_marker in out, "audio_ttfp should use SECONDS_BUCKETS containing le=0.05"
+
+    def test_audio_underrun_uses_seconds_fast_buckets(self, mod: OmniModalityMetrics) -> None:
+        stage, replica = "talker_fast", "0"
+        mod.observe_audio_underrun(stage, replica, 0.005)
+        out = generate_latest(REGISTRY).decode()
+        # SECONDS_FAST_BUCKETS includes le=0.001 which is absent from SECONDS_BUCKETS.
+        fast_marker = f'{defs.AUDIO_UNDERRUN_S}_bucket{{le="0.001"'
+        assert fast_marker in out, "audio_underrun_s should use SECONDS_FAST_BUCKETS containing le=0.001"

tests/metrics/test_prometheus.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+import re
+
+import pytest
+from prometheus_client import REGISTRY, CollectorRegistry, generate_latest
+
+from vllm_omni.metrics import OmniPrometheusMetrics
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODEL = "test-model"
+
+_PIPELINE_METRICS = [
+    "vllm:omni_num_requests_running",
+    "vllm:omni_num_requests_waiting",
+    "vllm:omni_requests_success",
+    "vllm:omni_e2e_request_latency_s",
+]
+
+
+@pytest.fixture(scope="module")
+def registry() -> CollectorRegistry:
+    return REGISTRY
+
+
+@pytest.fixture(scope="module")
+def prom() -> OmniPrometheusMetrics:
+    return OmniPrometheusMetrics(model_name=_MODEL)
+
+
+@pytest.fixture(scope="module")
+def scrape_output(prom: OmniPrometheusMetrics, registry: CollectorRegistry) -> str:
+    # Two natural completions (stop) + one length-cap + one failure (abort)
+    # exercise three distinct finished_reason buckets in the merged Counter.
+    prom.request_succeeded(e2e_seconds=1.5, finished_reason="stop")
+    prom.request_succeeded(e2e_seconds=2.0, finished_reason="stop")
+    prom.request_succeeded(e2e_seconds=3.0, finished_reason="length")
+    prom.request_failed()  # → finished_reason="abort"
+    prom.set_running(5)
+    prom.set_waiting(2)
+    return generate_latest(registry).decode()
+
+
+def _sample_value(output: str, metric_line: str) -> float | None:
+    for line in output.splitlines():
+        if line.startswith(metric_line):
+            return float(line.split()[-1])
+    return None
+
+
+class TestMetricObservation:
+    def test_all_metric_families_present(self, scrape_output: str) -> None:
+        for name in _PIPELINE_METRICS:
+            assert f"# HELP {name}" in scrape_output, f"missing metric family: {name}"
+
+    def test_counter_values(self, scrape_output: str) -> None:
+        # Per-reason buckets sourced from the merged completion Counter.
+        stop = _sample_value(
+            scrape_output,
+            f'vllm:omni_requests_success_total{{finished_reason="stop",model_name="{_MODEL}"}}',
+        )
+        assert stop == 2.0
+
+        length = _sample_value(
+            scrape_output,
+            f'vllm:omni_requests_success_total{{finished_reason="length",model_name="{_MODEL}"}}',
+        )
+        assert length == 1.0
+
+        abort = _sample_value(
+            scrape_output,
+            f'vllm:omni_requests_success_total{{finished_reason="abort",model_name="{_MODEL}"}}',
+        )
+        assert abort == 1.0
+
+    def test_gauge_values(self, scrape_output: str) -> None:
+        running = _sample_value(
+            scrape_output,
+            f'vllm:omni_num_requests_running{{model_name="{_MODEL}"}}',
+        )
+        assert running == 5.0
+
+        waiting = _sample_value(
+            scrape_output,
+            f'vllm:omni_num_requests_waiting{{model_name="{_MODEL}"}}',
+        )
+        assert waiting == 2.0
+
+    def test_histogram_counts(self, scrape_output: str) -> None:
+        # 3 successful completions (stop x2 + length x1) all observe e2e;
+        # the 1 failed completion only increments the Counter without
+        # observing the latency histogram, so the count stays at 3.
+        e2e_count = _sample_value(
+            scrape_output,
+            f'vllm:omni_e2e_request_latency_s_count{{model_name="{_MODEL}"}}',
+        )
+        assert e2e_count == 3.0
+
+
+class TestLabelCorrectness:
+    def test_pipeline_metrics_carry_model_name(self, scrape_output: str) -> None:
+        for name in _PIPELINE_METRICS:
+            pattern = rf'^{re.escape(name)}.*model_name="{re.escape(_MODEL)}"'
+            assert re.search(pattern, scrape_output, re.MULTILINE), f"{name} missing model_name label"
+
+    def test_no_legacy_engine_label(self, scrape_output: str) -> None:
+        assert 'engine="' not in scrape_output
+
+    def test_no_legacy_seconds_or_ms_families(self, scrape_output: str) -> None:
+        # Renamed: *_time_ms → *_s; *_time_seconds dropped.
+        for legacy in (
+            "vllm:omni_request_queue_time_seconds",
+            "vllm:omni_e2e_request_latency_seconds",
+        ):
+            assert legacy not in scrape_output, f"legacy family {legacy} still registered"
+
+
+class TestScrapeOutput:
+    def test_omni_metrics_in_default_registry(self, scrape_output: str) -> None:
+        for name in _PIPELINE_METRICS:
+            assert name in scrape_output
+
+    def test_process_metrics_in_default_registry(self, scrape_output: str) -> None:
+        # vllm:* metrics require a full PrometheusStatLogger with VllmConfig
+        # and are registered by the Orchestrator at server startup. Verifying
+        # their presence is covered by integration tests. Here we confirm the
+        # default registry is being scraped by checking for process_* metrics
+        # from the Python prometheus_client runtime.
+        assert "process_" in scrape_output
+
+
+class TestRequestLifecycleGauges:
+    """Regression tests for the running/waiting gauge race at request finalize.
+
+    Before the fix, the per-stage publish in OmniBase._process_single_result
+    ran while the completed request was still in self.request_states even
+    though the orchestrator had already decremented _running_counter — so a
+    single completed request briefly exposed running=0, total=1 → waiting=1,
+    which stuck on /metrics until another request triggered a refresh.
+
+    Two-layer fix: _process_single_result excludes the finalizing request
+    from `total`, and _log_summary_and_cleanup republishes gauges after the
+    pop as a fallback. These tests pin the post-cleanup state.
+    """
+
+    def test_running_and_waiting_zero_after_request_completes(self, registry: CollectorRegistry) -> None:
+        from types import SimpleNamespace
+
+        from vllm_omni.entrypoints.omni_base import OmniBase
+        from vllm_omni.metrics.prometheus import OmniPrometheusMetrics, OmniRequestCounter
+
+        obj = object.__new__(OmniBase)
+        obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter())
+        obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test")
+        obj.request_states = {}
+        obj.log_stats = False
+
+        # Simulate request lifecycle: start (counter 0→1, dict {} → {req}),
+        # then finalize (counter 1→0, dict still holds the request because
+        # _log_summary_and_cleanup hasn't run yet — this is the exact race
+        # window the fix addresses).
+        obj.engine._running_counter.increment()
+        obj.request_states["req-1"] = SimpleNamespace(metrics=SimpleNamespace(e2e_done={"req-1"}))
+        obj.engine._running_counter.decrement()
+
+        obj._log_summary_and_cleanup("req-1")
+
+        assert obj.engine._running_counter.value == 0
+        assert len(obj.request_states) == 0
+        out = generate_latest(registry).decode()
+        assert _sample_value(out, 'vllm:omni_num_requests_running{model_name="lifecycle-test"}') == 0.0
+        assert _sample_value(out, 'vllm:omni_num_requests_waiting{model_name="lifecycle-test"}') == 0.0
+
+    def test_gauges_reflect_remaining_requests_after_one_completes(self, registry: CollectorRegistry) -> None:
+        from types import SimpleNamespace
+
+        from vllm_omni.entrypoints.omni_base import OmniBase
+        from vllm_omni.metrics.prometheus import OmniPrometheusMetrics, OmniRequestCounter
+
+        obj = object.__new__(OmniBase)
+        obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter())
+        obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test-2")
+        obj.request_states = {}
+        obj.log_stats = False
+
+        # Two in flight, one finalizes — running should report 1, waiting 0.
+        obj.engine._running_counter.increment()
+        obj.engine._running_counter.increment()
+        obj.request_states["req-1"] = SimpleNamespace(metrics=SimpleNamespace(e2e_done={"req-1"}))
+        obj.request_states["req-2"] = SimpleNamespace(metrics=SimpleNamespace(e2e_done=set()))
+        obj.engine._running_counter.decrement()
+
+        obj._log_summary_and_cleanup("req-1")
+
+        out = generate_latest(registry).decode()
+        assert _sample_value(out, 'vllm:omni_num_requests_running{model_name="lifecycle-test-2"}') == 1.0
+        assert _sample_value(out, 'vllm:omni_num_requests_waiting{model_name="lifecycle-test-2"}') == 0.0

tests/metrics/test_stat_logger.py

@@ -0,0 +1,456 @@
+from __future__ import annotations
+
+import pytest
+from prometheus_client import CollectorRegistry, generate_latest
+
+from vllm_omni.metrics.stat_logger import (
+    _ENGINE_INDEX_MAP,
+    OmniPrometheusStatLogger,
+    _engine_to_stage_replica,
+    _OmniKVConnectorProm,
+    _OmniPerfMetricsProm,
+    _OmniSpecDecodingProm,
+    _RelabelCounter,
+    _RelabelGauge,
+    _RelabelHistogram,
+    _rewrite_labelnames,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _isolate_engine_map():
+    """Each test gets a clean _ENGINE_INDEX_MAP."""
+    _ENGINE_INDEX_MAP.clear()
+    yield
+    _ENGINE_INDEX_MAP.clear()
+
+
+@pytest.fixture
+def registry() -> CollectorRegistry:
+    return CollectorRegistry()
+
+
+# ---------------------------------------------------------------------------
+# _rewrite_labelnames
+# ---------------------------------------------------------------------------
+
+
+class TestRewriteLabelnames:
+    def test_engine_at_end(self):
+        assert _rewrite_labelnames(["model_name", "engine"]) == [
+            "model_name",
+            "stage",
+            "replica",
+        ]
+
+    def test_engine_in_middle(self):
+        # Upstream uses `labelnames + ["reason"]` etc., putting engine in middle.
+        assert _rewrite_labelnames(["model_name", "engine", "reason"]) == [
+            "model_name",
+            "stage",
+            "replica",
+            "reason",
+        ]
+
+    def test_no_engine_label(self):
+        # Unaffected (e.g. omni's own families that don't use engine).
+        assert _rewrite_labelnames(["model_name"]) == ["model_name"]
+
+    def test_tuple_input_returns_tuple(self):
+        out = _rewrite_labelnames(("model_name", "engine"))
+        assert isinstance(out, tuple)
+        assert out == ("model_name", "stage", "replica")
+
+    def test_none_passthrough(self):
+        assert _rewrite_labelnames(None) is None
+
+
+# ---------------------------------------------------------------------------
+# _engine_to_stage_replica
+# ---------------------------------------------------------------------------
+
+
+class TestEngineToStageReplica:
+    def test_int_engine_value(self):
+        # Mirrors upstream `.labels(engine=idx, ...)` with int (loggers.py:510).
+        _ENGINE_INDEX_MAP[7] = ("talker", "1")
+        assert _engine_to_stage_replica(7) == ("talker", "1")
+
+    def test_str_engine_value(self):
+        # Mirrors upstream `metrics_info["engine"] = str(idx)` (loggers.py:1055).
+        _ENGINE_INDEX_MAP[2] = ("thinker", "0")
+        assert _engine_to_stage_replica("2") == ("thinker", "0")
+
+    def test_missing_engine_idx_raises(self):
+        # Empty map → fail-fast rather than emit a wrong (stage, replica).
+        with pytest.raises(KeyError):
+            _engine_to_stage_replica(999)
+
+
+# ---------------------------------------------------------------------------
+# Wrapper class behavior
+# ---------------------------------------------------------------------------
+
+
+class TestRelabelGauge:
+    def test_labelnames_rewritten_at_creation(self, registry):
+        g = _RelabelGauge(
+            name="omni_test_gauge",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        assert g._labelnames == ("model_name", "stage", "replica")
+
+    def test_labels_kwarg_translated(self, registry):
+        _ENGINE_INDEX_MAP[5] = ("diffusion", "0")
+        g = _RelabelGauge(
+            name="omni_test_gauge_kwarg",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        g.labels(engine=5, model_name="qwen-omni").set(42.0)
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_gauge_kwarg{model_name="qwen-omni",replica="0",stage="diffusion"} 42.0' in out
+
+    def test_labels_positional_passthrough(self, registry):
+        # Double-rewrite guard: the per_engine_labelvalues setter rewrites
+        # the values to 3-tuple [model_name, stage, replica] BEFORE
+        # create_metric_per_engine fans them into .labels(*values). The mixin
+        # must detect that args length already matches the rewritten
+        # labelnames and pass through, otherwise it would re-interpret
+        # args[engine_label_index] as an engine_idx and splice (stage, replica)
+        # again, blowing label count to 4.
+        g = _RelabelGauge(
+            name="omni_test_gauge_pos",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        g.labels("qwen-omni", "thinker", "0").set(7.0)
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_gauge_pos{model_name="qwen-omni",replica="0",stage="thinker"} 7.0' in out
+
+    def test_multiprocess_mode_kwarg_passthrough(self, registry):
+        # Upstream creates Gauges with multiprocess_mode="mostrecent" — must not
+        # be eaten by our mixin.
+        g = _RelabelGauge(
+            name="omni_test_gauge_mp",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            multiprocess_mode="mostrecent",
+            registry=registry,
+        )
+        assert g._multiprocess_mode == "mostrecent"
+
+
+class TestRelabelCounter:
+    def test_labelnames_rewritten(self, registry):
+        c = _RelabelCounter(
+            name="omni_test_counter",
+            documentation="test",
+            labelnames=["model_name", "engine", "finished_reason"],
+            registry=registry,
+        )
+        assert c._labelnames == (
+            "model_name",
+            "stage",
+            "replica",
+            "finished_reason",
+        )
+
+    def test_labels_kwarg_translated(self, registry):
+        _ENGINE_INDEX_MAP[2] = ("thinker", "0")
+        c = _RelabelCounter(
+            name="omni_test_counter_kwarg",
+            documentation="test",
+            labelnames=["model_name", "engine", "finished_reason"],
+            registry=registry,
+        )
+        c.labels(engine=2, model_name="m", finished_reason="stop").inc(3)
+
+        out = generate_latest(registry).decode()
+        assert (
+            'omni_test_counter_kwarg_total{finished_reason="stop",model_name="m",replica="0",stage="thinker"} 3.0'
+            in out
+        )
+
+
+class TestRelabelHistogram:
+    def test_labelnames_rewritten(self, registry):
+        h = _RelabelHistogram(
+            name="omni_test_histo",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        assert h._labelnames == ("model_name", "stage", "replica")
+
+    def test_labels_kwarg_translated_and_observe(self, registry):
+        _ENGINE_INDEX_MAP[0] = ("talker", "0")
+        h = _RelabelHistogram(
+            name="omni_test_histo_obs",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        h.labels(engine=0, model_name="m").observe(0.5)
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_histo_obs_count{model_name="m",replica="0",stage="talker"} 1.0' in out
+
+    def test_no_engine_label_unaffected(self, registry):
+        # Families without engine label (e.g. omni-side own metrics) pass through.
+        h = _RelabelHistogram(
+            name="omni_test_no_engine",
+            documentation="test",
+            labelnames=["model_name"],
+            registry=registry,
+        )
+        assert h._labelnames == ("model_name",)
+        h.labels(model_name="m").observe(1.0)
+
+
+# ---------------------------------------------------------------------------
+# Positional .labels() with engine value (loggers.py:646, 679)
+# ---------------------------------------------------------------------------
+
+
+class TestPositionalEngine:
+    def test_positional_engine_at_middle_index(self, registry):
+        # Mirrors `counter_prompt_tokens_by_source.labels(model_name, str(idx), source)`.
+        # Family original labelnames = ["model_name", "engine", "source"].
+        _ENGINE_INDEX_MAP[5] = ("talker", "0")
+        c = _RelabelCounter(
+            name="omni_test_pos_mid",
+            documentation="test",
+            labelnames=["model_name", "engine", "source"],
+            registry=registry,
+        )
+        c.labels("m", "5", "decoder").inc(2)
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_pos_mid_total{model_name="m",replica="0",source="decoder",stage="talker"} 2.0' in out
+
+    def test_positional_engine_with_int_value(self, registry):
+        # Defensive: positional form may also receive an int (we accept both).
+        _ENGINE_INDEX_MAP[3] = ("thinker", "1")
+        c = _RelabelCounter(
+            name="omni_test_pos_int",
+            documentation="test",
+            labelnames=["model_name", "engine", "reason"],
+            registry=registry,
+        )
+        c.labels("m", 3, "stop").inc()
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_pos_int_total{model_name="m",reason="stop",replica="1",stage="thinker"} 1.0' in out
+
+
+# ---------------------------------------------------------------------------
+# String-form engine kwarg (loggers.py:1056 info_gauge)
+# ---------------------------------------------------------------------------
+
+
+class TestStrEngineKwarg:
+    def test_engine_kwarg_str_form(self, registry):
+        # Mirrors `info_gauge.labels(**metrics_info)` where metrics_info["engine"]="0".
+        _ENGINE_INDEX_MAP[0] = ("thinker", "0")
+        g = _RelabelGauge(
+            name="omni_test_info",
+            documentation="test",
+            labelnames=["cache_size", "engine"],
+            multiprocess_mode="mostrecent",
+            registry=registry,
+        )
+        # Upstream pattern: pass everything as kwargs from the metrics_info dict.
+        g.labels(cache_size="big", engine="0").set(1)
+
+        out = generate_latest(registry).decode()
+        assert 'omni_test_info{cache_size="big",replica="0",stage="thinker"} 1.0' in out
+
+
+# ---------------------------------------------------------------------------
+# Child metric does not re-trigger relabel logic
+# ---------------------------------------------------------------------------
+
+
+class TestChildNoRecursion:
+    def test_child_set_does_not_relookup(self, registry):
+        # Once .labels() returns a child, subsequent .set()/.inc() must not
+        # consult _ENGINE_INDEX_MAP again. We verify by clearing the map
+        # AFTER labels() and proving .set() still works.
+        _ENGINE_INDEX_MAP[4] = ("diffusion", "0")
+        g = _RelabelGauge(
+            name="omni_test_child",
+            documentation="test",
+            labelnames=["model_name", "engine"],
+            registry=registry,
+        )
+        child = g.labels(engine=4, model_name="m")
+        _ENGINE_INDEX_MAP.clear()  # would break a second .labels() lookup
+        child.set(99.0)  # but set() is on the bound child — no map needed
+
+        # Re-populate so generate_latest doesn't trip on anything else.
+        _ENGINE_INDEX_MAP[4] = ("diffusion", "0")
+        out = generate_latest(registry).decode()
+        assert 'omni_test_child{model_name="m",replica="0",stage="diffusion"} 99.0' in out
+
+
+# ---------------------------------------------------------------------------
+# OmniPrometheusStatLogger — focused on the wrap mechanics (full
+# PrometheusStatLogger init requires a real VllmConfig and is exercised by
+# the orchestrator integration tests).
+# ---------------------------------------------------------------------------
+
+
+class TestOmniPrometheusStatLogger:
+    def test_class_slots_point_to_wrappers(self):
+        # Upstream's __init__ uses self._gauge_cls(...) etc. when constructing
+        # families; class-level slot override is how we inject the relabel logic.
+        assert OmniPrometheusStatLogger._gauge_cls is _RelabelGauge
+        assert OmniPrometheusStatLogger._counter_cls is _RelabelCounter
+        assert OmniPrometheusStatLogger._histogram_cls is _RelabelHistogram
+
+    def test_per_engine_labelvalues_setter_rewrites_to_3tuple(self):
+        # Construct via __new__ to skip the upstream PrometheusStatLogger __init__
+        # (which needs a real VllmConfig). We only verify the property descriptor.
+        sl = OmniPrometheusStatLogger.__new__(OmniPrometheusStatLogger)
+        sl._stage_replica_map = {
+            0: ("thinker", "0"),
+            1: ("talker", "0"),
+            2: ("talker", "1"),
+        }
+
+        # Mirror upstream's loggers.py:433 assignment shape.
+        sl.per_engine_labelvalues = {
+            0: ["my-model", "0"],
+            1: ["my-model", "1"],
+            2: ["my-model", "2"],
+        }
+
+        # Getter should return the 3-tuple form for downstream
+        # create_metric_per_engine consumers.
+        assert sl.per_engine_labelvalues == {
+            0: ["my-model", "thinker", "0"],
+            1: ["my-model", "talker", "0"],
+            2: ["my-model", "talker", "1"],
+        }
+
+    def test_per_engine_labelvalues_getter_returns_internal_dict(self):
+        sl = OmniPrometheusStatLogger.__new__(OmniPrometheusStatLogger)
+        sl._stage_replica_map = {0: ("thinker", "0")}
+        sl._omni_per_engine_labelvalues = {0: ["m", "thinker", "0"]}
+        assert sl.per_engine_labelvalues == {0: ["m", "thinker", "0"]}
+
+    def test_stage_replica_map_property_exposed(self):
+        sl = OmniPrometheusStatLogger.__new__(OmniPrometheusStatLogger)
+        srm = {0: ("thinker", "0"), 1: ("diffusion", "0")}
+        sl._stage_replica_map = srm
+        assert sl.stage_replica_map is srm
+
+    def test_init_populates_engine_index_map(self):
+        # Simulate the bookkeeping portion of __init__ (clear + update) without
+        # calling super, since super needs a real VllmConfig.
+        _ENGINE_INDEX_MAP[99] = ("stale", "stale")  # leftover from prior
+        srm = {0: ("thinker", "0"), 1: ("talker", "0")}
+
+        # Manually invoke the bookkeeping the way __init__ does it.
+        _ENGINE_INDEX_MAP.clear()
+        _ENGINE_INDEX_MAP.update(srm)
+
+        assert dict(_ENGINE_INDEX_MAP) == srm
+        assert 99 not in _ENGINE_INDEX_MAP  # old entry was cleared
+
+
+# ---------------------------------------------------------------------------
+# Helper-class wraps for upstream's spec_decoding / kv_connector /
+# perf_metrics sub-collectors. Without these, OmniPrometheusStatLogger
+# crashes at startup with `Incorrect label count` because each helper builds
+# its internal Counter/Gauge/Histogram families with raw 2-element labelnames
+# (passed via constructor arg) while consuming the rewritten 3-element
+# per_engine_labelvalues from the property descriptor.
+# ---------------------------------------------------------------------------
+
+
+class TestHelperClassWraps:
+    def test_perf_metrics_wrap_routes_through_relabel_counter(self):
+        assert _OmniPerfMetricsProm._counter_cls is _RelabelCounter
+
+    def test_spec_decoding_wrap_routes_through_relabel_counter(self):
+        assert _OmniSpecDecodingProm._counter_cls is _RelabelCounter
+
+    def test_kv_connector_wrap_routes_through_all_three_relabel_classes(self):
+        # KVConnector lets each connector build any of Gauge/Counter/Histogram,
+        # so all three slots must be intercepted.
+        assert _OmniKVConnectorProm._gauge_cls is _RelabelGauge
+        assert _OmniKVConnectorProm._counter_cls is _RelabelCounter
+        assert _OmniKVConnectorProm._histogram_cls is _RelabelHistogram
+
+    def test_omni_logger_slots_point_to_helper_subclasses(self):
+        # Upstream's PrometheusStatLogger.__init__ instantiates each sub-helper
+        # via `self._<name>_cls(...)`, so the slot overrides on the omni
+        # subclass are what routes through to the relabel mixin.
+        assert OmniPrometheusStatLogger._perf_metrics_cls is _OmniPerfMetricsProm
+        assert OmniPrometheusStatLogger._spec_decoding_cls is _OmniSpecDecodingProm
+        assert OmniPrometheusStatLogger._kv_connector_cls is _OmniKVConnectorProm
+
+
+# ---------------------------------------------------------------------------
+# Double-rewrite guard. The mixin's positional-args path used to
+# unconditionally splice (stage, replica) at engine_label_index. Because
+# per_engine_labelvalues is rewritten to 3-tuples *before* feeding into
+# create_metric_per_engine, that splice would run a second time on the
+# already-rewritten values and blow the label count to 4. The guard now
+# detects len(args) == len(self._labelnames) and short-circuits to passthrough.
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def fresh_registry():
+    return CollectorRegistry()
+
+
+class TestDoubleRewriteGuard:
+    def test_pre_rewritten_3tuple_passes_through(self, fresh_registry):
+        # 2-label original → 3-label rewritten family. Caller passes 3 values
+        # (the rewritten shape) and they should land verbatim, not get
+        # re-spliced.
+        _ENGINE_INDEX_MAP.clear()
+        _ENGINE_INDEX_MAP[0] = ("0", "0")
+        _ENGINE_INDEX_MAP[1] = ("1", "0")
+        g = _RelabelGauge(
+            name="dr_pre_rewritten",
+            documentation="t",
+            labelnames=["model_name", "engine"],
+            registry=fresh_registry,
+        )
+        # 3 positional args matching the rewritten 3-label family.
+        g.labels("m", "1", "0").set(42)
+        out = generate_latest(fresh_registry).decode()
+        assert 'dr_pre_rewritten{model_name="m",replica="0",stage="1"} 42.0' in out
+
+    def test_legacy_2tuple_with_extra_label_still_splices(self, fresh_registry):
+        # 3-label original (engine in middle) → 4-label rewritten family.
+        # Caller passes 3 values matching the ORIGINAL labelnames (the
+        # gauge_waiting_by_reason / counter_request_success pattern from
+        # upstream loggers.py:646, 679). The mixin must splice
+        # (stage, replica) at engine's position to reach the 4-label family.
+        _ENGINE_INDEX_MAP.clear()
+        _ENGINE_INDEX_MAP[1] = ("1", "0")
+        c = _RelabelCounter(
+            name="dr_legacy_with_extra",
+            documentation="t",
+            labelnames=["model_name", "engine", "reason"],
+            registry=fresh_registry,
+        )
+        # 3 positional args matching ORIGINAL labelnames (model_name,
+        # engine_str, reason).
+        c.labels("m", "1", "stop").inc(3)
+        out = generate_latest(fresh_registry).decode()
+        assert 'dr_legacy_with_extra_total{model_name="m",reason="stop",replica="0",stage="1"} 3.0' in out

tests/metrics/test_transfer.py

@@ -0,0 +1,346 @@
+from __future__ import annotations
+
+import pytest
+from prometheus_client import REGISTRY, generate_latest
+
+from vllm_omni.metrics import definitions as defs
+from vllm_omni.metrics.stats import OrchestratorAggregator, StageRequestStats, StageStats
+from vllm_omni.metrics.transfer import OmniTransferMetrics
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+_MODEL = "test-transfer-model"
+
+
+@pytest.fixture(scope="module")
+def tx() -> OmniTransferMetrics:
+    return OmniTransferMetrics(model_name=_MODEL)
+
+
+def _sample_value(output: str, line_prefix: str) -> float | None:
+    for line in output.splitlines():
+        if line.startswith(line_prefix):
+            return float(line.split()[-1])
+    return None
+
+
+_EXPECTED_FAMILIES = [
+    defs.TRANSFER_SIZE_BYTES,
+    defs.TRANSFER_TX_S,
+    defs.TRANSFER_RX_S,
+    defs.TRANSFER_IN_FLIGHT_S,
+]
+
+
+# ---------------------------------------------------------------------------
+# Family registration
+# ---------------------------------------------------------------------------
+
+
+class TestRegistration:
+    def test_all_four_families_present(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_size(0, 0, 1, 0, 1024)
+        tx.observe_tx_time(0, 0, 1, 0, 0.005)
+        tx.observe_rx_time(0, 0, 1, 0, 0.008)
+        tx.observe_in_flight_time(0, 0, 1, 0, 0.002)
+
+        out = generate_latest(REGISTRY).decode()
+        for name in _EXPECTED_FAMILIES:
+            assert f"# HELP {name}" in out, f"missing family: {name}"
+
+    def test_legacy_ms_family_names_dropped(self, tx: OmniTransferMetrics) -> None:
+        out = generate_latest(REGISTRY).decode()
+        for legacy in (
+            "vllm:omni_transfer_tx_time_ms",
+            "vllm:omni_transfer_rx_decode_time_ms",
+            "vllm:omni_transfer_in_flight_time_ms",
+        ):
+            assert legacy not in out, f"legacy family {legacy} still registered"
+
+
+# ---------------------------------------------------------------------------
+# Observe APIs
+# ---------------------------------------------------------------------------
+
+
+class TestObserveSize:
+    def test_size_observed_with_correct_labels(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_size(2, 0, 3, 1, 65536)
+        out = generate_latest(REGISTRY).decode()
+        prefix = (
+            f"{defs.TRANSFER_SIZE_BYTES}_sum"
+            f'{{from_replica="0",from_stage="2",model_name="{_MODEL}",'
+            f'to_replica="1",to_stage="3"}}'
+        )
+        assert _sample_value(out, prefix) == 65536.0
+
+
+class TestObserveTxTime:
+    def test_tx_time_observed_in_seconds(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_tx_time(2, 1, 3, 0, 0.0125)
+        out = generate_latest(REGISTRY).decode()
+        prefix = (
+            f"{defs.TRANSFER_TX_S}_sum"
+            f'{{from_replica="1",from_stage="2",model_name="{_MODEL}",'
+            f'to_replica="0",to_stage="3"}}'
+        )
+        assert _sample_value(out, prefix) == pytest.approx(0.0125)
+
+
+class TestObserveRxTime:
+    def test_rx_time_observed_in_seconds(self, tx: OmniTransferMetrics) -> None:
+        # Use a (from, to) combo not shared with TestRegistration's setup
+        # observe on the rx_s Histogram, otherwise the per-test sum picks
+        # up accumulated values from earlier observes.
+        tx.observe_rx_time(3, 0, 4, 0, 0.0042)
+        out = generate_latest(REGISTRY).decode()
+        prefix = (
+            f"{defs.TRANSFER_RX_S}_sum"
+            f'{{from_replica="0",from_stage="3",model_name="{_MODEL}",'
+            f'to_replica="0",to_stage="4"}}'
+        )
+        assert _sample_value(out, prefix) == pytest.approx(0.0042)
+
+
+class TestObserveInFlightTime:
+    def test_in_flight_time_observed_in_seconds(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_in_flight_time(3, 0, 4, 0, 0.0017)
+        out = generate_latest(REGISTRY).decode()
+        prefix = (
+            f"{defs.TRANSFER_IN_FLIGHT_S}_sum"
+            f'{{from_replica="0",from_stage="3",model_name="{_MODEL}",'
+            f'to_replica="0",to_stage="4"}}'
+        )
+        assert _sample_value(out, prefix) == pytest.approx(0.0017)
+
+
+# ---------------------------------------------------------------------------
+# Multi (from, to) cardinality
+# ---------------------------------------------------------------------------
+
+
+class TestCardinality:
+    def test_multiple_edges_produce_independent_series(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_size(5, 0, 6, 0, 100)
+        tx.observe_size(5, 0, 6, 1, 200)
+        tx.observe_size(5, 1, 6, 0, 300)
+
+        out = generate_latest(REGISTRY).decode()
+
+        prefix_a = (
+            f"{defs.TRANSFER_SIZE_BYTES}_sum"
+            f'{{from_replica="0",from_stage="5",model_name="{_MODEL}",'
+            f'to_replica="0",to_stage="6"}}'
+        )
+        prefix_b = (
+            f"{defs.TRANSFER_SIZE_BYTES}_sum"
+            f'{{from_replica="0",from_stage="5",model_name="{_MODEL}",'
+            f'to_replica="1",to_stage="6"}}'
+        )
+        prefix_c = (
+            f"{defs.TRANSFER_SIZE_BYTES}_sum"
+            f'{{from_replica="1",from_stage="5",model_name="{_MODEL}",'
+            f'to_replica="0",to_stage="6"}}'
+        )
+        assert _sample_value(out, prefix_a) == 100.0
+        assert _sample_value(out, prefix_b) == 200.0
+        assert _sample_value(out, prefix_c) == 300.0
+
+
+# ---------------------------------------------------------------------------
+# Bucket selection
+# ---------------------------------------------------------------------------
+
+
+class TestBucketSelection:
+    def test_size_uses_bytes_buckets(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_size(7, 0, 8, 0, 4096)
+        out = generate_latest(REGISTRY).decode()
+        # prometheus_client renders bucket boundary as float (e.g. ``1024.0``)
+        # regardless of the declared int in BYTES_BUCKETS.
+        marker = f'{defs.TRANSFER_SIZE_BYTES}_bucket{{from_replica="0",from_stage="7",le="1024.0"'
+        assert marker in out
+
+    def test_time_families_use_seconds_fast_buckets(self, tx: OmniTransferMetrics) -> None:
+        tx.observe_tx_time(7, 0, 8, 0, 0.001)
+        out = generate_latest(REGISTRY).decode()
+        # SECONDS_FAST_BUCKETS has le=0.001 — distinctive vs SECONDS_BUCKETS (no 0.001).
+        marker = f'{defs.TRANSFER_TX_S}_bucket{{from_replica="0",from_stage="7",le="0.001"'
+        assert marker in out
+
+
+# ---------------------------------------------------------------------------
+# OrchestratorAggregator emit hook
+# ---------------------------------------------------------------------------
+
+
+class _StubTransferEmitter:
+    """Records every observe_* call so the hook routing can be asserted
+    without standing up a Prometheus registry."""
+
+    def __init__(self) -> None:
+        self.calls: list[tuple] = []
+
+    def observe_size(self, fs, fr, ts, tr, n):
+        self.calls.append(("observe_size", fs, fr, ts, tr, n))
+
+    def observe_tx_time(self, fs, fr, ts, tr, t):
+        self.calls.append(("observe_tx_time", fs, fr, ts, tr, t))
+
+    def observe_rx_time(self, fs, fr, ts, tr, t):
+        self.calls.append(("observe_rx_time", fs, fr, ts, tr, t))
+
+    def observe_in_flight_time(self, fs, fr, ts, tr, t):
+        self.calls.append(("observe_in_flight_time", fs, fr, ts, tr, t))
+
+
+def _make_stats(stage_id, request_id, *, rx_decode=0.0, rx_in_flight=0.0, rx_bytes=0):
+    """Minimal StageRequestStats for record_transfer_rx input."""
+    return StageRequestStats(
+        batch_id=1,
+        batch_size=1,
+        num_tokens_in=0,
+        num_tokens_out=0,
+        stage_gen_time_ms=0.0,
+        rx_transfer_bytes=rx_bytes,
+        rx_decode_time_ms=rx_decode,
+        rx_in_flight_time_ms=rx_in_flight,
+        stage_stats=StageStats(),
+        stage_id=stage_id,
+        request_id=request_id,
+    )
+
+
+class TestEmitHookTx:
+    def test_record_transfer_tx_emits_size_and_tx_time_in_seconds(self):
+        emitter = _StubTransferEmitter()
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=emitter,
+            replica_resolver=lambda s, rid: {0: 1, 1: 0}.get(s),
+        )
+        agg.record_transfer_tx(
+            from_stage=0,
+            to_stage=1,
+            request_id="r-tx-1",
+            size_bytes=2048,
+            tx_time_ms=7.5,
+            used_shm=False,
+        )
+        # tx_time_ms=7.5 -> 0.0075s at the emit boundary.
+        assert emitter.calls == [
+            ("observe_size", 0, 1, 1, 0, 2048),
+            ("observe_tx_time", 0, 1, 1, 0, pytest.approx(0.0075)),
+        ]
+
+    def test_record_transfer_tx_no_emit_when_emitter_none(self):
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=None,
+            replica_resolver=lambda s, rid: 0,
+        )
+        evt = agg.record_transfer_tx(
+            from_stage=0,
+            to_stage=1,
+            request_id="r-tx-2",
+            size_bytes=128,
+            tx_time_ms=1.0,
+            used_shm=True,
+        )
+        # Underlying accumulator keeps ms; only the Prometheus emit converts.
+        assert evt is not None
+        assert evt.size_bytes == 128
+        assert evt.tx_time_ms == 1.0
+
+    def test_record_transfer_tx_no_emit_when_resolver_returns_none(self):
+        emitter = _StubTransferEmitter()
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=emitter,
+            replica_resolver=lambda s, rid: None,
+        )
+        agg.record_transfer_tx(
+            from_stage=0,
+            to_stage=1,
+            request_id="r-tx-3",
+            size_bytes=512,
+            tx_time_ms=2.0,
+            used_shm=False,
+        )
+        assert emitter.calls == []
+
+    def test_record_transfer_tx_no_emit_when_one_side_resolves_to_none(self):
+        emitter = _StubTransferEmitter()
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=emitter,
+            replica_resolver=lambda s, rid: 0 if s == 0 else None,
+        )
+        agg.record_transfer_tx(
+            from_stage=0,
+            to_stage=1,
+            request_id="r-tx-4",
+            size_bytes=64,
+            tx_time_ms=0.5,
+            used_shm=False,
+        )
+        assert emitter.calls == []
+
+
+class TestEmitHookRx:
+    def test_record_transfer_rx_emits_rx_and_in_flight_in_seconds(self):
+        emitter = _StubTransferEmitter()
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=emitter,
+            replica_resolver=lambda s, rid: {0: 1, 1: 0}.get(s),
+        )
+        stats = _make_stats(stage_id=1, request_id="r-rx-1", rx_decode=4.2, rx_in_flight=1.7)
+        agg.record_transfer_rx(stats)
+        assert emitter.calls == [
+            ("observe_rx_time", 0, 1, 1, 0, pytest.approx(0.0042)),
+            ("observe_in_flight_time", 0, 1, 1, 0, pytest.approx(0.0017)),
+        ]
+
+    def test_record_transfer_rx_skips_stage_zero(self):
+        emitter = _StubTransferEmitter()
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+            transfer_emitter=emitter,
+            replica_resolver=lambda s, rid: 0,
+        )
+        stats = _make_stats(stage_id=0, request_id="r-rx-2", rx_decode=4.2)
+        agg.record_transfer_rx(stats)
+        assert emitter.calls == []
+
+    def test_record_transfer_rx_no_emit_when_emitter_none(self):
+        agg = OrchestratorAggregator(
+            num_stages=3,
+            log_stats=False,
+            wall_start_ts=0.0,
+            final_stage_id_for_e2e=2,
+        )
+        stats = _make_stats(stage_id=1, request_id="r-rx-3", rx_decode=1.0, rx_in_flight=0.5)
+        evt = agg.record_transfer_rx(stats)
+        assert evt is not None
+        assert evt.rx_decode_time_ms == 1.0
+        assert evt.in_flight_time_ms == 0.5

vllm_omni/benchmarks/metrics/metrics.py

@@ -7,6 +7,8 @@ from vllm.benchmarks.lib.endpoint_request_func import RequestFuncOutput
 from vllm.benchmarks.serve import MILLISECONDS_TO_SECONDS_CONVERSION, TERM_PLOTLIB_AVAILABLE, BenchmarkMetrics, TaskType
 from vllm.tokenizers import TokenizerLike
 
+from vllm_omni.metrics import definitions as defs
+
 
 @dataclass
 class MultiModalsBenchmarkMetrics(BenchmarkMetrics):
@@ -109,17 +111,17 @@ def process_one_metric(
         "tpot": "Time per Output Token (excl. 1st token)",
         "itl": "Inter-token Latency",
         "e2el": "End-to-end Latency",
-        "audio_ttfp": "Time to First Packet",
-        "audio_rtf": "Real Time Factor",
-        "audio_duration": "Audio Duration",
-        "audio_underrun": "Streaming Audio Underrun",
+        defs.AUDIO_TTFP: "Time to First Packet",
+        defs.AUDIO_RTF: "Real Time Factor",
+        defs.AUDIO_DURATION: "Audio Duration",
+        defs.AUDIO_UNDERRUN: "Streaming Audio Underrun",
     }
 
     header = metric_header_map.get(metric_attribute_name, metric_attribute_name)
     print("{s:{c}^{n}}".format(s=header, n=50, c="-"))
 
-    is_audio_rtf = metric_attribute_name == "audio_rtf"
-    is_audio_duration_or_underrun = metric_attribute_name in ("audio_duration", "audio_underrun")
+    is_audio_rtf = metric_attribute_name == defs.AUDIO_RTF
+    is_audio_duration_or_underrun = metric_attribute_name in (defs.AUDIO_DURATION, defs.AUDIO_UNDERRUN)
 
     suffix = "_ms"
     unit_suffix = " (ms)"
@@ -216,12 +218,12 @@ def calculate_metrics(
             all_tpots.append(tpot)
             itls += outputs[i].itl
             ttfts.append(outputs[i].ttft)
-            audio_ttfps.append(getattr(outputs[i], "audio_ttfp", 0.0))
-            audio_rtfs.append(getattr(outputs[i], "audio_rtf", 0.0))
-            audio_duration.append(getattr(outputs[i], "audio_duration", 0.0))
-            audio_frames.append(getattr(outputs[i], "audio_frames", 0.0))
-            audio_underruns.append(getattr(outputs[i], "audio_underrun_s", 0.0))
-            audio_continuity_ok.append(bool(getattr(outputs[i], "audio_continuity_ok", True)))
+            audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0))
+            audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
+            audio_duration.append(getattr(outputs[i], defs.AUDIO_DURATION, 0.0))
+            audio_frames.append(getattr(outputs[i], defs.AUDIO_FRAMES, 0.0))
+            audio_underruns.append(getattr(outputs[i], f"{defs.AUDIO_UNDERRUN}_s", 0.0))
+            audio_continuity_ok.append(bool(getattr(outputs[i], defs.AUDIO_CONTINUITY_OK, True)))
             e2els.append(outputs[i].latency)
             input_audio_duration += outputs[i].input_audio_duration
             completed += 1

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -1,16 +1,20 @@
 from __future__ import annotations
 
 import os
+import time
 from typing import Any
 
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.metrics.stats import SchedulerStats
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.core.sched.output import OmniChunkRecvHandle, OmniSchedulerOutput
 
 logger = init_logger(__name__)
 
+_STATS_INTERVAL_S = 1.0
+
 # Upper bound on how long a request may sit in full-payload-input wait
 # (the state ``OmniSchedulingCoordinator`` records via ``_waiting_since``)
 # before the scheduler force-fails it.  Defends against stuck consumer-side
@@ -145,3 +149,10 @@ class OmniSchedulerMixin:
             finished_requests_needing_kv_transfer=finished_requests_needing_kv_transfer or {},
             pending_input_registrations=pending_input_registrations,
         )
+
+    def make_stats(self, *args, **kwargs) -> SchedulerStats | None:
+        now = time.monotonic()
+        if now - getattr(self, "_last_stats_time", 0.0) < _STATS_INTERVAL_S:
+            return None
+        self._last_stats_time = now
+        return super().make_stats(*args, **kwargs)

vllm_omni/engine/async_omni_engine.py

@@ -110,6 +110,7 @@ from vllm_omni.entrypoints.utils import (
     load_and_resolve_stage_configs,
 )
 from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.metrics.prometheus import OmniRequestCounter
 from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
@@ -259,6 +260,12 @@ class AsyncOmniEngine:
         **kwargs: Additional arguments
     """
 
+    # Class-level defaults so tests that bypass __init__ via object.__new__
+    # don't AttributeError when stage-init / forward paths touch these attrs.
+    _log_stats: bool = False
+    _coordinator_runtime: Any = None
+    _transfer_emitter: Any = None
+
     def __init__(
         self,
         model: str,
@@ -267,11 +274,23 @@ class AsyncOmniEngine:
         init_timeout: int = 600,
         diffusion_batch_size: int = 1,
         single_stage_mode: bool = False,
+        transfer_emitter: Any = None,
+        log_stats: bool = False,
         **kwargs: Any,
     ) -> None:
         self.model = model
         self.diffusion_batch_size = diffusion_batch_size
         startup_timeout = int(init_timeout)
+        # Forwarded into Orchestrator so its _forward_to_next_stage path can
+        # emit per-edge transfer_tx_s / transfer_size_bytes histograms.
+        # Optional: when None, Orchestrator silently skips TX emit (existing
+        # RX path still works via OrchestratorAggregator).
+        self._transfer_emitter = transfer_emitter
+        # Drives upstream EngineCore + scheduler stats production. When False
+        # the engine skips SchedulerStats / IterationStats; the per-(stage,
+        # 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
 
         logger.info(f"[AsyncOmniEngine] Initializing with model {model}")
 
@@ -357,6 +376,7 @@ class AsyncOmniEngine:
         self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None
         self._weak_finalizer: weakref.finalize | None = None
         self._rpc_lock = threading.Lock()
+        self._running_counter = OmniRequestCounter()
 
         logger.info(f"[AsyncOmniEngine] Launching Orchestrator thread with {self.num_stages} stages")
 
@@ -905,7 +925,7 @@ class AsyncOmniEngine:
                                         launch_omni_core_engines(
                                             vllm_config=vllm_config,
                                             executor_class=executor_class,
-                                            log_stats=False,
+                                            log_stats=self._log_stats,
                                             omni_master_server=self._omni_master_server,
                                             stage_id=plan.metadata.stage_id,
                                             stage_config=stage_cfg,
@@ -917,7 +937,7 @@ class AsyncOmniEngine:
                                     addresses, proc, handshake_address = spawn_stage_core(
                                         vllm_config=vllm_config,
                                         executor_class=executor_class,
-                                        log_stats=False,
+                                        log_stats=self._log_stats,
                                     )
                                 logger.info(
                                     "[AsyncOmniEngine] Stage %s engine launch started",
@@ -1229,7 +1249,7 @@ class AsyncOmniEngine:
             if plan.replicas[0].metadata.stage_type != "diffusion":
                 stage_vllm_config = plan.replicas[0].stage_vllm_config
                 assert stage_vllm_config is not None
-                output_processor = build_llm_stage_output_processor(plan, stage_vllm_config)
+                output_processor = build_llm_stage_output_processor(plan, stage_vllm_config, log_stats=self._log_stats)
 
             stage_pools.append(
                 StagePool(
@@ -1374,9 +1394,12 @@ class AsyncOmniEngine:
                 stage_pools=self.stage_pools,
                 async_chunk=self.async_chunk,
                 pd_config=pd_config,
+                running_counter=self._running_counter,
                 coordinator_pub_address=coordinator_pub_address,
                 load_balancer_factory=load_balancer_factory,
                 remote_replica_factory=remote_replica_factory,
+                transfer_emitter=self._transfer_emitter,
+                log_stats=self._log_stats,
             )
             if not startup_future.done():
                 startup_future.set_result(asyncio.get_running_loop())

vllm_omni/engine/messages.py

@@ -80,6 +80,7 @@ class OutputMessage(EngineQueueMessage, kw_only=True):
     type: Literal["output"] = "output"
     request_id: str
     stage_id: int
+    replica_id: int | None = None
     engine_outputs: OmniRequestOutput
     metrics: StageRequestMetrics | None = None
     finished: bool
@@ -90,6 +91,7 @@ class StageMetricsMessage(EngineQueueMessage, kw_only=True):
     type: Literal["stage_metrics"] = "stage_metrics"
     request_id: str
     stage_id: int
+    replica_id: int | None = None
     metrics: StageRequestMetrics
     stage_submit_ts: float | None = None
 

vllm_omni/engine/orchestrator.py

@@ -29,6 +29,7 @@ from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine import EngineCoreOutputs
 from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.distributed.omni_coordinator import (
     LoadBalancer,
@@ -54,6 +55,8 @@ from vllm_omni.engine.messages import (
 )
 from vllm_omni.engine.serialization import serialize_additional_information
 from vllm_omni.engine.stage_pool import StagePool
+from vllm_omni.metrics.prometheus import OmniRequestCounter
+from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger
 from vllm_omni.outputs import OmniRequestOutput
 
 # Factory signature for building a head-side stage client for a
@@ -155,6 +158,12 @@ class Orchestrator:
     _WATCH_REPLICA_INTERVAL_S: float = 0.5
     _WATCH_REPLICA_IDLE_INTERVAL_S: float = 1.0
 
+    # Class-level defaults so tests that bypass __init__ via object.__new__
+    # don't AttributeError when transfer / counter emit paths access them.
+    _running_counter: OmniRequestCounter | None = None
+    _transfer_emitter: Any = None
+    _stat_logger: OmniPrometheusStatLogger | None = None
+
     def __init__(
         self,
         request_async_queue: janus.AsyncQueue[EngineQueueMessage],
@@ -164,9 +173,12 @@ class Orchestrator:
         *,
         async_chunk: bool = False,
         pd_config: dict[str, Any] | None = None,
+        running_counter: OmniRequestCounter | None = None,
         coordinator_pub_address: str | None = None,
         load_balancer_factory: Callable[[], LoadBalancer] | None = None,
         remote_replica_factory: RemoteReplicaFactory | None = None,
+        transfer_emitter: Any = None,
+        log_stats: bool = False,
     ) -> None:
         self.request_async_queue = request_async_queue
         self.output_async_queue = output_async_queue
@@ -186,6 +198,8 @@ class Orchestrator:
             self._pd_bootstrap_addr = pd_config.get("bootstrap_addr")
             self._pd_prefill_engine_id = pd_config.get("prefill_engine_id")
         self.request_states: dict[str, OrchestratorRequestState] = {}
+        self._init_metrics_state(stage_pools, running_counter, transfer_emitter, log_stats=log_stats)
+
         self._cfg_tracker = CfgCompanionTracker()
 
         self._shutdown_event = asyncio.Event()
@@ -215,6 +229,71 @@ class Orchestrator:
                 pool.attach_hub(self._hub)
                 pool.attach_load_balancer(factory())
 
+    def _init_metrics_state(
+        self,
+        stage_pools: list[StagePool],
+        running_counter: OmniRequestCounter | None,
+        transfer_emitter: Any,
+        log_stats: bool = False,
+    ) -> None:
+        """Wire up all metric-related orchestrator state.
+
+        Sets ``self._running_counter`` and ``self._transfer_emitter``
+        (both optional, used by request-add / forward paths), builds the
+        ``(stage_id, replica_id) ↔ engine_idx`` lookup used at record() time,
+        and best-effort constructs the ``OmniPrometheusStatLogger`` wrap
+        that exposes ~37 upstream ``vllm:*`` families with per-(stage,
+        replica) labels. Failure to build the wrap is logged and metrics
+        are simply disabled — orchestrator construction continues so unit
+        tests with a minimal ``vllm_config`` still pass.
+
+        ``log_stats=False`` short-circuits the wrap entirely so the
+        ~65 upstream ``vllm:*`` families are not registered in the
+        Prometheus default registry at all. The per-step record() path
+        already no-ops on ``scheduler_stats is None`` (which is what
+        the upstream scheduler returns when its own log_stats is False),
+        so this gate is mainly to keep the ``/metrics`` surface clean
+        when the user did not request stats.
+        """
+        self._running_counter = running_counter
+        self._transfer_emitter = transfer_emitter
+
+        # Flat engine_idx ↔ (stage, replica) maps. The reverse map is
+        # consulted at record() time to translate the orchestrator's
+        # (stage_id, replica_id) loop variables into an engine_idx the
+        # underlying PrometheusStatLogger can address.
+        stage_replica_map: dict[int, tuple[str, str]] = {}
+        self._stage_replica_to_engine_idx: dict[tuple[int, int], int] = {}
+        flat_idx = 0
+        for stage_id, pool in enumerate(stage_pools):
+            for replica_id in range(pool.num_replicas):
+                stage_replica_map[flat_idx] = (str(stage_id), str(replica_id))
+                self._stage_replica_to_engine_idx[(stage_id, replica_id)] = flat_idx
+                flat_idx += 1
+
+        if not log_stats:
+            self._stat_logger = None
+            return
+
+        vllm_config_for_stats = next(
+            (p.stage_vllm_config for p in stage_pools if p.stage_vllm_config is not None),
+            None,
+        )
+        if vllm_config_for_stats is None:
+            self._stat_logger = None
+            return
+        try:
+            self._stat_logger = OmniPrometheusStatLogger(
+                vllm_config=vllm_config_for_stats,
+                stage_replica_map=stage_replica_map,
+            )
+        except Exception:
+            # Minimal vllm_config in unit-test contexts can lack fields the
+            # upstream PrometheusStatLogger expects. Skip wrap rather than
+            # break orchestrator construction.
+            logger.exception("[Orchestrator] OmniPrometheusStatLogger init failed; metrics wrap disabled")
+            self._stat_logger = None
+
     async def run(self) -> None:
         """Main entry point for the Orchestrator event loop."""
         logger.info("[Orchestrator] Starting event loop")
@@ -393,6 +472,8 @@ class Orchestrator:
             mm_features=getattr(prompt, "mm_features", None),
         )
         self.request_states[request_id] = req_state
+        if self._running_counter is not None:
+            self._running_counter.increment()
         req_state.streaming.enabled = bool(getattr(prompt, "resumable", False))
         req_state.stage_submit_ts[stage_id] = _time.time()
         enqueue_ts = msg.enqueue_ts
@@ -603,7 +684,24 @@ class Orchestrator:
                                     "new_prompt_len_snapshot",
                                     None,
                                 )
-                            raw_output = await pool.process_llm_raw_outputs(replica_id, raw_outputs)
+                            # OmniSchedulerMixin.make_stats() already throttles
+                            # per-scheduler at 1 Hz, so raw_outputs.scheduler_stats
+                            # being non-None means this replica passed its own gate.
+                            # A second global throttle here would drop stats for
+                            # other (stage, replica) pairs in the same 1s window.
+                            record_stats = self._stat_logger is not None and raw_outputs.scheduler_stats is not None
+                            iteration_stats = IterationStats() if record_stats else None
+                            raw_output = await pool.process_llm_raw_outputs(
+                                replica_id,
+                                raw_outputs,
+                                iteration_stats=iteration_stats,
+                            )
+                            if record_stats:
+                                self._stat_logger.record(
+                                    raw_outputs.scheduler_stats,
+                                    iteration_stats,
+                                    engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)],
+                                )
                         except asyncio.CancelledError:
                             raise
                         except EngineDeadError as e:
@@ -681,7 +779,7 @@ class Orchestrator:
                 )
                 stage_metrics.pipeline_timings = dict(req_state.pipeline_timings)
 
-            await self._route_output(stage_id, output, req_state, stage_metrics)
+            await self._route_output(stage_id, replica_id, output, req_state, stage_metrics)
 
     async def _handle_stage_error(self, stage_id: int, output: Any) -> None:
         """Emit a frontend-visible error and clean up request state."""
@@ -713,7 +811,8 @@ class Orchestrator:
         self._release_request_bindings(request_ids)
         for request_id in request_ids:
             self._pd_kv_params.pop(request_id, None)
-            self.request_states.pop(request_id, None)
+            if self.request_states.pop(request_id, None) is not None and self._running_counter is not None:
+                self._running_counter.decrement()
 
     def _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any:
         """Attach CFG companion ids to diffusion sampling params when needed."""
@@ -735,6 +834,7 @@ class Orchestrator:
     async def _route_output(
         self,
         stage_id: int,
+        replica_id: int,
         output: Any,
         req_state: OrchestratorRequestState,
         stage_metrics: Any,
@@ -757,6 +857,7 @@ class Orchestrator:
                 OutputMessage(
                     request_id=req_id,
                     stage_id=stage_id,
+                    replica_id=replica_id,
                     engine_outputs=output,
                     metrics=stage_metrics,
                     finished=finished and stage_id == req_state.final_stage_id,
@@ -768,6 +869,7 @@ class Orchestrator:
                 StageMetricsMessage(
                     request_id=req_id,
                     stage_id=stage_id,
+                    replica_id=replica_id,
                     metrics=stage_metrics,
                     stage_submit_ts=submit_ts,
                 )
@@ -797,6 +899,7 @@ class Orchestrator:
                     stage_id,
                     output,
                     req_state,
+                    src_replica_id=replica_id,
                     is_streaming_session=req_state.streaming.enabled,
                     is_final_update=False,
                 )
@@ -807,6 +910,7 @@ class Orchestrator:
                         stage_id,
                         output,
                         req_state,
+                        src_replica_id=replica_id,
                         is_streaming_session=True,
                         is_final_update=True,
                     )
@@ -918,6 +1022,41 @@ class Orchestrator:
         )
         return sp
 
+    def _emit_tx_edge(
+        self,
+        *,
+        from_stage: int,
+        from_replica: int,
+        to_stage: int,
+        to_pool: StagePool,
+        request_id: str,
+        tx_ms: float,
+    ) -> None:
+        """Emit per-edge transfer_tx_s + transfer_size_bytes histograms.
+
+        ``tx_ms`` is the orchestrator-side wall-clock spent in ``next_pool.
+        submit_*`` (serialize + queue submit to the receiving worker). Best-
+        effort size_bytes left at 0 — orchestrator doesn't have a cheap handle
+        on the serialized payload size; a follow-up can plumb that from the
+        connector adapter.
+        """
+        if self._transfer_emitter is None:
+            return
+        to_replica = to_pool.get_bound_replica_id(request_id)
+        if to_replica is None:
+            return
+        try:
+            self._transfer_emitter.observe_size(from_stage, from_replica, to_stage, to_replica, 0)
+            self._transfer_emitter.observe_tx_time(from_stage, from_replica, to_stage, to_replica, tx_ms / 1000.0)
+        except Exception:
+            logger.debug(
+                "[Orchestrator] transfer_tx emit failed for edge %d->%d req=%s",
+                from_stage,
+                to_stage,
+                request_id,
+                exc_info=True,
+            )
+
     async def _forward_to_next_stage(
         self,
         req_id: str,
@@ -925,6 +1064,7 @@ class Orchestrator:
         output: Any,
         req_state: OrchestratorRequestState,
         *,
+        src_replica_id: int | None = None,
         is_streaming_session: bool = False,
         is_final_update: bool = False,
     ) -> None:
@@ -937,6 +1077,7 @@ class Orchestrator:
         next_stage_resumable = is_streaming_session and not is_final_update
         already_submitted = self._next_stage_already_submitted(src_stage_id, req_state)
         requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False)
+        _t_submit_start = _time.perf_counter()
 
         if next_pool.stage_type == "diffusion":
             companion_outputs = self._cfg_tracker.pop_companion_outputs(req_id)
@@ -1024,6 +1165,15 @@ class Orchestrator:
                     params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
                 )
             req_state.stage_submit_ts[next_logical] = _time.time()
+            _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0
+            self._emit_tx_edge(
+                from_stage=src_stage_id,
+                from_replica=src_replica_id if src_replica_id is not None else 0,
+                to_stage=next_logical,
+                to_pool=next_pool,
+                request_id=req_id,
+                tx_ms=_tx_ms,
+            )
             return
 
         # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params
@@ -1063,6 +1213,15 @@ class Orchestrator:
                     await next_pool.submit_initial(req_id, req_state, request, prompt_text=None)
 
             req_state.stage_submit_ts[next_logical] = _time.time()
+            _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0
+            self._emit_tx_edge(
+                from_stage=src_stage_id,
+                from_replica=src_replica_id if src_replica_id is not None else 0,
+                to_stage=next_logical,
+                to_pool=next_pool,
+                request_id=req_id,
+                tx_ms=_tx_ms,
+            )
             return
 
         if req_state.pd_prefill_multimodal_output is not None:
@@ -1107,6 +1266,15 @@ class Orchestrator:
                 await next_pool.submit_initial(req_id, req_state, request, prompt_text=None)
 
         req_state.stage_submit_ts[next_logical] = _time.time()
+        _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0
+        self._emit_tx_edge(
+            from_stage=src_stage_id,
+            from_replica=src_replica_id if src_replica_id is not None else 0,
+            to_stage=next_logical,
+            to_pool=next_pool,
+            request_id=req_id,
+            tx_ms=_tx_ms,
+        )
 
     async def _prewarm_async_chunk_stages(
         self,
@@ -1131,6 +1299,7 @@ class Orchestrator:
             params = req_state.sampling_params_list[next_stage_id]
 
             req_state.stage_submit_ts[next_stage_id] = _time.time()
+            _t_submit_start = _time.perf_counter()
 
             if next_pool.stage_type == "diffusion":
                 await next_pool.submit_initial(
@@ -1174,6 +1343,20 @@ class Orchestrator:
                 request.external_req_id = request.request_id
                 await next_pool.submit_initial(request_id, req_state, request, prompt_text=None)
 
+            # async_chunk pre-submit fires per stage edge (N-1 -> N). Source
+            # replica is stage 0's bound replica (single-replica thinker in
+            # all current configs); fall back to 0 if unknown.
+            _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0
+            src_replica = self.stage_pools[next_stage_id - 1].get_bound_replica_id(request_id)
+            self._emit_tx_edge(
+                from_stage=next_stage_id - 1,
+                from_replica=src_replica if src_replica is not None else 0,
+                to_stage=next_stage_id,
+                to_pool=next_pool,
+                request_id=request_id,
+                tx_ms=_tx_ms,
+            )
+
     def _build_kv_sender_info(
         self,
         sender_stage_ids: list[int],

vllm_omni/engine/stage_init_utils.py

@@ -831,8 +831,18 @@ def build_vllm_config(
     return vllm_config, executor_class
 
 
-def build_llm_stage_output_processor(plan: LogicalStageInitPlan, stage_vllm_config: Any) -> Any | None:
-    """Build one output processor per logical LLM stage."""
+def build_llm_stage_output_processor(
+    plan: LogicalStageInitPlan,
+    stage_vllm_config: Any,
+    log_stats: bool = False,
+) -> Any | None:
+    """Build one output processor per logical LLM stage.
+
+    ``log_stats`` controls whether the processor populates per-request
+    IterationStats (consumed by the Prometheus wrap). Default False matches
+    the upstream MultimodalOutputProcessor default and respects the
+    --log-stats CLI flag plumbed through AsyncOmniEngine.
+    """
 
     metadata = plan.replicas[0].metadata
     if stage_vllm_config.model_config.skip_tokenizer_init:
@@ -843,7 +853,7 @@ def build_llm_stage_output_processor(plan: LogicalStageInitPlan, stage_vllm_conf
         )
     return MultimodalOutputProcessor(
         tokenizer=tokenizer,
-        log_stats=False,
+        log_stats=log_stats,
         engine_core_output_type=metadata.engine_output_type,
     )
 

vllm_omni/engine/stage_pool.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, cast
 
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreOutputs
+from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.distributed.omni_coordinator import (
     LoadBalancer,
@@ -510,6 +511,7 @@ class StagePool:
             stage_gen_time_ms=stage_gen_time_ms,
             batch_id=batch_id,
             batch_size=1,
+            replica_id=replica_id,
             rx_decode_time_ms=0.0,
             rx_transfer_bytes=0,
             rx_in_flight_time_ms=0.0,
@@ -648,6 +650,7 @@ class StagePool:
         self,
         replica_id: int,
         raw_outputs: EngineCoreOutputs,
+        iteration_stats: IterationStats | None = None,
     ) -> list[Any]:
         """Run the shared LLM output processor on one raw poll result."""
         raw_client = self.clients[replica_id]
@@ -658,7 +661,7 @@ class StagePool:
         processed = processor.process_outputs(
             raw_outputs.outputs,
             raw_outputs.timestamp,
-            None,
+            iteration_stats,
         )
 
         if processed.reqs_to_abort:

vllm_omni/entrypoints/async_omni.py

@@ -160,6 +160,20 @@ class AsyncOmni(EngineClient, OmniBase):
                 renderer = renderer_from_config(vllm_config)
             self.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin)
 
+    def _resolve_transfer_replica(self, stage_id: int, request_id: str) -> int | None:
+        """Look up the sticky-routed replica for (stage_id, request_id).
+
+        Used as the ``replica_resolver`` callback by ``OrchestratorAggregator``
+        to label transfer_* metrics without plumbing replica ids through
+        ``TransferEdgeStats`` / ``StageRequestStats`` / connector adapters.
+        Returns None when stage_id is out of range or the request hasn't been
+        bound to a replica yet — the metric emit then defensive-skips.
+        """
+        pools = getattr(self.engine, "stage_pools", None)
+        if pools is None or not (0 <= stage_id < len(pools)):
+            return None
+        return pools[stage_id].get_bound_replica_id(request_id)
+
     def _get_comprehension_stage_index(self) -> int | None:
         fallback_idx: int | None = None
         for idx, stage_client in enumerate(self.engine.stage_clients):
@@ -331,6 +345,8 @@ class AsyncOmni(EngineClient, OmniBase):
                 self.log_stats,
                 wall_start_ts,
                 final_stage_id_for_e2e,
+                transfer_emitter=getattr(self, "transfer_metrics", None),
+                replica_resolver=self._resolve_transfer_replica,
             )
 
             req_state = ClientRequestState(
@@ -338,6 +354,7 @@ class AsyncOmni(EngineClient, OmniBase):
                 external_request_id=external_request_id,
             )
             req_state.metrics = metrics
+            req_state.request_arrival_ts = wall_start_ts
             self.request_states[request_id] = req_state
 
             # PD disaggregation: modify prefill-stage sampling params per request
@@ -387,10 +404,12 @@ class AsyncOmni(EngineClient, OmniBase):
         except (asyncio.CancelledError, GeneratorExit):
             if input_stream_task is not None and not input_stream_task.done():
                 input_stream_task.cancel()
+            self._fire_failure_counter_if_alive(request_id)
             await self._abort_internal_requests(request_id)
             logger.info(f"[AsyncOmni] Request {request_id} aborted.")
             raise
         except Exception as e:
+            self._fire_failure_counter_if_alive(request_id)
             await self._abort_internal_requests(request_id)
             logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}")
             raise

vllm_omni/entrypoints/client_request_state.py

@@ -17,3 +17,23 @@ class ClientRequestState:
         self.stage_id: int | None = None
         self.queue = queue if queue is not None else asyncio.Queue()
         self.metrics: OrchestratorAggregator | None = None
+        # Wall-clock time at which the user's request arrived in the engine
+        # entrypoint. Set in async_omni.generate() before the orchestrator
+        # accepts the request. Used as the t0 anchor for audio_ttfp.
+        self.request_arrival_ts: float = 0.0
+        # Wall-clock time at which the first audio packet was observed for
+        # this request. None means the streaming hook hasn't fired yet.
+        # Used as the once-per-request guard for audio_ttfp_s emit.
+        self.first_audio_ts: float | None = None
+        # Per-chunk timeline (seconds since request_arrival_ts) and PCM byte
+        # counts for the audio streaming response. Populated by the streaming
+        # endpoint on every audio.chunk emit; consumed at request finalize to
+        # compute audio_underrun_s and audio_continuity_ok_total.
+        self.audio_chunk_arrivals_s: list[float] = []
+        self.audio_chunk_bytes: list[int] = []
+        self.audio_sample_rate: int | None = None
+        # Stage / replica that produced the audio packets — captured at the
+        # first-packet hook so the finalize-time emit can label correctly
+        # without re-querying stage_pools.
+        self.audio_emit_stage_id: int | None = None
+        self.audio_emit_replica_id: int | None = None

vllm_omni/entrypoints/omni_base.py

@@ -23,7 +23,10 @@ from vllm_omni.engine.messages import (
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
 from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e
+from vllm_omni.metrics.modality import OmniModalityMetrics, observe_modality_at_finalize
+from vllm_omni.metrics.prometheus import OmniPrometheusMetrics
 from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
+from vllm_omni.metrics.transfer import OmniTransferMetrics
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -170,6 +173,10 @@ class OmniBase(PDDisaggregationMixin):
         self.tts_batch_max_items: int = kwargs.pop("tts_batch_max_items", 32)
 
         logger.info("[%s] Initializing with model %s", self.__class__.__name__, model)
+        # Construct transfer_metrics first so we can hand it to AsyncOmniEngine
+        # (which forwards it to the Orchestrator background thread for
+        # TX-side emit; see Orchestrator._forward_to_next_stage).
+        self.transfer_metrics = OmniTransferMetrics(model_name=model, log_stats=log_stats)
         st = time.time()
         self.engine = AsyncOmniEngine(
             model=model,
@@ -177,6 +184,8 @@ class OmniBase(PDDisaggregationMixin):
             init_timeout=init_timeout,
             stage_init_timeout=stage_init_timeout,
             diffusion_batch_size=diffusion_batch_size,
+            transfer_emitter=self.transfer_metrics,
+            log_stats=log_stats,
             **kwargs,
         )
         self._shutdown_called = False
@@ -190,6 +199,8 @@ class OmniBase(PDDisaggregationMixin):
         self.async_chunk = bool(getattr(self.engine, "async_chunk", False))
 
         self.request_states: dict[str, ClientRequestState] = {}
+        self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats)
+        self.mod_metrics = OmniModalityMetrics(model_name=model, log_stats=log_stats)
 
         self.default_sampling_params_list = self.engine.default_sampling_params_list
         if not self.output_modalities:
@@ -272,11 +283,29 @@ class OmniBase(PDDisaggregationMixin):
             raise ValueError(f"Expected {self.num_stages} sampling params, got {len(normalized)}")
         return normalized
 
+    def _fire_failure_counter_if_alive(self, request_id: str) -> None:
+        """Fire the abort/exception bucket of requests_success_total.
+
+        Called from cancel / exception paths in async_omni.generate() BEFORE
+        _abort_internal_requests pops request_states — that method resolves
+        the internal id by dict lookup, so popping first would no-op it. We
+        keep this counter fire separate from _log_summary_and_cleanup (which
+        pops) so the abort path can still find the state to clean up.
+        """
+        req_state = self.request_states.get(request_id)
+        prom = getattr(self, "prom_metrics", None)
+        if req_state is None or req_state.metrics is None or prom is None:
+            return
+        if str(request_id) not in req_state.metrics.e2e_done:
+            prom.request_failed()
+
     def _log_summary_and_cleanup(self, request_id: str) -> None:
         req_state = self.request_states.get(request_id)
         try:
             if req_state is None or req_state.metrics is None:
                 return
+            if str(request_id) not in req_state.metrics.e2e_done:
+                self.prom_metrics.request_failed()
             if self.log_stats:
                 # Emit per-request orchestrator timing (including e2e_total_ms)
                 # before dropping request state.
@@ -289,6 +318,16 @@ class OmniBase(PDDisaggregationMixin):
             )
         finally:
             self.request_states.pop(request_id, None)
+            # Republish gauges so any stale value left by the per-stage
+            # publish in _process_single_result (which runs while the request
+            # is still in self.request_states) is corrected after the pop.
+            prom = getattr(self, "prom_metrics", None)
+            counter = getattr(getattr(self, "engine", None), "_running_counter", None)
+            if prom is not None:
+                total = len(self.request_states)
+                running = counter.value if counter is not None else total
+                prom.set_running(running)
+                prom.set_waiting(max(0, total - running))
 
     def _compute_final_stage_id(self, output_modalities: list[str] | None) -> int:
         return get_final_stage_id_for_e2e(
@@ -435,6 +474,8 @@ class OmniBase(PDDisaggregationMixin):
         if not stage_meta.final_output:
             return None
 
+        output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type)
+
         try:
             rid_key = str(req_id)
             if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done and finished:
@@ -443,10 +484,41 @@ class OmniBase(PDDisaggregationMixin):
                     req_id,
                     req_start_ts.get(req_id, wall_start_ts),
                 )
+                e2e_seconds = now - req_start_ts.get(req_id, wall_start_ts)
+                # Extract finished_reason from upstream CompletionOutput so
+                # the per-reason completion Counter is labelled correctly.
+                completion_outputs = getattr(engine_outputs, "outputs", None) or []
+                fr = (getattr(completion_outputs[0], "finish_reason", None) if completion_outputs else None) or "stop"
+                self.prom_metrics.request_succeeded(
+                    e2e_seconds,
+                    finished_reason=fr,
+                )
+
+                # Modality observe inside the same finalize guard so it fires
+                # once per request and inherits the try/except isolation.
+                observe_modality_at_finalize(
+                    self.mod_metrics,
+                    output_type=output_type,
+                    stage_id=stage_id,
+                    replica_id=result.replica_id,
+                    stage_metrics=_m,
+                    engine_outputs=engine_outputs,
+                )
         except Exception:
             logger.exception("[%s] Finalize request handling error", self.__class__.__name__)
 
-        output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type)
+        # When this result finalizes the request, the orchestrator has
+        # already decremented _running_counter but _log_summary_and_cleanup
+        # hasn't popped self.request_states yet — exclude the finalizing
+        # request from `total` so waiting doesn't read 1 and stay stuck
+        # there until the next request arrives.
+        counter = getattr(self.engine, "_running_counter", None)
+        is_finalizing = finished and stage_id == final_stage_id_for_e2e
+        total = max(0, len(self.request_states) - (1 if is_finalizing else 0))
+        running = counter.value if counter is not None else total
+        self.prom_metrics.set_running(running)
+        self.prom_metrics.set_waiting(max(0, total - running))
+
         images = getattr(engine_outputs, "images", []) if output_type == "image" else []
         return OmniRequestOutput(
             request_id=req_id or "",

vllm_omni/entrypoints/openai/serving_chat.py

@@ -20,6 +20,11 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse
 from vllm_omni.entrypoints.utils import coerce_param_message_types
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
+from vllm_omni.metrics import definitions as _metric_defs
+from vllm_omni.metrics.modality import (
+    observe_audio_first_packet,
+    observe_audio_streaming_finalize,
+)
 
 try:
     import soundfile
@@ -110,6 +115,7 @@ from vllm_omni.entrypoints.openai.utils import (
 )
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate
 
 logger = init_logger(__name__)
 
@@ -1032,6 +1038,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         include_usage, include_continuous_usage = should_include_usage(stream_options, self.enable_force_include_usage)
 
         last_metrics: dict[str, Any] | None = None
+        # Hold a strong reference to the audio request state so the
+        # streaming-finalize hook below survives the inner generator's
+        # _log_summary_and_cleanup, which pops request_states[request_id]
+        # before this outer block runs.
+        req_state_audio_ref: Any = None
         try:
             async for omni_res in result_generator:
                 final_output_type = omni_res.final_output_type
@@ -1605,8 +1616,53 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         yield f"data: {data}\n\n"
 
                 elif final_output_type == "audio":
+                    # Observe audio_ttfp_s on first audio packet for this request_id
+                    # (once-per-request guard via first_audio_ts). The same hook
+                    # also captures (stage, replica) for the streaming-continuity
+                    # emit at request finalize. self.engine_client.request_states
+                    # is keyed by the internal UUID-suffixed id (set by
+                    # AsyncOmni.generate via _get_unique_request_id), but the
+                    # `request_id` we have here is the external (user-visible) id,
+                    # so resolve via the external_request_id field.
+                    req_state = next(
+                        (s for s in self.engine_client.request_states.values() if s.external_request_id == request_id),
+                        None,
+                    )
+                    if req_state is not None and req_state_audio_ref is None:
+                        req_state_audio_ref = req_state
+                    now_ts = time.time()
+                    if req_state is not None and req_state.first_audio_ts is None:
+                        req_state.first_audio_ts = now_ts
+                        stage_pools = getattr(self.engine_client.engine, "stage_pools", None)
+                        # The orchestrator binds requests by their internal id,
+                        # not the user-visible external id, so look up the
+                        # replica with req_state.request_id (internal).
+                        replica_id = (
+                            stage_pools[omni_res.stage_id].get_bound_replica_id(req_state.request_id)
+                            if stage_pools is not None and 0 <= omni_res.stage_id < len(stage_pools)
+                            else None
+                        )
+                        req_state.audio_emit_stage_id = omni_res.stage_id
+                        req_state.audio_emit_replica_id = replica_id
+                        observe_audio_first_packet(
+                            self.engine_client.mod_metrics,
+                            stage_id=omni_res.stage_id,
+                            replica_id=replica_id,
+                            arrival_ts=req_state.request_arrival_ts,
+                            now_ts=now_ts,
+                        )
+
                     role = self.get_chat_request_role(request)
                     choices_data = self._create_audio_choice(omni_res, role, request, stream=True)
+                    # Record per-chunk PCM byte count + arrival timestamp for
+                    # audio_underrun_s / audio_continuity_ok_total at finalize.
+                    if req_state is not None and req_state.request_arrival_ts > 0:
+                        chunk_bytes = audio_chunk_pcm_bytes(omni_res)
+                        if chunk_bytes > 0:
+                            req_state.audio_chunk_arrivals_s.append(max(now_ts - req_state.request_arrival_ts, 0.0))
+                            req_state.audio_chunk_bytes.append(chunk_bytes)
+                            if req_state.audio_sample_rate is None:
+                                req_state.audio_sample_rate = audio_chunk_sample_rate(omni_res)
                     chunk = OmniChatCompletionStreamResponse(
                         id=request_id,
                         object=chunk_object_type,
@@ -1627,6 +1683,24 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     logger.warning(f"Unsupported streaming final output type: {final_output_type}")
                     continue
 
+            # Emit audio_underrun_s + audio_continuity_ok_total once per
+            # request after the audio chunk stream is exhausted. The
+            # captured reference (req_state_audio_ref) outlives the inner
+            # _log_summary_and_cleanup that pops request_states[request_id].
+            if (
+                req_state_audio_ref is not None
+                and req_state_audio_ref.audio_chunk_arrivals_s
+                and req_state_audio_ref.audio_emit_replica_id is not None
+            ):
+                observe_audio_streaming_finalize(
+                    self.engine_client.mod_metrics,
+                    stage_id=req_state_audio_ref.audio_emit_stage_id or 0,
+                    replica_id=req_state_audio_ref.audio_emit_replica_id,
+                    chunk_arrival_times_s=req_state_audio_ref.audio_chunk_arrivals_s,
+                    chunk_bytes=req_state_audio_ref.audio_chunk_bytes,
+                    sample_rate=(req_state_audio_ref.audio_sample_rate or _metric_defs.DEFAULT_AUDIO_SAMPLE_RATE),
+                )
+
             # once the final token is handled, if stream_options.include_usage
             # is sent, send the usage
             if include_usage:

vllm_omni/metrics/__init__.py

@@ -1,7 +1,10 @@
+from .prometheus import OmniPrometheusMetrics, OmniRequestCounter
 from .stats import OrchestratorAggregator, StageRequestStats, StageStats
 from .utils import count_tokens_from_outputs
 
 __all__ = [
+    "OmniPrometheusMetrics",
+    "OmniRequestCounter",
     "OrchestratorAggregator",
     "StageStats",
     "StageRequestStats",

vllm_omni/metrics/definitions.py

@@ -0,0 +1,221 @@
+"""Single source of truth for vLLM-Omni Prometheus + bench CLI metric naming.
+
+Consumed by:
+- vllm_omni.metrics.prometheus (server-side /metrics pipeline families)
+- vllm_omni.metrics.modality (audio families)
+- vllm_omni.metrics.transfer (cross-stage transfer families)
+- vllm_omni.benchmarks.metrics.metrics (bench CLI MultiModalsBenchmarkMetrics)
+
+Naming conventions for the ``vllm_omni:*`` families exposed here:
+time-bearing metrics use the ``_s`` suffix (values in seconds), counters use
+``_total`` (auto-suffixed by the prometheus client), sizes use ``_bytes``.
+"""
+
+# vllm:omni_ avoids upstream's unregister_vllm_metrics() stripping, which
+# removes every collector whose ``_name`` does not start with ``vllm``.
+METRIC_PREFIX = "vllm:omni_"
+
+
+# ============================================================================
+# Bench-side stems (also used as RequestFuncOutput attribute names)
+# ============================================================================
+AUDIO_TTFP = "audio_ttfp"
+AUDIO_DURATION = "audio_duration"
+AUDIO_RTF = "audio_rtf"
+AUDIO_FRAMES = "audio_frames"
+AUDIO_UNDERRUN = "audio_underrun"
+AUDIO_CONTINUITY_OK = "audio_continuity_ok"
+AUDIO_SKIPPED_REQUESTS = "audio_skipped_requests"
+
+
+# ============================================================================
+# Pipeline-level metric families (request counts + e2e latency)
+# ============================================================================
+NUM_REQUESTS_RUNNING = METRIC_PREFIX + "num_requests_running"
+NUM_REQUESTS_WAITING = METRIC_PREFIX + "num_requests_waiting"
+E2E_REQUEST_LATENCY_S = METRIC_PREFIX + "e2e_request_latency_s"
+
+# Per-finished_reason Counter; finished_reason ∈ {stop, length, abort, ...}.
+# Aborts include client disconnect / cancellation paths. Counter auto-suffixes
+# ``_total`` at exposition time.
+REQUESTS_SUCCESS = METRIC_PREFIX + "requests_success"
+
+
+# ============================================================================
+# Audio family (per-stage + per-replica audio path metrics)
+# ============================================================================
+AUDIO_TTFP_S = METRIC_PREFIX + AUDIO_TTFP + "_s"
+AUDIO_DURATION_S = METRIC_PREFIX + AUDIO_DURATION + "_s"
+AUDIO_RTF_METRIC = METRIC_PREFIX + AUDIO_RTF
+AUDIO_FRAMES_METRIC = METRIC_PREFIX + AUDIO_FRAMES
+AUDIO_UNDERRUN_S = METRIC_PREFIX + AUDIO_UNDERRUN + "_s"
+AUDIO_CONTINUITY_OK_METRIC = METRIC_PREFIX + AUDIO_CONTINUITY_OK
+AUDIO_SKIPPED_REQUESTS_METRIC = METRIC_PREFIX + AUDIO_SKIPPED_REQUESTS
+
+
+# ============================================================================
+# Cross-stage Transfer family (per-physical-hop TX/RX/in-flight timings)
+# ============================================================================
+TRANSFER_SIZE_BYTES = METRIC_PREFIX + "transfer_size_bytes"
+TRANSFER_TX_S = METRIC_PREFIX + "transfer_tx_s"
+TRANSFER_RX_S = METRIC_PREFIX + "transfer_rx_s"
+TRANSFER_IN_FLIGHT_S = METRIC_PREFIX + "transfer_in_flight_s"
+
+
+# ============================================================================
+# Label sets
+# ============================================================================
+PIPELINE_LABELS = ("model_name",)
+SUCCESS_LABELS = ("model_name", "finished_reason")
+
+# Per-stage / per-replica label set used by the audio families and by the
+# OmniPrometheusStatLogger wrap which relabels upstream ``engine`` into
+# ``stage`` + ``replica``.
+STAGE_LABELS = ("model_name", "stage", "replica")
+
+# Audio continuity Counter carries an extra ``threshold_ms`` label so multiple
+# threshold buckets can be tracked simultaneously. The ``_ms`` suffix names a
+# numeric threshold *value* in ms, not a time-bearing metric.
+AUDIO_CONTINUITY_LABELS = ("model_name", "stage", "replica", "threshold_ms")
+
+# Audio skipped-requests Counter carries a `reason` label so the silent-loss
+# path (e.g. code2wav rejecting malformed codec input) can be distinguished
+# from other "200 OK + empty audio" cases.
+AUDIO_SKIPPED_LABELS = ("model_name", "stage", "replica", "reason")
+
+# Cross-stage transfer label set. Each observation is one physical hop from
+# (from_stage, from_replica) to (to_stage, to_replica).
+TRANSFER_LABELS = (
+    "model_name",
+    "from_stage",
+    "from_replica",
+    "to_stage",
+    "to_replica",
+)
+
+
+# ============================================================================
+# Histogram buckets
+# ============================================================================
+# Seconds bucket for e2e / generation / TTFP-style metrics that range from
+# ~10 ms to several minutes.
+SECONDS_BUCKETS = (
+    0.05,
+    0.1,
+    0.25,
+    0.5,
+    1.0,
+    2.0,
+    5.0,
+    10.0,
+    20.0,
+    30.0,
+    60.0,
+    120.0,
+    300.0,
+)
+
+# Seconds bucket for fine-grained metrics (cross-stage transfer + audio
+# underrun) that need millisecond-level resolution.
+SECONDS_FAST_BUCKETS = (
+    0.001,
+    0.005,
+    0.01,
+    0.025,
+    0.05,
+    0.1,
+    0.25,
+    0.5,
+    1.0,
+    2.5,
+    5.0,
+    10.0,
+    60.0,
+)
+
+# RTF SLO red line is 1.0 — TTS must generate faster than playback.
+RTF_BUCKETS = (
+    0.1,
+    0.25,
+    0.5,
+    0.75,
+    0.9,
+    1.0,
+    1.25,
+    1.5,
+    2.0,
+    5.0,
+    10.0,
+)
+
+# Bytes bucket for transfer payload size.
+BYTES_BUCKETS = (
+    1024,
+    4096,
+    16384,
+    65536,
+    262144,
+    1048576,
+    4194304,
+    16777216,
+    67108864,
+    268435456,
+)
+
+
+# ============================================================================
+# Audio-continuity defaults
+# ============================================================================
+# Default underrun threshold — kept aligned with the bench-side default and
+# the commonly-cited "audible gap" threshold for streaming TTS.
+AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S = 0.1
+
+
+# ============================================================================
+# Formula helpers (shared by server-side observe and bench-side calculation)
+# ============================================================================
+def compute_audio_rtf(stage_gen_time_s: float, audio_duration_s: float) -> float:
+    """RTF = stage_gen_time / audio_content_duration.
+
+    SLO red line < 1 — must generate faster than content plays back to stream.
+    Returns 0.0 when audio_duration_s is non-positive (caller decides whether
+    to observe; we don't want to divide by zero or emit negative samples).
+    """
+    if audio_duration_s <= 0:
+        return 0.0
+    return stage_gen_time_s / audio_duration_s
+
+
+# ============================================================================
+# Audio sample-rate resolution
+# ============================================================================
+# Most common across vllm-omni talker variants (cosyvoice3, omnivoice,
+# qwen3_tts, mimo_audio). voxcpm2 uses 48000, stable_audio 44100,
+# ming_flash 16000 — these models populate multimodal_output["audio_sample_rate"]
+# at runtime so this default only kicks in when the field is missing.
+DEFAULT_AUDIO_SAMPLE_RATE = 24000
+
+_SAMPLE_RATE_KEYS = ("audio_sample_rate", "sample_rate", "sampling_rate", "sr")
+
+
+def resolve_audio_sample_rate(multimodal_output: dict | None) -> int:
+    """Extract audio sample_rate from a multimodal_output dict, with fallbacks.
+
+    Tries the same key chain as serving_chat.py's audio response path so
+    /metrics audio_duration_s = audio_frames / sample_rate stays consistent
+    with what the OpenAI streaming endpoint reports back to clients.
+    Returns DEFAULT_AUDIO_SAMPLE_RATE when no usable value is present.
+    """
+    if not multimodal_output:
+        return DEFAULT_AUDIO_SAMPLE_RATE
+    for key in _SAMPLE_RATE_KEYS:
+        raw = multimodal_output.get(key)
+        if raw is None:
+            continue
+        try:
+            value = int(raw)
+        except (TypeError, ValueError):
+            continue
+        if value > 0:
+            return value
+    return DEFAULT_AUDIO_SAMPLE_RATE

vllm_omni/metrics/modality.py

@@ -0,0 +1,303 @@
+"""OmniModalityMetrics — per-modality Prometheus families (audio path only).
+
+7 audio business-semantic metric families. Text-path metrics (TTFT / ITL /
+TPOT / e2e) are NOT here — they come from the upstream
+``vllm:*{stage="thinker", ...}`` families exposed via the
+``OmniPrometheusStatLogger`` wrap.
+
+Contents:
+- Audio family declarations (Histograms + Counters)
+- ``OmniModalityMetrics``: label-bound observe API for the audio family
+- ``observe_modality_at_finalize``: dispatcher called from omni_base's e2e
+  finalize hook; currently handles the audio path only.
+- ``observe_audio_first_packet``: TTFP emit from the streaming SSE first
+  audio packet.
+- ``observe_audio_streaming_finalize``: emits ``audio_underrun_s`` +
+  ``audio_continuity_ok_total`` at SSE close using accumulated per-chunk
+  arrival timestamps.
+- ``_extract_mm_output`` / ``_count_audio_frames``: shape-tolerant helpers
+  for the heterogeneous multimodal_output payloads emitted by different
+  audio pipelines.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from prometheus_client import Counter, Histogram
+
+from vllm_omni.metrics import definitions as defs
+
+_stage_labels = list(defs.STAGE_LABELS)
+
+
+# ----------------------------------------------------------------------------
+# Audio family
+# ----------------------------------------------------------------------------
+_audio_ttfp_family = Histogram(
+    defs.AUDIO_TTFP_S,
+    "Time from request arrival to first audio packet/frame, in seconds.",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_BUCKETS,
+)
+_audio_duration_family = Histogram(
+    defs.AUDIO_DURATION_S,
+    "Generated audio content duration in seconds (audio_frames / sample_rate).",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_BUCKETS,
+)
+_audio_rtf_family = Histogram(
+    defs.AUDIO_RTF_METRIC,
+    "Audio real-time factor (stage_gen_time_s / audio_duration_s); streaming TTS requires < 1.",
+    labelnames=_stage_labels,
+    buckets=defs.RTF_BUCKETS,
+)
+_audio_frames_family = Counter(
+    defs.AUDIO_FRAMES_METRIC,
+    "Cumulative audio frame count; per-model rate (not summable across models). Throughput recovered via rate().",
+    labelnames=_stage_labels,
+)
+_audio_underrun_family = Histogram(
+    defs.AUDIO_UNDERRUN_S,
+    "Per-request worst-case player-deficit in seconds (max time the player "
+    "ran out of buffered audio). > 0 indicates listener experienced silent gaps.",
+    labelnames=_stage_labels,
+    buckets=defs.SECONDS_FAST_BUCKETS,
+)
+_audio_continuity_ok_family = Counter(
+    defs.AUDIO_CONTINUITY_OK_METRIC,
+    "Incremented when the request's worst underrun stayed below threshold_ms. "
+    "Pair with requests_success_total to compute streaming-continuity health rate.",
+    labelnames=list(defs.AUDIO_CONTINUITY_LABELS),
+)
+_audio_skipped_family = Counter(
+    defs.AUDIO_SKIPPED_REQUESTS_METRIC,
+    "Silent-loss counter — code2wav rejected malformed codec input and returned 200 OK with empty audio.",
+    labelnames=list(defs.AUDIO_SKIPPED_LABELS),
+)
+
+
+class OmniModalityMetrics:
+    """Per-modality observe API. Stage/replica are passed at observe time
+    because a single OmniModalityMetrics instance per pipeline serves all
+    stage+replica combinations.
+    """
+
+    def __init__(self, model_name: str, log_stats: bool = True) -> None:
+        self._model_name = model_name
+        self._log_stats = log_stats
+
+    # ---- Audio ------------------------------------------------------------
+
+    def observe_audio_ttfp(self, stage: str, replica: str, ttfp_seconds: float) -> None:
+        if not self._log_stats:
+            return
+        _audio_ttfp_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(ttfp_seconds)
+
+    def observe_audio_duration(self, stage: str, replica: str, duration_seconds: float) -> None:
+        if not self._log_stats:
+            return
+        _audio_duration_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(
+            duration_seconds
+        )
+
+    def observe_audio_rtf(self, stage: str, replica: str, rtf: float) -> None:
+        if not self._log_stats:
+            return
+        _audio_rtf_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(rtf)
+
+    def inc_audio_frames(self, stage: str, replica: str, n_frames: int) -> None:
+        if not self._log_stats or n_frames <= 0:
+            return
+        _audio_frames_family.labels(model_name=self._model_name, stage=stage, replica=replica).inc(n_frames)
+
+    def observe_audio_underrun(self, stage: str, replica: str, underrun_s: float) -> None:
+        if not self._log_stats:
+            return
+        _audio_underrun_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(
+            max(underrun_s, 0.0)
+        )
+
+    def inc_audio_continuity_ok(self, stage: str, replica: str, threshold_ms: int) -> None:
+        if not self._log_stats:
+            return
+        _audio_continuity_ok_family.labels(
+            model_name=self._model_name,
+            stage=stage,
+            replica=replica,
+            threshold_ms=str(int(threshold_ms)),
+        ).inc()
+
+    def inc_audio_skipped(self, stage: str, replica: str, reason: str) -> None:
+        if not self._log_stats:
+            return
+        _audio_skipped_family.labels(
+            model_name=self._model_name,
+            stage=stage,
+            replica=replica,
+            reason=reason or "unknown",
+        ).inc()
+
+
+def _extract_mm_output(engine_outputs: Any) -> dict[str, Any]:
+    """Return the multimodal_output dict regardless of where it's nested.
+
+    Three shapes seen in the wild:
+      * ``engine_outputs.multimodal_output`` — synthesized on OmniRequestOutput
+        for some pipelines (often empty for AR audio)
+      * ``engine_outputs.outputs[0].multimodal_output`` — vllm CompletionOutput
+        nesting (where actual qwen3-omni audio data lives)
+      * neither — returns ``{}``
+    """
+    mm = getattr(engine_outputs, "multimodal_output", None)
+    if isinstance(mm, dict) and mm:
+        return mm
+    outs = getattr(engine_outputs, "outputs", None)
+    if outs:
+        nested = getattr(outs[0], "multimodal_output", None)
+        if isinstance(nested, dict):
+            return nested
+    return {}
+
+
+def _count_audio_frames(mm_out: dict[str, Any]) -> int:
+    """Sum the per-tensor sample count of audio chunks in mm_out["audio"].
+
+    Returns the total number of audio frames (samples) across all chunks.
+    For multi-dim tensors (e.g. shape [channels, samples]) the last axis is
+    treated as the sample dim; for 1-D tensors the only axis is the sample
+    dim; scalars count as 1.
+    """
+    audio_chunks = mm_out.get("audio") if isinstance(mm_out, dict) else None
+    if audio_chunks is None:
+        return 0
+    chunks = audio_chunks if isinstance(audio_chunks, list) else [audio_chunks]
+    n = 0
+    for t in chunks:
+        try:
+            ndim = getattr(t, "ndim", 0)
+            shape = getattr(t, "shape", None)
+            if ndim == 0 or shape is None or len(shape) == 0:
+                n += 1
+            else:
+                n += int(shape[-1])
+        except Exception:
+            continue
+    return n
+
+
+def observe_modality_at_finalize(
+    mod_metrics: OmniModalityMetrics,
+    *,
+    output_type: str | None,
+    stage_id: int,
+    replica_id: int | None,
+    stage_metrics: Any,
+    engine_outputs: Any,
+) -> None:
+    """Route audio-path observations for a finalized request.
+
+    Used by ``omni_base._process_single_result`` inside the e2e_done finalize
+    guard so it fires once per request. Skips text path (covered by upstream
+    ``vllm:*{stage="thinker", ...}``) and any case where required inputs are
+    missing — caller should not need to pre-validate.
+
+    audio_ttfp is intentionally NOT observed here; it's emitted by the
+    streaming hook at first-packet time, not at finalize.
+    """
+    if replica_id is None or stage_metrics is None or output_type is None:
+        return
+    if output_type != "audio":
+        return
+
+    stage_label = str(stage_id)
+    replica_label = str(replica_id)
+    gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0
+    mm_out = _extract_mm_output(engine_outputs)
+
+    sample_rate = defs.resolve_audio_sample_rate(mm_out)
+    # `stage_metrics.audio_generated_frames` is the legacy per-chunk
+    # accumulator field on StageRequestStats. No production path currently
+    # fills it, so the fallback below is the live source — but we leave the
+    # field lookup in place in case the accumulator gets re-wired upstream.
+    n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0)
+    if n_frames == 0:
+        n_frames = _count_audio_frames(mm_out)
+    mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames)
+    duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0
+    if duration_s > 0:
+        mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s)
+        mod_metrics.observe_audio_rtf(
+            stage_label,
+            replica_label,
+            defs.compute_audio_rtf(gen_time_s, duration_s),
+        )
+    else:
+        # Request completed (finish_reason ∈ {stop, length} — error paths
+        # don't reach finalize) but no audio samples were produced. Covers
+        # silent `return None` skips in the talker→code2wav stage
+        # processors and the `parsed.append((0,0))` malformed-length path
+        # in qwen3-tts code2wav. raise-paths surface via the upstream
+        # vllm:request_success_total{finished_reason="error"} channel and
+        # never reach this branch.
+        mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data")
+    # audio_underrun / continuity are emitted from the streaming path in
+    # observe_audio_streaming_finalize; finalize is too late for the
+    # per-chunk timeline they need.
+
+
+def observe_audio_first_packet(
+    mod_metrics: OmniModalityMetrics,
+    *,
+    stage_id: int,
+    replica_id: int | None,
+    arrival_ts: float,
+    now_ts: float,
+) -> None:
+    """Observe audio_ttfp_s on a request's first audio packet.
+
+    Caller is responsible for the once-per-request guard (e.g. checking
+    ``ClientRequestState.first_audio_ts is None``) so this function fires at
+    most once per request_id. Defensive-skips when ``replica_id`` or
+    ``arrival_ts`` is insufficient — both can legitimately be missing in error
+    paths and we'd rather drop the sample than emit a wrong (stage, replica).
+    """
+    if replica_id is None or arrival_ts <= 0:
+        return
+    ttfp = max(now_ts - arrival_ts, 0.0)
+    mod_metrics.observe_audio_ttfp(str(stage_id), str(replica_id), ttfp)
+
+
+def observe_audio_streaming_finalize(
+    mod_metrics: OmniModalityMetrics,
+    *,
+    stage_id: int,
+    replica_id: int | None,
+    chunk_arrival_times_s: list[float],
+    chunk_bytes: list[int],
+    sample_rate: int,
+    threshold_s: float = defs.AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S,
+) -> None:
+    """Emit audio_underrun_s + audio_continuity_ok_total at request end.
+
+    Reuses the math from ``vllm_omni.benchmarks.audio_continuity`` so the
+    server-side and bench-side definitions stay aligned. Caller is responsible
+    for collecting per-chunk arrival timestamps and byte sizes during the
+    streaming response.
+    """
+    if replica_id is None or not chunk_arrival_times_s:
+        return
+    # Local import to keep the bench module optional at import time.
+    from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats
+
+    stats = compute_continuity_stats(
+        chunk_arrival_times_s=chunk_arrival_times_s,
+        chunk_bytes=chunk_bytes,
+        sample_rate=sample_rate,
+        threshold_s=threshold_s,
+    )
+    stage_label = str(stage_id)
+    replica_label = str(replica_id)
+    mod_metrics.observe_audio_underrun(stage_label, replica_label, stats.max_underrun_s)
+    if stats.is_continuous:
+        mod_metrics.inc_audio_continuity_ok(stage_label, replica_label, int(threshold_s * 1000))

vllm_omni/metrics/prometheus.py

@@ -0,0 +1,91 @@
+from prometheus_client import Counter, Gauge, Histogram
+
+from vllm_omni.metrics import definitions as defs
+
+_labelnames = list(defs.PIPELINE_LABELS)
+
+_running_family = Gauge(
+    defs.NUM_REQUESTS_RUNNING,
+    "Number of requests currently running across all pipeline stages.",
+    labelnames=_labelnames,
+)
+_waiting_family = Gauge(
+    defs.NUM_REQUESTS_WAITING,
+    "Number of requests waiting to be scheduled.",
+    labelnames=_labelnames,
+)
+_completion_family = Counter(
+    defs.REQUESTS_SUCCESS,
+    "Total requests by completion reason "
+    "(stop / length / abort / ...). Aborts cover client-disconnect / "
+    "cancellation paths in addition to upstream FinishReason.ABORT.",
+    labelnames=list(defs.SUCCESS_LABELS),
+)
+_e2e_latency_family = Histogram(
+    defs.E2E_REQUEST_LATENCY_S,
+    "Pipeline-global end-to-end request latency in seconds (user arrival to complete response).",
+    labelnames=_labelnames,
+    buckets=defs.SECONDS_BUCKETS,
+)
+
+
+class OmniPrometheusMetrics:
+    """Label-bound wrapper around the raw Prometheus metrics.
+
+    Metric collectors use the ``vllm:omni_`` prefix to avoid being
+    removed by upstream vLLM's ``unregister_vllm_metrics()``, which
+    strips every collector whose ``_name`` contains ``"vllm"``.
+    """
+
+    def __init__(self, model_name: str, log_stats: bool = True) -> None:
+        self._model_name = model_name
+        self._log_stats = log_stats
+        self._running = _running_family.labels(model_name=model_name)
+        self._waiting = _waiting_family.labels(model_name=model_name)
+        self._e2e_latency = _e2e_latency_family.labels(model_name=model_name)
+
+    def set_running(self, n: int) -> None:
+        if not self._log_stats:
+            return
+        self._running.set(n)
+
+    def set_waiting(self, n: int) -> None:
+        if not self._log_stats:
+            return
+        self._waiting.set(n)
+
+    def request_succeeded(
+        self,
+        e2e_seconds: float,
+        finished_reason: str = "stop",
+    ) -> None:
+        if not self._log_stats:
+            return
+        _completion_family.labels(
+            model_name=self._model_name,
+            finished_reason=finished_reason,
+        ).inc()
+        self._e2e_latency.observe(e2e_seconds)
+
+    def request_failed(self) -> None:
+        if not self._log_stats:
+            return
+        # Pipeline-level "fail" maps to the upstream FinishReason.ABORT bucket;
+        # a single counter family now covers both normal stops and aborts.
+        _completion_family.labels(
+            model_name=self._model_name,
+            finished_reason="abort",
+        ).inc()
+
+
+class OmniRequestCounter:
+    """Running-request counter written by the orchestrator thread, read by the client thread."""
+
+    def __init__(self) -> None:
+        self.value = 0
+
+    def increment(self) -> None:
+        self.value += 1
+
+    def decrement(self) -> None:
+        self.value = max(0, self.value - 1)

vllm_omni/metrics/stat_logger.py

@@ -0,0 +1,243 @@
+"""OmniPrometheusStatLogger — wrap upstream PrometheusStatLogger.
+
+Rewrites the upstream ``engine`` single-label scheme into a ``stage`` +
+``replica`` two-label scheme so the ~37 ``vllm:*`` metric families gain
+per-(stage, replica) visibility for multi-replica deployments.
+
+Contents:
+- ``_ENGINE_INDEX_MAP``: process-wide engine_idx → (stage_name, replica_id)
+  lookup, populated by ``OmniPrometheusStatLogger.__init__``.
+- ``_RelabelMixin``: rewrites ``labelnames`` at family creation and translates
+  ``.labels()`` calls; applied via ``_RelabelGauge`` / ``_RelabelCounter`` /
+  ``_RelabelHistogram``.
+- ``_OmniPerfMetricsProm`` / ``_OmniSpecDecodingProm`` / ``_OmniKVConnectorProm``:
+  helper-class wraps so the upstream sub-collectors construct their internal
+  families through the relabel mixin too.
+- ``OmniPrometheusStatLogger``: the subclass that wires everything together
+  and rewrites ``per_engine_labelvalues`` from 2-tuple to 3-tuple at setter
+  time.
+"""
+
+from __future__ import annotations
+
+from prometheus_client import Counter, Gauge, Histogram
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorProm
+from vllm.v1.metrics.loggers import PrometheusStatLogger
+from vllm.v1.metrics.perf import PerfMetricsProm
+from vllm.v1.spec_decode.metrics import SpecDecodingProm
+
+# Process-wide translation table written by OmniPrometheusStatLogger at init.
+# Keys are flat engine_idx values (as upstream PrometheusStatLogger sees them);
+# values are the (stage_name, replica_id_str) tuple we expose as labels.
+#
+# Module-level rather than per-instance because the wrapper metric classes are
+# constructed by upstream's __init__ and never get a back-reference to the
+# StatLogger that owns them. vLLM runs a single Orchestrator/StatLogger per
+# process, so a module global is safe; tests isolate by .clear()ing first.
+_ENGINE_INDEX_MAP: dict[int, tuple[str, str]] = {}
+
+
+def _rewrite_labelnames(labelnames):
+    """Replace `engine` in ``labelnames`` with (`stage`, `replica`) in place.
+
+    Preserves ordering (so ``["model_name", "engine", "reason"]`` becomes
+    ``["model_name", "stage", "replica", "reason"]``) and the original
+    container type (list vs tuple).
+    """
+    if labelnames is None:
+        return labelnames
+    seq = list(labelnames)
+    if "engine" not in seq:
+        return labelnames
+    out: list[str] = []
+    for name in seq:
+        if name == "engine":
+            out.extend(("stage", "replica"))
+        else:
+            out.append(name)
+    return type(labelnames)(out) if not isinstance(labelnames, list) else out
+
+
+def _engine_to_stage_replica(engine_value) -> tuple[str, str]:
+    """Look up (stage, replica) for an engine_idx, accepting int or str input.
+
+    Upstream emits engine values in two flavors:
+    - int form, e.g. ``gauge_engine_sleep_state.labels(engine=idx, ...)`` (loggers.py:510)
+    - str form, e.g. ``info_gauge.labels(**metrics_info)`` where ``metrics_info["engine"] = str(idx)`` (loggers.py:1055)
+
+    Raises ``KeyError`` when the value is missing from the map — fail-fast is
+    preferable to silently emitting series under a wrong (stage, replica).
+    """
+    key = int(engine_value) if isinstance(engine_value, str) else engine_value
+    return _ENGINE_INDEX_MAP[key]
+
+
+class _RelabelMixin:
+    """Mixin: rewrite ``labelnames`` at family creation and ``.labels()`` calls.
+
+    Handles all four upstream forms encountered in
+    ``vllm.v1.metrics.loggers.PrometheusStatLogger``:
+
+    1. ``.labels(engine=idx, ...)`` kwarg with int engine (loggers.py:510)
+    2. ``.labels(model_name, str(idx), source)`` positional with str engine
+       (loggers.py:646, 679)
+    3. ``.labels(**metrics_info)`` kwarg with str engine (loggers.py:1056)
+    4. Families without an ``engine`` label — passthrough (e.g. lora_info)
+
+    Drops into upstream's ``_gauge_cls`` / ``_counter_cls`` / ``_histogram_cls``
+    class slots.
+    """
+
+    def __init__(self, *args, **kwargs):
+        # Remember where `engine` sat in the original labelnames so positional
+        # `.labels()` calls can splice (stage, replica) at the right offset.
+        labelnames = kwargs.get("labelnames")
+        if labelnames is not None:
+            original = list(labelnames)
+            self._engine_label_index = original.index("engine") if "engine" in original else -1
+            kwargs["labelnames"] = _rewrite_labelnames(labelnames)
+        else:
+            self._engine_label_index = -1
+        super().__init__(*args, **kwargs)
+
+    def labels(self, *args, **kwargs):
+        if self._engine_label_index >= 0:
+            if args:
+                # Positional form. There are TWO upstream patterns:
+                #
+                # (a) Pre-rewritten path: create_metric_per_engine fans
+                #     `per_engine_labelvalues` (already a 3-tuple
+                #     [model_name, stage, replica] thanks to the property-
+                #     descriptor setter on OmniPrometheusStatLogger) into
+                #     `metric.labels(*values)`. len(args) matches the
+                #     rewritten label set already, so just pass through.
+                #
+                # (b) Legacy 2-tuple path: upstream sites like
+                #     `counter_request_success.labels(model_name, str(idx),
+                #     str(reason))` pass values shaped to the *original*
+                #     labelnames (engine still present at idx). Here
+                #     len(args) is short by 1 — splice (stage, replica)
+                #     in place of the engine value at engine_label_index.
+                if len(args) == len(self._labelnames):
+                    return super().labels(*args, **kwargs)
+                idx = self._engine_label_index
+                if idx < len(args):
+                    stage, replica = _engine_to_stage_replica(args[idx])
+                    args = (*args[:idx], stage, replica, *args[idx + 1 :])
+            elif "engine" in kwargs:
+                stage, replica = _engine_to_stage_replica(kwargs.pop("engine"))
+                kwargs["stage"] = stage
+                kwargs["replica"] = replica
+        return super().labels(*args, **kwargs)
+
+
+class _RelabelGauge(_RelabelMixin, Gauge):
+    pass
+
+
+class _RelabelCounter(_RelabelMixin, Counter):
+    pass
+
+
+class _RelabelHistogram(_RelabelMixin, Histogram):
+    pass
+
+
+# ----------------------------------------------------------------------------
+# Helper-class wraps for the three sub-metric collectors that upstream
+# PrometheusStatLogger constructs in its __init__ (loggers.py:438-446):
+#
+#     self.spec_decoding_prom = self._spec_decoding_cls(...)
+#     self.kv_connector_prom = self._kv_connector_cls(...)
+#     self.perf_metrics_prom = self._perf_metrics_cls(...)
+#
+# Each helper receives raw `labelnames` as a constructor argument and uses
+# its own class-level `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots
+# to build internal Counter/Gauge/Histogram families. The slot overrides on
+# OmniPrometheusStatLogger only reach families created via *its* slots, so
+# the helpers would otherwise still construct 2-label families and then hit
+# `Incorrect label count` when create_metric_per_engine feeds the rewritten
+# 3-element per_engine_labelvalues. Subclassing each helper and overriding
+# its slots routes the relabel mixin through to the helper-internal families
+# too. The helper kept seeing the OLD 2-element labelnames param, but that
+# is fine because the wrapper rewrites it at family-creation time.
+# ----------------------------------------------------------------------------
+
+
+class _OmniPerfMetricsProm(PerfMetricsProm):
+    _counter_cls = _RelabelCounter
+
+
+class _OmniSpecDecodingProm(SpecDecodingProm):
+    _counter_cls = _RelabelCounter
+
+
+class _OmniKVConnectorProm(KVConnectorProm):
+    _gauge_cls = _RelabelGauge
+    _counter_cls = _RelabelCounter
+    _histogram_cls = _RelabelHistogram
+
+
+class OmniPrometheusStatLogger(PrometheusStatLogger):
+    """Wrap upstream PrometheusStatLogger to expose per-(stage, replica) labels.
+
+    Replaces the upstream single ``engine`` label with two labels ``stage`` and
+    ``replica`` so the ~37 ``vllm:*`` metric families gain per-replica
+    visibility for multi-replica deployments.
+
+    The orchestrator builds ``stage_replica_map`` from the static stage_pools
+    config; flat engine_idx values map 1:1 to (stage_name, replica_id) tuples.
+    Dynamic add/remove of replicas at runtime is intentionally not supported —
+    the map is built once at construction and never mutated afterward.
+    """
+
+    # Inject our wrapper metric classes into upstream's class-level slots so
+    # every ~37 family is created with `engine` rewritten to `stage`+`replica`.
+    _gauge_cls = _RelabelGauge
+    _counter_cls = _RelabelCounter
+    _histogram_cls = _RelabelHistogram
+    # Inject helper-class wraps too so the perf / spec-decoding / kv-connector
+    # sub-collectors get the same labelname rewrite and don't crash with
+    # `Incorrect label count` when create_metric_per_engine fans out the
+    # rewritten 3-element per_engine_labelvalues over their internal families.
+    _perf_metrics_cls = _OmniPerfMetricsProm
+    _spec_decoding_cls = _OmniSpecDecodingProm
+    _kv_connector_cls = _OmniKVConnectorProm
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        stage_replica_map: dict[int, tuple[str, str]],
+    ) -> None:
+        self._stage_replica_map = stage_replica_map
+        # Populate the process-level translation table that wrapper metric
+        # classes consult on every `.labels()` call. Cleared first so a
+        # second OmniPrometheusStatLogger in the same process (e.g. tests,
+        # orchestrator restart) starts from a clean slate.
+        _ENGINE_INDEX_MAP.clear()
+        _ENGINE_INDEX_MAP.update(stage_replica_map)
+        super().__init__(
+            vllm_config=vllm_config,
+            engine_indexes=list(stage_replica_map.keys()),
+        )
+
+    @property
+    def stage_replica_map(self) -> dict[int, tuple[str, str]]:
+        return self._stage_replica_map
+
+    @property
+    def per_engine_labelvalues(self) -> dict[int, list[object]]:
+        return self._omni_per_engine_labelvalues
+
+    @per_engine_labelvalues.setter
+    def per_engine_labelvalues(self, value: dict[int, list[object]]) -> None:
+        # Upstream sets {idx: [model_name, str(idx)]} (loggers.py:433); we drop
+        # the engine str and append (stage, replica) so labelvalues match the
+        # 3-element labelnames our wrapper classes produce.
+        rewritten: dict[int, list[object]] = {}
+        for idx, vals in value.items():
+            model_name = vals[0]
+            stage, replica = self._stage_replica_map[idx]
+            rewritten[idx] = [model_name, stage, replica]
+        self._omni_per_engine_labelvalues = rewritten

vllm_omni/metrics/stats.py

@@ -5,12 +5,15 @@ from collections import defaultdict
 from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from vllm.logger import init_logger
 
 from vllm_omni.metrics.utils import _build_field_defs, _build_row, _format_table
 
+if TYPE_CHECKING:
+    from vllm_omni.metrics.transfer import OmniTransferMetrics
+
 logger = init_logger(__name__)
 
 
@@ -36,6 +39,7 @@ class StageRequestStats:
     rx_in_flight_time_ms: float
     stage_stats: StageStats
     stage_id: int | None = None
+    replica_id: int | None = None
     final_output_type: str | None = None
     request_id: str | None = None
     postprocess_time_ms: float = 0.0
@@ -121,6 +125,9 @@ class OrchestratorAggregator:
         log_stats: bool,
         wall_start_ts: float,
         final_stage_id_for_e2e: dict[str, int] | int,
+        *,
+        transfer_emitter: OmniTransferMetrics | None = None,
+        replica_resolver: Callable[[int, str], int | None] | None = None,
     ) -> None:
         self.num_stages = int(num_stages)
         self.log_stats = bool(log_stats)
@@ -131,6 +138,20 @@ class OrchestratorAggregator:
             tuple[int, int, str], TransferEdgeStats
         ] = {}  # Key: (from_stage, to_stage, request_id)
         self.e2e_events: list[RequestE2EStats] = []
+        # Emit per-physical-transfer Histogram observations to Prometheus
+        # alongside the existing TransferEdgeStats accumulation. Both deps
+        # are optional so OrchestratorAggregator stays usable in contexts
+        # that don't have a Prometheus registry (e.g. unit tests).
+        self._transfer_emitter = transfer_emitter
+        self._replica_resolver = replica_resolver
+        # Snapshot of (stage_id -> replica_id) per request, captured at
+        # ``record_transfer_rx`` time from the receiving stage's own
+        # ``StageRequestStats.replica_id``. Survives pool-side
+        # ``release_binding`` so the transfer emit can resolve label values
+        # even when the downstream stage's binding has already been released
+        # by the orchestrator's cleanup path (e.g. final-stage cleanup races
+        # the omni_base finalize emit).
+        self._replica_cache: dict[tuple[int, str], int] = {}
 
     def init_run_state(self, wall_start_ts: float) -> None:
         # Per-run aggregates and timing state
@@ -191,6 +212,14 @@ class OrchestratorAggregator:
             evt.size_bytes += int(size_bytes)
             evt.tx_time_ms += float(tx_time_ms)
             evt.used_shm = evt.used_shm or bool(used_shm)
+            # Emit per-physical-transfer Histogram observations to Prometheus.
+            self._emit_transfer_tx(
+                from_stage=int(from_stage),
+                to_stage=int(to_stage),
+                request_id=str(request_id),
+                size_bytes=int(size_bytes),
+                tx_time_ms=float(tx_time_ms),
+            )
             return evt
         except Exception:
             return None
@@ -205,6 +234,11 @@ class OrchestratorAggregator:
             from_stage = int(stats.stage_id) - 1
             to_stage = int(stats.stage_id)
             rid_key = str(stats.request_id)
+            # Snapshot the receiving (stage, replica) binding before the
+            # pool may release it on cleanup. Used by _resolve_edge_replicas
+            # as a fallback when get_bound_replica_id has already cleared.
+            if stats.replica_id is not None:
+                self._replica_cache[(to_stage, rid_key)] = int(stats.replica_id)
             evt = self._get_or_create_transfer_event(from_stage, to_stage, rid_key)
             # Accumulate rx metrics
             if evt.size_bytes == 0:
@@ -212,10 +246,79 @@ class OrchestratorAggregator:
                 evt.size_bytes = int(stats.rx_transfer_bytes)
             evt.rx_decode_time_ms += float(stats.rx_decode_time_ms)
             evt.in_flight_time_ms += float(stats.rx_in_flight_time_ms)
+            # Emit per-physical-receive Histogram observations to Prometheus.
+            self._emit_transfer_rx(
+                from_stage=from_stage,
+                to_stage=to_stage,
+                request_id=rid_key,
+                rx_decode_time_ms=float(stats.rx_decode_time_ms),
+                in_flight_time_ms=float(stats.rx_in_flight_time_ms),
+            )
             return evt
         except Exception:
             return None
 
+    # ------------------------------------------------------------------
+    # Prometheus emit hooks. Both helpers are no-ops when transfer_emitter
+    # or replica_resolver is None, or when the resolver cannot find a
+    # (stage_id, request_id) -> replica_id mapping. We deliberately
+    # fail-safe (skip) rather than emit a series with a wrong or invented
+    # replica label.
+    # ------------------------------------------------------------------
+
+    def _resolve_edge_replicas(self, from_stage: int, to_stage: int, request_id: str) -> tuple[int, int] | None:
+        # Prefer the per-request replica snapshot captured by
+        # record_transfer_rx (survives pool binding release); fall back to the
+        # resolver callback (pool.get_bound_replica_id) when the cache
+        # misses, e.g. for the TX side that fires before the receiving stage
+        # has recorded its stats.
+        from_r = self._replica_cache.get((from_stage, request_id))
+        to_r = self._replica_cache.get((to_stage, request_id))
+        if (from_r is None or to_r is None) and self._replica_resolver is not None:
+            if from_r is None:
+                from_r = self._replica_resolver(from_stage, request_id)
+            if to_r is None:
+                to_r = self._replica_resolver(to_stage, request_id)
+        if from_r is None or to_r is None:
+            return None
+        return int(from_r), int(to_r)
+
+    def _emit_transfer_tx(
+        self,
+        *,
+        from_stage: int,
+        to_stage: int,
+        request_id: str,
+        size_bytes: int,
+        tx_time_ms: float,
+    ) -> None:
+        if self._transfer_emitter is None:
+            return
+        replicas = self._resolve_edge_replicas(from_stage, to_stage, request_id)
+        if replicas is None:
+            return
+        from_r, to_r = replicas
+        self._transfer_emitter.observe_size(from_stage, from_r, to_stage, to_r, size_bytes)
+        self._transfer_emitter.observe_tx_time(from_stage, from_r, to_stage, to_r, tx_time_ms / 1000.0)
+
+    def _emit_transfer_rx(
+        self,
+        *,
+        from_stage: int,
+        to_stage: int,
+        request_id: str,
+        rx_decode_time_ms: float,
+        in_flight_time_ms: float,
+    ) -> None:
+        if self._transfer_emitter is None:
+            return
+        replicas = self._resolve_edge_replicas(from_stage, to_stage, request_id)
+        if replicas is None:
+            return
+        from_r, to_r = replicas
+        self._transfer_emitter.observe_rx_time(from_stage, from_r, to_stage, to_r, rx_decode_time_ms / 1000.0)
+        self._transfer_emitter.observe_in_flight_time(from_stage, from_r, to_stage, to_r, in_flight_time_ms / 1000.0)
+
     def record_audio_generated_frames(
         self,
         output_to_yield: Any,
@@ -462,6 +565,10 @@ class OrchestratorAggregator:
             ),
         )
         self.e2e_events.append(per_req_record)
+        # Drop per-request replica snapshots once the request is finalized.
+        if self._replica_cache:
+            for cached_key in [k for k in self._replica_cache if k[1] == rid_key]:
+                self._replica_cache.pop(cached_key, None)
 
     def build_and_log_summary(self) -> dict[str, Any]:
         if not self.log_stats:

vllm_omni/metrics/transfer.py

@@ -0,0 +1,145 @@
+"""OmniTransferMetrics — cross-stage transfer Prometheus families.
+
+Four families with ``{model_name, from_stage, from_replica, to_stage,
+to_replica}`` labels. Each ``observe_*()`` call corresponds to one physical
+transfer event (one chunk hop from a sender replica to a receiver replica),
+so the Histogram tracks the distribution of physical transfers, not
+request-aggregated sums.
+
+Data source: ``vllm_omni.metrics.stats.TransferEdgeStats`` accumulators in
+``OrchestratorAggregator.record_transfer_tx`` / ``record_transfer_rx``. The
+emit hook lives in stats.py; this module only registers the families and
+exposes the typed observe API.
+"""
+
+from __future__ import annotations
+
+from prometheus_client import Histogram
+
+from vllm_omni.metrics import definitions as defs
+
+_labelnames = list(defs.TRANSFER_LABELS)
+
+
+# ----------------------------------------------------------------------------
+# TX-side families (observed when record_transfer_tx fires)
+# ----------------------------------------------------------------------------
+_transfer_size_bytes_family = Histogram(
+    defs.TRANSFER_SIZE_BYTES,
+    "Per-transfer payload size in bytes (one observation per physical hop).",
+    labelnames=_labelnames,
+    buckets=defs.BYTES_BUCKETS,
+)
+_transfer_tx_family = Histogram(
+    defs.TRANSFER_TX_S,
+    "Sender-side time in seconds (serialize + submit to connector).",
+    labelnames=_labelnames,
+    buckets=defs.SECONDS_FAST_BUCKETS,
+)
+
+
+# ----------------------------------------------------------------------------
+# RX-side families (observed when record_transfer_rx fires)
+# ----------------------------------------------------------------------------
+_transfer_rx_family = Histogram(
+    defs.TRANSFER_RX_S,
+    "Receiver-side time in seconds (recv + deserialize).",
+    labelnames=_labelnames,
+    buckets=defs.SECONDS_FAST_BUCKETS,
+)
+_transfer_in_flight_family = Histogram(
+    defs.TRANSFER_IN_FLIGHT_S,
+    "Network in-flight time in seconds (TX done -> RX recv start).",
+    labelnames=_labelnames,
+    buckets=defs.SECONDS_FAST_BUCKETS,
+)
+
+
+class OmniTransferMetrics:
+    """Per-(from, to) replica observe API for cross-stage transfers.
+
+    A single instance per pipeline; ``model_name`` is bound at init and
+    every observe call carries it in the label set. Stage/replica are
+    passed at observe time because the same instance serves all
+    (from_stage, from_replica) -> (to_stage, to_replica) edges.
+    """
+
+    def __init__(self, model_name: str, log_stats: bool = True) -> None:
+        self._model_name = model_name
+        self._log_stats = log_stats
+
+    # ---- TX side (record_transfer_tx hook) -------------------------------
+
+    def observe_size(
+        self,
+        from_stage: int,
+        from_replica: int,
+        to_stage: int,
+        to_replica: int,
+        size_bytes: int,
+    ) -> None:
+        if not self._log_stats:
+            return
+        _transfer_size_bytes_family.labels(
+            model_name=self._model_name,
+            from_stage=str(from_stage),
+            from_replica=str(from_replica),
+            to_stage=str(to_stage),
+            to_replica=str(to_replica),
+        ).observe(size_bytes)
+
+    def observe_tx_time(
+        self,
+        from_stage: int,
+        from_replica: int,
+        to_stage: int,
+        to_replica: int,
+        tx_time_s: float,
+    ) -> None:
+        if not self._log_stats:
+            return
+        _transfer_tx_family.labels(
+            model_name=self._model_name,
+            from_stage=str(from_stage),
+            from_replica=str(from_replica),
+            to_stage=str(to_stage),
+            to_replica=str(to_replica),
+        ).observe(tx_time_s)
+
+    # ---- RX side (record_transfer_rx hook) -------------------------------
+
+    def observe_rx_time(
+        self,
+        from_stage: int,
+        from_replica: int,
+        to_stage: int,
+        to_replica: int,
+        rx_time_s: float,
+    ) -> None:
+        if not self._log_stats:
+            return
+        _transfer_rx_family.labels(
+            model_name=self._model_name,
+            from_stage=str(from_stage),
+            from_replica=str(from_replica),
+            to_stage=str(to_stage),
+            to_replica=str(to_replica),
+        ).observe(rx_time_s)
+
+    def observe_in_flight_time(
+        self,
+        from_stage: int,
+        from_replica: int,
+        to_stage: int,
+        to_replica: int,
+        in_flight_time_s: float,
+    ) -> None:
+        if not self._log_stats:
+            return
+        _transfer_in_flight_family.labels(
+            model_name=self._model_name,
+            from_stage=str(from_stage),
+            from_replica=str(from_replica),
+            to_stage=str(to_stage),
+            to_replica=str(to_replica),
+        ).observe(in_flight_time_s)

vllm_omni/patch.py

@@ -1,3 +1,4 @@
+import logging
 import sys
 from functools import cached_property
 
@@ -129,6 +130,43 @@ for module_name, module in list(sys.modules.items()):
     if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest:
         module.EngineCoreRequest = OmniEngineCoreRequest
 
+# =============================================================================
+# Patch unregister_vllm_metrics to skip vllm:omni_* collectors
+# =============================================================================
+# WHY: Upstream unregister_vllm_metrics() uses `"vllm" in collector._name` to
+# strip collectors before each new PrometheusStatLogger registers, which is
+# how it avoids "Duplicated timeseries" when the same process spawns multiple
+# engines / orchestrators.  But its substring match also wipes our vllm:omni_*
+# families, so we must replace it with a scoped version that keeps ours.
+#
+# REMOVAL: Remove this patch once upstream vLLM adds
+# _STAT_LOGGER_METRIC_NAMES to vllm.v1.metrics.prometheus and scopes
+# unregister_vllm_metrics() to that set.  Track:
+# https://github.com/vllm-project/vllm/pull/42331
+import vllm.v1.metrics.prometheus as _vllm_prometheus  # noqa: E402
+
+_logger = logging.getLogger(__name__)
+
+
+def _scoped_unregister_vllm_metrics():
+    """Drop upstream vllm:* collectors but preserve vllm:omni_* (ours)."""
+    from prometheus_client import REGISTRY
+
+    for collector in list(REGISTRY._collector_to_names):
+        name = getattr(collector, "_name", "")
+        if "vllm" not in name:
+            continue
+        if name.startswith("vllm:omni_") or name.startswith("vllm_omni"):
+            continue
+        REGISTRY.unregister(collector)
+
+
+_vllm_prometheus.unregister_vllm_metrics = _scoped_unregister_vllm_metrics
+_logger.warning(
+    "Monkey-patched unregister_vllm_metrics() to scope drops to non-omni vllm:* collectors. "
+    "Remove this patch once vLLM adds _STAT_LOGGER_METRIC_NAMES."
+)
+
 
 # Patch: add qwen3_omni_moe to vllm's chat template fallback registry.
 # Qwen/Qwen3-Omni-30B-A3B-Instruct stores its chat_template in a standalone

vllm_omni/utils/audio.py

@@ -7,6 +7,9 @@ import numpy as np
 import torch
 from torchaudio.functional import melscale_fbanks
 
+from vllm_omni.metrics import definitions as _metric_defs
+from vllm_omni.outputs import OmniRequestOutput
+
 
 def mel_filter_bank(
     sr: int,
@@ -66,3 +69,39 @@ def peak_normalize(
         return audio
     target = 10.0 ** (db_level / 20.0)
     return audio * (target / peak)
+
+
+def audio_chunk_pcm_bytes(omni_res: OmniRequestOutput) -> int:
+    """Best-effort PCM byte count of the last audio chunk in ``omni_res``.
+
+    Used by the audio-streaming continuity tracker to size the player buffer.
+    Returns 0 when the chunk shape can't be interpreted — caller drops the
+    sample rather than recording a wrong byte count.
+    """
+    try:
+        final_res = omni_res.request_output
+        mm_output = final_res.outputs[0].multimodal_output
+        audio_data = mm_output.get("audio")
+        if isinstance(audio_data, list):
+            if not audio_data:
+                return 0
+            audio_tensor = audio_data[-1]
+        else:
+            audio_tensor = audio_data
+        if audio_tensor is None:
+            return 0
+        n_samples = int(audio_tensor.numel() if hasattr(audio_tensor, "numel") else audio_tensor.size)
+        # PCM s16le mono → 2 bytes per sample. This matches what
+        # _create_audio_choice serialises via CreateAudio (response_format="wav").
+        return max(n_samples, 0) * 2
+    except Exception:
+        return 0
+
+
+def audio_chunk_sample_rate(omni_res: OmniRequestOutput) -> int:
+    """Resolve audio sample rate for the request's audio stream."""
+    try:
+        mm_output = omni_res.request_output.outputs[0].multimodal_output
+    except Exception:
+        return _metric_defs.DEFAULT_AUDIO_SAMPLE_RATE
+    return _metric_defs.resolve_audio_sample_rate(mm_output)