[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>

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])

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, :]

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]