From d530938b5b44dda96d189f86f603fd565d6dc972 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Tue, 8 Sep 2026 06:37:00 +0800 Subject: [PATCH] [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> --- .../moss_tts/test_local_depth_rope_cache.py | 15 +++- .../moss_tts/modeling_moss_tts_local_depth.py | 80 +++++++++++++------ 2 files changed, 68 insertions(+), 27 deletions(-) diff --git a/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py b/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py index 90c386ee0..566ba4595 100644 --- a/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py +++ b/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) diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py index e2d46ce2c..ba19c02d1 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py +++ b/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