[tts][bugfix] Fix CosyVoice3 sampling and stage handoff

[akshatvishu]

Purpose

Fixes #6416.

CosyVoice3 voice cloning could pass invalid probabilities to torch.multinomial, which killed stage 0 during decoding. Fixing that crash allowed the same supported request path to reach three failures that had been hidden by the sampler crash and test skips. BF16 tensors could not cross the shared stage serializer, serialized prompt conditioning was decoded as flat keys and dropped, and the causal HiFT STFT window could stay on CPU while its input was on the GPU.

Fix

Sampling and attention geometry. The CosyVoice sampler now uses vLLM's random_sample instead of torch.multinomial. It excludes NaN and both infinities after request masks and logit processors run, and it raises a clear error if a row has no finite token. The RAS fallback clones the selected score before masking it, so the only legal token can be restored without adding another Python device check or changing random number generator advancement. The outer CosyVoice config also exposes num_key_value_heads = 2, which matches the released nested Qwen config. Before this change, the Qwen layers and KV cache used 2 KV heads while vLLM attention metadata fell back to 14. The field is fixed because this integration supports one released CosyVoice3 checkpoint whose other backbone dimensions are already fixed in the same config. Valid scores keep the same categorical distribution, but a fixed seed can produce a different token sequence because vLLM and torch.multinomial consume random numbers differently. The finite row check also adds one device synchronization per sampler batch.

Stage handoff. The shared serializer now sends tensor storage through a uint8 view and restores the tensor with its recorded PyTorch dtype and shape. The change matches vLLM's raw tensor byte pattern, preserves BF16 exactly, and keeps the existing wire bytes for previously supported dtypes. CosyVoice now uses the shared decoder, which rebuilds nested keys such as embed.speech_token. The removed local decoder returned dotted keys, so serialized prompt conditioning was silently dropped. Decoding keeps one CPU byte copy, as the previous NumPy path did.

Vocoder device state. The causal HiFT STFT window is a nonpersistent buffer, so normal model moves include it without adding a checkpoint key. _stft and _istft also move and retain the window when its device differs from the input. The use site check is needed because vLLM's dummy loader does not call the model weight loading hook. The first mismatched call performs one device copy, and later calls reuse the moved buffer.

Online coverage. The three CosyVoice tests are marked core_model because blocking AMD CI runs them with dummy weights. The same tests load real weights and use Whisper quality checks when they run with --run-level advanced_model or full_model. ROCm requests use a 600 second client timeout because a cold AITER MHA kernel build took 301 seconds on one MI300X. CUDA keeps the existing 120 second default.

Test Plan

HIP_VISIBLE_DEVICES=0 COSYVOICE3_TRT=0 \
  .venv/bin/python -m pytest -o addopts='' -s -v \
  tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py \
  -m slow --run-level advanced_model

HIP_VISIBLE_DEVICES=0 COSYVOICE3_TRT=0 \
  .venv/bin/python -m pytest -o addopts='' -s -v \
  tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py \
  -m slow --run-level core_model

.venv/bin/python -m pytest -o addopts='' -q \
  tests/engine/test_data_entry_keys.py \
  tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py \
  tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py \
  tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py \
  tests/model_executor/models/cosyvoice3/test_cosyvoice3_config.py \
  --run-level core_model

vLLM Version: 0.27.0

vLLM-Omni Commit: ee889a0

Test Result

The test environment used one MI300X with 192 GB VRAM, ROCm 7.2.3, Python 3.12.13, BF16, and model snapshot 29e01c4e8d000f4bcd70751be16fa94bf3d85a18.

The real weight test passed all three Chinese sync, Chinese streaming, and English sync cases with the default stage 0 TRITON_ATTN backend:

3 passed, 14 warnings in 455.19s
Whisper small transcript similarity: 1.000 for all three requests

The targeted real weight test also passed with stage 0 forced to ROCM_AITER_FA through --stage-overrides '{"0":{"attention_backend":"ROCM_AITER_FA"}}':

