[Bugfix] Fix Qwen3-Omni streaming crops in mixed-length batches

[dshah1333]

Purpose

Fix the audio interval selected for shorter Qwen3-Omni requests in a padded streaming batch.

chunked_decode_streaming calculates each request's decoder trim by subtracting the whole batch's waveform length from that request's nominal length. When another request has a longer codec window, the subtraction becomes negative and clamps to zero. The shorter request then discards too much overlap and retains samples produced from padding.

For the pinned Qwen3-Omni checkpoint, the decoder returns frames * 1920 - 555 samples. A 29-frame request with four context frames, batched beside a 50-frame request, currently selects [7680:55680] instead of [7125:55125]. Both slices have 48,000 samples, so a duration-only check misses the shift.

Measure the trim from the decoded batch window, then use each request's valid start and end. Singleton and equal-length slices stay unchanged. Non-async decoding, the shared graph wrapper, sampling, and deployment defaults are untouched.

This follows the tail compensation in #4706. #4525 fixed graph-bucket padding, while this fixes padding between requests. The old batch test assumed padding supplied valid future context. The replacement checks values, lengths, output count, batch order, and a singleton → mixed batch → singleton schedule.

Reproduction

In a matching Linux vLLM-Omni development environment:

python -m pytest -q \
  tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py

With the final tests and only the production crop restored to baseline: 22 failed, 26 passed. With this fix: 48 passed. The regression needs the normal runtime imports but no model weights or GPU computation.

Test Plan

• Run the focused regression and adjacent Qwen/Code2Wav graph tests.
• Compare old and corrected crops on the same loaded pretrained decoder, including eager and actual batch-two graph execution.
• Run the full three-stage server at concurrency one and four, retaining streamed audio and checking chunk lengths and singleton PCM.
• Run changed-file pre-commit hooks and a separate read-only adversarial review.

vLLM Version: 0.28.0, Torch 2.13.0+cu130, Transformers 5.14.1, NVIDIA driver 580.173.02, 2x A100-80GB.

vLLM-Omni Commit: 99d0a5602bb22aa64ce886ab3434c14179ac9b50, based on 6fb7b0a0. Linux imported-module and GPU results were collected at 9464231f against ab561485 before updating the submission branch. The patch and affected model, packer, stage processor, graph wrapper, generation runner and test files are unchanged. Fresh extracted-method red/green checks and changed-file hooks passed after the update. Linux/GPU runs were not repeated after it.

Test Result

Check	Baseline	Fix
Focused Linux regression	22 failed, 26 passed	48 passed
Adjacent Qwen model and graph suite	Not run	126 passed
Pretrained FP32/BF16 causal-prefix comparisons	48/48 fail	48/48 bitwise exact
Full serving, c1 and c4	16/16 completed, one oversized first chunk	16/16 completed, correct chunk counts
Singleton decoded PCM	Reference	8/8 bitwise identical

The actual baseline server emitted a 7,680-sample first chunk instead of 7,125, an extra 555 samples at 24 kHz. All candidate first chunks contained 7,125 samples. The controlled GPU test holds packed width at 50 and uses width 64 for graph capture, comparing a short request with the valid prefix of a full-width decode containing genuine future frames. Constructor/dispatch scaffolding is lightweight, while the real imported packer and fully weighted decoder execute.

All applicable changed-file pre-commit hooks passed. A separate read-only Codex review found no introduced blocker after checking trim arithmetic, final chunks, graph double trimming, metadata fallback, dtype/lifetime and retained evidence. The test counts overlap and should not be added together.

Reproduction scripts, measured results, and full qualifications

Additional controls and limits

• An extracted-method CPU sweep passed 823,296 cases over request/companion widths 1–64, every valid context length, three trim values, and eager/graph-wrapper doubles. This supplements the normal Linux regression.
• Default cross-shape FP32/BF16 pointwise gates also failed healthy equal-width controls. Disabling only TF32 made all 36 FP32 candidate comparisons pass the original tolerance, while all baseline cases still failed. No production precision setting was changed. General BF16 cross-shape parity is not established.
• Concurrent serving takes were not identical codec/schedule replays. The fixed-code component tests provide the causal crop comparison. The server runs provide symptom and integration evidence.
• All 32 audio takes were transcribed. Exact normalized ASR matches were 4/8 for both singleton cohorts, 8/8 for baseline c4, and 7/8 for candidate c4. Singleton PCM is unchanged. This PR makes no broad speech-quality improvement or non-inferiority claim.
• The primary GPU comparison tensors were retained. The companion-poison test retained its executed equality verdict, not the auxiliary output tensor. Other accelerator backends, FP16/quantized checkpoints and PP capture are outside this validation.

AI assistance: Used Codex for the fix, regression tests, validation scripts, and PR text, plus a separate read-only Codex review.

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

[dshah1333]

Self-review:

• Traced codec windows through the async producer, padded batch packer, decoder and output crop. The old calculation mixes a request's nominal length with the longer batch output. Checked the 1920*T - 555 length contract against the reference convolution stack and confirmed both crop bounds need the correction.
• Checked singleton/equal-length behavior, partial final chunks, zero-trim controls, batch order, metadata fallback and graph-buffer lifetime. The graph wrapper removes capture-bucket padding first, so this per-request crop does not double-trim. The change adds no state, copies, synchronization, API or sampling policy.
• The final Linux regression fails 22 cases on the old crop and passes all 48 with the fix. The adjacent suite passes 126 tests. Pretrained FP32/BF16 prefix checks pass 48/48 exactly, and all eight singleton serving PCM outputs match baseline. A separate read-only Codex review found no introduced blocker.
• Verified the patch and affected execution/test files are unchanged after updating to current main. Fresh extracted-method checks and all applicable changed-file hooks passed afterward. The PR separates the earlier GPU runs, concurrent integration results and failed numerical/ASR controls, and makes no general speech-quality claim.

Scripts and measured data. AI assistance is disclosed in the PR body.

[vllm-omni-review-bot]

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

Module owners: @tzhouam @fake0fan @Gaohan123

Routing: @tzhouam via module of the changed files, CODEOWNERS; @fake0fan via module of the changed files; @Gaohan123 via module of the changed files

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

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 3f238fe6ccb4 produced:

• Priority: high. Prompt maintainer attention is suggested.

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

[linyueqian]

Approving at 3f238fe6. The old slice measured each row's decoder trim as code_seq_len * upsample - batch_wav.shape[-1], which is only right for the row that spans the whole decoded window; a shorter row padded with codec zeros got a smaller or zero trim and an end index past its valid samples, so in a mixed-length streaming batch the shorter request's chunk was shifted and included samples decoded from padding. Measuring the trim once from the decoded window and slicing each row to [left_context * upsample - tail, code_seq_len * upsample - tail] gives every row the same treatment a singleton decode gets, and the parametrised test compares batched rows against independent decodes by value across eager and CUDA-graph paths, both trims, both orderings and several context sizes, which is the right oracle. Singleton and equal-length batches are unchanged. Reviewed statically; fork head, no PR code executed. No lane had run on this head, so I added ready; merge on green.