diff --git a/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py b/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py new file mode 100644 index 000000000..85c0e8b56 --- /dev/null +++ b/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID +from vllm_omni.model_executor.models.mimo_audio.mimo_audio_code2wav import ( + AudioStreamerConfig, + MiMoAudioToken2WavForConditionalGenerationVLLM, + flat_codec_group_element_count, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_GROUP = 4 +_AC = 8 +_GROUP_WIDTH = flat_codec_group_element_count(_GROUP, _AC) +_FTP = 2 * 2 * 240 # frames_per_token from mocked tokenizer.config + + +def _codes_ns(empty: int = 555, eostm: int = 666): + return SimpleNamespace(empty=empty, eostm=eostm) + + +def _make_dummy_code_tensor() -> torch.Tensor: + """Pad-only talker dummy; matches _check_dummy_code_tensor.""" + t = torch.zeros(flat_codec_group_element_count(_GROUP, _AC), dtype=torch.long) + t = t.view(_GROUP, _AC + 1) + t[:, 0] = TALKER_CODEC_PAD_TOKEN_ID + return t.view(-1) + + +def _make_valid_flat_codes(num_groups: int = 1, *, empty_id: int = 555) -> torch.Tensor: + """Flat layout that extract_audio_code_tensor accepts (text col == empty).""" + parts = [] + for _ in range(num_groups): + g = torch.zeros(_GROUP, _AC + 1, dtype=torch.long) + g[0, 0] = empty_id + g[:, 1:] = torch.randint(1, 50, (_GROUP, _AC)) + parts.append(g.reshape(-1)) + return torch.cat(parts) + + +def _make_invalid_flat_immediate_eostm(eostm_id: int = 666) -> torch.Tensor: + g = torch.zeros(_GROUP, _AC + 1, dtype=torch.long) + g[0, 0] = eostm_id + return g.reshape(-1) + + +def _minimal_model(): + """Avoid __init__ (HF tokenizer paths); only fields used by _batch_decode_waveforms.""" + model = object.__new__(MiMoAudioToken2WavForConditionalGenerationVLLM) + model.device = torch.device("cpu") + model.config = SimpleNamespace(group_size=_GROUP, audio_channels=_AC) + model.streamer_config = AudioStreamerConfig(group_size=_GROUP, audio_channels=_AC) + model.codes = _codes_ns() + + decode_vq = Mock( + side_effect=lambda audio_codes: torch.ones( + audio_codes.shape[1], + 7, + dtype=torch.float32, + device=audio_codes.device, + ) + ) + decoder = Mock() + + audio_tok = SimpleNamespace( + encoder=SimpleNamespace(decode_vq=decode_vq), + decoder=decoder, + config=SimpleNamespace(avg_pooler=2, stride_size=2, hop_length=240), + ) + model._tokenizer_service = SimpleNamespace(audio_tokenizer=audio_tok) + return model, audio_tok + + +def test_batch_decode_waveforms_empty_input_list(): + """Empty input list returns a single zero-length float32 tensor on model device.""" + model, _ = _minimal_model() + out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, []) + assert len(out) == 1 + assert out[0].dtype == torch.float32 + assert out[0].numel() == 0 + assert out[0].device == model.device + + +def test_batch_decode_waveforms_single_vs_multiple_decoder_shapes(): + """Single and multi-request batches produce correctly shaped packed hidden states and trimmed waveforms.""" + model, audio_tok = _minimal_model() + decoder = audio_tok.decoder + + # Single valid request: decoder output rank-3 for double squeeze path + flat = _make_valid_flat_codes(1) + decoder.return_value = torch.ones(1, 1, 4 * _FTP + 100, dtype=torch.float32) + out1 = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, [flat]) + decoder.assert_called_once() + packed_hs, input_lengths = decoder.call_args[0] + assert packed_hs.shape == (4, 7) # T from decode_vq mock + assert input_lengths.tolist() == [4] + assert out1[0].shape == (4 * _FTP,) + assert out1[0].dtype == torch.float32 + + decoder.reset_mock() + a = _make_valid_flat_codes(1) + b = _make_valid_flat_codes(2) + decoder.return_value = torch.ones(2, 1, 8 * _FTP + 100, dtype=torch.float32) + out2 = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, [a, b]) + decoder.assert_called_once() + packed_hs2, input_lengths2 = decoder.call_args[0] + assert packed_hs2.shape == (4 + 8, 7) # 1 group -> T=4; 2 groups -> T=8 + assert input_lengths2.tolist() == [4, 8] + assert len(out2) == 2 + assert out2[0].shape == (4 * _FTP,) + assert out2[1].shape == (8 * _FTP,) + + +def test_batch_decode_waveforms_mixed_valid_invalid_requests(): + """Mixed valid and invalid requests: invalid slots get empty tensors, valid slots get decoded waveforms.""" + model, audio_tok = _minimal_model() + valid_a = _make_valid_flat_codes(1) + valid_b = _make_valid_flat_codes(1) + dummy = _make_dummy_code_tensor() + bad_extract = _make_invalid_flat_immediate_eostm(eostm_id=model.codes.eostm) + too_short = torch.tensor([1, 2, 3], dtype=torch.long) + + audio_tok.decoder.return_value = torch.ones(2, 1, 5000, dtype=torch.float32) + + inputs = [ + None, + torch.tensor([], dtype=torch.long), + dummy, + bad_extract, + too_short, + valid_a, + valid_b, + ] + out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, inputs) + + assert len(out) == len(inputs) + for i in range(5): + assert out[i].numel() == 0, f"index {i} should be empty" + assert out[5].shape == (4 * _FTP,) + assert out[6].shape == (4 * _FTP,) + audio_tok.decoder.assert_called_once() + packed_hs, input_lengths = audio_tok.decoder.call_args[0] + assert packed_hs.shape[0] == 8 + assert input_lengths.tolist() == [4, 4] + + +def test_batch_decode_waveforms_all_invalid_returns_per_request_empty(): + """All-invalid batch skips decoder entirely and returns empty tensors for every slot.""" + model, audio_tok = _minimal_model() + out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms( + model, + [None, _make_dummy_code_tensor(), torch.tensor([], dtype=torch.long)], + ) + assert len(out) == 3 + assert all(t.numel() == 0 for t in out) + audio_tok.decoder.assert_not_called() + + +def test_batch_decode_waveforms_output_shape_trim_when_decoder_returns_extra_samples(): + """Decoder output longer than valid_len is trimmed to the exact expected waveform length.""" + model, audio_tok = _minimal_model() + flat = _make_valid_flat_codes(1) + # Longer than valid_len so branch wav = wav[:valid_len] runs + audio_tok.decoder.return_value = torch.ones(1, 1, 10_000, dtype=torch.float32) + out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, [flat]) + assert out[0].dim() == 1 + assert out[0].numel() == 4 * _FTP + assert out[0].dtype == torch.float32 + + +def test_batch_decode_waveforms_multi_request_trims_each_row_when_decoder_returns_extra(): + """Else-branch split: per-request wav[:valid_len] when decoder pads each batch row.""" + model, audio_tok = _minimal_model() + a = _make_valid_flat_codes(1) + b = _make_valid_flat_codes(2) + audio_tok.decoder.return_value = torch.ones(2, 1, 10_000, dtype=torch.float32) + out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, [a, b]) + assert len(out) == 2 + assert out[0].shape == (4 * _FTP,) + assert out[1].shape == (8 * _FTP,) + assert out[0].dtype == torch.float32 + assert out[1].dtype == torch.float32 + + +def test_batch_decode_waveforms_valid_only_at_edges_maps_to_correct_indices(): + """Tensor packing order must match valid_indices when invalid requests are in the middle.""" + model, audio_tok = _minimal_model() + first = _make_valid_flat_codes(1) + last = _make_valid_flat_codes(2) + inputs = [ + first, + None, + _make_dummy_code_tensor(), + last, + ] + audio_tok.decoder.return_value = torch.ones(2, 1, 10_000, dtype=torch.float32) + out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, inputs) + assert len(out) == 4 + assert out[0].numel() == 4 * _FTP + assert out[1].numel() == 0 + assert out[2].numel() == 0 + assert out[3].numel() == 8 * _FTP + packed_hs, input_lengths = audio_tok.decoder.call_args[0] + assert packed_hs.shape[0] == 12 + assert input_lengths.tolist() == [4, 8] + + +def test_batch_decode_waveforms_output_shapes_1d_float32_for_all_slots(): + """Every slot is a 1-D float32 vector (empty or waveform), matching downstream expectations.""" + model, audio_tok = _minimal_model() + inputs = [_make_valid_flat_codes(1), None, _make_valid_flat_codes(1)] + audio_tok.decoder.return_value = torch.ones(2, 1, 5000, dtype=torch.float32) + out = MiMoAudioToken2WavForConditionalGenerationVLLM._batch_decode_waveforms(model, inputs) + assert len(out) == 3 + for i, t in enumerate(out): + assert t.dim() == 1, f"slot {i}" + assert t.dtype == torch.float32, f"slot {i}" + assert out[0].numel() == 4 * _FTP + assert out[1].numel() == 0 + assert out[2].numel() == 4 * _FTP diff --git a/tests/model_executor/stage_input_processors/test_mimo_audio_flush_remaining_codes.py b/tests/model_executor/stage_input_processors/test_mimo_audio_flush_remaining_codes.py new file mode 100644 index 000000000..4807ac627 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_mimo_audio_flush_remaining_codes.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.mimo_audio import _flush_remaining_codes + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _sentinel(): + return {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)} + + +def test_flush_remaining_codes_when_no_codes_accumulated_missing_request_id(): + """No entry for request_id: treat as empty, return finished sentinel with empty codes.""" + tm = SimpleNamespace(code_prompt_token_ids={}) + out = _flush_remaining_codes(tm, "missing", chunk_size=3, left_context_size=3) + assert out["code_predictor_codes"] == _sentinel()["code_predictor_codes"] + assert out["finished"].equal(_sentinel()["finished"]) + + +def test_flush_remaining_codes_when_no_codes_accumulated_empty_list(): + """Explicit empty accumulation list returns the same sentinel.""" + tm = SimpleNamespace(code_prompt_token_ids={"r": []}) + out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3) + assert out["code_predictor_codes"] == [] + assert out["finished"].item() is True + + +def test_flush_remaining_codes_partial_chunk_remaining(): + """ + length=7, chunk_size=3 -> remainder 1 frame in last chunk; context window uses + left_context + that partial (3+1), capped by length -> last 4 scalar entries. + """ + tm = SimpleNamespace( + code_prompt_token_ids={"r": [[1], [2], [3], [4], [5], [6], [7]]}, + ) + out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3) + assert out["finished"].item() is True + assert out["code_predictor_codes"] == [4, 5, 6, 7] + + +def test_flush_remaining_codes_when_length_is_exact_multiple_of_chunk_size(): + """length % chunk_size == 0 uses full last chunk as context_length (not remainder).""" + tm = SimpleNamespace( + code_prompt_token_ids={"r": [[1], [2], [3], [4], [5], [6]]}, + ) + out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3) + # context_length = chunk_size = 3, end_index = min(6, 6) -> all 6 + assert out["code_predictor_codes"] == [1, 2, 3, 4, 5, 6] + + +@pytest.mark.parametrize( + "length,chunk_size,left_context,expected_end_index", + [ + (2, 3, 3, 2), # chunk_length=2 -> context 2, min(2,5)=2 + (5, 3, 3, 5), # chunk_length=2 -> min(5,5)=5 + (7, 3, 3, 4), # chunk_length=1 -> min(7,4)=4 + (6, 3, 3, 6), # chunk_length=0 -> context 3, min(6,6)=6 + (10, 3, 3, 4), # chunk_length=1 -> min(10,4)=4 + (1, 5, 10, 1), # chunk_length=1 -> min(1,11)=1 + ], +) +def test_flush_remaining_codes_context_window_end_index( + length: int, chunk_size: int, left_context: int, expected_end_index: int +) -> None: + """Mirror _flush_remaining_codes context_length and end_index rules.""" + accumulated = [[i] for i in range(length)] + tm = SimpleNamespace(code_prompt_token_ids={"r": accumulated}) + out = _flush_remaining_codes(tm, "r", chunk_size=chunk_size, left_context_size=left_context) + expected_flat = list(range(length - expected_end_index, length)) + assert out["code_predictor_codes"] == expected_flat + assert out["finished"].item() is True diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py index 9528f0bfb..9acb81bce 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py @@ -814,11 +814,13 @@ class MiMoAudioForConditionalGeneration( ) ) - audio_tensor = self.generate_audio(code) + audio_result = self.token2wav(codes=code, **kwargs) + if isinstance(audio_result, OmniOutput): + return audio_result return OmniOutput( text_hidden_states=None, multimodal_outputs={ - "model_outputs": audio_tensor.reshape(1, -1) if audio_tensor is not None else audio_tensor + "model_outputs": audio_result.reshape(1, -1) if audio_result is not None else audio_result }, ) diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py index 495598bf7..7c2e87d5f 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py @@ -4,6 +4,7 @@ import os import time from collections.abc import Iterable from dataclasses import dataclass +from typing import Any import torch import torchaudio @@ -20,11 +21,21 @@ from vllm.v1.sample.sampler import Sampler from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID, MiMoAudioConfig from vllm_omni.model_executor.models.mimo_audio.modeling_audio_tokenizer import MiMoAudioTokenizer +from vllm_omni.model_executor.models.output_templates import OmniOutput logger = logging.getLogger(__name__) -DUMMY_CODE_SHAPE = 36 +def flat_codec_group_element_count(group_size: int, audio_channels: int) -> int: + """Flat token count for one MiMo talker codec group on the code2wav wire. + + The talker flattens a ``(group_size, audio_channels + 1)`` layout in column-major + order: one leading text/special column per row plus ``audio_channels`` RVQ code + columns. Each group therefore occupies ``group_size * (audio_channels + 1)`` + consecutive ids. This matches ``group_width`` in :func:`extract_audio_code_tensor` + and is used to count how many such groups are present in a flat ``codes`` tensor. + """ + return group_size * (audio_channels + 1) class MiMoAudioTokenizerWorker: @@ -308,7 +319,7 @@ def extract_audio_code_tensor( if flat_codes.numel() == 0: return None - group_width = group_size * (audio_channels + 1) + group_width = flat_codec_group_element_count(group_size, audio_channels) usable = (flat_codes.numel() // group_width) * group_width if usable == 0: return None @@ -377,6 +388,9 @@ def get_tokenizer_worker( class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): """Decode MiMo audio codes to waveform for the code2wav stage.""" + have_multimodal_outputs = True + enable_update_additional_information = True + def __init__(self, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -429,6 +443,23 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): * getattr(audio_tokenizer_config, "hop_length", 240) ) * self.config.group_size + connector_cfg = getattr(vllm_config.model_config, "stage_connector_config", None) + extra_cfg = ( + ( + connector_cfg.get("extra", connector_cfg) + if isinstance(connector_cfg, dict) + else getattr(connector_cfg, "extra", None) + ) + if connector_cfg + else None + ) + self._codec_chunk_frames = int(extra_cfg.get("codec_chunk_frames", 3)) if isinstance(extra_cfg, dict) else 3 + if self._codec_chunk_frames <= 0: + raise ValueError(f"codec_chunk_frames must be positive, got {self._codec_chunk_frames}") + self._codec_left_context_frames = ( + int(extra_cfg.get("codec_left_context_frames", 3)) if isinstance(extra_cfg, dict) else 3 + ) + def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], @@ -438,6 +469,61 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): # Decoder has no trainable weights to load in this stage. return set() + @staticmethod + def _split_flat_codes_for_requests( + ids: torch.Tensor, + seq_token_counts: list[int] | None, + runtime_additional_information: list[dict[str, Any]] | None, + ) -> list[torch.Tensor]: + """Split flat codec token ids per request. + + Prefer ``code_flat_numel`` from async-chunk / connector payload (same idea as + qwen3 omni/tts passing decode metadata in additional_information). Fall back + to runner-provided ``seq_token_counts`` (no forward_context / ubatch slicing). + """ + n = ids.numel() + if n == 0: + return [ids] + + if runtime_additional_information and all( + isinstance(info.get("code_flat_numel"), int) and int(info["code_flat_numel"]) > 0 + for info in runtime_additional_information + ): + sizes = [int(info["code_flat_numel"]) for info in runtime_additional_information] + if sum(sizes) == n: + parts: list[torch.Tensor] = [] + offset = 0 + for sz in sizes: + parts.append(ids[offset : offset + sz]) + offset += sz + return parts + + if seq_token_counts is not None and len(seq_token_counts) > 1: + boundaries = [0] + for count in seq_token_counts: + boundaries.append(boundaries[-1] + count) + return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))] + + return [ids] + + @staticmethod + def _mimo_codec_runtime_lists( + num_req: int, + runtime_additional_information: list[dict[str, Any]] | None, + ) -> tuple[list[int | None], list[int | None]]: + """Per-request ``left_context_size`` / ``codec_chunk_frames`` from runtime buffer.""" + left_frames: list[int | None] = [None] * num_req + chunk_frames: list[int | None] = [None] * num_req + if not runtime_additional_information: + return left_frames, chunk_frames + for i in range(min(num_req, len(runtime_additional_information))): + info = runtime_additional_information[i] + if "left_context_size" in info: + left_frames[i] = int(info["left_context_size"]) + if "codec_chunk_frames" in info: + chunk_frames[i] = int(info["codec_chunk_frames"]) + return left_frames, chunk_frames + def chunked_decode_streaming( self, codes: torch.Tensor, @@ -453,9 +539,12 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): wav_chunk = self._decode_waveform_from_codes(codes) if wav_chunk.numel() == 0: return wav_chunk - # Infer number of codec frames: flat format has 36 elements per frame (9*4 with prepend) num_flat = codes.numel() if codes.ndim == 1 else codes.shape[-1] - num_chunks = num_flat // 36 + elt_per_group = flat_codec_group_element_count( + self.streamer_config.group_size, + self.streamer_config.audio_channels, + ) + num_chunks = num_flat // elt_per_group if num_chunks <= chunk_size: context_size = 0 else: @@ -469,50 +558,65 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): self, input_ids: torch.Tensor | None = None, codes: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, **kwargs, - ) -> torch.Tensor | IntermediateTensors: - # Support both input_ids (from vLLM) and codes (from mimo_audio.py) + ) -> OmniOutput | torch.Tensor | IntermediateTensors: + if runtime_additional_information is None: + runtime_additional_information = kwargs.get("model_intermediate_buffer") or kwargs.get( + "runtime_additional_information" + ) code_tensor = codes if codes is not None else input_ids + empty = torch.zeros((0,), dtype=torch.float32, device=self.device) - if code_tensor is None: - raise ValueError("Either input_ids or codes must be provided") - - if code_tensor.numel() == 0: - raise ValueError("code_tensor is empty.") - - if getattr(self.vllm_config.model_config, "async_chunk", False): - waveform = self.chunked_decode_streaming(code_tensor, chunk_size=10, left_context_size=10) - else: - waveform = self._decode_waveform_from_codes(code_tensor) - - # For warmup/dummy run: if waveform is empty, return dummy hidden_states - # _dummy_run expects a tensor that can be processed by extract_multimodal_outputs - if waveform.numel() == 0: - # Return dummy hidden_states with shape [seq_len, hidden_size] - # This is needed for warmup to avoid NoneType errors - hidden_size = getattr(self.config, "hidden_size", None) - if hidden_size is None: - # Fallback: try to get from vllm_config.model_config - hidden_size = getattr( - self.vllm_config.model_config, - "hidden_size", - 4096, # Default fallback - ) - - # Return [seq_len=1, hidden_size] tensor for warmup - dummy_hidden = torch.zeros( - (1, hidden_size), - dtype=torch.bfloat16, - device=self.device, + if code_tensor is None or code_tensor.numel() == 0: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty]}, ) - return dummy_hidden + is_async_chunk = getattr(self.vllm_config.model_config, "async_chunk", False) + ids = code_tensor.reshape(-1).to(dtype=torch.long) + request_ids_list = self._split_flat_codes_for_requests( + ids, kwargs.get("seq_token_counts"), runtime_additional_information + ) + per_left, per_chunk = self._mimo_codec_runtime_lists(len(request_ids_list), runtime_additional_information) - if waveform.numel() == 1 and waveform.item() == 0: - # Designed for Text output only - return None + is_capturing = torch.cuda.is_current_stream_capturing() + if is_capturing: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty] * len(request_ids_list)}, + ) - return waveform + if is_async_chunk: + audios = self._batch_chunked_decode_streaming( + request_ids_list, + default_chunk_size=self._codec_chunk_frames, + default_left_context_size=self._codec_left_context_frames, + per_req_left_context_frames=per_left, + per_req_codec_chunk_frames=per_chunk, + ) + else: + audios = self._batch_decode_waveforms(request_ids_list) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audios}, + ) + + def make_omni_output(self, model_output: torch.Tensor | OmniOutput, **kwargs) -> OmniOutput: + """Convert raw model output to OmniOutput if needed.""" + if isinstance(model_output, OmniOutput): + return model_output + empty = torch.zeros((0,), dtype=torch.float32, device=self.device) + if model_output is None: + return OmniOutput(text_hidden_states=None, multimodal_outputs={"model_outputs": [empty]}) + if isinstance(model_output, torch.Tensor): + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [model_output.to(dtype=torch.float32).reshape(-1)]}, + ) + raise TypeError(f"Unexpected model output type: {type(model_output)}") def _get_full_code_sequence( self, @@ -545,8 +649,194 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): return code_tensor + @torch.inference_mode() + def _batch_decode_waveforms( + self, + request_codes_list: list[torch.Tensor], + ) -> list[torch.Tensor]: + """Batch-decode multiple requests' codes in a single decoder forward pass. + + Instead of calling _decode_waveform_from_codes per request (which incurs + 4 GPU↔CPU round-trips each), this method: + 1. Extracts audio_codes for all valid requests on CPU (cheap). + 2. Runs quantizer.decode_vq (embedding lookup) for each on GPU. + 3. Packs all hidden-states into one tensor and calls + decoder(packed_hs, input_lengths) once. + 4. Splits the output waveforms back to per-request tensors. + """ + empty = torch.zeros((0,), dtype=torch.float32, device=self.device) + num_req = len(request_codes_list) + if num_req == 0: + return [empty] + + tokenizer = self._tokenizer_service.audio_tokenizer + group_size = self.streamer_config.group_size + audio_channels = self.streamer_config.audio_channels + + hidden_list: list[torch.Tensor] = [] + lengths: list[int] = [] + valid_indices: list[int] = [] + + for i, req_codes in enumerate(request_codes_list): + if req_codes is None or req_codes.numel() == 0: + continue + if self._check_dummy_code_tensor(req_codes): + continue + + flat_codes = req_codes.detach().to(torch.long).cpu() + audio_codes = extract_audio_code_tensor( + flat_codes, + group_size, + audio_channels, + self.codes, + ) + if audio_codes is None or audio_codes.numel() == 0: + continue + + hs = tokenizer.encoder.decode_vq(audio_codes.to(self.device)) + hidden_list.append(hs) + lengths.append(hs.size(0)) + valid_indices.append(i) + + if not hidden_list: + return [empty] * num_req + + if len(hidden_list) == 1: + packed_hs = hidden_list[0] + else: + packed_hs = torch.cat(hidden_list, dim=0) + input_lengths = torch.tensor(lengths, device=self.device) + + recon_wav = tokenizer.decoder(packed_hs, input_lengths) + + cfg = tokenizer.config + frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length + + result: list[torch.Tensor] = [empty] * num_req + if len(valid_indices) == 1: + wav = recon_wav.squeeze(0).squeeze(0) + valid_len = lengths[0] * frames_per_token + if wav.numel() > valid_len: + wav = wav[:valid_len] + result[valid_indices[0]] = wav.to(dtype=torch.float32).reshape(-1) + else: + for j, idx in enumerate(valid_indices): + wav = recon_wav[j].squeeze(0) + valid_len = lengths[j] * frames_per_token + if wav.numel() > valid_len: + wav = wav[:valid_len] + result[idx] = wav.to(dtype=torch.float32).reshape(-1) + + return result + + @torch.inference_mode() + def _batch_chunked_decode_streaming( + self, + request_codes_list: list[torch.Tensor], + default_chunk_size: int, + default_left_context_size: int, + *, + per_req_left_context_frames: list[int | None] | None = None, + per_req_codec_chunk_frames: list[int | None] | None = None, + ) -> list[torch.Tensor]: + """Batch version of chunked_decode_streaming for async_chunk mode.""" + empty = torch.zeros((0,), dtype=torch.float32, device=self.device) + num_req = len(request_codes_list) + if num_req == 0: + return [empty] + + tokenizer = self._tokenizer_service.audio_tokenizer + group_size = self.streamer_config.group_size + audio_channels = self.streamer_config.audio_channels + + hidden_list: list[torch.Tensor] = [] + lengths: list[int] = [] + valid_indices: list[int] = [] + context_sizes: list[int] = [] + + for i, req_codes in enumerate(request_codes_list): + if req_codes is None or req_codes.numel() == 0: + continue + if self._check_dummy_code_tensor(req_codes): + continue + + flat_codes = req_codes.detach().to(torch.long).cpu() + audio_codes = extract_audio_code_tensor( + flat_codes, + group_size, + audio_channels, + self.codes, + ) + if audio_codes is None or audio_codes.numel() == 0: + continue + + hs = tokenizer.encoder.decode_vq(audio_codes.to(self.device)) + hidden_list.append(hs) + lengths.append(hs.size(0)) + valid_indices.append(i) + + num_flat = req_codes.numel() if req_codes.ndim == 1 else req_codes.shape[-1] + elt_per_group = flat_codec_group_element_count(group_size, audio_channels) + num_chunks = num_flat // elt_per_group + chunk_sz = default_chunk_size + if per_req_codec_chunk_frames is not None and i < len(per_req_codec_chunk_frames): + c = per_req_codec_chunk_frames[i] + if c is not None: + chunk_sz = c + if ( + per_req_left_context_frames is not None + and i < len(per_req_left_context_frames) + and per_req_left_context_frames[i] is not None + ): + ctx = int(per_req_left_context_frames[i]) + else: + ctx = 0 if num_chunks <= chunk_sz else default_left_context_size + context_sizes.append(ctx) + + if not hidden_list: + return [empty] * num_req + + if len(hidden_list) == 1: + packed_hs = hidden_list[0] + else: + packed_hs = torch.cat(hidden_list, dim=0) + input_lengths = torch.tensor(lengths, device=self.device) + + recon_wav = tokenizer.decoder(packed_hs, input_lengths) + + cfg = tokenizer.config + frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length + + result: list[torch.Tensor] = [empty] * num_req + if len(valid_indices) == 1: + wav = recon_wav.squeeze(0).squeeze(0) + valid_len = lengths[0] * frames_per_token + if wav.numel() > valid_len: + wav = wav[:valid_len] + drop = context_sizes[0] * self.total_upsample + if drop > 0 and drop < wav.numel(): + wav = wav[drop:] + elif drop >= wav.numel(): + wav = wav[:0] + result[valid_indices[0]] = wav.to(dtype=torch.float32).reshape(-1) + else: + for j, idx in enumerate(valid_indices): + wav = recon_wav[j].squeeze(0) + valid_len = lengths[j] * frames_per_token + if wav.numel() > valid_len: + wav = wav[:valid_len] + drop = context_sizes[j] * self.total_upsample + if drop > 0 and drop < wav.numel(): + wav = wav[drop:] + elif drop >= wav.numel(): + wav = wav[:0] + result[idx] = wav.to(dtype=torch.float32).reshape(-1) + + return result + def _check_dummy_code_tensor(self, code_tensor: torch.Tensor) -> bool: - if code_tensor is not None and code_tensor.numel() == DUMMY_CODE_SHAPE: + expected = flat_codec_group_element_count(self.config.group_size, self.config.audio_channels) + if code_tensor is not None and code_tensor.numel() == expected: code_groups = code_tensor.view(self.config.group_size, self.config.audio_channels + 1) return ( (code_groups[:, 0] == TALKER_CODEC_PAD_TOKEN_ID).all() and (code_groups[:, 1:].sum() == 0).all() diff --git a/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml index 04f7e656e..b3c6bbbaf 100644 --- a/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml +++ b/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml @@ -24,6 +24,8 @@ stage_args: max_model_len: 8192 max_num_batched_tokens: 8192 custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav_async_chunk + output_connectors: + to_stage_1: connector_of_shared_memory is_comprehension: true final_output: true final_output_type: text @@ -57,6 +59,8 @@ stage_args: max_model_len: 8192 max_num_batched_tokens: 8192 engine_input_source: [ 0 ] + input_connectors: + from_stage_0: connector_of_shared_memory is_comprehension: false final_output: true final_output_type: audio @@ -67,3 +71,26 @@ stage_args: max_tokens: 18192 seed: 42 detokenize: false + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + + connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 365536 + codec_streaming: true + connector_get_sleep_s: 0.001 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 3 + codec_left_context_frames: 3 + + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py index 2b347ab2a..71573c0e2 100644 --- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py +++ b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py @@ -50,6 +50,58 @@ def _make_finished_sentinel() -> dict[str, Any]: return {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)} +def _flush_remaining_codes( + transfer_manager: Any, + request_id: str, + chunk_size: int, + left_context_size: int, +) -> dict[str, Any]: + """Flush any accumulated but unsent codes when the request finishes.""" + accumulated = transfer_manager.code_prompt_token_ids.get(request_id, []) + if not accumulated: + return _make_finished_sentinel() + + length = len(accumulated) + chunk_length = length % chunk_size + context_length = chunk_length if chunk_length != 0 else chunk_size + end_index = min(length, left_context_size + context_length) + + # Align with qwen3_omni talker2code2wav_async_chunk: decoder strip uses explicit frame count. + left_ctx_frames = max(0, min(length - context_length, left_context_size)) + flat_codes = torch.tensor(accumulated[-end_index:]).reshape(-1).tolist() + + return { + "code_predictor_codes": flat_codes, + "left_context_size": left_ctx_frames, + "codec_chunk_frames": chunk_size, + "codec_left_context_frames": left_context_size, + "code_flat_numel": len(flat_codes), + "finished": torch.tensor(True, dtype=torch.bool), + } + + +def _is_codes_empty(codes: Any) -> bool: + """Check whether code_predictor_codes should be treated as empty / invalid.""" + if codes is None: + return True + if isinstance(codes, torch.Tensor): + return codes.numel() == 0 or not codes.any() + if hasattr(codes, "__len__") and len(codes) == 0: + return True + t = torch.tensor(codes, dtype=torch.long) if not isinstance(codes, torch.Tensor) else codes + return not t.any() + + +def _to_code_tensor(codes: Any) -> torch.Tensor | None: + """Convert codes to a (B, 1, 8, 4) long tensor, or return None if shape is invalid.""" + code_tensor = codes.to(torch.long) if isinstance(codes, torch.Tensor) else torch.tensor(codes, dtype=torch.long) + if code_tensor.ndim == 3: + code_tensor = code_tensor.unsqueeze(0) + if code_tensor.ndim != 4 or code_tensor.shape[-2:] != (8, 4): + return None + return code_tensor + + def llm2code2wav_async_chunk( transfer_manager: Any, pooling_output: dict[str, Any], @@ -62,57 +114,39 @@ def llm2code2wav_async_chunk( Accumulates codes in connector per request_id, returns payload only when chunk_size is full or request is finished; returns None when waiting. """ - if "code_predictor_codes" not in pooling_output: + + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + chunk_size = int(cfg.get("codec_chunk_frames", 3)) + left_context_size = int(cfg.get("codec_left_context_frames", 3)) + + request_id = getattr(request, "external_req_id", None) + + codes = pooling_output.get("code_predictor_codes") + + if _is_codes_empty(codes): if is_finished: - return _make_finished_sentinel() + return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size) return None - code_predictor_codes = pooling_output["code_predictor_codes"] - - if code_predictor_codes is None: + code_tensor = _to_code_tensor(codes) + if code_tensor is None: if is_finished: - return _make_finished_sentinel() - return None - if isinstance(code_predictor_codes, torch.Tensor): - if code_predictor_codes.numel() == 0: - if is_finished: - return _make_finished_sentinel() - return None - elif hasattr(code_predictor_codes, "__len__"): - if len(code_predictor_codes) == 0: - return None - - if isinstance(code_predictor_codes, torch.Tensor): - if not code_predictor_codes.any(): - if is_finished: - return _make_finished_sentinel() - return None - code_tensor = code_predictor_codes.to(torch.long) - else: - code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long) - if not code_tensor.any(): - if is_finished: - return _make_finished_sentinel() - return None - - if code_tensor.ndim == 3: - code_tensor = code_tensor.unsqueeze(0) - if code_tensor.ndim != 4 or code_tensor.shape[-2:] != (8, 4): + return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size) return None pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4, device=code_tensor.device, dtype=code_tensor.dtype) - code_final = prepend_and_flatten_colmajor(code_tensor, pad_vec) - code_list = code_final.tolist() + code_list = prepend_and_flatten_colmajor(code_tensor, pad_vec).tolist() + if sum(code_list) == 0: if is_finished: - return _make_finished_sentinel() + return _flush_remaining_codes(transfer_manager, request_id, chunk_size, left_context_size) return None - request_id = getattr(request, "external_req_id", None) if request_id is None: return None - chunk_size = left_context_size = 10 transfer_manager.code_prompt_token_ids[request_id].append(code_list) length = len(transfer_manager.code_prompt_token_ids[request_id]) chunk_length = length % chunk_size @@ -121,14 +155,17 @@ def llm2code2wav_async_chunk( context_length = chunk_length if chunk_length != 0 else chunk_size end_index = min(length, left_context_size + context_length) + left_ctx_frames = max(0, min(length - context_length, left_context_size)) + flat_codes = torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]).reshape(-1).tolist() - info = { - "code_predictor_codes": ( - torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]).reshape(-1).tolist() - ), + return { + "code_predictor_codes": flat_codes, + "left_context_size": left_ctx_frames, + "codec_chunk_frames": chunk_size, + "codec_left_context_frames": left_context_size, + "code_flat_numel": len(flat_codes), "finished": torch.tensor(is_finished, dtype=torch.bool), } - return info def llm2code2wav(