[Perf][PersonaPlex] Batch live duplex appends and replay Mimi codec frames from CUDA graphs (#8356)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

recipes/NVIDIA/PersonaPlex.md

@@ -84,8 +84,10 @@ duplex (`is_enabled()` is unconditionally true). Turn-based models
 
 ## Hardware Support
 
-Verified on GPU (Hopper-class). The serving path is plain PyTorch eager plus the
-native modules, so other CUDA GPUs with enough memory are expected to work.
+Verified on GPU (Hopper-class). The shipped deploy config runs the Stage 0
+talker with CUDA graphs and replays the streaming Mimi encode (Stage 0) and
+decode (Stage 1) frame steps from model-local CUDA graphs, bitwise equal to
+eager. Other CUDA GPUs with enough memory are expected to work.
 
 ## GPU
 
@@ -137,9 +139,11 @@ python tests/e2e/online_serving/personaplex_realtime_duplex.py \
 
 #### Notes
 
-- Realtime budget is 80 ms/frame. Verified eager per-tick latency at four
-  concurrent sessions is ~70-74 ms on this hardware class, i.e. all four
-  conversations stay realtime; single-session has comfortable headroom.
+- Realtime budget is 80 ms/frame. The shipped config admits two sessions; on a
+  141 GB card, 64 concurrent 30 s sessions stayed realtime after raising
+  `duplex_session.max_sessions` and both stages' `max_num_seqs` (see the
+  sizing notes in `vllm_omni/deploy/personaplex.yaml`). Longer sessions cost
+  more per step as their context fills.
 - Decoding is greedy only (temperature/top-k knobs are intentionally not
   exposed): greedy is what the parity gates pin against the reference
   implementation, and sampling amplifies sub-bit numeric drift into divergence.

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

@@ -17,6 +17,7 @@ from vllm.sampling_params import SamplingParams
 
 from tests.e2e.online_serving import personaplex_realtime_duplex as e2e_driver
 from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.duplex.config import DuplexSessionConfig
 from vllm_omni.engine.duplex.contracts import DuplexFence
 from vllm_omni.engine.duplex.plugin import (
@@ -26,6 +27,7 @@ from vllm_omni.engine.duplex.plugin import (
     load_duplex_plugin,
     validate_duplex_plugin_sampling,
 )
+from vllm_omni.model_executor.models.personaplex.configuration_personaplex import PersonaPlexConfig
 from vllm_omni.model_executor.models.personaplex.duplex import stage0
 from vllm_omni.model_executor.models.personaplex.duplex.config import DEFAULT_PERSONA, FRAME_SIZE, SAMPLE_RATE
 from vllm_omni.model_executor.models.personaplex.duplex.data_plane import PersonaPlexDataPlaneSession
@@ -106,6 +108,36 @@ def test_personaplex_deploy_is_duplex_and_propagates_capacity_to_all_model_stage
     assert "personaplex_codec_max_sessions" not in deploy.connectors["connector_of_shared_memory"]["extra"]
 
 
+def test_personaplex_deploy_enables_mimi_cuda_graphs_on_both_stage_configs(tmp_path: Path) -> None:
+    # The checkpoint ships an empty config.json; each stage's engine args build
+    # its PersonaPlexConfig from that plus the deploy YAML's hf_overrides.
+    (tmp_path / "config.json").write_text("{}")
+    stages = merge_pipeline_deploy(PERSONAPLEX_PIPELINE, load_deploy_config(DEPLOY_PATH))
+
+    def resolve(model_arch: str, hf_overrides: dict[str, object] | None) -> PersonaPlexConfig:
+        engine_args = OmniEngineArgs(
+            model=str(tmp_path),
+            model_arch=model_arch,
+            hf_overrides=hf_overrides,
+            skip_tokenizer_init=True,
+            trust_remote_code=True,
+        )
+        config = engine_args.create_model_config().hf_config
+        assert isinstance(config, PersonaPlexConfig)
+        return config
+
+    assert [stage.yaml_engine_args["model_arch"] for stage in stages] == [
+        "PersonaPlexTalkerForConditionalGeneration",
+        "PersonaPlexCode2Wav",
+    ]
+    for stage in stages:
+        model_arch = stage.yaml_engine_args["model_arch"]
+        assert stage.yaml_engine_args["hf_overrides"] == {"mimi_cuda_graphs": True}
+        assert resolve(model_arch, dict(stage.yaml_engine_args["hf_overrides"])).mimi_cuda_graphs is True
+        # Without the deploy override the flag keeps its default.
+        assert resolve(model_arch, None).mimi_cuda_graphs is False
+
+
 # --------------------------------------------------------------------------- #
 # Session policy                                                              #
 # --------------------------------------------------------------------------- #

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

@@ -0,0 +1,312 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Batched live-append preparation must match the per-request path exactly."""
+
+import base64
+from dataclasses import dataclass
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.personaplex.duplex.policy import (
+    SILENCE_TOKENS,
+    SINE_TOKENS,
+)
+from vllm_omni.model_executor.models.personaplex.duplex.stage0 import (
+    PersonaPlexStage0DuplexRuntime,
+)
+from vllm_omni.model_executor.models.personaplex.personaplex_talker import (
+    PersonaPlexTalkerForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_WIDTH = 17  # text row + 8 agent rows + 8 user rows
+_PREFILL_LEN = 18  # 2 voice rows + 2 x 6 silence frames + 3 persona tokens + the live frame
+
+
+class _FakeCodec:
+    """Shared streaming encoder whose codes encode the row's frame count."""
+
+    def __init__(self) -> None:
+        self.encode_calls = 0
+        self.frames: list[int] = []
+
+    def streaming_init(self, batch_size: int) -> None:
+        self.frames = [0] * batch_size
+
+    def encode_frame(self, pcm: torch.Tensor, active: torch.Tensor) -> torch.Tensor:
+        del pcm
+        self.encode_calls += 1
+        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] * 16 + torch.arange(8) + 1
+        return codes
+
+    def reset_slot(self, row: int) -> None:
+        self.frames[row] = 0
+
+
+@dataclass(frozen=True)
+class _TalkerConfig:
+    num_audio_codebooks: int = 16
+    audio_vocab_size: int = 2048
+
+
+class _EmbeddingTalker:
+    """Stage model whose input embedding is the token stack itself, so embeds expose the layout."""
+
+    def __init__(self) -> None:
+        self.device = torch.device("cpu")
+        self.dtype = torch.float32
+        self.config = _TalkerConfig()
+
+    def input_embeddings(self, stack: torch.Tensor) -> torch.Tensor:
+        return stack.squeeze(-1).to(torch.float32)
+
+    def _build_prefill_embed(self, tokens, offset, span, device, silence=None, user_sine=None):
+        del offset, silence, user_sine
+        values = tokens[:span].to(device=device, dtype=torch.float32)
+        return values[:, None].expand(-1, _WIDTH).contiguous()
+
+    _build_frame_embed = PersonaPlexTalkerForConditionalGeneration._build_frame_embed
+
+
+def _depformer(text_token, hidden, *, audio_tokens, audio_provided, num_steps):
+    """Deterministic stand-in whose codes depend on every input it is given."""
+    del hidden
+    assert audio_provided.dtype == torch.bool
+    base = audio_tokens[:, 8 : 8 + num_steps] * 7 + audio_provided[:, :num_steps].long()
+    return (base + text_token[:, None] + torch.arange(num_steps)) % 2048
+
+
+def _talker(max_sessions: int) -> PersonaPlexTalkerForConditionalGeneration:
+    runtime = PersonaPlexStage0DuplexRuntime(
+        _EmbeddingTalker(),
+        model_path="/unused",
+        device="cpu",
+        codec=_FakeCodec(),
+        max_sessions=max_sessions,
+        tokenizer=lambda _text: [7, 8, 9],
+        voice_loader=lambda _voice: {"embeddings": torch.arange(2 * _WIDTH, dtype=torch.float32).reshape(2, _WIDTH)},
+    )
+    talker = PersonaPlexTalkerForConditionalGeneration.__new__(PersonaPlexTalkerForConditionalGeneration)
+    torch.nn.Module.__init__(talker)
+    talker._personaplex_duplex_stage0_runtime = runtime
+    talker._dtype = torch.float32
+    talker.mtp_hidden_size = _WIDTH
+    talker.num_active_codebooks = 8
+    talker.depformer = _depformer
+    return talker
+
+
+def _append(session_id: str, seq: int, epoch: int = 0) -> dict:
+    pcm = np.zeros(1920, dtype="<f4")
+    return {
+        "data_plane": True,
+        "session_id": session_id,
+        "epoch": epoch,
+        "seq": seq,
+        "payload": {
+            "format": "pcm_f32le",
+            "sample_rate_hz": 24000,
+            "audio": base64.b64encode(pcm.tobytes()).decode("ascii"),
+        },
+        "runtime_config": {"personaplex_voice_prompt": "NATF2.pt", "personaplex_persona": "Be concise."},
+    }
+
+
+@dataclass
+class _StepResult:
+    embeds: dict[str, torch.Tensor]
+    audio_tokens: torch.Tensor
+    audio_provided: torch.Tensor
+    codes: torch.Tensor
+    agents: dict[tuple[str, int], torch.Tensor]
+
+
+def _run_step(
+    talker: PersonaPlexTalkerForConditionalGeneration,
+    step: list[tuple[str, dict, int]],
+    *,
+    batched: bool,
+    text_base: int,
+) -> _StepResult:
+    """One scheduler step: batch preprocess, per-request preprocess, then depformer sampling.
+
+    ``step`` holds ``(request_id, append, prompt_len)``; a prompt length of
+    ``_PREFILL_LEN`` marks a first append scheduled as one full prefill span.
+    """
+    runtime = talker._personaplex_duplex_stage0_runtime
+    infos = {}
+    spans = {}
+    for request_id, duplex, prompt_len in step:
+        spans[request_id] = prompt_len if prompt_len == _PREFILL_LEN else 1
+        infos[request_id] = {
+            "duplex": duplex,
+            "request_id": request_id,
+            "duplex_prompt_len": prompt_len,
+            "duplex_token_offset": prompt_len - spans[request_id],
+        }
+    if batched:
+        talker.preprocess_batch(req_ids=list(infos), model_intermediate_buffer=infos, device=torch.device("cpu"))
+    else:
+        runtime.encode_appends([info["duplex"] for info in infos.values()])
+    embeds, req_infos = {}, []
+    for request_id, info in infos.items():
+        input_ids = torch.zeros(spans[request_id], dtype=torch.long)
+        _, embed, update = talker.preprocess(input_ids, None, _omni_is_prefill=True, **info)
+        embeds[request_id] = embed
+        req_infos.append(update)
+    received = {}
+
+    def depformer(text_token, hidden, **kwargs):
+        received.update(kwargs)
+        return _depformer(text_token, hidden, **kwargs)
+
+    talker.depformer = depformer
+    req_ids = list(infos)
+    codes = talker.post_sample_talker_mtp(
+        input_ids=torch.arange(len(req_ids)) + text_base,
+        hidden_states=torch.zeros((len(req_ids), _WIDTH)),
+        req_ids=req_ids,
+        req_infos=req_infos,
+    )
+    return _StepResult(
+        embeds=embeds,
+        audio_tokens=received["audio_tokens"],
+        audio_provided=received["audio_provided"],
+        codes=codes,
+        agents={key: state.last_agent_codes.clone() for key, state in runtime.sessions.items()},
+    )
+
+
+def _assert_same(batched: _StepResult, reference: _StepResult) -> None:
+    assert batched.embeds.keys() == reference.embeds.keys()
+    for request_id, embed in batched.embeds.items():
+        assert torch.equal(embed, reference.embeds[request_id]), request_id
+    assert torch.equal(batched.audio_tokens, reference.audio_tokens)
+    assert torch.equal(batched.audio_provided, reference.audio_provided)
+    assert torch.equal(batched.codes, reference.codes)
+    assert batched.agents.keys() == reference.agents.keys()
+    for key, agent in batched.agents.items():
+        assert torch.equal(agent, reference.agents[key]), key
+
+
+def _device_prepared(talker, key: tuple[str, int]) -> bool:
+    state = talker._personaplex_duplex_stage0_runtime.sessions[key]
+    return state.prepared.depformer_audio_tokens_device is not None
+
+
+def test_mixed_prefill_and_live_steps_match_the_per_request_path() -> None:
+    steps = [
+        [("a", _append("a", 1), 18), ("b", _append("b", 1), 18)],
+        # A first append joins a step of live appends.
+        [("a", _append("a", 2), 19), ("b", _append("b", 2), 19), ("c", _append("c", 1), 18)],
+        [("a", _append("a", 3), 20), ("b", _append("b", 3), 20), ("c", _append("c", 2), 19)],
+        # Only one session appends; the others keep their state.
+        [("c", _append("c", 3), 20)],
+        [("a", _append("a", 4), 21), ("c", _append("c", 4), 21)],
+    ]
+    batched, reference = _talker(max_sessions=4), _talker(max_sessions=4)
+
+    for index, step in enumerate(steps):
+        text_base = 100 + 10 * index
+        result = _run_step(batched, step, batched=True, text_base=text_base)
+        _assert_same(result, _run_step(reference, step, batched=False, text_base=text_base))
+
+    assert _device_prepared(batched, ("a", 0)) and _device_prepared(batched, ("c", 0))
+    assert not _device_prepared(reference, ("a", 0))
+    # Batched live appends publish no host teacher-forcing copies; the per-request path still does.
+    host_fields = {"pplex_depformer_audio_tokens", "pplex_depformer_audio_provided"}
+    batched_info = batched._personaplex_duplex_stage0_runtime.sessions[("a", 0)].prepared.info_update
+    reference_info = reference._personaplex_duplex_stage0_runtime.sessions[("a", 0)].prepared.info_update
+    assert host_fields.isdisjoint(batched_info)
+    assert host_fields <= reference_info.keys()
+    runtime = batched._personaplex_duplex_stage0_runtime
+    assert runtime._shared_codec().encode_calls == len(steps)
+    assert len(runtime.sessions[("a", 0)].user_history_device) == 2
+    # Live embeds carry the causal user delay: cb0 from the previous frame, cb1..7 from the one before.
+    a4 = result.embeds["a"][0]
+    assert a4[9].item() == 3 * 16 + 1
+    assert a4[10:17].tolist() == [2 * 16 + k + 1 for k in range(1, 8)]
+
+
+def test_second_append_reads_sine_for_the_missing_user_history() -> None:
+    talker = _talker(max_sessions=1)
+
+    def state_of():
+        return talker._personaplex_duplex_stage0_runtime.sessions[("a", 0)]
+
+    first = _run_step(talker, [("a", _append("a", 1), 18)], batched=True, text_base=100)
+    # The first append forces agent cb1..7 to silence; cb0 is sampled.
+    assert state_of().last_agent_codes.tolist() == [first.codes[0, 0].item(), *SILENCE_TOKENS[1:]]
+    result = _run_step(talker, [("a", _append("a", 2), 19)], batched=True, text_base=101)
+    # Live appends force only the user codebooks, so the sampled agent frame is kept as is.
+    assert torch.equal(state_of().last_agent_codes, result.codes[0])
+
+    embed = result.embeds["a"][0]
+    assert embed[0].item() == 100
+    assert embed[9].item() == 16 + 1
+    assert embed[10:17].tolist() == list(SINE_TOKENS[1:])
+    assert result.audio_tokens[0].tolist() == [*SILENCE_TOKENS, 2 * 16 + 1, *(16 + k + 1 for k in range(1, 8))]
+    assert result.audio_provided[0].tolist() == [False] * 8 + [True] * 8
+
+
+def test_cancel_restart_and_recycled_slot_match_the_per_request_path() -> None:
+    steps = [
+        [("a0", _append("a", 1), 18), ("b", _append("b", 1), 18)],
+        [("a0", _append("a", 2), 19), ("b", _append("b", 2), 19)],
+        # A cancel restarts session a as epoch 1 while epoch 0's aborted request
+        # still has an append in this step, next to b's live append.
+        [("a0", _append("a", 3), 20), ("a1", _append("a", 1, epoch=1), 18), ("b", _append("b", 3), 20)],
+        [("a1", _append("a", 2, epoch=1), 19), ("b", _append("b", 4), 21)],
+    ]
+    batched, reference = _talker(max_sessions=2), _talker(max_sessions=2)
+
+    for index, step in enumerate(steps):
+        text_base = 100 + 10 * index
+        result = _run_step(batched, step, batched=True, text_base=text_base)
+        _assert_same(result, _run_step(reference, step, batched=False, text_base=text_base))
+
+    runtime = batched._personaplex_duplex_stage0_runtime
+    assert sorted(runtime.sessions) == [("a", 1), ("b", 0)]
+    assert _device_prepared(batched, ("a", 1))
+    # The restarted epoch starts from an empty history on its recycled row.
+    assert result.embeds["a1"][0][10:17].tolist() == list(SINE_TOKENS[1:])
+
+    # A new session on the row of a closed one also starts from a clean history.
+    for talker in (batched, reference):
+        talker.on_requests_finished({"b"})
+    followups = [
+        [("c", _append("c", 1), 18)],
+        [("c", _append("c", 2), 19), ("a1", _append("a", 3, epoch=1), 20)],
+    ]
+    for index, step in enumerate(followups):
+        text_base = 200 + 10 * index
+        result = _run_step(batched, step, batched=True, text_base=text_base)
+        _assert_same(result, _run_step(reference, step, batched=False, text_base=text_base))
+    assert result.embeds["c"][0][10:17].tolist() == list(SINE_TOKENS[1:])
+
+
+def test_prepare_live_appends_ignores_empty_first_and_unencoded_appends() -> None:
+    talker = _talker(max_sessions=2)
+    runtime = talker._personaplex_duplex_stage0_runtime
+
+    assert runtime.prepare_live_appends([]) == 0
+    talker.preprocess_batch(req_ids=["x"], model_intermediate_buffer={"x": {}}, device=torch.device("cpu"))
+    assert runtime.sessions == {}
+
+    runtime.encode_appends([_append("a", 1)])
+    assert runtime.prepare_live_appends([_append("a", 1)]) == 0  # the first append keeps the prefill path
+    _run_step(talker, [("a", _append("a", 1), 18)], batched=False, text_base=100)
+    # Not encoded in this step: left to prepare_append.
+    assert runtime.prepare_live_appends([_append("a", 2)]) == 0
+    runtime.encode_appends([_append("a", 2)])
+    assert runtime.prepare_live_appends([_append("a", 2), _append("a", 2)]) == 1
+    # Already prepared: a repeat in a later step is a no-op.
+    assert runtime.prepare_live_appends([_append("a", 2)]) == 0
+    assert runtime.prepare_append(_append("a", 2), prompt_len=19, request_id="a").prompt_offset == 18

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

