[Perf][Qwen3-TTS]Cached Incremental Decode

[sphinxkkkbc]

Purpose

Part W5 of #4855.

The existing ref_context_cache stores reference codec tokens but not decoder state, causing every async Code2Wav invocation to repeatedly decode previous context.

This PR introduces request-local incremental decoding for Qwen3-TTS Code2Wav. After the first chunk, async requests decode only newly completed codec frames using bounded cached context. Both ICL and x-vec-only requests are supported through shared eager and CUDA Graph phase grouping.

Files Modified

The core implementation changes and their test coverage are concentrated in the following files:

File	Diff	Role
segmented_graph_wrapper.py	+1,201	Implements phase-aware CUDA Graph capture and replay.
modeling_qwen3_tts_tokenizer_v2.py	+678 / −107	Implements incremental decoding and the common decode interface shared by eager and CUDA Graph backends.
qwen3_tts_code2wav.py	+127 / −213	Replaces request-level full decoding with delta-frame batching and maintains decoder state by request ID.
stage_input_processors/qwen3_tts.py	+25 / −12	Sends only newly generated codec frames for async chunks while retaining reference context in the first chunk.
test_qwen3_tts_incremental_decode.py	+1,104	Covers decode correctness, graph regressions, routing, cache lifecycle, batching, rolling, and final-chunk trimming.

The main implementation and test changes are concentrated in these five files. The detailed contract changes are described below. If needed, this PR can be split into smaller review units.

Contract Changes

1. For async_chunk=true, the Talker-to-Code2Wav connector uses a delta-input protocol:

   • The first ICL chunk contains the bounded reference prefix and initial codec chunk.
   • Subsequent ICL chunks contain only newly completed codec frames.
   • Prefixless/xvec requests follow the same protocol without a reference prefix.

2. Code2Wav maintains request-local decoder state keyed by request ID and releases it when the request finishes. The cached state contains:

   • quantized suffix frames needed for pre_conv boundary reconstruction;
   • the reference Transformer KV prefix;
   • the bounded rolling suffix window;
   • downstream hidden frames required by the upsampling and decoder stack.

3. Stateful eager and CUDA Graph execution share the same batched_request_decode() grouping policy. Requests are grouped by mode (icl or xvec), phase (first, growing, or rolling), and target suffix shape. Outputs and updated caches are restored to their original request slots.

4. Code2Wav constructs one right-padded [B, Q, F] transport batch and passes true codec lengths together with list[request_state] for stateful decoding or None for stateless decoding. Padding and trimming remain decoder responsibilities:

   • Variable-length ICL prefixes are left-padded, with hidden-state and attention masks excluding padded positions.
   • A partial final suffix delta may be right-padded to reuse the next captured phase; causality protects valid output positions.
   • The decoder returns only waveform samples produced from newly decoded frames and records each row's valid length in _last_output_audio_length.
   • Cache state produced from a padded final chunk is not published because the request finishes immediately.

5. Async CUDA Graph shapes are derived from initial_codec_chunk_frames, codec_chunk_frames, and optional codec_chunk_ramp. Non-async shapes continue to use decode_cudagraph_capture_sizes, while captured batch sizes use decode_cudagraph_batch_sizes.

6. Stateful dummy runs use synthetic request IDs to initialize caches during warmup and external CUDA Graph capture. Dummy-run replay and fallback events are excluded from runtime CUDA Graph statistics.

7. decode_batch_max_size limits temporary batched-cache allocation for eager stateful phase groups. 0 leaves eager batching unlimited and does not affect CUDA Graph replay batch sizes.

The non-async/stateless path remains supported and continues to use full-sequence chunked decoding without requiring request IDs.

CUDA Graph Shapes

The segmented wrapper captures graphs for ICL and xvec first chunks, every reachable growing suffix phase, and stable rolling decode.

With the default configuration:

initial_codec_chunk_frames: 1
codec_chunk_frames: 25
sliding_window: 72

the suffix transitions remain unchanged:

