[Bugfix][Qwen3-Omni] Fix Code2Wav CUDA-graph output length surplus (#4466)

[linyueqian]

Summary

Addresses the "noticeable lag in the middle" reported in #4466 when CUDA graph is enabled in Qwen3-Omni Code2Wav (stage 2, enforce_eager: false).

Root cause

CUDAGraphDecoderWrapper._decode trimmed each graph/compiled replay output to the nominal actual_size * total_upsample. Qwen3-Omni's Code2Wav forward returns fewer samples than nominal (a fixed causal-conv trim constant; empirically S * 1920 - 555). For a chunk whose size is not an exact capture bucket (zero-padded up to the bucket), trimming to nominal keeps that constant's worth of stale buffer-tail samples.

In streaming (async_chunk: true) the ramp-up chunks each leak this surplus, adding a fixed ~46 ms of extra samples near the start that the chunk concat turns into a one-time timing slip, heard as the "lag in the middle". Steady-state chunks are unaffected because chunk + left_context is always a capture bucket by construction, so the effect does not accumulate.

This was partially addressed by #3687 for the non-streaming chunked_decode_with_cudagraph path, but _decode itself (used by the default streaming path) was never fixed.

Fix

Trim relative to the captured buffer length instead of the nominal length, so the graph/compiled output matches the eager forward exactly. This is a no-op for decoders whose forward returns the nominal length (e.g. Qwen3-TTS, where the captured length already equals nominal).

Verification

On an H200-class GPU with Qwen3-Omni-30B-A3B-Instruct (vllm 0.22.0), serving the stock config (CUDA graph) vs an eager stage-2 variant on the same inputs. Stages 0 and 1 are identical and seeded, so the diff isolates the Code2Wav graph:

• Before: cuda-graph output is +1110 samples (+46 ms) vs eager, with a +23 ms cross-correlation lag step at a chunk seam (cross-correlation 1.000 elsewhere, i.e. identical content that is merely time-shifted).
• After: cuda-graph output is sample-identical to eager (length delta 0, lag flat 0).

This matches the reporter's attached files exactly: their cuda_graph track is +1110 samples vs eager with the same +23 ms slip.

Tests: adds regression tests for both the graph and compiled paths using a short-output decoder (fail before this change, pass after). Full wrapper suite: 65 passed.

Note on #4309

#4309 bundles two symptoms. The "plosive/lag between chunks from CUDA graph" part is this same surplus. The separate speaker-drift part tracks the CUDA toolkit version (the reporter found that moving from CUDA 13 to 12.9 fixes it) and did not reproduce on the H200-class box used here, so it is not addressed by this PR.

@amy-why-3459 would you be able to pull this branch and confirm the lag is gone on your setup? The change is contained to the Code2Wav CUDA graph trim, so the rest of the pipeline is unchanged.

[amy-why-3459]

@zsl24 PTAL

[hsliuustc0106]

Bugfix looks correct: trim to captured buffer length instead of nominal.

[Sy0307]

Reviewed the Code2Wav CUDA graph length fix.

The core fix looks technically correct. The old _decode() path trimmed graph replay output to actual_size * total_upsample, which silently assumes that decoder output length is nominal. That assumption is false for Qwen3-Omni Code2Wav, whose decoder returns a shorter causal-conv-trimmed output. For padded bucket replay, this preserved stale buffer tail and explains the +1110 samples / 46.25 ms surplus in #4466.

The new logic:

drop = (padded_size - actual_size) * self.decoder.total_upsample
actual_out_len = max(0, static_output.shape[-1] - drop)
return static_output[..., :actual_out_len]

is the right invariant: trim relative to the captured output's true length, not the nominal length. This keeps Qwen3-TTS as a no-op case while making short-output decoders match eager.

I also checked the added regression coverage. The new ShortOutputDecoder tests cover both:

• normal CUDA graph replay path;
• compiled CUDA graph path.

I ran the wrapper suite on A100 with vLLM 0.23:

65 passed

I also ran an end-to-end Qwen3-Omni graph/eager comparison using current main + this PR on A100:

graph: 82815 samples, 3.450625s, no mid-gap
eager: 82815 samples, 3.450625s, no mid-gap
sample_delta: 0
duration_delta: 0.0

This directly addresses the #4466 failure mode where graph previously produced +1110 samples / +46.25 ms versus eager.

One caveat worth keeping visible: _batched_chunked_decode() still assumes decode_fn(chunk) returns exactly chunk_len * total_upsample. The PR documents that short-output decoders such as Qwen3-Omni should not use that path, and current Qwen3-Omni streaming goes through the per-request decode path, so this is not a blocker. But if someone later routes Qwen3-Omni Code2Wav into batched_chunked_decode_with_cudagraph(), the same class of length bug would come back, just wearing a different hat. Please keep that invariant explicit.

Net: the fix is narrow, well targeted, and the tests cover the previously broken replay paths. No request-changes from me. Notice fix the CI issue.

[hsliuustc0106]

Thanks for the targeted fix. One nearby edge case still seems worth covering before this lands: Qwen3-Omni can pack multiple Code2Wav requests into a padded [B, 16, max_len] tensor when seq_token_counts is present, then chunked_decode_streaming() slices each row with code_seq_len * total_upsample.

Could you add a small regression test for a multi-request, variable-length Qwen3-Omni Code2Wav batch using a short-output decoder, or note why that path cannot hit CUDA graph replay in practice? The current new tests cover the single-request padded graph/compiled replay case well; this would close the closest remaining sibling case where one shorter row in a padded batch could still leak or over-slice relative to eager length.