[Feature] PersonaPlex (Moshi-based full-duplex S2S): native vLLM port + duplex serving (#4771)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Signed-off-by: Sy0307 <24841857+Sy0307@users.noreply.github.com>
Co-authored-by: Sy0307 <24841857+Sy0307@users.noreply.github.com>

examples/offline_inference/personaplex/personaplex_offline.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline PersonaPlex inference through the vllm-omni fullduplex backend.
+
+This is the runnable end-to-end proof for the integration: it drives a user WAV
+through :class:`PersonaPlexEngine` + :class:`PersonaPlexSession` (our wrappers
+around the Moshi/Mimi lockstep loop) and writes the agent's audio + inner-monologue
+text. It mirrors ``python -m moshi.offline`` but goes through our ``FrameStepper``
+seam, so a green run validates that the session/adapter plumbing is faithful.
+
+Prereqs: HF access to ``nvidia/personaplex-7b-v1`` (``export HF_TOKEN=...``). The
+stack is moshi-free (native temporal + depformer + streaming Mimi via
+``transformers``); decoding is greedy.
+
+Example:
+    python examples/offline_inference/personaplex/personaplex_offline.py \\
+        --input-wav input_assistant.wav --output-wav out.wav \\
+        --voice-prompt NATF2.pt --persona "You are a wise and friendly teacher."
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+
+import numpy as np
+
+from vllm_omni.experimental.fullduplex.personaplex import (
+    PersonaPlexConfig,
+    PersonaPlexEngine,
+    PersonaPlexSession,
+)
+from vllm_omni.experimental.fullduplex.personaplex.config import DEFAULT_PERSONA
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--input-wav", required=True, help="User-side input WAV")
+    parser.add_argument("--output-wav", required=True, help="Agent audio output WAV")
+    parser.add_argument("--output-text", default=None, help="Optional JSON for agent text")
+    parser.add_argument("--voice-prompt", default="NATF2.pt", help="Voice clone (bundled .pt or .wav)")
+    parser.add_argument("--persona", default=DEFAULT_PERSONA, help="System role text")
+    parser.add_argument("--hf-repo", default="nvidia/personaplex-7b-v1")
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--cpu-offload", action="store_true")
+    args = parser.parse_args()
+
+    import sphn  # local import: only the runnable driver needs the audio I/O dep
+
+    config = PersonaPlexConfig(
+        hf_repo=args.hf_repo,
+        voice_prompt=args.voice_prompt,
+        persona=args.persona,
+        device=args.device,
+        cpu_offload=args.cpu_offload,
+    )
+
+    engine = PersonaPlexEngine(config).load()
+    session = PersonaPlexSession(engine, config)
+    session.open()
+
+    # Load exactly like the reference (moshi.lm.load_audio): native read then
+    # explicit resample, so greedy runs stay bit-faithful to `python -m moshi.offline`.
+    user_pcm, src_sr = sphn.read(args.input_wav)
+    user_pcm = sphn.resample(user_pcm, src_sample_rate=src_sr, dst_sample_rate=config.sample_rate)
+    audio, text = session.run(np.asarray(user_pcm[0], dtype=np.float32))
+
+    sphn.write_wav(args.output_wav, audio, config.sample_rate)
+    print(f"wrote {audio.shape[0] / config.sample_rate:.2f}s of agent audio to {args.output_wav}")
+    print(f"agent text: {text!r}")
+    if args.output_text:
+        with open(args.output_text, "w") as f:
+            json.dump({"text": text}, f, ensure_ascii=False)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/personaplex/README.md

@@ -0,0 +1,84 @@
+# PersonaPlex full-duplex online serving
+
+Serve [`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
+(a Moshi-based full-duplex speech-to-speech model) with the native vLLM-Omni engine.
+
+The server hosts the **official PersonaPlex web client** at `/` (auto-downloaded from
+the model repo) and implements its WebSocket protocol at `/api/chat`, so you talk to the
+model live in the browser — mic in, agent speech + inner-monologue text out.
+
+> Experimental. Requires a GPU and Hugging Face access to the gated repo
+> (`HF_TOKEN` with access to `nvidia/personaplex-7b-v1`).
+
+## Start the server
+
+```bash
+HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 bash run_server.sh
+# or directly:
+HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 python -m \
+    vllm_omni.experimental.fullduplex.personaplex.serving.server --port 8124
+```
+
+## Talk to it
+
+- **Browser (recommended):** open `http://localhost:8124/` and allow the microphone.
+  Use headphones so the agent does not hear itself.
+- **Headless:** stream a 24 kHz mono WAV and save the reply:
+
+  ```bash
+  python duplex_client.py --url ws://localhost:8124/api/chat --input user.wav --out reply.wav
+  ```
+
+## Unified Realtime API
+
+The default `vllm_omni/deploy/personaplex.yaml` also enables main's engine-owned
+full-duplex control plane:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
+  /path/to/personaplex-7b-v1 \
+  --omni \
+  --stage-configs-path vllm_omni/deploy/personaplex.yaml
+```
+
+Validate the actual `/v1/realtime?duplex=1` scheduler path with paced 24 kHz
+PCM, two concurrent sessions, overflow admission, per-session slot recycling,
+and non-silent output:
+
+```bash
+python tests/e2e/online_serving/personaplex_realtime_duplex.py \
+  --model /path/to/personaplex-7b-v1 \
+  --input-wav /path/to/speech.wav \
+  --output-dir /tmp/personaplex-realtime-duplex
+```
+
+This is distinct from the standalone `/api/chat` compatibility server above.
+The unified endpoint advertises `supports_barge_in=false`: overlapping speech
+is native model behavior, but destructive output interruption and model-state
+rewind have not been validated for PersonaPlex.
+
+## Protocol (`/api/chat`)
+
+Binary WebSocket messages, first byte is a tag (identical to `moshi.server`):
+
+| Direction | Tag | Payload |
+| --- | --- | --- |
+| server → client | `\x00` | ready handshake (after the system-prompt prefill) |
+| client → server | `\x01` | Opus-encoded mic audio (24 kHz) |
+| server → client | `\x01` | Opus-encoded agent audio |
+| server → client | `\x02` | UTF-8 inner-monologue text |
+
+Connect with query params `text_prompt` (persona) and `voice_prompt` (voice file).
+A raw-PCM endpoint (`/v1/audio/duplex`, JSON `open` + float32 frames) is also available
+for clients that do not want an Opus dependency.
+
+## Notes
+
+- **Run the client near the server.** Real-time 80 ms audio is sensitive to network
+  latency/jitter; over a high-latency remote link playback can stutter regardless of
+  engine speed. On localhost it is smooth.
+- Greedy decoding by default. One conversation per server by default (a second
+  concurrent connection is rejected with WS close 1013 until the first ends);
+  pass `--batch-size N` to host N concurrent conversations on one engine
+  (elastic batching; connections beyond N are rejected until a slot frees).
+- Full runbook: [`recipes/NVIDIA/PersonaPlex.md`](../../../recipes/NVIDIA/PersonaPlex.md).

examples/online_serving/personaplex/duplex_client.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Headless client for the PersonaPlex duplex server (official Moshi protocol).
+
+Streams a 24 kHz mono WAV to ``/api/chat`` exactly as the browser client does
+(Opus over WebSocket, tags ``\\x00`` ready / ``\\x01`` audio / ``\\x02`` text),
+collects the agent's reply audio + inner-monologue text, and writes a WAV.
+
+The browser UI at ``http://<host>:<port>/`` is the primary demo; this script is
+for headless testing / scripting.
+
+    python duplex_client.py --url ws://localhost:8124/api/chat \
+        --input user.wav --out reply.wav
+
+Deps: ``websockets``, ``sphn``, ``soundfile``, ``numpy``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+
+import numpy as np
+import soundfile as sf
+import sphn
+
+SR = 24000  # PersonaPlex / Mimi codec rate
+
+
+async def converse(url: str, pcm: np.ndarray) -> tuple[np.ndarray, str]:
+    import websockets
+
+    reader = sphn.OpusStreamReader(SR)  # decode agent audio
+    writer = sphn.OpusStreamWriter(SR)  # encode our mic audio
+    text: list[str] = []
+    ready = asyncio.Event()
+    last = {"t": 0.0}
+
+    async with websockets.connect(url, max_size=None) as ws:
+
+        async def receive() -> None:
+            loop = asyncio.get_event_loop()
+            while True:
+                try:
+                    msg = await ws.recv()
+                except Exception:
+                    return
+                if not isinstance(msg, (bytes, bytearray)) or not msg:
+                    continue
+                tag, payload = msg[0], msg[1:]
+                if tag == 0:
+                    ready.set()
+                elif tag == 1:
+                    reader.append_bytes(bytes(payload))
+                    last["t"] = loop.time()
+                elif tag == 2:
+                    text.append(payload.decode("utf8", errors="replace"))
+
+        recv_task = asyncio.create_task(receive())
+        await asyncio.wait_for(ready.wait(), timeout=120)  # waits for system-prompt prefill
+
+        frame = 1920  # 80 ms
+        for i in range(0, len(pcm), frame):
+            chunk = pcm[i : i + frame]
+            if len(chunk) < frame:
+                chunk = np.concatenate([chunk, np.zeros(frame - len(chunk), dtype=np.float32)])
+            writer.append_pcm(chunk)
+            data = writer.read_bytes()
+            if data:
+                await ws.send(b"\x01" + data)
+            await asyncio.sleep(frame / SR)  # pace at real time
+
+        loop = asyncio.get_event_loop()
+        last["t"] = loop.time()
+        while loop.time() - last["t"] < 2.5:  # drain until the agent stops talking
+            await asyncio.sleep(0.1)
+        recv_task.cancel()
+
+    chunks: list[np.ndarray] = []
+    while True:
+        out = reader.read_pcm()
+        if out is None or out.shape[-1] == 0:
+            break
+        chunks.append(np.asarray(out, dtype=np.float32))
+    audio = np.concatenate(chunks) if chunks else np.zeros(0, dtype=np.float32)
+    return audio, "".join(text)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--url", default="ws://localhost:8124/api/chat")
+    ap.add_argument("--input", required=True, help="24 kHz mono WAV of the user turn")
+    ap.add_argument("--out", default="personaplex_reply.wav")
+    ap.add_argument("--tail", type=float, default=10.0, help="trailing silence (s) for the agent reply")
+    ap.add_argument("--persona", default="")
+    ap.add_argument("--voice", default="NATF2.pt")
+    args = ap.parse_args()
+
+    wav, sr = sf.read(args.input, dtype="float32")
+    if wav.ndim > 1:
+        wav = wav.mean(axis=1)
+    if sr != SR:
+        raise SystemExit(f"input must be {SR} Hz mono; got {sr} Hz. Resample it first.")
+    wav = np.concatenate([wav, np.zeros(int(args.tail * SR), dtype=np.float32)])
+
+    url = f"{args.url}?text_prompt={args.persona}&voice_prompt={args.voice}"
+    audio, text = asyncio.run(converse(url, wav))
+    sf.write(args.out, audio if audio.size else np.zeros(SR, dtype=np.float32), SR)
+    print(f"agent reply: {audio.size / SR:.2f}s -> {args.out}")
+    print(f"inner-monologue: {text.strip()!r}")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/personaplex/run_server.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# PersonaPlex full-duplex online serving.
+#
+# Serves the OFFICIAL PersonaPlex web client at / (auto-downloaded from
+# nvidia/personaplex-7b-v1) and the Moshi WebSocket protocol at /api/chat,
+# backed by the native vLLM-Omni engine. Requires HF access to the gated repo.
+
+HOST="${HOST:-0.0.0.0}"
+PORT="${PORT:-8124}"
+
+echo "Starting PersonaPlex duplex server on ${HOST}:${PORT}..."
+echo "Open http://localhost:${PORT}/ for the official web client (run it near the server)."
+
+python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
+    --host "$HOST" --port "$PORT"

recipes/NVIDIA/PersonaPlex.md

@@ -0,0 +1,174 @@
+# PersonaPlex
+
+> Live full-duplex speech-to-speech (talk to the model in your browser, it listens
+> while it speaks) with persona and voice control, on the native vLLM-Omni engine.
+
+## Summary
+
+- Vendor: NVIDIA
+- Model: [`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
+  (gated; a Moshi finetune)
+- Task: Full-duplex speech-to-speech. 24 kHz mic audio in, agent speech +
+  inner-monologue text out, in 80 ms lockstep. Persona (role text) and voice
+  (zero-shot clone) are set per session.
+- Mode: Live duplex WebSocket server with the official browser client, single or
+  multi-session (`--batch-size`), plus an offline WAV-in / WAV-out example
+- Maintainer: [`@linyueqian`](https://github.com/linyueqian)
+
+## When to use this recipe
+
+Use this recipe to run `nvidia/personaplex-7b-v1` as a live voice agent on one GPU.
+The integration is **moshi-free**: the temporal transformer, depformer, and
+streaming Mimi codec are all native vLLM-Omni modules (Mimi runs on
+`transformers.MimiModel`), so no vendored fork is installed. Decoding is greedy,
+matching the reference implementation frame for frame on the golden replays used
+as the acceptance gate.
+
+PersonaPlex is the first Moshi-class (pure-lockstep) model in the experimental
+full-duplex framework (`vllm_omni/experimental/fullduplex/`), and the first with
+elastic multi-session batching: one engine hosts `--batch-size` concurrent
+conversations, callers join or leave on any 80 ms tick without disturbing live
+slots (greedy isolation is bit-exact).
+
+## References
+
+- Serving package:
+  [`vllm_omni/experimental/fullduplex/personaplex/`](../../vllm_omni/experimental/fullduplex/personaplex/)
+- Native model modules (temporal / depformer / streaming Mimi):
+  [`vllm_omni/model_executor/models/personaplex/`](../../vllm_omni/model_executor/models/personaplex/)
+- Online example (server + headless client):
+  [`examples/online_serving/personaplex/`](../../examples/online_serving/personaplex/)
+- Offline example:
+  [`examples/offline_inference/personaplex/personaplex_offline.py`](../../examples/offline_inference/personaplex/personaplex_offline.py)
+- Staged offline pipeline (talker -> Mimi code2wav, async-chunk streaming):
+  [`vllm_omni/deploy/personaplex.yaml`](../../vllm_omni/deploy/personaplex.yaml)
+- Integration PR: [vllm-project/vllm-omni#4771](https://github.com/vllm-project/vllm-omni/pull/4771)
+- Framework context: [RFC #3745](https://github.com/vllm-project/vllm-omni/issues/3745)
+  (duplex adapter patterns), [#1335](https://github.com/vllm-project/vllm-omni/issues/1335)
+  (full-duplex target), [PR #3907](https://github.com/vllm-project/vllm-omni/pull/3907)
+  (fullduplex core contracts this adapter conforms to)
+- Upstream model card:
+  [`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
+
+## Why it is different from MiniCPM-o / JoyVL duplex
+
+| | MiniCPM-o 4.5 (#3907) | JoyVL | **PersonaPlex (this)** |
+|---|---|---|---|
+| Cadence | 1 s chunk groups | ~1 fps frames | **80 ms lockstep (12.5 Hz)** |
+| Turn control | learned `⟨listen⟩`/`⟨speak⟩` | `</silence>`/`</response>` | **none, pure lockstep** |
+| Per step | variable-length token group | text decision | **1 user frame in -> 1 agent frame + 1 text token** |
+| Barge-in | at chunk boundary | n/a | **native (model always hears the user)** |
+| Session state | chunk-group KV | per-tick HTTP | **persistent ring KV + streaming Mimi state** |
+
+This is the lockstep ("parallel-frame joint") shape of the duplex adapter
+patterns: the adapter declares `DuplexCapability.continuous = True`, and
+`core.DuplexRuntime` runs ONE eternal response that consumes input frames as they
+arrive and drains on close. The flag is a small, default-off lifecycle mode, so
+turn-based adapters (JoyVL, MiniCPM-o) are unaffected.
+
+## Architecture (Moshi RQ-Transformer, all native)
+
+- **Mimi codec** 24 kHz, 12.5 Hz, 1920 samples/frame, 8 active codebooks
+  (card 2048). Streaming encode + decode via `transformers.MimiModel` with the
+  PersonaPlex reference checkpoint weights.
+- **Temporal transformer** (Helium backbone) 4096-d / 32 layers / 32 heads,
+  sliding window 3000 frames (a 4-minute rolling context) held in a ring KV cache
+  with per-slot recycle.
+- **Depformer** 1024-d / 6 layers / 16 heads, autoregressive over 16 codebooks
+  per frame (8 vocoded), reset each frame.
+- Per-frame token column `[B, 17, 1]`: row 0 = inner-monologue text,
+  rows 1-8 = agent audio, rows 9-16 = user audio.
+- Persona + voice are injected once at session open through the same lockstep
+  step: voice clone forces the agent stream from reference-audio Mimi codes;
+  persona forces the text stream from `<system> ... <system>` tokens.
+
+## 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.
+
+## GPU
+
+### 1x 141 GB Hopper-class GPU (H20-class, verified)
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+ (CI covers 3.11 / 3.12)
+- vLLM-Omni: PR #4771 branch or current `main` once merged
+- Extra dependency for the live server: `pip install sphn` (Opus framing;
+  `aiohttp` already ships with vLLM)
+- `export HF_TOKEN=...` with access to the gated
+  `nvidia/personaplex-7b-v1` repo (accept the license on the model page)
+
+First run auto-downloads from the model repo: `model.safetensors`,
+`tokenizer_spm_32k_3.model`, `voices.tgz`, and the Mimi reference checkpoint;
+`kyutai/mimi` is fetched via `transformers` for the codec module graph.
+
+#### Command
+
+Live duplex server (official browser client served at `/`):
+
+```bash
+HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 \
+python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
+    --port 8124 --voice NATF2.pt
+```
+
+Multi-session on the same GPU (elastic batching, callers join/leave any tick):
+
+```bash
+HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 \
+python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
+    --port 8124 --batch-size 4
+```
+
+Then open `http://localhost:8124/` and allow the microphone (use headphones so
+the agent does not hear itself). Headless alternative:
+
+```bash
+python examples/online_serving/personaplex/duplex_client.py \
+    --url ws://localhost:8124/api/chat --input user.wav --out reply.wav
+```
+
+Offline WAV-in / WAV-out (no server):
+
+```bash
+python examples/offline_inference/personaplex/personaplex_offline.py \
+    --input-wav question.wav --output-wav reply.wav --output-text reply.json \
+    --voice-prompt NATF2.pt --persona "You are a wise and friendly teacher."
+```
+
+#### Verification
+
+```bash
+# GPU-free contract tests (stubbed FrameStepper)
+pytest tests/e2e/features/fullduplex/test_personaplex_adapter.py \
+       tests/e2e/features/fullduplex/test_personaplex_session.py \
+       tests/e2e/features/fullduplex/test_personaplex_batched.py -q
+
+# GPU e2e: batched isolation + slot recycle (coherence check needs a real
+# question WAV via PPLEX_QUESTION_WAV)
+pytest tests/e2e/features/fullduplex/test_personaplex_elastic.py -q
+```
+
+A green offline run producing a coherent spoken reply (and inner-monologue text
+in `reply.json`) is the quickest end-to-end check.
+
+#### Notes
+
+- Realtime budget is 80 ms/frame. Verified eager per-tick latency at
+  `--batch-size 4` is ~70-74 ms on this hardware class, i.e. all four
+  conversations stay realtime; single-session has comfortable headroom.
+- 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.
+- With `--batch-size N`, connections beyond N are rejected until a slot frees;
+  a recycled slot is bit-exact with a fresh engine for the same inputs.
+- Voice prompts: bundled `voices.tgz` names (`NATF*`/`NATM*` natural,
+  `VARF*`/`VARM*` varied) or a path to your own `.pt`/`.wav`.
+- Run the browser client near the server. The 80 ms cadence is sensitive to
+  network jitter; over a high-latency link playback can stutter regardless of
+  engine speed. On localhost it is smooth.
+- Sliding window is 3000 frames (~4 min). The ring KV recycles beyond that;
+  very long sessions keep running with a rolling context.

tests/config/test_omni_config.py

@@ -81,6 +81,21 @@ def _from_pipeline_key(
     )
 
 
+def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
+    omni_config = _from_pipeline_key("personaplex")
+
+    assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2]
+
+
+def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None:
+    deploy_path = tmp_path / "personaplex-turn.yaml"
+    deploy_path.write_text("session_mode: turn\n", encoding="utf-8")
+
+    omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path))
+
+    assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1]
+
+
 @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
 def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str):
     pipeline = _resolve_pipeline_or_skip(model_type)
@@ -105,6 +120,7 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
         assert omni_stage.runtime_config.num_replicas == legacy_stage.yaml_runtime.get("num_replicas", 1)
 
         engine_args = legacy_stage.yaml_engine_args
+        assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1)
         assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
         assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
         assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
@@ -379,6 +395,7 @@ def test_runtime_config_fields_match_rfc_runtime_scope():
 def test_sub_config_fields_match_rfc_scopes():
     assert {f.name for f in fields(OmniStageModelConfig)} == {
         "active_stream_window",
+        "duplex_max_sessions",
         "enable_sleep_mode",
         "default_sampling_params",
         "subtalker_sampling_params",

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -110,6 +110,9 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
     session.status = RequestStatus.RUNNING
     session.resumable = True
     session.num_computed_tokens = len(session.prompt_token_ids)
+    # schedule() has counted this step's token as in flight; update_from_output()
+    # must settle it before handling the segment boundary.
+    session.num_in_flight_tokens = 1
 
     sched = MagicMock()
     sched.requests = {session.request_id: session}
@@ -118,12 +121,20 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
         is_done_receiving_chunks=lambda _request_id: True,
         segment_finished_requests={session.request_id},
     )
-    sched._handle_stopped_request.return_value = False
+
+    def rearm_resumable_request(request: Request) -> bool:
+        request.status = RequestStatus.WAITING
+        return False
+
+    sched._handle_stopped_request.side_effect = rearm_resumable_request
     sched.running = [session]
     sched.waiting = MagicMock()
     sched.skipped_waiting = MagicMock()
     sched.structured_output_manager.should_advance.return_value = False
-    sched._pending_finish_reqs = []
+    # Async scheduling can observe the segment boundary in the next schedule()
+    # before this model output is applied and defer the same request for finish.
+    # update_from_output() must not re-arm the resumable request twice.
+    sched._pending_finish_reqs = [session]
     sched.recompute_kv_load_failures = False
     sched.connector = None
     sched.kv_cache_manager.take_events.return_value = None
@@ -155,6 +166,8 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
     output = outputs[session.client_index].outputs[0]
     assert output.finish_reason is not None
     assert output.is_segment_finished is True
+    sched._handle_stopped_request.assert_called_once_with(session)
+    assert sched._pending_finish_reqs == []
 
 
 def test_async_chunk_resumable_stop_rearms_connector_polling() -> None:

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -390,6 +390,105 @@ def test_send_single_request_respects_processor_receiver_boundary(build_adapter,
     assert sent_payload.meta.is_segment_finished.item() is False
 
 
+def test_send_single_request_personaplex_pending_frame_is_not_segment_boundary(
+    build_adapter,
+):
+    from vllm_omni.model_executor.stage_input_processors.personaplex import (
+        talker2code2wav_async_chunk,
+    )
+
+    adapter, connector = build_adapter(
+        stage_id=0,
+        connector_extra={
+            "initial_codec_chunk_frames": 1,
+            "codec_chunk_frames": 5,
+        },
+    )
+    request = _req(
+        "req-personaplex",
+        RequestStatus.WAITING,
+        external_req_id="ext-personaplex",
+    )
+    request.resumable = True
+    request.additional_information = {
+        "codes": {
+            "audio": torch.arange(8, dtype=torch.long).reshape(1, 8),
+        }
+    }
+    adapter.custom_process_next_stage_input_func = talker2code2wav_async_chunk
+
+    adapter._send_single_request(
+        {
+            "multimodal_output": None,
+            "request": request,
+            "is_finished": False,
+            "is_segment_finished": True,
+        }
+    )
+
+    sent_payload = connector.put.call_args.kwargs["data"]
+    assert sent_payload.codes is None
+    assert sent_payload.meta.finished.item() is False
+    assert sent_payload.meta.is_segment_finished.item() is False
+
+
+def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter):
+    from vllm_omni.model_executor.stage_input_processors.personaplex import (
+        talker2code2wav_async_chunk,
+    )
+
+    adapter, _ = build_adapter(
+        stage_id=0,
+        connector_extra={
+            "initial_codec_chunk_frames": 1,
+            "codec_chunk_frames": 5,
+        },
+    )
+    request = _req(
+        "req-personaplex-first",
+        RequestStatus.WAITING,
+        external_req_id="ext-personaplex-reused",
+    )
+    request.resumable = True
+    request.additional_information = {
+        "codes": {
+            "audio": torch.arange(8, dtype=torch.long).reshape(1, 8),
+        }
+    }
+
+    first = talker2code2wav_async_chunk(
+        adapter,
+        multimodal_output=None,
+        request=request,
+        is_finished=True,
+    )
+    assert first is not None
+    assert first.codes is None
+
+    adapter.cleanup_sender(request.external_req_id)
+
+    replacement = _req(
+        "req-personaplex-replacement",
+        RequestStatus.WAITING,
+        external_req_id=request.external_req_id,
+    )
+    replacement.resumable = True
+    replacement.additional_information = {
+        "codes": {
+            "audio": torch.arange(8, 16, dtype=torch.long).reshape(1, 8),
+        }
+    }
+    second = talker2code2wav_async_chunk(
+        adapter,
+        multimodal_output=None,
+        request=replacement,
+        is_finished=True,
+    )
+
+    assert second is not None
+    assert second.codes is None
+
+
 def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
     adapter, _ = build_adapter(stage_id=1)
     request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")

tests/e2e/features/fullduplex/personaplex/test_import_boundary.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import subprocess
+import sys
+
+
+def test_batched_runtime_import_does_not_require_optional_opus_dependency() -> None:
+    script = """
+import builtins
+
+real_import = builtins.__import__
+
+def guarded_import(name, *args, **kwargs):
+    if name == "sphn":
+        raise ModuleNotFoundError("blocked optional dependency: sphn")
+    return real_import(name, *args, **kwargs)
+
+builtins.__import__ = guarded_import
+from vllm_omni.experimental.fullduplex.personaplex.serving.batched import BatchedSessionManager
+assert BatchedSessionManager is not None
+"""
+    result = subprocess.run(
+        [sys.executable, "-c", script],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+
+    assert result.returncode == 0, result.stderr

tests/e2e/features/fullduplex/personaplex/test_stage0_runtime.py

@@ -0,0 +1,303 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import base64
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.experimental.fullduplex.personaplex.policy import (
+    SILENCE_TOKENS,
+    SINE_TOKENS,
+)
+from vllm_omni.experimental.fullduplex.personaplex.stage0 import (
+    PersonaPlexStage0DuplexRuntime,
+)
+from vllm_omni.model_executor.models.personaplex.personaplex_talker import (
+    PersonaPlexTalkerForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeCodec:
+    def __init__(self) -> None:
+        self.encode_calls = 0
+        self.reset_calls = 0
+
+    def streaming_init(self, batch_size: int) -> None:
+        assert batch_size == 1
+
+    def encode_frame(self, pcm: torch.Tensor) -> torch.Tensor:
+        assert pcm.shape == (1, 1920)
+        self.encode_calls += 1
+        return torch.full((1, 8), self.encode_calls, dtype=torch.long)
+
+    def reset_streaming(self) -> None:
+        self.reset_calls += 1
+
+
+class _FakeTalker:
+    def __init__(self) -> None:
+        self.device = torch.device("cpu")
+        self.dtype = torch.float32
+        self.frame_calls: list[dict[str, torch.Tensor | None]] = []
+
+    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, 4).contiguous()
+
+    def _build_frame_embed(
+        self,
+        text_token,
+        last_agent,
+        prev_agent,
+        device,
+        user_d0=None,
+        user_d1=None,
+    ):
+        self.frame_calls.append(
+            {
+                "text_token": text_token.clone(),
+                "last_agent": None if last_agent is None else last_agent.clone(),
+                "prev_agent": None if prev_agent is None else prev_agent.clone(),
+                "user_d0": None if user_d0 is None else user_d0.clone(),
+                "user_d1": None if user_d1 is None else user_d1.clone(),
+            }
+        )
+        value = float(user_d0[0].item()) if user_d0 is not None else -1.0
+        return torch.full((1, 4), value, device=device)
+
+
+def _duplex_info(*, seq: int, session_id: str = "session", incarnation: int = 1):
+    pcm = np.zeros(1920, dtype="<f4")
+    return {
+        "data_plane": True,
+        "session_id": session_id,
+        "incarnation": incarnation,
+        "epoch": 0,
+        "seq": seq,
+        "payload": {
+            "format": "pcm_f32le",
+            "sample_rate_hz": 24000,
+            "audio": base64.b64encode(pcm.tobytes()).decode("ascii"),
+        },
+        "runtime_config": {
+            "personaplex_model_path": "/unused",
+            "personaplex_voice_prompt": "NATF2.pt",
+            "personaplex_persona": "Be concise.",
+        },
+    }
+
+
+def _runtime(*codecs: _FakeCodec) -> PersonaPlexStage0DuplexRuntime:
+    voice_embeddings = torch.arange(8, dtype=torch.float32).reshape(2, 1, 1, 4)
+    codec_args = {"codec": codecs[0]}
+    if len(codecs) > 1:
+        available = iter(codecs)
+        codec_args = {
+            "codec_factory": lambda: next(available),
+            "max_sessions": len(codecs),
+        }
+    return PersonaPlexStage0DuplexRuntime(
+        _FakeTalker(),
+        model_path="/unused",
+        device="cpu",
+        **codec_args,
+        tokenizer=lambda _text: [7, 8, 9],
+        voice_loader=lambda _voice: {
+            "embeddings": voice_embeddings,
+            "cache": torch.zeros((1, 17, 4), dtype=torch.long),
+        },
+    )
+
+
+def test_first_append_prepends_voice_and_persona_once() -> None:
+    codec = _FakeCodec()
+    runtime = _runtime(codec)
+
+    first = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
+    second = runtime.prepare_append(_duplex_info(seq=2), prompt_len=18)
+
+    assert first.prefill_applied is True
+    assert first.prompt_offset == 0
+    assert first.user_codes.shape == (1, 8)
+    assert first.inputs_embeds.shape == (18, 4)
+    assert second.prefill_applied is False
+    assert second.prompt_offset == 17
+    assert second.user_codes.shape == (2, 8)
+    assert second.inputs_embeds.shape == (1, 4)
+    assert codec.encode_calls == 2
+
+
+@pytest.mark.parametrize(
+    ("append_count", "expected_tokens", "expected_provided"),
+    [
+        (
+            1,
+            [*SILENCE_TOKENS, 1, *SINE_TOKENS[1:]],
+            [False, *([True] * 15)],
+        ),
+        (
+            2,
+            [*SILENCE_TOKENS, 2, *([1] * 7)],
+            [*([False] * 8), *([True] * 8)],
+        ),
+    ],
+)
+def test_prepare_append_preserves_native_depformer_teacher_forcing(
+    append_count: int,
+    expected_tokens: list[int],
+    expected_provided: list[bool],
+) -> None:
+    runtime = _runtime(_FakeCodec())
+
+    result = None
+    for seq in range(1, append_count + 1):
+        result = runtime.prepare_append(_duplex_info(seq=seq), prompt_len=18)
+
+    assert result is not None
+    assert result.info_update["pplex_depformer_audio_tokens"].tolist() == expected_tokens
+    assert result.info_update["pplex_depformer_audio_provided"].tolist() == expected_provided
+
+
+def test_frame_embed_uses_previous_effective_agent_frame() -> None:
+    fake_talker = SimpleNamespace(
+        config=SimpleNamespace(
+            num_audio_codebooks=16,
+            audio_vocab_size=2048,
+        ),
+        input_embeddings=lambda stack: stack,
+    )
+    last_agent = torch.arange(10, 18)
+
+    embeds = PersonaPlexTalkerForConditionalGeneration._build_frame_embed(
+        fake_talker,
+        torch.tensor([3]),
+        last_agent,
+        torch.arange(8),
+        torch.device("cpu"),
+    )
+
+    assert embeds.shape == (1, 17)
+    assert torch.equal(embeds[0, 1:9], last_agent)
+
+
+def test_repeated_append_identity_does_not_advance_codec() -> None:
+    codec = _FakeCodec()
+    runtime = _runtime(codec)
+
+    info = _duplex_info(seq=1)
+    first = runtime.prepare_append(info, prompt_len=18, request_id="req")
+    retry = runtime.prepare_append(info, prompt_len=18, request_id="req")
+
+    assert codec.encode_calls == 1
+    assert torch.equal(first.inputs_embeds, retry.inputs_embeds)
+    assert torch.equal(first.user_codes, retry.user_codes)
+
+
+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)
+    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))
+    assert torch.equal(second_call["user_d1"], sine)
+
+    third_call = runtime.stage_model.frame_calls[2]
+    assert third_call["text_token"].tolist() == [102]
+    assert torch.equal(third_call["last_agent"], second_agent)
+    assert torch.equal(third_call["prev_agent"], second_agent)
+    assert torch.equal(third_call["user_d0"], torch.full((8,), 2, dtype=torch.long))
+    assert torch.equal(third_call["user_d1"], torch.full((8,), 1, dtype=torch.long))
+
+
+def test_decoded_pcm_is_writable_for_torch_zero_copy() -> None:
+    pcm = PersonaPlexStage0DuplexRuntime._decode_pcm(_duplex_info(seq=1)["payload"])
+
+    assert pcm.flags.writeable
+
+
+def _prepare_two_sessions(
+    runtime: PersonaPlexStage0DuplexRuntime,
+) -> tuple[object, object, object, object]:
+    first_1 = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
+    second_1 = runtime.prepare_append(
+        _duplex_info(seq=1, session_id="other", incarnation=2),
+        prompt_len=18,
+    )
+    first_2 = runtime.prepare_append(_duplex_info(seq=2), prompt_len=19)
+    second_2 = runtime.prepare_append(
+        _duplex_info(seq=2, session_id="other", incarnation=2),
+        prompt_len=19,
+    )
+    return first_1, second_1, first_2, second_2
+
+
+def test_live_sessions_keep_independent_streaming_encoders() -> None:
+    first_codec = _FakeCodec()
+    second_codec = _FakeCodec()
+    runtime = _runtime(first_codec, second_codec)
+
+    first_1, second_1, first_2, second_2 = _prepare_two_sessions(runtime)
+
+    assert first_1.user_codes[:, 0].tolist() == [1]
+    assert second_1.user_codes[:, 0].tolist() == [1]
+    assert first_2.user_codes[:, 0].tolist() == [1, 2]
+    assert second_2.user_codes[:, 0].tolist() == [1, 2]
+
+
+def test_stage0_session_capacity_fails_before_codec_state_is_shared() -> None:
+    runtime = _runtime(_FakeCodec(), _FakeCodec())
+    _prepare_two_sessions(runtime)
+
+    with pytest.raises(RuntimeError, match="capacity 2"):
+        runtime.prepare_append(
+            _duplex_info(seq=1, session_id="third", incarnation=3),
+            prompt_len=18,
+        )
+
+
+def test_close_session_resets_and_reuses_released_codec() -> None:
+    first_codec = _FakeCodec()
+    second_codec = _FakeCodec()
+    runtime = _runtime(first_codec, second_codec)
+    _prepare_two_sessions(runtime)
+
+    runtime.close_session("session", 1)
+
+    assert first_codec.reset_calls == 1
+    assert second_codec.reset_calls == 0
+    replacement = runtime.prepare_append(
+        _duplex_info(seq=1, session_id="replacement", incarnation=4),
+        prompt_len=18,
+    )
+    assert replacement.user_codes[:, 0].tolist() == [3]

tests/e2e/features/fullduplex/personaplex/test_streaming_code2wav.py

@@ -0,0 +1,316 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.personaplex import (
+    personaplex_code2wav,
+    personaplex_mimi,
+)
+from vllm_omni.model_executor.models.personaplex.personaplex_code2wav import (
+    PersonaPlexCode2Wav,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeStreamingMimi(nn.Module):
+    def __init__(self, stream_value: int = 0) -> None:
+        super().__init__()
+        self.stream_value = stream_value
+        self.decode_frame_calls = 0
+        self.reset_calls = 0
+
+    def decode(self, codes: torch.Tensor, return_dict: bool = True):
+        assert return_dict
+        frames = int(codes.shape[-1])
+        return SimpleNamespace(audio_values=torch.arange(frames * 4, dtype=torch.float32).reshape(1, 1, -1))
+
+    def streaming_init(self, batch_size: int) -> None:
+        assert batch_size == 1
+
+    def decode_frame(self, codes: torch.Tensor) -> torch.Tensor:
+        assert codes.shape == (1, 2)
+        self.decode_frame_calls += 1
+        value = self.stream_value + self.decode_frame_calls
+        return torch.full((1, 4), float(value), dtype=torch.float32)
+
+    def reset_streaming(self) -> None:
+        self.reset_calls += 1
+
+
+class _FakeChunkedMimi(_FakeStreamingMimi):
+    def __init__(self) -> None:
+        super().__init__()
+        self.decode_frames_calls = 0
+        self.decode_frames_shapes: list[tuple[int, ...]] = []
+
+    def decode_frames(self, codes: torch.Tensor) -> torch.Tensor:
+        self.decode_frames_calls += 1
+        self.decode_frames_shapes.append(tuple(codes.shape))
+        return torch.arange(codes.shape[-1] * 4, dtype=torch.float32).reshape(1, -1)
+
+
+class _FakeMimiModel(nn.Module):
+    def __init__(self, _config) -> None:
+        super().__init__()
+        self.encoder_transformer = nn.Linear(1, 1, bias=False)
+        self.decoder_transformer = nn.Linear(1, 1, bias=False)
+        self.quantizer = nn.Linear(1, 1, bias=False)
+        self.encoder = SimpleNamespace(layers=[])
+        self.downsample = SimpleNamespace(conv=nn.Conv1d(1, 1, 1))
+        self.upsample = SimpleNamespace(conv=nn.ConvTranspose1d(1, 1, 1))
+        self.decoder = SimpleNamespace(layers=[])
+
+    def load_state_dict(self, _state_dict, strict: bool = True):
+        assert not strict
+        return SimpleNamespace(
+            missing_keys=[
+                "encoder_transformer.weight",
+                "decoder_transformer.weight",
+            ],
+            unexpected_keys=[],
+        )
+
+
+class _FakeMimiTransformer(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(1))
+
+    def load_weights(self, _state_dict, _prefix: str) -> int:
+        return 80
+
+
+def test_moshi_checkpoint_mapping_drops_replaced_transformer_weights(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    mapper = getattr(personaplex_mimi, "_map_moshi_codec_weights", None)
+    assert callable(mapper), "PersonaPlex Mimi must load its bundled codec weights without Hugging Face"
+
+    source = {
+        "encoder.model.3.conv.conv.weight": torch.ones(1),
+        "downsample.conv.conv.conv.weight": torch.ones(1),
+        "quantizer.rvq_first.vq.layers.0._codebook.embedding_sum": torch.ones(1),
+        "quantizer.rvq_rest.vq.layers.3._codebook._initialized": torch.ones(1),
+        "encoder_transformer.transformer.layers.0.norm1.weight": torch.ones(1),
+    }
+
+    mapped = mapper(source)
+
+    assert set(mapped) == {
+        "encoder.layers.3.conv.weight",
+        "downsample.conv.weight",
+        "quantizer.semantic_residual_vector_quantizer.layers.0.codebook.embed_sum",
+        "quantizer.acoustic_residual_vector_quantizer.layers.3.codebook.initialized",
+    }
+    assert mapped["encoder.layers.3.conv.weight"] is source["encoder.model.3.conv.conv.weight"]
+    monkeypatch.setattr("transformers.MimiConfig", lambda: object())
+    monkeypatch.setattr("transformers.MimiModel", _FakeMimiModel)
+    monkeypatch.setattr("safetensors.torch.load_file", lambda *_args, **_kwargs: {})
+    monkeypatch.setattr(personaplex_mimi, "_MimiStreamingTransformer", _FakeMimiTransformer)
+
+    codec = personaplex_mimi.PersonaPlexMimiCodec(
+        checkpoint="/unused",
+        device="cpu",
+    )
+
+    assert not hasattr(codec.model, "encoder_transformer")
+    assert not hasattr(codec.model, "decoder_transformer")
+    assert not any(
+        name.startswith(("model.encoder_transformer.", "model.decoder_transformer."))
+        for name, _ in codec.named_parameters()
+    )
+
+
+def test_mimi_full_stream_reset_reuses_all_state_storage_and_clears_offsets() -> None:
+    codec = personaplex_mimi.PersonaPlexMimiCodec.__new__(personaplex_mimi.PersonaPlexMimiCodec)
+    nn.Module.__init__(codec)
+    codec.device = torch.device("cpu")
+    codec.dtype = torch.float32
+    codec._enc_stages = [("conv", personaplex_mimi._StreamConv1d(nn.Conv1d(1, 1, 3)))]
+    codec._dec_stages = []
+    codec._downsample = personaplex_mimi._StreamConv1d(nn.Conv1d(1, 1, 3))
+    codec._upsample = personaplex_mimi._StreamConvTr1d(nn.ConvTranspose1d(1, 1, 3))
+    codec.encoder_transformer = personaplex_mimi._MimiStreamingTransformer(
+        num_layers=1,
+        dim=4,
+        num_heads=1,
+        context=3,
+    )
+    codec.decoder_transformer = personaplex_mimi._MimiStreamingTransformer(
+        num_layers=1,
+        dim=4,
+        num_heads=1,
+        context=3,
+    )
+    codec.streaming_init(batch_size=2)
+
+    conv_states = list(codec._conv_states())
+    conv_buffers = [
+        state.prev if isinstance(state, personaplex_mimi._StreamConv1d) else state.partial for state in conv_states
+    ]
+    transformers = [codec.encoder_transformer, codec.decoder_transformer]
+    kv_states = [kv for transformer in transformers for kv in transformer._kv]
+    state_tensors = [
+        *conv_buffers,
+        *(kv.cache for kv in kv_states),
+        *(kv.end_offset for kv in kv_states),
+        *(kv.start_offset for kv in kv_states),
+        *(transformer._offset for transformer in transformers),
+    ]
+    for tensor in state_tensors:
+        tensor.fill_(1)
+    for state in conv_states:
+        state._fresh.zero_()
+    storage = [tensor.data_ptr() for tensor in state_tensors]
+
+    codec.reset_streaming()
+
+    assert storage == [tensor.data_ptr() for tensor in state_tensors]
+    cleared = [
+        *conv_buffers,
+        *(kv.end_offset for kv in kv_states),
+        *(kv.start_offset for kv in kv_states),
+        *(transformer._offset for transformer in transformers),
+    ]
+    assert all(not tensor.any() for tensor in cleared)
+    assert all(state._fresh.all() for state in conv_states)
+
+
+def _model(*, max_sessions: int = 1) -> tuple[PersonaPlexCode2Wav, _FakeStreamingMimi]:
+    mimi_config = SimpleNamespace(num_codebooks=2, sample_rate=24000, samples_per_frame=4, mimi_name=None)
+    config = SimpleNamespace(mimi_config=mimi_config, mimi_name=None)
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            model="/unused",
+            hf_config=config,
+            duplex_max_sessions=max_sessions,
+        ),
+        device_config=SimpleNamespace(device="cpu"),
+    )
+    model = PersonaPlexCode2Wav(vllm_config=vllm_config)
+    mimi = _FakeStreamingMimi()
+    model.mimi = mimi
+    model._mimi_device = torch.device("cpu")
+    return model, mimi
+
+
+def _codes(frames: int, *, start: int = 0) -> torch.Tensor:
+    return torch.stack([torch.arange(start + offset, start + offset + frames) for offset in (0, 100)]).reshape(-1)
+
+
+def _audio(output) -> torch.Tensor:
+    return output.multimodal_outputs["model_outputs"][0]
+
+
+@pytest.mark.parametrize("second_codes", [_codes(3), _codes(1, start=100)])
+def test_resumable_codes_emit_only_new_pcm(second_codes: torch.Tensor) -> None:
+    model, mimi = _model()
+
+    first = model(input_ids=_codes(2), request_ids=["req"])
+    second = model(input_ids=second_codes, request_ids=["req"])
+
+    assert _audio(first).numel() == 8
+    assert _audio(second).numel() == 4
+    assert mimi.decode_frame_calls == 3
+
+
+@pytest.mark.parametrize(
+    ("frames", "expected_shapes"),
+    [
+        (5, [(1, 2, 5)]),
+        (12, [(1, 2, 5), (1, 2, 5), (1, 2, 2)]),
+    ],
+)
+def test_streaming_decode_uses_bounded_multi_frame_calls(
+    frames: int,
+    expected_shapes: list[tuple[int, ...]],
+) -> None:
+    model, _ = _model()
+    mimi = _FakeChunkedMimi()
+    model.mimi = mimi
+
+    output = model(input_ids=_codes(frames), request_ids=["req"])
+
+    assert _audio(output).numel() == frames * 4
+    assert mimi.decode_frames_shapes == expected_shapes
+    assert mimi.decode_frame_calls == 0
+
+
+def test_profile_inputs_skip_decode_while_malformed_online_input_warns(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    model, mimi = _model()
+    runtime_info = [{"meta": {"personaplex_dummy_profile": True}}]
+    warnings: list[tuple] = []
+    monkeypatch.setattr(personaplex_code2wav.logger, "warning", lambda *args: warnings.append(args))
+
+    output = model(
+        input_ids=torch.arange(3),
+        runtime_additional_information=runtime_info,
+    )
+
+    assert _audio(output).numel() == 0
+    assert mimi.decode_frame_calls == 0
+    assert warnings == []
+
+    output = model(input_ids=torch.arange(3))
+
+    assert _audio(output).numel() == 0
+    assert mimi.decode_frame_calls == 0
+    assert len(warnings) == 1
+    assert "not divisible by" in warnings[0][0]
+    assert model.get_dummy_runtime_additional_information(2) == [
+        {"meta": {"personaplex_dummy_profile": True}},
+        {"meta": {"personaplex_dummy_profile": True}},
+    ]
+
+
+def test_request_id_falls_back_to_runtime_information() -> None:
+    model, _ = _model()
+    info = [{"request_id": "runtime-req"}]
+
+    model(input_ids=_codes(2), runtime_additional_information=info)
+    second = model(input_ids=_codes(3), runtime_additional_information=info)
+
+    assert _audio(second).numel() == 4
+
+
+def test_decoder_slot_lifecycle_isolated_capacity_and_reuse() -> None:
+    model, _ = _model(max_sessions=2)
+    first = _FakeStreamingMimi(stream_value=10)
+    second = _FakeStreamingMimi(stream_value=20)
+    model._set_mimi_codecs([first, second])
+    assert model._max_codec_sessions == 2
+
+    initial = model(
+        input_ids=torch.cat([_codes(1), _codes(1)]),
+        request_ids=["first", "second"],
+        seq_token_counts=[2, 2],
+    )
+    continued = model(
+        input_ids=torch.cat([_codes(1, start=10), _codes(1, start=20)]),
+        request_ids=["first", "second"],
+        seq_token_counts=[2, 2],
+    )
+
+    initial_audio = initial.multimodal_outputs["model_outputs"]
+    continued_audio = continued.multimodal_outputs["model_outputs"]
+    assert initial_audio[0].tolist() == [11.0] * 4
+    assert initial_audio[1].tolist() == [21.0] * 4
+    assert continued_audio[0].tolist() == [12.0] * 4
+    assert continued_audio[1].tolist() == [22.0] * 4
+    assert model._request_codec_slots == {"first": 0, "second": 1}
+    with pytest.raises(RuntimeError, match="decoder capacity 2 is exhausted"):
+        model._codec_for_request("overflow")
+
+    model.on_requests_finished(["first"])
+    assert first.reset_calls == 1
+    assert second.reset_calls == 0
+    replacement = model(input_ids=_codes(1), request_ids=["replacement"])
+    assert _audio(replacement).tolist() == [13.0] * 4

tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py

@@ -0,0 +1,364 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import base64
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+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.experimental.fullduplex.engine.duplex_runtime import (
+    DuplexInputMode,
+)
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+from vllm_omni.experimental.fullduplex.personaplex.data_plane import (
+    PersonaPlexDataPlaneContext,
+    PersonaPlexDataPlaneSession,
+)
+from vllm_omni.experimental.fullduplex.personaplex.input import (
+    PersonaPlexPcmAppendBuffer,
+)
+from vllm_omni.experimental.fullduplex.personaplex.runtime_extension import (
+    PersonaPlexDuplexRuntimeExtension,
+)
+from vllm_omni.experimental.fullduplex.personaplex.serving_adapter import (
+    PersonaPlexServingRuntimeAdapter,
+)
+from vllm_omni.model_executor.models.personaplex.pipeline import (
+    PERSONAPLEX_PIPELINE,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _pcm_payload(samples: np.ndarray, *, sample_rate_hz: int = 24000) -> dict[str, object]:
+    samples = np.ascontiguousarray(samples, dtype="<f4")
+    return {
+        "type": "audio",
+        "format": "pcm_f32le",
+        "sample_rate_hz": sample_rate_hz,
+        "audio": base64.b64encode(samples.tobytes()).decode("ascii"),
+    }
+
+
+def _audio_client(
+    frame_count: int = 10,
+    *,
+    voiced_frames: int = 0,
+) -> SimpleNamespace:
+    frames = np.full((frame_count, e2e_driver.FRAME_SAMPLES), 7, dtype="<i2")
+    frames[:voiced_frames] = 1000
+    raw = frames.tobytes()
+    event = {
+        "type": "response.audio.delta",
+        "response_id": "response-1",
+        "sample_rate_hz": e2e_driver.SAMPLE_RATE_HZ,
+        "metadata": {
+            "vllm_omni": {
+                "runtime_impl": "scheduler_data_plane",
+                "uses_model_runner_scheduler": True,
+                "runner_kv_backed": True,
+            }
+        },
+    }
+    events = SimpleNamespace(
+        events=[event],
+        response_audio={"response-1": [raw]},
+        audio_bytes=lambda: raw,
+    )
+    return SimpleNamespace(events=events)
+
+
+def test_realtime_audio_frame_stats_separate_voiced_and_silent_frames() -> None:
+    silent = np.zeros(e2e_driver.FRAME_SAMPLES, dtype="<i2")
+    voiced = np.full(e2e_driver.FRAME_SAMPLES, 1000, dtype="<i2")
+
+    stats = e2e_driver._audio_frame_stats(
+        silent.tobytes() + voiced.tobytes(),
+        input_frames=3,
+        voiced_frame_rms_threshold=1e-4,
+    )
+
+    assert stats == {
+        "output_frames": 2,
+        "voiced_frames": 1,
+        "silent_frames": 1,
+        "frame_deficit": 1,
+        "frame_coverage_ratio": pytest.approx(2 / 3),
+    }
+
+
+def test_realtime_audio_frame_stats_reject_partial_codec_frame() -> None:
+    partial_frame = np.zeros(e2e_driver.FRAME_SAMPLES - 1, dtype="<i2")
+
+    with pytest.raises(AssertionError, match="whole 80 ms codec frames"):
+        e2e_driver._audio_frame_stats(
+            partial_frame.tobytes(),
+            input_frames=1,
+            voiced_frame_rms_threshold=1e-4,
+        )
+
+
+@pytest.mark.parametrize("frame_count", [10, 40])
+def test_e2e_driver_rejects_inaudible_sessions(frame_count: int) -> None:
+    args = SimpleNamespace(
+        max_frame_deficit=4,
+        voiced_frame_rms_threshold=1e-3,
+        min_voiced_frames=5,
+    )
+
+    with pytest.raises(AssertionError, match="audible speech"):
+        e2e_driver._session_result(
+            _audio_client(frame_count),
+            input_frames=frame_count,
+            args=args,
+            minimum_chunks=1,
+        )
+
+
+def test_e2e_driver_uses_absolute_audible_floor() -> None:
+    args = SimpleNamespace(
+        max_frame_deficit=4,
+        voiced_frame_rms_threshold=1e-3,
+        min_voiced_frames=5,
+    )
+
+    e2e_driver._session_result(
+        _audio_client(200, voiced_frames=5),
+        input_frames=200,
+        args=args,
+        minimum_chunks=1,
+    )
+
+
+def test_e2e_driver_defaults_replacement_to_full_input() -> None:
+    args = e2e_driver.parse_args(
+        [
+            "--model",
+            "/model",
+            "--input-wav",
+            "/input.wav",
+        ]
+    )
+
+    assert args.replacement_frames == 0
+    assert args.min_voiced_frames == 5
+
+
+def test_e2e_driver_records_and_validates_input_sha256(tmp_path: Path) -> None:
+    input_wav = tmp_path / "input.wav"
+    input_wav.write_bytes(b"fixture-wav")
+
+    identity = e2e_driver._input_identity(
+        input_wav,
+        expected_sha256="bade4b3c163edde390ff391207d34a887257d8a9cc3b621cc8c618b6e6761304",
+    )
+
+    assert identity == {
+        "path": str(input_wav.resolve()),
+        "sha256": "bade4b3c163edde390ff391207d34a887257d8a9cc3b621cc8c618b6e6761304",
+    }
+
+
+def test_e2e_driver_rejects_unexpected_input_sha256(tmp_path: Path) -> None:
+    input_wav = tmp_path / "input.wav"
+    input_wav.write_bytes(b"fixture-wav")
+
+    with pytest.raises(ValueError, match="input WAV SHA-256 mismatch"):
+        e2e_driver._input_identity(input_wav, expected_sha256="0" * 64)
+
+
+@pytest.mark.parametrize(
+    ("attribute", "expected"),
+    [
+        (
+            "duplex_runtime_extension",
+            "vllm_omni.experimental.fullduplex.personaplex.runtime_extension.PersonaPlexDuplexRuntimeExtension",
+        ),
+        (
+            "duplex_serving_adapter",
+            "vllm_omni.experimental.fullduplex.personaplex.serving_adapter.PersonaPlexServingRuntimeAdapter",
+        ),
+    ],
+)
+def test_personaplex_pipeline_registers_unified_duplex_components(
+    attribute: str,
+    expected: str,
+) -> None:
+    assert PERSONAPLEX_PIPELINE.duplex_control_enabled is True
+    assert getattr(PERSONAPLEX_PIPELINE, attribute) == expected
+
+
+@pytest.mark.parametrize(
+    ("engine_arg", "expected"),
+    [
+        ("skip_tokenizer_init", True),
+        ("enable_prefix_caching", False),
+    ],
+)
+def test_personaplex_stage_engine_args(
+    engine_arg: str,
+    expected: bool,
+) -> None:
+    deploy_path = Path(__file__).parents[5] / "vllm_omni" / "deploy" / "personaplex.yaml"
+    stages = merge_pipeline_deploy(
+        PERSONAPLEX_PIPELINE,
+        load_deploy_config(deploy_path),
+    )
+
+    assert [stage.yaml_engine_args.get(engine_arg) for stage in stages] == [
+        expected,
+        expected,
+    ]
+
+
+def test_personaplex_duplex_capacity_is_propagated_to_all_model_stages() -> None:
+    deploy_path = Path(__file__).parents[5] / "vllm_omni" / "deploy" / "personaplex.yaml"
+    deploy = load_deploy_config(deploy_path)
+    stages = merge_pipeline_deploy(PERSONAPLEX_PIPELINE, deploy)
+
+    assert deploy.duplex_session.max_sessions == 2
+    assert [stage.yaml_engine_args.get("duplex_max_sessions") for stage in stages] == [2, 2]
+    assert "personaplex_codec_max_sessions" not in deploy.connectors["connector_of_shared_memory"]["extra"]
+
+
+def test_personaplex_capabilities_are_honest() -> None:
+    single = PersonaPlexServingRuntimeAdapter(lambda *_: None).capabilities(max_sessions=1)
+    multi = PersonaPlexServingRuntimeAdapter(lambda *_: None).capabilities(max_sessions=2)
+
+    assert multi.implementation_level == "model_native_duplex"
+    assert multi.input_modes == ["append_audio_chunk"]
+    assert multi.chunk_period_ms == 80
+    assert single.supports_multi_session is False
+    assert multi.supports_multi_session is True
+    assert multi.supports_multi_session_same_replica is True
+    assert multi.supports_barge_in is False
+
+
+def test_pcm_buffer_emits_one_80ms_frame_transactionally() -> None:
+    buffer = PersonaPlexPcmAppendBuffer()
+    reservation = buffer.prepare_append(
+        _pcm_payload(np.arange(1920, dtype=np.float32)),
+        operation_id="op-1",
+        chunk_period_ms=80,
+        allow_emit=True,
+    )
+
+    assert reservation is not None
+    assert reservation.byte_count == 1920 * 4
+    assert buffer.pending_byte_count == 0
+    reservation.rollback()
+    assert buffer.pending_byte_count == 1920 * 4
+
+
+def test_pcm_buffer_rejects_changed_sample_rate() -> None:
+    buffer = PersonaPlexPcmAppendBuffer()
+    assert (
+        buffer.prepare_append(
+            _pcm_payload(np.zeros(100, np.float32)),
+            operation_id="op-1",
+            chunk_period_ms=80,
+            allow_emit=False,
+        )
+        is None
+    )
+
+    with pytest.raises(ValueError, match="sample_rate"):
+        buffer.prepare_append(
+            _pcm_payload(np.zeros(100, np.float32), sample_rate_hz=16000),
+            operation_id="op-2",
+            chunk_period_ms=80,
+            allow_emit=False,
+        )
+
+
+def test_pcm_buffer_rejects_non_finite_samples() -> None:
+    buffer = PersonaPlexPcmAppendBuffer()
+
+    with pytest.raises(ValueError, match="finite"):
+        buffer.prepare_append(
+            _pcm_payload(np.array([0.0, np.nan], np.float32)),
+            operation_id="op-1",
+            chunk_period_ms=80,
+            allow_emit=False,
+        )
+
+
+def test_runtime_extension_builds_one_frame_scheduler_append() -> None:
+    extension = PersonaPlexDuplexRuntimeExtension()
+    defaults = (
+        SamplingParams(temperature=0.8, top_k=10, max_tokens=10),
+        SamplingParams(max_tokens=1024),
+    )
+
+    configured = extension.configure_sampling_params(runtime_config={}, defaults=defaults)
+    assert configured[0].temperature == 0.0
+    assert configured[0].top_k == 1
+    assert configured[0].max_tokens == 1
+    assert configured[1] is defaults[1]
+
+    plan = extension.plan_append(
+        request_id="req",
+        fence=DuplexFence("session", incarnation=2, epoch=3),
+        session_config={"instructions": "Be concise."},
+        runtime_config={"personaplex_prefill_slots": 4},
+        seq=1,
+        turn_seq=1,
+        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+        payload=_pcm_payload(np.zeros(1920, np.float32)),
+        final=False,
+        sampling_params=configured[0],
+    )
+    assert len(plan.prompt["prompt_token_ids"]) == 5
+    duplex = plan.prompt["model_intermediate_buffer"]["duplex"]
+    assert duplex["session_id"] == "session"
+    assert duplex["incarnation"] == 2
+    assert duplex["epoch"] == 3
+    assert duplex["seq"] == 1
+    assert duplex["data_plane"] is True
+
+
+def test_projector_emits_only_cumulative_audio_and_text_suffixes() -> None:
+    encoded_sizes: list[int] = []
+
+    def encode_audio(audio, _sample_rate, _response_format, _speed):
+        size = int(np.asarray(audio, dtype=np.float32).size)
+        encoded_sizes.append(size)
+        return f"audio-{size}"
+
+    projector = PersonaPlexDataPlaneSession(encode_audio)
+    projector.begin_request("req")
+    context = PersonaPlexDataPlaneContext(response_format="wav")
+
+    first = SimpleNamespace(
+        request_id="req",
+        outputs=[SimpleNamespace(text="he", multimodal_output={})],
+        multimodal_output={"model_outputs": np.arange(4, dtype=np.float32), "sr": 24000},
+        finished=False,
+    )
+    second = SimpleNamespace(
+        request_id="req",
+        outputs=[SimpleNamespace(text="hello", multimodal_output={})],
+        multimodal_output={"model_outputs": np.arange(6, dtype=np.float32), "sr": 24000},
+        finished=False,
+    )
+
+    projected = list(
+        projector.project(
+            {"data_plane_outputs": [first, second]},
+            context=context,
+        )
+    )
+
+    assert encoded_sizes == [4, 2]
+    assert [item["text"] for item in projected] == ["he", "llo"]
+    assert [item["audio_data"] for item in projected] == ["audio-4", "audio-2"]
+    assert all(item["sample_rate_hz"] == 24000 for item in projected)
+    assert all(item["end_of_turn"] is False for item in projected)

tests/e2e/features/fullduplex/test_core_contracts.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Contract tests for the fullduplex core (PR #3907 alignment).
+
+Covers the DuplexFence identity value, the session identity rules (epoch only
+advances on barge-in, response_index once per begin_response), stale-output
+dropping in the turn lifecycle. The lockstep (single-eternal-response)
+lifecycle is model-owned (see
+tests/e2e/features/fullduplex/test_personaplex_adapter.py).
+"""
+
+import asyncio
+import dataclasses
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.core import protocol as ev
+from vllm_omni.experimental.fullduplex.core.adapter import (
+    DuplexAdapter,
+    DuplexCapability,
+    OutputChunk,
+)
+from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime
+from vllm_omni.experimental.fullduplex.core.session import (
+    DuplexSession,
+    DuplexSessionConfig,
+    DuplexState,
+)
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_fence_is_frozen_value_identity():
+    f = DuplexFence(session_id="s1")
+    assert (f.epoch, f.turn_id, f.response_seq) == (0, 0, 0)
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        f.epoch = 1  # type: ignore[misc]
+    assert DuplexFence("s1", 1, 2, 3) == DuplexFence("s1", 1, 2, 3)
+    assert DuplexFence("s1", 1, 2, 3) != DuplexFence("s1", 2, 2, 3)
+
+
+def test_session_identity_rules():
+    s = DuplexSession(session_id="test", config=DuplexSessionConfig())
+    r1, e1 = s.begin_response()
+    r2, e2 = s.begin_response()
+    assert r2 == r1 + 1, "response index advances exactly once per begin"
+    assert e1 == e2 == 0, "epoch does not move on response begin"
+    s.end_response()
+    assert s.epoch == 0, "epoch does not move on response end"
+    new_epoch = s.barge_in()
+    assert new_epoch == 1, "epoch advances only on barge-in"
+    assert s.is_stale(e1) and not s.is_stale(new_epoch)
+
+
+class _TurnAdapter(DuplexAdapter):
+    """Emits chunks until cancelled; counts respond() lifecycles."""
+
+    def __init__(self):
+        self.responses = 0
+        self.barge_ins = 0
+
+    def capabilities(self) -> DuplexCapability:
+        return DuplexCapability(
+            input_modalities=frozenset({"audio"}),
+            output_modalities=frozenset({"audio"}),
+            proactive=True,
+        )
+
+    async def on_input(self, session, modality, data) -> None:
+        pass
+
+    def should_respond(self, session) -> bool:
+        return True
+
+    async def on_barge_in(self, session) -> None:
+        self.barge_ins += 1
+
+    async def respond(self, session):
+        self.responses += 1
+        for i in range(100):
+            await asyncio.sleep(0.001)
+            yield OutputChunk("audio", i)
+
+
+async def _drive(runtime, events):
+    async def gen():
+        for e in events:
+            await asyncio.sleep(0.005)
+            yield e
+
+    out = []
+
+    async def emit(e):
+        out.append(e)
+
+    await runtime.run(gen(), emit)
+    return out
+
+
+def test_turn_mode_supersede_and_stale_drop():
+    adapter = _TurnAdapter()
+    session = DuplexSession(session_id="test", config=DuplexSessionConfig(proactive=True))
+    runtime = DuplexRuntime(session, adapter)
+    out = asyncio.run(
+        _drive(
+            runtime,
+            [
+                {"type": ev.RESPONSE_CREATE},
+                {"type": ev.RESPONSE_CREATE},  # supersedes the first
+                {"type": ev.CLOSE},
+            ],
+        )
+    )
+    assert adapter.responses == 2
+    assert adapter.barge_ins == 1, "superseding an in-flight response barges in"
+    created = [e for e in out if e["type"] == ev.RESPONSE_CREATED]
+    done = [e for e in out if e["type"] == ev.RESPONSE_DONE]
+    assert len(created) == 2
+    assert len(done) == 1, "stale first response must not emit done"
+    assert done[0]["response_index"] == created[1]["response_index"]
+    assert session.state is DuplexState.CLOSED

tests/e2e/features/fullduplex/test_personaplex_adapter.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex adapter driven through the model-owned lockstep runtime.
+
+PersonaPlexDuplexRuntime starts ONE response for the whole session, streams one
+agent frame per user frame, and drains on close (core stays verbatim upstream;
+the lockstep lifecycle is model policy) — all with a GPU-free stub stepper.
+"""
+
+import asyncio
+
+import numpy as np
+import pytest
+
+from vllm_omni.experimental.fullduplex.core import protocol as ev
+from vllm_omni.experimental.fullduplex.core.session import (
+    DuplexSession,
+    DuplexSessionConfig,
+    DuplexState,
+)
+from vllm_omni.experimental.fullduplex.personaplex.adapter import (
+    PersonaPlexDuplexAdapter,
+    PersonaPlexDuplexRuntime,
+)
+from vllm_omni.experimental.fullduplex.personaplex.config import FRAME_SIZE
+from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class StubStepper:
+    """Deterministic FrameStepper: frame i -> audio filled with i, text 't{i}' on odds."""
+
+    sample_rate = 24000
+    frame_size = FRAME_SIZE
+
+    def __init__(self) -> None:
+        self.opened: tuple[str | None, str | None] | None = None
+        self.steps = 0
+
+    def open_session(self, voice_prompt=None, persona=None) -> None:
+        self.opened = (voice_prompt, persona)
+        self.steps = 0
+
+    def step(self, user_pcm):
+        assert user_pcm.shape[0] == self.frame_size
+        self.steps += 1
+        i = self.steps
+        return FrameOutput(
+            audio=np.full(self.frame_size, float(i), dtype=np.float32),
+            text=(f"t{i}" if i % 2 else None),
+        )
+
+
+class RecordingAdapter(PersonaPlexDuplexAdapter):
+    def __init__(self, stepper: StubStepper) -> None:
+        super().__init__(stepper)
+        self.closed = False
+
+    async def on_close(self, session: DuplexSession) -> None:
+        self.closed = True
+        await super().on_close(session)
+
+
+def _collector():
+    out: list[dict] = []
+
+    async def emit(event: dict) -> None:
+        out.append(event)
+
+    return out, emit
+
+
+async def _feed(events):
+    for e in events:
+        yield e
+
+
+async def _raising_feed(frame):
+    yield {"type": ev.INPUT_APPEND, "modality": "audio", "data": frame}
+    await asyncio.sleep(0)
+    raise RuntimeError("input failed")
+
+
+def _respond_tasks() -> list[asyncio.Task]:
+    current = asyncio.current_task()
+    return [
+        task
+        for task in asyncio.all_tasks()
+        if task is not current and "_respond_once" in getattr(task.get_coro(), "__qualname__", "")
+    ]
+
+
+def _session():
+    cfg = DuplexSessionConfig(
+        input_modalities=("audio",),
+        output_modalities=("audio", "text"),
+    )
+    return DuplexSession("s", cfg)
+
+
+@pytest.mark.asyncio
+async def test_lockstep_one_agent_frame_per_user_frame():
+    stub = StubStepper()
+    adapter = PersonaPlexDuplexAdapter(stub, voice_prompt="NATF2.pt", persona="be terse")
+    rt = PersonaPlexDuplexRuntime(_session(), adapter)
+    out, emit = _collector()
+
+    frame = np.zeros(FRAME_SIZE, dtype=np.float32)
+    events = [{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame} for _ in range(3)]
+    events.append({"type": ev.CLOSE})
+    await rt.run(_feed(events), emit)
+
+    types = [e["type"] for e in out]
+    # exactly one response for the whole session, drained cleanly on close
+    assert types.count(ev.RESPONSE_CREATED) == 1
+    assert types.count(ev.RESPONSE_DONE) == 1
+    audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
+    text = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "text"]
+    assert len(audio) == 3  # one agent frame per user frame
+    assert [int(e["data"][0]) for e in audio] == [1, 2, 3]  # lockstep order
+    assert text == ["t1", "t3"]
+    assert stub.steps == 3
+    assert stub.opened == ("NATF2.pt", "be terse")
+
+
+@pytest.mark.asyncio
+async def test_unaligned_chunks_are_reframed_to_80ms():
+    stub = StubStepper()
+    adapter = PersonaPlexDuplexAdapter(stub)
+    rt = PersonaPlexDuplexRuntime(_session(), adapter)
+    out, emit = _collector()
+
+    # 2.5 frames in one chunk + 0.5 frame in another => 3 whole frames
+    events = [
+        {"type": ev.INPUT_APPEND, "modality": "audio", "data": np.zeros(FRAME_SIZE * 5 // 2, dtype=np.float32)},
+        {"type": ev.INPUT_APPEND, "modality": "audio", "data": np.zeros(FRAME_SIZE // 2, dtype=np.float32)},
+        {"type": ev.CLOSE},
+    ]
+    await rt.run(_feed(events), emit)
+    audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
+    assert len(audio) == 3
+    assert stub.steps == 3
+
+
+@pytest.mark.asyncio
+async def test_close_without_input_is_clean():
+    adapter = PersonaPlexDuplexAdapter(StubStepper())
+    rt = PersonaPlexDuplexRuntime(_session(), adapter)
+    out, emit = _collector()
+    await rt.run(_feed([{"type": ev.CLOSE}]), emit)
+    # no input -> response never started -> no created/done, no error
+    assert not any(e["type"] == ev.ERROR for e in out)
+    assert not any(e["type"] == ev.RESPONSE_CREATED for e in out)
+
+
+@pytest.mark.asyncio
+async def test_iterator_eof_without_close_drains_and_stops():
+    # Transport drop: the input iterator ends WITHOUT an explicit CLOSE event.
+    # The runtime must still signal on_close so the eternal response drains
+    # instead of blocking on its inbox forever.
+    stub = StubStepper()
+    adapter = PersonaPlexDuplexAdapter(stub)
+    rt = PersonaPlexDuplexRuntime(_session(), adapter)
+    out, emit = _collector()
+    frame = np.zeros(FRAME_SIZE, dtype=np.float32)
+    events = [{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame}]
+    await asyncio.wait_for(rt.run(_feed(events), emit), timeout=5.0)
+    types = [e["type"] for e in out]
+    assert types.count(ev.RESPONSE_DONE) == 1
+    assert stub.steps == 1
+
+
+@pytest.mark.asyncio
+async def test_input_exception_closes_response_and_runtime() -> None:
+    stub = StubStepper()
+    adapter = RecordingAdapter(stub)
+    rt = PersonaPlexDuplexRuntime(_session(), adapter)
+    out, emit = _collector()
+    frame = np.zeros(FRAME_SIZE, dtype=np.float32)
+
+    with pytest.raises(RuntimeError, match="input failed"):
+        await rt.run(_raising_feed(frame), emit)
+
+    leaked = _respond_tasks()
+    for task in leaked:
+        task.cancel()
+    if leaked:
+        await asyncio.gather(*leaked, return_exceptions=True)
+
+    assert adapter.closed is True
+    assert rt.session.state is DuplexState.CLOSED
+    assert not leaked
+
+
+@pytest.mark.asyncio
+async def test_default_session_config_gets_lockstep_lifecycle():
+    # The lifecycle comes from the model-owned runtime class, not a session
+    # flag: a plain default-constructed session config runs lockstep.
+    stub = StubStepper()
+    adapter = PersonaPlexDuplexAdapter(stub)
+    session = DuplexSession(
+        "s",
+        DuplexSessionConfig(input_modalities=("audio",), output_modalities=("audio", "text")),
+    )
+    rt = PersonaPlexDuplexRuntime(session, adapter)
+    out, emit = _collector()
+    frame = np.zeros(FRAME_SIZE, dtype=np.float32)
+    events = [{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame}, {"type": ev.CLOSE}]
+    await rt.run(_feed(events), emit)
+    types = [e["type"] for e in out]
+    assert types.count(ev.RESPONSE_CREATED) == 1
+    assert types.count(ev.RESPONSE_DONE) == 1
+    assert stub.steps == 1
+
+
+@pytest.mark.asyncio
+async def test_partial_tail_is_padded_and_stepped_on_close():
+    # A trailing sub-frame chunk is zero-padded and stepped at close (same
+    # contract as PersonaPlexSession.flush) rather than silently dropped.
+    stub = StubStepper()
+    adapter = PersonaPlexDuplexAdapter(stub)
+    rt = PersonaPlexDuplexRuntime(_session(), adapter)
+    out, emit = _collector()
+    events = [
+        {"type": ev.INPUT_APPEND, "modality": "audio", "data": np.ones(FRAME_SIZE // 2, dtype=np.float32)},
+        {"type": ev.CLOSE},
+    ]
+    await rt.run(_feed(events), emit)
+    audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
+    assert len(audio) == 1
+    assert stub.steps == 1
+
+
+@pytest.mark.asyncio
+async def test_single_eternal_response_across_barge_in():
+    # Moved from the core contract tests when the lockstep lifecycle became
+    # model-owned: ONE response for the whole session, RESPONSE_CANCEL bumps the
+    # epoch fence (playback cursor) but the stream keeps flowing, and close
+    # drains via on_close.
+    stub = StubStepper()
+    adapter = PersonaPlexDuplexAdapter(stub)
+    session = _session()
+    rt = PersonaPlexDuplexRuntime(session, adapter)
+    out, emit = _collector()
+
+    frame = np.zeros(FRAME_SIZE, dtype=np.float32)
+    events = [
+        {"type": ev.INPUT_APPEND, "modality": "audio", "data": frame},
+        {"type": ev.RESPONSE_CANCEL},  # barge-in: epoch bumps, stream continues
+        {"type": ev.INPUT_APPEND, "modality": "audio", "data": frame},
+        {"type": ev.CLOSE},
+    ]
+    await rt.run(_feed(events), emit)
+
+    types = [e["type"] for e in out]
+    assert types.count(ev.RESPONSE_CREATED) == 1, "exactly one eternal response"
+    assert types.count(ev.RESPONSE_DONE) == 1
+    audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
+    assert len(audio) == 2, "stream keeps flowing across barge-in"
+    assert session.epoch == 1, "cancel still advanced the epoch fence"
+    assert stub.steps == 2

tests/e2e/features/fullduplex/test_personaplex_batched.py

@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GPU-free tests for the batched PersonaPlex slot manager (stub engine)."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput, PrefillStep
+from vllm_omni.experimental.fullduplex.personaplex.serving.batched import (
+    BatchedSessionManager,
+    SlotPhase,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+FRAME = 4
+
+
+class StubEngine:
+    """Records the slot-API call sequence; echoes each row's pcm as its output."""
+
+    frame_size = FRAME
+
+    def __init__(self, batch_size: int = 4, prefill_len: int = 2) -> None:
+        self.config = SimpleNamespace(batch_size=batch_size)
+        self.prefill_len = prefill_len
+        self.recycled: list[int] = []
+        self.voice_ended: list[int] = []
+        self.prefill_done: list[int] = []
+        self.steps: list[tuple[np.ndarray, dict | None]] = []
+
+    def prefill_steps(self, persona=None):
+        steps = [PrefillStep("sacrifice"), PrefillStep("voice", embedding=None), PrefillStep("voice_end")]
+        steps += [PrefillStep("tokens", moshi_tokens=None, text_token=3, user_sine=True)] * self.prefill_len
+        return steps
+
+    def recycle_slot(self, b):
+        self.recycled.append(b)
+
+    def voice_end(self, b):
+        self.voice_ended.append(b)
+
+    def end_prefill(self, b):
+        self.prefill_done.append(b)
+
+    def step_batch(self, pcm, prefill=None):
+        self.steps.append((pcm.copy(), dict(prefill) if prefill else None))
+        return [FrameOutput(audio=pcm[b].copy(), text=None) for b in range(self.config.batch_size)]
+
+
+def test_all_idle_does_not_step():
+    eng = StubEngine()
+    mgr = BatchedSessionManager(eng)
+    assert mgr.tick() is False
+    assert eng.steps == []
+
+
+def test_pristine_slot_is_live_immediately():
+    eng = StubEngine()
+    mgr = BatchedSessionManager(eng)
+    slot, ready_now = mgr.acquire()
+    assert ready_now is True
+    assert slot.phase is SlotPhase.LIVE
+    assert eng.recycled == []
+
+
+def test_custom_persona_prefills_and_recycles_on_tick_thread():
+    eng = StubEngine(prefill_len=2)
+    mgr = BatchedSessionManager(eng)
+    slot, ready_now = mgr.acquire(persona="custom")
+    assert ready_now is False
+    assert slot.phase is SlotPhase.PREFILL
+    assert eng.recycled == []  # deferred: reset must run between engine steps
+
+    ready = []
+    slot.on_ready = lambda: ready.append(True)
+
+    mgr.tick()  # recycle + sacrifice
+    assert eng.recycled == [slot.index]
+    assert eng.steps[-1][1][slot.index].kind == "sacrifice"
+
+    mgr.tick()  # voice frame
+    assert eng.steps[-1][1][slot.index].kind == "voice"
+    assert eng.voice_ended == []
+
+    mgr.tick()  # voice_end executed BETWEEN ticks, then first tokens frame
+    assert eng.voice_ended == [slot.index]
+    assert eng.steps[-1][1][slot.index].kind == "tokens"
+
+    mgr.tick()  # second tokens frame
+    assert not ready
+    mgr.tick()  # prefill exhausted -> live
+    assert eng.prefill_done == [slot.index]
+    assert slot.phase is SlotPhase.LIVE
+    assert ready == [True]
+
+
+def test_pristine_slot_after_ticking_requires_recycle():
+    # step_batch advances EVERY row each tick, so a pristine idle row that has
+    # lived through ticks has drifted from its boot prefill and must be
+    # recycled before a client may join it.
+    eng = StubEngine()
+    mgr = BatchedSessionManager(eng)
+    first, ready_now = mgr.acquire()
+    assert ready_now is True  # claimed before any tick: truly pristine
+    mgr.tick()
+    second, ready_now2 = mgr.acquire()
+    assert second.index != first.index
+    assert ready_now2 is False
+    assert second.phase is SlotPhase.PREFILL
+
+
+def test_used_slot_requires_recycle_even_with_default_persona():
+    eng = StubEngine()
+    mgr = BatchedSessionManager(eng)
+    slot, ready_now = mgr.acquire()
+    assert ready_now
+    mgr.release(slot)
+    slot2, ready_now2 = mgr.acquire()
+    assert slot2.index == slot.index
+    assert ready_now2 is False
+    assert slot2.phase is SlotPhase.PREFILL
+
+
+def test_acquire_returns_none_when_full():
+    eng = StubEngine(batch_size=2)
+    mgr = BatchedSessionManager(eng)
+    assert mgr.acquire() is not None
+    assert mgr.acquire() is not None
+    assert mgr.acquire() is None
+
+
+def test_feed_frames_epoch_guard_and_backlog_cap():
+    eng = StubEngine()
+    mgr = BatchedSessionManager(eng, max_buffered_frames=3)
+    slot, _ = mgr.acquire()
+    mgr.feed(slot, np.arange(FRAME * 2, dtype=np.float32), slot.epoch)
+    assert len(slot.inq) == 2
+    mgr.feed(slot, np.zeros(3, dtype=np.float32), slot.epoch)  # partial tail buffered
+    assert len(slot.inq) == 2 and slot.pending.shape[0] == 3
+    mgr.feed(slot, np.zeros(FRAME * 5, dtype=np.float32), slot.epoch)
+    assert len(slot.inq) == 3  # oldest dropped beyond the cap
+    mgr.feed(slot, np.zeros(FRAME, dtype=np.float32), slot.epoch - 1)  # stale epoch ignored
+    assert len(slot.inq) == 3
+
+
+def test_live_slot_consumes_input_and_emits_output():
+    eng = StubEngine()
+    mgr = BatchedSessionManager(eng)
+    slot, _ = mgr.acquire()
+    got = []
+    slot.on_output = got.append
+    frame = np.full(FRAME, 0.5, dtype=np.float32)
+    mgr.feed(slot, frame, slot.epoch)
+    mgr.tick()
+    assert len(got) == 1
+    np.testing.assert_array_equal(got[0].audio, frame)
+    # queue drained: next tick feeds zeros
+    mgr.tick()
+    np.testing.assert_array_equal(got[1].audio, np.zeros(FRAME, dtype=np.float32))
+
+
+def test_release_mid_flight_stops_emission():
+    eng = StubEngine()
+    mgr = BatchedSessionManager(eng)
+    slot, _ = mgr.acquire()
+    got = []
+    slot.on_output = got.append
+    mgr.tick()
+    assert len(got) == 1
+    mgr.release(slot)
+    # another slot keeps the batch non-idle
+    other, _ = mgr.acquire()  # same slot index, now PREFILL (used)
+    assert other.phase is SlotPhase.PREFILL
+    mgr.tick()
+    assert len(got) == 1  # no further output for the released epoch

tests/e2e/features/fullduplex/test_personaplex_elastic.py

@@ -0,0 +1,240 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GPU e2e for PersonaPlex elastic batching (bit-exact isolation + recycle).
+
+Greedy decoding makes every cross-slot leak observable as a bit difference:
+
+1. Isolation (bit-exact): recycling one slot mid-flight leaves every other
+   slot's output bit-identical to a run that never recycled.
+2. Replay exactness: a recycled slot's opening live ticks match a pristine boot
+   at numerical-noise level (a broken KV mask / staging roll garbles frame 0).
+3. Coherence (behavioral): a recycled slot fed a real spoken question produces
+   a real reply in its inner-monologue text stream.
+
+Run (moshi-free native stack; ~30 GB GPU + HF access to the gated repo; set
+PPLEX_QUESTION_WAV to a spoken-question wav to enable the coherence test):
+    pytest tests/e2e/features/fullduplex/test_personaplex_elastic.py -v -s
+
+Bench (B sweep, ms/tick vs the 80 ms budget):
+    python tests/e2e/features/fullduplex/test_personaplex_elastic.py --bench 2 4 8 16
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+
+# slow + H100 + omni routes this module into the weekly Buildkite E2E group
+# ("Omni · H100" in .buildkite/cuda/test-weekly.yml: -m "slow and H100 and
+# omni" --run-level full_model); the ready/merge CPU sweeps exclude it by mark.
+pytestmark = [pytest.mark.full_model, pytest.mark.omni, pytest.mark.slow, pytest.mark.H100]
+
+if not torch.cuda.is_available():  # noqa: E402
+    pytest.skip("needs CUDA", allow_module_level=True)
+
+from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig  # noqa: E402
+from vllm_omni.experimental.fullduplex.personaplex.runtime import (  # noqa: E402
+    PersonaPlexEngine,
+)
+
+B = 4
+LIVE_TICKS = 40  # live frames compared per claim
+RECYCLE_AT = 10
+
+
+def _make_engine() -> PersonaPlexEngine:
+    cfg = PersonaPlexConfig(batch_size=B)  # native stepper is greedy-only
+    try:
+        return PersonaPlexEngine(cfg).load()
+    except Exception as exc:  # gated-repo access is environment-dependent
+        name = type(exc).__name__
+        if "Gated" in name or "401" in str(exc) or "403" in str(exc):
+            pytest.skip(f"no access to the gated nvidia/personaplex-7b-v1 repo: {name}")
+        raise
+
+
+@pytest.fixture(scope="module")
+def engine() -> PersonaPlexEngine:
+    return _make_engine()
+
+
+def _slot_stream(slot: int, tick: int, frame: int) -> np.ndarray:
+    """Deterministic per-slot input: distinct low-amplitude tones."""
+    sr = 24000
+    freq = 180.0 + 90.0 * slot
+    t0 = tick * frame
+    t = (np.arange(frame) + t0) / sr
+    return (0.12 * np.sin(2 * np.pi * freq * t)).astype(np.float32)
+
+
+def _run(engine, total_ticks: int, recycle: bool):
+    """Run the lockstep loop; optionally recycle slot 1 at RECYCLE_AT.
+
+    Returns per-slot lists of (audio | None) per live tick index, where the
+    recycled slot's list restarts its input stream at its post-prefill tick.
+    """
+    frame = engine.frame_size
+    outs: list[list[np.ndarray | None]] = [[] for _ in range(B)]
+    stream_tick = [0] * B  # per-slot conversation clock driving its input
+    prefill = None
+    tick = 0
+    while tick < total_ticks:
+        if recycle and tick == RECYCLE_AT and prefill is None:
+            engine.recycle_slot(1)
+            prefill = list(engine.prefill_steps(None))
+        pcm = np.zeros((B, frame), dtype=np.float32)
+        step_for_1 = None
+        if prefill:
+            while prefill and prefill[0].kind == "voice_end":
+                engine.voice_end(1)
+                prefill.pop(0)
+            if prefill:
+                step_for_1 = prefill.pop(0)
+            else:
+                engine.end_prefill(1)
+                prefill = None
+                stream_tick[1] = 0  # restart slot 1's input stream from its own t=0
+        for b in range(B):
+            if b == 1 and step_for_1 is not None:
+                continue  # engine overrides the user channel for prefill rows
+            pcm[b] = _slot_stream(b, stream_tick[b], frame)
+            stream_tick[b] += 1
+        frame_outs = engine.step_batch(pcm, {1: step_for_1} if step_for_1 is not None else None)
+        for b in range(B):
+            if b == 1 and step_for_1 is not None:
+                continue  # prefill tick: nothing conversational to record
+            frame_out = frame_outs[b]
+            outs[b].append(None if frame_out is None else frame_out.audio)
+        tick += 1
+    return outs
+
+
+def test_recycle_isolation_and_fresh_equivalence(engine):
+    engine.open_batch()
+    prefill_len = sum(1 for s in engine.prefill_steps(None) if s.kind != "voice_end")
+    total = RECYCLE_AT + prefill_len + LIVE_TICKS
+
+    base = _run(engine, total, recycle=False)
+
+    engine.open_batch()  # deterministic re-boot (greedy, fixed voice/persona)
+    rec = _run(engine, total, recycle=True)
+
+    # Claim 1 -- isolation: untouched slots bit-identical to the no-recycle run.
+    for b in (0, 2, 3):
+        assert len(base[b]) == total and len(rec[b]) == total
+        for i, (x, y) in enumerate(zip(base[b], rec[b])):
+            assert (x is None) == (y is None), f"slot {b} tick {i} validity diverged"
+            if x is not None:
+                assert np.array_equal(x, y), f"slot {b} tick {i} audio diverged (recycle leaked)"
+
+    # Claim 2 -- replay exactness: the recycled slot's opening live ticks match a
+    # pristine boot at numerical-noise level. A broken KV mask / staging-ring
+    # roll / Mimi reset garbles the audio from the very first frame (diffs of
+    # 0.1-1.0); a correct replay differs only by decoder-state rounding (~1e-4).
+    #
+    # Full-length bitwise identity is NOT expected, for two structural reasons:
+    # the replayed system prompt sits at a different ABSOLUTE RoPE offset (bf16
+    # rounds differently there), and the recycled Mimi ENCODER cannot restore a
+    # fresh stream's start padding (a batch-shared scalar), so from the first
+    # non-silent user frame it emits different-but-equally-valid codes for the
+    # same audio. Once any near-tie flips a token, the two equally valid
+    # trajectories legitimately diverge; coherence of the recycled slot is
+    # asserted behaviorally in test_recycled_slot_converses.
+    fresh = [a for a in base[1] if a is not None][:LIVE_TICKS]
+    reborn = [a for a in rec[1][RECYCLE_AT:] if a is not None][:LIVE_TICKS]
+    assert len(reborn) >= LIVE_TICKS // 2, "recycled slot produced too few live frames"
+    diffs = [float(np.max(np.abs(x - y))) for x, y in zip(fresh, reborn)]
+    first_flip = next((i for i, d in enumerate(diffs) if d > 5e-2), len(diffs))
+    print(
+        f"[elastic] replay exactness: {first_flip}/{len(diffs)} opening ticks at noise level "
+        f"(head diffs: {[f'{d:.1e}' for d in diffs[:6]]})"
+    )
+    assert first_flip >= 4, (
+        f"recycled slot diverged from fresh behavior at tick {first_flip} "
+        f"(diffs {diffs[: first_flip + 2]}); the system-prompt replay is broken"
+    )
+    assert max(diffs[:4]) < 2e-2, f"reborn audio not at noise level: {max(diffs[:4])}"
+
+
+def test_recycled_slot_converses(engine):
+    """Behavioral guarantee: a recycled slot holds a coherent conversation.
+
+    Feed a real spoken question into a slot that was recycled mid-flight (while
+    other slots keep running); the model's own inner-monologue text stream must
+    produce a real reply. This is the end-user property the encoder-state
+    caveats above cannot break.
+    """
+    import os
+
+    import sphn
+
+    wav_path = os.environ.get("PPLEX_QUESTION_WAV")
+    if not wav_path:
+        pytest.skip("set PPLEX_QUESTION_WAV to a spoken-question wav to run")
+    wav, _ = sphn.read(wav_path, sample_rate=24000)
+    wav = wav[0].astype(np.float32)
+    frame = engine.frame_size
+    pad = (-len(wav)) % frame
+    wav = np.concatenate([wav, np.zeros(pad, dtype=np.float32)])
+    q_frames = [wav[i : i + frame] for i in range(0, len(wav), frame)]
+
+    engine.open_batch()
+    pcm = np.zeros((B, frame), dtype=np.float32)
+    for _ in range(5):  # everyone live briefly
+        engine.step_batch(pcm)
+    engine.recycle_slot(1)
+    prefill = list(engine.prefill_steps(None))
+    while prefill:
+        while prefill and prefill[0].kind == "voice_end":
+            engine.voice_end(1)
+            prefill.pop(0)
+        if not prefill:
+            break
+        engine.step_batch(pcm, {1: prefill.pop(0)})
+    engine.end_prefill(1)
+
+    text = ""
+    total = len(q_frames) + 120  # question + ~10 s to answer
+    for i in range(total):
+        pcm = np.zeros((B, frame), dtype=np.float32)
+        if i < len(q_frames):
+            pcm[1] = q_frames[i]
+        outs = engine.step_batch(pcm)
+        if outs[1] is not None and outs[1].text:
+            text += outs[1].text
+    alpha = sum(c.isalpha() for c in text)
+    print(f"[elastic] recycled slot inner monologue: {text!r}")
+    assert alpha >= 15, f"recycled slot produced no coherent reply: {text!r}"
+
+
+def _bench(batch_sizes: list[int]) -> None:
+    import time
+
+    for bs in batch_sizes:
+        cfg = PersonaPlexConfig(batch_size=max(bs, 2))
+        eng = PersonaPlexEngine(cfg).load()
+        eng.open_batch()
+        frame = eng.frame_size
+        pcm = np.zeros((cfg.batch_size, frame), dtype=np.float32)
+        for _ in range(10):  # warm
+            eng.step_batch(pcm)
+        torch.cuda.synchronize()
+        t0 = time.monotonic()
+        n = 100
+        for _ in range(n):
+            eng.step_batch(pcm)
+        torch.cuda.synchronize()
+        ms = (time.monotonic() - t0) / n * 1000
+        mem = torch.cuda.max_memory_allocated() / 2**30
+        print(f"B={cfg.batch_size:3d}  {ms:6.1f} ms/tick  RTF/stream {ms / 80:.3f}  {mem:.1f} GiB")
+        del eng
+        torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sizes = [int(x) for x in sys.argv[2:]] if len(sys.argv) > 2 and sys.argv[1] == "--bench" else [2, 4, 8]
+    _bench(sizes)

tests/e2e/features/fullduplex/test_personaplex_server_lease.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Single-session lease teardown: a cancelled engine call must be drained.
+
+Cancelling an asyncio task blocked on a worker-thread engine call returns
+immediately while the thread keeps running (executor futures cannot be
+interrupted). The server therefore tracks the in-flight future on the session
+state and waits it out before releasing the capacity-one lease; otherwise a
+fast reconnect could call ``open()`` on the shared engine while the previous
+connection's final ``step()`` is still mutating it.
+"""
+
+import asyncio
+import contextlib
+import sys
+import threading
+import types
+
+import pytest
+
+try:
+    import sphn  # noqa: F401
+except ImportError:
+    # These lease tests never encode/decode Opus. Keep the optional serving
+    # dependency from hiding dependency-free teardown regressions.
+    sys.modules["sphn"] = types.ModuleType("sphn")
+
+from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig
+from vllm_omni.experimental.fullduplex.personaplex.serving.server import DuplexServer
+from vllm_omni.experimental.fullduplex.personaplex.session import PersonaPlexServingSessionState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.asyncio
+async def test_cancelled_engine_call_is_drained_before_lease_release():
+    server = DuplexServer(PersonaPlexConfig())  # engine never loaded: stub calls only
+    state = PersonaPlexServingSessionState()
+
+    started = threading.Event()
+    release = threading.Event()
+    finished = threading.Event()
+
+    def slow_engine_op():
+        started.set()
+        release.wait(timeout=5.0)
+        finished.set()
+        return []
+
+    task = asyncio.create_task(server._engine_call(state, slow_engine_op))
+    await asyncio.to_thread(started.wait, 5.0)
+
+    # Cancel while the worker thread is mid-call: the await returns at once,
+    # the thread keeps running, and its completion event stays tracked on the
+    # state (the asyncio future itself reads as cancelled/done here, which is
+    # exactly why the drain must not rely on it).
+    task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await task
+    assert state.inflight is not None
+    assert not state.inflight.is_set()
+
+    # The drain must genuinely wait for the worker, not just return.
+    drain = asyncio.create_task(server._drain_inflight(state))
+    await asyncio.sleep(0.05)
+    assert not drain.done()
+
+    release.set()
+    await asyncio.wait_for(drain, timeout=5.0)
+    assert finished.is_set()
+    assert state.inflight.is_set()
+
+
+@pytest.mark.asyncio
+async def test_completed_engine_call_clears_inflight():
+    server = DuplexServer(PersonaPlexConfig())
+    state = PersonaPlexServingSessionState()
+    result = await server._engine_call(state, lambda: "ok")
+    assert result == "ok"
+    assert state.inflight is None
+    await server._drain_inflight(state)  # no-op on a clean state
+
+
+@pytest.mark.asyncio
+async def test_drain_timeout_keeps_lease_until_worker_finishes():
+    server = DuplexServer(PersonaPlexConfig())
+    server._active = True
+    state = PersonaPlexServingSessionState()
+    state.inflight = threading.Event()
+
+    drained = await server._release_lease_after_drain(state, timeout=0.01)
+
+    assert drained is False
+    assert server._active is True
+
+    state.inflight.set()
+    for _ in range(100):
+        if not server._active:
+            break
+        await asyncio.sleep(0.01)
+    assert server._active is False

tests/e2e/features/fullduplex/test_personaplex_session.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lockstep-contract tests for the PersonaPlex session driver (no GPU / no moshi).
+
+A stub :class:`FrameStepper` stands in for the real Moshi backend so we can assert
+the framing + 1-frame-in / 1-frame-out invariants without weights.
+"""
+
+import numpy as np
+import pytest
+
+from vllm_omni.experimental.fullduplex.personaplex.config import FRAME_SIZE, PersonaPlexConfig
+from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput, FrameStepper
+from vllm_omni.experimental.fullduplex.personaplex.session import PersonaPlexSession
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class StubStepper:
+    """Deterministic FrameStepper: frame i -> audio filled with i, text 't{i}' on odds."""
+
+    sample_rate = 24000
+    frame_size = FRAME_SIZE
+
+    def __init__(self) -> None:
+        self.opened: tuple[str | None, str | None] | None = None
+        self.steps = 0
+
+    def open_session(self, voice_prompt=None, persona=None) -> None:
+        self.opened = (voice_prompt, persona)
+        self.steps = 0
+
+    def step(self, user_pcm):
+        assert user_pcm.shape[0] == self.frame_size
+        self.steps += 1
+        i = self.steps
+        return FrameOutput(
+            audio=np.full(self.frame_size, float(i), dtype=np.float32),
+            text=(f"t{i}" if i % 2 else None),
+        )
+
+
+def test_stub_satisfies_protocol():
+    assert isinstance(StubStepper(), FrameStepper)
+
+
+def test_feed_emits_one_output_per_complete_frame():
+    s = PersonaPlexSession(StubStepper(), PersonaPlexConfig())
+    s.open(voice_prompt="NATF2.pt", persona="be terse")
+    out = s.feed(np.zeros(FRAME_SIZE * 3, dtype=np.float32))
+    assert len(out) == 3
+    # lockstep order preserved: frame i -> audio filled with i
+    assert [int(o.audio[0]) for o in out] == [1, 2, 3]
+    assert [o.text for o in out] == ["t1", None, "t3"]
+
+
+def test_partial_frames_are_buffered_until_complete():
+    stub = StubStepper()
+    s = PersonaPlexSession(stub, PersonaPlexConfig())
+    s.open()
+    assert s.feed(np.zeros(FRAME_SIZE // 2, dtype=np.float32)) == []  # nothing yet
+    assert stub.steps == 0
+    out = s.feed(np.zeros(FRAME_SIZE, dtype=np.float32))  # 1.5 frames total -> 1 frame
+    assert len(out) == 1
+    assert stub.steps == 1
+
+
+def test_flush_pads_the_tail():
+    stub = StubStepper()
+    s = PersonaPlexSession(stub, PersonaPlexConfig())
+    s.open()
+    s.feed(np.ones(FRAME_SIZE + 10, dtype=np.float32))  # 1 frame + 10-sample tail
+    assert stub.steps == 1
+    tail = s.flush()
+    assert len(tail) == 1 and stub.steps == 2
+    assert s.flush() == []  # nothing left
+
+
+def test_run_concatenates_audio_and_joins_text():
+    s = PersonaPlexSession(StubStepper(), PersonaPlexConfig())
+    s.open()
+    audio, text = s.run(np.zeros(FRAME_SIZE * 4, dtype=np.float32))
+    assert audio.shape[0] == FRAME_SIZE * 4
+    assert text == "t1t3"  # frames 1..4, text on odds
+
+
+def test_open_forwards_voice_and_persona():
+    stub = StubStepper()
+    s = PersonaPlexSession(stub, PersonaPlexConfig())
+    s.open(voice_prompt="NATM1.pt", persona="pirate")
+    assert stub.opened == ("NATM1.pt", "pirate")
+
+
+def test_feed_before_open_raises():
+    s = PersonaPlexSession(StubStepper(), PersonaPlexConfig())
+    with pytest.raises(RuntimeError):
+        s.feed(np.zeros(FRAME_SIZE, dtype=np.float32))

tests/e2e/online_serving/personaplex_realtime_duplex.py

@@ -0,0 +1,492 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Strict PersonaPlex validation for the unified OpenAI Realtime duplex path."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import hashlib
+import json
+import math
+import uuid
+import wave
+from collections.abc import Sequence
+from pathlib import Path
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+import numpy as np
+
+from vllm_omni.experimental.fullduplex.client import (
+    RealtimeDuplexClient,
+    wait_for,
+    write_pcm16_wav,
+)
+
+SAMPLE_RATE_HZ = 24_000
+FRAME_SAMPLES = 1_920
+FRAME_PERIOD_S = FRAME_SAMPLES / SAMPLE_RATE_HZ
+
+
+def _input_identity(
+    path: Path,
+    *,
+    expected_sha256: str | None,
+) -> dict[str, str]:
+    resolved = path.resolve()
+    digest = hashlib.sha256()
+    with resolved.open("rb") as input_file:
+        for chunk in iter(lambda: input_file.read(1024 * 1024), b""):
+            digest.update(chunk)
+    actual = digest.hexdigest()
+    if expected_sha256 is not None and actual != expected_sha256.lower():
+        raise ValueError(
+            f"input WAV SHA-256 mismatch: expected={expected_sha256.lower()}, actual={actual}, path={resolved}"
+        )
+    return {"path": str(resolved), "sha256": actual}
+
+
+def _realtime_url(base_url: str, model: str, session_id: str) -> str:
+    parts = urlsplit(base_url)
+    query = dict(parse_qsl(parts.query, keep_blank_values=True))
+    query.update(duplex="1", model=model, autostart="0", session_id=session_id)
+    return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
+
+
+def _read_wav_as_float32(path: Path) -> np.ndarray:
+    with wave.open(str(path), "rb") as wav_file:
+        channels = wav_file.getnchannels()
+        sample_width = wav_file.getsampwidth()
+        source_rate = wav_file.getframerate()
+        frames = wav_file.readframes(wav_file.getnframes())
+    dtypes = {1: np.uint8, 2: np.dtype("<i2"), 4: np.dtype("<i4")}
+    if sample_width not in dtypes:
+        raise ValueError(f"unsupported WAV sample width: {sample_width}")
+    pcm = np.frombuffer(frames, dtype=dtypes[sample_width]).astype(np.float32)
+    if sample_width == 1:
+        pcm = (pcm - 128.0) / 128.0
+    else:
+        pcm /= float(1 << (sample_width * 8 - 1))
+    if channels > 1:
+        pcm = pcm.reshape(-1, channels).mean(axis=1)
+    if source_rate != SAMPLE_RATE_HZ and pcm.size:
+        output_samples = max(1, round(pcm.size * SAMPLE_RATE_HZ / source_rate))
+        pcm = np.interp(
+            np.linspace(0, pcm.size - 1, output_samples),
+            np.arange(pcm.size),
+            pcm,
+        ).astype(np.float32)
+    pcm = np.ascontiguousarray(pcm, dtype="<f4")
+    if not pcm.size or not np.isfinite(pcm).all():
+        raise ValueError("input WAV must contain finite audio")
+    return pcm
+
+
+def _events(client: RealtimeDuplexClient, event_type: str) -> list[dict[str, object]]:
+    return [event for event in client.events.events if event.get("type") == event_type]
+
+
+async def _open_session(
+    args: argparse.Namespace,
+    *,
+    session_id: str,
+    persona: str,
+    expect_error: bool = False,
+) -> tuple[RealtimeDuplexClient, dict[str, object]]:
+    client = RealtimeDuplexClient(_realtime_url(args.url, args.model, session_id))
+    await client.__aenter__()
+    await client.send(
+        {
+            "type": "session.update",
+            "session": {
+                "session_id": session_id,
+                "model": args.model,
+                "modalities": ["audio", "text"],
+                "input_audio_format": "pcm_f32le",
+                "output_audio_format": "pcm16",
+                "voice": args.voice,
+                "instructions": persona,
+                "turn_detection": None,
+                "extra_body": {"auto_response": True},
+            },
+        }
+    )
+    event_type = "error" if expect_error else "session.created"
+    await wait_for(
+        lambda: client.events.count(event_type) > 0,
+        timeout_s=args.timeout_s,
+        label=event_type,
+    )
+    return client, _events(client, event_type)[-1]
+
+
+async def _close_session(client: RealtimeDuplexClient, *, timeout_s: float) -> None:
+    await client.send({"type": "session.close"})
+    await wait_for(
+        lambda: client.events.count("session.closed") > 0,
+        timeout_s=timeout_s,
+        label="session.closed",
+    )
+    await client.__aexit__(None, None, None)
+
+
+async def _stream_frames(
+    client: RealtimeDuplexClient,
+    pcm: np.ndarray,
+    *,
+    max_frames: int | None = None,
+) -> int:
+    frame_count = math.ceil(pcm.size / FRAME_SAMPLES)
+    if max_frames is not None:
+        if max_frames < 0:
+            raise ValueError(f"max_frames must be non-negative, got {max_frames}")
+        if max_frames > 0:
+            frame_count = min(frame_count, max_frames)
+    for seq in range(frame_count):
+        frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
+        frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size))
+        frame = np.ascontiguousarray(frame, dtype="<f4")
+        await client.send(
+            {
+                "type": "input_audio_buffer.append",
+                "audio": base64.b64encode(frame).decode("ascii"),
+                "format": "pcm_f32le",
+                "sample_rate_hz": SAMPLE_RATE_HZ,
+                "duration_ms": 80,
+                "audio_end_ms": (seq + 1) * 80,
+                "is_speech": bool(np.sqrt(np.mean(np.square(frame))) >= 1e-4),
+            }
+        )
+        await asyncio.sleep(FRAME_PERIOD_S)
+    return frame_count
+
+
+def _audio_frame_stats(
+    raw: bytes,
+    *,
+    input_frames: int,
+    voiced_frame_rms_threshold: float,
+) -> dict[str, int | float]:
+    pcm = np.frombuffer(raw, dtype="<i2")
+    if pcm.size % FRAME_SAMPLES:
+        raise AssertionError(
+            "PersonaPlex output must contain whole 80 ms codec frames: "
+            f"samples={pcm.size}, samples_per_frame={FRAME_SAMPLES}"
+        )
+    output_frames = pcm.size // FRAME_SAMPLES
+    frames = pcm.astype(np.float32).reshape(output_frames, FRAME_SAMPLES) / 32768.0
+    voiced_frames = int(np.count_nonzero(np.sqrt(np.mean(np.square(frames), axis=1)) >= voiced_frame_rms_threshold))
+    return {
+        "output_frames": int(output_frames),
+        "voiced_frames": voiced_frames,
+        "silent_frames": int(output_frames - voiced_frames),
+        "frame_deficit": int(input_frames - output_frames),
+        "frame_coverage_ratio": float(output_frames / input_frames),
+    }
+
+
+def _response_ids(client: RealtimeDuplexClient) -> set[str]:
+    return {
+        str(event["response_id"])
+        for event in _events(client, "response.audio.delta")
+        if isinstance(event.get("response_id"), str)
+    }
+
+
+def _session_result(
+    client: RealtimeDuplexClient,
+    *,
+    input_frames: int,
+    args: argparse.Namespace,
+    minimum_chunks: int,
+) -> tuple[bytes, set[str], dict[str, object]]:
+    chunks = sum(len(items) for items in client.events.response_audio.values())
+    raw = client.events.audio_bytes()
+    if chunks < minimum_chunks or not raw or len(raw) % 2:
+        raise AssertionError(f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}")
+    pcm = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
+    rms = float(np.sqrt(np.mean(np.square(pcm))))
+    # This low floor only catches corrupt, non-finite, or effectively zero PCM.
+    # Audible speech is checked per frame below with the higher 1e-3 default.
+    if not np.isfinite(pcm).all() or rms <= 1e-5:
+        raise AssertionError(f"output audio is non-finite or silent: samples={pcm.size}, rms={rms}")
+    rates = {
+        int(event["sample_rate_hz"])
+        for event in _events(client, "response.audio.delta")
+        if isinstance(event.get("sample_rate_hz"), int)
+    }
+    if rates != {SAMPLE_RATE_HZ}:
+        raise AssertionError(f"unexpected output sample rates: {sorted(rates)}")
+    stats = _audio_frame_stats(
+        raw,
+        input_frames=input_frames,
+        voiced_frame_rms_threshold=args.voiced_frame_rms_threshold,
+    )
+    if not 0 <= int(stats["frame_deficit"]) <= args.max_frame_deficit:
+        raise AssertionError(f"invalid synchronous frame accounting: input={input_frames}, stats={stats}")
+    min_voiced_frames = int(args.min_voiced_frames)
+    if min_voiced_frames < 1:
+        raise ValueError(f"min_voiced_frames must be positive, got {min_voiced_frames}")
+    if int(stats["voiced_frames"]) < min_voiced_frames:
+        raise AssertionError(
+            f"output contained insufficient audible speech: minimum_voiced_frames={min_voiced_frames}, stats={stats}"
+        )
+    runtime_metadata = [
+        metadata["vllm_omni"]
+        for event in _events(client, "response.audio.delta")
+        if isinstance((metadata := event.get("metadata")), dict) and isinstance(metadata.get("vllm_omni"), dict)
+    ]
+    if not any(
+        item.get("runtime_impl") == "scheduler_data_plane"
+        and item.get("uses_model_runner_scheduler") is True
+        and item.get("runner_kv_backed") is True
+        for item in runtime_metadata
+    ):
+        raise AssertionError(f"audio events did not prove the scheduler data plane: {runtime_metadata}")
+    response_ids = _response_ids(client)
+    if len(response_ids) != 1:
+        raise AssertionError(f"continuous stream used multiple responses: {sorted(response_ids)}")
+    return (
+        raw,
+        response_ids,
+        {
+            **stats,
+            "audio_chunks": chunks,
+            "audio_samples": len(raw) // 2,
+            "audio_rms": rms,
+            "response_ids": sorted(response_ids),
+        },
+    )
+
+
+def _capabilities(created: dict[str, object]) -> dict[str, object]:
+    session = created.get("session")
+    capabilities = session.get("capabilities") if isinstance(session, dict) else None
+    if not isinstance(capabilities, dict):
+        raise AssertionError(f"session.created omitted capabilities: {created}")
+    return capabilities
+
+
+def _save(output_dir: Path, name: str, client: RealtimeDuplexClient, raw: bytes) -> None:
+    (output_dir / f"{name}-events.jsonl").write_text(
+        "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events),
+        encoding="utf-8",
+    )
+    write_pcm16_wav(output_dir / f"{name}-output.wav", raw, sample_rate_hz=SAMPLE_RATE_HZ)
+
+
+async def run(args: argparse.Namespace) -> dict[str, object]:
+    input_path = Path(args.input_wav)
+    input_identity = _input_identity(
+        input_path,
+        expected_sha256=args.expected_input_sha256,
+    )
+    pcm = _read_wav_as_float32(input_path)
+    if args.tail_s > 0:
+        pcm = np.concatenate([pcm, np.zeros(round(args.tail_s * SAMPLE_RATE_HZ), dtype=np.float32)])
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    ids = {name: f"personaplex-{name}-{uuid.uuid4().hex}" for name in ("primary", "secondary", "replacement")}
+    primary, created = await _open_session(args, session_id=ids["primary"], persona=args.persona)
+    secondary, secondary_created = await _open_session(
+        args,
+        session_id=ids["secondary"],
+        persona=args.secondary_persona,
+    )
+    capabilities = _capabilities(created)
+    if _capabilities(secondary_created) != capabilities:
+        raise AssertionError("concurrent sessions returned different capabilities")
+    expected_capabilities = {
+        "implementation_level": "model_native_duplex",
+        "chunk_period_ms": 80,
+        "supports_multi_session": True,
+        "supports_multi_session_same_replica": True,
+        "supports_barge_in": False,
+    }
+    if any(capabilities.get(key) != value for key, value in expected_capabilities.items()):
+        raise AssertionError(f"unexpected PersonaPlex capabilities: {capabilities}")
+
+    overflow, error = await _open_session(
+        args,
+        session_id=f"personaplex-overflow-{uuid.uuid4().hex}",
+        persona=args.persona,
+        expect_error=True,
+    )
+    error_body = error.get("error")
+    overflow_code = error_body.get("code") if isinstance(error_body, dict) else error.get("code")
+    await overflow.__aexit__(None, None, None)
+    if overflow_code != "resource_exhausted":
+        raise AssertionError(f"overflow session was not rejected: {error}")
+
+    primary_frames, secondary_frames = await asyncio.gather(
+        _stream_frames(primary, pcm),
+        _stream_frames(secondary, pcm),
+    )
+    await asyncio.gather(
+        wait_for(
+            lambda: len(primary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit,
+            timeout_s=args.timeout_s,
+            label="primary frame coverage",
+        ),
+        wait_for(
+            lambda: (
+                len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit
+            ),
+            timeout_s=args.timeout_s,
+            label="secondary frame coverage",
+        ),
+    )
+    await asyncio.sleep(args.drain_s)
+    primary_audio, primary_response_ids, primary_stats = _session_result(
+        primary,
+        input_frames=primary_frames,
+        args=args,
+        minimum_chunks=args.minimum_audio_chunks,
+    )
+    secondary_audio, secondary_response_ids, _ = _session_result(
+        secondary,
+        input_frames=secondary_frames,
+        args=args,
+        minimum_chunks=args.minimum_audio_chunks,
+    )
+    if primary_response_ids & secondary_response_ids:
+        raise AssertionError("concurrent sessions shared a response id")
+    _save(output_dir, "primary", primary, primary_audio)
+    await _close_session(primary, timeout_s=args.timeout_s)
+
+    replacement, replacement_created = await _open_session(
+        args,
+        session_id=ids["replacement"],
+        persona=args.replacement_persona,
+    )
+    if _capabilities(replacement_created) != capabilities:
+        raise AssertionError("replacement session returned different capabilities")
+    continuation_frames, replacement_frames = await asyncio.gather(
+        _stream_frames(secondary, pcm, max_frames=args.replacement_frames),
+        _stream_frames(replacement, pcm, max_frames=args.replacement_frames),
+    )
+    secondary_total_frames = secondary_frames + continuation_frames
+    await asyncio.gather(
+        wait_for(
+            lambda: (
+                len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES)
+                >= secondary_total_frames - args.max_frame_deficit
+            ),
+            timeout_s=args.timeout_s,
+            label="survivor frame coverage after slot recycle",
+        ),
+        wait_for(
+            lambda: (
+                len(replacement.events.audio_bytes()) // (2 * FRAME_SAMPLES)
+                >= replacement_frames - args.max_frame_deficit
+            ),
+            timeout_s=args.timeout_s,
+            label="replacement frame coverage",
+        ),
+    )
+    await asyncio.sleep(args.drain_s)
+    secondary_audio, secondary_after_ids, secondary_stats = _session_result(
+        secondary,
+        input_frames=secondary_total_frames,
+        args=args,
+        minimum_chunks=args.minimum_audio_chunks + 1,
+    )
+    replacement_audio, replacement_response_ids, replacement_stats = _session_result(
+        replacement,
+        input_frames=replacement_frames,
+        args=args,
+        minimum_chunks=1,
+    )
+    if secondary_after_ids != secondary_response_ids:
+        raise AssertionError("survivor response changed while another session slot was recycled")
+    if replacement_response_ids & (primary_response_ids | secondary_response_ids):
+        raise AssertionError("replacement session reused another session's response id")
+    _save(output_dir, "secondary", secondary, secondary_audio)
+    _save(output_dir, "replacement", replacement, replacement_audio)
+    await asyncio.gather(
+        _close_session(secondary, timeout_s=args.timeout_s),
+        _close_session(replacement, timeout_s=args.timeout_s),
+    )
+
+    errors = primary.events.errors() + secondary.events.errors() + replacement.events.errors()
+    result = {
+        "ok": not errors,
+        "model": args.model,
+        "input": {**input_identity, "tail_s": args.tail_s},
+        "capabilities": capabilities,
+        "overflow_error_code": overflow_code,
+        "primary": {"session_id": ids["primary"], "input_frames": primary_frames, **primary_stats},
+        "secondary": {
+            "session_id": ids["secondary"],
+            "input_frames": secondary_total_frames,
+            "continuation_frames_after_recycle": continuation_frames,
+            **secondary_stats,
+        },
+        "replacement": {
+            "session_id": ids["replacement"],
+            "input_frames": replacement_frames,
+            **replacement_stats,
+        },
+        "errors": errors,
+        "output_dir": str(output_dir),
+    }
+    (output_dir / "result.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
+    if errors:
+        raise AssertionError(f"Realtime sessions emitted errors: {errors}")
+    return result
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1")
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--input-wav", required=True)
+    parser.add_argument(
+        "--expected-input-sha256",
+        help="Fail before opening sessions when the input WAV does not match this SHA-256.",
+    )
+    parser.add_argument("--output-dir", default="/tmp/personaplex-realtime-duplex")
+    parser.add_argument("--voice", default="NATF2.pt")
+    parser.add_argument("--persona", default="You are a concise and helpful assistant.")
+    parser.add_argument("--secondary-persona", default="You are a calm and precise assistant.")
+    parser.add_argument("--replacement-persona", default="You are a cheerful travel assistant.")
+    parser.add_argument("--tail-s", type=float, default=2.0)
+    parser.add_argument("--drain-s", type=float, default=2.0)
+    parser.add_argument("--timeout-s", type=float, default=180.0)
+    parser.add_argument(
+        "--max-frame-deficit",
+        type=int,
+        default=4,
+        help=(
+            "Maximum unflushed tail frames. PersonaPlex Code2Wav emits five-frame "
+            "chunks, so a session cutoff can leave at most four pending frames."
+        ),
+    )
+    parser.add_argument("--voiced-frame-rms-threshold", type=float, default=1e-3)
+    parser.add_argument(
+        "--min-voiced-frames",
+        type=int,
+        default=5,
+        help=(
+            "Require at least one normal five-frame Code2Wav chunk (400 ms) above "
+            "the per-frame RMS threshold; this detects inaudible output, not speech quality."
+        ),
+    )
+    parser.add_argument("--minimum-audio-chunks", type=int, default=2)
+    parser.add_argument(
+        "--replacement-frames",
+        type=int,
+        default=0,
+        help="Frames sent after slot reuse; 0 replays the full input so the replacement must produce audible output.",
+    )
+    return parser.parse_args(argv)
+
+
+def main() -> None:
+    print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()

tests/model_executor/stage_input_processors/test_personaplex.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.personaplex.personaplex_talker import (
+    PersonaPlexTalkerForConditionalGeneration,
+)
+from vllm_omni.model_executor.stage_input_processors.personaplex import (
+    talker2code2wav_async_chunk,
+    talker2code2wav_full_payload,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _raw_agent_codes() -> torch.Tensor:
+    return torch.arange(16, dtype=torch.long).reshape(2, 8)
+
+
+@pytest.mark.parametrize("source", ["pooling_output", "additional_information_cpu"])
+def test_full_payload_accepts_worker_and_cached_payload_sources(source: str) -> None:
+    audio = _raw_agent_codes()
+    payload = talker2code2wav_full_payload(
+        pooling_output={"codes": {"audio": audio}} if source == "pooling_output" else None,
+        request=(
+            SimpleNamespace()
+            if source == "pooling_output"
+            else SimpleNamespace(additional_information_cpu={"codes": {"audio": audio}})
+        ),
+    )
+
+    expected = torch.cat([audio[:-1, :1], audio[1:, 1:]], dim=1).reshape(-1)
+    assert torch.equal(payload.codes.audio, expected)
+
+
+def test_async_chunk_keeps_delay_tail_across_resumable_segments() -> None:
+    manager = SimpleNamespace(
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "initial_codec_chunk_frames": 1,
+                    "codec_chunk_frames": 5,
+                }
+            }
+        )
+    )
+    request = SimpleNamespace(
+        request_id="req",
+        external_req_id="req",
+        resumable=True,
+        is_finished=lambda: True,
+        additional_information=None,
+    )
+    first_frame = torch.arange(8, dtype=torch.long).reshape(1, 8)
+    second_frame = torch.arange(8, 16, dtype=torch.long).reshape(1, 8)
+
+    request.additional_information = {"codes": {"audio": first_frame}}
+    first = talker2code2wav_async_chunk(
+        manager,
+        multimodal_output=None,
+        request=request,
+        is_finished=True,
+    )
+    request.additional_information = {"codes": {"audio": second_frame}}
+    second = talker2code2wav_async_chunk(
+        manager,
+        multimodal_output=None,
+        request=request,
+        is_finished=True,
+    )
+
+    assert first is not None
+    assert first.codes is None
+    assert first.meta is not None
+    assert first.meta.finished.item() is False
+    assert first.meta.is_segment_finished.item() is False
+    expected = torch.cat([first_frame[:, :1], second_frame[:, 1:]], dim=1).reshape(-1)
+    assert torch.equal(second.codes.audio, expected)
+    assert manager.request_payload["req"]["personaplex_frames"][0].equal(second_frame.reshape(-1))
+
+
+def test_post_sample_talker_mtp_uses_current_temporal_state() -> None:
+    received: dict[str, torch.Tensor] = {}
+    recorded: list[tuple[str, torch.Tensor, torch.Tensor]] = []
+
+    def depformer(
+        text_token: torch.Tensor,
+        hidden: torch.Tensor,
+        *,
+        audio_tokens: torch.Tensor | None = None,
+        audio_provided: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        received["text_token"] = text_token
+        received["hidden"] = hidden
+        received["audio_tokens"] = audio_tokens
+        received["audio_provided"] = audio_provided
+        return torch.arange(16, dtype=torch.long).reshape(1, 16)
+
+    model = SimpleNamespace(
+        _dtype=torch.float32,
+        depformer=depformer,
+        _duplex_stage0_runtime=lambda: SimpleNamespace(
+            record_sample=lambda *, request_id, text_token, agent_codes: recorded.append(
+                (request_id, text_token.clone(), agent_codes.clone())
+            )
+        ),
+    )
+    method = getattr(PersonaPlexTalkerForConditionalGeneration, "post_sample_talker_mtp", None)
+    assert callable(method)
+
+    codes = method(
+        model,
+        input_ids=torch.tensor([101]),
+        hidden_states=torch.arange(4, dtype=torch.float32).reshape(1, 4),
+        req_ids=["r1"],
+        req_infos=[
+            {
+                "duplex": {"data_plane": True},
+                "pplex_depformer_audio_tokens": torch.arange(16),
+                "pplex_depformer_audio_provided": torch.arange(16) > 0,
+            }
+        ],
+    )
+
+    assert codes.shape == (1, 16)
+    assert received["text_token"].tolist() == [101]
+    assert received["hidden"].shape == (1, 1, 4)
+    assert torch.equal(received["audio_tokens"], torch.arange(16).reshape(1, 16))
+    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])

tests/worker/test_gpu_ar_model_runner.py

@@ -1,4 +1,5 @@
 import asyncio
+import copy
 from collections.abc import Sequence
 from types import SimpleNamespace
 
@@ -32,6 +33,160 @@ def _make_runner(engine_output_type: str | None, downstream_req_ids: set[str]) -
     return runner
 
 
+def _mtp_runner(*, async_scheduling: bool, buffers: dict[str, dict]) -> tuple[GPUARModelRunner, dict[str, object]]:
+    received: dict[str, object] = {}
+
+    def post_sample_talker_mtp(*, input_ids, hidden_states, req_ids, req_infos):
+        received.update(
+            input_ids=input_ids.clone(),
+            hidden_states=hidden_states.clone(),
+            req_ids=list(req_ids),
+            req_infos=copy.deepcopy(req_infos),
+        )
+        return torch.tensor([[11, 12, 13]], dtype=torch.long)
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.use_async_scheduling = async_scheduling
+    runner.model = SimpleNamespace(post_sample_talker_mtp=post_sample_talker_mtp)
+    runner.requests = {req_id: SimpleNamespace() for req_id in buffers}
+    runner.model_intermediate_buffer = buffers
+    return runner, received
+
+
+def test_post_sample_talker_mtp_uses_current_sample_and_hidden() -> None:
+    runner, received = _mtp_runner(
+        async_scheduling=False,
+        buffers={
+            "ready": {"duplex": {"data_plane": True}},
+            "partial-prefill": {
+                "duplex": {"data_plane": True},
+                "codes": {"audio": torch.tensor([[99]])},
+            },
+            "offline": {},
+        },
+    )
+
+    multimodal = GPUARModelRunner._run_post_sample_talker_mtp(
+        runner,
+        req_ids=["ready", "partial-prefill", "offline"],
+        valid_sampled_token_ids=[[101], [], [102]],
+        sampled_token_ids=torch.tensor([[101], [0], [102]], dtype=torch.long),
+        invalid_req_indices=[],
+        sample_hidden_states=torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
+        ),
+        multimodal_outputs={
+            "codes": {
+                "audio": torch.tensor(
+                    [[70, 71, 72], [80, 81, 82], [90, 91, 92]],
+                    dtype=torch.long,
+                )
+            },
+            "meta": {"source": "temporal"},
+        },
+    )
+
+    assert received["req_ids"] == ["ready"]
+    assert received["input_ids"].tolist() == [101]
+    assert received["hidden_states"].tolist() == [[1.0, 2.0]]
+    assert received["req_infos"] == [{"duplex": {"data_plane": True}}]
+    assert runner.model_intermediate_buffer["ready"]["codes"]["audio"].tolist() == [[11, 12, 13]]
+    # A partial-prefill row has no sampled token yet. It must not replay the
+    # prior frame's codes into the current connector payload.
+    audio = multimodal["codes"]["audio"]
+    assert audio[0].tolist() == [[11, 12, 13]]
+    assert audio[1].numel() == 0
+    # Mixed batches must retain outputs produced by non-duplex rows.
+    assert audio[2].tolist() == [[90, 91, 92]]
+    assert multimodal["meta"] == {"source": "temporal"}
+
+
+def test_post_sample_talker_mtp_uses_gpu_token_with_async_scheduling() -> None:
+    runner, received = _mtp_runner(
+        async_scheduling=True,
+        buffers={
+            "ready": {"duplex": {"data_plane": True}},
+            "discarded-prefill": {"duplex": {"data_plane": True}},
+        },
+    )
+
+    multimodal = GPUARModelRunner._run_post_sample_talker_mtp(
+        runner,
+        req_ids=["ready", "discarded-prefill"],
+        valid_sampled_token_ids=[],
+        sampled_token_ids=torch.tensor([[101], [102]], dtype=torch.long),
+        invalid_req_indices=[1],
+        sample_hidden_states=torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0]],
+        ),
+        multimodal_outputs=None,
+    )
+
+    assert received["req_ids"] == ["ready"]
+    assert received["input_ids"].tolist() == [101]
+    assert received["hidden_states"].tolist() == [[1.0, 2.0]]
+    assert received["req_infos"] == [{"duplex": {"data_plane": True}}]
+    assert multimodal["codes"]["audio"][0].tolist() == [[11, 12, 13]]
+    assert multimodal["codes"]["audio"][1].numel() == 0
+
+
+def test_post_sample_talker_mtp_skips_shape_validation_without_duplex_rows() -> None:
+    runner = object.__new__(GPUARModelRunner)
+    runner.use_async_scheduling = False
+    runner.model = SimpleNamespace(
+        post_sample_talker_mtp=lambda **_: pytest.fail("hook must not run"),
+    )
+    runner.model_intermediate_buffer = {"offline": {}}
+    multimodal = {"codes": {"audio": torch.tensor([[7, 8]], dtype=torch.long)}}
+
+    result = GPUARModelRunner._run_post_sample_talker_mtp(
+        runner,
+        req_ids=["offline"],
+        valid_sampled_token_ids=[[101, 102]],
+        sampled_token_ids=torch.tensor([[101, 102]], dtype=torch.long),
+        invalid_req_indices=[],
+        sample_hidden_states=torch.tensor([[1.0, 2.0]]),
+        multimodal_outputs=multimodal,
+    )
+
+    assert result is multimodal
+
+
+@pytest.mark.parametrize(
+    ("sampled_token_ids", "error_match"),
+    [
+        (torch.tensor([101], dtype=torch.long), "rows do not match requests"),
+        (
+            torch.tensor([[101, 102]], dtype=torch.long),
+            "does not support speculative token rows",
+        ),
+    ],
+)
+def test_post_sample_talker_mtp_rejects_invalid_selected_token_shape(
+    sampled_token_ids: torch.Tensor,
+    error_match: str,
+) -> None:
+    runner = object.__new__(GPUARModelRunner)
+    runner.use_async_scheduling = False
+    runner.model = SimpleNamespace(
+        post_sample_talker_mtp=lambda **_: pytest.fail("hook must not run"),
+    )
+    runner.model_intermediate_buffer = {
+        "duplex": {"duplex": {"data_plane": True}},
+    }
+
+    with pytest.raises(ValueError, match=error_match):
+        GPUARModelRunner._run_post_sample_talker_mtp(
+            runner,
+            req_ids=["duplex"],
+            valid_sampled_token_ids=[[101]],
+            sampled_token_ids=sampled_token_ids,
+            invalid_req_indices=[],
+            sample_hidden_states=torch.tensor([[1.0, 2.0]]),
+            multimodal_outputs=None,
+        )
+
+
 def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatch):
     requested = {"temperature": 0.7, "top_p": 0.8, "top_k": 17}
     request = OpenAICreateSpeechRequest.model_validate(

vllm_omni/config/model.py

@@ -122,6 +122,7 @@ class OmniModelConfig(ModelConfig):
     retains_state_across_chunks: bool = False
     # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin.
     active_stream_window: int = 0
+    duplex_max_sessions: int = 1
     model_stage: str = "thinker"
     model_arch: str | None = None
     worker_type: str | None = None

vllm_omni/config/omni_config.py

@@ -258,6 +258,7 @@ class OmniStageModelConfig:
     """Per-stage model behavior."""
 
     active_stream_window: int = Field(default=0, ge=0)
+    duplex_max_sessions: int = Field(default=1, ge=1)
     enable_sleep_mode: bool = False
     default_sampling_params: dict[str, Any] | None = None
     subtalker_sampling_params: dict[str, Any] | None = None
@@ -968,7 +969,12 @@ def _build_common_stage_config_kwargs(
     return (
         {
             "stage_pipeline_config": topology,
-            "model_config": _build_model_config(topology, stage_deploy, engine.model),
+            "model_config": _build_model_config(
+                topology,
+                stage_deploy,
+                engine.model,
+                duplex_max_sessions=(deploy.duplex_session.max_sessions if deploy.session_mode == "duplex" else 1),
+            ),
             "load_config": _build_load_config(engine.load),
             "cache_config": _build_cache_config(deploy, engine.cache),
             "scheduler_config": _build_scheduler_config(deploy, engine.scheduler),
@@ -1121,6 +1127,8 @@ def _build_model_config(
     topology: StagePipelineConfig,
     stage_deploy: StageDeployConfig | None,
     engine: _ModelEngineOverrides,
+    *,
+    duplex_max_sessions: int,
 ) -> OmniStageModelConfig:
     default_sampling_params = _stage_sampling_params(stage_deploy, topology)
     kwargs = _config_kwargs(engine)
@@ -1132,6 +1140,7 @@ def _build_model_config(
         kwargs["tokenizer_subdir"] = topology.tokenizer_subdir
     return OmniStageModelConfig(
         default_sampling_params=default_sampling_params,
+        duplex_max_sessions=duplex_max_sessions,
         **kwargs,
     )
 

vllm_omni/config/pipeline_registry.py

@@ -86,6 +86,7 @@ from vllm_omni.model_executor.models.moss_tts.pipeline import (
 )
 from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE
 from vllm_omni.model_executor.models.omnivoice.pipeline import OMNIVOICE_PIPELINE
+from vllm_omni.model_executor.models.personaplex.pipeline import PERSONAPLEX_PIPELINE
 from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
     QWEN2_5_OMNI_PIPELINE,
     QWEN2_5_OMNI_THINKER_ONLY_PIPELINE,
@@ -113,6 +114,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "aura_omni": AURA_OMNI_PIPELINE,
     "qwen2_5_omni": QWEN2_5_OMNI_PIPELINE,
     "qwen2_5_omni_thinker_only": QWEN2_5_OMNI_THINKER_ONLY_PIPELINE,
+    "personaplex": PERSONAPLEX_PIPELINE,
     "qwen3_omni_moe": resolve_qwen3_omni_pipeline,
     "qwen3_tts": QWEN3_TTS_PIPELINE,
     "step_audio_2": STEP_AUDIO2_PIPELINE,

vllm_omni/config/stage_config.py

@@ -857,6 +857,11 @@ def _build_engine_args(
     # Materialize the resolved pipeline-wide async_chunk value into every
     # stage so explicit False overrides do not get lost downstream.
     engine_args["async_chunk"] = bool(deploy.async_chunk)
+    if deploy.session_mode == "duplex":
+        # The engine admission limit is also the authoritative capacity for
+        # model-owned streaming state. Propagate it to every stage instead of
+        # making individual models duplicate the value in connector extras.
+        engine_args["duplex_max_sessions"] = deploy.duplex_session.max_sessions
     if ps.omni_kv_config:
         engine_args["omni_kv_config"] = dict(ps.omni_kv_config)
     return engine_args

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -506,7 +506,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Finish async_chunk requests that schedule() collected because their
         # upstream completed with no remaining codec tokens.
         for request in self._pending_finish_reqs:
-            if request.is_finished():
+            if request in stopped_running_reqs or request.is_finished():
                 continue
             request.status = RequestStatus.FINISHED_STOPPED
             finish_reason = request.get_finished_reason()

vllm_omni/deploy/personaplex.yaml

@@ -0,0 +1,90 @@
+# PersonaPlex deploy: talker -> code2wav (Mimi) via shared-memory chunk streaming.
+#
+# PersonaPlex is a Moshi finetune (full-duplex S2S), served offline as a 2-stage
+# audio->audio pipeline that reuses the Qwen3-TTS staged deploy layout. Fields
+# omitted from a stage fall back to StageDeployConfig defaults (see
+# vllm_omni/config/stage_config.py).
+#
+# Notable choices for this model:
+#   * Stage 0 (talker) keeps a small max_num_batched_tokens for async-chunk
+#     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.
+pipeline: personaplex
+session_mode: duplex
+duplex_session:
+  max_sessions: 2
+
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      # Mimi runs at 12.5 Hz (80 ms / frame). Stream a few frames per chunk.
+      codec_chunk_frames: 5
+      codec_left_context_frames: 10
+      # Emit only the first audio chunk early, then return to codec_chunk_frames.
+      initial_codec_chunk_frames: 1
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 10
+    gpu_memory_utilization: 0.4
+    trust_remote_code: true
+    skip_tokenizer_init: true
+    # Prompt IDs are placeholders; voice/persona/audio semantics live in
+    # request-specific embeddings, so token-ID prefix caching is invalid.
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    # Temporal transformer sliding_window is 3000 frames; keep model len within.
+    max_model_len: 3000
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      # Greedy talker decoding: the HeliumModel temporal matches Moshi within cos
+      # 0.999 (not bit-exact), and chaotic sampling amplifies that into divergence.
+      # Greedy is reliably coherent; raise temperature only once the temporal is
+      # bit-exact.
+      temperature: 0.0
+      top_k: 1
+      max_tokens: 3000
+      seed: 42
+      repetition_penalty: 1.0
+
+  - stage_id: 1
+    max_num_seqs: 10
+    gpu_memory_utilization: 0.3
+    enforce_eager: true
+    trust_remote_code: true
+    skip_tokenizer_init: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    # Codec prefill length is num_codebooks * num_frames; keep ample headroom.
+    max_num_batched_tokens: 65536
+    max_model_len: 65536
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0
+
+platforms:
+  npu:
+    stages:
+      - stage_id: 0
+        enforce_eager: false

vllm_omni/engine/arg_utils.py

@@ -23,6 +23,11 @@ _ARCH_TO_MODEL_TYPE: dict[str, str] = {
     "IndexTTS2S2MelDecoder": "indextts2",
     "IndexTTS2TalkerForConditionalGeneration": "indextts2",
     "OmniVoiceModel": "omnivoice",
+    # PersonaPlex ships an empty config.json, so create_model_config() must patch
+    # model_type=personaplex from the arch name or the staged pipeline can't load
+    # the HF config without manual overrides.
+    "PersonaPlexTalkerForConditionalGeneration": "personaplex",
+    "PersonaPlexCode2Wav": "personaplex",
     "VoxCPM2TalkerForConditionalGeneration": "voxcpm2",
 }
 
@@ -48,6 +53,9 @@ def _register_omni_hf_configs() -> None:
             MossTTSLocalConfig,
             MossTTSRealtimeConfig,
         )
+        from vllm_omni.model_executor.models.personaplex.configuration_personaplex import (
+            PersonaPlexConfig,
+        )
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
             Qwen3TTSConfig,
         )
@@ -74,6 +82,7 @@ def _register_omni_hf_configs() -> None:
         ("moss_tts_local", MossTTSLocalConfig),
         ("moss_tts_realtime", MossTTSRealtimeConfig),
         ("qwen3_tts", Qwen3TTSConfig),
+        ("personaplex", PersonaPlexConfig),
         ("cosyvoice3", CosyVoice3Config),
         ("glm_tts", GLMTTSConfig),
         ("omnivoice", OmniVoiceConfig),
@@ -164,6 +173,9 @@ class OmniEngineArgs(EngineArgs):
     # Must be declared here so engine_args dict propagation does not silently
     # drop the value when constructing OmniEngineArgs from kwargs.
     active_stream_window: int = 0
+    # Engine admission is the single source of truth for per-stage pools of
+    # model-owned streaming state (codec, decoder, and similar resources).
+    duplex_max_sessions: int = 1
     omni_kv_config: dict | None = None
     quantization_config: Any | None = None
     force_cutlass_fp8: bool | None = None
@@ -344,6 +356,7 @@ class OmniEngineArgs(EngineArgs):
             async_chunk=self.async_chunk,
             retains_state_across_chunks=self.retains_state_across_chunks,
             active_stream_window=self.active_stream_window,
+            duplex_max_sessions=self.duplex_max_sessions,
             model_stage=self.model_stage,
             model_arch=self.model_arch,
             worker_type=self.worker_type,

vllm_omni/experimental/fullduplex/README.md

@@ -1,10 +1,12 @@
 # Experimental Full-Duplex Runtime
 
-This package contains two experimental integrations:
+This package contains three experimental integrations:
 
 - the existing JoyVL framework and example integration;
 - the MiniCPM-o 4.5 native audio path used by `/v1/duplex` and
-  `/v1/realtime?duplex=1`.
+  `/v1/realtime?duplex=1`;
+- the PersonaPlex lockstep speech-to-speech path (browser demo + batched
+  serving, see `recipes/NVIDIA/PersonaPlex.md`).
 
 For the MiniCPM active runtime path, lifecycle invariants, capability boundary,
 Realtime response contract, and validation scope, see [`DESIGN.md`](DESIGN.md).
@@ -15,14 +17,39 @@ To run JoyVL, see
 ## Package boundaries
 
 ```text
-core/       existing JoyVL framework and experimental compatibility exports
-engine/     AsyncOmni/orchestrator scheduler data-plane adapter
-openai/     WebSocket transport, Realtime projection, and audio codecs
-minicpmo45/ MiniCPM input framing, policy, compatibility, and Stage0 state
-joyvl/      JoyVL model-specific integration
+core/        model-agnostic duplex contracts (adapter, session, turn runtime)
+engine/      AsyncOmni/orchestrator scheduler data-plane adapter
+openai/      WebSocket transport, Realtime projection, and audio codecs
+minicpmo45/  MiniCPM input framing, policy, compatibility, and Stage0 state
+joyvl/       JoyVL model-specific integration
+personaplex/ PersonaPlex lockstep engine, model-owned runtime, and serving
 ```
 
-MiniCPM does not run through the removed experimental `core.DuplexRuntime`
+MiniCPM does not run through the experimental `core.DuplexRuntime`
 facade. Its active path uses the `openai` session controller, the experimental
 engine contracts, the standard scheduler/model runners, and an injected
 MiniCPM-specific runtime extension from `minicpmo45/runtime.py`.
+
+`personaplex/` is a Moshi-class, pure-lockstep speech-to-speech model on the
+`core/` contracts. It keeps `core/` untouched: the lockstep lifecycle (ONE
+eternal, frame-clocked response that drains on close, instead of the turn-style
+start/cancel-per-trigger one) is model policy and lives in the model package as
+`PersonaPlexDuplexRuntime`, mirroring the model-owned runtime shape of the
+MiniCPM-o duplex work. Its runnable serving path is `personaplex/serving/`
+(single-session lease or `--batch-size` elastic slots) over
+`personaplex/session.py` (lockstep driver).
+
+## Adding a full-duplex model on the core contracts
+
+The seam is `core.DuplexAdapter`. `core/` owns the session lifecycle,
+epoch-based barge-in, playback cursor, and the event protocol; you implement
+only model policy.
+
+1. Create a sibling package `vllm_omni/experimental/fullduplex/<model>/`; keep
+   model-specific code there and do not touch `core/`.
+2. Implement one `DuplexAdapter` (`capabilities` / `on_input` / `respond`; the
+   rest have defaults). Turn-based models run through `core.DuplexRuntime`
+   unchanged; a model needing a different lifecycle carries its own runtime in
+   its package (see `personaplex/adapter.py::PersonaPlexDuplexRuntime`).
+3. Promote a helper from a model package up into `core/` only once a second
+   model actually needs it.

vllm_omni/experimental/fullduplex/personaplex/DESIGN.md

@@ -0,0 +1,259 @@
+# PersonaPlex Unified Full-Duplex Design
+
+## Status and target
+
+This design adapts PersonaPlex PR #4771 to the unified full-duplex runtime on
+top of:
+
+- vLLM-Omni `origin/main`: `67c54777bb22e9e7e08fdf7c47a64f06b566fc47`
+- PersonaPlex PR head: `477fb7c225f0c06991bc8aa55eadbd908ba282e4`
+
+The target is the engine-native path:
+
+```text
+/v1/realtime?duplex=1
+  -> OpenAI Realtime session actor
+  -> DuplexRequestClient
+  -> AsyncOmni correlated RPC
+  -> DuplexControlPlane
+  -> resumable Stage 0 request
+  -> PersonaPlex Talker
+  -> streaming PersonaPlex Code2Wav
+  -> PersonaPlex data-plane projector
+  -> response.audio.delta + response.audio_transcript.delta
+```
+
+The standalone `/api/chat` and `/v1/audio/duplex` server remains a compatibility
+surface. It is not evidence that the unified engine path works.
+
+## Why configuration-only enablement is invalid
+
+The current staged pipeline is explicitly turn based. Its Talker reads
+`pplex_user_codes`, `pplex_prefill_text`, and `pplex_silence_codes`, but no
+production staged input path writes those fields. The voice prompt, persona
+prefill, and streaming Mimi state live only in the standalone
+`PersonaPlexEngine`.
+
+Setting only the following fields would therefore advertise an endpoint whose
+model never receives the live microphone stream:
+
+```python
+duplex_control_enabled = True
+duplex_runtime_extension = "..."
+duplex_serving_adapter = "..."
+```
+
+The adapter must supply a real scheduler data plane, not just endpoint
+capabilities.
+
+## Supported scope
+
+The unified implementation supports:
+
+- up to two engine-owned sessions on one replica;
+- 24 kHz mono float PCM input;
+- one 1920-sample, 80 ms model frame per physical append unit;
+- continuous user input while assistant audio is generated or played;
+- bundled `.pt` voice prompts and a session persona;
+- greedy text and depformer sampling, matching the current PersonaPlex port;
+- `/v1/realtime?duplex=1`;
+- engine lease close, disconnect cleanup, reconnect after cleanup, and explicit
+  response cancellation without cross-session state reuse.
+
+The implementation does not claim:
+
+- more than two simultaneous PersonaPlex sessions on one replica;
+- arbitrary WAV voice cloning;
+- turn-based `response.create` semantics for an otherwise continuous model;
+- scheduler migration of live codec state between replicas;
+- exact output equality with the standalone engine after different scheduling
+  boundaries.
+
+The capability payload derives multi-session support from the configured
+session limit. The shipped two-session deployment reports
+`supports_multi_session=true` and `supports_multi_session_same_replica=true`.
+`duplex_session.max_sessions` is the only capacity source: config resolution
+propagates it to every stage as `duplex_max_sessions`, and both Mimi pools read
+that model-config value. Connector extras do not carry a second model-specific
+capacity knob that could drift from engine admission.
+It still reports `supports_barge_in=false`: the generic epoch fence can suppress
+stale transport output, but neither PersonaPlex nor the current MiniCPM-o 4.5
+adapter proves that model-owned streaming state can be destructively rewound or
+restarted at a playback cursor. Continuous overlapping speech is model-native
+duplex behavior, not by itself a barge-in contract.
+
+## Components and ownership
+
+### Serving adapter
+
+`PersonaPlexServingRuntimeAdapter` owns only serving-side state:
+
+- a transactional PCM append buffer;
+- validation of `voice_prompt` and `instructions`;
+- public capabilities;
+- PersonaPlex data-plane output projection.
+
+It accepts 24 kHz `pcm_f32le`, groups client packets into whole 1920-sample
+frames, zero-pads only the final residual, and rolls a reservation back when an
+engine append fails. It never loads CUDA weights and never encodes user audio.
+
+Client input cannot provide local filesystem paths. A voice is a bundled
+basename such as `NATF2.pt`; the worker resolves it under the local model
+checkpoint. `instructions` is the persona string.
+
+### Runtime extension
+
+`PersonaPlexDuplexRuntimeExtension` is pure model policy. It:
+
+- configures greedy Stage 0 sampling and bounded segment lengths;
+- maps each accepted PCM append to a scheduler prompt;
+- places immutable session identity, append sequence, PCM payload, voice, and
+  persona under `model_intermediate_buffer["duplex"]`;
+- reserves exactly one scheduler prompt slot per encoded Mimi frame, plus the
+  first-append voice/persona prefill length;
+- never performs model inference or owns session state.
+
+The extension returns no turn/listen decision. PersonaPlex is an always-clocked
+model, so visible audio/text comes from the final stage data plane.
+
+### Stage 0 streaming runtime
+
+The Talker owns a `PersonaPlexStage0DuplexRuntime` helper, analogous to
+MiniCPM-o's Stage 0 helper but with PersonaPlex lockstep semantics.
+
+For each admitted session it owns:
+
+- streaming Mimi encoder convolution and transformer state;
+- the selected voice embedding bundle;
+- persona tokenization and prefill embeddings;
+- the prior user code frame needed by the one-frame acoustic delay;
+- append identity used to make a retried scheduler update idempotent.
+
+The first append builds this ordered prefill:
+
+```text
+voice embeddings
+  -> encoded silence
+  -> <system> persona <system> tokens with encoded silence
+  -> encoded silence
+  -> first live user frame
+```
+
+Later appends encode only new 1920-sample frames. The helper returns the
+per-frame user codes and prompt embeddings through the request's
+`model_intermediate_buffer`. The normal vLLM runner remains authoritative for
+attention metadata, block tables, KV allocation, scheduling, and sampling.
+
+The Talker must distinguish resumable prompt prefill from decode positions.
+Prompt rows consume the exact prepared embeddings; sampled decode rows continue
+to use the existing delayed agent/user frame construction. No code path may
+fall back to an all-initial user stream for a duplex request.
+
+Cleanup is keyed by the full `(session_id, incarnation)` identity. Every live
+session has an independent Mimi encoder instance; encoder convolution/KV state
+is never shared between asynchronously scheduled sessions. A finished or
+aborted scheduler request resets and returns only that session's encoder.
+
+### Stage 1 streaming decoder
+
+The current `PersonaPlexCode2Wav` calls one-shot `MimiModel.decode` and is not
+CUDA-graph safe. Unified duplex uses eager Stage 1 and maintains an independent
+streaming Mimi decoder for every active request. Decoder ownership is keyed by
+the stable Stage 1 request id and released by `on_requests_finished`; a mixed
+batch must never advance another request's convolution or transformer state.
+
+Each Stage 0 segment emits de-delayed agent codebooks. Stage 1 decodes only the
+new code frames, emits only the new PCM suffix, and resets state when the
+request is closed. Connector chunk boundaries retain the final raw code frame
+needed to de-delay the next chunk.
+
+The deploy default sets Stage 1 `enforce_eager: true`; a default configuration
+that fails during CUDA graph capture is not an acceptable deployment profile.
+
+### Data-plane projector
+
+`PersonaPlexDataPlaneSession` converts cumulative or delta Stage 1 output into
+model-neutral native results:
+
+```python
+{
+    "stage_role": "tts",
+    "data_plane_request_id": request_id,
+    "text": text_delta,
+    "audio_data": encoded_audio_delta,
+    "audio_format": response_format,
+    "sample_rate_hz": 24000,
+    "audio_duration_ms": delta_duration,
+    "end_of_turn": False,
+}
+```
+
+It owns per-request audio and text cursors so a cumulative output cannot replay
+old audio. PersonaPlex keeps one visible response open while continuous output
+arrives. Session close or cancellation terminates that response through the
+generic Realtime lifecycle; a codec segment finishing is not a conversational
+turn boundary.
+
+## Error and lifecycle contracts
+
+- Unsupported sample rate, malformed base64, non-finite PCM, invalid voice
+  basename, or changed format fails before scheduler submission.
+- Append is prepare/submit/commit. Failure rolls back the exact reserved PCM
+  bytes and does not advance the model frame cursor.
+- A repeated `operation_id` must not encode or submit the same frame twice.
+- Input iterator exceptions execute the same cleanup as explicit close.
+- Cancellation does not release the engine lease until the stage request and
+  any in-flight codec operation are actually finished.
+- A bounded cleanup timeout returns a cleanup error and keeps the session in
+  the closing admission set; it must not make the slot available while work
+  still mutates shared state.
+- New sessions cannot observe the previous voice, persona, PCM tail, Mimi
+  convolution state, or Talker delayed code frame.
+
+The legacy `PersonaPlexDuplexRuntime.run()` and standalone server drain path
+must follow the same exception-safe rule while they remain in the tree.
+
+## Testing and acceptance
+
+### Contract tests
+
+Tests first cover:
+
+- pipeline registration enables the control plane and selects both PersonaPlex
+  adapters;
+- the capability payload is two-session, 80 ms, append-only, and honest;
+- PCM reservation commit/rollback, partial-frame flush, invalid input, and
+  operation idempotency;
+- runtime prompt fields and exact token budgeting;
+- first-append voice/persona prefill followed by live user codes;
+- later appends retain per-session Mimi state and do not replay prefill;
+- interleaved Stage 0 and Stage 1 work preserves independent codec histories;
+- output projection emits only audio/text deltas;
+- close, exception, timeout, and reconnect cleanup;
+- ordinary non-duplex imports do not load PersonaPlex modules.
+
+### Remote H20 validation
+
+Validation runs in an isolated remote worktree using the ModelScope
+`nv-community/personaplex-7b-v1` checkpoint and its Mimi dependency.
+
+The required evidence is:
+
+1. default `personaplex.yaml` reaches ready without a local eager override;
+2. `/health` returns 200;
+3. `/v1/realtime?duplex=1` reports `model_native_duplex`, `chunk_period_ms=80`,
+   and two-session admission;
+4. paced 24 kHz PCM appends produce finite, non-silent 24 kHz audio deltas and
+   text deltas;
+5. microphone input continues during assistant output without cancelling the
+   scheduler request;
+6. two paced sessions simultaneously produce independent non-empty audio, and
+   a third session is rejected with `resource_exhausted`;
+7. closing either session frees only its scheduler request and GPU codec state,
+   after which a replacement session can use a different persona without state
+   leakage;
+8. malformed input returns one typed error and does not poison the next append;
+9. all focused unit tests and `git diff --check` pass.
+
+Audio that is empty, all zero, non-finite, or only a protocol `listen` event is
+not a successful end-to-end result.

vllm_omni/experimental/fullduplex/personaplex/__init__.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex full-duplex backend for the experimental fullduplex framework.
+
+PersonaPlex (``nvidia/personaplex-7b-v1``) is a Moshi finetune: a pure-lockstep
+speech-to-speech model. This package plugs it into the model-agnostic
+``fullduplex/core`` seam without touching the engine:
+
+- :class:`PersonaPlexConfig`   immutable session config (voice / persona / sampling)
+- :class:`PersonaPlexEngine`   the moshi-free lockstep engine (``FrameStepper``;
+  single session and B concurrent slots with elastic per-slot recycle)
+- :class:`PersonaPlexSession`  lockstep driver (the runnable serving primitive)
+- :class:`PersonaPlexDuplexAdapter`  the ``core.DuplexAdapter`` for lockstep audio
+- :class:`PersonaPlexDuplexRuntime`  the model-owned lockstep lifecycle
+  (start-once / drain-on-close; ``core`` stays verbatim upstream)
+"""
+
+from vllm_omni.experimental.fullduplex.personaplex.adapter import (
+    PersonaPlexDuplexAdapter,
+    PersonaPlexDuplexRuntime,
+)
+from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig
+from vllm_omni.experimental.fullduplex.personaplex.engine import (
+    FrameOutput,
+    FrameStepper,
+    PersonaPlexEngine,
+)
+from vllm_omni.experimental.fullduplex.personaplex.runtime import PrefillStep
+from vllm_omni.experimental.fullduplex.personaplex.session import PersonaPlexSession
+
+__all__ = [
+    "FrameOutput",
+    "FrameStepper",
+    "PersonaPlexConfig",
+    "PersonaPlexDuplexAdapter",
+    "PersonaPlexDuplexRuntime",
+    "PersonaPlexEngine",
+    "PersonaPlexSession",
+    "PrefillStep",
+]

vllm_omni/experimental/fullduplex/personaplex/adapter.py

@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex duplex adapter + model-owned lockstep runtime.
+
+PersonaPlex is a *continuous* (frame-clocked, no turn-taking) model: one eternal
+``respond()`` that pulls one 80 ms user frame at a time and emits one agent
+frame + text token. The shared ``core.DuplexRuntime`` implements the turn
+lifecycle (start/cancel per trigger), which does not fit lockstep, so the
+lifecycle lives HERE as :class:`PersonaPlexDuplexRuntime` — model policy stays
+in the model package, mirroring #3907's model-owned minicpmo45 runtime, and
+``core`` remains verbatim upstream.
+
+Stepping happens off the event loop (``asyncio.to_thread``, as Moshi's own
+server does) so the input side stays live — that is what makes barge-in work
+without any explicit flush.
+
+Depends only on the ``FrameStepper`` protocol, so it is exercised by GPU-free
+tests with a stub stepper; ``PersonaPlexEngine`` is the real backend.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from typing import Any
+
+import numpy as np
+
+from vllm_omni.experimental.fullduplex.core import protocol as ev
+from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter, DuplexCapability, OutputChunk
+from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime, Emit
+from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexState
+from vllm_omni.experimental.fullduplex.personaplex.engine import FrameStepper
+from vllm_omni.experimental.fullduplex.personaplex.session import MAX_PENDING_INPUT, offer_realtime
+
+_CLOSE = object()  # sentinel: drain the inbox and end the eternal response
+
+
+class PersonaPlexDuplexAdapter(DuplexAdapter):
+    """Drive a :class:`FrameStepper` as a continuous full-duplex response."""
+
+    def __init__(
+        self,
+        stepper: FrameStepper,
+        *,
+        voice_prompt: str | None = None,
+        persona: str | None = None,
+    ) -> None:
+        self._stepper = stepper
+        self._voice_prompt = voice_prompt
+        self._persona = persona
+        self._frame_size = stepper.frame_size
+        # Bounded, drop-oldest: input arriving faster than the stepper can run
+        # must shed stale frames (realtime), not grow memory + latency unbounded.
+        self._inbox: asyncio.Queue[Any] = asyncio.Queue(maxsize=MAX_PENDING_INPUT)
+        self._pending = np.zeros(0, dtype=np.float32)
+        self._opened = False
+
+    def capabilities(self) -> DuplexCapability:
+        return DuplexCapability(
+            input_modalities=frozenset({"audio"}),
+            output_modalities=frozenset({"audio", "text"}),
+            proactive=True,
+        )
+
+    async def on_input(self, session: DuplexSession, modality: str, data: Any) -> None:
+        """Frame incoming PCM into exact 80 ms windows and queue them for the loop."""
+        if modality != "audio":
+            return
+        samples = np.ascontiguousarray(data, dtype=np.float32).reshape(-1)
+        self._pending = np.concatenate([self._pending, samples]) if self._pending.size else samples
+        n = self._frame_size
+        while self._pending.shape[0] >= n:
+            frame, self._pending = self._pending[:n], self._pending[n:]
+            offer_realtime(self._inbox, frame)
+
+    async def respond(self, session: DuplexSession) -> AsyncIterator[OutputChunk]:
+        """The eternal lockstep loop: one user frame in, one agent frame + text out."""
+        if not self._opened:
+            await asyncio.to_thread(self._stepper.open_session, self._voice_prompt, self._persona)
+            self._opened = True
+        while True:
+            frame = await self._inbox.get()
+            if frame is _CLOSE:
+                break
+            out = await asyncio.to_thread(self._stepper.step, frame)
+            if out.audio is not None:
+                yield OutputChunk("audio", out.audio)
+            if out.text:
+                yield OutputChunk("text", out.text)
+
+    async def on_close(self, session: DuplexSession) -> None:
+        # Model-owned drain signal (not part of the core DuplexAdapter ABC):
+        # PersonaPlexDuplexRuntime calls this on every exit path. Pad + flush a
+        # partial final frame (same contract as PersonaPlexSession.flush()) so
+        # the tail of the user's audio is stepped, not dropped.
+        if self._pending.size:
+            pad = self._frame_size - self._pending.shape[0]
+            frame = np.concatenate([self._pending, np.zeros(pad, dtype=np.float32)])
+            self._pending = np.zeros(0, dtype=np.float32)
+            offer_realtime(self._inbox, frame)
+        offer_realtime(self._inbox, _CLOSE)
+
+    async def on_barge_in(self, session: DuplexSession) -> None:
+        # No-op: Moshi-class models perceive the user continuously, so an interrupt
+        # needs no downstream flush — the next frames simply steer the model.
+        return
+
+
+class PersonaPlexDuplexRuntime(DuplexRuntime):
+    """Model-owned lockstep lifecycle: start once, never supersede, drain on exit.
+
+    ``core.DuplexRuntime`` runs the turn lifecycle — every trigger starts a new
+    response that supersedes and cancels the previous one. A lockstep model has
+    exactly ONE response for the whole session: it consumes user frames as they
+    arrive and only ends when the input side goes away. Keeping this here (not
+    as a flag on the shared runtime) mirrors #3907's model-owned minicpmo45
+    runtime: model policy lives in the model package.
+    """
+
+    async def run(self, inputs: AsyncIterator[dict], emit: Emit) -> None:
+        task: asyncio.Task | None = None
+        try:
+            async for event in inputs:
+                etype = event.get("type")
+                if etype == ev.INPUT_APPEND:
+                    modality = event.get("modality", "")
+                    if modality not in self._capabilities.input_modalities:
+                        await emit(ev.error(f"unsupported input modality: {modality}"))
+                        continue
+                    await self.adapter.on_input(self.session, modality, event.get("data"))
+                    self.session.state = DuplexState.LISTENING
+                    task = self._ensure_started(task, emit)
+                elif etype in (ev.INPUT_COMMIT, ev.RESPONSE_CREATE):
+                    task = self._ensure_started(task, emit)
+                elif etype == ev.RESPONSE_CANCEL:
+                    # Barge-in is native for lockstep models: bump the epoch for the
+                    # playback cursor but keep the single response streaming.
+                    await self._barge_in()
+                    await emit(ev.cancelled(self.session.response_index))
+                elif etype == ev.PLAYBACK_ACK:
+                    await self.adapter.on_playback_ack(self.session, int(event.get("cursor", 0)))
+                elif etype == ev.CLOSE:
+                    break
+        finally:
+            # Drain on EVERY exit path: CLOSE, iterator EOF, transport errors,
+            # and cancellation must not leave the eternal response blocked on
+            # its inbox or the session outside CLOSED.
+            try:
+                if task is not None and not task.done():
+                    await self.adapter.on_close(self.session)
+                    await task
+            finally:
+                self.session.state = DuplexState.CLOSED
+
+    def _ensure_started(self, task: asyncio.Task | None, emit: Emit) -> asyncio.Task:
+        # Start the single eternal response exactly once; never supersede it.
+        if task is None or task.done():
+            return asyncio.create_task(self._respond_once(emit))
+        return task
+
+    async def _respond_once(self, emit: Emit) -> None:
+        # One response for the whole session. No staleness checks: the loop ends
+        # only when the adapter's respond() returns (after on_close drains its
+        # input).
+        response_index, _ = self.session.begin_response()
+        await emit(ev.created(response_index))
+        try:
+            async for chunk in self.adapter.respond(self.session):
+                await emit(ev.delta(response_index, chunk.modality, chunk.data))
+        except asyncio.CancelledError:
+            raise
+        except Exception as err:
+            await emit(ev.error(f"response failed: {err}"))
+        finally:
+            await emit(ev.done(response_index))
+            self.session.end_response()

vllm_omni/experimental/fullduplex/personaplex/config.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Configuration for the PersonaPlex full-duplex backend.
+
+PersonaPlex (``nvidia/personaplex-7b-v1``) is a Moshi finetune: a pure-lockstep
+full-duplex speech-to-speech model running at the Mimi codec frame rate
+(12.5 Hz / 80 ms). One config drives both the offline driver and the duplex
+adapter. Defaults mirror the PersonaPlex reference loop.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# Mimi codec constants (loaders.py: SAMPLE_RATE / FRAME_RATE). Fixed by the
+# pretrained checkpoint; exposed here only so the rest of the package never
+# hard-codes them.
+SAMPLE_RATE = 24000
+FRAME_RATE = 12.5
+FRAME_SIZE = int(SAMPLE_RATE / FRAME_RATE)  # 1920 samples per 80 ms frame
+
+# Default assistant persona shipped with PersonaPlex (offline.py:335).
+DEFAULT_PERSONA = "You are a wise and friendly teacher. Answer questions or provide advice in a clear and engaging way."
+
+
+@dataclass(frozen=True)
+class PersonaPlexConfig:
+    """Immutable session configuration for a PersonaPlex conversation.
+
+    Attributes:
+        hf_repo: HuggingFace repo holding the weights, Mimi codec and tokenizer.
+        voice_prompt: Voice-clone reference. Either a bundled basename
+            (``"NATF2.pt"`` / ``"NATM1.pt"`` from ``voices.tgz``) or a path to a
+            ``.pt`` embedding bundle or a reference ``.wav``.
+        persona: System role text; injected as ``<system> ... <system>`` into the
+            inner-monologue stream at session start.
+        device: Torch device for the backend (``"cuda"`` / ``"cpu"``).
+        cpu_offload: Offload LM layers to CPU when GPU memory is tight (needs
+            ``accelerate``).
+        batch_size: Concurrent conversation slots sharing one engine. ``1`` is
+            the single-session path; ``> 1`` enables elastic batching with
+            per-slot recycle for new callers.
+
+    Note: the native stepper decodes greedily (argmax) for both the text head
+    and the depformer, so there are no sampling knobs here yet. Temperature /
+    top-k / seed fields will be added if and when a sampling path is wired in.
+    """
+
+    hf_repo: str = "nvidia/personaplex-7b-v1"
+    voice_prompt: str = "NATF2.pt"
+    persona: str = DEFAULT_PERSONA
+
+    device: str = "cuda"
+    cpu_offload: bool = False
+
+    batch_size: int = 1
+
+    @property
+    def sample_rate(self) -> int:
+        return SAMPLE_RATE
+
+    @property
+    def frame_size(self) -> int:
+        return FRAME_SIZE

vllm_omni/experimental/fullduplex/personaplex/data_plane.py

@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Mapping
+from dataclasses import dataclass
+
+import numpy as np
+
+from vllm_omni.experimental.fullduplex.engine.contracts import (
+    duplex_resource_request_belongs_to_session,
+)
+
+EncodeAudio = Callable[[object, int, str, float | None], str | None]
+
+
+@dataclass(frozen=True, slots=True)
+class PersonaPlexDataPlaneContext:
+    epoch: int = 0
+    turn_id: int = 0
+    active_response_turn_id: int | None = None
+    active_response_id: str | None = None
+    auto_responds: bool = True
+    response_format: str = "wav"
+    speed: float | None = None
+    modalities: tuple[str, ...] = ("audio", "text")
+
+
+@dataclass(slots=True)
+class _RequestCursor:
+    audio_samples: int = 0
+    text: str = ""
+    terminal: bool = False
+
+
+class PersonaPlexDataPlaneSession:
+    """Project cumulative staged PersonaPlex output into Realtime deltas."""
+
+    def __init__(self, encode_audio: EncodeAudio) -> None:
+        self._encode_audio = encode_audio
+        self._requests: dict[str, _RequestCursor] = {}
+
+    def begin_request(self, request_id: str) -> None:
+        self._requests.setdefault(request_id, _RequestCursor()).terminal = False
+
+    def is_terminal(self, request_id: str | None) -> bool:
+        if request_id is None:
+            return False
+        state = self._requests.get(request_id)
+        return state is not None and state.terminal
+
+    def mark_terminal(self, request_id: str) -> None:
+        self._requests.setdefault(request_id, _RequestCursor()).terminal = True
+
+    def close_stream(self, request_id: str) -> None:
+        self._requests.pop(request_id, None)
+
+    def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None:
+        if active_request_id is not None:
+            self._requests.pop(active_request_id, None)
+        for request_id in list(self._requests):
+            if duplex_resource_request_belongs_to_session(request_id, session_id):
+                self._requests.pop(request_id, None)
+
+    def project(
+        self,
+        result: object,
+        *,
+        context: PersonaPlexDataPlaneContext | None = None,
+    ) -> Iterator[dict[str, object]]:
+        if not isinstance(result, dict):
+            return
+        outputs = result.get("data_plane_outputs")
+        if not isinstance(outputs, list):
+            return
+        context = context or PersonaPlexDataPlaneContext()
+        for output in outputs:
+            projected = self._project_output(output, context=context)
+            if projected is not None:
+                yield projected
+
+    def _project_output(
+        self,
+        output: object,
+        *,
+        context: PersonaPlexDataPlaneContext,
+    ) -> dict[str, object] | None:
+        output, completion = _unwrap_output(output)
+        request_id = getattr(output, "request_id", None)
+        if not isinstance(request_id, str) or not request_id:
+            request_id = None
+        state = self._requests.setdefault(request_id, _RequestCursor()) if request_id is not None else _RequestCursor()
+        multimodal = _multimodal_output(output, completion)
+        audio = _audio_value(multimodal)
+        audio_delta = _slice_audio(audio, state.audio_samples)
+        audio_samples = _num_samples(audio)
+        if audio_samples is not None:
+            state.audio_samples = audio_samples
+
+        text = _text_value(multimodal, completion)
+        text_delta = _text_delta(text, state.text)
+        if text:
+            state.text = text
+
+        sample_rate_hz = _sample_rate(multimodal)
+        encoded = self._encode_audio(audio_delta, sample_rate_hz, context.response_format, context.speed)
+        if not encoded and not text_delta:
+            return None
+        delta_samples = _num_samples(audio_delta) or 0
+        return {
+            "supported": True,
+            "stage_role": "tts",
+            "is_listen": False,
+            "data_plane_request_id": request_id,
+            "text": text_delta,
+            "audio_data": encoded or "",
+            "audio_format": context.response_format,
+            "sample_rate_hz": sample_rate_hz,
+            "audio_duration_ms": round(delta_samples * 1000 / max(1, sample_rate_hz)),
+            "end_of_turn": False,
+            "uses_model_runner_scheduler": True,
+            "runner_kv_backed": True,
+            "runtime_impl": "scheduler_data_plane",
+            "owned_runtime": False,
+        }
+
+
+def _unwrap_output(output: object) -> tuple[object, object | None]:
+    inner = getattr(output, "request_output", None)
+    if inner is not None and inner is not output:
+        output = inner
+    outputs = getattr(output, "outputs", None)
+    completion = outputs[0] if isinstance(outputs, list) and outputs else None
+    return output, completion
+
+
+def _multimodal_output(output: object, completion: object | None) -> dict[str, object]:
+    for candidate in (
+        getattr(output, "multimodal_output", None),
+        getattr(completion, "multimodal_output", None) if completion is not None else None,
+    ):
+        if isinstance(candidate, Mapping) and candidate:
+            return dict(candidate)
+    return {}
+
+
+def _audio_value(multimodal: dict[str, object]) -> object | None:
+    value = next(
+        (multimodal[key] for key in ("audio", "model_outputs", "latent") if key in multimodal),
+        None,
+    )
+    if isinstance(value, list) and len(value) == 1:
+        return value[0]
+    return value
+
+
+def _text_value(multimodal: dict[str, object], completion: object | None) -> str:
+    for candidate in (
+        multimodal.get("text"),
+        multimodal.get("llm_output_text"),
+        getattr(completion, "text", None) if completion is not None else None,
+    ):
+        if isinstance(candidate, str) and candidate:
+            return candidate
+    return ""
+
+
+def _text_delta(text: str, previous: str) -> str:
+    if not text:
+        return ""
+    if text == previous:
+        return ""
+    if previous and text.startswith(previous):
+        return text[len(previous) :]
+    return text
+
+
+def _slice_audio(audio: object | None, offset: int) -> object | None:
+    samples = _num_samples(audio)
+    if samples is None or samples <= 0:
+        return None
+    if offset <= 0 or samples < offset:
+        return audio
+    if samples == offset:
+        return None
+    try:
+        import torch
+
+        if isinstance(audio, torch.Tensor):
+            return audio.reshape(-1)[offset:].contiguous()
+    except Exception:
+        pass
+    return np.asarray(audio, dtype=np.float32).reshape(-1)[offset:]
+
+
+def _num_samples(audio: object | None) -> int | None:
+    if audio is None:
+        return None
+    try:
+        import torch
+
+        if isinstance(audio, torch.Tensor):
+            return int(audio.numel())
+    except Exception:
+        pass
+    try:
+        return int(np.asarray(audio, dtype=np.float32).size)
+    except Exception:
+        return None
+
+
+def _sample_rate(multimodal: dict[str, object]) -> int:
+    value = multimodal.get("sr", multimodal.get("sample_rate_hz", 24000))
+    if isinstance(value, list) and value:
+        value = value[0]
+    if hasattr(value, "item"):
+        try:
+            value = value.item()
+        except Exception:
+            value = 24000
+    return int(value) if isinstance(value, int | float) else 24000
+
+
+__all__ = [
+    "PersonaPlexDataPlaneContext",
+    "PersonaPlexDataPlaneSession",
+]

vllm_omni/experimental/fullduplex/personaplex/engine.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex duplex protocol home (moshi-free).
+
+Defines the narrow ``FrameStepper`` seam (one 80 ms user frame in, one agent
+frame + text piece out) that the session driver, the duplex adapter and the
+GPU-free tests depend on, and re-exports the native lockstep engine that
+implements it. The implementation lives in ``runtime.py`` and is built
+entirely from vllm-omni-native components (streaming Helium temporal, input
+embeddings, depformer, streaming Mimi codec) — the external ``moshi`` package
+is no longer used.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class FrameOutput:
+    """Result of stepping the model by exactly one 80 ms frame.
+
+    Attributes:
+        audio: Agent PCM for this frame (``frame_size`` mono float32 samples), or
+            ``None`` during the model's initial delay/warmup frames.
+        text: Inner-monologue text piece for this frame, or ``None`` when the
+            model emitted a padding token (silence between words).
+    """
+
+    audio: NDArray[np.float32] | None
+    text: str | None
+
+
+@runtime_checkable
+class FrameStepper(Protocol):
+    """The single seam between PersonaPlex and the duplex framework.
+
+    A stepper consumes one user audio frame and returns one agent frame. It owns
+    all conversation state (KV caches / streaming buffers); callers never re-feed
+    history. ``PersonaPlexEngine`` is the real implementation; tests pass a stub.
+    """
+
+    sample_rate: int
+    frame_size: int
+
+    def open_session(self, voice_prompt: str | None = None, persona: str | None = None) -> None:
+        """Reset streaming state and inject the voice clone + persona prompt."""
+        ...
+
+    def step(self, user_pcm: NDArray[np.float32]) -> FrameOutput:
+        """Advance the conversation by one frame of exactly ``frame_size`` samples."""
+        ...
+
+
+from vllm_omni.experimental.fullduplex.personaplex.runtime import (  # noqa: E402
+    PersonaPlexEngine,
+    PrefillStep,
+)
+
+__all__ = [
+    "FrameOutput",
+    "FrameStepper",
+    "PersonaPlexEngine",
+    "PrefillStep",
+]

vllm_omni/experimental/fullduplex/personaplex/input.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import base64
+import binascii
+
+import numpy as np
+
+_SAMPLE_RATE_HZ = 24000
+_FRAME_SAMPLES = 1920
+_BYTES_PER_SAMPLE = 4
+
+
+class PersonaPlexPcmAppendReservation:
+    __slots__ = (
+        "_active",
+        "_owner",
+        "_raw",
+        "operation_id",
+        "payload",
+    )
+
+    def __init__(
+        self,
+        *,
+        owner: PersonaPlexPcmAppendBuffer,
+        operation_id: str,
+        payload: dict[str, object] | None,
+        raw: bytes,
+    ) -> None:
+        self._owner = owner
+        self.operation_id = operation_id
+        self.payload = payload
+        self._raw = raw
+        self._active = True
+
+    @property
+    def active(self) -> bool:
+        return self._active
+
+    @property
+    def byte_count(self) -> int:
+        return len(self._raw)
+
+    def commit(self) -> None:
+        self._owner._commit_reservation(self)
+
+    def rollback(self) -> None:
+        self._owner._rollback_reservation(self)
+
+
+class PersonaPlexPcmAppendBuffer:
+    """Transactionally frame 24 kHz float PCM into PersonaPlex 80 ms units."""
+
+    def __init__(self) -> None:
+        self._buffer = bytearray()
+        self._sample_rate_hz: int | None = None
+        self._reservations: list[PersonaPlexPcmAppendReservation] = []
+
+    @property
+    def pending_byte_count(self) -> int:
+        return len(self._buffer)
+
+    def has_pending(self) -> bool:
+        return bool(self._buffer)
+
+    def has_reserved(self) -> bool:
+        return any(reservation.active for reservation in self._reservations)
+
+    def clear_force_listen(self) -> None:
+        return
+
+    def clear(self) -> None:
+        for reservation in self._reservations:
+            reservation._active = False
+        self._reservations.clear()
+        self._buffer.clear()
+        self._sample_rate_hz = None
+
+    def prepare_append(
+        self,
+        payload: dict[str, object],
+        *,
+        operation_id: str,
+        chunk_period_ms: int,
+        allow_emit: bool,
+    ) -> PersonaPlexPcmAppendReservation | None:
+        if chunk_period_ms != 80:
+            raise ValueError("PersonaPlex chunk_period_ms must be 80")
+        if any(reservation.active and reservation.operation_id == operation_id for reservation in self._reservations):
+            raise ValueError(f"PersonaPlex duplicate active operation_id: {operation_id}")
+
+        raw = self._decode_payload(payload)
+        self._buffer.extend(raw)
+        if not allow_emit or len(self._buffer) < _FRAME_SAMPLES * _BYTES_PER_SAMPLE:
+            return None
+        return self._reserve_frame(payload, operation_id=operation_id, flush=False)
+
+    def prepare_commit(
+        self,
+        *,
+        operation_id: str,
+        chunk_period_ms: int,
+    ) -> PersonaPlexPcmAppendReservation:
+        if chunk_period_ms != 80:
+            raise ValueError("PersonaPlex chunk_period_ms must be 80")
+        if not self._buffer:
+            reservation = PersonaPlexPcmAppendReservation(
+                owner=self,
+                operation_id=operation_id,
+                payload=None,
+                raw=b"",
+            )
+            self._reservations.append(reservation)
+            return reservation
+        payload: dict[str, object] = {
+            "type": "audio",
+            "format": "pcm_f32le",
+            "sample_rate_hz": self._sample_rate_hz or _SAMPLE_RATE_HZ,
+            "audio": "",
+            "final": True,
+        }
+        return self._reserve_frame(payload, operation_id=operation_id, flush=True)
+
+    def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None:
+        reservation = self.prepare_commit(
+            operation_id="personaplex-flush",
+            chunk_period_ms=chunk_period_ms,
+        )
+        payload = reservation.payload
+        reservation.commit()
+        return payload
+
+    def _decode_payload(self, payload: dict[str, object]) -> bytes:
+        if payload.get("format") != "pcm_f32le":
+            raise ValueError("PersonaPlex input format must be pcm_f32le")
+        sample_rate_hz = payload.get("sample_rate_hz")
+        if sample_rate_hz != _SAMPLE_RATE_HZ:
+            raise ValueError("PersonaPlex sample_rate_hz must be 24000")
+        if self._sample_rate_hz is not None and self._sample_rate_hz != sample_rate_hz:
+            raise ValueError("PersonaPlex sample_rate_hz changed within a session")
+        audio = payload.get("audio")
+        if not isinstance(audio, str):
+            raise ValueError("PersonaPlex audio must be base64 pcm_f32le")
+        try:
+            raw = base64.b64decode(audio, validate=True)
+        except (binascii.Error, ValueError) as exc:
+            raise ValueError("PersonaPlex audio is not valid base64") from exc
+        if len(raw) % _BYTES_PER_SAMPLE:
+            raise ValueError("PersonaPlex pcm_f32le byte length must be divisible by four")
+        samples = np.frombuffer(raw, dtype="<f4")
+        if not np.isfinite(samples).all():
+            raise ValueError("PersonaPlex pcm_f32le samples must be finite")
+        self._sample_rate_hz = int(sample_rate_hz)
+        return raw
+
+    def _reserve_frame(
+        self,
+        source_payload: dict[str, object],
+        *,
+        operation_id: str,
+        flush: bool,
+    ) -> PersonaPlexPcmAppendReservation:
+        frame_bytes = _FRAME_SAMPLES * _BYTES_PER_SAMPLE
+        if flush:
+            consumed_bytes = min(len(self._buffer), frame_bytes)
+        else:
+            consumed_bytes = frame_bytes
+        raw = bytes(self._buffer[:consumed_bytes])
+        del self._buffer[:consumed_bytes]
+        encoded_raw = raw + b"\x00" * (frame_bytes - consumed_bytes)
+        payload = dict(source_payload)
+        payload["type"] = "audio"
+        payload["format"] = "pcm_f32le"
+        payload["sample_rate_hz"] = _SAMPLE_RATE_HZ
+        payload["audio"] = base64.b64encode(encoded_raw).decode("ascii")
+        reservation = PersonaPlexPcmAppendReservation(
+            owner=self,
+            operation_id=operation_id,
+            payload=payload,
+            raw=raw,
+        )
+        self._reservations.append(reservation)
+        return reservation
+
+    def _commit_reservation(self, reservation: PersonaPlexPcmAppendReservation) -> None:
+        if not reservation.active:
+            return
+        reservation._active = False
+        self._reservations.remove(reservation)
+
+    def _rollback_reservation(self, reservation: PersonaPlexPcmAppendReservation) -> None:
+        if not reservation.active:
+            return
+        try:
+            index = self._reservations.index(reservation)
+        except ValueError:
+            reservation._active = False
+            return
+        rolled_back = self._reservations[index:]
+        restored = b"".join(item._raw for item in rolled_back if item.active)
+        self._buffer[:0] = restored
+        for item in rolled_back:
+            item._active = False
+        del self._reservations[index:]
+
+
+__all__ = [
+    "PersonaPlexPcmAppendBuffer",
+    "PersonaPlexPcmAppendReservation",
+]

vllm_omni/experimental/fullduplex/personaplex/policy.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex frame/token contract (the model *policy*, no engine state).
+
+This mirrors the ``minicpmo45`` model-package layout from PR #3907: the token
+contract lives in ``policy.py``, the model-owned runtime in ``runtime.py``.
+Everything here is a constant or a pure function of the checkpoint's lockstep
+conventions; nothing touches weights, devices, or streaming state.
+
+Frame layout per 80 ms tick: 17 token rows -- row 0 inner-monologue text,
+rows 1-8 agent audio codebooks, rows 9-16 user audio codebooks.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+# Inner-monologue token ids that carry no displayable text (EPAD, PAD).
+NON_TEXT_TOKEN_IDS = frozenset({0, 3})
+ZERO_TEXT_TOKEN = 3
+# Mimi token constants for the hybrid system prompt (agent silence / user sine).
+SILENCE_TOKENS = (948, 243, 1178, 546, 1736, 1030, 1978, 2008)
+SINE_TOKENS = (430, 1268, 381, 1611, 1095, 1495, 56, 472)
+# LM initial ("beginning of stream") token ids: audio rows use the codec
+# cardinality, the text row uses the text vocabulary size.
+AUDIO_INITIAL_TOKEN = 2048
+TEXT_INITIAL_TOKEN = 32000
+TEXT_VOCAB_SIZE = 32000
+AUDIO_SILENCE_FRAME_CNT = 6  # 0.5 s at 12.5 Hz
+
+
+@dataclass(frozen=True)
+class PrefillStep:
+    """One tick of a recycled slot's system-prompt replay (see batched serving)."""
+
+    kind: str  # "sacrifice" | "voice" | "voice_end" | "tokens"
+    embedding: Any = None  # [1, 1, H] tensor for kind == "voice"
+    moshi_tokens: Any = None  # [8] tensor for kind == "tokens"
+    text_token: int | None = None
+    user_sine: bool = False
+
+
+def wrap_with_system_tags(text: str) -> str:
+    t = text.strip()
+    return t if t.startswith("<system>") else f"<system> {t} <system>"
+
+
+__all__ = [
+    "AUDIO_INITIAL_TOKEN",
+    "AUDIO_SILENCE_FRAME_CNT",
+    "NON_TEXT_TOKEN_IDS",
+    "PrefillStep",
+    "SILENCE_TOKENS",
+    "SINE_TOKENS",
+    "TEXT_INITIAL_TOKEN",
+    "TEXT_VOCAB_SIZE",
+    "ZERO_TEXT_TOKEN",
+    "wrap_with_system_tags",
+]

vllm_omni/experimental/fullduplex/personaplex/runtime.py

@@ -0,0 +1,473 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Moshi-free PersonaPlex lockstep engine (native temporal + depformer + Mimi).
+
+This is the real-time duplex backend built entirely from vllm-omni-native
+components: :class:`PersonaPlexTemporalStreaming` (ring-KV Helium),
+:class:`PersonaPlexInputEmbeddings`, :class:`PersonaPlexDepformer` and
+:class:`PersonaPlexMimiCodec`. It reimplements the reference lockstep semantics
+(staging ring, per-codebook delays, warmup gates, hybrid system-prompt prefill,
+per-slot elastic recycle) as first-class policy code, so the external ``moshi``
+package is no longer needed at runtime.
+
+Lockstep contract per 80 ms tick (17 token rows: text, 8 agent codebooks, 8
+user codebooks; delays ``[0, 1x8-ish from checkpoint]``):
+
+1. Mimi-encode the user frame and write it (delayed) into the staging ring.
+2. Read the model input at ``(o - 1) % CT``, embed, step the temporal.
+3. Greedy text token; depformer expands the 16 audio rows (teacher-forced where
+   the staging ring already holds provided values).
+4. Write sampled rows at ``o % CT``; gather the de-delayed output frame; decode
+   agent rows 1..8 through Mimi (rows still in warmup are sanitized to silence
+   tokens: the LM's initial audio token id equals the codec cardinality and is
+   out of range for the codebook tables).
+
+Elastic recycle mirrors the verified design from the moshi-hosted path: the
+temporal ring KV masks the recycled row (plus the one LM "sacrifice tick"
+write), the staging row is wiped to initial tokens, and the system prompt
+(voice embeddings, silence spacer, persona tokens, silence) is replayed per-row
+while other slots keep conversing.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import tarfile
+import threading
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig
+from vllm_omni.experimental.fullduplex.personaplex.policy import (
+    AUDIO_INITIAL_TOKEN,
+    AUDIO_SILENCE_FRAME_CNT,
+    NON_TEXT_TOKEN_IDS,
+    SILENCE_TOKENS,
+    SINE_TOKENS,
+    TEXT_INITIAL_TOKEN,
+    TEXT_VOCAB_SIZE,
+    ZERO_TEXT_TOKEN,
+    PrefillStep,
+    wrap_with_system_tags,
+)
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+logger = logging.getLogger(__name__)
+
+
+class PersonaPlexEngine:
+    """Frame-lockstep PersonaPlex on native components only (no moshi).
+
+    One class serves both surfaces: the single-session ``FrameStepper`` protocol
+    (``open_session`` / ``step``) and the batched slot API consumed by
+    ``serving/batched.py`` (``open_batch`` / ``prefill_steps`` / ``recycle_slot``
+    / ``voice_end`` / ``end_prefill`` / ``step_batch``).
+    """
+
+    def __init__(self, config: PersonaPlexConfig | None = None) -> None:
+        self.config = config or PersonaPlexConfig()
+        self.sample_rate = self.config.sample_rate
+        self.frame_size = self.config.frame_size
+        # Serializes the single-session FrameStepper surface (open_session/step)
+        # across worker threads: a cancelled asyncio task cannot stop a running
+        # engine call, so a racing open() must wait for the in-flight step.
+        # The batched surface is serialized by BatchedSessionManager instead.
+        self._fs_lock = threading.Lock()
+        self._loaded = False
+        self._opened = False
+        self._tokenizer = None
+        self._codec = None
+        self._temporal = None
+        self._emb = None
+        self._dep = None
+        self._voice_embeddings = None  # [T, 1, 1, H]
+        self._voice_cache = None  # [1, K, CT]
+        self._loaded_voice = None  # name/path of the currently loaded voice
+        # lockstep state (built at open time)
+        self._cache = None  # [B, K, CT] long
+        self._provided = None  # [B, K, CT] bool
+        self._initial = None  # [K] long
+        self._delays = None  # [K] long (python list)
+        self._phys = 0
+        self._local = None  # [B] cpu long
+
+    # -- loading --------------------------------------------------------------
+
+    def load(self) -> PersonaPlexEngine:
+        if self._loaded:
+            return self
+        import sentencepiece
+        import torch
+        from huggingface_hub import hf_hub_download
+        from safetensors.torch import load_file
+
+        from vllm_omni.model_executor.models.personaplex.configuration_personaplex import (
+            PersonaPlexConfig as _PPModelConfig,
+        )
+        from vllm_omni.model_executor.models.personaplex.personaplex_depformer import (
+            PersonaPlexDepformer,
+        )
+        from vllm_omni.model_executor.models.personaplex.personaplex_embeddings import (
+            PersonaPlexInputEmbeddings,
+        )
+        from vllm_omni.model_executor.models.personaplex.personaplex_mimi import (
+            PersonaPlexMimiCodec,
+        )
+        from vllm_omni.model_executor.models.personaplex.personaplex_temporal import (
+            PersonaPlexTemporalStreaming,
+        )
+
+        cfg = self.config
+        B = cfg.batch_size
+        logger.info("PersonaPlex native: loading Mimi codec")
+        self._codec = PersonaPlexMimiCodec(device=cfg.device)
+        self._codec.streaming_init(B)
+
+        tok_path = hf_hub_download(cfg.hf_repo, "tokenizer_spm_32k_3.model")
+        self._tokenizer = sentencepiece.SentencePieceProcessor(tok_path)
+
+        logger.info("PersonaPlex native: loading temporal + depformer + embeddings")
+        sd = load_file(hf_hub_download(cfg.hf_repo, "model.safetensors"), device=cfg.device)
+        dtype = torch.bfloat16
+        self._temporal = PersonaPlexTemporalStreaming().to(cfg.device, dtype)
+        self._temporal.load_weights(sd)
+        self._temporal.eval()
+        self._temporal.streaming_init(B)
+
+        ppcfg = _PPModelConfig()
+        self._emb = PersonaPlexInputEmbeddings(ppcfg).to(cfg.device, dtype)
+        self._emb.load_weights(sd)
+        self._emb.eval()
+        self._dep = PersonaPlexDepformer(ppcfg.depformer_config, temporal_hidden_size=4096, text_card=32000).to(
+            cfg.device, dtype
+        )
+        self._dep.load_weights(sd)
+        self._dep.eval()
+        del sd
+        torch.accelerator.empty_cache()
+
+        K = 17
+        self._delays = [0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1]  # text, agent cb0..7, user cb0..7
+        self._initial = torch.tensor(
+            [TEXT_INITIAL_TOKEN] + [AUDIO_INITIAL_TOKEN] * 16, dtype=torch.long, device=cfg.device
+        )
+        ct = max(self._delays) + 3
+        self._cache = torch.full((B, K, ct), AUDIO_INITIAL_TOKEN, dtype=torch.long, device=cfg.device)
+        self._provided = torch.zeros((B, K, ct), dtype=torch.bool, device=cfg.device)
+        self._local = torch.zeros(B, dtype=torch.long)
+        self._sine = torch.tensor(SINE_TOKENS, dtype=torch.long, device=cfg.device)
+        self._silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long, device=cfg.device)
+
+        self._load_voice(cfg.voice_prompt)
+        self._loaded = True
+        return self
+
+    def _load_voice(self, voice: str) -> None:
+        import torch
+
+        if voice == self._loaded_voice:
+            return
+        path = self._resolve_voice_prompt(voice)
+        if not path.endswith(".pt"):
+            raise ValueError(f"native stack needs a .pt voice-embedding bundle; got {path!r}")
+        state = torch.load(path, map_location=self.config.device, weights_only=True)
+        self._voice_embeddings = state["embeddings"].to(self.config.device)
+        self._voice_cache = state["cache"].to(self.config.device)
+        self._loaded_voice = voice
+
+    def _resolve_voice_prompt(self, voice: str) -> str:
+        from huggingface_hub import hf_hub_download
+
+        if os.path.exists(voice):
+            return voice
+        tgz = Path(hf_hub_download(self.config.hf_repo, "voices.tgz"))
+        vdir = tgz.parent / "voices"
+        if not vdir.exists():
+            with tarfile.open(tgz, "r:gz") as tar:
+                tar.extractall(path=tgz.parent)  # noqa: S202 - trusted first-party bundle
+        cand = vdir / voice
+        if not cand.exists():
+            raise FileNotFoundError(f"voice prompt {voice!r} not found in voices.tgz")
+        return str(cand)
+
+    # -- lockstep core ----------------------------------------------------------
+
+    def _reset_all_streaming(self) -> None:
+        self._codec.reset_streaming()
+        self._temporal.reset_streaming()
+        self._cache.fill_(AUDIO_INITIAL_TOKEN)
+        self._cache[:, 0, :] = TEXT_INITIAL_TOKEN
+        self._provided.fill_(False)
+        self._phys = 0
+        self._local.zero_()
+
+    def _tick(
+        self,
+        user_codes,  # [B, 8] long or None (prefill rows overridden separately)
+        forced_moshi=None,  # [B, 8] long (values used on force_mask rows)
+        forced_text=None,  # [B] long
+        force_mask=None,  # [B] bool (device)
+        embed_override=None,  # ([B, 1, H], [B] bool rows)
+    ):
+        """One lockstep step. Returns (ran, out[B,17] | None)."""
+        import torch
+
+        B, K, CT = self._cache.shape
+        o = self._phys
+        lo = self._local
+
+        if user_codes is not None:
+            for q in range(8):
+                k = 9 + q
+                wp = (o + self._delays[k]) % CT
+                self._cache[:, k, wp] = user_codes[:, q]
+                self._provided[:, k, wp] = True
+        if forced_moshi is not None and force_mask is not None:
+            for q in range(8):
+                k = 1 + q
+                wp = (o + self._delays[k]) % CT
+                self._cache[force_mask, k, wp] = forced_moshi[force_mask, q]
+                self._provided[force_mask, k, wp] = True
+        if forced_text is not None and force_mask is not None:
+            wp = (o + self._delays[0]) % CT
+            self._cache[force_mask, 0, wp] = forced_text[force_mask]
+            self._provided[force_mask, 0, wp] = True
+
+        for k, delay in enumerate(self._delays):
+            warm = lo <= delay  # [B] cpu bool
+            if bool(warm.any()):
+                wd = warm.to(self._cache.device)
+                pos = o % CT
+                self._cache[:, k, pos] = torch.where(wd, self._initial[k].expand(B), self._cache[:, k, pos])
+                self._provided[:, k, pos] = self._provided[:, k, pos] | wd
+
+        if o == 0:
+            # Genuine cold start: no prior physical frame to read as model input.
+            self._cache[:, :, 0] = self._initial.view(1, K).expand(B, K)
+            self._phys += 1
+            self._local += 1
+            return False, None
+
+        mip = (o - 1) % CT
+        tp = o % CT
+        input_ = self._cache[:, :, mip]  # [B, K]
+        target_ = self._cache[:, :, tp]
+        provided_ = self._provided[:, :, tp]
+
+        emb = self._emb(input_.unsqueeze(-1))  # [B, 1, H]
+        if embed_override is not None:
+            ov, rows = embed_override
+            emb = torch.where(rows.view(B, 1, 1), ov.to(emb.dtype), emb)
+        hidden, text_logits = self._temporal.step(emb)
+
+        sampled_text = text_logits.float().view(B, -1).argmax(-1)  # greedy
+        next_text = torch.where(provided_[:, 0], target_[:, 0], sampled_text)
+        audio = self._dep(next_text, hidden, audio_tokens=target_[:, 1:], audio_provided=provided_[:, 1:])  # [B, 16]
+
+        self._provided[:, :, mip] = False
+        self._cache[:, 0, tp] = torch.where(~provided_[:, 0], sampled_text, self._cache[:, 0, tp])
+        self._cache[:, 1:, tp] = torch.where(~provided_[:, 1:], audio, self._cache[:, 1:, tp])
+
+        max_delay = max(self._delays)
+        valid = lo > max_delay
+        if not bool(valid.any()):
+            self._phys += 1
+            self._local += 1
+            return True, None
+
+        idx = torch.tensor(
+            [(self._phys - max_delay + d) % CT for d in self._delays],
+            device=self._cache.device,
+            dtype=torch.long,
+        )
+        out = self._cache.gather(2, idx.view(1, K, 1).expand(B, K, 1)).squeeze(-1)  # [B, K]
+        self._last_valid = valid.clone()
+        self._phys += 1
+        self._local += 1
+        return True, out
+
+    def _boot_prefill(self, persona: str | None) -> None:
+        """Batch-wide hybrid system prompt: voice -> cache restore -> silence ->
+        persona -> silence (all rows identical)."""
+        import torch
+
+        B = self.config.batch_size
+        dev = self.config.device
+        dummy_user = self._initial[9:].view(1, 8).expand(B, 8)
+        dummy_moshi = self._initial[1:9].view(1, 8).expand(B, 8)
+        zero_text = torch.full((B,), ZERO_TEXT_TOKEN, dtype=torch.long, device=dev)
+        all_rows = torch.ones(B, dtype=torch.bool, device=dev)
+
+        for e in self._voice_embeddings:
+            ov = e.view(1, 1, -1).expand(B, 1, -1)
+            while True:
+                ran, _ = self._tick(dummy_user, dummy_moshi, zero_text, all_rows, embed_override=(ov, all_rows))
+                if ran:
+                    break
+        self._cache.copy_(self._voice_cache.expand_as(self._cache))
+
+        sine = self._sine.view(1, 8).expand(B, 8)
+        silence = self._silence.view(1, 8).expand(B, 8)
+        text = persona if persona is not None else self.config.persona
+        toks = self._tokenizer.encode(wrap_with_system_tags(text)) if text else []
+        for phase_toks in (
+            [ZERO_TEXT_TOKEN] * AUDIO_SILENCE_FRAME_CNT,
+            toks,
+            [ZERO_TEXT_TOKEN] * AUDIO_SILENCE_FRAME_CNT,
+        ):
+            for t in phase_toks:
+                tt = torch.full((B,), int(t), dtype=torch.long, device=dev)
+                self._tick(sine, silence, tt, all_rows)
+        self._codec.reset_streaming()
+
+    # -- batched surface (mirrors the previous engine's API) --------------------
+
+    def open_batch(self) -> None:
+        if not self._loaded:
+            self.load()
+        self._reset_all_streaming()
+        self._boot_prefill(None)
+        self._opened = True
+        logger.info("PersonaPlex native: %d slots prefilled", self.config.batch_size)
+
+    def prefill_steps(self, persona: str | None = None) -> list[PrefillStep]:
+        zero = self._silence
+        silence = [
+            PrefillStep("tokens", moshi_tokens=zero, text_token=ZERO_TEXT_TOKEN, user_sine=True)
+            for _ in range(AUDIO_SILENCE_FRAME_CNT)
+        ]
+        steps: list[PrefillStep] = [PrefillStep("sacrifice")]
+        for e in self._voice_embeddings:
+            steps.append(PrefillStep("voice", embedding=e.view(1, 1, -1)))
+        steps.append(PrefillStep("voice_end"))
+        steps.extend(silence)
+        text = persona if persona is not None else self.config.persona
+        if text:
+            for tok in self._tokenizer.encode(wrap_with_system_tags(text)):
+                steps.append(PrefillStep("tokens", moshi_tokens=zero, text_token=int(tok), user_sine=True))
+        steps.extend(silence)
+        return steps
+
+    def recycle_slot(self, b: int) -> None:
+        self._temporal.reset_slot(b)
+        self._temporal.bump_slot_start(b)  # LM sacrifice-tick semantics
+        self._local[b] = 0
+        self._cache[b] = self._initial.view(-1, 1).expand(-1, self._cache.shape[2])
+        self._provided[b] = False
+
+    def voice_end(self, b: int) -> None:
+        import torch
+
+        ct = self._cache.shape[2]
+        shift = int((self._phys - int(self._local[b])) % ct)
+        self._cache[b] = torch.roll(self._voice_cache[0], shifts=shift, dims=-1)
+
+    def end_prefill(self, b: int) -> None:
+        self._codec.reset_slot(b)
+
+    def step_batch(self, user_pcm: NDArray[np.float32], prefill: dict[int, PrefillStep] | None = None):
+        import torch
+
+        if not self._opened:
+            raise RuntimeError("open_batch() must be called before step_batch()")
+        cfg = self.config
+        B = cfg.batch_size
+        pcm = np.ascontiguousarray(user_pcm, dtype=np.float32).reshape(B, self.frame_size)
+        codes = self._codec.encode_frame(torch.from_numpy(pcm).to(cfg.device)).long()
+        return self._step_from_codes(codes, prefill)
+
+    def step_codes(self, user_codes, prefill: dict[int, PrefillStep] | None = None):
+        """Drive the lockstep loop from precomputed user codes (golden replays).
+
+        Kept public deliberately: bit-exact golden-replay verification must
+        bypass the Mimi encoder, whose streaming conv state makes encoded
+        codes run-order dependent (valid but not reproducible token-for-token).
+        """
+        return self._step_from_codes(user_codes.to(self.config.device).long(), prefill)
+
+    def _step_from_codes(self, codes, prefill):
+        import torch
+
+        from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput
+
+        cfg = self.config
+        B = cfg.batch_size
+        dev = cfg.device
+        forced_moshi = forced_text = force_mask = embed_override = None
+        if prefill:
+            force_mask = torch.zeros(B, dtype=torch.bool, device=dev)
+            embed_rows = torch.zeros(B, dtype=torch.bool, device=dev)
+            emb_full = None
+            forced_moshi = self._initial[1:9].view(1, 8).expand(B, 8).clone()
+            forced_text = torch.full((B,), ZERO_TEXT_TOKEN, dtype=torch.long, device=dev)
+            for b, stp in prefill.items():
+                if stp.kind == "voice_end":
+                    raise ValueError("voice_end is not a tick; call engine.voice_end(b)")
+                codes[b] = self._sine if stp.user_sine else self._initial[9:]
+                if stp.kind == "voice":
+                    embed_rows[b] = True
+                    if emb_full is None:
+                        emb_full = torch.zeros(B, 1, stp.embedding.shape[-1], dtype=stp.embedding.dtype, device=dev)
+                    emb_full[b] = stp.embedding[0]
+                    force_mask[b] = True  # dummy staging like step_embeddings
+                else:
+                    force_mask[b] = True
+                    if stp.kind == "tokens":
+                        forced_moshi[b] = stp.moshi_tokens
+                        forced_text[b] = int(stp.text_token)
+            if bool(embed_rows.any()):
+                embed_override = (emb_full, embed_rows)
+
+        ran, out = self._tick(codes, forced_moshi, forced_text, force_mask, embed_override)
+        self.last_out_tokens = None if out is None else out.detach().cpu()
+        if not ran or out is None:
+            return [None] * B
+
+        valid = self._last_valid  # [B] cpu bool
+        emit = [bool(valid[b]) and (not prefill or b not in prefill) for b in range(B)]
+        audio_tokens = out[:, 1:9].clone()
+        if not all(emit):
+            hide = torch.tensor([not e for e in emit], device=dev).view(B, 1)
+            audio_tokens = torch.where(hide, self._silence.view(1, 8).expand(B, 8), audio_tokens)
+        pcm_out = self._codec.decode_frame(audio_tokens)  # [B, frame]
+        outs = []
+        for b in range(B):
+            if not emit[b]:
+                outs.append(None)
+                continue
+            audio = pcm_out[b].float().detach().cpu().numpy().astype(np.float32)
+            outs.append(FrameOutput(audio=audio, text=self._decode_text(int(out[b, 0].item()))))
+        return outs
+
+    # -- single-session FrameStepper surface -------------------------------------
+
+    def open_session(self, voice_prompt: str | None = None, persona: str | None = None) -> None:
+        with self._fs_lock:
+            if not self._loaded:
+                self.load()
+            # Always resolve a voice: falling back to the configured default here
+            # keeps a previous session's custom voice from leaking into this one.
+            self._load_voice(voice_prompt or self.config.voice_prompt)
+            self._reset_all_streaming()
+            self._boot_prefill(persona)
+            self._opened = True
+
+    def step(self, user_pcm: NDArray[np.float32]):
+        with self._fs_lock:
+            pcm = np.ascontiguousarray(user_pcm, dtype=np.float32).reshape(-1)
+            if pcm.shape[0] != self.frame_size:
+                raise ValueError(f"expected {self.frame_size} samples per frame, got {pcm.shape[0]}")
+            outs = self.step_batch(pcm.reshape(1, -1))
+        from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput
+
+        return outs[0] if outs[0] is not None else FrameOutput(audio=None, text=None)
+
+    def _decode_text(self, token_id: int) -> str | None:
+        if token_id in NON_TEXT_TOKEN_IDS or token_id >= TEXT_VOCAB_SIZE:
+            return None
+        return self._tokenizer.id_to_piece(token_id).replace("▁", " ")

vllm_omni/experimental/fullduplex/personaplex/runtime_extension.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import base64
+import binascii
+from typing import Any
+
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.experimental.fullduplex.engine.duplex_runtime import (
+    DuplexAppendPlan,
+    DuplexInputMode,
+    DuplexOutputDecision,
+)
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+
+_FRAME_BYTES = 1920 * 4
+
+
+def _validated_frame_payload(payload: object) -> dict[str, object]:
+    if not isinstance(payload, dict):
+        raise ValueError("PersonaPlex duplex append payload must be a mapping")
+    if payload.get("format") != "pcm_f32le":
+        raise ValueError("PersonaPlex duplex append format must be pcm_f32le")
+    if payload.get("sample_rate_hz") != 24000:
+        raise ValueError("PersonaPlex duplex append sample_rate_hz must be 24000")
+    audio = payload.get("audio")
+    if not isinstance(audio, str):
+        raise ValueError("PersonaPlex duplex append audio must be base64 pcm_f32le")
+    try:
+        raw = base64.b64decode(audio, validate=True)
+    except (binascii.Error, ValueError) as exc:
+        raise ValueError("PersonaPlex duplex append audio is not valid base64") from exc
+    if len(raw) != _FRAME_BYTES:
+        raise ValueError("PersonaPlex duplex append must contain exactly 1920 samples")
+    return dict(payload)
+
+
+class PersonaPlexDuplexRuntimeExtension:
+    """PersonaPlex policy for the engine-owned resumable duplex request."""
+
+    def configure_sampling_params(
+        self,
+        *,
+        runtime_config: dict[str, Any],
+        defaults: tuple[object, ...],
+    ) -> tuple[object, ...]:
+        del runtime_config
+        if not defaults:
+            return defaults
+        configured = list(defaults)
+        stage0 = defaults[0]
+        if isinstance(stage0, SamplingParams):
+            stage0 = stage0.clone()
+            stage0.temperature = 0.0
+            stage0.top_k = 1
+            stage0.max_tokens = 1
+            configured[0] = stage0
+        return tuple(configured)
+
+    def plan_append(
+        self,
+        *,
+        request_id: str,
+        fence: DuplexFence,
+        session_config: dict[str, Any],
+        runtime_config: dict[str, Any],
+        seq: int,
+        turn_seq: int,
+        mode: DuplexInputMode,
+        payload: object,
+        final: bool,
+        sampling_params: object,
+    ) -> DuplexAppendPlan:
+        del sampling_params
+        if mode is not DuplexInputMode.APPEND_AUDIO_CHUNK:
+            raise ValueError(f"PersonaPlex does not support duplex input mode {mode.value!r}")
+        normalized_payload = _validated_frame_payload(payload)
+        prefill_slots = runtime_config.get("personaplex_prefill_slots", 0)
+        try:
+            prefill_slots = max(0, int(prefill_slots))
+        except (TypeError, ValueError) as exc:
+            raise ValueError("personaplex_prefill_slots must be a non-negative integer") from exc
+        prompt_slots = 1 + (prefill_slots if seq <= 1 else 0)
+        return DuplexAppendPlan(
+            prompt={
+                "prompt_token_ids": [0] * prompt_slots,
+                "model_intermediate_buffer": {
+                    "request_id": request_id,
+                    "global_request_id": [fence.session_id],
+                    "duplex": {
+                        "data_plane": True,
+                        "fence": fence,
+                        "session_id": fence.session_id,
+                        "incarnation": fence.incarnation,
+                        "epoch": fence.epoch,
+                        "turn_id": fence.turn_id,
+                        "response_seq": fence.response_seq,
+                        "seq": seq,
+                        "turn_seq": turn_seq,
+                        "mode": mode.value,
+                        "payload": normalized_payload,
+                        "final": final,
+                        "session_config": dict(session_config),
+                        "runtime_config": dict(runtime_config),
+                        "scheduler_token_budget": prompt_slots,
+                    },
+                },
+            }
+        )
+
+    def decide_output(
+        self,
+        *,
+        stage_id: int,
+        final_stage_id: int,
+        segment_finished: bool,
+        segment_token_ids: tuple[int, ...],
+        segment_output_metadata: dict[str, Any],
+        output: object,
+    ) -> DuplexOutputDecision | None:
+        del stage_id, final_stage_id, segment_finished
+        del segment_token_ids, segment_output_metadata, output
+        return None
+
+
+__all__ = ["PersonaPlexDuplexRuntimeExtension"]

vllm_omni/experimental/fullduplex/personaplex/serving/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex full-duplex WebSocket serving (PCM in / PCM + text out)."""
+
+from typing import Any
+
+__all__ = ["create_app", "main"]
+
+
+def __getattr__(name: str) -> Any:
+    if name not in __all__:
+        raise AttributeError(name)
+    from vllm_omni.experimental.fullduplex.personaplex.serving import server
+
+    return getattr(server, name)

vllm_omni/experimental/fullduplex/personaplex/serving/batched.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Slot manager + lockstep tick loop for batched PersonaPlex serving.
+
+Pure-python orchestration over a :class:`PersonaPlexEngine`-shaped object
+(injected, so tests run GPU-free with a stub): a fixed table of B conversation
+slots, one shared wall-clock 80 ms tick that advances every non-idle slot, and
+per-slot recycle prefill for new callers.
+
+Threading model: ``tick()`` runs on a dedicated stepping thread (the engine call
+is synchronous CUDA work); ``acquire``/``release``/``feed`` are called from the
+asyncio thread. All slot state is guarded by one lock; per-frame output is
+handed off via each slot's ``on_output`` callback (called on the stepping
+thread -- the server bridges to asyncio with ``call_soon_threadsafe``).
+"""
+
+from __future__ import annotations
+
+import enum
+import logging
+import threading
+import time
+from collections import deque
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class SlotPhase(enum.Enum):
+    IDLE = "idle"
+    PREFILL = "prefill"
+    LIVE = "live"
+
+
+@dataclass
+class Slot:
+    index: int
+    phase: SlotPhase = SlotPhase.IDLE
+    used: bool = False  # a used slot must be recycled before reuse
+    needs_recycle: bool = False  # engine.reset executed by the tick thread only
+    epoch: int = 0
+    pending: np.ndarray = field(default_factory=lambda: np.zeros(0, dtype=np.float32))
+    inq: deque = field(default_factory=deque)
+    prefill: deque = field(default_factory=deque)
+    on_output: Callable[[Any], None] | None = None  # FrameOutput, stepping thread
+    on_ready: Callable[[], None] | None = None  # prefill finished, stepping thread
+
+
+class BatchedSessionManager:
+    """Own the slot table and drive one engine tick per call to :meth:`tick`."""
+
+    def __init__(self, engine, max_buffered_frames: int = 25) -> None:
+        self.engine = engine
+        self.frame_size = engine.frame_size
+        self.batch_size = engine.config.batch_size
+        self.max_buffered_frames = max_buffered_frames
+        self.slots = [Slot(index=i) for i in range(self.batch_size)]
+        self.lock = threading.Lock()
+        # step_batch advances EVERY row each tick (lockstep), so once any tick has
+        # run, "pristine" idle rows have silently drifted from their boot prefill
+        # and must be recycled before a client may join them.
+        self._ticked = False
+
+    # -- connection side (asyncio thread) ------------------------------------
+
+    def acquire(self, persona: str | None = None) -> tuple[Slot, bool] | None:
+        """Claim an idle slot. Returns ``(slot, ready_now)`` or ``None`` if full.
+
+        A pristine slot (never used since boot, and claimed before the first
+        engine tick) still carries the default voice + persona from
+        ``open_batch`` and is live immediately when no custom persona is
+        requested. Otherwise the slot is recycled and its system prompt
+        replayed over the next ticks (``ready_now == False``) -- including
+        pristine rows after ticking has started, because lockstep stepping
+        advances every row whether or not a client is attached.
+        """
+        with self.lock:
+            slot = next((s for s in self.slots if s.phase is SlotPhase.IDLE), None)
+            if slot is None:
+                return None
+            slot.epoch += 1
+            slot.pending = np.zeros(0, dtype=np.float32)
+            slot.inq.clear()
+            slot.prefill.clear()
+            if not slot.used and persona is None and not self._ticked:
+                slot.phase = SlotPhase.LIVE
+                slot.used = True
+                return slot, True
+            # The engine reset itself must run BETWEEN engine steps, i.e. on the
+            # stepping thread -- only mark it here; tick() executes it.
+            slot.prefill.extend(self.engine.prefill_steps(persona))
+            slot.needs_recycle = True
+            slot.phase = SlotPhase.PREFILL
+            slot.used = True
+            return slot, False
+
+    def release(self, slot: Slot) -> None:
+        with self.lock:
+            slot.phase = SlotPhase.IDLE
+            slot.epoch += 1
+            slot.on_output = None
+            slot.on_ready = None
+            slot.inq.clear()
+            slot.prefill.clear()
+            slot.pending = np.zeros(0, dtype=np.float32)
+
+    def feed(self, slot: Slot, pcm: np.ndarray, epoch: int) -> None:
+        """Buffer incoming PCM and carve it into exact engine frames."""
+        with self.lock:
+            if slot.epoch != epoch or slot.phase is SlotPhase.IDLE:
+                return
+            samples = np.ascontiguousarray(pcm, dtype=np.float32).reshape(-1)
+            slot.pending = np.concatenate([slot.pending, samples]) if slot.pending.size else samples
+            n = self.frame_size
+            while slot.pending.shape[0] >= n:
+                frame, slot.pending = slot.pending[:n], slot.pending[n:]
+                slot.inq.append(frame)
+                while len(slot.inq) > self.max_buffered_frames:
+                    slot.inq.popleft()  # realtime: drop the oldest backlog
+
+    # -- stepping side (dedicated thread) ------------------------------------
+
+    def tick(self) -> bool:
+        """Advance every non-idle slot one frame. Returns False when all idle."""
+        with self.lock:
+            if all(s.phase is SlotPhase.IDLE for s in self.slots):
+                return False
+            self._ticked = True
+            pcm = np.zeros((self.batch_size, self.frame_size), dtype=np.float32)
+            prefill: dict[int, Any] = {}
+            emitters: list[tuple[Slot, int]] = []
+            for slot in self.slots:
+                if slot.phase is SlotPhase.LIVE:
+                    if slot.inq:
+                        pcm[slot.index] = slot.inq.popleft()
+                    emitters.append((slot, slot.epoch))
+                elif slot.phase is SlotPhase.PREFILL:
+                    if slot.needs_recycle:
+                        self.engine.recycle_slot(slot.index)
+                        slot.needs_recycle = False
+                    while slot.prefill and slot.prefill[0].kind == "voice_end":
+                        self.engine.voice_end(slot.index)
+                        slot.prefill.popleft()
+                    if slot.prefill:
+                        prefill[slot.index] = slot.prefill.popleft()
+                    else:
+                        self.engine.end_prefill(slot.index)
+                        slot.phase = SlotPhase.LIVE
+                        emitters.append((slot, slot.epoch))
+                        if slot.on_ready is not None:
+                            slot.on_ready()
+
+        outs = self.engine.step_batch(pcm, prefill or None)
+
+        with self.lock:
+            for slot, epoch in emitters:
+                out = outs[slot.index]
+                if out is None or slot.epoch != epoch or slot.on_output is None:
+                    continue
+                slot.on_output(out)
+        return True
+
+    def run_forever(self, stop: threading.Event, frame_seconds: float = 0.08) -> None:
+        """Wall-clock-paced stepping loop (run on a dedicated thread)."""
+        next_t = time.monotonic()
+        while not stop.is_set():
+            did = self.tick()
+            if not did:
+                time.sleep(frame_seconds / 4)
+                next_t = time.monotonic()
+                continue
+            next_t += frame_seconds
+            delay = next_t - time.monotonic()
+            if delay > 0:
+                time.sleep(delay)
+            elif delay < -1.0:
+                # Fell far behind realtime (e.g. B too large for the GPU); resync
+                # rather than bursting to catch up.
+                logger.warning("batched tick loop %.1fs behind realtime; resyncing", -delay)
+                next_t = time.monotonic()

vllm_omni/experimental/fullduplex/personaplex/serving/server.py

@@ -0,0 +1,539 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex full-duplex server, compatible with the OFFICIAL Moshi web client.
+
+Serves the official PersonaPlex web UI (``dist.tgz`` from ``nvidia/personaplex-7b-v1``,
+the same bundle ``moshi.server`` downloads) and implements its WebSocket protocol at
+``/api/chat`` -- but driven by our native-component engine instead of moshi's server.
+
+Built on **aiohttp**, mirroring ``moshi.server`` exactly: aiohttp writes+drains each
+WebSocket frame promptly, whereas uvicorn's ``websockets`` backend coalesces many small
+(~80 ms Opus) frames into a send buffer, which makes browser playback choppy even when
+the engine is real-time. aiohttp ships as a moshi dependency, so this adds nothing new.
+
+Moshi protocol (binary WS messages, first byte = tag):
+  server -> client  ``\\x00``                 handshake (ready, after system prompts)
+  client -> server  ``\\x01`` + opus bytes    mic audio (Opus @ 24kHz)
+  server -> client  ``\\x01`` + opus bytes    agent audio (Opus @ 24kHz)
+  server -> client  ``\\x02`` + utf8          inner-monologue text
+Query params on connect: ``text_prompt`` (persona), ``voice_prompt`` (voice file).
+
+Run (auto-downloads the official UI; open http://<host>:8124/):
+    HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 python -m \\
+        vllm_omni.experimental.fullduplex.personaplex.serving.server --port 8124
+
+A raw-PCM endpoint (``/v1/audio/duplex``, JSON ``open`` + float32 frames) remains for
+tooling/tests that do not want an Opus dependency.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import contextlib
+import functools
+import json
+import logging
+import tarfile
+import threading
+from pathlib import Path
+
+import numpy as np
+import sphn
+from aiohttp import WSMsgType, web
+from huggingface_hub import hf_hub_download
+
+from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig
+from vllm_omni.experimental.fullduplex.personaplex.engine import PersonaPlexEngine
+from vllm_omni.experimental.fullduplex.personaplex.serving.batched import BatchedSessionManager
+from vllm_omni.experimental.fullduplex.personaplex.session import (
+    PersonaPlexServingSessionState,
+    PersonaPlexSession,
+    offer_realtime,
+)
+
+logger = logging.getLogger(__name__)
+_HF_REPO = "nvidia/personaplex-7b-v1"
+
+
+def _opus_encode(writer: sphn.OpusStreamWriter, pcm: np.ndarray) -> bytes:
+    """Encode one PCM frame to Opus bytes, across sphn API versions.
+
+    sphn 0.2.x returns the encoded bytes straight from ``append_pcm``; the older
+    0.1.x buffers and requires a ``read_bytes`` drain. Handle both.
+    """
+    out = writer.append_pcm(np.ascontiguousarray(pcm, dtype=np.float32))
+    if out is None and hasattr(writer, "read_bytes"):
+        out = writer.read_bytes()
+    return out or b""
+
+
+def _opus_decode(reader: sphn.OpusStreamReader, data: bytes) -> np.ndarray | None:
+    """Decode Opus bytes to PCM, across sphn API versions (see ``_opus_encode``)."""
+    out = reader.append_bytes(data)
+    if out is None and hasattr(reader, "read_pcm"):
+        out = reader.read_pcm()
+    if out is None or getattr(out, "shape", (0,))[-1] == 0:
+        return None
+    return np.asarray(out, dtype=np.float32)
+
+
+def _official_web_dir() -> Path | None:
+    """Download + extract the official PersonaPlex web client (dist.tgz)."""
+    try:
+        tgz = Path(hf_hub_download(_HF_REPO, "dist.tgz"))
+        dist = tgz.parent / "dist"
+        if not dist.exists():
+            with tarfile.open(tgz, "r:gz") as tar:
+                tar.extractall(path=tgz.parent)  # noqa: S202 - trusted first-party bundle
+        return dist if dist.exists() else None
+    except Exception as exc:  # network / auth / layout
+        logger.warning("official web client unavailable (%s); / will return an error note", exc)
+        return None
+
+
+class DuplexServer:
+    """Owns the loaded engine; serves one duplex conversation at a time.
+
+    A capacity-one lease covers the entire WebSocket lifetime on BOTH endpoints:
+    the engine is a single shared streaming state, so a second concurrent
+    connection would reset and corrupt the live conversation. Extra connections
+    are rejected with 1013 (try again later). Engine work (``open``/``feed``,
+    synchronous CUDA) runs on a worker thread so the event loop stays free for
+    WebSocket traffic and close handling.
+    """
+
+    def __init__(self, config: PersonaPlexConfig) -> None:
+        self.config = config
+        self.engine = PersonaPlexEngine(config)
+        self._active = False  # capacity-one session lease (event-loop confined)
+        self._lease_release_tasks: set[asyncio.Task] = set()
+
+    def load(self) -> None:
+        self.engine.load()
+
+    async def _lease(self, ws: web.WebSocketResponse) -> bool:
+        """Try to claim the single-session lease; reject the connection if held."""
+        if self._active:
+            await ws.close(code=1013, message=b"another conversation is active")
+            return False
+        self._active = True
+        return True
+
+    async def _engine_call(self, state: PersonaPlexServingSessionState, fn, *args, **kwargs):
+        """Run one synchronous engine call on a worker thread, tracked for drain.
+
+        Cancelling the awaiting task does NOT stop the worker thread, and it
+        marks the asyncio wrapper future cancelled (``done()``) even while the
+        thread keeps running, so the future itself cannot signal completion.
+        The worker instead sets a plain ``threading.Event`` when the call truly
+        finishes; :meth:`_drain_inflight` waits on that event before the lease
+        is released. Without this, a fast reconnect could acquire the lease and
+        call ``open()`` while the previous connection's final ``step()`` is
+        still mutating the shared engine on its thread.
+        """
+        call = functools.partial(fn, *args, **kwargs)
+        finished = threading.Event()
+
+        def run_and_signal():
+            try:
+                return call()
+            finally:
+                finished.set()
+
+        state.inflight = finished
+        result = await asyncio.get_running_loop().run_in_executor(None, run_and_signal)
+        state.inflight = None  # on cancellation, the event stays referenced for drain
+        return result
+
+    @staticmethod
+    async def _drain_inflight(state: PersonaPlexServingSessionState, timeout: float = 5.0) -> bool:
+        """Wait out any engine call still running on a worker thread."""
+        finished = state.inflight
+        if finished is None or finished.is_set():
+            return True
+        return bool(await asyncio.to_thread(finished.wait, timeout))
+
+    async def _release_lease_when_finished(self, state: PersonaPlexServingSessionState) -> None:
+        finished = state.inflight
+        if finished is not None and not finished.is_set():
+            await asyncio.to_thread(finished.wait)
+        self._active = False
+
+    async def _release_lease_after_drain(
+        self,
+        state: PersonaPlexServingSessionState,
+        *,
+        timeout: float = 5.0,
+    ) -> bool:
+        """Release now when drained, otherwise retain the lease until completion."""
+        if await self._drain_inflight(state, timeout):
+            self._active = False
+            return True
+
+        logger.error(
+            "PersonaPlex engine cleanup exceeded %.1fs; retaining the session lease until the worker thread finishes",
+            timeout,
+        )
+        task = asyncio.create_task(self._release_lease_when_finished(state))
+        self._lease_release_tasks.add(task)
+        task.add_done_callback(self._lease_release_tasks.discard)
+        return False
+
+    # ---- official Moshi protocol (Opus over aiohttp WS) ---------------------
+
+    async def handle_chat(self, request: web.Request) -> web.WebSocketResponse:
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        if not await self._lease(ws):
+            return ws
+        state = PersonaPlexServingSessionState()
+        try:
+            persona = (request.query.get("text_prompt") or "").strip() or None
+            voice = (request.query.get("voice_prompt") or "").strip() or None
+            sr = self.engine.sample_rate
+            session = PersonaPlexSession(self.engine, self.config)
+            # System-prompt prefill is seconds of GPU work; keep it off the loop.
+            await self._engine_call(state, session.open, voice_prompt=voice, persona=persona)
+
+            reader = sphn.OpusStreamReader(sr)
+            writer = sphn.OpusStreamWriter(sr)
+            await ws.send_bytes(b"\x00")  # ready handshake
+
+            async def recv_loop() -> None:
+                try:
+                    async for msg in ws:
+                        if msg.type == WSMsgType.BINARY:
+                            data = msg.data
+                            if data and data[0] == 1:
+                                pcm = _opus_decode(reader, data[1:])
+                                if pcm is not None:
+                                    offer_realtime(state.pcm_queue, pcm)
+                        elif msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.ERROR):
+                            break
+                finally:
+                    state.closed = True
+
+            async def proc_loop() -> None:
+                # Drain decoded PCM and step the model, streaming agent Opus + text
+                # out. The synchronous engine step runs on a worker thread.
+                while not state.closed:
+                    try:
+                        pcm = await asyncio.wait_for(state.pcm_queue.get(), timeout=0.25)
+                    except asyncio.TimeoutError:
+                        continue
+                    for frame_out in await self._engine_call(state, session.feed, pcm):
+                        if frame_out.audio is not None:
+                            opus = _opus_encode(writer, frame_out.audio)
+                            if opus:
+                                await ws.send_bytes(b"\x01" + opus)
+                        if frame_out.text:
+                            await ws.send_bytes(b"\x02" + frame_out.text.encode("utf8"))
+
+            tasks = [asyncio.create_task(recv_loop()), asyncio.create_task(proc_loop())]
+            try:
+                _, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+                for t in pending:
+                    t.cancel()
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await t
+            finally:
+                with contextlib.suppress(Exception):
+                    await ws.close()
+        finally:
+            # A cancelled feed/open may still be running on its worker thread;
+            # wait it out BEFORE freeing the lease, or a fast reconnect races
+            # the shared engine (CUDA illegal access / cross-session bleed).
+            await self._release_lease_after_drain(state)
+        return ws
+
+    # ---- simple raw-PCM protocol (no Opus, for tests) ----------------------
+
+    async def handle_raw(self, request: web.Request) -> web.WebSocketResponse:
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        if not await self._lease(ws):
+            return ws
+        session = PersonaPlexSession(self.engine, self.config)
+        state = PersonaPlexServingSessionState()
+        opened = False
+        try:
+            async for msg in ws:
+                if msg.type == WSMsgType.TEXT:
+                    req = json.loads(msg.data)
+                    if req.get("type") == "open":
+                        await self._engine_call(
+                            state, session.open, voice_prompt=req.get("voice"), persona=req.get("persona")
+                        )
+                        opened = True
+                        await ws.send_json({"type": "ready"})
+                    elif req.get("type") == "close":
+                        if opened:
+                            await self._emit(ws, await self._engine_call(state, session.flush))
+                        await ws.send_json({"type": "done"})
+                        break
+                elif msg.type == WSMsgType.BINARY:
+                    if not opened:
+                        await self._engine_call(state, session.open)
+                        opened = True
+                    pcm = np.frombuffer(msg.data, dtype=np.float32)
+                    await self._emit(ws, await self._engine_call(state, session.feed, pcm))
+                elif msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.ERROR):
+                    break
+        finally:
+            # Same as handle_chat: wait out any worker-thread engine call
+            # before freeing the lease (see _engine_call).
+            await self._release_lease_after_drain(state)
+            with contextlib.suppress(Exception):
+                await ws.close()
+        return ws
+
+    @staticmethod
+    async def _emit(ws: web.WebSocketResponse, outputs: list) -> None:
+        for frame_out in outputs:
+            if frame_out.audio is not None:
+                await ws.send_bytes(np.ascontiguousarray(frame_out.audio, dtype=np.float32).tobytes())
+            if frame_out.text:
+                await ws.send_json({"type": "text", "text": frame_out.text})
+
+
+class BatchedDuplexServer:
+    """Serve ``batch_size`` concurrent duplex conversations on one engine.
+
+    Same wire protocols as :class:`DuplexServer`; each connection claims a slot
+    from the shared :class:`BatchedSessionManager`, whose stepping thread ticks
+    all live slots in lockstep at 12.5 Hz wall clock. A connection with a custom
+    ``text_prompt`` (or reusing a previously used slot) waits while the slot's
+    system prompt is replayed per-row (voice embeddings + persona), then gets the
+    ready handshake. Per-connection custom voices are not supported in batched
+    mode (fixed default voice; ``voice_prompt`` is ignored with a warning).
+    """
+
+    def __init__(self, config: PersonaPlexConfig) -> None:
+        self.config = config
+        self.engine = PersonaPlexEngine(config)
+        self.manager: BatchedSessionManager | None = None
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def load(self) -> None:
+        self.engine.load()
+        self.engine.open_batch()
+        self.manager = BatchedSessionManager(self.engine)
+        self._thread = threading.Thread(
+            target=self.manager.run_forever, args=(self._stop,), daemon=True, name="pplex-tick"
+        )
+        self._thread.start()
+
+    def shutdown(self) -> None:
+        self._stop.set()
+        if self._thread is not None:
+            self._thread.join(timeout=2.0)
+
+    async def _acquire(self, ws: web.WebSocketResponse, persona: str | None):
+        """Claim a slot and wait until it is live.
+
+        Returns ``PersonaPlexServingSessionState`` (slot + bounded output queue)
+        or ``None`` when the server is full. Exception-safe: if this coroutine
+        is cancelled (client gone mid-prefill) or the wait fails, the slot is
+        released instead of leaking in the PREFILL/LIVE phase.
+        """
+        acquired = self.manager.acquire(persona=persona)
+        if acquired is None:
+            await ws.close(code=1013, message=b"all conversation slots busy")
+            return None
+        slot, ready_now = acquired
+        try:
+            loop = asyncio.get_running_loop()
+            state = PersonaPlexServingSessionState(slot=slot, epoch=slot.epoch)
+            ready = asyncio.Event()
+            with self.manager.lock:
+                # Bounded handoff: the tick thread produces one frame per 80 ms;
+                # a client that stops reading gets drop-oldest, not unbounded RAM.
+                slot.on_output = lambda frame_out: loop.call_soon_threadsafe(offer_realtime, state.out_queue, frame_out)
+                if ready_now:
+                    ready.set()
+                else:
+                    slot.on_ready = lambda: loop.call_soon_threadsafe(ready.set)
+            await ready.wait()
+        except BaseException:
+            self.manager.release(slot)
+            raise
+        return state
+
+    async def handle_chat(self, request: web.Request) -> web.WebSocketResponse:
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        persona = (request.query.get("text_prompt") or "").strip() or None
+        if (request.query.get("voice_prompt") or "").strip():
+            logger.warning("batched mode ignores per-connection voice_prompt (fixed default voice)")
+        state = await self._acquire(ws, persona)
+        if state is None:
+            return ws
+        try:
+            sr = self.engine.sample_rate
+            reader = sphn.OpusStreamReader(sr)
+            writer = sphn.OpusStreamWriter(sr)
+            await ws.send_bytes(b"\x00")  # ready handshake
+
+            async def recv_loop() -> None:
+                try:
+                    async for msg in ws:
+                        if msg.type == WSMsgType.BINARY:
+                            data = msg.data
+                            if data and data[0] == 1:
+                                pcm = _opus_decode(reader, data[1:])
+                                if pcm is not None:
+                                    self.manager.feed(state.slot, pcm, state.epoch)
+                        elif msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.ERROR):
+                            break
+                finally:
+                    state.closed = True
+
+            async def out_loop() -> None:
+                while not state.closed:
+                    try:
+                        frame_out = await asyncio.wait_for(state.out_queue.get(), timeout=0.25)
+                    except asyncio.TimeoutError:
+                        continue
+                    if frame_out.audio is not None:
+                        msg = _opus_encode(writer, frame_out.audio)
+                        if msg:
+                            await ws.send_bytes(b"\x01" + msg)
+                    if frame_out.text:
+                        await ws.send_bytes(b"\x02" + frame_out.text.encode("utf8"))
+
+            tasks = [asyncio.create_task(recv_loop()), asyncio.create_task(out_loop())]
+            _, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+            for t in pending:
+                t.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await t
+        finally:
+            self.manager.release(state.slot)
+            with contextlib.suppress(Exception):
+                await ws.close()
+        return ws
+
+    async def handle_raw(self, request: web.Request) -> web.WebSocketResponse:
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        state: PersonaPlexServingSessionState | None = None
+        sender: asyncio.Task | None = None
+
+        async def pump(outq: asyncio.Queue) -> None:
+            # Forward slot output continuously (the tick loop produces in realtime).
+            while True:
+                frame_out = await outq.get()
+                if frame_out.audio is not None:
+                    await ws.send_bytes(np.ascontiguousarray(frame_out.audio, dtype=np.float32).tobytes())
+                if frame_out.text:
+                    await ws.send_json({"type": "text", "text": frame_out.text})
+
+        try:
+            async for msg in ws:
+                if msg.type == WSMsgType.TEXT:
+                    req = json.loads(msg.data)
+                    if req.get("type") == "open":
+                        if state is not None:
+                            # One slot lease per connection; a second `open` would
+                            # silently leak the first slot and its pump.
+                            await ws.send_json({"type": "error", "error": "session already open"})
+                            continue
+                        state = await self._acquire(ws, (req.get("persona") or "").strip() or None)
+                        if state is None:
+                            return ws
+                        sender = asyncio.create_task(pump(state.out_queue))
+                        await ws.send_json({"type": "ready"})
+                    elif req.get("type") == "close":
+                        await asyncio.sleep(0.3)  # let in-flight frames flush via pump
+                        await ws.send_json({"type": "done"})
+                        break
+                elif msg.type == WSMsgType.BINARY and state is not None:
+                    self.manager.feed(state.slot, np.frombuffer(msg.data, dtype=np.float32), state.epoch)
+                elif msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.ERROR):
+                    break
+        finally:
+            # Release FIRST and unconditionally: a pump that already failed (e.g.
+            # peer reset mid-send) re-raises on await and must not leak the slot.
+            if state is not None:
+                self.manager.release(state.slot)
+            if sender is not None:
+                sender.cancel()
+                try:
+                    await sender
+                except asyncio.CancelledError:
+                    pass
+                except Exception:
+                    logger.debug("raw output pump ended with an error", exc_info=True)
+            with contextlib.suppress(Exception):
+                await ws.close()
+        return ws
+
+
+def create_app(config: PersonaPlexConfig) -> web.Application:
+    server = BatchedDuplexServer(config) if config.batch_size > 1 else DuplexServer(config)
+    web_dir = _official_web_dir()
+
+    async def _startup(_app: web.Application) -> None:
+        logger.info("PersonaPlex duplex server: loading engine (batch_size=%d)", config.batch_size)
+        server.load()
+        logger.info("PersonaPlex duplex server: ready (official web: %s)", bool(web_dir))
+
+    async def _cleanup(_app: web.Application) -> None:
+        if isinstance(server, BatchedDuplexServer):
+            server.shutdown()
+
+    app = web.Application()
+    app.on_startup.append(_startup)
+    app.on_cleanup.append(_cleanup)
+
+    async def health(_request: web.Request) -> web.Response:
+        return web.json_response({"status": "ok"})
+
+    app.router.add_get("/health", health)
+    app.router.add_get("/api/chat", server.handle_chat)
+    app.router.add_get("/v1/audio/duplex", server.handle_raw)
+
+    # Serve the official PersonaPlex web client (explicit routes are matched first).
+    if web_dir is not None:
+
+        async def index(_request: web.Request) -> web.FileResponse:
+            return web.FileResponse(web_dir / "index.html")
+
+        app.router.add_get("/", index)
+        app.router.add_static("/", path=str(web_dir), follow_symlinks=True, name="web")
+    else:
+
+        async def index_missing(_request: web.Request) -> web.Response:
+            return web.json_response({"error": f"official web client unavailable; check access to {_HF_REPO}"})
+
+        app.router.add_get("/", index_missing)
+
+    return app
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--host", default="0.0.0.0")
+    ap.add_argument("--port", type=int, default=8124)
+    ap.add_argument("--persona", default=None)
+    ap.add_argument("--voice", default="NATF2.pt")
+    ap.add_argument(
+        "--batch-size",
+        type=int,
+        default=1,
+        help="Concurrent conversation slots on one engine (elastic batching when > 1).",
+    )
+    args = ap.parse_args()
+    logging.basicConfig(level=logging.INFO)
+    cfg = PersonaPlexConfig(
+        voice_prompt=args.voice,
+        batch_size=args.batch_size,
+        **({"persona": args.persona} if args.persona else {}),
+    )
+    web.run_app(create_app(cfg), host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

vllm_omni/experimental/fullduplex/personaplex/serving_adapter.py

@@ -0,0 +1,233 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, Mapping
+from copy import deepcopy
+from dataclasses import dataclass, field
+from pathlib import PurePath
+from typing import Any
+
+from vllm_omni.experimental.fullduplex.openai.protocol import (
+    DuplexCapabilities,
+)
+from vllm_omni.experimental.fullduplex.openai.runtime_adapter import (
+    ServingRuntimeConfigError,
+)
+from vllm_omni.experimental.fullduplex.personaplex.config import DEFAULT_PERSONA
+from vllm_omni.experimental.fullduplex.personaplex.data_plane import (
+    PersonaPlexDataPlaneContext,
+    PersonaPlexDataPlaneSession,
+)
+from vllm_omni.experimental.fullduplex.personaplex.input import (
+    PersonaPlexPcmAppendBuffer,
+)
+
+EncodeAudio = Callable[[object, int, str, float | None], str | None]
+
+_PRIVATE_RUNTIME_CONFIG_KEYS = frozenset(
+    {
+        "personaplex_prefill_slots",
+        "personaplex_model_path",
+    }
+)
+
+
+@dataclass(slots=True)
+class PersonaPlexServingSessionState:
+    audio_buffer: PersonaPlexPcmAppendBuffer = field(default_factory=PersonaPlexPcmAppendBuffer)
+    input_since_commit: bool = False
+    speech_since_commit: bool = False
+    committed_audio_payload: dict[str, object] | None = None
+    committed_audio_operation_id: str | None = None
+    committed_audio_reserved_bytes: int = 0
+    deferred_response_create: bool = False
+    deferred_precreate_response: bool = False
+    data_plane_task: asyncio.Task[None] | None = None
+    data_plane_restart_requested: bool = False
+    continuation_owner_id: str | None = None
+    continuation_units: int = 0
+    pending_silence_task: asyncio.Task[bool] | None = None
+    pending_silence_owner_id: str | None = None
+    silence_continuation_scheduler: Callable[..., Awaitable[bool]] | None = None
+
+    def retain_committed_audio(
+        self,
+        payload: dict[str, object],
+        *,
+        operation_id: str | None,
+        reserved_bytes: int = 0,
+    ) -> None:
+        self.committed_audio_payload = payload
+        self.committed_audio_operation_id = operation_id
+        self.committed_audio_reserved_bytes += max(0, int(reserved_bytes))
+
+    def clear_committed_audio(self) -> int:
+        reserved_bytes = self.committed_audio_reserved_bytes
+        self.committed_audio_payload = None
+        self.committed_audio_operation_id = None
+        self.committed_audio_reserved_bytes = 0
+        self.deferred_response_create = False
+        self.deferred_precreate_response = False
+        return reserved_bytes
+
+    def clear_continuation(self) -> None:
+        self.continuation_owner_id = None
+        self.continuation_units = 0
+        self.pending_silence_task = None
+        self.pending_silence_owner_id = None
+
+
+class PersonaPlexServingRuntimeAdapter:
+    adapter_id = "personaplex"
+    clean_response_done_prefix = ""
+    interrupted_tts_prefix = ""
+    private_runtime_config_keys = _PRIVATE_RUNTIME_CONFIG_KEYS
+
+    def __init__(self, encode_audio: EncodeAudio) -> None:
+        self.session_states: dict[str, PersonaPlexServingSessionState] = {}
+        self.data_plane = PersonaPlexDataPlaneSession(encode_audio)
+
+    def create_session_state(self) -> PersonaPlexServingSessionState:
+        return PersonaPlexServingSessionState()
+
+    def session_state(self, session_id: str) -> PersonaPlexServingSessionState:
+        return self.session_states.setdefault(session_id, self.create_session_state())
+
+    def remove_session_state(self, session_id: str) -> None:
+        self.session_states.pop(session_id, None)
+
+    @staticmethod
+    def is_enabled(config: object) -> bool:
+        del config
+        return True
+
+    @staticmethod
+    def capabilities(*, max_sessions: int) -> DuplexCapabilities:
+        supports_multi_session = max_sessions > 1
+        return DuplexCapabilities(
+            supports_model_native_turn_policy=True,
+            supports_external_turn_signal=False,
+            supports_client_commit=False,
+            supports_barge_in=False,
+            supports_playback_ack=True,
+            supports_input_append=True,
+            supports_replace_latest_chunk=False,
+            supports_reencode_context=False,
+            supports_rollback_to_checkpoint=False,
+            supports_turn_commit_only=False,
+            supports_model_internal_state=True,
+            supports_stage_resumption=True,
+            supports_core_resumable_request=True,
+            supports_stage_connector_handoff=True,
+            supports_independent_io_streams=True,
+            supports_realtime_endpoint=True,
+            supports_multi_session=supports_multi_session,
+            supports_multi_session_same_replica=supports_multi_session,
+            supports_session_lease=True,
+            supports_session_resume=False,
+            session_admission_mode="engine_managed",
+            supports_audio_truncate=False,
+            requires_model_runner_kv=True,
+            requires_native_stage_role=True,
+            implementation_level="model_native_duplex",
+            adapter_patterns=["scheduler_data_plane"],
+            input_modes=["append_audio_chunk"],
+            signal_sources=["model_native", "client_event"],
+            stage_handoff_transport="scheduler_data_plane",
+            chunk_period_ms=80,
+            target_barge_in_latency_ms=None,
+        )
+
+    @staticmethod
+    def validate_client_extra_body(extra_body: object) -> None:
+        if not isinstance(extra_body, dict):
+            return
+        private = sorted(_PRIVATE_RUNTIME_CONFIG_KEYS.intersection(extra_body))
+        if private:
+            raise ServingRuntimeConfigError("PersonaPlex runtime configuration is server-owned: " + ", ".join(private))
+
+    @classmethod
+    async def prepare_runtime_config(
+        cls,
+        config: object,
+        *,
+        model_config: Any,
+    ) -> dict[str, object]:
+        extra_body = getattr(config, "extra_body", None)
+        cls.validate_client_extra_body(extra_body)
+        model_path = getattr(model_config, "model", None)
+        if not isinstance(model_path, str) or not model_path:
+            raise ServingRuntimeConfigError("PersonaPlex model path is unavailable")
+        voice = cls._voice_name(getattr(config, "voice", None))
+        instructions = getattr(config, "instructions", None) or DEFAULT_PERSONA
+        from vllm_omni.experimental.fullduplex.personaplex.stage0 import (
+            personaplex_prefill_slots,
+        )
+
+        try:
+            prefill_slots = await asyncio.to_thread(
+                personaplex_prefill_slots,
+                model_path,
+                voice,
+                str(instructions),
+            )
+        except Exception as exc:
+            raise ServingRuntimeConfigError(f"PersonaPlex voice/persona prefill could not be prepared: {exc}") from exc
+        return {
+            "personaplex_model_path": model_path,
+            "personaplex_voice_prompt": voice,
+            "personaplex_persona": str(instructions),
+            "personaplex_prefill_slots": prefill_slots,
+        }
+
+    @classmethod
+    def runtime_config_for_update(
+        cls,
+        config: object,
+        current: Mapping[str, object],
+    ) -> dict[str, object]:
+        cls.validate_client_extra_body(getattr(config, "extra_body", None))
+        runtime_config = deepcopy(dict(current))
+        runtime_config["personaplex_voice_prompt"] = cls._voice_name(getattr(config, "voice", None))
+        runtime_config["personaplex_persona"] = str(getattr(config, "instructions", None) or DEFAULT_PERSONA)
+        return runtime_config
+
+    @staticmethod
+    def data_plane_context(
+        *,
+        epoch: int,
+        turn_id: int,
+        active_response_turn_id: int | None,
+        active_response_id: str | None,
+        auto_responds: bool,
+        response_format: str,
+        speed: float | None,
+        modalities: tuple[str, ...],
+    ) -> PersonaPlexDataPlaneContext:
+        return PersonaPlexDataPlaneContext(
+            epoch=epoch,
+            turn_id=turn_id,
+            active_response_turn_id=active_response_turn_id,
+            active_response_id=active_response_id,
+            auto_responds=auto_responds,
+            response_format=response_format,
+            speed=speed,
+            modalities=modalities,
+        )
+
+    @staticmethod
+    def _voice_name(value: object) -> str:
+        voice = value if isinstance(value, str) and value else "NATF2.pt"
+        path = PurePath(voice)
+        if path.name != voice or path.suffix != ".pt" or any(part == ".." for part in path.parts):
+            raise ServingRuntimeConfigError("PersonaPlex voice must be a bundled .pt basename")
+        return voice
+
+
+__all__ = [
+    "PersonaPlexServingRuntimeAdapter",
+    "PersonaPlexServingSessionState",
+]

vllm_omni/experimental/fullduplex/personaplex/session.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex lockstep session driver — the runnable serving primitive.
+
+Mirrors JoyVL's ``serving/session.py``: a thin, framework-free driver that owns a
+``FrameStepper`` and turns arbitrary-length PCM chunks from a client into exact
+80 ms frames, stepping the model one frame at a time. This is the path that
+actually runs end to end (offline driver and any WS server build on it); the
+``DuplexAdapter`` is a separate demonstration of the generic framework seam.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig
+from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput, FrameStepper
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+# Realtime queue bounds: 80 ms/frame, so 50 frames = 4 s of backlog. Anything
+# beyond that means the peer cannot keep up and old frames are already useless.
+MAX_PENDING_INPUT = 50
+MAX_PENDING_OUTPUT = 50
+
+
+def offer_realtime(queue: asyncio.Queue, item: Any) -> None:
+    """Enqueue with a drop-oldest policy (realtime streams must never grow unbounded)."""
+    while True:
+        try:
+            queue.put_nowait(item)
+            return
+        except asyncio.QueueFull:
+            with contextlib.suppress(asyncio.QueueEmpty):
+                queue.get_nowait()
+
+
+@dataclass
+class PersonaPlexServingSessionState:
+    """Mutable per-connection serving state (the minicpmo45 ``session.py`` role).
+
+    One instance per WebSocket connection, for both server modes: the
+    single-session server uses ``pcm_queue`` (decoded mic PCM awaiting the
+    engine); the batched server uses ``out_queue`` + ``slot``/``epoch`` (engine
+    output awaiting send, and the claimed slot lease). ``closed`` tears both
+    pump loops down.
+    """
+
+    pcm_queue: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=MAX_PENDING_INPUT))
+    out_queue: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=MAX_PENDING_OUTPUT))
+    closed: bool = False
+    slot: Any = None  # batched mode: the claimed serving/batched.py Slot
+    epoch: int = 0  # batched mode: slot epoch at claim time (stale-guard)
+    # Single-session mode: completion event of the engine call currently running
+    # on a worker thread. Task cancellation cannot stop that thread (and marks
+    # the asyncio future cancelled while the thread runs on), so the server
+    # waits on this event before releasing its lease (DuplexServer._engine_call).
+    inflight: Any = None
+
+
+class PersonaPlexSession:
+    """Drive one PersonaPlex conversation in strict frame lockstep.
+
+    A session buffers incoming PCM and emits exactly one :class:`FrameOutput` per
+    completed ``frame_size`` window. Partial tails are held until enough samples
+    arrive (or ``flush()`` pads them out at end of stream).
+    """
+
+    def __init__(self, stepper: FrameStepper, config: PersonaPlexConfig | None = None) -> None:
+        self.stepper = stepper
+        self.config = config or PersonaPlexConfig()
+        self.frame_size = stepper.frame_size
+        self._pending = np.zeros(0, dtype=np.float32)
+        self._opened = False
+
+    def open(self, voice_prompt: str | None = None, persona: str | None = None) -> None:
+        """Inject the voice clone + persona and arm the conversation."""
+        self.stepper.open_session(voice_prompt=voice_prompt, persona=persona)
+        self._pending = np.zeros(0, dtype=np.float32)
+        self._opened = True
+
+    def feed(self, pcm: NDArray[np.float32]) -> list[FrameOutput]:
+        """Append PCM and step every newly completed frame.
+
+        Returns one :class:`FrameOutput` per frame consumed from the buffer this
+        call (possibly empty if fewer than ``frame_size`` samples are buffered).
+        """
+        if not self._opened:
+            raise RuntimeError("PersonaPlexSession.open() must be called before feed()")
+        samples = np.ascontiguousarray(pcm, dtype=np.float32).reshape(-1)
+        self._pending = np.concatenate([self._pending, samples]) if self._pending.size else samples
+
+        outputs: list[FrameOutput] = []
+        n = self.frame_size
+        while self._pending.shape[0] >= n:
+            frame, self._pending = self._pending[:n], self._pending[n:]
+            outputs.append(self.stepper.step(frame))
+        return outputs
+
+    def flush(self) -> list[FrameOutput]:
+        """Zero-pad and step any buffered tail; call once at end of input."""
+        if self._pending.shape[0] == 0:
+            return []
+        pad = self.frame_size - self._pending.shape[0]
+        frame = np.concatenate([self._pending, np.zeros(pad, dtype=np.float32)])
+        self._pending = np.zeros(0, dtype=np.float32)
+        return [self.stepper.step(frame)]
+
+    def run(self, pcm: NDArray[np.float32]) -> tuple[NDArray[np.float32], str]:
+        """Convenience for offline use: feed a whole utterance, return (audio, text).
+
+        Concatenates every emitted agent frame and joins the inner-monologue text.
+        """
+        outputs = self.feed(pcm) + self.flush()
+        chunks = [o.audio for o in outputs if o.audio is not None]
+        audio = np.concatenate(chunks) if chunks else np.zeros(0, dtype=np.float32)
+        text = "".join(o.text for o in outputs if o.text)
+        return audio, text

vllm_omni/experimental/fullduplex/personaplex/stage0.py

@@ -0,0 +1,467 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import base64
+import binascii
+import io
+import tarfile
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from vllm_omni.experimental.fullduplex.personaplex.policy import (
+    AUDIO_SILENCE_FRAME_CNT,
+    SILENCE_TOKENS,
+    SINE_TOKENS,
+    ZERO_TEXT_TOKEN,
+    wrap_with_system_tags,
+)
+
+_FRAME_SAMPLES = 1920
+
+
+@dataclass(slots=True)
+class PersonaPlexStage0PreparedAppend:
+    input_ids: Any
+    inputs_embeds: Any
+    user_codes: Any
+    info_update: dict[str, Any]
+    prefill_applied: bool
+    prompt_offset: int
+
+
+@dataclass(slots=True)
+class PersonaPlexStage0SessionState:
+    session_id: str
+    incarnation: int
+    user_codes: Any | None = None
+    last_text_token: Any | None = None
+    last_agent_codes: Any | None = None
+    prefill_slots: int = 0
+    prepared_identity: tuple[int, int] | None = None
+    sampled_identity: tuple[int, int] | None = None
+    prepared: PersonaPlexStage0PreparedAppend | None = None
+    last_seq: int = 0
+    request_ids: set[str] = field(default_factory=set)
+    codec: Any | None = None
+
+
+def _tokenizer_path(model_path: str) -> Path:
+    path = Path(model_path) / "tokenizer_spm_32k_3.model"
+    if not path.is_file():
+        raise FileNotFoundError(f"PersonaPlex tokenizer not found: {path}")
+    return path
+
+
+def load_personaplex_tokenizer(model_path: str):
+    import sentencepiece
+
+    tokenizer = sentencepiece.SentencePieceProcessor(str(_tokenizer_path(model_path)))
+    return lambda text: list(tokenizer.encode(text))
+
+
+def load_personaplex_voice_state(model_path: str, voice: str) -> dict[str, Any]:
+    import torch
+
+    root = Path(model_path)
+    for candidate in (root / "voices" / voice, root / voice):
+        if candidate.is_file():
+            state = torch.load(candidate, map_location="cpu", weights_only=True)
+            if isinstance(state, dict):
+                return state
+            raise ValueError(f"PersonaPlex voice bundle is not a mapping: {candidate}")
+
+    archive = root / "voices.tgz"
+    if archive.is_file():
+        with tarfile.open(archive, "r:gz") as tar:
+            member = next(
+                (item for item in tar.getmembers() if item.isfile() and Path(item.name).name == voice),
+                None,
+            )
+            if member is not None:
+                extracted = tar.extractfile(member)
+                if extracted is None:
+                    raise FileNotFoundError(f"PersonaPlex voice prompt {voice!r} is unreadable")
+                state = torch.load(io.BytesIO(extracted.read()), map_location="cpu", weights_only=True)
+                if isinstance(state, dict):
+                    return state
+                raise ValueError(f"PersonaPlex voice bundle {voice!r} is not a mapping")
+    raise FileNotFoundError(f"PersonaPlex bundled voice prompt {voice!r} was not found under {model_path!r}")
+
+
+def personaplex_prefill_slots(model_path: str, voice: str, persona: str) -> int:
+    state = load_personaplex_voice_state(model_path, voice)
+    embeddings = state.get("embeddings")
+    if not hasattr(embeddings, "shape") or len(embeddings.shape) < 1:
+        raise ValueError(f"PersonaPlex voice prompt {voice!r} has no embeddings")
+    tokenizer = load_personaplex_tokenizer(model_path)
+    persona_tokens = tokenizer(wrap_with_system_tags(persona)) if persona else []
+    return int(embeddings.shape[0]) + 2 * AUDIO_SILENCE_FRAME_CNT + len(persona_tokens)
+
+
+class PersonaPlexStage0DuplexRuntime:
+    """Own per-session streaming Mimi encoders and first-append prefill."""
+
+    def __init__(
+        self,
+        stage_model: Any,
+        *,
+        model_path: str,
+        device: str,
+        codec: Any | None = None,
+        codec_factory: Callable[[], Any] | None = None,
+        max_sessions: int = 1,
+        tokenizer=None,
+        voice_loader=None,
+    ) -> None:
+        if max_sessions <= 0:
+            raise ValueError("PersonaPlex Stage 0 max_sessions must be positive")
+        self.stage_model = stage_model
+        self.model_path = model_path
+        self.device = device
+        self.max_sessions = max_sessions
+        self._codec_factory = codec_factory
+        self._free_codecs: list[Any] = []
+        self._tokenizer = tokenizer
+        self._voice_loader = voice_loader
+        self.sessions: dict[tuple[str, int], PersonaPlexStage0SessionState] = {}
+        self.request_sessions: dict[str, tuple[str, int]] = {}
+        if codec is not None:
+            codec.streaming_init(1)
+            self._free_codecs.append(codec)
+
+    def prepare_append(
+        self,
+        duplex: dict[str, Any],
+        *,
+        prompt_len: int,
+        request_id: str | None = None,
+    ) -> PersonaPlexStage0PreparedAppend:
+        import torch
+
+        session_id = duplex.get("session_id")
+        if not isinstance(session_id, str) or not session_id:
+            raise ValueError("PersonaPlex duplex append requires session_id")
+        incarnation = _coerce_non_negative_int(duplex.get("incarnation"), "incarnation")
+        epoch = _coerce_non_negative_int(duplex.get("epoch"), "epoch")
+        seq = _coerce_positive_int(duplex.get("seq"), "seq")
+        identity = (epoch, seq)
+        key = (session_id, incarnation)
+
+        state = self.sessions.get(key)
+        if state is None:
+            if len(self.sessions) >= self.max_sessions:
+                raise RuntimeError(f"PersonaPlex Stage 0 session capacity {self.max_sessions} is exhausted")
+            state = PersonaPlexStage0SessionState(
+                session_id=session_id,
+                incarnation=incarnation,
+                codec=self._acquire_codec(),
+            )
+            self.sessions[key] = state
+        if request_id:
+            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
+        if seq <= state.last_seq:
+            raise ValueError(f"PersonaPlex duplex append seq must increase: last={state.last_seq}, got={seq}")
+
+        pcm = self._decode_pcm(duplex.get("payload"))
+        codec = state.codec
+        if codec is None:
+            raise RuntimeError("PersonaPlex Stage 0 session has no streaming encoder")
+        encoded = codec.encode_frame(torch.from_numpy(pcm).reshape(1, _FRAME_SAMPLES))
+        user_frame = encoded.detach().to(dtype=torch.long, device="cpu").reshape(1, -1)
+        if user_frame.shape[1] < 8:
+            raise RuntimeError(
+                f"PersonaPlex Mimi encoder returned {user_frame.shape[1]} codebooks, expected at least 8"
+            )
+        user_frame = user_frame[:, :8].contiguous()
+        state.user_codes = user_frame if state.user_codes is None else torch.cat([state.user_codes, user_frame], dim=0)
+
+        runtime_config = duplex.get("runtime_config")
+        runtime_config = dict(runtime_config) if isinstance(runtime_config, dict) else {}
+        first_append = state.last_seq == 0
+        device, dtype = self._model_device_dtype()
+        silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long, device=device)
+        sine = torch.tensor(SINE_TOKENS, dtype=torch.long, device=device)
+        text_token = state.last_text_token
+        if text_token is None:
+            text_token = torch.tensor([ZERO_TEXT_TOKEN], dtype=torch.long, device=device)
+        last_agent = state.last_agent_codes if state.last_agent_codes is not None else silence
+        user_frame_count = state.user_codes.shape[0]
+        # Match the native lockstep ring: the temporal input sees the previous
+        # effective agent frame. User cb0 trails the frame being appended by one
+        # tick and user cb1..7 trail it by two ticks. Before live user frames
+        # fill those slots, text-prompt prefill leaves encoded sine on user rows.
+        user_d0 = state.user_codes[-2].to(device) if user_frame_count > 1 else sine
+        user_d1 = state.user_codes[-3].to(device) if user_frame_count > 2 else sine
+        depformer_audio_tokens = torch.cat(
+            [
+                silence,
+                user_frame[0, :1].to(device),
+                user_d0[1:8],
+            ]
+        )
+        depformer_audio_provided = torch.tensor(
+            [
+                False,
+                *([first_append] * 7),
+                *([True] * 8),
+            ],
+            dtype=torch.bool,
+            device=device,
+        )
+        live_embed = self.stage_model._build_frame_embed(
+            text_token,
+            last_agent,
+            last_agent,
+            device,
+            user_d0=user_d0,
+            user_d1=user_d1,
+        )
+        if first_append:
+            voice = runtime_config.get("personaplex_voice_prompt", "NATF2.pt")
+            persona = runtime_config.get("personaplex_persona", "")
+            if not isinstance(voice, str) or not voice:
+                raise ValueError("PersonaPlex runtime voice prompt is invalid")
+            if not isinstance(persona, str):
+                raise ValueError("PersonaPlex runtime persona is invalid")
+            voice_state = self._load_voice(voice)
+            voice_embeddings = voice_state.get("embeddings")
+            if not isinstance(voice_embeddings, torch.Tensor) or voice_embeddings.numel() == 0:
+                raise ValueError(f"PersonaPlex voice prompt {voice!r} has no embeddings")
+            voice_embeddings = voice_embeddings.reshape(-1, voice_embeddings.shape[-1]).to(
+                device=device,
+                dtype=dtype,
+            )
+            tokenizer = self._load_tokenizer()
+            persona_tokens = tokenizer(wrap_with_system_tags(persona)) if persona else []
+            prefill_tokens = torch.tensor(
+                [
+                    *([ZERO_TEXT_TOKEN] * AUDIO_SILENCE_FRAME_CNT),
+                    *persona_tokens,
+                    *([ZERO_TEXT_TOKEN] * AUDIO_SILENCE_FRAME_CNT),
+                ],
+                dtype=torch.long,
+                device=device,
+            )
+            token_prefill = self.stage_model._build_prefill_embed(
+                prefill_tokens,
+                0,
+                int(prefill_tokens.numel()),
+                device,
+                silence,
+                user_sine=sine,
+            )
+            prefill_embeds = torch.cat(
+                [
+                    voice_embeddings,
+                    token_prefill.to(dtype=dtype),
+                ],
+                dim=0,
+            )
+            state.prefill_slots = int(prefill_embeds.shape[0])
+            full_embeds = torch.cat([prefill_embeds, live_embed.to(dtype=dtype)], dim=0)
+        else:
+            full_embeds = live_embed.to(dtype=dtype)
+
+        prepared_len = int(full_embeds.shape[0])
+        prompt_offset = int(prompt_len) - prepared_len
+        if prompt_offset < 0:
+            raise ValueError(
+                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,
+                "incarnation": incarnation,
+                "epoch": epoch,
+                "seq": seq,
+            },
+        }
+        prepared = 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
+        state.prepared = prepared
+        state.last_seq = seq
+        return prepared
+
+    def record_sample(
+        self,
+        *,
+        request_id: str,
+        text_token: Any,
+        agent_codes: Any,
+    ) -> None:
+        """Commit one sampled temporal frame for the next live append."""
+        import torch
+
+        key = self.request_sessions.get(request_id)
+        if key is None:
+            raise KeyError(f"PersonaPlex Stage 0 request is not attached to a live session: {request_id}")
+        state = self.sessions.get(key)
+        if state is None or state.prepared_identity is None:
+            raise RuntimeError(f"PersonaPlex Stage 0 request has no prepared append: {request_id}")
+        if state.sampled_identity == state.prepared_identity:
+            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)
+        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()
+        state.last_text_token = text.detach()
+        state.sampled_identity = state.prepared_identity
+
+    def close_request(self, request_id: str) -> None:
+        key = self.request_sessions.pop(request_id, None)
+        if key is None:
+            return
+        state = self.sessions.get(key)
+        if state is None:
+            return
+        state.request_ids.discard(request_id)
+        if not state.request_ids:
+            self.close_session(*key)
+
+    def close_session(self, session_id: str, incarnation: int) -> None:
+        key = (session_id, incarnation)
+        state = self.sessions.pop(key, None)
+        if state is None:
+            return
+        for request_id in state.request_ids:
+            self.request_sessions.pop(request_id, None)
+        if state.codec is not None:
+            state.codec.reset_streaming()
+            self._free_codecs.append(state.codec)
+            state.codec = None
+
+    def _acquire_codec(self):
+        if self._free_codecs:
+            return self._free_codecs.pop()
+        if self._codec_factory is not None:
+            codec = self._codec_factory()
+            codec.streaming_init(1)
+            return codec
+        from vllm_omni.model_executor.models.personaplex.personaplex_mimi import (
+            PersonaPlexMimiCodec,
+        )
+
+        checkpoint = Path(self.model_path) / "tokenizer-e351c8d8-checkpoint125.safetensors"
+        codec = PersonaPlexMimiCodec(
+            checkpoint=str(checkpoint) if checkpoint.is_file() else None,
+            device=self.device,
+        )
+        codec.streaming_init(1)
+        return codec
+
+    def _load_tokenizer(self):
+        if self._tokenizer is None:
+            self._tokenizer = load_personaplex_tokenizer(self.model_path)
+        return self._tokenizer
+
+    def _load_voice(self, voice: str) -> dict[str, Any]:
+        if self._voice_loader is not None:
+            state = self._voice_loader(voice)
+        else:
+            state = load_personaplex_voice_state(self.model_path, voice)
+        if not isinstance(state, dict):
+            raise ValueError(f"PersonaPlex voice prompt {voice!r} is not a mapping")
+        return state
+
+    def _model_device_dtype(self):
+        import torch
+
+        try:
+            parameter = next(self.stage_model.parameters())
+            return parameter.device, parameter.dtype
+        except Exception:
+            device = getattr(self.stage_model, "device", torch.device(self.device))
+            dtype = getattr(self.stage_model, "dtype", torch.float32)
+            return torch.device(device), dtype
+
+    @staticmethod
+    def _decode_pcm(payload: object) -> np.ndarray:
+        if not isinstance(payload, dict):
+            raise ValueError("PersonaPlex duplex payload must be a mapping")
+        if payload.get("format") != "pcm_f32le" or payload.get("sample_rate_hz") != 24000:
+            raise ValueError("PersonaPlex Stage 0 requires 24 kHz pcm_f32le")
+        audio = payload.get("audio")
+        if not isinstance(audio, str):
+            raise ValueError("PersonaPlex Stage 0 requires base64 audio")
+        try:
+            raw = base64.b64decode(audio, validate=True)
+        except (binascii.Error, ValueError) as exc:
+            raise ValueError("PersonaPlex Stage 0 audio is not valid base64") from exc
+        samples = np.frombuffer(raw, dtype="<f4")
+        if samples.size != _FRAME_SAMPLES:
+            raise ValueError(f"PersonaPlex Stage 0 requires {_FRAME_SAMPLES} samples per append")
+        if not np.isfinite(samples).all():
+            raise ValueError("PersonaPlex Stage 0 samples must be finite")
+        return np.ascontiguousarray(samples, dtype=np.float32).copy()
+
+
+def _coerce_non_negative_int(value: object, name: str) -> int:
+    try:
+        result = int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"PersonaPlex duplex {name} must be an integer") from exc
+    if result < 0:
+        raise ValueError(f"PersonaPlex duplex {name} must be non-negative")
+    return result
+
+
+def _coerce_positive_int(value: object, name: str) -> int:
+    result = _coerce_non_negative_int(value, name)
+    if result <= 0:
+        raise ValueError(f"PersonaPlex duplex {name} must be positive")
+    return result
+
+
+__all__ = [
+    "PersonaPlexStage0DuplexRuntime",
+    "PersonaPlexStage0PreparedAppend",
+    "PersonaPlexStage0SessionState",
+    "load_personaplex_tokenizer",
+    "load_personaplex_voice_state",
+    "personaplex_prefill_slots",
+]

vllm_omni/model_executor/models/personaplex/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""vllm-omni integration for PersonaPlex (a Moshi finetune, full-duplex S2S).
+
+Offline/batch runs go through a two-stage audio->audio pipeline (talker ->
+code2wav); real-time conversation is served over the duplex WebSocket path.
+"""
+
+from vllm_omni.model_executor.models.personaplex.configuration_personaplex import (
+    PersonaPlexConfig,
+    PersonaPlexDepformerConfig,
+    PersonaPlexMimiConfig,
+)
+
+__all__ = [
+    "PersonaPlexConfig",
+    "PersonaPlexDepformerConfig",
+    "PersonaPlexMimiConfig",
+]

vllm_omni/model_executor/models/personaplex/configuration_helium.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Configuration for the PersonaPlex Helium temporal transformer."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from transformers import PretrainedConfig
+
+__all__ = ["HeliumConfig"]
+
+
+class HeliumConfig(PretrainedConfig):
+    """Minimal HF config for the Moshi temporal LM backbone.
+
+    The defaults are measured from the PersonaPlex checkpoint rather than inferred
+    from Moshi's higher-level kwargs.
+    """
+
+    model_type = "helium"
+    keys_to_ignore_at_inference = ("past_key_values",)
+
+    def __init__(
+        self,
+        hidden_size: int = 4096,
+        num_hidden_layers: int = 32,
+        num_attention_heads: int = 32,
+        head_dim: int = 128,
+        num_key_value_heads: int = 32,
+        intermediate_size: int = 11264,
+        max_position_embeddings: int = 3000,
+        rope_theta: float = 10000.0,
+        rms_norm_eps: float = 1e-8,
+        vocab_size: int = 32000,
+        sliding_window: int = 3000,
+        hidden_act: str = "silu",
+        attention_bias: bool = False,
+        mlp_bias: bool = False,
+        tie_word_embeddings: bool = False,
+        use_cache: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = head_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.vocab_size = vocab_size
+        self.sliding_window = sliding_window
+        self.hidden_act = hidden_act
+        self.attention_bias = attention_bias
+        self.mlp_bias = mlp_bias
+        self.use_cache = use_cache

vllm_omni/model_executor/models/personaplex/configuration_personaplex.py

@@ -0,0 +1,278 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Configuration for PersonaPlex (a Moshi finetune; 2-stage audio->audio pipeline).
+
+PersonaPlex is a staged AR speech model composed of:
+
+* a *temporal transformer* (a Llama-variant Helium backbone) that predicts, per
+  frame, the next text token and the codebook-0 audio token;
+* a *depformer* (a small per-step transformer) that, conditioned on the temporal
+  hidden state, predicts the remaining audio codebooks;
+* the Mimi neural audio codec, which turns the audio codebooks into 24 kHz PCM.
+
+The temporal-transformer hyperparameters live in
+:class:`~vllm_omni.model_executor.models.personaplex.configuration_helium.HeliumConfig`
+(already used by the talker the lead is building). This module reuses it as the
+``temporal_config`` sub-config so the config tree carries a single source of
+truth, and adds two further sub-configs (``depformer_config`` and
+``mimi_config``).
+
+All defaults are *measured* from the PersonaPlex checkpoint, whose own
+``config.json`` is empty; do not treat the Moshi higher-level kwargs as
+authoritative.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+from vllm_omni.model_executor.models.personaplex.configuration_helium import (
+    HeliumConfig,
+)
+
+logger = logging.get_logger(__name__)
+
+__all__ = [
+    "PersonaPlexConfig",
+    "PersonaPlexDepformerConfig",
+    "PersonaPlexMimiConfig",
+    "HeliumConfig",
+]
+
+
+class PersonaPlexDepformerConfig(PretrainedConfig):
+    r"""Configuration for the PersonaPlex depformer (per-step code predictor).
+
+    The depformer is a small autoregressive transformer that runs ``dep_q``
+    inner steps per temporal frame, conditioned on the temporal hidden state,
+    to predict the audio codebooks. Only the first ``num_active_codebooks``
+    codebooks (``cb 0..7``) are decoded to PCM by Mimi; the remaining codebooks
+    up to ``dep_q`` are predicted but not vocoded.
+
+    Args:
+        hidden_size (`int`, *optional*, defaults to 1024):
+            Dimension of the depformer hidden representations.
+        num_hidden_layers (`int`, *optional*, defaults to 6):
+            Number of depformer transformer layers.
+        num_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads per depformer layer.
+        head_dim (`int`, *optional*, defaults to 64):
+            Per-head attention dimension (``hidden_size // num_attention_heads``).
+        max_position_embeddings (`int`, *optional*, defaults to 8):
+            The depformer context length (``ctx``): one position per inner step.
+        dep_q (`int`, *optional*, defaults to 16):
+            Number of audio codebooks the depformer predicts per frame.
+        num_active_codebooks (`int`, *optional*, defaults to 8):
+            Number of leading codebooks actually decoded to PCM by Mimi.
+        card (`int`, *optional*, defaults to 2048):
+            Per-codebook cardinality (audio vocab size).
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-8):
+            The epsilon used by the fp32 RMS normalization layers.
+        hidden_act (`str`, *optional*, defaults to `"silu"`):
+            SwiGLU gate activation.
+        attention_bias (`bool`, *optional*, defaults to `False`):
+            Whether attention projections carry a bias.
+        mlp_bias (`bool`, *optional*, defaults to `False`):
+            Whether the MLP projections carry a bias.
+    """
+
+    model_type = "personaplex_depformer"
+    keys_to_ignore_at_inference = ("past_key_values",)
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        num_hidden_layers: int = 6,
+        num_attention_heads: int = 16,
+        head_dim: int = 64,
+        num_key_value_heads: int = 16,
+        intermediate_size: int = 2816,
+        max_position_embeddings: int = 8,
+        dep_q: int = 16,
+        num_active_codebooks: int = 8,
+        card: int = 2048,
+        rope_theta: float = 10000.0,
+        rms_norm_eps: float = 1e-8,
+        hidden_act: str = "silu",
+        attention_bias: bool = False,
+        mlp_bias: bool = False,
+        tie_word_embeddings: bool = False,
+        use_cache: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = head_dim
+        self.num_key_value_heads = num_key_value_heads
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.dep_q = dep_q
+        self.num_active_codebooks = num_active_codebooks
+        self.card = card
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.hidden_act = hidden_act
+        self.attention_bias = attention_bias
+        self.mlp_bias = mlp_bias
+        self.use_cache = use_cache
+
+
+class PersonaPlexMimiConfig(PretrainedConfig):
+    r"""Configuration for the Mimi neural audio codec used by PersonaPlex.
+
+    The Mimi weights and module live in the external ``moshi`` package; this
+    config only carries the scalars the vllm-omni serving layer needs to size
+    buffers and compute audio durations. The actual decoder is instantiated by
+    :class:`PersonaPlexCode2Wav` via ``moshi.models.loaders.get_mimi``.
+
+    Args:
+        sample_rate (`int`, *optional*, defaults to 24000):
+            Output PCM sample rate in Hz.
+        frame_rate (`float`, *optional*, defaults to 12.5):
+            Mimi codec frame rate in Hz (one frame every 80 ms).
+        samples_per_frame (`int`, *optional*, defaults to 1920):
+            PCM samples produced per codec frame (``sample_rate / frame_rate``).
+        num_codebooks (`int`, *optional*, defaults to 8):
+            Number of active audio codebooks decoded to PCM (``cb 0..7``).
+        card (`int`, *optional*, defaults to 2048):
+            Per-codebook cardinality.
+        num_channels (`int`, *optional*, defaults to 1):
+            Number of output audio channels (mono).
+        mimi_name (`str`, *optional*):
+            Filename of the Mimi weight checkpoint inside the model repo. When
+            ``None``, :class:`PersonaPlexCode2Wav` falls back to
+            ``moshi.models.loaders.MIMI_NAME``.
+    """
+
+    model_type = "personaplex_mimi"
+
+    def __init__(
+        self,
+        sample_rate: int = 24000,
+        frame_rate: float = 12.5,
+        samples_per_frame: int = 1920,
+        num_codebooks: int = 8,
+        card: int = 2048,
+        num_channels: int = 1,
+        mimi_name: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.sample_rate = sample_rate
+        self.frame_rate = frame_rate
+        self.samples_per_frame = samples_per_frame
+        self.num_codebooks = num_codebooks
+        self.card = card
+        self.num_channels = num_channels
+        self.mimi_name = mimi_name
+
+
+class PersonaPlexConfig(PretrainedConfig):
+    r"""Top-level configuration for ``PersonaPlexTalkerForConditionalGeneration``.
+
+    Mirrors the Qwen3-TTS config layout: a top-level config holding sub-configs
+    for each component. The temporal transformer config is the *text config*
+    that vLLM consumes (it exposes ``hidden_size`` / ``num_attention_heads`` and
+    drives the talker's vLLM Llama backbone).
+
+    Args:
+        temporal_config (`dict` or `HeliumConfig`, *optional*):
+            The temporal-transformer (Helium) backbone config.
+        depformer_config (`dict` or `PersonaPlexDepformerConfig`, *optional*):
+            The depformer (per-step code predictor) config.
+        mimi_config (`dict` or `PersonaPlexMimiConfig`, *optional*):
+            The Mimi codec config.
+        text_vocab_size (`int`, *optional*, defaults to 32000):
+            Text / ``lm_head`` vocabulary size.
+        text_embedding_rows (`int`, *optional*, defaults to 32001):
+            Number of rows in the text embedding table (one extra padding row).
+        audio_vocab_size (`int`, *optional*, defaults to 2048):
+            Per-codebook audio cardinality (``card``).
+        num_audio_codebooks (`int`, *optional*, defaults to 16):
+            Total number of audio codebooks (``n_q``).
+        mimi_name (`str`, *optional*):
+            Convenience mirror of ``mimi_config.mimi_name``; if set, it
+            overrides the value carried inside ``mimi_config``.
+    """
+
+    model_type = "personaplex"
+    sub_configs = {
+        "temporal_config": HeliumConfig,
+        "depformer_config": PersonaPlexDepformerConfig,
+        "mimi_config": PersonaPlexMimiConfig,
+    }
+
+    def __init__(
+        self,
+        temporal_config: dict | HeliumConfig | None = None,
+        depformer_config: dict | PersonaPlexDepformerConfig | None = None,
+        mimi_config: dict | PersonaPlexMimiConfig | None = None,
+        text_vocab_size: int = 32000,
+        text_embedding_rows: int = 32001,
+        audio_vocab_size: int = 2048,
+        num_audio_codebooks: int = 16,
+        mimi_name: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if temporal_config is None:
+            temporal_config = {}
+            logger.info("temporal_config is None. Initializing Helium backbone with default values")
+        if depformer_config is None:
+            depformer_config = {}
+            logger.info("depformer_config is None. Initializing depformer with default values")
+        if mimi_config is None:
+            mimi_config = {}
+            logger.info("mimi_config is None. Initializing Mimi codec config with default values")
+
+        self.temporal_config = self._coerce(temporal_config, HeliumConfig)
+        self.depformer_config = self._coerce(depformer_config, PersonaPlexDepformerConfig)
+        self.mimi_config = self._coerce(mimi_config, PersonaPlexMimiConfig)
+
+        super().__init__(**kwargs)
+
+        # The PersonaPlex checkpoint's config.json is empty (no ``architectures``);
+        # supply the primary (stage-0 talker) arch so vLLM's ModelConfig resolves a
+        # generate-capable model. The code2wav stage arch comes from the pipeline.
+        if not getattr(self, "architectures", None):
+            self.architectures = ["PersonaPlexTalkerForConditionalGeneration"]
+
+        self.text_vocab_size = text_vocab_size
+        self.text_embedding_rows = text_embedding_rows
+        self.audio_vocab_size = audio_vocab_size
+        self.num_audio_codebooks = num_audio_codebooks
+
+        # ``mimi_name`` is allowed both at the top level (convenient for
+        # hf_overrides / a flat config.json) and inside ``mimi_config``. Keep
+        # them in sync, preferring an explicit top-level value.
+        if mimi_name is not None:
+            self.mimi_config.mimi_name = mimi_name
+        self.mimi_name = self.mimi_config.mimi_name
+
+    @staticmethod
+    def _coerce(value: Any, config_cls: type[PretrainedConfig]) -> PretrainedConfig:
+        if isinstance(value, config_cls):
+            return value
+        if isinstance(value, dict):
+            return config_cls(**value)
+        raise TypeError(f"Expected dict or {config_cls.__name__}, got {type(value).__name__}")
+
+    @property
+    def sample_rate(self) -> int:
+        """Output PCM sample rate in Hz (delegates to the Mimi config)."""
+        return int(self.mimi_config.sample_rate)
+
+    def get_text_config(self, **kwargs: Any) -> PretrainedConfig:
+        # vLLM expects the text config to expose hidden_size / num_attention_heads.
+        # For PersonaPlex the temporal-transformer (Helium) config is the text
+        # model config that drives the talker's vLLM Llama backbone.
+        return self.temporal_config

vllm_omni/model_executor/models/personaplex/modeling_helium.py

@@ -0,0 +1,548 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""vLLM-native PersonaPlex Helium temporal transformer.
+
+The module maps the Moshi temporal LM backbone onto vLLM's Llama-style decoder
+components. It intentionally does not include Moshi's text/audio input
+embeddings or depformer; PersonaPlex feeds temporal `inputs_embeds` directly.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from itertools import islice
+
+import torch
+from torch import nn
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
+from vllm.model_executor.models.utils import (
+    PPMissingLayer,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backend import AttentionType
+
+from vllm_omni.model_executor.models.personaplex.configuration_helium import (
+    HeliumConfig,
+)
+
+__all__ = ["HeliumForCausalLM", "HeliumModel"]
+
+
+class HeliumRMSNorm(RMSNorm):
+    """Moshi-compatible RMSNorm.
+
+    Moshi stores the parameter as `alpha` with shape `[1, 1, hidden]` and
+    computes variance and the alpha multiply in fp32 before casting back to the
+    input dtype. vLLM's parameter is named `weight`; the loader squeezes alpha.
+    """
+
+    def _forward_f32(self, x: torch.Tensor) -> torch.Tensor:
+        output_dtype = x.dtype
+        x_f32 = x.float()
+        variance = x_f32.pow(2).mean(dim=-1, keepdim=True)
+        x_f32 = x_f32 * torch.rsqrt(variance + self.variance_epsilon)
+        if self.has_weight:
+            x_f32 = x_f32 * self.weight.float()
+        return x_f32.to(output_dtype)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            return self._forward_f32(x)
+        residual = residual + x
+        return self._forward_f32(residual), residual
+
+
+class HeliumMLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        if hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported.")
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+def _get_helium_rope(config: HeliumConfig):
+    rope_kwargs = {
+        "head_size": config.head_dim,
+        "max_position": config.max_position_embeddings,
+        "is_neox_style": False,
+        "rope_parameters": {"base": config.rope_theta},
+    }
+    try:
+        return get_rope(**rope_kwargs)
+    except TypeError:
+        # vLLM 0.23 uses `rope_parameters`; this fallback is for older local
+        # source trees whose get_rope still takes rotary_dim/base/rope_scaling.
+        return get_rope(
+            head_size=config.head_dim,
+            rotary_dim=config.head_dim,
+            max_position=config.max_position_embeddings,
+            base=config.rope_theta,
+            is_neox_style=False,
+        )
+
+
+class HeliumAttention(nn.Module):
+    def __init__(
+        self,
+        config: HeliumConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        attn_type: str = AttentionType.DECODER,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = config.head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+        self.rotary_emb = _get_helium_rope(config)
+        self.attn = Attention(
+            self.num_heads,
+            self.head_dim,
+            self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            per_layer_sliding_window=config.sliding_window,
+            prefix=f"{prefix}.attn",
+            attn_type=attn_type,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class HeliumDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+        config: HeliumConfig | None = None,
+    ) -> None:
+        super().__init__()
+        config = config or vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = HeliumAttention(
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=config.num_key_value_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+            attn_type=AttentionType.DECODER,
+        )
+        self.mlp = HeliumMLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = HeliumRMSNorm(
+            config.hidden_size,
+            eps=config.rms_norm_eps,
+        )
+        self.post_attention_layernorm = HeliumRMSNorm(
+            config.hidden_size,
+            eps=config.rms_norm_eps,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+
+        hidden_states, residual = self.post_attention_layernorm(
+            hidden_states,
+            residual,
+        )
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+def _flatten_inputs(
+    hidden_states: torch.Tensor,
+    positions: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int] | None]:
+    if hidden_states.dim() != 3:
+        return hidden_states, positions.reshape(-1), None
+
+    batch, seq_len, hidden = hidden_states.shape
+    flat_hidden = hidden_states.reshape(batch * seq_len, hidden)
+    if positions.dim() == 1:
+        if positions.numel() == seq_len:
+            flat_positions = positions.unsqueeze(0).expand(batch, seq_len).reshape(-1)
+        elif positions.numel() == batch * seq_len:
+            flat_positions = positions.reshape(-1)
+        else:
+            raise ValueError(
+                "1D positions must have length S or B*S when inputs_embeds is "
+                f"3D, got {positions.numel()} for B={batch}, S={seq_len}."
+            )
+    else:
+        flat_positions = positions.reshape(-1)
+    return flat_hidden, flat_positions, (batch, seq_len)
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        "positions": 0,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+    }
+)
+class HeliumModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+        layer_type: type[nn.Module] = HeliumDecoderLayer,
+        config: HeliumConfig | None = None,
+    ) -> None:
+        super().__init__()
+        # ``config`` lets a composite model (e.g. the PersonaPlex talker) build the
+        # Helium backbone from its temporal sub-config while the engine's top-level
+        # hf_config is the composite PersonaPlexConfig. Defaults to the engine's.
+        config = config or vllm_config.model_config.hf_config
+        self.config = config
+        self.quant_config = vllm_config.quant_config
+        self.vocab_size = config.vocab_size
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: layer_type(vllm_config=vllm_config, prefix=prefix, config=config),
+            prefix=f"{prefix}.layers",
+        )
+        if get_pp_group().is_last_rank:
+            self.norm = HeliumRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"],
+            config.hidden_size,
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        raise ValueError(
+            "Helium temporal forward requires inputs_embeds; Moshi token "
+            "embeddings are intentionally not part of this vLLM module."
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        unflatten_shape: tuple[int, int] | None = None
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is None:
+                raise ValueError("HeliumModel.forward requires inputs_embeds.")
+            hidden_states, positions, unflatten_shape = _flatten_inputs(
+                inputs_embeds,
+                positions,
+            )
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+            positions = positions.reshape(-1)
+
+        for layer in islice(self.layers, self.start_layer, self.end_layer):
+            hidden_states, residual = layer(
+                positions,
+                hidden_states,
+                residual,
+            )
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+        if unflatten_shape is not None:
+            batch, seq_len = unflatten_shape
+            hidden_states = hidden_states.reshape(batch, seq_len, -1)
+        return hidden_states
+
+
+class HeliumForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+    embedding_modules = {
+        "lm_head": "output_embeddings",
+    }
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+        layer_type: type[nn.Module] = HeliumDecoderLayer,
+    ) -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.config = config
+        self.quant_config = vllm_config.quant_config
+        self.model = HeliumModel(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+            layer_type=layer_type,
+        )
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=self.quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+        else:
+            self.lm_head = PPMissingLayer()
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        return self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        if hidden_states.dim() == 3:
+            batch, seq_len, hidden = hidden_states.shape
+            flat_hidden = hidden_states.reshape(batch * seq_len, hidden)
+            logits = self.logits_processor(self.lm_head, flat_hidden)
+            if logits is None:
+                return None
+            return logits.reshape(batch, seq_len, -1)
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]] | Mapping[str, torch.Tensor],
+    ) -> set[str]:
+        if isinstance(weights, Mapping):
+            weights = weights.items()
+
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            mapped = self._load_moshi_weight(name, loaded_weight, params_dict)
+            loaded_params.update(mapped)
+
+        return loaded_params
+
+    def _load_moshi_weight(
+        self,
+        name: str,
+        loaded_weight: torch.Tensor,
+        params_dict: dict[str, nn.Parameter],
+    ) -> set[str]:
+        if name == "out_norm.alpha":
+            return self._load_direct(
+                "model.norm.weight",
+                loaded_weight.squeeze(),
+                params_dict,
+            )
+        if name == "text_linear.weight":
+            return self._load_direct("lm_head.weight", loaded_weight, params_dict)
+
+        prefix = "transformer.layers."
+        if not name.startswith(prefix):
+            return set()
+
+        rest = name.removeprefix(prefix)
+        layer_index, _, suffix = rest.partition(".")
+        if not layer_index.isdigit() or not suffix:
+            return set()
+
+        base = f"model.layers.{layer_index}"
+        if suffix == "self_attn.in_proj_weight":
+            q_weight, k_weight, v_weight = loaded_weight.chunk(3, dim=0)
+            param_name = f"{base}.self_attn.qkv_proj.weight"
+            self._load_shard(param_name, q_weight, "q", params_dict)
+            self._load_shard(param_name, k_weight, "k", params_dict)
+            self._load_shard(param_name, v_weight, "v", params_dict)
+            return {param_name}
+        if suffix == "self_attn.out_proj.weight":
+            return self._load_direct(
+                f"{base}.self_attn.o_proj.weight",
+                loaded_weight,
+                params_dict,
+            )
+        if suffix == "gating.linear_in.weight":
+            gate_weight, up_weight = loaded_weight.chunk(2, dim=0)
+            param_name = f"{base}.mlp.gate_up_proj.weight"
+            self._load_shard(param_name, gate_weight, 0, params_dict)
+            self._load_shard(param_name, up_weight, 1, params_dict)
+            return {param_name}
+        if suffix == "gating.linear_out.weight":
+            return self._load_direct(
+                f"{base}.mlp.down_proj.weight",
+                loaded_weight,
+                params_dict,
+            )
+        if suffix == "norm1.alpha":
+            return self._load_direct(
+                f"{base}.input_layernorm.weight",
+                loaded_weight.squeeze(),
+                params_dict,
+            )
+        if suffix == "norm2.alpha":
+            return self._load_direct(
+                f"{base}.post_attention_layernorm.weight",
+                loaded_weight.squeeze(),
+                params_dict,
+            )
+        return set()
+
+    def _load_direct(
+        self,
+        name: str,
+        loaded_weight: torch.Tensor,
+        params_dict: dict[str, nn.Parameter],
+    ) -> set[str]:
+        if name not in params_dict or is_pp_missing_parameter(name, self):
+            return set()
+        param = params_dict[name]
+        weight_loader = getattr(param, "weight_loader", default_weight_loader)
+        weight_loader(param, loaded_weight)
+        return {name}
+
+    def _load_shard(
+        self,
+        name: str,
+        loaded_weight: torch.Tensor,
+        shard_id: str | int,
+        params_dict: dict[str, nn.Parameter],
+    ) -> None:
+        if name not in params_dict or is_pp_missing_parameter(name, self):
+            return
+        param = params_dict[name]
+        weight_loader = param.weight_loader
+        weight_loader(param, loaded_weight, shard_id)

vllm_omni/model_executor/models/personaplex/personaplex_code2wav.py

@@ -0,0 +1,433 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage-1 Code2Wav model for PersonaPlex (Moshi finetune).
+
+This is the ``LLM_GENERATION`` codec stage of the 2-stage PersonaPlex pipeline.
+It consumes the per-frame audio codebooks produced by the talker (stage 0) and
+the depformer (built by the lead), and turns them into 24 kHz PCM by calling the
+external Mimi neural codec from the ``moshi`` package.
+
+Layout contract (mirrors ``Qwen3TTSCode2Wav``):
+
+* Per request, ``input_ids`` holds a flat codebook-major codec sequence
+  ``[k * F]`` where ``k`` is the number of *active* codebooks Mimi decodes
+  (``num_codebooks``, i.e. the ``cb 0..7`` slice) and ``F`` is the number of
+  codec frames. The talker emits a 17-row token stack per frame; the input
+  processor (built by the lead) keeps only rows ``1:9`` (the 8 PCM-bearing audio
+  codebooks) and flattens them codebook-major before this stage.
+* ``forward(...)`` returns an :class:`OmniOutput` whose ``multimodal_outputs``
+  carries ``{"model_outputs": [wav_per_request], "sr": [sr_per_request]}`` — the
+  exact shape ``Qwen3TTSCode2Wav`` returns, so the downstream audio packer is
+  unchanged.
+
+The Mimi decoder is the transformers ``MimiModel`` (kyutai/mimi weights), not
+in vLLM's safetensors loader, so they are loaded eagerly in ``load_weights``
+part of the vLLM weights iterator (the codec owns its own checkpoint; see the
+experimental fullduplex package for the same pattern).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.forward_context import get_forward_context, is_forward_context_available
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+logger = init_logger(__name__)
+
+_MIMI_DECODE_BATCH_FRAMES = 5
+
+
+def _codec_ids_from_payload_or_input(
+    input_ids: torch.Tensor,
+    runtime_info: Mapping[str, Any] | None,
+) -> torch.Tensor:
+    """Prefer connector-delivered codec ids over scheduler placeholders."""
+    if isinstance(runtime_info, Mapping):
+        codes = runtime_info.get("codes")
+        if isinstance(codes, Mapping):
+            audio = codes.get("audio")
+            if isinstance(audio, torch.Tensor) and audio.numel() > 0:
+                return audio.reshape(-1).to(device=input_ids.device, dtype=torch.long)
+            if isinstance(audio, (list, tuple)) and audio:
+                return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1)
+    return input_ids.reshape(-1).to(dtype=torch.long)
+
+
+class PersonaPlexCode2Wav(nn.Module):
+    """Stage-1 Code2Wav model for PersonaPlex (GenerationModelRunner).
+
+    Wraps the external Mimi decoder. The vLLM generation runner only calls a
+    handful of methods on this module: ``embed_input_ids`` (dummy embeddings),
+    ``compute_logits`` (none -- this stage never samples), ``forward`` (the
+    actual codec->PCM decode), ``make_omni_output`` (output normalization), and
+    ``load_weights`` (eager Mimi construction).
+    """
+
+    input_modalities = "audio"
+    requires_request_ids = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.model_path = vllm_config.model_config.model
+        self.config = vllm_config.model_config.hf_config
+
+        # Runner-facing capability flags, matching Qwen3TTSCode2Wav so the
+        # GenerationModelRunner drives this stage the same way.
+        self.have_multimodal_outputs = True
+        self.has_preprocess = False
+        self.has_postprocess = False
+        self.enable_update_additional_information = True
+        self.requires_raw_input_tokens = True
+
+        mimi_cfg = getattr(self.config, "mimi_config", None)
+        # Number of *active* codebooks Mimi decodes to PCM (cb 0..7).
+        self._num_codebooks = int(getattr(mimi_cfg, "num_codebooks", 8))
+        self._output_sample_rate = int(getattr(mimi_cfg, "sample_rate", 24000))
+        self._samples_per_frame = int(getattr(mimi_cfg, "samples_per_frame", 1920))
+        self._mimi_name = getattr(mimi_cfg, "mimi_name", None) or getattr(self.config, "mimi_name", None)
+        self._max_codec_sessions = int(getattr(vllm_config.model_config, "duplex_max_sessions", 1))
+
+        # The Mimi module is constructed in load_weights() (it owns its own
+        # weight format) and assigned here so vLLM's memory profiler can see it.
+        self.mimi: nn.Module | None = None
+        self._additional_mimi = nn.ModuleList()
+        self._mimi_device: torch.device | None = None
+        self._request_codes: dict[str, torch.Tensor] = {}
+        self._request_codec_slots: dict[str, int] = {}
+
+    # ------------------------------------------------------------------
+    # Runner-facing no-op / placeholder hooks (mirror Qwen3TTSCode2Wav).
+    # ------------------------------------------------------------------
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        # This stage ignores token embeddings; keep a stable dummy embedding so
+        # the vLLM runner has a valid hidden-state placeholder.
+        if input_ids.numel() == 0:
+            return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32)
+        return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor | OmniOutput,
+        sampling_metadata: Any = None,
+    ) -> None:
+        # Code2Wav never samples; the runner short-circuits when logits are None.
+        return None
+
+    def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]:
+        return [{"meta": {"personaplex_dummy_profile": True}} for _ in range(num_reqs)]
+
+    def _split_request_ids(
+        self,
+        ids: torch.Tensor,
+        seq_token_counts: list[int] | None = None,
+    ) -> list[torch.Tensor]:
+        """Split concatenated ``input_ids`` into per-request segments.
+
+        Uses ``seq_token_counts`` (injected by the runner via model_kwargs) when
+        available, falling back to forward-context ``ubatch_slices`` when
+        micro-batching is active. Returns ``[ids]`` for single-request batches.
+        Mirrors ``Qwen3TTSCode2Wav._split_request_ids``.
+        """
+        if seq_token_counts is not None and len(seq_token_counts) > 1:
+            boundaries = [0]
+            for count in seq_token_counts:
+                boundaries.append(boundaries[-1] + count)
+            n = ids.numel()
+            return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))]
+        if is_forward_context_available():
+            slices = get_forward_context().ubatch_slices
+            if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices):
+                boundaries = [0]
+                for s in slices:
+                    boundaries.append(boundaries[-1] + s)
+                return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)]
+        return [ids]
+
+    # ------------------------------------------------------------------
+    # Decode.
+    # ------------------------------------------------------------------
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: Any = None,
+        inputs_embeds: torch.Tensor | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        """Decode flat codebook-major codec ids into PCM via Mimi.
+
+        ``input_ids`` per request is ``[k * F]`` (codebook-major), where ``k``
+        is ``self._num_codebooks``. The connector may send either a new delta
+        chunk or the cumulative prefix of a resumable request. Request-local
+        code history identifies the new suffix, and the streaming Mimi decoder
+        consumes each new frame exactly once.
+        """
+        sr_val = int(self._output_sample_rate)
+        sr_tensor = torch.tensor(sr_val, dtype=torch.int32)
+        empty = torch.zeros((0,), dtype=torch.float32)
+
+        if input_ids is None or input_ids.numel() == 0:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]},
+            )
+
+        ids = input_ids.reshape(-1).to(dtype=torch.long)
+        request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
+        num_req = len(request_ids_list)
+        state_ids = self._resolve_request_ids(
+            num_req,
+            kwargs.get("request_ids"),
+            runtime_additional_information,
+        )
+
+        if self.mimi is None:
+            raise RuntimeError("PersonaPlexCode2Wav.forward called before Mimi was loaded in load_weights().")
+
+        k = int(self._num_codebooks)
+        device = self._mimi_device or ids.device
+        audios: list[torch.Tensor] = [empty] * num_req
+        srs = [sr_tensor] * num_req
+
+        for i, req_ids in enumerate(request_ids_list):
+            runtime_info = (
+                runtime_additional_information[i]
+                if runtime_additional_information is not None and i < len(runtime_additional_information)
+                else None
+            )
+            meta = runtime_info.get("meta") if isinstance(runtime_info, Mapping) else None
+            if isinstance(meta, Mapping) and meta.get("personaplex_dummy_profile") is True:
+                continue
+            flat = _codec_ids_from_payload_or_input(req_ids, runtime_info)
+            n = int(flat.numel())
+            if n == 0 or n % k != 0:
+                if n > 0:
+                    logger.warning(
+                        "PersonaPlex Code2Wav input_ids length %d not divisible by "
+                        "num_codebooks %d; skipping malformed request.",
+                        n,
+                        k,
+                    )
+                continue
+            frames = n // k
+            codes_kf = flat.reshape(k, frames)
+            state_id = state_ids[i]
+            delta_kf = self._new_code_suffix(state_id, codes_kf)
+            if delta_kf.shape[1] == 0:
+                continue
+            wav = self._decode_streaming_frames(state_id, delta_kf.to(device=device))
+            if wav.numel() > 0:
+                audios[i] = wav.to(dtype=torch.float32).reshape(-1)
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": audios, "sr": srs},
+        )
+
+    @staticmethod
+    def _resolve_request_ids(
+        count: int,
+        request_ids: object,
+        runtime_additional_information: list[dict[str, Any]] | None,
+    ) -> list[str | None]:
+        if isinstance(request_ids, list | tuple):
+            if len(request_ids) != count:
+                raise ValueError(
+                    "PersonaPlex Code2Wav request id count does not match inputs: "
+                    f"request_ids={len(request_ids)}, inputs={count}"
+                )
+            return [str(request_id) for request_id in request_ids]
+
+        infos = runtime_additional_information or []
+        resolved: list[str | None] = []
+        for index in range(count):
+            info = infos[index] if index < len(infos) else None
+            request_id = None
+            if isinstance(info, Mapping):
+                request_id = info.get("request_id")
+                if request_id is None:
+                    meta = info.get("meta")
+                    if isinstance(meta, Mapping):
+                        request_id = meta.get("request_id")
+            resolved.append(str(request_id) if request_id is not None else None)
+        return resolved
+
+    def _new_code_suffix(self, request_id: str | None, codes_kf: torch.Tensor) -> torch.Tensor:
+        if request_id is None:
+            return codes_kf
+
+        incoming = codes_kf.detach().to(device="cpu", dtype=torch.long)
+        previous = self._request_codes.get(request_id)
+        if (
+            previous is not None
+            and incoming.shape[1] >= previous.shape[1]
+            and torch.equal(incoming[:, : previous.shape[1]], previous)
+        ):
+            delta = incoming[:, previous.shape[1] :]
+            self._request_codes[request_id] = incoming
+            return delta
+
+        self._request_codes[request_id] = incoming if previous is None else torch.cat([previous, incoming], dim=1)
+        return incoming
+
+    def _decode_streaming_frames(self, request_id: str | None, codes_kf: torch.Tensor) -> torch.Tensor:
+        codec, ephemeral = self._codec_for_request(request_id)
+        if not hasattr(codec, "decode_frame"):
+            raise RuntimeError("PersonaPlex Code2Wav requires a streaming Mimi decoder")
+
+        decode_frames = getattr(codec, "decode_frames", None)
+        chunks: list[torch.Tensor] = []
+        for start in range(0, codes_kf.shape[1], _MIMI_DECODE_BATCH_FRAMES):
+            frame_batch = codes_kf[:, start : start + _MIMI_DECODE_BATCH_FRAMES]
+            if frame_batch.shape[1] > 1 and callable(decode_frames):
+                chunks.append(self._flatten_wav(decode_frames(frame_batch.unsqueeze(0))))
+            else:
+                chunks.extend(
+                    self._flatten_wav(codec.decode_frame(frame_batch[:, frame].unsqueeze(0)))
+                    for frame in range(frame_batch.shape[1])
+                )
+        wav = torch.cat(chunks, dim=0) if chunks else codes_kf.new_empty(0, dtype=torch.float32)
+        if ephemeral:
+            codec.reset_streaming()
+        return wav
+
+    def _set_mimi_codecs(self, codecs: list[nn.Module]) -> None:
+        if not codecs:
+            raise ValueError("PersonaPlex Code2Wav requires at least one Mimi decoder")
+        self.mimi = codecs[0]
+        self._additional_mimi = nn.ModuleList(codecs[1:])
+        self._request_codec_slots.clear()
+
+    def _mimi_codecs(self) -> list[nn.Module]:
+        if self.mimi is None:
+            return []
+        return [self.mimi, *self._additional_mimi]
+
+    def _codec_for_request(self, request_id: str | None) -> tuple[nn.Module, bool]:
+        codecs = self._mimi_codecs()
+        if not codecs:
+            raise RuntimeError("PersonaPlexCode2Wav.forward called before Mimi was loaded in load_weights().")
+        occupied = set(self._request_codec_slots.values())
+        if request_id is not None:
+            slot = self._request_codec_slots.get(request_id)
+            if slot is None:
+                slot = next((index for index in range(len(codecs)) if index not in occupied), None)
+                if slot is None:
+                    raise RuntimeError(f"PersonaPlex Code2Wav decoder capacity {len(codecs)} is exhausted")
+                codecs[slot].streaming_init(1)
+                self._request_codec_slots[request_id] = slot
+            return codecs[slot], False
+
+        slot = next((index for index in range(len(codecs)) if index not in occupied), None)
+        if slot is None:
+            raise RuntimeError(f"PersonaPlex Code2Wav decoder capacity {len(codecs)} is exhausted")
+        codecs[slot].streaming_init(1)
+        return codecs[slot], True
+
+    def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
+        for request_id in finished_req_ids:
+            state_id = str(request_id)
+            self._request_codes.pop(state_id, None)
+            slot = self._request_codec_slots.pop(state_id, None)
+            codecs = self._mimi_codecs()
+            if slot is not None and slot < len(codecs):
+                codecs[slot].reset_streaming()
+
+    @staticmethod
+    def _flatten_wav(wav: torch.Tensor) -> torch.Tensor:
+        """Normalize Mimi's decode output to a 1D PCM tensor.
+
+        Mimi may return ``[B, samples]`` or ``[B, C, samples]`` depending on the
+        moshi version; PersonaPlex is mono (``C == 1``), so collapse to 1D.
+        """
+        if wav.dim() == 3:
+            # [B, C, T] -> take batch 0, channel 0.
+            return wav[0, 0]
+        if wav.dim() == 2:
+            # [B, T] -> batch 0.
+            return wav[0]
+        return wav.reshape(-1)
+
+    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput:
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+
+        if isinstance(model_outputs, tuple) and len(model_outputs) == len(OmniOutput._fields):
+            return OmniOutput(*model_outputs)
+
+        if not (isinstance(model_outputs, tuple) and len(model_outputs) == 2):
+            raise TypeError(
+                "PersonaPlexCode2Wav expected OmniOutput, OmniOutput tuple, "
+                f"or (audio_tensor, sr) outputs, got {type(model_outputs)}"
+            )
+
+        audio_tensor, sr = model_outputs
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={
+                "model_outputs": audio_tensor,
+                "sr": sr,
+            },
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Construct and load the external Mimi decoder.
+
+        The primary vLLM weights iterator carries no Code2Wav parameters (Mimi
+        owns its own checkpoint format), so it is drained and the Mimi module is
+        built from the moshi package's loader.
+        """
+        # Drain the primary iterator so callers don't hang on an unconsumed
+        # generator; none of these weights belong to this stage.
+        for _ in weights:
+            pass
+
+        device = self.vllm_config.device_config.device
+        from vllm_omni.model_executor.models.personaplex.personaplex_mimi import (
+            PersonaPlexMimiCodec,
+        )
+
+        checkpoint = Path(self.model_path) / (self._mimi_name or "tokenizer-e351c8d8-checkpoint125.safetensors")
+        codecs = [
+            PersonaPlexMimiCodec(
+                checkpoint=str(checkpoint) if checkpoint.is_file() else None,
+                device=str(device),
+            ).eval()
+            for _ in range(self._max_codec_sessions)
+        ]
+        self._set_mimi_codecs(codecs)
+        self._mimi_device = torch.device(str(device))
+        reported_sr = getattr(codecs[0].model.config, "sampling_rate", None)
+        if reported_sr is not None:
+            self._output_sample_rate = int(reported_sr)
+        bytes_per_mib = 1024**2
+        per_slot_weight_mib = (
+            sum(parameter.numel() * parameter.element_size() for parameter in codecs[0].parameters()) / bytes_per_mib
+        )
+        total_weight_mib = (
+            sum(parameter.numel() * parameter.element_size() for codec in codecs for parameter in codec.parameters())
+            / bytes_per_mib
+        )
+        logger.info(
+            "PersonaPlex Code2Wav loaded %d Mimi stream(s): "
+            "per_slot_weight_mib=%.2f total_weight_mib=%.2f "
+            "num_codebooks=%d sample_rate=%d samples_per_frame=%d",
+            len(codecs),
+            per_slot_weight_mib,
+            total_weight_mib,
+            self._num_codebooks,
+            self._output_sample_rate,
+            self._samples_per_frame,
+        )
+        return {name for name, _ in self.named_parameters() if name.startswith(("mimi.", "_additional_mimi."))}

vllm_omni/model_executor/models/personaplex/personaplex_depformer.py

@@ -0,0 +1,366 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex depformer: the per-step audio code predictor.
+
+The depformer is the small autoregressive head that, conditioned on one temporal
+hidden state, predicts the ``dep_q`` audio codebooks of a single Mimi frame. It
+is a faithful eager port of Moshi's depformer (``moshi.models.lm`` +
+``moshi.modules.transformer``), kept deliberately outside vLLM's paged-attention
+engine: its attention KV is reset every frame and only spans the ``dep_q`` inner
+steps, so paging buys nothing here. This mirrors how the existing TTS code
+predictors live alongside, not inside, the AR talker.
+
+Faithful-port details (all measured from ``nvidia/personaplex-7b-v1``):
+
+* **Per-step weights** (``weights_per_step``): each inner step ``t`` uses its own
+  attention in/out projection and its own gating MLP, selected by index (Moshi's
+  ``multi_linear``). There are ``dep_q`` weight sets per layer.
+* **No positional embedding** (``depformer_pos_emb="none"``): no RoPE, no sin.
+* **Causal attention over all inner steps**: Moshi forces the depformer's
+  ``context`` to None, so each step attends to every prior inner step (KV
+  capacity == ``weights_per_step``); the KV cache is rebuilt every temporal frame.
+* **fp32 RMSNorm** (``rms_norm_f32``, eps 1e-8) with the weight stored as
+  ``alpha`` of shape ``[1, 1, dim]``.
+* **SiLU gating** MLP (Moshi ``ActivationGating``): ``linear_in`` projects to
+  ``2 * hidden``; the halves are combined as ``silu(a) * b`` then ``linear_out``.
+* **ScaledEmbedding** with ``zero_idx = -1`` (the special token maps to the zero
+  vector); audio table is ``card + 1`` rows, text table ``text_card + 1``.
+
+The checkpoint ships ``dep_q = 8`` weight sets; the production loader (Moshi's
+``get_moshi_lm``) overrides ``dep_q = 16`` and fills sets 8..15 by copying 0..7.
+:meth:`PersonaPlexDepformer.load_weights` replicates that expansion so the module
+loads directly from the raw ``model.safetensors`` as well as from an
+already-expanded Moshi instance.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Mapping
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm_omni.model_executor.models.personaplex.configuration_personaplex import (
+    PersonaPlexDepformerConfig,
+)
+
+__all__ = ["PersonaPlexDepformer"]
+
+# Moshi's special "no token / zero embedding" sentinel (``LMModel.zero_token_id``).
+_ZERO_IDX = -1
+
+# A builder takes the source state dict and returns the tensor for one target
+# param, or None when the source lacks it.
+_Builder = Callable[[dict], "torch.Tensor | None"]
+
+
+def _rms_norm_f32(x: torch.Tensor, alpha: torch.Tensor, eps: float) -> torch.Tensor:
+    """Moshi ``_rms_norm`` with the fp32 accumulation path.
+
+    Variance and the alpha multiply run in float32 before casting back, matching
+    ``moshi.modules.transformer._rms_norm`` exactly (``dtype=torch.float``).
+    """
+    x_dtype = x.dtype
+    x_f32 = x.float()
+    var = eps + x_f32.pow(2).mean(dim=-1, keepdim=True)
+    y = x_f32 * (alpha.float() * torch.rsqrt(var))
+    return y.to(x_dtype)
+
+
+class _ScaledEmbedding(nn.Embedding):
+    """``nn.Embedding`` matching Moshi's ``ScaledEmbedding`` (no norm, zero_idx).
+
+    Inputs equal to ``zero_idx`` produce the zero vector; all other inputs are
+    clamped to ``>= 0`` before the lookup. ``norm_emb`` is False for PersonaPlex,
+    so no embedding norm is applied.
+    """
+
+    def __init__(self, num_embeddings: int, embedding_dim: int, zero_idx: int = _ZERO_IDX) -> None:
+        super().__init__(num_embeddings, embedding_dim)
+        self.zero_idx = zero_idx
+
+    def forward(self, idx: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
+        is_zero = idx == self.zero_idx
+        y = super().forward(idx.clamp(min=0))
+        zero = torch.zeros(1, dtype=y.dtype, device=y.device)
+        return torch.where(is_zero[..., None], zero, y)
+
+
+class _DepformerLayer(nn.Module):
+    """One depformer transformer layer with per-step attention + gating.
+
+    Attention and gating weights carry a leading ``dep_q`` axis; the active inner
+    step selects the slice (Moshi ``multi_linear`` / ``gating[t]``). KV is supplied
+    per call so the caller controls the per-frame reset and the sliding window.
+    """
+
+    def __init__(self, config: PersonaPlexDepformerConfig) -> None:
+        super().__init__()
+        dim = config.hidden_size
+        self.dim = dim
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim
+        self.dep_q = config.dep_q
+        self.eps = config.rms_norm_eps
+
+        # Per-step fused QKV and output projections, stored exactly as Moshi does:
+        # ``in_proj_weight`` is [dep_q * 3*dim, dim]; ``out_proj`` is [dep_q*dim, dim].
+        self.in_proj_weight = nn.Parameter(torch.empty(self.dep_q * 3 * dim, dim))
+        self.out_proj_weight = nn.Parameter(torch.empty(self.dep_q * dim, dim))
+
+        # fp32 RMSNorm weights (Moshi ``alpha``, shape [1, 1, dim]).
+        self.norm1_alpha = nn.Parameter(torch.ones(1, 1, dim))
+        self.norm2_alpha = nn.Parameter(torch.ones(1, 1, dim))
+
+        # Per-step SiLU gating MLP. ``linear_in`` -> [2*hidden, dim]; ``linear_out`` -> [dim, hidden].
+        hidden = config.intermediate_size
+        self.gating_in = nn.Parameter(torch.empty(self.dep_q, 2 * hidden, dim))
+        self.gating_out = nn.Parameter(torch.empty(self.dep_q, dim, hidden))
+
+    def _in_proj(self, step: int) -> torch.Tensor:
+        return self.in_proj_weight.view(self.dep_q, 3 * self.dim, self.dim)[step]
+
+    def _out_proj(self, step: int) -> torch.Tensor:
+        return self.out_proj_weight.view(self.dep_q, self.dim, self.dim)[step]
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        step: int,
+        kv: dict[str, torch.Tensor | None],
+    ) -> torch.Tensor:
+        """Run inner step ``step`` for a ``[B, 1, dim]`` input, updating ``kv`` in place."""
+        # --- self-attention (per-step weights, causal sliding window, no RoPE) ---
+        residual = x
+        h = _rms_norm_f32(x, self.norm1_alpha, self.eps).squeeze(1)  # [B, dim]
+        qkv = F.linear(h, self._in_proj(step))  # [B, 3*dim]
+        q, k, v = qkv.split(self.dim, dim=-1)
+        b = q.shape[0]
+        q = q.view(b, self.num_heads, 1, self.head_dim)
+        k = k.view(b, self.num_heads, 1, self.head_dim)
+        v = v.view(b, self.num_heads, 1, self.head_dim)
+
+        if kv["k"] is None:
+            k_hist, v_hist = k, v
+        else:
+            k_hist = torch.cat([kv["k"], k], dim=2)
+            v_hist = torch.cat([kv["v"], v], dim=2)
+        kv["k"], kv["v"] = k_hist, v_hist
+
+        # Moshi forces the depformer's ``context`` to None (lm.py: kwargs_dep[
+        # "context"] = None), so attention spans ALL prior inner steps -- the KV
+        # capacity is ``weights_per_step`` (== dep_q), never windowed. The query is
+        # the newest step and every key is causal w.r.t. it, so no mask is needed.
+        attn = F.scaled_dot_product_attention(q, k_hist, v_hist)  # [B, H, 1, Dh]
+        attn = attn.reshape(b, self.dim)
+        out = F.linear(attn, self._out_proj(step)).unsqueeze(1)  # [B, 1, dim]
+        x = residual + out
+
+        # --- feed-forward (per-step SiLU gating) ---
+        residual = x
+        h = _rms_norm_f32(x, self.norm2_alpha, self.eps).squeeze(1)  # [B, dim]
+        gated = F.linear(h, self.gating_in[step])  # [B, 2*hidden]
+        a, b_half = gated.chunk(2, dim=-1)
+        update = F.linear(F.silu(a) * b_half, self.gating_out[step]).unsqueeze(1)
+        return residual + update
+
+
+class PersonaPlexDepformer(nn.Module):
+    """Moshi depformer as a standalone per-frame code predictor.
+
+    Call :meth:`forward` once per temporal frame with the frame's text token and
+    temporal hidden state; it runs ``dep_q`` inner AR steps (greedy by default,
+    matching Moshi's ``use_sampling=False``) and returns the ``dep_q`` audio codes.
+    """
+
+    def __init__(
+        self,
+        config: PersonaPlexDepformerConfig,
+        temporal_hidden_size: int,
+        text_card: int = 32000,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.dep_q = config.dep_q
+        self.card = config.card
+        dim = config.hidden_size
+
+        # Per-codebook projection of the temporal hidden state (``depformer_multi_linear``).
+        self.depformer_in = nn.ModuleList([nn.Linear(temporal_hidden_size, dim, bias=False) for _ in range(self.dep_q)])
+        # Step-0 conditions on the text token; steps 1..dep_q-1 on the previous audio code.
+        self.depformer_text_emb = _ScaledEmbedding(text_card + 1, dim)
+        self.depformer_emb = nn.ModuleList([_ScaledEmbedding(self.card + 1, dim) for _ in range(self.dep_q - 1)])
+        self.layers = nn.ModuleList([_DepformerLayer(config) for _ in range(config.num_hidden_layers)])
+        # Per-codebook output heads (Moshi ``linears``).
+        self.linears = nn.ModuleList([nn.Linear(dim, self.card, bias=False) for _ in range(self.dep_q)])
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        text_token: torch.Tensor,
+        transformer_out: torch.Tensor,
+        audio_tokens: torch.Tensor | None = None,
+        audio_provided: torch.Tensor | None = None,
+        return_logits: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """Predict the ``dep_q`` audio codes for one frame.
+
+        Args:
+            text_token: ``[B]`` sampled text token for the frame.
+            transformer_out: ``[B, 1, temporal_hidden]`` temporal hidden state.
+            audio_tokens: optional ``[B, dep_q]`` teacher-forcing codes.
+            audio_provided: optional ``[B, dep_q]`` bool mask; where True the
+                corresponding ``audio_tokens`` entry overrides the sampled code as
+                the next-step input (partial teacher forcing, as in Moshi).
+            return_logits: also return the per-step logits ``[B, dep_q, card]``.
+
+        Returns:
+            ``[B, dep_q]`` long tensor of sampled (greedy) audio codes, plus the
+            per-step logits when ``return_logits`` is set.
+        """
+        if transformer_out.dim() != 3 or transformer_out.shape[1] != 1:
+            raise ValueError(f"transformer_out must be [B, 1, H]; got {tuple(transformer_out.shape)}")
+        prev = text_token
+        kv = [{"k": None, "v": None} for _ in self.layers]
+        codes: list[torch.Tensor] = []
+        logits_all: list[torch.Tensor] = []
+        for step in range(self.dep_q):
+            x = self.depformer_in[step](transformer_out)  # [B, 1, dim]
+            if step == 0:
+                cond = self.depformer_text_emb(prev.unsqueeze(1))
+            else:
+                cond = self.depformer_emb[step - 1](prev.unsqueeze(1))
+            x = x + cond
+            for li, layer in enumerate(self.layers):
+                x = layer(x, step, kv[li])
+            logits = self.linears[step](x).squeeze(1)  # [B, card]
+            sampled = logits.float().argmax(dim=-1)  # greedy [B]
+            codes.append(sampled)
+            if return_logits:
+                logits_all.append(logits.float())
+            if audio_provided is not None and audio_tokens is not None:
+                prev = torch.where(audio_provided[:, step], audio_tokens[:, step], sampled)
+            else:
+                prev = sampled
+        stacked = torch.stack(codes, dim=1)  # [B, dep_q]
+        if return_logits:
+            return stacked, torch.stack(logits_all, dim=1)  # [B, dep_q, card]
+        return stacked
+
+    # ------------------------------------------------------------------
+    # Weight loading
+    # ------------------------------------------------------------------
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]] | Mapping[str, torch.Tensor],
+    ) -> set[str]:
+        """Load Moshi depformer weights, expanding ``dep_q`` 8 -> 16 if needed.
+
+        Accepts the raw checkpoint names (``depformer.*``, ``depformer_in.N``,
+        ``depformer_emb.N``, ``depformer_text_emb``, ``linears.N``). Per-step
+        modules present only for steps 0..7 are copied to 8..15, and the fused
+        attention projections are tiled, exactly as Moshi's ``get_moshi_lm`` does.
+        """
+        src = dict(weights.items() if isinstance(weights, Mapping) else weights)
+        src = {k: v for k, v in src.items() if self._is_depformer_key(k)}
+        params = dict(self.named_parameters())
+        loaded: set[str] = set()
+
+        for tgt, build in self._weight_plan().items():
+            tensor = build(src)
+            if tensor is None:
+                continue
+            param = params[tgt]
+            if tensor.shape != param.shape:
+                raise ValueError(
+                    f"shape mismatch for {tgt}: checkpoint {tuple(tensor.shape)} vs param {tuple(param.shape)}"
+                )
+            param.data.copy_(tensor.to(param.dtype))
+            loaded.add(tgt)
+        return loaded
+
+    @staticmethod
+    def _is_depformer_key(name: str) -> bool:
+        return name.startswith(("depformer.", "depformer_in.", "depformer_emb.", "depformer_text_emb", "linears."))
+
+    def _weight_plan(self) -> dict[str, _Builder]:
+        """Map each of this module's params to a builder over the source state dict."""
+        dim = self.config.hidden_size
+
+        def step_or_copy(prefix: str, step: int, suffix: str = "weight") -> _Builder:
+            # cb 8..15 fall back to cb (step-8) when absent (Moshi copy_missing_weights).
+            def build(src: dict[str, torch.Tensor]) -> torch.Tensor | None:
+                key = f"{prefix}.{step}.{suffix}"
+                if key in src:
+                    return src[key]
+                if step >= 8:
+                    return src.get(f"{prefix}.{step - 8}.{suffix}")
+                return None
+
+            return build
+
+        def expand_attn(layer: int, kind: str, per_step_rows: int) -> _Builder:
+            # ``in_proj_weight`` [steps*per_step_rows, dim] / ``out_proj.weight`` likewise;
+            # tile an 8-step checkpoint block up to dep_q (Moshi concat-copy expansion).
+            def build(src: dict[str, torch.Tensor]) -> torch.Tensor | None:
+                key = f"depformer.layers.{layer}.self_attn.{kind}"
+                t = src.get(key)
+                if t is None:
+                    return None
+                steps = t.shape[0] // per_step_rows
+                if steps < self.dep_q:
+                    reps = (self.dep_q + steps - 1) // steps
+                    t = t.repeat(reps, 1)[: self.dep_q * per_step_rows]
+                return t
+
+            return build
+
+        plan: dict[str, _Builder] = {}
+        # depformer_in / linears / depformer_emb / text_emb
+        for step in range(self.dep_q):
+            plan[f"depformer_in.{step}.weight"] = step_or_copy("depformer_in", step)
+            plan[f"linears.{step}.weight"] = step_or_copy("linears", step)
+        plan["depformer_text_emb.weight"] = lambda src: src.get("depformer_text_emb.weight")
+        for step in range(self.dep_q - 1):
+            plan[f"depformer_emb.{step}.weight"] = step_or_copy("depformer_emb", step)
+
+        # per-layer attention + norms + gating
+        for layer in range(self.config.num_hidden_layers):
+            lp = f"depformer.layers.{layer}"
+            plan[f"layers.{layer}.in_proj_weight"] = expand_attn(layer, "in_proj_weight", 3 * dim)
+            plan[f"layers.{layer}.out_proj_weight"] = expand_attn(layer, "out_proj.weight", dim)
+            plan[f"layers.{layer}.norm1_alpha"] = self._reshape_alpha(f"{lp}.norm1.alpha")
+            plan[f"layers.{layer}.norm2_alpha"] = self._reshape_alpha(f"{lp}.norm2.alpha")
+            plan[f"layers.{layer}.gating_in"] = self._stack_gating(lp, "linear_in")
+            plan[f"layers.{layer}.gating_out"] = self._stack_gating(lp, "linear_out")
+        return plan
+
+    @staticmethod
+    def _reshape_alpha(key: str) -> _Builder:
+        def build(src: dict[str, torch.Tensor]) -> torch.Tensor | None:
+            t = src.get(key)
+            return None if t is None else t.reshape(1, 1, -1)
+
+        return build
+
+    def _stack_gating(self, layer_prefix: str, which: str) -> _Builder:
+        """Stack the dep_q per-step gating weights into one [dep_q, ...] tensor.
+
+        Steps 8..15 fall back to step (n-8) when the checkpoint only ships 8
+        (Moshi ``copy_missing_weights``).
+        """
+
+        def build(src: dict[str, torch.Tensor]) -> torch.Tensor | None:
+            slabs: list[torch.Tensor] = []
+            for step in range(self.dep_q):
+                key = f"{layer_prefix}.gating.{step}.{which}.weight"
+                if key not in src and step >= 8:
+                    key = f"{layer_prefix}.gating.{step - 8}.{which}.weight"
+                t = src.get(key)
+                if t is None:
+                    return None
+                slabs.append(t)
+            return torch.stack(slabs, dim=0)
+
+        return build

vllm_omni/model_executor/models/personaplex/personaplex_embeddings.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex input embeddings (Moshi ``embed_codes``).
+
+The temporal transformer (:class:`HeliumForCausalLM`) consumes precomputed
+``inputs_embeds`` rather than token ids. Those embeddings come from Moshi's
+``embed_codes``: a per-frame ``[B, 1 + n_q, S]`` token stack (row 0 = text,
+rows 1..n_q = audio codebooks) is mapped through one text embedding plus ``n_q``
+audio-codebook embeddings and summed into ``[B, S, hidden]``.
+
+This module owns exactly those embedding tables (Moshi ``text_emb`` and
+``emb.0..n_q-1``), which ``modeling_helium`` deliberately omits. The talker
+(Task 2) uses it to build the temporal ``inputs_embeds`` for both the system /
+voice-prompt prefill and the per-frame decode step.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.personaplex.configuration_personaplex import (
+    PersonaPlexConfig,
+)
+from vllm_omni.model_executor.models.personaplex.personaplex_depformer import (
+    _ScaledEmbedding,
+)
+
+__all__ = ["PersonaPlexInputEmbeddings"]
+
+
+class PersonaPlexInputEmbeddings(nn.Module):
+    """Moshi ``embed_codes``: summed text + audio-codebook embeddings.
+
+    Args:
+        config: the top-level PersonaPlex config (supplies the temporal hidden
+            size, the audio cardinality / codebook count, and the text table size).
+    """
+
+    def __init__(self, config: PersonaPlexConfig) -> None:
+        super().__init__()
+        dim = config.temporal_config.hidden_size
+        self.num_audio_codebooks = config.num_audio_codebooks  # n_q (16)
+        # Text table has one extra row for the special token (text_card + 1).
+        self.text_emb = _ScaledEmbedding(config.text_embedding_rows, dim)
+        # Each audio table is card + 1 rows; ``embed_codes`` uses all n_q codebooks.
+        self.audio_emb = nn.ModuleList(
+            [_ScaledEmbedding(config.audio_vocab_size + 1, dim) for _ in range(self.num_audio_codebooks)]
+        )
+
+    def forward(self, sequence: torch.Tensor) -> torch.Tensor:
+        """Map a ``[B, 1 + n_q, S]`` token stack to ``[B, S, hidden]``.
+
+        Row 0 is the text stream; rows ``1..n_q`` are the audio codebooks (in
+        Moshi's layout the audio rows carry both the agent and user streams).
+        """
+        if sequence.dim() != 3 or sequence.shape[1] != self.num_audio_codebooks + 1:
+            raise ValueError(f"sequence must be [B, {self.num_audio_codebooks + 1}, S]; got {tuple(sequence.shape)}")
+        emb = self.text_emb(sequence[:, 0])  # [B, S, hidden]
+        for cb in range(self.num_audio_codebooks):
+            emb = emb + self.audio_emb[cb](sequence[:, cb + 1])
+        return emb
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]] | Mapping[str, torch.Tensor],
+    ) -> set[str]:
+        """Load Moshi ``text_emb`` and ``emb.0..n_q-1`` embedding tables."""
+        src = dict(weights.items() if isinstance(weights, Mapping) else weights)
+        params = dict(self.named_parameters())
+        loaded: set[str] = set()
+
+        plan: dict[str, str] = {"text_emb.weight": "text_emb.weight"}
+        for cb in range(self.num_audio_codebooks):
+            plan[f"audio_emb.{cb}.weight"] = f"emb.{cb}.weight"
+
+        for tgt, key in plan.items():
+            if key not in src:
+                continue
+            param = params[tgt]
+            tensor = src[key]
+            if tensor.shape != param.shape:
+                raise ValueError(
+                    f"shape mismatch for {tgt}: checkpoint {tuple(tensor.shape)} vs param {tuple(param.shape)}"
+                )
+            param.data.copy_(tensor.to(param.dtype))
+            loaded.add(tgt)
+        return loaded

vllm_omni/model_executor/models/personaplex/personaplex_mimi.py

@@ -0,0 +1,413 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Streaming Mimi codec for PersonaPlex (moshi-free).
+
+Frame-clocked duplex needs a codec that encodes/decodes exactly one 80 ms frame
+per call with state carried across the conversation. The ``moshi`` package
+provided that; this module removes the dependency by combining:
+
+- **transformers ``MimiModel``** (``kyutai/mimi`` — the same checkpoint family
+  PersonaPlex ships) for the quantizer and every SEANet conv weight. Its
+  streaming support only covers the encoder convs, so conv streaming is done
+  here instead, uniformly.
+- **Our own streaming wrappers** mirroring the Moshi reference semantics (MIT),
+  verified against recorded reference outputs:
+  * ``Conv1d``: left-context carry of ``effective_kernel - stride`` samples,
+    zero-initialized at stream start (``pad_mode="constant"``).
+  * ``ConvTranspose1d``: overlap-add tail carry of ``kernel - stride`` output
+    samples, with the double-counted bias subtracted on merge.
+  * transformer: the encoder/decoder transformers use a 250-position sliding
+    context over a ring KV with absolute-offset RoPE. Hugging Face's cache path
+    diverges once the window engages (position 250), so the transformers are
+    reimplemented here on the same ring-KV design as
+    ``personaplex_temporal.py`` and loaded directly from the PersonaPlex
+    checkpoint's fused layout (LayerNorm + per-layer LayerScale + GELU FFN).
+
+All per-stream state is ``[B, ...]`` with per-row reset (``reset_slot``), so the
+codec composes with elastic slot recycling in batched duplex serving.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.personaplex.personaplex_temporal import (
+    _apply_rope,
+    _RingKV,
+)
+
+DEFAULT_HF_REPO = "kyutai/mimi"
+FRAME_SIZE = 1920
+CODEBOOKS = 8
+
+
+def _map_moshi_codec_weights(
+    state_dict: dict[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+    """Map bundled Moshi codec keys to ``transformers.MimiModel`` keys."""
+
+    mapped: dict[str, torch.Tensor] = {}
+    transformer_prefixes = ("encoder_transformer.", "decoder_transformer.")
+    for name, tensor in state_dict.items():
+        if name.startswith(transformer_prefixes):
+            continue
+
+        if name.startswith("encoder.model."):
+            target = name.replace("encoder.model.", "encoder.layers.", 1)
+            target = target.replace(".conv.conv.", ".conv.")
+        elif name.startswith("decoder.model."):
+            target = name.replace("decoder.model.", "decoder.layers.", 1)
+            target = target.replace(".convtr.convtr.", ".conv.")
+            target = target.replace(".conv.conv.", ".conv.")
+        elif name.startswith("downsample.conv.conv.conv."):
+            target = name.replace("downsample.conv.conv.conv.", "downsample.conv.", 1)
+        elif name.startswith("upsample.convtr.convtr.convtr."):
+            target = name.replace("upsample.convtr.convtr.convtr.", "upsample.conv.", 1)
+        elif name.startswith("quantizer.rvq_first."):
+            target = name.replace(
+                "quantizer.rvq_first.",
+                "quantizer.semantic_residual_vector_quantizer.",
+                1,
+            )
+        elif name.startswith("quantizer.rvq_rest."):
+            target = name.replace(
+                "quantizer.rvq_rest.",
+                "quantizer.acoustic_residual_vector_quantizer.",
+                1,
+            )
+        else:
+            raise KeyError(f"unrecognized PersonaPlex Mimi checkpoint key: {name}")
+
+        target = target.replace(".vq.layers.", ".layers.")
+        target = target.replace("._codebook._initialized", ".codebook.initialized")
+        target = target.replace("._codebook.cluster_usage", ".codebook.cluster_usage")
+        target = target.replace("._codebook.embedding_sum", ".codebook.embed_sum")
+        mapped[target] = tensor
+    return mapped
+
+
+class _StreamConv1d:
+    """Moshi ``RawStreamingConv1d``: causal left-context carry per call.
+
+    ``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.
+    """
+
+    def __init__(self, conv: nn.Conv1d, pad_mode: str = "constant") -> None:
+        self.conv = conv
+        self.kernel = (conv.kernel_size[0] - 1) * conv.dilation[0] + 1
+        self.stride = conv.stride[0]
+        self.pad_mode = pad_mode
+        self.prev: torch.Tensor | None = None
+        self._fresh: torch.Tensor | None = None
+
+    def reset(self, batch_size: int, device, dtype) -> None:
+        pad = self.kernel - self.stride
+        self.prev = torch.zeros(batch_size, self.conv.in_channels, pad, device=device, dtype=dtype)
+        self._fresh = torch.ones(batch_size, dtype=torch.bool)
+
+    def reset_slot(self, b: int) -> None:
+        self.prev[b].zero_()
+        self._fresh[b] = True
+
+    def __call__(self, x: torch.Tensor) -> torch.Tensor:
+        if self.pad_mode == "replicate" and bool(self._fresh.any()):
+            pad = self.prev.shape[-1]
+            edge = x[..., 0:1].expand(-1, -1, pad)
+            fresh = self._fresh.to(x.device).view(-1, 1, 1)
+            self.prev = torch.where(fresh, edge.to(self.prev.dtype), self.prev)
+        self._fresh[:] = False
+        x = torch.cat([self.prev, x], dim=-1)
+        t = x.shape[-1]
+        num_frames = max(0, (t - self.kernel) // self.stride + 1)
+        self.prev = x[..., num_frames * self.stride :]
+        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."""
+
+    def __init__(self, conv: nn.ConvTranspose1d) -> None:
+        self.conv = conv
+        self.kernel = conv.kernel_size[0]
+        self.stride = conv.stride[0]
+        self.partial: torch.Tensor | None = None
+
+    def reset(self, batch_size: int, device, dtype) -> None:
+        self.partial = torch.zeros(
+            batch_size, self.conv.out_channels, self.kernel - self.stride, device=device, dtype=dtype
+        )
+        self._fresh = torch.ones(batch_size, dtype=torch.bool)
+
+    def reset_slot(self, b: int) -> None:
+        self.partial[b].zero_()
+        self._fresh[b] = True
+
+    def __call__(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.conv(x)
+        length = out.shape[-1]
+        tail = self.kernel - self.stride
+        pt = self.partial.shape[-1]
+        merge = self.partial.clone()
+        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] = 0.0
+            self._fresh[:] = False
+        out[..., :pt] += merge
+        self.partial = out[..., length - tail :].clone()
+        return out[..., : length - tail]
+
+
+class _MimiTransformerLayer(nn.Module):
+    """Moshi mimi transformer layer: LayerNorm, LayerScale, GELU FFN, no biases."""
+
+    def __init__(self, dim: int = 512, num_heads: int = 8, ffn: int = 2048) -> None:
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.in_proj_weight = nn.Parameter(torch.empty(3 * dim, dim))
+        self.out_proj_weight = nn.Parameter(torch.empty(dim, dim))
+        self.linear1 = nn.Parameter(torch.empty(ffn, dim))
+        self.linear2 = nn.Parameter(torch.empty(dim, ffn))
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.scale1 = nn.Parameter(torch.empty(dim))
+        self.scale2 = nn.Parameter(torch.empty(dim))
+
+    def forward(self, x: torch.Tensor, kv: _RingKV, offset: torch.Tensor, context: int) -> torch.Tensor:
+        B, T, _ = x.shape
+        h = self.norm1(x)
+        qkv = F.linear(h, self.in_proj_weight)
+        qkv = qkv.view(B, T, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        q, k = _apply_rope(q, k, offset)
+        keys, values, pos_k = kv.complete(k, v)
+        pos_k = pos_k.view(pos_k.shape[0], 1, pos_k.shape[1])
+        pos_q = offset + torch.arange(T, device=q.device, dtype=torch.long).view(1, -1, 1)
+        delta = pos_q - pos_k
+        attn_bias = (pos_k >= 0) & (delta >= 0) & (delta < context)
+        attn = F.scaled_dot_product_attention(q, keys, values, attn_bias.unsqueeze(1), dropout_p=0.0)
+        attn = attn.transpose(1, 2).reshape(B, T, self.dim)
+        x = x + self.scale1 * F.linear(attn, self.out_proj_weight)
+        h = self.norm2(x)
+        h = F.linear(F.gelu(F.linear(h, self.linear1)), self.linear2)
+        return x + self.scale2 * h
+
+
+class _MimiStreamingTransformer(nn.Module):
+    """The 8-layer mimi encoder/decoder transformer as a stateful stepper."""
+
+    def __init__(self, num_layers: int = 8, dim: int = 512, num_heads: int = 8, context: int = 250) -> None:
+        super().__init__()
+        self.context = context
+        self.layers = nn.ModuleList([_MimiTransformerLayer(dim, num_heads) for _ in range(num_layers)])
+        self._kv: list[_RingKV] | None = None
+        self._offset: torch.Tensor | None = None
+
+    def streaming_init(self, batch_size: int) -> None:
+        p = next(self.parameters())
+        heads = self.layers[0].num_heads
+        hd = self.layers[0].head_dim
+        self._kv = [_RingKV(batch_size, heads, hd, self.context, p.device, p.dtype) for _ in self.layers]
+        self._offset = torch.zeros(1, device=p.device, dtype=torch.long)
+
+    def reset_streaming(self) -> None:
+        for kv in self._kv:
+            kv.reset()
+        self._offset.zero_()
+
+    def reset_slot(self, b: int) -> None:
+        for kv in self._kv:
+            kv.reset_slot(b)
+
+    def step(self, x: torch.Tensor) -> torch.Tensor:
+        """``x`` is ``[B, T, dim]`` (T = positions this frame, typically 2)."""
+        for layer, kv in zip(self.layers, self._kv):
+            x = layer(x, kv, self._offset, self.context)
+        self._offset.add_(x.shape[1])
+        return x
+
+    def load_weights(self, state_dict: dict[str, torch.Tensor], prefix: str) -> int:
+        loaded = 0
+        with torch.no_grad():
+            for i, layer in enumerate(self.layers):
+                base = f"{prefix}.transformer.layers.{i}"
+                pairs = [
+                    (layer.in_proj_weight, f"{base}.self_attn.in_proj_weight"),
+                    (layer.out_proj_weight, f"{base}.self_attn.out_proj.weight"),
+                    (layer.linear1, f"{base}.linear1.weight"),
+                    (layer.linear2, f"{base}.linear2.weight"),
+                    (layer.norm1.weight, f"{base}.norm1.weight"),
+                    (layer.norm1.bias, f"{base}.norm1.bias"),
+                    (layer.norm2.weight, f"{base}.norm2.weight"),
+                    (layer.norm2.bias, f"{base}.norm2.bias"),
+                    (layer.scale1, f"{base}.layer_scale_1.scale"),
+                    (layer.scale2, f"{base}.layer_scale_2.scale"),
+                ]
+                for param, name in pairs:
+                    param.data.copy_(state_dict[name].reshape(param.shape).to(param.dtype))
+                    loaded += 1
+        return loaded
+
+
+def _walk_seanet(layers) -> list[tuple[str, object]]:
+    """Wrap a Hugging Face Mimi SEANet layer list with streaming conv state."""
+    stages: list[tuple[str, object]] = []
+    for layer in layers:
+        kind = type(layer).__name__
+        if kind == "MimiConv1d":
+            stages.append(("conv", _StreamConv1d(layer.conv)))
+        elif kind == "MimiConvTranspose1d":
+            stages.append(("convtr", _StreamConvTr1d(layer.conv)))
+        elif kind == "ELU":
+            stages.append(("act", layer))
+        elif kind == "MimiResnetBlock":
+            block = (
+                layer.block[0],
+                _StreamConv1d(layer.block[1].conv),
+                layer.block[2],
+                _StreamConv1d(layer.block[3].conv),
+            )
+            stages.append(("res", block))
+        else:  # pragma: no cover - unexpected layer type
+            raise ValueError(f"unhandled Mimi SEANet layer: {kind}")
+    return stages
+
+
+class PersonaPlexMimiCodec(nn.Module):
+    """Streaming Mimi encode/decode at one 80 ms frame per call (moshi-free)."""
+
+    def __init__(self, hf_repo: str = DEFAULT_HF_REPO, checkpoint: str | None = None, device: str = "cuda") -> None:
+        super().__init__()
+        from huggingface_hub import hf_hub_download
+        from safetensors.torch import load_file
+        from transformers import MimiConfig, MimiModel
+
+        self.device = torch.device(device)
+
+        # The PersonaPlex repo ships the reference mimi checkpoint in the moshi
+        # fused layout. Build the matching Transformers graph locally, then load
+        # every codec tensor from that bundled checkpoint. Only the two HF
+        # transformer stacks remain absent because the streaming implementations
+        # below replace them with the checkpoint's fused QKV layout.
+        if checkpoint is None:
+            checkpoint = hf_hub_download(
+                "nvidia/personaplex-7b-v1",
+                "tokenizer-e351c8d8-checkpoint125.safetensors",
+            )
+        sd = load_file(checkpoint, device=str(self.device))
+        self.model = MimiModel(MimiConfig())
+        codec_state = _map_moshi_codec_weights(sd)
+        incompatible = self.model.load_state_dict(codec_state, strict=False)
+        expected_missing = {
+            name
+            for name in self.model.state_dict()
+            if name.startswith(("encoder_transformer.", "decoder_transformer."))
+        }
+        if set(incompatible.missing_keys) != expected_missing or incompatible.unexpected_keys:
+            raise RuntimeError(
+                "PersonaPlex Mimi checkpoint did not exactly cover the local codec graph: "
+                f"missing={sorted(set(incompatible.missing_keys) - expected_missing)}, "
+                f"unexpected={sorted(incompatible.unexpected_keys)}"
+            )
+        # The fused streaming transformers below replace these unloaded HF
+        # stacks. Drop their randomly initialized weights before moving Mimi to
+        # the target device.
+        del self.model.encoder_transformer, self.model.decoder_transformer
+        self.model = self.model.to(self.device).eval()
+        self.dtype = next(self.model.parameters()).dtype
+        self.encoder_transformer = _MimiStreamingTransformer().to(self.device, self.dtype)
+        self.decoder_transformer = _MimiStreamingTransformer().to(self.device, self.dtype)
+        n_enc = self.encoder_transformer.load_weights(sd, "encoder_transformer")
+        n_dec = self.decoder_transformer.load_weights(sd, "decoder_transformer")
+        assert n_enc == n_dec == 80, (n_enc, n_dec)
+        del sd
+
+        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
+
+    # -- streaming state ------------------------------------------------------
+
+    def _conv_states(self):
+        for _, stage in (*self._enc_stages, *self._dec_stages):
+            if isinstance(stage, (_StreamConv1d, _StreamConvTr1d)):
+                yield stage
+            elif isinstance(stage, tuple):
+                yield stage[1]
+                yield stage[3]
+        yield self._downsample
+        yield self._upsample
+
+    def streaming_init(self, batch_size: int) -> None:
+        self._batch_size = batch_size
+        for s in self._conv_states():
+            s.reset(batch_size, self.device, self.dtype)
+        self.encoder_transformer.streaming_init(batch_size)
+        self.decoder_transformer.streaming_init(batch_size)
+
+    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)
+        self.encoder_transformer.reset_streaming()
+        self.decoder_transformer.reset_streaming()
+
+    def reset_slot(self, b: int) -> None:
+        for s in self._conv_states():
+            s.reset_slot(b)
+        self.encoder_transformer.reset_slot(b)
+        self.decoder_transformer.reset_slot(b)
+
+    # -- per-frame codec -------------------------------------------------------
+
+    @staticmethod
+    def _run_stages(x: torch.Tensor, stages) -> torch.Tensor:
+        for kind, stage in stages:
+            if kind == "res":
+                act0, conv1, act2, conv3 = stage
+                x = x + conv3(act2(conv1(act0(x))))
+            else:
+                x = stage(x)
+        return x
+
+    @torch.no_grad()
+    def encode_frame(self, pcm: torch.Tensor) -> torch.Tensor:
+        """``[B, frame_size]`` float PCM -> ``[B, 8]`` codes."""
+        x = pcm.to(self.device, self.dtype).view(-1, 1, FRAME_SIZE)
+        x = self._run_stages(x, self._enc_stages)
+        x = self.encoder_transformer.step(x.transpose(1, 2)).transpose(1, 2)
+        x = self._downsample(x)
+        codes = self.model.quantizer.encode(x)  # [Q, B, T]
+        return codes[:CODEBOOKS, :, 0].transpose(0, 1).contiguous()
+
+    @torch.no_grad()
+    def decode_frame(self, codes: torch.Tensor) -> torch.Tensor:
+        """``[B, 8]`` codes -> ``[B, frame_size]`` float PCM."""
+        emb = self.model.quantizer.decode(codes.to(self.device).view(-1, CODEBOOKS, 1))
+        emb = self._upsample(emb)
+        emb = self.decoder_transformer.step(emb.transpose(1, 2)).transpose(1, 2)
+        x = self._run_stages(emb, self._dec_stages)
+        return x[:, 0, :]
+
+    @torch.no_grad()
+    def decode_frames(self, codes: torch.Tensor) -> torch.Tensor:
+        """``[B, 8, F]`` codes -> ``[B, F * frame_size]`` float PCM."""
+        emb = self.model.quantizer.decode(codes.to(self.device))
+        emb = self._upsample(emb)
+        emb = self.decoder_transformer.step(emb.transpose(1, 2)).transpose(1, 2)
+        x = self._run_stages(emb, self._dec_stages)
+        return x[:, 0, :]

vllm_omni/model_executor/models/personaplex/personaplex_talker.py

@@ -0,0 +1,639 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex talker: the temporal transformer as a vLLM-native omni AR stage.
+
+This is the stage-0 (``LLM_AR``) model the OmniGPUModelRunner drives. It composes
+three pieces, each verified in isolation against Moshi:
+
+* the Helium temporal transformer (:class:`HeliumModel`) on vLLM paged attention,
+  consuming per-frame ``inputs_embeds`` and producing the per-frame hidden state;
+* the input embeddings (:class:`PersonaPlexInputEmbeddings`, ``embed_codes``) that
+  build those ``inputs_embeds`` from the delayed 17-row token stack;
+* the depformer (:class:`PersonaPlexDepformer`) that, conditioned on the temporal
+  hidden state and the sampled text token, predicts the per-frame audio codes.
+
+Per-frame protocol (OmniGPUModelRunner, gpu_model_runner.py):
+
+1. ``compute_logits`` produces the text logits; the engine samples the text token.
+2. ``preprocess`` (per-request, with that request's ``additional_information``)
+   carries Moshi's acoustic-delay cache and the precomputed user-audio code stream,
+   builds the base ``inputs_embeds`` for the current frame, and exposes the
+   previous frame's temporal hidden + text-step embedding via ``mtp_inputs``.
+3. ``talker_mtp`` (batched, stateless) runs the depformer to predict the agent
+   codes and finishes the next frame's ``inputs_embeds``; the codes are stored
+   under ``talker_mtp_output_key=("codes","audio")`` for the Mimi code2wav stage.
+
+Phase 1 is turn-based: the user-audio rows (9..16) come from a precomputed Mimi
+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
+
+import torch
+from torch import nn
+from vllm.config import VllmConfig
+from vllm.distributed import get_pp_group
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix
+from vllm.sequence import IntermediateTensors
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.personaplex.configuration_personaplex import (
+    PersonaPlexConfig,
+)
+from vllm_omni.model_executor.models.personaplex.modeling_helium import HeliumModel
+from vllm_omni.model_executor.models.personaplex.personaplex_depformer import (
+    PersonaPlexDepformer,
+)
+from vllm_omni.model_executor.models.personaplex.personaplex_embeddings import (
+    PersonaPlexInputEmbeddings,
+)
+
+__all__ = ["PersonaPlexTalkerForConditionalGeneration"]
+
+
+class PersonaPlexTalkerForConditionalGeneration(nn.Module):
+    """vLLM-native PersonaPlex talker (temporal transformer + depformer).
+
+    Plain ``nn.Module`` (not ``SupportsPP``): inheriting a vLLM Protocol puts
+    ``Protocol`` in the MRO and breaks vLLM's runtime ``isinstance`` check for
+    ``VllmModelForTextGeneration``, which would misclassify the talker as a
+    non-generate model. Qwen3-TTS's talker is plain ``nn.Module`` for the same
+    reason; PersonaPlex does not need pipeline parallelism.
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        self.vllm_config = vllm_config
+        config: PersonaPlexConfig = vllm_config.model_config.hf_config  # type: ignore[assignment]
+        self.config = config
+        self.temporal_config = config.temporal_config
+        hidden = self.temporal_config.hidden_size
+        self._dtype = getattr(vllm_config.model_config, "dtype", torch.bfloat16)
+
+        # Temporal backbone on vLLM paged attention (consumes inputs_embeds).
+        self.model = HeliumModel(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+            config=self.temporal_config,
+        )
+        # Text head (Moshi text_linear -> lm_head), vocab = text_vocab_size (32000).
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                config.text_vocab_size,
+                hidden,
+                quant_config=vllm_config.quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+            )
+        else:
+            self.lm_head = PPMissingLayer()
+        self.logits_processor = LogitsProcessor(config.text_vocab_size)
+        self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
+
+        # Verified custom components: embed_codes + depformer.
+        self.input_embeddings = PersonaPlexInputEmbeddings(config)
+        self.depformer = PersonaPlexDepformer(
+            config.depformer_config,
+            temporal_hidden_size=hidden,
+            text_card=config.text_vocab_size,
+        )
+
+        # Omni AR runner contract.
+        self.have_multimodal_outputs = True
+        self.has_preprocess = True
+        self.has_postprocess = True  # capture each frame's hidden for the next depformer step
+        self.requires_full_prefix_cached_hidden_states = False
+        # Keep the per-frame "last" hidden on GPU (avoids a CPU round-trip each step).
+        self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("hidden_states", "last")}
+        self.mtp_hidden_size = hidden
+        self.talker_mtp_output_key = ("codes", "audio")
+        # dep_q audio codebooks per frame; only cb 0..num_active are vocoded.
+        self.dep_q = config.depformer_config.dep_q
+        self.num_active_codebooks = config.depformer_config.num_active_codebooks
+
+    # ------------------------------------------------------------------
+    # Core forward / logits
+    # ------------------------------------------------------------------
+    def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
+        """Placeholder embedding for vLLM's ``VllmModel`` protocol.
+
+        The talker is driven by precomputed ``inputs_embeds`` (built in
+        ``preprocess`` via ``embed_codes``); ``input_ids`` are only in-vocab
+        bookkeeping placeholders. Return a zero embedding of the temporal hidden
+        size — the runner replaces it with the real per-frame ``inputs_embeds``.
+        """
+        return torch.zeros(
+            (input_ids.shape[0], self.mtp_hidden_size),
+            device=input_ids.device,
+            dtype=self._dtype,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **_: Any,
+    ) -> torch.Tensor | IntermediateTensors:
+        return self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor | OmniOutput,
+        sampling_metadata: Any = None,
+    ) -> torch.Tensor | None:
+        if isinstance(hidden_states, OmniOutput):
+            hidden_states = hidden_states.text_hidden_states
+        if hidden_states is None:
+            return None
+        if hidden_states.dim() == 3:
+            b, s, h = hidden_states.shape
+            logits = self.logits_processor(self.lm_head, hidden_states.reshape(b * s, h))
+            return None if logits is None else logits.reshape(b, s, -1)
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+        hidden = model_outputs
+        info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or []
+        audio_codes_list: list[torch.Tensor] = []
+        for info in info_dicts:
+            if not isinstance(info, dict):
+                continue
+            ac = info.get("codes", {}).get("audio")
+            if isinstance(ac, torch.Tensor):
+                audio_codes_list.append(ac)
+        if not audio_codes_list:
+            return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
+        audio_codes = torch.cat(audio_codes_list, dim=0)
+        hidden = hidden[: int(audio_codes.shape[0])]
+        return OmniOutput(text_hidden_states=hidden, multimodal_outputs={"codes": {"audio": audio_codes}})
+
+    # ------------------------------------------------------------------
+    # Per-frame omni AR protocol: preprocess (stateful) + talker_mtp (batched)
+    # ------------------------------------------------------------------
+    def _initial_frame_embed(self, device: torch.device) -> torch.Tensor:
+        """``embed_codes`` of Moshi's initial token frame (text+audio SOS)."""
+        text_init = self.config.text_vocab_size  # text_initial_token_id
+        audio_init = self.config.audio_vocab_size  # card == initial audio token
+        n_q = self.config.num_audio_codebooks
+        stack = torch.empty((1, 1 + n_q, 1), dtype=torch.long, device=device)
+        stack[:, 0] = text_init
+        stack[:, 1:] = audio_init
+        return self.input_embeddings(stack).reshape(1, -1)  # [1, hidden]
+
+    def _build_frame_embed(
+        self,
+        text_token: torch.Tensor,
+        last_agent: torch.Tensor | None,
+        prev_agent: torch.Tensor | None,
+        device: torch.device,
+        user_d0: torch.Tensor | None = None,
+        user_d1: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Full ``inputs_embeds`` for a decode frame (Moshi cache read at offset-1).
+
+        Acoustic-delay decomposition (delays ``[0, 0,1x7, 0,1x7]``): after a
+        depformer step, Moshi writes every unforced agent codebook into the
+        current cache target. The next temporal input therefore reads the whole
+        previous effective agent frame. User writes are delayed by codebook, so
+        user cb0 and cb1..7 still come from ``user_d0`` and ``user_d1``.
+        Unfilled positions default to the initial token (card).
+        """
+        n_q = self.config.num_audio_codebooks  # 16 (8 agent + 8 user)
+        n_user = n_q // 2  # 8
+        audio_init = self.config.audio_vocab_size  # card == initial audio token (Moshi cache default)
+        # No-history / not-yet-generated delayed positions = the initial token (card),
+        # NOT 0 (verified against Moshi's per-frame input stack).
+        stack = torch.full((1, 1 + n_q, 1), audio_init, dtype=torch.long, device=device)
+        stack[:, 0] = text_token.reshape(1)
+        # The full effective agent frame was stored at the previous cache target
+        # after applying any teacher forcing. The next cache read sees all eight
+        # codebooks from that frame.
+        if last_agent is not None and last_agent.numel() >= n_user:
+            stack[0, 1 : 1 + n_user, 0] = last_agent.reshape(-1)[:n_user].to(device)
+        # user cb0 (row 9) delay-0; user cb1..7 (rows 10..16) delay-1.
+        user_base = 1 + n_user  # row index 9
+        if user_d0 is not None and user_d0.numel() >= 1:
+            stack[0, user_base, 0] = user_d0.reshape(-1)[0].to(device)
+        if user_d1 is not None and user_d1.numel() >= n_user:
+            stack[0, user_base + 1 : user_base + n_user, 0] = user_d1.reshape(-1)[1:n_user].to(device)
+        return self.input_embeddings(stack).reshape(1, -1)  # [1, hidden]
+
+    def _build_prefill_embed(
+        self,
+        prefill_text: torch.Tensor,
+        offset: int,
+        span: int,
+        device: torch.device,
+        silence: torch.Tensor | None = None,
+        user_sine: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Per-frame embeds for the persona system-prompt prefill.
+
+        Each frame forces the persona/silence text token (row 0). The agent rows
+        use Mimi-encoded silence, while the user rows optionally use the encoded
+        sine frame from PersonaPlex's native text-prompt prefill. Without an
+        explicit user frame, both sides fall back to silence.
+        """
+        n_q = self.config.num_audio_codebooks
+        n_user = n_q // 2
+        zero_text = 3  # Moshi LMGen.zero_text_code (silence/pad text)
+        sil = (
+            silence.reshape(-1).to(device) if isinstance(silence, torch.Tensor) and silence.numel() >= n_user else None
+        )
+        user = (
+            user_sine.reshape(-1).to(device)
+            if isinstance(user_sine, torch.Tensor) and user_sine.numel() >= n_user
+            else sil
+        )
+        total = int(prefill_text.numel())
+        rows = []
+        for i in range(span):
+            pos = offset + i
+            text_tok = int(prefill_text[pos].item()) if pos < total else zero_text
+            stack = torch.zeros((1, 1 + n_q, 1), dtype=torch.long, device=device)
+            stack[:, 0] = text_tok
+            if sil is not None:
+                stack[0, 1 : 1 + n_user, 0] = sil[:n_user]  # agent rows = encoded silence
+            if user is not None:
+                stack[0, 1 + n_user : 1 + 2 * n_user, 0] = user[:n_user]
+            rows.append(self.input_embeddings(stack).reshape(1, -1))
+        return torch.cat(rows, dim=0)  # [span, hidden]
+
+    @staticmethod
+    def _user_frame(info: dict[str, Any], frame_idx: int) -> torch.Tensor | None:
+        """Fetch the Mimi-encoded user codes for ``frame_idx`` (or None / out of range)."""
+        uc = info.get("pplex_user_codes")
+        if uc is None:
+            return None
+        if not isinstance(uc, torch.Tensor):
+            uc = torch.as_tensor(uc, dtype=torch.long)
+        if uc.ndim != 2 or frame_idx < 0 or frame_idx >= uc.shape[0]:
+            return None
+        return uc[frame_idx]
+
+    @staticmethod
+    def _last_agent_codes(info: dict[str, Any]) -> torch.Tensor | None:
+        codes = info.get("codes")
+        if isinstance(codes, dict):
+            ac = codes.get("audio")
+            if isinstance(ac, torch.Tensor) and ac.numel() > 0:
+                return (ac[-1] if ac.ndim == 2 else ac).reshape(-1)
+        return None
+
+    def preprocess(
+        self,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor | None,
+        **info_dict: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
+        """Build the current frame's ``inputs_embeds`` and the mtp inputs.
+
+        Prefill emits the initial-token frame embedding; each decode step builds the
+        delayed base frame (everything but agent cb0, which ``talker_mtp`` adds).
+        """
+        additional = info_dict.get("additional_information")
+        if isinstance(additional, dict):
+            merged = {k: v for k, v in info_dict.items() if k != "additional_information"}
+            for k, v in additional.items():
+                merged.setdefault(k, v)
+            info_dict = merged
+        meta = info_dict.get("meta", {}) if isinstance(info_dict.get("meta"), dict) else {}
+
+        device = input_ids.device
+        span = int(input_ids.shape[0])
+        is_prefill_raw = info_dict.get("_omni_is_prefill")
+        if isinstance(is_prefill_raw, bool):
+            is_prefill = is_prefill_raw
+        else:
+            is_prefill = span > 1
+
+        duplex = info_dict.get("duplex")
+        if isinstance(duplex, dict) and duplex.get("data_plane") is True and is_prefill:
+            prompt_len_raw = info_dict.get("duplex_prompt_len", span)
+            try:
+                prompt_len = int(prompt_len_raw)
+            except (TypeError, ValueError):
+                prompt_len = span
+            prepared = self._duplex_stage0_runtime().prepare_append(
+                duplex,
+                prompt_len=prompt_len,
+                request_id=(str(info_dict["request_id"]) if isinstance(info_dict.get("request_id"), str) else None),
+            )
+            offset_raw = info_dict.get("duplex_token_offset", 0)
+            try:
+                offset = max(0, int(offset_raw))
+            except (TypeError, ValueError):
+                offset = 0
+            local_offset = offset - prepared.prompt_offset
+            if local_offset < 0:
+                raise ValueError(
+                    "PersonaPlex scheduled span precedes the current append: "
+                    f"offset={offset}, append_offset={prepared.prompt_offset}, "
+                    f"span={span}, prompt={prompt_len}"
+                )
+            req_embeds = prepared.inputs_embeds[local_offset : local_offset + span].to(
+                device=device,
+                dtype=self._dtype,
+            )
+            req_input_ids = prepared.input_ids[local_offset : local_offset + span].to(device=device)
+            if req_embeds.shape[0] != span or req_input_ids.shape[0] != span:
+                raise ValueError(
+                    "PersonaPlex duplex prompt slice is shorter than the scheduled span: "
+                    f"offset={offset}, span={span}, prompt={prompt_len}"
+                )
+            return req_input_ids, req_embeds, prepared.info_update
+
+        zero_hidden = torch.zeros((1, self.mtp_hidden_size), device=device, dtype=self._dtype)
+
+        prefill_text = info_dict.get("pplex_prefill_text")
+        prefill_len = int(prefill_text.numel()) if isinstance(prefill_text, torch.Tensor) else 0
+
+        if is_prefill:
+            offset = max(0, int(info_dict.get("_omni_num_computed_tokens", 0) or 0))
+            if isinstance(prefill_text, torch.Tensor) and prefill_text.numel() > 0:
+                # Persona system-prompt prefill (step_system_prompts analog).
+                silence = info_dict.get("pplex_silence_codes")
+                emb = self._build_prefill_embed(prefill_text, offset, span, device, silence)
+            else:
+                emb = self._initial_frame_embed(device)
+                if span > 1:
+                    emb = emb.expand(span, -1).contiguous()
+            ids_out = input_ids.clone()
+            info_update = {
+                "meta": {
+                    "pplex_frame": int(meta.get("pplex_frame", 0)) + span,
+                    "pplex_prefill_len": prefill_len,
+                },
+                "mtp_inputs": (zero_hidden, zero_hidden),
+            }
+            return ids_out, emb, info_update
+
+        # Decode: input_ids == the text token sampled last step (delay-0 text).
+        text_token = input_ids.reshape(-1)[:1]
+        # Agent acoustic delay: cb0 (delay 0) = gen[t-1] = the last stored codes;
+        # cb1..7 (delay 1) = gen[t-2] = the codes carried from the previous frame.
+        last_agent = self._last_agent_codes(info_dict)
+        prev_agent = info_dict.get("pplex_prev_agent")
+        if prev_agent is not None and not isinstance(prev_agent, torch.Tensor):
+            prev_agent = torch.as_tensor(prev_agent, dtype=torch.long)
+        # Real user stream (Phase-1 turn-based): user cb0 delay-0, cb1..7 delay-1.
+        # Index into the user stream relative to the first decode frame (after the
+        # persona prefill), so the user audio aligns with the agent's response.
+        prefill_len = int(meta.get("pplex_prefill_len", prefill_len) or 0)
+        # decode_frame = pplex_frame - prefill_len (the prefill frame already advanced
+        # pplex_frame, so this is d+1 at decode step d). Verified vs Moshi's per-frame
+        # stack: user cb0 (delay 0) = enc[decode_frame-1] = enc[d]; cb1..7 (delay 1) =
+        # enc[decode_frame-2]. (The agent stream lags one more: agent cb0 = gen[d-1].)
+        decode_frame = max(0, int(meta.get("pplex_frame", 0)) - prefill_len)
+        user_d0 = self._user_frame(info_dict, decode_frame - 1)
+        user_d1 = self._user_frame(info_dict, decode_frame - 2)
+        base = self._build_frame_embed(text_token, last_agent, prev_agent, device, user_d0, user_d1)
+        text_step = self.input_embeddings.text_emb(text_token.reshape(1, 1)).reshape(1, -1)
+        # The previous frame's temporal hidden (written by postprocess into
+        # hidden_states["last"]) conditions this frame's depformer. Falls back to
+        # zeros only on the very first decode step (no prior hidden yet).
+        hs = info_dict.get("hidden_states", {}) if isinstance(info_dict.get("hidden_states"), dict) else {}
+        last_hidden = hs.get("last")
+        if isinstance(last_hidden, torch.Tensor) and last_hidden.numel() > 0:
+            last_hidden = last_hidden.reshape(1, -1).to(device=device, dtype=self._dtype)
+        else:
+            last_hidden = zero_hidden
+        info_update = {
+            "meta": {"pplex_frame": int(meta.get("pplex_frame", 0)) + 1},
+            "mtp_inputs": (last_hidden, text_step),
+        }
+        # Carry this frame's gen[t-1] forward; next frame reads it as gen[t-2] (cb1..7).
+        if last_agent is not None:
+            info_update["pplex_prev_agent"] = last_agent.detach().to(torch.long).cpu()
+        return input_ids, base, info_update
+
+    def _duplex_stage0_runtime(self):
+        runtime = getattr(self, "_personaplex_duplex_stage0_runtime", None)
+        if runtime is not None:
+            return runtime
+        from vllm_omni.experimental.fullduplex.personaplex.stage0 import (
+            PersonaPlexStage0DuplexRuntime,
+        )
+
+        model_path = str(getattr(self.vllm_config.model_config, "model", ""))
+        device = str(next(self.parameters()).device)
+        runtime = PersonaPlexStage0DuplexRuntime(
+            self,
+            model_path=model_path,
+            device=device,
+            max_sessions=int(getattr(self.vllm_config.model_config, "duplex_max_sessions", 1)),
+        )
+        self._personaplex_duplex_stage0_runtime = runtime
+        return runtime
+
+    def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
+        runtime = getattr(self, "_personaplex_duplex_stage0_runtime", None)
+        if runtime is None:
+            return
+        for request_id in finished_req_ids:
+            runtime.close_request(request_id)
+
+    def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]:
+        """Capture this frame's last hidden for the next step's depformer (mtp)."""
+        if hidden_states is None or hidden_states.numel() == 0:
+            return {}
+        return {"hidden_states": {"last": hidden_states[-1, :].detach()}}
+
+    def talker_mtp(
+        self,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        text_step: torch.Tensor,
+        **kwargs: Any,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Run the depformer for the frame and finish its ``inputs_embeds``.
+
+        Returns ``(inputs_embeds, audio_codes[B, dep_q])``. ``audio_codes`` are
+        stored under ``("codes","audio")``; ``inputs_embeds`` (base from preprocess
+        + the fresh agent cb0 embedding, the delay-0 acoustic token) feeds the next
+        temporal forward.
+        """
+        bsz = int(input_ids.shape[0])
+        dtype = self._dtype
+        text_token = input_ids.reshape(bsz).to(torch.long)
+        hidden = last_talker_hidden.reshape(bsz, 1, -1).to(dtype)
+
+        codes = self.depformer(text_token, hidden)  # [B, dep_q] == gen[t]
+
+        # The frame's inputs_embeds is built fully in preprocess (Moshi cache read at
+        # offset-1, with agent cb0 = gen[t-1]); gen[t] feeds the NEXT frame, so pass
+        # the embed through unchanged and just emit this frame's codes.
+        inputs_embeds = input_embeds.reshape(bsz, -1).to(dtype)
+        return inputs_embeds, codes.to(torch.long)
+
+    def post_sample_talker_mtp(
+        self,
+        *,
+        input_ids: torch.Tensor,
+        hidden_states: torch.Tensor,
+        req_ids: list[str],
+        req_infos: list[dict[str, Any]],
+    ) -> torch.Tensor:
+        """Generate depformer codes for a one-token resumable duplex segment.
+
+        The normal runner invokes ``talker_mtp`` at the start of the next decode
+        step. PersonaPlex's unified duplex request instead appends one audio
+        frame and stops after one sampled text token, so there is no next decode
+        step. Run the same depformer dependency immediately from the current
+        sampled text token and temporal hidden state.
+        """
+        bsz = int(input_ids.shape[0])
+        if len(req_infos) != bsz:
+            raise ValueError(
+                f"PersonaPlex depformer request information does not match batch: {len(req_infos)} != {bsz}"
+            )
+        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:
+            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):
+                raise ValueError("PersonaPlex duplex depformer teacher-forcing state is missing")
+            tokens = tokens.reshape(-1)
+            provided = provided.reshape(-1)
+            if tokens.shape != provided.shape:
+                raise ValueError(
+                    "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,
+            ),
+        ).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
+
+    # ------------------------------------------------------------------
+    # Weight loading
+    # ------------------------------------------------------------------
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Route the Moshi checkpoint into the three components.
+
+        * ``transformer.*`` / ``out_norm.alpha`` -> the Helium temporal backbone
+          (same q/k-split + gate/up-split + alpha-squeeze map as HeliumForCausalLM).
+        * ``text_linear.weight`` -> ``lm_head``.
+        * ``emb.*`` / ``text_emb.weight`` -> input embeddings.
+        * ``depformer*`` / ``linears.*`` -> depformer.
+        """
+        weights = list(weights)
+        params = dict(self.named_parameters(remove_duplicate=False))
+        loaded: set[str] = set()
+
+        emb_w: list[tuple[str, torch.Tensor]] = []
+        dep_w: list[tuple[str, torch.Tensor]] = []
+        for name, w in weights:
+            if name.startswith(("emb.", "text_emb.")):
+                emb_w.append((name, w))
+            elif name.startswith(("depformer.", "depformer_in.", "depformer_emb.", "depformer_text_emb", "linears.")):
+                dep_w.append((name, w))
+            elif name == "text_linear.weight":
+                self._load_direct("lm_head.weight", w, params, loaded)
+            else:
+                loaded |= self._load_temporal(name, w, params)
+
+        # Delegate to the verified component loaders, prefixed into this module.
+        for sub, sub_w in (("input_embeddings", emb_w), ("depformer", dep_w)):
+            module = getattr(self, sub)
+            for tgt in module.load_weights(sub_w):
+                loaded.add(f"{sub}.{tgt}")
+        return loaded
+
+    def _load_temporal(
+        self,
+        name: str,
+        loaded_weight: torch.Tensor,
+        params: dict[str, nn.Parameter],
+    ) -> set[str]:
+        out: set[str] = set()
+        if name == "out_norm.alpha":
+            self._load_direct("model.norm.weight", loaded_weight.squeeze(), params, out)
+            return out
+        prefix = "transformer.layers."
+        if not name.startswith(prefix):
+            return out
+        rest = name.removeprefix(prefix)
+        layer_index, _, suffix = rest.partition(".")
+        if not layer_index.isdigit() or not suffix:
+            return out
+        base = f"model.layers.{layer_index}"
+        if suffix == "self_attn.in_proj_weight":
+            q, k, v = loaded_weight.chunk(3, dim=0)
+            pname = f"{base}.self_attn.qkv_proj.weight"
+            self._load_shard(pname, q, "q", params)
+            self._load_shard(pname, k, "k", params)
+            self._load_shard(pname, v, "v", params)
+            out.add(pname)
+        elif suffix == "self_attn.out_proj.weight":
+            self._load_direct(f"{base}.self_attn.o_proj.weight", loaded_weight, params, out)
+        elif suffix == "gating.linear_in.weight":
+            gate, up = loaded_weight.chunk(2, dim=0)
+            pname = f"{base}.mlp.gate_up_proj.weight"
+            self._load_shard(pname, gate, 0, params)
+            self._load_shard(pname, up, 1, params)
+            out.add(pname)
+        elif suffix == "gating.linear_out.weight":
+            self._load_direct(f"{base}.mlp.down_proj.weight", loaded_weight, params, out)
+        elif suffix == "norm1.alpha":
+            self._load_direct(f"{base}.input_layernorm.weight", loaded_weight.squeeze(), params, out)
+        elif suffix == "norm2.alpha":
+            self._load_direct(f"{base}.post_attention_layernorm.weight", loaded_weight.squeeze(), params, out)
+        return out
+
+    @staticmethod
+    def _load_direct(
+        name: str,
+        loaded_weight: torch.Tensor,
+        params: dict[str, nn.Parameter],
+        loaded: set[str],
+    ) -> None:
+        if name not in params:
+            return
+        param = params[name]
+        weight_loader = getattr(param, "weight_loader", default_weight_loader)
+        weight_loader(param, loaded_weight)
+        loaded.add(name)
+
+    @staticmethod
+    def _load_shard(
+        name: str,
+        loaded_weight: torch.Tensor,
+        shard_id: str | int,
+        params: dict[str, nn.Parameter],
+    ) -> None:
+        if name not in params:
+            return
+        param = params[name]
+        param.weight_loader(param, loaded_weight, shard_id)

vllm_omni/model_executor/models/personaplex/personaplex_temporal.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Streaming Helium temporal transformer for PersonaPlex (moshi-free, plain torch).
+
+Frame-clocked duplex needs a stateful per-frame forward: one 80 ms step consumes
+one summed frame embedding and returns the temporal hidden + text logits, with a
+sliding-window KV carried across the conversation. vLLM's engine path serves the
+offline/staged flavor of this model (``modeling_helium.HeliumForCausalLM``); this
+module is the REAL-TIME flavor: no engine context, a fixed-capacity ring KV, and
+static shapes so the whole step is CUDA-graphable.
+
+Semantics mirror the Moshi streaming transformer exactly (MIT; the reference for
+the ``nvidia/personaplex-7b-v1`` checkpoint), op for op, so greedy decoding is
+bit-comparable against recorded reference outputs:
+
+- fp32 RMSNorm with ``alpha`` weights (norm1/norm2/out_norm)
+- fused in-proj attention, interleaved (non-neox) RoPE applied at the ABSOLUTE
+  offset in fp32, relative-position causal mask over a ring KV with
+  ``context``-window truncation
+- SiLU gating MLP with fused ``linear_in`` (gate/up) and ``linear_out``
+- ``transformer_out`` is the POST-``out_norm`` hidden; ``text_logits`` is
+  ``text_linear`` of that hidden
+
+Weights load directly from the moshi checkpoint's fused layout
+(``transformer.layers.{i}.self_attn.in_proj_weight`` etc.), keeping the exact
+computation order of the reference.
+
+Per-slot elastic recycle (concurrent serving) uses the same per-row
+``start_offset`` masking verified for the moshi-hosted path: writes stay uniform
+because every slot ticks in lockstep; ``reset_slot`` only moves that row's
+valid-window start. The LM-specific "+1 sacrifice tick" bump is the CALLER's
+job (see the fullduplex PersonaPlex stage0), matching the split introduced when
+the +1 was found to be wrong for codec streams.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.personaplex.personaplex_depformer import _rms_norm_f32
+
+
+def _apply_rope(q: torch.Tensor, k: torch.Tensor, offset: torch.Tensor, max_period: float = 10_000.0):
+    """Interleaved RoPE at absolute ``offset``, fp32 rotation (moshi ``apply_rope``)."""
+    B, H, T, D = q.shape
+    ds = torch.arange(D // 2, device=q.device, dtype=torch.float32)
+    freqs = torch.exp(ds * (-math.log(max_period) * 2 / D))
+    ts = offset.float() + torch.arange(T, device=q.device, dtype=torch.float32)
+    ts = ts.view(1, -1, 1)
+
+    dims = q.shape[:-1]
+    q = q.view(*dims, D // 2, 2)
+    k = k.view(*dims, D // 2, 2)
+    qr, qi = q[..., 0].float(), q[..., 1].float()
+    kr, ki = k[..., 0].float(), k[..., 1].float()
+    rotr = torch.cos(freqs * ts)
+    roti = torch.sin(freqs * ts)
+    qor = qr * rotr - qi * roti
+    qoi = qr * roti + qi * rotr
+    kor = kr * rotr - ki * roti
+    koi = kr * roti + ki * rotr
+    dtype = q.dtype
+    qo = torch.stack([qor.to(dtype), qoi.to(dtype)], dim=-1)
+    ko = torch.stack([kor.to(dtype), koi.to(dtype)], dim=-1)
+    return qo.view(*dims, D), ko.view(*dims, D)
+
+
+class _RingKV:
+    """Fixed-capacity KV ring with per-row valid-window start (elastic recycle).
+
+    Uniform writes (all rows tick together) + a per-row mask via absolute
+    positions; with ``start_offset == 0`` everywhere this is the plain streaming
+    cache. All shapes static -> CUDA-graph safe.
+    """
+
+    def __init__(self, batch_size: int, num_heads: int, dim_per_head: int, capacity: int, device, dtype):
+        self.capacity = capacity
+        self.cache = torch.zeros((2, batch_size, num_heads, capacity, dim_per_head), device=device, dtype=dtype)
+        self.end_offset = torch.zeros(1, device=device, dtype=torch.long)
+        self.start_offset = torch.zeros(batch_size, device=device, dtype=torch.long)
+
+    def reset(self) -> None:
+        self.end_offset.zero_()
+        self.start_offset.zero_()
+
+    def reset_slot(self, b: int) -> None:
+        # Mask everything written so far for row b; the row's next write is its
+        # first visible entry. (LM sacrifice-tick +1 is applied by the caller.)
+        self.start_offset[b] = self.end_offset.clone()
+
+    def bump_slot_start(self, b: int) -> None:
+        self.start_offset[b] += 1
+
+    def complete(self, k: torch.Tensor, v: torch.Tensor):
+        B, H, T, D = k.shape
+        indexes = torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype) + self.end_offset
+        indexes = indexes % self.capacity
+        self.cache[0].index_copy_(2, indexes, k)
+        self.cache[1].index_copy_(2, indexes, v)
+        self.end_offset.add_(T)
+
+        idx = torch.arange(self.capacity, device=self.end_offset.device, dtype=torch.long)
+        invalid = idx >= self.end_offset
+        end_index = self.end_offset % self.capacity
+        delta = idx - end_index
+        # `delta <= 0` (not `< 0`) is moshi's exact convention (transformer.py
+        # RingKVCache.complete). It labels the just-past-newest slot as the future
+        # write position, so once the ring has wrapped the single oldest in-window
+        # cell is excluded and the effective window is capacity-1. This is inherited
+        # verbatim from the reference and only shows after the window fills (Helium
+        # ~3000 frames / 240 s); it costs one frame out of thousands. Do NOT change
+        # this to `< 0`: it would diverge from moshi and break greedy bit-parity.
+        positions = torch.where(delta <= 0, self.end_offset + delta, self.end_offset + delta - self.capacity)
+        positions = torch.where(invalid, torch.full_like(positions, -1), positions)
+        positions = positions.view(1, -1)  # [1, capacity]
+        below = positions < self.start_offset.view(-1, 1)  # [B, capacity]
+        positions = torch.where(below, torch.full_like(positions, -1), positions)
+        return self.cache[0], self.cache[1], positions
+
+
+class _TemporalLayer(nn.Module):
+    """One Helium decoder layer in the checkpoint's fused layout."""
+
+    def __init__(self, dim: int, num_heads: int, hidden: int) -> None:
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.in_proj_weight = nn.Parameter(torch.empty(3 * dim, dim))
+        self.out_proj_weight = nn.Parameter(torch.empty(dim, dim))
+        self.gating_in = nn.Parameter(torch.empty(2 * hidden, dim))
+        self.gating_out = nn.Parameter(torch.empty(dim, hidden))
+        self.norm1_alpha = nn.Parameter(torch.ones(1, 1, dim))
+        self.norm2_alpha = nn.Parameter(torch.ones(1, 1, dim))
+
+    def forward(self, x: torch.Tensor, kv: _RingKV, offset: torch.Tensor, context: int) -> torch.Tensor:
+        B, T, _ = x.shape
+        h = _rms_norm_f32(x, self.norm1_alpha, 1e-8)
+        qkv = F.linear(h, self.in_proj_weight)
+        # Layout mirrors moshi's einops "b t (p h d) -> p b h t d" exactly, so the
+        # downstream kernels see identical strides (bit-level replay agreement).
+        qkv = qkv.view(B, T, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        q, k = _apply_rope(q, k, offset)
+
+        keys, values, pos_k = kv.complete(k, v)
+        pos_k = pos_k.view(pos_k.shape[0], 1, pos_k.shape[1])  # [B, 1, cap]
+        pos_q = offset + torch.arange(T, device=q.device, dtype=torch.long).view(1, -1, 1)
+        delta = pos_q - pos_k
+        attn_bias = (pos_k >= 0) & (delta >= 0) & (delta < context)
+        attn_bias = attn_bias.unsqueeze(1)  # [B, 1, T, cap]
+        attn = F.scaled_dot_product_attention(q, keys, values, attn_bias, dropout_p=0.0)
+        attn = attn.transpose(1, 2).reshape(B, T, self.dim)
+        x = x + F.linear(attn, self.out_proj_weight)
+
+        h = _rms_norm_f32(x, self.norm2_alpha, 1e-8)
+        a, b = F.linear(h, self.gating_in).chunk(2, dim=-1)
+        return x + F.linear(F.silu(a) * b, self.gating_out)
+
+
+class PersonaPlexTemporalStreaming(nn.Module):
+    """The 7B Helium temporal backbone as a stateful per-frame stepper.
+
+    ``step(frame_embedding [B, 1, dim]) -> (transformer_out [B, 1, dim],
+    text_logits [B, 1, 1, text_card])``, matching moshi's
+    ``LMModel.forward_embeddings`` contract so the existing native depformer and
+    input embeddings compose unchanged.
+    """
+
+    def __init__(
+        self,
+        dim: int = 4096,
+        num_layers: int = 32,
+        num_heads: int = 32,
+        hidden: int = 11264,
+        context: int = 3000,
+        text_card: int = 32000,
+        max_period: float = 10_000.0,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.context = context
+        self.max_period = max_period
+        self.layers = nn.ModuleList([_TemporalLayer(dim, num_heads, hidden) for _ in range(num_layers)])
+        self.out_norm_alpha = nn.Parameter(torch.ones(1, 1, dim))
+        self.text_linear = nn.Parameter(torch.empty(text_card, dim))
+        self._kv: list[_RingKV] | None = None
+        self._offset: torch.Tensor | None = None
+
+    # -- streaming state ----------------------------------------------------
+
+    def streaming_init(self, batch_size: int) -> None:
+        p = next(self.parameters())
+        heads = self.layers[0].num_heads
+        head_dim = self.layers[0].head_dim
+        # Capacity = context window (the mask truncates at `context` anyway).
+        self._kv = [_RingKV(batch_size, heads, head_dim, self.context, p.device, p.dtype) for _ in self.layers]
+        self._offset = torch.zeros(1, device=p.device, dtype=torch.long)
+
+    def reset_streaming(self) -> None:
+        assert self._kv is not None, "call streaming_init first"
+        for kv in self._kv:
+            kv.reset()
+        self._offset.zero_()
+
+    def reset_slot(self, b: int) -> None:
+        assert self._kv is not None
+        for kv in self._kv:
+            kv.reset_slot(b)
+
+    def bump_slot_start(self, b: int) -> None:
+        """LM sacrifice-tick semantics: mask the one spurious post-recycle write."""
+        assert self._kv is not None
+        for kv in self._kv:
+            kv.bump_slot_start(b)
+
+    # -- per-frame step -------------------------------------------------------
+
+    @torch.no_grad()
+    def step(self, frame_embedding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        assert self._kv is not None, "call streaming_init first"
+        x = frame_embedding
+        for layer, kv in zip(self.layers, self._kv):
+            x = layer(x, kv, self._offset, self.context)
+        self._offset.add_(x.shape[1])
+        out = _rms_norm_f32(x, self.out_norm_alpha, 1e-8)
+        text_logits = F.linear(out, self.text_linear)
+        return out, text_logits[:, None]
+
+    # -- weights --------------------------------------------------------------
+
+    def load_weights(self, state_dict: dict[str, torch.Tensor]) -> int:
+        """Load from the moshi checkpoint layout (fused, exact tensor reuse)."""
+        loaded = 0
+        with torch.no_grad():
+            for i, layer in enumerate(self.layers):
+                base = f"transformer.layers.{i}"
+                pairs = [
+                    (layer.in_proj_weight, f"{base}.self_attn.in_proj_weight"),
+                    (layer.out_proj_weight, f"{base}.self_attn.out_proj.weight"),
+                    (layer.gating_in, f"{base}.gating.linear_in.weight"),
+                    (layer.gating_out, f"{base}.gating.linear_out.weight"),
+                    (layer.norm1_alpha, f"{base}.norm1.alpha"),
+                    (layer.norm2_alpha, f"{base}.norm2.alpha"),
+                ]
+                for param, name in pairs:
+                    src = state_dict[name]
+                    param.data.copy_(src.reshape(param.shape).to(param.dtype))
+                    loaded += 1
+            self.out_norm_alpha.data.copy_(state_dict["out_norm.alpha"].reshape(self.out_norm_alpha.shape))
+            self.text_linear.data.copy_(state_dict["text_linear.weight"].to(self.text_linear.dtype))
+            loaded += 2
+        return loaded

vllm_omni/model_executor/models/personaplex/pipeline.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""PersonaPlex pipeline: Talker (AR decode -> Mimi codebooks) -> Code2Wav (codebooks -> 24 kHz PCM).
+
+PersonaPlex is a Moshi finetune (full-duplex speech-to-speech). For offline/batch
+runs it is served as a 2-stage vllm-omni audio->audio pipeline, reusing the
+Qwen3-TTS staged topology:
+
+* Stage 0 (``personaplex``) is the AR talker: the Helium temporal transformer
+  plus the depformer (both built by the lead). It emits the per-frame audio
+  codebooks as latents.
+* Stage 1 (``code2wav``) wraps the external Mimi codec and decodes the active
+  codebooks (``cb 0..7``) into PCM.
+
+The inter-stage input processors named below live in
+``vllm_omni.model_executor.stage_input_processors.personaplex`` and are part of
+the talker<->code2wav seam; they are built by the lead alongside the talker.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+# Talker<->Code2Wav input processors (built by lead; coupled to talker output).
+_PROC = "vllm_omni.model_executor.stage_input_processors.personaplex"
+
+PERSONAPLEX_PIPELINE = PipelineConfig(
+    model_type="personaplex",
+    # Pipeline-level default; the code2wav stage overrides per-stage below.
+    model_arch="PersonaPlexTalkerForConditionalGeneration",
+    default_deploy_config_name="personaplex.yaml",
+    duplex_runtime_extension=(
+        "vllm_omni.experimental.fullduplex.personaplex.runtime_extension.PersonaPlexDuplexRuntimeExtension"
+    ),
+    duplex_serving_adapter=(
+        "vllm_omni.experimental.fullduplex.personaplex.serving_adapter.PersonaPlexServingRuntimeAdapter"
+    ),
+    duplex_control_enabled=True,
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="personaplex",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            # built by lead: talker emits a 17-row Mimi token stack per frame;
+            # these processors keep the PCM-bearing rows (cb 0..7) and flatten
+            # them codebook-major for the Code2Wav stage.
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
+            custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload",
+            sampling_constraints={
+                "detokenize": False,
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="code2wav",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="PersonaPlexCode2Wav",
+            # In sync (non-async-chunk) mode the only pre-stage input proc is a
+            # length-only placeholder; the bulk codec payload ships via the
+            # worker connector from stage 0's full-payload producer.
+            # built by lead.
+            sync_process_input_func=f"{_PROC}.talker2code2wav_token_only",
+            sampling_constraints={"detokenize": True},
+        ),
+    ),
+)

vllm_omni/model_executor/models/registry.py

@@ -164,6 +164,18 @@ _OMNI_MODELS = {
         "qwen3_tts_code2wav",
         "Qwen3TTSCode2Wav",
     ),
+    ## personaplex (Moshi finetune, full-duplex S2S; 2-stage audio->audio offline pipeline)
+    # built by lead: personaplex_talker module + PersonaPlexTalkerForConditionalGeneration.
+    "PersonaPlexTalkerForConditionalGeneration": (
+        "personaplex",
+        "personaplex_talker",
+        "PersonaPlexTalkerForConditionalGeneration",
+    ),
+    "PersonaPlexCode2Wav": (
+        "personaplex",
+        "personaplex_code2wav",
+        "PersonaPlexCode2Wav",
+    ),
     ## higgs-audio v2
     "HiggsAudioV2ForConditionalGeneration": (
         "higgs_audio_v2",

vllm_omni/model_executor/stage_input_processors/personaplex.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Talker -> Code2Wav input processors for PersonaPlex.
+
+The talker (stage 0) emits, per frame, the ``dep_q`` depformer audio codes under
+``("codes","audio")``. Only the leading ``num_active_codebooks`` (the agent's
+``cb 0..7``) are decoded to PCM by Mimi; the trailing codebooks are the user
+stream and are not vocoded. These processors take the accumulated per-frame agent
+codes ``[F, dep_q]``, keep ``cb 0..7``, and flatten them codebook-major
+(``[8 * F]``) — the exact layout :class:`PersonaPlexCode2Wav` consumes.
+
+Mirrors the Qwen3-TTS processors (sync ``full_payload`` + ``token_only``; an
+async-chunk variant for the streaming path), but with PersonaPlex's agent-codebook
+slice instead of Qwen3-TTS's residual layout.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from vllm_omni.data_entry_keys import (
+    CodesStruct,
+    MetaStruct,
+    OmniPayloadStruct,
+)
+
+_NUM_ACTIVE_CODEBOOKS = 8  # agent cb 0..7 (the PCM-bearing rows)
+
+
+def _empty_finished_payload() -> OmniPayloadStruct:
+    return OmniPayloadStruct(
+        codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+        meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+    )
+
+
+def _agent_codes_to_codebook_major(audio: torch.Tensor) -> torch.Tensor:
+    """``[F, dep_q]`` raw depformer agent codes -> de-delayed flat codebook-major.
+
+    The talker emits the raw per-frame depformer codes ``gen[t]`` (cb 0..7). Mimi
+    needs them ACOUSTICALLY DE-DELAYED to a common time step (Moshi agent delays
+    ``[0, 1, 1, 1, 1, 1, 1, 1]``): acoustic frame t's cb_k is predicted at step
+    ``t + delay[k]``, so output frame t = ``[gen[t][0], gen[t+1][1:8]]`` (cb0 from
+    the current step, cb1..7 from the NEXT step, since the delayed codebooks lag).
+    Without this the codebooks are misaligned by one frame and Mimi decodes garble.
+    The last frame has no successor for cb1..7, so it is dropped (the delay warmup).
+    """
+    if audio.ndim != 2 or audio.shape[0] < 2:
+        return torch.empty(0, dtype=torch.long)
+    audio = audio.to(torch.long)
+    k = min(_NUM_ACTIVE_CODEBOOKS, int(audio.shape[1]))
+    agent = audio[:, :k]
+    valid = (agent >= 0).all(dim=1)
+    agent = agent[valid]
+    if agent.shape[0] < 2:
+        return torch.empty(0, dtype=torch.long)
+    # De-delay: cb0 from frame t, cb1..7 from frame t+1 (drop the last frame).
+    cb0 = agent[:-1, 0:1]  # [F-1, 1]
+    cb_rest = agent[1:, 1:k]  # [F-1, k-1]
+    dd = torch.cat([cb0, cb_rest], dim=1)  # [F-1, k] de-delayed
+    # [F-1, k] -> [k, F-1] -> flat [k * (F-1)] (codebook-major), as Code2Wav expects.
+    return dd.transpose(0, 1).contiguous().reshape(-1)
+
+
+def talker2code2wav_token_only(
+    source_outputs: list,
+    prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+) -> list:
+    """Sync ``process_engine_inputs``: build the code2wav placeholder inputs.
+
+    Returns one :class:`OmniTokensPrompt` per finished talker request, with
+    ``prompt_token_ids`` sized to the flat codebook-major codec length
+    (``num_active_codebooks * num_agent_frames``). The actual codec ids are
+    delivered via the worker connector payload from ``talker2code2wav_full_payload``.
+    """
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    del prompt, _requires_multimodal_data
+    inputs: list = []
+    for talker_output in source_outputs:
+        if not getattr(talker_output, "finished", False):
+            continue
+        output = talker_output.outputs[0]
+        # The per-request output is the "latent" (engine_output_type="latent"); the
+        # codes ship via the connector full_payload. Size the placeholder from the
+        # generated token count (one AR token == one Mimi frame). prompt was 1 frame.
+        token_ids = getattr(output, "cumulative_token_ids", None) or getattr(output, "token_ids", None) or []
+        n_frames = max(len(token_ids) - 1, 0)
+        prompt_len = _NUM_ACTIVE_CODEBOOKS * n_frames
+        inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * prompt_len,
+                additional_information=None,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return inputs
+
+
+def talker2code2wav_full_payload(
+    transfer_manager: Any = None,
+    pooling_output: Any = None,
+    request: Any = None,
+    is_finished: bool = False,
+    **kwargs: Any,
+) -> OmniPayloadStruct:
+    """Producer: collect the talker's accumulated agent codes -> Code2Wav input.
+
+    Called by the connector with (transfer_manager, pooling_output, request,
+    is_finished). Prefer the request's additional_information payload under
+    ("codes","audio") (talker_mtp_output_key), while accepting pooling_output
+    and the legacy multimodal_output keyword as compatibility fallbacks.
+    """
+    del transfer_manager, is_finished
+
+    def _codes_from(src: Any) -> torch.Tensor | None:
+        if not isinstance(src, dict):
+            return None
+        nested = src.get("codes")
+        audio = nested.get("audio") if isinstance(nested, dict) else None
+        return audio if audio is not None else src.get("codes.audio")
+
+    audio = None
+    for source in (
+        getattr(request, "additional_information", None),
+        getattr(request, "additional_information_cpu", None),
+        pooling_output,
+        # Temporary compatibility shim for the two producer call contracts.
+        # Remove after https://github.com/vllm-project/vllm-omni/issues/4872
+        # provides validated full-payload and async-chunk processor protocols.
+        kwargs.get("multimodal_output"),
+    ):
+        audio = _codes_from(source)
+        if audio is not None:
+            break
+    if not isinstance(audio, torch.Tensor) or audio.numel() == 0:
+        return _empty_finished_payload()
+    flat = _agent_codes_to_codebook_major(audio)
+    if flat.numel() == 0:
+        return _empty_finished_payload()
+    return OmniPayloadStruct(
+        codes=CodesStruct(audio=flat),
+        meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+    )
+
+
+def talker2code2wav_async_chunk(
+    transfer_manager: Any,
+    multimodal_output: Any,
+    request: Any,
+    is_finished: bool = False,
+) -> OmniPayloadStruct | None:
+    """Streaming: accumulate per-frame agent codes, emit a codebook-major chunk.
+
+    Minimal fixed-chunk variant (no left-context / ref-code complexity, which
+    PersonaPlex does not use). Frames are buffered on the transfer manager until a
+    chunk's worth is ready (or the request finishes), then flushed.
+    """
+    request_id = getattr(request, "external_req_id", getattr(request, "request_id", "?"))
+    # The adapter passes ``is_finished=True`` for both a resumable segment
+    # boundary and the terminal request boundary. PersonaPlex must preserve the
+    # delayed cb1..7 tail across the former; only a non-resumable stop flushes
+    # the stream.
+    finished = bool(is_finished and not getattr(request, "resumable", False))
+    request_payload = getattr(transfer_manager, "request_payload", None)
+    if not isinstance(request_payload, dict):
+        request_payload = {}
+        transfer_manager.request_payload = request_payload
+    state = request_payload.setdefault(request_id, {})
+    frames = state.setdefault("personaplex_frames", [])
+
+    # Codes live in the server-side request's additional_information under
+    # ("codes","audio") (talker_mtp_output_key), not in multimodal_output (latent).
+    def _codes_from(src: Any) -> torch.Tensor | None:
+        if isinstance(src, dict):
+            nested = src.get("codes")
+            a = nested.get("audio") if isinstance(nested, dict) else None
+            return a if a is not None else src.get("codes.audio")
+        return None
+
+    # Explicit None fallback: `a or b` would evaluate bool(a) on a multi-element
+    # Tensor and raise "Boolean value of Tensor ... is ambiguous".
+    audio = _codes_from(getattr(request, "additional_information", None))
+    if audio is None:
+        audio = _codes_from(multimodal_output)
+    if isinstance(audio, torch.Tensor) and audio.numel() > 0:
+        a = audio if audio.ndim == 2 else audio.reshape(1, -1)
+        frames.append(a[-1].to(torch.long).cpu())  # latest frame's codes
+
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+    chunk = int(cfg.get("codec_chunk_frames", 25))
+    initial_chunk = int(cfg.get("initial_codec_chunk_frames") or 0)
+    if chunk <= 0 or initial_chunk < 0:
+        raise ValueError(
+            "PersonaPlex codec chunk sizes must be positive/non-negative: "
+            f"codec_chunk_frames={chunk}, initial_codec_chunk_frames={initial_chunk}"
+        )
+    target_frames = initial_chunk if not state.get("personaplex_emitted") and initial_chunk > 0 else chunk
+
+    # De-delay needs one successor raw frame: N output acoustic frames require
+    # N + 1 raw depformer rows.
+    available_frames = max(0, len(frames) - 1)
+    if available_frames < target_frames and not finished:
+        # Each full-duplex input frame is a resumable stage-0 segment. Returning
+        # None would make the generic chunk adapter synthesize a
+        # segment-finished marker, wake Code2Wav with its one-token placeholder,
+        # and discard the buffered de-delay tail. Explicitly keep this transport
+        # chunk non-terminal until enough successor frames exist.
+        pending = torch.tensor(False, dtype=torch.bool)
+        return OmniPayloadStruct(
+            meta=MetaStruct(
+                finished=pending,
+                is_segment_finished=pending,
+            )
+        )
+    emit_frames = available_frames if finished else target_frames
+    if emit_frames <= 0:
+        if finished:
+            request_payload.pop(request_id, None)
+            return _empty_finished_payload()
+        return None
+
+    stacked = torch.stack(frames[: emit_frames + 1], dim=0)  # [F+1, dep_q]
+    flat = _agent_codes_to_codebook_major(stacked)
+    if finished:
+        request_payload.pop(request_id, None)
+    else:
+        # Row ``emit_frames`` is the successor used by the last emitted frame
+        # and the cb0 source for the next frame.
+        state["personaplex_frames"] = frames[emit_frames:]
+        state["personaplex_emitted"] = True
+    return OmniPayloadStruct(
+        codes=CodesStruct(audio=flat),
+        meta=MetaStruct(finished=torch.tensor(bool(finished), dtype=torch.bool)),
+    )
+
+
+__all__ = [
+    "talker2code2wav_token_only",
+    "talker2code2wav_full_payload",
+    "talker2code2wav_async_chunk",
+]

vllm_omni/worker/gpu_ar_model_runner.py

@@ -1516,6 +1516,120 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             return None
         return hidden_states_cpu[start:end]
 
+    def _run_post_sample_talker_mtp(
+        self,
+        *,
+        req_ids: list[str],
+        valid_sampled_token_ids: list[list[int]],
+        sampled_token_ids: torch.Tensor,
+        invalid_req_indices: list[int],
+        sample_hidden_states: torch.Tensor,
+        multimodal_outputs: Any,
+    ) -> Any:
+        """Run an opt-in talker MTP with the token sampled in this step.
+
+        Some frame-synchronous talkers need the current temporal hidden state
+        and its newly sampled text token to produce the same frame's codec
+        codes. A resumable ``max_tokens=1`` segment has no later decode
+        preprocess in which to do that work, so those models expose
+        ``post_sample_talker_mtp`` instead.
+        """
+        model = getattr(self, "model", None)
+        hook = getattr(model, "post_sample_talker_mtp", None)
+        if not callable(hook):
+            return multimodal_outputs
+
+        selected_indices: list[int] = []
+        selected_req_ids: list[str] = []
+        selected_req_infos: list[dict[str, Any]] = []
+        duplex_indices: list[int] = []
+        invalid_indices = set(invalid_req_indices)
+        use_async_scheduling = bool(getattr(self, "use_async_scheduling", False))
+        for idx, req_id in enumerate(req_ids):
+            req_info = self.model_intermediate_buffer.get(req_id)
+            duplex = req_info.get("duplex") if isinstance(req_info, dict) else None
+            if not isinstance(duplex, dict) or duplex.get("data_plane") is not True:
+                continue
+            duplex_indices.append(idx)
+            if use_async_scheduling:
+                if idx in invalid_indices:
+                    continue
+            else:
+                token_ids = valid_sampled_token_ids[idx] if idx < len(valid_sampled_token_ids) else []
+                if not token_ids:
+                    continue
+                if len(token_ids) != 1:
+                    raise ValueError(
+                        "post_sample_talker_mtp requires exactly one sampled token "
+                        f"per request, got {len(token_ids)} for {req_id}"
+                    )
+            selected_indices.append(idx)
+            selected_req_ids.append(req_id)
+            selected_req_infos.append(req_info)
+
+        if not selected_indices:
+            return multimodal_outputs
+        if sampled_token_ids.ndim != 2 or sampled_token_ids.shape[0] < len(req_ids):
+            raise ValueError(
+                "post_sample_talker_mtp sampled-token rows do not match requests: "
+                f"shape={tuple(sampled_token_ids.shape)}, requests={len(req_ids)}"
+            )
+        if sampled_token_ids.shape[1] != 1:
+            raise ValueError(
+                "post_sample_talker_mtp does not support speculative token rows: "
+                f"shape={tuple(sampled_token_ids.shape)}"
+            )
+
+        empty_audio = torch.empty(0, dtype=torch.long, device=sample_hidden_states.device)
+        merged = dict(multimodal_outputs) if isinstance(multimodal_outputs, dict) else {}
+        existing_codes = merged.get("codes")
+        codes_payload = dict(existing_codes) if isinstance(existing_codes, dict) else {}
+        existing_audio = codes_payload.get("audio")
+        if isinstance(existing_audio, list) and len(existing_audio) == len(req_ids):
+            per_request_audio = list(existing_audio)
+        elif (
+            isinstance(existing_audio, torch.Tensor)
+            and existing_audio.ndim > 0
+            and existing_audio.shape[0] == len(req_ids)
+        ):
+            per_request_audio = [existing_audio[idx : idx + 1] for idx in range(len(req_ids))]
+        else:
+            per_request_audio = [empty_audio for _ in req_ids]
+        for idx in duplex_indices:
+            per_request_audio[idx] = empty_audio
+
+        indices = torch.tensor(
+            selected_indices,
+            dtype=torch.long,
+            device=sample_hidden_states.device,
+        )
+        input_ids = sampled_token_ids.index_select(0, indices).reshape(-1).to(torch.long)
+        codes = hook(
+            input_ids=input_ids,
+            hidden_states=sample_hidden_states.index_select(0, indices),
+            req_ids=selected_req_ids,
+            req_infos=selected_req_infos,
+        )
+        if not isinstance(codes, torch.Tensor) or codes.ndim != 2:
+            raise TypeError(f"post_sample_talker_mtp must return a rank-2 Tensor, got {type(codes).__name__}")
+        if codes.shape[0] != len(selected_req_ids):
+            raise ValueError(
+                f"post_sample_talker_mtp row count does not match requests: {codes.shape[0]} != {len(selected_req_ids)}"
+            )
+        for row, (idx, req_id) in enumerate(zip(selected_indices, selected_req_ids, strict=True)):
+            request_codes = codes[row : row + 1]
+            per_request_audio[idx] = request_codes
+            self._update_intermediate_buffer(
+                req_id,
+                {"codes": {"audio": request_codes}},
+            )
+
+        # Empty rows are intentional: they suppress replay of the previous
+        # frame while a chunked prefill has not sampled a token yet.
+        codes_payload["audio"] = per_request_audio
+        merged["codes"] = codes_payload
+        return merged
+
     def _build_multimodal_outputs(
         self,
         per_req_payloads: list[dict[str, object] | None] | None,
@@ -2036,6 +2150,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 scheduler_output.total_num_scheduled_tokens,
             )
 
+        multimodal_outputs = self._run_post_sample_talker_mtp(
+            req_ids=req_ids_output_copy,
+            valid_sampled_token_ids=valid_sampled_token_ids,
+            sampled_token_ids=sampler_output.sampled_token_ids,
+            invalid_req_indices=invalid_req_indices,
+            sample_hidden_states=sample_hidden_states,
+            multimodal_outputs=multimodal_outputs,
+        )
+
         if propose_drafts_after_bookkeeping:
             # ngram and other speculative decoding methods use the sampled
             # tokens on the CPU, so they are run after bookkeeping.