From bfe8967ca29bf4253a30177f053759f0b6cab158 Mon Sep 17 00:00:00 2001 From: boatman <1930807094@qq.com> Date: Sat, 29 Aug 2026 12:45:24 +0800 Subject: [PATCH] [Bugfix][StepAudio2] Fix async-chunk metadata mismatch and boundary glitches (#5917) Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com> --- .../test_step_audio2_token2wav_async_chunk.py | 21 +++++++++++++++---- .../step_audio2/step_audio2_token2wav.py | 19 ++++++++++++----- 2 files changed, 31 insertions(+), 9 deletions(-) diff --git a/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py b/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py index 050a5b724..53a90452e 100644 --- a/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py +++ b/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -76,8 +77,8 @@ def test_step_audio2_token2wav_async_chunk_batch_guard(): model._forward_async_chunk( input_ids=torch.tensor([1, 2, 3]), runtime_additional_information=[ - {"left_context_size": 0}, - {"left_context_size": 1}, + {"meta": {"left_context_size": 0}}, + {"meta": {"left_context_size": 1}}, ], ) @@ -116,7 +117,7 @@ def test_step_audio2_token2wav_async_chunk_last_chunk_resets_state(monkeypatch): out = model.forward( input_ids=torch.tensor([10, 11, 12]), positions=torch.tensor([0, 1, 2]), - runtime_additional_information=[{"left_context_size": 1}], + runtime_additional_information=[{"meta": {"left_context_size": 1}}], ) audio_list = out.multimodal_outputs["model_outputs"] @@ -126,12 +127,24 @@ def test_step_audio2_token2wav_async_chunk_last_chunk_resets_state(monkeypatch): assert calls == {"setup": 1, "stream": 1, "reset": 1} +def test_step_audio2_token2wav_routes_non_final_chunk_to_async_forward(): + model = _build_model() + model._forward_async_chunk = MagicMock() + + model.forward( + input_ids=torch.tensor([10, 11, 12]), + positions=torch.tensor([0, 1, 2]), + runtime_additional_information=[{"meta": {"left_context_size": 0}}], + ) + model._forward_async_chunk.assert_called_once() + + def test_step_audio2_token2wav_async_chunk_empty_eof_returns_zero_chunk(): model = _build_model() out = model._forward_async_chunk( input_ids=torch.tensor([], dtype=torch.int64), - runtime_additional_information=[{"left_context_size": 1}], + runtime_additional_information=[{"meta": {"left_context_size": 1}}], ) audio_list = out.multimodal_outputs["model_outputs"] 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 6e9236766..e9ca78b21 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 @@ -44,6 +44,14 @@ from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( logger = init_logger(__name__) +def _get_left_context_size(info: dict[str, Any]) -> Any | None: + """Return the async-chunk marker from a nested metadata payload.""" + meta = info.get("meta") + if not isinstance(meta, dict): + return None + return meta.get("left_context_size") + + def fade_in_out( fade_in_mel: torch.Tensor, fade_out_mel: torch.Tensor, @@ -540,7 +548,8 @@ class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP): additional: Additional information dict (deprecated) runtime_additional_information: Per-request dicts from the async_chunk processor. Each dict contains ``audio_tokens`` - (list[int]) and ``last_chunk`` (bool). + (list[int]) and ``meta.left_context_size`` as the + last-chunk marker (0 = non-final, 1 = final). **kwargs: Other kwargs Returns: @@ -553,7 +562,7 @@ class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP): # [{}]), so a bare truthiness check would incorrectly enter this # branch in synchronous mode. if runtime_additional_information and any( - "left_context_size" in info for info in runtime_additional_information + _get_left_context_size(info) is not None for info in runtime_additional_information ): return self._forward_async_chunk(input_ids, runtime_additional_information, **kwargs) @@ -622,15 +631,15 @@ class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP): async_chunk payload to ``request.additional_information``). """ # --- batch=1 guard --- - batch_size = sum(1 for info in runtime_additional_information if "left_context_size" in info) + batch_size = sum(1 for info in runtime_additional_information if _get_left_context_size(info) is not None) if batch_size != 1: raise RuntimeError( f"Token2Wav async_chunk only supports batch=1, got {batch_size}. " "Batch>1 requires framework support for batch_req_ids." ) - info = next(info for info in runtime_additional_information if "left_context_size" in info) - last_chunk = info.get("left_context_size", 0) == 1 + info = next(info for info in runtime_additional_information if _get_left_context_size(info) is not None) + last_chunk = _get_left_context_size(info) == 1 # --- Manage single stream state --- # If the previous request was preempted (setup_done but not finished),