[Enhancement] Custom chunk_size for mimo-audio model (#1964)

Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>

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

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