Uw KV-cache heeft geen bit-probleem. Het heeft een geometrie-probleem.

Bij identieke 2-bit precisie kan één beslissing over welke as je kwantiseert, de benchmarkscore doen schommelen van 2,88 tot 63,53. Sleutels en waarden hebben tegenovergestelde behandeling nodig — en de reden ligt in de aandachtvergelijking, niet in de hardware.

Neem Llama-2-13B. Groepeer de sleutel-waarde-cache door een kwantiseringsgroepgrootte van 32 in twee bits, terwijl alles anders op zijn plaats blijft — hetzelfde model, hetzelfde bit-budget, dezelfde groepsgroottes, dezelfde benchmarks.

Afhankelijk van één enkel element van een implementatiebeslissing, resulteert de CoQA-nauwkeurigheid in 2,88 of 63,53. De score met volledige precisie is 66,37.

De beslissing gaat niet over hoeveel totale bits worden gebruikt. De vraag is eenvoudigweg welke as je kiest om te groeperen wanneer je elke schaalfactor berekent? Wanneer je besluit om het kanaal als groeperingsdimensie te gebruiken (sleutels) en het token als groeperingsdimensie (waarden), kom je ergens binnen vier punten van de volledige precisieprestaties. Als je een van deze keuzes omdraait, ervaar je een kwaliteitsverlies. Als je beide van deze keuzes omdraait, werkt het model niet meer.

Vier manieren om dezelfde 2 bits uit te geven aan dezelfde cache. Resultaten van de KIVI-ablatie op Llama-2-13B bij groepsgrootte 32.

Kwantiseren wordt meestal gezien als één knop: 8 bits, 4 bits, 2 bits, met een gladde nauwkeurigheidskosten. Binnen de KV-cache is het niet zo. Het is het kiezen van coördinatensystemen, en verschillende systemen zijn van toepassing op sleutels en waarden. Dit artikel legt uit waarom. Kort samengevat: kwantiseringsfout hangt af van het bereik van waarden binnen groepen; sleutels en waarden hebben een heel andere structuur; en mensen trippen vaak op omdat je de juiste as niet kunt afleiden uit de waardeverdeling. Je moet kijken naar hoe de fout verandert nadat aandacht het heeft verbruikt. Dat geeft een algemeen principe voor het comprimeren van tussenliggende activaties en een goede reden om aan te nemen dat reconstructiefout niet noodzakelijkerwijs een proxy voor kwaliteit is.

Waarom de KV-cache hier een probleem heeft

Tijdens de generatiefase slaat een transformer alle sleutel-waardeprojectie (KV)-gegevens van tokens die het eerder heeft verwerkt op in een cache, zodat het deze gegevens niet opnieuw hoeft te berekenen. Deze cache groeit lineair met contextlengte en batchgrootte. Uiteindelijk zal dit ertoe leiden dat de cache groter wordt dan het model zelf.

Deze toename in groei kan gemakkelijk worden geïdentificeerd wanneer men naar het geheugengebruik van verschillende delen van het model kijkt. In de KVQuant-analyse van LLaMA-7B, zijn gewichten verantwoordelijk voor ongeveer 98 procent van het geheugen bij een sequentielengte van 512, met activaties bij 2 procent. Bij 128K-context keert het verhouding om tot ongeveer 16 procent gewichten en 84 procent KV-cache. Wanneer we naar een analyse van OPT-175B kijken, die wordt aangehaald door de KIVI-auteurs, vonden ze soortgelijke resultaten. Specifiek, bij een batchgrootte van 512 met een 512-tokenprompt, bereikt de KV-cache 1,2 TB — meerdere keren de grootte van de modelgewichten.

Capaciteit is echter maar de helft van het probleem hier. De GPU moet de hele KV-cache vanaf het apparaatgeheugen lezen voor elk enkel token dat het genereert. Dit betekent dat terwijl de GPU de KV-cache uitleest, de compute-kernen inactief blijven. Door de algehele grootte van de cache te verkleinen, neemt zowel de beschikbare verwerkingsspanning toe als de tijd die wordt besteed aan het wachten op gegevensoverdrachten.

Wat kwantiseringsfout eigenlijk is

