[Bugfix] Fix /v1/audio/speech usage token accounting for Qwen3-TTS (#4646) (#4673)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Sy03 <1370724210@qq.com>

docs/serving/speech_api.md

@@ -122,7 +122,8 @@ Content-Type: application/json
 | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
 | `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
 | `non_streaming_mode` | bool | null | Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. |
-| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
+| `stream_format` | string | null | Streaming output format. `"audio"` streams raw audio bytes as they are decoded; `"sse"` streams OpenAI `speech.audio.*` Server-Sent Events. Omit (null) for a single non-streaming binary response. See [Response Format](#response-format). |
+| `stream` | bool | false | Legacy streaming switch, equivalent to `stream_format="audio"`. Requires `response_format="pcm"` or `"wav"`. Speed adjustment is not supported when streaming. |
 
 **Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian.
 
@@ -136,7 +137,60 @@ Content-Type: application/json
 
 ### Response Format
 
-Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
+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.
+
+**Raw audio stream** (`stream=true` or `stream_format="audio"`). Streams raw
+audio bytes (PCM or WAV) as they are decoded.
+
+Both streaming modes (`stream_format="audio"` and `"sse"`) require
+`response_format="pcm"` or `"wav"`, and `speed` must be `1.0` (or omitted).
+
+**SSE stream** (`stream_format="sse"`). Streams [OpenAI `speech.audio.*` Server-Sent
+Events](https://platform.openai.com/docs/api-reference/audio-streaming). Set
+`stream_format="sse"` and leave the legacy `stream` unset: `stream=true` takes
+precedence and forces the raw-audio path, so the response would be audio bytes with
+no `speech.audio.*` events (and no terminal `usage`). Each event has an `event:`
+line and a JSON `data:` line:
+
+- `speech.audio.delta` — a base64 audio chunk:
+
+    ```json
+    { "type": "speech.audio.delta", "audio": "<base64>", "response_format": "pcm" }
+    ```
+
+- `speech.audio.done` — terminal event, carrying token `usage`:
+
+    ```json
+    {
+        "type": "speech.audio.done",
+        "usage": {
+            "input_tokens": 119,
+            "output_tokens": 77,
+            "total_tokens": 196,
+            "input_token_details": { "text_tokens": 18, "audio_tokens": 101 }
+        }
+    }
+    ```
+
+- `speech.audio.error` — emitted instead of `speech.audio.done` if generation fails:
+
+    ```json
+    { "type": "speech.audio.error", "error": { "message": "...", "type": "server_error", "param": null, "code": 500 } }
+    ```
+
+The `usage` object on `speech.audio.done` is the same shape returned per item by the
+[batch endpoint](#batch-speech-generation):
+
+- `input_tokens` = `text_tokens` + `audio_tokens`
+    - `text_tokens`: tokens of the synthesized text (`input` plus `instructions`)
+    - `audio_tokens`: reference-audio codec frames, non-zero only for in-context
+      voice cloning (Base task); `0` for CustomVoice/VoiceDesign or x-vector-only
+- `output_tokens`: generated codec tokens
+- `total_tokens` = `input_tokens` + `output_tokens`
 
 ### Voices Endpoint
 
@@ -490,7 +544,13 @@ Each item in the `items` array requires only `input` (the text). All other field
             "index": 0,
             "status": "success",
             "audio_data": "<base64-encoded audio>",
-            "media_type": "audio/wav"
+            "media_type": "audio/wav",
+            "usage": {
+                "input_tokens": 119,
+                "output_tokens": 77,
+                "total_tokens": 196,
+                "input_token_details": { "text_tokens": 18, "audio_tokens": 101 }
+            }
         },
         {
             "index": 1,
@@ -504,6 +564,18 @@ Each item in the `items` array requires only `input` (the text). All other field
 }
 ```
 
+Each successful item carries a `usage` object (errored items omit it):
+
+- `input_tokens` = `text_tokens` + `audio_tokens`
+    - `text_tokens`: tokens of the synthesized text (`input` plus `instructions`)
+    - `audio_tokens`: reference-audio codec frames, non-zero only for in-context
+      voice cloning (Base task); `0` for CustomVoice/VoiceDesign or x-vector-only
+- `output_tokens`: generated codec tokens
+- `total_tokens` = `input_tokens` + `output_tokens`
+
+This is the same `usage` object emitted on the terminal `speech.audio.done` event
+of the single endpoint's [SSE stream](#response-format) (`stream_format="sse"`).
+
 ### Examples
 
 **Basic batch with shared defaults:**

tests/entrypoints/openai_api/test_serving_speech.py

@@ -2535,6 +2535,61 @@ def test_api_server_create_speech_batch_without_handler_returns_404(mocker: Mock
     )
 
 
+def test_api_server_create_speech_batch_omits_null_fields(mocker: MockerFixture):
+    # The batch response must omit optional null fields rather than serialize them
+    # as null (issue #4646 follow-up): errored items drop usage/audio_data/media_type,
+    # successful items drop error. This is the shape documented in speech_api.md.
+    from vllm_omni.entrypoints.openai.protocol.audio import (
+        BatchSpeechResponse,
+        SpeechBatchItemResult,
+        SpeechInputTokenDetails,
+        SpeechTokenUsage,
+    )
+
+    handler = mocker.MagicMock()
+    handler.create_speech_batch = mocker.AsyncMock(
+        return_value=BatchSpeechResponse(
+            id="speech-batch-test",
+            results=[
+                SpeechBatchItemResult(
+                    index=0,
+                    status="success",
+                    audio_data="YWJj",
+                    media_type="audio/wav",
+                    usage=SpeechTokenUsage(
+                        input_tokens=119,
+                        output_tokens=77,
+                        total_tokens=196,
+                        input_token_details=SpeechInputTokenDetails(text_tokens=18, audio_tokens=101),
+                    ),
+                ),
+                SpeechBatchItemResult(index=1, status="error", error="Input text cannot be empty"),
+            ],
+            total=2,
+            succeeded=1,
+            failed=1,
+        )
+    )
+    raw_request = _make_api_server_request(handler, path="/v1/audio/speech/batch")
+    request = BatchSpeechRequest(items=[SpeechBatchItem(input="hi"), SpeechBatchItem(input="")])
+
+    response = asyncio.run(api_server_module.create_speech_batch(request, raw_request))
+
+    assert isinstance(response, JSONResponse)
+    assert response.status_code == 200
+    body = json.loads(response.body)
+    success, errored = body["results"][0], body["results"][1]
+    # Successful item carries usage and drops the null `error`.
+    assert success["usage"]["total_tokens"] == 196
+    assert success["usage"]["input_token_details"] == {"text_tokens": 18, "audio_tokens": 101}
+    assert "error" not in success
+    # Errored item drops usage/audio_data/media_type instead of serializing null.
+    assert "usage" not in errored
+    assert "audio_data" not in errored
+    assert "media_type" not in errored
+    assert errored["error"] == "Input text cannot be empty"
+
+
 def test_api_server_create_audio_generate_without_handler_returns_404(mocker: MockerFixture):
     fake_base = _patch_api_server_base(mocker)
     raw_request = _make_api_server_request(None, path="/v1/audio/generate")

tests/entrypoints/openai_api/test_speech_usage.py

@@ -0,0 +1,220 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for speech token-usage accounting (issue #4646).
+
+These exercise the pure usage logic without a model/server: input-token
+breakdown (text + reference-audio gating) and output-token accumulation.
+"""
+
+from dataclasses import dataclass
+
+import pytest
+
+from vllm_omni.entrypoints.openai.speech_usage import (
+    SpeechOutputTokenCounter,
+    build_speech_usage,
+    gate_audio_tokens,
+    qwen3_tts_input_token_details,
+)
+
+
+# A deterministic stand-in for a real tokenizer: 1 token per whitespace word.
+def word_count(text: str) -> int:
+    return len(text.split())
+
+
+# --- input breakdown: gating ------------------------------------------------
+
+
+def test_gate_audio_tokens_base_in_context_counts_ref_frames():
+    assert (
+        gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=None, ref_code_length=120)
+        == 120
+    )
+
+
+def test_gate_audio_tokens_customvoice_is_zero():
+    assert (
+        gate_audio_tokens(
+            task_type="CustomVoice", x_vector_only_mode=False, icl_mode_override=None, ref_code_length=120
+        )
+        == 0
+    )
+
+
+def test_gate_audio_tokens_x_vector_only_is_zero_even_with_ref():
+    # x-vector cloning inserts NO codec frames, so audio_tokens must be 0
+    # even though ref_code_length is populated.
+    assert (
+        gate_audio_tokens(task_type="Base", x_vector_only_mode=True, icl_mode_override=None, ref_code_length=120) == 0
+    )
+
+
+def test_gate_audio_tokens_icl_override_wins():
+    assert (
+        gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=False, ref_code_length=120) == 0
+    )
+    assert (
+        gate_audio_tokens(task_type="Base", x_vector_only_mode=True, icl_mode_override=True, ref_code_length=80) == 80
+    )
+
+
+def test_gate_audio_tokens_icl_override_ignored_for_non_base():
+    # icl_mode_override is honored only for the Base task; CustomVoice stays 0
+    # even if such a flag leaks into its params (mirrors the estimator).
+    assert (
+        gate_audio_tokens(
+            task_type="CustomVoice", x_vector_only_mode=False, icl_mode_override=True, ref_code_length=120
+        )
+        == 0
+    )
+
+
+def test_gate_audio_tokens_bad_ref_length_is_zero():
+    assert (
+        gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=None, ref_code_length="oops")
+        == 0
+    )
+
+
+# --- input breakdown: Qwen3-TTS reproduces issue #4646 observation ----------
+
+
+def test_customvoice_text_scales_with_input_no_audio():
+    short = qwen3_tts_input_token_details(
+        input_text="hi there",
+        instructions=None,
+        tts_params={"task_type": ["CustomVoice"]},
+        count_text_tokens=word_count,
+    )
+    long = qwen3_tts_input_token_details(
+        input_text="hi there this is a much longer sentence",
+        instructions=None,
+        tts_params={"task_type": ["CustomVoice"]},
+        count_text_tokens=word_count,
+    )
+    assert short.audio_tokens == 0 and long.audio_tokens == 0
+    assert long.text_tokens > short.text_tokens  # scales with input
+
+
+def test_base_icl_text_tracks_input_audio_tracks_ref():
+    # The core #4646 fix: text_tokens reflect the input (not dropped), and
+    # audio_tokens track the reference audio independently.
+    params = {"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [100]}
+    short = qwen3_tts_input_token_details(
+        input_text="hi", instructions=None, tts_params=params, count_text_tokens=word_count
+    )
+    long = qwen3_tts_input_token_details(
+        input_text="hi there friend", instructions=None, tts_params=params, count_text_tokens=word_count
+    )
+    # text now varies with input (the bug was that it was dropped for Base)
+    assert long.text_tokens > short.text_tokens
+    # audio fixed by the reference, independent of input text
+    assert short.audio_tokens == long.audio_tokens == 100
+
+
+def test_instructions_count_toward_text():
+    without = qwen3_tts_input_token_details(
+        input_text="hello world",
+        instructions=None,
+        tts_params={"task_type": ["CustomVoice"]},
+        count_text_tokens=word_count,
+    )
+    with_instr = qwen3_tts_input_token_details(
+        input_text="hello world",
+        instructions="speak slowly and warmly",
+        tts_params={"task_type": ["CustomVoice"]},
+        count_text_tokens=word_count,
+    )
+    assert with_instr.text_tokens == without.text_tokens + word_count("speak slowly and warmly")
+
+
+# --- build usage ------------------------------------------------------------
+
+
+def test_build_speech_usage_aggregates_and_totals():
+    details = qwen3_tts_input_token_details(
+        input_text="hi there",
+        instructions=None,
+        tts_params={"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [100]},
+        count_text_tokens=word_count,
+    )
+    usage = build_speech_usage(details, output_tokens=250)
+    assert usage.input_tokens == details.text_tokens + 100
+    assert usage.output_tokens == 250
+    assert usage.total_tokens == usage.input_tokens + 250
+    assert usage.input_token_details.text_tokens == 2
+    assert usage.input_token_details.audio_tokens == 100
+
+
+def test_batch_item_result_serializes_usage():
+    # Guards the batch per-item usage carrier (and that it is optional/None on error).
+    from vllm_omni.entrypoints.openai.protocol.audio import SpeechBatchItemResult
+
+    details = qwen3_tts_input_token_details(
+        input_text="hello world",
+        instructions=None,
+        tts_params={"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [50]},
+        count_text_tokens=word_count,
+    )
+    item = SpeechBatchItemResult(
+        index=0,
+        status="success",
+        audio_data="x",
+        media_type="audio/wav",
+        usage=build_speech_usage(details, output_tokens=30),
+    )
+    dumped = item.model_dump()
+    assert dumped["usage"]["input_tokens"] == 2 + 50
+    assert dumped["usage"]["output_tokens"] == 30
+    assert dumped["usage"]["input_token_details"] == {"text_tokens": 2, "audio_tokens": 50}
+    # Errored items carry no usage.
+    assert SpeechBatchItemResult(index=1, status="error", error="boom").model_dump()["usage"] is None
+
+
+# --- output token accumulation ----------------------------------------------
+# The generated codec-token count is read from per-stage pipeline metrics on the
+# final output: res.metrics["stage_metrics"][<stage>]["num_tokens_out"]. Only the
+# AR stage (0) is non-zero; code2wav (stage 1) reports 0.
+
+
+@dataclass
+class _FakeRes:
+    metrics: dict
+
+
+def _final_res(stage0_out, stage1_out=0):
+    return _FakeRes(
+        metrics={
+            "stage_metrics": {
+                "0": {"num_tokens_in": 110, "num_tokens_out": stage0_out},
+                "1": {"num_tokens_in": 0, "num_tokens_out": stage1_out},
+            }
+        }
+    )
+
+
+def test_output_tokens_from_stage_metrics():
+    acc = SpeechOutputTokenCounter()
+    # Intermediate outputs carry empty metrics -> contribute nothing.
+    acc.observe(_FakeRes(metrics={}))
+    acc.observe(_FakeRes(metrics={}))
+    # Final output carries stage metrics; stage-0 num_tokens_out is the count.
+    acc.observe(_final_res(stage0_out=77))
+    assert acc.total() == 77
+
+
+def test_output_tokens_takes_max_across_stages():
+    acc = SpeechOutputTokenCounter()
+    acc.observe(_final_res(stage0_out=42, stage1_out=0))
+    assert acc.total() == 42
+
+
+def test_output_tokens_zero_without_metrics():
+    acc = SpeechOutputTokenCounter()
+    acc.observe(_FakeRes(metrics={}))
+    acc.observe(_FakeRes(metrics={"stage_metrics": {}}))
+    assert acc.total() == 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

vllm_omni/entrypoints/openai/api_server.py

@@ -1317,7 +1317,10 @@ async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request)
         result = await handler.create_speech_batch(request)
         if isinstance(result, ErrorResponse):
             return _error_response_to_json_response(result)
-        return JSONResponse(content=result.model_dump())
+        # exclude_none so optional per-item fields are omitted rather than
+        # serialized as null: errored items drop `usage`/`audio_data`/`media_type`,
+        # successful items drop `error`. Matches the documented batch response shape.
+        return JSONResponse(content=result.model_dump(exclude_none=True))
     except (EngineGenerateError, EngineDeadError) as exc:
         return _create_engine_error_json_response(raw_request, exc)
     except ValueError as e:

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -421,12 +421,65 @@ class BatchSpeechRequest(BaseModel):
     non_streaming_mode: bool | None = None
 
 
+class SpeechInputTokenDetails(BaseModel):
+    """Per-modality breakdown of the speech request's *input* tokens.
+
+    The aggregate ``input_tokens`` on :class:`SpeechTokenUsage` is the sum of
+    these. We surface the split (rather than one opaque number) for the same
+    reason OpenAI's realtime/chat usage does: text and audio inputs are billed
+    and reasoned about differently, and folding them together is misleading.
+
+    Fields:
+        text_tokens: Tokens of the text to synthesize. This is ``input`` plus
+            ``instructions`` (style/emotion prompt), because both are tokenized
+            into the model prefill. This is the number that should scale with
+            how much text the caller asked to speak.
+        audio_tokens: Reference-audio codec frames used as voice-cloning
+            conditioning. NON-ZERO only when in-context voice cloning is
+            actually active (Qwen3-TTS ``task_type='Base'`` ICL). It is 0 for
+            CustomVoice/VoiceDesign and for x-vector-only cloning, because those
+            paths put no reference codec frames into the prefill. See issue
+            #4646: this is the value that previously leaked into ``prompt_tokens``
+            and made Base usage look independent of the input text.
+    """
+
+    text_tokens: int = 0
+    audio_tokens: int = 0
+
+
+class SpeechTokenUsage(BaseModel):
+    """Token usage for a speech (TTS) request.
+
+    Field naming follows OpenAI's documented ``speech.audio.done`` event
+    (``input_tokens``/``output_tokens``/``total_tokens``), NOT chat's
+    ``prompt_tokens``/``completion_tokens``.
+
+    input_tokens  = text_tokens + audio_tokens   (see SpeechInputTokenDetails)
+    output_tokens = generated codec/audio tokens (stage-0 decode steps)
+    total_tokens  = input_tokens + output_tokens
+
+    IMPORTANT: ``input_tokens`` is computed from the *semantic* request inputs
+    (tokenized text + reference-audio frames), NOT from ``len(prompt_token_ids)``.
+    For staged TTS models that engine prompt is a ``[1] * prefill_len``
+    placeholder whose length mirrors the model prefill, so it is not a faithful
+    input-token count (issue #4646).
+    """
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_tokens: int = 0
+    input_token_details: SpeechInputTokenDetails = Field(default_factory=SpeechInputTokenDetails)
+
+
 class SpeechBatchItemResult(BaseModel):
     index: int
     status: Literal["success", "error"]
     audio_data: str | None = None
     media_type: str | None = None
     error: str | None = None
+    # Per-item token usage (input text + reference-audio conditioning, and
+    # generated audio tokens). None when the item errored before generation.
+    usage: SpeechTokenUsage | None = None
 
 
 class BatchSpeechResponse(BaseModel):

vllm_omni/entrypoints/openai/serving_speech.py

@@ -43,6 +43,13 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
     OpenAICreateSpeechRequest,
     SpeechBatchItem,
     SpeechBatchItemResult,
+    SpeechInputTokenDetails,
+    SpeechTokenUsage,
+)
+from vllm_omni.entrypoints.openai.speech_usage import (
+    SpeechOutputTokenCounter,
+    build_speech_usage,
+    qwen3_tts_input_token_details,
 )
 from vllm_omni.entrypoints.openai.tts_adapters import (
     SpeechServingContext,
@@ -895,6 +902,69 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e)
             return 2048
 
+    def _get_usage_text_tokenizer(self):
+        """Return a text tokenizer for counting input-text usage tokens.
+
+        Prefer the per-model tokenizer already loaded for prompt-length
+        estimation (`_tts_tokenizer`, which is the *correct* tokenizer for the
+        active model). Fall back to a lazily-loaded, cached generic tokenizer
+        for models that never populate `_tts_tokenizer`. Returns None if no
+        tokenizer can be obtained (usage then reports text_tokens=0).
+        """
+        if self._tts_tokenizer is not None:
+            return self._tts_tokenizer
+        if getattr(self, "_usage_text_tokenizer", None) is None:
+            try:
+                from transformers import AutoTokenizer
+
+                self._usage_text_tokenizer = AutoTokenizer.from_pretrained(
+                    self.engine_client.model_config.model, trust_remote_code=True
+                )
+            except Exception as e:  # pragma: no cover - environment dependent
+                logger.warning("Usage: could not load a text tokenizer (%s); text_tokens will be 0", e)
+                self._usage_text_tokenizer = None
+        return self._usage_text_tokenizer
+
+    def _count_usage_text_tokens(self, text: str) -> int:
+        """Token count of `text` using the model's text tokenizer (0 on failure)."""
+        if not text:
+            return 0
+        tok = self._get_usage_text_tokenizer()
+        if tok is None:
+            return 0
+        try:
+            return len(tok(text, padding=False)["input_ids"])
+        except Exception:
+            return 0
+
+    def _compute_speech_input_details(
+        self, request: OpenAICreateSpeechRequest, tts_params: dict[str, Any]
+    ) -> SpeechInputTokenDetails:
+        """Input-token breakdown (text + reference-audio) for a speech request.
+
+        Counts `input` (+ `instructions`) as text tokens, and reference-audio
+        codec frames as audio tokens *only* when in-context voice cloning is
+        active (see `qwen3_tts_input_token_details` / `gate_audio_tokens`). The
+        audio gating reads Qwen3-TTS `tts_params` conventions; other TTS models
+        do not set those keys, so they degrade cleanly to text-only counts.
+        """
+        return qwen3_tts_input_token_details(
+            input_text=request.input,
+            instructions=request.instructions,
+            tts_params=tts_params or {},
+            count_text_tokens=self._count_usage_text_tokens,
+        )
+
+    def _build_speech_usage(
+        self,
+        request: OpenAICreateSpeechRequest,
+        tts_params: dict[str, Any],
+        output_tokens: int,
+    ) -> SpeechTokenUsage:
+        """Assemble the full usage object (input breakdown + generated tokens)."""
+        details = self._compute_speech_input_details(request, tts_params)
+        return build_speech_usage(details, output_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 (
@@ -2635,6 +2705,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         raw_request: Request | None = None,
         request_start_s: float | None = None,
         include_sample_rate: bool = False,
+        usage_acc: SpeechOutputTokenCounter | None = None,
     ):
         """Generate audio chunks for streaming response.
 
@@ -2661,6 +2732,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         try:
             async for res in generator:
+                # Tally generated codec tokens for usage (reads per-stage metrics
+                # off the final output; a cheap early-return on every other res).
+                if usage_acc is not None:
+                    usage_acc.observe(res)
                 audio_output, audio_key = self._extract_audio_output(res)
                 if audio_key is None:
                     continue
@@ -2786,6 +2861,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         response_format: str = "pcm",
         raw_request: Request | None = None,
         request_start_s: float | None = None,
+        request: OpenAICreateSpeechRequest | None = None,
+        tts_params: dict[str, Any] | None = None,
     ):
         """Generate OpenAI-style SSE events with base64 audio deltas.
 
@@ -2793,7 +2870,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         carries the base64 chunk in ``audio`` (not ``delta`` — that is the
         Realtime API ``response.audio.delta`` convention, a different event).
         See https://platform.openai.com/docs/api-reference/audio-streaming.
+
+        The terminal ``speech.audio.done`` event carries a ``usage`` object
+        (``input_tokens``/``output_tokens``/``total_tokens`` + a per-modality
+        ``input_token_details`` breakdown), matching OpenAI's documented
+        ``speech.audio.done`` schema. ``output_tokens`` is accumulated from the
+        stage-0 deltas as they stream (see ``SpeechOutputTokenCounter``);
+        ``input_tokens`` is computed from the request text + reference audio.
         """
+        usage_acc = SpeechOutputTokenCounter()
         try:
             async for chunk in self._generate_audio_chunks(
                 generator,
@@ -2801,6 +2886,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 response_format,
                 raw_request=raw_request,
                 request_start_s=request_start_s,
+                usage_acc=usage_acc,
             ):
                 payload = {
                     "type": "speech.audio.delta",
@@ -2809,7 +2895,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 }
                 data = json.dumps(payload, separators=(",", ":"))
                 yield f"event: speech.audio.delta\ndata: {data}\n\n"
-            done = json.dumps({"type": "speech.audio.done"}, separators=(",", ":"))
+            done_payload: dict[str, Any] = {"type": "speech.audio.done"}
+            if request is not None:
+                # Streaming path: output_tokens = sum of stage-0 deltas.
+                usage = self._build_speech_usage(request, tts_params or {}, usage_acc.total())
+                done_payload["usage"] = usage.model_dump()
+            done = json.dumps(done_payload, separators=(",", ":"))
             yield f"event: speech.audio.done\ndata: {done}\n\n"
         except asyncio.CancelledError:
             raise
@@ -3650,8 +3741,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request: OpenAICreateSpeechRequest,
         base64_encode: bool = False,
         request_id: str | None = None,
+        usage_out: list[SpeechTokenUsage] | None = None,
     ) -> tuple[bytes | str, str]:
-        request_id, generator, _ = await self._prepare_speech_generation(request, request_id=request_id)
+        # ``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.
+        request_id, generator, bytes_tts_params = await self._prepare_speech_generation(request, request_id=request_id)
         artifact_ready = False
 
         try:
@@ -3664,8 +3760,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             moss_sample_rate: int | None = None
 
             final_output: OmniRequestOutput | None = None
+            # Non-streaming is FINAL_ONLY, so the stage-0 output carries the full
+            # token sequence; the counter records its length for output_tokens.
+            usage_acc = SpeechOutputTokenCounter()
             async for res in generator:
                 final_output = res
+                usage_acc.observe(res)
                 if not is_moss:
                     continue
                 try:
@@ -3749,6 +3849,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             audio_response: AudioResponse = self.create_audio(audio_obj)
             self._mark_ref_audio_artifact_ready_for_request(request_id)
             artifact_ready = True
+            if usage_out is not None:
+                usage_out.append(self._build_speech_usage(request, bytes_tts_params or {}, usage_acc.total()))
             return audio_response.audio_data, audio_response.media_type
         finally:
             if not artifact_ready:
@@ -3983,7 +4085,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 if error is not None:
                     return error
 
-                _, generator, _ = await self._prepare_speech_generation(request, request_id=request_id)
+                _, generator, sse_tts_params = await self._prepare_speech_generation(request, request_id=request_id)
                 return StreamingResponse(
                     self._generate_audio_sse_events(
                         generator,
@@ -3991,6 +4093,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         response_format,
                         raw_request=raw_request,
                         request_start_s=request_start_s,
+                        request=request,
+                        tts_params=sse_tts_params,
                     ),
                     media_type="text/event-stream",
                 )
@@ -4108,8 +4212,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             validation_error = self._validate_tts_request(req)
             if validation_error is not None:
                 return SpeechBatchItemResult(index=idx, status="error", error=validation_error)
+            usage_box: list[SpeechTokenUsage] = []
             try:
-                audio_data, media_type = await self._generate_audio_bytes(req, base64_encode=True)
+                audio_data, media_type = await self._generate_audio_bytes(req, base64_encode=True, usage_out=usage_box)
             except Exception as e:
                 logger.exception("Batch item %d failed: %s", idx, e)
                 return SpeechBatchItemResult(index=idx, status="error", error=str(e))
@@ -4118,6 +4223,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 status="success",
                 audio_data=audio_data,
                 media_type=media_type,
+                usage=usage_box[0] if usage_box else None,
             )
 
         results = await asyncio.gather(

vllm_omni/entrypoints/openai/speech_usage.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Token-usage accounting for the Speech (``/v1/audio/speech``) API.
+
+Why this module exists (issue #4646)
+------------------------------------
+For staged TTS models the engine prompt is a PLACEHOLDER: the serving layer
+builds ``prompt_token_ids = [1] * prefill_len`` and lets the model rebuild the
+real conditioning (text / ref_audio / ref_text) into ``inputs_embeds`` later.
+So ``len(prompt_token_ids)`` is NOT a faithful count of what the caller sent.
+
+For Qwen3-TTS specifically the placeholder length mirrors the model prefill:
+  * CustomVoice/VoiceDesign: the full input text is embedded in the prefill, so
+    the placeholder length scales with the input text.
+  * Base in-context voice cloning: the prefill embeds only ``codec_bos`` + the
+    reference-audio codec frames; the input text is consumed incrementally
+    during DECODE, not prefill. So the placeholder length tracks the *reference
+    audio* and is independent of the input text.
+
+That is why ``usage.prompt_tokens`` (== ``len(prompt_token_ids)``) looked wrong
+for Base: it counted reference-audio frames instead of the synthesized text.
+
+This module computes usage from the *semantic* inputs instead:
+
+    input_tokens  = text_tokens + audio_tokens
+        text_tokens  -> tokens of ``input`` (+ ``instructions``); the text to speak
+        audio_tokens -> reference-audio codec frames used as voice-clone
+                        conditioning, counted ONLY when in-context cloning is
+                        actually active (see ``gate_audio_tokens``)
+    output_tokens = generated codec/audio tokens (stage-0 decode steps)
+    total_tokens  = input_tokens + output_tokens
+
+Naming (``input_tokens``/``output_tokens``) follows OpenAI's ``speech.audio.done``
+event; the ``input_token_details`` breakdown follows OpenAI's realtime/chat
+convention of never folding audio into an opaque text count.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    SpeechInputTokenDetails,
+    SpeechTokenUsage,
+)
+
+
+def _first(value: Any, default: Any = None) -> Any:
+    """Unwrap the singleton-list convention used by ``tts_params``.
+
+    ``tts_params`` wraps scalars in 1-element lists (e.g. ``task_type=["Base"]``)
+    because the model side batches per request. This returns the inner scalar.
+    """
+    if isinstance(value, (list, tuple)):
+        return value[0] if value else default
+    return value if value is not None else default
+
+
+def gate_audio_tokens(
+    *,
+    task_type: str | None,
+    x_vector_only_mode: bool,
+    icl_mode_override: bool | None,
+    ref_code_length: Any,
+) -> int:
+    """Reference-audio codec frames that actually enter the prefill.
+
+    Mirrors the ``in_context_mode`` decision in
+    ``Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information``:
+    reference codec frames are prefilled ONLY for Base in-context voice cloning.
+    They are NOT prefilled for:
+      * CustomVoice / VoiceDesign (no reference audio at all), or
+      * Base ``x_vector_only_mode`` (the reference audio is reduced to a single
+        speaker embedding vector; no per-frame codec context is inserted).
+
+    Counting ``ref_code_length`` outside the in-context path would re-introduce
+    the issue #4646 inconsistency in the other direction, so we gate it here.
+    """
+    in_context = (task_type == "Base") and not x_vector_only_mode
+    # An explicit per-request ``voice_clone_prompt.icl_mode`` wins if present,
+    # but ONLY for the Base task -- the estimator consults ``icl_mode`` inside
+    # its ``task_type == "Base"`` branch only, so CustomVoice / VoiceDesign must
+    # never have audio counted even if such a flag leaks into their params.
+    if task_type == "Base" and icl_mode_override is not None:
+        in_context = bool(icl_mode_override)
+    if not in_context:
+        return 0
+    try:
+        return max(0, int(ref_code_length)) if ref_code_length is not None else 0
+    except (TypeError, ValueError):
+        return 0
+
+
+def qwen3_tts_input_token_details(
+    *,
+    input_text: str,
+    instructions: str | None,
+    tts_params: dict[str, Any],
+    count_text_tokens: Callable[[str], int],
+) -> SpeechInputTokenDetails:
+    """Compute the input-token breakdown for a Qwen3-TTS request.
+
+    ``count_text_tokens`` tokenizes a string with the model's *text* tokenizer
+    and returns the token count. ``tts_params`` is the finalized param dict from
+    the adapter build (it already carries ``ref_code_length`` for ICL clones and
+    the resolved ``task_type`` / ``x_vector_only_mode``).
+    """
+    # Text tokens = the text to synthesize plus the style/emotion instructions,
+    # because Qwen3-TTS tokenizes and prepends the instruction block too.
+    text = input_text or ""
+    instr = instructions or ""
+    text_tokens = count_text_tokens(text) if text else 0
+    if instr.strip():
+        text_tokens += count_text_tokens(instr)
+
+    # Audio tokens = reference codec frames, gated to the in-context clone path.
+    voice_clone_prompt = _first(tts_params.get("voice_clone_prompt"), None)
+    icl_override = None
+    if isinstance(voice_clone_prompt, dict):
+        icl_flag = voice_clone_prompt.get("icl_mode")
+        if isinstance(icl_flag, bool):
+            icl_override = icl_flag
+    audio_tokens = gate_audio_tokens(
+        task_type=_first(tts_params.get("task_type"), "CustomVoice"),
+        x_vector_only_mode=bool(_first(tts_params.get("x_vector_only_mode"), False)),
+        icl_mode_override=icl_override,
+        ref_code_length=_first(tts_params.get("ref_code_length"), None),
+    )
+    return SpeechInputTokenDetails(text_tokens=int(text_tokens), audio_tokens=int(audio_tokens))
+
+
+def build_speech_usage(details: SpeechInputTokenDetails, output_tokens: int) -> SpeechTokenUsage:
+    """Assemble the final usage object from the input breakdown + output count."""
+    input_tokens = int(details.text_tokens) + int(details.audio_tokens)
+    out = max(0, int(output_tokens))
+    return SpeechTokenUsage(
+        input_tokens=input_tokens,
+        output_tokens=out,
+        total_tokens=input_tokens + out,
+        input_token_details=details,
+    )
+
+
+@dataclass
+class SpeechOutputTokenCounter:
+    """Accumulates the count of generated output (codec) tokens off the stream.
+
+    Why not count ``res.outputs[0].token_ids``? In the staged TTS pipeline the
+    token-generating stage (stage 0, the AR talker) is NOT surfaced to the audio
+    consumer loop — only the downstream code2wav (audio) outputs reach it, and
+    those carry an empty ``token_ids`` (their payload is a waveform, not tokens).
+
+    The generated codec-token count is instead reported by the pipeline itself
+    in per-stage metrics on the final (finished) output:
+    ``res.metrics['stage_metrics'][<stage>]['num_tokens_out']``. Only the AR
+    stage has a non-zero ``num_tokens_out`` (code2wav reports 0), so we take the
+    max across stages. Tracking the max is safe because the metrics appear on
+    the terminal output (intermediate outputs carry empty metrics).
+    """
+
+    output_tokens: int = 0
+
+    def observe(self, res: Any) -> None:
+        metrics = getattr(res, "metrics", None)
+        if not isinstance(metrics, dict):
+            return
+        stage_metrics = metrics.get("stage_metrics")
+        if not isinstance(stage_metrics, dict):
+            return
+        for stage in stage_metrics.values():
+            if not isinstance(stage, dict):
+                continue
+            n = stage.get("num_tokens_out")
+            if isinstance(n, int) and n > self.output_tokens:
+                self.output_tokens = n
+
+    def total(self) -> int:
+        """Generated codec/audio token count (max stage ``num_tokens_out``)."""
+        return self.output_tokens