example-9d59422e9cb3.py

 1import torch
 2
 3
 4class KVCache:
 5    def __init__(self) -> None:
 6        self.k: torch.Tensor | None = None
 7        self.v: torch.Tensor | None = None
 8
 9    @torch.no_grad()
10    def step(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
11        # One token per call: q, k, v have shape (1, dim).
12        self.k = k if self.k is None else torch.cat([self.k, k], dim=0)
13        self.v = v if self.v is None else torch.cat([self.v, v], dim=0)
14        scores = (q @ self.k.T) / q.shape[-1] ** 0.5
15        return scores.softmax(dim=-1) @ self.v
16
17
18torch.manual_seed(3)
19q, k, v = [torch.randn(5, 8) for _ in range(3)]
20cache = KVCache()
21incremental = torch.cat([
22    cache.step(q[t:t+1], k[t:t+1], v[t:t+1])
23    for t in range(5)
24])
25
26# Compare against all tokens at once with a causal mask.
27scores = (q @ k.T) / 8 ** 0.5
28future = torch.ones(5, 5, dtype=torch.bool).triu(1)
29full = scores.masked_fill(future, float("-inf")).softmax(-1) @ v
30torch.testing.assert_close(incremental, full)
31assert cache.k is not None
32print(cache.k.shape)  # torch.Size([5, 8])