[Model] Optimize MOSS-TTS Local batched execution and streaming codec

[Sy0307]

MOSS-TTS Local repeatedly evaluates the growing 12-codebook prefix for each audio frame and performs per-request output and codec work that can be batched. This change reuses frame-local K/V, batches owned output snapshots, and accelerates streaming codec attention and terminal chunks while preserving request boundaries and output lengths.

For Local requests with matching ref_audio and ref_text, the adapter now builds the official continuation prompt and emits newly generated audio. Audio-only requests retain generation mode. Long English Seed-TTS outputs are scored in full instead of silently truncating ASR input at 30 seconds.

Changes

• Reuse frame-local K/V with dynamic compilation; capture mixed batches up to 512 tokens in the H200 recipe.
• Batch tensor snapshots with owned storage before graph buffers can be reused, and avoid unnecessary intermediate-stage output materialization.
• Add the Triton streaming-attention backend and coalesce compatible terminal codec chunks while preserving valid lengths and resetting completed slots.
• Keep the Local reference-transcript fix in main's extracted adapter implementation; consolidate its four continuation/fallback cases in one parameterized test.
• Reconcile nullable return, payload, and metric annotations in touched interfaces so the current type checks pass. No new sampling or stopping policy is introduced.

Quality validation

Compared main 968f3e37 plus the same transcript-adaptation fix against PR 28dd49a5, on one H200 GPU 6, c256, capacity 256 and 32 GiB KV. Two complete 1,088-item English launches per arm ran in main/PR/PR/main order. All generated clips, including anomalies, remain in the metrics.

Condition	WER, runs	Mean WER	Mean UTMOS	Outputs >30s
main + transcript fix	1.683%, 1.331%	1.507%	4.0147	0/2176
PR	1.440%, 1.733%	1.587%	4.0132	1/2176

The paired utterance-bootstrap PR-minus-main intervals cross zero: WER +0.0795 percentage points, 95% CI[-0.1851, +0.3846]; UTMOS-0.00154,95% CI[-0.01490, +0.01158]. These are conditional on the measured launches, not proof of zero rare failures.

A follow-up weighted stress test retained 2,048 requests per arm, including 1,024 repetitions of the sentence responsible for the extreme PR output:

Condition	WER	UTMOS	Outputs >30s
main + transcript fix	1.150%	4.0019	1/2048
PR	1.189%	4.0009	1/2048

Main reproduced the same “The team colors were teal, red and black.” 163.76s failure without these performance changes. A suspected 11-word omission was independently transcribed in full by Whisper-small using the original audio without a target-text prompt; the original Whisper-large-v3 benchmark scores were not selectively corrected. Earlier focused HF and two full HF runs also reproduced native-model repetition/premature stopping. The extreme case is therefore not unique to this PR, but the same 163.76-second case has not been reproduced in official HF. Its root cause remains unresolved; main reproduction alone does not rule out a shared integration issue. This PR does not claim to eliminate long-output failures.

The primary comparisons use identical frozen CPU reference codes, with an actual-encoder equality check on first use in each service. Two-token cache warmups are separate from the measured requests with a 2,048-token generation budget. Main uses its SDPA codec; PR exercises the added Triton backend. Scoring uses full-length Whisper-large-v3 with EnglishTextNormalizer and UTMOS with verified PCM scaling. No failed or long clips are discarded, and the focused stress set is not a population failure-rate estimate.

Performance validation

Fresh end-to-end comparison: PR 28dd49a5 versus current main ab561485 plus only the same reference-transcript adapter fix. One physical H200, both stages on that card, c256, capacity 256/stage, 32 GiB talker KV, 512-token talker graph buckets and identical codec batch buckets. Main uses its SDPA codec; PR enables the introduced Triton backend. Payload hashes, software versions, CPU affinity and all runtime source hashes were verified. Current main contains four commits after the PR merge base; this is a comparison of these frozen versions, not a per-change ablation.

Metric	main + transcript fix	PR	Change
Steady throughput	209.91 audio-s/s	299.87 audio-s/s	+42.85%
Throughput including fill and drain	203.84 audio-s/s	293.10 audio-s/s	+43.79%
First audio P95	787.87 ms	632.74 ms	-19.69%
Completion P95	9.01 s	6.22 s	-30.99%

Three independent server launches per arm, two 90-second trials per launch. Per-launch mean throughput (audio-s/s): main 209.14 / 210.54 / 210.05; PR 299.31 / 299.99 / 300.29. The launch-mean spans are 0.66% and 0.33%, respectively. PR launches ran first, followed by main launches on the same card; the order does not eliminate temporal drift through interleaving. Results are specific to c256 and do not establish gains at every concurrency.

