diff --git a/tests/model_executor/models/personaplex/duplex/test_plugin.py b/tests/model_executor/models/personaplex/duplex/test_plugin.py index 56a90e4b5..72cb6d9aa 100644 --- a/tests/model_executor/models/personaplex/duplex/test_plugin.py +++ b/tests/model_executor/models/personaplex/duplex/test_plugin.py @@ -101,6 +101,8 @@ def test_personaplex_deploy_is_duplex_and_propagates_capacity_to_all_model_stage assert deploy.session_mode == "duplex" assert deploy.duplex_session.max_sessions == 2 assert [stage.yaml_engine_args.get("duplex_max_sessions") for stage in stages] == [2, 2] + # Lockstep frames of all sessions must share a step (see the deploy comment). + assert stages[0].yaml_engine_args.get("async_scheduling") is False assert "personaplex_codec_max_sessions" not in deploy.connectors["connector_of_shared_memory"]["extra"] 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 61c6e6c1b..40368c51a 100644 --- a/tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py +++ b/tests/model_executor/models/personaplex/duplex/test_stage0_runtime.py @@ -13,8 +13,10 @@ from vllm_omni.model_executor.models.personaplex.duplex.policy import ( SINE_TOKENS, ) from vllm_omni.model_executor.models.personaplex.duplex.stage0 import ( + PersonaPlexStage0CapacityError, PersonaPlexStage0DuplexRuntime, PersonaPlexStage0PreparedAppend, + PersonaPlexStage0StaleEpochError, ) from vllm_omni.model_executor.models.personaplex.personaplex_talker import ( PersonaPlexTalkerForConditionalGeneration, @@ -24,20 +26,30 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _FakeCodec: + """Shared streaming encoder whose code for a row is that row's frame count.""" + def __init__(self) -> None: self.encode_calls = 0 - self.reset_calls = 0 + self.reset_slots: list[int] = [] + self.frames: list[int] = [] def streaming_init(self, batch_size: int) -> None: - assert batch_size == 1 + self.frames = [0] * batch_size - def encode_frame(self, pcm: torch.Tensor) -> torch.Tensor: - assert pcm.shape == (1, 1920) + def encode_frame(self, pcm: torch.Tensor, active: torch.Tensor) -> torch.Tensor: + assert pcm.shape == (len(self.frames), 1920) + assert active.shape == (len(self.frames),) self.encode_calls += 1 - return torch.full((1, 8), self.encode_calls, dtype=torch.long) + 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] + return codes - def reset_streaming(self) -> None: - self.reset_calls += 1 + def reset_slot(self, row: int) -> None: + self.reset_slots.append(row) + self.frames[row] = 0 class _FakeTalker: @@ -101,20 +113,14 @@ def _duplex_info(*, seq: int, session_id: str = "session", epoch: int = 0): } -def _runtime(*codecs: _FakeCodec) -> PersonaPlexStage0DuplexRuntime: +def _runtime(codec: _FakeCodec, max_sessions: int = 1) -> PersonaPlexStage0DuplexRuntime: voice_embeddings = torch.arange(8, dtype=torch.float32).reshape(2, 1, 1, 4) - codec_args: dict[str, object] = {"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, + codec=codec, + max_sessions=max_sessions, tokenizer=lambda _text: [7, 8, 9], voice_loader=lambda _voice: { "embeddings": voice_embeddings, @@ -268,9 +274,7 @@ def _prepare_two_sessions( def test_live_sessions_keep_independent_streaming_encoders() -> None: - first_codec = _FakeCodec() - second_codec = _FakeCodec() - runtime = _runtime(first_codec, second_codec) + runtime = _runtime(_FakeCodec(), max_sessions=2) first_1, second_1, first_2, second_2 = _prepare_two_sessions(runtime) @@ -281,7 +285,7 @@ def test_live_sessions_keep_independent_streaming_encoders() -> None: def test_stage0_session_capacity_fails_before_codec_state_is_shared() -> None: - runtime = _runtime(_FakeCodec(), _FakeCodec()) + runtime = _runtime(_FakeCodec(), max_sessions=2) _prepare_two_sessions(runtime) with pytest.raises(RuntimeError, match="capacity 2"): @@ -291,21 +295,23 @@ def test_stage0_session_capacity_fails_before_codec_state_is_shared() -> None: ) -def test_close_session_resets_and_reuses_released_codec() -> None: - first_codec = _FakeCodec() - second_codec = _FakeCodec() - runtime = _runtime(first_codec, second_codec) +def test_close_session_resets_only_its_row_and_reuses_it() -> None: + codec = _FakeCodec() + runtime = _runtime(codec, max_sessions=2) _prepare_two_sessions(runtime) + closed_slot = runtime.sessions[("session", 0)].slot + other_slot = runtime.sessions[("other", 0)].slot runtime.close_session("session", 0) - assert first_codec.reset_calls == 1 - assert second_codec.reset_calls == 0 + assert codec.reset_slots == [closed_slot] replacement = runtime.prepare_append( _duplex_info(seq=1, session_id="replacement"), prompt_len=18, ) - assert replacement.user_codes[:, 0].tolist() == [3] + assert runtime.sessions[("replacement", 0)].slot == closed_slot + assert replacement.user_codes[:, 0].tolist() == [1] + assert codec.frames[other_slot] == 2 def test_a_new_epoch_replays_the_prefill_and_recycles_the_codec() -> None: @@ -324,8 +330,8 @@ def test_a_new_epoch_replays_the_prefill_and_recycles_the_codec() -> None: assert restarted.user_codes.shape == (1, 8) assert list(runtime.sessions) == [("session", 1)] assert runtime.request_sessions == {"req-e1": ("session", 1)} - assert codec.reset_calls == 1 - assert runtime.sessions[("session", 1)].codec is codec + assert codec.reset_slots == [0] + assert runtime.sessions[("session", 1)].slot == 0 def test_a_late_finish_of_the_old_epoch_request_does_not_close_the_new_state() -> None: @@ -341,11 +347,164 @@ def test_a_late_finish_of_the_old_epoch_request_does_not_close_the_new_state() - def test_stage0_capacity_counts_live_epochs_not_superseded_ones() -> None: - runtime = _runtime(_FakeCodec(), _FakeCodec()) + runtime = _runtime(_FakeCodec(), max_sessions=2) runtime.prepare_append(_duplex_info(seq=1), prompt_len=18) runtime.prepare_append(_duplex_info(seq=1, session_id="other"), prompt_len=18) - # Restarting one session must not need a third codec. + # Restarting one session must not need a third encoder row. runtime.prepare_append(_duplex_info(seq=1, epoch=1), prompt_len=18) assert sorted(runtime.sessions) == [("other", 0), ("session", 1)] + + +def test_encode_appends_batches_sessions_into_one_encoder_call() -> None: + codec = _FakeCodec() + runtime = _runtime(codec, max_sessions=4) + + runtime.encode_appends([_duplex_info(seq=1), _duplex_info(seq=1, session_id="other")]) + first = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18) + other = runtime.prepare_append(_duplex_info(seq=1, session_id="other"), prompt_len=18) + + assert codec.encode_calls == 1 + assert first.user_codes[:, 0].tolist() == [1] + assert other.user_codes[:, 0].tolist() == [1] + + +def test_encode_appends_leaves_rows_without_a_new_append_untouched() -> None: + codec = _FakeCodec() + runtime = _runtime(codec, max_sessions=2) + runtime.prepare_append(_duplex_info(seq=1), prompt_len=18) + runtime.prepare_append(_duplex_info(seq=1, session_id="other"), prompt_len=18) + other_slot = runtime.sessions[("other", 0)].slot + + runtime.encode_appends([_duplex_info(seq=2)]) + runtime.prepare_append(_duplex_info(seq=2), prompt_len=19) + + assert codec.frames[other_slot] == 1 + assert runtime.sessions[("session", 0)].user_codes[:, 0].tolist() == [1, 2] + + +def test_encode_appends_skips_an_identity_that_was_already_encoded_or_prepared() -> None: + codec = _FakeCodec() + runtime = _runtime(codec) + + # A chunked first prefill puts the same append in two scheduler steps. + runtime.encode_appends([_duplex_info(seq=1)]) + runtime.encode_appends([_duplex_info(seq=1)]) + runtime.prepare_append(_duplex_info(seq=1), prompt_len=18) + runtime.encode_appends([_duplex_info(seq=1)]) + retry = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18) + + assert codec.encode_calls == 1 + assert retry.user_codes[:, 0].tolist() == [1] + + +def test_encode_appends_leaves_over_capacity_appends_to_prepare_append() -> None: + codec = _FakeCodec() + runtime = _runtime(codec) + + runtime.encode_appends([_duplex_info(seq=1), _duplex_info(seq=1, session_id="other")]) + + assert list(runtime.sessions) == [("session", 0)] + with pytest.raises(RuntimeError, match="capacity 1"): + runtime.prepare_append(_duplex_info(seq=1, session_id="other"), prompt_len=18) + + +def _restart_overlap(runtime: PersonaPlexStage0DuplexRuntime) -> tuple[dict, dict]: + """Epoch 0 of a session is live, then a cancel restarts it as epoch 1 while + epoch 0's aborted request still has an append in the same scheduler step.""" + runtime.prepare_append(_duplex_info(seq=1), prompt_len=18, request_id="req-e0") + runtime.prepare_append(_duplex_info(seq=2), prompt_len=19, request_id="req-e0") + return _duplex_info(seq=3), _duplex_info(seq=1, epoch=1) + + +@pytest.mark.parametrize("max_sessions", [1, 2]) +@pytest.mark.parametrize("encode_old_first", [True, False]) +@pytest.mark.parametrize("prepare_old_first", [True, False]) +def test_cancel_overlap_full_processing_order_never_re_leases_the_aborted_epoch( + max_sessions: int, encode_old_first: bool, prepare_old_first: bool +) -> None: + codec = _FakeCodec() + runtime = _runtime(codec, max_sessions=max_sessions) + old, new = _restart_overlap(runtime) + + # One scheduler step: batched encode, then per-request prepare, then sampling. + runtime.encode_appends([old, new] if encode_old_first else [new, old]) + order = [("req-e0", old, 19), ("req-e1", new, 18)] + restarted = None + for request_id, duplex, prompt_len in order if prepare_old_first else order[::-1]: + if request_id == "req-e0": + with pytest.raises(PersonaPlexStage0StaleEpochError): + runtime.prepare_append(duplex, prompt_len=prompt_len, request_id=request_id) + 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.close_request("req-e0") # the engine's late finish of the aborted request + + assert list(runtime.sessions) == [("session", 1)] + assert len(runtime._free_slots) == max_sessions - 1 + assert restarted is not None and restarted.user_codes[:, 0].tolist() == [1] + assert codec.encode_calls == 3 + assert runtime.sessions[("session", 1)].sampled_identity == (1, 1) + + +def test_codec_init_failure_propagates_after_one_attempt() -> None: + attempts = 0 + + def failing_codec(): + nonlocal attempts + attempts += 1 + raise RuntimeError("CUDA out of memory while building the Mimi encoder") + + runtime = PersonaPlexStage0DuplexRuntime( + _FakeTalker(), + model_path="/unused", + device="cpu", + codec_factory=failing_codec, + max_sessions=16, + tokenizer=lambda _text: [7, 8, 9], + voice_loader=lambda _voice: {}, + ) + appends = [_duplex_info(seq=1, session_id=f"s{i}") for i in range(16)] + + with pytest.raises(RuntimeError, match="out of memory"): + runtime.encode_appends(appends) + assert attempts == 1 + assert not runtime.sessions + + +def test_capacity_error_is_a_dedicated_exception() -> None: + runtime = _runtime(_FakeCodec(), max_sessions=1) + runtime.prepare_append(_duplex_info(seq=1), prompt_len=18) + + with pytest.raises(PersonaPlexStage0CapacityError, match="capacity 1"): + runtime.prepare_append(_duplex_info(seq=1, session_id="other"), prompt_len=18) + + +def test_talker_gives_the_aborted_epoch_a_neutral_row_in_the_same_step() -> None: + codec = _FakeCodec() + runtime = _runtime(codec, max_sessions=1) + old, new = _restart_overlap(runtime) + talker = PersonaPlexTalkerForConditionalGeneration.__new__(PersonaPlexTalkerForConditionalGeneration) + torch.nn.Module.__init__(talker) + talker._personaplex_duplex_stage0_runtime = runtime + talker._dtype = torch.float32 + talker.mtp_hidden_size = 4 + + infos = { + "req-e0": {"duplex": old, "request_id": "req-e0", "duplex_prompt_len": 19, "duplex_token_offset": 0}, + "req-e1": {"duplex": new, "request_id": "req-e1", "duplex_prompt_len": 18, "duplex_token_offset": 0}, + } + talker.preprocess_batch(req_ids=list(infos), model_intermediate_buffer=infos, device=torch.device("cpu")) + outs = { + req: talker.preprocess(torch.zeros(1, dtype=torch.long), None, _omni_is_prefill=True, **info) + for req, info in infos.items() + } + + _, stale_embeds, stale_info = outs["req-e0"] + assert stale_info["duplex"] == {"stage0_stale": True} + assert torch.count_nonzero(stale_embeds) == 0 + assert stale_info["pplex_depformer_audio_tokens"].numel() == 16 + assert outs["req-e1"][2]["duplex"]["epoch"] == 1 + assert list(runtime.sessions) == [("session", 1)] diff --git a/tests/model_executor/models/personaplex/test_depformer_num_steps.py b/tests/model_executor/models/personaplex/test_depformer_num_steps.py new file mode 100644 index 000000000..934bd9dbf --- /dev/null +++ b/tests/model_executor/models/personaplex/test_depformer_num_steps.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch + +from vllm_omni.model_executor.models.personaplex.configuration_personaplex import ( + PersonaPlexDepformerConfig, +) +from vllm_omni.model_executor.models.personaplex.personaplex_depformer import PersonaPlexDepformer + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _tiny_depformer() -> PersonaPlexDepformer: + torch.manual_seed(0) + config = PersonaPlexDepformerConfig( + hidden_size=32, + num_hidden_layers=2, + num_attention_heads=4, + head_dim=8, + num_key_value_heads=4, + intermediate_size=48, + dep_q=16, + num_active_codebooks=8, + card=64, + ) + model = PersonaPlexDepformer(config, temporal_hidden_size=24, text_card=100) + for param in model.parameters(): + torch.nn.init.normal_(param, std=0.2) + return model.eval() + + +def test_num_steps_matches_prefix_of_full_run_with_teacher_forcing(): + model = _tiny_depformer() + text = torch.tensor([3, 7, 11]) + hidden = torch.randn(3, 1, 24) + tokens = torch.randint(0, 64, (3, 16)) + provided = torch.zeros(3, 16, dtype=torch.bool) + provided[0, 1:8] = True + provided[:, 8:] = True + + full, full_logits = model(text, hidden, tokens, provided, return_logits=True) + head, head_logits = model(text, hidden, tokens, provided, return_logits=True, num_steps=8) + + assert full.shape == (3, 16) + assert head.shape == (3, 8) + assert torch.equal(head, full[:, :8]) + assert torch.equal(head_logits, full_logits[:, :8]) + + +@pytest.mark.parametrize("num_steps", [0, 17]) +def test_num_steps_out_of_range_is_rejected(num_steps): + model = _tiny_depformer() + with pytest.raises(ValueError, match="num_steps"): + model(torch.tensor([1]), torch.randn(1, 1, 24), num_steps=num_steps) 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 fce42f40c..0fa216e14 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 @@ -367,3 +367,30 @@ def test_mimi_conv_carries_preserve_inactive_rows(kind: str) -> None: else: assert torch.equal(_stream_carry(batched)[row], carry_before[row]) assert torch.equal(batched._fresh[row], fresh_before[row]) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_recycled_mimi_transformer_row_matches_a_fresh_stream(): + from vllm_omni.model_executor.models.personaplex.personaplex_mimi import _MimiStreamingTransformer + + torch.manual_seed(0) + used = _MimiStreamingTransformer(num_layers=2, dim=16, num_heads=2, context=8) + for param in used.parameters(): + torch.nn.init.normal_(param, std=0.1) + fresh = _MimiStreamingTransformer(num_layers=2, dim=16, num_heads=2, context=8) + fresh.load_state_dict(used.state_dict()) + used.streaming_init(2) + fresh.streaming_init(2) + both = torch.tensor([True, True]) + + # Wrap row 0's ring (6 frames x 2 positions > context 8) before recycling it. + for _ in range(6): + used.step(torch.randn(2, 2, 16), both) + used.reset_slot(0) + + frames = [torch.randn(2, 2, 16) for _ in range(5)] + for x in frames: + recycled = used.step(x, both) + expected = fresh.step(x, both) + torch.testing.assert_close(recycled[0], expected[0], rtol=0, atol=0) diff --git a/tests/model_executor/models/personaplex/test_talker_omni_output.py b/tests/model_executor/models/personaplex/test_talker_omni_output.py new file mode 100644 index 000000000..62acb6ef8 --- /dev/null +++ b/tests/model_executor/models/personaplex/test_talker_omni_output.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.personaplex.personaplex_talker import ( + PersonaPlexTalkerForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_make_omni_output_keeps_token_rows_for_mixed_prefill_and_live_rows(): + # One live session carries one cached audio row; a newly admitted session + # prefills 150 rows in the same step. The runner indexes the hidden states + # with token-space logits indices, so the last prefill row must survive. + hidden = torch.randn(151, 16) + live_codes = torch.zeros(1, 8, dtype=torch.long) + out = PersonaPlexTalkerForConditionalGeneration.make_omni_output( + None, + hidden, + model_intermediate_buffer=[{"codes": {"audio": live_codes}}, {}], + ) + + assert isinstance(out, OmniOutput) + assert out.text_hidden_states.shape[0] == 151 + logits_indices = torch.tensor([0, 150]) + assert out.text_hidden_states[logits_indices].shape == (2, 16) + assert torch.equal(out.multimodal_outputs["codes"]["audio"], live_codes) diff --git a/tests/model_executor/stage_input_processors/test_personaplex.py b/tests/model_executor/stage_input_processors/test_personaplex.py index 37f464201..6521fdd96 100644 --- a/tests/model_executor/stage_input_processors/test_personaplex.py +++ b/tests/model_executor/stage_input_processors/test_personaplex.py @@ -93,15 +93,18 @@ def test_post_sample_talker_mtp_uses_current_temporal_state() -> None: *, audio_tokens: torch.Tensor | None = None, audio_provided: torch.Tensor | None = None, + num_steps: int | 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) + received["num_steps"] = num_steps + return torch.arange(num_steps or 16, dtype=torch.long).reshape(1, -1) model = SimpleNamespace( _dtype=torch.float32, + num_active_codebooks=8, depformer=depformer, _duplex_stage0_runtime=lambda: SimpleNamespace( record_sample=lambda *, request_id, text_token, agent_codes: recorded.append( @@ -126,7 +129,9 @@ def test_post_sample_talker_mtp_uses_current_temporal_state() -> None: ], ) - assert codes.shape == (1, 16) + # Only the vocoded agent codebooks are drawn on the duplex path. + assert received["num_steps"] == 8 + assert codes.shape == (1, 8) 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)) diff --git a/vllm_omni/deploy/personaplex.yaml b/vllm_omni/deploy/personaplex.yaml index 65cdec6d6..567d0fd69 100644 --- a/vllm_omni/deploy/personaplex.yaml +++ b/vllm_omni/deploy/personaplex.yaml @@ -12,6 +12,10 @@ # 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. +# * 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. pipeline: personaplex session_mode: duplex duplex_session: @@ -38,7 +42,12 @@ stages: # 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 + # Every live session appends one frame per 80 ms tick as a one-token + # resumable request. With async scheduling the next step is scheduled + # before a session's previous frame is parked, so a session can join only + # every other step and concurrent sessions split into alternating batches. + # Synchronous scheduling puts every waiting frame in one step. + async_scheduling: false max_num_batched_tokens: 512 # Temporal transformer sliding_window is 3000 frames; keep model len within. max_model_len: 3000 diff --git a/vllm_omni/model_executor/models/personaplex/duplex/stage0.py b/vllm_omni/model_executor/models/personaplex/duplex/stage0.py index 5deecafdb..dfeca1878 100644 --- a/vllm_omni/model_executor/models/personaplex/duplex/stage0.py +++ b/vllm_omni/model_executor/models/personaplex/duplex/stage0.py @@ -27,6 +27,14 @@ from vllm_omni.model_executor.models.personaplex.duplex.policy import ( _FRAME_SAMPLES = FRAME_SIZE +class PersonaPlexStage0CapacityError(RuntimeError): + """Every streaming encoder row is leased by a live session.""" + + +class PersonaPlexStage0StaleEpochError(RuntimeError): + """The append belongs to an epoch that a newer epoch of the same session has superseded.""" + + @dataclass(slots=True) class PersonaPlexStage0PreparedAppend: input_ids: Any @@ -52,7 +60,9 @@ class PersonaPlexStage0SessionState: prepared: PersonaPlexStage0PreparedAppend | None = None last_seq: int = 0 request_ids: set[str] = field(default_factory=set) - codec: Any | None = None + slot: int | None = None + encoded_identity: tuple[int, int] | None = None + encoded_frame: Any | None = None def _tokenizer_path(model_path: str) -> Path: @@ -124,7 +134,13 @@ def personaplex_prefill_slots(model_path: str, voice: str, persona: str) -> int: class PersonaPlexStage0DuplexRuntime: - """Own per-session streaming Mimi encoders and first-append prefill.""" + """Own the shared streaming Mimi encoder and each session's first-append prefill. + + One encoder holds ``max_sessions`` streaming rows; a live ``(session, epoch)`` + leases one row for its lifetime. ``encode_appends`` encodes every new append + of a scheduler step in one batched call (rows without a new append are + inactive and keep their state), and ``prepare_append`` consumes the result. + """ def __init__( self, @@ -145,14 +161,61 @@ class PersonaPlexStage0DuplexRuntime: self.device = device self.max_sessions = max_sessions self._codec_factory = codec_factory - self._free_codecs: list[Any] = [] + self._codec: Any | None = None + self._free_slots: list[int] = list(reversed(range(max_sessions))) self._tokenizer = tokenizer self._voice_loader = voice_loader self.sessions: dict[tuple[str, int], PersonaPlexStage0SessionState] = {} self.request_sessions: dict[str, tuple[str, int]] = {} + # 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() if codec is not None: - codec.streaming_init(1) - self._free_codecs.append(codec) + codec.streaming_init(max_sessions) + self._codec = codec + + def encode_appends(self, appends: list[dict[str, Any]]) -> None: + """Encode the new user frame of each append in one batched encoder call. + + Called once per scheduler step before the per-request ``prepare_append`` + calls. An append whose ``(epoch, seq)`` is already encoded or prepared + (a chunked first prefill spans several steps) is skipped, so a row's + streaming state advances exactly once per frame. Appends that cannot be + admitted are left for ``prepare_append`` to reject. + """ + parsed: list[tuple[str, int, int, dict[str, Any]]] = [] + for duplex in appends: + try: + parsed.append((*_append_identity(duplex), duplex)) + except ValueError: + continue + # A cancel can put a session's aborted epoch and its restarted epoch in + # one step. Only the newest epoch is live: admitting it closes the old + # one, so the old append must neither be encoded nor re-leased. + newest: dict[str, int] = {} + for session_id, epoch, _, _ in parsed: + newest[session_id] = max(epoch, newest.get(session_id, epoch)) + for session_id, epoch in self.sessions: + if session_id in newest: + newest[session_id] = max(epoch, newest[session_id]) + rows: list[tuple[PersonaPlexStage0SessionState, tuple[int, int], np.ndarray]] = [] + for session_id, epoch, seq, duplex in parsed: + if epoch != newest[session_id]: + continue + try: + state = self._session_state(session_id, epoch) + except PersonaPlexStage0CapacityError: + # Left for prepare_append to reject. Any other error (e.g. the + # shared codec failing to initialize) propagates immediately. + continue + identity = (epoch, seq) + if identity in (state.prepared_identity, state.encoded_identity) or seq <= state.last_seq: + continue + if any(row[0] is state for row in rows): + continue + rows.append((state, identity, self._decode_pcm(duplex.get("payload")))) + if rows: + self._encode_rows(rows) def prepare_append( self, @@ -163,30 +226,15 @@ class PersonaPlexStage0DuplexRuntime: ) -> 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") - epoch = _coerce_non_negative_int(duplex.get("epoch"), "epoch") - seq = _coerce_positive_int(duplex.get("seq"), "seq") + session_id, epoch, seq = _append_identity(duplex) identity = (epoch, seq) key = (session_id, epoch) - - state = self.sessions.get(key) - if state is None: - # A newer epoch supersedes the session's earlier lockstep state: - # the engine aborted that request, but its finish notification - # may still be in flight, so release it here rather than let the - # two epochs share the codec budget. - for stale_key in [k for k in self.sessions if k[0] == session_id and k[1] < epoch]: - self.close_session(*stale_key) - 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, - epoch=epoch, - codec=self._acquire_codec(), - ) - self.sessions[key] = state + try: + state = self._session_state(session_id, epoch) + except PersonaPlexStage0StaleEpochError: + if request_id: + self._stale_requests.add(request_id) + raise if request_id: state.request_ids.add(request_id) self.request_sessions[request_id] = key @@ -195,17 +243,11 @@ class PersonaPlexStage0DuplexRuntime: 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() + if state.encoded_identity != identity: + self._encode_rows([(state, identity, self._decode_pcm(duplex.get("payload")))]) + user_frame = state.encoded_frame + 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) runtime_config = duplex.get("runtime_config") @@ -342,6 +384,8 @@ class PersonaPlexStage0DuplexRuntime: """Commit one sampled temporal frame for the next live append.""" import torch + if request_id in self._stale_requests: + return 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}") @@ -376,6 +420,7 @@ class PersonaPlexStage0DuplexRuntime: state.sampled_identity = state.prepared_identity def close_request(self, request_id: str) -> None: + self._stale_requests.discard(request_id) key = self.request_sessions.pop(request_id, None) if key is None: return @@ -393,28 +438,75 @@ class PersonaPlexStage0DuplexRuntime: 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 + if state.slot is not None: + self._shared_codec().reset_slot(state.slot) + self._free_slots.append(state.slot) + state.slot = None - def _acquire_codec(self): - if self._free_codecs: - return self._free_codecs.pop() + def _session_state(self, session_id: str, epoch: int) -> PersonaPlexStage0SessionState: + key = (session_id, epoch) + state = self.sessions.get(key) + if state is not None: + return state + # A newer epoch supersedes the session's earlier lockstep state: the + # engine aborted that request, but its finish notification may still + # be in flight, so release it here rather than let the two epochs share + # the encoder budget. + # The reverse also holds: once a newer epoch is live, an append of an + # older epoch is a leftover of the aborted request and must not re-lease + # a row (at capacity it would fail the whole step). + newer = [k[1] for k in self.sessions if k[0] == session_id and k[1] > epoch] + if newer: + raise PersonaPlexStage0StaleEpochError( + f"PersonaPlex Stage 0 epoch {epoch} of session {session_id} is superseded by epoch {max(newer)}" + ) + for stale_key in [k for k in self.sessions if k[0] == session_id and k[1] < epoch]: + self.close_session(*stale_key) + if len(self.sessions) >= self.max_sessions or not self._free_slots: + raise PersonaPlexStage0CapacityError( + f"PersonaPlex Stage 0 session capacity {self.max_sessions} is exhausted" + ) + self._shared_codec() + state = PersonaPlexStage0SessionState(session_id=session_id, epoch=epoch, slot=self._free_slots.pop()) + self.sessions[key] = state + return state + + def _encode_rows(self, rows: list[tuple[PersonaPlexStage0SessionState, tuple[int, int], np.ndarray]]) -> None: + import torch + + codec = self._shared_codec() + pcm = torch.zeros((self.max_sessions, _FRAME_SAMPLES), dtype=torch.float32) + active = torch.zeros((self.max_sessions,), dtype=torch.bool) + for state, _, samples in rows: + assert state.slot is not None + pcm[state.slot] = torch.from_numpy(samples) + active[state.slot] = True + encoded = codec.encode_frame(pcm, active) + # 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_identity = identity + + def _shared_codec(self): + if self._codec is not None: + return self._codec 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, - ) + else: + 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) + 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(self.max_sessions) + self._codec = codec return codec def _load_tokenizer(self): @@ -453,6 +545,15 @@ class PersonaPlexStage0DuplexRuntime: return pcm_f32le_samples(raw) +def _append_identity(duplex: dict[str, Any]) -> tuple[str, int, int]: + session_id = duplex.get("session_id") + if not isinstance(session_id, str) or not session_id: + raise ValueError("PersonaPlex duplex append requires session_id") + epoch = _coerce_non_negative_int(duplex.get("epoch"), "epoch") + seq = _coerce_positive_int(duplex.get("seq"), "seq") + return session_id, epoch, seq + + def _coerce_non_negative_int(value: object, name: str) -> int: try: result = int(value) diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_depformer.py b/vllm_omni/model_executor/models/personaplex/personaplex_depformer.py index 2a9a12905..011d74b9f 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_depformer.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_depformer.py @@ -203,8 +203,9 @@ class PersonaPlexDepformer(nn.Module): audio_tokens: torch.Tensor | None = None, audio_provided: torch.Tensor | None = None, return_logits: bool = False, + num_steps: int | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """Predict the ``dep_q`` audio codes for one frame. + """Predict the first ``num_steps`` (default ``dep_q``) audio codes for one frame. Args: text_token: ``[B]`` sampled text token for the frame. @@ -213,19 +214,26 @@ class PersonaPlexDepformer(nn.Module): 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]``. + return_logits: also return the per-step logits ``[B, num_steps, card]``. + num_steps: stop after this many inner steps. The loop is causal, so the + first ``num_steps`` codes are identical to a full ``dep_q`` run; + the trailing steps predict the user's codebooks, which the serving + path never consumes. Returns: - ``[B, dep_q]`` long tensor of sampled (greedy) audio codes, plus the + ``[B, num_steps]`` 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)}") + steps = self.dep_q if num_steps is None else int(num_steps) + if not 1 <= steps <= self.dep_q: + raise ValueError(f"num_steps must be in [1, {self.dep_q}]; got {num_steps}") 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): + for step in range(steps): x = self.depformer_in[step](transformer_out) # [B, 1, dim] if step == 0: cond = self.depformer_text_emb(prev.unsqueeze(1)) @@ -243,9 +251,9 @@ class PersonaPlexDepformer(nn.Module): prev = torch.where(audio_provided[:, step], audio_tokens[:, step], sampled) else: prev = sampled - stacked = torch.stack(codes, dim=1) # [B, dep_q] + stacked = torch.stack(codes, dim=1) # [B, steps] if return_logits: - return stacked, torch.stack(logits_all, dim=1) # [B, dep_q, card] + return stacked, torch.stack(logits_all, dim=1) # [B, steps, card] return stacked # ------------------------------------------------------------------ diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py index 356694380..af56c7969 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py @@ -249,8 +249,11 @@ class _MimiStreamingTransformer(nn.Module): self._offset.zero_() def reset_slot(self, b: int) -> None: + # A recycled row restarts at position 0, exactly like a fresh stream, + # instead of carrying its predecessor's absolute RoPE positions. for kv in self._kv: - kv.reset_slot(b) + kv.reset_row(b) + self._offset[b] = 0 def step(self, x: torch.Tensor, active: torch.Tensor) -> torch.Tensor: """``x`` is ``[B, T, dim]`` (T = positions this frame, typically 2).""" diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_talker.py b/vllm_omni/model_executor/models/personaplex/personaplex_talker.py index 0ea3e4f12..5ccc48f7d 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_talker.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_talker.py @@ -173,7 +173,9 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module): 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])] + # Keep every token row: the runner indexes this tensor with token-space + # logits indices, and a step can mix a new session's multi-row prefill + # with one-row live appends, so the audio row count is not the token count. return OmniOutput(text_hidden_states=hidden, multimodal_outputs={"codes": {"audio": audio_codes}}) # ------------------------------------------------------------------ @@ -323,11 +325,18 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module): 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), + from vllm_omni.model_executor.models.personaplex.duplex.stage0 import ( + PersonaPlexStage0StaleEpochError, ) + + try: + 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), + ) + except PersonaPlexStage0StaleEpochError: + return self._stale_append_passthrough(input_ids, span) offset_raw = info_dict.get("duplex_token_offset", 0) try: offset = max(0, int(offset_raw)) @@ -435,6 +444,52 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module): self._personaplex_duplex_stage0_runtime = runtime return runtime + def _stale_append_passthrough( + self, input_ids: torch.Tensor, span: int + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Neutral inputs for a request of a superseded epoch that is still in this step. + + The engine already aborted it and discards its output; it only has to keep + the batch shapes valid without touching the live session's encoder row. + """ + from vllm_omni.model_executor.models.personaplex.duplex.policy import SILENCE_TOKENS + + device = input_ids.device + silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long) + embeds = torch.zeros((span, self.mtp_hidden_size), device=device, dtype=self._dtype) + return ( + input_ids, + embeds, + { + "pplex_depformer_audio_tokens": torch.cat([silence, silence]), + "pplex_depformer_audio_provided": torch.zeros(2 * silence.numel(), dtype=torch.bool), + "duplex": {"stage0_stale": True}, + }, + ) + + def preprocess_batch( + self, + *, + req_ids: list[str], + model_intermediate_buffer: dict[str, dict[str, Any]], + device: torch.device, + ) -> None: + """Encode every live duplex append of this step in one shared-encoder call.""" + del device + appends: list[dict[str, Any]] = [] + for req_id in req_ids: + info = model_intermediate_buffer.get(req_id) + if not isinstance(info, dict): + continue + duplex = info.get("duplex") + if not isinstance(duplex, dict): + additional = info.get("additional_information") + duplex = additional.get("duplex") if isinstance(additional, dict) else None + if isinstance(duplex, dict) and duplex.get("data_plane") is True: + appends.append(duplex) + if appends: + self._duplex_stage0_runtime().encode_appends(appends) + 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: @@ -526,6 +581,7 @@ class PersonaPlexTalkerForConditionalGeneration(nn.Module): 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): diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py b/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py index 28759dafb..38887917a 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py @@ -92,6 +92,12 @@ class _RingKV: # first visible entry. (LM sacrifice-tick +1 is applied by the caller.) self.start_offset[b] = self.end_offset[b] + def reset_row(self, b: int) -> None: + # Restart row b at position 0. Every cached entry of the row sits at or + # past the new end offset, so all of them are masked until overwritten. + self.end_offset[b] = 0 + self.start_offset[b] = 0 + def bump_slot_start(self, b: int) -> None: self.start_offset[b] += 1