[Perf][TTS] step_audio2: keep streaming tokens on-device and batch audio-feature length syncs

[JuanPZuluaga]

Purpose

Two small changes that cut host syncs on the hot path, with no change to outputs:

• token2wav streaming passed audio tokens through a Python list every chunk (input_ids.flatten().cpu().tolist(), then rebuilt a device tensor). Now it passes the tensor straight through and stays on device.
• The thinker's audio-feature split sliced with one 0-d CUDA tensor per item, which syncs once per item. Now it reads all the lengths to host once with a single .tolist().

Numbers

Measured with a device-to-host (D2H) sync counter, greedy decoding, output checked byte-for-byte.

• Speech-in path (ASR, 8 concurrent audio to text requests): D2H syncs drop from 17 to 2 (−88%). Transcripts identical before and after.
• Text-to-speech path: unchanged. The audio encoder only runs when there's audio input, so text to speech never reaches this code checks.

Test Plan

vLLM Version:

vLLM-Omni Commit:

Test Result

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)

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

[linyueqian]

@QiuMike can you take a look?

[lishunyang12]

I like keeping this on-device 👍 — the one thing I'm not sure about: when audio_tokens already arrives as int32 on self.device, .to(...) is a no-op and token_tensor ends up aliasing the caller's input_ids storage. Today the ids come in as int64 so the dtype cast copies, but if inference_chunk ever writes into token in place that'd quietly corrupt input_ids. A trailing .contiguous() (or .clone()) would make it robust regardless of the incoming dtype.

[JuanPZuluaga]

Added .clone() so token_tensor never aliases the caller's input_ids, regardless of incoming dtype.

[lishunyang12]

Nice N-syncs-to-1 win here. The one thing I want to confirm is that _process_audio_input never runs inside a CUDA-graph capture region — .tolist() is a hard host sync. The old per-item slice synced too so this is no regression, I just want to be sure that assumption still holds on this path.

[JuanPZuluaga]

Correct: _process_audio_input runs in the eager prefill path, not inside a CUDA-graph capture region, so the .tolist() host sync is safe; I added a comment noting this.

[JuanPZuluaga]

Fixed and added some comments @lishunyang12

[JuanPZuluaga]

@lishunyang12 please let me know if there's anything else to be done here.
thanks.

[linyueqian]

@vllm-omni-review-bot

[vllm-omni-review-bot]

Omni ReviewBot review

The on-device streaming tensor path preserves the prior [1, T] int32 input shape and non-aliasing behavior, while the batched .tolist() conversion preserves equivalent feature slicing. No concrete correctness defects found.

No actionable findings.

[linyueqian]

Ran an A/B on an idle L20X box (vLLM 0.26.0, step_audio_2_async_chunk.yaml, Step-Audio-2-mini, TP=2). Baseline is pr5067^2 = 8e44e780, so the two legs differ by exactly this PR's 17 lines.

leg	median	min	output
baseline 8e44e780	2.87s	2.83s	316844 B, 3/3 ok
PR head ce3bcada	2.88s	2.84s	316844 B, 3/3 ok

No regression, and no measurable speedup on this workload — but that's a statement about my probe, not about the patch. Text→speech never enters _process_audio_input, so the thinker .tolist() change wasn't exercised at all, and max_num_seqs: 1 in that config means no batching pressure on the D2H syncs the token2wav change removes. If you have a workload where these showed up in a profile (audio-in requests? concurrency? long chunks?), I'm happy to re-run against that — it would make the win visible instead of me reporting a null result.

Code reads fine to me: audio_tokens as a device tensor has no list-assuming consumer left in _forward_async_chunk (only numel() and the pass-through), and the .clone() guarding against aliasing input_ids when the ids already arrive int32-on-device is a good catch.

Unrelated to this PR, in case it's useful: serving step_audio2 from a clean env needed librosa, soundfile, s3tokenizer, hyperpyyaml and step-audio2 on top of the base install, and the failure modes were misleading — first a pydantic ValidationError: Model architectures [...] failed to be inspected, then a server that reports startup complete and only dies on the first request with ImportError: no such class as cosyvoice2.flow.flow.CausalMaskedDiffWithXvec. Only some of those are declared (in the dev extra).

[linyueqian]

Verified by A/B against pr5067^2: no regression (2.88s vs 2.87s median, 3/3 ok both legs, identical output size). Flat on my workload rather than faster, but the probe was text-to-speech so it never entered _process_audio_input, and max_num_seqs: 1 means no batching pressure on the syncs this removes — so that's a limit of my measurement, not evidence against the change. Code reads clean: the device-tensor audio_tokens has no list-assuming consumer left, and the .clone() against aliasing input_ids is a good catch.