[Model] Reuse frame-local KV in MOSS-TTS Local

Cache each position once while generating the audio codebooks, preserving frame-local RoPE positions and the full-prefix reference path. Share compiled batch and position shapes to avoid specialization-limit fallback. Add one compact parity regression covering frame reuse.

Signed-off-by: Sy03 <1370724210@qq.com>

tests/model_executor/models/moss_tts/test_local_depth_rope_cache.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
 """CPU parity tests for MOSS-TTS Local Depth RoPE caching."""
 
 import pytest
@@ -42,3 +42,16 @@ def test_rope_cache_matches_reference_and_reuses_storage() -> None:
     # Smaller slices retain the fixed allocation used by CUDA Graph.
     attn.prepare_rope_cache(6, device, torch.float32)
     assert attn._rope_cos_cache.data_ptr() == cache_ptr
+
+
+def test_frame_local_kv_matches_full_prefix_and_reuses_slots() -> None:
+    model = MossTTSLocalDepthTransformer(_test_config()).eval()
+    model.h[0].attn.prepare_rope_cache(12, torch.device("cpu"), torch.float32)
+    cache = (torch.empty(2, 4, 12, 8), torch.empty(2, 4, 12, 8))
+    with torch.inference_mode():
+        for _ in range(2):
+            inputs = torch.randn(2, 12, 32)
+            for position in range(12):
+                expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:]
+                actual = model._forward_prefix(inputs[:, position : position + 1], cache, position)
+                torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5)

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License").
@@ -101,8 +103,13 @@ class _MossTTSLocalAttention(nn.Module):
         odd = x[..., 1::2]
         return torch.stack((-odd, even), dim=-1).reshape_as(x)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        """``hidden_states``: ``(B, S, H)``. Re-prefills with a fresh causal mask over ``[0, S)`` every call."""
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        position: int = 0,
+    ) -> torch.Tensor:
+        """Run a causal prefix, or one new token using frame-local K/V."""
         batch_size, seq_len, _ = hidden_states.shape
         qkv = self.c_attn(hidden_states)
         query, key, value = qkv.split(self.embed_dim, dim=-1)
@@ -110,14 +117,22 @@ class _MossTTSLocalAttention(nn.Module):
         key = key.view(batch_size, seq_len, self.n_head, self.head_dim)
         value = value.view(batch_size, seq_len, self.n_head, self.head_dim)
 
-        cos, sin = self._rope_cos_sin(seq_len, hidden_states.device, hidden_states.dtype)
+        cos, sin = self._rope_cos_sin(position + seq_len, hidden_states.device, hidden_states.dtype)
+        cos, sin = cos[:, position:], sin[:, position:]
         query = query * cos + self._rotate_half(query) * sin
         key = key * cos + self._rotate_half(key) * sin
 
         query = query.transpose(1, 2)
         key = key.transpose(1, 2)
         value = value.transpose(1, 2)
-        attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=True)
+        if kv_cache is not None:
+            assert seq_len == 1, "Cached Local Depth execution consumes one token at a time"
+            k_cache, v_cache = kv_cache
+            k_cache[:, :, position : position + 1].copy_(key)
+            v_cache[:, :, position : position + 1].copy_(value)
+            key = k_cache[:, :, : position + 1]
+            value = v_cache[:, :, : position + 1]
+        attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=kv_cache is None)
         attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
         return self.c_proj(attn_output)
 
@@ -141,8 +156,13 @@ class _MossTTSLocalBlock(nn.Module):
         self.ln_2 = nn.LayerNorm(hidden_size, eps=eps)
         self.mlp = _MossTTSLocalMLP(hidden_size, inner_size)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = hidden_states + self.attn(self.ln_1(hidden_states))
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        position: int = 0,
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.ln_1(hidden_states), kv_cache, position)
         hidden_states = hidden_states + self.mlp(self.ln_2(hidden_states))
         return hidden_states
 
@@ -156,11 +176,9 @@ class MossTTSLocalDepthTransformer(nn.Module):
         and codebook-0's head (``audio_lm_heads[0]``) simultaneously.
       - codebooks 1..n_vq-1 are sampled sequentially: each sampled code is
         re-embedded (``audio_embeddings[c]``) and appended as the next
