Loading your learning…
Loading your learning…
At each decoding step, a new query compares against all available keys and combines their values. With a causal mask, an earlier token’s key and value never depend on later tokens, so they stay valid as new tokens arrive. Input per call: query, key, and value [B, H, 1, D]. The cache grows to [B, H, T, D]. Output is only the newest token’s contextual features [B, H, 1, D].
import torchfrom torch import nn class KVCache: def __init__(self): self.keys = None self.values = None def append(self, keys, values): # Each tensor: [B, H, new_tokens, D]. self.keys = keys if self.keys is None else torch.cat([self.keys, keys], dim=2) self.values = values if self.values is None else torch.cat([self.values, values], dim=2) def cached_attention(query, new_keys, new_values, cache): # One new token per call: query, new_keys, new_values [B, H, 1, D]. cache.append(new_keys, new_values) scores = query @ cache.keys.transpose(-2, -1) * query.shape[-1] ** -0.5 weights = scores.softmax(dim=-1) # [B, H, 1, cached_tokens] return weights @ cache.values # [B, H, 1, D] if __name__ == "__main__": torch.manual_seed(0) q, k, v = [torch.randn(2, 2, 3, 4) for _ in range(3)] cache = KVCache() with torch.no_grad(): outputs = [cached_attention(q[:, :, t:t+1], k[:, :, t:t+1], v[:, :, t:t+1], cache) for t in range(3)] print("Outputs:", torch.cat(outputs, dim=2).shape) # [2, 2, 3, 4] print("Cache:", cache.keys.shape) # [2, 2, 3, 4]The newest query asks what past content matters now. Old queries are not needed to compute its output. Store keys for comparisons and values for the weighted mixture. Create a new cache for each sequence batch and each attention layer. This cache is request state, not learned model parameters.
from torch import nn class KVCache: def __init__(self): self.keys = None self.values = None def append(self, keys, values):dim=2 is the token axis in [B, H, T, D]. The first call stores one token; later calls concatenate one more. Batch size, head count, head dimension, dtype, and device must stay compatible. Concatenation copies the growing cache, which is simple but inefficient for long sequences.
self.values = None def append(self, keys, values): # Each tensor: [B, H, new_tokens, D]. self.keys = keys if self.keys is None else torch.cat([self.keys, keys], dim=2) self.values = values if self.values is None else torch.cat([self.values, values], dim=2) def cached_attention(query, new_keys, new_values, cache):