From 6d50348ef19444f06cd1293f8afa235c385f2b0f Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:59:41 -0700 Subject: [PATCH] [Perf][Qwen3-TTS] Skip per-step hidden-state D2H via hidden pooler payload opt-out (#4879) Signed-off-by: Yueqian Lin --- .../model_executor/models/qwen3_tts/qwen3_tts_talker.py | 6 ++++++ vllm_omni/worker/gpu_ar_model_runner.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index 347c76dfa..2d28a0730 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -309,6 +309,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # tensor read; postprocess receives the tail-only slice instead, which # avoids ~18 ms merge + ~6 ms write per step (Sy0307 profile, #3665). self.requires_full_prefix_cached_hidden_states = False + # Stage 1 (code2wav) consumes runtime audio codes, not hidden states, + # so the inter-stage pooler payload never needs a CPU hidden-states + # view. Opting out lets the runner skip the per-step blocking + # ``hidden_states[:n].to("cpu")`` (measured at 24% of stage-0 on-CPU + # samples at c=64 on H20). + self.omni_pooler_payload_include_hidden = False # ``codes.audio`` is only needed for future prefix-hit reconstruction # after a request has produced codec rows. Keep per-step rows on GPU and # materialize the CPU OmniTensorPrefixCache entry once at completion. diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index a5e4a9bd1..59bfbb855 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -1309,7 +1309,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # prefix caching but the downstream pooler payload path still # needs a CPU hidden-states view. Materialize it synchronously # in that case; the legacy behavior is preserved. - if hs_for_cache is None: + if hs_for_cache is None and self._model_omni_pooler_payload_include_hidden(): hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous() slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu self.omni_prefix_cache.schedule_async_write(