Gör inference billig: KV cache, batching och kvantisering

För att generera en token tar en decoder-only transformer hela sekvensen hittills, kör den genom varje lager och läser probability distribution från den sista positionen. Sedan lägger den till vald token och gör om det. Den beskrivningen är korrekt, och det är vad den långsamma körningen gör.

Den är också enormt slösaktig, och orsaken är causal mask från Kapitel 9. Position 7:s key- och value-vektorer beräknas från position 7:s input och positionerna före den. När position 8 kommer kan position 7 inte se den — det är vad causal betyder — så position 7:s key och value är exakt samma tal som tidigare. Den långsamma körningen räknar om dem ändå, vid varje steg.

Så lagra dem. Det lagret är key-value cache, den enskilt mest betydelsefulla optimeringen i språkmodell-serving:

out = model(prompt_ids, use_cache=True)          # prefill: the whole prompt
past = out.past_key_values                        
nxt = out.logits[:, -1].argmax(-1, keepdim=True)

for _ in range(n - 1):
    out = model(nxt, past_key_values=past, use_cache=True)   
    past = out.past_key_values                                
    nxt = out.logits[:, -1].argmax(-1, keepdim=True)

Titta på vad som matas in i modellen inne i loopen: nxt, en token. Inte sekvensen. Den nya token:s query gör attention mot varje cachelagrad key, och de cachelagrade keys skulle aldrig ändras. Det här är inte en approximation — kontrollen med identisk output ovan är poängen. Cache byter inte kvalitet mot hastighet; den tar bort redundant aritmetik.

För att se skalningen tydligt, skala bort transformer och ta tiden på ett enda attention head med $d = 64$, ett generation-steg beräknat på båda sätten:

Kolumnen längst till höger är orsaken. Omräkning bygger hela $n \times n$ attention-matrisen vid varje steg — $O(n^2)$ från rutan om asymptotisk notation i Kapitel 9, betald en gång per token. Med cache bygger du en $1 \times n$-rad i stället: vid 4 096 tokens, 67 MB scores mot 16 KB.

Att räkna multiply-accumulates i stället för millisekunder tar bort maskinen ur argumentet. För att generera $T$ tokens från kall start:

Per steg är den cachelagrade versionen linjär i context och den ocachelagrade kvadratisk; summerat över en generation, $O(T^2)$ mot $O(T^3)$, med en ratio som växer utan gräns. Den niofaldiga skillnaden i inledningen mättes över 48 tokens — kortare än tabellens första rad.

Cache ändrar också vad som måste finnas i minnet. På en laptop-GPU med 8 GB som genererar 256 tokens i fp16, med allocatorns topp minus resident weights:

8,3 gånger mer minne, använt för att producera samma tokens långsammare. Det här är löftet från Kapitel 5, som anländer från ett oväntat håll: där behövde reverse-mode autodiff hålla varje intermediate vid liv för backward pass, och activations dominerade träningsminnet. Vid inference finns inget backward pass och inget att behålla för det — så det som dominerar minnet i stället är cache, och det är ett medvetet val snarare än en oundviklig kostnad.

Titta igen på den snabba körningen: dess första token betedde sig inte som de övriga fyrtiosju.

prefill, 40 prompt tokens : 1.0224 s   ->  25.6 ms per token
decode,  47 steps         : 0.1665 s mean per step

Prompt kostade 25,6 ms per token och varje genererad token kostade 166 ms. Samma model, samma hårdvara, samma weights, en sexfaldig skillnad per token — och den går åt det håll de flesta inte väntar sig. Prompt är den billiga delen. Generation delas i två faser med genuint olika fysik:

Båda halvorna kan mätas. Prefill, ett pass över $P$ tokens:

Decode, en token mot en cache på $C$:

Läs den andra tabellen två gånger. Att gå från 16 tokens context till 1 024 — sextiofyra gånger mer historik att göra attention över — ändrade kostnaden för ett steg med ingenting mätbart. Attention mot cache är verkligt arbete, men det dränks av den fasta kostnaden för att dra en halv miljard weights genom minnesbussen för att producera en vektor. Den fasta kostnaden är orsaken till allt i nästa avsnitt.