@@ -245,23 +245,23 @@ def test_repeated_append_identity_does_not_advance_codec() -> None:
 def test_next_append_uses_prior_sample_and_causally_delayed_user_frame() -> None:
     runtime = _runtime(_FakeCodec())
 
-    runtime.prepare_append(_duplex_info(seq=1), prompt_len=18, request_id="req")
-    first_agent = torch.arange(8, dtype=torch.long)
-    runtime.record_sample(request_id="req", text_token=torch.tensor(101), agent_codes=first_agent)
-    runtime.prepare_append(_duplex_info(seq=2), prompt_len=19, request_id="req")
-    second_agent = torch.arange(10, 18, dtype=torch.long)
-    runtime.record_sample(request_id="req", text_token=torch.tensor(102), agent_codes=second_agent)
-    runtime.prepare_append(_duplex_info(seq=3), prompt_len=20, request_id="req")
-
     silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long)
     sine = torch.tensor(SINE_TOKENS, dtype=torch.long)
+    runtime.prepare_append(_duplex_info(seq=1), prompt_len=18, request_id="req")
+    # The first append teacher-forces agent cb1..7 to silence.
+    expected_first_effective = torch.cat([torch.tensor([0]), silence[1:]])
+    runtime.record_sample(request_id="req", text_token=torch.tensor(101), effective_codes=expected_first_effective)
+    runtime.prepare_append(_duplex_info(seq=2), prompt_len=19, request_id="req")
+    second_agent = torch.arange(10, 18, dtype=torch.long)
+    runtime.record_sample(request_id="req", text_token=torch.tensor(102), effective_codes=second_agent)
+    runtime.prepare_append(_duplex_info(seq=3), prompt_len=20, request_id="req")
+
     first_call = runtime.stage_model.frame_calls[0]
     assert all(torch.equal(first_call[key], silence) for key in ("last_agent", "prev_agent"))
     assert all(torch.equal(first_call[key], sine) for key in ("user_d0", "user_d1"))
 
     second_call = runtime.stage_model.frame_calls[1]
     assert second_call["text_token"].tolist() == [101]
-    expected_first_effective = torch.cat([first_agent[:1], silence[1:]])
     assert torch.equal(second_call["last_agent"], expected_first_effective)
     assert torch.equal(second_call["prev_agent"], expected_first_effective)
     assert torch.equal(second_call["user_d0"], torch.full((8,), 1, dtype=torch.long))
@@ -468,7 +468,7 @@ def test_cancel_overlap_full_processing_order_never_re_leases_the_aborted_epoch(
         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.record_sample(request_id=request_id, text_token=5, effective_codes=list(range(8)))
     runtime.close_request("req-e0")  # the engine's late finish of the aborted request
 
     assert list(runtime.sessions) == [("session", 1)]

tests/model_executor/models/personaplex/test_mimi_cudagraph.py

@@ -0,0 +1,597 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fixed-buffer streaming state and CUDA graph replay of the PersonaPlex Mimi codec."""
+
+import gc
+import weakref
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.personaplex.personaplex_code2wav import _MIMI_DECODE_BATCH_FRAMES
+from vllm_omni.model_executor.models.personaplex.personaplex_mimi import (
+    CODEBOOKS,
+    FRAME_SIZE,
+    PersonaPlexMimiCodec,
+    _MimiStreamingTransformer,
+    _StreamConv1d,
+    _StreamConvTr1d,
+)
+
+pytestmark = pytest.mark.core_model
+
+SEED = 4321
+
+
+def _mask(*rows: bool, device: torch.device | str = "cpu") -> torch.Tensor:
+    return torch.tensor(rows, dtype=torch.bool, device=device)
+
+
+def _random_codec(device: torch.device, batch_size: int) -> PersonaPlexMimiCodec:
+    """The real Mimi graph with random weights (no checkpoint download)."""
+    from transformers import MimiConfig, MimiModel
+
+    torch.manual_seed(SEED)
+    codec = PersonaPlexMimiCodec.__new__(PersonaPlexMimiCodec)
+    nn.Module.__init__(codec)
+    codec.device = device
+    model = MimiModel(MimiConfig())
+    del model.encoder_transformer, model.decoder_transformer
+    with torch.no_grad():
+        for name, buffer in model.quantizer.named_buffers():
+            if name.endswith("embed_sum"):
+                buffer.normal_()
+    codec.model = model.to(device).eval()
+    codec.dtype = torch.float32
+    codec.encoder_transformer = _MimiStreamingTransformer().to(device)
+    codec.decoder_transformer = _MimiStreamingTransformer().to(device)
+    for transformer in (codec.encoder_transformer, codec.decoder_transformer):
+        for parameter in transformer.parameters():
+            nn.init.normal_(parameter, std=0.02)
+    codec._init_streaming_stages()
+    codec.streaming_init(batch_size)
+    return codec
+
+
+def _state_tensors(codec: PersonaPlexMimiCodec) -> list[torch.Tensor]:
+    tensors: list[torch.Tensor] = []
+    for state in codec._conv_states():
+        tensors.append(state.prev if isinstance(state, _StreamConv1d) else state.partial)
+        tensors.append(state._fresh)
+    for transformer in (codec.encoder_transformer, codec.decoder_transformer):
+        tensors.append(transformer._offset)
+        for kv in transformer._kv:
+            tensors.extend((kv.cache, kv.end_offset, kv.start_offset))
+    return tensors
+
+
+def _snapshot(codec: PersonaPlexMimiCodec) -> list[torch.Tensor]:
+    return [tensor.clone() for tensor in _state_tensors(codec)]
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("kind", ["conv", "conv_replicate", "convtr"])
+def test_stream_carries_update_in_place(kind: str) -> None:
+    torch.manual_seed(SEED)
+    if kind == "convtr":
+        stream = _StreamConvTr1d(nn.ConvTranspose1d(2, 3, kernel_size=4, stride=2))
+        samples = 3
+    else:
+        pad_mode = "replicate" if kind == "conv_replicate" else "constant"
+        stream = _StreamConv1d(nn.Conv1d(2, 3, kernel_size=5, stride=2), pad_mode=pad_mode)
+        samples = 6
+    stream.reset(3, torch.device("cpu"), torch.float32)
+    carry = stream.prev if isinstance(stream, _StreamConv1d) else stream.partial
+    fresh = stream._fresh
+    pointers = (carry.data_ptr(), fresh.data_ptr())
+
+    with torch.no_grad():
+        for active in (_mask(True, False, True), _mask(False, True, True), _mask(True, True, True)):
+            stream(torch.randn(3, 2, samples), active)
+            current = stream.prev if isinstance(stream, _StreamConv1d) else stream.partial
+            assert current is carry and stream._fresh is fresh
+            assert (current.data_ptr(), stream._fresh.data_ptr()) == pointers
+    assert carry.abs().sum() > 0
+    assert not fresh.any()
+
+    stream.reset_slot(1)
+    assert not carry[1].any() and fresh[1]
+    assert carry[0].abs().sum() > 0 and not fresh[0]
+
+    stream.reset_all()
+    assert not carry.any() and fresh.all()
+    assert (carry.data_ptr(), fresh.data_ptr()) == pointers
+
+
+@pytest.mark.cpu
+def test_streaming_init_at_same_batch_size_resets_in_place() -> None:
+    codec = _random_codec(torch.device("cpu"), batch_size=2)
+    pcm = torch.randn(2, FRAME_SIZE)
+    first = codec.encode_frame(pcm)
+    codec.decode_frame(first)
+    pointers = [tensor.data_ptr() for tensor in _state_tensors(codec)]
+
+    codec.streaming_init(2)
+
+    assert pointers == [tensor.data_ptr() for tensor in _state_tensors(codec)]
+    assert torch.equal(codec.encode_frame(pcm), first)
+
+    codec.streaming_init(3)
+    assert codec._batch_size == 3
+    assert codec.encode_frame(torch.randn(3, FRAME_SIZE)).shape == (3, CODEBOOKS)
+
+
+@pytest.mark.cpu
+def test_capture_off_cuda_stays_eager() -> None:
+    codec = _random_codec(torch.device("cpu"), batch_size=2)
+    assert codec.capture_cuda_graphs() == []
+    assert codec._cuda_graphs == {}
+    assert codec.encode_frame(torch.randn(2, FRAME_SIZE)).shape == (2, CODEBOOKS)
+
+
+@pytest.mark.cpu
+def test_quantizer_decode_matches_transformers() -> None:
+    codec = _random_codec(torch.device("cpu"), batch_size=2)
+    codes = torch.randint(0, 2048, (2, CODEBOOKS, 3))
+    with torch.no_grad():
+        expected = codec.model.quantizer.decode(codes)
+        actual = codec._quantizer_decode(codes)
+    assert torch.equal(actual, expected)
+
+
+def _drive(codec: PersonaPlexMimiCodec, frames: int, recycle_at: int, device: torch.device):
+    """Encode/decode ``frames`` frames under a mixed active schedule with a row recycle."""
+    batch_size = codec._batch_size
+    generator = torch.Generator(device="cpu").manual_seed(SEED)
+    schedule = [_mask(*[(step + row) % 3 != 0 for row in range(batch_size)], device=device) for step in range(frames)]
+    codes_out, pcm_out, inactive_ok = [], [], True
+    for step, active in enumerate(schedule):
+        if step == recycle_at:
+            codec.reset_slot(1)
+        pcm = 0.1 * torch.randn(batch_size, FRAME_SIZE, generator=generator)
+        before = _snapshot(codec)
+        codes = codec.encode_frame(pcm, active)
+        wav = codec.decode_frame(codes, active)
+        after = _state_tensors(codec)
+        for row, is_active in enumerate(active.tolist()):
+            if is_active:
+                continue
+            for old, new in zip(before, after):
+                if new.dim() > 0 and new.shape[0] == batch_size:
+                    inactive_ok &= torch.equal(old[row], new[row])
+                elif new.dim() > 1 and new.shape[1] == batch_size:
+                    inactive_ok &= torch.equal(old[:, row], new[:, row])
+        codes_out.append(codes)
+        pcm_out.append(wav)
+    return torch.stack(codes_out), torch.stack(pcm_out), inactive_ok
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_graph_replay_is_bitwise_equal_to_eager_across_recycle() -> None:
+    device = torch.device("cuda")
+    codec = _random_codec(device, batch_size=3)
+    eager_codes, eager_pcm, eager_inactive_ok = _drive(codec, frames=12, recycle_at=7, device=device)
+    assert eager_inactive_ok
+    assert torch.isfinite(eager_pcm).all()
+
+    codec.streaming_init(3)
+    pointers = [tensor.data_ptr() for tensor in _state_tensors(codec)]
+    assert codec.capture_cuda_graphs(decode_frame_counts=(1, 3)) == ["decode_f1", "decode_f3", "encode"]
+    assert pointers == [tensor.data_ptr() for tensor in _state_tensors(codec)]
+
+    graph_codes, graph_pcm, graph_inactive_ok = _drive(codec, frames=12, recycle_at=7, device=device)
+
+    assert codec._cuda_graphs["encode"].replays == 12
+    assert codec._cuda_graphs["decode_f1"].replays == 12
+    assert graph_inactive_ok
+    assert torch.equal(graph_codes, eager_codes)
+    assert torch.equal(graph_pcm, eager_pcm)
+
+    # A full reset in place (a new stream on the same codec) replays the same
+    # frames from a fresh state.
+    codec.streaming_init(3)
+    again_codes, again_pcm, _ = _drive(codec, frames=12, recycle_at=7, device=device)
+    assert torch.equal(again_codes, eager_codes)
+    assert torch.equal(again_pcm, eager_pcm)
+    assert pointers == [tensor.data_ptr() for tensor in _state_tensors(codec)]
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("batch_size", [1, 3])
+def test_multi_frame_decode_graph_matches_eager_chunks(batch_size: int) -> None:
+    # Code2Wav records the first single-frame delta and the full chunk, on one
+    # B=1 decoder per session. B=3 runs the same graphs on shared rows, where
+    # every call carries all rows and only some are active.
+    chunk = _MIMI_DECODE_BATCH_FRAMES
+    partial = chunk - 2
+    assert partial not in (1, chunk)
+    device = torch.device("cuda")
+    codec = _random_codec(device, batch_size=batch_size)
+    generator = torch.Generator(device="cpu").manual_seed(SEED)
+    # First frame alone, full chunks, a partial chunk with no graph, then a
+    # single frame and a full chunk again on the state the eager step left.
+    frame_counts = [1, chunk, chunk, partial, 1, chunk]
+    row_masks = [(True, True, False), (True, False, True), (True, True, True), (False, True, True)]
+    steps = [
+        (frames, None if batch_size == 1 else _mask(*row_masks[step % len(row_masks)], device=device))
+        for step, frames in enumerate(frame_counts)
+    ]
+    chunks = [torch.randint(0, 2048, (batch_size, CODEBOOKS, frames), generator=generator) for frames, _ in steps]
+
+    eager = [codec.decode_frames(chunk_codes, active) for chunk_codes, (_, active) in zip(chunks, steps)]
+    codec.streaming_init(batch_size)
+    assert codec.capture_cuda_graphs(encode=False, decode_frame_counts=(1, chunk)) == ["decode_f1", f"decode_f{chunk}"]
+    graphed = [codec.decode_frames(chunk_codes, active) for chunk_codes, (_, active) in zip(chunks, steps)]
+
+    # The partial chunk has no graph, so it is the one step that ran eagerly.
+    assert codec._cuda_graphs["decode_f1"].replays == 2
+    assert codec._cuda_graphs[f"decode_f{chunk}"].replays == 3
+    assert f"decode_f{partial}" not in codec._cuda_graphs
+    for (frames, active), expected, actual in zip(steps, eager, graphed):
+        assert actual.shape == (batch_size, frames * FRAME_SIZE)
+        rows = slice(None) if active is None else active
+        assert torch.equal(actual[rows], expected[rows])
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_failed_capture_warns_with_the_error_and_stays_eager(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.model_executor.models.personaplex import personaplex_mimi_cudagraph
+
+    device = torch.device("cuda")
+    codec = _random_codec(device, batch_size=2)
+    pcm = 0.1 * torch.randn(2, FRAME_SIZE, generator=torch.Generator(device="cpu").manual_seed(SEED))
+    expected_codes = codec.encode_frame(pcm)
+    expected_pcm = codec.decode_frame(expected_codes)
+    codec.streaming_init(2)
+
+    class _FailingCapture:
+        def __enter__(self) -> None:
+            pass
+
+        def __exit__(self, *_args) -> None:
+            raise RuntimeError("CUDA error: out of memory")
+
+    warnings: list[tuple[tuple, dict]] = []
+    monkeypatch.setattr(torch.cuda, "graph", lambda *_args, **_kwargs: _FailingCapture())
+    monkeypatch.setattr(
+        personaplex_mimi_cudagraph.logger,
+        "warning",
+        lambda *args, **kwargs: warnings.append((args, kwargs)),
+    )
+
+    assert codec.capture_cuda_graphs(encode=True, decode_frame_counts=(1,)) == []
+
+    assert codec._cuda_graphs == {}
+    assert len(warnings) == 1
+    args, kwargs = warnings[0]
+    assert "requested" in args[0] and "capture failed" in args[0] and "eagerly" in args[0]
+    assert args[1:3] == ("encode/decode_f1", 2)
+    assert "Traceback (most recent call last):" in args[3]
+    assert "RuntimeError: CUDA error: out of memory" in args[3]
+    assert not kwargs
+    # The warmup frames are rolled back, so the eager codec starts a fresh stream.
+    assert torch.equal(codec.encode_frame(pcm), expected_codes)
+    assert torch.equal(codec.decode_frame(expected_codes), expected_pcm)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_capture_does_not_hide_errors_that_are_not_cuda_failures(monkeypatch: pytest.MonkeyPatch) -> None:
+    codec = _random_codec(torch.device("cuda"), batch_size=2)
+
+    def _broken(*_args, **_kwargs):
+        raise TypeError("not a capture failure")
+
+    monkeypatch.setattr(codec, "_encode_frame_eager", _broken)
+    with pytest.raises(TypeError, match="not a capture failure"):
+        codec.capture_cuda_graphs(encode=True, decode_frame_counts=())
+    assert codec._cuda_graphs == {}
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_graph_capture_leaves_a_fresh_stream_and_resize_drops_graphs() -> None:
+    device = torch.device("cuda")
+    codec = _random_codec(device, batch_size=2)
+    codec.capture_cuda_graphs(encode=True, decode_frame_counts=())
+    assert list(codec._cuda_graphs) == ["encode"]
+    for state in codec._conv_states():
+        carry = state.prev if isinstance(state, _StreamConv1d) else state.partial
+        assert not carry.any() and state._fresh.all()
+    for transformer in (codec.encoder_transformer, codec.decoder_transformer):
+        assert not transformer._offset.any()
+        assert all(not kv.end_offset.any() for kv in transformer._kv)
+
+    # Resizing reallocates the state, so graphs recorded at the old size are dropped.
+    codec.streaming_init(4)
+    assert codec._cuda_graphs == {}
+    assert codec.encode_frame(torch.zeros(4, FRAME_SIZE)).shape == (4, CODEBOOKS)
+
+
+class _RecordingCodec:
+    """Stage 0 shared-encoder stand-in that records the graph capture request."""
+
+    def __init__(self) -> None:
+        self.calls: list[tuple[str, object]] = []
+
+    def streaming_init(self, batch_size: int) -> None:
+        self.calls.append(("streaming_init", batch_size))
+
+    def capture_cuda_graphs(self, *, encode: bool, decode_frame_counts: tuple[int, ...]) -> list[str]:
+        self.calls.append(("capture", {"encode": encode, "decode_frame_counts": decode_frame_counts}))
+        return ["encode"]
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("enabled", [True, False])
+def test_stage0_captures_the_shared_encoder_once_when_enabled(enabled: bool) -> None:
+    from vllm_omni.model_executor.models.personaplex.duplex.stage0 import PersonaPlexStage0DuplexRuntime
+
+    codec = _RecordingCodec()
+    runtime = PersonaPlexStage0DuplexRuntime(
+        object(),
+        model_path="/unused",
+        device="cpu",
+        codec_factory=lambda: codec,
+        max_sessions=4,
+        codec_cuda_graphs=enabled,
+    )
+
+    assert runtime._shared_codec() is codec
+    assert runtime._shared_codec() is codec
+    expected: list[tuple[str, object]] = [("streaming_init", 4)]
+    if enabled:
+        expected.append(("capture", {"encode": True, "decode_frame_counts": ()}))
+    assert codec.calls == expected
+
+
+@pytest.mark.cpu
+def test_resize_after_capture_warns_that_graphs_are_dropped(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.model_executor.models.personaplex import personaplex_mimi
+
+    warnings: list[tuple] = []
+    monkeypatch.setattr(personaplex_mimi.logger, "warning", lambda *args: warnings.append(args))
+    codec = _random_codec(torch.device("cpu"), batch_size=2)
+    codec._cuda_graphs = {"encode": object()}
+
+    codec.streaming_init(2)
+    assert "encode" in codec._cuda_graphs and not warnings
+
+    codec.streaming_init(3)
+    assert codec._cuda_graphs == {}
+    assert len(warnings) == 1 and warnings[0][1:] == (2, 3, "encode")
+
+
+@pytest.fixture
+def capture_faults(monkeypatch: pytest.MonkeyPatch):
+    """Exercise capture orchestration and error boundaries with CPU tensors."""
+    from vllm_omni.model_executor.models.personaplex import personaplex_mimi_cudagraph as graphs
+
+    state = SimpleNamespace(
+        errors={},
+        counts={},
+        events=[],
+        warnings=[],
+        phase="setup",
+        capturing=False,
+        graph_refs=[],
+        capture_refs=[],
+        output_refs=[],
+        check_released=False,
+    )
+    state.original_stream = state.stream = object()
+
+    def record(event):
+        state.events.append(event)
+        state.counts[event] = state.counts.get(event, 0) + 1
+        error = state.errors.get((event, state.counts[event]), state.errors.get(event))
+        if error is not None:
+            raise error
+
+    def eager(name):
+        def step(frame, active):
+            state.phase = "capture" if state.capturing else "warmup"
+            record(f"{name}_{state.phase}")
+            if state.capturing:
+                output = frame.clone()
+                state.output_refs.append(weakref.ref(output))
+                return output
+            return frame
+
+        return step
+
+    def reset():
+        assert state.stream is state.original_stream
+        if state.check_released:
+            assert all(ref() is None for ref in state.graph_refs + state.capture_refs + state.output_refs)
+        state.phase = "reset"
+        record("reset")
+
+    codec = SimpleNamespace(
+        device="cuda",
+        dtype=torch.float32,
+        _batch_size=2,
+        _encode_frame_eager=eager("encode"),
+        _decode_frames_eager=eager("decode"),
+        reset_streaming=reset,
+    )
+
+    class Graph:
+        def __init__(self):
+            record("construct")
+            state.graph_refs.append(weakref.ref(self))
+
+    class Capture:
+        def __init__(self, graph):
+            self.graph = graph
+            state.capture_refs.append(weakref.ref(self))
+
+        def __enter__(self):
+            record("capture_setup_sync")
+            record("capture_setup_empty_cache")
+            state.stream = object()
+            state.capturing = True
+            record("capture_begin")
+
+        def __exit__(self, exc_type, exc, traceback):
+            state.capturing = False
+            # Match torch.cuda.graph: capture_end can raise before restoration,
+            # and can replace a body exception with capture-invalidated.
+            record("capture_end")
+            state.stream = state.original_stream
+
+    @contextmanager
+    def stream_context(stream):
+        previous = state.stream
+        state.stream = stream
+        try:
+            yield
+        finally:
+            state.stream = previous
+
+    zeros, ones = torch.zeros, torch.ones
+    monkeypatch.setattr(torch, "zeros", lambda *args, **kwargs: zeros(*args, **(kwargs | {"device": "cpu"})))
+    monkeypatch.setattr(torch, "ones", lambda *args, **kwargs: ones(*args, **(kwargs | {"device": "cpu"})))
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "graph_pool_handle", lambda: (0, 0))
+    monkeypatch.setattr(torch.cuda, "CUDAGraph", Graph)
+    monkeypatch.setattr(torch.cuda, "graph", lambda graph, **kwargs: Capture(graph))
+    monkeypatch.setattr(torch.cuda, "current_stream", lambda device: state.stream)
+    monkeypatch.setattr(torch.cuda, "stream", stream_context)
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda device: record(f"sync_{state.phase}"))
+    monkeypatch.setattr(graphs.logger, "warning", lambda *args, **kwargs: state.warnings.append((args, kwargs)))
+    state.capture = lambda: graphs.capture_mimi_frame_graphs(
+        codec,
+        encode=True,
+        decode_frame_counts=(1,),
+        warmup_iters=1,
+    )
+    return state
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("step", ["encode_warmup", "decode_warmup"])
+@pytest.mark.parametrize("error_type", [RuntimeError, torch.cuda.OutOfMemoryError])
+def test_eager_warmup_failures_propagate_before_capture(capture_faults, step, error_type):
+    error = error_type("injected eager tensor failure")
+    capture_faults.errors[step] = error
+    with pytest.raises(error_type) as raised:
+        capture_faults.capture()
+    assert raised.value is error
+    assert capture_faults.counts.get("construct", 0) == 0
+    assert capture_faults.counts["reset"] == 1
+    assert not capture_faults.warnings
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "point",
+    [
+        "sync_warmup",
+        "construct",
+        "capture_setup_sync",
+        "capture_setup_empty_cache",
+        "capture_begin",
+        "sync_capture",
+        "sync_reset",
+        "reset",
+    ],
+)
+def test_device_and_reset_failures_are_never_capture_fallbacks(capture_faults, point):
+    # Even an otherwise recoverable error class must propagate at these boundaries.
+    error = torch.cuda.OutOfMemoryError("injected device or reset failure")
+    capture_faults.errors[point] = error
+    with pytest.raises(torch.cuda.OutOfMemoryError) as raised:
+        capture_faults.capture()
+    assert raised.value is error
+    assert not capture_faults.warnings
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "message",
+    [
+        "tensor size mismatch",
+        "CUDA error: an illegal memory access was encountered",
+        "CUDA error: device-side assert triggered",
+        "CUDA error: operation failed due to a previous error during capture",
+    ],
+)
+def test_unknown_or_fatal_capture_errors_propagate(capture_faults, message):
+    error = RuntimeError(message)
+    capture_faults.errors["encode_capture"] = error
+    with pytest.raises(RuntimeError) as raised:
+        capture_faults.capture()
+    assert raised.value is error
+    assert capture_faults.stream is capture_faults.original_stream
+    assert capture_faults.counts["reset"] == 1
+    assert not capture_faults.warnings
+
+
+@pytest.mark.cpu
+def test_recoverable_capture_failure_discards_partial_graphs_after_all_warmups(capture_faults):
+    error = torch.cuda.OutOfMemoryError("capture allocation failed")
+    capture_faults.errors[("capture_end", 2)] = error
+    capture_faults.check_released = True
+    was_enabled = gc.isenabled()
+    gc.disable()
+    try:
+        assert capture_faults.capture() == {}
+    finally:
+        if was_enabled:
+            gc.enable()
+    assert capture_faults.events.index("decode_warmup") < capture_faults.events.index("construct")
+    assert capture_faults.counts["capture_end"] == 2  # an earlier graph completed
+    assert capture_faults.counts["sync_reset"] == 1
+    assert capture_faults.stream is capture_faults.original_stream
+    assert len(capture_faults.warnings) == 1
+    assert "capture allocation failed" in capture_faults.warnings[0][0][-1]
+    assert "Traceback (most recent call last):" in capture_faults.warnings[0][0][-1]
+    assert not capture_faults.warnings[0][1]
+    assert error.__traceback__ is None
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("recoverable", [True, False])
+def test_capture_end_does_not_hide_the_original_failure(capture_faults, recoverable):
+    message = "CUDA error: operation not permitted when stream is capturing" if recoverable else "tensor size mismatch"
+    original = RuntimeError(message)
+    capture_faults.errors["encode_capture"] = original
+    invalidated = RuntimeError("CUDA error: operation failed due to a previous error during capture")
+    capture_faults.errors["capture_end"] = invalidated
+    if recoverable:
+        capture_faults.check_released = True
+        was_enabled = gc.isenabled()
+        gc.disable()
+        try:
+            assert capture_faults.capture() == {}
+        finally:
+            if was_enabled:
+                gc.enable()
+        assert len(capture_faults.warnings) == 1
+        assert original.__traceback__ is None and invalidated.__traceback__ is None
+    else:
+        with pytest.raises(RuntimeError) as raised:
+            capture_faults.capture()
+        assert raised.value is invalidated and raised.value.__context__ is original
+        assert not capture_faults.warnings
+    assert capture_faults.stream is capture_faults.original_stream
+    assert capture_faults.counts["reset"] == 1
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("point", ["reset", "sync_reset"])
+def test_capture_fallback_requires_successful_reset(capture_faults, point):
+    capture_faults.errors["capture_end"] = torch.cuda.OutOfMemoryError("capture allocation failed")
+    error = RuntimeError("codec could not reset")
+    capture_faults.errors[point] = error
+    with pytest.raises(RuntimeError) as raised:
+        capture_faults.capture()
+    assert raised.value is error
+    assert not capture_faults.warnings

tests/model_executor/models/personaplex/test_ring_per_row_offsets.py

@@ -122,6 +122,8 @@ def _make_codec_stub(mocker: MockerFixture) -> PersonaPlexMimiCodec:
         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])
+    # The codec decodes through its own RVQ sum; route it to the stub.
+    codec._quantizer_decode = quantizer.decode
 
     codec.device = torch.device("cpu")
     codec.dtype = torch.float32

tests/model_executor/stage_input_processors/test_personaplex.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 
@@ -100,16 +100,18 @@ def test_post_sample_talker_mtp_uses_current_temporal_state() -> None:
         received["audio_tokens"] = audio_tokens
         received["audio_provided"] = audio_provided
         received["num_steps"] = num_steps
-        return torch.arange(num_steps or 16, dtype=torch.long).reshape(1, -1)
+        return 100 + torch.arange(num_steps or 16, dtype=torch.long).reshape(1, -1)
 
     model = SimpleNamespace(
         _dtype=torch.float32,
         num_active_codebooks=8,
         depformer=depformer,
+        _depformer_teacher_forcing=PersonaPlexTalkerForConditionalGeneration._depformer_teacher_forcing,
         _duplex_stage0_runtime=lambda: SimpleNamespace(
-            record_sample=lambda *, request_id, text_token, agent_codes: recorded.append(
-                (request_id, text_token.clone(), agent_codes.clone())
-            )
+            prepared_depformer_state=lambda _request_id: None,
+            record_sample=lambda *, request_id, text_token, effective_codes: recorded.append(
+                (request_id, text_token.clone(), effective_codes.clone())
+            ),
         ),
     )
     method = getattr(PersonaPlexTalkerForConditionalGeneration, "post_sample_talker_mtp", None)
@@ -137,4 +139,6 @@ def test_post_sample_talker_mtp_uses_current_temporal_state() -> None:
     assert torch.equal(received["audio_tokens"], torch.arange(16).reshape(1, 16))
     assert torch.equal(received["audio_provided"], (torch.arange(16) > 0).reshape(1, 16))
     assert [(row[0], row[1].item()) for row in recorded] == [("r1", 101)]
-    assert torch.equal(recorded[0][2], codes[0])
+    # Teacher-forced agent codebooks (1..7 here) replace the sampled ones.
+    assert codes[0].tolist() == list(range(100, 108))
+    assert recorded[0][2].tolist() == [100, 1, 2, 3, 4, 5, 6, 7]

vllm_omni/deploy/personaplex.yaml

@@ -10,12 +10,19 @@
 #     first-chunk latency (not correctness).
 #   * Stage 1 (code2wav) wraps the external Mimi codec. Each codec frame is
 #     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.
+#   * Stage 0 runs cudagraph by default; Stage 1 (Mimi) runs with
+#     enforce_eager because the external decoder performs host-to-device copies
+#     that the runner cannot capture. The streaming Mimi encode (Stage 0) and
+#     decode (Stage 1) frame steps instead replay model-local CUDA graphs
+#     (hf_overrides.mimi_cuda_graphs), bitwise equal to eager. The graphs add
+#     about 80 MiB on Stage 0 and 5 MiB per decoder slot on Stage 1.
 #   * 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.
+#     max_num_seqs together, and size Stage 0 gpu_memory_utilization for
+#     N x 1.5 GiB of KV: on a 141 GB card, utilization 0.75 holds about 60 full
+#     windows. Measured there (Stage 1 at 0.15), 64 concurrent 30 s sessions
+#     (375 frames each) stayed realtime. Longer sessions also cost more per
+#     step as their context fills.
 pipeline: personaplex
 session_mode: duplex
 duplex_session:
@@ -52,6 +59,8 @@ stages:
     # Temporal transformer sliding_window is 3000 frames; keep model len within.
     max_model_len: 3000
     devices: "0"
+    hf_overrides:
+      mimi_cuda_graphs: true
     output_connectors:
       to_stage_1: connector_of_shared_memory
     default_sampling_params:
@@ -77,6 +86,8 @@ stages:
     max_num_batched_tokens: 65536
     max_model_len: 65536
     devices: "0"
+    hf_overrides:
+      mimi_cuda_graphs: true
     input_connectors:
       from_stage_0: connector_of_shared_memory
     default_sampling_params:

vllm_omni/model_executor/models/personaplex/configuration_personaplex.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Configuration for PersonaPlex (a Moshi finetune; 2-stage audio->audio pipeline).
 
 PersonaPlex is a staged AR speech model composed of:
@@ -203,6 +203,10 @@ class PersonaPlexConfig(PretrainedConfig):
         mimi_name (`str`, *optional*):
             Convenience mirror of ``mimi_config.mimi_name``; if set, it
             overrides the value carried inside ``mimi_config``.
+        mimi_cuda_graphs (`bool`, *optional*, defaults to `False`):
+            Replay the streaming Mimi per-frame encode (Stage 0) and decode
+            (Stage 1) from model-local CUDA graphs. Can be enabled through
+            ``hf_overrides``.
     """
 
     model_type = "personaplex"
@@ -222,6 +226,7 @@ class PersonaPlexConfig(PretrainedConfig):
         audio_vocab_size: int = 2048,
         num_audio_codebooks: int = 16,
         mimi_name: str | None = None,
+        mimi_cuda_graphs: bool = False,
         **kwargs: Any,
     ) -> None:
         if temporal_config is None:
@@ -257,6 +262,7 @@ class PersonaPlexConfig(PretrainedConfig):
         if mimi_name is not None:
             self.mimi_config.mimi_name = mimi_name
         self.mimi_name = self.mimi_config.mimi_name
+        self.mimi_cuda_graphs = mimi_cuda_graphs
 
     @staticmethod
     def _coerce(value: Any, config_cls: type[PretrainedConfig]) -> PretrainedConfig:

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

@@ -9,9 +9,10 @@ from collections.abc import Callable
 from dataclasses import dataclass, field
 from functools import lru_cache
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
+from vllm.logger import init_logger
 
 from vllm_omni.model_executor.common.audio.pcm import pcm_f32le_samples
 from vllm_omni.model_executor.common.duplex.payload import decode_pcm_f32le_payload
@@ -24,8 +25,13 @@ from vllm_omni.model_executor.models.personaplex.duplex.policy import (
     wrap_with_system_tags,
 )
 
+if TYPE_CHECKING:
+    import torch
+
 _FRAME_SAMPLES = FRAME_SIZE
 
+logger = init_logger(__name__)
+
 
 class PersonaPlexStage0CapacityError(RuntimeError):
     """Every streaming encoder row is leased by a live session."""
@@ -43,6 +49,9 @@ class PersonaPlexStage0PreparedAppend:
     info_update: dict[str, Any]
     prefill_applied: bool
     prompt_offset: int
