From 71748388f820a9b71d7b73954b99f3c83593c0a8 Mon Sep 17 00:00:00 2001 From: Gao Han Date: Mon, 31 Aug 2026 00:43:25 +0800 Subject: [PATCH] [Bugfix] Recover Omni benchmark TPOT from Stage 0 metrics (#6818) Signed-off-by: Gao Han --- tests/benchmarks/patch/test_patch.py | 103 +++++++++++++++++++++++++++ vllm_omni/benchmarks/patch/patch.py | 46 ++++++++++++ 2 files changed, 149 insertions(+) diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index e40e8f954..ad33a03ab 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/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 # ============================================================================ diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 217467a2c..fd52c197b 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/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