[arun-elr]

#495 has landed, so this is now rebased on main and the diff is only this PR's commits.
Fixes #599

What does this PR do?

CosyVoice3's flow node runs its 10-step Euler solve eagerly, about 145 ms a call at streamed-hop sizes whatever the size (41 ms replayed; likely launch-bound, not profiled), which every request pays. The flow node now replays the solve as piecewise CUDA graphs the engine captures, one per frame bucket and batch size.

• cosyvoice3/components/flow.py splits the solve into prepare_solve (pads) and a fixed-shape solve_prepared, so nothing host-side is in the captured part; eager output is unchanged.
• cosyvoice3/submodules.py returns one PiecewiseBatchedConfig per frame bucket and replays the smallest bucket and batch size that fit, else runs eagerly; the engine captures (AGENTS.md invariant 1).
• cosyvoice3/config.py: 20 buckets, 100 to 4,200 frames, batch 1 to 8 within 8,192 batch-frames (97 graphs). flow_capture_frame_buckets, flow_capture_batch_sizes and flow_capture_batch_frames can be set in model_kwargs; flow_capture_frame_buckets: [] captures nothing and every solve runs eagerly, as on [Model] Add Fun-CosyVoice3: zero-shot voice-cloning TTS with streamed 24 kHz audio #495.

One H200, tts_speech_bench, #495 vs this PR: 8.7 to 14.9 audio-s/s at c1 and 18.9 to 34.7 at c32, first audio p50 306 to 172 ms at c1. Capture is on by default and costs 4.8 GiB of GPU memory and about 72 s of startup (94 s vs 21 s to /health); model_kwargs: {flow_capture_frame_buckets: []} turns it off for a deployment that cares more about either.

How was it tested?