Dessa två faser är ursprunget till de två tal varje serving-system rapporterar. Time to first token är i princip prefill, och den växer med prompt, vilket är varför en lång konversation känns långsam att starta. Tokens per second är $1/\text{decode step}$, och den är ungefär konstant, vilket är varför svaret sedan flyter jämnt. En chat som startar långsamt och sedan streamar jämnt är inte ett renderingstrick. Det är de här två tabellerna.

Cache byter aritmetik mot minne, och minnet den vill ha är inte litet. För varje token i context håller varje lager en key-vektor och en value-vektor per key-value head:

2:an är för keys och values; allt annat är arkitekturen. För modellen som mäts genom hela kapitlet — 24 lager, 14 query heads, 2 key-value heads, head-dimension 64 — är det i fp16 $2 \times 24 \times 2 \times 64 \times 2 = 12{,}288$ bytes per token.

Formler i det här fältet har en vana att vara fel med en faktor två, så kontrollera den mot allocatorn i stället för att tro på den:

KV cache tensors per layer: (1, 2, 295, 64) float16
measured: 3,624,960 bytes for 295 tokens = 12,288 bytes/token
formula : 2 * 24 * 2 * 64 * 2                = 12,288 bytes/token

Exakt, och den förblir exakt över varje form som prövades:

De sista tre raderna förtjänar en andra titt. Trettiotvå användare med 2 048 tokens var, sextiofyra med 1 024, etthundratjugoåtta med 512 — cache är 768 MB i varje fall, eftersom alla tre håller 65 536 tokens. Cache beror bara på det totala antalet tokens som är resident, inte på hur de är fördelade mellan användare. Det faktumet är grunden för avsnittet om batching.

Kapitel 9 introducerade multi-query och grouped-query attention och sköt upp orsaken till det här kapitlet. Orsaken är den formeln, och särskilt $H_{kv}$ i den.

Standard multi-head attention ger varje query head sina egna key och value heads. Modellen här har 14 query heads; med full multi-head attention skulle dess cache vara $2 \times 24 \times 14 \times 64 \times 2 = 86{,}016$ bytes per token — 84 KB i stället för 12 KB, exakt sju gånger mer, förhållandet mellan query heads och key-value heads.

Multi-query attention¹ tar det här till gränsen: alla query heads delar ett enda key-value head. Grouped-query attention² är kompromissen som vann — en handfull key-value heads, där varje delas av en grupp query heads — eftersom MQA:s kvalitetsförlust var verklig och GQA:s inte är det. Ingen av dem köper någon aritmetik. De finns för att dela den formeln med ett heltal, och de spreds över branschen i samma ögonblick som långa context gjorde cache till den bindande begränsningen.

Vilket den gör, snabbt. För en 7B-klass model med 32 lager och 8 key-value heads med dimension 128 är cache 128 KB per token i fp16:

Modellens egna weights är 13,0 GB i fp16, siffran i tabellen i slutet av kapitlet. Så vid ett 128 000-token context är en användares cache större än modellen. Det här är aritmetiken som Kapitel 16 gör om till pengar, och det är därför en lång konversation inte bara är långsam — den upptar en fast del av en maskin så länge request lever.

Decode är memory-bound: weights dras genom bussen för att producera en token, och de aritmetiska enheterna står sysslolösa. Lägg därför mer arbete i samma steg. Kör flera requests samtidigt, och weights, lästa en gång, tjänar dem alla. Mätt på samma model, där varje request håller en 64-token cache och decodar en token:

Läs de två kolumnerna till höger mot varandra, för de är hela poängen. Att gå från en request till sexton multiplicerar throughput med 6,0 och väntetiden för varje enskild request med 2,67. Batch gjorde servern bättre och varje användare sämre.

Det är inte en bug som kan trimmas bort; det är själva trade-offen, och den har ett namn på varje sida. Latency är vad en person som väntar på ett svar upplever. Throughput är vad fakturan delas med. Ingen inställning förbättrar båda.

Notera också var det slutar. Från 16 till 32 ökar throughput med 9 % medan latency nästan fördubblas: steget har slutat vara memory-bound och blivit compute-bound, och efter den knäpunkten köper batch ingenting. Varje deployment har en sådan knäpunkt; dess läge måste mätas på din, men dess existens behöver inte det.

