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

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

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