Din KV-cache har inte ett bitproblem. Den har ett geometriproblem.

Vid identisk 2-bitars precision kan ett beslut om vilken axel man kvantiserar längs svänga en benchmark-poäng från 2,88 till 63,53. Nycklar och värden behöver motsatt behandling – och orsaken ligger i uppmärksamhets ekvationen, inte i hårdvaran.

Ta Llama-2-13B. Gruppera dess nyckel-värde-cache efter en kvantiseringgruppsstorlek på 32 till två bitar, medan allt annat lämnas på plats – samma modell, samma bitbudget, samma gruppstorlekar, samma benchmark.

Beroende på ett enda element i en implementeringsbeslut, resulterar CoQA-precisionen i antingen 2,88 eller 63,53. Poängen med full precision är 66,37.

Beslutet handlar inte om hur många totala bitar som används. Frågan är enkelt vilken axel man väljer att gruppera över när man beräknar varje skalfaktor? När man bestämmer sig för att använda kanalen som grupperingsdimension (nycklar) och token som grupperingsdimension (värden), hamnar man någonstans inom fyra poäng från full precision-prestanda. Om man vänder på något av dessa val, upplever man en förlust av kvalitet. Om man vänder på båda dessa val, fungerar inte modellen längre.

Fyra sätt att spendera samma 2 bitar på samma cache. Resultat från KIVI-ablationen på Llama-2-13B vid gruppstorlek 32.

Kvantisering anses vanligtvis vara bara en ratt: 8 bitar, 4 bitar, 2 bitar, med en jämn kostnad för precision. Inuti KV-cachen är det inte så. Det handlar om att välja koordinatsystem, och olika system gäller för nycklar och värden. Den här artikeln förklarar varför. Kortfattat: kvantiseringfel beror på värdets omfång inom grupper; nycklar och värden har mycket olika struktur; och människor trippar ofta upp för att man inte kan härleda rätt axel från värdefördelning över huvud taget. Man måste titta på hur felet förändras efter att uppmärksamheten har konsumerat det. Det ger en allmän princip för att komprimera mellanliggande aktiveringar och en bra anledning att ifrågasätta rekonstruktionsfelet som en proxy för kvalitet.

Varför KV-cachen är där detta beter sig

Under genereringsfasen lagrar en transformer all nyckel-värde-projektionsdata (KV) för token som den tidigare har bearbetat i en cache så att den inte behöver beräkna om denna data igen. Den här cachen växer linjärt med kontextlängd och batchstorlek. Till slut kommer detta att resultera i att cachen växer större än modellen själv.

Denna ökning av tillväxt kan lätt identifieras när man tittar på minnesanvändningen för olika delar av modellen. I KVQuant-analysen av LLaMA-7B står vikter för cirka 98 procent av minnet vid en sekvenslängd på 512, med aktiveringar på 2 procent. Vid 128K kontext inverteras förhållandet till cirka 16 procent vikter och 84 procent KV-cache. När vi tittar på en analys av OPT-175B som citeras av KIVI-författarna, fann de liknande resultat. Specifikt, vid en batchstorlek på 512 med en 512-token-prompt, når KV-cachen 1,2 TB – flera gånger större än modellvikterna.

Men kapacitet är bara halva problemet här. GPU:n måste läsa hela KV-cachen från enhetsminnet för varje enskild token som genereras. Det betyder att medan GPU:n läser ut KV-cachen, sitter beräkningskärnorna inaktiva. Som sådan ökar minskningen av den totala cachen storleken inte bara den tillgängliga bearbetningshuvudrummet, utan minskar också den tid som tillbringas med att vänta på dataöverföringar.

Vad kvantiseringfel egentligen består av

Uniform integer-kvantisering är matematiskt sett enkelt. För en grupp av tal, registrerar man det minsta talet som en nollpunkt och delar sedan omfånget för den gruppen med antalet nivåer som kan representeras för att få en stegstorlek. Man avrundar sedan varje element till närmaste steg. Två omedelbara resultat följer. Först är felet per element begränsat till halva steget. Andra, steget är gruppens omfång dividerat med 2^B – 1. Vid 2 bitar har man bara 4 nivåer för att täcka den spridning som finns inom gruppen. Så ett element som är hundra gånger större jämfört med sina grannar presterar inte bara dåligt. Det inflatar steget för alla andra element som delar samma grupp, och alla blir grövre tillsammans. Grupp är enheten för skada. Valet av axel betyder att bestämma vilka element som lider tillsammans. Att formulera frågan på ett annat sätt, är det inte längre “hur många bitar kan jag unna mig att offra?” utan “var är extrema värden och kan jag isolera dem?”

Nycklar: Avvikarna bor i fasta kanaler

Stora språkmodeller innehåller aktiveringar som är ovanligt stora jämfört med de flesta aktiveringar. Sun och kollegor katalogiserade dessa mycket stora aktiveringar över olika modellfamiljer: i Mixtral 8x7B är den största magnituden nära 7000 medan medianfunktionen är runt 0,3 – ungefär fyra storleksordningar isär. Dessa är mycket sällsynta; de förblir fasta i dimensioner som sällan ändras med indata, och de är inte tillfälliga. De fungerar som implicita bias, och de är vad som fokuserar uppmärksamheten på bara några få token: uppmärksamhetssänkor beteende. I nyckelcachen är denna struktur mycket tydlig: specifika kanaler bär mycket stora magnituder konsekvent över varje token i en sekvens. Gruppera efter token, och varje grupp innehåller dessa avvikande kanaler, så varje grupps steget bestäms av avvikarna, och alla vanliga kanaler betalar för det. Gruppera efter kanaler, och de avvikande kanalerna bildar sina egna grupper. Deras inre omfång är stort men självinneslutande; de vanliga kanalerna lämnas ensamma. Resultaten stämmer. Genomsnittligt över lager och huvuden på Llama-2-13B rapporterar KIVI nyckelrekonstruktionsfel på 13,67 under per-token-gruppering mot 4,55 per-kanal, och – viktigare – uppmärksamhetsscore-fel på 47,00 mot 9,60. Kvantisering av nycklar per token producerar ungefär fem gånger större score-fel. Poängen stämmer med meningsfulla mått för nycklar; kanal-kvantisering utmärker sig på båda fronter.

Värden: Där intuitionen bryter

Värdecachen visar inte en kanal-avvikande mönster. Den verkar vara ganska platt. På egen hand, enligt omfångsargumentet, kunde vi förvänta oss att antingen av dessa axlar skulle producera en liknande kvalitet på komprimering.

De gör inte det. Oavsett hur nyckelhanteringen implementeras (2,80- och 2,88-resultaten), kollapsar komprimering per-kanal värden modellen.

Och här är fällan: om du mätte denna förlust med hjälp av det råa rekonstruktionsfelet på den ursprungliga tensorn för vilken varje värde komprimerades, ser per-kanal-värde-kvantisering faktiskt ut att vara något bättre, med 3,73 mot 4,57. Om du validerade din komprimering på det uppenbara sättet, skulle du välja konfigurationen som förstör modellen.

Värdecache-kvantiseringsfel på Llama-2-13B, mätt på två sätt. Lagringsmåttet och konsumtionsmåttet stämmer inte överens med mer än en storleksordning.

Lösningen är att värdecachen aldrig läses direkt. Den konsumeras av en matrisprodukt: uppmärksamhetsutdata är en vägd summa av värdevektorer över token, med softmax-uppmärksamhetsskör som vikter. På grund av detta är det relevanta felet det som introduceras under denna process och inte inom tensorerna själva. Mätt i termer av uppmärksamhetsutdata, var felet i omvänd ordning. Relativt fel rapporterat av KIVI för uppmärksamhetsutdata på grund av per-token-värdevektor-kvantisering var 3,55 jämfört med 49,89 för per-kanal-kvantisering – över fjorton gånger högre för vad som verkade vara det bättre valet baserat på hur väl det var komprimerat.