Det naiva sättet att batcha är att samla $B$ requests, köra dem tillsammans och returnera när alla är klara. Men de blir inte klara samtidigt: vissa svar är tjugo tokens och andra femhundra. En fast batch kör tills dess längsta medlem är klar, och varje färdig request fortsätter att uppta sin plats och bidra med padding tills dess.

Ta 64 requests med en realistisk snedfördelning av output-längder — median 18 tokens, längsta 231, totalt 1 874 — och simulera båda policyerna med den uppmätta per-steg-kostnaden för åtta platser:

Throughput förbättras med 1,6x. Genomsnittlig latency förbättras med mer än tio gånger, eftersom en request som blev klar på fyra steg under statisk batching ändå väntar på en granne med 231 tokens innan någon får höra den.

Continuous batching³ är lösningen, och den är så enkel som den låter: batch är inte en grupp utan en uppsättning platser, och en plats som frigörs tar in nästa köade request på allra nästa steg. Schemaläggaren arbetar på granulariteten en token snarare än en request. Varje serving-stack i produktion gör detta nu.

Den har en andra halva, och det är cache. Platser som kommer och går lämnar cache-minnet fragmenterat, och att reservera varje plats för dess största möjliga context slösar bort det mesta av reservationen. PagedAttention⁴ lånar svaret från operativsystem: lagra cache i block med fast storlek med en blocktabell per sekvens, så att en sekvens cache kan vara fysiskt utspridd men logiskt sammanhängande — vilket också låter två sekvenser med delat prefix dela blocken som håller det. Det är vad vLLM bygger på, och varför en serving engine är en minnesallocator med en transformer fastsatt.

Den andra halvan av notan är själva weights. En halv miljard parameters på fyra bytes var är 1,98 GB; på två bytes, 0,99 GB; på en byte, 0,49 GB. Färre bits per weight krymper modellen på disk, krymper den i minnet och — eftersom decode är bandwidth-bound — gör varje steg snabbare, eftersom färre bytes måste flyttas.

Det enklaste schemat är symmetrisk absolutmax-quantization, och det ryms på tre rader:

qmax  = 2 ** (bits - 1) - 1
scale = W.abs().max() / qmax                        
Wq    = torch.round(W / scale).clamp(-qmax - 1, qmax)
W_hat = Wq * scale                                  # dequantized

Välj en scale så att den största weight mappas till det största heltalet, dela, avrunda, lagra heltalen och scale. Rekonstruera genom att multiplicera tillbaka. Inget med det är smart, och det fungerar — ända tills det inte gör det.

Mätt på modellens verkliga weights: alla 168 projection matrices, 357,8 miljoner parameters, relative error $\lVert W - \hat{W}\rVert / \lVert W \rVert$:

Den fjärde raden är kollapsen. Ett relative error på 0,99 på den sämsta matrisen betyder att rekonstruktionen i princip inte behåller något av originalet — matrisen har ersatts av brus med ungefär rätt storleksordning. Orsaken syns i samma experiment på en enda matris:

model.layers.12.mlp.down_proj.weight   (896 x 4864)
mean |w| 0.01386   std 0.01822   max |w| 0.43945   max/std 24.1
weights beyond 6 sigma: 692 of 4,358,144   (0.016 %)

En weight av sextusen ligger bortom sex standardavvikelser, och den största är 24 bort. Med en enda scale för hela matrisen sätter den ena weight stegstorleken för alla 4,3 miljoner. Vid 8 bits finns 256 steg och den typiska weight hamnar fortfarande på ett meningsfullt. Vid 4 bits finns 16, de yttersta reserveras för ett värde som nästan ingenting har, och de vanliga weights — vilket är allihop — avrundas till två eller tre distinkta nivåer.

Allt efter den raden är samma reparation i olika granularitet: ge scale ett mindre territorium. Per output-rad delar felet med 3,4; per grupp om 128 intilliggande weights delar det igen. Kostnaden är bokföring — en 16-bit scale per grupp om 128 är $4 + 16/128 = 4.125$ bits per weight i stället för 4 — och det köper tillbaka större delen av glappet.