Each launch warms all 123 references using normal-length requests, primes c256 for 20 seconds and drains, then runs two trials with 10-second fill, 90-second measurement and normal completion of outstanding requests. Each measured window has three adjacent 30-second subwindows; those are not independent launches. The fixed workload cycles 128 Seed-TTS EN texts with matching reference audio and transcripts, official sampling, and max_new_tokens=1024. Throughput counts actual received 48 kHz stereo s16 PCM bytes / 192000 / elapsed seconds. Full-run throughput includes fill and drain; latency percentiles pool requests started inside measured windows. The shared CPU reference-code cache and first-encode equality check described above are used only to avoid repeating untimed reference encoding.

All requests complete normally; client/server completion counts agree. There are no request errors, new reference-cache stores during measurement, or JIT warnings recorded by the monitor after formal timing begins. Timing records retain every long output: main has 7/25,735 outputs above 30 seconds, PR 10/36,218; each arm has 7 outputs near the 1024-token budget, with maximum audio duration 81.84 seconds. These counts include fill/drain but exclude warmup. This performance rerun saves timing and byte-count records, not PCM waveforms; it adds no new WER/UTMOS or quality-passing claim. The separate quality results above use a 2048-token budget and saved audio.

Fresh real-weight Local Depth microbenchmark at batch 256, including all 12 audio codebooks, measured warmed CUDA-graph replays in counterbalanced order (four measurements per variant, 50 replays each): 8.739 ms to 4.185 ms per frame batch, 52.1% lower time. This is a submodule measurement, not a claim of 2x end-to-end throughput. Individual optimization gains are not added together.

Checks

• Refreshed on main; merge conflict in the extracted MOSS adapter resolved.
• Full changed-file pre-commit checks pass, including mypy, formatting and SPDX; no checker suppression added.
• Local regression groups: 77 tests including CUDA, 64 entrypoint tests, 62 AR-runner CPU tests passed. These are executed cases, not 203 new tests.
• DCO/identity and whitespace checks pass. The earlier submission/cancellation, backpressure and SHM runtime fixes remain outside this PR.

With matching Linux vLLM/PyTorch and pytest installed, CPU adapter validation is:

python -m pytest tests/entrypoints/openai_api/test_tts_adapter.py -q
python -m pytest tests/entrypoints/openai_api/test_tts_adapter.py -m 'core_model and cpu' --run-level core_model

CUDA regression coverage additionally exercises the streaming attention, terminal batching and owned snapshot cases. Real-weight validation used model revision be7766a6735b98bd793f7c79fb720b4d0f5d13b8, vLLM 0.28.0, PyTorch 2.13.0/CUDA 13, Transformers 5.15.1 and Triton 3.7.1.

CI follow-up

Current head incorporates upstream commit 09b720b0 (#7314), the project quarantine for Qwen3-Omni Server VAD issue #7279. Build 14872 failed that exact test with input_backpressure; this incorporates the existing upstream handling rather than changing MOSS or claiming to fix Server VAD. All runtime files under vllm_omni/ remain byte-identical to the measured head 28dd49a5.

The other failed job, MiniMax-H3 L3, passed its FL2VA and REF2VA cases before its log stopped during Turbo LoRA DLO initialization. Its exit status was -1, without a Python traceback or agent-reported completion; the available log does not establish OOM or a code regression. It is retained for verification in the new CI build, without speculative model or test changes.

[hsliuustc0106]

This PR touches vllm_omni/model_executor/, tests/model_executor/, vllm_omni/entrypoints/, vllm_omni/worker/, .buildkite/cuda/ (18 files). Based on CODEOWNERS coverage of the changed files, the most-related reviewers appear to be:

@tzhouam @linyueqian @NickCao

Could one of you take a look when you get a chance? Thanks!

[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/ar_runtime.md, docs/design/module/model_integration.md.

Module owners: @tzhouam @Gaohan123 @fake0fan

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

@Sy0307, 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 dedb04522567 produced:

• Priority: high. Prompt maintainer attention is suggested.

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

[vllm-omni-review-bot]

Omni ReviewBot: three questions on the performance claim

@Sy0307 this PR reads as a performance or value claim:

• title: [Model] Optimize MOSS-TTS Local batched execution and streaming codec

Before the full evidence checklist, three short questions:

1. Bottleneck — what is the current bottleneck, and which profile, trace or per-stage measurement shows it?
2. Value — what does the change buy the user or the system (latency, throughput, memory, cost), and at which workload?
3. A/B or ablation — is there a same-workload, same-head/config comparison that isolates each main claim on its own? For stacked optimizations, one number per item rather than a blended delta.

When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution.