From aad30908772a8188c17997341d62768dc0bab3b7 Mon Sep 17 00:00:00 2001 From: Linze Shi Date: Mon, 5 Oct 2026 11:52:41 -0700 Subject: [PATCH] [Model][PersonaPlex] Vectorize prefill embedding construction (#7481) Signed-off-by: Linze-Shi --- .../personaplex/duplex/test_stage0_runtime.py | 29 ++++ .../personaplex/test_prefill_embeddings.py | 149 ++++++++++++++++++ .../models/personaplex/personaplex_talker.py | 23 ++- 3 files changed, 188 insertions(+), 13 deletions(-) create mode 100644 tests/model_executor/models/personaplex/test_prefill_embeddings.py diff --git a/tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py b/tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py index 40368c51a..e4ff44bdd 100644 --- a/tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py +++ b/tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py @@ -8,6 +8,8 @@ import numpy as np import pytest import torch +from tests.model_executor.models.personaplex.test_prefill_embeddings import _make_talker, _reference_prefill +from vllm_omni.model_executor.models.personaplex.configuration_personaplex import PersonaPlexConfig from vllm_omni.model_executor.models.personaplex.duplex.policy import ( SILENCE_TOKENS, SINE_TOKENS, @@ -147,6 +149,33 @@ def test_first_append_prepends_voice_and_persona_once() -> None: assert codec.encode_calls == 2 +@torch.inference_mode() +def test_first_append_uses_real_prefill_embeddings_once(monkeypatch, mocker) -> None: + runtime = _runtime(_FakeCodec()) + device = torch.device("cpu") + talker = _make_talker(device, torch.float32, PersonaPlexConfig(temporal_config={"hidden_size": 4})) + prefill = mocker.spy(talker, "_build_prefill_embed") + # Exercise real prefill while retaining the codec/live-frame fixture. + monkeypatch.setattr(runtime.stage_model, "_build_prefill_embed", prefill) + first = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18) + second = runtime.prepare_append(_duplex_info(seq=2), prompt_len=18) + expected = _reference_prefill( + talker, + torch.tensor([3] * 6 + [7, 8, 9] + [3] * 6), + 0, + 15, + device, + torch.tensor(SILENCE_TOKENS), + torch.tensor(SINE_TOKENS), + ) + prefill.assert_called_once() + assert torch.equal(first.inputs_embeds[:2], torch.arange(8, dtype=torch.float32).reshape(2, 4)) + assert torch.equal(first.inputs_embeds[2:17].view(torch.uint8), expected.view(torch.uint8)) + assert first.inputs_embeds.shape == (18, 4) + assert second.inputs_embeds.shape == (1, 4) + assert second.prefill_applied is False + + @pytest.mark.parametrize( ("append_count", "expected_tokens", "expected_provided"), [ diff --git a/tests/model_executor/models/personaplex/test_prefill_embeddings.py b/tests/model_executor/models/personaplex/test_prefill_embeddings.py new file mode 100644 index 000000000..56d2dd458 --- /dev/null +++ b/tests/model_executor/models/personaplex/test_prefill_embeddings.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.personaplex.configuration_personaplex import PersonaPlexConfig +from vllm_omni.model_executor.models.personaplex.personaplex_embeddings import PersonaPlexInputEmbeddings +from vllm_omni.model_executor.models.personaplex.personaplex_talker import PersonaPlexTalkerForConditionalGeneration + +pytestmark = pytest.mark.core_model + + +def _make_talker( + device: torch.device, dtype: torch.dtype, config: PersonaPlexConfig | None = None +) -> PersonaPlexTalkerForConditionalGeneration: + talker = PersonaPlexTalkerForConditionalGeneration.__new__(PersonaPlexTalkerForConditionalGeneration) + nn.Module.__init__(talker) + talker.config = config or PersonaPlexConfig( + temporal_config={"hidden_size": 32}, + text_vocab_size=31, + text_embedding_rows=32, + audio_vocab_size=16, + ) + talker.input_embeddings = PersonaPlexInputEmbeddings(talker.config) + generator = torch.Generator().manual_seed(7389) + with torch.no_grad(): + for parameter in talker.parameters(): + parameter.copy_(torch.randn(parameter.shape, generator=generator)) + return talker.to(device=device, dtype=dtype).eval() + + +def _reference_prefill( + talker: PersonaPlexTalkerForConditionalGeneration, + prefill_text: torch.Tensor, + offset: int, + span: int, + device: torch.device, + silence: torch.Tensor | None = None, + user_sine: torch.Tensor | None = None, +) -> torch.Tensor: + """Original loop from main@13b85c56.""" + n_q = talker.config.num_audio_codebooks + n_user = n_q // 2 + sil = silence.reshape(-1).to(device) if isinstance(silence, torch.Tensor) and silence.numel() >= n_user else None + user = ( + user_sine.reshape(-1).to(device) if isinstance(user_sine, torch.Tensor) and user_sine.numel() >= n_user else sil + ) + rows = [] + for i in range(span): + pos = offset + i + text_tok = int(prefill_text[pos].item()) if pos < prefill_text.numel() else 3 + stack = torch.zeros((1, 1 + n_q, 1), dtype=torch.long, device=device) + stack[:, 0] = text_tok + if sil is not None: + stack[0, 1 : 1 + n_user, 0] = sil[:n_user] + if user is not None: + stack[0, 1 + n_user : 1 + 2 * n_user, 0] = user[:n_user] + rows.append(talker.input_embeddings(stack).reshape(1, -1)) + return torch.cat(rows, dim=0) + + +@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)]) +def device(request): + if request.param == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA PyTorch and a visible GPU are required") + return torch.device(request.param) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + ("length", "offset", "span"), + [(5, 0, 1), (5, 1, 3), (5, 3, 4), (5, 5, 3), (5, 8, 2), (0, 0, 3), (67, 3, 64)], + ids=["single", "offset", "partial-pad", "at-end", "past-end", "empty-text", "long"], +) +@pytest.mark.parametrize( + ("silence_size", "user_size"), + [(None, None), (8, None), (8, 8), (None, 8), (7, 7), (8, 7), (10, 10)], + ids=["absent", "silence-fallback", "both", "user-only", "short", "short-user", "long-codes"], +) +@torch.inference_mode() +def test_prefill_matches_loop_bitwise(device, dtype, length, offset, span, silence_size, user_size): + talker = _make_talker(device, dtype) + # Strided inputs include special token rows. + text = (torch.arange(2 * length, device=device) % 33 - 1)[::2] + silence = None if silence_size is None else ((torch.arange(2 * silence_size, device=device) - 1) % 18 - 1)[::2] + user = None if user_size is None else (torch.arange(2 * user_size, device=device) % 17)[::2] + originals = [value.clone() for value in (text, silence, user) if value is not None] + expected = _reference_prefill(talker, text, offset, span, device, silence, user) + actual = talker._build_prefill_embed(text, offset, span, device, silence, user) + assert actual.shape == (span, 32) + assert actual.dtype == dtype + assert actual.device == expected.device + assert actual.is_contiguous() + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + for value, original in zip((v for v in (text, silence, user) if v is not None), originals): + assert torch.equal(value, original) + + +@torch.inference_mode() +def test_prefill_manual_layout_and_padding(device): + talker = _make_talker(device, torch.float32) + # Encode table identity and token ID in the weights. + for table_index, table in enumerate([talker.input_embeddings.text_emb, *talker.input_embeddings.audio_emb]): + values = torch.arange(table.num_embeddings, device=device) + 100 * table_index + table.weight.copy_(values[:, None].expand_as(table.weight)) + actual = talker._build_prefill_embed( + torch.tensor([10, 20, 30], device=device), + 1, + 4, + device, + torch.arange(1, 9, device=device), + torch.arange(9, 17, device=device), + ) + # sum(100 * k + k for k in 1..16) = 13736. + expected = torch.tensor([13756, 13766, 13739, 13739], dtype=torch.float32, device=device) + assert torch.equal(actual, expected[:, None].expand(-1, 32)) + + +@torch.inference_mode() +def test_prefill_joins_chunked_slices(device): + talker = _make_talker(device, torch.bfloat16) + text = torch.arange(23, device=device) + silence = torch.arange(8, device=device) + full = talker._build_prefill_embed(text, 0, 27, device, silence) + pieces = [ + talker._build_prefill_embed(text, offset, span, device, silence) for offset, span in [(0, 7), (7, 13), (20, 7)] + ] + assert torch.equal(full.view(torch.uint8), torch.cat(pieces).view(torch.uint8)) + + +@torch.inference_mode() +def test_prefill_uses_one_embedding_call_without_scalar_reads(device, monkeypatch): + talker = _make_talker(device, torch.float32) + text = torch.arange(12, device=device) + shapes = [] + handle = talker.input_embeddings.register_forward_pre_hook(lambda _module, args: shapes.append(args[0].shape)) + + def reject_scalar_read(_self): + raise AssertionError("prefill must not read device scalars into Python") + + try: + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "item", reject_scalar_read) + talker._build_prefill_embed(text, 0, 12, device) + finally: + handle.remove() + assert shapes == [(1, 17, 12)] diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_talker.py b/vllm_omni/model_executor/models/personaplex/personaplex_talker.py index 5ccc48f7d..14374fa57 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_talker.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_talker.py @@ -256,19 +256,16 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module): if isinstance(user_sine, torch.Tensor) and user_sine.numel() >= n_user else sil ) - total = int(prefill_text.numel()) - rows = [] - for i in range(span): - pos = offset + i - text_tok = int(prefill_text[pos].item()) if pos < total else zero_text - stack = torch.zeros((1, 1 + n_q, 1), dtype=torch.long, device=device) - stack[:, 0] = text_tok - if sil is not None: - stack[0, 1 : 1 + n_user, 0] = sil[:n_user] # agent rows = encoded silence - if user is not None: - stack[0, 1 + n_user : 1 + 2 * n_user, 0] = user[:n_user] - rows.append(self.input_embeddings(stack).reshape(1, -1)) - return torch.cat(rows, dim=0) # [span, hidden] + + stack = torch.zeros((1, 1 + n_q, span), dtype=torch.long, device=device) + stack[:, 0] = zero_text + text = prefill_text[offset : offset + span] + stack[0, 0, : text.numel()] = text.to(device) + if sil is not None: + stack[0, 1 : 1 + n_user] = sil[:n_user, None] + if user is not None: + stack[0, 1 + n_user : 1 + 2 * n_user] = user[:n_user, None] + return self.input_embeddings(stack).reshape(span, -1) # [span, hidden] @staticmethod def _user_frame(info: dict[str, Any], frame_idx: int) -> torch.Tensor | None: