[Bugfix] Recover Omni benchmark TPOT from Stage 0 metrics (#6818)

Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>

tests/benchmarks/patch/test_patch.py

@@ -9,6 +9,7 @@ import asyncio
 import base64
 import json
 import time
+from argparse import Namespace
 from types import SimpleNamespace
 
 import pytest
@@ -20,6 +21,7 @@ from vllm_omni.benchmarks.patch.patch import (
     _apply_stage0_token_timings,
     async_request_openai_chat_omni_completions,
     async_request_openai_realtime_duplex,
+    should_request_stage_metrics,
 )
 from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector
 
@@ -274,6 +276,107 @@ def create_sse_chunk(data_dict):
     return f"data: {json.dumps(data_dict)}\n\n".encode()
 
 
+def test_chat_text_timing_metrics_request_stage_metrics():
+    args = Namespace(
+        backend="openai-chat-omni",
+        percentile_metrics="ttft,tpot,itl,e2el",
+        print_stage=False,
+        extra_body={},
+    )
+
+    assert should_request_stage_metrics(args) is True
+
+
+@pytest.mark.asyncio
+async def test_bundled_first_text_chunk_uses_stage0_token_timings(mocker: MockerFixture):
+    """Engine timings recover TPOT when every text token arrives together."""
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=10,
+        output_len=20,
+    )
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "ABCD"}}],
+                "modality": "text",
+                "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14},
+                "metrics": {
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 4,
+                            "vllm_itls_ms": [10.0, 11.0, 12.0],
+                            "vllm_tpot_ms": 11.0,
+                        }
+                    }
+                },
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+    mock_response = MockResponse(200, chunks)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    assert output.success is True
+    assert output.output_tokens == 4
+    assert output.itl == pytest.approx([0.010, 0.011, 0.012])
+    assert output.text_latency - output.ttft == pytest.approx(0.033)
+    assert output.tpot_measured is True
+
+
+@pytest.mark.asyncio
+async def test_positive_client_text_timings_take_precedence_over_stage0(mocker: MockerFixture):
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=10,
+        output_len=20,
+    )
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "A"}}],
+                "modality": "text",
+                "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "B"}}],
+                "modality": "text",
+                "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
+                "metrics": {
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 2,
+                            "vllm_itls_ms": [1000.0],
+                            "vllm_tpot_ms": 1000.0,
+                        }
+                    }
+                },
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+    mock_response = MockResponse(200, chunks, delay_between_chunks=0.01)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    assert output.success is True
+    assert len(output.itl) == 1
+    assert 0.0 < output.itl[0] < 0.1
+
+
 # ============================================================================
 # output_tokens Tests
 # ============================================================================

vllm_omni/benchmarks/patch/patch.py

@@ -119,6 +119,17 @@ def should_request_stage_metrics(args: Any) -> bool:
     if backend in _IMAGE_STAGE_METRICS_BACKENDS:
         return True
 
+    percentile_metrics = getattr(args, "percentile_metrics", ()) or ()
+    if isinstance(percentile_metrics, str):
+        percentile_metrics = percentile_metrics.split(",")
+    selected_metrics = {str(metric).strip().lower() for metric in percentile_metrics}
+    if backend == "openai-chat-omni" and selected_metrics.intersection({"tpot", "itl"}):
+        # Client receive timestamps normally provide these metrics. Under
+        # enough load, however, every text token can be coalesced into one
+        # event-loop read, leaving no positive interval. Request the engine's
+        # Stage 0 timings so that path still has an authoritative measurement.
+        return True
+
     extra_body = getattr(args, "extra_body", None) or {}
     modalities = extra_body.get("modalities") if isinstance(extra_body, dict) else None
     return backend == "openai-chat-omni" and "image" in (modalities or [])
@@ -854,6 +865,33 @@ def _update_output_stage_metrics_from_payload(
         output.stage_metrics.update(stage_snapshot)
 
 
+def _apply_chat_stage0_token_timings(output: MixRequestFuncOutput) -> bool:
+    """Apply native Stage 0 timings from a chat response snapshot."""
+    stage_metrics = output.stage_metrics
+    if not isinstance(stage_metrics, dict):
+        return False
+    stage0 = stage_metrics.get("0")
+    if not isinstance(stage0, dict):
+        return False
+
+    output_tokens = coerce_positive_int_scalar(stage0.get(defs.NUM_TOKENS_OUT))
+    expected_output_tokens = coerce_positive_int_scalar(output.output_tokens)
+    if output_tokens is None or expected_output_tokens is None or output_tokens != expected_output_tokens:
+        return False
+
+    return _apply_stage0_token_timings(
+        output,
+        [
+            {
+                "output_token_count": output_tokens,
+                "itls_ms": stage0.get(defs.VLLM_ITLS_MS),
+                "tpot_ms": stage0.get(defs.VLLM_TPOT_MS),
+            }
+        ],
+        expected_output_tokens=expected_output_tokens,
+    )
+
+
 def _image_metrics_from_stage_metrics(metrics: dict[str, Any] | None) -> tuple[int, float, int, float]:
     if not isinstance(metrics, dict):
         return 0, 0.0, 0, 0.0
@@ -1121,6 +1159,14 @@ async def async_request_openai_chat_omni_completions(
 
                     output.latency = timestamp - st
                     output.generated_text = generated_text
+                    if output.output_tokens > 1 and not any(
+                        isinstance(value, int | float)
+                        and not isinstance(value, bool)
+                        and np.isfinite(value)
+                        and value > 0
+                        for value in output.itl
+                    ):
+                        _apply_chat_stage0_token_timings(output)
                     if output.itl:
                         # Align text_latency with ITL so TPOT formula and
                         # mean(ITL) are consistent.  Do NOT infer output_tokens