+    # Device copies of the depformer teacher-forcing state (batched path only).
+    depformer_audio_tokens_device: torch.Tensor | None = None
+    depformer_audio_provided_device: torch.Tensor | None = None
 
 
 @dataclass(slots=True)
@@ -63,6 +72,10 @@ class PersonaPlexStage0SessionState:
     slot: int | None = None
     encoded_identity: tuple[int, int] | None = None
     encoded_frame: Any | None = None
+    # Device copies of the encoded frame and of the last two consumed user
+    # frames (newest last), so live appends are built without host round trips.
+    encoded_frame_device: torch.Tensor | None = None
+    user_history_device: list[torch.Tensor] = field(default_factory=list)
 
 
 def _tokenizer_path(model_path: str) -> Path:
@@ -151,6 +164,7 @@ class PersonaPlexStage0DuplexRuntime:
         codec: Any | None = None,
         codec_factory: Callable[[], Any] | None = None,
         max_sessions: int = 1,
+        codec_cuda_graphs: bool = False,
         tokenizer=None,
         voice_loader=None,
     ) -> None:
@@ -161,6 +175,7 @@ class PersonaPlexStage0DuplexRuntime:
         self.device = device
         self.max_sessions = max_sessions
         self._codec_factory = codec_factory
+        self._codec_cuda_graphs = codec_cuda_graphs
         self._codec: Any | None = None
         self._free_slots: list[int] = list(reversed(range(max_sessions)))
         self._tokenizer = tokenizer
@@ -170,6 +185,7 @@ class PersonaPlexStage0DuplexRuntime:
         # 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()
+        self._constants_cache: dict[str, torch.Tensor] | None = None
         if codec is not None:
             codec.streaming_init(max_sessions)
             self._codec = codec
@@ -239,13 +255,22 @@ class PersonaPlexStage0DuplexRuntime:
             state.request_ids.add(request_id)
             self.request_sessions[request_id] = key
         if state.prepared_identity == identity and state.prepared is not None:
-            return state.prepared
+            prepared = state.prepared
+            if prepared.prompt_offset < 0:
+                prepared.prompt_offset = int(prompt_len) - int(prepared.inputs_embeds.shape[0])
+                if prepared.prompt_offset < 0:
+                    raise ValueError(
+                        "PersonaPlex scheduler prompt reservation mismatch: "
+                        f"reserved={prompt_len}, prepared={prepared.inputs_embeds.shape[0]}"
+                    )
+            return prepared
         if seq <= state.last_seq:
             raise ValueError(f"PersonaPlex duplex append seq must increase: last={state.last_seq}, got={seq}")
 
         if state.encoded_identity != identity:
             self._encode_rows([(state, identity, self._decode_pcm(duplex.get("payload")))])
         user_frame = state.encoded_frame
+        self._push_user_history(state)
         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)
@@ -344,44 +369,188 @@ class PersonaPlexStage0DuplexRuntime:
                 f"PersonaPlex scheduler prompt reservation mismatch: reserved={prompt_len}, prepared={prepared_len}"
             )
         input_ids = torch.zeros((prepared_len,), dtype=torch.long, device=device)
-        info_update = {
-            "pplex_user_codes": state.user_codes,
-            "pplex_silence_codes": silence.detach().cpu(),
-            "pplex_depformer_audio_tokens": depformer_audio_tokens.detach().cpu(),
-            "pplex_depformer_audio_provided": depformer_audio_provided.detach().cpu(),
-            "meta": {
-                "pplex_frame": state.prefill_slots + int(state.user_codes.shape[0]),
-                "pplex_prefill_len": state.prefill_slots,
-            },
-            "duplex": {
-                "stage0_prepared": True,
-                "prefill_applied": first_append,
-                "session_id": session_id,
-                "epoch": epoch,
-                "seq": seq,
-            },
-        }
-        prepared = PersonaPlexStage0PreparedAppend(
+        info_update = self._prepared_info_update(
+            state, seq, silence=silence.detach().cpu(), prefill_applied=first_append
+        )
+        # This path has no device teacher-forcing copies, so the talker reads
+        # these host tensors from the request information.
+        info_update["pplex_depformer_audio_tokens"] = depformer_audio_tokens.detach().cpu()
+        info_update["pplex_depformer_audio_provided"] = depformer_audio_provided.detach().cpu()
+        return self._store_prepared(
+            state,
+            seq,
+            PersonaPlexStage0PreparedAppend(
                 input_ids=input_ids,
                 inputs_embeds=full_embeds,
                 user_codes=state.user_codes,
                 info_update=info_update,
                 prefill_applied=first_append,
                 prompt_offset=prompt_offset,
+            ),
         )
-        state.prepared_identity = identity
+
+    @staticmethod
+    def _prepared_info_update(
+        state: PersonaPlexStage0SessionState,
+        seq: int,
+        *,
+        silence: torch.Tensor,
+        prefill_applied: bool,
+    ) -> dict[str, Any]:
+        """Request-information update shared by the per-request and batched append paths."""
+        return {
+            "pplex_user_codes": state.user_codes,
+            "pplex_silence_codes": silence,
+            "meta": {
+                "pplex_frame": state.prefill_slots + int(state.user_codes.shape[0]),
+                "pplex_prefill_len": state.prefill_slots,
+            },
+            "duplex": {
+                "stage0_prepared": True,
+                "prefill_applied": prefill_applied,
+                "session_id": state.session_id,
+                "epoch": state.epoch,
+                "seq": seq,
+            },
+        }
+
+    @staticmethod
+    def _store_prepared(
+        state: PersonaPlexStage0SessionState,
+        seq: int,
+        prepared: PersonaPlexStage0PreparedAppend,
+    ) -> PersonaPlexStage0PreparedAppend:
+        """Record ``prepared`` as the session's current append and advance its sequence."""
+        state.prepared_identity = (state.epoch, seq)
         state.prepared = prepared
         state.last_seq = seq
         return prepared
 
+    @staticmethod
+    def _push_user_history(state: PersonaPlexStage0SessionState) -> None:
+        """Move the consumed encoded frame into the device history (last two frames)."""
+        if state.encoded_frame_device is None:
+            raise RuntimeError("PersonaPlex Stage 0 append has no encoded user frame")
+        state.user_history_device.append(state.encoded_frame_device)
+        del state.user_history_device[:-2]
+        state.encoded_frame_device = None
+
+    def _constants(self) -> dict[str, torch.Tensor]:
+        """Live-append constant frames on the device (plus host silence), built once."""
+        if self._constants_cache is not None:
+            return self._constants_cache
+        import torch
+
+        device, _ = self._model_device_dtype()
+        silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long)
+        self._constants_cache = {
+            "silence": silence,
+            "silence_device": silence.to(device),
+            "sine_device": torch.tensor(SINE_TOKENS, dtype=torch.long, device=device),
+            "live_provided_device": torch.tensor([False] * 8 + [True] * 8, dtype=torch.bool, device=device),
+            "zero_text_device": torch.tensor([ZERO_TEXT_TOKEN], dtype=torch.long, device=device),
+        }
+        return self._constants_cache
+
+    def prepare_live_appends(self, appends: list[dict[str, Any]]) -> int:
+        """Prepare every one-frame live append of this step with one batched embedding call.
+
+        Only appends whose frame was encoded this step and that are not the
+        first append of their epoch (which carries the voice/persona prefill)
+        take this path; the rest are left to :meth:`prepare_append`. Results are
+        stored exactly as :meth:`prepare_append` stores them, so its per-request
+        call returns the cached preparation. Returns the number prepared.
+        """
+        import torch
+
+        batch: dict[tuple[str, int], tuple[PersonaPlexStage0SessionState, int]] = {}
+        for duplex in appends:
+            try:
+                session_id, epoch, seq = _append_identity(duplex)
+            except ValueError:
+                continue
+            key = (session_id, epoch)
+            state = self.sessions.get(key)
+            # Only an append encoded in this step and not yet consumed: its seq
+            # is past last_seq, so it has not been prepared either.
+            if state is None or state.last_seq == 0 or seq <= state.last_seq or key in batch:
+                continue
+            if state.encoded_identity != (epoch, seq):
+                continue
+            batch[key] = (state, seq)
+        if not batch:
+            return 0
+        c = self._constants()
+        device, dtype = self._model_device_dtype()
+        sine_d, silence_d = c["sine_device"], c["silence_device"]
+        texts, agents, d0s, d1s, currents = [], [], [], [], []
+        for state, _ in batch.values():
+            history = state.user_history_device  # consumed frames, newest last
+            texts.append(
+                state.last_text_token.reshape(1) if state.last_text_token is not None else c["zero_text_device"]
+            )
+            agents.append(state.last_agent_codes[:8] if state.last_agent_codes is not None else silence_d)
+            # The frame being appended is not in the history yet: user cb0 reads
+            # the previous frame and user cb1..7 the one before it.
+            d0s.append(history[-1] if len(history) >= 1 else sine_d)
+            d1s.append(history[-2] if len(history) >= 2 else sine_d)
+            currents.append(state.encoded_frame_device)
+        n = len(batch)
+        text = torch.stack(texts)
+        agent = torch.stack(agents)
+        user_d0 = torch.stack(d0s)
+        user_d1 = torch.stack(d1s)
+        current = torch.stack(currents)
+        stack = torch.cat([text, agent, user_d0[:, :1], user_d1[:, 1:8]], dim=1).unsqueeze(-1)
+        embeds = self.stage_model.input_embeddings(stack).reshape(n, 1, -1).to(dtype=dtype)
+        dep_tokens_device = torch.cat([silence_d.expand(n, 8), current[:, :1], user_d0[:, 1:8]], dim=1)
+        input_ids = torch.zeros((n, 1), dtype=torch.long, device=device)
+        for row, (state, seq) in enumerate(batch.values()):
+            user_frame = state.encoded_frame
+            self._push_user_history(state)
+            state.encoded_identity = None
+            state.encoded_frame = None
+            state.user_codes = torch.cat([state.user_codes, user_frame], dim=0)
+            # The talker reads teacher forcing from the device copies below, so
+            # no host tokens or mask are built for these appends.
+            self._store_prepared(
+                state,
+                seq,
+                PersonaPlexStage0PreparedAppend(
+                    input_ids=input_ids[row],
+                    inputs_embeds=embeds[row],
+                    user_codes=state.user_codes,
+                    info_update=self._prepared_info_update(state, seq, silence=c["silence"], prefill_applied=False),
+                    prefill_applied=False,
+                    # Resolved by prepare_append from the scheduler's prompt length.
+                    prompt_offset=-1,
+                    depformer_audio_tokens_device=dep_tokens_device[row],
+                    depformer_audio_provided_device=c["live_provided_device"],
+                ),
+            )
+        return n
+
+    def prepared_depformer_state(self, request_id: str) -> tuple[torch.Tensor, torch.Tensor] | None:
+        """Device teacher-forcing tensors of the request's current append, when prepared on the device."""
+        key = self.request_sessions.get(request_id)
+        state = self.sessions.get(key) if key is not None else None
+        prepared = state.prepared if state is not None else None
+        if prepared is None or prepared.depformer_audio_tokens_device is None:
+            return None
+        return prepared.depformer_audio_tokens_device, prepared.depformer_audio_provided_device
+
     def record_sample(
         self,
         *,
         request_id: str,
         text_token: Any,
-        agent_codes: Any,
+        effective_codes: Any,
     ) -> None:
