From b91c907694f96a3fb9da03d4bbdc83fbcd3a40a4 Mon Sep 17 00:00:00 2001 From: WangChou Lu Date: Fri, 7 Jul 2023 03:48:08 +0800 Subject: [PATCH] Avoid rearranging all caches (#1483) * avoid rearranging all kv_caches * avoid calculating the same kv_cache from cross attn * Update decoding.py * linter fix --------- Co-authored-by: Jong Wook Kim --- whisper/decoding.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/whisper/decoding.py b/whisper/decoding.py index 457ee7c..ecd98a4 100644 --- a/whisper/decoding.py +++ b/whisper/decoding.py @@ -146,6 +146,10 @@ class PyTorchInference(Inference): self.kv_cache = {} self.hooks = [] + key_modules = [block.attn.key for block in self.model.decoder.blocks] + value_modules = [block.attn.value for block in self.model.decoder.blocks] + self.kv_modules = key_modules + value_modules + def logits(self, tokens: Tensor, audio_features: Tensor) -> Tensor: if not self.kv_cache: self.kv_cache, self.hooks = self.model.install_kv_cache_hooks() @@ -164,9 +168,10 @@ class PyTorchInference(Inference): self.hooks = [] def rearrange_kv_cache(self, source_indices): - for module, tensor in self.kv_cache.items(): - # update the key/value cache to contain the selected sequences - self.kv_cache[module] = tensor[source_indices].detach() + if source_indices != list(range(len(source_indices))): + for module in self.kv_modules: + # update the key/value cache to contain the selected sequences + self.kv_cache[module] = self.kv_cache[module][source_indices].detach() class SequenceRanker: @@ -668,7 +673,6 @@ class DecodingTask: return languages, lang_probs def _main_loop(self, audio_features: Tensor, tokens: Tensor): - assert audio_features.shape[0] == tokens.shape[0] n_batch = tokens.shape[0] sum_logprobs: Tensor = torch.zeros(n_batch, device=audio_features.device) no_speech_probs = [np.nan] * n_batch @@ -721,8 +725,7 @@ class DecodingTask: ) ] - # repeat the audio & text tensors by the group size, for beam search or best-of-n sampling - audio_features = audio_features.repeat_interleave(self.n_group, dim=0) + # repeat text tensors by the group size, for beam search or best-of-n sampling tokens = tokens.repeat_interleave(self.n_group, dim=0).to(audio_features.device) # call the main sampling loop