[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>

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