-        """Commit one sampled temporal frame for the next live append."""
+        """Commit one sampled temporal frame for the next live append.
+
+        ``effective_codes`` are the frame's eight agent codes with the append's
+        depformer teacher forcing already applied, which the next append reads.
+        """
         import torch
 
         if request_id in self._stale_requests:
@@ -396,26 +565,12 @@ class PersonaPlexStage0DuplexRuntime:
             return
 
         text = torch.as_tensor(text_token, device=self._model_device_dtype()[0], dtype=torch.long).reshape(-1)
-        codes = torch.as_tensor(agent_codes, device=text.device, dtype=torch.long).reshape(-1)
+        codes = torch.as_tensor(effective_codes, device=text.device, dtype=torch.long).reshape(-1)
         if text.numel() != 1:
             raise ValueError(f"PersonaPlex Stage 0 expected one sampled text token, got {text.numel()}")
-        if codes.numel() < 8:
-            raise ValueError(f"PersonaPlex Stage 0 expected at least 8 agent codes, got {codes.numel()}")
-
-        effective_codes = codes[:8].clone()
-        prepared = state.prepared
-        if prepared is None:
-            raise RuntimeError(f"PersonaPlex Stage 0 request has no prepared payload: {request_id}")
-        target = prepared.info_update.get("pplex_depformer_audio_tokens")
-        provided = prepared.info_update.get("pplex_depformer_audio_provided")
-        if not isinstance(target, torch.Tensor) or not isinstance(provided, torch.Tensor):
-            raise RuntimeError("PersonaPlex Stage 0 prepared payload has no depformer teacher-forcing state")
-        target = target.reshape(-1).to(device=effective_codes.device, dtype=torch.long)
-        provided = provided.reshape(-1).to(device=effective_codes.device, dtype=torch.bool)
-        if target.numel() < 8 or provided.numel() < 8:
-            raise RuntimeError("PersonaPlex Stage 0 depformer teacher-forcing state is shorter than one agent frame")
-        effective_codes = torch.where(provided[:8], target[:8], effective_codes)
-        state.last_agent_codes = effective_codes.detach()
+        if codes.numel() != 8:
+            raise ValueError(f"PersonaPlex Stage 0 expected 8 effective agent codes, got {codes.numel()}")
+        state.last_agent_codes = codes.detach()
         state.last_text_token = text.detach()
         state.sampled_identity = state.prepared_identity
 
@@ -482,12 +637,14 @@ class PersonaPlexStage0DuplexRuntime:
             pcm[state.slot] = torch.from_numpy(samples)
             active[state.slot] = True
         encoded = codec.encode_frame(pcm, active)
+        codes_device = encoded.detach()[:, :8].to(dtype=torch.long)
         # 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_frame_device = codes_device[state.slot]
             state.encoded_identity = identity
 
     def _shared_codec(self):
@@ -506,6 +663,15 @@ class PersonaPlexStage0DuplexRuntime:
                 device=self.device,
             )
         codec.streaming_init(self.max_sessions)
+        if self._codec_cuda_graphs:
+            # Rows are leased and recycled in place, so one graph at
+            # max_sessions with the active mask as input serves every step.
+            captured = codec.capture_cuda_graphs(encode=True, decode_frame_counts=())
+            logger.info(
+                "PersonaPlex Stage 0 Mimi encoder %s at %d rows",
+                "replays a CUDA graph" if captured else "runs eagerly",
+                self.max_sessions,
+            )
         self._codec = codec
         return codec
 

vllm_omni/model_executor/models/personaplex/personaplex_code2wav.py

@@ -406,6 +406,28 @@ class PersonaPlexCode2Wav(nn.Module):
             ).eval()
             for _ in range(self._max_codec_sessions)
         ]
+        if getattr(self.config, "mimi_cuda_graphs", False):
+            # Each decoder is leased per request and reset in place, so its B=1
+            # graphs stay valid. Deltas arrive as the first single frame and then
+            # full chunks, so those two frame counts are recorded; other counts
+            # run eagerly. The graphs replay one at a time on this stage's stream
+            # and copy their output, so they share one pool.
+            pool = torch.cuda.graph_pool_handle() if torch.device(str(device)).type == "cuda" else None
+            captured = 0
+            for codec in codecs:
+                codec.streaming_init(1)
+                captured += bool(
+                    codec.capture_cuda_graphs(
+                        encode=False,
+                        decode_frame_counts=(1, _MIMI_DECODE_BATCH_FRAMES),
+                        pool=pool,
+                    )
+                )
+            logger.info(
+                "PersonaPlex Code2Wav replays Mimi decode from CUDA graphs on %d/%d stream(s)",
+                captured,
+                len(codecs),
+            )
         self._set_mimi_codecs(codecs)
         self._mimi_device = torch.device(str(device))
         reported_sr = getattr(codecs[0].model.config, "sampling_rate", None)

vllm_omni/model_executor/models/personaplex/personaplex_mimi.py