1  + 25 -> 26
26 + 25 -> 51
51 + 25 -> 76
72 + 25 -> 97

When a chunk ramp is configured, capture shapes are derived from its actual deltas. For example:

codec_chunk_ramp: [4, 4, 8, 16, 25]
codec_chunk_frames: 25

produces:

first chunk: 4
4  + 4  -> 8
8  + 8  -> 16
16 + 16 -> 32
32 + 25 -> 57
57 + 25 -> 82
72 + 25 -> 97  # stable rolling phase

Intermediate ramp phases therefore use their real delta shapes instead of being padded to codec_chunk_frames. Default graph shapes and runtime behavior are unchanged.

Compatibility Notes

Cached decoding is numerically close to full decoding but is not guaranteed to be bit-exact. Differences in input shape, memory layout, and contiguity may produce small floating-point differences.

Test Plan

pytest -v \
  tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py \
  tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py \
  tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

Test file	Result	Coverage
test_qwen3_tts_incremental_decode.py	34 passed	Numerical correctness; cache initialization and graph replay; routing and slot restoration; padding and trimming
test_qwen3_tts_code2wav.py	22 passed	Unified padded batching, state handling, stateless compatibility, and decoder-owned trimming
test_qwen3_tts_async_chunk.py	101 passed, 1 local failure	First-chunk reference context and follow-up delta-input contract

The single local failure is test_parse_ramp_warns_on_tail_mismatch. The expected warning is emitted but is not captured by the local caplog configuration. The test and logging path are unchanged by this PR.

Benchmark

Benchmarked on single 3090-24GB

vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
    --deploy-config vllm-omni/vllm_omni/deploy/qwen3_tts.yaml \
    --omni --port 8091

python3 vllm-omni/benchmarks/tts/bench_tts.py \
  --host 127.0.0.1 \
  --port 8091 \
  --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
  --locale en \
  --dataset-path /path/to/seedtts_testset \
  --num-prompts 32 \
  --concurrency 1 2 4 8 16 \
  --wer-eval \
  --output-dir ./results/qwen3-tts-12hz-1.7b-base

Both baseline and cached runs used batch-size-1 graph capture. Additional experiments with graph buckets [1, 2, 4] showed little difference up to concurrency 8, while batch size 4 or larger improved performance at concurrency 16 and above.

Concurrency	Baseline RTF	Cached RTF	Baseline TTFP	Cached TTFP	Throughput improvement
1	0.244	0.222	142 ms	143 ms	+10.3%
2	0.296	0.250	189 ms	160 ms	+20.2%
4	0.407	0.308	288 ms	203 ms	+31.1%
8	0.609	0.431	440 ms	364 ms	+40.4%
16	1.071	0.718	1049 ms	837 ms	+34.1%

Seed-TTS evaluation of 32 cached-path requests produced a mean WER of 0.0089, with no request, PCM capture, or ASR failures.

• vLLM version: 0.26.0
• vLLM-Omni commit: b8edbc8d23d24bf310f671a5cafe6b4d8949973e

cc @Sy0307 @linyueqian

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)

[sphinxkkkbc]

Resolved the review comments and revised the relevant test.

All tests passed except for test_parse_ramp_warns_on_tail_mismatch, as also mentioned in descriptions. This test was not modified by this PR, and the local failure appears to be related to log capture—the expected warning is emitted but not recognized by the assertion, rather than indicating an issue in the runtime path. I can investigate and fix it separately if needed.

Test command:

pytest -v \
  tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py \
  tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py \
  tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

Result: 159 passed, 1 failed.

@linyueqian Please feel free to take another look. Thanks!

[sphinxkkkbc]

Update ffe679e and 08c3ba1:

CI exposed a missing dummy-run guard that did not reproduce in my local serving tests for test_qwen3_tts_base and test_qwen3_tts_customvoice. The CI configuration uses the dummy weight loader, which bypasses load_weights(). As a result, the inner Code2Wav CUDA Graph wrapper is not initialized, while the outer vLLM CUDA Graph capture still invokes Code2Wav.forward(). The dummy request was then incorrectly routed to _decode_xvec_first_chunk() instead of _forward_exact(), causing CUDA Graph capture to fail.

