From ae430dddc0db6c12054b9935e6df7e0b708a32d4 Mon Sep 17 00:00:00 2001 From: LOGO127 Date: Tue, 6 Oct 2026 02:49:33 +0800 Subject: [PATCH] [Model][PersonaPlex] Encode only consumed Mimi codebooks (#7399) Signed-off-by: luozijian --- .../personaplex/benchmark_mimi_quantizer.py | 104 ++++++++++++++++ .../personaplex/test_mimi_encode_codebooks.py | 111 ++++++++++++++++++ .../personaplex/test_ring_per_row_offsets.py | 4 +- .../models/personaplex/personaplex_mimi.py | 2 +- 4 files changed, 219 insertions(+), 2 deletions(-) create mode 100644 tests/model_executor/models/personaplex/benchmark_mimi_quantizer.py create mode 100644 tests/model_executor/models/personaplex/test_mimi_encode_codebooks.py diff --git a/tests/model_executor/models/personaplex/benchmark_mimi_quantizer.py b/tests/model_executor/models/personaplex/benchmark_mimi_quantizer.py new file mode 100644 index 000000000..50a2a6c9d --- /dev/null +++ b/tests/model_executor/models/personaplex/benchmark_mimi_quantizer.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Compare full and 8-codebook Mimi RVQ on CUDA without model checkpoints. + +Run: python tests/model_executor/models/personaplex/benchmark_mimi_quantizer.py + +This measures the HF quantizer on random centroids and synthetic latents, not +full codec/LM latency, audio quality, or serving concurrency. Timings include +host launches and synchronized CUDA execution. No CPU fallback is allowed. +""" + +import hashlib +import inspect +import json +import statistics +from datetime import datetime, timezone + +import torch +import transformers +from torch.utils.benchmark import Timer +from transformers import MimiConfig +from transformers.models.mimi.modeling_mimi import MimiSplitResidualVectorQuantizer + + +def main() -> None: + # Explicit device allocation below fails rather than falling back to CPU. + torch.set_num_threads(1) + torch.manual_seed(314159) + config = MimiConfig() + assert config.num_quantizers == 32 and config.num_semantic_quantizers == 1 + quantizer = MimiSplitResidualVectorQuantizer(config).eval().to(device="cuda:0") + for rvq in (quantizer.semantic_residual_vector_quantizer, quantizer.acoustic_residual_vector_quantizer): + for layer in rvq.layers: + layer.codebook.embed_sum.normal_() + result = { + "observed_at": datetime.now(timezone.utc).isoformat(), + "device": str(next(quantizer.parameters()).device), + "torch": torch.__version__, + "cuda_runtime": torch.version.cuda, + "transformers": transformers.__version__, + "dtype": "float32", + "scope": "HF Mimi split RVQ only; random centroids and synthetic latents; no pretrained weights or serving", + "config": { + key: getattr(config, key) + for key in ( + "hidden_size", + "vector_quantization_hidden_dimension", + "codebook_dim", + "codebook_size", + "num_quantizers", + "num_semantic_quantizers", + ) + }, + "source_sha256": hashlib.sha256(inspect.getsource(MimiSplitResidualVectorQuantizer).encode()).hexdigest(), + "warmup_calls_per_path": 16, + "rounds": 6, + "calls_per_round": 16, + "timing": "torch.utils.benchmark.Timer; synchronizes accelerators and adds its own warmup", + } + batches = [] + for batch in (1, 8, 32): + inputs = [torch.randn(batch, config.hidden_size, 1, device="cuda") for _ in range(32)] + for x in inputs: + full = quantizer.encode(x) + prefix = quantizer.encode(x, num_quantizers=8) + assert torch.equal(full[:8], prefix), "8-codebook prefix differs from full RVQ" + for _ in range(16): + quantizer.encode(inputs[0]) + quantizer.encode(inputs[0], num_quantizers=8) + timers = { + count: Timer( + stmt="for x in inputs: quantizer.encode(x, num_quantizers=count)", + globals={"quantizer": quantizer, "inputs": inputs[:16], "count": count}, + num_threads=1, + ) + for count in (32, 8) + } + timings: dict[int, list[float]] = {32: [], 8: []} + for repeat in range(6): + for count in (32, 8) if repeat % 2 == 0 else (8, 32): + timings[count].append(timers[count].timeit(number=1).mean * 1000 / 16) + full_ms = statistics.median(timings[32]) + prefix_ms = statistics.median(timings[8]) + row = { + "batch": batch, + "frames_per_call": 1, + "parity_cases": len(inputs), + "prefix_bit_identical": True, + "full_ms": timings[32], + "prefix_ms": timings[8], + "full_median_ms": full_ms, + "prefix_median_ms": prefix_ms, + "module_speedup": full_ms / prefix_ms, + } + batches.append(row) + print("BATCH " + json.dumps(row), flush=True) + result["batches"] = batches + print("RESULT " + json.dumps(result), flush=True) + + +if __name__ == "__main__": + with torch.no_grad(): + main() diff --git a/tests/model_executor/models/personaplex/test_mimi_encode_codebooks.py b/tests/model_executor/models/personaplex/test_mimi_encode_codebooks.py new file mode 100644 index 000000000..e3d7cfb8e --- /dev/null +++ b/tests/model_executor/models/personaplex/test_mimi_encode_codebooks.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Checkpoint-free coverage using real streaming layers and the HF Mimi RVQ.""" + +import pytest +import torch +from pytest_mock import MockerFixture +from torch import nn +from transformers import MimiConfig +from transformers.models.mimi.modeling_mimi import MimiSplitResidualVectorQuantizer + +from vllm_omni.model_executor.models.personaplex import personaplex_mimi as mimi + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _codec(batch_size: int) -> mimi.PersonaPlexMimiCodec: + # Use a small real front end rather than loading pretrained codec weights. + # Keep all 32 codebooks, with nonzero centroids, to exercise prefix parity. + with torch.random.fork_rng(devices=[]): + torch.manual_seed(17) + codec = mimi.PersonaPlexMimiCodec.__new__(mimi.PersonaPlexMimiCodec) + nn.Module.__init__(codec) + codec.device = torch.device("cpu") + codec.dtype = torch.float32 + codec.model = nn.Module() + config = MimiConfig( + hidden_size=8, + vector_quantization_hidden_dimension=8, + codebook_dim=8, + codebook_size=16, + num_quantizers=32, + num_semantic_quantizers=1, + ) + codec.model.quantizer = MimiSplitResidualVectorQuantizer(config).eval() + for rvq in ( + codec.model.quantizer.semantic_residual_vector_quantizer, + codec.model.quantizer.acoustic_residual_vector_quantizer, + ): + for layer in rvq.layers: + layer.codebook.embed_sum.normal_() + codec._enc_stages = [("conv", mimi._StreamConv1d(nn.Conv1d(1, 8, mimi.FRAME_SIZE + 2, stride=mimi.FRAME_SIZE)))] + codec._dec_stages = [] + codec._downsample = mimi._StreamConv1d(nn.Conv1d(8, 8, 3), pad_mode="replicate") + codec._upsample = mimi._StreamConvTr1d(nn.ConvTranspose1d(8, 8, 3)) + for name in ("encoder_transformer", "decoder_transformer"): + transformer = mimi._MimiStreamingTransformer(num_layers=1, dim=8, num_heads=2, context=3) + transformer.layers = nn.ModuleList([mimi._MimiTransformerLayer(dim=8, num_heads=2, ffn=16)]) + for parameter in transformer.parameters(): + nn.init.uniform_(parameter, -0.2, 0.2) + setattr(codec, name, transformer) + codec.streaming_init(batch_size) + return codec.eval() + + +@pytest.mark.parametrize("batch_size", [1, 2, 9]) +def test_encode_frame_computes_only_consumed_codebooks(batch_size: int, mocker: MockerFixture) -> None: + codec = _codec(batch_size) + quantizer = codec.model.quantizer + layers = [ + *quantizer.semantic_residual_vector_quantizer.layers, + *quantizer.acoustic_residual_vector_quantizer.layers, + ] + layer_spies = [mocker.spy(layer, "encode") for layer in layers] + quantizer_spy = mocker.spy(quantizer, "encode") + pcm = torch.randn(batch_size, mimi.FRAME_SIZE, generator=torch.Generator().manual_seed(23)) + + actual = codec.encode_frame(pcm) + counts = [spy.call_count for spy in layer_spies] + latents = quantizer_spy.call_args.args[0] + reference = quantizer.encode(latents, num_quantizers=32)[: mimi.CODEBOOKS, :, 0].transpose(0, 1).contiguous() + + assert actual.shape == (batch_size, mimi.CODEBOOKS) + assert actual.dtype == torch.long + assert actual.is_contiguous() + assert torch.equal(actual, reference) + assert actual.unique().numel() > 1 # Do not pass through all-zero codebooks. + assert counts == [1] * mimi.CODEBOOKS + [0] * (32 - mimi.CODEBOOKS) + assert quantizer.max_num_quantizers == 32 # Encoding must not mutate decoder capacity. + + +@pytest.mark.parametrize("batch_size", [1, 2, 9]) +def test_encode_prefix_survives_ring_wrap_and_slot_reset(batch_size: int, monkeypatch: pytest.MonkeyPatch) -> None: + candidate = _codec(batch_size) + reference = _codec(batch_size) + encode_all = reference.model.quantizer.encode + + def full_encode(embeddings: torch.Tensor, num_quantizers: int | None = None) -> torch.Tensor: + # Execute the pre-optimization 32-codebook path on the same HF quantizer. + return encode_all(embeddings, num_quantizers=32) + + monkeypatch.setattr(reference.model.quantizer, "encode", full_encode) + generator = torch.Generator().manual_seed(29) + for frame in range(9): + if frame == 4: + candidate.reset_slot(batch_size - 1) + reference.reset_slot(batch_size - 1) + elif frame == 7: + candidate.reset_streaming() + reference.reset_streaming() + pcm = torch.randn(batch_size, mimi.FRAME_SIZE, generator=generator) + assert torch.equal(candidate.encode_frame(pcm), reference.encode_frame(pcm)) + assert torch.equal(candidate.encoder_transformer._offset, reference.encoder_transformer._offset) + for actual, expected in zip(candidate.encoder_transformer._kv, reference.encoder_transformer._kv): + assert torch.equal(actual.end_offset, expected.end_offset) + assert torch.equal(actual.start_offset, expected.start_offset) + for actual, expected in zip(candidate._conv_states(), reference._conv_states()): + actual_buffer = actual.prev if isinstance(actual, mimi._StreamConv1d) else actual.partial + expected_buffer = expected.prev if isinstance(expected, mimi._StreamConv1d) else expected.partial + assert torch.equal(actual_buffer, expected_buffer) diff --git a/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py b/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py index 0fa216e14..c60b5c197 100644 --- a/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py +++ b/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py @@ -118,7 +118,9 @@ def _make_codec_stub(mocker: MockerFixture) -> PersonaPlexMimiCodec: nn.Module.__init__(codec) quantizer = mocker.Mock(spec=["encode", "decode"]) - quantizer.encode.side_effect = lambda x: torch.zeros(8, x.shape[0], x.shape[-1], dtype=torch.long) + quantizer.encode.side_effect = lambda x, num_quantizers=8: torch.zeros( + num_quantizers, x.shape[0], x.shape[-1], dtype=torch.long + ) quantizer.decode.side_effect = lambda codes: torch.zeros(codes.shape[0], 1, codes.shape[-1]) codec.device = torch.device("cpu") diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py index af56c7969..81729dbdf 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py @@ -421,7 +421,7 @@ class PersonaPlexMimiCodec(nn.Module): x = self._run_stages(x, self._enc_stages, active) x = self.encoder_transformer.step(x.transpose(1, 2), active).transpose(1, 2) x = self._downsample(x, active) - codes = self.model.quantizer.encode(x) # [Q, B, T] + codes = self.model.quantizer.encode(x, num_quantizers=CODEBOOKS) # [Q, B, T] return codes[:CODEBOOKS, :, 0].transpose(0, 1).contiguous() @torch.no_grad()