Rebased on main 18b0859 (ee60508): ruff check . clean; CPU Core (ci.yml paths w/o the TP lockstep sim, -n 4): 8,632 passed, 935 skipped, 40 failed, the same 40 as on main 18b0859 in this venv (37 Rust-runtime cases that need a newer mstar_rust build than this venv's, 3 need ftfy). The only change in the rebase: the CPU capture stand-in in test_cosyvoice3_flow_capture.py takes the node_name that #810 now passes to capture_into_graph.

On d5f118f (before the rebase, on #495's head 1d346b8, which is on main a3d73a0): ruff check . clean. CPU Core with ci.yml's paths incl. the TP lockstep sim (pytest -n 4 --dist worksteal, shared venv): 6788 passed, 1200 skipped, 6 failed; the same 6 fail on main a3d73a0 in that venv (they need prometheus_client, soundfile or ftfy), so none are this PR's. python -m pytest test/modular/test_cosyvoice3_flow_capture.py test/modular/test_cosyvoice3_model.py test/integration/test_cosyvoice3_real_weights.py: 134 passed, 11 skipped on CPU; on 7e99011, 145 passed on one H200 with the checkpoint, the CUDA capture test included (bf16 and float32). New: test_capture_fields_come_from_model_kwargs and test_a_deployment_can_turn_capture_off (both fail before the model_kwargs change).

Perf A/B on one H200, benchmark.tts_speech_bench, #495 / this PR / #631 rotated a b c c b a a b c with 3 server runs each, then 2 runs with flow_capture_frame_buckets: []: the numbers in the chart, 0 errors in 1,377 requests per arm, ranges over the runs never overlap; capture off matches #495 within spread (8.8 / 13.9 / 16.3 / 18.8 audio-s/s). Table and commands in Performance. Measured on 437bcbf (#495) and 0847206 (this PR); the heads since add main's #807 (chat attachments), the store-error wording in _resolve_ref_audio and a GPU-test fix, none on the measured path.

E2E compare matrix's cosyvoice3 clone case against #495's head (one H200): no regressions, same WER (2.54%, 24/24 identical transcripts) and speaker similarity (0.958), c=1 latency 761 to 413 ms, c=max throughput 11.1 to 22.8 audio-s/s (Performance).

Not tested

A served float32 flow estimator (the GPU test covers its capture); other GPUs; a smaller bucket ladder for small GPUs (the fields are settable now, but no smaller ladder was measured).

Depends on and merge order

Depends on: #495, now on main. Nothing else to wait for.

Design

• Why padding is safe: prepare_solve right-pads a batch to the bucket's frame count. Embeddings past a row's tokens are zeroed, the DiT's convolutions are left-padded (causal), attention masks every key past a row's valid frames (chunk_attention_mask), and row_keep zeroes the attention output of padded query frames. So the frames of a row do not read the padding; in exact arithmetic a padded solve gives each row the frames of its own solve. In floating point, another shape can pick other kernels, which is the same spread the eager path already has when it pads a short row to a longer one in a batch (Correctness).
• What is captured: solve_prepared, the chunk mask, the noise slice and the 10 Euler steps with guidance, at the default step count. It reads nothing on the host. prepare_solve (embedding, look-ahead, prompt mel, speaker projection) stays eager and writes the static inputs.
• Routing: FlowSubmodule._solve takes the smallest bucket that holds the batch's longest row, then the runner pads the batch to the smallest captured batch size. A padding row has zero inputs and valid clamped to 1, so its attention stays finite; its output is dropped. A solve past the largest bucket (4,200 frames, about 84 s of prompt plus speech), a batch size that bucket did not capture, or another n_timesteps runs the same code eagerly.
• Ladder: buckets are whole 25-token attention chunks (a streamed hop ends on one). Batch sizes 1, 2, 3, 4, 6 and 8 are captured for a bucket while batch * frames <= 8192. Below about 3,000 batch-frames an eager solve takes about 145 ms whatever the size, likely bound by its kernel launches (a replay of the same kernels is 41 ms at 250 frames); past that it is bound by compute, and replay saves 1 to 13% (Performance), so finer buckets there would cost startup and memory for little.
• Turning it off: the three capture fields come from model_kwargs (sorted, repeats dropped, positive integers; a frame budget of 0 is allowed). An empty bucket list, or a frame budget of 0, leaves get_piecewise_accelerator_graph_configs empty, so the engine captures nothing and _solve runs every solve eagerly; the served numbers with it off match [Model] Add Fun-CosyVoice3: zero-shot voice-cloning TTS with streamed 24 kHz audio #495 within run-to-run spread (table above). The default stays on.
• Precision: the regions capture under the node's _without_tf32, as the eager solve runs, and the flow node opts out of autocast, so capture and eager run the same precision.
• AGENTS.md invariant 1: nothing under mstar/model/ captures, syncs or manages streams; the node returns PiecewiseBatchedConfigs from get_piecewise_accelerator_graph_configs and calls runner.run, as Kokoro does. The runner's output is a copy, so a mel outlives the next replay.

Correctness

Served, one H200, the same clips as Performance (27 sentences of t2s.txt, --voice zero_shot --seed 7). WER is benchmark.tts_speech_wer (whisper-large-v3-turbo); speaker similarity is the cosine of each clip's CAMPPlus x-vector (the one the flow is conditioned on) to the reference clip's. Measured on d7aa2c0 and 11c30c2 (older heads of #495 and this PR; the solve and routing are unchanged since):

Check	Before (d7aa2c0)	After (11c30c2)
WER, c1 clips	1.95% (22 of 1,131 words)	1.95% (22 of 1,131)
WER, c8 clips	2.12% (24 of 1,131)	2.12% (24 of 1,131)
Speaker similarity, c1, mean (min)	0.7409 (0.6763)	0.7411 (0.6811)
Speaker similarity, c8, mean over the two starts	0.7436, 0.7432	0.7437, 0.7441
c1 clips byte identical across the build's two server starts	27 of 27	27 of 27
c1 streamed vs non-streamed (--extra stream=false)	27 of 27 byte identical	27 of 27 byte identical

Before vs after at c1: 23 of 27 clips are byte identical; the other 4 have the same length (same tokens) and differ in the flow's rounding, as below. On the current heads, with #495's look-ahead change, a review run found 30 of 32 clips byte identical to #495 at the same WER; the differences on the older heads were in short clips' last chunk only (log-mel SNR 26 dB or better). At c8 the batch composition varies between starts, so even before vs before leaves 6 of 27 clips different; WER and speaker similarity are the check there.

Flow solve parity, on 11c30c2, real weights, bf16 estimator, fixed speech tokens and the model's fixed noise:

Comparison	Result
Replay vs the eager solve of the same padded inputs (9 single rows, batches of 3 and 8)	bit identical, 11 of 11
Replay vs the eager solve at the row's own length	bit identical for 7 of 9 rows; the other 2: max abs 0.59 and 1.03, mel SNR 38.0 and 37.8 dB
Replayed batches of 3 and 8 vs each row solved alone	min mel SNR 36.3 and 38.7 dB
Yardstick: eager alone vs eager batched with a longer row (the padding #495 already does)	bit identical for 4 of 9; others max abs 0.39 to 0.85, mel SNR 38.2 to 40.8 dB
Scale: bf16 vs float32 estimator, one row	max abs 2.07, mel SNR 31.8 dB

The eager path itself is unchanged: solve_rows on #495's head and on this branch give bit identical mels (CPU, float32, 3 and 10 steps, padded batches with final and non-final rows). The GPU test test_captured_solve_equals_the_eager_solve_of_the_same_padded_inputs is the in-repo form of the first row (bf16 and float32 estimators, random weights); it passes on one H200 on the head.

Performance

One H200 (torch 2.12.1+cu130, spot VM), one server per run, configs/cosyvoice3.yaml plus model_kwargs.voices_dir holding the CosyVoice Mandarin prompt clip as the zero_shot preset, MSTAR_SHM_ARENA=1, servers started and stopped with .claude/skills/benchmarking/scripts/server.sh. Arms rotated #495 (437bcbf), this PR (0847206), #631 (bc022bd) as a b c c b a a b c, 3 server runs each, then 2 runs of this PR with flow_capture_frame_buckets: []. Each cell is the median over server runs of the harness's median. Inputs: t2s.txt (27 sentences) for c1, the same file 3 times (81 lines) for c8 and 6 times (162 lines) for c16 and c32, so every concurrency gets at least 5x its count in requests.

B="python -m benchmark.tts_speech_bench --engine mstar --url http://127.0.0.1:18800 --model cosyvoice3 --voice zero_shot --seed 7"
$B --sentences t2s.txt    --concurrency 1  --repeats 2 --warmup 3  --out c1.json
$B --sentences t2s_x3.txt --concurrency 8  --repeats 1 --warmup 16 --out c8.json
$B --sentences t2s_x6.txt --concurrency 16 --repeats 1 --warmup 32 --out c16.json
$B --sentences t2s_x6.txt --concurrency 32 --repeats 1 --warmup 64 --out c32.json

c	arm	first audio p50 (range)	p95	RTF	audio-s/s (range)
1	#495	306 (299 to 314)	368	0.122	8.7 (8.5 to 8.9)
1	this PR	172 (171 to 173)	226	0.071	14.9 (14.8 to 14.9)
1	capture off	303 (301 to 305)	371	0.120	8.8
8	#495	1454 (1440 to 1494)	2010	0.597	14.1 (13.7 to 14.3)
8	this PR	710 (677 to 729)	926	0.296	28.6 (28.5 to 29.0)
8	capture off	1507 (1487 to 1526)	2033	0.610	13.9
16	#495	2531 (2466 to 2589)	3706	1.028	16.5 (15.8 to 16.8)
16	this PR	1272 (1267 to 1301)	1843	0.526	32.0 (32.0 to 32.4)
16	capture off	2513 (2475 to 2552)	3666	1.037	16.3
32	#495	7967 (7464 to 7995)	10868	1.785	18.9 (18.3 to 19.5)
32	this PR	4315 (4194 to 4343)	5903	0.972	34.7 (34.2 to 34.7)
32	capture off	7818	11215	1.797	18.8

Total audio per pass matched within 0.4% between #495 and this PR, so audio-s/s is not a length artifact. Startup to /health (server.sh's poll): 21 s on #495, 90 to 94 s with capture, 21 s with it off; GPU memory after startup (nvidia-smi): 6.2 GiB on #495 and with capture off, 11.0 GiB with capture. The c1 gain is larger than in the earlier run below because #495's look-ahead change (cuDNN off for those convolutions) has since landed underneath.

Earlier H200 run (d7aa2c0 vs 11c30c2, before #495's look-ahead change, ABBA, 2 server starts per arm, 3 repeats each): c1 first audio p50 454 to 333 ms, c8 1309 to 953 ms, c16 3218 to 2109 ms, c32 8879 to 6070 ms; throughput 7.3 to 10.1, 11.8 to 17.7, 12.9 to 19.8, 15.5 to 22.8 audio-s/s.

Flow solve microbenchmark, on 11c30c2, real weights, bf16 estimator, TF32 off as the node runs it: the eager 10-step solve (wall time to GPU completion) against a replay of the same solve captured as one CUDA graph, median of 5 (3 for the largest):

Batch	Frames	Eager (ms)	Replay (ms)	Speedup	Replay vs eager
1	250	145	41	3.54x	bit identical
8	250	147	108	1.36x	bit identical
1	500	144	51	2.81x	bit identical
4	500	146	112	1.31x	bit identical
1	1000	146	76	1.90x	bit identical
2	1000	146	121	1.21x	bit identical
8	1000	443	414	1.07x	bit identical
1	2000	157	138	1.13x	bit identical
4	2000	512	483	1.06x	bit identical
1	4200	360	329	1.10x	bit identical
8	4200	2477	2446	1.01x	bit identical

Up to about 3,000 batch-frames the eager solve takes about 145 ms whatever its size, while a replay of the same kernels takes 41 to 138 ms: kernel launches are the likely cost (a hypothesis; not profiled). That is the regime streamed hops live in (prompt plus 25 to 100 settled tokens, a few hundred frames). The scripts behind that table, the parity table and the capture cost are in this comment.

E2E compare matrix (#652's cosyvoice3 clone case, 24 clones, c=1 and c=max, 3 passes), one H200, #495's head as base; startup 22 / 92 s, peak VRAM 15 / 20 GB: no regressions.

python -m benchmark.eval compare --base b0ae44174 --head 7e9901142 --models cosyvoice3 --hosts <h200 host>

Model	Modality	Correct	Errors, base / head	c=1 latency, base to head	c=max throughput, base to head	Quality, base to head	Verdict
cosyvoice3	tts (voice clone, streamed)	yes, WER 2.54% to 2.54%, 24/24 identical transcripts	0 / 0	761 ms to 413 ms (+45.7%)	11.1 to 22.8 audio s/s (+106.6%)	UTMOS 3.98 to 3.98; speaker similarity 0.958 to 0.958	faster

Tests in detail

test/modular/test_cosyvoice3_flow_capture.py (CPU tests drive the real PiecewiseAcceleratorGraphRunner with capture replaced by an eager rerun of the region):

• a solve padded past its rows (by 0, 1, 37 frames and to twice the length) gives every row the frames of its own solve; a bucket shorter than a row is refused;
• the regions: batch sizes within the batch-frame budget and the node's batch cap, static input shapes and dtypes, no region past the flow's noise;
• the node replays the smallest bucket and batch size that hold a batch (including a batch of 3 on the batch-4 graph), and runs eagerly past the largest bucket, for a batch size a bucket did not capture, and for another step count, with the same mels as eager each time;
• a streamed request's two hops through the buckets equal one eager stream, and a replay's mel outlives the next replay;
• the capture fields come from model_kwargs (sorted, deduplicated, bad values refused), and an empty bucket list or a zero frame budget captures no region while the solve gives the same mel as a replay;
• CUDA only: real capture and replay of every region against the eager solve of the same padded inputs (bit identical at a captured batch size, within tolerance when a padding row changes the batch shape), with replays of the other regions sharing the pool in between, for bf16 and float32 estimators. It now passes main's warmup_and_capture its plan (None, every bucket); the old call raised a TypeError on a GPU while CPU CI skipped it.

Resources

State or op	Resource	Notes
20 piecewise regions, 97 graphs, captured at startup	engine graph pool shared by the node's regions	0.88 GiB pool (largest bucket captured first; 1.01 GiB smallest first)
Static inputs and outputs per graph	device memory outside the pool	0.21 GiB in total
All of it	GPU memory	4.8 GiB more after startup on the served H200 (11.0 vs 6.2 GiB); the rest beyond the pool and buffers is outside the caching allocator, consistent with the instantiated graphs themselves (not separately measured)
Startup	server /health	90 to 94 s with capture, 21 s without (flow_capture_frame_buckets: [])

The flow node declares no engine resources, so the regions declare no step. Pool and buffer sizes measured on 11c30c2 with capture_mem.py from the linked comment, which builds the flow node as the server does and captures its regions through PiecewiseAcceleratorGraphRunner.

Checklist

• ruff check . passes
• Added or updated tests / docs where relevant

AI assistance: written with AI coding agent

[arun-elr]

Scripts behind the Correctness and Performance numbers (not committed; run from the repo root with PYTHONPATH=$PWD on one GPU, the checkpoint in the HF cache).

Scripts

solve_bench.py: flow solve microbenchmark, eager vs one captured graph: python solve_bench.py 1,2,4,8 250,500,1000,2000,4200

"""Eager vs CUDA-graph-replayed CosyVoice3 flow solve, per (batch, frames).

Scratch microbenchmark for #599 (not committed). Real flow weights, bf16
estimator, TF32 off as the flow node runs it. Prints host launch time of the
eager solve, its GPU-complete time, the replay time, and capture memory.
"""

import json
import sys
import time

import torch

from mstar.model.chatterbox.loader import resolve_snapshot
from mstar.model.cosyvoice3.components.dit import chunk_attention_mask
from mstar.model.cosyvoice3.components.flow import CosyVoice3Flow
from mstar.model.cosyvoice3.config import CosyVoice3Config

REPO = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"


def load_flow(device):
    cfg = CosyVoice3Config()
    snap = resolve_snapshot(REPO)
    flow = CosyVoice3Flow(cfg.flow).to(device)
    sd = torch.load(f"{snap}/{cfg.flow_weights}", map_location="cpu", weights_only=True, mmap=True)
    flow.load_weights((k, v.to(device)) for k, v in sd.items())
    flow.set_estimator_dtype(torch.bfloat16)
    return flow.eval(), cfg


def solve(flow, cfg, mu, cond, spks, valid, streaming, steps=10):
    b, _, frames = mu.shape
    mask = chunk_attention_mask(valid, frames, cfg.flow.dit.static_chunk_size, streaming)
    noise = flow.noise[:, :, :frames].expand(b, -1, -1).contiguous()
    return flow.decoder.solve(mu, mask, spks, cond, noise, steps)


def main():
    torch.set_float32_matmul_precision("highest")
    torch.backends.cudnn.allow_tf32 = False
    dev = torch.device("cuda")
    flow, cfg = load_flow(dev)
    batches = [int(x) for x in sys.argv[1].split(",")] if len(sys.argv) > 1 else [1, 2, 4, 8]
    frames_list = [int(x) for x in sys.argv[2].split(",")] if len(sys.argv) > 2 else [256, 512, 768, 1024, 1536, 2048, 3072, 4096]
    out = []
    g = torch.Generator(device=dev).manual_seed(0)
    with torch.no_grad():
        for frames in frames_list:
            for b in batches:
                mu = torch.randn(b, 80, frames, device=dev, generator=g)
                cond = torch.zeros(b, 80, frames, device=dev)
                cond[:, :, :174] = torch.randn(b, 80, 174, device=dev, generator=g)
                spks = torch.randn(b, 80, device=dev, generator=g)
                valid = torch.full((b,), frames, dtype=torch.long, device=dev)
                streaming = torch.ones(b, dtype=torch.bool, device=dev)
                args = (flow, cfg, mu, cond, spks, valid, streaming)
                for _ in range(2):
                    ref = solve(*args)
                torch.cuda.synchronize()
                n = 5 if b * frames <= 8192 else 3
                launch, total = [], []
                for _ in range(n):
                    torch.cuda.synchronize()
                    t0 = time.perf_counter()
                    solve(*args)
                    t1 = time.perf_counter()
                    torch.cuda.synchronize()
                    t2 = time.perf_counter()
                    launch.append(t1 - t0)
                    total.append(t2 - t0)
                torch.cuda.empty_cache()
                before = torch.cuda.memory_reserved()
                pool = torch.cuda.graph_pool_handle()
                graph = torch.cuda.CUDAGraph()
                s = torch.cuda.Stream()
                s.wait_stream(torch.cuda.current_stream())
                with torch.cuda.stream(s):
                    solve(*args)
                torch.cuda.current_stream().wait_stream(s)
                torch.cuda.synchronize()
                tc0 = time.perf_counter()
                with torch.cuda.graph(graph, pool=pool):
                    static_out = solve(*args)
                torch.cuda.synchronize()
                capture_s = time.perf_counter() - tc0
                pool_mb = (torch.cuda.memory_reserved() - before) / 2**20
                graph.replay()
                torch.cuda.synchronize()
                exact = torch.equal(static_out, ref)
                maxdiff = (static_out - ref).abs().max().item()
                replay = []
                for _ in range(n):
                    torch.cuda.synchronize()
                    t0 = time.perf_counter()
                    graph.replay()
                    torch.cuda.synchronize()
                    replay.append(time.perf_counter() - t0)
                row = {
                    "b": b, "frames": frames,
                    "eager_launch_ms": 1000 * sorted(launch)[n // 2],
                    "eager_total_ms": 1000 * sorted(total)[n // 2],
                    "replay_ms": 1000 * sorted(replay)[n // 2],
                    "capture_s": capture_s, "pool_mb": pool_mb,
                    "replay_equals_eager": exact, "max_abs_diff": maxdiff,
                }
                out.append(row)
                print(json.dumps(row), flush=True)
                del graph, static_out, pool
                torch.cuda.empty_cache()
    print("peak_allocated_mb", torch.cuda.max_memory_allocated() / 2**20)


if __name__ == "__main__":
    main()

parity.py: captured vs eager flow solves on real weights: python parity.py <reference clip.wav>

"""Captured vs eager CosyVoice3 flow solves on real weights and real speech tokens (#599).

Builds the flow node exactly as the server does (model.get_submodule), captures
its regions through the engine's PiecewiseCudaGraphRunner, and compares on
fixed tokens and the model's fixed noise:

  replay  vs eager on the same padded inputs        (same kernels: bits)
  replay  vs eager at the row's own length          (padding to the bucket)
  yardstick: eager alone vs eager batched with a longer row (the padding the
             eager path already applies to every batch)
  scale:  bf16 estimator vs float32 estimator, same inputs

Prints max |diff| and mel SNR (dB) per case; capture time and pool memory.
"""

import json
import math
import sys
import time
from unittest import mock

import soundfile as sf
import torch

from mstar.engine.cuda_graph_runner import PiecewiseCudaGraphRunner
from mstar.model.chatterbox.components.audio_frontend import resample
from mstar.model.cosyvoice3.components.flow import FlowRow
from mstar.model.cosyvoice3.cosyvoice3_model import CosyVoice3Model
from mstar.model.cosyvoice3.submodules import _without_tf32, flow_region


def snr(ref, x):
    err = (ref - x).double().pow(2).sum().item()
    return math.inf if err == 0 else 10 * math.log10(ref.double().pow(2).sum().item() / err)


def stats(ref, x):
    return {"max_abs": (ref - x).abs().max().item(), "snr_db": snr(ref, x), "bitexact": torch.equal(ref, x)}


def main():
    dev = "cuda"
    model = CosyVoice3Model(model_path_hf="FunAudioLLM/Fun-CosyVoice3-0.5B-2512")
    node = model.get_submodule("flow", dev)
    flow = node.flow
    wav, sr = sf.read(sys.argv[1], dtype="float32")
    wav24 = resample(torch.from_numpy(wav).to(dev), sr, 24000)
    voice = node.encoder.voice(wav24)
    p = voice.num_prompt_tokens
    # real speech tokens: the clip's own, tiled
    speech = voice.prompt_tokens.repeat(8).to(dev)
    torch.cuda.synchronize()
    free0 = torch.cuda.mem_get_info()[0]
    t0 = time.perf_counter()
    pool = torch.cuda.graph_pool_handle()
    runners = {}
    for label, config in node.get_piecewise_cuda_graph_configs(torch.device(dev, 0), torch.float32).items():
        r = PiecewiseCudaGraphRunner(label=label, config=config, resources={}, step_runner=mock.Mock(),
                                     device=torch.device(dev, 0), autocast_dtype=None, num_slots=1, memory_pool=pool)
        r.warmup_and_capture()
        runners[label] = r
    torch.cuda.synchronize()
    capture_s = time.perf_counter() - t0
    graphs = sum(len(r._graphs) for r in runners.values())
    print(json.dumps({"prompt_tokens": p, "regions": len(runners), "graphs": graphs, "capture_s": capture_s,
                      "capture_mem_gib": (free0 - torch.cuda.mem_get_info()[0]) / 2**30}), flush=True)

    def solve_bucketed(rows):
        frames = max(flow.valid_frames(r) for r in rows)
        bucket = next(b for b in model.config.flow_capture_frame_buckets if b >= frames)
        inputs = flow.prepare_solve(rows, frames=bucket)
        out = runners[flow_region(bucket)].run(static_inputs=inputs._asdict(), real_bs=len(rows))["mel"]
        with _without_tf32():
            eager_same = flow.solve_prepared(inputs, 10)
        return bucket, flow.row_mels(rows, out), flow.row_mels(rows, eager_same)

    results = []
    # a streamed request's hops (prompt aligned to a chunk), then odd final lengths
    first = -(-p // 25) * 25 - p + 25
    settles = [first, first + 50, first + 150, first + 250, 37, 61, 118, 203, 290]
    with _without_tf32():
        for settle in settles:
            final = settle not in settles[:4]
            n = settle if final else settle + 3
            row = FlowRow(voice, speech[:n], settle, final=final, streaming=True)
            bucket, (rep,), (same,) = solve_bucketed([row])
            exact = flow.solve_rows([row], 10)[0]
            longer = FlowRow(voice, speech[: settle + 120], settle + 120, final=True, streaming=True)
            batched = flow.solve_rows([row, longer], 10)[0]
            results.append({"settle": settle, "frames": flow.valid_frames(row), "bucket": bucket,
                            "replay_vs_eager_same_pad": stats(same, rep),
                            "replay_vs_eager_exact": stats(exact, rep),
                            "yardstick_eager_alone_vs_batched": stats(exact, batched)})
            print(json.dumps(results[-1]), flush=True)
        # batches: 3 rows (replays the batch-4 graph if 3 is not captured), 8 rows
        for k in (3, 8):
            rows = [FlowRow(voice, speech[: 40 + 23 * i], 40 + 23 * i, final=True, streaming=True) for i in range(k)]
            bucket, reps, sames = solve_bucketed(rows)
            exacts = [flow.solve_rows([r], 10)[0] for r in rows]
            entry = {"batch": k, "bucket": bucket,
                     "replay_vs_eager_same_pad_bitexact": all(torch.equal(a, b) for a, b in zip(sames, reps)),
                     "replay_vs_eager_exact_min_snr_db": min(snr(e, r) for e, r in zip(exacts, reps))}
            results.append(entry)
            print(json.dumps(entry), flush=True)
        # scale: bf16 vs float32 estimator on one row
        row = FlowRow(voice, speech[:203], 203, final=True, streaming=True)
        bf16 = flow.solve_rows([row], 10)[0]
        flow.decoder.estimator.float()
        fp32 = flow.solve_rows([row], 10)[0]
        flow.decoder.estimator.to(torch.bfloat16)
        print(json.dumps({"scale_bf16_vs_fp32_estimator": stats(fp32, bf16)}), flush=True)


if __name__ == "__main__":
    main()

capture_mem.py: capture time, graph pool and static buffers of the flow node's regions: python capture_mem.py desc

"""Capture cost of the flow node's regions (#599): time, graph pool, static buffers.

usage: capture_mem.py <asc|desc>   (capture order of the frame buckets)
"""

import json
import sys
import time
from unittest import mock

import torch

from mstar.engine.cuda_graph_runner import PiecewiseCudaGraphRunner
from mstar.model.cosyvoice3.cosyvoice3_model import CosyVoice3Model


def main():
    order = sys.argv[1]
    dev = torch.device("cuda", 0)
    model = CosyVoice3Model(model_path_hf="FunAudioLLM/Fun-CosyVoice3-0.5B-2512")
    node = model.get_submodule("flow", "cuda")
    configs = node.get_piecewise_cuda_graph_configs(dev, torch.float32)
    labels = sorted(configs, key=lambda k: configs[k].seq_len, reverse=order == "desc")
    torch.cuda.synchronize()
    torch.cuda.empty_cache()
    free0 = torch.cuda.mem_get_info()[0]
    alloc0 = torch.cuda.memory_allocated()
    t0 = time.perf_counter()
    pool = torch.cuda.graph_pool_handle()
    runners = []
    for label in labels:
        r = PiecewiseCudaGraphRunner(label=label, config=configs[label], resources={}, step_runner=mock.Mock(),
                                     device=dev, autocast_dtype=None, num_slots=1, memory_pool=pool)
        r.warmup_and_capture()
        runners.append(r)
    torch.cuda.synchronize()
    capture_s = time.perf_counter() - t0
    free1 = torch.cuda.mem_get_info()[0]
    torch.cuda.empty_cache()
    free2 = torch.cuda.mem_get_info()[0]
    pools: dict = {}
    for seg in torch.cuda.memory_snapshot():
        key = str(seg.get("segment_pool_id"))
        pools[key] = pools.get(key, 0) + seg["total_size"]
    static = sum(t.numel() * t.element_size() for r in runners for d in r._graphs.values()
                 for t in (*d.static_inputs.values(), *d.static_outputs.values()))
    gib = 2**30
    print(json.dumps({
        "order": order, "graphs": sum(len(r._graphs) for r in runners), "capture_s": round(capture_s, 1),
        "free_drop_gib": round((free0 - free1) / gib, 2), "free_drop_after_empty_cache_gib": round((free0 - free2) / gib, 2),
        "static_buffers_gib": round(static / gib, 3), "allocated_growth_gib": round((torch.cuda.memory_allocated() - alloc0) / gib, 3),
        "segments_by_pool_gib": {k: round(v / gib, 3) for k, v in pools.items()},
    }), flush=True)


if __name__ == "__main__":
    main()

spk_sim.py: speaker similarity and clip-by-clip identity between runs: python spk_sim.py <ref.wav> before=<wav dir> after=<wav dir>

"""Speaker similarity of served clips to the reference clip, and clip-by-clip
comparison between runs.

Speaker embedding: CosyVoice3's own CAMPPlus x-vector (the one the flow is
conditioned on), computed as ``ReferenceEncoder.voice`` does: 16 kHz Kaldi
fbank -> CAMPPlus. Cosine similarity of each clip to the reference clip.

usage: spk_sim.py <ref.wav> <name>=<wav dir> [<name>=<wav dir> ...] [--out f.json]
"""

import json
import sys
from pathlib import Path

import numpy as np
import soundfile as sf
import torch

from mstar.model.chatterbox.components.audio import kaldi_fbank_80
from mstar.model.chatterbox.components.audio_frontend import resample
from mstar.model.cosyvoice3.cosyvoice3_model import CosyVoice3Model


def load(path: Path) -> tuple[torch.Tensor, int]:
    wav, sr = sf.read(str(path), dtype="float32", always_2d=True)
    return torch.from_numpy(wav.mean(axis=1)), sr


def main():
    args = sys.argv[1:]
    out = None
    if "--out" in args:
        i = args.index("--out")
        out = args[i + 1]
        del args[i : i + 2]
    ref_path, dirs = Path(args[0]), dict(a.split("=", 1) for a in args[1:])
    model = CosyVoice3Model(model_path_hf="FunAudioLLM/Fun-CosyVoice3-0.5B-2512")
    campplus = model._reference_encoder("cuda").campplus

    @torch.no_grad()
    def embed(wav: torch.Tensor, sr: int) -> torch.Tensor:
        wav16 = resample(wav.cuda(), sr, 16000)
        return torch.nn.functional.normalize(campplus(kaldi_fbank_80(wav16)[None])[0], dim=0)

    ref = embed(*load(ref_path))
    report = {}
    clips = {}
    for name, d in dirs.items():
        files = sorted(Path(d).glob("*.wav"))
        sims = {}
        for f in files:
            wav, sr = load(f)
            clips.setdefault(f.name, {})[name] = wav
            sims[f.name] = float(embed(wav, sr) @ ref)
        report[name] = {"n": len(sims), "mean": float(np.mean(list(sims.values()))),
                        "min": float(np.min(list(sims.values()))), "per_clip": sims}
        print(f"{name}: {len(sims)} clips, speaker similarity mean {report[name]['mean']:.4f} "
              f"min {report[name]['min']:.4f}")
    names = list(dirs)
    pairs = {}
    for a_i, a in enumerate(names):
        for b in names[a_i + 1 :]:
            same = diff_len = 0
            snr = []
            for _, by in clips.items():
                if a not in by or b not in by:
                    continue
                x, y = by[a], by[b]
                if x.shape != y.shape:
                    diff_len += 1
                    continue
                if torch.equal(x, y):
                    same += 1
                    continue
                err = (x - y).pow(2).sum().item()
                snr.append(10 * np.log10(x.pow(2).sum().item() / max(err, 1e-20)))
            pairs[f"{a}~{b}"] = {"identical": same, "different_length": diff_len,
                                 "same_length_differing": len(snr),
                                 "snr_db_median": float(np.median(snr)) if snr else None,
                                 "snr_db_min": float(np.min(snr)) if snr else None}
            print(f"{a} vs {b}: {pairs[f'{a}~{b}']}")
    report["pairs"] = pairs
    if out:
        Path(out).write_text(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()

eager_same.py: eager solve_rows of one checkout, for a bitwise comparison across checkouts (CPU): python eager_same.py out.pt

"""Eager flow solves of one checkout, saved for a bitwise comparison across checkouts (#599).

usage: PYTHONPATH=<tree> python eager_same.py <out.pt>   (run from <tree>)
"""

import sys

import torch

from mstar.model.cosyvoice3.components.flow import CosyVoice3Flow, FlowRow, ReferenceVoice
from mstar.model.cosyvoice3.config import CosyVoice3Config, DiTConfig, FlowConfig

torch.set_num_threads(2)
flow_cfg = FlowConfig(vocab_size=40, pre_lookahead_channels=32, spk_embed_dim=16, noise_frames=2000,
                      dit=DiTConfig(dim=32, depth=2, heads=2, dim_head=16, conv_pos_kernel=5, conv_pos_groups=4,
                                    time_freq_dim=16))
cfg = CosyVoice3Config(flow=flow_cfg, silent_tokens=(1, 2), max_speech_tokens=512)
torch.manual_seed(0)
flow = CosyVoice3Flow(cfg.flow)
with torch.no_grad():
    for p in flow.parameters():
        p.normal_(0, 0.2)
flow = flow.float().eval()


def voice(n, seed):
    g = torch.Generator().manual_seed(seed)
    return ReferenceVoice(prompt_tokens=torch.randint(0, 40, (n,), generator=g),
                          prompt_feat=torch.randn(2 * n, 80, generator=g), embedding=torch.randn(16, generator=g))


def tokens(n, seed):
    return torch.randint(3, 40, (n,), generator=torch.Generator().manual_seed(seed))


batches = [
    [FlowRow(voice(10, 1), tokens(60, 1), 40, final=False, streaming=True)],
    [FlowRow(voice(23, 2), tokens(31, 2), 31, final=True, streaming=True),
     FlowRow(voice(5, 3), tokens(12, 3), 12, final=True, streaming=False),
     FlowRow(voice(17, 4), tokens(80, 4), 75, final=False, streaming=True)],
]
out = [mel for rows in batches for steps in (3, 10) for mel in flow.solve_rows(rows, steps)]
torch.save(out, sys.argv[1])
print(len(out), "mels")

[github-actions]

Automated review

The PR splits CosyVoice3's flow solve into an eager prepare_solve, which pads the batch, and a fixed-shape solve_prepared. The flow node then declares one PiecewiseBatchedConfig per frame bucket: 20 buckets and 97 graphs, which matches the ladder in config.py. FlowSubmodule._solve replays the smallest bucket that holds the batch through the engine's PiecewiseCudaGraphRunner. It falls back to eager past the largest bucket, for a batch size the bucket did not capture, or when the step count differs from the captured one. The code matches the description.

Findings: none. I traced the replay path through PiecewiseCudaGraphRunner and found nothing that breaks:

• The batched key seq_len * padded_bs agrees between capture and _resolve.
• The runner zeroes the padded rows and clamp_min(1) keeps their attention finite.
• Outputs are captured into buffers outside the shared pool, and PiecewiseOutput clones them.
• n_timesteps routing matches the step count baked into _solve_capture.
• Precision is set by _without_tf32 together with disable_autocast, so the runner gets autocast_dtype=None.
• compile defaults to False.
• With no declare_step, the runner keeps one slot, so the async worker is not involved.
• Nothing under mstar/model/ captures graphs or synchronizes (invariant 1).

Scope: as the description asks, I reviewed only this PR's commit: components/flow.py, the flow section of submodules.py, the capture fields in config.py, test/modular/test_cosyvoice3_flow_capture.py and test/cpu-core.txt. I did not review the rest of the diff, which is the stacked #495: the LM, DiT, HiFT, ONNX and text components, the model file, the API-server adapters and speech chunking, the Chatterbox and CLI edits, and the other CosyVoice3 tests.

Description:

• The H200 GPU runs don't name the torch/CUDA build. AGENTS.md asks for it.
• The flow-parity table was measured on 19a7ff9, not the head. The description says so and explains that the head only changes capture order. The GPU parity test is reported on the head.

---

_{Advisory, not a merge gate. Rules: AGENTS.md. Reply to discuss.}

[arun-elr]

The H200 GPU runs don't name the torch/CUDA build.

Added to the Performance toggle: torch 2.12.1+cu130 (CUDA 13.0, cuDNN 9.2), driver 580.173.02, the same build for every number in the PR.

The flow-parity table was measured on ..., not the head.

Reran it on the head: every value is unchanged, and the table is now labeled with the head.

[NSagan271]

🤖 Automated live test, AI-generated (Claude Code). Not checked line by line by a human. Numbers come from one allocation, so treat small deltas as provisional.

Live test: [CosyVoice3][Perf] Replay the flow's Euler solve as piecewise CUDA graphs per frame bucket

Works: throughput +50% at c1 and +75% at c8, TTFA p50 down 15–35%, WER unchanged; absolute numbers are from a heavily host-loaded node and not comparable to the PR's.

1x H200, configs/cosyvoice3.yaml + voices_dir (the CosyVoice repo's Mandarin zero-shot clip), SHM tensor transport · PR head c35a2d6d vs #495 head be299053 · head now d7008823: rebased onto #495's newer head (which adds the cuDNN bypass for the look-ahead convolutions, so both arms would now carry it), plus capture-order and capture-ladder commits; not retested

Claims

Claim	Result
c8 throughput 11.8 → 17.7 audio-s/s (+49%)	✅ 6.4 / 6.7 / 6.7 → 11.2 / 12.3 / 11.9 (+76%)
c1 first audio p50 454 → 333 ms	✅ 1072 / 785 / 817 → 816 / 682 / 600 ms (median −22%; one base start overlaps one PR start)
c8 first audio p50 1.31 → 0.95 s	✅ 4.34 / 3.34 / 3.51 → 2.53 / 2.23 / 2.37 s
WER unchanged	✅ Whisper-turbo, c1: 2.30% on all 6 starts (26/1131 words each); c8: base 2.30 / 2.56 / 2.30%, PR 1.68 / 2.65 / 1.77%
Capture adds ~63 s of startup	✅ to /health: base 103 / 46 / 51 s, PR 207 / 154 / 170 s (+~110 s here, host-loaded)
c16 / c32 gains	➖ not run

Performance

tts_speech_bench --voice zero_shot --seed 7, one burst of c per repeat, median of repeats per start (c1: 2, c8: 3), 3 starts per arm in order base, PR, PR, base, base, PR. Values per start.

	c1	c8
audio-s/s, base	3.8 / 4.4 / 4.3	6.4 / 6.7 / 6.7
audio-s/s, PR	5.9 / 6.6 / 6.7	11.2 / 12.3 / 11.9
TTFA p50 ms, base	1072 / 785 / 817	4344 / 3342 / 3514
TTFA p50 ms, PR	816 / 682 / 600	2530 / 2230 / 2370
RTF, base	0.297 / 0.256 / 0.260	1.294 / 1.206 / 1.194
RTF, PR	0.194 / 0.173 / 0.168	0.755 / 0.680 / 0.709

• Throughput ranges don't overlap at either concurrency; output volume matched (407–409 audio-s per repeat on every start).
• The node's 1-min load was 400–535 on 256 cores (other users) throughout, so every number here is 2–3x below the PR's. That the solve is launch-bound (the PR's microbenchmark) fits a larger relative gain under host contention; hypothesis, not measured.
• The first start of each arm shared the node's other GPUs with untimed servers and is its slowest; without them, c1 audio-s/s is base 4.4 / 4.3 vs PR 6.7. The PR's third c8 value is the median of 2 complete repeats.

Commands

mstar serve cosyvoice3 --config cosyvoice3_r6.yaml   # configs/cosyvoice3.yaml + model_kwargs.voices_dir
python -m benchmark.tts_speech_bench --engine mstar --url ... --model cosyvoice3 --voice zero_shot --seed 7 \
  --sentences benchmark/assets/t2s.txt --concurrency 1 --repeats 2 --warmup 2 --save-audio-dir ...
... --concurrency 8 --repeats 3 --warmup 8
python -m benchmark.tts_speech_wer --audio-dir <dir> --sentences benchmark/assets/t2s.txt

Other

• Not tested: c16, c32; the CUDA cases of test_cosyvoice3_flow_capture.py (CPU run: 104 passed, 3 skipped with test_cosyvoice3_model.py); GPU memory cost; float32 flow estimator; speaker similarity.

[merceod]

Reviewed c35a2d6 and 6ee4473 against #495 on a klone H200 (clock-limited node); 0a6e6e5 is the same patch rebased. Bucket-edge check with real weights: replay equals eager on the same padded inputs 120/120, the eager fallback triggers past 4,200 frames and for uncaptured batch sizes, the eager path is bit-identical to #495 48/48. Served on the current heads, 30 of 32 clips are byte-identical to #495 and the WER is identical; on the older heads short clips differed in their last chunk only (log-mel SNR 26 dB or better), which the description's "23/27 identical" undersells for short sentences.

Timing, 3 interleaved launches per arm, medians with ranges that never overlap: c=1 first audio p50 489 to 234 ms and RTF 0.30 to 0.14, c=8 first audio 3,428 to 1,357 ms, RTF 1.49 to 0.55, audio per second 5.2 to 14.0. With this PR a c=8 stream keeps up with playback on this node, without it it does not. The costs: startup 66 to 84 s against 169 to 205 s, and 11.3 GB at ready against 6.4 GB (97 graphs, 4.9 GiB, the description says 4.9).

One note on the premise rather than the code: dc248c1 (now in #495) attributes about 155 ms per call to cuDNN's float32 plan on the look-ahead convolutions and turns cuDNN off there, while this PR's description attributes the eager solve cost to kernel launches. The A/B above ran on the heads with that change, so the gain stands either way, but the description's explanation is dated.

Not run: a served request past the largest bucket with ignore_eos (the offline check covers that path), a float32 flow.

[arun-elr]

Thanks for the review and the live test.

One note on the premise rather than the code: [...] attributes about 155 ms per call to cuDNN's float32 plan on the look-ahead convolutions and turns cuDNN off there, while this PR's description attributes the eager solve cost to kernel launches. [...] the description's explanation is dated.

Updated the description; no code change. The look-ahead convolutions run in the eager prepare_solve, outside the captured solve, so their cost was in both arms of the PR's own served table, which predates that change; the description now says so and cites the A/B on heads with it (c1 first audio 489 to 234 ms). For the solve itself it now states what was measured, about 145 ms eager against 41 ms for a replay of the same kernels at 250 frames, and labels kernel launches as the likely cause, not profiled.

on the older heads short clips differed in their last chunk only (log-mel SNR 26 dB or better), which the description's "23/27 identical" undersells for short sentences.

The Correctness toggle now gives 30 of 32 byte-identical on the current heads, and that the older heads' differences were in short clips' last chunk only.

| Capture adds ~63 s of startup | ✅ to /health: base 103 / 46 / 51 s, PR 207 / 154 / 170 s (+~110 s here, host-loaded) |

• Real cost, host-dependent: 63 s is from an idle H200; the review measured a similar +100 to 120 s on a clock-limited node. The description now gives 63 s on that H200 and up to 120 s on busier nodes, and whether capture stays on by default at that cost is its open question for maintainers.

| c16 / c32 gains | ➖ not run |

• Noise: not run here; the description's c16 and c32 rows (one H200, 6 repeats per arm, ranges not overlapping) stand, and the review's c1 and c8 A/B on the current heads agrees in direction. The tested head and the current head have the same own commits; the rebase onto [Model] Add Fun-CosyVoice3: zero-shot voice-cloning TTS with streamed 24 kHz audio #495's newer head adds its look-ahead change to both arms.

[arun-elr]

#262 renames: the flow node's override is now get_piecewise_accelerator_graph_configs (main's engine calls only that name, so the old one captured nothing and every solve ran eagerly with no error) [...] GPU numbers predate this and were not rerun (no GPU in this pass).

Served check on one H200 that the renamed flow override takes effect: this head merged onto current main (clean merge), mstar serve cosyvoice3 with configs/cosyvoice3.yaml and a preset voice, MSTAR_REQUIRE_ACCELERATOR_GRAPHS=1, the same requests as the check on #495. Counts come from a local probe on the runners' replay paths and the eager solve_rows (not committed).

• Boot: healthy in 95 s; flow piecewise: 97 of 97 buckets (20 frame buckets from 100 to 4,200, batch 1 to 8 within 8,192 batch-frames) and the LM's 6 decode buckets, no capture failures. The flow capture took 1,044 MiB against 890 predicted (logged as a warning); GPU memory after boot 11.3 GB (6.4 GB on [Model] Add Fun-CosyVoice3: zero-shot voice-cloning TTS with streamed 24 kHz audio #495).
• Replay: 133 flow solves over the requests, all 133 replays of a captured bucket (8 frame buckets from 250 to 1,200, batch 1 to 8), no eager solve; 3,073 LM decode steps, all replayed. On [Model] Add Fun-CosyVoice3: zero-shot voice-cloning TTS with streamed 24 kHz audio #495 the same requests ran 130 eager solves.
• Audio: the same lengths as on [Model] Add Fun-CosyVoice3: zero-shot voice-cloning TTS with streamed 24 kHz audio #495 (long Korean, kana and Russian 82.0, 64.7 and 64.4 s; Chinese 71.8 against 72.6 s), English WER 2.4% (as on [Model] Add Fun-CosyVoice3: zero-shot voice-cloning TTS with streamed 24 kHz audio #495), Korean CER 1.2%, Russian WER 3.7%, Chinese complete (checked on its last 20 s); a streamed request's first bytes in 0.22 s (0.37 s on [Model] Add Fun-CosyVoice3: zero-shot voice-cloning TTS with streamed 24 kHz audio #495).
• Malformed: invalid JSON 422, empty input 400, a string seed 422, an unknown voice 400; the next request served normally.
• No problem from the capture. The garbled last kana word reported on [Model] Add Fun-CosyVoice3: zero-shot voice-cloning TTS with streamed 24 kHz audio #495 is the same here, from the LM.

[arun-elr]

@vasilevklart rebased on main now that #495 is in, so the diff is just this PR's commits (plus one test stub that takes #810's node_name). CI's green, ok to merge?

[vasilevklart]

This doubles main on 1x H200 (c1 8.2 → 14.9 audio-s/s, first audio 313 → 164 ms, same WER), but vllm-omni 0.32.0 on the same GPU and texts is 2.2-2.9x ahead at c8/c32 (68 / 102 audio-s/s vs 30 / 36) with lower WER (1.87% vs 2.77%) so the gap opens with concurrency (inline)

[vasilevklart]

Rows with the same step count solve together here, but the flow batch is capped at flow_max_batch_size = 8 and each row's tokens still sync to the host in prepare_inputs (:649)

[arun-elr]

Rows with the same step count solve together here, but the flow batch is capped at flow_max_batch_size = 8 and each row's tokens still sync to the host in prepare_inputs (:649)

@vasilevklart thanks for running it against vllm-omni. Those two look like the main candidates for the c8/c32 gap, and both are on main as they are here: this PR captures batch sizes 1 to 8 under that cap (fewer for the larger frame buckets, within 8,192 batch-frames) and leaves prepare_inputs alone. Raising the cap alone would run batches past 8 eagerly; capturing them too costs more of the 4.8 GiB and of the startup.

I'd keep this PR to the capture and work on the gap in a follow-up:

1. one host copy of the speech tokens for the whole flow batch instead of one per request, by moving the readback and filtering from prepare_inputs into the batched preprocess;
2. a higher flow batch cap, with the captured batch sizes and frame budget sized for it;

measured on one H200 at c8 and c32 against vllm-omni 0.32.0 on your texts. OK to merge this one as is, or would you rather have 1 here?

[vasilevklart]

lgtm, 1 and 2 as a follow-up sounds good

[vasilevklart]

ThankS!