@@ -29,9 +29,12 @@ codec composes with elastic slot recycling in batched duplex serving.
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from vllm.logger import init_logger
 
 from vllm_omni.model_executor.models.personaplex.personaplex_temporal import (
     _apply_rope,
@@ -42,13 +45,17 @@ DEFAULT_HF_REPO = "kyutai/mimi"
 FRAME_SIZE = 1920
 CODEBOOKS = 8
 
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    from vllm_omni.model_executor.models.personaplex.personaplex_mimi_cudagraph import MimiFrameGraph
+
 
 def _normalize_active(active: torch.Tensor | None, all_active: torch.Tensor) -> torch.Tensor:
-    # Stage 0 and Code2Wav currently use one B=1 codec per session and omit
-    # `active`, so preserve that API by treating None as all rows active. When
-    # they switch to shared B>1 codecs, their batch builders must pass bool[B]:
+    # Stage 0 shares one codec with a row per session slot and passes bool[B]:
     # True advances that row's streaming state; False keeps an absent or padded
-    # row's offsets and convolution carries unchanged.
+    # row's offsets and convolution carries unchanged. Code2Wav leases one B=1
+    # codec per request and omits `active`, which means all rows are active.
     if active is None:
         return all_active
     if active.shape != all_active.shape:
@@ -107,6 +114,10 @@ class _StreamConv1d:
     ``pad_mode`` controls the stream-start left padding: SEANet convs use zeros
     (``constant``); the down/upsample resamplers use ``replicate`` (the first
     real sample), marked per row so elastic slot recycling re-primes correctly.
+
+    The carry and the fresh-row flags are allocated once in ``reset`` and only
+    ever updated in place, so a CUDA graph recorded over ``__call__`` keeps
+    reading and writing the live state across frames, row resets and sessions.
     """
 
     def __init__(self, conv: nn.Conv1d, pad_mode: str = "constant") -> None:
@@ -122,29 +133,36 @@ class _StreamConv1d:
         self.prev = torch.zeros(batch_size, self.conv.in_channels, pad, device=device, dtype=dtype)
         self._fresh = torch.ones(batch_size, dtype=torch.bool, device=device)
 
+    def reset_all(self) -> None:
+        self.prev.zero_()
+        self._fresh.fill_(True)
+
     def reset_slot(self, b: int) -> None:
         self.prev[b].zero_()
         self._fresh[b] = True
 
     def __call__(self, x: torch.Tensor, active: torch.Tensor) -> torch.Tensor:
+        active_view = active.view(-1, 1, 1)
         if self.pad_mode == "replicate":
             pad = self.prev.shape[-1]
             edge = x[..., 0:1].expand(-1, -1, pad)
             fresh = (self._fresh & active).view(-1, 1, 1)
-            self.prev = torch.where(fresh, edge.to(self.prev.dtype), self.prev)
-        self._fresh[active] = False
+            self.prev.copy_(torch.where(fresh, edge.to(self.prev.dtype), self.prev))
+        self._fresh.logical_and_(~active)
         x = torch.cat([self.prev, x], dim=-1)
         t = x.shape[-1]
         num_frames = max(0, (t - self.kernel) // self.stride + 1)
-        prev = x[..., num_frames * self.stride :]
-        self.prev = torch.where(active.view(-1, 1, 1), prev, self.prev)
+        self.prev.copy_(torch.where(active_view, x[..., num_frames * self.stride :], self.prev))
         if num_frames == 0:
             return x.new_zeros(x.shape[0], self.conv.out_channels, 0)
         return self.conv(x[..., : (num_frames - 1) * self.stride + self.kernel])
 
 
 class _StreamConvTr1d:
-    """Moshi ``RawStreamingConvTranspose1d``: overlap-add tail carry per call."""
+    """Moshi ``RawStreamingConvTranspose1d``: overlap-add tail carry per call.
+
+    Like ``_StreamConv1d``, the tail carry is updated in place only.
+    """
 
     def __init__(self, conv: nn.ConvTranspose1d) -> None:
         self.conv = conv
@@ -159,6 +177,10 @@ class _StreamConvTr1d:
         )
         self._fresh = torch.ones(batch_size, dtype=torch.bool, device=device)
 
+    def reset_all(self) -> None:
+        self.partial.zero_()
+        self._fresh.fill_(True)
+
     def reset_slot(self, b: int) -> None:
         self.partial[b].zero_()
         self._fresh[b] = True
@@ -168,17 +190,17 @@ class _StreamConvTr1d:
         length = out.shape[-1]
         tail = self.kernel - self.stride
         pt = self.partial.shape[-1]
-        merge = self.partial.clone()
+        merge = self.partial
         if self.conv.bias is not None:
             # The carried tail already includes the bias; the fresh output adds
             # it again, so subtract one copy -- except on a row's very first
             # frame, where the carry is zeros by construction.
             merge = merge - self.conv.bias[:, None]
-            merge[self._fresh & active] = 0.0
-            self._fresh[active] = False
+            first = (self._fresh & active).view(-1, 1, 1)
+            merge = torch.where(first, torch.zeros_like(merge), merge)
+            self._fresh.logical_and_(~active)
         out[..., :pt] += merge
-        partial = out[..., length - tail :].clone()
-        self.partial = torch.where(active.view(-1, 1, 1), partial, self.partial)
+        self.partial.copy_(torch.where(active.view(-1, 1, 1), out[..., length - tail :], self.partial))
         return out[..., : length - tail]
 
 
@@ -358,13 +380,16 @@ class PersonaPlexMimiCodec(nn.Module):
         n_dec = self.decoder_transformer.load_weights(sd, "decoder_transformer")
         assert n_enc == n_dec == 80, (n_enc, n_dec)
         del sd
+        self._init_streaming_stages()
 
+    def _init_streaming_stages(self) -> None:
         m = self.model
         self._enc_stages = _walk_seanet(m.encoder.layers)
         self._downsample = _StreamConv1d(m.downsample.conv, pad_mode="replicate")
         self._upsample = _StreamConvTr1d(m.upsample.conv)
         self._dec_stages = _walk_seanet(m.decoder.layers)
         self._batch_size: int | None = None
+        self._cuda_graphs: dict[str, MimiFrameGraph] = {}
         self._all_active: torch.Tensor
 
     # -- streaming state ------------------------------------------------------
@@ -380,7 +405,22 @@ class PersonaPlexMimiCodec(nn.Module):
         yield self._upsample
 
     def streaming_init(self, batch_size: int) -> None:
+        # Re-initializing at the same batch size resets in place: the state
+        # tensors keep their storage, so captured frame graphs stay valid.
+        if getattr(self, "_batch_size", None) == batch_size:
+            self.reset_streaming()
+            return
+        if getattr(self, "_cuda_graphs", None):
+            # The state is reallocated below, so graphs recorded at the old
+            # batch size would read stale tensors; drop them and say so.
+            logger.warning(
+                "PersonaPlex Mimi stream resized from %s to %d rows; dropping its CUDA graphs (%s)",
+                self._batch_size,
+                batch_size,
+                "/".join(sorted(self._cuda_graphs)),
+            )
         self._batch_size = batch_size
+        self._cuda_graphs = {}
         self._all_active = torch.ones(batch_size, dtype=torch.bool, device=self.device)
         for s in self._conv_states():
             s.reset(batch_size, self.device, self.dtype)
@@ -390,8 +430,7 @@ class PersonaPlexMimiCodec(nn.Module):
     def reset_streaming(self) -> None:
         assert self._batch_size is not None
         for state in self._conv_states():
-            for b in range(self._batch_size):
-                state.reset_slot(b)
+            state.reset_all()
         self.encoder_transformer.reset_streaming()
         self.decoder_transformer.reset_streaming()
 
@@ -401,6 +440,42 @@ class PersonaPlexMimiCodec(nn.Module):
         self.encoder_transformer.reset_slot(b)
         self.decoder_transformer.reset_slot(b)
 
+    def capture_cuda_graphs(
+        self,
+        *,
+        encode: bool = True,
+        decode_frame_counts: tuple[int, ...] = (1,),
+        warmup_iters: int = 2,
+        pool: tuple[int, int] | None = None,
+    ) -> list[str]:
+        """Record the per-frame codec steps at the streaming batch size.
+
+        ``encode`` records ``encode_frame``; each ``F`` in ``decode_frame_counts``
+        records ``decode_frames`` over ``F`` frames (``decode_frame`` is ``F=1``).
+        Every row is recorded active; the ``active`` mask is a graph input, so
+        one graph serves any subset of live rows. Warmup frames run through the
+        real streaming state, which is then reset in place, so the codec is
+        fresh afterwards. ``pool`` is a CUDA graph memory pool to share with
+        other codecs replayed on the same stream; by default the codec gets a
+        private one. Returns the captured graph names; on a non-CUDA device the
+        codec stays eager, and so it does after a recoverable capture failure,
+        which is logged with its traceback. Warmup, device execution and reset
+        failures propagate instead of leaving an apparently usable codec.
+        """
+        from vllm_omni.model_executor.models.personaplex.personaplex_mimi_cudagraph import (
+            capture_mimi_frame_graphs,
+        )
+
+        assert self._batch_size is not None, "call streaming_init first"
+        self._cuda_graphs = capture_mimi_frame_graphs(
+            self,
+            encode=encode,
+            decode_frame_counts=decode_frame_counts,
+            warmup_iters=warmup_iters,
+            pool=pool,
+        )
+        return sorted(self._cuda_graphs)
+
     # -- per-frame codec -------------------------------------------------------
 
     @staticmethod
@@ -413,32 +488,68 @@ class PersonaPlexMimiCodec(nn.Module):
                 x = stage(x, active) if kind in {"conv", "convtr"} else stage(x)
         return x
 
+    def _frame_graph(self, name: str) -> MimiFrameGraph | None:
+        graph = self._cuda_graphs.get(name)
+        if graph is None or torch.cuda.is_current_stream_capturing():
+            return None
+        return graph
+
     @torch.no_grad()
     def encode_frame(self, pcm: torch.Tensor, active: torch.Tensor | None = None) -> torch.Tensor:
         """``[B, frame_size]`` float PCM -> ``[B, 8]`` codes."""
-        x = pcm.to(self.device, self.dtype).view(-1, 1, FRAME_SIZE)
         active = _normalize_active(active, self._all_active)
+        graph = self._frame_graph("encode")
+        if graph is not None:
+            return graph.replay(pcm.reshape(-1, FRAME_SIZE), active)
+        return self._encode_frame_eager(pcm, active)
+
+    def _encode_frame_eager(self, pcm: torch.Tensor, active: torch.Tensor) -> torch.Tensor:
+        x = pcm.to(self.device, self.dtype).view(-1, 1, FRAME_SIZE)
         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, num_quantizers=CODEBOOKS)  # [Q, B, T]
         return codes[:CODEBOOKS, :, 0].transpose(0, 1).contiguous()
 
+    def _quantizer_decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """``MimiSplitResidualVectorQuantizer.decode`` without its host-side scalar.
+
+        The Hugging Face RVQ decode seeds its sum with ``torch.tensor(0.0)``, a
+        host-to-device copy that cannot be recorded in a CUDA graph. This
+        seeds it with an on-device zero instead; the arithmetic is unchanged.
+        """
+        quantizer = self.model.quantizer
+        split = quantizer.num_semantic_quantizers
+
+        def rvq_decode(rvq, rvq_codes: torch.Tensor) -> torch.Tensor:
+            out = torch.zeros((), device=rvq_codes.device)
+            for layer, indices in zip(rvq.layers, rvq_codes.transpose(0, 1)):
+                out = out + layer.decode(indices)
+            if rvq.output_proj is not None:
+                out = rvq.output_proj(out)
+            return out
+
+        out = rvq_decode(quantizer.semantic_residual_vector_quantizer, codes[:, :split])
+        if codes.shape[1] > split:
+            out += rvq_decode(quantizer.acoustic_residual_vector_quantizer, codes[:, split:])
+        return out
+
     @torch.no_grad()
     def decode_frame(self, codes: torch.Tensor, active: torch.Tensor | None = None) -> torch.Tensor:
         """``[B, 8]`` codes -> ``[B, frame_size]`` float PCM."""
-        emb = self.model.quantizer.decode(codes.to(self.device).view(-1, CODEBOOKS, 1))
-        active = _normalize_active(active, self._all_active)
-        emb = self._upsample(emb, active)
-        emb = self.decoder_transformer.step(emb.transpose(1, 2), active).transpose(1, 2)
-        x = self._run_stages(emb, self._dec_stages, active)
-        return x[:, 0, :]
+        return self.decode_frames(codes.reshape(-1, CODEBOOKS, 1), active)
 
     @torch.no_grad()
     def decode_frames(self, codes: torch.Tensor, active: torch.Tensor | None = None) -> torch.Tensor:
         """``[B, 8, F]`` codes -> ``[B, F * frame_size]`` float PCM."""
-        emb = self.model.quantizer.decode(codes.to(self.device))
         active = _normalize_active(active, self._all_active)
+        graph = self._frame_graph(f"decode_f{codes.shape[-1]}")
+        if graph is not None:
+            return graph.replay(codes, active)
+        return self._decode_frames_eager(codes, active)
+
+    def _decode_frames_eager(self, codes: torch.Tensor, active: torch.Tensor) -> torch.Tensor:
+        emb = self._quantizer_decode(codes.to(self.device))
         emb = self._upsample(emb, active)
         emb = self.decoder_transformer.step(emb.transpose(1, 2), active).transpose(1, 2)
         x = self._run_stages(emb, self._dec_stages, active)

vllm_omni/model_executor/models/personaplex/personaplex_mimi_cudagraph.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Replay the PersonaPlex streaming Mimi frame steps from CUDA graphs.
+
+One 80 ms Mimi encode or decode frame is a few hundred small kernels (SEANet
+convs, an 8-layer streaming transformer and the RVQ), so eager execution is
+launch bound. Every streaming state tensor of ``PersonaPlexMimiCodec`` (conv
+carries, fresh-row flags, ring KV caches and offsets) is allocated once by
+``streaming_init`` and only updated in place, and row resets are in place too,
+so a graph recorded at the streaming batch size stays valid across frames,
+session slot recycling and full stream resets.
+
+The ``active`` row mask is a graph input: one graph at the shared batch size
+serves any subset of live rows, and inactive rows keep their state exactly as
+in eager execution.
+"""
+
+from __future__ import annotations
+
+import traceback
+from collections.abc import Callable
+from contextlib import ExitStack
+from typing import TYPE_CHECKING
+
+import torch
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm_omni.model_executor.models.personaplex.personaplex_mimi import PersonaPlexMimiCodec
+
+logger = init_logger(__name__)
+
+__all__ = ["MimiFrameGraph", "capture_mimi_frame_graphs"]
+
+
+def _is_recoverable_capture_error(error: RuntimeError) -> bool:
+    """Allow only known capture limitations, never an unexplained CUDA error."""
+    pending: list[BaseException] = [error]
+    seen: set[int] = set()
+    known_failure = False
+    while pending:
+        current = pending.pop()
+        if id(current) in seen:
+            continue
+        seen.add(id(current))
+        message = str(current).split("\n", 1)[0]
+        if isinstance(current, torch.cuda.OutOfMemoryError) or (
+            isinstance(current, RuntimeError)
+            and (
+                message.startswith("CUDA out of memory.")
+                or message
+                in {
+                    "CUDA error: out of memory",
+                    "CUDA error: operation not permitted when stream is capturing",
+                }
+            )
+        ):
+            known_failure = True
+        elif not (
+            isinstance(current, RuntimeError)
+            and message == "CUDA error: operation failed due to a previous error during capture"
+        ):
+            return False
+        # capture_end can replace a body exception with capture-invalidated.
+        # The original failure must also be recoverable, even if suppressed.
+        pending.extend(exc for exc in (current.__cause__, current.__context__) if exc is not None)
+    return known_failure
+
+
+def _capture_error_details(error: RuntimeError) -> str:
+    """Preserve diagnostics without retaining capture frames or graph pools."""
+    details = "".join(traceback.format_exception(error))
+    pending: list[BaseException] = [error]
+    seen: set[int] = set()
+    while pending:
+        current = pending.pop()
+        if id(current) in seen:
+            continue
+        seen.add(id(current))
+        if current.__traceback__ is not None:
+            # ExitStack frames can retain their own traceback in exception tuples.
+            traceback.clear_frames(current.__traceback__)
+        current.__traceback__ = None
+        pending.extend(exc for exc in (current.__cause__, current.__context__) if exc is not None)
+    return details
+
+
+class MimiFrameGraph:
+    """A recorded ``(frame, active) -> output`` codec step with static buffers."""
+
+    def __init__(
+        self,
+        graph: torch.cuda.CUDAGraph,
+        static_input: torch.Tensor,
+        static_active: torch.Tensor,
+        static_output: torch.Tensor,
+        eager: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+    ) -> None:
+        self.graph = graph
+        self.static_input = static_input
+        self.static_active = static_active
+        self.static_output = static_output
+        self.eager = eager
+        self.replays = 0
+        self._warned = False
+
+    def replay(self, frame: torch.Tensor, active: torch.Tensor) -> torch.Tensor:
+        if frame.shape != self.static_input.shape or active.shape != self.static_active.shape:
+            if not self._warned:
+                self._warned = True
+                logger.warning(
+                    "PersonaPlex Mimi graph input shape %s does not match capture shape %s; running eagerly",
+                    tuple(frame.shape),
+                    tuple(self.static_input.shape),
+                )
+            return self.eager(frame, active)
+        self.static_input.copy_(frame)
+        self.static_active.copy_(active)
+        self.graph.replay()
+        self.replays += 1
+        # The static output is overwritten by the next replay.
+        return self.static_output.clone()
+
+
+def capture_mimi_frame_graphs(
+    codec: PersonaPlexMimiCodec,
+    *,
+    encode: bool,
+    decode_frame_counts: tuple[int, ...],
+    warmup_iters: int = 2,
+    pool: tuple[int, int] | None = None,
+) -> dict[str, MimiFrameGraph]:
+    """Capture the codec's encode and ``F``-frame decode steps.
+
+    Returns an empty dict off CUDA or after a recoverable capture failure.
+    Warmup, device execution and stream-reset failures always propagate.
+    """
+    from vllm_omni.model_executor.models.personaplex.personaplex_mimi import CODEBOOKS, FRAME_SIZE
+
+    device = torch.device(codec.device)
+    if device.type != "cuda" or not torch.cuda.is_available():
+        logger.info("PersonaPlex Mimi CUDA graphs skipped on device %s", device)
+        return {}
+    batch_size = codec._batch_size
+    specs: list[tuple[str, Callable, torch.Tensor]] = []
+    if encode:
+        specs.append(
+            (
+                "encode",
+                codec._encode_frame_eager,
+                torch.zeros(batch_size, FRAME_SIZE, device=device, dtype=codec.dtype),
+            )
+        )
+    for frames in sorted(set(decode_frame_counts)):
+        specs.append(
+            (
+                f"decode_f{frames}",
+                codec._decode_frames_eager,
+                torch.zeros(batch_size, CODEBOOKS, frames, device=device, dtype=torch.long),
+            )
+        )
+    # A private pool by default: these graphs are captured outside the model
+    # runner's own capture sequence, so they must not reuse its pool blocks.
+    if pool is None:
+        pool = torch.cuda.graph_pool_handle()
+    graphs: dict[str, MimiFrameGraph] = {}
+    warmed_specs: list[tuple[str, Callable, torch.Tensor, torch.Tensor]] = []
+    capture_error_details: str | None = None
+    try:
+        with torch.no_grad():
+            # Validate every eager step before capture can select the fallback.
+            for name, eager, static_input in specs:
+                static_active = torch.ones(batch_size, dtype=torch.bool, device=device)
+                for _ in range(max(warmup_iters, 1)):
+                    eager(static_input, static_active)
+                torch.accelerator.synchronize(device)
+                warmed_specs.append((name, eager, static_input, static_active))
+            for name, eager, static_input, static_active in warmed_specs:
+                # capture_end may raise before torch.cuda.graph restores its
+                # stream. Always restore our caller's stream on that path.
+                with torch.cuda.stream(torch.cuda.current_stream(device)):
+                    graph = torch.cuda.CUDAGraph()
+                    # __enter__ synchronizes and clears the CUDA cache. A
+                    # failure there is a setup failure, not a capture fallback.
+                    capture = torch.cuda.graph(graph, pool=pool, capture_error_mode="thread_local")
+                    capture.__enter__()
+                    try:
+                        # Other worker threads may issue CUDA calls while this records.
+                        with ExitStack() as stack:
+                            stack.push(capture)
+                            static_output = eager(static_input, static_active)
+                    except RuntimeError as error:
+                        if not _is_recoverable_capture_error(error):
+                            raise
+                        capture_error_details = _capture_error_details(error)
+                        graphs.clear()
+                        # Release partial graphs and graph-owned output tensors
+                        # before reset, including the last loop iteration.
+                        graph = static_output = capture = None
+                        break
+                # A device execution failure is not an optional capture failure.
+                torch.accelerator.synchronize(device)
+                graphs[name] = MimiFrameGraph(graph, static_input, static_active, static_output, eager)
+    finally:
+        # Warmup frames advanced the real streaming state; restore a fresh
+        # stream without reallocating the tensors the graphs point at.
+        codec.reset_streaming()
+    # Reset enqueues device work too. Do not report usable graphs or an eager
+    # fallback unless that work completed successfully.
+    torch.accelerator.synchronize(device)
+    if capture_error_details is not None:
+        logger.warning(
+            "PersonaPlex Mimi CUDA graphs were requested (%s at batch size %d) but capture failed; "
+            "the codec runs eagerly, which lowers realtime session capacity\n%s",
+            "/".join(name for name, _, _ in specs),
+            batch_size,
+            capture_error_details,
+        )
+    if graphs:
+        logger.debug("Captured PersonaPlex Mimi %s graph(s) at batch size %d", "/".join(graphs), batch_size)
+    return graphs

vllm_omni/model_executor/models/personaplex/personaplex_talker.py

@@ -30,7 +30,7 @@ encode of the input WAV (built in ``preprocess``); live duplex is Phase 2.
 from __future__ import annotations
 
 from collections.abc import Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 from torch import nn
@@ -54,6 +54,11 @@ from vllm_omni.model_executor.models.personaplex.personaplex_embeddings import (
     PersonaPlexInputEmbeddings,
 )
 
+if TYPE_CHECKING:
+    from vllm_omni.model_executor.models.personaplex.duplex.stage0 import (
+        PersonaPlexStage0DuplexRuntime,
+    )
+
 __all__ = ["PersonaPlexTalkerForConditionalGeneration"]
 
 
@@ -437,6 +442,7 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module):
             model_path=model_path,
             device=device,
             max_sessions=int(getattr(self.vllm_config.model_config, "duplex_max_sessions", 1)),