NF4 angriper det från andra hållet.⁵ Nivåerna behöver inte ligga jämnt. Weights inom ett block är ungefär normalfördelade, så välj de sexton nivåerna som kvantilerna i en normalfördelning: täta nära noll där weights faktiskt finns, glesa i svansarna där de inte gör det. Samma fyra bits, samma block-scaling, vid ett mindre block — 4,25 bits per weight mot group-128:s 4,125 — och det uppmätta felet faller från 0,1323 till 0,0952, 28 % lägre. En del av det är det finare blocket och resten är att lägga nivåerna där massan finns, och att separera de två skulle kräva en tredje rad.

Kapitel 2:s floating-point-ruta slutade med ett löfte: att det här kapitlet skulle kvantisera weights till 8 och 4 bits och hitta en handfull outlier features som vägrade pressas ihop. Här är de, och de förklarar varför ”bara avrunda talen” aldrig skulle fungera på activations.

Weights ovan betedde sig illa. Activations är i en annan liga. Ta en vanlig 84-token prompt, fånga residual stream vid varje lager och mät den största magnituden som var och en av de 896 dimensionerna når:

Dimension 62 når 1 579,6 medan mediandimensionen aldrig överstiger 1,6. Det är inte en slump från en token eller ett lager: samma dimension finns där vid lager 4 och är fortfarande där vid lager 20, med nästan samma värde. Det här är outlier features,⁶ och de är systematiska — en egenskap hos den tränade modellen, inte hos input.

Histogrammet över dessa 896 per-dimension maxima vid lager 16 gör formen omisskännlig:

     0 -      1 | ######################################## 254
     1 -      2 | ######################################## 283
     2 -      4 | ######################################## 226
     4 -      8 | ######################################## 93
     8 -     16 | ##################                       18
    16 -     32 | #########                                9
    32 -     64 | #######                                  7
    64 -    128 | #####                                    5
   128 -    256 |                                          0
   256 -    512 |                                          0
   512 -   1024 |                                          0
  1024 -   4096 | #                                        1

Niohundra dimensioner i en prydlig hög under 8, ingenting alls i tre oktaver, sedan en enda dimension ensam längst bort. Kvantisera nu den tensorn till INT8 och räkna vad som händer:

Fjorton nivåer av 256. Scale sattes av 1 579,6, så varje steg är 12,44 brett, och den typiska activation — medianmagnitud 0,26, nittionionde percentil 2,51 — har ingenstans att landa. Per dimension är det ännu skarpare:

single tensor-wide scale = 12.4378
  dim 826 (max |h| = 4.77):  1 distinct level out of 256
  dim 336 (max |h| = 1.62):  1 distinct level out of 256
  dim  96 (max |h| = 0.69):  1 distinct level out of 256

after excluding the top 4 dimensions, scale = 0.5749  (22x smaller)
  dim 826: 8 levels    dim 336: 4 levels    dim  96: 3 levels

En nivå. Hela dimensionen, varje token, kvantiserad till samma tal. Åtta bits tilldelades och ungefär noll användes, och modellen som läser dessa activations får en konstant.

Den mätningen är motiveringen för varje teknik som människor faktiskt använder:

Nästan varje artikel om quantization stannar vid föregående avsnitt: den förklarar metoden, citerar ett compression ratio och påstår att kvaliteten är ”largely preserved”. Kapitel 4 handlade om att inte lura sig själv, så låt oss ta reda på det.

Samma model, weights kvantiserade på plats med varje scheme, sedan tre mätningar: perplexity på 2 048 tokens undanhållen engelsk prosa — här utkastet till den här kursen, vilket är varför repositoryt ersätter den med en fast public-domain-bok och skriver ut en tabell med samma form men andra siffror — ett batteri med 16 korta faktafrågor med kända svar under greedy decoding, och andelen tokens där den kvantiserade modellen håller med full-precision-modellen givet identisk context.

Fyra saker i den tabellen är värda att säga rakt ut.