-        position, re-prefilling the block over the growing (<=n_vq) sequence
-        with a fresh causal mask each call -- mathematically identical to
-        incremental KV-cache decoding since attention is strictly causal and
-        only the last position is ever read, but avoids any cache plumbing
-        given the trivially short sequence length.
+        position. K/V for earlier positions are reused within the frame, so
+        each position's projections and MLP run once. A new frame overwrites
+        every cache position before reading it, including CUDA Graph replay.
     """
 
     def __init__(self, gpt2_config, hidden_size: int | None = None) -> None:
@@ -174,11 +192,13 @@ class MossTTSLocalDepthTransformer(nn.Module):
         self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps)
         self._compiled_forward_prefix = None
 
-    def _forward_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor:
-        hidden_states = seq_embeds
-        for block in self.h:
-            hidden_states = block(hidden_states)
-        return self.ln_f(hidden_states)
+    def _forward_prefix(
+        self,
+        seq_embeds: torch.Tensor,
+        kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        position: int = 0,
+    ) -> torch.Tensor:
+        return self.ln_f(self.h[0](seq_embeds, kv_cache, position))
 
     def setup_compile(self) -> None:
         if self._compiled_forward_prefix is not None:
@@ -189,14 +209,18 @@ class MossTTSLocalDepthTransformer(nn.Module):
             return
         self._compiled_forward_prefix = torch.compile(
             self._forward_prefix,
-            dynamic=False,
+            # Share batch/position shapes instead of exhausting Dynamo's
+            # recompilation budget while warming request-size graph buckets.
+            dynamic=True,
             options={"epilogue_fusion": False},
         )
-        logger.info("MOSS-TTS local depth prefix enabled with torch.compile")
+        logger.info("MOSS-TTS local depth frame-local KV execution enabled with torch.compile")
 
-    def _run_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor:
+    def _run_prefix(
+        self, seq_embeds: torch.Tensor, kv_cache: tuple[torch.Tensor, torch.Tensor], position: int
+    ) -> torch.Tensor:
         forward_prefix = self._compiled_forward_prefix or self._forward_prefix
-        return forward_prefix(seq_embeds)
+        return forward_prefix(seq_embeds, kv_cache, position)
 
     @torch.no_grad()
     def generate_frame(
@@ -240,9 +264,13 @@ class MossTTSLocalDepthTransformer(nn.Module):
         for block in self.h:
             block.attn.prepare_rope_cache(n_vq, backbone_last_hidden.device, dtype)
 
-        embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype)
-        embeds[:, 0, :] = backbone_last_hidden.to(dtype)
-        hidden = self._run_prefix(embeds[:, :1, :])
+        attn = self.h[0].attn
+        cache_shape = (batch_size, attn.n_head, n_vq, attn.head_dim)
+        kv_cache = (
+            backbone_last_hidden.new_empty(cache_shape, dtype=dtype),
+            backbone_last_hidden.new_empty(cache_shape, dtype=dtype),
+        )
+        hidden = self._run_prefix(backbone_last_hidden[:, None, :].to(dtype), kv_cache, 0)
         local_hidden = hidden[:, 0, :]
 
         binary_logits = local_text_lm_head(local_hidden).float()
@@ -292,9 +320,9 @@ class MossTTSLocalDepthTransformer(nn.Module):
             codes[:, channel_index] = channel_token
 
             if channel_index + 1 < n_vq:
-                embeds[:, channel_index + 1, :] = audio_embeddings[channel_index](channel_token).to(dtype)
-                hidden = self._run_prefix(embeds[:, : channel_index + 2, :])
-                local_hidden = hidden[:, channel_index + 1, :]
+                embeds = audio_embeddings[channel_index](channel_token).to(dtype)[:, None, :]
+                hidden = self._run_prefix(embeds, kv_cache, channel_index + 1)
+                local_hidden = hidden[:, 0, :]
 
         return should_continue, codes