1 passed, 14 warnings in 182.99s
Whisper small transcript similarity: 1.000
AITER specialization: gqa_ratio=7, head_size=64

The AITER specialization confirms that the loaded attention implementation used 14 query heads and 2 KV heads. The CPU config regression separately checks that the outer vLLM metadata contract resolves the same 2 KV heads. The dummy weight test passed all three cases after the STFT use site fix:

3 passed, 14 warnings in 788.60s

The focused CPU suite passed on the latest head:

79 passed, 15 warnings in 3.78s

Pre-commit passed for all 11 changed files. The earlier H20 review was performed at 38eb12bf, before the RAS, finite row, STFT use site, and attention geometry fixes. The latest head still needs an H20 or H100 rerun to verify the CUDA Flash Attention path before the CUDA test coverage is considered resolved. An MI300X cannot provide that proof because ROCm uses different attention backends, kernels, and graph implementations. The MI300X results therefore validate only the ROCm Triton and ROCm AITER paths.

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)

[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.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/model_integration.md.

Module owners: @tzhouam @gcanlin

@akshatvishu, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[akshatvishu]

@vllm-omni-review-bot

[linyueqian]

Reviewed at 38eb12bf008efea8c57a8aef4f1cad68d15397fd against base 5730af0a95299387c233e12dfbdbe728aace0e58.

I ran the e2e on one H20 (vLLM 0.27.0, torch 2.13.0+cu132, bf16, CUDA graphs on stage 0), comparing base with only the three skip marks removed against this PR head:

	base + unskip	PR head
Stage 0	dies: probability tensor contains either inf, nan or element < 0, device-side assert (#6416 reproduced)	survives
Talker logit rows	n/a	81 sampler calls, 46 entirely NaN
HTTP	500, Stage-0 has no live replica	200, WAV returned
Whisper transcript	n/a	さぁさっ, さぁさっ, empty

The PR does stop the crash. It does not fix what produces it.

[blocking] The vocoder device fix does not apply under load_format: dummy

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

stft_window only reaches the GPU via self.hift.to(device) at cosyvoice3_code2wav.py:318, which sits inside code2wav.load_weights() and is reached only through the model's load_weights hook at cosyvoice3.py:1197. Under load_format: dummy vLLM never calls load_weights, and persistent=False keeps the buffer out of state_dict(), so dummy initialization does not move it either. The buffer stays on CPU.

Measured on this head at --run-level core_model:

RuntimeError: stft input and window must be on the same device
             but got self on cuda:0 and window on cpu
-> StageEngineCoreProc_stage1_replica0 dead -> HTTP 500 -> 3 failed

Same commit at --run-level advanced_model (real weights): stage 1 healthy, audio produced.

This matters because the PR removes @pytest.mark.skip from all three tests, and .buildkite/amd/test-amd-ready.yml:196 runs this file as -m "slow" --run-level "core_model" on mi300_1 with grade: Blocking. tests/helpers/stage_config.py:670 stamps load_format: dummy at that run level. So merging as-is turns that step from 3 skipped and green into 3 failed and red on every subsequent PR.

Smallest fix: make device placement independent of the weight-loading path. The PR intentionally dropped upstream CosyVoice's per-call window.to(device) in favour of the buffer, and that is the optimization that breaks here. Placing the window on the input device inside _stft/_istft when it does not already match keeps the buffer benefit and removes the dependency on load_weights having run.

[blocking] nan=0.0 turns the #6416 crash into uniformly random tokens

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py:571

I instrumented _sanitize_sampling_scores on this head. Of 81 rows entering the talker sampler for test_voice_clone_zh_001, 46 arrived entirely NaN. nan_to_num(nan=0.0) flattens an all-NaN row to all zeros, which is a uniform distribution over the full 151923-token vocabulary, so 57% of decode steps select a uniformly random token. That is the measured cause of the unintelligible output above, and the reason the server now returns 200 with garbage rather than failing.

The NaN producer is untouched by this PR. Either diagnose it, or use nan=float("-inf") so an invalid logit is excluded rather than promoted, with an explicit failure when no finite candidate remains.

For the greedy branch nan=0.0 changes nothing, since torch.argmax already returns the NaN index (measured: raw 0, after nan_to_num(nan=0.0) 0, with nan=-inf 1). The choice only has an effect on the sampling path, where it keeps the broken token as a live candidate.

[high] The _ras_sample_one restore guard is dead code, and its test passes on the bug

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py:646

original_score = weighted_scores[top_id]
weighted_scores[top_id] = float("-inf")
weighted_scores[top_id] = torch.where(
    torch.isfinite(weighted_scores).any(), weighted_scores[top_id], original_score)

Integer indexing returns a 0-dim view, so line 647 also sets original_score to -inf and torch.where reads -inf in both branches. When top_id was the only unmasked token the row becomes all -inf, softmax is all NaN, and random_sample's argmax returns index 0, which is a token the mask had forbidden.

Reproduced through the real CosyVoice3Model.sample(): with allowed_token_ids_mask leaving only token 1 (or 2) legal and RAS rejecting it, sample() returns 0.

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py:428 places the sole finite token at index 0 and asserts 0, which is exactly the buggy output, so it passes without exercising the guard. Moving that token to index 1 fails.

Fix: original_score = weighted_scores[top_id].clone(), and return top_id directly when nothing else is finite instead of sampling from a NaN distribution. The instrumented run entered this fallback on 22 of 81 steps, so the path is live.

[medium] The advanced_model marker selects nothing for this file

No pipeline step selects test_cosyvoice3_tts_expansion.py with -m advanced_model. The real coverage is .buildkite/cuda/test-weekly.yml:91, the TTS · H100 step, which runs -m "slow and H100 and tts" --run-level "full_model". hardware_marks at tests/helpers/mark.py:165 applies marks for every platform in res unconditionally, so these tests carry H100, cuda, MI325 and rocm.

That weekly H100 step is where the Whisper similarity assertion actually runs, and it is where #4644 recorded 0.81, 0.62 and 0.03. The test plan here covers only MI300X, so the un-skip lands unvalidated on the hardware that produced the original failures. Worth either adding an advanced_model step or dropping the marker so the selection story is unambiguous.

Verified clean

• The serializer round-trips bit-exactly for float64, bfloat16, float16, float32, int64, int32, bool and uint8, including 0-dim and empty tensors. I also stress-tested the torch.frombuffer(bytearray(...)).view(dtype) path for alignment across 10 dtypes and 305 sizes (3050 round-trips, no failures, data_ptr % itemsize == 0 throughout; CPython's minimum bytearray alignment measured at 16 bytes against a maximum torch itemsize of 8). Wire bytes are unchanged for pre-existing dtypes, and the pattern already matches the mooncake/mori/yuanrong connectors.
• random_sample's in-place div_ is safe at both call sites, since _nucleus_sample_one and _ras_sample_one each pass a freshly allocated tensor.
• persistent=False is safe with load_state_dict(..., strict=True), and CausalHiFTGenerator.__init__ calling torch.nn.Module.__init__ directly does mean the base class register_buffer never ran, so the change itself is correct.

Worth adding to the description

Replacing _decode_additional_information with the shared deserialize_additional_information fixes a bug the PR does not claim. The old helper returned flat dotted keys such as embed.speech_token, so info.get("embed", {}) at stage_input_processors/cosyvoice3.py:126 always returned {} and prompt conditioning was silently dropped on the serialized path. That is a plausible explanation for the streaming reference bleed and the 0.03 English similarity reported in #4644, and it deserves a line in the description.

Caveats on the measurements

H20 is not H100 or MI325, and vllm_omni/deploy/cosyvoice3.yaml states its defaults are tuned for H100. TensorRT was absent in my environment, so campplus fell back to ONNX-Runtime, which is the same path your COSYVOICE3_TRT=0 run takes. What is hardware-independent is the code reading behind each finding. What the run adds is that base reproduces #6416's assert verbatim on this GPU, and that the head replaces it with 46 all-NaN rows out of 81.

[akshatvishu]

Thanks for the detailed review @linyueqian; I have updated the branch with the following changes:

• STFT window placement. _stft and _istft now move and retain the window when its device does not match the input. The dummy loader never calls load_weights, so placement must also happen where the window is used. The first mismatch causes one device copy and later calls reuse the moved buffer.

• Logits and sampling. The sampler now excludes NaN and both infinities and it raises an error if no finite token remains. It also uses vLLM's random_sample instead of torch.multinomial. The finite row check adds one device synchronization per batch, but it prevents the server from returning random speech for an invalid row.

• RAS restore. The saved score is cloned before the selected token is masked, so the mask cannot also change the saved value. The test now puts the only legal repeated token at a nonzero index, which catches the old fallback to token 0. I kept the restore check on the device to avoid another synchronization and preserve random number generator advancement.

• Test markers. The tests now use core_model because the blocking AMD job runs this file with dummy weights. --run-level advanced_model still loads real weights and runs the Whisper quality check.

• KV heads. The outer config now sets num_key_value_heads = 2, matching the official nested Qwen config and loaded layers. Without the field, vLLM metadata fell back to 14 KV heads. The field duplicates the released checkpoint value because vLLM reads the outer config before the nested Qwen config is loaded.

• Stage handoff. The shared serializer now transfers BF16 tensors through a raw uint8 view, which preserves their bytes and payload size. CosyVoice also uses the shared decoder. The old local decoder returned flat keys such as embed.speech_token, so the caller silently dropped prompt conditioning. These fixes stay in this PR because the same request reached this stage boundary after the stage 0 crash was fixed.

• ROCm timeout. One MI300X run spent 301 seconds building the AITER MHA kernel, so ROCm uses a 600 second test timeout while other platforms keep 120 seconds. Cached runs do not pay the same build cost.

Measurement caveats

• Hardware scope. The review ran on an H20. The deploy file has an H100 tuning comment, and the ROCm test marker names MI325. The base run reproduced the exact issue 6416 CUDA assertion on that H20, while commit 38eb12bf produced 46 all-NaN rows out of 81. Those counts apply only to that commit and H20 run. The dummy loader, RAS view, and KV head findings come from the code paths and do not depend on those counts.

• TensorRT scope. TensorRT was unavailable in the H20 review, so campplus used ONNX Runtime. Our MI300X runs set COSYVOICE3_TRT=0 and used the same campplus path. Matching campplus removes one difference between the runs, but it does not validate CUDA Flash Attention or CUDA graph behavior.

Testing and validation

• The focused CPU suite passed 79 tests, and pre-commit passed for all changed files.
• One MI300X passed three real weight TRITON_ATTN tests and one forced ROCM_AITER_FA test. Whisper similarity was 1.000 in all four tests.
• CUDA Flash Attention and CUDA graphs remain unverified. An H20 or H100 run on the latest commit is required before the PR can claim that the original CUDA failure is fixed or enable the CUDA test.

The branch still conflicts with current main. I will merge current main and keep issue 6416 skipped only on CUDA until a current H20 or H100 run passes.

[linyueqian]

Both blocking findings are resolved. Re-reviewed at ad499969a308679676a0ad933d0d2168234950c2 against merge base dc72559d5.

CUDA evidence you asked for

You noted that an H20 or H100 run on the latest commit was required before the PR could claim the original CUDA failure is fixed. I ran it on one H20 (vLLM 0.27.0, torch 2.13.0+cu132, bf16, CUDA graphs on stage 0), with _SKIP_CUDA_ISSUE_6416 patched out locally so the CUDA path actually executes, and with the talker sampler instrumented.

Run level	Weights	Result
advanced_model	real, Whisper check on	3 passed
core_model	dummy, matches the AMD ready job	3 passed

Sampler counters across all six server runs: 0 rows with NaN out of 201 / 201 / 108, 0 rows with +inf, 0 rows with no finite logit, and the new ValueError guard never fired.

Transcripts at advanced_model:

zh_001  收到好友从远方寄来的生日礼物那份意外的惊喜与深深的感动让我热泪盈眶   exact
zh_002  收到好友从远方进来的生日礼物那份意外的惊喜与深深的感动让我热泪盈眶   one homophone, 寄来 -> 进来
en_001  Hello, this is a voice cloning test with English text.              exact

For comparison, the same GPU on 38eb12bf gave 46 all-NaN rows out of 81 and transcribed さぁさっ, さぁさっ, and empty. #4644 recorded 0.81, 0.62 and 0.03 on H100. The streaming case is the one #4644 flagged for leaking reference text, and there is no leak now, just a Whisper homophone.

On this evidence the CUDA skipif can come off. I would keep it in this PR if you would rather land the fixes first and enable CUDA in a follow-up, but the data supports enabling it now.

On the root cause

num_key_value_heads = 2 is the fix that matters. I confirmed the checkpoint's nested CosyVoice-BlankEN/config.json declares num_attention_heads: 14 and num_key_value_heads: 2, while the outer CosyVoice3Config declared only the 14, so vLLM's engine metadata sized the KV cache for 14 KV heads against layers built with 2. That explains the NaN logits, and it explains why #6416 showed up with enforce_eager: false. The fact that the finite check never fires in any of the six runs is the confirmation that it is a safety net rather than the fix.

_get_stft_window on the base HiFTGenerator is the right place for the device fix, and core_model passing is the direct proof: that path was dead before, since hift.to(device) only runs inside load_weights, which the dummy loader never calls.

Non-blocking

1. cosyvoice3.py:687, logits.masked_fill_(~finite_logits, float("-inf")) also maps +inf to -inf, which turns a maximally likely token into a forbidden one. I measured zero +inf rows so this is theoretical, but torch.isposinf deserves separate handling from NaN if it is ever reachable.

2. cosyvoice3.py:686, raising ValueError inside sample() kills the EngineCore for every in-flight request, not just the offending one. Failing loudly beats returning random speech, so this is an improvement over the previous behaviour, but a per-request abort would be better if the runner supports it. Worth a follow-up rather than a change here.

3. num_key_value_heads is hardcoded next to the other Qwen2-0.5B values. That matches the file's existing style, but since the failure mode is silent NaN rather than a load error, reading it from the nested config (AutoConfig is already imported) or asserting the two agree at construction would stop this recurring on a future checkpoint.

Note for maintainers, not this PR

#6454 landed in main as dc72559d5 and is byte-identical to this PR's original register_buffer(..., persistent=False) hunk, without the _get_stft_window half. So main currently carries the incomplete fix and the stft input and window must be on the same device failure under load_format: dummy is live there, masked only because these tests are still skipped. This PR is what closes it. Worth landing reasonably promptly for that reason.

Caveats

H20 is not H100 or MI325, and vllm_omni/deploy/cosyvoice3.yaml states its defaults are H100 tuned. TensorRT was unavailable, so campplus used ONNX Runtime, which is the same path your COSYVOICE3_TRT=0 runs take. CUDA graphs and the CUDA attention backend were exercised, since stage 0 runs with enforce_eager: false and this is where the original assert fired.

LGTM.

[akshatvishu]

Thanks for the detailed re-review and H20 validation @linyueqian ! In the current PR, I've only removed _SKIP_CUDA_ISSUE_6416 and retained the +inf handling (PyTorch softmax invalidates rows with +inf, matching CosyVoice and vLLM). I'll handle the sampling failure and attention geometry assertions in follow-ups.

Also, could you re-run CI? (in-case my last push undid the first ci check)