[Bugfix] Fix Qwen3-TTS word timestamps with async chunking

[Bezdarnost]

Purpose

Fixes #7505. /v1/audio/speech requests with word_timestamps=true can hang when Qwen3-TTS runs with the bundled asynchronous chunk configuration and a forced-aligner stage.

The aligner consumes a completed waveform through code2wav2aligner; it is not a chunk-connector consumer. Previously it inherited pipeline-wide async chunking, so the orchestrator pre-submitted a placeholder and the scheduler waited for connector chunks. The completed audio was not forwarded through the synchronous stage-input processor.

Allow a stage to opt out of pipeline-wide async chunking, and disable it for the injected pooling aligner. Resolve the same mode in legacy and structured configuration, including processor selection. The shared connector predicate now excludes stages with async chunking disabled, so the orchestrator waits for completed audio before submitting the aligner. Talker and Code2Wav retain asynchronous streaming.

Test Plan

CPU regressions cover injected-stage configuration with both deploy and CLI async settings, legacy/structured parity, deferred aligner submission, and completion after both final outputs. The speech test client now forwards word_timestamps and validates nonempty, ordered intervals from X-Word-Timestamps.

GPU regressions exercise two concurrent timestamped WAV requests and ordinary streaming speech on the same server. They use the existing weekly slow and L4 and tts sweep, with full model weights; no new Buildkite job is needed.

Reproduction with sufficient GPU memory:

vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --deploy-config vllm_omni/deploy/qwen3_tts.yaml --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B --trust-remote-code

In a second terminal:

curl --max-time 120 -D timestamps.headers -o speech.wav http://localhost:8000/v1/audio/speech -H 'Content-Type: application/json' -d '{"model":"Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice","input":"The weather is nice today, perfect for a walk in the park.","voice":"vivian","language":"English","response_format":"wav","word_timestamps":true}'

The response should finish with WAV audio and a populated X-Word-Timestamps header. The regression fixture below uses smaller memory/batch settings for a single GPU.

python -m pytest -q -o addopts= tests/config/test_forced_aligner_injection.py tests/config/test_omni_config.py tests/engine/test_orchestrator_kv_sender_info.py tests/model_executor/stage_input_processors/test_forced_aligner_stage.py tests/utils/test_forced_aligner.py tests/entrypoints/openai_api/test_serving_speech_stream.py tests/helpers/tests/test_assertions.py
VLLM_USE_FLASHINFER_SAMPLER=0 OMP_NUM_THREADS=4 python -m pytest -s -v -o addopts= tests/e2e/online_serving/test_qwen3_tts_word_timestamps.py -m 'slow and L4 and tts' --run-level full_model
pre-commit run

vLLM Version: 0.29.0 (current source installation requirement)

vLLM-Omni Commit: 01f4479 + this change

Test Result

• CPU regression suite: 741 passed, 2 skipped. On unmodified base commit 01f4479, the new configuration/routing regressions produce 4 failures; they pass with this fix.
• New full-model GPU tests: 2 passed on one RTX 5080 (16 GB), including two concurrent timestamped requests and the streaming control. CUDA graphs were enabled, and the harness's speech transcription checks passed.
• Additional live API checks passed for single/concurrent timestamped requests, raw WAV streaming without timestamps, and a timestamped request after client cancellation. Each timestamped sample returned 12 word intervals for the 12-word input.
• All applicable changed-file pre-commit hooks passed, including mypy and CI marker checks; git diff --check passed.

GPU environment: Python 3.12.11, PyTorch 2.13.0, CUDA 13, vLLM 0.29.0, Transformers 5.14.1, NVIDIA driver 595.91.07. The regression fixture retains the bundled async connector settings and limits sequence/batch sizes and KV cache to fit a single GPU. The L4 CI target was not run locally; the same CI marker command passed on RTX 5080. No throughput or multi-GPU performance claim is made.

AI assistance: Used Codex to investigate the issue, implement the fix and regression tests, run validation, and prepare the PR description.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-15T01:17:27.574531Z	565654f	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/vllm_omni_config.md, docs/design/module/entrypoints.md, docs/design/module/stage_runtime.md.

Module owners: @lishunyang12 @alex-jw-brooks @fake0fan

Routing: @lishunyang12 via module of the changed files, CODEOWNERS; @alex-jw-brooks via module of the changed files; @fake0fan via module of the changed files

@Bezdarnost, 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.

[Bezdarnost]

This PR appears to belong to: docs/design/module/vllm_omni_config.md, docs/design/module/entrypoints.md, docs/design/module/stage_runtime.md.

Module owners: @lishunyang12 @alex-jw-brooks @fake0fan

Routing: @lishunyang12 via module of the changed files, CODEOWNERS; @alex-jw-brooks via module of the changed files; @fake0fan via module of the changed files

@Bezdarnost, 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.

Author self-review:

