[Bugfix] Fix benchmark Total input tokens for multimodal requests (#2540) (#2549)

Signed-off-by: Dnoob <dxpouo@gmail.com>

tests/benchmarks/metrics/test_metrics.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for metrics.py
+"""
+
+import pytest
+from vllm.benchmarks.serve import TaskType
+
+from vllm_omni.benchmarks.metrics.metrics import calculate_metrics
+from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
+
+
+def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutput:
+    """Build a minimal successful MixRequestFuncOutput for metrics aggregation."""
+    output = MixRequestFuncOutput()
+    output.success = True
+    output.prompt_len = prompt_len
+    output.output_tokens = output_tokens
+    output.generated_text = "x" * output_tokens
+    output.ttft = 0.1
+    output.text_latency = 1.0
+    output.latency = 1.0
+    output.start_time = 0.0
+    output.itl = [0.1] * max(output_tokens - 1, 0)
+    output.audio_ttfp = 0.0
+    output.audio_rtf = 0.0
+    output.audio_duration = 0.0
+    output.audio_frames = 0
+    output.input_audio_duration = 0.0
+    output.error = ""
+    return output
+
+
+# ============================================================================
+# total_input Tests
+# ============================================================================
+
+
+def test_total_input_aggregated_from_output_prompt_len():
+    """Test that total_input sums outputs[i].prompt_len, not input_requests[i].prompt_len."""
+    outputs = [_make_output(4992), _make_output(3000)]
+
+    metrics, _ = calculate_metrics(
+        input_requests=[],
+        outputs=outputs,
+        dur_s=10.0,
+        tokenizer=None,
+        selected_percentiles=[99.0],
+        goodput_config_dict={},
+        task_type=TaskType.GENERATION,
+        selected_percentile_metrics=[],
+        max_concurrency=None,
+        request_rate=float("inf"),
+        benchmark_duration=10.0,
+    )
+
+    assert metrics.total_input == 7992, (
+        "total_input should aggregate from outputs[i].prompt_len to reflect the true multimodal input token count"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

tests/benchmarks/patch/test_patch.py

@@ -574,5 +574,59 @@ class TestTextLatencyAttribute:
         )
 
 
+# ============================================================================
+# prompt_len Tests
+# ============================================================================
+
+
+@pytest.mark.asyncio
+async def test_prompt_len_assigned_from_usage(mocker: MockerFixture):
+    # Arrange: request claims prompt_len=100, but server reports 4992 (multimodal).
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=100,
+        output_len=20,
+    )
+
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "Hello"}}],
+                "modality": "text",
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": " world"}}],
+                "modality": "text",
+            }
+        ),
+        # Final usage chunk emitted because stream_options.include_usage=True.
+        create_sse_chunk(
+            {
+                "choices": [],
+                "usage": {"prompt_tokens": 4992, "completion_tokens": 2, "total_tokens": 4994},
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+
+    mock_response = MockResponse(200, chunks)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+    # Act
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    # Assert
+    assert output.success is True
+    assert output.prompt_len == 4992, (
+        "prompt_len should be overridden by usage.prompt_tokens to reflect the true multimodal input token count"
+    )
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

vllm_omni/benchmarks/metrics/metrics.py

@@ -185,7 +185,7 @@ def calculate_metrics(
                     # Note : this may inflate the output token count slightly
                     output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
             actual_output_lens.append(output_len)
-            total_input += input_requests[i].prompt_len
+            total_input += outputs[i].prompt_len
             tpot = 0
             if output_len > 1:
                 latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft

vllm_omni/benchmarks/patch/patch.py

@@ -190,6 +190,10 @@ async def async_request_openai_chat_omni_completions(
                                 if metrics := data.get("metrics"):
                                     output.output_tokens = metrics.get("num_tokens_out", 0)
 
+                                if usage := data.get("usage"):
+                                    if (pt := usage.get("prompt_tokens")) is not None:
+                                        output.prompt_len = pt
+
                     output.latency = timestamp - st
                     output.generated_text = generated_text
                     if generated_audio is not None: