Heuristic:Openai Whisper KV Cache Optimization
| Knowledge Sources | |
|---|---|
| Domains | Decoding, Optimization |
| Last Updated | 2025-06-25 00:00 GMT |
Overview
Key-Value caching optimization for autoregressive decoding that avoids redundant recomputation of attention keys and values, reducing per-token inference cost from O(n) to O(1).
Description
During autoregressive decoding, the Whisper decoder generates tokens one at a time. Without caching, each new token requires recomputing attention over all previous tokens. KV caching stores the key and value projections from previous positions, so only the new token's projections need to be computed at each step. Whisper implements this via forward hooks on the key and value Linear layers that automatically append new outputs to a persistent cache dictionary.
Usage
Use this heuristic when optimizing decoder inference speed. KV caching is enabled automatically by `PyTorchInference` during decoding. After the first forward pass (which processes all initial tokens), subsequent steps process only the last token and reuse cached keys/values. Cross-attention keys and values are computed once per segment and reused for all decoding steps.
The Insight (Rule of Thumb)
- Action: Use `model.install_kv_cache_hooks()` to enable KV caching during autoregressive decoding. After the first pass, pass only the last token to the decoder.
- Value: Reduces per-step computation from processing all N tokens to processing 1 token.
- Trade-off: Increases memory usage linearly with sequence length (stores cached tensors). Cache must be properly cleaned up after decoding to avoid memory leaks.
- Key detail: For beam search, the cache must be rearranged when beams are reordered, using `rearrange_kv_cache(source_indices)`.
Reasoning
In Transformer self-attention, the key and value projections for earlier positions do not change as new tokens are generated. Caching them avoids O(n^2) total computation across the full sequence, reducing it to O(n). For cross-attention, the audio features are constant across the entire segment, so their keys and values only need to be computed once. This is critical for practical inference speed.
Code evidence from `whisper/decoding.py:155-163`:
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()
if tokens.shape[-1] > self.initial_token_length:
# only need to use the last token except in the first forward pass
tokens = tokens[:, -1:]
return self.model.decoder(tokens, audio_features, kv_cache=self.kv_cache)
Hook-based cache mechanism from `whisper/model.py:327-333`:
def save_to_cache(module, _, output):
if module not in cache or output.shape[1] > self.dims.n_text_ctx:
# save as-is, for the first token or cross attention
cache[module] = output
else:
cache[module] = torch.cat([cache[module], output], dim=1).detach()
return cache[module]
Cache reordering for beam search from `whisper/decoding.py:172-176`:
def rearrange_kv_cache(self, source_indices):
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()