Din KV-cache har ikke et bit-problem. Den har et geometri-problem.

Ved identisk 2-bit præcision kan én beslutning om, hvilken akse du kvantiserer langs, svinger en benchmark-score fra 2,88 til 63,53. Nøgler og værdier kræver modsatrettede behandlinger – og årsagen ligger i opmærksomheds-ligningen, ikke i hardwaren.

Tag Llama-2-13B. Grupper dine nøgle-værdi-cache efter en kvantiseringsgruppe-størrelse på 32 i to bit, mens du efterlader alt andet på plads – samme model, samme bit-budget, samme gruppe-størrelser, samme benchmarks.

Afhængigt af ét enkelt element i en implementeringsbeslutning resulterer CoQA-nøjagtighed i enten 2,88 eller 63,53. Scorerne ved fuld præcision er 66,37.

Beslutningen handler ikke om, hvor mange bits der bruges i alt. Spørgsmålet er simpelthen, hvilken akse vælger du at gruppere på, når du beregner hver skala-faktor? Når du vælger at bruge kanalen som din gruppe-dimension (nøgler) og tokenet som din gruppe-dimension (værdier), ender du med en score inden for fire punkter fra fuld præcision. Hvis du flipper enten af disse valg, oplever du en kvalitets-tab. Hvis du flipper begge disse valg, fungerer modellen ikke længere.

Fire måder at bruge de samme 2 bit på samme cache. Resultater fra KIVI-ablation på Llama-2-13B ved gruppe-størrelse 32.

Kvantisering anses normalt for at være bare én knap: 8 bit, 4 bit, 2 bit, med en glidende nøjagtigheds-omkostning. Inden for KV-cache er det ikke sådan. Det handler om at vælge koordinat-systemer, og forskellige systemer gælder for nøgler og værdier. Denne artikel forklarer hvorfor. Kort sagt: kvantiserings-fejl afhænger af værdi-rækkevidden inden for grupper; nøgler og værdier har meget forskellig struktur; og mennesker tripper ofte over, fordi du ikke kan aflede den rigtige akse fra værdi-fordelingen overhovedet. Du må se på, hvordan fejlen ændrer sig efter, at opmærksomhed har forbrugt den. Dette giver en generel princip for at komprimere mellemlege-aktiveringer og en god grund til at tvivle på rekonstruktions-fejl som en proxy for kvalitet.

Hvorfor KV-cachen er, hvor dette bider

Under genererings-fasen gemmer en transformer alle nøgle-værdi-projektions-data (KV) for token, som den tidligere har behandlet, i en cache, så den ikke behøver at regne denne data igen. Denne cache vokser lineært med kontekst-længde og batch-størrelse. Til sidst vil dette resultere i, at cachen vokser større end modellen selv.

Denne vækst kan let identificeres, når man ser på hukommelses-forbrug af forskellige dele af modellen. I KVQuant-analyse af LLaMA-7B, udgør vægte omtrent 98 procent af hukommelsen ved en sekvens-længde på 512, med aktiveringer på 2 procent. Ved 128K kontekst inverterer forholdet til omkring 16 procent vægte og 84 procent KV-cache. Når vi ser på en analyse af OPT-175B citeret af KIVI-forfatterne, fandt de lignende resultater. Specifikt, ved en batch-størrelse på 512 med en 512-token prompt, når KV-cachen 1,2 TB – flere gange størrelsen på model-vægtene.

Men kapacitet er kun halvdelen af problemet her. GPU’en skal læse hele KV-cachen fra enhedshukommelsen for hver enkelt token, den genererer. Dette betyder, at mens GPU’en læser ud KV-cachen, sidder beregnings-kernerne inaktive. Som sådan reducerer reducering af cachens samlede størrelse både den tilgængelige beregnings-hovedrum og tiden, der bruges på at vente på data-overførsler.

Hvad kvantiserings-fejl faktisk består af

Uniform integer-kvantisering er matematisk set ret enkelt. For en gruppe tal, optegner du det mindste tal som et nulpunkt og derefter dividerer gruppen af tal med antallet af niveauer, der kan repræsenteres, for at få en skridt-størrelse. Du afrunder derefter hvert element til det nærmeste skridt. To umiddelbare resultater følger. Først er fejlen per element begrænset af halv en skridt-størrelse. Anden, skridt-størrelsen er gruppen af tal divideret med 2^B – 1. Ved 2 bit har du kun 4 niveauer til at dække, hvad som helst der findes inden for gruppen. Så et element, der er hundred gange større i forhold til sine naboer, fungerer ikke kun dårligt. Det øger skridt-størrelsen for alle andre elementer, der deler samme gruppe, og alle bliver grovere sammen. Gruppe er enheden for skade. At vælge en akse betyder at bestemme, hvilke elementer lider sammen. At formulere spørgsmålet anderledes, er det ikke længere “hvor mange bit kan jeg af med?” men “hvor er ekstreme værdier, og kan jeg isolere dem?”

Nøgler: Outlierne bor i faste kanaler

Store sprogmodeller indeholder aktiveringer, der er usædvanligt store i forhold til de fleste aktiveringer. Sun og kolleger katalogiserede disse meget store aktiveringer over forskellige modelfamilier: i Mixtral 8x7B er den største størrelse nær 7000, mens median-feature-størrelsen er omkring 0,3 – omtrent fire størrelsesordener fra hinanden. Disse er meget sjældne; de forbliver faste i dimensioner, der sjældent ændrer sig med input, og de er ikke tilfældige. De fungerer som implicit bias, og de er, hvad der fokuserer opmærksomhed på kun få token: opmærksomhed-sænkning-adfærd. I nøgle-cachen er denne struktur meget tydelig: bestemte kanaler bærer meget store størrelser konsekvent over hver token i en sekvens. Grupperer du efter token, indeholder hver gruppe disse outlier-kanaler, så hver gruppes skridt-størrelse er fastsat af outlier-kanalerne, og alle almindelige kanaler betaler for det. Grupperer du efter kanaler, danner outlier-kanalerne deres egne grupper. Deres interne rækkevidde er stor, men selv-indsat; de almindelige kanaler bliver alene. Resultaterne matcher. Gennemsnitligt over lag og hoveder på Llama-2-13B rapporterer KIVI nøgle-rekonstruktions-fejl på 13,67 under per-token-gruppering mod 4,55 per-kanal, og – vigtigere – opmærksomheds-score-fejl på 47,00 mod 9,60. Kvantisering af nøgler per token producerer omtrent fem gange større score-fejl. Scorerne er enige med meningsfulde metrikker for nøgler; kanal-kvantisering excellerer på begge fronter.

Værdier: Hvor intuitionen fejler

Værdi-cachen viser ikke en kanal-outlier-mønster. Den synes at være ret flad. På sig selv, efter rækkevidde-argumentet, kunne vi forvente, at enten af disse akser ville producere en lignende kvalitet af kompression.

De gør det ikke. Uanset, hvordan nøgle-håndtering er implementeret (2,80 og 2,88 resultaterne), kollapser komprimering per-kanal værdier modellen.

Og her er fælden: hvis du målte denne tab ved hjælp af den rene rekonstruktions-fejl på den oprindelige tensor for hver værdi, der blev komprimeret, ser per-kanal-værdi-kvantisering faktisk lidt bedre ud, med 3,73 mod 4,57. Hvis du validerede din kompression på den åbenlyse måde, ville du vælge den konfiguration, der ødelægger modellen.

Værdi-cache kvantiserings-fejl på Llama-2-13B, målt på to måder. Den gemte-tensor-metrik og den forbrugte-udgangs-metrik er uenige med mere end en størrelsesorden.

Løsningen er, at værdi-cachen aldrig læses direkte. Den forbruges af en matrix-produkt: opmærksomheds-udgangen er en vægtet sum af værdi-vektorer over token, med softmax-opmærksomheds-scorer som vægt. Fordi af denne, er den relevante fejl den, der introduceres under denne proces, og ikke inden for tensorerne selv. Målt i forhold til opmærksomheds-udgangen, var den relative fejl rapporteret af KIVI for opmærksomheds-udgangen på grund af per-token-værdi-vektor-kvantisering 3,55 sammenlignet med 49,89 for per-kanal-kvantisering – over 14 gange højere for, hvad der så ud til at være det bedre valg baseret på, hvor godt det var komprimeret.

