diff --git a/recipes/NVIDIA/PersonaPlex.md b/recipes/NVIDIA/PersonaPlex.md index 65c6453ff..ee199aec9 100644 --- a/recipes/NVIDIA/PersonaPlex.md +++ b/recipes/NVIDIA/PersonaPlex.md @@ -84,8 +84,10 @@ duplex (`is_enabled()` is unconditionally true). Turn-based models ## Hardware Support -Verified on GPU (Hopper-class). The serving path is plain PyTorch eager plus the -native modules, so other CUDA GPUs with enough memory are expected to work. +Verified on GPU (Hopper-class). The shipped deploy config runs the Stage 0 +talker with CUDA graphs and replays the streaming Mimi encode (Stage 0) and +decode (Stage 1) frame steps from model-local CUDA graphs, bitwise equal to +eager. Other CUDA GPUs with enough memory are expected to work. ## GPU @@ -137,9 +139,11 @@ python tests/e2e/online_serving/personaplex_realtime_duplex.py \ #### Notes -- Realtime budget is 80 ms/frame. Verified eager per-tick latency at four - concurrent sessions is ~70-74 ms on this hardware class, i.e. all four - conversations stay realtime; single-session has comfortable headroom. +- Realtime budget is 80 ms/frame. The shipped config admits two sessions; on a + 141 GB card, 64 concurrent 30 s sessions stayed realtime after raising + `duplex_session.max_sessions` and both stages' `max_num_seqs` (see the + sizing notes in `vllm_omni/deploy/personaplex.yaml`). Longer sessions cost + more per step as their context fills. - Decoding is greedy only (temperature/top-k knobs are intentionally not exposed): greedy is what the parity gates pin against the reference implementation, and sampling amplifies sub-bit numeric drift into divergence. diff --git a/tests/model_executor/models/personaplex/duplex/test_plugin.py b/tests/model_executor/models/personaplex/duplex/test_plugin.py index 72cb6d9aa..d409f098d 100644 --- a/tests/model_executor/models/personaplex/duplex/test_plugin.py +++ b/tests/model_executor/models/personaplex/duplex/test_plugin.py @@ -17,6 +17,7 @@ from vllm.sampling_params import SamplingParams from tests.e2e.online_serving import personaplex_realtime_duplex as e2e_driver from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy +from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.engine.duplex.config import DuplexSessionConfig from vllm_omni.engine.duplex.contracts import DuplexFence from vllm_omni.engine.duplex.plugin import ( @@ -26,6 +27,7 @@ from vllm_omni.engine.duplex.plugin import ( load_duplex_plugin, validate_duplex_plugin_sampling, ) +from vllm_omni.model_executor.models.personaplex.configuration_personaplex import PersonaPlexConfig from vllm_omni.model_executor.models.personaplex.duplex import stage0 from vllm_omni.model_executor.models.personaplex.duplex.config import DEFAULT_PERSONA, FRAME_SIZE, SAMPLE_RATE from vllm_omni.model_executor.models.personaplex.duplex.data_plane import PersonaPlexDataPlaneSession @@ -106,6 +108,36 @@ def test_personaplex_deploy_is_duplex_and_propagates_capacity_to_all_model_stage assert "personaplex_codec_max_sessions" not in deploy.connectors["connector_of_shared_memory"]["extra"] +def test_personaplex_deploy_enables_mimi_cuda_graphs_on_both_stage_configs(tmp_path: Path) -> None: + # The checkpoint ships an empty config.json; each stage's engine args build + # its PersonaPlexConfig from that plus the deploy YAML's hf_overrides. + (tmp_path / "config.json").write_text("{}") + stages = merge_pipeline_deploy(PERSONAPLEX_PIPELINE, load_deploy_config(DEPLOY_PATH)) + + def resolve(model_arch: str, hf_overrides: dict[str, object] | None) -> PersonaPlexConfig: + engine_args = OmniEngineArgs( + model=str(tmp_path), + model_arch=model_arch, + hf_overrides=hf_overrides, + skip_tokenizer_init=True, + trust_remote_code=True, + ) + config = engine_args.create_model_config().hf_config + assert isinstance(config, PersonaPlexConfig) + return config + + assert [stage.yaml_engine_args["model_arch"] for stage in stages] == [ + "PersonaPlexTalkerForConditionalGeneration", + "PersonaPlexCode2Wav", + ] + for stage in stages: + model_arch = stage.yaml_engine_args["model_arch"] + assert stage.yaml_engine_args["hf_overrides"] == {"mimi_cuda_graphs": True} + assert resolve(model_arch, dict(stage.yaml_engine_args["hf_overrides"])).mimi_cuda_graphs is True + # Without the deploy override the flag keeps its default. + assert resolve(model_arch, None).mimi_cuda_graphs is False + + # --------------------------------------------------------------------------- # # Session policy # # --------------------------------------------------------------------------- # diff --git a/tests/model_executor/models/personaplex/duplex/test_stage0_batched_append.py b/tests/model_executor/models/personaplex/duplex/test_stage0_batched_append.py new file mode 100644 index 000000000..64e3874d8 --- /dev/null +++ b/tests/model_executor/models/personaplex/duplex/test_stage0_batched_append.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Batched live-append preparation must match the per-request path exactly.""" + +import base64 +from dataclasses import dataclass + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.models.personaplex.duplex.policy import ( + SILENCE_TOKENS, + SINE_TOKENS, +) +from vllm_omni.model_executor.models.personaplex.duplex.stage0 import ( + PersonaPlexStage0DuplexRuntime, +) +from vllm_omni.model_executor.models.personaplex.personaplex_talker import ( + PersonaPlexTalkerForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_WIDTH = 17 # text row + 8 agent rows + 8 user rows +_PREFILL_LEN = 18 # 2 voice rows + 2 x 6 silence frames + 3 persona tokens + the live frame + + +class _FakeCodec: + """Shared streaming encoder whose codes encode the row's frame count.""" + + def __init__(self) -> None: + self.encode_calls = 0 + self.frames: list[int] = [] + + def streaming_init(self, batch_size: int) -> None: + self.frames = [0] * batch_size + + def encode_frame(self, pcm: torch.Tensor, active: torch.Tensor) -> torch.Tensor: + del pcm + self.encode_calls += 1 + codes = torch.zeros((len(self.frames), 8), dtype=torch.long) + for row, is_active in enumerate(active.tolist()): + if is_active: + self.frames[row] += 1 + codes[row] = self.frames[row] * 16 + torch.arange(8) + 1 + return codes + + def reset_slot(self, row: int) -> None: + self.frames[row] = 0 + + +@dataclass(frozen=True) +class _TalkerConfig: + num_audio_codebooks: int = 16 + audio_vocab_size: int = 2048 + + +class _EmbeddingTalker: + """Stage model whose input embedding is the token stack itself, so embeds expose the layout.""" + + def __init__(self) -> None: + self.device = torch.device("cpu") + self.dtype = torch.float32 + self.config = _TalkerConfig() + + def input_embeddings(self, stack: torch.Tensor) -> torch.Tensor: + return stack.squeeze(-1).to(torch.float32) + + def _build_prefill_embed(self, tokens, offset, span, device, silence=None, user_sine=None): + del offset, silence, user_sine + values = tokens[:span].to(device=device, dtype=torch.float32) + return values[:, None].expand(-1, _WIDTH).contiguous() + + _build_frame_embed = PersonaPlexTalkerForConditionalGeneration._build_frame_embed + + +def _depformer(text_token, hidden, *, audio_tokens, audio_provided, num_steps): + """Deterministic stand-in whose codes depend on every input it is given.""" + del hidden + assert audio_provided.dtype == torch.bool + base = audio_tokens[:, 8 : 8 + num_steps] * 7 + audio_provided[:, :num_steps].long() + return (base + text_token[:, None] + torch.arange(num_steps)) % 2048 + + +def _talker(max_sessions: int) -> PersonaPlexTalkerForConditionalGeneration: + runtime = PersonaPlexStage0DuplexRuntime( + _EmbeddingTalker(), + model_path="/unused", + device="cpu", + codec=_FakeCodec(), + max_sessions=max_sessions, + tokenizer=lambda _text: [7, 8, 9], + voice_loader=lambda _voice: {"embeddings": torch.arange(2 * _WIDTH, dtype=torch.float32).reshape(2, _WIDTH)}, + ) + talker = PersonaPlexTalkerForConditionalGeneration.__new__(PersonaPlexTalkerForConditionalGeneration) + torch.nn.Module.__init__(talker) + talker._personaplex_duplex_stage0_runtime = runtime + talker._dtype = torch.float32 + talker.mtp_hidden_size = _WIDTH + talker.num_active_codebooks = 8 + talker.depformer = _depformer + return talker + + +def _append(session_id: str, seq: int, epoch: int = 0) -> dict: + pcm = np.zeros(1920, dtype=" _StepResult: + """One scheduler step: batch preprocess, per-request preprocess, then depformer sampling. + + ``step`` holds ``(request_id, append, prompt_len)``; a prompt length of + ``_PREFILL_LEN`` marks a first append scheduled as one full prefill span. + """ + runtime = talker._personaplex_duplex_stage0_runtime + infos = {} + spans = {} + for request_id, duplex, prompt_len in step: + spans[request_id] = prompt_len if prompt_len == _PREFILL_LEN else 1 + infos[request_id] = { + "duplex": duplex, + "request_id": request_id, + "duplex_prompt_len": prompt_len, + "duplex_token_offset": prompt_len - spans[request_id], + } + if batched: + talker.preprocess_batch(req_ids=list(infos), model_intermediate_buffer=infos, device=torch.device("cpu")) + else: + runtime.encode_appends([info["duplex"] for info in infos.values()]) + embeds, req_infos = {}, [] + for request_id, info in infos.items(): + input_ids = torch.zeros(spans[request_id], dtype=torch.long) + _, embed, update = talker.preprocess(input_ids, None, _omni_is_prefill=True, **info) + embeds[request_id] = embed + req_infos.append(update) + received = {} + + def depformer(text_token, hidden, **kwargs): + received.update(kwargs) + return _depformer(text_token, hidden, **kwargs) + + talker.depformer = depformer + req_ids = list(infos) + codes = talker.post_sample_talker_mtp( + input_ids=torch.arange(len(req_ids)) + text_base, + hidden_states=torch.zeros((len(req_ids), _WIDTH)), + req_ids=req_ids, + req_infos=req_infos, + ) + return _StepResult( + embeds=embeds, + audio_tokens=received["audio_tokens"], + audio_provided=received["audio_provided"], + codes=codes, + agents={key: state.last_agent_codes.clone() for key, state in runtime.sessions.items()}, + ) + + +def _assert_same(batched: _StepResult, reference: _StepResult) -> None: + assert batched.embeds.keys() == reference.embeds.keys() + for request_id, embed in batched.embeds.items(): + assert torch.equal(embed, reference.embeds[request_id]), request_id + assert torch.equal(batched.audio_tokens, reference.audio_tokens) + assert torch.equal(batched.audio_provided, reference.audio_provided) + assert torch.equal(batched.codes, reference.codes) + assert batched.agents.keys() == reference.agents.keys() + for key, agent in batched.agents.items(): + assert torch.equal(agent, reference.agents[key]), key + + +def _device_prepared(talker, key: tuple[str, int]) -> bool: + state = talker._personaplex_duplex_stage0_runtime.sessions[key] + return state.prepared.depformer_audio_tokens_device is not None + + +def test_mixed_prefill_and_live_steps_match_the_per_request_path() -> None: + steps = [ + [("a", _append("a", 1), 18), ("b", _append("b", 1), 18)], + # A first append joins a step of live appends. + [("a", _append("a", 2), 19), ("b", _append("b", 2), 19), ("c", _append("c", 1), 18)], + [("a", _append("a", 3), 20), ("b", _append("b", 3), 20), ("c", _append("c", 2), 19)], + # Only one session appends; the others keep their state. + [("c", _append("c", 3), 20)], + [("a", _append("a", 4), 21), ("c", _append("c", 4), 21)], + ] + batched, reference = _talker(max_sessions=4), _talker(max_sessions=4) + + for index, step in enumerate(steps): + text_base = 100 + 10 * index + result = _run_step(batched, step, batched=True, text_base=text_base) + _assert_same(result, _run_step(reference, step, batched=False, text_base=text_base)) + + assert _device_prepared(batched, ("a", 0)) and _device_prepared(batched, ("c", 0)) + assert not _device_prepared(reference, ("a", 0)) + # Batched live appends publish no host teacher-forcing copies; the per-request path still does. + host_fields = {"pplex_depformer_audio_tokens", "pplex_depformer_audio_provided"} + batched_info = batched._personaplex_duplex_stage0_runtime.sessions[("a", 0)].prepared.info_update + reference_info = reference._personaplex_duplex_stage0_runtime.sessions[("a", 0)].prepared.info_update + assert host_fields.isdisjoint(batched_info) + assert host_fields <= reference_info.keys() + runtime = batched._personaplex_duplex_stage0_runtime + assert runtime._shared_codec().encode_calls == len(steps) + assert len(runtime.sessions[("a", 0)].user_history_device) == 2 + # Live embeds carry the causal user delay: cb0 from the previous frame, cb1..7 from the one before. + a4 = result.embeds["a"][0] + assert a4[9].item() == 3 * 16 + 1 + assert a4[10:17].tolist() == [2 * 16 + k + 1 for k in range(1, 8)] + + +def test_second_append_reads_sine_for_the_missing_user_history() -> None: + talker = _talker(max_sessions=1) + + def state_of(): + return talker._personaplex_duplex_stage0_runtime.sessions[("a", 0)] + + first = _run_step(talker, [("a", _append("a", 1), 18)], batched=True, text_base=100) + # The first append forces agent cb1..7 to silence; cb0 is sampled. + assert state_of().last_agent_codes.tolist() == [first.codes[0, 0].item(), *SILENCE_TOKENS[1:]] + result = _run_step(talker, [("a", _append("a", 2), 19)], batched=True, text_base=101) + # Live appends force only the user codebooks, so the sampled agent frame is kept as is. + assert torch.equal(state_of().last_agent_codes, result.codes[0]) + + embed = result.embeds["a"][0] + assert embed[0].item() == 100 + assert embed[9].item() == 16 + 1 + assert embed[10:17].tolist() == list(SINE_TOKENS[1:]) + assert result.audio_tokens[0].tolist() == [*SILENCE_TOKENS, 2 * 16 + 1, *(16 + k + 1 for k in range(1, 8))] + assert result.audio_provided[0].tolist() == [False] * 8 + [True] * 8 + + +def test_cancel_restart_and_recycled_slot_match_the_per_request_path() -> None: + steps = [ + [("a0", _append("a", 1), 18), ("b", _append("b", 1), 18)], + [("a0", _append("a", 2), 19), ("b", _append("b", 2), 19)], + # A cancel restarts session a as epoch 1 while epoch 0's aborted request + # still has an append in this step, next to b's live append. + [("a0", _append("a", 3), 20), ("a1", _append("a", 1, epoch=1), 18), ("b", _append("b", 3), 20)], + [("a1", _append("a", 2, epoch=1), 19), ("b", _append("b", 4), 21)], + ] + batched, reference = _talker(max_sessions=2), _talker(max_sessions=2) + + for index, step in enumerate(steps): + text_base = 100 + 10 * index + result = _run_step(batched, step, batched=True, text_base=text_base) + _assert_same(result, _run_step(reference, step, batched=False, text_base=text_base)) + + runtime = batched._personaplex_duplex_stage0_runtime + assert sorted(runtime.sessions) == [("a", 1), ("b", 0)] + assert _device_prepared(batched, ("a", 1)) + # The restarted epoch starts from an empty history on its recycled row. + assert result.embeds["a1"][0][10:17].tolist() == list(SINE_TOKENS[1:]) + + # A new session on the row of a closed one also starts from a clean history. + for talker in (batched, reference): + talker.on_requests_finished({"b"}) + followups = [ + [("c", _append("c", 1), 18)], + [("c", _append("c", 2), 19), ("a1", _append("a", 3, epoch=1), 20)], + ] + for index, step in enumerate(followups): + text_base = 200 + 10 * index + result = _run_step(batched, step, batched=True, text_base=text_base) + _assert_same(result, _run_step(reference, step, batched=False, text_base=text_base)) + assert result.embeds["c"][0][10:17].tolist() == list(SINE_TOKENS[1:]) + + +def test_prepare_live_appends_ignores_empty_first_and_unencoded_appends() -> None: + talker = _talker(max_sessions=2) + runtime = talker._personaplex_duplex_stage0_runtime + + assert runtime.prepare_live_appends([]) == 0 + talker.preprocess_batch(req_ids=["x"], model_intermediate_buffer={"x": {}}, device=torch.device("cpu")) + assert runtime.sessions == {} + + runtime.encode_appends([_append("a", 1)]) + assert runtime.prepare_live_appends([_append("a", 1)]) == 0 # the first append keeps the prefill path + _run_step(talker, [("a", _append("a", 1), 18)], batched=False, text_base=100) + # Not encoded in this step: left to prepare_append. + assert runtime.prepare_live_appends([_append("a", 2)]) == 0 + runtime.encode_appends([_append("a", 2)]) + assert runtime.prepare_live_appends([_append("a", 2), _append("a", 2)]) == 1 + # Already prepared: a repeat in a later step is a no-op. + assert runtime.prepare_live_appends([_append("a", 2)]) == 0 + assert runtime.prepare_append(_append("a", 2), prompt_len=19, request_id="a").prompt_offset == 18 diff --git a/tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py b/tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py index e4ff44bdd..197c92b90 100644 --- a/tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py +++ b/tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py @@ -245,23 +245,23 @@ def test_repeated_append_identity_does_not_advance_codec() -> None: def test_next_append_uses_prior_sample_and_causally_delayed_user_frame() -> None: runtime = _runtime(_FakeCodec()) - runtime.prepare_append(_duplex_info(seq=1), prompt_len=18, request_id="req") - first_agent = torch.arange(8, dtype=torch.long) - runtime.record_sample(request_id="req", text_token=torch.tensor(101), agent_codes=first_agent) - runtime.prepare_append(_duplex_info(seq=2), prompt_len=19, request_id="req") - second_agent = torch.arange(10, 18, dtype=torch.long) - runtime.record_sample(request_id="req", text_token=torch.tensor(102), agent_codes=second_agent) - runtime.prepare_append(_duplex_info(seq=3), prompt_len=20, request_id="req") - silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long) sine = torch.tensor(SINE_TOKENS, dtype=torch.long) + runtime.prepare_append(_duplex_info(seq=1), prompt_len=18, request_id="req") + # The first append teacher-forces agent cb1..7 to silence. + expected_first_effective = torch.cat([torch.tensor([0]), silence[1:]]) + runtime.record_sample(request_id="req", text_token=torch.tensor(101), effective_codes=expected_first_effective) + runtime.prepare_append(_duplex_info(seq=2), prompt_len=19, request_id="req") + second_agent = torch.arange(10, 18, dtype=torch.long) + runtime.record_sample(request_id="req", text_token=torch.tensor(102), effective_codes=second_agent) + runtime.prepare_append(_duplex_info(seq=3), prompt_len=20, request_id="req") + first_call = runtime.stage_model.frame_calls[0] assert all(torch.equal(first_call[key], silence) for key in ("last_agent", "prev_agent")) assert all(torch.equal(first_call[key], sine) for key in ("user_d0", "user_d1")) second_call = runtime.stage_model.frame_calls[1] assert second_call["text_token"].tolist() == [101] - expected_first_effective = torch.cat([first_agent[:1], silence[1:]]) assert torch.equal(second_call["last_agent"], expected_first_effective) assert torch.equal(second_call["prev_agent"], expected_first_effective) assert torch.equal(second_call["user_d0"], torch.full((8,), 1, dtype=torch.long)) @@ -468,7 +468,7 @@ def test_cancel_overlap_full_processing_order_never_re_leases_the_aborted_epoch( else: restarted = runtime.prepare_append(duplex, prompt_len=prompt_len, request_id=request_id) for request_id in ("req-e0", "req-e1"): - runtime.record_sample(request_id=request_id, text_token=5, agent_codes=list(range(8))) + runtime.record_sample(request_id=request_id, text_token=5, effective_codes=list(range(8))) runtime.close_request("req-e0") # the engine's late finish of the aborted request assert list(runtime.sessions) == [("session", 1)] diff --git a/tests/model_executor/models/personaplex/test_mimi_cudagraph.py b/tests/model_executor/models/personaplex/test_mimi_cudagraph.py new file mode 100644 index 000000000..934968466 --- /dev/null +++ b/tests/model_executor/models/personaplex/test_mimi_cudagraph.py @@ -0,0 +1,597 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fixed-buffer streaming state and CUDA graph replay of the PersonaPlex Mimi codec.""" + +import gc +import weakref +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.personaplex.personaplex_code2wav import _MIMI_DECODE_BATCH_FRAMES +from vllm_omni.model_executor.models.personaplex.personaplex_mimi import ( + CODEBOOKS, + FRAME_SIZE, + PersonaPlexMimiCodec, + _MimiStreamingTransformer, + _StreamConv1d, + _StreamConvTr1d, +) + +pytestmark = pytest.mark.core_model + +SEED = 4321 + + +def _mask(*rows: bool, device: torch.device | str = "cpu") -> torch.Tensor: + return torch.tensor(rows, dtype=torch.bool, device=device) + + +def _random_codec(device: torch.device, batch_size: int) -> PersonaPlexMimiCodec: + """The real Mimi graph with random weights (no checkpoint download).""" + from transformers import MimiConfig, MimiModel + + torch.manual_seed(SEED) + codec = PersonaPlexMimiCodec.__new__(PersonaPlexMimiCodec) + nn.Module.__init__(codec) + codec.device = device + model = MimiModel(MimiConfig()) + del model.encoder_transformer, model.decoder_transformer + with torch.no_grad(): + for name, buffer in model.quantizer.named_buffers(): + if name.endswith("embed_sum"): + buffer.normal_() + codec.model = model.to(device).eval() + codec.dtype = torch.float32 + codec.encoder_transformer = _MimiStreamingTransformer().to(device) + codec.decoder_transformer = _MimiStreamingTransformer().to(device) + for transformer in (codec.encoder_transformer, codec.decoder_transformer): + for parameter in transformer.parameters(): + nn.init.normal_(parameter, std=0.02) + codec._init_streaming_stages() + codec.streaming_init(batch_size) + return codec + + +def _state_tensors(codec: PersonaPlexMimiCodec) -> list[torch.Tensor]: + tensors: list[torch.Tensor] = [] + for state in codec._conv_states(): + tensors.append(state.prev if isinstance(state, _StreamConv1d) else state.partial) + tensors.append(state._fresh) + for transformer in (codec.encoder_transformer, codec.decoder_transformer): + tensors.append(transformer._offset) + for kv in transformer._kv: + tensors.extend((kv.cache, kv.end_offset, kv.start_offset)) + return tensors + + +def _snapshot(codec: PersonaPlexMimiCodec) -> list[torch.Tensor]: + return [tensor.clone() for tensor in _state_tensors(codec)] + + +@pytest.mark.cpu +@pytest.mark.parametrize("kind", ["conv", "conv_replicate", "convtr"]) +def test_stream_carries_update_in_place(kind: str) -> None: + torch.manual_seed(SEED) + if kind == "convtr": + stream = _StreamConvTr1d(nn.ConvTranspose1d(2, 3, kernel_size=4, stride=2)) + samples = 3 + else: + pad_mode = "replicate" if kind == "conv_replicate" else "constant" + stream = _StreamConv1d(nn.Conv1d(2, 3, kernel_size=5, stride=2), pad_mode=pad_mode) + samples = 6 + stream.reset(3, torch.device("cpu"), torch.float32) + carry = stream.prev if isinstance(stream, _StreamConv1d) else stream.partial + fresh = stream._fresh + pointers = (carry.data_ptr(), fresh.data_ptr()) + + with torch.no_grad(): + for active in (_mask(True, False, True), _mask(False, True, True), _mask(True, True, True)): + stream(torch.randn(3, 2, samples), active) + current = stream.prev if isinstance(stream, _StreamConv1d) else stream.partial + assert current is carry and stream._fresh is fresh + assert (current.data_ptr(), stream._fresh.data_ptr()) == pointers + assert carry.abs().sum() > 0 + assert not fresh.any() + + stream.reset_slot(1) + assert not carry[1].any() and fresh[1] + assert carry[0].abs().sum() > 0 and not fresh[0] + + stream.reset_all() + assert not carry.any() and fresh.all() + assert (carry.data_ptr(), fresh.data_ptr()) == pointers + + +@pytest.mark.cpu +def test_streaming_init_at_same_batch_size_resets_in_place() -> None: + codec = _random_codec(torch.device("cpu"), batch_size=2) + pcm = torch.randn(2, FRAME_SIZE) + first = codec.encode_frame(pcm) + codec.decode_frame(first) + pointers = [tensor.data_ptr() for tensor in _state_tensors(codec)] + + codec.streaming_init(2) + + assert pointers == [tensor.data_ptr() for tensor in _state_tensors(codec)] + assert torch.equal(codec.encode_frame(pcm), first) + + codec.streaming_init(3) + assert codec._batch_size == 3 + assert codec.encode_frame(torch.randn(3, FRAME_SIZE)).shape == (3, CODEBOOKS) + + +@pytest.mark.cpu +def test_capture_off_cuda_stays_eager() -> None: + codec = _random_codec(torch.device("cpu"), batch_size=2) + assert codec.capture_cuda_graphs() == [] + assert codec._cuda_graphs == {} + assert codec.encode_frame(torch.randn(2, FRAME_SIZE)).shape == (2, CODEBOOKS) + + +@pytest.mark.cpu +def test_quantizer_decode_matches_transformers() -> None: + codec = _random_codec(torch.device("cpu"), batch_size=2) + codes = torch.randint(0, 2048, (2, CODEBOOKS, 3)) + with torch.no_grad(): + expected = codec.model.quantizer.decode(codes) + actual = codec._quantizer_decode(codes) + assert torch.equal(actual, expected) + + +def _drive(codec: PersonaPlexMimiCodec, frames: int, recycle_at: int, device: torch.device): + """Encode/decode ``frames`` frames under a mixed active schedule with a row recycle.""" + batch_size = codec._batch_size + generator = torch.Generator(device="cpu").manual_seed(SEED) + schedule = [_mask(*[(step + row) % 3 != 0 for row in range(batch_size)], device=device) for step in range(frames)] + codes_out, pcm_out, inactive_ok = [], [], True + for step, active in enumerate(schedule): + if step == recycle_at: + codec.reset_slot(1) + pcm = 0.1 * torch.randn(batch_size, FRAME_SIZE, generator=generator) + before = _snapshot(codec) + codes = codec.encode_frame(pcm, active) + wav = codec.decode_frame(codes, active) + after = _state_tensors(codec) + for row, is_active in enumerate(active.tolist()): + if is_active: + continue + for old, new in zip(before, after): + if new.dim() > 0 and new.shape[0] == batch_size: + inactive_ok &= torch.equal(old[row], new[row]) + elif new.dim() > 1 and new.shape[1] == batch_size: + inactive_ok &= torch.equal(old[:, row], new[:, row]) + codes_out.append(codes) + pcm_out.append(wav) + return torch.stack(codes_out), torch.stack(pcm_out), inactive_ok + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_graph_replay_is_bitwise_equal_to_eager_across_recycle() -> None: + device = torch.device("cuda") + codec = _random_codec(device, batch_size=3) + eager_codes, eager_pcm, eager_inactive_ok = _drive(codec, frames=12, recycle_at=7, device=device) + assert eager_inactive_ok + assert torch.isfinite(eager_pcm).all() + + codec.streaming_init(3) + pointers = [tensor.data_ptr() for tensor in _state_tensors(codec)] + assert codec.capture_cuda_graphs(decode_frame_counts=(1, 3)) == ["decode_f1", "decode_f3", "encode"] + assert pointers == [tensor.data_ptr() for tensor in _state_tensors(codec)] + + graph_codes, graph_pcm, graph_inactive_ok = _drive(codec, frames=12, recycle_at=7, device=device) + + assert codec._cuda_graphs["encode"].replays == 12 + assert codec._cuda_graphs["decode_f1"].replays == 12 + assert graph_inactive_ok + assert torch.equal(graph_codes, eager_codes) + assert torch.equal(graph_pcm, eager_pcm) + + # A full reset in place (a new stream on the same codec) replays the same + # frames from a fresh state. + codec.streaming_init(3) + again_codes, again_pcm, _ = _drive(codec, frames=12, recycle_at=7, device=device) + assert torch.equal(again_codes, eager_codes) + assert torch.equal(again_pcm, eager_pcm) + assert pointers == [tensor.data_ptr() for tensor in _state_tensors(codec)] + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("batch_size", [1, 3]) +def test_multi_frame_decode_graph_matches_eager_chunks(batch_size: int) -> None: + # Code2Wav records the first single-frame delta and the full chunk, on one + # B=1 decoder per session. B=3 runs the same graphs on shared rows, where + # every call carries all rows and only some are active. + chunk = _MIMI_DECODE_BATCH_FRAMES + partial = chunk - 2 + assert partial not in (1, chunk) + device = torch.device("cuda") + codec = _random_codec(device, batch_size=batch_size) + generator = torch.Generator(device="cpu").manual_seed(SEED) + # First frame alone, full chunks, a partial chunk with no graph, then a + # single frame and a full chunk again on the state the eager step left. + frame_counts = [1, chunk, chunk, partial, 1, chunk] + row_masks = [(True, True, False), (True, False, True), (True, True, True), (False, True, True)] + steps = [ + (frames, None if batch_size == 1 else _mask(*row_masks[step % len(row_masks)], device=device)) + for step, frames in enumerate(frame_counts) + ] + chunks = [torch.randint(0, 2048, (batch_size, CODEBOOKS, frames), generator=generator) for frames, _ in steps] + + eager = [codec.decode_frames(chunk_codes, active) for chunk_codes, (_, active) in zip(chunks, steps)] + codec.streaming_init(batch_size) + assert codec.capture_cuda_graphs(encode=False, decode_frame_counts=(1, chunk)) == ["decode_f1", f"decode_f{chunk}"] + graphed = [codec.decode_frames(chunk_codes, active) for chunk_codes, (_, active) in zip(chunks, steps)] + + # The partial chunk has no graph, so it is the one step that ran eagerly. + assert codec._cuda_graphs["decode_f1"].replays == 2 + assert codec._cuda_graphs[f"decode_f{chunk}"].replays == 3 + assert f"decode_f{partial}" not in codec._cuda_graphs + for (frames, active), expected, actual in zip(steps, eager, graphed): + assert actual.shape == (batch_size, frames * FRAME_SIZE) + rows = slice(None) if active is None else active + assert torch.equal(actual[rows], expected[rows]) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_failed_capture_warns_with_the_error_and_stays_eager(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.model_executor.models.personaplex import personaplex_mimi_cudagraph + + device = torch.device("cuda") + codec = _random_codec(device, batch_size=2) + pcm = 0.1 * torch.randn(2, FRAME_SIZE, generator=torch.Generator(device="cpu").manual_seed(SEED)) + expected_codes = codec.encode_frame(pcm) + expected_pcm = codec.decode_frame(expected_codes) + codec.streaming_init(2) + + class _FailingCapture: + def __enter__(self) -> None: + pass + + def __exit__(self, *_args) -> None: + raise RuntimeError("CUDA error: out of memory") + + warnings: list[tuple[tuple, dict]] = [] + monkeypatch.setattr(torch.cuda, "graph", lambda *_args, **_kwargs: _FailingCapture()) + monkeypatch.setattr( + personaplex_mimi_cudagraph.logger, + "warning", + lambda *args, **kwargs: warnings.append((args, kwargs)), + ) + + assert codec.capture_cuda_graphs(encode=True, decode_frame_counts=(1,)) == [] + + assert codec._cuda_graphs == {} + assert len(warnings) == 1 + args, kwargs = warnings[0] + assert "requested" in args[0] and "capture failed" in args[0] and "eagerly" in args[0] + assert args[1:3] == ("encode/decode_f1", 2) + assert "Traceback (most recent call last):" in args[3] + assert "RuntimeError: CUDA error: out of memory" in args[3] + assert not kwargs + # The warmup frames are rolled back, so the eager codec starts a fresh stream. + assert torch.equal(codec.encode_frame(pcm), expected_codes) + assert torch.equal(codec.decode_frame(expected_codes), expected_pcm) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_capture_does_not_hide_errors_that_are_not_cuda_failures(monkeypatch: pytest.MonkeyPatch) -> None: + codec = _random_codec(torch.device("cuda"), batch_size=2) + + def _broken(*_args, **_kwargs): + raise TypeError("not a capture failure") + + monkeypatch.setattr(codec, "_encode_frame_eager", _broken) + with pytest.raises(TypeError, match="not a capture failure"): + codec.capture_cuda_graphs(encode=True, decode_frame_counts=()) + assert codec._cuda_graphs == {} + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_graph_capture_leaves_a_fresh_stream_and_resize_drops_graphs() -> None: + device = torch.device("cuda") + codec = _random_codec(device, batch_size=2) + codec.capture_cuda_graphs(encode=True, decode_frame_counts=()) + assert list(codec._cuda_graphs) == ["encode"] + for state in codec._conv_states(): + carry = state.prev if isinstance(state, _StreamConv1d) else state.partial + assert not carry.any() and state._fresh.all() + for transformer in (codec.encoder_transformer, codec.decoder_transformer): + assert not transformer._offset.any() + assert all(not kv.end_offset.any() for kv in transformer._kv) + + # Resizing reallocates the state, so graphs recorded at the old size are dropped. + codec.streaming_init(4) + assert codec._cuda_graphs == {} + assert codec.encode_frame(torch.zeros(4, FRAME_SIZE)).shape == (4, CODEBOOKS) + + +class _RecordingCodec: + """Stage 0 shared-encoder stand-in that records the graph capture request.""" + + def __init__(self) -> None: + self.calls: list[tuple[str, object]] = [] + + def streaming_init(self, batch_size: int) -> None: + self.calls.append(("streaming_init", batch_size)) + + def capture_cuda_graphs(self, *, encode: bool, decode_frame_counts: tuple[int, ...]) -> list[str]: + self.calls.append(("capture", {"encode": encode, "decode_frame_counts": decode_frame_counts})) + return ["encode"] + + +@pytest.mark.cpu +@pytest.mark.parametrize("enabled", [True, False]) +def test_stage0_captures_the_shared_encoder_once_when_enabled(enabled: bool) -> None: + from vllm_omni.model_executor.models.personaplex.duplex.stage0 import PersonaPlexStage0DuplexRuntime + + codec = _RecordingCodec() + runtime = PersonaPlexStage0DuplexRuntime( + object(), + model_path="/unused", + device="cpu", + codec_factory=lambda: codec, + max_sessions=4, + codec_cuda_graphs=enabled, + ) + + assert runtime._shared_codec() is codec + assert runtime._shared_codec() is codec + expected: list[tuple[str, object]] = [("streaming_init", 4)] + if enabled: + expected.append(("capture", {"encode": True, "decode_frame_counts": ()})) + assert codec.calls == expected + + +@pytest.mark.cpu +def test_resize_after_capture_warns_that_graphs_are_dropped(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.model_executor.models.personaplex import personaplex_mimi + + warnings: list[tuple] = [] + monkeypatch.setattr(personaplex_mimi.logger, "warning", lambda *args: warnings.append(args)) + codec = _random_codec(torch.device("cpu"), batch_size=2) + codec._cuda_graphs = {"encode": object()} + + codec.streaming_init(2) + assert "encode" in codec._cuda_graphs and not warnings + + codec.streaming_init(3) + assert codec._cuda_graphs == {} + assert len(warnings) == 1 and warnings[0][1:] == (2, 3, "encode") + + +@pytest.fixture +def capture_faults(monkeypatch: pytest.MonkeyPatch): + """Exercise capture orchestration and error boundaries with CPU tensors.""" + from vllm_omni.model_executor.models.personaplex import personaplex_mimi_cudagraph as graphs + + state = SimpleNamespace( + errors={}, + counts={}, + events=[], + warnings=[], + phase="setup", + capturing=False, + graph_refs=[], + capture_refs=[], + output_refs=[], + check_released=False, + ) + state.original_stream = state.stream = object() + + def record(event): + state.events.append(event) + state.counts[event] = state.counts.get(event, 0) + 1 + error = state.errors.get((event, state.counts[event]), state.errors.get(event)) + if error is not None: + raise error + + def eager(name): + def step(frame, active): + state.phase = "capture" if state.capturing else "warmup" + record(f"{name}_{state.phase}") + if state.capturing: + output = frame.clone() + state.output_refs.append(weakref.ref(output)) + return output + return frame + + return step + + def reset(): + assert state.stream is state.original_stream + if state.check_released: + assert all(ref() is None for ref in state.graph_refs + state.capture_refs + state.output_refs) + state.phase = "reset" + record("reset") + + codec = SimpleNamespace( + device="cuda", + dtype=torch.float32, + _batch_size=2, + _encode_frame_eager=eager("encode"), + _decode_frames_eager=eager("decode"), + reset_streaming=reset, + ) + + class Graph: + def __init__(self): + record("construct") + state.graph_refs.append(weakref.ref(self)) + + class Capture: + def __init__(self, graph): + self.graph = graph + state.capture_refs.append(weakref.ref(self)) + + def __enter__(self): + record("capture_setup_sync") + record("capture_setup_empty_cache") + state.stream = object() + state.capturing = True + record("capture_begin") + + def __exit__(self, exc_type, exc, traceback): + state.capturing = False + # Match torch.cuda.graph: capture_end can raise before restoration, + # and can replace a body exception with capture-invalidated. + record("capture_end") + state.stream = state.original_stream + + @contextmanager + def stream_context(stream): + previous = state.stream + state.stream = stream + try: + yield + finally: + state.stream = previous + + zeros, ones = torch.zeros, torch.ones + monkeypatch.setattr(torch, "zeros", lambda *args, **kwargs: zeros(*args, **(kwargs | {"device": "cpu"}))) + monkeypatch.setattr(torch, "ones", lambda *args, **kwargs: ones(*args, **(kwargs | {"device": "cpu"}))) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "graph_pool_handle", lambda: (0, 0)) + monkeypatch.setattr(torch.cuda, "CUDAGraph", Graph) + monkeypatch.setattr(torch.cuda, "graph", lambda graph, **kwargs: Capture(graph)) + monkeypatch.setattr(torch.cuda, "current_stream", lambda device: state.stream) + monkeypatch.setattr(torch.cuda, "stream", stream_context) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda device: record(f"sync_{state.phase}")) + monkeypatch.setattr(graphs.logger, "warning", lambda *args, **kwargs: state.warnings.append((args, kwargs))) + state.capture = lambda: graphs.capture_mimi_frame_graphs( + codec, + encode=True, + decode_frame_counts=(1,), + warmup_iters=1, + ) + return state + + +@pytest.mark.cpu +@pytest.mark.parametrize("step", ["encode_warmup", "decode_warmup"]) +@pytest.mark.parametrize("error_type", [RuntimeError, torch.cuda.OutOfMemoryError]) +def test_eager_warmup_failures_propagate_before_capture(capture_faults, step, error_type): + error = error_type("injected eager tensor failure") + capture_faults.errors[step] = error + with pytest.raises(error_type) as raised: + capture_faults.capture() + assert raised.value is error + assert capture_faults.counts.get("construct", 0) == 0 + assert capture_faults.counts["reset"] == 1 + assert not capture_faults.warnings + + +@pytest.mark.cpu +@pytest.mark.parametrize( + "point", + [ + "sync_warmup", + "construct", + "capture_setup_sync", + "capture_setup_empty_cache", + "capture_begin", + "sync_capture", + "sync_reset", + "reset", + ], +) +def test_device_and_reset_failures_are_never_capture_fallbacks(capture_faults, point): + # Even an otherwise recoverable error class must propagate at these boundaries. + error = torch.cuda.OutOfMemoryError("injected device or reset failure") + capture_faults.errors[point] = error + with pytest.raises(torch.cuda.OutOfMemoryError) as raised: + capture_faults.capture() + assert raised.value is error + assert not capture_faults.warnings + + +@pytest.mark.cpu +@pytest.mark.parametrize( + "message", + [ + "tensor size mismatch", + "CUDA error: an illegal memory access was encountered", + "CUDA error: device-side assert triggered", + "CUDA error: operation failed due to a previous error during capture", + ], +) +def test_unknown_or_fatal_capture_errors_propagate(capture_faults, message): + error = RuntimeError(message) + capture_faults.errors["encode_capture"] = error + with pytest.raises(RuntimeError) as raised: + capture_faults.capture() + assert raised.value is error + assert capture_faults.stream is capture_faults.original_stream + assert capture_faults.counts["reset"] == 1 + assert not capture_faults.warnings + + +@pytest.mark.cpu +def test_recoverable_capture_failure_discards_partial_graphs_after_all_warmups(capture_faults): + error = torch.cuda.OutOfMemoryError("capture allocation failed") + capture_faults.errors[("capture_end", 2)] = error + capture_faults.check_released = True + was_enabled = gc.isenabled() + gc.disable() + try: + assert capture_faults.capture() == {} + finally: + if was_enabled: + gc.enable() + assert capture_faults.events.index("decode_warmup") < capture_faults.events.index("construct") + assert capture_faults.counts["capture_end"] == 2 # an earlier graph completed + assert capture_faults.counts["sync_reset"] == 1 + assert capture_faults.stream is capture_faults.original_stream + assert len(capture_faults.warnings) == 1 + assert "capture allocation failed" in capture_faults.warnings[0][0][-1] + assert "Traceback (most recent call last):" in capture_faults.warnings[0][0][-1] + assert not capture_faults.warnings[0][1] + assert error.__traceback__ is None + + +@pytest.mark.cpu +@pytest.mark.parametrize("recoverable", [True, False]) +def test_capture_end_does_not_hide_the_original_failure(capture_faults, recoverable): + message = "CUDA error: operation not permitted when stream is capturing" if recoverable else "tensor size mismatch" + original = RuntimeError(message) + capture_faults.errors["encode_capture"] = original + invalidated = RuntimeError("CUDA error: operation failed due to a previous error during capture") + capture_faults.errors["capture_end"] = invalidated + if recoverable: + capture_faults.check_released = True + was_enabled = gc.isenabled() + gc.disable() + try: + assert capture_faults.capture() == {} + finally: + if was_enabled: + gc.enable() + assert len(capture_faults.warnings) == 1 + assert original.__traceback__ is None and invalidated.__traceback__ is None + else: + with pytest.raises(RuntimeError) as raised: + capture_faults.capture() + assert raised.value is invalidated and raised.value.__context__ is original + assert not capture_faults.warnings + assert capture_faults.stream is capture_faults.original_stream + assert capture_faults.counts["reset"] == 1 + + +@pytest.mark.cpu +@pytest.mark.parametrize("point", ["reset", "sync_reset"]) +def test_capture_fallback_requires_successful_reset(capture_faults, point): + capture_faults.errors["capture_end"] = torch.cuda.OutOfMemoryError("capture allocation failed") + error = RuntimeError("codec could not reset") + capture_faults.errors[point] = error + with pytest.raises(RuntimeError) as raised: + capture_faults.capture() + assert raised.value is error + assert not capture_faults.warnings diff --git a/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py b/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py index c60b5c197..2e410852d 100644 --- a/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py +++ b/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py @@ -122,6 +122,8 @@ def _make_codec_stub(mocker: MockerFixture) -> PersonaPlexMimiCodec: num_quantizers, x.shape[0], x.shape[-1], dtype=torch.long ) quantizer.decode.side_effect = lambda codes: torch.zeros(codes.shape[0], 1, codes.shape[-1]) + # The codec decodes through its own RVQ sum; route it to the stub. + codec._quantizer_decode = quantizer.decode codec.device = torch.device("cpu") codec.dtype = torch.float32 diff --git a/tests/model_executor/stage_input_processors/test_personaplex.py b/tests/model_executor/stage_input_processors/test_personaplex.py index 6521fdd96..5407b57f5 100644 --- a/tests/model_executor/stage_input_processors/test_personaplex.py +++ b/tests/model_executor/stage_input_processors/test_personaplex.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace @@ -100,16 +100,18 @@ def test_post_sample_talker_mtp_uses_current_temporal_state() -> None: received["audio_tokens"] = audio_tokens received["audio_provided"] = audio_provided received["num_steps"] = num_steps - return torch.arange(num_steps or 16, dtype=torch.long).reshape(1, -1) + return 100 + torch.arange(num_steps or 16, dtype=torch.long).reshape(1, -1) model = SimpleNamespace( _dtype=torch.float32, num_active_codebooks=8, depformer=depformer, + _depformer_teacher_forcing=PersonaPlexTalkerForConditionalGeneration._depformer_teacher_forcing, _duplex_stage0_runtime=lambda: SimpleNamespace( - record_sample=lambda *, request_id, text_token, agent_codes: recorded.append( - (request_id, text_token.clone(), agent_codes.clone()) - ) + prepared_depformer_state=lambda _request_id: None, + record_sample=lambda *, request_id, text_token, effective_codes: recorded.append( + (request_id, text_token.clone(), effective_codes.clone()) + ), ), ) method = getattr(PersonaPlexTalkerForConditionalGeneration, "post_sample_talker_mtp", None) @@ -137,4 +139,6 @@ def test_post_sample_talker_mtp_uses_current_temporal_state() -> None: assert torch.equal(received["audio_tokens"], torch.arange(16).reshape(1, 16)) assert torch.equal(received["audio_provided"], (torch.arange(16) > 0).reshape(1, 16)) assert [(row[0], row[1].item()) for row in recorded] == [("r1", 101)] - assert torch.equal(recorded[0][2], codes[0]) + # Teacher-forced agent codebooks (1..7 here) replace the sampled ones. + assert codes[0].tolist() == list(range(100, 108)) + assert recorded[0][2].tolist() == [100, 1, 2, 3, 4, 5, 6, 7] diff --git a/vllm_omni/deploy/personaplex.yaml b/vllm_omni/deploy/personaplex.yaml index 567d0fd69..ec9084bcc 100644 --- a/vllm_omni/deploy/personaplex.yaml +++ b/vllm_omni/deploy/personaplex.yaml @@ -10,12 +10,19 @@ # first-chunk latency (not correctness). # * Stage 1 (code2wav) wraps the external Mimi codec. Each codec frame is # 8 active codebooks (cb 0..7) -> 1920 PCM samples @ 24 kHz / 12.5 Hz. -# * Stage 0 runs cudagraph by default; Stage 1 (Mimi) is eager because the -# external decoder performs host-to-device copies that cannot be captured. +# * Stage 0 runs cudagraph by default; Stage 1 (Mimi) runs with +# enforce_eager because the external decoder performs host-to-device copies +# that the runner cannot capture. The streaming Mimi encode (Stage 0) and +# decode (Stage 1) frame steps instead replay model-local CUDA graphs +# (hf_overrides.mimi_cuda_graphs), bitwise equal to eager. The graphs add +# about 80 MiB on Stage 0 and 5 MiB per decoder slot on Stage 1. # * Each live session holds one 3000-frame KV window (~1.5 GiB) on Stage 0. # To serve more sessions, raise duplex_session.max_sessions and both stages' -# max_num_seqs together, and size Stage 0 gpu_memory_utilization for the KV; -# on a 141 GB card, 16 sessions at Stage 0 utilization 0.75 stay realtime. +# max_num_seqs together, and size Stage 0 gpu_memory_utilization for +# N x 1.5 GiB of KV: on a 141 GB card, utilization 0.75 holds about 60 full +# windows. Measured there (Stage 1 at 0.15), 64 concurrent 30 s sessions +# (375 frames each) stayed realtime. Longer sessions also cost more per +# step as their context fills. pipeline: personaplex session_mode: duplex duplex_session: @@ -52,6 +59,8 @@ stages: # Temporal transformer sliding_window is 3000 frames; keep model len within. max_model_len: 3000 devices: "0" + hf_overrides: + mimi_cuda_graphs: true output_connectors: to_stage_1: connector_of_shared_memory default_sampling_params: @@ -77,6 +86,8 @@ stages: max_num_batched_tokens: 65536 max_model_len: 65536 devices: "0" + hf_overrides: + mimi_cuda_graphs: true input_connectors: from_stage_0: connector_of_shared_memory default_sampling_params: diff --git a/vllm_omni/model_executor/models/personaplex/configuration_personaplex.py b/vllm_omni/model_executor/models/personaplex/configuration_personaplex.py index e43e12168..e16e9165e 100644 --- a/vllm_omni/model_executor/models/personaplex/configuration_personaplex.py +++ b/vllm_omni/model_executor/models/personaplex/configuration_personaplex.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Configuration for PersonaPlex (a Moshi finetune; 2-stage audio->audio pipeline). PersonaPlex is a staged AR speech model composed of: @@ -203,6 +203,10 @@ class PersonaPlexConfig(PretrainedConfig): mimi_name (`str`, *optional*): Convenience mirror of ``mimi_config.mimi_name``; if set, it overrides the value carried inside ``mimi_config``. + mimi_cuda_graphs (`bool`, *optional*, defaults to `False`): + Replay the streaming Mimi per-frame encode (Stage 0) and decode + (Stage 1) from model-local CUDA graphs. Can be enabled through + ``hf_overrides``. """ model_type = "personaplex" @@ -222,6 +226,7 @@ class PersonaPlexConfig(PretrainedConfig): audio_vocab_size: int = 2048, num_audio_codebooks: int = 16, mimi_name: str | None = None, + mimi_cuda_graphs: bool = False, **kwargs: Any, ) -> None: if temporal_config is None: @@ -257,6 +262,7 @@ class PersonaPlexConfig(PretrainedConfig): if mimi_name is not None: self.mimi_config.mimi_name = mimi_name self.mimi_name = self.mimi_config.mimi_name + self.mimi_cuda_graphs = mimi_cuda_graphs @staticmethod def _coerce(value: Any, config_cls: type[PretrainedConfig]) -> PretrainedConfig: diff --git a/vllm_omni/model_executor/models/personaplex/duplex/stage0.py b/vllm_omni/model_executor/models/personaplex/duplex/stage0.py index dfeca1878..895e3e644 100644 --- a/vllm_omni/model_executor/models/personaplex/duplex/stage0.py +++ b/vllm_omni/model_executor/models/personaplex/duplex/stage0.py @@ -9,9 +9,10 @@ from collections.abc import Callable from dataclasses import dataclass, field from functools import lru_cache from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np +from vllm.logger import init_logger from vllm_omni.model_executor.common.audio.pcm import pcm_f32le_samples from vllm_omni.model_executor.common.duplex.payload import decode_pcm_f32le_payload @@ -24,8 +25,13 @@ from vllm_omni.model_executor.models.personaplex.duplex.policy import ( wrap_with_system_tags, ) +if TYPE_CHECKING: + import torch + _FRAME_SAMPLES = FRAME_SIZE +logger = init_logger(__name__) + class PersonaPlexStage0CapacityError(RuntimeError): """Every streaming encoder row is leased by a live session.""" @@ -43,6 +49,9 @@ class PersonaPlexStage0PreparedAppend: info_update: dict[str, Any] prefill_applied: bool prompt_offset: int + # Device copies of the depformer teacher-forcing state (batched path only). + depformer_audio_tokens_device: torch.Tensor | None = None + depformer_audio_provided_device: torch.Tensor | None = None @dataclass(slots=True) @@ -63,6 +72,10 @@ class PersonaPlexStage0SessionState: slot: int | None = None encoded_identity: tuple[int, int] | None = None encoded_frame: Any | None = None + # Device copies of the encoded frame and of the last two consumed user + # frames (newest last), so live appends are built without host round trips. + encoded_frame_device: torch.Tensor | None = None + user_history_device: list[torch.Tensor] = field(default_factory=list) def _tokenizer_path(model_path: str) -> Path: @@ -151,6 +164,7 @@ class PersonaPlexStage0DuplexRuntime: codec: Any | None = None, codec_factory: Callable[[], Any] | None = None, max_sessions: int = 1, + codec_cuda_graphs: bool = False, tokenizer=None, voice_loader=None, ) -> None: @@ -161,6 +175,7 @@ class PersonaPlexStage0DuplexRuntime: self.device = device self.max_sessions = max_sessions self._codec_factory = codec_factory + self._codec_cuda_graphs = codec_cuda_graphs self._codec: Any | None = None self._free_slots: list[int] = list(reversed(range(max_sessions))) self._tokenizer = tokenizer @@ -170,6 +185,7 @@ class PersonaPlexStage0DuplexRuntime: # Requests of a superseded epoch that still reached this step; they are # finished by the engine and must not lease a row or record a sample. self._stale_requests: set[str] = set() + self._constants_cache: dict[str, torch.Tensor] | None = None if codec is not None: codec.streaming_init(max_sessions) self._codec = codec @@ -239,13 +255,22 @@ class PersonaPlexStage0DuplexRuntime: state.request_ids.add(request_id) self.request_sessions[request_id] = key if state.prepared_identity == identity and state.prepared is not None: - return state.prepared + prepared = state.prepared + if prepared.prompt_offset < 0: + prepared.prompt_offset = int(prompt_len) - int(prepared.inputs_embeds.shape[0]) + if prepared.prompt_offset < 0: + raise ValueError( + "PersonaPlex scheduler prompt reservation mismatch: " + f"reserved={prompt_len}, prepared={prepared.inputs_embeds.shape[0]}" + ) + return prepared if seq <= state.last_seq: raise ValueError(f"PersonaPlex duplex append seq must increase: last={state.last_seq}, got={seq}") if state.encoded_identity != identity: self._encode_rows([(state, identity, self._decode_pcm(duplex.get("payload")))]) user_frame = state.encoded_frame + self._push_user_history(state) state.encoded_identity = None state.encoded_frame = None state.user_codes = user_frame if state.user_codes is None else torch.cat([state.user_codes, user_frame], dim=0) @@ -344,44 +369,188 @@ class PersonaPlexStage0DuplexRuntime: f"PersonaPlex scheduler prompt reservation mismatch: reserved={prompt_len}, prepared={prepared_len}" ) input_ids = torch.zeros((prepared_len,), dtype=torch.long, device=device) - info_update = { + info_update = self._prepared_info_update( + state, seq, silence=silence.detach().cpu(), prefill_applied=first_append + ) + # This path has no device teacher-forcing copies, so the talker reads + # these host tensors from the request information. + info_update["pplex_depformer_audio_tokens"] = depformer_audio_tokens.detach().cpu() + info_update["pplex_depformer_audio_provided"] = depformer_audio_provided.detach().cpu() + return self._store_prepared( + state, + seq, + PersonaPlexStage0PreparedAppend( + input_ids=input_ids, + inputs_embeds=full_embeds, + user_codes=state.user_codes, + info_update=info_update, + prefill_applied=first_append, + prompt_offset=prompt_offset, + ), + ) + + @staticmethod + def _prepared_info_update( + state: PersonaPlexStage0SessionState, + seq: int, + *, + silence: torch.Tensor, + prefill_applied: bool, + ) -> dict[str, Any]: + """Request-information update shared by the per-request and batched append paths.""" + return { "pplex_user_codes": state.user_codes, - "pplex_silence_codes": silence.detach().cpu(), - "pplex_depformer_audio_tokens": depformer_audio_tokens.detach().cpu(), - "pplex_depformer_audio_provided": depformer_audio_provided.detach().cpu(), + "pplex_silence_codes": silence, "meta": { "pplex_frame": state.prefill_slots + int(state.user_codes.shape[0]), "pplex_prefill_len": state.prefill_slots, }, "duplex": { "stage0_prepared": True, - "prefill_applied": first_append, - "session_id": session_id, - "epoch": epoch, + "prefill_applied": prefill_applied, + "session_id": state.session_id, + "epoch": state.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 + + @staticmethod + def _store_prepared( + state: PersonaPlexStage0SessionState, + seq: int, + prepared: PersonaPlexStage0PreparedAppend, + ) -> PersonaPlexStage0PreparedAppend: + """Record ``prepared`` as the session's current append and advance its sequence.""" + state.prepared_identity = (state.epoch, seq) state.prepared = prepared state.last_seq = seq return prepared + @staticmethod + def _push_user_history(state: PersonaPlexStage0SessionState) -> None: + """Move the consumed encoded frame into the device history (last two frames).""" + if state.encoded_frame_device is None: + raise RuntimeError("PersonaPlex Stage 0 append has no encoded user frame") + state.user_history_device.append(state.encoded_frame_device) + del state.user_history_device[:-2] + state.encoded_frame_device = None + + def _constants(self) -> dict[str, torch.Tensor]: + """Live-append constant frames on the device (plus host silence), built once.""" + if self._constants_cache is not None: + return self._constants_cache + import torch + + device, _ = self._model_device_dtype() + silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long) + self._constants_cache = { + "silence": silence, + "silence_device": silence.to(device), + "sine_device": torch.tensor(SINE_TOKENS, dtype=torch.long, device=device), + "live_provided_device": torch.tensor([False] * 8 + [True] * 8, dtype=torch.bool, device=device), + "zero_text_device": torch.tensor([ZERO_TEXT_TOKEN], dtype=torch.long, device=device), + } + return self._constants_cache + + def prepare_live_appends(self, appends: list[dict[str, Any]]) -> int: + """Prepare every one-frame live append of this step with one batched embedding call. + + Only appends whose frame was encoded this step and that are not the + first append of their epoch (which carries the voice/persona prefill) + take this path; the rest are left to :meth:`prepare_append`. Results are + stored exactly as :meth:`prepare_append` stores them, so its per-request + call returns the cached preparation. Returns the number prepared. + """ + import torch + + batch: dict[tuple[str, int], tuple[PersonaPlexStage0SessionState, int]] = {} + for duplex in appends: + try: + session_id, epoch, seq = _append_identity(duplex) + except ValueError: + continue + key = (session_id, epoch) + state = self.sessions.get(key) + # Only an append encoded in this step and not yet consumed: its seq + # is past last_seq, so it has not been prepared either. + if state is None or state.last_seq == 0 or seq <= state.last_seq or key in batch: + continue + if state.encoded_identity != (epoch, seq): + continue + batch[key] = (state, seq) + if not batch: + return 0 + c = self._constants() + device, dtype = self._model_device_dtype() + sine_d, silence_d = c["sine_device"], c["silence_device"] + texts, agents, d0s, d1s, currents = [], [], [], [], [] + for state, _ in batch.values(): + history = state.user_history_device # consumed frames, newest last + texts.append( + state.last_text_token.reshape(1) if state.last_text_token is not None else c["zero_text_device"] + ) + agents.append(state.last_agent_codes[:8] if state.last_agent_codes is not None else silence_d) + # The frame being appended is not in the history yet: user cb0 reads + # the previous frame and user cb1..7 the one before it. + d0s.append(history[-1] if len(history) >= 1 else sine_d) + d1s.append(history[-2] if len(history) >= 2 else sine_d) + currents.append(state.encoded_frame_device) + n = len(batch) + text = torch.stack(texts) + agent = torch.stack(agents) + user_d0 = torch.stack(d0s) + user_d1 = torch.stack(d1s) + current = torch.stack(currents) + stack = torch.cat([text, agent, user_d0[:, :1], user_d1[:, 1:8]], dim=1).unsqueeze(-1) + embeds = self.stage_model.input_embeddings(stack).reshape(n, 1, -1).to(dtype=dtype) + dep_tokens_device = torch.cat([silence_d.expand(n, 8), current[:, :1], user_d0[:, 1:8]], dim=1) + input_ids = torch.zeros((n, 1), dtype=torch.long, device=device) + for row, (state, seq) in enumerate(batch.values()): + user_frame = state.encoded_frame + self._push_user_history(state) + state.encoded_identity = None + state.encoded_frame = None + state.user_codes = torch.cat([state.user_codes, user_frame], dim=0) + # The talker reads teacher forcing from the device copies below, so + # no host tokens or mask are built for these appends. + self._store_prepared( + state, + seq, + PersonaPlexStage0PreparedAppend( + input_ids=input_ids[row], + inputs_embeds=embeds[row], + user_codes=state.user_codes, + info_update=self._prepared_info_update(state, seq, silence=c["silence"], prefill_applied=False), + prefill_applied=False, + # Resolved by prepare_append from the scheduler's prompt length. + prompt_offset=-1, + depformer_audio_tokens_device=dep_tokens_device[row], + depformer_audio_provided_device=c["live_provided_device"], + ), + ) + return n + + def prepared_depformer_state(self, request_id: str) -> tuple[torch.Tensor, torch.Tensor] | None: + """Device teacher-forcing tensors of the request's current append, when prepared on the device.""" + key = self.request_sessions.get(request_id) + state = self.sessions.get(key) if key is not None else None + prepared = state.prepared if state is not None else None + if prepared is None or prepared.depformer_audio_tokens_device is None: + return None + return prepared.depformer_audio_tokens_device, prepared.depformer_audio_provided_device + def record_sample( self, *, request_id: str, text_token: Any, - agent_codes: Any, + effective_codes: Any, ) -> None: - """Commit one sampled temporal frame for the next live append.""" + """Commit one sampled temporal frame for the next live append. + + ``effective_codes`` are the frame's eight agent codes with the append's + depformer teacher forcing already applied, which the next append reads. + """ import torch if request_id in self._stale_requests: @@ -396,26 +565,12 @@ class PersonaPlexStage0DuplexRuntime: return text = torch.as_tensor(text_token, device=self._model_device_dtype()[0], dtype=torch.long).reshape(-1) - codes = torch.as_tensor(agent_codes, device=text.device, dtype=torch.long).reshape(-1) + codes = torch.as_tensor(effective_codes, device=text.device, dtype=torch.long).reshape(-1) if text.numel() != 1: raise ValueError(f"PersonaPlex Stage 0 expected one sampled text token, got {text.numel()}") - if codes.numel() < 8: - raise ValueError(f"PersonaPlex Stage 0 expected at least 8 agent codes, got {codes.numel()}") - - effective_codes = codes[:8].clone() - prepared = state.prepared - if prepared is None: - raise RuntimeError(f"PersonaPlex Stage 0 request has no prepared payload: {request_id}") - target = prepared.info_update.get("pplex_depformer_audio_tokens") - provided = prepared.info_update.get("pplex_depformer_audio_provided") - if not isinstance(target, torch.Tensor) or not isinstance(provided, torch.Tensor): - raise RuntimeError("PersonaPlex Stage 0 prepared payload has no depformer teacher-forcing state") - target = target.reshape(-1).to(device=effective_codes.device, dtype=torch.long) - provided = provided.reshape(-1).to(device=effective_codes.device, dtype=torch.bool) - if target.numel() < 8 or provided.numel() < 8: - raise RuntimeError("PersonaPlex Stage 0 depformer teacher-forcing state is shorter than one agent frame") - effective_codes = torch.where(provided[:8], target[:8], effective_codes) - state.last_agent_codes = effective_codes.detach() + if codes.numel() != 8: + raise ValueError(f"PersonaPlex Stage 0 expected 8 effective agent codes, got {codes.numel()}") + state.last_agent_codes = codes.detach() state.last_text_token = text.detach() state.sampled_identity = state.prepared_identity @@ -482,12 +637,14 @@ class PersonaPlexStage0DuplexRuntime: pcm[state.slot] = torch.from_numpy(samples) active[state.slot] = True encoded = codec.encode_frame(pcm, active) + codes_device = encoded.detach()[:, :8].to(dtype=torch.long) # One device-to-host copy per step, whatever the number of sessions. codes = encoded.detach().to(dtype=torch.long, device="cpu") if codes.shape[-1] < 8: raise RuntimeError(f"PersonaPlex Mimi encoder returned {codes.shape[-1]} codebooks, expected at least 8") for state, identity, _ in rows: state.encoded_frame = codes[state.slot : state.slot + 1, :8].clone() + state.encoded_frame_device = codes_device[state.slot] state.encoded_identity = identity def _shared_codec(self): @@ -506,6 +663,15 @@ class PersonaPlexStage0DuplexRuntime: device=self.device, ) codec.streaming_init(self.max_sessions) + if self._codec_cuda_graphs: + # Rows are leased and recycled in place, so one graph at + # max_sessions with the active mask as input serves every step. + captured = codec.capture_cuda_graphs(encode=True, decode_frame_counts=()) + logger.info( + "PersonaPlex Stage 0 Mimi encoder %s at %d rows", + "replays a CUDA graph" if captured else "runs eagerly", + self.max_sessions, + ) self._codec = codec return codec diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_code2wav.py b/vllm_omni/model_executor/models/personaplex/personaplex_code2wav.py index 63456b4f9..dfda7b390 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_code2wav.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_code2wav.py @@ -406,6 +406,28 @@ class PersonaPlexCode2Wav(nn.Module): ).eval() for _ in range(self._max_codec_sessions) ] + if getattr(self.config, "mimi_cuda_graphs", False): + # Each decoder is leased per request and reset in place, so its B=1 + # graphs stay valid. Deltas arrive as the first single frame and then + # full chunks, so those two frame counts are recorded; other counts + # run eagerly. The graphs replay one at a time on this stage's stream + # and copy their output, so they share one pool. + pool = torch.cuda.graph_pool_handle() if torch.device(str(device)).type == "cuda" else None + captured = 0 + for codec in codecs: + codec.streaming_init(1) + captured += bool( + codec.capture_cuda_graphs( + encode=False, + decode_frame_counts=(1, _MIMI_DECODE_BATCH_FRAMES), + pool=pool, + ) + ) + logger.info( + "PersonaPlex Code2Wav replays Mimi decode from CUDA graphs on %d/%d stream(s)", + captured, + len(codecs), + ) self._set_mimi_codecs(codecs) self._mimi_device = torch.device(str(device)) reported_sr = getattr(codecs[0].model.config, "sampling_rate", None) diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py index 81729dbdf..3f605c859 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py @@ -29,9 +29,12 @@ codec composes with elastic slot recycling in batched duplex serving. from __future__ import annotations +from typing import TYPE_CHECKING + import torch import torch.nn as nn import torch.nn.functional as F +from vllm.logger import init_logger from vllm_omni.model_executor.models.personaplex.personaplex_temporal import ( _apply_rope, @@ -42,13 +45,17 @@ DEFAULT_HF_REPO = "kyutai/mimi" FRAME_SIZE = 1920 CODEBOOKS = 8 +logger = init_logger(__name__) + +if TYPE_CHECKING: + from vllm_omni.model_executor.models.personaplex.personaplex_mimi_cudagraph import MimiFrameGraph + def _normalize_active(active: torch.Tensor | None, all_active: torch.Tensor) -> torch.Tensor: - # Stage 0 and Code2Wav currently use one B=1 codec per session and omit - # `active`, so preserve that API by treating None as all rows active. When - # they switch to shared B>1 codecs, their batch builders must pass bool[B]: + # Stage 0 shares one codec with a row per session slot and passes bool[B]: # True advances that row's streaming state; False keeps an absent or padded - # row's offsets and convolution carries unchanged. + # row's offsets and convolution carries unchanged. Code2Wav leases one B=1 + # codec per request and omits `active`, which means all rows are active. if active is None: return all_active if active.shape != all_active.shape: @@ -107,6 +114,10 @@ class _StreamConv1d: ``pad_mode`` controls the stream-start left padding: SEANet convs use zeros (``constant``); the down/upsample resamplers use ``replicate`` (the first real sample), marked per row so elastic slot recycling re-primes correctly. + + The carry and the fresh-row flags are allocated once in ``reset`` and only + ever updated in place, so a CUDA graph recorded over ``__call__`` keeps + reading and writing the live state across frames, row resets and sessions. """ def __init__(self, conv: nn.Conv1d, pad_mode: str = "constant") -> None: @@ -122,29 +133,36 @@ class _StreamConv1d: self.prev = torch.zeros(batch_size, self.conv.in_channels, pad, device=device, dtype=dtype) self._fresh = torch.ones(batch_size, dtype=torch.bool, device=device) + def reset_all(self) -> None: + self.prev.zero_() + self._fresh.fill_(True) + def reset_slot(self, b: int) -> None: self.prev[b].zero_() self._fresh[b] = True def __call__(self, x: torch.Tensor, active: torch.Tensor) -> torch.Tensor: + active_view = active.view(-1, 1, 1) if self.pad_mode == "replicate": pad = self.prev.shape[-1] edge = x[..., 0:1].expand(-1, -1, pad) fresh = (self._fresh & active).view(-1, 1, 1) - self.prev = torch.where(fresh, edge.to(self.prev.dtype), self.prev) - self._fresh[active] = False + self.prev.copy_(torch.where(fresh, edge.to(self.prev.dtype), self.prev)) + self._fresh.logical_and_(~active) x = torch.cat([self.prev, x], dim=-1) t = x.shape[-1] num_frames = max(0, (t - self.kernel) // self.stride + 1) - prev = x[..., num_frames * self.stride :] - self.prev = torch.where(active.view(-1, 1, 1), prev, self.prev) + self.prev.copy_(torch.where(active_view, x[..., num_frames * self.stride :], self.prev)) if num_frames == 0: return x.new_zeros(x.shape[0], self.conv.out_channels, 0) return self.conv(x[..., : (num_frames - 1) * self.stride + self.kernel]) class _StreamConvTr1d: - """Moshi ``RawStreamingConvTranspose1d``: overlap-add tail carry per call.""" + """Moshi ``RawStreamingConvTranspose1d``: overlap-add tail carry per call. + + Like ``_StreamConv1d``, the tail carry is updated in place only. + """ def __init__(self, conv: nn.ConvTranspose1d) -> None: self.conv = conv @@ -159,6 +177,10 @@ class _StreamConvTr1d: ) self._fresh = torch.ones(batch_size, dtype=torch.bool, device=device) + def reset_all(self) -> None: + self.partial.zero_() + self._fresh.fill_(True) + def reset_slot(self, b: int) -> None: self.partial[b].zero_() self._fresh[b] = True @@ -168,17 +190,17 @@ class _StreamConvTr1d: length = out.shape[-1] tail = self.kernel - self.stride pt = self.partial.shape[-1] - merge = self.partial.clone() + merge = self.partial if self.conv.bias is not None: # The carried tail already includes the bias; the fresh output adds # it again, so subtract one copy -- except on a row's very first # frame, where the carry is zeros by construction. merge = merge - self.conv.bias[:, None] - merge[self._fresh & active] = 0.0 - self._fresh[active] = False + first = (self._fresh & active).view(-1, 1, 1) + merge = torch.where(first, torch.zeros_like(merge), merge) + self._fresh.logical_and_(~active) out[..., :pt] += merge - partial = out[..., length - tail :].clone() - self.partial = torch.where(active.view(-1, 1, 1), partial, self.partial) + self.partial.copy_(torch.where(active.view(-1, 1, 1), out[..., length - tail :], self.partial)) return out[..., : length - tail] @@ -358,13 +380,16 @@ class PersonaPlexMimiCodec(nn.Module): n_dec = self.decoder_transformer.load_weights(sd, "decoder_transformer") assert n_enc == n_dec == 80, (n_enc, n_dec) del sd + self._init_streaming_stages() + def _init_streaming_stages(self) -> None: m = self.model self._enc_stages = _walk_seanet(m.encoder.layers) self._downsample = _StreamConv1d(m.downsample.conv, pad_mode="replicate") self._upsample = _StreamConvTr1d(m.upsample.conv) self._dec_stages = _walk_seanet(m.decoder.layers) self._batch_size: int | None = None + self._cuda_graphs: dict[str, MimiFrameGraph] = {} self._all_active: torch.Tensor # -- streaming state ------------------------------------------------------ @@ -380,7 +405,22 @@ class PersonaPlexMimiCodec(nn.Module): yield self._upsample def streaming_init(self, batch_size: int) -> None: + # Re-initializing at the same batch size resets in place: the state + # tensors keep their storage, so captured frame graphs stay valid. + if getattr(self, "_batch_size", None) == batch_size: + self.reset_streaming() + return + if getattr(self, "_cuda_graphs", None): + # The state is reallocated below, so graphs recorded at the old + # batch size would read stale tensors; drop them and say so. + logger.warning( + "PersonaPlex Mimi stream resized from %s to %d rows; dropping its CUDA graphs (%s)", + self._batch_size, + batch_size, + "/".join(sorted(self._cuda_graphs)), + ) self._batch_size = batch_size + self._cuda_graphs = {} self._all_active = torch.ones(batch_size, dtype=torch.bool, device=self.device) for s in self._conv_states(): s.reset(batch_size, self.device, self.dtype) @@ -390,8 +430,7 @@ class PersonaPlexMimiCodec(nn.Module): def reset_streaming(self) -> None: assert self._batch_size is not None for state in self._conv_states(): - for b in range(self._batch_size): - state.reset_slot(b) + state.reset_all() self.encoder_transformer.reset_streaming() self.decoder_transformer.reset_streaming() @@ -401,6 +440,42 @@ class PersonaPlexMimiCodec(nn.Module): self.encoder_transformer.reset_slot(b) self.decoder_transformer.reset_slot(b) + def capture_cuda_graphs( + self, + *, + encode: bool = True, + decode_frame_counts: tuple[int, ...] = (1,), + warmup_iters: int = 2, + pool: tuple[int, int] | None = None, + ) -> list[str]: + """Record the per-frame codec steps at the streaming batch size. + + ``encode`` records ``encode_frame``; each ``F`` in ``decode_frame_counts`` + records ``decode_frames`` over ``F`` frames (``decode_frame`` is ``F=1``). + Every row is recorded active; the ``active`` mask is a graph input, so + one graph serves any subset of live rows. Warmup frames run through the + real streaming state, which is then reset in place, so the codec is + fresh afterwards. ``pool`` is a CUDA graph memory pool to share with + other codecs replayed on the same stream; by default the codec gets a + private one. Returns the captured graph names; on a non-CUDA device the + codec stays eager, and so it does after a recoverable capture failure, + which is logged with its traceback. Warmup, device execution and reset + failures propagate instead of leaving an apparently usable codec. + """ + from vllm_omni.model_executor.models.personaplex.personaplex_mimi_cudagraph import ( + capture_mimi_frame_graphs, + ) + + assert self._batch_size is not None, "call streaming_init first" + self._cuda_graphs = capture_mimi_frame_graphs( + self, + encode=encode, + decode_frame_counts=decode_frame_counts, + warmup_iters=warmup_iters, + pool=pool, + ) + return sorted(self._cuda_graphs) + # -- per-frame codec ------------------------------------------------------- @staticmethod @@ -413,32 +488,68 @@ class PersonaPlexMimiCodec(nn.Module): x = stage(x, active) if kind in {"conv", "convtr"} else stage(x) return x + def _frame_graph(self, name: str) -> MimiFrameGraph | None: + graph = self._cuda_graphs.get(name) + if graph is None or torch.cuda.is_current_stream_capturing(): + return None + return graph + @torch.no_grad() def encode_frame(self, pcm: torch.Tensor, active: torch.Tensor | None = None) -> torch.Tensor: """``[B, frame_size]`` float PCM -> ``[B, 8]`` codes.""" - x = pcm.to(self.device, self.dtype).view(-1, 1, FRAME_SIZE) active = _normalize_active(active, self._all_active) + graph = self._frame_graph("encode") + if graph is not None: + return graph.replay(pcm.reshape(-1, FRAME_SIZE), active) + return self._encode_frame_eager(pcm, active) + + def _encode_frame_eager(self, pcm: torch.Tensor, active: torch.Tensor) -> torch.Tensor: + x = pcm.to(self.device, self.dtype).view(-1, 1, FRAME_SIZE) x = self._run_stages(x, self._enc_stages, active) x = self.encoder_transformer.step(x.transpose(1, 2), active).transpose(1, 2) x = self._downsample(x, active) codes = self.model.quantizer.encode(x, num_quantizers=CODEBOOKS) # [Q, B, T] return codes[:CODEBOOKS, :, 0].transpose(0, 1).contiguous() + def _quantizer_decode(self, codes: torch.Tensor) -> torch.Tensor: + """``MimiSplitResidualVectorQuantizer.decode`` without its host-side scalar. + + The Hugging Face RVQ decode seeds its sum with ``torch.tensor(0.0)``, a + host-to-device copy that cannot be recorded in a CUDA graph. This + seeds it with an on-device zero instead; the arithmetic is unchanged. + """ + quantizer = self.model.quantizer + split = quantizer.num_semantic_quantizers + + def rvq_decode(rvq, rvq_codes: torch.Tensor) -> torch.Tensor: + out = torch.zeros((), device=rvq_codes.device) + for layer, indices in zip(rvq.layers, rvq_codes.transpose(0, 1)): + out = out + layer.decode(indices) + if rvq.output_proj is not None: + out = rvq.output_proj(out) + return out + + out = rvq_decode(quantizer.semantic_residual_vector_quantizer, codes[:, :split]) + if codes.shape[1] > split: + out += rvq_decode(quantizer.acoustic_residual_vector_quantizer, codes[:, split:]) + return out + @torch.no_grad() def decode_frame(self, codes: torch.Tensor, active: torch.Tensor | None = None) -> torch.Tensor: """``[B, 8]`` codes -> ``[B, frame_size]`` float PCM.""" - emb = self.model.quantizer.decode(codes.to(self.device).view(-1, CODEBOOKS, 1)) - active = _normalize_active(active, self._all_active) - emb = self._upsample(emb, active) - emb = self.decoder_transformer.step(emb.transpose(1, 2), active).transpose(1, 2) - x = self._run_stages(emb, self._dec_stages, active) - return x[:, 0, :] + return self.decode_frames(codes.reshape(-1, CODEBOOKS, 1), active) @torch.no_grad() def decode_frames(self, codes: torch.Tensor, active: torch.Tensor | None = None) -> torch.Tensor: """``[B, 8, F]`` codes -> ``[B, F * frame_size]`` float PCM.""" - emb = self.model.quantizer.decode(codes.to(self.device)) active = _normalize_active(active, self._all_active) + graph = self._frame_graph(f"decode_f{codes.shape[-1]}") + if graph is not None: + return graph.replay(codes, active) + return self._decode_frames_eager(codes, active) + + def _decode_frames_eager(self, codes: torch.Tensor, active: torch.Tensor) -> torch.Tensor: + emb = self._quantizer_decode(codes.to(self.device)) emb = self._upsample(emb, active) emb = self.decoder_transformer.step(emb.transpose(1, 2), active).transpose(1, 2) x = self._run_stages(emb, self._dec_stages, active) diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_mimi_cudagraph.py b/vllm_omni/model_executor/models/personaplex/personaplex_mimi_cudagraph.py new file mode 100644 index 000000000..fefffd87b --- /dev/null +++ b/vllm_omni/model_executor/models/personaplex/personaplex_mimi_cudagraph.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Replay the PersonaPlex streaming Mimi frame steps from CUDA graphs. + +One 80 ms Mimi encode or decode frame is a few hundred small kernels (SEANet +convs, an 8-layer streaming transformer and the RVQ), so eager execution is +launch bound. Every streaming state tensor of ``PersonaPlexMimiCodec`` (conv +carries, fresh-row flags, ring KV caches and offsets) is allocated once by +``streaming_init`` and only updated in place, and row resets are in place too, +so a graph recorded at the streaming batch size stays valid across frames, +session slot recycling and full stream resets. + +The ``active`` row mask is a graph input: one graph at the shared batch size +serves any subset of live rows, and inactive rows keep their state exactly as +in eager execution. +""" + +from __future__ import annotations + +import traceback +from collections.abc import Callable +from contextlib import ExitStack +from typing import TYPE_CHECKING + +import torch +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm_omni.model_executor.models.personaplex.personaplex_mimi import PersonaPlexMimiCodec + +logger = init_logger(__name__) + +__all__ = ["MimiFrameGraph", "capture_mimi_frame_graphs"] + + +def _is_recoverable_capture_error(error: RuntimeError) -> bool: + """Allow only known capture limitations, never an unexplained CUDA error.""" + pending: list[BaseException] = [error] + seen: set[int] = set() + known_failure = False + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + message = str(current).split("\n", 1)[0] + if isinstance(current, torch.cuda.OutOfMemoryError) or ( + isinstance(current, RuntimeError) + and ( + message.startswith("CUDA out of memory.") + or message + in { + "CUDA error: out of memory", + "CUDA error: operation not permitted when stream is capturing", + } + ) + ): + known_failure = True + elif not ( + isinstance(current, RuntimeError) + and message == "CUDA error: operation failed due to a previous error during capture" + ): + return False + # capture_end can replace a body exception with capture-invalidated. + # The original failure must also be recoverable, even if suppressed. + pending.extend(exc for exc in (current.__cause__, current.__context__) if exc is not None) + return known_failure + + +def _capture_error_details(error: RuntimeError) -> str: + """Preserve diagnostics without retaining capture frames or graph pools.""" + details = "".join(traceback.format_exception(error)) + pending: list[BaseException] = [error] + seen: set[int] = set() + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + if current.__traceback__ is not None: + # ExitStack frames can retain their own traceback in exception tuples. + traceback.clear_frames(current.__traceback__) + current.__traceback__ = None + pending.extend(exc for exc in (current.__cause__, current.__context__) if exc is not None) + return details + + +class MimiFrameGraph: + """A recorded ``(frame, active) -> output`` codec step with static buffers.""" + + def __init__( + self, + graph: torch.cuda.CUDAGraph, + static_input: torch.Tensor, + static_active: torch.Tensor, + static_output: torch.Tensor, + eager: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + ) -> None: + self.graph = graph + self.static_input = static_input + self.static_active = static_active + self.static_output = static_output + self.eager = eager + self.replays = 0 + self._warned = False + + def replay(self, frame: torch.Tensor, active: torch.Tensor) -> torch.Tensor: + if frame.shape != self.static_input.shape or active.shape != self.static_active.shape: + if not self._warned: + self._warned = True + logger.warning( + "PersonaPlex Mimi graph input shape %s does not match capture shape %s; running eagerly", + tuple(frame.shape), + tuple(self.static_input.shape), + ) + return self.eager(frame, active) + self.static_input.copy_(frame) + self.static_active.copy_(active) + self.graph.replay() + self.replays += 1 + # The static output is overwritten by the next replay. + return self.static_output.clone() + + +def capture_mimi_frame_graphs( + codec: PersonaPlexMimiCodec, + *, + encode: bool, + decode_frame_counts: tuple[int, ...], + warmup_iters: int = 2, + pool: tuple[int, int] | None = None, +) -> dict[str, MimiFrameGraph]: + """Capture the codec's encode and ``F``-frame decode steps. + + Returns an empty dict off CUDA or after a recoverable capture failure. + Warmup, device execution and stream-reset failures always propagate. + """ + from vllm_omni.model_executor.models.personaplex.personaplex_mimi import CODEBOOKS, FRAME_SIZE + + device = torch.device(codec.device) + if device.type != "cuda" or not torch.cuda.is_available(): + logger.info("PersonaPlex Mimi CUDA graphs skipped on device %s", device) + return {} + batch_size = codec._batch_size + specs: list[tuple[str, Callable, torch.Tensor]] = [] + if encode: + specs.append( + ( + "encode", + codec._encode_frame_eager, + torch.zeros(batch_size, FRAME_SIZE, device=device, dtype=codec.dtype), + ) + ) + for frames in sorted(set(decode_frame_counts)): + specs.append( + ( + f"decode_f{frames}", + codec._decode_frames_eager, + torch.zeros(batch_size, CODEBOOKS, frames, device=device, dtype=torch.long), + ) + ) + # A private pool by default: these graphs are captured outside the model + # runner's own capture sequence, so they must not reuse its pool blocks. + if pool is None: + pool = torch.cuda.graph_pool_handle() + graphs: dict[str, MimiFrameGraph] = {} + warmed_specs: list[tuple[str, Callable, torch.Tensor, torch.Tensor]] = [] + capture_error_details: str | None = None + try: + with torch.no_grad(): + # Validate every eager step before capture can select the fallback. + for name, eager, static_input in specs: + static_active = torch.ones(batch_size, dtype=torch.bool, device=device) + for _ in range(max(warmup_iters, 1)): + eager(static_input, static_active) + torch.accelerator.synchronize(device) + warmed_specs.append((name, eager, static_input, static_active)) + for name, eager, static_input, static_active in warmed_specs: + # capture_end may raise before torch.cuda.graph restores its + # stream. Always restore our caller's stream on that path. + with torch.cuda.stream(torch.cuda.current_stream(device)): + graph = torch.cuda.CUDAGraph() + # __enter__ synchronizes and clears the CUDA cache. A + # failure there is a setup failure, not a capture fallback. + capture = torch.cuda.graph(graph, pool=pool, capture_error_mode="thread_local") + capture.__enter__() + try: + # Other worker threads may issue CUDA calls while this records. + with ExitStack() as stack: + stack.push(capture) + static_output = eager(static_input, static_active) + except RuntimeError as error: + if not _is_recoverable_capture_error(error): + raise + capture_error_details = _capture_error_details(error) + graphs.clear() + # Release partial graphs and graph-owned output tensors + # before reset, including the last loop iteration. + graph = static_output = capture = None + break + # A device execution failure is not an optional capture failure. + torch.accelerator.synchronize(device) + graphs[name] = MimiFrameGraph(graph, static_input, static_active, static_output, eager) + finally: + # Warmup frames advanced the real streaming state; restore a fresh + # stream without reallocating the tensors the graphs point at. + codec.reset_streaming() + # Reset enqueues device work too. Do not report usable graphs or an eager + # fallback unless that work completed successfully. + torch.accelerator.synchronize(device) + if capture_error_details is not None: + logger.warning( + "PersonaPlex Mimi CUDA graphs were requested (%s at batch size %d) but capture failed; " + "the codec runs eagerly, which lowers realtime session capacity\n%s", + "/".join(name for name, _, _ in specs), + batch_size, + capture_error_details, + ) + if graphs: + logger.debug("Captured PersonaPlex Mimi %s graph(s) at batch size %d", "/".join(graphs), batch_size) + return graphs diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_talker.py b/vllm_omni/model_executor/models/personaplex/personaplex_talker.py index 14374fa57..178ecba67 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_talker.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_talker.py @@ -30,7 +30,7 @@ encode of the input WAV (built in ``preprocess``); live duplex is Phase 2. from __future__ import annotations from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import torch from torch import nn @@ -54,6 +54,11 @@ from vllm_omni.model_executor.models.personaplex.personaplex_embeddings import ( PersonaPlexInputEmbeddings, ) +if TYPE_CHECKING: + from vllm_omni.model_executor.models.personaplex.duplex.stage0 import ( + PersonaPlexStage0DuplexRuntime, + ) + __all__ = ["PersonaPlexTalkerForConditionalGeneration"] @@ -437,6 +442,7 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module): model_path=model_path, device=device, max_sessions=int(getattr(self.vllm_config.model_config, "duplex_max_sessions", 1)), + codec_cuda_graphs=bool(getattr(self.config, "mimi_cuda_graphs", False)), ) self._personaplex_duplex_stage0_runtime = runtime return runtime @@ -485,7 +491,9 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module): if isinstance(duplex, dict) and duplex.get("data_plane") is True: appends.append(duplex) if appends: - self._duplex_stage0_runtime().encode_appends(appends) + runtime = self._duplex_stage0_runtime() + runtime.encode_appends(appends) + runtime.prepare_live_appends(appends) def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: runtime = getattr(self, "_personaplex_duplex_stage0_runtime", None) @@ -551,9 +559,48 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module): ) text_token = input_ids.reshape(bsz).to(torch.long) hidden = hidden_states.reshape(bsz, 1, -1).to(self._dtype) - audio_tokens: list[torch.Tensor] = [] - audio_provided: list[torch.Tensor] = [] - for info in req_infos: + runtime = self._duplex_stage0_runtime() + audio_tokens, audio_provided = self._depformer_teacher_forcing(runtime, req_ids, req_infos, hidden.device) + codes = self.depformer( + text_token, + hidden, + audio_tokens=audio_tokens, + audio_provided=audio_provided, + num_steps=self.num_active_codebooks, + ).to(torch.long) + # The next append reads the agent frame with the teacher-forced + # codebooks applied; merge it once for the whole batch. + effective = torch.where(audio_provided[:, :8], audio_tokens[:, :8], codes[:, :8]) + for row, request_id in enumerate(req_ids): + runtime.record_sample( + request_id=request_id, + text_token=text_token[row], + effective_codes=effective[row], + ) + return codes + + @staticmethod + def _depformer_teacher_forcing( + runtime: PersonaPlexStage0DuplexRuntime, + req_ids: list[str], + req_infos: list[dict[str, Any]], + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Stack each request's depformer teacher-forcing tokens and mask on ``device``. + + Live appends batched in ``preprocess_batch`` already hold device copies. + The rest (first appends, aborted epochs) carry host tensors in their + request information and are moved in one copy. + """ + tokens_rows: list[torch.Tensor] = [] + provided_rows: list[torch.Tensor] = [] + host_rows: list[int] = [] + for row, (request_id, info) in enumerate(zip(req_ids, req_infos)): + prepared = runtime.prepared_depformer_state(request_id) + if prepared is not None: + tokens_rows.append(prepared[0]) + provided_rows.append(prepared[1]) + continue tokens = info.get("pplex_depformer_audio_tokens") provided = info.get("pplex_depformer_audio_provided") if not isinstance(tokens, torch.Tensor) or not isinstance(provided, torch.Tensor): @@ -565,29 +612,18 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module): "PersonaPlex depformer teacher-forcing token/mask shapes differ: " f"{tuple(tokens.shape)} != {tuple(provided.shape)}" ) - audio_tokens.append(tokens) - audio_provided.append(provided) - codes = self.depformer( - text_token, - hidden, - audio_tokens=torch.stack(audio_tokens).to( - device=hidden.device, - dtype=torch.long, - ), - audio_provided=torch.stack(audio_provided).to( - device=hidden.device, - dtype=torch.bool, - ), - num_steps=self.num_active_codebooks, - ).to(torch.long) - runtime = self._duplex_stage0_runtime() - for row, request_id in enumerate(req_ids): - runtime.record_sample( - request_id=request_id, - text_token=text_token[row], - agent_codes=codes[row], - ) - return codes + host_rows.append(row) + tokens_rows.append(tokens) + provided_rows.append(provided) + if host_rows: + host_tokens = torch.stack([tokens_rows[row] for row in host_rows]).to(device=device, dtype=torch.long) + host_provided = torch.stack([provided_rows[row] for row in host_rows]).to(device=device, dtype=torch.bool) + if len(host_rows) == len(req_ids): + return host_tokens, host_provided + for index, row in enumerate(host_rows): + tokens_rows[row] = host_tokens[index] + provided_rows[row] = host_provided[index] + return torch.stack(tokens_rows), torch.stack(provided_rows) # ------------------------------------------------------------------ # Weight loading