[Perf] [Qwen3-TTS] Keep audio_codes and last_talker_hidden on GPU to eliminate per-step sync stalls

[DomBrown]

Purpose

Keep audio_codes and last_talker_hidden on GPU in the Qwen3-TTS talker's per-step intermediate buffer, eliminating two cudaStreamSynchronize stalls per AR decode step.

Previously, both tensors were copied to CPU immediately after production via .to("cpu"), which blocked the CPU for ~2ms each while waiting for all pending GPU kernels to drain. Neither tensor is consumed on CPU during the decode loop — audio_codes is only used by make_omni_output (torch.cat on GPU) and last_talker_hidden is immediately transferred back to GPU in the next step's preprocess. The D2H transfer for audio_codes is deferred to sample_tokens, where hidden_states.to("cpu") already synchronizes the stream, making the subsequent copy zero-cost.

This uses the existing gpu_resident_buffer_keys mechanism (already used by Qwen3-Omni for last_talker_hidden) to opt specific keys into GPU-resident storage via _update_intermediate_buffer.

Test Plan

Performance benchmark

Test Result

Tested on NVIDIA RTX Pro 6000 Blackwell Server Edition

Command:

vllm serve Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml --omni --port 8000 --trust-remote-code

time vllm-omni bench serve \
    --backend openai-audio-speech \
    --endpoint /v1/audio/speech \
    --model "$MODEL" \
    --dataset-name hf \
    --dataset-path philschmid/mt-bench \
    --skip-chat-template \
    --num-prompts 10 \
    --max_concurrency 1 \
    --percentile-metrics e2el,audio_rtf,audio_duration \
    --extra-body '{"voice": "Vivian", "instructions": "Speak with great enthusiasm", "language": "English"}'

Before

============ Serving Benchmark Result ============
Successful requests:                     10        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  142.50    
Request throughput (req/s):              0.07      
Peak concurrent requests:                2.00      
----------------End-to-end Latency----------------
Mean E2EL (ms):                          14249.69  
Median E2EL (ms):                        9697.36   
P99 E2EL (ms):                           38141.95  
================== Text Result ===================
Total input tokens:                      785       
Total generated tokens:                  0         
Output token throughput (tok/s):         0.00      
Peak output token throughput (tok/s):    1.00      
Peak concurrent requests:                2.00      
Total Token throughput (tok/s):          5.51      
================== Audio Result ==================
Total audio duration generated(s):       246.96    
Total audio frames generated:            5927040   
Audio throughput(audio duration/s):      1.73      
-----------------Real Time Factor-----------------
Mean AUDIO_RTF:                          0.58      
Median AUDIO_RTF:                        0.58      
P99 AUDIO_RTF:                           0.59      
------------------Audio Duration------------------
Mean AUDIO_DURATION (s):                 24.70     
Median AUDIO_DURATION (s):               16.80     
P99 AUDIO_DURATION (s):                  66.04     
==================================================

After

============ Serving Benchmark Result ============
Successful requests:                     10        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  131.47    
Request throughput (req/s):              0.08      
Peak concurrent requests:                2.00      
----------------End-to-end Latency----------------
Mean E2EL (ms):                          13146.68  
Median E2EL (ms):                        8844.79   
P99 E2EL (ms):                           38730.45  
================== Text Result ===================
Total input tokens:                      785       
Total generated tokens:                  0         
Output token throughput (tok/s):         0.00      
Peak output token throughput (tok/s):    1.00      
Peak concurrent requests:                2.00      
Total Token throughput (tok/s):          5.97      
================== Audio Result ==================
Total audio duration generated(s):       267.36    
Total audio frames generated:            6416640   
Audio throughput(audio duration/s):      2.03      
-----------------Real Time Factor-----------------
Mean AUDIO_RTF:                          0.49      
Median AUDIO_RTF:                        0.49      
P99 AUDIO_RTF:                           0.50      
------------------Audio Duration------------------
Mean AUDIO_DURATION (s):                 26.74     
Median AUDIO_DURATION (s):               18.04     
P99 AUDIO_DURATION (s):                  79.38     
==================================================

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 73aed6091a

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[DomBrown]

Looks like some of this overlapped with #1852

[linyueqian]

Tested locally on A100-80GB with Qwen3-TTS-12Hz-0.6B-CustomVoice, single-request sequential benchmark (5 runs after warmup):

Metric	main (with #1852)	This PR	Delta
E2EL mean	2037 ms	1772 ms	-13%
E2EL median	1976 ms	1767 ms	-11%
RTF mean	0.510	0.445	-13%
RTF median	0.508	0.436	-14%

Confirms the improvement. Already cleanly rebased on top of #1852. The talker-side change is just adding "audio_codes" to the existing gpu_resident_buffer_keys set. The bulk of the value comes from the runner-side changes (deferred D2H for code_predictor_codes in gpu_model_runner.py and hoisted multimodal copy in gpu_ar_model_runner.py).

LGTM.

[lishunyang12]

LGTM