Forklaringen er opmærksomheds-sparthed, som de målte til 84,3 procent. Flertallet af den information, der er indeholdt i udgangen, kan tilskrives en lille mængde meget vigtige token. Per-token-kvantisering begrænser hver tokens fejl til det token, så fejl på unyttige token bliver multipliceret med næsten-nul-opmærksomheds-vægte og forsvinder effektivt. Per-kanal-kvantisering spreder hver tokens fejl over en delt kanal-skala, så dårligt repræsenterede token forurener repræsentationen af dem, der er vigtige. Spartheden, der gør opmærksomhed effektiv, er den samme egenskab, der gør per-token-kvantisering sikker.

Den overførbare lære er bredere end KV-cachen: mål kompressions-fejl, hvor tensor er forbrugt, ikke hvor den er gemt. En implicit antagelse, der er lavet af rekonstruktions-fejl, er, at hver komponent af en tensor har lige stor vægt, når det bidrager til den endelige udgang. Opmærksomhed gør ikke implicit. Enhver downstream-operation, der vægter, lukker eller sparsommer sin input, bryder denne antagelse. Læsere, der er bekendt med min tidligere artikel om blindspots i evaluering-metrikker i retrieval-systemer, vil genkende, at disse resultater er lignende med tidligere beskrevne fejl: let beregnelige metrikker, der rapporterer om noget andet end, hvad der var tiltænkt.

Rotary-embedding komplicerer nøglerne

Der er nogle problemer med at bruge Rotary Position Embeddings (RoPE). RoPE roterer par af kanaler baseret på den relative position af hver token. Denne blanding opløser delvist den faste-kanal-struktur, der gjorde per-kanal-nøgle-kvantisering til at fungere i første omgang – en outlier-kanal bliver roteret ind i sine naboer, og naboerne arver rækkevidden. KVQuants svar er orden: kvantisér nøgler, før rotationen anvendes, og anvend RoPE efter dekvantisering. Ved siden af per-kanal-nøgle-kvantisering, ikke-uniforme datatyper og isolering af en lille fraktion af outlier, får dette dem under 0,1 perplexity-degradering ved 3 bit, og muliggør servering af LLaMA-7B op til 1 million token af kontekst på en enkelt A100-80GB.

Det er også vigtigt at forstå niveauet af impact fra RoPE. Forfatterne af artiklen “RotateKV” rapporterede en øgning på 145 procent i kvantiserings-fejl, når RoPE blev tilføjet, og bemærkede, at outlier-kanaler forskellige sig over opmærksomheds-hoveder – hvilket er, hvorfor anvendelse af en delt rotations-matrix overalt er utilstrækkelig, og hoved-adaptiver rotationer gør bedre.

System-skatten og hvorfor det ikke er en detalje

Per-token-kvantisering passer godt til afkodning. Hver token ankommer; du kvantiserer det, tilføjer det til sekvensen (langs token-dimensionen), intet andet flytter sig.

Men per-kanal-kvantisering passer ikke. Da en kanals statistik spænder over token, der endnu ikke er blevet genereret, kan du ikke beregne en skala-faktor, når en token kommer ind. KIVIs workaround er at holde de seneste token – op til 128 – i fuld præcision i en residual-buffer, og kvantiserer i grupper, når nok har akkumuleret.

Som det sker, bliver residual-bufferen bærende, snarere end bare en tilfældig ting. På GSM8K med Llama-2-7B, scorer fuld præcision 13,50. Fuld kvantiseret til 2 bit med de korrekte akser, scorer 5,76. Samme akser og samme bit, plus residual-bufferen af nyligt producerede token i fuld præcision, scorer 12,74. En glide-vindue af nyligt producerede token i fuld præcision vil gendanne meget af, hvad der var tabt på grund af aggressiv kvantisering på svære multi-trin-problemer – hvilket ville have mening, hvis vi overvejer, hvilke token der blev opmærksomt af en kæde af aritmetiske operationer.

Der er en betydelig fordel ved at gøre alle disse ting korrekt – som KIVI rapporterer, 2,6 gange mindre peak-hukommelses-forbrug for Llama-2-7B, hvilket tillader batch-størrelser op til 4 gange større, samt 2,35 til 3,47 gange bedre gennemstrømning på en real-world service-opgave.

Hvad du skal gøre med dette

