[Edge][WP5] Add codec-token streaming mode with on-device decoding

* qwen3_tts_talker_only pipeline (single terminal LLM_AR stage,
  final_output_type "latent") + vllm_omni/deploy/qwen3_tts_talker_only.yaml;
  a terminal talker publishes its per-step codes.audio rows to the client
  (omni_client_multimodal_output_keys).
* /v1/audio/speech/stream: session.config.output_mode="codec_tokens" streams
  codec.start (JSON) + binary frames (<IHB seq/n_frames/flags + int16 rows,
  vllm_omni/entrypoints/openai/codec_stream.py) + codec.done; new input.cancel
  message aborts an in-flight generation for both codec and PCM sessions
  (single reader task + session inbox replaces direct socket reads); optional
  session.config.seed; audio.done gains "cancelled" only when it happened.
* serving_speech._generate_codec_chunks aligns the cumulative code tensor to
  output token ids (drops prefill placeholder/EOS rows) and emits deltas.
* TTS adapters: supports_codec_stream / codec_stream_spec (Qwen3-TTS: 16
  codebooks, 2048 ids, 12.5 Hz, 24 kHz, decoder qwen3_tts_tokenizer_v2_12hz).
* examples/online_serving/text_to_speech/qwen3_tts/edge_decoder_client.py
  decodes locally with Qwen3TTSTokenizerV2Decoder (ramp schedule, 72-frame
  context, cancel/reset), codec_parity_check.py and
  benchmarks/tts/codec_stream_e2e.py compare against the two-stage server.
* docs/serving/codec_token_streaming.md; CPU unit tests for the wire format,
  alignment, session config and handler (frames, unsupported adapter, cancel);
  L3 e2e tests/e2e/online_serving/test_qwen3_tts_codec_stream.py.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