Förklaringen är uppmärksamhetssparsitet, som de mätte till 84,3 procent. Majoriteten av den information som innehålls i utdata kan tillskrivas ett fåtal mycket viktiga token. Per-token-kvantisering begränsar varje tokens fel till den token, så fel på oviktiga token multipliceras med nästan noll uppmärksamhetsskör och försvinner effektivt. Per-kanal-kvantisering sprider varje tokens fel över en delad kanalskala, så illa representerade token förorenar representationen av de som är viktiga. Sparsiteten som gör uppmärksamheten effektiv är samma egenskap som gör per-token-kvantisering säker.

Den överförbara lärdomen är bredare än KV-cachen: mät komprimeringsfel där tensorn konsumeras, inte där den lagras. En implicit antagande som görs av rekonstruktionsfelet är att varje komponent av en tensor har lika stor vikt när den bidrar till den slutliga utdata. Uppmärksamhet gör det inte explicit. Någon nedströmsoperation som väger, stänger av eller sprider sin indata bryter mot den antaganden. Läsare som är bekanta med min tidigare artikel om blindspots i utvärderingsmått i hämtningsystem kommer att känna igen att dessa resultat är liknande tidigare beskrivna misslyckanden: lättberäknade mått som rapporterar på något annat än vad som avsågs.

Rotary-inkodningar komplicerar nycklarna

Det finns vissa problem med att använda Rotary Position Embeddings (RoPE). RoPE roterar par av kanaler baserat på den relativa positionen för varje token. Den här blandningen löser delvis upp den fasta kanalsstruktur som gjorde per-kanal-nyckel-kvantisering arbete från första början – en avvikande kanal roteras in i sina grannar, och grannarna ärver omfånget. KVQuants svar är ordning: kvantisera nycklar innan rotationen appliceras, och applicera RoPE efter dekvantisering. Tillsammans med per-kanal-nyckel-kvantisering, icke-uniforma datatyper och isolering av en liten fraktion av avvikare, får detta dem under 0,1 perplexitetsdegradering vid 3 bitar, och möjliggör att serva LLaMA-7B upp till 1 miljon token i kontext på en enda A100-80GB.

Det är också viktigt att förstå den nivå av påverkan från RoPE. Författarna till artikeln “RotateKV” rapporterade en ökning på 145 procent i kvantiseringfel när RoPE lades till, och noterade att avvikande kanaler skiljer sig åt mellan uppmärksamhets-huvuden – vilket är varför tillämpning av en delad rotationsmatris överallt är otillräcklig, och huvud-adaptiva rotationer fungerar bättre.

Systemskatten och varför den inte är en detalj

Per-token-kvantisering passar avkodning bra. Varje token anländer; du kvantiserar den, lägger till den i sekvensen (längs token-dimensionen), ingenting annat flyttar.

Men per-kanal-kvantisering passar inte. Eftersom en kanals statistik spänner över token som inte har genererats ännu, kan du inte beräkna en skalfaktor när en token kommer in. KIVI:s lösning är att behålla de senaste token – upp till 128 – i full precision i en restbuffert, och kvantisera i grupper när tillräckligt många har ackumulerats.

Som det händer, blir restbufferten bärande, snarare än bara en tillfällig sak. På GSM8K med Llama-2-7B, full precision-poängen 13,50. Fullt kvantiserad till 2 bitar med rätt axlar, poängen 5,76. Samma axlar och samma bitar, plus restbufferten med nyligen producerade token i full precision, poängen 12,74. Ett glidande fönster av nyligen producerade token i full precision kommer att återställa mycket av vad som förlorades på grund av aggressiv kvantisering på svåra multi-stegsproblem – vilket skulle vara logiskt om vi överväger vilka token som uppmärksammades av en kedja av aritmetiska operationer.

Det finns en betydande fördel med att göra allt detta korrekt – som KIVI rapporterar, 2,6 gånger mindre toppminnesanvändning för Llama-2-7B, vilket möjliggör batchstorlekar upp till 4 gånger större, samt 2,35 till 3,47 gånger bättre genomströmning på en verklig serviceuppgift.

Vad man ska göra med detta

