From 33a93718f62b452a733ac3e698d8c9267078250c Mon Sep 17 00:00:00 2001 From: darknight054 Date: Fri, 20 Feb 2026 19:15:34 +0530 Subject: [PATCH] Changed the implementation to fuse the transformers (#85) --- pocket_tts/models/tts_model.py | 49 ++++++--- pocket_tts/modules/mimi_transformer.py | 144 +----------------------- pocket_tts/modules/transformer.py | 145 +++++++++++++++---------- 3 files changed, 124 insertions(+), 214 deletions(-) diff --git a/pocket_tts/models/tts_model.py b/pocket_tts/models/tts_model.py index 63a0444..7ef3146 100644 --- a/pocket_tts/models/tts_model.py +++ b/pocket_tts/models/tts_model.py @@ -336,21 +336,26 @@ class TTSModel(nn.Module): def _flow_lm_current_end(self, model_state: dict) -> int: for module_state in model_state.values(): - current_end = module_state.get("current_end") - if current_end is not None: - return int(current_end.shape[0]) + offset = module_state.get("offset") + if offset is not None: + return int(offset.view(-1)[0].item()) raise ValueError( - "Could not find current_end in model state, please open an issue " + "Could not find offset in model state, please open an issue " "at https://github.com/kyutai-labs/pocket-tts/issues" ) @torch.no_grad - def _decode_audio_worker(self, latents_queue: queue.Queue, result_queue: queue.Queue): + def _decode_audio_worker( + self, + latents_queue: queue.Queue, + result_queue: queue.Queue, + mimi_sequence_length: int, + mimi_steps_per_latent: int, + ): """Worker thread function for decoding audio latents from queue with immediate streaming.""" try: audio_chunks = [] - mimi_context = self.config.mimi.transformer.context - mimi_state = init_states(self.mimi, batch_size=1, sequence_length=mimi_context) + mimi_state = init_states(self.mimi, batch_size=1, sequence_length=mimi_sequence_length) while True: latent = latents_queue.get() if latent is None: @@ -361,7 +366,7 @@ class TTSModel(nn.Module): t = time.monotonic() audio_frame = self.mimi.decode_from_latent(quantized, mimi_state) - increment_steps(self.mimi, mimi_state, increment=16) + increment_steps(self.mimi, mimi_state, increment=mimi_steps_per_latent) audio_frame_duration = audio_frame.shape[2] / self.config.mimi.sample_rate # We could log the timings here. logger.debug( @@ -538,13 +543,21 @@ class TTSModel(nn.Module): if copy_state: model_state = copy.deepcopy(model_state) + prepared = self.flow_lm.conditioner.prepare(text_to_generate) + token_count = prepared.tokens.shape[1] + max_gen_len = self._estimate_max_gen_len(token_count) + mimi_steps_per_latent = int(self.mimi.encoder_frame_rate / self.mimi.frame_rate) + mimi_sequence_length = max_gen_len * mimi_steps_per_latent + # Set up multithreaded generation and decoding latents_queue = queue.Queue() result_queue = queue.Queue() # Start decoder worker thread decoder_thread = threading.Thread( - target=self._decode_audio_worker, args=(latents_queue, result_queue), daemon=True + target=self._decode_audio_worker, + args=(latents_queue, result_queue, mimi_sequence_length, mimi_steps_per_latent), + daemon=True, ) logger.info("starting timer now!") t_generating = time.monotonic() @@ -553,7 +566,8 @@ class TTSModel(nn.Module): # Generate latents and add them to queue (decoder processes them in parallel) self._generate( model_state=model_state, - text_to_generate=text_to_generate, + prepared=prepared, + max_gen_len=max_gen_len, frames_after_eos=frames_after_eos, latents_queue=latents_queue, result_queue=result_queue, @@ -600,14 +614,13 @@ class TTSModel(nn.Module): def _generate( self, model_state: dict, - text_to_generate: str, + prepared: TokenizedText, + max_gen_len: int, frames_after_eos: int, latents_queue: queue.Queue, result_queue: queue.Queue, ): - prepared = self.flow_lm.conditioner.prepare(text_to_generate) token_count = prepared.tokens.shape[1] - max_gen_len = self._estimate_max_gen_len(token_count) current_end = self._flow_lm_current_end(model_state) required_len = current_end + token_count + max_gen_len self._expand_kv_cache(model_state, sequence_length=required_len) @@ -880,5 +893,13 @@ def _import_model_state(source: str | Path) -> dict[str, dict[str, torch.Tensor] for key in f.keys(): module_name, tensor_key = key.split("/") result.setdefault(module_name, {}) - result[module_name][tensor_key] = f.get_tensor(key) + if tensor_key == "current_end": + # we used the shape[0] as step index before for torch.compile() compatibility, + # but it's not needed anymore + tensor = f.get_tensor(key) + result[module_name]["offset"] = torch.full( + (1,), fill_value=tensor.shape[0], dtype=torch.long, device=tensor.device + ) + else: + result[module_name][tensor_key] = f.get_tensor(key) return result diff --git a/pocket_tts/modules/mimi_transformer.py b/pocket_tts/modules/mimi_transformer.py index 0a18319..1b6e3cd 100644 --- a/pocket_tts/modules/mimi_transformer.py +++ b/pocket_tts/modules/mimi_transformer.py @@ -1,139 +1,14 @@ -from typing import NamedTuple - import torch import torch.nn as nn -from einops import rearrange from torch.nn import functional as F from typing_extensions import Self from pocket_tts.modules.layer_scale import LayerScale from pocket_tts.modules.rope import RotaryEmbedding -from pocket_tts.modules.stateful_module import StatefulModule from pocket_tts.modules.transformer import StreamingMultiheadAttention from pocket_tts.utils.config import FlowLMTransformerConfig -class KVCacheResult(NamedTuple): - keys: torch.Tensor - values: torch.Tensor - positions: torch.Tensor - - @staticmethod - def from_kv(keys: torch.Tensor, values: torch.Tensor) -> "KVCacheResult": - B, H, T, D = keys.shape - assert tuple(values.shape[:-1]) == (B, H, T) - positions = torch.arange(T, device=keys.device, dtype=torch.long) - return KVCacheResult(keys, values, positions.expand(B, -1)) - - -def complete( - cache: torch.Tensor, end_offset: torch.Tensor, k: torch.Tensor, v: torch.Tensor -) -> KVCacheResult: - capacity = cache.shape[3] - assert k.shape[:-1] == v.shape[:-1], (k.shape, v.shape) - B, H, T, D = k.shape - assert T > 0 - indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype) - indexes = indexes + end_offset.view(-1, 1) - indexes = indexes % capacity - # indexes is [B, T] - # k is [B, H, T, D] - # cache is [B, H, T', D] - this_indexes = indexes.view(B, 1, T, 1) - this_indexes = this_indexes.expand(-1, H, T, D) - cache[0].scatter_(2, this_indexes, k) - cache[1].scatter_(2, this_indexes, v) - - keys = cache[0] - values = cache[1] - - indexes = torch.arange(capacity, device=end_offset.device, dtype=torch.long) - - # end_index correspond to the actual index where the last value was written. - last_offset = end_offset.view(-1, 1) + T - 1 - end_index = last_offset % capacity - delta = indexes - end_index - - positions = torch.where(delta <= 0, last_offset + delta, last_offset + delta - capacity) - end_offset[:] = end_offset + T - invalid = indexes >= end_offset.view(-1, 1) - positions = torch.where(invalid, torch.full_like(positions, -1), positions) - - return KVCacheResult(keys, values, positions) - - -class MimiStreamingMultiheadAttention(StatefulModule): - def __init__(self, embed_dim: int, num_heads: int, context: int, rope: RotaryEmbedding): - super().__init__() - - self.embed_dim = embed_dim - self.context = context - self.rope = rope - self.num_heads = num_heads - out_dim = 3 * embed_dim - - self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False) - self.in_proj = nn.Linear(embed_dim, out_dim, bias=False) - - def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]: - dim_per_head = self.embed_dim // self.num_heads - - state = {} - device = self.in_proj.weight.device - state["offset"] = torch.zeros(batch_size, dtype=torch.long, device=device) - state["cache"] = torch.zeros( - (2, batch_size, self.num_heads, sequence_length, dim_per_head), device=device - ) - state["end_offset"] = torch.zeros(batch_size, dtype=torch.long, device=device) - return state - - def increment_step(self, state, increment: int = 1): - state["offset"] += increment - - def _complete_kv(self, k, v, model_state: dict | None) -> KVCacheResult: - if model_state is None: - return KVCacheResult.from_kv(k, v) - else: - layer_state = self.get_state(model_state) - return complete(layer_state["cache"], layer_state["end_offset"], k, v) - - def forward(self, query: torch.Tensor, model_state: dict | None) -> torch.Tensor: - B, T = query.shape[:2] - - if model_state is None: - offset = torch.zeros(B, device=query.device, dtype=torch.long) - else: - offset = self.get_state(model_state)["offset"] - - projected = self.in_proj(query) - - q, k, v = rearrange(projected, "b t (p h d) -> p b h t d", p=3, h=self.num_heads) - - # Permute from [b, h, t, d] to [b, t, h, d] for rope - q = q.permute(0, 2, 1, 3) - k = k.permute(0, 2, 1, 3) - q, k = self.rope(q, k, offset) - # Permute back from [b, t, h, d] to [b, h, t, d] - q = q.permute(0, 2, 1, 3) - k = k.permute(0, 2, 1, 3) - - k, v, pos_k = self._complete_kv(k, v, model_state) - pos_k = pos_k[:, None] - pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view( - -1, 1 - ) - delta = pos_q - pos_k - attn_bias = (pos_k >= 0) & (delta >= 0) - attn_bias = attn_bias & (delta < self.context) - attn_bias = attn_bias[:, None] - - x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0) - - x = rearrange(x, "b h t d -> b t (h d)") - x = self.out_proj(x) - return x - - class StreamingTransformerLayer(nn.Module): def __init__( self, @@ -143,21 +18,11 @@ class StreamingTransformerLayer(nn.Module): context: int | None, rope: RotaryEmbedding, layer_scale: float | None = None, - attention_kind: str = "mimi", ): super().__init__() - # Redefine self_attn to our streaming multi-head attention - if attention_kind == "mimi": - # TODO: we should actually use StreamingMultiheadAttention here and add context window - # support. And we should then delete MimiStreamingMultiheadAttention. - # The implementation is really close. - self.self_attn = MimiStreamingMultiheadAttention( - context=context, rope=rope, embed_dim=d_model, num_heads=num_heads - ) - else: - self.self_attn = StreamingMultiheadAttention( - rope=rope, embed_dim=d_model, num_heads=num_heads - ) + self.self_attn = StreamingMultiheadAttention( + rope=rope, embed_dim=d_model, num_heads=num_heads, context=context + ) self.norm1 = nn.LayerNorm(d_model, eps=1e-5) self.norm2 = nn.LayerNorm(d_model, eps=1e-5) @@ -199,7 +64,6 @@ class StreamingTransformer(nn.Module): dim_feedforward: int | list[int] = 2048, context: int | None = None, max_period: float = 10_000.0, - kind: str = "mimi", ): super().__init__() assert d_model % num_heads == 0 @@ -217,7 +81,6 @@ class StreamingTransformer(nn.Module): context=context, rope=self.rope, layer_scale=layer_scale, - attention_kind=kind, ) ) @@ -230,7 +93,6 @@ class StreamingTransformer(nn.Module): num_layers=config.num_layers, dim_feedforward=dim_feedforward, max_period=float(config.max_period), - kind="flow_lm", ) def forward(self, x: torch.Tensor, model_state: dict | None): diff --git a/pocket_tts/modules/transformer.py b/pocket_tts/modules/transformer.py index 2b8278c..4cf18e6 100644 --- a/pocket_tts/modules/transformer.py +++ b/pocket_tts/modules/transformer.py @@ -7,28 +7,81 @@ from pocket_tts.modules.stateful_module import StatefulModule def complete_kv( - cache: torch.Tensor, current_end: torch.Tensor, k: torch.Tensor, v: torch.Tensor + cache: torch.Tensor, offset: torch.Tensor, k: torch.Tensor, v: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - current_end = current_end.shape[0] + if offset.numel() > 1 and not torch.all(offset == offset.view(-1)[0]): + raise ValueError("Linear cache offset must be identical across batch.") + offset_value = int(offset.view(-1)[0].item()) - cache[0, :, current_end : current_end + k.shape[1]] = k - cache[1, :, current_end : current_end + v.shape[1]] = v - valid = cache[:, :, : current_end + k.shape[1]] + cache[0, :, offset_value : offset_value + k.shape[1]] = k + cache[1, :, offset_value : offset_value + v.shape[1]] = v + valid = cache[:, :, : offset_value + k.shape[1]] return valid[0], valid[1] -def _materialize_causal_mask( - shape: tuple[int, ...], shift: int, device: str | torch.device = "cpu" +def _build_attention_mask( + pos_q: torch.Tensor, pos_k: torch.Tensor, context: int | None ) -> torch.Tensor: - dtype = torch.float32 + delta = pos_q[:, :, None] - pos_k[:, None, :] + mask = (pos_k[:, None, :] >= 0) & (delta >= 0) + if context is not None: + mask = mask & (delta < context) + return mask[:, None] - num_queries, num_keys = shape[-2:] - shift = num_keys - num_queries - tensor = torch.full(shape, dtype=dtype, fill_value=1, device=device) - mask = torch.tril(tensor, diagonal=shift).to(dtype) - mask = torch.log(mask) - return mask.to(dtype) +# Per-layer streaming state schemas (returned by init_state and stored in model_state): +# - Linear cache (FlowLM / full causal): +# - offset: torch.long[B] # absolute time index for the next write / RoPE offset +# # (batch must be in sync) +# - cache: torch.[dtype][2, B, T, H, D] # K/V stored contiguously along T (allocated capacity) + + +class _LinearKVCacheBackend: + requires_state = True + + def __init__(self, num_heads: int, dim_per_head: int): + self.num_heads = num_heads + self.dim_per_head = dim_per_head + + def init_state( + self, batch_size: int, sequence_length: int, device: torch.device, dtype: torch.dtype + ) -> dict[str, torch.Tensor]: + return dict( + offset=torch.zeros(batch_size, dtype=torch.long, device=device), + cache=torch.full( + (2, batch_size, sequence_length, self.num_heads, self.dim_per_head), + float("NaN"), + device=device, + dtype=dtype, + ), + ) + + def increment_step(self, state: dict[str, torch.Tensor], increment: int) -> None: + state["offset"] += increment + + def rope_offset( + self, state: dict[str, torch.Tensor] | None, batch_size: int, device: torch.device + ) -> torch.Tensor: + if state is None: + return torch.zeros((), dtype=torch.long, device=device) + return state["offset"].view(-1)[0] + + def append_and_get( + self, k: torch.Tensor, v: torch.Tensor, state: dict[str, torch.Tensor] | None + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if state is None: + k_attn = k.permute(0, 2, 1, 3) + v_attn = v.permute(0, 2, 1, 3) + pos_k = torch.arange(k_attn.shape[2], device=k_attn.device, dtype=torch.long) + pos_k = pos_k.view(1, -1).expand(k_attn.shape[0], -1) + offset = torch.zeros(k_attn.shape[0], device=k_attn.device, dtype=torch.long) + return k_attn, v_attn, pos_k, offset + cache_k, cache_v = complete_kv(state["cache"], state["offset"], k, v) + k_attn = cache_k.permute(0, 2, 1, 3) + v_attn = cache_v.permute(0, 2, 1, 3) + pos_k = torch.arange(k_attn.shape[2], device=k_attn.device, dtype=torch.long) + pos_k = pos_k.view(1, -1).expand(k_attn.shape[0], -1) + return k_attn, v_attn, pos_k, state["offset"] class StreamingMultiheadAttention(StatefulModule): @@ -44,12 +97,17 @@ class StreamingMultiheadAttention(StatefulModule): dtype (torch.dtype, optional): dtype to use. """ - def __init__(self, embed_dim: int, num_heads: int, rope: RotaryEmbedding): + def __init__( + self, embed_dim: int, num_heads: int, rope: RotaryEmbedding, context: int | None = None + ): super().__init__() self.embed_dim = embed_dim self.rope = rope self.num_heads = num_heads + self.context = context + self.dim_per_head = embed_dim // num_heads + self._cache_backend = _LinearKVCacheBackend(self.num_heads, self.dim_per_head) out_dim = embed_dim num_kv = num_heads @@ -59,62 +117,31 @@ class StreamingMultiheadAttention(StatefulModule): self.in_proj = nn.Linear(embed_dim, mult * out_dim, bias=False) self.out_proj = nn.Linear(embed_dim, mult * embed_dim, bias=False) - def _get_mask(self, shape: tuple[int, int], shift: int, device: torch.device) -> torch.Tensor: - return _materialize_causal_mask(shape, shift=shift, device=device) - def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]: - dim_per_head = self.embed_dim // self.num_heads - initial_current_end = torch.zeros((0,)).to(self.in_proj.weight.device) - return dict( - current_end=initial_current_end, - cache=torch.full( - (2, batch_size, sequence_length, self.num_heads, dim_per_head), - float("NaN"), - device=self.in_proj.weight.device, - dtype=self.in_proj.weight.dtype, - ), - ) + device = self.in_proj.weight.device + dtype = self.in_proj.weight.dtype + return self._cache_backend.init_state(batch_size, sequence_length, device, dtype) def increment_step(self, state: dict, increment: int = 1): - new_size = state["current_end"].shape[0] + increment - state["current_end"] = torch.zeros((new_size,)).to(state["current_end"].device) - - def _complete_kv(self, k, v, state: dict | None): - k, v = complete_kv(state["cache"], state["current_end"], k, v) - return k, v - - def _apply_rope(self, query: torch.Tensor, key: torch.Tensor, state: dict | None): - # Apply rope embeddings to query and key tensors. - streaming_offset = self._streaming_offset(state) - return self.rope(query, key, offset=streaming_offset) - - def _streaming_offset(self, state: dict | None) -> torch.Tensor | int: - return state["current_end"].shape[0] - - def check_model_state(self, model_state: dict): - if model_state is None: - raise ValueError("model_state must be provided") - return self.get_state(model_state) + self._cache_backend.increment_step(state, increment) def forward(self, query: torch.Tensor, model_state: dict | None): - state = self.check_model_state(model_state) + state = None if model_state is None else self.get_state(model_state) projected = self.in_proj(query) # Reshape from (b, t, p*h*d) to (b, t, p, h, d) where p=3, h=num_heads b, t, _ = projected.shape - d = self.embed_dim // self.num_heads + d = self.dim_per_head packed = projected.view(b, t, 3, self.num_heads, d) q, k, v = torch.unbind(packed, dim=2) - q, k = self._apply_rope(q, k, state) - k, v = self._complete_kv(k, v, state) + rope_offset = self._cache_backend.rope_offset(state, b, q.device) + q, k = self.rope(q, k, offset=rope_offset) + q = q.transpose(1, 2) - mask_shape = (query.shape[1], query.shape[1] + state["current_end"].shape[0]) - shift = state["current_end"].shape[0] - - attn_mask = self._get_mask(mask_shape, shift=shift, device=q.device) - - q, k, v = [x.transpose(1, 2) for x in (q, k, v)] - x = F.scaled_dot_product_attention(q, k, v, attn_mask) + k_attn, v_attn, pos_k, offset = self._cache_backend.append_and_get(k, v, state) + pos_q = offset.view(-1, 1) + torch.arange(t, device=q.device, dtype=torch.long).view(1, -1) + attn_mask = _build_attention_mask(pos_q, pos_k, self.context) + x = F.scaled_dot_product_attention(q, k_attn, v_attn, attn_mask, dropout_p=0.0) x = x.transpose(1, 2) # Reshape from (b, t, h, d) to (b, t, h*d) b, t, h, d = x.shape