examples/online_serving/text_to_speech/qwen3_tts/codec_parity_check.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+"""Parity check: server-decoded PCM vs client-decoded PCM from codec tokens.
+
+Run the full 2-stage server and the talker-only server one after another (or
+on two ports) with ``enforce_eager`` on stage 0 so per-request seeds are
+reproducible, then:
+
+  python codec_parity_check.py --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \\
+      --full-port 8000 --talker-port 8001 --seed 1234 --text "..."
+
+Reports the length-trimmed SNR (dB) between the two waveforms and the codec
+frame count. Both requests use the same ``seed``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import io
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+
+def snr_db(ref: np.ndarray, test: np.ndarray) -> float:
+    n = min(ref.size, test.size)
+    if n == 0:
+        return float("nan")
+    ref, test = ref[:n].astype(np.float64), test[:n].astype(np.float64)
+    noise = ref - test
+    p_sig = float(np.sum(ref**2)) + 1e-12
+    p_noise = float(np.sum(noise**2)) + 1e-12
+    return 10.0 * np.log10(p_sig / p_noise)
+
+
+def fetch_full_pcm(host: str, port: int, model: str, text: str, voice: str, seed: int) -> np.ndarray:
+    import requests
+
+    body = {
+        "model": model,
+        "input": text,
+        "voice": voice,
+        "response_format": "wav",
+        "seed": seed,
+    }
+    r = requests.post(f"http://{host}:{port}/v1/audio/speech", json=body, timeout=600)
+    r.raise_for_status()
+    import soundfile as sf
+
+    data, sr = sf.read(io.BytesIO(r.content), dtype="float32")
+    return data.reshape(-1)
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", required=True)
+    p.add_argument("--host", default="localhost")
+    p.add_argument("--full-port", type=int, default=8000)
+    p.add_argument("--talker-port", type=int, default=8001)
+    p.add_argument("--voice", default="vivian")
+    p.add_argument("--text", default="Parity check between server decoding and on-device decoding.")
+    p.add_argument("--seed", type=int, default=1234)
+    p.add_argument("--device", default="cpu")
+    p.add_argument("--full-wav", default=None, help="use a previously saved full-server WAV instead of requesting one")
+    p.add_argument("--out-dir", default=None)
+    args = p.parse_args(argv)
+
+    import edge_decoder_client as client
+
+    out_dir = Path(args.out_dir) if args.out_dir else None
+    if out_dir:
+        out_dir.mkdir(parents=True, exist_ok=True)
+    if args.full_wav:
+        import soundfile as sf
+
+        ref, _ = sf.read(args.full_wav, dtype="float32")
+        ref = ref.reshape(-1)
+    else:
+        ref = fetch_full_pcm(args.host, args.full_port, args.model, args.text, args.voice, args.seed)
+    client_args = argparse.Namespace(
+        host=args.host,
+        port=args.talker_port,
+        model=args.model,
+        text=args.text,
+        voice=args.voice,
+        task_type="CustomVoice",
+        language=None,
+        instructions=None,
+        max_new_tokens=None,
+        device=args.device,
+        left_context=72,
+        ramp=None,
+        seed=args.seed,
+        out=str(out_dir / "client.wav") if out_dir else None,
+        json=str(out_dir / "client.json") if out_dir else None,
+    )
+    result = asyncio.run(client.run(client_args))
+    import soundfile as sf
+
+    test, _ = sf.read(client_args.out, dtype="float32") if client_args.out else (np.zeros(0, dtype=np.float32), 0)
+    report = {
+        "ref_samples": int(ref.size),
+        "client_samples": int(test.size),
+        "snr_db": snr_db(ref, test),
+        "length_ratio": (test.size / ref.size) if ref.size else None,
+        "client_ttfa_ms": result.get("ttfa_ms"),
+        "client_rtf": result.get("rtf"),
+        "frames": result.get("total_frames"),
+    }
+    print(json.dumps(report, indent=1))
+    if out_dir:
+        (out_dir / "parity.json").write_text(json.dumps(report, indent=1))
+        sf.write(str(out_dir / "full.wav"), ref, 24000, format="WAV")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

examples/online_serving/text_to_speech/qwen3_tts/edge_decoder_client.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+"""Codec-token streaming client with on-device decoding (Qwen3-TTS).
+
+Connects to ``/v1/audio/speech/stream`` of a **talker-only** deployment
+(``--deploy-config qwen3_tts_talker_only.yaml``), requests
+``output_mode="codec_tokens"``, receives 16-codebook frames at 12.5 Hz and
+decodes them locally with ``Qwen3TTSTokenizerV2Decoder`` (from the model's
+``speech_tokenizer/`` folder) using the same chunk schedule / 72-frame left
+context as the server-side Code2Wav stage. Measures TTFA (first PCM out),
+per-chunk decode time, playback-start latency and RTF, and writes a WAV.
+
+Example:
+  python edge_decoder_client.py --host localhost --port 8000 \\
+      --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice --voice vivian \\
+      --text "Hello from the edge." --device cpu --out out.wav
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+sys.path.insert(0, str(REPO_ROOT / "benchmarks" / "tts"))
+
+from vllm_omni.entrypoints.openai.codec_stream import FLAG_EOS, FLAG_RESET, unpack_codec_frame  # noqa: E402
+from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp  # noqa: E402
+
+DEFAULT_RAMP = [2, 4, 8, 16, 25]
+
+
+# ------------------------------------------------------------------ decoder
+class LocalDecoder:
+    """Stateful chunked Qwen3-TTS decoder mirroring the Code2Wav stage."""
+
+    def __init__(
+        self, model: str, device: str = "cpu", left_context: int = 72, chunk_frames: int = 25, dtype: str = "float32"
+    ):
+        from safetensors.torch import load_file
+
+        from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
+            Qwen3TTSTokenizerV2Config,
+        )
+        from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+            Qwen3TTSTokenizerV2Decoder,
+        )
+
+        model_dir = self._resolve_dir(model)
+        cfg = Qwen3TTSTokenizerV2Config.from_pretrained(str(model_dir), subfolder="speech_tokenizer")
+        self.decoder = Qwen3TTSTokenizerV2Decoder._from_config(cfg.decoder_config)
+        state = load_file(str(model_dir / "speech_tokenizer" / "model.safetensors"))
+        dec_state = {k[len("decoder.") :]: v for k, v in state.items() if k.startswith("decoder.")}
+        missing, unexpected = self.decoder.load_state_dict(dec_state, strict=False)
+        if unexpected:
+            raise RuntimeError(f"unexpected decoder keys: {unexpected[:5]}")
+        if missing:
+            print(f"[decoder] warning: {len(missing)} missing keys (first: {missing[:3]})", file=sys.stderr)
+        self.decoder.eval().to(device=device, dtype=getattr(torch, dtype))
+        if hasattr(self.decoder, "precompute_snake_caches"):
+            self.decoder.precompute_snake_caches()
+        if hasattr(self.decoder, "_incremental_chunk_frames"):
+            self.decoder._incremental_chunk_frames = chunk_frames
+        self.device = device
+        self.num_quantizers = int(cfg.decoder_config.num_quantizers)
+        self.sample_rate = int(cfg.output_sample_rate)
+        self.left_context = left_context
+        self.chunk_frames = chunk_frames
+        self.reset()
+
+    @staticmethod
+    def _resolve_dir(model: str) -> Path:
+        p = Path(model)
+        if p.is_dir():
+            return p
+        from huggingface_hub import snapshot_download
+
+        return Path(snapshot_download(model, local_files_only=True))
+
+    def reset(self) -> None:
+        self.state: dict[str, Any] = {"prefix_frames": 0}
+        self.frames_decoded = 0
+
+    @torch.inference_mode()
+    def decode(self, codes: np.ndarray) -> np.ndarray:
+        """Decode ``[n_frames, Q]`` new frames, carrying state across calls; returns float32 PCM."""
+        if codes.shape[0] == 0:
+            return np.zeros(0, dtype=np.float32)
+        qf = torch.as_tensor(codes.T, dtype=torch.long, device=self.device)  # [Q, F]
+        wavs = self.decoder.batched_chunked_decode(
+            [qf],
+            [int(qf.shape[-1])],
+            caches=[self.state],
+            chunk_size=self.chunk_frames,
+            left_context_size=self.left_context,
+            max_batch_size=1,
+        )
+        wav = wavs[0]
+        if wav.dim() == 2 and wav.shape[0] == 1:
+            wav = wav[0]
+        self.frames_decoded += int(codes.shape[0])
+        return wav.float().cpu().numpy().reshape(-1)
+
+
+# ------------------------------------------------------------------- client
+async def run(args: argparse.Namespace) -> dict[str, Any]:
+    import websockets
+
+    ramp = args.ramp or DEFAULT_RAMP
+    schedule = parse_chunk_ramp({"codec_chunk_ramp": ramp}, steady=ramp[-1]) or ramp
+    steady = schedule[-1]
+    decoder = LocalDecoder(args.model, device=args.device, left_context=args.left_context, chunk_frames=steady)
+    url = f"ws://{args.host}:{args.port}/v1/audio/speech/stream"
+    config: dict[str, Any] = {
+        "type": "session.config",
+        "model": args.model,
+        "output_mode": "codec_tokens",
+        "stream_audio": True,
+        "response_format": "pcm",
+    }
+    for key in ("voice", "task_type", "language", "instructions", "max_new_tokens", "seed"):
+        val = getattr(args, key, None)
+        if val is not None:
+            config[key] = val
+    pcm_parts: list[np.ndarray] = []
+    chunk_log: list[dict[str, Any]] = []
+    pending = np.zeros((0, decoder.num_quantizers), dtype=np.int64)
+    chunk_idx = 0
+    t0 = None
+    t_first_pcm = None
+    codec_start = None
+    total_frames = 0
+    async with websockets.connect(url, max_size=64 * 1024 * 1024) as ws:
+        await ws.send(json.dumps(config))
+        await ws.send(json.dumps({"type": "input.text", "text": args.text}))
+        t0 = time.perf_counter()
+        await ws.send(json.dumps({"type": "input.done"}))
+        while True:
+            raw = await ws.recv()
+            now = time.perf_counter()
+            if isinstance(raw, (bytes, bytearray)):
+                seq, n, flags, arr = unpack_codec_frame(bytes(raw), decoder.num_quantizers)
+                if flags & FLAG_RESET:
+                    decoder.reset()
+                    pending = pending[:0]
+                total_frames += n
+                if n:
+                    pending = np.concatenate([pending, arr.astype(np.int64)], axis=0)
+                eos = bool(flags & FLAG_EOS)
+                target = schedule[min(chunk_idx, len(schedule) - 1)]
+                while pending.shape[0] >= target or (eos and pending.shape[0] > 0):
+                    take = pending[: min(target, steady)] if not eos else pending[:steady]
+                    pending = pending[take.shape[0] :]
+                    td = time.perf_counter()
+                    pcm = decoder.decode(take)
+                    t_dec = time.perf_counter() - td
+                    if t_first_pcm is None and pcm.size:
+                        t_first_pcm = time.perf_counter()
+                    pcm_parts.append(pcm)
+                    chunk_log.append(
+                        {
+                            "idx": chunk_idx,
+                            "t_ms": (time.perf_counter() - t0) * 1e3,
+                            "arrival_ms": (now - t0) * 1e3,
+                            "frames": int(take.shape[0]),
+                            "samples": int(pcm.size),
+                            "decode_ms": t_dec * 1e3,
+                        }
+                    )
+                    chunk_idx += 1
+                    target = schedule[min(chunk_idx, len(schedule) - 1)]
+                if eos:
+                    continue
+            else:
+                msg = json.loads(raw)
+                mtype = msg.get("type")
+                if mtype == "codec.start":
+                    codec_start = msg
+                elif mtype in ("codec.done", "audio.done"):
+                    pass
+                elif mtype == "session.done":
+                    break
+                elif mtype == "error":
+                    raise RuntimeError(f"server error: {msg.get('message')}")
+        await ws.send(json.dumps({"type": "session.close"}))
+    import stream_latency_metrics as metrics
+
+    sr = decoder.sample_rate
+    audio = np.concatenate(pcm_parts) if pcm_parts else np.zeros(0, dtype=np.float32)
+    total_s = (time.perf_counter() - t0) if t0 else 0.0
+    result = {
+        "codec_start": codec_start,
+        "schedule": schedule,
+        "total_frames": total_frames,
+        "ttfa_ms": ((t_first_pcm - t0) * 1e3) if t_first_pcm else None,
+        "playback_start_ms": metrics.playback_start_ms(chunk_log, sr),
+        "stall_at_ttfa_ms": metrics.stall_at_ttfa_ms(chunk_log, sr),
+        "audio_s": audio.size / sr,
+        "total_s": total_s,
+        "rtf": (total_s / (audio.size / sr)) if audio.size else None,
+        "decode_ms_total": sum(c["decode_ms"] for c in chunk_log),
+        "chunks": chunk_log,
+        "device": args.device,
+    }
+    if args.out:
+        import soundfile as sf
+
+        sf.write(args.out, audio, sr, format="WAV")
+        result["wav"] = args.out
+    if args.json:
+        Path(args.json).write_text(json.dumps(result, indent=1))
+    print(json.dumps({k: v for k, v in result.items() if k != "chunks"}, indent=1))
+    return result
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--host", default="localhost")
+    p.add_argument("--port", type=int, default=8000)
+    p.add_argument(
+        "--model", required=True, help="HF id or local dir (decoder weights are read from speech_tokenizer/)"
+    )
+    p.add_argument("--text", default="Hello, this is a codec token streaming test from the edge decoder client.")
+    p.add_argument("--voice", default="vivian")
+    p.add_argument("--task-type", dest="task_type", default="CustomVoice")
+    p.add_argument("--language", default=None)
+    p.add_argument("--instructions", default=None)
+    p.add_argument("--max-new-tokens", dest="max_new_tokens", type=int, default=None)
+    p.add_argument("--seed", type=int, default=None)
+    p.add_argument("--device", default="cpu")
+    p.add_argument("--left-context", type=int, default=72)
+    p.add_argument("--ramp", type=int, nargs="*", default=None, help="chunk schedule in frames, default 2 4 8 16 25")
+    p.add_argument("--out", default=None)
+    p.add_argument("--json", default=None)
+    args = p.parse_args(argv)
+    asyncio.run(run(args))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())