KV-cache

I en decoder-only-transformer projicerer hvert attention-lag hvert token til en query-, en key- og en value-vektor. For at generere token t skal dets query sammenlignes med keys og values for alle tokens 1…t. Da kausal maskering betyder, at tidligere tokens' keys og values aldrig ændrer sig, kan de beregnes én gang og gemmes. Inferens deles derfor i to faser: prefill, der behandler hele prompten parallelt og skriver dens K- og V-tensorer i cachen, og decode, der frembringer ét token pr. trin, tilføjer én ny K/V-post pr. lag og læser hele cachen. Uden cachen skulle hvert trin genberegne attention-input for hele præfikset, så lineært arbejde pr. token blev til kvadratisk samlet arbejde.

Hukommelsesforbruget er let at estimere: bytes pr. token = 2 (K og V) × antal lag × antal KV-heads × head-dimension × bytes pr. element. For en model af typen Llama-2-70B (80 lag, 8 KV-heads, head-dimension 128) med 16-bit-præcision giver det omkring 320 KiB pr. token, så en enkelt kontekst på 32.000 tokens fylder cirka 10 GiB, og en batch af sådanne forespørgsler kan fylde mere end selve vægtene. Fordi decode skal læse hele cachen fra HBM i hvert trin, gør lange kontekster genereringen begrænset af hukommelsesbåndbredden og begrænser, hvor mange sekvenser en GPU kan betjene samtidig.

Det meste af det nyere arbejde med arkitektur og systemer går efter netop denne omkostning. Multi-query attention (Shazeer, 2019) deler ét K/V-head mellem alle query-heads, og grouped-query attention (Ainslie m.fl., 2023) bruger et lille antal fælles grupper, hvilket mindsker cachen med forholdet mellem query-heads og KV-heads. DeepSeek-V2's multi-head latent attention gemmer en komprimeret latent repræsentation i stedet for fulde K/V. Sliding-window attention begrænser cachen til de seneste W tokens, og selve cachen kan kvantiseres til 8 bit eller mindre, så man bytter lidt nøjagtighed for kapacitet. På systemsiden gemmer vLLM's PagedAttention (Kwon m.fl., 2023) cachen i blokke af fast størrelse, der slås op via en bloktabel ligesom sider i virtuel hukommelse; det fjerner det meste af den fragmentering, der opstår, når man reserverer sammenhængende buffere på forhånd, og gør det muligt at dele fælles præfikser mellem sekvenser med copy-on-write.

Flere misforståelser går igen. KV-cachen hører til den enkelte forespørgsel og lever kun, mens der genereres, medmindre et serving-system bevidst gemmer den; gør det det på tværs af forespørgsler, hedder funktionen prefix caching eller prompt caching. Den er ikke en cache af svar, så identiske spørgsmål bliver ikke gratis. Når cacheblokke smides ud eller swappes under hukommelsespres, må de genberegnes, eller forespørgsler må sættes på pause, hvilket ses som spidser i latensen. Endelig skaber en cache, der deles mellem kunder, en timing-sidekanal, så serving-systemer med flere lejere bør isolere genbrug af cache pr. kunde.