Uniforme gehele getallenkwantiseren is eenvoudige wiskunde. Voor een groep getallen, wordt het kleinste getal als nulpunt opgenomen en vervolgens de bereik van die groep gedeeld door het aantal niveaus dat kan worden weergegeven om een stapgrootte te krijgen. Vervolgens wordt elk element afgerond tot het dichtstbijzijnde niveau. Twee onmiddellijke resultaten volgen. Ten eerste is de fout per element begrensd door de helft van een stap. Ten tweede is de stapgrootte het bereik van de groep gedeeld door 2ᴮ − 1. Bij 2 bits hebt u slechts 4 niveaus om het bereik dat binnen die groep bestaat te dekken. Een element dat honderd keer groter is in vergelijking met zijn buren, presteert niet alleen slecht. Het zet de stapgrootte voor alle andere elementen die dezelfde groep delen op, en allemaal worden ze samen grover. De groep is de eenheid van schade. Het kiezen van een as betekent beslissen welke elementen samen lijden. De vraag anders geformuleerd, is het niet langer “hoeveel bits kan ik me veroorloven?” maar “waar zijn extreme waarden en kan ik ze isoleren?”

Sleutels: de uitschieters leven in vaste kanalen

Grote taalmodellen bevatten activaties die ongewoon groot zijn in vergelijking met de meeste activaties. Sun en collega’s hebben deze zeer grote activaties in kaart gebracht voor verschillende modellenfamilies: in Mixtral 8x7B, is de grootste magnitude in de buurt van 7000, terwijl de mediaan van de kenmerkmagnitude ongeveer 0,3 is — ongeveer vier orden van grootte uit elkaar. Deze zijn zeer zeldzaam; ze blijven consistent in dimensies die zelden veranderen met invoer, en ze zijn niet toevallig. Ze fungeren als impliciete biases, en ze zijn wat de aandacht op een paar tokens richt: aandachtsink-gedrag. In de sleutelcache is deze structuur zeer duidelijk: specifieke kanalen dragen consistent grote magnitudes over elk token in een sequentie. Groepeer over tokens, en elke groep bevat die outlier-kanalen, dus elke groep stapgrootte wordt bepaald door de outliers, en alle normale kanalen betalen ervoor. Groepeer over kanalen, en de outlier-kanalen vormen hun eigen groepen. Hun interne bereik is groot, maar zelfbeperkt; de normale kanalen worden met rust gelaten. Resultaten komen overeen. Gemiddeld over lagen en hoofden op Llama-2-13B, rapporteert KIVI een sleutelreconstructiefout van 13,67 onder per-token-groepering tegen 4,55 per kanaal, en — meer belangrijk — een aandachtscorefout van 47,00 tegen 9,60. Kwantiseren van sleutels per token produceert ongeveer vijf keer de scorefout. Scores komen overeen met betekenisvolle metrics voor sleutels; kanaal-kwantisatie blinkt uit op beide fronten.

Waarden: waar de intuïtie breekt

De waarde-cache vertoont geen kanaal-outlierpatroon. Het lijkt redelijk plat te zijn. Op zichzelf, door het bereikargument, zouden we kunnen verwachten dat een van deze assen een soortgelijke kwaliteit van compressie zal produceren.

Ze doen dat niet. Ongeacht hoe sleutelbeheer wordt geïmplementeerd (de resultaten 2,80 en 2,88), comprimeert per-kanaalwaarden de modellencollaps.

En hier is de catch: als je deze verlies meet met de brute reconstructiefout op de oorspronkelijke tensor waarvoor elke waarde is gecomprimeerd, ziet per-kanaalwaardekwantiseren er eigenlijk een beetje beter uit, met 3,73 tegen 4,57. Als je je compressie op de voor de hand liggende manier valideert, zou je de configuratie kiezen die het model vernietigt.

Waarde-cache kwantiseringsfout op Llama-2-13B, gemeten op twee manieren. De opgeslagen-tensor-metric en de verbruikte-uitvoer-metric komen niet overeen met meer dan een orde van grootte.

De oplossing is dat de waarde-cache nooit direct wordt gelezen. Het wordt verbruikt door een matrixproduct: de aandachtsuitvoer is een gewogen som van waardevectoren over tokens, met softmax-aandachtscores als gewichten. Omdat van dit, is de relevante fout die wordt geïntroduceerd tijdens dit proces en niet binnen de tensors zelf. Gemeten in termen van de aandachtsuitvoer, was de volgorde volledig omgekeerd. Relatieve fout gerapporteerd door KIVI voor de aandachtsuitvoer als gevolg van per-token-waardevector-kwantisatie was 3,55 in vergelijking met 49,89 voor per-kanaal-kwantisatie — meer dan veertien keer hoger voor wat leek op de betere keuze op basis van hoe goed het was gecomprimeerd.

