for t in range(max_tokens):
    logits = model(full_sequence)  # Recompute ALL attention!

Step 1: "Hello"
       Q1 → Attn(K1, V1) → output1
       → Cache: [K1, V1]

Step 2: "Hello world"
       Q2 → Attn([K1,K2], [V1,V2]) → output2
       → Cache: [K1,K2, V1,V2]

Step 3: ...
       → Only compute new Q, K, V

class MiniGPT(nn.Module):
    def __init__(self, vocab_size, n_embd=64, n_head=4, n_layer=4, block_size=128):
        super().__init__()
        self.block_size = block_size
        self.n_embd = n_embd
        self.n_head = n_head
        self.n_layer = n_layer

        self.token_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.blocks = nn.ModuleList([TransformerBlock(n_embd, n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)

    def forward(self, idx, targets=None, past_kv=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)

        tok_emb = self.token_emb(idx)
        pos_emb = self.pos_emb(pos)
        x = tok_emb + pos_emb

        new_kv = []
        for i, block in enumerate(self.blocks):
            cache = past_kv[i] if past_kv else None
            x, cache = block(x, cache)
            new_kv.append(cache)

        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))

        return logits, loss, new_kv

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        self.n_head = n_head
        self.d_k = n_embd // n_head
        self.Wq = nn.Linear(n_embd, n_embd)
        self.Wk = nn.Linear(n_embd, n_embd)
        self.Wv = nn.Linear(n_embd, n_embd)
        self.Wo = nn.Linear(n_embd, n_embd)

    def forward(self, x, past_kv=None):
        B, T, C = x.shape
        q = self.Wq(x).view(B, T, self.n_head, self.d_k).transpose(1, 2)  # (B,H,T,d)

        if past_kv is not None:
            k_past, v_past = past_kv
            k = torch.cat([k_past, self.Wk(x)], dim=2)
            v = torch.cat([v_past, self.Wv(x)], dim=2)
        else:
            k = self.Wk(x).view(B, T, self.n_head, self.d_k).transpose(1, 2)
            v = self.Wv(x).view(B, T, self.n_head, self.d_k).transpose(1, 2)

        # Scaled dot-product
        att = (q @ k.transpose(-2, -1)) * (1.0 / (self.d_k ** 0.5))

        # Causal mask (only for current T)
        if past_kv is None:
            mask = torch.tril(torch.ones(T, T, device=x.device))
        else:
            past_len = k.size(2) - T
            mask = torch.ones(T, k.size(2), device=x.device)
            mask = torch.tril(mask, diagonal=past_len)
        att = att.masked_fill(mask == 0, float('-inf'))

        att = F.softmax(att, dim=-1)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.Wo(y)

        return y, (k, v)

class TransformerBlock(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head)
        self.ln2 = nn.LayerNorm(n_embd)
        self.ff = nn.Sequential(
            nn.Linear(n_embd, 4*n_embd),
            nn.GELU(),
            nn.Linear(4*n_embd, n_embd)
        )

    def forward(self, x, past_kv=None):
        attn_out, new_kv = self.attn(self.ln1(x), past_kv)
        x = x + attn_out
        x = x + self.ff(self.ln2(x))
        return x, new_kv

@torch.no_grad()
def generate(model, idx, max_new_tokens=100, cache=None):
    model.eval()
    for _ in range(max_new_tokens):
        # Only pass last token if cache exists
        idx_cond = idx if cache is None else idx[:, -1:]
        logits, _, cache = model(idx_cond, past_kv=cache)
        logits = logits[:, -1, :]
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, 1)
        idx = torch.cat([idx, idx_next], dim=1)
    return idx, cache

import time

model = MiniGPT(vocab_size=65)  # TinyShakespeare
context = torch.tensor([[0, 1, 2, 3]], dtype=torch.long)  # dummy

# Warmup
generate(model, context, 10)

# No cache
start = time.time()
for _ in range(10):
    generate(model, context, 100)
no_cache = time.time() - start

# With cache
cache = None
start = time.time()
for _ in range(10):
    _, cache = generate(model, context, 100, cache)
    cache = None  # reset
with_cache = time.time() - start

print(f"No Cache:  {no_cache:.3f}s")
print(f"With Cache: {with_cache:.3f}s")
print(f"Speedup: {no_cache/with_cache:.1f}x")

Layer 1: K: (B, H, T, d) → grows with T
         V: (B, H, T, d)

Total memory: O(L × H × T × d)

# Keep only last N tokens
def prune_cache(cache, keep_last=512):
    if cache is None: return None
    return [(k[:, :, -keep_last:], v[:, :, -keep_last:]) for k, v in cache]

@torch.no_grad()
def generate_stream(model, prompt, max_tokens=200):
    idx = torch.tensor(encode(prompt), dtype=torch.long).unsqueeze(0)
    cache = None
    generated = prompt

    for _ in range(max_tokens):
        logits, _, cache = model(idx[:, -1:] if cache else idx, past_kv=cache)
        next_token = torch.multinomial(F.softmax(logits[:, -1, :], dim=-1), 1)
        idx = torch.cat([idx, next_token], dim=1)
        generated += decode([next_token.item()])
        print(decode([next_token.item()]), end='', flush=True)
        if next_token.item() == 0: break  # EOS
    return generated

import torch
import torch.nn as nn
import torch.nn.functional as F

# [CausalSelfAttention, TransformerBlock, MiniGPT with cache]

# Speed test
model = MiniGPT(vocab_size=65)
context = torch.tensor([[0]], dtype=torch.long)

# With cache
cache = None
start = time.time()
for _ in range(100):
    _, cache = model.generate(context, 50, cache)
    cache = None
print(f"100 tokens: {time.time()-start:.3f}s")