This is the scenario that the original codes.shape[-1] == 4 check was intended to handle. Added explicit dummy/capture guards covering both PIECEWISE:decoder.forward() and FULL:Code2Wav.forward(). Corresponding regression tests were added for both paths.

In addition, this update removes the redundant:
exact per-request waveform → pad to a batch tensor → slice back to the exact length
conversion after batched decoding.

The updated output contracts are:

• Stateful decoding returns list[torch.Tensor] with an exact-length waveform for each request.
• For stateless decoding, the helper removes both right-side CUDA Graph/batch padding and the decoder-internal 25-frame chunk overlap, returning an exact-length waveform for each request. Code2Wav.forward() only removes the scheduler-provided leading context once from the assembled waveform.
• Trimming such as ICL reference-prefix removal and suffix decoder-context removal remains inside the decoder.

Feedback and suggestions regarding this contract change are welcome.

The single-request stateful path has also been consolidated into _batched_request_decode(). The redundant mixed graph/eager _decode_icl_prefix() and _decode_suffix() entrypoint was removed. Batch size 1 now uses the same xvec-prefix, ICL-prefix, suffix, and eager-fallback routing as larger batches.

Unit tests:

pytest -v \
  tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py \
  tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py \
  tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py \
  tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

All functional tests passed. The only remaining local failure is the previously mentioned log-capture issue.

E2E tests:

pytest -v \
  tests/e2e/online_serving/test_qwen3_tts_customvoice.py \
  tests/e2e/online_serving/test_qwen3_tts_base.py

Both passed.

I also reran the benchmark, and the numbers reported in the PR description remain unchanged. These changes primarily consolidate the decode routing, remove redundant output padding, and add robust CUDA Graph dummy-run fallbacks. No behavioral or performance regression was observed.

[linyueqian]

Re-reviewed at e43c98aec. Five of my six items are cleanly fixed, and I re-ran the end-to-end check on hardware. Two things still need a change.

The reference-audio leak is fixed — measured, not just read

I re-ran the same A/B I used to find it: H20, Qwen3-TTS-12Hz-1.7B-Base, vLLM 0.26.0, async_chunk: false, voice=clone, 8.0800s reference, merge-base vs this head.

arm	correlation of the leading 8.08s vs. the reference
b3f4fbf9e (base)	−0.0124 / +0.0012
e43c98aec (before your fix)	+0.9010 / +0.9010
e43c98aec (now)	+0.0005 / −0.0038

Durations no longer differ by the reference length either. #4421 stays fixed.

Thanks for the other four as well — and particularly for replacing the codes.shape[-1] == 4 shortcut with an explicit _is_dummy_run flag rather than just moving the magic number around. I only asked where the 4 came from; you turned it into a real contract.

1. left_context_size is indexed with the wrong index

row walks the compacted valid_codes_qf, but left_context_size is still sized and filled by original batch position:

279:  left_context_size = [0] * len(request_ids_list)   # original batch index
384:  valid_indices.append(i)                          # the mapping, already here
459:  for row in range(len(valid_codes_qf)):           # compacted index
466:  start = left_context_size[row] * self._total_upsample

Any request skipped earlier in the loop (empty codec payload, or a length not divisible by num_quantizers) shifts every later row. A valid ICL row then gets someone else's trim — too small and the reference audio comes back, too large and generated speech is cut. valid_indices[row] is the fix; it is already built for exactly this.

My A/B was single-request, so it cannot see this one. It needs a batch with a skipped row.

2. Aborted requests never free their decoder state

on_requests_finished pops the scheduler's request id, but _decoder_state_cache is keyed on meta.request_id from the payload — the external id. chunk_transfer_adapter.py keeps a request_ids_mapping precisely because those two can differ. An aborted request produces no terminal payload, so this pop is the only cleanup path, and it misses.