1. Använd aldrig en kvantisator för båda. Använd olika kvantisatorer för nycklar (per-kanal) och för värden (per-token). En pipeline som applicerar en enda kvantisator på “KV-cachen” har förmodligen redan offrat mest möjlig kvalitet när man använder ett litet antal bitar för att representera varje värde.
2. Kvantisera nycklar innan RoPE. Detta är en fråga om korrekthet snarare än en fråga om preferens.
3. Lagra ett fullt precisionsfönster av nyligen genererade token. Även om lagring av ett sådant fönster tar mycket lite minne jämfört med hur stor en cache kan vara, är det just detta område som genererar mycket av precisionen för svåra uppgifter.
4. Validera inte på rekonstruktionsfel. Validera alltid baserat på uppmärksamhetsutdata eller baserat på slutuppgiftsprestanda. Lagringsmåttet är inte bara brusigt – för värden pekar det åt fel håll.
5. Validera inte på kortkontext, flervalsfrågor. KIVI-författarna undvek medvetet slutna uppgifter som MMLU för denna utvärdering, eftersom en enda avkodningssteg som läser utdata-logiterna knappt övade cachen alls. Någon utvärdering som inte bygger upp en cache över tid och sedan utför generering från den kommer aldrig att kunna observera de misslyckanden som är inbyggda i systemdesignen.

Vart arbetet är på väg

Även om det fortfarande finns lite att göra när det gäller den geometriska naturen hos problemet, studerar många forskare sätt på vilka avvikande kanaler fördelas bland olika transformerhuvuden, och hur hårdvarubegränsningar påverkar vilka grupperingar som är billigast: InnerQ viker in kanalvis nyckelnormalisering i nyckel- och frågevikterna under förifyllning. Därför uppstår ingen extra överhuvudet under körning. Dessutom lagrar InnerQ högprecisionsfönster för både nyligen genererade token och uppmärksamhetssänktoken. Genom att göra så eliminerar InnerQ möjligheten för avvikande kanaler i sänkkanalen att förorena grannkanalerna.

Andra föreslår att istället för att lagra hela cachen, borde vi lagra tillräckligt med information för att kunna återskapa nyckeln och/eller värdet på begäran från en mindre cachad representation.

Slutligen är det viktigt att komma ihåg att precision är inte den enda parametern som kvantisering påverkar. Nyligen publicerad forskning visade på en försämring av justering som resulterade från kvantisering av KV-cachar. Dessutom dokumenterade denna forskning en försämring av justering även i produktionsvLLM-tjänstmiljöer som använder FP8-cachar tillsammans med en träningsfri återställningsprotokoll som återställde upp till 97 procent av vad som förlorades i termer av justering. Som sådan, även om en konfiguration kan behålla sina benchmark-resultat, betyder det inte nödvändigtvis att den behåller alla andra relevanta parametrar man bryr sig om.

Den allmänna principen

Idén om kvantisering har ramats in som en “precisionsbudget”: hur många bitar kan jag unna mig att offra? KV-cachen visar att den mer användbara frågan är strukturell. Precision allokeras i grupper; gruppen är enheten för skada, och axeln man grupperar efter bestämmer vilka element som delar sin öde. Den korrekta axeln är den som din tensor konsumeras på, dvs. hur du använder din tensor och INTE hur din tensor ser ut när den lagras i minnet. Nycklar används via en punktproduktberäkning mot frågan. En enda korrupt kanal kommer att förgifta alla poäng. Värden konsumeras genom en sparse-vägd-genomsnittsberäkning över token. Därför är en enda korrupt token bara vägd ut. Två tensorer med identiska dimensioner och genererade av två på varandra följande lager behandlas olika. Det är värt att fråga om varje aktivering man planerar att komprimera: vilken operation kontraherar bort detta, och respekterar min gruppering det? Rage-beräkning över token. Därför är en enda korrupt token bara vägd ut. Två tensorer med identiska dimensioner och genererade av två på varandra följande lager behandlas olika. Det är värt att fråga om varje aktivering man planerar att komprimera: vilken operation kontraherar bort detta, och respekterar min gruppering det?