From 12a5f6fbf3ac08ffeec8cb4a2fc05344b8d1e4cf Mon Sep 17 00:00:00 2001 From: Juan Pablo Zuluaga <46724788+JuanPZuluaga@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:24:25 +0200 Subject: [PATCH] [Perf][TTS] step_audio2: keep streaming tokens on-device and batch audio-feature length syncs (#5067) Signed-off-by: JuanPZuluaga Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .../models/step_audio2/step_audio2_thinker.py | 6 +++++- .../models/step_audio2/step_audio2_token2wav.py | 16 ++++++++++++---- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py index c92f23978..b56229cf6 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py @@ -758,7 +758,11 @@ class StepAudio2ThinkerForConditionalGeneration(nn.Module, SupportsMultiModal, S audio_feature_lens = (audio_lens - 1) // 2 + 1 - audio_feature_list = [audio_features[i, : audio_feature_lens[i]] for i in range(audio_features.size(0))] + # Single host sync; slicing with per-item 0-d cuda tensors syncs once per item. + # Safe: _process_audio_input runs in the eager multimodal-encoder prefill path, + # never inside a decode CUDA-graph capture, so this .tolist() host sync is fine. + audio_feature_lens_cpu = audio_feature_lens.tolist() + audio_feature_list = [audio_features[i, : audio_feature_lens_cpu[i]] for i in range(audio_features.size(0))] return audio_feature_list diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py index b116b6bcd..9666a4f5a 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py @@ -364,7 +364,7 @@ class StepAudio2Token2WavCore(nn.Module): def stream_chunk_for( self, - audio_tokens: list[int], + audio_tokens: list[int] | torch.Tensor, prompt_wav: str, last_chunk: bool, state: _StreamState, @@ -377,7 +377,14 @@ class StepAudio2Token2WavCore(nn.Module): self.cache[prompt_wav] = self._prepare_prompt(prompt_wav) _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav] - token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device) + if isinstance(audio_tokens, torch.Tensor): + # Already on device; avoids the D2H->H2D round-trip of rebuilding from a list. + # clone() so token_tensor never aliases the caller's input_ids: when the ids + # already arrive int32-on-device the .to() is a no-op and would share storage, + # so an in-place write inside inference_chunk could corrupt input_ids. + token_tensor = audio_tokens.reshape(1, -1).to(device=self.device, dtype=torch.int32).clone() + else: + token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device) with torch.amp.autocast( str(self.device.type), @@ -651,10 +658,11 @@ class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP): state = self._stream_states[0] # --- Extract audio tokens --- - audio_tokens = input_ids.flatten().cpu().tolist() + # Kept on device; stream_chunk_for takes the tensor directly. + audio_tokens = input_ids.flatten() # Empty chunk (e.g. EOF with no audio tokens) - if not audio_tokens: + if audio_tokens.numel() == 0: if state.setup_done: self.token2wav.reset_stream_for(state) else: