Optimize CosyVoice3 Stage1 flow batching

[gerayking]

Purpose

Add opt-in CosyVoice3 Stage-1 flow batching for concurrent streaming requests.

Main changes:

• Batch Stage-1 streaming flow items inside CosyVoice3Model.forward() when multiple requests are scheduled together and COSYVOICE3_BATCH_FLOW=1.
• Batch the flow-matching mel path while keeping HiFT/vocoder execution and cache state per request.
• Support ragged codec-token lengths with padding, per-row lengths, and attention masks.
• Preserve positional request/result alignment with an explicit completeness invariant.
• Log group-size distribution and enable detailed profiler scopes only under the debug toggle.
• Remove the experimental bf16/fp16 varlen attention path after its measured performance regression.

Runtime controls

• COSYVOICE3_BATCH_FLOW=1: enable cross-request Stage-1 flow batching. It is disabled by default.
• COSYVOICE3_BATCH_FLOW_DEBUG=1: log batching group statistics and enable detailed Stage-1 profiler scopes. It is disabled by default.

Batching can produce small waveform differences because batched GEMM shapes change floating-point execution. It remains opt-in until WER/SIM or equivalent quality measurements justify changing the default.

Test Plan

Validation completed locally:

python3 -m py_compile \
  vllm_omni/model_executor/models/cosyvoice3/runtime.py \
  vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py \
  vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py \
  vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py \
  vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py \
  tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py \
  tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

git diff --check upstream/main..HEAD

Added coverage:

• Verify batching is opt-in and debug is disabled by default.
• Verify the float32 attention path receives the 2D per-row mask without Q/K/V dtype casting.
• Verify padded codec-token batches preserve per-row lengths.
• Add an L4-marked numerical regression test that runs a real lightweight DiT/CFM flow and compares ragged batched mels with per-request mels using length-consistent noise.

The local machine does not have PyTorch, pytest, or a GPU, so dependency-complete unit and GPU execution is delegated to PR CI. The branch is rebased onto the latest main and is conflict-free.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[Sy0307]

Please solve the conflicts and the DCO issue. And I found it has much benefit for CosyVoice3 throughput. Wish we can merge it as soon as possible. Also cc @linyueqian

[gerayking]

Please solve the conflicts and the DCO issue. And I found it has much benefit for CosyVoice3 throughput. Wish we can merge it as soon as possible. Also cc @linyueqian

thanks, I'll fix conflict soon

[gerayking]

@Sy0307 solved conflict.

[linyueqian]

Reviewed at 7b62e32 against base dd527b4. I ran the batching path on GPU against the real Fun-CosyVoice3-0.5B-2512 weights rather than only reading it.

The batching math is correct. No blockers from me. Notes inline, then what I could not measure.

Verified on GPU

I patched the CFM noise draw to be deterministic and length-consistent, so a batched run and a per-request run see the same z. Run-to-run is then bit-exact, which makes every delta below real rather than sampling noise.

• Output lengths match batched vs per-request in every case I tried: equal-length B=2, ragged B=2 (40/25 codec tokens), ragged B=4 (50/40/30/25), finalize=True, and token_offset 10/5. The valid_mel / trim_mel / pre_lookahead_len slicing reproduces what flow_model.inference returns unbatched.
• The new per-row mask is load-bearing and it works. Forcing _normalize_attention_mask to return None makes a padded row drift by rel_mean=1.27; with the mask in place it is 0.021. Padding does not leak into valid positions.
• No row-position bias: two identical items in one batch give bit-identical audio (max|d| = 0.0).
• Real CosyVoice3Model.forward, 2 concurrent requests across 3 streaming chunks: identical per-chunk sample counts (7680 / 11520 / 18240) and identical totals with COSYVOICE3_BATCH_FLOW=1 vs 0. The streaming vocoder-cache handoff and the audio stitching both survive the loop restructure.

What I could not measure

The only GPU I could reach was a 4xL20X pod where all four cards sit at 100% utilization from another tenant, and any GPU process of mine gets killed at about 45 seconds. That took out every server boot, the offline e2e test, and all throughput runs. So there is no throughput number, no online-serving e2e, and no WER/SIM from me.

Four paired base-then-PR runs of a single streaming chunk came out at -0.9%, -0.2%, -0.3% and +1.4%, while base-to-base drift over the same session was +1.9%. So no regression is detectable on the default batch=1 path, at roughly +/-2% resolution. At 2.65s per flow call on that box any fixed CPU overhead is diluted below visibility, so the profiler-scope cost I flag inline stays unmeasured.

Two failures reproduce on dd527b4 without this PR, so they are not yours: tests/e2e/offline_inference/test_cosyvoice3_expansion.py errors with Diffusers pipeline index not found, and tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py already fails 1 and errors 7 when run CPU-only despite its cpu marker.

[andyluo7]

Thanks for unblocking the ROCm lane. I checked the exact target logs on 1cfa7b7b:

• AMD #11310, Diffusion Model Test collected the six accelerator tests in this file and skipped all six under the updated shared marker.
• CUDA #14545, Diffusion Model Test executed all six successfully.

Only four of those tests call the bundled paged flash-attention API. test_the_decode_graph_captures_into_the_shared_platform_pool and test_a_second_request_replays_the_first_capture exercise the graph lifecycle without that kernel and should remain runnable on ROCm. Applying the capability check to the shared cuda_only marker therefore removes useful ROCm coverage. The literal 64 also duplicates a capability parameter instead of using the test's existing constants.

#6935 splits the markers: the four paged-kernel tests use the existing paged_decode_supported(...) capability probe, while the two graph-only tests retain the broader accelerator guard. It contains no ROCm/CUDA behavior branch. Could you please drop this unrelated SenseNova hunk in favor of #6935, or use the same split-marker structure here?

[andyluo7]

Follow-up with exact-head evidence from #6935: AMD #11313 skipped exactly the four paged-kernel cases and passed both graph-only cases; CUDA #14547 executed and passed all six. This confirms the split-marker approach preserves useful ROCm coverage while leaving supported CUDA unchanged. Please drop the broad shared-marker SenseNova hunk here in favor of #6935.

[vllm-omni-review-bot]

Omni ReviewBot: three questions on the performance claim

@gerayking this PR reads as a performance or value claim:

• title: Optimize CosyVoice3 Stage1 flow batching

Before the full evidence checklist, three short questions:

1. Bottleneck — what is the current bottleneck, and which profile, trace or per-stage measurement shows it?
2. Value — what does the change buy the user or the system (latency, throughput, memory, cost), and at which workload?
3. A/B or ablation — is there a same-workload, same-head/config comparison that isolates each main claim on its own? For stacked optimizations, one number per item rather than a blended delta.

When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution.

[hsliuustc0106]

Omni ReviewBot: three questions on the performance claim

@gerayking this PR reads as a performance or value claim:

• title: Optimize CosyVoice3 Stage1 flow batching

Before the full evidence checklist, three short questions:

1. Bottleneck — what is the current bottleneck, and which profile, trace or per-stage measurement shows it?
2. Value — what does the change buy the user or the system (latency, throughput, memory, cost), and at which workload?
3. A/B or ablation — is there a same-workload, same-head/config comparison that isolates each main claim on its own? For stacked optimizations, one number per item rather than a blended delta.

When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution.

@gerayking please help reply the reviews