INT8 gjort rätt är gratis. Per-row INT8 får 22,96 mot referensens 23,08 — ett gap på en del på tvåhundra, vilket är brus och bör läsas som ”identiskt”. Åt vilket håll bruset pekar är inte stabilt: på repositoryts public-domain-korpus kommer samma två schemes ut som 22,24 mot 22,18: halva avståndet, och åt andra hållet. Det håller med full-precision-modellen på 142 av 144 genererade tokens. En fjärdedel av minnet mot fp32-referensen, hälften mot fp16 som du faktiskt skulle deploya, och ingen detekterbar kostnad. INT8 gjort vårdslöst är nästan gratis också: en scale per matris kostar 0,5 perplexity-poäng och inga svar i batteriet. Åtta bits är förlåtande nog för att granulariteten knappt spelar roll, vilket är exakt varför människor generaliserar från INT8 till INT4 och gör sig illa.

INT4 med en scale per tensor förstör modellen. Perplexity 365 miljoner: inte försämrad, utplånad. Granularitet är sedan hela spelet — per-tensor 365 416 000, per-row 46,18, per-group-of-128 31,08, NF4 24,55. Samma fyra bits per weight, en faktor på femton miljoner mellan sämst och bäst.

Perplexity är ett grovt instrument och batteriet ett ännu grövre. Mellan NF4 och group-128 INT4 är perplexity-gapet 6,5 poäng och batteriet skiljer med en fråga — och Kapitel 4:s confidence interval säger att en fråga av sexton inte särskiljer någonting alls. Det finns en skarpare demonstration än intervallet: kör samma batteri med modellens standard-repetition penalty avstängd, vilket är vad greedy decoding faktiskt betyder, och de två raderna byter plats. En fråga av sexton är inte en liten effekt, det är ingen effekt. Kapitel 8:s varning gäller också: perplexity är bara jämförbar mellan modeller som delar tokenizer, så ett tal från någon annans genomgång kan inte jämföras med ditt.

Agreement-kolumnen är den skarpaste av de tre, och nästan gratis: kör full-precision-modellen greedily, och fråga sedan den kvantiserade, vid varje position, vad den skulle ha valt givet samma prefix. Den har 144 oberoende observationer i stället för 16, behöver ingen ground truth och försämras mjukt där batteriet försämras i hopp. Den är också exakt den storhet nästa avsnitt behöver.

Det här är löftet Kapitel 1 gav om det här kapitlet, levererat enligt plan: matematiken säger att en 4-bit model är möjlig, och engineering avgör om den är användbar.

Kapitel 12 annonserade detta och lämnade notan här.

Idén kommer rakt ur prefill/decode-delningen. Att verifiera en föreslagen sekvens om $\gamma$ tokens kostar ett forward pass över $\gamma$ positioner — en matrix-matrix product, knappt dyrare än passet över en. Alltså:

Output distribution är oförändrad. Med greedy decoding är det uppenbart — en token accepteras bara om target skulle ha producerat den. Med sampling kräver det en modifierad acceptance rule, och Leviathan et al. bevisar att den resulterande distributionen är exakt target:s.¹⁰ Det här är den andra exakta optimeringen i kapitlet.

Allt hänger därför på acceptance rate $\alpha$, som kan mätas — det är agreement-kolumnen ovan, vilket är varför den beräknades där. Med varje kvantiserad model som draft för full-precision target, över 144 genererade positioner:

Det förväntade antalet tokens som accepteras per verification pass, vid draft-längd $\gamma$, är

och den netto-speedup delar det med draft:ens egen kostnad, en bråkdel $c$ av target per token:

Den feta posten är den du ska minnas: speculative decoding kan göra generation långsammare. Vid 30 % acceptance med en draft som kostar en femtedel av target betalar du för fem forward passes och behåller 1,4 tokens. Den sista kolumnen är den andra fällan — en längre draft hjälper bara när acceptance är hög, eftersom svansen på en $\gamma$-token-gissning nästan aldrig nås. Vid 90 % acceptance är $\gamma = 8$ värt 3,40x och vid 30 % är det värt 0,79x: samma konfiguration, en vinst eller förlust beroende på ett tal uppmätt på din trafik.

Quantization krymper en model genom att lagra samma funktion i färre bits. Distillation krymper den genom att träna en mindre model att imitera en större¹¹ — en idé som föregår deep learning med nästan ett decennium.¹²

