From e788ef6e76a93e8fe73c3072064f985ea0b50889 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Tue, 4 Aug 2026 20:10:33 -0700 Subject: [PATCH] [Feature] PersonaPlex (Moshi-based full-duplex S2S): native vLLM port + duplex serving (#4771) Signed-off-by: Yueqian Lin Signed-off-by: Sy0307 <24841857+Sy0307@users.noreply.github.com> Co-authored-by: Sy0307 <24841857+Sy0307@users.noreply.github.com> --- .../personaplex/personaplex_offline.py | 77 +++ examples/online_serving/personaplex/README.md | 84 +++ .../personaplex/duplex_client.py | 114 ++++ .../online_serving/personaplex/run_server.sh | 15 + recipes/NVIDIA/PersonaPlex.md | 174 +++++ tests/config/test_omni_config.py | 17 + ...mni_generation_scheduler_update_session.py | 17 +- .../test_chunk_transfer_adapter.py | 99 +++ .../personaplex/test_import_boundary.py | 30 + .../personaplex/test_stage0_runtime.py | 303 +++++++++ .../personaplex/test_streaming_code2wav.py | 316 +++++++++ .../personaplex/test_unified_runtime.py | 364 ++++++++++ .../fullduplex/test_core_contracts.py | 122 ++++ .../fullduplex/test_personaplex_adapter.py | 262 +++++++ .../fullduplex/test_personaplex_batched.py | 178 +++++ .../fullduplex/test_personaplex_elastic.py | 240 +++++++ .../test_personaplex_server_lease.py | 101 +++ .../fullduplex/test_personaplex_session.py | 97 +++ .../personaplex_realtime_duplex.py | 492 ++++++++++++++ .../test_personaplex.py | 135 ++++ tests/worker/test_gpu_ar_model_runner.py | 155 +++++ vllm_omni/config/model.py | 1 + vllm_omni/config/omni_config.py | 11 +- vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/config/stage_config.py | 5 + .../core/sched/omni_generation_scheduler.py | 2 +- vllm_omni/deploy/personaplex.yaml | 90 +++ vllm_omni/engine/arg_utils.py | 13 + vllm_omni/experimental/fullduplex/README.md | 43 +- .../fullduplex/personaplex/DESIGN.md | 259 +++++++ .../fullduplex/personaplex/__init__.py | 40 ++ .../fullduplex/personaplex/adapter.py | 177 +++++ .../fullduplex/personaplex/config.py | 64 ++ .../fullduplex/personaplex/data_plane.py | 228 +++++++ .../fullduplex/personaplex/engine.py | 74 ++ .../fullduplex/personaplex/input.py | 213 ++++++ .../fullduplex/personaplex/policy.py | 60 ++ .../fullduplex/personaplex/runtime.py | 473 +++++++++++++ .../personaplex/runtime_extension.py | 129 ++++ .../personaplex/serving/__init__.py | 15 + .../fullduplex/personaplex/serving/batched.py | 184 +++++ .../fullduplex/personaplex/serving/server.py | 539 +++++++++++++++ .../fullduplex/personaplex/serving_adapter.py | 233 +++++++ .../fullduplex/personaplex/session.py | 124 ++++ .../fullduplex/personaplex/stage0.py | 467 +++++++++++++ .../models/personaplex/__init__.py | 19 + .../personaplex/configuration_helium.py | 62 ++ .../personaplex/configuration_personaplex.py | 278 ++++++++ .../models/personaplex/modeling_helium.py | 548 +++++++++++++++ .../personaplex/personaplex_code2wav.py | 433 ++++++++++++ .../personaplex/personaplex_depformer.py | 366 ++++++++++ .../personaplex/personaplex_embeddings.py | 90 +++ .../models/personaplex/personaplex_mimi.py | 413 +++++++++++ .../models/personaplex/personaplex_talker.py | 639 ++++++++++++++++++ .../personaplex/personaplex_temporal.py | 257 +++++++ .../models/personaplex/pipeline.py | 75 ++ vllm_omni/model_executor/models/registry.py | 12 + .../stage_input_processors/personaplex.py | 248 +++++++ vllm_omni/worker/gpu_ar_model_runner.py | 123 ++++ 59 files changed, 10389 insertions(+), 12 deletions(-) create mode 100644 examples/offline_inference/personaplex/personaplex_offline.py create mode 100644 examples/online_serving/personaplex/README.md create mode 100644 examples/online_serving/personaplex/duplex_client.py create mode 100755 examples/online_serving/personaplex/run_server.sh create mode 100644 recipes/NVIDIA/PersonaPlex.md create mode 100644 tests/e2e/features/fullduplex/personaplex/test_import_boundary.py create mode 100644 tests/e2e/features/fullduplex/personaplex/test_stage0_runtime.py create mode 100644 tests/e2e/features/fullduplex/personaplex/test_streaming_code2wav.py create mode 100644 tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py create mode 100644 tests/e2e/features/fullduplex/test_core_contracts.py create mode 100644 tests/e2e/features/fullduplex/test_personaplex_adapter.py create mode 100644 tests/e2e/features/fullduplex/test_personaplex_batched.py create mode 100644 tests/e2e/features/fullduplex/test_personaplex_elastic.py create mode 100644 tests/e2e/features/fullduplex/test_personaplex_server_lease.py create mode 100644 tests/e2e/features/fullduplex/test_personaplex_session.py create mode 100644 tests/e2e/online_serving/personaplex_realtime_duplex.py create mode 100644 tests/model_executor/stage_input_processors/test_personaplex.py create mode 100644 vllm_omni/deploy/personaplex.yaml create mode 100644 vllm_omni/experimental/fullduplex/personaplex/DESIGN.md create mode 100644 vllm_omni/experimental/fullduplex/personaplex/__init__.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/adapter.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/config.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/data_plane.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/engine.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/input.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/policy.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/runtime.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/runtime_extension.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/serving/__init__.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/serving/batched.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/serving/server.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/serving_adapter.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/session.py create mode 100644 vllm_omni/experimental/fullduplex/personaplex/stage0.py create mode 100644 vllm_omni/model_executor/models/personaplex/__init__.py create mode 100644 vllm_omni/model_executor/models/personaplex/configuration_helium.py create mode 100644 vllm_omni/model_executor/models/personaplex/configuration_personaplex.py create mode 100644 vllm_omni/model_executor/models/personaplex/modeling_helium.py create mode 100644 vllm_omni/model_executor/models/personaplex/personaplex_code2wav.py create mode 100644 vllm_omni/model_executor/models/personaplex/personaplex_depformer.py create mode 100644 vllm_omni/model_executor/models/personaplex/personaplex_embeddings.py create mode 100644 vllm_omni/model_executor/models/personaplex/personaplex_mimi.py create mode 100644 vllm_omni/model_executor/models/personaplex/personaplex_talker.py create mode 100644 vllm_omni/model_executor/models/personaplex/personaplex_temporal.py create mode 100644 vllm_omni/model_executor/models/personaplex/pipeline.py create mode 100644 vllm_omni/model_executor/stage_input_processors/personaplex.py diff --git a/examples/offline_inference/personaplex/personaplex_offline.py b/examples/offline_inference/personaplex/personaplex_offline.py new file mode 100644 index 000000000..ea011b8a9 --- /dev/null +++ b/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() diff --git a/examples/online_serving/personaplex/README.md b/examples/online_serving/personaplex/README.md new file mode 100644 index 000000000..cf105dd6c --- /dev/null +++ b/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). diff --git a/examples/online_serving/personaplex/duplex_client.py b/examples/online_serving/personaplex/duplex_client.py new file mode 100644 index 000000000..bb1eaba2b --- /dev/null +++ b/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://:/`` 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() diff --git a/examples/online_serving/personaplex/run_server.sh b/examples/online_serving/personaplex/run_server.sh new file mode 100755 index 000000000..53f6f0886 --- /dev/null +++ b/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" diff --git a/recipes/NVIDIA/PersonaPlex.md b/recipes/NVIDIA/PersonaPlex.md new file mode 100644 index 000000000..62d04a499 --- /dev/null +++ b/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⟩` | ``/`` | **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 ` ... ` 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. diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 9650471d6..6a1b263d2 100644 --- a/tests/config/test_omni_config.py +++ b/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", diff --git a/tests/core/sched/test_omni_generation_scheduler_update_session.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py index 922cc5eb6..0f654715a 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/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: diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index e1a1e36f9..b1ff5db94 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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") diff --git a/tests/e2e/features/fullduplex/personaplex/test_import_boundary.py b/tests/e2e/features/fullduplex/personaplex/test_import_boundary.py new file mode 100644 index 000000000..463c0c98e --- /dev/null +++ b/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 diff --git a/tests/e2e/features/fullduplex/personaplex/test_stage0_runtime.py b/tests/e2e/features/fullduplex/personaplex/test_stage0_runtime.py new file mode 100644 index 000000000..3bdc4456b --- /dev/null +++ b/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=" 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] diff --git a/tests/e2e/features/fullduplex/personaplex/test_streaming_code2wav.py b/tests/e2e/features/fullduplex/personaplex/test_streaming_code2wav.py new file mode 100644 index 000000000..abe2146b6 --- /dev/null +++ b/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 diff --git a/tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py b/tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py new file mode 100644 index 000000000..f4a3c994e --- /dev/null +++ b/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=" SimpleNamespace: + frames = np.full((frame_count, e2e_driver.FRAME_SAMPLES), 7, dtype=" None: + silent = np.zeros(e2e_driver.FRAME_SAMPLES, dtype=" None: + partial_frame = np.zeros(e2e_driver.FRAME_SAMPLES - 1, dtype=" 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) diff --git a/tests/e2e/features/fullduplex/test_core_contracts.py b/tests/e2e/features/fullduplex/test_core_contracts.py new file mode 100644 index 000000000..2a514b9f2 --- /dev/null +++ b/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 diff --git a/tests/e2e/features/fullduplex/test_personaplex_adapter.py b/tests/e2e/features/fullduplex/test_personaplex_adapter.py new file mode 100644 index 000000000..458333627 --- /dev/null +++ b/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 diff --git a/tests/e2e/features/fullduplex/test_personaplex_batched.py b/tests/e2e/features/fullduplex/test_personaplex_batched.py new file mode 100644 index 000000000..f2ee73729 --- /dev/null +++ b/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 diff --git a/tests/e2e/features/fullduplex/test_personaplex_elastic.py b/tests/e2e/features/fullduplex/test_personaplex_elastic.py new file mode 100644 index 000000000..453d90153 --- /dev/null +++ b/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) diff --git a/tests/e2e/features/fullduplex/test_personaplex_server_lease.py b/tests/e2e/features/fullduplex/test_personaplex_server_lease.py new file mode 100644 index 000000000..451ffb714 --- /dev/null +++ b/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 diff --git a/tests/e2e/features/fullduplex/test_personaplex_session.py b/tests/e2e/features/fullduplex/test_personaplex_session.py new file mode 100644 index 000000000..f838de4ff --- /dev/null +++ b/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)) diff --git a/tests/e2e/online_serving/personaplex_realtime_duplex.py b/tests/e2e/online_serving/personaplex_realtime_duplex.py new file mode 100644 index 000000000..62ad9e972 --- /dev/null +++ b/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(" 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=" 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="= 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="= 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=" 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() diff --git a/tests/model_executor/stage_input_processors/test_personaplex.py b/tests/model_executor/stage_input_processors/test_personaplex.py new file mode 100644 index 000000000..37f464201 --- /dev/null +++ b/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]) diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index e8c40c690..116b5ee4d 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/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( diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 683c6bd52..d172347f7 100644 --- a/vllm_omni/config/model.py +++ b/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 diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 0207b0132..0f523eb48 100644 --- a/vllm_omni/config/omni_config.py +++ b/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, ) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index bd929bbc3..d702f1ace 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/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, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 649e5e4cf..4831e6f92 100644 --- a/vllm_omni/config/stage_config.py +++ b/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 diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 8a165e2a1..e1632764c 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/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() diff --git a/vllm_omni/deploy/personaplex.yaml b/vllm_omni/deploy/personaplex.yaml new file mode 100644 index 000000000..9f1da7e63 --- /dev/null +++ b/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 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index c9a4f44f5..fbc2bec3a 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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, diff --git a/vllm_omni/experimental/fullduplex/README.md b/vllm_omni/experimental/fullduplex/README.md index 5359edf87..e08c5ee1c 100644 --- a/vllm_omni/experimental/fullduplex/README.md +++ b/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//`; 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. diff --git a/vllm_omni/experimental/fullduplex/personaplex/DESIGN.md b/vllm_omni/experimental/fullduplex/personaplex/DESIGN.md new file mode 100644 index 000000000..623f08353 --- /dev/null +++ b/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 + -> persona 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. diff --git a/vllm_omni/experimental/fullduplex/personaplex/__init__.py b/vllm_omni/experimental/fullduplex/personaplex/__init__.py new file mode 100644 index 000000000..bf6c711bb --- /dev/null +++ b/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", +] diff --git a/vllm_omni/experimental/fullduplex/personaplex/adapter.py b/vllm_omni/experimental/fullduplex/personaplex/adapter.py new file mode 100644 index 000000000..1508931f8 --- /dev/null +++ b/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() diff --git a/vllm_omni/experimental/fullduplex/personaplex/config.py b/vllm_omni/experimental/fullduplex/personaplex/config.py new file mode 100644 index 000000000..1cfa4b736 --- /dev/null +++ b/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 `` ... `` 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 diff --git a/vllm_omni/experimental/fullduplex/personaplex/data_plane.py b/vllm_omni/experimental/fullduplex/personaplex/data_plane.py new file mode 100644 index 000000000..96e1f9e20 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/experimental/fullduplex/personaplex/engine.py b/vllm_omni/experimental/fullduplex/personaplex/engine.py new file mode 100644 index 000000000..488676aa1 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/experimental/fullduplex/personaplex/input.py b/vllm_omni/experimental/fullduplex/personaplex/input.py new file mode 100644 index 000000000..2fe57442a --- /dev/null +++ b/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=" 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", +] diff --git a/vllm_omni/experimental/fullduplex/personaplex/policy.py b/vllm_omni/experimental/fullduplex/personaplex/policy.py new file mode 100644 index 000000000..8e6062d0d --- /dev/null +++ b/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("") else f" {t} " + + +__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", +] diff --git a/vllm_omni/experimental/fullduplex/personaplex/runtime.py b/vllm_omni/experimental/fullduplex/personaplex/runtime.py new file mode 100644 index 000000000..885b36aa5 --- /dev/null +++ b/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("▁", " ") diff --git a/vllm_omni/experimental/fullduplex/personaplex/runtime_extension.py b/vllm_omni/experimental/fullduplex/personaplex/runtime_extension.py new file mode 100644 index 000000000..036103a43 --- /dev/null +++ b/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"] diff --git a/vllm_omni/experimental/fullduplex/personaplex/serving/__init__.py b/vllm_omni/experimental/fullduplex/personaplex/serving/__init__.py new file mode 100644 index 000000000..c3044a80d --- /dev/null +++ b/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) diff --git a/vllm_omni/experimental/fullduplex/personaplex/serving/batched.py b/vllm_omni/experimental/fullduplex/personaplex/serving/batched.py new file mode 100644 index 000000000..26abd92a8 --- /dev/null +++ b/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() diff --git a/vllm_omni/experimental/fullduplex/personaplex/serving/server.py b/vllm_omni/experimental/fullduplex/personaplex/serving/server.py new file mode 100644 index 000000000..fcb857628 --- /dev/null +++ b/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://: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() diff --git a/vllm_omni/experimental/fullduplex/personaplex/serving_adapter.py b/vllm_omni/experimental/fullduplex/personaplex/serving_adapter.py new file mode 100644 index 000000000..66244ca68 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/experimental/fullduplex/personaplex/session.py b/vllm_omni/experimental/fullduplex/personaplex/session.py new file mode 100644 index 000000000..b3a476f12 --- /dev/null +++ b/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 diff --git a/vllm_omni/experimental/fullduplex/personaplex/stage0.py b/vllm_omni/experimental/fullduplex/personaplex/stage0.py new file mode 100644 index 000000000..a11c5dc40 --- /dev/null +++ b/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=" 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", +] diff --git a/vllm_omni/model_executor/models/personaplex/__init__.py b/vllm_omni/model_executor/models/personaplex/__init__.py new file mode 100644 index 000000000..b4e3d49cd --- /dev/null +++ b/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", +] diff --git a/vllm_omni/model_executor/models/personaplex/configuration_helium.py b/vllm_omni/model_executor/models/personaplex/configuration_helium.py new file mode 100644 index 000000000..1158637ac --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/personaplex/configuration_personaplex.py b/vllm_omni/model_executor/models/personaplex/configuration_personaplex.py new file mode 100644 index 000000000..e43e12168 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/personaplex/modeling_helium.py b/vllm_omni/model_executor/models/personaplex/modeling_helium.py new file mode 100644 index 000000000..e4ac8d885 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_code2wav.py b/vllm_omni/model_executor/models/personaplex/personaplex_code2wav.py new file mode 100644 index 000000000..a5712aa40 --- /dev/null +++ b/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."))} diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_depformer.py b/vllm_omni/model_executor/models/personaplex/personaplex_depformer.py new file mode 100644 index 000000000..2a9a12905 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_embeddings.py b/vllm_omni/model_executor/models/personaplex/personaplex_embeddings.py new file mode 100644 index 000000000..18c502f98 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py new file mode 100644 index 000000000..48df20fab --- /dev/null +++ b/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, :] diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_talker.py b/vllm_omni/model_executor/models/personaplex/personaplex_talker.py new file mode 100644 index 000000000..5e56e7a7c --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py b/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py new file mode 100644 index 000000000..717e4d505 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/personaplex/pipeline.py b/vllm_omni/model_executor/models/personaplex/pipeline.py new file mode 100644 index 000000000..1865d4b3a --- /dev/null +++ b/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}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 2e1662ea0..fd76dadd8 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/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", diff --git a/vllm_omni/model_executor/stage_input_processors/personaplex.py b/vllm_omni/model_executor/stage_input_processors/personaplex.py new file mode 100644 index 000000000..71f33ba33 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 5d66a2405..b2437ca02 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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.