• Checked that the injected aligner opts out of async chunking while Talker and Code2Wav keep streaming; pipeline-wide async_chunk=false still disables async for every stage.
• Checked configuration/CLI precedence, deferred aligner submission until the complete waveform is available, and request completion after both audio and timestamps.
• Validation: 741 CPU tests passed (2 skipped), both full-model GPU regressions passed on RTX 5080 with CUDA graphs, and all applicable changed-file pre-commit hooks passed. Confirmed that the new regression tests fail on the original code and pass with this fix.
• Codex assisted with implementation, tests, validation, and this comment.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 565654f38e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 216ba3636648 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[linyueqian]

Requesting changes at 701706f6, narrowly. The core of this is right and I traced it end to end: the per-stage async_chunk opt-out reaches every site that classifies a chunk consumer (stage_receives_chunks in the orchestrator and the chunk adapter, the prewarm placeholder path, the scheduler-side adapter creation), the injected aligner resolves to False on both the legacy and the structured path, the edge validator is exempt on the Code2Wav-to-aligner edge because Code2Wav declares no async producer and correctly rejects opting out one end of the talker-to-Code2Wav edge, all 42 shipped async_chunk: true deploys still validate, and an explicit per-stage False survives a CLI --async-chunk. The non-streaming word-timestamps hang in #7505 is fixed and the tests pin it.

The two things I want changed are inline: the streaming endpoint now returns wrong timestamps rather than hanging, because Code2Wav runs in DELTA mode and the aligner only ever sees the last delta; and the new per-stage flag lets a stage-0 opt-out silently turn off downstream chunking because the engine still derives the pipeline flag from stage 0. Both are small. Reviewed statically against the diff with a two-model cross-check; fork head, no PR code executed. I added ready so the lane runs on this head; I will re-review as soon as the push lands.

[linyueqian]

This branch went into conflict with main after #7413 merged; git merge-tree shows the conflict is confined to tests/helpers/tests/test_assertions.py, with tests/helpers/assertions.py, omni_config.py and stage_config.py auto-merging cleanly. When you push the fixes for the two inline findings above (streaming word_timestamps under DELTA and the stage-0 opt-out), please rebase onto current main in the same push so the lane runs against the tree it would merge into. The branch already carries ready, which means a push does not start a new build on its own; ping here after the push and I will re-fire it.

[Bezdarnost]

@linyueqian Rebased onto 4392af5c and addressed both findings:

• Streaming timestamp requests retain cumulative audio for the aligner; the response emits only the newly generated samples, so audio is not repeated.
• The engine enables chunking when any resolved stage needs it, rather than reading only stage 0. Global async=false still disables every stage.

Validation: 1,145 CPU tests passed (2 skipped), and all 3 full-model GPU tests passed on RTX 5080 with CUDA graphs, including WebSocket streaming with full-sentence timestamps. The new focused regressions fail before these fixes and pass afterward.

Local pre-commit passes except mypy: current main already defines _create_default_diffusion_stage_cfg twice in omni_engine_base.py (no-redef). I reproduced the same failure before my review fixes. Intel build 8342 failed during checkout with No space left on device, before tests ran.

Could you re-fire CI and re-review this update? Codex assisted with the fixes and validation.

[linyueqian]

Approving at 216ba363, which is baf1197a plus a merge of current main (5aadcbb0) with no change to the PR's own hunks. Both items from my first pass are fixed at the source rather than papered over, and each carries a regression that fails before the fix.

Streaming word timestamps: a word_timestamps request that streams now clones its sampling params and switches them to RequestOutputKind.CUMULATIVE, so the aligner stage receives the complete waveform at the terminal output instead of the last DELTA chunk, and the WAV emitter compensates by tracking emitted_samples and slicing only the new tail out of each cumulative tensor before resampling, so the client still hears each sample once. The list-mode branch already emitted only the new tail via prev_count, so the two modes now agree. test_cumulative_audio_terminal_output_supplies_full_aligner_waveform and test_streaming_timestamps_retain_complete_audio pin the two halves, and test_async_chunk_streaming_word_timestamps covers it end to end on the WebSocket path. Non-timestamp streaming keeps DELTA and is untouched.

Stage-0 opt-out: the engine's async_chunk is now true when any resolved stage needs chunking rather than mirroring stage 0, and the per-stage async_chunk: bool | None opt-out is resolved through one helper (resolve_stage_async_chunk) that both the engine-args materialisation and the processor selection use, so the two can no longer disagree. validate_stage_async_chunk_edges then rejects a pipeline whose producer and consumer sit on different sides of a chunked connector edge with an actionable error, which is the case that would previously have started and then stalled. test_engine_async_chunk_includes_downstream_stages, test_stage_async_chunk_opt_out_matches_legacy_config and test_async_chunk_rejects_mismatched_connector_edge cover the three paths.

The rebase onto 4392af5c resolved the test_assertions.py conflict, and the follow-up merge of main leaves the head zero commits behind with no file overlap. Validation: static read of the delta between the two heads plus the new tests; no PR code executed. I re-fired the general lane at this head after the push and will merge on green; the author reports 1,145 CPU tests and the three full-model GPU tests passing locally on an RTX 5080.