[BugFix][Metrics] Fix per-stage time and token stats in metrics API (#4974)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

docs/design/metrics.md

@@ -148,7 +148,13 @@ Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, retur
 
 `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.
+The orchestrator side does not add its own throttle on top. Scheduler gauges are
+recorded when `raw_outputs.scheduler_stats is not None` (i.e. this replica's
+scheduler passed its own 1 Hz gate), while output-batch request/token metrics
+are recorded from `IterationStats` whenever a non-empty output batch is
+processed. 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
 

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -21,6 +21,7 @@ import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
 from vllm.sampling_params import SamplingParams
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.engine import EngineCoreEventType
+from vllm.v1.metrics.stats import PrefillStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
@@ -93,6 +94,17 @@ def _make_update(**kwargs) -> StreamingUpdate:
     return StreamingUpdate(**defaults)
 
 
+def test_generation_scheduler_records_prefill_stats_for_metrics() -> None:
+    request = _make_request(prompt_token_ids=[1, 2, 3, 4])
+    request.prefill_stats = PrefillStats()
+
+    OmniGenerationScheduler._record_prefill_stats(request)
+
+    assert request.prefill_stats.num_prompt_tokens == 4
+    assert request.prefill_stats.num_computed_tokens == 4
+    assert request.prefill_stats.num_cached_tokens == 0
+
+
 def test_resumable_generation_stop_marks_segment_boundary() -> None:
     session = _make_request(request_id="req-generation-segment")
     session.status = RequestStatus.RUNNING

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -10,8 +10,10 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
+from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
 from vllm.v1.request import RequestStatus
 
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
 from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt
 from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
@@ -37,8 +39,10 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
         external_req_id=external_req_id or req_id,
         status=status,
         prompt_token_ids=[],
+        num_prompt_tokens=0,
         num_computed_tokens=0,
         num_output_placeholders=0,
+        prefill_stats=None,
         additional_information=None,
         is_finished=lambda: status == RequestStatus.FINISHED_STOPPED,
     )
@@ -485,6 +489,48 @@ def test_load_poll_ar_request_additional_information_concats_tensors(build_adapt
     assert request.additional_information["meta"]["finished"].item() is True
 
 
+def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter):
+    adapter, connector = build_adapter(stage_id=2, model_mode="generation")
+    request = _req("req-later-chunk", RequestStatus.WAITING, external_req_id="ext-later-chunk")
+    request.prefill_stats = PrefillStats()
+    adapter.request_ids_mapping[request.request_id] = request.external_req_id
+
+    connector.get.return_value = (
+        {
+            "codes": {"audio": torch.tensor([7, 8], dtype=torch.long)},
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+        },
+        8,
+    )
+    assert adapter._poll_single_request(request) is True
+    OmniGenerationScheduler._record_prefill_stats(request)
+    first_chunk_stats = request.prefill_stats
+    request.prefill_stats = None
+
+    connector.get.return_value = (
+        {
+            "codes": {"audio": torch.tensor([9, 10, 11], dtype=torch.long)},
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+        },
+        8,
+    )
+    assert adapter._poll_single_request(request) is True
+    assert isinstance(request.prefill_stats, PrefillStats)
+    OmniGenerationScheduler._record_prefill_stats(request)
+    second_chunk_stats = request.prefill_stats
+
+    assert first_chunk_stats is not None
+    assert first_chunk_stats.num_prompt_tokens == 2
+    assert second_chunk_stats is not None
+    assert second_chunk_stats.num_prompt_tokens == 3
+
+    prompt_token_stats = PromptTokenStats()
+    prompt_token_stats.update_from_output(first_chunk_stats)
+    prompt_token_stats.update_from_output(second_chunk_stats)
+    assert prompt_token_stats.total == 5
+    assert prompt_token_stats.computed == 5
+
+
 def test_sender_only_adapter_does_not_park_or_clear_requests(build_adapter):
     adapter, _ = build_adapter(stage_id=1, connector_extra={"role": "sender"})
     request = _req("req-1", RequestStatus.WAITING)

tests/engine/test_async_omni_engine_stage_init.py

@@ -172,6 +172,41 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     )
 
 
+def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.stage_configs = [types.SimpleNamespace()]
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.single_stage_mode = False
+    engine.diffusion_batch_size = 1
+    engine.async_chunk = False
+    engine.tokenizer = None
+    engine._single_stage_id_filter = None
+    engine._omni_master_address = None
+    engine._omni_master_port = None
+    engine._omni_dp_size_local = 1
+    engine._omni_heartbeat_timeout = 30.0
+    engine._omni_lb_policy = "random"
+    engine.request_queue = types.SimpleNamespace()
+    engine._log_stats = True
+
+    captured: dict[str, object] = {}
+    runtime = types.SimpleNamespace(stage_pools=[], initialize=lambda: None)
+
+    def _capture_create_stage_runtime(**kwargs):
+        captured.update(kwargs)
+        return runtime
+
+    monkeypatch.setattr(engine_mod, "create_stage_runtime", _capture_create_stage_runtime)
+
+    engine._initialize_stages(stage_init_timeout=7)
+
+    assert captured["stage_init_timeout"] == 7
+    assert captured["log_stats"] is True
+
+
 def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -342,6 +377,95 @@ def test_stage_runtime_initializes_stage_pools(monkeypatch):
     assert runtime.stage_pools[1].output_processor is stage1_output_processor
 
 
+def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+        log_stats=True,
+    )
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    plan = _make_llm_plan(0, stage_id=0, vllm_config=cfg).replicas[0]
+    plan.engine_args_dict = {}
+
+    captured: dict[str, object] = {}
+    addresses = types.SimpleNamespace(
+        inputs=["tcp://127.0.0.1:1"],
+        outputs=["tcp://127.0.0.1:2"],
+        frontend_stats_publish_address=None,
+    )
+    resources = types.SimpleNamespace(manager=object(), coordinator=None, addresses=addresses)
+    stage_client = types.SimpleNamespace()
+
+    @contextlib.contextmanager
+    def _capture_launch_stage_replica(**kwargs):
+        captured["log_stats"] = kwargs["log_stats"]
+        yield resources
+
+    monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: [])
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", _capture_launch_stage_replica)
+    monkeypatch.setattr(
+        runtime_mod.StageEngineCoreClientBase,
+        "make_async_mp_client",
+        lambda **kwargs: (captured.__setitem__("client_log_stats", kwargs["log_stats"]) or stage_client),
+    )
+
+    assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client
+    assert captured["log_stats"] is True
+    assert captured["client_log_stats"] is True
+
+
+def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+        log_stats=True,
+    )
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=cfg)
+    stage_client = types.SimpleNamespace(
+        stage_type="llm",
+        is_comprehension=False,
+        final_output=True,
+        final_output_type=None,
+        default_sampling_params=types.SimpleNamespace(),
+    )
+    captured: dict[str, object] = {}
+    output_processor = object()
+
+    def _capture_build_llm_stage_output_processor(plan, stage_vllm_config, *, log_stats=False):
+        captured["plan"] = plan
+        captured["stage_vllm_config"] = stage_vllm_config
+        captured["log_stats"] = log_stats
+        return output_processor
+
+    monkeypatch.setattr(
+        runtime_mod,
+        "build_llm_stage_output_processor",
+        _capture_build_llm_stage_output_processor,
+    )
+
+    pools = runtime._assemble_stage_pools([stage_plan], {0: [stage_client]})
+
+    assert pools[0].output_processor is output_processor
+    assert captured == {
+        "plan": stage_plan,
+        "stage_vllm_config": cfg,
+        "log_stats": True,
+    }
+
+
 def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch):
     import vllm_omni.engine.stage_runtime as runtime_mod
 
@@ -537,10 +661,11 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
         stage_configs=[],
         model="dummy-model",
         config_path="dummy-config",
-        stage_init_timeout=302,
+        stage_init_timeout=1,
         diffusion_batch_size=1,
         async_chunk=False,
     )
+    stage_init_timeout = 302
 
     fake_vllm_config = types.SimpleNamespace()
     fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None)
@@ -587,14 +712,14 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
     )
 
     try:
-        runtime._initialize_local_llm_replica(plan, 302)
+        runtime._initialize_local_llm_replica(plan, stage_init_timeout)
     finally:
         if prev_device_env is None:
             os.environ.pop(device_env_var, None)
         else:
             os.environ[device_env_var] = prev_device_env
 
-    assert captured_timeout == 302
+    assert captured_timeout == stage_init_timeout
 
 
 def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_path):

tests/engine/test_orchestrator.py

@@ -15,6 +15,7 @@ import pytest
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
@@ -145,7 +146,7 @@ class FakeStageClient:
         try:
             return self._engine_core_outputs.get_nowait()
         except queue.Empty:
-            return SimpleNamespace(outputs=[])
+            return SimpleNamespace(outputs=[], scheduler_stats=None, finished_requests=None)
 
     def get_diffusion_output_nowait(self):
         try:
@@ -253,12 +254,40 @@ class FakeOutputProcessor:
         return None
 
 
+class IterationStatsOutputProcessor(FakeOutputProcessor):
+    def process_outputs(self, *_args, **kwargs):
+        iteration_stats = kwargs.get("iteration_stats")
+        if iteration_stats is None and len(_args) >= 3:
+            iteration_stats = _args[2]
+        assert iteration_stats is not None
+        iteration_stats.num_generation_tokens += 3
+        return super().process_outputs(*_args, **kwargs)
+
+
+class RecordingOutputProcessor(FakeOutputProcessor):
+    def __init__(self, *, request_outputs: list[object] | None = None) -> None:
+        super().__init__(request_outputs=request_outputs)
+        self.process_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = []
+
+    def process_outputs(self, *args, **kwargs):
+        self.process_calls.append((args, kwargs))
+        return super().process_outputs(*args, **kwargs)
+
+
+class RecordingStatLogger:
+    def __init__(self) -> None:
+        self.records: list[tuple[Any, Any, int]] = []
+
+    def record(self, scheduler_stats, iteration_stats, *, engine_idx: int = 0) -> None:
+        self.records.append((scheduler_stats, iteration_stats, engine_idx))
+
+
 def _sampling_params(max_tokens: int = 4) -> SamplingParams:
     return SamplingParams(max_tokens=max_tokens)
 
 
 def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace:
-    return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None)
+    return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None)
 
 
 def _build_request_output(
@@ -329,6 +358,7 @@ def _build_harness(
     output_processors: list[object] | None = None,
     stage_vllm_configs: list[object] | None = None,
     async_chunk: bool = False,
+    log_stats: bool = False,
     stage_pools: list[StagePool] | None = None,
 ) -> OrchestratorFixture:
     """Build an Orchestrator test harness.
@@ -364,6 +394,7 @@ def _build_harness(
                 rpc_async_queue=rpc_queue.async_q,
                 stage_pools=stage_pools,
                 async_chunk=async_chunk,
+                log_stats=log_stats,
             )
             ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue))
             await orchestrator.run()
@@ -2137,38 +2168,143 @@ async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator
         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."
+@pytest.mark.asyncio
+async def test_orchestrator_records_iteration_stats_without_scheduler_stats(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = IterationStatsOutputProcessor(
+        request_outputs=[_build_request_output("req-stats", token_ids=[1], finished=True)]
     )
-    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."
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    stat_logger = RecordingStatLogger()
+    orchestrator_fixture.orchestrator._stat_logger = stat_logger
+    request = SimpleNamespace(request_id="req-stats", prompt_token_ids=[1, 2])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-stats",
+            prompt=request,
+            original_prompt={"prompt": "hello"},
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+        )
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+
+        stage0.push_engine_core_outputs(_engine_core_outputs("raw-without-scheduler-stats", 1.0))
+
+        await _wait_for(lambda: len(stat_logger.records) == 1)
+        scheduler_stats, iteration_stats, engine_idx = stat_logger.records[0]
+        assert scheduler_stats is None
+        assert iteration_stats is not None
+        assert iteration_stats.num_generation_tokens == 3
+        assert engine_idx == 0
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_records_scheduler_stats_without_outputs(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    stat_logger = RecordingStatLogger()
+    orchestrator_fixture.orchestrator._stat_logger = stat_logger
+    scheduler_stats = SimpleNamespace(num_running_reqs=0, kv_cache_usage=0.5)
+
+    try:
+        stage0.push_engine_core_outputs(
+            SimpleNamespace(
+                outputs=[],
+                timestamp=1.0,
+                scheduler_stats=scheduler_stats,
+                finished_requests=None,
+            )
+        )
+
+        await _wait_for(lambda: len(stat_logger.records) == 1)
+        recorded_scheduler_stats, iteration_stats, engine_idx = stat_logger.records[0]
+        assert recorded_scheduler_stats is scheduler_stats
+        assert iteration_stats is None
+        assert engine_idx == 0
+        assert processor.process_calls[0][0][2] is None
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_does_not_build_iteration_stats_for_finished_only_batch(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    stat_logger = RecordingStatLogger()
+    orchestrator_fixture.orchestrator._stat_logger = stat_logger
+
+    try:
+        stage0.push_engine_core_outputs(
+            SimpleNamespace(
+                outputs=[],
+                timestamp=1.0,
+                scheduler_stats=None,
+                finished_requests={"req-finished"},
+            )
+        )
+
+        await _wait_for(lambda: bool(processor.process_calls))
+        assert processor.process_calls[0][0][2] is None
+        assert stat_logger.records == []
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_does_not_build_iteration_stats_without_stat_logger(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    orchestrator_fixture.orchestrator._stat_logger = None
+
+    try:
+        stage0.push_engine_core_outputs(_engine_core_outputs("raw-output", 1.0))
+
+        await _wait_for(lambda: bool(processor.process_calls))
+        assert processor.process_calls[0][0][2] is None
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_preserves_none_iteration_stats() -> None:
+    client = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    pool = StagePool(
+        0,
+        [client],
+        output_processor=processor,
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
     )
+    raw_outputs = SimpleNamespace(outputs=["raw"], timestamp=1.0, scheduler_stats=None)
+
+    await pool.process_llm_raw_outputs(0, raw_outputs, iteration_stats=None)
+
+    assert processor.process_calls[0][0][2] is None
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_process_llm_raw_outputs_mutates_iteration_stats() -> None:
+    client = FakeStageClient(stage_type="llm", final_output=True)
+    processor = IterationStatsOutputProcessor()
+    pool = StagePool(
+        0,
+        [client],
+        output_processor=processor,
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+    raw_outputs = SimpleNamespace(outputs=["raw"], timestamp=1.0, scheduler_stats=None)
+    iteration_stats = IterationStats()
+
+    await pool.process_llm_raw_outputs(0, raw_outputs, iteration_stats=iteration_stats)
+
+    assert iteration_stats.num_generation_tokens == 3
 
 
 @pytest.mark.asyncio

tests/engine/test_output_processor.py

@@ -10,9 +10,11 @@ import pytest
 import torch
 from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
-from vllm.v1.engine import FinishReason
+from vllm.v1.engine import EngineCoreEvent, EngineCoreEventType, FinishReason
 from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
+from vllm.v1.metrics.stats import IterationStats, PrefillStats
 
+from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.outputs import output_processor
 from vllm_omni.outputs.output_modality import OutputModality, OutputModalityNames
 from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState
@@ -616,3 +618,48 @@ def test_no_detokenizer_final_only():
     result = s.make_request_output([], None, FinishReason.STOP, None)
     assert result is not None
     assert AUDIO in result.outputs[0].multimodal_output
+
+
+def test_mm_only_outputs_update_iteration_stats():
+    processor = MultimodalOutputProcessor(
+        tokenizer=None,
+        log_stats=True,
+        engine_core_output_type=AUDIO,
+    )
+    state = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
+    processor.request_states[state.request_id] = state
+    processor.external_req_ids[state.external_req_id].append(state.request_id)
+
+    prefill_stats = PrefillStats()
+    prefill_stats.set(
+        num_prompt_tokens=4,
+        num_local_cached_tokens=0,
+        num_external_cached_tokens=0,
+    )
+    output = OmniEngineCoreOutput(
+        request_id=state.request_id,
+        new_token_ids=[10, 11],
+        finish_reason=FinishReason.STOP,
+        events=[
+            EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0),
+            EngineCoreEvent(EngineCoreEventType.SCHEDULED, 2.0),
+        ],
+        prefill_stats=prefill_stats,
+        multimodal_output={AUDIO: torch.ones(1, 8)},
+    )
+    iteration_stats = IterationStats()
+
+    processor.process_outputs(
+        [output],
+        engine_core_timestamp=3.0,
+        iteration_stats=iteration_stats,
+    )
+
+    assert iteration_stats.num_prompt_tokens == 4
+    assert iteration_stats.num_generation_tokens == 2
+    assert len(iteration_stats.time_to_first_tokens_iter) == 1
+    assert len(iteration_stats.finished_requests) == 1
+    finished = iteration_stats.finished_requests[0]
+    assert finished.finish_reason == FinishReason.STOP
+    assert finished.num_prompt_tokens == state.prompt_len
+    assert finished.num_generation_tokens == 2

tests/engine/test_single_stage_mode.py

@@ -778,6 +778,7 @@ class TestSingleStageReplicaInitialization:
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",
             omni_master_port=26000,
+            log_stats=True,
         )
         runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer)
         runtime._omni_master_server.get_stage_config.return_value = {"stage_id": 1, "stage_type": "llm"}
@@ -811,10 +812,17 @@ class TestSingleStageReplicaInitialization:
         sentinel_client = SimpleNamespace()
 
         mock_connect = mocker.patch.object(runtime_mod, "connect_remote_engine_cores", side_effect=_fake_connect)
+        client_kwargs: dict[str, Any] = {}
+
+        def _capture_make_async_mp_client(**kwargs):
+            client_kwargs.update(kwargs)
+            events.append("attach")
+            return sentinel_client
+
         mocker.patch.object(
             StageEngineCoreClientBase,
             "make_async_mp_client",
-            side_effect=lambda **_: (events.append("attach"), sentinel_client)[1],
+            side_effect=_capture_make_async_mp_client,
         )
 
         result = runtime._initialize_remote_replica(plan, stage_init_timeout=60)
@@ -824,6 +832,7 @@ class TestSingleStageReplicaInitialization:
         assert mock_connect.call_args.kwargs["vllm_config"].parallel_config.data_parallel_size_local == 0
         assert mock_connect.call_args.kwargs["stage_id"] == 1
         assert mock_connect.call_args.kwargs["replica_id"] == 0
+        assert client_kwargs["log_stats"] is True
         assert events == ["enter", "exit", "attach"]
 
     def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture):

tests/entrypoints/test_stream_finish_reason.py

@@ -17,6 +17,7 @@ Key invariants tested:
 """
 
 import json
+import time
 from unittest.mock import MagicMock
 
 import pytest
@@ -80,28 +81,37 @@ def _make_audio_omni_output(
     request_id: str = "test-req",
     index: int = 0,
     num_prompt_tokens: int = 3,
+    stage_id: int | None = None,
+    replica_id: int | None = None,
+    audio_samples: int = 0,
 ) -> OmniRequestOutput:
     """Build an OmniRequestOutput for audio (no torch dependency)."""
+    completion = CompletionOutput(
+        index=index,
+        text="",
+        token_ids=[],
+        cumulative_logprob=0.0,
+        logprobs=None,
+        finish_reason="stop",
+        stop_reason=None,
+    )
+    if audio_samples > 0:
+        completion.multimodal_output = {
+            "audio": [MagicMock(numel=MagicMock(return_value=audio_samples))],
+            "sr": 24000,
+        }
     res = RequestOutput(
         request_id=request_id,
         prompt="test",
         prompt_token_ids=list(range(num_prompt_tokens)),
         prompt_logprobs=None,
-        outputs=[
-            CompletionOutput(
-                index=index,
-                text="",
-                token_ids=[],
-                cumulative_logprob=0.0,
-                logprobs=None,
-                finish_reason="stop",
-                stop_reason=None,
-            )
-        ],
+        outputs=[completion],
         finished=True,
     )
     return OmniRequestOutput(
         request_id=request_id,
+        stage_id=stage_id,
+        replica_id=replica_id,
         final_output_type="audio",
         request_output=res,
         finished=True,
@@ -321,6 +331,74 @@ async def test_single_modality_audio_only_one_stop():
     assert finish_reasons[-1] == "stop"
 
 
+@pytest.mark.parametrize(
+    ("output_replica_id", "pool_replica_id", "expected_replica_id", "pool_lookup"),
+    [
+        (7, None, 7, False),
+        (None, 5, 5, True),
+    ],
+)
+@pytest.mark.asyncio
+async def test_streaming_audio_metrics_resolve_replica_id(
+    monkeypatch,
+    output_replica_id,
+    pool_replica_id,
+    expected_replica_id,
+    pool_lookup,
+):
+    import vllm_omni.entrypoints.openai.serving_chat as serving_chat_mod
+    from vllm_omni.entrypoints.client_request_state import ClientRequestState
+
+    serving_chat = _build_serving_chat()
+    request = _make_request(modalities=["audio"])
+    req_state = ClientRequestState(
+        request_id="internal-req",
+        external_request_id="test-req",
+    )
+    req_state.request_arrival_ts = time.time() - 1.0
+    serving_chat.engine_client.request_states = {"internal-req": req_state}
+    serving_chat.engine_client.mod_metrics = object()
+    serving_chat.engine_client.engine = MagicMock()
+    serving_chat.engine_client.engine.stage_pools = [MagicMock(), MagicMock(), MagicMock()]
+    serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.return_value = pool_replica_id
+
+    first_packet_calls = []
+    finalize_calls = []
+
+    def fake_observe_audio_first_packet(*args, **kwargs):
+        first_packet_calls.append(kwargs)
+
+    def fake_observe_audio_streaming_finalize(*args, **kwargs):
+        finalize_calls.append(kwargs)
+
+    monkeypatch.setattr(serving_chat_mod, "observe_audio_first_packet", fake_observe_audio_first_packet)
+    monkeypatch.setattr(serving_chat_mod, "observe_audio_streaming_finalize", fake_observe_audio_streaming_finalize)
+
+    async def result_generator():
+        yield _make_audio_omni_output(stage_id=2, replica_id=output_replica_id, audio_samples=2400)
+
+    await _collect_stream(
+        serving_chat.chat_completion_stream_generator(
+            request=request,
+            result_generator=result_generator(),
+            request_id="test-req",
+            model_name="test-model",
+            conversation=[],
+            tokenizer=MagicMock(),
+            request_metadata=MagicMock(),
+        )
+    )
+
+    if pool_lookup:
+        serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.assert_called_once_with("internal-req")
+    else:
+        serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.assert_not_called()
+    assert first_packet_calls[0]["stage_id"] == 2
+    assert first_packet_calls[0]["replica_id"] == expected_replica_id
+    assert finalize_calls[0]["stage_id"] == 2
+    assert finalize_calls[0]["replica_id"] == expected_replica_id
+
+
 # ---------------------------------------------------------------------------
 # Tests: fallback stop chunk when declared modality is not produced
 # ---------------------------------------------------------------------------

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -56,6 +56,24 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             )
         self._latest_omni_connector_output: OmniConnectorOutput | None = None
 
+    @staticmethod
+    def _record_prefill_stats(request: Request) -> None:
+        """Mirror upstream first-prefill prompt accounting for generation stages.
+
+        Async-chunk generation stages can grow ``prompt_token_ids`` after the
+        first schedule; until this path tracks final chunk lengths, this records
+        the first scheduled prompt snapshot. Keep this in sync with vLLM
+        Scheduler.schedule() if this fast path gains prefix-cache or external-KV
+        accounting.
+        """
+        if request.prefill_stats is None:
+            return
+        request.prefill_stats.set(
+            num_prompt_tokens=request.num_prompt_tokens,
+            num_local_cached_tokens=0,
+            num_external_cached_tokens=0,
+        )
+
     def _handle_stopped_request(self, request: Request) -> bool:
         if (
             request.resumable
@@ -147,6 +165,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 break
             if self.log_stats:
                 request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
+            if num_computed_tokens == 0:
+                self._record_prefill_stats(request)
             req_to_new_blocks[request.request_id] = new_blocks
             num_scheduled_tokens[request.request_id] = num_new_tokens
             cached_prompt_token_ids[request.request_id] = request.prompt_token_ids
@@ -213,6 +233,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             self.running.append(request)
             if self.log_stats:
                 request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
+            if request.num_computed_tokens == 0:
+                self._record_prefill_stats(request)
 
             req_to_new_blocks[request.request_id] = new_blocks
             num_scheduled_tokens[request.request_id] = num_new_tokens

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -7,6 +7,7 @@ from collections.abc import Callable, Mapping
 from typing import Any
 
 import torch
+from vllm.v1.metrics.stats import PrefillStats
 from vllm.v1.request import Request, RequestStatus
 
 from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload
@@ -109,6 +110,12 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         num_placeholders = int(getattr(request, "num_output_placeholders", 0) or 0)
         return max(0, num_computed - num_placeholders)
 
+    @staticmethod
+    def _refresh_generation_chunk_prefill_state(request: Request) -> None:
+        request.num_prompt_tokens = len(request.prompt_token_ids)
+        if getattr(request, "prefill_stats", None) is None:
+            request.prefill_stats = PrefillStats()
+
     @classmethod
     def create_connector(cls, model_config: Any):
         connector_config = getattr(model_config, "stage_connector_config", None)
@@ -249,15 +256,18 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 new_ids = payload_data.get("codes", {}).get("audio")
                 has_tensor_codes = isinstance(new_ids, torch.Tensor)
                 use_tensor_codes = has_tensor_codes and new_ids.ndim >= 2
+                prompt_token_ids: list[int]
                 if use_tensor_codes:
-                    request.prompt_token_ids = [0] if new_ids.numel() > 0 else []
+                    prompt_token_ids = [0] if new_ids.numel() > 0 else []
                 elif has_tensor_codes:
                     new_ids = new_ids.tolist()
+                    prompt_token_ids = new_ids
                 elif new_ids is None:
                     new_ids = []
-                    request.prompt_token_ids = new_ids
-                if not use_tensor_codes:
-                    request.prompt_token_ids = new_ids
+                    prompt_token_ids = new_ids
+                else:
+                    prompt_token_ids = new_ids
+                request.prompt_token_ids = prompt_token_ids
                 prev_info = getattr(request, "additional_information", None)
                 info = dict(prev_info) if isinstance(prev_info, dict) else {}
                 for key, value in payload_data.items():
@@ -293,6 +303,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                     # chunk as ready, otherwise Stage1 can consume before the
                     # first DAC frame arrives.
                     return False
+                self._refresh_generation_chunk_prefill_state(request)
 
             # Mark as finished for consumption
             self._finished_load_reqs.add(req_id)

vllm_omni/engine/async_omni_engine.py

@@ -331,6 +331,7 @@ class AsyncOmniEngine:
             omni_heartbeat_timeout=self._omni_heartbeat_timeout,
             omni_lb_policy=self._omni_lb_policy,
             request_queue=self.request_queue,
+            log_stats=self._log_stats,
         )
         self._runtime.initialize()
 

vllm_omni/engine/orchestrator.py

@@ -362,6 +362,7 @@ class Orchestrator:
         self.async_chunk = bool(async_chunk)
         self.num_stages = len(stage_pools)
         self.stage_pools: list[StagePool] = stage_pools
+        self.log_stats = log_stats
         self._orch_monitor = create_orch_monitor(
             enabled=enable_orch_monitor,
             replica_sampler=self._sample_replica_metrics,
@@ -446,7 +447,10 @@ class Orchestrator:
         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.
+        when the user did not request stats. When stats are enabled, record()
+        must still receive IterationStats even on ticks where scheduler_stats
+        is None; upstream PrometheusStatLogger records token/request metrics
+        from iteration_stats independently of scheduler gauges.
         """
         self._running_counter = running_counter
         self._transfer_emitter = transfer_emitter
@@ -937,19 +941,17 @@ class Orchestrator:
                                 )
                                 if req_state.streaming.enabled:
                                     await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state)
-                            # 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
+                            iteration_stats = (
+                                IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None
+                            )
                             raw_output = await pool.process_llm_raw_outputs(
                                 replica_id,
                                 raw_outputs,
                                 iteration_stats=iteration_stats,
                             )
-                            if record_stats:
+                            if self._stat_logger is not None and (
+                                raw_outputs.scheduler_stats is not None or iteration_stats is not None
+                            ):
                                 self._stat_logger.record(
                                     raw_outputs.scheduler_stats,
                                     iteration_stats,

vllm_omni/engine/stage_pool.py

@@ -1081,7 +1081,11 @@ class StagePool:
     async def _poll_stage_raw(self, client: StagePoolLLMClient) -> EngineCoreOutputs | None:
         """Pull raw EngineCoreOutputs from a stage replica without processing."""
         outputs = await client.get_output_async()
-        if not outputs.outputs:
+        # Keep scheduler-only / finished-only batches. Omni schedulers mirror
+        # upstream by emitting SchedulerStats on throttled ticks even when no
+        # request output is produced, and dropping those batches loses KV/queue
+        # gauges for that interval.
+        if not outputs.outputs and outputs.scheduler_stats is None and not outputs.finished_requests:
             return None
         return outputs
 
@@ -1097,7 +1101,6 @@ class StagePool:
             return []
         client = cast(StagePoolLLMClient, raw_client)
         processor = self.output_processor
-        iteration_stats = IterationStats()
         processed = processor.process_outputs(
             raw_outputs.outputs,
             raw_outputs.timestamp,

vllm_omni/engine/stage_runtime.py

@@ -126,6 +126,7 @@ class StageRuntime:
         diffusion_batch_size: int,
         async_chunk: bool,
         tokenizer: str | None = None,
+        log_stats: bool = False,
     ) -> None:
         self._stage_configs = stage_configs
         self._model = model
@@ -134,6 +135,7 @@ class StageRuntime:
         self._diffusion_batch_size = diffusion_batch_size
         self._async_chunk = async_chunk
         self._tokenizer = tokenizer
+        self._log_stats = log_stats
         self._num_stages = len(stage_configs)
 
         # Populated by initialize()
@@ -575,7 +577,7 @@ class StageRuntime:
                 with launch_stage_replica(
                     vllm_config=vllm_config,
                     executor_class=executor_class,
-                    log_stats=False,
+                    log_stats=self._log_stats,
                     stage_id=plan.metadata.stage_id,
                     replica_id=plan.replica_id,
                     stage_config=plan.stage_cfg,
@@ -594,6 +596,7 @@ class StageRuntime:
             stage_client = StageEngineCoreClientBase.make_async_mp_client(
                 vllm_config=vllm_config,
                 executor_class=executor_class,
+                log_stats=self._log_stats,
                 metadata=plan.metadata,
                 client_addresses=self._client_addresses_from_zmq(resources.addresses),
                 engine_manager=resources.manager,
@@ -704,7 +707,11 @@ class StageRuntime:
                 stage_vllm_config = plan.replicas[0].stage_vllm_config
                 if stage_vllm_config is None:
                     raise RuntimeError(f"Stage {plan.stage_id} is missing vllm_config")
-                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(
@@ -742,10 +749,11 @@ class DistStageRuntime(StageRuntime):
         stage_init_timeout: int,
         diffusion_batch_size: int,
         async_chunk: bool,
-        tokenizer: str | None = None,
         single_stage_id_filter: int | None,
         omni_master_address: str,
         omni_master_port: int,
+        tokenizer: str | None = None,
+        log_stats: bool = False,
         omni_dp_size_local: int = 1,
         omni_heartbeat_timeout: float = 30.0,
         omni_lb_policy: str = "random",
@@ -759,6 +767,7 @@ class DistStageRuntime(StageRuntime):
             diffusion_batch_size=diffusion_batch_size,
             async_chunk=async_chunk,
             tokenizer=tokenizer,
+            log_stats=log_stats,
         )
         self._single_stage_id_filter = single_stage_id_filter
         self._omni_master_address = omni_master_address
@@ -1051,6 +1060,7 @@ class DistStageRuntime(StageRuntime):
             client = StageEngineCoreClientBase.make_async_mp_client(
                 vllm_config=vllm_config,
                 executor_class=ctx.executor_class,
+                log_stats=self._log_stats,
                 metadata=metadata,
                 client_addresses=client_addresses,
                 engine_manager=resources.manager,
@@ -1089,6 +1099,7 @@ def create_stage_runtime(
     omni_heartbeat_timeout: float = 30.0,
     omni_lb_policy: str = "random",
     request_queue: janus.Queue[EngineQueueMessage] | None = None,
+    log_stats: bool = False,
 ) -> StageRuntime:
     """Factory: select StageRuntime or DistStageRuntime."""
     if single_stage_mode:
@@ -1102,6 +1113,7 @@ def create_stage_runtime(
             diffusion_batch_size=diffusion_batch_size,
             async_chunk=async_chunk,
             tokenizer=tokenizer,
+            log_stats=log_stats,
             single_stage_id_filter=single_stage_id_filter,
             omni_master_address=omni_master_address,
             omni_master_port=omni_master_port,
@@ -1118,4 +1130,5 @@ def create_stage_runtime(
         diffusion_batch_size=diffusion_batch_size,
         async_chunk=async_chunk,
         tokenizer=tokenizer,
+        log_stats=log_stats,
     )

vllm_omni/entrypoints/omni_base.py

@@ -612,6 +612,7 @@ class OmniBase(PDDisaggregationMixin):
             request_id=req_id or "",
             finished=finished,
             stage_id=stage_id,
+            replica_id=result.replica_id,
             final_output_type=output_type,
             request_output=engine_outputs,
             images=images,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -2062,15 +2062,19 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     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
-                        )
+                        replica_id = getattr(omni_res, "replica_id", None)
+                        if replica_id is None:
+                            stage_pools = getattr(self.engine_client.engine, "stage_pools", None)
+                            # Fallback for older outputs. The orchestrator binds
+                            # requests by internal id, but can release that
+                            # binding before the API sees the final audio chunk,
+                            # so prefer the OutputMessage replica copied onto
+                            # OmniRequestOutput above.
+                            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(

vllm_omni/outputs/__init__.py

@@ -83,6 +83,7 @@ class OmniRequestOutput:
         request_id: Unique identifier for this request
         finished: Whether generation is complete
         stage_id: Identifier of the stage that produced this output (pipeline mode)
+        replica_id: Identifier of the stage replica that produced this output
         final_output_type: Type of output ("text", "image", "audio", "latents")
         request_output: The underlying RequestOutput from the stage (pipeline mode)
         images: List of generated PIL images (diffusion mode)
@@ -96,6 +97,7 @@ class OmniRequestOutput:
 
     # Pipeline stage fields
     stage_id: int | None = None
+    replica_id: int | None = None
     final_output_type: str = "text"
     request_output: RequestOutput | None = None
 
@@ -154,6 +156,7 @@ class OmniRequestOutput:
         stage_id: int,
         final_output_type: str,
         request_output: RequestOutput,
+        replica_id: int | None = None,
     ) -> "OmniRequestOutput":
         """Create output from pipeline stage.
 
@@ -168,6 +171,7 @@ class OmniRequestOutput:
         return cls(
             request_id=getattr(request_output, "request_id", ""),
             stage_id=stage_id,
+            replica_id=replica_id,
             final_output_type=final_output_type,
             request_output=request_output,
             finished=True,

vllm_omni/outputs/output_processor.py

@@ -523,7 +523,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 upstream_outputs.append(eco)
 
         # Handle multimodal-only outputs (generation stages) locally.
-        mm_request_outputs = self._process_mm_only_outputs(mm_only_outputs)
+        mm_request_outputs = self._process_mm_only_outputs(
+            mm_only_outputs,
+            engine_core_timestamp=engine_core_timestamp,
+            iteration_stats=iteration_stats,
+        )
 
         # Delegate text/pooling outputs to upstream.
         processed = super().process_outputs(
@@ -537,6 +541,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
     def _process_mm_only_outputs(
         self,
         engine_core_outputs: list[EngineCoreOutput],
+        *,
+        engine_core_timestamp: float | None = None,
+        iteration_stats: IterationStats | None = None,
     ) -> list[OmniRequestOutput | PoolingRequestOutput]:
         """Handle outputs from generation stages that have no detokenizer.
 
@@ -554,7 +561,15 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             stop_reason = eco.stop_reason
             kv_transfer_params = eco.kv_transfer_params
             routed_experts = eco.routed_experts
-            req_state.num_cached_tokens = eco.num_cached_tokens
+            self._update_stats_from_output(
+                req_state,
+                eco,
+                engine_core_timestamp,
+                iteration_stats,
+            )
+            prefill_stats = getattr(eco, "prefill_stats", None)
+            if prefill_stats is not None:
+                req_state.num_cached_tokens = prefill_stats.num_cached_tokens
             req_state.is_prefilling = False
 
             is_non_final_audio_chunk = (
@@ -578,6 +593,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             is_segment_finished = bool(getattr(eco, "is_segment_finished", False))
             if finish_reason is not None and not is_segment_finished and not is_non_final_audio_chunk:
                 self._finish_request(req_state)
+                self._update_stats_from_finished(
+                    req_state,
+                    finish_reason,
+                    iteration_stats,
+                )
         return request_outputs
 
     def _update_stats_from_output(