Loading your learning…
Loading your learning…
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.
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] class CachedSelfAttention(nn.Module): def __init__(self, dim, heads): super().__init__() if dim % heads: raise ValueError("dim must divide evenly across heads") self.heads = heads self.qkv = nn.Linear(dim, 3 * dim, bias=False) self.output = nn.Linear(dim, dim, bias=False) def forward(self, x, cache=None): batch, length, dim = x.shape if cache is not None and length != 1: raise ValueError("Cached decoding accepts one new token per call") qkv = self.qkv(x).view(batch, length, 3, self.heads, dim // self.heads) q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) if cache is not None: mixed = cached_attention(q, k, v, cache) else: scores = q @ k.transpose(-2, -1) * q.shape[-1] ** -0.5 future = torch.ones(length, length, device=x.device, dtype=torch.bool).triu(1) weights = scores.masked_fill(future, float("-inf")).softmax(dim=-1) mixed = weights @ v mixed = mixed.transpose(1, 2).reshape(batch, length, dim) return self.output(mixed) if __name__ == "__main__": torch.manual_seed(0) attention = CachedSelfAttention(dim=8, heads=2).eval() x = torch.randn(2, 4, 8) # [B=2, T=4, E=8] cache = KVCache() # Fresh cache for this batch and layer. with torch.no_grad(): full = attention(x) incremental = torch.cat([attention(x[:, t:t+1], cache) for t in range(x.shape[1])], dim=1) torch.testing.assert_close(full, incremental) print("Full and cached outputs agree:", incremental.shape)CachedSelfAttention owns learned projections; KVCache is supplied separately. That separation keeps one user’s history from leaking into another request and allows a fresh cache without rebuilding the network. Each layer in a larger transformer needs its own cache.
return weights @ cache.values # [B, H, 1, D] class CachedSelfAttention(nn.Module): def __init__(self, dim, heads): super().__init__() if dim % heads: raise ValueError("dim must divide evenly across heads") self.heads = heads self.qkv = nn.Linear(dim, 3 * dim, bias=False) self.output = nn.Linear(dim, dim, bias=False) def forward(self, x, cache=None):qkv maps [B, T, E] to [B, T, 3E]. Reshaping and permuting yields Q/K/V [B, H, T, D], where D=E/H. With a cache, this happens only for the new token. The explicit length check prevents silently using the single-token algorithm on a multi-token chunk.
def forward(self, x, cache=None): batch, length, dim = x.shape if cache is not None and length != 1: raise ValueError("Cached decoding accepts one new token per call") qkv = self.qkv(x).view(batch, length, 3, self.heads, dim // self.heads) q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) if cache is not None: mixed = cached_attention(q, k, v, cache)