KV-cache

När modellen genererar token t+1 räknar den attention över alla tidigare tokens. Utan cache räknas K och V om för hela sekvensen varje gång — men de är identiska med förra gången, eftersom tidigare tokens inte ändras.

KV-cache: spara K och V per lager och huvud. Vid nästa token beräknas bara den nya positionens K och V och läggs till.

Effekt: varje nytt token kostar O(t) i stället för O(t²). För ett svar på 500 tokens är det skillnaden mellan sekunder och minuter.

Priset är minne. Cachen växer linjärt med kontextlängden och blir ofta det som begränsar hur många samtidiga användare en GPU klarar.

import numpy as np

def softmax(z):
    z = z - z.max(-1, keepdims=True); e = np.exp(z); return e / e.sum(-1, keepdims=True)

class CachadAttention:
    def __init__(self, d):
        self.d = d; self.K = None; self.V = None
    def steg(self, q, k, v):
        """Ett nytt token: q, k, v har form (d,). Returnerar attention-utdata."""
        self.K = k[None] if self.K is None else np.vstack([self.K, k])
        self.V = v[None] if self.V is None else np.vstack([self.V, v])
        poang = (self.K @ q) / np.sqrt(self.d)       # (t,)
        return softmax(poang) @ self.V                # (d,)

rng = np.random.default_rng(0); d = 8
c = CachadAttention(d)
for t in range(5):
    ut = c.steg(rng.normal(size=d), rng.normal(size=d), rng.normal(size=d))
print(c.K.shape, ut.shape)      # (5, 8) (8,)

Minnesformel:

bytes = 2 (K och V) × lager × kv_huvuden × d_head × sekvenslängd × batch × bytes_per_tal

En 7B-modell (32 lager, 32 huvuden, d_head 128, fp16) vid 8 192 tokens: 2 · 32 · 32 · 128 · 8 192 · 2 B ≈ 4,3 GB — för en användare. Därför GQA (färre KV-huvuden), kvantiserad cache och sliding window.