Combined with the cache now being a warning rather than a bound, that is an unbounded leak of GPU decoder state. I raised this as a soft question last time; on a second look it deserves to block. Either key the cache on the internal id, or translate before popping.

Two external reviewers independently flagged this one as the most serious item, which is what changed my mind about its severity.

[linyueqian]

row is the compacted index over valid_codes_qf, but left_context_size is [0] * len(request_ids_list) — original batch positions. The loop above continues past empty and malformed requests without appending, so once one is skipped every later row reads a neighbour's trim.

Too small and the reference audio leaks again; too large and generated speech gets cut. valid_indices is already populated at line 384 for exactly this mapping:

Suggested change

	start = left_context_size[row] * self._total_upsample
	start = left_context_size[valid_indices[row]] * self._total_upsample

Worth a regression test with a batch where the first request has an empty codec payload — the single-request A/B I ran cannot reach this.

[linyueqian]

This pops the scheduler's request id, but the cache is keyed on meta.request_id — the external id from the payload. chunk_transfer_adapter.py maintains a request_ids_mapping because those two are allowed to differ.

An aborted request emits no terminal payload, so the finished or segment_finished path in forward() never runs and this is the only cleanup left. When the ids differ, the pop misses and the entry stays. _decoder_state_cache_warn_entries logs but does not bound, so nothing stops it growing.

Either key the cache on the internal id, or translate here before popping. I asked about this as a question last round rather than a blocker; two external reviewers both rated it the most serious item and I think they are right.

[sphinxkkkbc]

Thanks for catching the potential issues that were not exposed by my local tests, @linyueqian. I’ve added regression tests for both issues.

Regarding the external/internal request IDs, I previously verified that the cache key and finished-request ID were identical in my test setup, so I did not treat this as a blocker. But as you pointed out, the two IDs are not guaranteed to be identical and may drift when request ID mapping is involved. The decoder state cache is now explicitly keyed by the scheduler’s internal request ID, matching finished_req_ids and ensuring aborted requests are cleaned up correctly.

[linyueqian]

Approving at c63c2efc0. Both remaining items are fixed, and I re-ran the end-to-end check on the new head rather than carrying over the old result.

No reference-audio leak on this head. Same H20 A/B, control arm is current main (12a5f6fb), not the stale merge-base I used the first time:

arm	leading-8.08s correlation vs. the reference
12a5f6fb (main)	−0.0034 / +0.0028
c63c2efc0 (this PR)	+0.0005 / −0.0038

For contrast, the same probe read +0.9010 / +0.9010 before you fixed it. #4421 stays closed.

One thing worth stating plainly: the pr arm's output bytes are identical to my previous run against e43c98aec. That is expected, not a stale run — with a single valid request valid_indices[0] == 0, so the index fix is an identity transform on this path. My A/B therefore does not exercise fix 1 at all. Your test_qwen3_tts_code2wav.py additions do, which is the right place for it.

Fix 1 — left_context_size[valid_indices[row]] at :483. Exactly right.

Fix 2 — keying the cache on the runner-side id is better than what I suggested. I said "key on the internal id or translate before popping" and missed that direct forward() calls and CUDA Graph dummy runs have no runner id at all; a straight swap would have left them without a key. The cache_request_ids copy with a payload-id fallback handles that. Thanks for thinking it through further than the review did.

For the record, the six items I raised and where they landed:

	fix
reference audio leaked into stateless output	left_context_size parsing restored, trim under request_states is None
uninitialized tail on short rows	decoder returns a list, so no padded region to trim
is_segment_finished never evicted state	finished or segment_finished
codes.shape[-1] == 4 magic number	explicit _is_dummy_run — you turned a question into a contract
initial_chunk_frames hardcoded to 1	plumbed from config
unbounded state cache + id mismatch	internal-id key, plus a size warning

Nice work on the perf numbers too — +34–40% at c8–c16 with WER 0.0089 is a real result, and the design held up under three rounds of review without needing to be restructured.