+            codec_cuda_graphs=bool(getattr(self.config, "mimi_cuda_graphs", False)),
         )
         self._personaplex_duplex_stage0_runtime = runtime
         return runtime
@@ -485,7 +491,9 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module):
             if isinstance(duplex, dict) and duplex.get("data_plane") is True:
                 appends.append(duplex)
         if appends:
-            self._duplex_stage0_runtime().encode_appends(appends)
+            runtime = self._duplex_stage0_runtime()
+            runtime.encode_appends(appends)
+            runtime.prepare_live_appends(appends)
 
     def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
         runtime = getattr(self, "_personaplex_duplex_stage0_runtime", None)
@@ -551,9 +559,48 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module):
             )
         text_token = input_ids.reshape(bsz).to(torch.long)
         hidden = hidden_states.reshape(bsz, 1, -1).to(self._dtype)
-        audio_tokens: list[torch.Tensor] = []
-        audio_provided: list[torch.Tensor] = []
-        for info in req_infos:
+        runtime = self._duplex_stage0_runtime()
+        audio_tokens, audio_provided = self._depformer_teacher_forcing(runtime, req_ids, req_infos, hidden.device)
+        codes = self.depformer(
+            text_token,
+            hidden,
+            audio_tokens=audio_tokens,
+            audio_provided=audio_provided,
+            num_steps=self.num_active_codebooks,
+        ).to(torch.long)
+        # The next append reads the agent frame with the teacher-forced
+        # codebooks applied; merge it once for the whole batch.
+        effective = torch.where(audio_provided[:, :8], audio_tokens[:, :8], codes[:, :8])
+        for row, request_id in enumerate(req_ids):
+            runtime.record_sample(
+                request_id=request_id,
+                text_token=text_token[row],
+                effective_codes=effective[row],
+            )
+        return codes
+
+    @staticmethod
+    def _depformer_teacher_forcing(
+        runtime: PersonaPlexStage0DuplexRuntime,
+        req_ids: list[str],
+        req_infos: list[dict[str, Any]],
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Stack each request's depformer teacher-forcing tokens and mask on ``device``.
+
+        Live appends batched in ``preprocess_batch`` already hold device copies.
+        The rest (first appends, aborted epochs) carry host tensors in their
+        request information and are moved in one copy.
+        """
+        tokens_rows: list[torch.Tensor] = []
+        provided_rows: list[torch.Tensor] = []
+        host_rows: list[int] = []
+        for row, (request_id, info) in enumerate(zip(req_ids, req_infos)):
+            prepared = runtime.prepared_depformer_state(request_id)
+            if prepared is not None:
+                tokens_rows.append(prepared[0])
+                provided_rows.append(prepared[1])
+                continue
             tokens = info.get("pplex_depformer_audio_tokens")
             provided = info.get("pplex_depformer_audio_provided")
             if not isinstance(tokens, torch.Tensor) or not isinstance(provided, torch.Tensor):
@@ -565,29 +612,18 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module):
                     "PersonaPlex depformer teacher-forcing token/mask shapes differ: "
                     f"{tuple(tokens.shape)} != {tuple(provided.shape)}"
                 )
-            audio_tokens.append(tokens)
-            audio_provided.append(provided)
-        codes = self.depformer(
-            text_token,
-            hidden,
-            audio_tokens=torch.stack(audio_tokens).to(
-                device=hidden.device,
-                dtype=torch.long,
-            ),
-            audio_provided=torch.stack(audio_provided).to(
-                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):
-            runtime.record_sample(
-                request_id=request_id,
-                text_token=text_token[row],
-                agent_codes=codes[row],
-            )
-        return codes
+            host_rows.append(row)
+            tokens_rows.append(tokens)
+            provided_rows.append(provided)
+        if host_rows:
+            host_tokens = torch.stack([tokens_rows[row] for row in host_rows]).to(device=device, dtype=torch.long)
+            host_provided = torch.stack([provided_rows[row] for row in host_rows]).to(device=device, dtype=torch.bool)
+            if len(host_rows) == len(req_ids):
+                return host_tokens, host_provided
+            for index, row in enumerate(host_rows):
+                tokens_rows[row] = host_tokens[index]
+                provided_rows[row] = host_provided[index]
+        return torch.stack(tokens_rows), torch.stack(provided_rows)
 
     # ------------------------------------------------------------------
     # Weight loading