1. Brug aldrig én kvantiser til begge. Brug forskellige kvantisører til nøgler (per-kanal) og til værdier (per-token). En pipeline, der anvender en enkelt kvantiser til “KV-cachen”, har sandsynligvis allerede ofret det meste af den mulige kvalitet, når du bruger et lille antal bit til at repræsentere hver værdi.
2. Kvantiser nøgler, før RoPE. Dette er en spørgsmål om korrekthed snarere end en spørgsmål om præference.
3. Gem en fuld præcision-vindue af nyligt genererede token. Selvom gemme af et sådant vindue tager meget lidt hukommelse i forhold til, hvor stor en cache kan være, er det præcis dette område, der genererer meget af nøjagtigheden for svære opgaver.
4. Valider aldrig på rekonstruktions-fejl. Valider altid baseret på opmærksomheds-udgangen eller baseret på end-opgave-præstation. Lager-metrikken er ikke kun støj – for værdier peger den den forkerte vej.
5. Valider aldrig på kort-kontekst multiple-choice-benchmarks. KIVI-forfatterne undgår bevidst lukkede opgaver som MMLU til denne evaluering, fordi en enkelt afkodnings-trin læser output-logit ikke udnytter cachen overhovedet. Enhver evaluering, der ikke bygger en cache over tid og derefter udfører generation fra den, vil aldrig være i stand til at observere fejlene, der er indbygget i dit system-design.

Hvor arbejdet er på vej

Selvom der stadig er noget at gøre med den geometriske natur af problemet, studerer mange forskere måder, hvorpå outlier-kanaler er fordelt blandt de forskellige transformer-hoveder, og hvordan hardware-begrænsninger påvirker, hvilke grupperinger er billigst: InnerQ folder kanal-vis nøgle-normalisering ind i nøgle- og forespørgsels-vægt under prefill. Derfor incureres ingen ekstra omkostning ved runtime. Desuden gemmer InnerQ høj-præcision-vinduer for både nyligt genererede token og opmærksomheds-sænk-token. Ved at gøre dette eliminerer InnerQ muligheden for, at outlier i sink-kanalen forurener nabokanaler.

Andre foreslår, at i stedet for at gemme hele cachen, skal vi kun gemme nok information til at kunne genopbygge nøglen og/eller værdierne på krav fra en mindre cached-repræsentation.

Endelig er det vigtigt at huske, at nøjagtighed ikke er det eneste parameter, der påvirkes af kvantisering. Nyligt publiceret forskning demonstrerede alignment-degradering som følge af kvantisering af KV-cacher. Desuden dokumenterede denne forskning alignment-degradering, selv i produktions-vLLM-server-miljøer, der anvender FP8-cacher sammen med en trænings-fri genoprettelses-protokol, der genskabte op til 97 procent af, hvad der var tabt i forhold til alignment. Som sådan, selvom en konfiguration kan fastholde sine benchmark-resultater, betyder det ikke nødvendigvis, at den fastholder alle andre relevante parametre, du bekymrer dig om.

Den generelle princip

Idéen om kvantisering er blevet ramt som en “præcisions-budget”: hvor mange bit kan jeg af med? KV-cachen viser, at det mere nyttige spørgsmål er strukturelt. Præcision er allokeret i grupper; gruppen er enheden for skade, og aksen, du grupper efter, bestemmer, hvilke elementer deler deres skæbne. Den korrekte akse er den, på hvilken din tensor forbruges, dvs. den måde, du bruger din tensor på, og ikke, hvordan din tensor ser ud, når den er gemt i hukommelsen. Nøgler bruges via en punkt-produkt-beregning mod forespørgslen. En enkelt korrupt kanal vil forgifte alle scoringer. Værdier forbruges gennem en sparsom-vægtet-gennemsnits-beregning over token. Derfor er en enkelt korrupt token blot vægtet ud.

To tensorer af identiske dimensioner og genereret af to på hinanden følgende lag behandles forskelligt. Det er værd at spørge om enhver aktivering, du planlægger at komprimere: hvilken operation kontraherer dette væk, og respekterer min gruppering det? Rage-komputation over token. Derfor er en enkelt korrupt token blot vægtet ud. To tensorer af identiske dimensioner og genereret af to på hinanden følgende lag behandles forskelligt. Det er værd at spørge om enhver aktivering, du planlægger at komprimere: hvilken operation kontraherer dette væk, og respekterer min gruppering det?