Add speech token usage headers (#4499)

Signed-off-by: JackyLiu <lzwgre@126.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/serving/speech_api.md

@@ -141,7 +141,20 @@ The response shape depends on the streaming parameters:
 
 **Non-streaming (default).** With `stream=false` and no `stream_format`, returns the
 complete clip as binary audio data with an appropriate `Content-Type` header (e.g.
-`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported.
+`audio/wav`). Because the raw-bytes body has no JSON carrier, successful
+non-streaming responses from non-diffusion speech servers report usage through
+response headers:
+
+| Header | Description |
+| --- | --- |
+| `x-vllm-omni-input-tokens` | Total input tokens (`text_tokens` + `audio_tokens`). |
+| `x-vllm-omni-output-tokens` | Generated codec/audio tokens. |
+| `x-vllm-omni-total-tokens` | `input_tokens` + `output_tokens`. |
+| `x-vllm-omni-input-text-tokens` | Tokens from the synthesized text (`input` plus `instructions`). |
+| `x-vllm-omni-input-audio-tokens` | Reference-audio codec frames, non-zero only for in-context voice cloning. |
+
+Diffusion-mode speech servers route through a separate response path and do not
+emit these headers.
 
 **Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or
 WAV) as they are decoded.

tests/entrypoints/openai_api/test_serving_speech.py

@@ -33,6 +33,8 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
     OpenAICreateAudioGenerateRequest,
     OpenAICreateSpeechRequest,
     SpeechBatchItem,
+    SpeechInputTokenDetails,
+    SpeechTokenUsage,
     StreamingSpeechSessionConfig,
 )
 from vllm_omni.entrypoints.openai.serving_speech import (
@@ -155,6 +157,7 @@ class TestAudioMixin:
 # Helper to create mock model output for endpoint tests
 def create_mock_audio_output_for_test(
     request_id: str = "speech-mock-123",
+    metrics: dict | None = None,
 ) -> OmniRequestOutput:
     class MockCompletionOutput:
         def __init__(self, index: int = 0):
@@ -181,11 +184,16 @@ def create_mock_audio_output_for_test(
     audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples))
     mock_request_output = MockRequestOutput(request_id=request_id, audio_tensor=audio_tensor)
 
-    return OmniRequestOutput.from_stage_output(
+    if metrics is None:
+        metrics = {"stage_metrics": {"0": {"num_tokens_in": 0, "num_tokens_out": 5}}}
+
+    output = OmniRequestOutput.from_stage_output(
         mock_request_output,
         stage_id=0,
         final_output_type="audio",
     )
+    output.metrics = metrics
+    return output
 
 
 def _write_custom_voice_manifest(root: Path, *, model_type: str, voices: dict) -> None:
@@ -233,6 +241,7 @@ def test_app(mocker: MockerFixture):
         models=mock_models,
         request_logger=mock_request_logger,
     )
+    speech_server._tts_tokenizer = lambda text, padding=False: {"input_ids": text.split()}
 
     # Skip TTS validation in tests (mock doesn't set up supported_speakers)
     speech_server._validate_tts_request = mocker.MagicMock(return_value=None)
@@ -370,6 +379,37 @@ class TestSpeechAPI:
         assert response.headers["content-type"] == "audio/wav"
         assert len(response.content) > 0
 
+    def test_create_speech_includes_token_usage_and_detail_headers(self, client):
+        payload = {
+            "input": "Hello world",
+            "model": "tts-model",
+            "voice": "alloy",
+            "response_format": "wav",
+        }
+        response = client.post("/v1/audio/speech", json=payload)
+        assert response.status_code == 200
+        assert response.headers["x-vllm-omni-input-tokens"] == "2"
+        assert response.headers["x-vllm-omni-output-tokens"] == "5"
+        assert response.headers["x-vllm-omni-total-tokens"] == "7"
+        assert response.headers["x-vllm-omni-input-text-tokens"] == "2"
+        assert response.headers["x-vllm-omni-input-audio-tokens"] == "0"
+
+    def test_build_speech_usage_headers_uses_usage_and_detail_field_names(self):
+        usage = SpeechTokenUsage(
+            input_tokens=7,
+            output_tokens=11,
+            total_tokens=18,
+            input_token_details=SpeechInputTokenDetails(text_tokens=3, audio_tokens=4),
+        )
+        headers = OmniOpenAIServingSpeech._build_speech_usage_headers(usage)
+        assert headers == {
+            "X-VLLM-OMNI-INPUT-TOKENS": "7",
+            "X-VLLM-OMNI-OUTPUT-TOKENS": "11",
+            "X-VLLM-OMNI-TOTAL-TOKENS": "18",
+            "X-VLLM-OMNI-INPUT-TEXT-TOKENS": "3",
+            "X-VLLM-OMNI-INPUT-AUDIO-TOKENS": "4",
+        }
+
     def test_create_speech_mp3_format(self, client):
         payload = {
             "input": "Hello world",

vllm_omni/entrypoints/openai/serving_speech.py

@@ -80,6 +80,12 @@ from vllm_omni.utils.speaker_cache import (
 
 logger = init_logger(__name__)
 
+_SPEECH_USAGE_INPUT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TOKENS"
+_SPEECH_USAGE_OUTPUT_TOKENS_HEADER = "X-VLLM-OMNI-OUTPUT-TOKENS"
+_SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS"
+_SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TEXT-TOKENS"
+_SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS"
+
 # TTS Configuration
 #
 # The stage-key -> model-type mapping is NOT declared here: it is derived from
@@ -950,6 +956,24 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         details = self._compute_speech_input_details(request, tts_params)
         return build_speech_usage(details, output_tokens)
 
+    @staticmethod
+    def _build_speech_usage_headers(usage: SpeechTokenUsage | None) -> dict[str, str]:
+        """Map speech usage into non-streaming response headers.
+
+        Returns an empty dict when usage is unavailable, allowing callers to
+        merge these headers with other optional response headers.
+        """
+        if usage is None:
+            return {}
+
+        return {
+            _SPEECH_USAGE_INPUT_TOKENS_HEADER: str(usage.input_tokens),
+            _SPEECH_USAGE_OUTPUT_TOKENS_HEADER: str(usage.output_tokens),
+            _SPEECH_USAGE_TOTAL_TOKENS_HEADER: str(usage.total_tokens),
+            _SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER: str(usage.input_token_details.text_tokens),
+            _SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER: str(usage.input_token_details.audio_tokens),
+        }
+
     def _estimate_fish_ref_code_len(self, ref_audio: object) -> int | None:
         """Estimate Fish Speech semantic token length from raw reference audio."""
         from vllm_omni.model_executor.models.fish_speech.dac_utils import (
@@ -3170,7 +3194,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         # ``usage_out`` is an opt-in output channel: when a list is passed, the
         # computed SpeechTokenUsage is appended to it. The return stays a
         # 2-tuple so existing callers (and their test mocks) are unaffected;
-        # only the batch path, which surfaces per-item usage, opts in.
+        # batch and non-streaming response-header paths opt in when surfacing
+        # usage outside the raw audio body.
         request_id, generator, bytes_tts_params = await self._prepare_speech_generation(
             request, request_id=request_id, has_inline_ref_audio=has_inline_ref_audio
         )
@@ -3576,7 +3601,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 )
 
             collect: dict = {}
-            audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id, collect=collect)
+            usage_box: list[SpeechTokenUsage] = []
+            audio_bytes, media_type = await self._generate_audio_bytes(
+                request, request_id=request_id, usage_out=usage_box, collect=collect
+            )
             total_ms = (time.perf_counter() - request_start_s) * 1000.0
             logger.info(
                 "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d",
@@ -3584,7 +3612,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 total_ms,
                 len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else len(str(audio_bytes)),
             )
-            headers = {}
+            headers = self._build_speech_usage_headers(usage_box[0] if usage_box else None)
             if collect.get("word_timestamps") is not None:
                 # Default ensure_ascii keeps the header latin-1 encodable (non-ASCII words \uXXXX-escaped).
                 ts_json = json.dumps(collect["word_timestamps"])