From d7216ef5144866c2f039947e1cc396b1984fcf09 Mon Sep 17 00:00:00 2001 From: JackyLiu Date: Tue, 18 Aug 2026 06:53:06 +0800 Subject: [PATCH] Add speech token usage headers (#4499) Signed-off-by: JackyLiu Co-authored-by: Hongsheng Liu --- docs/serving/speech_api.md | 15 ++++++- .../openai_api/test_serving_speech.py | 42 ++++++++++++++++++- .../entrypoints/openai/serving_speech.py | 34 +++++++++++++-- 3 files changed, 86 insertions(+), 5 deletions(-) diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 0deea7f8c..11c943ca4 100644 --- a/docs/serving/speech_api.md +++ b/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. diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 934c0ac64..604eb54f3 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/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", diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 1b44d397d..a3069c38c 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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"])