KV-cache

De KV-cache is de plek waar een taalmodel zijn tussenresultaten bewaart terwijl het een antwoord schrijft. Het model schrijft stukje voor stukje, en elk nieuw stukje hangt af van alle tekst ervoor. De cache bewaart het werk dat al op die tekst is gedaan, zodat elke stap alleen het laatste stukje verwerkt. Antwoorden komen zo sneller, maar de cache kost geheugen: hoe langer de gesprekken, hoe minder gebruikers per server.

Hoe werkt het?

In het attention-mechanisme van een transformer krijgt elk token drie vectoren: een query, een key en een value (Vaswani e.a., 2017). Om het volgende token te voorspellen, vergelijkt het model de query van het laatste token met de keys van alle eerdere tokens, en weegt het met die scores hun values. Een token kijkt alleen terug, dus zijn key en value veranderen niet meer als ze eenmaal berekend zijn. Volgens de documentatie van Hugging Face kun je ze daarom bewaren en hergebruiken: elke stap berekent alleen voor het nieuwe token een key en een value, in elke laag.

Een verzoek verwerken gaat in twee fasen (Kwon e.a., 2023). In de promptfase verwerkt het model de hele prompt in één keer en vult het de cache. In de generatiefase voegt het één token per stap toe. Die stappen kunnen niet parallel lopen en kosten het grootste deel van de tijd die een verzoek nodig heeft. Het nieuwe token wordt nog steeds met elke opgeslagen key vergeleken, dus de attention-kosten per stap groeien lineair met de lengte van de tekst in plaats van kwadratisch.

Hoeveel geheugen kost het?

De cache bevat per token een key en een value voor elke laag. Kwon e.a. rekenen het voor bij het OPT-model met 13 miljard parameters: 2 (key en value) × 5.120 (verborgen dimensie) × 40 (lagen) × 2 bytes is 800 KB per token, of tot 1,6 GB voor een verzoek van 2.048 tokens. Op een NVIDIA A100 van 40 GB met een model van 13B ging ongeveer 65% van het geheugen naar de gewichten en bijna 30% naar de cache. De cache groeit met elk token en elk verzoek, en beperkt zo hoeveel verzoeken een GPU tegelijk aankan. Modellen als Mamba houden in plaats daarvan een toestand van vaste grootte bij (Lahoti e.a., 2026).

Hoe maak je de cache kleiner?

• Keys en values delen tussen attention-heads. Stap voor stap decoderen is vaak traag omdat de grote key- en value-tensoren steeds opnieuw geladen worden (Shazeer, 2019). In zijn multi-query attention (MQA) delen alle heads één key en één value, waardoor decoderen veel sneller gaat met een klein kwaliteitsverlies. Bij grouped-query attention (GQA) delen groepen heads ze, met een kwaliteit dicht bij gewone multi-head attention en een snelheid vergelijkbaar met MQA (Ainslie e.a., 2023). Meta gebruikte GQA voor de versies van 34B en 70B van Llama 2.
• Het geheugen in pagina's beheren. Oudere servingsystemen reserveerden per verzoek ruimte voor de maximale lengte, en in de metingen van Kwon e.a. bevatte maar 20,4% tot 38,2% van het cachegeheugen echte tokendata. Hun PagedAttention, de basis van vLLM, slaat de cache op in kleine blokken en verspilt bijna niets. Bij dezelfde latentie haalde vLLM 2 tot 4 keer de doorvoer van FasterTransformer en Orca.
• Minder tokens bewaren. H2O (Zhang e.a., 2023) bewaart de laatste tokens en de paar heavy hitters die het meest aan attention bijdragen. Met 20% heavy hitters haalde het tot 29 keer de doorvoer van DeepSpeed Zero-Inference en Hugging Face Accelerate.
• Minder bits gebruiken. vLLM kan de cache in 8 bits opslaan (FP8 KV cache), een vorm van kwantisatie.
• Hergebruiken. Beginnen verzoeken met dezelfde tekst, zoals een lange systeemprompt, dan kan vLLM de cache van dat gedeelde begin hergebruiken (automatic prefix caching). API-aanbieders bieden iets vergelijkbaars aan onder de naam prompt caching.

Voorbeeld

Stel: een bedrijf draait OPT-13B op één A100 van 40 GB voor zijn klantenservicechat. Bij 800 KB per token neemt een gesprek van 2.000 tokens zo'n 1,6 GB cache in beslag. De ongeveer 12 GB voor de cache (bijna 30% van 40 GB) biedt dan plaats aan zo'n zeven van die gesprekken tegelijk. Een cache in 8 bits halveert die 800 KB per token en verdubbelt dat aantal.