[Perf][PersonaPlex] Serve 16 concurrent realtime duplex sessions on one GPU (#8192)

tests/model_executor/models/personaplex/duplex/test_plugin.py

@@ -101,6 +101,8 @@ def test_personaplex_deploy_is_duplex_and_propagates_capacity_to_all_model_stage
     assert deploy.session_mode == "duplex"
     assert deploy.duplex_session.max_sessions == 2
     assert [stage.yaml_engine_args.get("duplex_max_sessions") for stage in stages] == [2, 2]
+    # Lockstep frames of all sessions must share a step (see the deploy comment).
+    assert stages[0].yaml_engine_args.get("async_scheduling") is False
     assert "personaplex_codec_max_sessions" not in deploy.connectors["connector_of_shared_memory"]["extra"]
 
 

tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py

@@ -13,8 +13,10 @@ from vllm_omni.model_executor.models.personaplex.duplex.policy import (
     SINE_TOKENS,
 )
 from vllm_omni.model_executor.models.personaplex.duplex.stage0 import (
+    PersonaPlexStage0CapacityError,
     PersonaPlexStage0DuplexRuntime,
     PersonaPlexStage0PreparedAppend,
+    PersonaPlexStage0StaleEpochError,
 )
 from vllm_omni.model_executor.models.personaplex.personaplex_talker import (
     PersonaPlexTalkerForConditionalGeneration,
@@ -24,20 +26,30 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _FakeCodec:
+    """Shared streaming encoder whose code for a row is that row's frame count."""
+
     def __init__(self) -> None:
         self.encode_calls = 0
-        self.reset_calls = 0
+        self.reset_slots: list[int] = []
+        self.frames: list[int] = []
 
     def streaming_init(self, batch_size: int) -> None:
-        assert batch_size == 1
+        self.frames = [0] * batch_size
 
-    def encode_frame(self, pcm: torch.Tensor) -> torch.Tensor:
-        assert pcm.shape == (1, 1920)
+    def encode_frame(self, pcm: torch.Tensor, active: torch.Tensor) -> torch.Tensor:
+        assert pcm.shape == (len(self.frames), 1920)
+        assert active.shape == (len(self.frames),)
         self.encode_calls += 1
-        return torch.full((1, 8), self.encode_calls, dtype=torch.long)
+        codes = torch.zeros((len(self.frames), 8), dtype=torch.long)
+        for row, is_active in enumerate(active.tolist()):
+            if is_active:
+                self.frames[row] += 1
+                codes[row] = self.frames[row]
+        return codes
 
-    def reset_streaming(self) -> None:
-        self.reset_calls += 1
+    def reset_slot(self, row: int) -> None:
+        self.reset_slots.append(row)
+        self.frames[row] = 0
 
 
 class _FakeTalker:
@@ -101,20 +113,14 @@ def _duplex_info(*, seq: int, session_id: str = "session", epoch: int = 0):
     }
 
 
-def _runtime(*codecs: _FakeCodec) -> PersonaPlexStage0DuplexRuntime:
+def _runtime(codec: _FakeCodec, max_sessions: int = 1) -> PersonaPlexStage0DuplexRuntime:
     voice_embeddings = torch.arange(8, dtype=torch.float32).reshape(2, 1, 1, 4)
-    codec_args: dict[str, object] = {"codec": codecs[0]}
-    if len(codecs) > 1:
-        available = iter(codecs)
-        codec_args = {
-            "codec_factory": lambda: next(available),
-            "max_sessions": len(codecs),
-        }
     return PersonaPlexStage0DuplexRuntime(
         _FakeTalker(),
         model_path="/unused",
         device="cpu",
-        **codec_args,
+        codec=codec,
+        max_sessions=max_sessions,
         tokenizer=lambda _text: [7, 8, 9],
         voice_loader=lambda _voice: {
             "embeddings": voice_embeddings,
@@ -268,9 +274,7 @@ def _prepare_two_sessions(
 
 
 def test_live_sessions_keep_independent_streaming_encoders() -> None:
-    first_codec = _FakeCodec()
-    second_codec = _FakeCodec()
-    runtime = _runtime(first_codec, second_codec)
+    runtime = _runtime(_FakeCodec(), max_sessions=2)
 
     first_1, second_1, first_2, second_2 = _prepare_two_sessions(runtime)
 
@@ -281,7 +285,7 @@ def test_live_sessions_keep_independent_streaming_encoders() -> None:
 
 
 def test_stage0_session_capacity_fails_before_codec_state_is_shared() -> None:
-    runtime = _runtime(_FakeCodec(), _FakeCodec())
+    runtime = _runtime(_FakeCodec(), max_sessions=2)
     _prepare_two_sessions(runtime)
 
     with pytest.raises(RuntimeError, match="capacity 2"):
@@ -291,21 +295,23 @@ def test_stage0_session_capacity_fails_before_codec_state_is_shared() -> None:
         )
 
 
-def test_close_session_resets_and_reuses_released_codec() -> None:
-    first_codec = _FakeCodec()
-    second_codec = _FakeCodec()
-    runtime = _runtime(first_codec, second_codec)
+def test_close_session_resets_only_its_row_and_reuses_it() -> None:
+    codec = _FakeCodec()
+    runtime = _runtime(codec, max_sessions=2)
     _prepare_two_sessions(runtime)
+    closed_slot = runtime.sessions[("session", 0)].slot
+    other_slot = runtime.sessions[("other", 0)].slot
 
     runtime.close_session("session", 0)
 
-    assert first_codec.reset_calls == 1
-    assert second_codec.reset_calls == 0
+    assert codec.reset_slots == [closed_slot]
     replacement = runtime.prepare_append(
         _duplex_info(seq=1, session_id="replacement"),
         prompt_len=18,
     )
-    assert replacement.user_codes[:, 0].tolist() == [3]
+    assert runtime.sessions[("replacement", 0)].slot == closed_slot
+    assert replacement.user_codes[:, 0].tolist() == [1]
+    assert codec.frames[other_slot] == 2
 
 
 def test_a_new_epoch_replays_the_prefill_and_recycles_the_codec() -> None:
@@ -324,8 +330,8 @@ def test_a_new_epoch_replays_the_prefill_and_recycles_the_codec() -> None:
     assert restarted.user_codes.shape == (1, 8)
     assert list(runtime.sessions) == [("session", 1)]
     assert runtime.request_sessions == {"req-e1": ("session", 1)}
-    assert codec.reset_calls == 1
-    assert runtime.sessions[("session", 1)].codec is codec
+    assert codec.reset_slots == [0]
+    assert runtime.sessions[("session", 1)].slot == 0
 
 
 def test_a_late_finish_of_the_old_epoch_request_does_not_close_the_new_state() -> None:
@@ -341,11 +347,164 @@ def test_a_late_finish_of_the_old_epoch_request_does_not_close_the_new_state() -
 
 
 def test_stage0_capacity_counts_live_epochs_not_superseded_ones() -> None:
-    runtime = _runtime(_FakeCodec(), _FakeCodec())
+    runtime = _runtime(_FakeCodec(), max_sessions=2)
     runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
     runtime.prepare_append(_duplex_info(seq=1, session_id="other"), prompt_len=18)
 
-    # Restarting one session must not need a third codec.
+    # Restarting one session must not need a third encoder row.
     runtime.prepare_append(_duplex_info(seq=1, epoch=1), prompt_len=18)
 
     assert sorted(runtime.sessions) == [("other", 0), ("session", 1)]
+
+
+def test_encode_appends_batches_sessions_into_one_encoder_call() -> None:
+    codec = _FakeCodec()
+    runtime = _runtime(codec, max_sessions=4)
+
+    runtime.encode_appends([_duplex_info(seq=1), _duplex_info(seq=1, session_id="other")])
+    first = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
+    other = runtime.prepare_append(_duplex_info(seq=1, session_id="other"), prompt_len=18)
+
+    assert codec.encode_calls == 1
+    assert first.user_codes[:, 0].tolist() == [1]
+    assert other.user_codes[:, 0].tolist() == [1]
+
+
+def test_encode_appends_leaves_rows_without_a_new_append_untouched() -> None:
+    codec = _FakeCodec()
+    runtime = _runtime(codec, max_sessions=2)
+    runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
+    runtime.prepare_append(_duplex_info(seq=1, session_id="other"), prompt_len=18)
+    other_slot = runtime.sessions[("other", 0)].slot
+
+    runtime.encode_appends([_duplex_info(seq=2)])
+    runtime.prepare_append(_duplex_info(seq=2), prompt_len=19)
+
+    assert codec.frames[other_slot] == 1
+    assert runtime.sessions[("session", 0)].user_codes[:, 0].tolist() == [1, 2]
+
+
+def test_encode_appends_skips_an_identity_that_was_already_encoded_or_prepared() -> None:
+    codec = _FakeCodec()
+    runtime = _runtime(codec)
+
+    # A chunked first prefill puts the same append in two scheduler steps.
+    runtime.encode_appends([_duplex_info(seq=1)])
+    runtime.encode_appends([_duplex_info(seq=1)])
+    runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
+    runtime.encode_appends([_duplex_info(seq=1)])
+    retry = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
+
+    assert codec.encode_calls == 1
+    assert retry.user_codes[:, 0].tolist() == [1]
+
+
+def test_encode_appends_leaves_over_capacity_appends_to_prepare_append() -> None:
+    codec = _FakeCodec()
+    runtime = _runtime(codec)
+
+    runtime.encode_appends([_duplex_info(seq=1), _duplex_info(seq=1, session_id="other")])
+
+    assert list(runtime.sessions) == [("session", 0)]
+    with pytest.raises(RuntimeError, match="capacity 1"):
+        runtime.prepare_append(_duplex_info(seq=1, session_id="other"), prompt_len=18)
+
+
+def _restart_overlap(runtime: PersonaPlexStage0DuplexRuntime) -> tuple[dict, dict]:
+    """Epoch 0 of a session is live, then a cancel restarts it as epoch 1 while
+    epoch 0's aborted request still has an append in the same scheduler step."""
+    runtime.prepare_append(_duplex_info(seq=1), prompt_len=18, request_id="req-e0")
+    runtime.prepare_append(_duplex_info(seq=2), prompt_len=19, request_id="req-e0")
+    return _duplex_info(seq=3), _duplex_info(seq=1, epoch=1)
+
+
+@pytest.mark.parametrize("max_sessions", [1, 2])
+@pytest.mark.parametrize("encode_old_first", [True, False])
+@pytest.mark.parametrize("prepare_old_first", [True, False])
+def test_cancel_overlap_full_processing_order_never_re_leases_the_aborted_epoch(
+    max_sessions: int, encode_old_first: bool, prepare_old_first: bool
+) -> None:
+    codec = _FakeCodec()
+    runtime = _runtime(codec, max_sessions=max_sessions)
+    old, new = _restart_overlap(runtime)
+
+    # One scheduler step: batched encode, then per-request prepare, then sampling.
+    runtime.encode_appends([old, new] if encode_old_first else [new, old])
+    order = [("req-e0", old, 19), ("req-e1", new, 18)]
+    restarted = None
+    for request_id, duplex, prompt_len in order if prepare_old_first else order[::-1]:
+        if request_id == "req-e0":
+            with pytest.raises(PersonaPlexStage0StaleEpochError):
+                runtime.prepare_append(duplex, prompt_len=prompt_len, request_id=request_id)
+        else:
+            restarted = runtime.prepare_append(duplex, prompt_len=prompt_len, request_id=request_id)
+    for request_id in ("req-e0", "req-e1"):
+        runtime.record_sample(request_id=request_id, text_token=5, agent_codes=list(range(8)))
+    runtime.close_request("req-e0")  # the engine's late finish of the aborted request
+
+    assert list(runtime.sessions) == [("session", 1)]
+    assert len(runtime._free_slots) == max_sessions - 1
+    assert restarted is not None and restarted.user_codes[:, 0].tolist() == [1]
+    assert codec.encode_calls == 3
+    assert runtime.sessions[("session", 1)].sampled_identity == (1, 1)
+
+
+def test_codec_init_failure_propagates_after_one_attempt() -> None:
+    attempts = 0
+
+    def failing_codec():
+        nonlocal attempts
+        attempts += 1
+        raise RuntimeError("CUDA out of memory while building the Mimi encoder")
+
+    runtime = PersonaPlexStage0DuplexRuntime(
+        _FakeTalker(),
+        model_path="/unused",
+        device="cpu",
+        codec_factory=failing_codec,
+        max_sessions=16,
+        tokenizer=lambda _text: [7, 8, 9],
+        voice_loader=lambda _voice: {},
+    )
+    appends = [_duplex_info(seq=1, session_id=f"s{i}") for i in range(16)]
+
+    with pytest.raises(RuntimeError, match="out of memory"):
+        runtime.encode_appends(appends)
+    assert attempts == 1
+    assert not runtime.sessions
+
+
+def test_capacity_error_is_a_dedicated_exception() -> None:
+    runtime = _runtime(_FakeCodec(), max_sessions=1)
+    runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
+
+    with pytest.raises(PersonaPlexStage0CapacityError, match="capacity 1"):
+        runtime.prepare_append(_duplex_info(seq=1, session_id="other"), prompt_len=18)
+
+
+def test_talker_gives_the_aborted_epoch_a_neutral_row_in_the_same_step() -> None:
+    codec = _FakeCodec()
+    runtime = _runtime(codec, max_sessions=1)
+    old, new = _restart_overlap(runtime)
+    talker = PersonaPlexTalkerForConditionalGeneration.__new__(PersonaPlexTalkerForConditionalGeneration)
+    torch.nn.Module.__init__(talker)
+    talker._personaplex_duplex_stage0_runtime = runtime
+    talker._dtype = torch.float32
+    talker.mtp_hidden_size = 4
+
+    infos = {
+        "req-e0": {"duplex": old, "request_id": "req-e0", "duplex_prompt_len": 19, "duplex_token_offset": 0},
+        "req-e1": {"duplex": new, "request_id": "req-e1", "duplex_prompt_len": 18, "duplex_token_offset": 0},
+    }
+    talker.preprocess_batch(req_ids=list(infos), model_intermediate_buffer=infos, device=torch.device("cpu"))
+    outs = {
+        req: talker.preprocess(torch.zeros(1, dtype=torch.long), None, _omni_is_prefill=True, **info)
+        for req, info in infos.items()
+    }
+
+    _, stale_embeds, stale_info = outs["req-e0"]
+    assert stale_info["duplex"] == {"stage0_stale": True}
+    assert torch.count_nonzero(stale_embeds) == 0
+    assert stale_info["pplex_depformer_audio_tokens"].numel() == 16
+    assert outs["req-e1"][2]["duplex"]["epoch"] == 1
+    assert list(runtime.sessions) == [("session", 1)]

tests/model_executor/models/personaplex/test_depformer_num_steps.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.personaplex.configuration_personaplex import (
+    PersonaPlexDepformerConfig,
+)
+from vllm_omni.model_executor.models.personaplex.personaplex_depformer import PersonaPlexDepformer
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _tiny_depformer() -> PersonaPlexDepformer:
+    torch.manual_seed(0)
+    config = PersonaPlexDepformerConfig(
+        hidden_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        head_dim=8,
+        num_key_value_heads=4,
+        intermediate_size=48,
+        dep_q=16,
+        num_active_codebooks=8,
+        card=64,
+    )
+    model = PersonaPlexDepformer(config, temporal_hidden_size=24, text_card=100)
+    for param in model.parameters():
+        torch.nn.init.normal_(param, std=0.2)
+    return model.eval()
+
+
+def test_num_steps_matches_prefix_of_full_run_with_teacher_forcing():
+    model = _tiny_depformer()
+    text = torch.tensor([3, 7, 11])
+    hidden = torch.randn(3, 1, 24)
+    tokens = torch.randint(0, 64, (3, 16))
+    provided = torch.zeros(3, 16, dtype=torch.bool)
+    provided[0, 1:8] = True
+    provided[:, 8:] = True
+
+    full, full_logits = model(text, hidden, tokens, provided, return_logits=True)
+    head, head_logits = model(text, hidden, tokens, provided, return_logits=True, num_steps=8)
+
+    assert full.shape == (3, 16)
+    assert head.shape == (3, 8)
+    assert torch.equal(head, full[:, :8])
+    assert torch.equal(head_logits, full_logits[:, :8])
+
+
+@pytest.mark.parametrize("num_steps", [0, 17])
+def test_num_steps_out_of_range_is_rejected(num_steps):
+    model = _tiny_depformer()
+    with pytest.raises(ValueError, match="num_steps"):
+        model(torch.tensor([1]), torch.randn(1, 1, 24), num_steps=num_steps)

tests/model_executor/models/personaplex/test_ring_per_row_offsets.py

@@ -367,3 +367,30 @@ def test_mimi_conv_carries_preserve_inactive_rows(kind: str) -> None:
             else:
                 assert torch.equal(_stream_carry(batched)[row], carry_before[row])
                 assert torch.equal(batched._fresh[row], fresh_before[row])
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_recycled_mimi_transformer_row_matches_a_fresh_stream():
+    from vllm_omni.model_executor.models.personaplex.personaplex_mimi import _MimiStreamingTransformer
+
+    torch.manual_seed(0)
+    used = _MimiStreamingTransformer(num_layers=2, dim=16, num_heads=2, context=8)
+    for param in used.parameters():
+        torch.nn.init.normal_(param, std=0.1)
+    fresh = _MimiStreamingTransformer(num_layers=2, dim=16, num_heads=2, context=8)
+    fresh.load_state_dict(used.state_dict())
+    used.streaming_init(2)
+    fresh.streaming_init(2)
+    both = torch.tensor([True, True])
+
+    # Wrap row 0's ring (6 frames x 2 positions > context 8) before recycling it.
+    for _ in range(6):
+        used.step(torch.randn(2, 2, 16), both)
+    used.reset_slot(0)
+
+    frames = [torch.randn(2, 2, 16) for _ in range(5)]
+    for x in frames:
+        recycled = used.step(x, both)
+        expected = fresh.step(x, both)
+        torch.testing.assert_close(recycled[0], expected[0], rtol=0, atol=0)

tests/model_executor/models/personaplex/test_talker_omni_output.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.personaplex.personaplex_talker import (
+    PersonaPlexTalkerForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_make_omni_output_keeps_token_rows_for_mixed_prefill_and_live_rows():
+    # One live session carries one cached audio row; a newly admitted session
+    # prefills 150 rows in the same step. The runner indexes the hidden states
+    # with token-space logits indices, so the last prefill row must survive.
+    hidden = torch.randn(151, 16)
+    live_codes = torch.zeros(1, 8, dtype=torch.long)
+    out = PersonaPlexTalkerForConditionalGeneration.make_omni_output(
+        None,
+        hidden,
+        model_intermediate_buffer=[{"codes": {"audio": live_codes}}, {}],
+    )
+
+    assert isinstance(out, OmniOutput)
+    assert out.text_hidden_states.shape[0] == 151
+    logits_indices = torch.tensor([0, 150])
+    assert out.text_hidden_states[logits_indices].shape == (2, 16)
+    assert torch.equal(out.multimodal_outputs["codes"]["audio"], live_codes)

tests/model_executor/stage_input_processors/test_personaplex.py

@@ -93,15 +93,18 @@ def test_post_sample_talker_mtp_uses_current_temporal_state() -> None:
         *,
         audio_tokens: torch.Tensor | None = None,
         audio_provided: torch.Tensor | None = None,
+        num_steps: int | None = None,
     ) -> torch.Tensor:
         received["text_token"] = text_token
         received["hidden"] = hidden
         received["audio_tokens"] = audio_tokens
         received["audio_provided"] = audio_provided
-        return torch.arange(16, dtype=torch.long).reshape(1, 16)
+        received["num_steps"] = num_steps
+        return torch.arange(num_steps or 16, dtype=torch.long).reshape(1, -1)
 
     model = SimpleNamespace(
         _dtype=torch.float32,
+        num_active_codebooks=8,
         depformer=depformer,
         _duplex_stage0_runtime=lambda: SimpleNamespace(
             record_sample=lambda *, request_id, text_token, agent_codes: recorded.append(
@@ -126,7 +129,9 @@ def test_post_sample_talker_mtp_uses_current_temporal_state() -> None:
         ],
     )
 
-    assert codes.shape == (1, 16)
+    # Only the vocoded agent codebooks are drawn on the duplex path.
+    assert received["num_steps"] == 8
+    assert codes.shape == (1, 8)
     assert received["text_token"].tolist() == [101]
     assert received["hidden"].shape == (1, 1, 4)
     assert torch.equal(received["audio_tokens"], torch.arange(16).reshape(1, 16))

vllm_omni/deploy/personaplex.yaml

@@ -12,6 +12,10 @@
 #     8 active codebooks (cb 0..7) -> 1920 PCM samples @ 24 kHz / 12.5 Hz.
 #   * Stage 0 runs cudagraph by default; Stage 1 (Mimi) is eager because the
 #     external decoder performs host-to-device copies that cannot be captured.
+#   * Each live session holds one 3000-frame KV window (~1.5 GiB) on Stage 0.
+#     To serve more sessions, raise duplex_session.max_sessions and both stages'
+#     max_num_seqs together, and size Stage 0 gpu_memory_utilization for the KV;
+#     on a 141 GB card, 16 sessions at Stage 0 utilization 0.75 stay realtime.
 pipeline: personaplex
 session_mode: duplex
 duplex_session:
@@ -38,7 +42,12 @@ stages:
     # Prompt IDs are placeholders; voice/persona/audio semantics live in
     # request-specific embeddings, so token-ID prefix caching is invalid.
     enable_prefix_caching: false
-    async_scheduling: true
+    # Every live session appends one frame per 80 ms tick as a one-token
+    # resumable request. With async scheduling the next step is scheduled
+    # before a session's previous frame is parked, so a session can join only
+    # every other step and concurrent sessions split into alternating batches.
+    # Synchronous scheduling puts every waiting frame in one step.
+    async_scheduling: false
     max_num_batched_tokens: 512
     # Temporal transformer sliding_window is 3000 frames; keep model len within.
     max_model_len: 3000

vllm_omni/model_executor/models/personaplex/duplex/stage0.py

@@ -27,6 +27,14 @@ from vllm_omni.model_executor.models.personaplex.duplex.policy import (
 _FRAME_SAMPLES = FRAME_SIZE
 
 
+class PersonaPlexStage0CapacityError(RuntimeError):
+    """Every streaming encoder row is leased by a live session."""
+
+
+class PersonaPlexStage0StaleEpochError(RuntimeError):
+    """The append belongs to an epoch that a newer epoch of the same session has superseded."""
+
+
 @dataclass(slots=True)
 class PersonaPlexStage0PreparedAppend:
     input_ids: Any
@@ -52,7 +60,9 @@ class PersonaPlexStage0SessionState:
     prepared: PersonaPlexStage0PreparedAppend | None = None
     last_seq: int = 0
     request_ids: set[str] = field(default_factory=set)
-    codec: Any | None = None
+    slot: int | None = None
+    encoded_identity: tuple[int, int] | None = None
+    encoded_frame: Any | None = None
 
 
 def _tokenizer_path(model_path: str) -> Path:
@@ -124,7 +134,13 @@ def personaplex_prefill_slots(model_path: str, voice: str, persona: str) -> int:
 
 
 class PersonaPlexStage0DuplexRuntime:
-    """Own per-session streaming Mimi encoders and first-append prefill."""
+    """Own the shared streaming Mimi encoder and each session's first-append prefill.
+
+    One encoder holds ``max_sessions`` streaming rows; a live ``(session, epoch)``
+    leases one row for its lifetime. ``encode_appends`` encodes every new append
+    of a scheduler step in one batched call (rows without a new append are
+    inactive and keep their state), and ``prepare_append`` consumes the result.
+    """
 
     def __init__(
         self,
@@ -145,14 +161,61 @@ class PersonaPlexStage0DuplexRuntime:
         self.device = device
         self.max_sessions = max_sessions
         self._codec_factory = codec_factory
-        self._free_codecs: list[Any] = []
+        self._codec: Any | None = None
+        self._free_slots: list[int] = list(reversed(range(max_sessions)))
         self._tokenizer = tokenizer
         self._voice_loader = voice_loader
         self.sessions: dict[tuple[str, int], PersonaPlexStage0SessionState] = {}
         self.request_sessions: dict[str, tuple[str, int]] = {}
+        # Requests of a superseded epoch that still reached this step; they are
+        # finished by the engine and must not lease a row or record a sample.
+        self._stale_requests: set[str] = set()
         if codec is not None:
-            codec.streaming_init(1)
-            self._free_codecs.append(codec)
+            codec.streaming_init(max_sessions)
+            self._codec = codec
+
+    def encode_appends(self, appends: list[dict[str, Any]]) -> None:
+        """Encode the new user frame of each append in one batched encoder call.
+
+        Called once per scheduler step before the per-request ``prepare_append``
+        calls. An append whose ``(epoch, seq)`` is already encoded or prepared
+        (a chunked first prefill spans several steps) is skipped, so a row's
+        streaming state advances exactly once per frame. Appends that cannot be
+        admitted are left for ``prepare_append`` to reject.
+        """
+        parsed: list[tuple[str, int, int, dict[str, Any]]] = []
+        for duplex in appends:
+            try:
+                parsed.append((*_append_identity(duplex), duplex))
+            except ValueError:
+                continue
+        # A cancel can put a session's aborted epoch and its restarted epoch in
+        # one step. Only the newest epoch is live: admitting it closes the old
+        # one, so the old append must neither be encoded nor re-leased.
+        newest: dict[str, int] = {}
+        for session_id, epoch, _, _ in parsed:
+            newest[session_id] = max(epoch, newest.get(session_id, epoch))
+        for session_id, epoch in self.sessions:
+            if session_id in newest:
+                newest[session_id] = max(epoch, newest[session_id])
+        rows: list[tuple[PersonaPlexStage0SessionState, tuple[int, int], np.ndarray]] = []
+        for session_id, epoch, seq, duplex in parsed:
+            if epoch != newest[session_id]:
+                continue
+            try:
+                state = self._session_state(session_id, epoch)
+            except PersonaPlexStage0CapacityError:
+                # Left for prepare_append to reject. Any other error (e.g. the
+                # shared codec failing to initialize) propagates immediately.
+                continue
+            identity = (epoch, seq)
+            if identity in (state.prepared_identity, state.encoded_identity) or seq <= state.last_seq:
+                continue
+            if any(row[0] is state for row in rows):
+                continue
+            rows.append((state, identity, self._decode_pcm(duplex.get("payload"))))
+        if rows:
+            self._encode_rows(rows)
 
     def prepare_append(
         self,
@@ -163,30 +226,15 @@ class PersonaPlexStage0DuplexRuntime:
     ) -> PersonaPlexStage0PreparedAppend:
         import torch
 
-        session_id = duplex.get("session_id")
-        if not isinstance(session_id, str) or not session_id:
-            raise ValueError("PersonaPlex duplex append requires session_id")
-        epoch = _coerce_non_negative_int(duplex.get("epoch"), "epoch")
-        seq = _coerce_positive_int(duplex.get("seq"), "seq")
+        session_id, epoch, seq = _append_identity(duplex)
         identity = (epoch, seq)
         key = (session_id, epoch)
-
-        state = self.sessions.get(key)
-        if state is None:
-            # A newer epoch supersedes the session's earlier lockstep state:
-            # the engine aborted that request, but its finish notification
-            # may still be in flight, so release it here rather than let the
-            # two epochs share the codec budget.
-            for stale_key in [k for k in self.sessions if k[0] == session_id and k[1] < epoch]:
-                self.close_session(*stale_key)
-            if len(self.sessions) >= self.max_sessions:
-                raise RuntimeError(f"PersonaPlex Stage 0 session capacity {self.max_sessions} is exhausted")
-            state = PersonaPlexStage0SessionState(
-                session_id=session_id,
-                epoch=epoch,
-                codec=self._acquire_codec(),
-            )
-            self.sessions[key] = state
+        try:
+            state = self._session_state(session_id, epoch)
+        except PersonaPlexStage0StaleEpochError:
+            if request_id:
+                self._stale_requests.add(request_id)
+            raise
         if request_id:
             state.request_ids.add(request_id)
             self.request_sessions[request_id] = key
@@ -195,17 +243,11 @@ class PersonaPlexStage0DuplexRuntime:
         if seq <= state.last_seq:
             raise ValueError(f"PersonaPlex duplex append seq must increase: last={state.last_seq}, got={seq}")
 
-        pcm = self._decode_pcm(duplex.get("payload"))
-        codec = state.codec
-        if codec is None:
-            raise RuntimeError("PersonaPlex Stage 0 session has no streaming encoder")
-        encoded = codec.encode_frame(torch.from_numpy(pcm).reshape(1, _FRAME_SAMPLES))
-        user_frame = encoded.detach().to(dtype=torch.long, device="cpu").reshape(1, -1)
-        if user_frame.shape[1] < 8:
-            raise RuntimeError(
-                f"PersonaPlex Mimi encoder returned {user_frame.shape[1]} codebooks, expected at least 8"
-            )
-        user_frame = user_frame[:, :8].contiguous()
+        if state.encoded_identity != identity:
+            self._encode_rows([(state, identity, self._decode_pcm(duplex.get("payload")))])
+        user_frame = state.encoded_frame
+        state.encoded_identity = None
+        state.encoded_frame = None
         state.user_codes = user_frame if state.user_codes is None else torch.cat([state.user_codes, user_frame], dim=0)
 
         runtime_config = duplex.get("runtime_config")
@@ -342,6 +384,8 @@ class PersonaPlexStage0DuplexRuntime:
         """Commit one sampled temporal frame for the next live append."""
         import torch
 
+        if request_id in self._stale_requests:
+            return
         key = self.request_sessions.get(request_id)
         if key is None:
             raise KeyError(f"PersonaPlex Stage 0 request is not attached to a live session: {request_id}")
@@ -376,6 +420,7 @@ class PersonaPlexStage0DuplexRuntime:
         state.sampled_identity = state.prepared_identity
 
     def close_request(self, request_id: str) -> None:
+        self._stale_requests.discard(request_id)
         key = self.request_sessions.pop(request_id, None)
         if key is None:
             return
@@ -393,28 +438,75 @@ class PersonaPlexStage0DuplexRuntime:
             return
         for request_id in state.request_ids:
             self.request_sessions.pop(request_id, None)
-        if state.codec is not None:
-            state.codec.reset_streaming()
-            self._free_codecs.append(state.codec)
-            state.codec = None
+        if state.slot is not None:
+            self._shared_codec().reset_slot(state.slot)
+            self._free_slots.append(state.slot)
+            state.slot = None
 
-    def _acquire_codec(self):
-        if self._free_codecs:
-            return self._free_codecs.pop()
+    def _session_state(self, session_id: str, epoch: int) -> PersonaPlexStage0SessionState:
+        key = (session_id, epoch)
+        state = self.sessions.get(key)
+        if state is not None:
+            return state
+        # A newer epoch supersedes the session's earlier lockstep state: the
+        # engine aborted that request, but its finish notification may still
+        # be in flight, so release it here rather than let the two epochs share
+        # the encoder budget.
+        # The reverse also holds: once a newer epoch is live, an append of an
+        # older epoch is a leftover of the aborted request and must not re-lease
+        # a row (at capacity it would fail the whole step).
+        newer = [k[1] for k in self.sessions if k[0] == session_id and k[1] > epoch]
+        if newer:
+            raise PersonaPlexStage0StaleEpochError(
+                f"PersonaPlex Stage 0 epoch {epoch} of session {session_id} is superseded by epoch {max(newer)}"
+            )
+        for stale_key in [k for k in self.sessions if k[0] == session_id and k[1] < epoch]:
+            self.close_session(*stale_key)
+        if len(self.sessions) >= self.max_sessions or not self._free_slots:
+            raise PersonaPlexStage0CapacityError(
+                f"PersonaPlex Stage 0 session capacity {self.max_sessions} is exhausted"
+            )
+        self._shared_codec()
+        state = PersonaPlexStage0SessionState(session_id=session_id, epoch=epoch, slot=self._free_slots.pop())
+        self.sessions[key] = state
+        return state
+
+    def _encode_rows(self, rows: list[tuple[PersonaPlexStage0SessionState, tuple[int, int], np.ndarray]]) -> None:
+        import torch
+
+        codec = self._shared_codec()
+        pcm = torch.zeros((self.max_sessions, _FRAME_SAMPLES), dtype=torch.float32)
+        active = torch.zeros((self.max_sessions,), dtype=torch.bool)
+        for state, _, samples in rows:
+            assert state.slot is not None
+            pcm[state.slot] = torch.from_numpy(samples)
+            active[state.slot] = True
+        encoded = codec.encode_frame(pcm, active)
+        # One device-to-host copy per step, whatever the number of sessions.
+        codes = encoded.detach().to(dtype=torch.long, device="cpu")
+        if codes.shape[-1] < 8:
+            raise RuntimeError(f"PersonaPlex Mimi encoder returned {codes.shape[-1]} codebooks, expected at least 8")
+        for state, identity, _ in rows:
+            state.encoded_frame = codes[state.slot : state.slot + 1, :8].clone()
+            state.encoded_identity = identity
+
+    def _shared_codec(self):
+        if self._codec is not None:
+            return self._codec
         if self._codec_factory is not None:
             codec = self._codec_factory()
-            codec.streaming_init(1)
-            return codec
-        from vllm_omni.model_executor.models.personaplex.personaplex_mimi import (
-            PersonaPlexMimiCodec,
-        )
+        else:
+            from vllm_omni.model_executor.models.personaplex.personaplex_mimi import (
+                PersonaPlexMimiCodec,
+            )
 
-        checkpoint = Path(self.model_path) / "tokenizer-e351c8d8-checkpoint125.safetensors"
-        codec = PersonaPlexMimiCodec(
-            checkpoint=str(checkpoint) if checkpoint.is_file() else None,
-            device=self.device,
-        )
-        codec.streaming_init(1)
+            checkpoint = Path(self.model_path) / "tokenizer-e351c8d8-checkpoint125.safetensors"
+            codec = PersonaPlexMimiCodec(
+                checkpoint=str(checkpoint) if checkpoint.is_file() else None,
+                device=self.device,
+            )
+        codec.streaming_init(self.max_sessions)
+        self._codec = codec
         return codec
 
     def _load_tokenizer(self):
@@ -453,6 +545,15 @@ class PersonaPlexStage0DuplexRuntime:
         return pcm_f32le_samples(raw)
 
 
+def _append_identity(duplex: dict[str, Any]) -> tuple[str, int, int]:
+    session_id = duplex.get("session_id")
+    if not isinstance(session_id, str) or not session_id:
+        raise ValueError("PersonaPlex duplex append requires session_id")
+    epoch = _coerce_non_negative_int(duplex.get("epoch"), "epoch")
+    seq = _coerce_positive_int(duplex.get("seq"), "seq")
+    return session_id, epoch, seq
+
+
 def _coerce_non_negative_int(value: object, name: str) -> int:
     try:
         result = int(value)

vllm_omni/model_executor/models/personaplex/personaplex_depformer.py

@@ -203,8 +203,9 @@ class PersonaPlexDepformer(nn.Module):
         audio_tokens: torch.Tensor | None = None,
         audio_provided: torch.Tensor | None = None,
         return_logits: bool = False,
+        num_steps: int | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """Predict the ``dep_q`` audio codes for one frame.
+        """Predict the first ``num_steps`` (default ``dep_q``) audio codes for one frame.
 
         Args:
             text_token: ``[B]`` sampled text token for the frame.
@@ -213,19 +214,26 @@ class PersonaPlexDepformer(nn.Module):
             audio_provided: optional ``[B, dep_q]`` bool mask; where True the
                 corresponding ``audio_tokens`` entry overrides the sampled code as
                 the next-step input (partial teacher forcing, as in Moshi).
-            return_logits: also return the per-step logits ``[B, dep_q, card]``.
+            return_logits: also return the per-step logits ``[B, num_steps, card]``.
+            num_steps: stop after this many inner steps. The loop is causal, so the
+                first ``num_steps`` codes are identical to a full ``dep_q`` run;
+                the trailing steps predict the user's codebooks, which the serving
+                path never consumes.
 
         Returns:
-            ``[B, dep_q]`` long tensor of sampled (greedy) audio codes, plus the
+            ``[B, num_steps]`` long tensor of sampled (greedy) audio codes, plus the
             per-step logits when ``return_logits`` is set.
         """
         if transformer_out.dim() != 3 or transformer_out.shape[1] != 1:
             raise ValueError(f"transformer_out must be [B, 1, H]; got {tuple(transformer_out.shape)}")
+        steps = self.dep_q if num_steps is None else int(num_steps)
+        if not 1 <= steps <= self.dep_q:
+            raise ValueError(f"num_steps must be in [1, {self.dep_q}]; got {num_steps}")
         prev = text_token
         kv = [{"k": None, "v": None} for _ in self.layers]
         codes: list[torch.Tensor] = []
         logits_all: list[torch.Tensor] = []
-        for step in range(self.dep_q):
+        for step in range(steps):
             x = self.depformer_in[step](transformer_out)  # [B, 1, dim]
             if step == 0:
                 cond = self.depformer_text_emb(prev.unsqueeze(1))
@@ -243,9 +251,9 @@ class PersonaPlexDepformer(nn.Module):
                 prev = torch.where(audio_provided[:, step], audio_tokens[:, step], sampled)
             else:
                 prev = sampled
-        stacked = torch.stack(codes, dim=1)  # [B, dep_q]
+        stacked = torch.stack(codes, dim=1)  # [B, steps]
         if return_logits:
-            return stacked, torch.stack(logits_all, dim=1)  # [B, dep_q, card]
+            return stacked, torch.stack(logits_all, dim=1)  # [B, steps, card]
         return stacked
 
     # ------------------------------------------------------------------

vllm_omni/model_executor/models/personaplex/personaplex_mimi.py

@@ -249,8 +249,11 @@ class _MimiStreamingTransformer(nn.Module):
         self._offset.zero_()
 
     def reset_slot(self, b: int) -> None:
+        # A recycled row restarts at position 0, exactly like a fresh stream,
+        # instead of carrying its predecessor's absolute RoPE positions.
         for kv in self._kv:
-            kv.reset_slot(b)
+            kv.reset_row(b)
+        self._offset[b] = 0
 
     def step(self, x: torch.Tensor, active: torch.Tensor) -> torch.Tensor:
         """``x`` is ``[B, T, dim]`` (T = positions this frame, typically 2)."""

vllm_omni/model_executor/models/personaplex/personaplex_talker.py

@@ -173,7 +173,9 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module):
         if not audio_codes_list:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
         audio_codes = torch.cat(audio_codes_list, dim=0)
-        hidden = hidden[: int(audio_codes.shape[0])]
+        # Keep every token row: the runner indexes this tensor with token-space
+        # logits indices, and a step can mix a new session's multi-row prefill
+        # with one-row live appends, so the audio row count is not the token count.
         return OmniOutput(text_hidden_states=hidden, multimodal_outputs={"codes": {"audio": audio_codes}})
 
     # ------------------------------------------------------------------
@@ -323,11 +325,18 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module):
                 prompt_len = int(prompt_len_raw)
             except (TypeError, ValueError):
                 prompt_len = span
-            prepared = self._duplex_stage0_runtime().prepare_append(
-                duplex,
-                prompt_len=prompt_len,
-                request_id=(str(info_dict["request_id"]) if isinstance(info_dict.get("request_id"), str) else None),
+            from vllm_omni.model_executor.models.personaplex.duplex.stage0 import (
+                PersonaPlexStage0StaleEpochError,
             )
+
+            try:
+                prepared = self._duplex_stage0_runtime().prepare_append(
+                    duplex,
+                    prompt_len=prompt_len,
+                    request_id=(str(info_dict["request_id"]) if isinstance(info_dict.get("request_id"), str) else None),
+                )
+            except PersonaPlexStage0StaleEpochError:
+                return self._stale_append_passthrough(input_ids, span)
             offset_raw = info_dict.get("duplex_token_offset", 0)
             try:
                 offset = max(0, int(offset_raw))
@@ -435,6 +444,52 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module):
         self._personaplex_duplex_stage0_runtime = runtime
         return runtime
 
+    def _stale_append_passthrough(
+        self, input_ids: torch.Tensor, span: int
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+        """Neutral inputs for a request of a superseded epoch that is still in this step.
+
+        The engine already aborted it and discards its output; it only has to keep
+        the batch shapes valid without touching the live session's encoder row.
+        """
+        from vllm_omni.model_executor.models.personaplex.duplex.policy import SILENCE_TOKENS
+
+        device = input_ids.device
+        silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long)
+        embeds = torch.zeros((span, self.mtp_hidden_size), device=device, dtype=self._dtype)
+        return (
+            input_ids,
+            embeds,
+            {
+                "pplex_depformer_audio_tokens": torch.cat([silence, silence]),
+                "pplex_depformer_audio_provided": torch.zeros(2 * silence.numel(), dtype=torch.bool),
+                "duplex": {"stage0_stale": True},
+            },
+        )
+
+    def preprocess_batch(
+        self,
+        *,
+        req_ids: list[str],
+        model_intermediate_buffer: dict[str, dict[str, Any]],
+        device: torch.device,
+    ) -> None:
+        """Encode every live duplex append of this step in one shared-encoder call."""
+        del device
+        appends: list[dict[str, Any]] = []
+        for req_id in req_ids:
+            info = model_intermediate_buffer.get(req_id)
+            if not isinstance(info, dict):
+                continue
+            duplex = info.get("duplex")
+            if not isinstance(duplex, dict):
+                additional = info.get("additional_information")
+                duplex = additional.get("duplex") if isinstance(additional, dict) else None
+            if isinstance(duplex, dict) and duplex.get("data_plane") is True:
+                appends.append(duplex)
+        if appends:
+            self._duplex_stage0_runtime().encode_appends(appends)
+
     def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
         runtime = getattr(self, "_personaplex_duplex_stage0_runtime", None)
         if runtime is None:
@@ -526,6 +581,7 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module):
                 device=hidden.device,
                 dtype=torch.bool,
             ),
+            num_steps=self.num_active_codebooks,
         ).to(torch.long)
         runtime = self._duplex_stage0_runtime()
         for row, request_id in enumerate(req_ids):

vllm_omni/model_executor/models/personaplex/personaplex_temporal.py

@@ -92,6 +92,12 @@ class _RingKV:
         # first visible entry. (LM sacrifice-tick +1 is applied by the caller.)
         self.start_offset[b] = self.end_offset[b]
 
+    def reset_row(self, b: int) -> None:
+        # Restart row b at position 0. Every cached entry of the row sits at or
+        # past the new end offset, so all of them are masked until overwritten.
+        self.end_offset[b] = 0
+        self.start_offset[b] = 0
+
     def bump_slot_start(self, b: int) -> None:
         self.start_offset[b] += 1