[Model][PersonaPlex] Vectorize prefill embedding construction (#7481)

Signed-off-by: Linze-Shi <linzeshi0@gmail.com>

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"),
     [

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)]

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 = []
+        stack = torch.zeros((1, 1 + n_q, span), dtype=torch.long, device=device)
-        for i in range(span):
+        stack[:, 0] = zero_text
-            pos = offset + i
+        text = prefill_text[offset : offset + span]
-            text_tok = int(prefill_text[pos].item()) if pos < total else zero_text
+        stack[0, 0, : text.numel()] = text.to(device)
-            stack = torch.zeros((1, 1 + n_q, 1), dtype=torch.long, device=device)
+        if sil is not None:
-            stack[:, 0] = text_tok
+            stack[0, 1 : 1 + n_user] = sil[:n_user, None]
-            if sil is not None:
+        if user is not None:
-                stack[0, 1 : 1 + n_user, 0] = sil[:n_user]  # agent rows = encoded silence
+            stack[0, 1 + n_user : 1 + 2 * n_user] = user[:n_user, None]
-            if user is not None:
+        return self.input_embeddings(stack).reshape(span, -1)  # [span, hidden]
-                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]
 
     @staticmethod
     def _user_frame(info: dict[str, Any], frame_idx: int) -> torch.Tensor | None: