[Perf][CosyVoice3] Bounded-window incremental HiFT vocoder streaming (#7521)

Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>

benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py

@@ -0,0 +1,517 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CosyVoice3 streaming e2e benchmark: concurrency x utterance-length matrix.
+
+Launches the server once per (concurrency, step-size) combo, streams TTS
+requests cycling through several utterance lengths, and reports
+TTFA/E2E/RTF/per-chunk latency (mean, stdev, range) per cell. "Step size" is
+``codec_chunk_frames``, the token-hop length the first emitted chunk waits
+for; later chunks wait for a doubling hop capped at ``4 * codec_chunk_frames``
+(see ``_token_hop_lengths``). Per-chunk latency is reconstructed from that
+same ramp rather than from raw transport read boundaries, which don't align
+to model chunks (see ``one_request``'s docstring).
+
+Running with no flags reproduces the default matrix. Per-request correctness
+is checked: non-silent, long enough, and (with a fixed ``--seed``)
+byte-identical PCM across requests for the same text.
+"""
+
+import argparse
+import json
+import os
+import signal
+import statistics
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+
+import httpx
+import numpy as np
+
+REPO = Path(__file__).resolve().parents[2]
+DEFAULT_MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
+
+DEFAULT_REF_AUDIO = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
+DEFAULT_REF_TEXT = "希望你以后能够做的比我还好呦。"
+
+DEFAULT_TEXTS = [
+    "Thank you so much for the birthday gift you sent me.",
+    "The weather is lovely today, with warm sunshine and a gentle breeze. "
+    "We decided to take a walk in the park to enjoy the spring flowers and the fresh air.",
+    "The weather is lovely today, with warm sunshine and a gentle breeze. We decided to take a walk in the "
+    "park to enjoy the spring flowers and the fresh air. Many people were out walking, children were "
+    "running and playing on the grass, and older folks sat on benches chatting with one another. We "
+    "strolled slowly along the path and came across a cherry tree in full bloom, its pink petals drifting "
+    "on the wind. It was so beautiful that we stopped to take some photos to remember the moment. Then we "
+    "continued on and reached a small lake, its water so clear you could see straight to the bottom, with "
+    "the blue sky and white clouds reflected on its surface. A few ducks paddled back and forth across the "
+    "water, occasionally letting out a soft quack. We sat on a bench by the lake for a while, enjoying the "
+    "peaceful afternoon. Time passed quickly, and before we knew it evening had arrived, with the setting "
+    "sun painting the sky in shades of gold. Reluctant to leave, we finally headed home, already planning "
+    "our next visit. On the way back we talked about everything we had seen and heard that day, feeling "
+    "wonderfully happy. It was truly an unforgettable weekend.",
+]
+
+# PCM16 mono @ 24 kHz (CosyVoice3 output rate)
+SAMPLE_RATE = 24000
+
+# Well below any real speech's RMS; catches all-zero/near-silent PCM that byte
+# count and duration checks alone pass through.
+MIN_RMS = 50.0
+
+# Overlay: only codec_chunk_frames is overridden; must live next to cosyvoice3.yaml (relative base_config)
+DEPLOY_TEMPLATE = """# Auto-generated by bench_cosyvoice3_concurrency_matrix.py -- do not edit.
+base_config: cosyvoice3.yaml
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: {chunk_frames}
+      codec_max_chunk_frames: {max_chunk_frames}
+      codec_stream_scale_factor: {stream_scale_factor}
+      codec_left_context_frames: 25
+      codec_vocab_size: 6561
+"""
+
+# CosyVoice3Config constants (transformers_utils/configs/cosyvoice3.py): tokens
+# are consumed at token_mel_ratio mel frames each, and HiFT emits
+# samples_per_mel PCM16 samples per mel frame. Pinned here (not left to the
+# server's defaults) since one_request's byte-boundary reconstruction below
+# depends on knowing them exactly.
+TOKEN_MEL_RATIO = 2
+SAMPLES_PER_MEL = 480
+STREAM_SCALE_FACTOR = 2
+
+
+def _stats(values: list[float]) -> dict[str, float]:
+    """Mean/median/stdev/p95/min/max/n for a list of per-request measurements."""
+    n = len(values)
+    if n == 0:
+        return {
+            "n": 0,
+            "mean": float("nan"),
+            "median": float("nan"),
+            "stdev": float("nan"),
+            "p95": float("nan"),
+            "min": float("nan"),
+            "max": float("nan"),
+        }
+    return {
+        "n": n,
+        "mean": statistics.mean(values),
+        "median": statistics.median(values),
+        "stdev": statistics.stdev(values) if n >= 2 else 0.0,
+        "p95": statistics.quantiles(values, n=20)[18] if n >= 2 else values[0],
+        "min": min(values),
+        "max": max(values),
+    }
+
+
+def vllm_bin() -> str:
+    for c in (REPO / ".venv/bin/vllm", Path(sys.executable).parent / "vllm"):
+        if c.is_file():
+            return str(c)
+    return "vllm"
+
+
+def wait_ready(url: str, proc: subprocess.Popen, log: Path, timeout: float) -> None:
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if proc.poll() is not None:
+            tail = log.read_text(errors="replace").splitlines()[-40:] if log.exists() else []
+            raise RuntimeError("server exited during startup:\n  " + "\n  ".join(tail))
+        try:
+            if httpx.get(f"{url}/health", timeout=2.0).status_code == 200:
+                return
+        except Exception:
+            pass
+        time.sleep(2.0)
+    raise TimeoutError(f"server not ready within {timeout}s (log: {log})")
+
+
+def _token_hop_lengths(chunk_frames: int, max_chunk_frames: int):
+    """The token-hop-length ramp talker2code2wav_async_chunk uses to decide
+    when to fire the next flow-solve/emission: chunk_frames, then doubling
+    (STREAM_SCALE_FACTOR) each step, capped at max_chunk_frames. This -- not
+    codec_chunk_frames alone -- is what actually paces emitted audio chunks:
+    with the defaults here, hops are chunk_frames, 2x, 4x, 4x, 4x, ...
+
+    The very first emission is consistently one token short of chunk_frames
+    (confirmed against real speech.audio.delta event sizes) even though the
+    ramp's own progression is computed from the nominal chunk_frames, not the
+    actual first hop -- so only the first yield gets the -1 adjustment.
+    """
+    yield chunk_frames - 1
+    hop = chunk_frames
+    while True:
+        hop = min(max_chunk_frames, max(chunk_frames, hop * STREAM_SCALE_FACTOR))
+        yield hop
+
+
+def one_request(client, url, model, text, ref_audio, ref_text, chunk_frames, seed=None):
+    """Run one streaming request, returning TTFA, E2E, per-chunk gaps, and audio bytes.
+
+    ``response_format=pcm`` + ``stream_format=audio`` streams raw bytes with no
+    framing, so a client read boundary alone is a transport artifact, not a
+    model chunk boundary. Reconstructed instead by replicating the server's
+    own token-hop-length ramp (see ``_token_hop_lengths``) to compute each
+    real chunk's expected byte size ahead of time, and slicing the
+    accumulated byte stream at those offsets -- independent of how the
+    transport happens to fragment the underlying reads. The final chunk
+    (finalize, which releases extra held-back frames) is whatever remains
+    when the stream closes, not a predicted hop.
+    """
+    payload = {
+        "model": model,
+        "input": text,
+        "response_format": "pcm",
+        "stream": True,
+        "stream_format": "audio",
+        "ref_audio": ref_audio,
+        "ref_text": ref_text,
+    }
+    if seed is not None:
+        payload["seed"] = seed
+    t0 = time.perf_counter()
+    ttfa = None
+    nbytes = 0
+    audio_bytes = bytearray()
+    chunk_gaps_ms = []  # inter-chunk arrival gaps (per model chunk), by step index
+    prev_t = None
+    step_idx = 0
+
+    bytes_per_token = TOKEN_MEL_RATIO * SAMPLES_PER_MEL * 2
+    hop_gen = _token_hop_lengths(chunk_frames, 4 * chunk_frames)
+    next_boundary = next(hop_gen) * bytes_per_token
+    last_crossed_bytes = 0
+    received = 0
+    last_now = t0
+
+    with client.stream("POST", f"{url}/v1/audio/speech", json=payload, headers={"Authorization": "Bearer EMPTY"}) as r:
+        if r.status_code != 200:
+            raise RuntimeError(f"{r.status_code}: {r.read()[:300]!r}")
+        for chunk in r.iter_bytes():
+            if not chunk:
+                continue
+            now = time.perf_counter()
+            if ttfa is None:
+                ttfa = now - t0
+            last_now = now
+            received += len(chunk)
+            nbytes += len(chunk)
+            audio_bytes += chunk
+            while received >= next_boundary:
+                if prev_t is not None:
+                    chunk_gaps_ms.append((step_idx, (now - prev_t) * 1e3))
+                    step_idx += 1
+                prev_t = now
+                last_crossed_bytes = next_boundary
+                next_boundary += next(hop_gen) * bytes_per_token
+
+    if received > last_crossed_bytes:
+        # Finalize chunk: releases extra held-back frames, so it doesn't
+        # match a predicted hop size -- its arrival is the stream closing.
+        if prev_t is not None:
+            chunk_gaps_ms.append((step_idx, (last_now - prev_t) * 1e3))
+            step_idx += 1
+
+    e2e = time.perf_counter() - t0
+    audio_s = nbytes / 2 / SAMPLE_RATE
+    return {
+        "ttfa_s": ttfa,
+        "e2e_s": e2e,
+        "audio_s": audio_s,
+        "rtf": e2e / audio_s if audio_s > 0 else float("nan"),
+        "chunk_gaps_ms": chunk_gaps_ms,
+        "nbytes": nbytes,
+        "audio_bytes": audio_bytes,
+    }
+
+
+def check_correctness(r: dict, min_audio_s: float) -> list[str]:
+    """Return a list of correctness violations (empty = pass)."""
+    problems = []
+    if r["nbytes"] <= 0:
+        problems.append("empty audio (0 bytes)")
+    if r["audio_s"] < min_audio_s:
+        problems.append(f"audio too short ({r['audio_s']:.2f}s < {min_audio_s:.2f}s)")
+    if r["nbytes"] > 0:
+        samples = np.frombuffer(bytes(r["audio_bytes"]), dtype=np.int16).astype(np.float64)
+        rms = float(np.sqrt(np.mean(samples**2)))
+        if rms < MIN_RMS:
+            problems.append(f"near-silent audio (rms={rms:.1f} < {MIN_RMS})")
+    return problems
+
+
+def drive(url, args, n, conc, label, results, lock, idx, chunk_frames):
+    """Run ``n`` requests at concurrency ``conc``, cycling through ``args.texts``."""
+    texts = args.texts
+
+    def worker():
+        with httpx.Client(timeout=600.0) as c:
+            while True:
+                with lock:
+                    i = idx[0]
+                    idx[0] += 1
+                if i >= n:
+                    return
+                text = texts[i % len(texts)]
+                try:
+                    r = one_request(
+                        c, url, args.model, text, args.ref_audio, args.ref_text, chunk_frames, seed=args.seed
+                    )
+                    r["text"] = text
+                    with lock:
+                        results.append(r)
+                except Exception as e:
+                    print(f"    [{label}] request {i} failed: {e}", flush=True)
+
+    ts = [threading.Thread(target=worker) for _ in range(conc)]
+    t0 = time.perf_counter()
+    for t in ts:
+        t.start()
+    for t in ts:
+        t.join()
+    return time.perf_counter() - t0
+
+
+def run_combo(args, conc: int, chunk_frames: int, out: Path) -> dict:
+    tag = f"c{conc}_chunk{chunk_frames}"
+    run_dir = out / tag
+    run_dir.mkdir(parents=True, exist_ok=True)
+    log = run_dir / "server.log"
+
+    deploy = REPO / "vllm_omni/deploy" / f"_autogen_steps_{tag}.yaml"
+    deploy.write_text(
+        DEPLOY_TEMPLATE.format(
+            chunk_frames=chunk_frames,
+            max_chunk_frames=4 * chunk_frames,
+            stream_scale_factor=STREAM_SCALE_FACTOR,
+        )
+    )
+
+    cmd = [
+        vllm_bin(),
+        "serve",
+        args.model,
+        "--omni",
+        "--deploy-config",
+        str(deploy),
+        "--host",
+        "127.0.0.1",
+        "--port",
+        str(args.port),
+        "--allowed-local-media-path",
+        str(REPO),
+    ]
+    env = dict(os.environ)
+    if args.cuda_visible_devices:
+        env["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices
+    if args.trt is not None:
+        env["COSYVOICE3_TRT"] = args.trt
+
+    print(f"\n=== concurrency={conc} step_size={chunk_frames} ===", flush=True)
+    print(f"  launching server (log: {log})", flush=True)
+    with open(log, "w") as lf:
+        proc = subprocess.Popen(cmd, env=env, stdout=lf, stderr=subprocess.STDOUT, start_new_session=True)
+    url = f"http://127.0.0.1:{args.port}"
+    try:
+        wait_ready(url, proc, log, args.server_timeout)
+        print("  server ready", flush=True)
+
+        results: list[dict] = []
+        lock = threading.Lock()
+        idx = [0]
+        if args.warmup:
+            print(f"  warmup ({args.warmup} req)", flush=True)
+            drive(url, args, args.warmup, min(conc, args.warmup), "warmup", results, lock, idx, chunk_frames)
+            results.clear()
+            idx[0] = 0
+
+        print(f"  running {args.num_prompts} req @ concurrency {conc}", flush=True)
+        wall = drive(url, args, args.num_prompts, conc, "bench", results, lock, idx, chunk_frames)
+
+        ok = [r for r in results if r.get("ttfa_s") is not None]
+        problems = []
+        for r in ok:
+            problems += check_correctness(r, args.min_audio_s)
+        # With a fixed seed, same-text requests must produce byte-identical PCM (per text group)
+        if args.seed is not None:
+            by_text: dict[str, bytes] = {}
+            for r in ok:
+                t = r.get("text")
+                ref = by_text.get(t)
+                if ref is None:
+                    by_text[t] = r.get("audio_bytes")
+                elif r.get("audio_bytes") != ref:
+                    problems.append(f"byte-identity violation for text {t!r}: seeded requests produced different PCM")
+        n_ok = len(ok)
+        total_audio_s = sum(r["audio_s"] for r in ok)
+        all_gaps = [g for r in ok for g in r["chunk_gaps_ms"]]  # (step_index, gap_ms) per request
+        gap_vals = [g for _, g in all_gaps]
+
+        # bucket gaps by chunk position within a request, to see any O(n) slope
+        by_idx: dict[int, list[float]] = {}
+        for step_index, g in all_gaps:
+            by_idx.setdefault(step_index, []).append(g)
+        step_latency_by_idx = {i: statistics.median(vals) for i, vals in sorted(by_idx.items())}
+
+        agg = {
+            "concurrency": conc,
+            "step_size": chunk_frames,
+            "ok": n_ok,
+            "errors": args.num_prompts - n_ok,
+            "correctness_problems": len(problems),
+            "correctness_sample": problems[:5],
+            "total_audio_s": total_audio_s,
+            "throughput_audio_s_per_s": total_audio_s / wall if wall > 0 else float("nan"),
+            "ttfa_s": _stats([r["ttfa_s"] for r in ok]),
+            "e2e_s": _stats([r["e2e_s"] for r in ok]),
+            "rtf": _stats([r["rtf"] for r in ok]),
+            "step_latency_ms": _stats(gap_vals),
+            "step_latency_ms_by_idx": step_latency_by_idx,
+            "num_steps": len(all_gaps),
+            # per-utterance-length breakdown; n per text = num_prompts / len(texts)
+            "by_text": {
+                t: {
+                    "audio_s": _stats([r["audio_s"] for r in ok if r.get("text") == t]),
+                    "e2e_s": _stats([r["e2e_s"] for r in ok if r.get("text") == t]),
+                    "rtf": _stats([r["rtf"] for r in ok if r.get("text") == t]),
+                }
+                for t in sorted({r.get("text") for r in ok})
+            },
+        }
+        print(json.dumps(agg, indent=2, ensure_ascii=False), flush=True)
+        return agg
+    finally:
+        print("  shutting down server", flush=True)
+        try:
+            os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
+        except Exception:
+            proc.terminate()
+        try:
+            proc.wait(timeout=60)
+        except subprocess.TimeoutExpired:
+            try:
+                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+            except Exception:
+                proc.kill()
+            proc.wait(timeout=20)
+        deploy.unlink(missing_ok=True)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", default=DEFAULT_MODEL)
+    p.add_argument(
+        "--text",
+        dest="texts",
+        nargs="+",
+        default=DEFAULT_TEXTS,
+        help="One or more utterance texts. With several, requests cycle through "
+        "them so a single server launch covers all utterance lengths; results "
+        "are bucketed per text.",
+    )
+    p.add_argument("--ref-audio", default=DEFAULT_REF_AUDIO)
+    p.add_argument("--ref-text", default=DEFAULT_REF_TEXT)
+    p.add_argument("--concurrency", type=int, nargs="+", default=[1, 8])
+    p.add_argument(
+        "--step-size",
+        type=int,
+        nargs="+",
+        default=[15],
+        help="codec_chunk_frames values to sweep (the streaming step size).",
+    )
+    p.add_argument("--num-prompts", type=int, default=30)
+    p.add_argument("--warmup", type=int, default=2)
+    p.add_argument("--port", type=int, default=8091)
+    p.add_argument(
+        "--server-timeout",
+        type=float,
+        default=1800.0,
+        help="Startup budget; first run also downloads weights and builds the TRT engine.",
+    )
+    p.add_argument("--output-dir", default="./perf_steps")
+    p.add_argument("--cuda-visible-devices", default="0")
+    p.add_argument(
+        "--trt",
+        choices=["0", "1"],
+        default=None,
+        help="Pin COSYVOICE3_TRT. 0=force torch estimator, 1=allow TensorRT. Default: unset.",
+    )
+    p.add_argument(
+        "--min-audio-s",
+        type=float,
+        default=1.0,
+        help="Minimum expected audio seconds per request for the correctness check.",
+    )
+    p.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Fixed seed for reproducible (byte-identical) audio. Pass --seed '' "
+        "equivalent (any non-int) to disable -- there is no unset sentinel here "
+        "by design, since the default sweep is meant to be exactly reproducible.",
+    )
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    out = Path(args.output_dir).resolve()
+    out.mkdir(parents=True, exist_ok=True)
+
+    all_results = []
+    failed = False
+    for conc in args.concurrency:
+        for chunk_frames in args.step_size:
+            try:
+                result = run_combo(args, conc, chunk_frames, out)
+                all_results.append(result)
+                if result["errors"] > 0 or result["correctness_problems"] > 0:
+                    failed = True
+            except Exception as e:
+                print(f"  combo c={conc} chunk={chunk_frames} FAILED: {e}", flush=True)
+                failed = True
+
+    out_path = out / "steps_results.json"
+    with open(out_path, "w", encoding="utf-8") as f:
+        json.dump({"model": args.model, "texts": args.texts, "runs": all_results}, f, indent=2, ensure_ascii=False)
+    print(f"\nResults written to {out_path}", flush=True)
+
+    print("\n=== SUMMARY ===")
+    header = (
+        f"{'conc':>4} {'step':>4} {'ok':>4} {'corr':>4} "
+        f"{'TTFA med(ms)':>12} {'step med(ms)':>12} {'step p95(ms)':>12} {'RTF mean+-sd':>14}"
+    )
+    print(header)
+    print("-" * len(header))
+    for r in all_results:
+        print(
+            f"{r['concurrency']:>4} {r['step_size']:>4} {r['ok']:>4} {r['correctness_problems']:>4} "
+            f"{r['ttfa_s']['median'] * 1e3:>12.1f} "
+            f"{r['step_latency_ms']['median']:>12.1f} {r['step_latency_ms']['p95']:>12.1f} "
+            f"{r['rtf']['mean']:>6.3f}+-{r['rtf']['stdev']:>5.3f}"
+        )
+        for t, stats in r.get("by_text", {}).items():
+            print(
+                f"      text {t[:24]!r}: e2e {stats['e2e_s']['mean'] * 1e3:8.1f}"
+                f"+-{stats['e2e_s']['stdev'] * 1e3:6.1f} ms  (n={stats['e2e_s']['n']}, "
+                f"audio {stats['audio_s']['mean']:.1f}s)"
+            )
+
+    if failed:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -66,14 +66,12 @@ dev = [
     "jiwer==4.0.0",
     "zhon==2.1.1",
     "zhconv==1.4.3",
-    "scipy==1.18.0",
     "funasr==1.3.29",
     "lpips==0.1.4", # for quantization fp8
     "orjson==3.11.9", # to make v1/embeddings API fast
     "voxcpm==2.0.3", # for VoxCPM2 path-locating in tests
     "torchdiffeq==0.2.5", # for covo_audio Test
     "FlagEmbedding==1.4.0", # for hunyuan_image3 tests
-    "s3tokenizer==0.3.0", # for tts test
     # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
     # whose top-level stepaudio2 package conflicts with this distribution.
     "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)

requirements/common.txt

@@ -41,3 +41,8 @@ pydantic-settings>=2.0.0
 gguf>=0.10.0
 # Required for running the Cosmos3 model with guardrails enabled (default behavior).
 cosmos-guardrail>=0.3.1
+# Required for CosyVoice3 (get_window/resample_poly). >=1.18.0 requires
+# Python >=3.12, breaking base installs on the project's supported 3.10/3.11;
+# these are long-stable APIs with no version-specific need for 1.18.
+scipy>=1.11.0
+s3tokenizer==0.3.0

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py

@@ -15,6 +15,7 @@ from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import AttentionConfig
 from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
     CausalHiFTGenerator,
+    HiFTGenerator,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -76,6 +77,39 @@ def test_causal_hift_stft_moves_window_to_input_device(causal_hift):
     assert causal_hift.stft_window.device == waveform.device
 
 
+class _MinimalF0Predictor(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.proj = nn.Conv1d(80, 1, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj(x).squeeze(1).abs()
+
+
+@pytest.mark.parametrize("sampling_rate", [22050, 24000])
+def test_noncausal_hift_inference_runs_with_real_sinegen(sampling_rate):
+    """Base (non-causal) HiFTGenerator with a real SineGen/SineGen2 (not a
+    test double) must still run: SourceModuleHnNSF/SineGen/SineGen2/_f02sine
+    all branch on `self.causal` to return the pre-streaming (shorter) tuple
+    contract for non-causal callers, and any one of them regressing back to
+    always returning the streaming-length tuple breaks this immediately.
+    """
+    hift = HiFTGenerator(
+        base_channels=32,
+        sampling_rate=sampling_rate,
+        upsample_rates=[8, 5, 3],
+        upsample_kernel_sizes=[16, 11, 7],
+        source_resblock_kernel_sizes=[7, 7, 11],
+        source_resblock_dilation_sizes=[[1, 3, 5]] * 3,
+        f0_predictor=_MinimalF0Predictor(),
+    ).eval()
+
+    speech, _ = hift.inference(torch.randn(1, 80, 20))
+
+    assert speech.shape[0] == 1
+    assert torch.isfinite(speech).all()
+
+
 def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypatch):
     class NpuTensor:
         device = SimpleNamespace(type="npu")
@@ -83,16 +117,17 @@ def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypa
     tensor = NpuTensor()
     stft_calls = []
     istft_calls = []
-    monkeypatch.setattr(
-        causal_hift,
-        "_stft_on_cpu",
-        lambda value: stft_calls.append(value) or ("real", "imag"),
-    )
-    monkeypatch.setattr(
-        causal_hift,
-        "_istft_on_cpu",
-        lambda magnitude, phase: istft_calls.append((magnitude, phase)) or "waveform",
-    )
+
+    def fake_stft(value):
+        stft_calls.append(value)
+        return ("real", "imag")
+
+    def fake_istft(magnitude, phase):
+        istft_calls.append((magnitude, phase))
+        return "waveform"
+
+    monkeypatch.setattr(causal_hift, "_stft_on_cpu", fake_stft)
+    monkeypatch.setattr(causal_hift, "_istft_on_cpu", fake_istft)
 
     assert causal_hift._stft(tensor) == ("real", "imag")
     assert causal_hift._istft(tensor, tensor) == "waveform"
@@ -627,7 +662,11 @@ def test_code2wav_forward_finalizes_hift_tail():
 
         def inference(self, speech_feat, finalize=True):
             self.finalize_calls.append(bool(finalize))
-            return torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype), None
+            return (
+                torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype),
+                None,
+                None,
+            )
 
     model = object.__new__(CosyVoice3Code2Wav)
     nn.Module.__init__(model)

tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py

@@ -0,0 +1,415 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Correctness + boundedness tests for the incremental (windowed) HiFT path.
+
+The streaming vocoder runs HiFT over a bounded mel window instead of the full
+cumulative history, carrying harmonic phase and noise-buffer offset across
+chunks.
+
+* **Correctness** — windowed output must match the full cumulative re-run.
+  With full history the match is bit-exact (proves the phase/noise carry is
+  right). With a bounded window, float64 still agrees to ~1e-16 (the window
+  is algebraically exact); float32 differs by ~2e-7 absolute (reassociation
+  rounding, not truncation), so tests assert a tight atol/rtol rather than
+  PCM16 byte-identity, since a sub-LSB perturbation can flip a boundary sample.
+* **Boundedness** — the per-chunk window must not grow with mel history.
+
+Tests are parametrized across REAL_CONFIG (24000 Hz -> SineGen2, what
+CosyVoice3 ships) and LEGACY_CONFIG (22050 Hz -> SineGen) so a class-selection
+regression in either path is caught.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
+    CausalConvRNNF0Predictor,
+    CausalHiFTGenerator,
+    SineGen,
+    SineGen2,
+    _wrapped_slice,
+)
+from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+CHUNK_LEN = 24
+TOTAL_MEL = 96
+SPM = 480  # samples per mel frame for the small test HiFT (8*5*3*4), matching the real model
+
+LONG_TOTAL_MEL = 400  # long enough that the 64-frame window actually truncates
+PCM16_LSB = 1.0 / 32767.0  # smallest representable PCM16 difference
+
+# Worst observed deviation is 1.97e-7 (~2x float32 eps); atol=1e-6 keeps 5x headroom.
+ATOL = 1e-6
+RTOL = 1e-5
+
+# What CosyVoice3 ships (-> SineGen2). nb_harmonics=8 is required, not a free
+# choice: SineGen2's causal buffers are hardcoded to dim 9 (= harmonic_num+1).
+REAL_CONFIG = {"sampling_rate": 24000, "nb_harmonics": 8}
+LEGACY_CONFIG = {"sampling_rate": 22050, "nb_harmonics": 4}  # -> SineGen, still shipped code
+
+CONFIGS = [
+    pytest.param(REAL_CONFIG, id="sinegen2_real_config"),
+    pytest.param(LEGACY_CONFIG, id="sinegen1_legacy_config"),
+]
+
+
+def _make_hift(config: dict = REAL_CONFIG) -> CausalHiFTGenerator:
+    """Small CausalHiFTGenerator; defaults to the config CosyVoice3 ships."""
+    torch.manual_seed(0)
+    f0_predictor = CausalConvRNNF0Predictor(num_class=1, in_channels=80, cond_channels=16)
+    return CausalHiFTGenerator(
+        in_channels=80,
+        base_channels=32,
+        nb_harmonics=config["nb_harmonics"],
+        sampling_rate=config["sampling_rate"],
+        upsample_rates=[8, 5, 3],
+        upsample_kernel_sizes=[16, 11, 7],
+        source_resblock_kernel_sizes=[7, 7, 11],
+        source_resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+        resblock_kernel_sizes=[3, 7, 11],
+        resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+        f0_predictor=f0_predictor,
+    ).eval()
+
+
+def test_real_config_resolves_to_sinegen2():
+    """Guard against silently testing the wrong SineGen class again."""
+    hift = _make_hift(REAL_CONFIG)
+    assert isinstance(hift.m_source.l_sin_gen, SineGen2)
+
+    hift_legacy = _make_hift(LEGACY_CONFIG)
+    assert isinstance(hift_legacy.m_source.l_sin_gen, SineGen)
+
+
+def _make_model(hift: CausalHiFTGenerator, window_len: int) -> CosyVoice3Code2Wav:
+    """Build a CosyVoice3Code2Wav shell wired to the given HiFT."""
+    model = object.__new__(CosyVoice3Code2Wav)
+    nn.Module.__init__(model)
+    model.hift = hift
+    model._hift_window_len = window_len
+    return model
+
+
+def _chunks(total_mel: int = TOTAL_MEL, chunk_len: int = CHUNK_LEN) -> list[torch.Tensor]:
+    torch.manual_seed(1)
+    mel = torch.randn(1, 80, total_mel)
+    return [mel[:, :, i : i + chunk_len] for i in range(0, total_mel, chunk_len)]
+
+
+def _variable_chunks(total_mel: int, first_len: int, step_len: int) -> list[torch.Tensor]:
+    """First chunk sized like a real prompt prefill, then small fixed-size steps."""
+    torch.manual_seed(1)
+    mel = torch.randn(1, 80, total_mel)
+    chunks = [mel[:, :, :first_len]]
+    pos = first_len
+    while pos < total_mel:
+        chunks.append(mel[:, :, pos : pos + step_len])
+        pos += step_len
+    return chunks
+
+
+def _full_reference(model: CosyVoice3Code2Wav, chunks: list[torch.Tensor]) -> torch.Tensor:
+    """Re-run HiFT over the full cumulative mel each chunk (pre-windowing behavior).
+
+    Uses non-finalize inference per chunk, matching the old streaming vocoder; a
+    single finalize pass would emit the look-right tail streaming holds back.
+    """
+    torch.manual_seed(0)
+    emitted = []
+    cache: dict[str, torch.Tensor] | None = None
+    for chunk in chunks:
+        cached = None if cache is None else cache.get("mel")
+        if cached is not None and cached.numel() > 0:
+            tts_mel = torch.cat([cached, chunk], dim=-1)
+        else:
+            tts_mel = chunk
+        speech, _, _ = model.hift.inference(speech_feat=tts_mel, finalize=False)
+        speech = speech.reshape(speech.shape[0], -1)
+        offset = 0 if cache is None else cache.get("speech_offset")
+        emitted.append(speech[:, offset:])
+        cache = {"mel": tts_mel.detach().cpu(), "speech_offset": int(speech.shape[-1])}
+    return torch.cat(emitted, dim=-1)
+
+
+def _pcm16(x: torch.Tensor) -> torch.Tensor:
+    """Quantize float audio to PCM16, the format the server streams to clients."""
+    return (x.clamp(-1, 1) * 32767).round().to(torch.int16)
+
+
+def _incremental(model: CosyVoice3Code2Wav, chunks: list[torch.Tensor]) -> torch.Tensor:
+    """Run the new windowed path via _stream_hift_from_feat."""
+    torch.manual_seed(0)
+    emitted = []
+    cache: dict[str, torch.Tensor] | None = None
+    for chunk in chunks:
+        speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False)
+        emitted.append(speech.reshape(speech.shape[0], -1))
+    return torch.cat(emitted, dim=-1)
+
+
+def _rel_divergence(a: torch.Tensor, b: torch.Tensor) -> float:
+    return ((a - b).abs().mean() / a.abs().mean().clamp_min(1e-6)).item()
+
+
+@pytest.mark.parametrize("config", CONFIGS)
+def test_incremental_hift_matches_full_reference_exactly_with_full_history(config):
+    """With a window covering the full history, the windowed path must match the
+    full cumulative re-run exactly. This proves the phase/noise carry is right."""
+    hift = _make_hift(config)
+    model = _make_model(hift, window_len=TOTAL_MEL)  # window >= history
+    chunks = _chunks()
+
+    full = _full_reference(model, chunks)
+    incr = _incremental(model, chunks)
+
+    assert full.shape == incr.shape, f"{full.shape} vs {incr.shape}"
+    assert _rel_divergence(full, incr) < 1e-6
+
+
+@pytest.mark.parametrize("config", CONFIGS)
+def test_incremental_hift_bounded_window_is_close(config):
+    """With a bounded window, the windowed path stays close to the full re-run;
+    only the truncated deep history diverges, and it is small."""
+    hift = _make_hift(config)
+    model = _make_model(hift, window_len=48)
+    chunks = _chunks()
+
+    full = _full_reference(model, chunks)
+    incr = _incremental(model, chunks)
+
+    assert full.shape == incr.shape
+    torch.testing.assert_close(incr, full, atol=ATOL, rtol=RTOL)
+
+
+@pytest.mark.parametrize("config", CONFIGS)
+def test_incremental_hift_matches_streaming_reference_within_tolerance(config):
+    """Windowed output matches the full-cumulative reference within float32 rounding.
+
+    The load-bearing correctness test: 400 mel frames far exceeds the 64-frame
+    window, so truncation genuinely exercises the phase/noise carry. Float64
+    agrees to ~1e-16, confirming the residual here is rounding, not truncation.
+    """
+    hift = _make_hift(config)
+    model = _make_model(hift, window_len=64)  # the real _hift_window_len
+    chunks = _chunks(total_mel=LONG_TOTAL_MEL)
+
+    trim = int(hift.f0_predictor.condnet[0].causal_padding)
+    assert LONG_TOTAL_MEL > 64 + trim + CHUNK_LEN  # window must actually truncate
+
+    full = _full_reference(model, chunks)
+    incr = _incremental(model, chunks)
+
+    assert full.shape == incr.shape, f"{full.shape} vs {incr.shape}"
+    torch.testing.assert_close(incr, full, atol=ATOL, rtol=RTOL)
+
+    max_dev = (full - incr).abs().max().item()  # stays far below one PCM16 step
+    assert max_dev < PCM16_LSB / 10, f"max deviation {max_dev:.3e} vs PCM16 LSB {PCM16_LSB:.3e}"
+
+
+def test_incremental_hift_window_is_bounded():
+    """Per-chunk window must not grow with cumulative mel history."""
+    hift = _make_hift()
+    model = _make_model(hift, window_len=32)
+    chunks = _chunks(total_mel=200, chunk_len=24)
+
+    trim = int(hift.f0_predictor.condnet[0].causal_padding)
+    f0_margin = int(hift.f0_predictor.left_context_frames)
+
+    cache = None
+    window_sizes = []
+    for chunk in chunks:
+        _, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False)
+        window_sizes.append(int(cache["mel"].shape[-1]))
+
+    assert max(window_sizes) <= 32 + trim + f0_margin + 24  # bounded regardless of utterance length
+    steady = max(window_sizes)
+    assert window_sizes.count(steady) >= 2  # steady-state value repeats once history exceeds window
+
+
+def test_incremental_hift_finalize_releases_tail():
+    """Finalize must emit the released look-right tail without crashing."""
+    hift = _make_hift()
+    model = _make_model(hift, window_len=32)
+    chunks = _chunks(total_mel=72, chunk_len=24)
+
+    cache = None
+    emitted = []
+    for i, chunk in enumerate(chunks):
+        finalize = i == len(chunks) - 1
+        speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize)
+        emitted.append(speech.reshape(speech.shape[0], -1))
+        if finalize:
+            assert cache is None
+
+    full = torch.cat(emitted, dim=-1)
+    assert full.shape[-1] > 0
+    assert torch.isfinite(full).all()
+
+
+def test_incremental_hift_emits_full_audio_length():
+    """Streamed output covers the full mel history at samples-per-mel resolution.
+
+    Guards against emission arithmetic dropping an upsample stage, which
+    silently truncates audio (e.g. 1/3 of expected length with [8,5,3]).
+    """
+    hift = _make_hift()
+    model = _make_model(hift, window_len=32)
+    chunks = _chunks(total_mel=72, chunk_len=24)
+
+    cache = None
+    emitted = []
+    for i, chunk in enumerate(chunks):
+        finalize = i == len(chunks) - 1
+        speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize)
+        emitted.append(speech.reshape(speech.shape[0], -1))
+        if finalize:
+            assert cache is None
+
+    full = torch.cat(emitted, dim=-1)
+    # total_mel * SPM, +slack for finalize's look-right tail, -slack for a dropped upsample stage
+    total_mel = sum(c.shape[-1] for c in chunks)
+    assert full.shape[-1] > total_mel * SPM * 0.8, f"{full.shape[-1]} vs {total_mel * SPM}"
+    assert full.shape[-1] < total_mel * SPM * 1.5, f"{full.shape[-1]} vs {total_mel * SPM}"
+
+
+@pytest.mark.parametrize("config", CONFIGS)
+def test_finalize_matches_single_pass_length_exactly(config):
+    """Streamed total (chunked, last=finalize) must exactly match a single
+    non-streaming finalize=True call over the whole mel -- not just approximately.
+
+    The finalize release must include `trim` (the F0 predictor's own held-back
+    frames), or the last 3 mel frames (1,440 samples / 60ms at 24kHz) of every
+    stream's audio go missing.
+    """
+    hift = _make_hift(config)
+    chunks = _chunks(total_mel=200, chunk_len=24)
+
+    torch.manual_seed(0)
+    full_mel = torch.cat(chunks, dim=-1)
+    ground_truth, _, _ = hift.inference(speech_feat=full_mel, finalize=True)
+    ground_truth = ground_truth.reshape(ground_truth.shape[0], -1)
+
+    hift2 = _make_hift(config)
+    model = _make_model(hift2, window_len=64)
+    torch.manual_seed(0)
+    cache = None
+    emitted = []
+    for i, chunk in enumerate(chunks):
+        finalize = i == len(chunks) - 1
+        speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize)
+        emitted.append(speech.reshape(speech.shape[0], -1))
+    streamed = torch.cat(emitted, dim=-1)
+
+    assert streamed.shape[-1] == ground_truth.shape[-1], (
+        f"{streamed.shape[-1]} vs {ground_truth.shape[-1]} (diff={ground_truth.shape[-1] - streamed.shape[-1]} samples)"
+    )
+    torch.testing.assert_close(streamed, ground_truth, atol=ATOL, rtol=RTOL)
+
+
+@pytest.mark.parametrize("config", CONFIGS)
+def test_incremental_hift_matches_reference_for_voiced_f0(config):
+    """Forces F0 above nsf_voiced_threshold (untrained weights never trigger voiced,
+    so no other test here does) and checks windowed output against the full-
+    cumulative reference, which exercises condnet's left-causal convs and
+    SineGen/SineGen2's phase integration under genuinely voiced harmonics.
+
+    Tolerance is looser than ATOL/RTOL: summing this sustained constant-F0 phase
+    in per-window chunks vs. one recompute reassociates float32 adds differently;
+    real speech's varying/unvoiced F0 keeps this far smaller in practice.
+    """
+    hift = _make_hift(config)
+    with torch.no_grad():
+        hift.f0_predictor.classifier.bias.fill_(150.0)  # clearly voiced (threshold is 10)
+    chunks = _chunks(total_mel=200, chunk_len=24)
+
+    full = _full_reference(_make_model(hift, window_len=400), chunks)
+    incr = _incremental(_make_model(hift, window_len=64), chunks)
+
+    assert full.shape == incr.shape
+    torch.testing.assert_close(incr, full, atol=2e-4, rtol=2e-4)
+
+
+@pytest.mark.parametrize("config", CONFIGS)
+@pytest.mark.parametrize("first_len,step_len", [(10, 6), (10, 3), (20, 3), (10, 1)])
+def test_incremental_hift_matches_reference_for_voiced_f0_small_chunks(config, first_len, step_len):
+    """Same as test_incremental_hift_matches_reference_for_voiced_f0, but with small,
+    sub-receptive-field chunk sizes (down to 1 mel frame) instead of the uniform
+    24-frame chunks used elsewhere. cache_state must carry the F0 predictor's own
+    left receptive field across calls regardless of how small a single chunk is,
+    or the next call's margin comes up short and substitutes zero-padding for real
+    history -- corrupting SineGen's phase integration the same way as an absent
+    margin entirely.
+    """
+    hift = _make_hift(config)
+    with torch.no_grad():
+        hift.f0_predictor.classifier.bias.fill_(150.0)  # clearly voiced (threshold is 10)
+    chunks = _variable_chunks(total_mel=200, first_len=first_len, step_len=step_len)
+
+    full = _full_reference(_make_model(hift, window_len=400), chunks)
+    incr = _incremental(_make_model(hift, window_len=64), chunks)
+
+    assert full.shape == incr.shape
+    torch.testing.assert_close(incr, full, atol=2e-4, rtol=2e-4)
+
+
+def test_wrapped_slice_wraps_position_deterministically():
+    """SineGen/SineGen2/SourceModuleHnNSF's noise buffers are a fixed 300s at
+    24kHz; a stream running past that must keep repeating the realization
+    (position-deterministic) rather than the slice coming back short and
+    failing to broadcast against the rest of that call's tensors.
+    """
+    buf = torch.arange(10).reshape(1, 10, 1)
+
+    assert _wrapped_slice(buf, 2, 3).flatten().tolist() == [2, 3, 4]
+    assert _wrapped_slice(buf, 8, 4).flatten().tolist() == [8, 9, 0, 1]
+    assert _wrapped_slice(buf, 12, 3).flatten().tolist() == [2, 3, 4]  # offset past the end
+    assert _wrapped_slice(buf, 5, 13).flatten().tolist() == [5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7]
+
+
+def test_incremental_hift_empty_middle_chunk_emits_nothing():
+    """An empty middle chunk (reachable from forward_streaming_batch when a
+    row's post-trim feat comes out empty near the lookahead boundary) must
+    emit zero samples and leave cache_state's phase_acc unchanged, not
+    re-emit the re-decoded overlap window as if it were new audio.
+    """
+    hift = _make_hift(REAL_CONFIG)
+    model = _make_model(hift, window_len=64)
+    chunks = _chunks(total_mel=96, chunk_len=24)
+    empty = chunks[0][:, :, :0]
+
+    cache = None
+    for chunk in chunks[:2]:
+        _, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False)
+    assert cache is not None
+    phase_acc_before = cache["phase_acc"]
+
+    speech, cache = model._stream_hift_from_feat(empty, cache_state=cache, finalize=False)
+
+    assert speech.shape[-1] == 0
+    if phase_acc_before is None:
+        assert cache["phase_acc"] is None
+    else:
+        torch.testing.assert_close(cache["phase_acc"], phase_acc_before)
+
+
+@pytest.mark.parametrize("config", CONFIGS)
+@pytest.mark.parametrize("mel_frames", [1, 2, 3, 4])
+def test_incremental_hift_short_finalize_chunk_does_not_crash(config, mel_frames):
+    """A one-token (or otherwise sub-trim) utterance that finalizes on its very
+    first chunk must not crash. On finalize the F0 predictor consumes the whole
+    window (it does not hold back `trim` frames), so the phase-carry index must
+    not be inflated by `trim` -- it would point past the end of a short window.
+    """
+    hift = _make_hift(config)
+    model = _make_model(hift, window_len=32)
+    chunk = torch.randn(1, 80, mel_frames)
+
+    speech, cache = model._stream_hift_from_feat(chunk, cache_state=None, finalize=True)
+
+    assert cache is None
+    assert speech.shape[0] == 1
+    assert speech.shape[-1] > 0
+    assert torch.isfinite(speech).all()

tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py

@@ -177,6 +177,8 @@ class _MiniDiT(nn.Module):
         cnn_cache_buffer: torch.Tensor | None = None,
         att_cache_buffer: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        assert cnn_cache_buffer is not None
+        assert att_cache_buffer is not None
         x = x.transpose(1, 2)
         x = self.in_proj(x)
         for b_idx in range(len(self.blocks)):

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py

@@ -98,6 +98,46 @@ class ResBlock(torch.nn.Module):
             remove_weight_norm(self.convs2[idx])
 
 
+def _carry_phase_at_boundary(
+    cum_total: torch.Tensor,
+    phase_acc: torch.Tensor | None,
+    next_overlap: int,
+    *,
+    dim: int,
+    trim: int = 0,
+) -> torch.Tensor | None:
+    """Accumulated phase (mod 1) carried into the next streaming window.
+
+    Shared by SineGen and SineGen2 so harmonic phase stays continuous across
+    chunk boundaries when HiFT runs over a bounded mel window. ``cum_total``
+    must already include ``phase_acc``; falls back to it unchanged if
+    ``next_overlap`` (+ ``trim``) covers the whole window.
+    """
+    length = cum_total.shape[dim]
+    if next_overlap >= length + trim:
+        return phase_acc
+    idx = length + trim - next_overlap - 1
+    return cum_total.narrow(dim, idx, 1) % 1
+
+
+def _wrapped_slice(buf: torch.Tensor, offset: int, length: int, dim: int = 1) -> torch.Tensor:
+    """``length`` elements of ``buf`` along ``dim`` starting at ``offset``,
+    wrapping (mod ``buf.shape[dim]``) instead of running past the end and
+    coming back short. Shared by SineGen/SineGen2/SourceModuleHnNSF's
+    position-indexed noise buffers, now that streams can outlast their 300s.
+    """
+    n = buf.shape[dim]
+    start = offset % n
+    parts = []
+    remaining = length
+    while remaining > 0:
+        take = min(remaining, n - start)
+        parts.append(buf.narrow(dim, start, take))
+        remaining -= take
+        start = 0
+    return parts[0] if len(parts) == 1 else torch.cat(parts, dim=dim)
+
+
 class SineGen(torch.nn.Module):
     """Definition of sine generator
     SineGen(samp_rate, harmonic_num = 0,
@@ -114,13 +154,24 @@ class SineGen(torch.nn.Module):
         segment is always sin(np.pi) or cos(0)
     """
 
-    def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0):
+    def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0, causal=False):
         super().__init__()
         self.sine_amp = sine_amp
         self.noise_std = noise_std
         self.harmonic_num = harmonic_num
         self.sampling_rate = samp_rate
         self.voiced_threshold = voiced_threshold
+        self.causal = causal
+        # Fixed per model instance like SineGen2's rand_ini, not redrawn per
+        # call, or each harmonic jumps to a new random phase every window.
+        self.phase_vec = Uniform(low=-np.pi, high=np.pi).sample(sample_shape=(1, harmonic_num + 1, 1))
+        self.phase_vec[:, 0, :] = 0
+        if causal:
+            self.register_buffer(
+                "noise_buf",
+                torch.randn(1, 300 * 24000, harmonic_num + 1),
+                persistent=False,
+            )
 
     def _f02uv(self, f0):
         # generate uv signal
@@ -128,22 +179,32 @@ class SineGen(torch.nn.Module):
         return uv
 
     @torch.no_grad()
-    def forward(self, f0):
+    def forward(self, f0, phase_acc=None, next_overlap=0, trim=0, noise_offset=0):
         """sine_tensor, uv = forward(f0)
         input F0: tensor(batchsize=1, dim=1, length)
                   f0 for unvoiced steps should be 0
         output sine_tensor: tensor(batchsize=1, length, dim)
         output uv: tensor(batchsize=1, length, 1)
+
+        ``next_overlap`` and ``trim`` are sample-indexed here (``f0`` arrives
+        post-``f0_upsamp``), unlike SineGen2 which downsamples to mel rate
+        before its own cumsum -- callers must scale both by samples-per-mel.
         """
         f0 = f0.transpose(1, 2)
         F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device)
         for i in range(self.harmonic_num + 1):
             F_mat[:, i : i + 1, :] = f0 * (i + 1) / self.sampling_rate
 
-        theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1)
-        u_dist = Uniform(low=-np.pi, high=np.pi)
-        phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device)
-        phase_vec[:, 0, :] = 0
+        cum = torch.cumsum(F_mat, dim=-1)
+        cum_total = cum if phase_acc is None else phase_acc + cum
+        theta_mat = 2 * np.pi * (cum_total % 1)
+        new_phase_acc = _carry_phase_at_boundary(cum_total, phase_acc, next_overlap, dim=-1, trim=trim)
+        if self.causal:
+            phase_vec = self.phase_vec.to(F_mat.device)
+        else:
+            u_dist = Uniform(low=-np.pi, high=np.pi)
+            phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device)
+            phase_vec[:, 0, :] = 0
 
         # generate sine waveforms
         sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec)
@@ -155,11 +216,16 @@ class SineGen(torch.nn.Module):
         #        std = self.sine_amp/3 -> max value ~ self.sine_amp
         # .       for voiced regions is self.noise_std
         noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
-        noise = noise_amp * torch.randn_like(sine_waves)
+        if self.training is False and self.causal is True:
+            noise = noise_amp * _wrapped_slice(self.noise_buf, noise_offset, sine_waves.shape[-1]).transpose(1, 2)
+        else:
+            noise = noise_amp * torch.randn_like(sine_waves)
 
         # first: set the unvoiced part to 0 by uv
         # then: additive noise
         sine_waves = sine_waves * uv + noise
+        if self.causal:
+            return sine_waves.transpose(1, 2), uv.transpose(1, 2), noise, new_phase_acc
         return sine_waves.transpose(1, 2), uv.transpose(1, 2), noise
 
 
@@ -215,7 +281,7 @@ class SineGen2(torch.nn.Module):
         uv = (f0 > self.voiced_threshold).type(torch.float32)
         return uv
 
-    def _f02sine(self, f0_values):
+    def _f02sine(self, f0_values, phase_acc=None, next_overlap=0, trim=0):
         """f0_values: (batchsize, length, dim)
         where dim indicates fundamental tone and overtones
         """
@@ -223,9 +289,12 @@ class SineGen2(torch.nn.Module):
         # because 2 * np.pi * n doesn't affect phase
         rad_values = (f0_values / self.sampling_rate) % 1
 
-        # initial phase noise (no noise for fundamental component)
+        # initial phase noise (no noise for fundamental component). Only applied
+        # on the genuine first chunk of a stream (see rationale above the class);
+        # phase_acc is None exactly when there is no carried-forward phase yet.
         if self.training is False and self.causal is True:
-            rad_values[:, 0, :] = rad_values[:, 0, :] + self.rand_ini.to(rad_values.device)
+            if phase_acc is None:
+                rad_values[:, 0, :] = rad_values[:, 0, :] + self.rand_ini.to(rad_values.device)
         else:
             rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device)
             rand_ini[:, 0] = 0
@@ -237,7 +306,10 @@ class SineGen2(torch.nn.Module):
                 rad_values.transpose(1, 2), scale_factor=1 / self.upsample_scale, mode="linear"
             ).transpose(1, 2)
 
-            phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
+            cum = torch.cumsum(rad_values, dim=1)
+            phase = cum if phase_acc is None else cum + phase_acc
+            new_phase_acc = _carry_phase_at_boundary(phase, phase_acc, next_overlap, dim=1, trim=trim)
+            phase = phase * 2 * np.pi
             phase = torch.nn.functional.interpolate(
                 phase.transpose(1, 2) * self.upsample_scale,
                 scale_factor=self.upsample_scale,
@@ -272,9 +344,10 @@ class SineGen2(torch.nn.Module):
 
             # get the sines
             sines = torch.cos(i_phase * 2 * np.pi)
-        return sines
+            new_phase_acc = None
+        return (sines, new_phase_acc) if self.causal else sines
 
-    def forward(self, f0):
+    def forward(self, f0, phase_acc=None, next_overlap=0, trim=0, noise_offset=0):
         """sine_tensor, uv = forward(f0)
         input F0: tensor(batchsize=1, length, dim=1)
                   f0 for unvoiced steps should be 0
@@ -284,8 +357,12 @@ class SineGen2(torch.nn.Module):
         # fundamental component
         fn = torch.multiply(f0, self.harmonic_ids)
 
-        # generate sine waveforms
-        sine_waves = self._f02sine(fn) * self.sine_amp
+        if self.causal:
+            sine_waves, new_phase_acc = self._f02sine(fn, phase_acc, next_overlap, trim)
+        else:
+            sine_waves = self._f02sine(fn)
+            new_phase_acc = None
+        sine_waves = sine_waves * self.sine_amp
 
         # generate uv signal
         uv = self._f02uv(f0)
@@ -295,14 +372,15 @@ class SineGen2(torch.nn.Module):
         # .       for voiced regions is self.noise_std
         noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
         if self.training is False and self.causal is True:
-            noise = noise_amp * self.sine_waves[:, : sine_waves.shape[1]].to(sine_waves.device)
+            # indexed by absolute sample position, not window-relative position
+            noise = noise_amp * _wrapped_slice(self.sine_waves, noise_offset, sine_waves.shape[1]).to(sine_waves.device)
         else:
             noise = noise_amp * torch.randn_like(sine_waves)
 
         # first: set the unvoiced part to 0 by uv
         # then: additive noise
         sine_waves = sine_waves * uv + noise
-        return sine_waves, uv, noise
+        return (sine_waves, uv, noise, new_phase_acc) if self.causal else (sine_waves, uv, noise)
 
 
 class SourceModuleHnNSF(torch.nn.Module):
@@ -338,10 +416,13 @@ class SourceModuleHnNSF(torch.nn.Module):
 
         self.sine_amp = sine_amp
         self.noise_std = add_noise_std
+        self.upsample_scale = upsample_scale
 
         # to produce sine waveforms
         if sinegen_type == "1":
-            self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold)
+            self.l_sin_gen = SineGen(
+                sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold, causal=causal
+            )
         else:
             self.l_sin_gen = SineGen2(
                 sampling_rate, upsample_scale, harmonic_num, sine_amp, add_noise_std, voiced_threshold, causal=causal
@@ -354,24 +435,38 @@ class SourceModuleHnNSF(torch.nn.Module):
         if causal is True:
             self.uv = torch.rand(1, 300 * 24000, 1)
 
-    def forward(self, x):
+    def forward(self, x, phase_acc=None, uv_offset=0, next_overlap=0, trim=0):
         """
         Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
         F0_sampled (batchsize, length, 1)
         Sine_source (batchsize, length, 1)
         noise_source (batchsize, length 1)
         """
-        # source for harmonic branch
         with torch.no_grad():
-            sine_wavs, uv, _ = self.l_sin_gen(x)
+            if not self.causal:
+                sine_wavs, uv, _ = self.l_sin_gen(x)
+                new_phase_acc = None
+            elif isinstance(self.l_sin_gen, SineGen):
+                # next_overlap/trim arrive in mel frames; SineGen needs samples.
+                sine_wavs, uv, _, new_phase_acc = self.l_sin_gen(
+                    x,
+                    phase_acc,
+                    next_overlap * self.upsample_scale,
+                    trim * self.upsample_scale,
+                    noise_offset=uv_offset,
+                )
+            else:
+                sine_wavs, uv, _, new_phase_acc = self.l_sin_gen(
+                    x, phase_acc, next_overlap, trim, noise_offset=uv_offset
+                )
         sine_merge = self.l_tanh(self.l_linear(sine_wavs))
 
         # source for noise branch, in the same shape as uv
         if self.training is False and self.causal is True:
-            noise = self.uv[:, : uv.shape[1]] * self.sine_amp / 3
+            noise = _wrapped_slice(self.uv, uv_offset, uv.shape[1]) * self.sine_amp / 3
         else:
             noise = torch.randn_like(uv) * self.sine_amp / 3
-        return sine_merge, noise, uv
+        return sine_merge, noise, uv, new_phase_acc
 
 
 class HiFTGenerator(nn.Module):
@@ -606,7 +701,7 @@ class HiFTGenerator(nn.Module):
         f0 = self.f0_predictor(speech_feat)
         # f0->source
         s = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
-        s, _, _ = self.m_source(s)
+        s, _, _, _ = self.m_source(s)
         s = s.transpose(1, 2)
         # mel+source->speech
         generated_speech = self.decode(x=speech_feat, s=s)
@@ -622,7 +717,7 @@ class HiFTGenerator(nn.Module):
         f0 = self.f0_predictor(speech_feat)
         # f0->source
         s = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
-        s, _, _ = self.m_source(s)
+        s, _, _, _ = self.m_source(s)
         s = s.transpose(1, 2)
         # use cache_source to avoid glitch
         if cache_source.shape[2] != 0:
@@ -785,14 +880,28 @@ class CausalHiFTGenerator(HiFTGenerator):
         return x
 
     @torch.inference_mode()
-    def inference(self, speech_feat: torch.Tensor, finalize: bool = True) -> torch.Tensor:
+    def inference(
+        self,
+        speech_feat: torch.Tensor,
+        finalize: bool = True,
+        phase_acc: torch.Tensor | None = None,
+        uv_offset: int = 0,
+        next_overlap: int = 0,
+        trim: int = 0,
+        f0_margin: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
         # mel->f0 NOTE f0_predictor precision is crucial for causal inference, move
         # self.f0_predictor to cpu if necessary
         self.f0_predictor.to("cpu")
         f0 = self.f0_predictor(speech_feat.cpu(), finalize=finalize).to(speech_feat)
+        if f0_margin > 0:
+            f0 = f0[:, f0_margin:]
+            speech_feat = speech_feat[:, :, f0_margin:]
         # f0->source
         s = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
-        s, _, _ = self.m_source(s)
+        s, _, _, new_phase_acc = self.m_source(
+            s, phase_acc=phase_acc, uv_offset=uv_offset, next_overlap=next_overlap, trim=trim
+        )
         s = s.transpose(1, 2)
         if finalize is True:
             generated_speech = self.decode(x=speech_feat, s=s, finalize=finalize)
@@ -800,7 +909,7 @@ class CausalHiFTGenerator(HiFTGenerator):
             generated_speech = self.decode(
                 x=speech_feat[:, :, : -self.f0_predictor.condnet[0].causal_padding], s=s, finalize=finalize
             )
-        return generated_speech, s
+        return generated_speech, s, new_phase_acc
 
 
 class CausalConv1dUpsample(torch.nn.Conv1d):
@@ -954,6 +1063,11 @@ class CausalConvRNNF0Predictor(nn.Module):
             nn.ELU(),
         )
         self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class)
+        self.left_context_frames = sum(
+            layer.causal_padding
+            for layer in self.condnet
+            if isinstance(layer, CausalConv1d) and layer.causal_type == "left"
+        )
 
     def forward(self, x: torch.Tensor, finalize: bool = True) -> torch.Tensor:
         if finalize is True:

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py

@@ -9,10 +9,8 @@ This module contains the code2wav (token-to-waveform) stage which uses:
 3. HiFiGAN vocoder for waveform synthesis
 """
 
-from __future__ import annotations
-
 from collections import Counter
-from typing import cast
+from typing import TypedDict, cast
 
 import numpy as np
 import torch
@@ -39,6 +37,21 @@ from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 logger = init_logger(__name__)
 
 
+class StreamingFlowItem(TypedDict, total=False):
+    """One entry of the batched-streaming item list passed to forward_streaming_batch."""
+
+    index: int
+    req_id: str | None
+    stream_finished: bool
+    token: torch.Tensor
+    prompt_token: torch.Tensor
+    prompt_feat: torch.Tensor
+    embedding: torch.Tensor
+    cache_state: dict[str, torch.Tensor] | None
+    token_offset_tokens: int
+    finalize: bool
+
+
 class CosyVoice3Code2Wav(nn.Module):
     """CosyVoice3 Code2Wav stage for token-to-waveform conversion.
 
@@ -119,6 +132,10 @@ class CosyVoice3Code2Wav(nn.Module):
         self.mel_cache_len = 20
         self.source_cache_len = int(self.mel_cache_len * 256)
         self.speech_window = np.hamming(2 * self.source_cache_len)
+        # Must cover decode()'s own causal receptive field, not just the F0
+        # margin; window_len=48 already passes
+        # test_incremental_hift_bounded_window_is_close, so 64 has headroom.
+        self._hift_window_len = 64
 
     @property
     def input_frame_rate(self) -> int:
@@ -226,40 +243,79 @@ class CosyVoice3Code2Wav(nn.Module):
         chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype)
 
         cached_mel = None if not cache_state else cache_state.get("mel")
-        speech_offset_obj = None if not cache_state else cache_state.get("speech_offset")
-        try:
-            speech_offset = int(speech_offset_obj) if speech_offset_obj is not None else 0
-        except (TypeError, ValueError):
-            speech_offset = 0
-
+        cached_offset = 0 if not cache_state else int(cache_state.get("mel_offset", 0))
+        window_len = self._hift_window_len
+        samples_per_mel = int(np.prod(self.hift.upsample_rates) * self.hift.istft_params["hop_len"])
+        trim = int(self.hift.f0_predictor.condnet[0].causal_padding)
+        f0_margin_frames = int(self.hift.f0_predictor.left_context_frames)
+        f0_margin = 0
         if isinstance(cached_mel, torch.Tensor) and cached_mel.numel() > 0:
             cached_mel = cached_mel.to(device=chunk_mel.device, dtype=chunk_mel.dtype)
-            tts_mel = torch.cat([cached_mel, chunk_mel], dim=-1) if chunk_mel.numel() > 0 else cached_mel
+            total_hist = cached_mel.shape[-1]
+            overlap = min(window_len + trim, total_hist)
+            window_mel = torch.cat([cached_mel[..., -overlap:], chunk_mel], dim=-1)
+            window_offset = cached_offset + total_hist - overlap
+            uv_offset = window_offset * samples_per_mel
+            f0_margin = min(f0_margin_frames, total_hist - overlap)
+            f0_input_mel = (
+                torch.cat(
+                    [cached_mel[..., total_hist - overlap - f0_margin : total_hist - overlap], window_mel], dim=-1
+                )
+                if f0_margin > 0
+                else window_mel
+            )
         else:
-            tts_mel = chunk_mel
+            window_mel = chunk_mel
+            f0_input_mel = chunk_mel
+            window_offset = 0
+            uv_offset = 0
 
-        if tts_mel.shape[-1] == 0:
+        phase_acc = None if not cache_state else cache_state.get("phase_acc")
+        if phase_acc is not None:
+            phase_acc = phase_acc.to(device=chunk_mel.device, dtype=chunk_mel.dtype)
+
+        if window_mel.shape[-1] == 0:
             tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype)
+            new_phase_acc = phase_acc
         else:
-            tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
+            next_overlap = min(window_len + trim, window_mel.shape[-1])
+            # On finalize the F0 predictor consumes the whole window (it does not
+            # hold back `trim` frames), so trim must not inflate the phase-carry
+            # index -- it would point past the end of a short window.
+            carry_trim = 0 if finalize else trim
+            tts_speech, _, new_phase_acc = self.hift.inference(
+                speech_feat=f0_input_mel,
+                finalize=finalize,
+                phase_acc=phase_acc,
+                uv_offset=uv_offset,
+                next_overlap=next_overlap,
+                trim=carry_trim,
+                f0_margin=f0_margin,
+            )
 
         tts_speech = tts_speech.reshape(tts_speech.shape[0], -1)
-        speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1])))
-        emitted_speech = tts_speech[:, speech_offset:]
+        new_samples = chunk_mel.shape[-1] * samples_per_mel
+        if finalize:
+            frames_withheld_per_chunk = trim + self.hift.conv_pre_look_right + 1
+            released = frames_withheld_per_chunk * samples_per_mel
+            emitted_speech = tts_speech[:, -(new_samples + released) :]
+        else:
+            emitted_speech = tts_speech[:, -new_samples:] if new_samples > 0 else tts_speech[:, :0]
 
         if finalize:
             return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), None
 
         new_state = {
-            "mel": tts_mel.detach().cpu().contiguous(),
-            "speech_offset": int(tts_speech.shape[-1]),
+            "mel": f0_input_mel.detach().cpu().contiguous(),
+            "mel_offset": window_offset - f0_margin,
+            "phase_acc": new_phase_acc.detach().cpu().contiguous() if new_phase_acc is not None else None,
         }
         return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state
 
     @torch.inference_mode()
     def forward_streaming_batch(
         self,
-        items: list[dict[str, object]],
+        items: list[StreamingFlowItem],
         *,
         n_timesteps: int = 10,
     ) -> list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]]:
@@ -270,20 +326,16 @@ class CosyVoice3Code2Wav(nn.Module):
         group and passed to the flow as per-row token lengths.
         """
         results: list[tuple[torch.Tensor, dict[str, torch.Tensor] | None] | None] = [None] * len(items)
-        groups: dict[tuple[int, int, int, bool], list[tuple[int, dict[str, object]]]] = {}
+        groups: dict[tuple[int, int, int, bool], list[tuple[int, StreamingFlowItem]]] = {}
         for index, item in enumerate(items):
-            token = item["token"]
-            prompt_token = item["prompt_token"]
-            prompt_feat = item["prompt_feat"]
-            embedding = item["embedding"]
-            assert isinstance(token, torch.Tensor)
-            assert isinstance(prompt_token, torch.Tensor)
-            assert isinstance(prompt_feat, torch.Tensor)
-            assert isinstance(embedding, torch.Tensor)
+            assert isinstance(item["token"], torch.Tensor)
+            assert isinstance(item["prompt_token"], torch.Tensor)
+            assert isinstance(item["prompt_feat"], torch.Tensor)
+            assert isinstance(item["embedding"], torch.Tensor)
             key = (
-                int(prompt_token.shape[1]),
-                int(prompt_feat.shape[1]),
-                int(embedding.shape[1]),
+                int(item["prompt_token"].shape[1]),
+                int(item["prompt_feat"].shape[1]),
+                int(item["embedding"].shape[1]),
                 bool(item.get("finalize", False)),
             )
             groups.setdefault(key, []).append((index, item))
@@ -305,11 +357,11 @@ class CosyVoice3Code2Wav(nn.Module):
             if len(group) == 1:
                 index, item = group[0]
                 result = self.forward_streaming(
-                    token=item["token"],  # type: ignore[arg-type]
-                    prompt_token=item["prompt_token"],  # type: ignore[arg-type]
-                    prompt_feat=item["prompt_feat"],  # type: ignore[arg-type]
-                    embedding=item["embedding"],  # type: ignore[arg-type]
-                    cache_state=item.get("cache_state"),  # type: ignore[arg-type]
+                    token=item["token"],
+                    prompt_token=item["prompt_token"],
+                    prompt_feat=item["prompt_feat"],
+                    embedding=item["embedding"],
+                    cache_state=item.get("cache_state"),
                     n_timesteps=n_timesteps,
                     token_offset_tokens=int(item.get("token_offset_tokens", 0)),
                     finalize=bool(item.get("finalize", False)),
@@ -318,15 +370,13 @@ class CosyVoice3Code2Wav(nn.Module):
                 continue
 
             token_tensors = [item["token"] for _, item in group]
-            assert all(isinstance(token, torch.Tensor) for token in token_tensors)
             token_lens = torch.tensor(
-                [int(token.shape[1]) for token in token_tensors],  # type: ignore[union-attr]
+                [int(token.shape[1]) for token in token_tensors],
                 dtype=torch.int32,
             )
             max_token_len = int(token_lens.max().item())
             padded_tokens = []
             for token in token_tensors:
-                assert isinstance(token, torch.Tensor)
                 if int(token.shape[1]) == max_token_len:
                     padded_tokens.append(token)
                 else:
@@ -337,9 +387,9 @@ class CosyVoice3Code2Wav(nn.Module):
                     )
                     padded_tokens.append(torch.cat([token, pad], dim=1))
             tokens = torch.cat(padded_tokens, dim=0)
-            prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0)  # type: ignore[list-item]
-            prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0)  # type: ignore[list-item]
-            embeddings = torch.cat([item["embedding"] for _, item in group], dim=0)  # type: ignore[list-item]
+            prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0)
+            prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0)
+            embeddings = torch.cat([item["embedding"] for _, item in group], dim=0)
             prompt_token_lens = torch.full((len(group),), prompt_tokens.shape[1], dtype=torch.int32)
             prompt_feat_lens = torch.full((len(group),), prompt_feats.shape[1], dtype=torch.int32)
             finalize = bool(group[0][1].get("finalize", False))
@@ -370,7 +420,7 @@ class CosyVoice3Code2Wav(nn.Module):
                     row_feat = row_feat[:, :, trim_mel:]
                 results[index] = self._stream_hift_from_feat(
                     row_feat,
-                    cache_state=item.get("cache_state"),  # type: ignore[arg-type]
+                    cache_state=item.get("cache_state"),
                     finalize=finalize,
                 )
 
@@ -443,7 +493,7 @@ class CosyVoice3Code2Wav(nn.Module):
                 dtype=tts_mel.dtype,
             )
         else:
-            tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=True)
+            tts_speech, _, _ = self.hift.inference(speech_feat=tts_mel, finalize=True)
 
         return tts_speech