Changed the implementation to fuse the transformers (#85)

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

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

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