Paper → implementation
See this paper in code.
Community-written snippets that make the paper’s ideas concrete. These are not official implementations from the paper’s authors.
Reuse the past while producing one new output
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].
3 4class KVCache:5 def __init__(self):6 self.keys = None7 self.values = None8 9 def append(self, keys, values):Wrap projections and verify cached decoding
This module accepts token features [B, T, E] and owns the Q/K/V and output projections. A full call uses a causal mask. A cached call accepts [B, 1, E] and reuses stored keys and values. The runnable check compares both paths on the same inputs and weights, making the cache’s purpose observable.
20 21class CachedSelfAttention(nn.Module):22 def __init__(self, dim, heads):23 super().__init__()24 if dim % heads:25 raise ValueError("dim must divide evenly across heads")26 self.heads = headsYour place is saved automatically. Reaching the end marks this paper as read.
Page 1