De verklaring is aandachtspariteit, die ze mat op 84,3 procent. De meerderheid van de informatie die is opgenomen in de uitvoer kan worden toegeschreven aan een klein aantal zeer belangrijke tokens. Per-token-kwantisatie beperkt elke tokenfout tot die token, dus fouten op onbelangrijke tokens worden vermenigvuldigd met near-nul aandachtsgewichten en verdwijnen effectief. Per-kanaal-kwantisatie smeert elke tokenfout over een gedeelde kanaalschaal, dus slecht weergegeven tokens verontreinigen de weergave van degenen die ertoe doen. De pariteit die aandacht efficiënt maakt, is hetzelfde eigenschap dat per-token-kwantisatie veilig maakt.

De overdraagbare les is breder dan de KV-cache: meet compressiefout waar de tensor wordt verbruikt, niet waar het wordt opgeslagen. Een impliciete veronderstelling die wordt gemaakt door reconstructiefout is dat elke component van een tensor een gelijk gewicht heeft bij het bijdragen aan de einduitvoer. Aandacht doet dit expliciet niet. Elke downstream-bewerking die invoer weegt, gate of sparsifieert, breekt deze veronderstelling. Lezers die bekend zijn met mijn vorige artikel over blindspots in evaluatiemetrics in retrieval-systemen, zullen deze resultaten herkennen als soortgelijk aan eerder beschreven fouten: gemakkelijk berekende metrics die rapporteren over iets anders dan de bedoeling was.

Rotary-embeddings compliceren de sleutels

Er zijn enkele problemen met het gebruik van Rotary Position Embeddings (RoPE). RoPE roteert paren van kanalen op basis van de relatieve positie van elk token. Die menging lost de vaste kanaalstructuur gedeeltelijk op die per-kanaal-sleutel-kwantisatie in de eerste plaats liet werken — een outlier-kanaal wordt geroteerd naar zijn buren, en de buren erven het bereik. KVQuant’s antwoord is ordening: kwantiseer sleutels voordat de rotatie wordt toegepast, en pas RoPE toe na dekwantisatie. Naast per-kanaal-sleutel-kwantisatie, niet-uniforme gegevenstypen en het isoleren van een kleine fractie van outliers, krijgen ze dit onder 0,1 perplexity-degradatie bij 3 bits, en maken het mogelijk om LLaMA-7B te serveren tot 1 miljoen tokens van context op één A100-80GB.

Het is ook belangrijk om het niveau van impact van RoPE te begrijpen. De auteurs van het artikel “RotateKV”, meldden een toename van 145 procent in kwantiseringsfouten zodra RoPE werd toegevoegd, en merkte op dat outlier-kanalen verschillen over aandachtskoppen — wat de reden is dat het toepassen van één gedeelde rotatiematrix overal onvoldoende is, en hoofd-adaptieve rotaties beter doen.

De systeemtaks, en waarom het geen detail is

Per-token-kwantisatie past goed bij decoderen. Elk token arriveert; je kwantiseert het, voegt het toe aan de sequentie (langs de tokendimensie), niets anders verplaatst.

Per-kanaal-kwantisatie past echter niet. Aangezien een kanaals statistiek over tokens strekt die nog niet zijn gegenereerd, kun je geen schaalfactor berekenen wanneer een token binnenkomt. KIVI’s workaround is om de meest recente tokens — tot 128 — in volledige precisie in een restbuffer op te slaan, en te kwantiseren in groepen zodra er genoeg zijn opgehoopt.

Het gebeurt dat de restbuffer belastbaar wordt, in plaats van alleen een incidenteel ding. Op GSM8K met Llama-2-7B, scoort volledige precisie 13,50. Volledig gekwantiseerd tot 2 bits met de juiste assen, scoort het 5,76. Dezelfde assen en dezelfde bits, plus de restbuffer van onlangs gegenereerde tokens in volledige precisie, scoren 12,74. Een schuifvenster van onlangs gegenereerde tokens in volledige precisie zal veel van wat verloren ging door agressieve kwantisatie op moeilijke multi-stap-problemen herstellen — wat zou kloppen als we overwegen welke tokens werden aandachtig bekeken door een keten van rekenkundige bewerkingen.

Er is een aanzienlijk voordeel om al deze dingen correct te doen — zoals KIVI rapporteert, 2,6 keer minder piekgeheugengebruik voor Llama-2-7B, waardoor batchgroottes tot 4 keer groter mogelijk zijn, evenals 2,35 tot 3,47 keer betere doorvoer op een real-world servicetaak.

