import torch


class KVCache:
    def __init__(self) -> None:
        self.k: torch.Tensor | None = None
        self.v: torch.Tensor | None = None

    @torch.no_grad()
    def step(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
        # One token per call: q, k, v have shape (1, dim).
        self.k = k if self.k is None else torch.cat([self.k, k], dim=0)
        self.v = v if self.v is None else torch.cat([self.v, v], dim=0)
        scores = (q @ self.k.T) / q.shape[-1] ** 0.5
        return scores.softmax(dim=-1) @ self.v


torch.manual_seed(3)
q, k, v = [torch.randn(5, 8) for _ in range(3)]
cache = KVCache()
incremental = torch.cat([
    cache.step(q[t:t+1], k[t:t+1], v[t:t+1])
    for t in range(5)
])

# Compare against all tokens at once with a causal mask.
scores = (q @ k.T) / 8 ** 0.5
future = torch.ones(5, 5, dtype=torch.bool).triu(1)
full = scores.masked_fill(future, float("-inf")).softmax(-1) @ v
torch.testing.assert_close(incremental, full)
assert cache.k is not None
print(cache.k.shape)  # torch.Size([5, 8])