[Edge][WP5] Codec stream diagnostics: per-request output/codes counters, client saves received codes, offline parity diag

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

benchmarks/tts/codec_parity_diag.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+"""Offline parity diagnosis for the codec-token stream (CPU only).
+
+Given a codec_stream_e2e.py output directory (client_<i>.codes.npy, client_<i>.wav,
+full_<i>.wav) decode the *received* codes with the exact (non-chunked) decoder and
+compare (a) against the client's chunked decode (tests the client's chunking) and
+(b) against the full server's PCM (tests token equality + the server's decode).
+Reports lag-aligned SNR and correlation.
+
+  python benchmarks/tts/codec_parity_diag.py --dir benchmarks/tts/edge_results/codec_stream/<ts>
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+
+
+def aligned_snr(ref: np.ndarray, test: np.ndarray, max_lag: int = 4000) -> tuple[float, float, int]:
+    n = min(ref.size, test.size)
+    ref, test = ref[:n].astype(np.float64), test[:n].astype(np.float64)
+    if n < 3 * max_lag:
+        lag = 0
+    else:
+        c = np.correlate(ref[max_lag:-max_lag], test, mode="valid")
+        lag = int(np.argmax(c)) - max_lag
+    if lag >= 0:
+        a, b = test[lag:], ref[: n - lag]
+    else:
+        a, b = test[: n + lag], ref[-lag:]
+    m = min(a.size, b.size)
+    a, b = a[:m], b[:m]
+    snr = 10 * np.log10((b**2).sum() / (((b - a) ** 2).sum() + 1e-12))
+    corr = float(np.corrcoef(a, b)[0, 1]) if m > 1 else 0.0
+    return float(snr), corr, lag
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--dir", required=True)
+    p.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice")
+    args = p.parse_args(argv)
+    from vllm_omni.edge.decoder_export import load_code2wav_decoder
+
+    dec = load_code2wav_decoder(args.model)
+    d = Path(args.dir)
+    report = []
+    for codes_path in sorted(d.glob("client_*.codes.npy")):
+        i = codes_path.name.split("_")[1].split(".")[0]
+        codes = np.load(codes_path)
+        if codes.size == 0:
+            continue
+        with torch.inference_mode():
+            exact = dec._forward_exact(torch.as_tensor(codes.T[None], dtype=torch.long)).reshape(-1).numpy()
+        client, _ = sf.read(d / f"client_{i}.wav", dtype="float32")
+        full, _ = sf.read(d / f"full_{i}.wav", dtype="float32")
+        s1, c1, l1 = aligned_snr(exact, client.reshape(-1))
+        s2, c2, l2 = aligned_snr(full.reshape(-1), exact)
+        row = {
+            "i": int(i),
+            "frames": int(codes.shape[0]),
+            "exact_vs_client_chunked": {"snr_db": s1, "corr": c1, "lag": l1},
+            "full_server_vs_exact": {"snr_db": s2, "corr": c2, "lag": l2},
+        }
+        report.append(row)
+        print(json.dumps(row))
+    (d / "parity_diag.json").write_text(json.dumps(report, indent=1))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

examples/online_serving/text_to_speech/qwen3_tts/edge_decoder_client.py

@@ -199,6 +199,7 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
         if val is not None:
             config[key] = val
     pcm_parts: list[np.ndarray] = []
+    codes_all: list[np.ndarray] = []
     chunk_log: list[dict[str, Any]] = []
     pending = np.zeros((0, decoder.num_quantizers), dtype=np.int64)
     chunk_idx = 0
@@ -222,6 +223,7 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
                 total_frames += n
                 if n:
                     pending = np.concatenate([pending, arr.astype(np.int64)], axis=0)
+                    codes_all.append(arr.astype(np.int64))
                 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):
@@ -283,6 +285,10 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
 
         sf.write(args.out, audio, sr, format="WAV")
         result["wav"] = args.out
+        # Raw received codec rows (frames x codebooks) for offline parity checks.
+        codes_path = str(Path(args.out).with_suffix(".codes.npy"))
+        np.save(codes_path, np.concatenate(codes_all, axis=0) if codes_all else np.zeros((0, 0), dtype=np.int64))
+        result["codes"] = codes_path
     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))

vllm_omni/entrypoints/openai/serving_speech.py

@@ -3293,8 +3293,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         token_ids: list[int] = []
         sent = 0
         finished = False
+        n_outputs = 0
+        n_with_codes = 0
+        first_codes_at = -1
         try:
             async for res in generator:
+                n_outputs += 1
                 if usage_acc is not None:
                     usage_acc.observe(res)
                 outputs = getattr(res, "outputs", None)
@@ -3308,6 +3312,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 codes_cum = self._extract_codec_output(res)
                 done = bool(getattr(res, "finished", False))
                 if codes_cum is not None:
+                    n_with_codes += 1
+                    if first_codes_at < 0:
+                        first_codes_at = n_outputs
                     aligned = align_codec_rows(codes_cum, token_ids, codebook_size)
                     delta = aligned[sent:]
                     if delta.shape[0] >= min_frames or (done and delta.shape[0] > 0):
@@ -3319,6 +3326,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     finished = True
                     yield None, True
         finally:
+            logger.info(
+                "[codec_stream] req=%s engine outputs=%d with codes=%d first codes at output #%d rows sent=%d",
+                request_id,
+                n_outputs,
+                n_with_codes,
+                first_codes_at,
+                sent,
+            )
             if usage_acc is not None and tts_params is not None and finished:
                 self._validate_tts_generation(tts_params, usage_acc)