Wat te doen met dit

1. Gebruik nooit één kwantiseerder voor beide. Gebruik verschillende kwantiseerders voor sleutels (per kanaal) en voor waarden (per token). Een pijplijn die één kwantiseerder toepast op “de KV-cache” heeft waarschijnlijk al het grootste deel van de mogelijke kwaliteit opgeofferd bij het gebruik van een klein aantal bits om elke waarde weer te geven.
2. Kwantiseer sleutels voordat RoPE. Dit is een kwestie van correctheid in plaats van een kwestie van voorkeur.
3. Sla een venster van onlangs gegenereerde tokens op in volledige precisie. Hoewel het opslaan van een dergelijk venster weinig geheugen in beslag neemt in vergelijking met hoe groot een cache kan zijn, is het precies dit gebied dat veel van de nauwkeurigheid voor moeilijke taken genereert.
4. Valideer nooit op reconstructiefout. Valideer altijd op basis van de aandachtsuitvoer of op basis van eindtaakprestaties. De opslagmetric is niet alleen luidruchtig — voor waarden wijst het de verkeerde kant op.
5. Valideer nooit op korte-context multiple-choice-benchmarks. De KIVI-auteurs vermijden expres gesloten taken zoals MMLU voor deze evaluatie, omdat een enkele decoderingsstap die uitvoerlogits leest, de cache amper gebruikt. Elke evaluatie die de cache niet over tijd opbouwt en vervolgens generatie uitvoert, zal nooit in staat zijn om de fouten inherent aan uw systeemontwerp te observeren.

Waar het werk naartoe gaat

Hoewel er nog wat te doen is met betrekking tot de geometrische aard van het probleem, blijven veel onderzoekers manieren bestuderen waarop outlier-kanalen worden verdeeld over de verschillende transformerhoofden, en hoe hardwarebeperkingen de goedkoopste groeperingen beïnvloeden: InnerQ, vouwt kanaalwijze sleutelnormalisatie op in de sleutel- en querygewichten tijdens de prefilleerfase. Hierdoor wordt geen extra overhead gegenereerd tijdens runtime. Bovendien slaat InnerQ hoge precisievensters op voor zowel onlangs gegenereerde tokens als aandachtsink-tokens. Door dit te doen, elimineert InnerQ de mogelijkheid voor outliers in de sink-kanaal om aangrenzende kanalen te verontreinigen.

Anderen stellen voor dat we, in plaats van de hele cache op te slaan, alleen genoeg informatie moeten opslaan om de sleutel en/of waarde(s) op aanvraag te hermaterialiseren uit een kleinere gecachte weergave.

Tenslotte is het belangrijk om te onthouden dat nauwkeurigheid niet de enige parameter is die kwantisatie beïnvloedt. Onlangs gepubliceerde onderzoeken toonden alignementdegradatie aan als gevolg van het kwantiseren van KV-caches. Bovendien documenteerde dit onderzoek alignementdegradatie, zelfs in productie vLLM-serviceomgevingen die FP8-caches gebruiken in combinatie met een trainingsvrije herstelprotocol dat tot 97 procent van wat verloren ging in termen van alignement herstelde. Als zodanig, terwijl een configuratie zijn benchmarkresultaten kan behouden, betekent dit niet noodzakelijkerwijs dat het alle andere relevante parameters behoudt waar je om geeft.

Het algemene principe

Het idee van kwantisatie is gevat als een “precisiebudget”: hoeveel bits kan ik me veroorloven op te offeren? De KV-cache toont aan dat de meer bruikbare vraag structureel is. Precisie wordt toegewezen in groepen; de groep is de eenheid van schade, en de as waarlangs je groepeert, bepaalt welke elementen hun lot delen. De juiste as is die waarlangs je tensor wordt verbruikt, d.w.z. de manier waarop je je tensor gebruikt en NIET hoe je tensor eruitziet wanneer deze in het geheugen wordt opgeslagen. Sleutels worden gebruikt via een dot-productberekening tegen de query. Een enkel beschadigd kanaal zal alle scores vergiftigen. Waarden worden verbruikt via een sparse-gewogen-gemiddelde-bewerking over tokens. Daarom wordt een enkel beschadigd token eenvoudig gewogen uit.

Twee tensors van identieke dimensies en gegenereerd door twee opeenvolgende lagen worden anders behandeld. Het is de moeite waard om te vragen over elke activatie die je van plan bent te comprimeren: wat voor bewerking contracteert dit, en respecteert mijn groepering dit?