Den subtila delen är vad studenten lär sig av. Inte det korrekta svaret: den kunde ha tränats direkt på det. Det läraren tillför är hela distributionen. Fråga modellen vad som följer på en fras och titta förbi argmax:

"She poured the milk into the"
  ' jug' 0.1355   ' cup' 0.1051   ' bowl' 0.0605   ' large' 0.0380   ' milk' 0.0360

Hard label säger jug och inget annat. Soft label säger jug, och också att cup var nästan lika bra, bowl plausibel, och large — ett adjektiv, en helt annan grammatisk fortsättning — fortfarande levande. Det är det ursprungliga argumentet: det här är en 7:a, men den liknar ganska mycket en 1:a, och likheten är information som hard label kastar bort.

Det är också därför distillation använder en temperature. Att dela logits med $T$ före softmax plattar ut distributionen och höjer den relativa vikten hos tvåorna: på den här frasen faller ration mellan top token och den tredje från 2,24 vid $T = 1$ till 1,50 vid $T = 2$ — kvadratroten av den första, vilket är vad det gör med en ratio att dela logits med två. Samma ordning, mer av förlustens attention på nästan-träffarna. Studentens gradient bär lärarens osäkerhet och inte bara dess dom.

Allt i det här kapitlet är nu en summa:

där $T$ är det totala antalet tokens resident över alla samtidiga requests. Tillämpat: 7B- och 70B-raderna antar 8 key-value heads med dimension 128, 13B-raden full multi-head attention med 40 heads, vilket är hur de generationerna av modeller byggdes — och det märks.

8 GB

16 GB

24 GB

Titta på 13B-raden i 8 GB-tabellen. Weights ryms — 6,2 GB av 8 — så med det vanliga sättet att prata ”körs en 13B model på ett 8 GB-kort”. Den har 337 tokens context, vilket inte är en konversation utan knappt en prompt. ”Ryms den?” är fel fråga. Den rätta är ”med hur mycket context, och för hur många användare samtidigt?”

Titta också på de två 16 GB int8-raderna. 7B får 65 378 tokens och 13B får 3 136 — en tjugofaldig skillnad från 5,6 GB extra weights, eftersom 13B här har multi-head attention och dess cache kostar 800 KB per token mot 7B:s 128 KB. Två modeller av liknande storlek, en oanvändbar för långt context, av en orsak som inte syns i någon model cards rubrik.

För tretton kapitel sedan var det här en perceptron med två weights och en bias. Nu är det en transformer som har designats, tränats, alignats, lärt sig att lägga compute på svåra frågor och serveras till en uppmätt kostnad per token — utan någon låda kvar oöppnad i den.

Det slutar här, och det slutar avsiktligt.

Kapitel 14 börjar med modellen någon annanstans. Inte i din process, inte i ditt minne, inte i en variabel du kan printa: på en maskin du inte administrerar, bakom en API key, en port och en nota. Allt som mättes här händer fortfarande — prefill körs fortfarande före första token, cache växer fortfarande med konversationen, batchen du är i tillhör fortfarande någon annan och avgör fortfarande din latency — men från och med nu observerar du det genom en stream av Server-Sent Events, en finish_reason och en HTTP 429 med en Retry-After header. Frågorna förändras med utsiktspunkten: inte hur beräknas denna gradient utan varför tredubblades min faktura. Det gör även språket, och Kapitel 14 förklarar den regeln i stället för att bara annonsera den — fram hit höll koden weights, gradients, logits och tokenizer-bytes; därifrån och framåt håller den en connection, ett retry, en cancellation och ackumulerat state. De tretton kapitlen bakom dig kastas inte bort vid övergången. De är beskrivningen av vad som kör på andra sidan porten.

Två utelämnanden är avsiktliga. FlashAttention (Dao et al., arXiv:2205.14135) är inte en annan attention — den beräknar samma funktion genom att tiling operationen så att $n \times n$ score matrix aldrig skrivs till minnet, vilket är varför de 67 MB i kapitlets andra tabell är mindre i praktiken än aritmetiken antyder. Och själva kernels delegeras: föreläsning 10 i Stanfords CS336 täcker inference systems på ett djup som detta inte försöker nå, och llama.cpp-repositoryt och GGUF-specifikationen är primärkällorna för CPU-sidan.