From a47c7e4ba984d3edbcafb392aef679728e1c3d81 Mon Sep 17 00:00:00 2001 From: darknight054 Date: Tue, 3 Feb 2026 17:02:46 +0530 Subject: [PATCH] Reduce KV cache memory (#51) Hi! This PR tackles the memory issues discussed in [Reduce Memory Usage](https://github.com/kyutai-labs/pocket-tts/issues/14) by removing the "always allocate 1000 steps" KV-cache behavior. --- pocket_tts/models/tts_model.py | 67 +++++++++++++++------------------- 1 file changed, 30 insertions(+), 37 deletions(-) diff --git a/pocket_tts/models/tts_model.py b/pocket_tts/models/tts_model.py index f66f956..1e71a96 100644 --- a/pocket_tts/models/tts_model.py +++ b/pocket_tts/models/tts_model.py @@ -1,5 +1,6 @@ import copy import logging +import math import os import queue import statistics @@ -55,6 +56,9 @@ VOICE_CLONING_UNSUPPORTED = ( class TTSModel(nn.Module): + _TOKENS_PER_SECOND_ESTIMATE = 3.0 + _GEN_SECONDS_PADDING = 2.0 + def __init__( self, flow_lm: FlowLMModel, @@ -279,33 +283,6 @@ class TTSModel(nn.Module): conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight) return conditioning - def _slice_kv_cache(self, model_state: dict, num_frames: int) -> None: - """Slice KV cache to only keep the first num_frames elements. - - This optimizes memory usage when caching voice states by discarding - unused cache capacity beyond the actual audio prompt length. - - Args: - model_state: The model state dict containing KV caches for all modules - num_frames: Number of frames to keep in the KV cache - """ - original_size = 0 - sliced_size = 0 - for module_name, module_state in model_state.items(): - if "cache" in module_state: - # KV cache has shape [2, batch_size, sequence_length, num_heads, dim_per_head] - cache = module_state["cache"] - original_size += cache.numel() * cache.element_size() - # Slice to keep only the first num_frames positions - module_state["cache"] = cache[:, :, :num_frames, :, :].clone() - sliced_size += module_state["cache"].numel() * module_state["cache"].element_size() - - memory_saved_mb = (original_size - sliced_size) / (1024 * 1024) - logger.info( - f"Sliced KV cache from {original_size / (1024 * 1024):.1f} MB to {sliced_size / (1024 * 1024):.1f} MB " - f"(saved {memory_saved_mb:.1f} MB)" - ) - def _expand_kv_cache(self, model_state: dict, sequence_length: int) -> None: """Expand KV cache back to full sequence_length for generation. @@ -339,12 +316,23 @@ class TTSModel(nn.Module): expanded_cache[:, :, :current_length, :, :] = cache module_state["cache"] = expanded_cache + 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]) + raise ValueError( + "Could not find current_end 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): """Worker thread function for decoding audio latents from queue with immediate streaming.""" try: audio_chunks = [] - mimi_state = init_states(self.mimi, batch_size=1, sequence_length=1000) + mimi_context = self.config.mimi.transformer.context + mimi_state = init_states(self.mimi, batch_size=1, sequence_length=mimi_context) while True: latent = latents_queue.get() if latent is None: @@ -503,9 +491,6 @@ class TTSModel(nn.Module): if copy_state: model_state = copy.deepcopy(model_state) - # Expand sliced KV caches back to full size for generation - self._expand_kv_cache(model_state, sequence_length=1000) - # Set up multithreaded generation and decoding latents_queue = queue.Queue() result_queue = queue.Queue() @@ -573,9 +558,12 @@ class TTSModel(nn.Module): latents_queue: queue.Queue, result_queue: queue.Queue, ): - gen_len_sec = len(text_to_generate.split()) * 1 + 2.0 - max_gen_len = int(gen_len_sec * 12.5) 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) with display_execution_time("Prompting text"): self._run_flow_lm_and_increment_step( @@ -713,17 +701,22 @@ class TTSModel(nn.Module): with display_execution_time("Encoding audio prompt"): prompt = self._encode_audio(audio_conditioning.unsqueeze(0).to(self.device)) - model_state = init_states(self.flow_lm, batch_size=1, sequence_length=1000) + model_state = init_states(self.flow_lm, batch_size=1, sequence_length=prompt.shape[1]) with display_execution_time("Prompting audio"): self._run_flow_lm_and_increment_step(model_state=model_state, audio_conditioning=prompt) - # Optimize memory by slicing KV cache to only keep frames from the audio prompt - num_audio_frames = prompt.shape[1] - self._slice_kv_cache(model_state, num_audio_frames) + logger.info( + "Size of the model state for audio prompt: %d MB", size_of_dict(model_state) // 1e6 + ) return model_state + def _estimate_max_gen_len(self, token_count: int) -> int: + gen_len_sec = token_count / self._TOKENS_PER_SECOND_ESTIMATE + self._GEN_SECONDS_PADDING + frame_rate = self.config.mimi.frame_rate + return math.ceil(gen_len_sec * frame_rate) + @torch.no_grad def save_audio_prompt( self,