From 825ab20b07268050015516f9a65c733d7d089d28 Mon Sep 17 00:00:00 2001 From: MrDongsls Date: Thu, 17 Sep 2026 16:50:40 +0800 Subject: [PATCH] [Model][PersonaPlex] add per-row RingKV offsets and active rows (#7670) Signed-off-by: Mr.Dong <1042542469@qq.com> Co-authored-by: Mr.Dong <1042542469@qq.com> --- .../personaplex/test_ring_per_row_offsets.py | 369 ++++++++++++++++++ .../models/personaplex/personaplex_mimi.py | 94 +++-- .../personaplex/personaplex_temporal.py | 76 ++-- 3 files changed, 483 insertions(+), 56 deletions(-) create mode 100644 tests/model_executor/models/personaplex/test_ring_per_row_offsets.py 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 new file mode 100644 index 000000000..fce42f40c --- /dev/null +++ b/tests/model_executor/models/personaplex/test_ring_per_row_offsets.py @@ -0,0 +1,369 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +from pytest_mock import MockerFixture + +from vllm_omni.model_executor.models.personaplex.personaplex_mimi import ( + PersonaPlexMimiCodec, + _MimiStreamingTransformer, + _StreamConv1d, + _StreamConvTr1d, +) +from vllm_omni.model_executor.models.personaplex.personaplex_temporal import ( + PersonaPlexTemporalStreaming, + _RingKV, +) + +pytestmark = pytest.mark.core_model + +SEED = 1234 +CUDA_DEVICE = torch.device("cuda") + + +def _mask(*rows: bool, device: torch.device | str = "cpu") -> torch.Tensor: + return torch.tensor(rows, dtype=torch.bool, device=device) + + +def _inputs(shape: tuple[int, ...], seed: int, device: torch.device) -> torch.Tensor: + generator = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(*shape, generator=generator).to(device) + + +def _make_mimi_transformer( + batch_size: int, + device: torch.device, + context: int = 6, +) -> _MimiStreamingTransformer: + torch.manual_seed(SEED) + transformer = _MimiStreamingTransformer( + num_layers=2, + dim=32, + num_heads=4, + context=context, + ).to(device) + for parameter in transformer.parameters(): + if parameter.dtype.is_floating_point: + nn.init.normal_(parameter, std=0.1) + transformer.streaming_init(batch_size) + return transformer + + +def _make_temporal( + batch_size: int, + device: torch.device, + context: int = 6, +) -> PersonaPlexTemporalStreaming: + torch.manual_seed(SEED) + temporal = PersonaPlexTemporalStreaming( + dim=16, + num_layers=2, + num_heads=4, + hidden=32, + context=context, + text_card=11, + ).to(device) + for parameter in temporal.parameters(): + if parameter.dtype.is_floating_point: + nn.init.normal_(parameter, std=0.1) + temporal.streaming_init(batch_size) + return temporal + + +def _assert_valid_ring_row_matches( + batched: _RingKV, + reference: _RingKV, + batched_positions: torch.Tensor, + reference_positions: torch.Tensor, + row: int, +) -> None: + assert torch.equal(batched_positions[row], reference_positions) + assert torch.equal(batched.end_offset[row], reference.end_offset[0]) + + # The active mask covers the physical write as well as offset advancement, + # so the complete per-row ring state remains singleton-identical. + assert torch.equal(batched.cache[:, row], reference.cache[:, 0]) + + +def _stream_carry(stream: _StreamConv1d | _StreamConvTr1d) -> torch.Tensor: + if isinstance(stream, _StreamConv1d): + assert stream.prev is not None + return stream.prev + assert stream.partial is not None + return stream.partial + + +def _make_passthrough_stage(mocker: MockerFixture): + stage = mocker.Mock(spec=["reset", "reset_slot"]) + stage.side_effect = lambda x, active: x + return stage + + +def _make_passthrough_transformer(mocker: MockerFixture): + transformer = mocker.Mock(spec=["streaming_init", "reset_streaming", "reset_slot", "step"]) + transformer.step.side_effect = lambda x, active: x + return transformer + + +def _active_args(call_args_list) -> torch.Tensor: + return torch.stack([call.args[1] for call in call_args_list]) + + +def _make_codec_stub(mocker: MockerFixture) -> PersonaPlexMimiCodec: + codec = PersonaPlexMimiCodec.__new__(PersonaPlexMimiCodec) + nn.Module.__init__(codec) + + quantizer = mocker.Mock(spec=["encode", "decode"]) + quantizer.encode.side_effect = lambda x: torch.zeros(8, x.shape[0], x.shape[-1], dtype=torch.long) + quantizer.decode.side_effect = lambda codes: torch.zeros(codes.shape[0], 1, codes.shape[-1]) + + codec.device = torch.device("cpu") + codec.dtype = torch.float32 + codec.model = SimpleNamespace(quantizer=quantizer) + codec._enc_stages = [] + codec._dec_stages = [] + codec._downsample = _make_passthrough_stage(mocker) + codec._upsample = _make_passthrough_stage(mocker) + codec.encoder_transformer = _make_passthrough_transformer(mocker) + codec.decoder_transformer = _make_passthrough_transformer(mocker) + codec.streaming_init(batch_size=2) + return codec + + +@pytest.mark.cpu +def test_codec_entrypoints_forward_active_mask(mocker: MockerFixture) -> None: + codec = _make_codec_stub(mocker) + active = _mask(True, False) + all_active = torch.ones_like(active) + + codec.encode_frame(torch.zeros(2, 1920), active) + codec.decode_frame(torch.zeros(2, 8), active) + codec.decode_frames(torch.zeros(2, 8, 3), active) + codec.decode_frames(torch.zeros(2, 8, 3), active=None) + + assert torch.equal(_active_args(codec._downsample.call_args_list), active[None]) + assert torch.equal(_active_args(codec._upsample.call_args_list), torch.stack((active, active, all_active))) + assert torch.equal(_active_args(codec.encoder_transformer.step.call_args_list), active[None]) + assert torch.equal( + _active_args(codec.decoder_transformer.step.call_args_list), + torch.stack((active, active, all_active)), + ) + + +@pytest.mark.cpu +def test_ring_mixed_offsets_match_singleton_streams() -> None: + batch_size, heads, head_dim, capacity, tokens = 3, 2, 3, 6, 2 + batched = _RingKV(batch_size, heads, head_dim, capacity, torch.device("cpu"), torch.float32) + singletons = [_RingKV(1, heads, head_dim, capacity, torch.device("cpu"), torch.float32) for _ in range(batch_size)] + reference_positions: list[torch.Tensor | None] = [None] * batch_size + active_schedule = [ + _mask(True, True, True), + _mask(True, False, False), + _mask(False, False, True), + _mask(True, True, False), + _mask(False, True, True), + _mask(True, False, True), + _mask(True, True, True), + ] * 2 + + torch.manual_seed(SEED) + for active in active_schedule: + keys = torch.randn(batch_size, heads, tokens, head_dim) + values = torch.randn_like(keys) + _, _, positions = batched.complete(keys, values, active) + + for row, is_active in enumerate(active.tolist()): + if is_active: + _, _, singleton_positions = singletons[row].complete( + keys[row : row + 1], + values[row : row + 1], + _mask(True), + ) + reference_positions[row] = singleton_positions[0].clone() + + assert reference_positions[row] is not None + _assert_valid_ring_row_matches( + batched, + singletons[row], + positions, + reference_positions[row], + row, + ) + + assert any(int(offset) > capacity for offset in batched.end_offset) + assert not torch.equal(batched.end_offset[0], batched.end_offset[1]) + + +@pytest.mark.cpu +def test_ring_slot_recycle_masks_old_history() -> None: + ring = _RingKV(2, 1, 1, 8, torch.device("cpu"), torch.float32) + for _ in range(4): + ring.complete(torch.randn(2, 1, 2, 1), torch.randn(2, 1, 2, 1), _mask(True, True)) + + old_end = ring.end_offset.clone() + ring.reset_slot(1) + _, _, positions = ring.complete( + torch.randn(2, 1, 2, 1), + torch.randn(2, 1, 2, 1), + _mask(False, True), + ) + visible = positions[1] >= 0 + assert torch.all(positions[1][visible] >= old_end[1]) + + ring.reset_slot(0) + ring.bump_slot_start(0) + _, _, positions = ring.complete( + torch.randn(2, 1, 2, 1), + torch.randn(2, 1, 2, 1), + _mask(True, False), + ) + visible = positions[0] >= 0 + assert torch.all(positions[0][visible] >= old_end[0] + 1) + + +@pytest.mark.cpu +def test_temporal_streaming_rejects_invalid_active_shape() -> None: + temporal = _make_temporal(2, torch.device("cpu")) + + with pytest.raises(ValueError, match=r"active must have shape \(2,\)"): + temporal.step(torch.zeros(2, 1, 16), torch.ones(1)) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mimi_transformer_mixed_offsets_match_singletons() -> None: + batch_size, tokens, dim = 3, 2, 32 + batched = _make_mimi_transformer(batch_size, CUDA_DEVICE) + singletons = [_make_mimi_transformer(1, CUDA_DEVICE) for _ in range(batch_size)] + active_schedule = [ + _mask(True, True, True, device=CUDA_DEVICE), + _mask(True, False, False, device=CUDA_DEVICE), + _mask(False, False, True, device=CUDA_DEVICE), + _mask(True, True, False, device=CUDA_DEVICE), + _mask(False, True, True, device=CUDA_DEVICE), + _mask(True, False, True, device=CUDA_DEVICE), + _mask(True, True, True, device=CUDA_DEVICE), + ] * 2 + + for step, active in enumerate(active_schedule): + inputs = _inputs((batch_size, tokens, dim), SEED + step, CUDA_DEVICE) + offsets_before = batched._offset.clone() + cache_before = [kv.cache.clone() for kv in batched._kv] + output = batched.step(inputs, active) + + for row, is_active in enumerate(active.tolist()): + if is_active: + expected = singletons[row].step( + inputs[row : row + 1], + _mask(True, device=CUDA_DEVICE), + ) + torch.testing.assert_close(output[row : row + 1], expected, rtol=0.0, atol=0.0) + assert torch.equal(batched._offset[row], singletons[row]._offset[0]) + for batched_kv, singleton_kv in zip(batched._kv, singletons[row]._kv): + assert torch.equal(batched_kv.end_offset[row], singleton_kv.end_offset[0]) + torch.testing.assert_close(batched_kv.cache[:, row], singleton_kv.cache[:, 0], rtol=0.0, atol=0.0) + else: + assert torch.equal(batched._offset[row], offsets_before[row]) + for kv, previous_cache in zip(batched._kv, cache_before): + assert torch.equal(kv.end_offset[row], offsets_before[row]) + assert torch.equal(kv.cache[:, row], previous_cache[:, row]) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_temporal_streaming_mixed_offsets_match_isolated_rows() -> None: + # Keep the reference shape at B=2 so exact comparisons use the same CUDA + # kernel shape while the target row runs independently. + batched = _make_temporal(2, CUDA_DEVICE) + references = [_make_temporal(2, CUDA_DEVICE) for _ in range(2)] + active_schedule = [ + _mask(True, True, device=CUDA_DEVICE), + _mask(True, False, device=CUDA_DEVICE), + _mask(False, False, device=CUDA_DEVICE), + _mask(False, True, device=CUDA_DEVICE), + _mask(True, True, device=CUDA_DEVICE), + _mask(True, False, device=CUDA_DEVICE), + _mask(False, True, device=CUDA_DEVICE), + _mask(True, True, device=CUDA_DEVICE), + ] * 2 + + for step, active in enumerate(active_schedule): + inputs = _inputs((2, 1, 16), SEED + 100 + step, CUDA_DEVICE) + offsets_before = batched._offset.clone() + cache_before = [kv.cache.clone() for kv in batched._kv] + output, logits = batched.step(inputs, active) + + for row, (reference, is_active) in enumerate(zip(references, active.tolist())): + if is_active: + reference_inputs = torch.zeros_like(inputs) + reference_inputs[row] = inputs[row] + reference_active = torch.zeros_like(active) + reference_active[row] = True + expected_output, expected_logits = reference.step(reference_inputs, reference_active) + torch.testing.assert_close(output[row : row + 1], expected_output[row : row + 1], rtol=0.0, atol=0.0) + torch.testing.assert_close(logits[row : row + 1], expected_logits[row : row + 1], rtol=0.0, atol=0.0) + assert torch.equal(batched._offset[row], reference._offset[row]) + for batched_kv, reference_kv in zip(batched._kv, reference._kv): + assert torch.equal(batched_kv.end_offset[row], reference_kv.end_offset[row]) + torch.testing.assert_close(batched_kv.cache[:, row], reference_kv.cache[:, row], rtol=0.0, atol=0.0) + else: + assert torch.equal(batched._offset[row], offsets_before[row]) + for kv, previous_cache in zip(batched._kv, cache_before): + assert torch.equal(kv.end_offset[row], offsets_before[row]) + assert torch.equal(kv.cache[:, row], previous_cache[:, row]) + + +@pytest.mark.parametrize("kind", ["conv", "convtr"]) +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_mimi_conv_carries_preserve_inactive_rows(kind: str) -> None: + if kind == "conv": + conv = nn.Conv1d(1, 2, kernel_size=3, stride=2, device=CUDA_DEVICE) + batched = _StreamConv1d(conv, pad_mode="replicate") + row0 = _StreamConv1d(conv, pad_mode="replicate") + row1 = _StreamConv1d(conv, pad_mode="replicate") + samples = 4 + else: + conv = nn.ConvTranspose1d(1, 2, kernel_size=4, stride=2, device=CUDA_DEVICE) + batched = _StreamConvTr1d(conv) + row0 = _StreamConvTr1d(conv) + row1 = _StreamConvTr1d(conv) + samples = 2 + + for stream, batch_size in ((batched, 2), (row0, 1), (row1, 1)): + stream.reset(batch_size, CUDA_DEVICE, torch.float32) + + singletons = [row0, row1] + active_schedule = [ + _mask(False, False, device=CUDA_DEVICE), + _mask(True, False, device=CUDA_DEVICE), + _mask(True, True, device=CUDA_DEVICE), + _mask(False, True, device=CUDA_DEVICE), + _mask(True, True, device=CUDA_DEVICE), + ] + + for step, active in enumerate(active_schedule): + inputs = _inputs((2, 1, samples), SEED + 200 + step, CUDA_DEVICE) + carry_before = _stream_carry(batched).clone() + fresh_before = batched._fresh.clone() + output = PersonaPlexMimiCodec._run_stages(inputs, [(kind, batched)], active) + + for row, is_active in enumerate(active.tolist()): + if is_active: + expected = PersonaPlexMimiCodec._run_stages( + inputs[row : row + 1], + [(kind, singletons[row])], + _mask(True, device=CUDA_DEVICE), + ) + torch.testing.assert_close(output[row : row + 1], expected, rtol=0.0, atol=0.0) + torch.testing.assert_close( + _stream_carry(batched)[row], _stream_carry(singletons[row])[0], rtol=0.0, atol=0.0 + ) + assert torch.equal(batched._fresh[row], singletons[row]._fresh[0]) + else: + assert torch.equal(_stream_carry(batched)[row], carry_before[row]) + assert torch.equal(batched._fresh[row], fresh_before[row]) diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py index 3b8813faf..356694380 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py @@ -43,6 +43,19 @@ FRAME_SIZE = 1920 CODEBOOKS = 8 +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]: + # True advances that row's streaming state; False keeps an absent or padded + # row's offsets and convolution carries unchanged. + if active is None: + return all_active + if active.shape != all_active.shape: + raise ValueError(f"active must have shape {tuple(all_active.shape)}, got {tuple(active.shape)}") + return active.to(device=all_active.device, dtype=torch.bool) + + def _map_moshi_codec_weights( state_dict: dict[str, torch.Tensor], ) -> dict[str, torch.Tensor]: @@ -107,23 +120,24 @@ class _StreamConv1d: def reset(self, batch_size: int, device, dtype) -> None: pad = self.kernel - self.stride self.prev = torch.zeros(batch_size, self.conv.in_channels, pad, device=device, dtype=dtype) - self._fresh = torch.ones(batch_size, dtype=torch.bool) + self._fresh = torch.ones(batch_size, dtype=torch.bool, device=device) def reset_slot(self, b: int) -> None: self.prev[b].zero_() self._fresh[b] = True - def __call__(self, x: torch.Tensor) -> torch.Tensor: - if self.pad_mode == "replicate" and bool(self._fresh.any()): + def __call__(self, x: torch.Tensor, active: torch.Tensor) -> torch.Tensor: + if self.pad_mode == "replicate": pad = self.prev.shape[-1] edge = x[..., 0:1].expand(-1, -1, pad) - fresh = self._fresh.to(x.device).view(-1, 1, 1) + fresh = (self._fresh & active).view(-1, 1, 1) self.prev = torch.where(fresh, edge.to(self.prev.dtype), self.prev) - self._fresh[:] = False + self._fresh[active] = False x = torch.cat([self.prev, x], dim=-1) t = x.shape[-1] num_frames = max(0, (t - self.kernel) // self.stride + 1) - self.prev = x[..., num_frames * self.stride :] + prev = x[..., num_frames * self.stride :] + self.prev = torch.where(active.view(-1, 1, 1), prev, 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]) @@ -137,18 +151,19 @@ class _StreamConvTr1d: self.kernel = conv.kernel_size[0] self.stride = conv.stride[0] self.partial: torch.Tensor | None = None + self._fresh: torch.Tensor | None = None def reset(self, batch_size: int, device, dtype) -> None: self.partial = torch.zeros( batch_size, self.conv.out_channels, self.kernel - self.stride, device=device, dtype=dtype ) - self._fresh = torch.ones(batch_size, dtype=torch.bool) + self._fresh = torch.ones(batch_size, dtype=torch.bool, device=device) def reset_slot(self, b: int) -> None: self.partial[b].zero_() self._fresh[b] = True - def __call__(self, x: torch.Tensor) -> torch.Tensor: + def __call__(self, x: torch.Tensor, active: torch.Tensor) -> torch.Tensor: out = self.conv(x) length = out.shape[-1] tail = self.kernel - self.stride @@ -159,10 +174,11 @@ class _StreamConvTr1d: # it again, so subtract one copy -- except on a row's very first # frame, where the carry is zeros by construction. merge = merge - self.conv.bias[:, None] - merge[self._fresh] = 0.0 - self._fresh[:] = False + merge[self._fresh & active] = 0.0 + self._fresh[active] = False out[..., :pt] += merge - self.partial = out[..., length - tail :].clone() + partial = out[..., length - tail :].clone() + self.partial = torch.where(active.view(-1, 1, 1), partial, self.partial) return out[..., : length - tail] @@ -183,16 +199,23 @@ class _MimiTransformerLayer(nn.Module): self.scale1 = nn.Parameter(torch.empty(dim)) self.scale2 = nn.Parameter(torch.empty(dim)) - def forward(self, x: torch.Tensor, kv: _RingKV, offset: torch.Tensor, context: int) -> torch.Tensor: + def forward( + self, + x: torch.Tensor, + kv: _RingKV, + offset: torch.Tensor, + context: int, + active: torch.Tensor, + ) -> torch.Tensor: B, T, _ = x.shape h = self.norm1(x) qkv = F.linear(h, self.in_proj_weight) qkv = qkv.view(B, T, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] q, k = _apply_rope(q, k, offset) - keys, values, pos_k = kv.complete(k, v) + keys, values, pos_k = kv.complete(k, v, active=active) pos_k = pos_k.view(pos_k.shape[0], 1, pos_k.shape[1]) - pos_q = offset + torch.arange(T, device=q.device, dtype=torch.long).view(1, -1, 1) + pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(1, -1, 1) delta = pos_q - pos_k attn_bias = (pos_k >= 0) & (delta >= 0) & (delta < context) attn = F.scaled_dot_product_attention(q, keys, values, attn_bias.unsqueeze(1), dropout_p=0.0) @@ -218,7 +241,7 @@ class _MimiStreamingTransformer(nn.Module): heads = self.layers[0].num_heads hd = self.layers[0].head_dim self._kv = [_RingKV(batch_size, heads, hd, self.context, p.device, p.dtype) for _ in self.layers] - self._offset = torch.zeros(1, device=p.device, dtype=torch.long) + self._offset = torch.zeros(batch_size, device=p.device, dtype=torch.long) def reset_streaming(self) -> None: for kv in self._kv: @@ -229,11 +252,11 @@ class _MimiStreamingTransformer(nn.Module): for kv in self._kv: kv.reset_slot(b) - def step(self, x: torch.Tensor) -> torch.Tensor: + def step(self, x: torch.Tensor, active: torch.Tensor) -> torch.Tensor: """``x`` is ``[B, T, dim]`` (T = positions this frame, typically 2).""" for layer, kv in zip(self.layers, self._kv): - x = layer(x, kv, self._offset, self.context) - self._offset.add_(x.shape[1]) + x = layer(x, kv, self._offset, self.context, active) + self._offset.add_(x.shape[1] * active.to(self._offset.dtype)) return x def load_weights(self, state_dict: dict[str, torch.Tensor], prefix: str) -> int: @@ -339,6 +362,7 @@ class PersonaPlexMimiCodec(nn.Module): self._upsample = _StreamConvTr1d(m.upsample.conv) self._dec_stages = _walk_seanet(m.decoder.layers) self._batch_size: int | None = None + self._all_active: torch.Tensor # -- streaming state ------------------------------------------------------ @@ -354,6 +378,7 @@ class PersonaPlexMimiCodec(nn.Module): def streaming_init(self, batch_size: int) -> None: self._batch_size = batch_size + 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) self.encoder_transformer.streaming_init(batch_size) @@ -376,39 +401,42 @@ class PersonaPlexMimiCodec(nn.Module): # -- per-frame codec ------------------------------------------------------- @staticmethod - def _run_stages(x: torch.Tensor, stages) -> torch.Tensor: + def _run_stages(x: torch.Tensor, stages, active: torch.Tensor) -> torch.Tensor: for kind, stage in stages: if kind == "res": act0, conv1, act2, conv3 = stage - x = x + conv3(act2(conv1(act0(x)))) + x = x + conv3(act2(conv1(act0(x), active)), active) else: - x = stage(x) + x = stage(x, active) if kind in {"conv", "convtr"} else stage(x) return x @torch.no_grad() - def encode_frame(self, pcm: torch.Tensor) -> torch.Tensor: + 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) - x = self._run_stages(x, self._enc_stages) - x = self.encoder_transformer.step(x.transpose(1, 2)).transpose(1, 2) - x = self._downsample(x) + active = _normalize_active(active, self._all_active) + 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) # [Q, B, T] return codes[:CODEBOOKS, :, 0].transpose(0, 1).contiguous() @torch.no_grad() - def decode_frame(self, codes: torch.Tensor) -> torch.Tensor: + 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)) - emb = self._upsample(emb) - emb = self.decoder_transformer.step(emb.transpose(1, 2)).transpose(1, 2) - x = self._run_stages(emb, self._dec_stages) + 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, :] @torch.no_grad() - def decode_frames(self, codes: torch.Tensor) -> torch.Tensor: + 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)) - emb = self._upsample(emb) - emb = self.decoder_transformer.step(emb.transpose(1, 2)).transpose(1, 2) - x = self._run_stages(emb, self._dec_stages) + 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, :] diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py b/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py index 717e4d505..28759dafb 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_temporal.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_temporal.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 """Streaming Helium temporal transformer for PersonaPlex (moshi-free, plain torch). Frame-clocked duplex needs a stateful per-frame forward: one 80 ms step consumes @@ -49,8 +49,8 @@ def _apply_rope(q: torch.Tensor, k: torch.Tensor, offset: torch.Tensor, max_peri B, H, T, D = q.shape ds = torch.arange(D // 2, device=q.device, dtype=torch.float32) freqs = torch.exp(ds * (-math.log(max_period) * 2 / D)) - ts = offset.float() + torch.arange(T, device=q.device, dtype=torch.float32) - ts = ts.view(1, -1, 1) + ts = offset.float().view(-1, 1) + torch.arange(T, device=q.device, dtype=torch.float32) + ts = ts.view(-1, 1, T, 1) dims = q.shape[:-1] q = q.view(*dims, D // 2, 2) @@ -80,7 +80,7 @@ class _RingKV: def __init__(self, batch_size: int, num_heads: int, dim_per_head: int, capacity: int, device, dtype): self.capacity = capacity self.cache = torch.zeros((2, batch_size, num_heads, capacity, dim_per_head), device=device, dtype=dtype) - self.end_offset = torch.zeros(1, device=device, dtype=torch.long) + self.end_offset = torch.zeros(batch_size, device=device, dtype=torch.long) self.start_offset = torch.zeros(batch_size, device=device, dtype=torch.long) def reset(self) -> None: @@ -90,23 +90,36 @@ class _RingKV: def reset_slot(self, b: int) -> None: # Mask everything written so far for row b; the row's next write is its # first visible entry. (LM sacrifice-tick +1 is applied by the caller.) - self.start_offset[b] = self.end_offset.clone() + self.start_offset[b] = self.end_offset[b] def bump_slot_start(self, b: int) -> None: self.start_offset[b] += 1 - def complete(self, k: torch.Tensor, v: torch.Tensor): + def complete(self, k: torch.Tensor, v: torch.Tensor, active: torch.Tensor): B, H, T, D = k.shape - indexes = torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype) + self.end_offset - indexes = indexes % self.capacity - self.cache[0].index_copy_(2, indexes, k) - self.cache[1].index_copy_(2, indexes, v) - self.end_offset.add_(T) + indexes = ( + torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype).view(1, -1) + + self.end_offset.view(-1, 1) + ) % self.capacity + idx4 = indexes.view(B, 1, T, 1).expand(-1, H, -1, D) + # Keep inactive rows completely inert. Once the ring is full, the + # physical future-write slot is also addressable by the position mask; + # writing it for an inactive row would therefore leak padded data into + # its next attention step. + active_view = active.view(B, 1, 1, 1) + old_k = self.cache[0].gather(2, idx4) + old_v = self.cache[1].gather(2, idx4) + k = torch.where(active_view, k, old_k) + v = torch.where(active_view, v, old_v) + self.cache[0].scatter_(2, idx4, k) + self.cache[1].scatter_(2, idx4, v) + self.end_offset.add_(T * active.to(self.end_offset.dtype)) idx = torch.arange(self.capacity, device=self.end_offset.device, dtype=torch.long) - invalid = idx >= self.end_offset - end_index = self.end_offset % self.capacity - delta = idx - end_index + end_offset = self.end_offset.view(-1, 1) + invalid = idx.view(1, -1) >= end_offset + end_index = end_offset % self.capacity + delta = idx.view(1, -1) - end_index # `delta <= 0` (not `< 0`) is moshi's exact convention (transformer.py # RingKVCache.complete). It labels the just-past-newest slot as the future # write position, so once the ring has wrapped the single oldest in-window @@ -114,14 +127,19 @@ class _RingKV: # verbatim from the reference and only shows after the window fills (Helium # ~3000 frames / 240 s); it costs one frame out of thousands. Do NOT change # this to `< 0`: it would diverge from moshi and break greedy bit-parity. - positions = torch.where(delta <= 0, self.end_offset + delta, self.end_offset + delta - self.capacity) + positions = torch.where(delta <= 0, end_offset + delta, end_offset + delta - self.capacity) positions = torch.where(invalid, torch.full_like(positions, -1), positions) - positions = positions.view(1, -1) # [1, capacity] below = positions < self.start_offset.view(-1, 1) # [B, capacity] positions = torch.where(below, torch.full_like(positions, -1), positions) return self.cache[0], self.cache[1], positions +def _normalize_temporal_active(active: torch.Tensor, reference: torch.Tensor) -> torch.Tensor: + if active.shape != reference.shape: + raise ValueError(f"active must have shape {tuple(reference.shape)}, got {tuple(active.shape)}") + return active.to(device=reference.device, dtype=torch.bool) + + class _TemporalLayer(nn.Module): """One Helium decoder layer in the checkpoint's fused layout.""" @@ -137,7 +155,14 @@ class _TemporalLayer(nn.Module): self.norm1_alpha = nn.Parameter(torch.ones(1, 1, dim)) self.norm2_alpha = nn.Parameter(torch.ones(1, 1, dim)) - def forward(self, x: torch.Tensor, kv: _RingKV, offset: torch.Tensor, context: int) -> torch.Tensor: + def forward( + self, + x: torch.Tensor, + kv: _RingKV, + offset: torch.Tensor, + context: int, + active: torch.Tensor, + ) -> torch.Tensor: B, T, _ = x.shape h = _rms_norm_f32(x, self.norm1_alpha, 1e-8) qkv = F.linear(h, self.in_proj_weight) @@ -147,9 +172,9 @@ class _TemporalLayer(nn.Module): q, k, v = qkv[0], qkv[1], qkv[2] q, k = _apply_rope(q, k, offset) - keys, values, pos_k = kv.complete(k, v) + keys, values, pos_k = kv.complete(k, v, active) pos_k = pos_k.view(pos_k.shape[0], 1, pos_k.shape[1]) # [B, 1, cap] - pos_q = offset + torch.arange(T, device=q.device, dtype=torch.long).view(1, -1, 1) + pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(1, -1, 1) delta = pos_q - pos_k attn_bias = (pos_k >= 0) & (delta >= 0) & (delta < context) attn_bias = attn_bias.unsqueeze(1) # [B, 1, T, cap] @@ -199,7 +224,7 @@ class PersonaPlexTemporalStreaming(nn.Module): head_dim = self.layers[0].head_dim # Capacity = context window (the mask truncates at `context` anyway). self._kv = [_RingKV(batch_size, heads, head_dim, self.context, p.device, p.dtype) for _ in self.layers] - self._offset = torch.zeros(1, device=p.device, dtype=torch.long) + self._offset = torch.zeros(batch_size, device=p.device, dtype=torch.long) def reset_streaming(self) -> None: assert self._kv is not None, "call streaming_init first" @@ -221,12 +246,17 @@ class PersonaPlexTemporalStreaming(nn.Module): # -- per-frame step ------------------------------------------------------- @torch.no_grad() - def step(self, frame_embedding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def step( + self, + frame_embedding: torch.Tensor, + active: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: assert self._kv is not None, "call streaming_init first" + active = _normalize_temporal_active(active, self._offset) x = frame_embedding for layer, kv in zip(self.layers, self._kv): - x = layer(x, kv, self._offset, self.context) - self._offset.add_(x.shape[1]) + x = layer(x, kv, self._offset, self.context, active) + self._offset.add_(x.shape[1] * active.to(self._offset.dtype)) out = _rms_norm_f32(x, self.out_norm_alpha, 1e-8) text_logits = F.linear(out, self.text_linear) return out, text_logits[:, None]