Ihr KV-Cache hat kein Bit-Problem. Er hat ein Geometrie-Problem.

Bei identischer 2-Bit-Präzision kann eine Entscheidung über die Achse, entlang derer quantifiziert wird, den Benchmark-Score von 2,88 auf 63,53 schwanken. Schlüssel und Werte benötigen entgegengesetzte Behandlung – und der Grund liegt in der Aufmerksamkeitsgleichung, nicht in der Hardware.

Nehmen wir Llama-2-13B. Gruppieren Sie den Schlüssel-Wert-Cache nach einer Quantisierungsgruppengröße von 32 in zwei Bits, während alles andere an seinem Platz bleibt – dasselbe Modell, derselbe Bit-Budget, dieselben Gruppengrößen, dieselben Benchmarks.

Je nach einer einzigen Implementierungsentscheidung ergeben sich bei CoQA-Genauigkeit entweder 2,88 oder 63,53. Der Score bei voller Präzision beträgt 66,37.

Die Entscheidung dreht sich nicht darum, wie viele Gesamtbits verwendet werden. Die Frage ist einfach, welche Achse Sie als Gruppierungsdimension für die Berechnung jedes Skalierungsfaktors wählen. Wenn Sie sich entscheiden, den Kanal als Gruppierungsdimension (Schlüssel) und den Token als Gruppierungsdimension (Werte) zu verwenden, landen Sie innerhalb von vier Punkten der Vollpräzisionsleistung. Wenn Sie eine dieser Entscheidungen ändern, erleben Sie einen Qualitätsverlust. Wenn Sie beide dieser Entscheidungen ändern, funktioniert das Modell nicht mehr.

Vier Möglichkeiten, die gleichen 2 Bits auf dem gleichen Cache auszugeben. Ergebnisse aus der KIVI-Ablation auf Llama-2-13B bei Gruppengröße 32.

Quantisierung wird normalerweise als eine einzige Einstellung betrachtet: 8 Bits, 4 Bits, 2 Bits, mit einer glatten Kosten-Genauigkeits-Beziehung. Innerhalb des KV-Caches ist es nicht so. Es geht darum, Koordinatensysteme zu wählen, und unterschiedliche Systeme gelten für Schlüssel und Werte. Dieser Artikel erklärt, warum. Kurz gesagt: Quantisierungsfehler hängen vom Wertebereich innerhalb der Gruppen ab; Schlüssel und Werte haben sehr unterschiedliche Strukturen; und Menschen stolpern oft, weil man die richtige Achse nicht aus der Verteilung der Werte ableiten kann. Man muss sich ansehen, wie sich der Fehler nach der Aufmerksamkeitsverarbeitung ändert. Das gibt einen allgemeinen Grundsatz für die Komprimierung von Zwischenaktivierungen und einen guten Grund, den Rekonstruktionsfehler als Qualitätsproxy in Frage zu stellen.

Warum der KV-Cache der Punkt ist, an dem dies beißt

Während der Generierungsphase speichert ein Transformer alle Schlüssel- und Werteprojektionsdaten (KV) der zuvor verarbeiteten Token in einem Cache, um diese Daten nicht erneut berechnen zu müssen. Dieser Cache wächst linear mit der Kontextlänge und der Batch-Größe. Schließlich wird dieser Cache größer als das Modell selbst.

Dieses Wachstum kann leicht identifiziert werden, wenn man den Speicherbedarf verschiedener Teile des Modells betrachtet. In der KVQuant-Analyse von LLaMA-7B machen Gewichte etwa 98 Prozent des Speichers bei einer Sequenzlänge von 512 aus, während Aktivierungen bei 2 Prozent liegen. Bei 128K-Kontext kehrt sich das Verhältnis um auf etwa 16 Prozent Gewichte und 84 Prozent KV-Cache. Wenn wir eine Analyse von OPT-175B betrachten, die von den KIVI-Autoren zitiert wird, finden wir ähnliche Ergebnisse. Insbesondere bei einer Batch-Größe von 512 mit einem 512-Token-Prompt erreicht der KV-Cache 1,2 TB – mehrere Male die Größe der Modellgewichte.

Allerdings ist die Kapazität nur die Hälfte des Problems hier. Der GPU muss den gesamten KV-Cache von der Gerätespeicher für jeden generierten Token lesen. Das bedeutet, dass der GPU beim Lesen des KV-Caches stillsteht, während die Recheneinheiten untätig bleiben. Durch die Reduzierung der Gesamtgröße des Caches erhöht sich die verfügbare Verarbeitungskapazität und reduziert sich die Zeit, die für Datenübertragungen aufgewendet wird.

Was der Quantisierungsfehler eigentlich ist

Die einheitliche ganzzahlige Quantisierung ist mathematisch einfach. Für eine Gruppe von Zahlen wird die kleinste Zahl als Nullpunkt aufgezeichnet und dann der Bereich der Gruppe durch die Anzahl der darstellbaren Stufen geteilt, um eine Schrittgröße zu erhalten. Jedes Element wird dann auf die nächste Stufe gerundet. Zwei sofortige Ergebnisse folgen. Erstens ist der Fehler pro Element durch die halbe Schrittgröße begrenzt. Zweitens ist die Schrittgröße der Bereich der Gruppe geteilt durch 2^B – 1. Bei 2 Bits haben wir nur 4 Stufen, um jeden Bereich abzudecken, der existiert. Ein Element, das hundertmal größer ist als seine Nachbarn, funktioniert nicht nur schlecht. Es bläht die Schrittgröße für alle anderen Elemente auf, die dieselbe Gruppe teilen, und alle werden zusammen gröber. Die Gruppe ist die Einheit des Schadens. Die Wahl der Achse bedeutet, zu entscheiden, welche Elemente zusammen leiden. Anders ausgedrückt, ist es nicht mehr “wie viele Bits kann ich mir leisten?” sondern “wo sind extreme Werte und kann ich sie isolieren?”

Schlüssel: Die Ausreißer leben in festen Kanälen

Große Sprachmodelle enthalten Aktivierungen, die im Vergleich zu den meisten Aktivierungen ungewöhnlich groß sind. Sun und Kollegen haben diese sehr großen Aktivierungen in verschiedenen Modellfamilien katalogisiert: in Mixtral 8x7B liegt die größte Größenordnung bei etwa 7000, während die Median-Feature-Größe etwa 0,3 beträgt – etwa vier Größenordnungen auseinander. Diese sind sehr selten; sie bleiben in Dimensionen, die sich selten mit der Eingabe ändern, und sie sind nicht zufällig. Sie wirken als implizite Vorurteile und lenken die Aufmerksamkeit auf nur wenige Token: Aufmerksamkeitssenken-Verhalten. Im Schlüssel-Cache ist diese Struktur sehr klar: bestimmte Kanäle tragen konsistent über jeden Token in einer Sequenz sehr große Größen. Gruppieren Sie entlang der Token, und jede Gruppe enthält diese Ausreißer-Kanäle, so dass jede Gruppen-Schrittgröße durch die Ausreißer bestimmt wird und alle normalen Kanäle dafür bezahlen. Gruppieren Sie entlang der Kanäle, und die Ausreißer-Kanäle bilden ihre eigenen Gruppen. Ihr interner Bereich ist groß, aber selbst enthalten; die normalen Kanäle bleiben unberührt. Die Ergebnisse stimmen überein. Im Durchschnitt über Schichten und Köpfe auf Llama-2-13B berichtet KIVI über einen Schlüssel-Rekonstruktionsfehler von 13,67 bei pro-Token-Gruppierung gegenüber 4,55 pro-Kanal, und – wichtiger noch – einem Aufmerksamkeits-Score-Fehler von 47,00 gegenüber 9,60. Die Quantisierung von Schlüsseln pro Token produziert etwa fünfmal den Score-Fehler. Die Scores stimmen mit sinnvollen Metriken für Schlüssel überein; Kanal-Quantisierung excelt auf beiden Seiten.

Werte: Wo die Intuition bricht

Der Wert-Cache zeigt kein Kanal-Ausreißer-Muster. Er erscheint ziemlich flach. Allein durch den Bereichs-Argument würde man erwarten, dass entweder dieser Achsen eine ähnliche Qualität der Komprimierung produziert.

Sie tun es nicht. Unabhängig davon, wie die Schlüsselverwaltung implementiert wird (die Ergebnisse 2,80 und 2,88), kollabiert die Komprimierung pro-Kanal-Werte das Modell.

Und hier ist der Haken: Wenn man diesen Verlust mit dem rohen Rekonstruktionsfehler auf dem ursprünglichen Tensor misst, für den jeder Wert komprimiert wurde, sieht die pro-Kanal-Wert-Quantisierung tatsächlich ein bisschen besser aus, bei 3,73 gegenüber 4,57. Wenn man die Komprimierung auf die übliche Weise validiert, würde man die Konfiguration wählen, die das Modell zerstört.

Wert-Cache-Quantisierungsfehler auf Llama-2-13B, auf zwei Arten gemessen. Die gespeicherte-Tensor-Metrik und die konsumierte-Ausgabe-Metrik stimmen um mehr als eine Größenordnung nicht überein.

Die Auflösung ist, dass der Wert-Cache nie direkt gelesen wird. Er wird durch ein Matrix-Produkt konsumiert: Die Aufmerksamkeitsausgabe ist eine gewichtete Summe von Wert-Vektoren über Token, mit Softmax-Aufmerksamkeitsscores als Gewichten. Da dies der Fall ist, ist der relevante Fehler der, der während dieses Prozesses eingeführt wird, und nicht innerhalb der Tensoren selbst. Gemessen in Bezug auf die Aufmerksamkeitsausgabe wurde der relative Fehler durch KIVI für die pro-Token-Wert-Vektor-Quantisierung mit 3,55 gegenüber 49,89 für pro-Kanal-Quantisierung berichtet – mehr als vierzehnmal höher für das, was aufgrund der Komprimierung die bessere Wahl zu sein schien.

Die Erklärung ist die Aufmerksamkeitssparheit, die sie mit 84,3 Prozent gemessen haben. Der größte Teil der Information, die in der Ausgabe enthalten ist, kann einigen wenigen sehr wichtigen Token zugeschrieben werden. Pro-Token-Quantisierung beschränkt jeden Tokens Fehler auf diesen Token, so dass Fehler auf unimportanten Tokenen durch nahezu Null-Aufmerksamkeitsscores multipliziert werden und effektiv verschwinden. Pro-Kanal-Quantisierung verteilt jeden Tokens Fehler über einen gemeinsamen Kanal-Skalierungs-Faktor, so dass schlecht dargestellte Tokenen die Darstellung der wichtigen Tokenen kontaminieren. Die Sparheit, die die Aufmerksamkeit effizient macht, ist dieselbe Eigenschaft, die pro-Token-Quantisierung sicher macht.

Die übertragbare Lektion ist breiter als der KV-Cache: Messen Sie den Komprimierungsfehler dort, wo der Tensor konsumiert wird, nicht dort, wo er gespeichert wird. Eine implizite Annahme, die durch den Rekonstruktionsfehler gemacht wird, ist, dass jeder Bestandteil eines Tensors gleiches Gewicht hat, wenn er zur endgültigen Ausgabe beiträgt. Aufmerksamkeit tut dies explizit nicht. Jede nachgelagerte Operation, die ihre Eingabe gewichtet, torpediert oder sparsam macht, bricht diese Annahme. Leser, die mit meinem vorherigen Artikel über Blindspots in Bewertungsmetriken in Retrieval-Systemen vertraut sind, werden erkennen, dass diese Ergebnisse ähnlich sind wie die zuvor beschriebenen Ausfälle: leicht berechnete Metriken, die über etwas anderes berichten, als beabsichtigt.

Rotary-Embeddings komplizieren die Schlüssel

Es gibt einige Probleme mit der Verwendung von Rotary-Position-Embeddings (RoPE). RoPE rotiert Paare von Kanälen basierend auf der relativen Position jedes Tokens. Diese Mischung löst die feste Kanal-Struktur teilweise auf, die die pro-Kanal-Schlüssel-Quantisierung zum Funktionieren brachte – ein Ausreißer-Kanal wird in seine Nachbarn rotiert, und die Nachbarn erben den Bereich. KVQuants Antwort ist die Reihenfolge: Quantisieren Sie die Schlüssel, bevor die Rotation angewendet wird, und wenden Sie RoPE nach der De-Quantisierung an. Neben der pro-Kanal-Schlüssel-Quantisierung, nicht-uniformen Datentypen und der Isolierung einer kleinen Anzahl von Ausreißern erhält man so weniger als 0,1 Perplexitäts-Abnahme bei 3 Bits und ermöglicht es, LLaMA-7B mit bis zu 1 Million Token-Kontext auf einem einzelnen A100-80GB zu bedienen.

Es ist auch wichtig, den Umfang des Einflusses von RoPE zu verstehen. Die Autoren des Papiers “RotateKV” berichteten über eine Erhöhung der Quantisierungsfehler um 145 Prozent, nachdem RoPE hinzugefügt wurde, und stellten fest, dass Ausreißer-Kanäle sich über Aufmerksamkeits-Köpfe unterscheiden – was der Grund dafür ist, dass die Anwendung einer gemeinsamen Rotationsmatrix überall unzureichend ist und head-adaptive Rotationen besser abschneiden.

Die Systemsteuer, und warum sie kein Detail ist

Pro-Token-Quantisierung eignet sich gut für die Dekodierung. Jedes Token kommt; Sie quantisieren es, fügen es der Sequenz hinzu (entlang der Token-Dimension), nichts anderes bewegt sich.

Pro-Kanal-Quantisierung hingegen passt nicht. Da die Statistik eines Kanals Token umfasst, die noch nicht generiert wurden, kann der Skalierungsfaktor nicht berechnet werden, wenn ein Token hereinkommt. KIVIs Workaround ist, die meisten recenten Token – bis zu 128 – in voller Präzision in einem Rest-Buffer zu halten und zu quantisieren, sobald genug angesammelt wurden.

Es stellt sich heraus, dass der Rest-Buffer lasttragend wird, anstatt nur eine Nebensache zu sein. Bei GSM8K mit Llama-2-7B erzielt die volle Präzision einen Score von 13,50. Vollständig quantifiziert auf 2 Bits mit den richtigen Achsen, erzielt es einen Score von 5,76. Dieselben Achsen und dieselben Bits, plus der Rest-Buffer der kürzlich produzierten Token in voller Präzision, erzielen einen Score von 12,74. Ein gleitendes Fenster kürzlich produzierter Token in voller Präzision kann viel von dem wiederherstellen, was durch aggressive Quantisierung bei schwierigen Mehrschritt-Problemen verloren ging – was sinnvoll wäre, wenn man bedenkt, welche Token von einer Kette arithmetischer Operationen aufmerksamkeitsmäßig verfolgt werden.

Es gibt einen erheblichen Vorteil, wenn man all dies richtig macht – wie KIVI berichtet, 2,6-mal weniger Spitzen-Speicherbedarf für Llama-2-7B, was Batch-Größen bis zu viermal größer ermöglicht, sowie 2,35- bis 3,47-mal bessere Durchsatzleistung bei einer realen Service-Aufgabe.

Was man damit tun kann

1. Verwenden Sie nie denselben Quantisierer für beide. Verwenden Sie unterschiedliche Quantisierer für Schlüssel (pro-Kanal) und Werte (pro-Token). Eine Pipeline, die einen einzelnen Quantisierer auf “den KV-Cache” anwendet, hat wahrscheinlich bereits die meisten möglichen Qualitäten aufgegeben, wenn sie eine kleine Anzahl von Bits verwendet, um jeden Wert darzustellen.
2. Quantisieren Sie Schlüssel vor RoPE. Dies ist eine Frage der Richtigkeit und nicht der Vorliebe.
3. Speichern Sie ein Fenster von kürzlich generierten Token in voller Präzision. Obwohl das Speichern eines solchen Fensters im Vergleich zu der Größe des Caches sehr wenig Speicher benötigt, ist es genau diese Fläche, die bei schwierigen Aufgaben viel Genauigkeit erzeugt.
4. Validieren Sie nicht auf Rekonstruktionsfehler. Validieren Sie immer auf der Grundlage der Aufmerksamkeitsausgabe oder der Endaufgabenleistung. Die Speicher-Metrik ist nicht nur laut – für Werte zeigt sie in die falsche Richtung.
5. Validieren Sie nicht auf kurze Kontext-Mehrfachwahl-Benchmarks. Die KIVI-Autoren vermeiden absichtlich geschlossene Aufgaben wie MMLU für diese Bewertung, da ein einzelner Dekodierungsschritt, der die Ausgabe-Logits liest, den Cache überhaupt nicht beansprucht. Jede Bewertung, die den Cache nicht im Laufe der Zeit aufbaut und dann nicht aus ihm generiert, kann die im System-Design inhärenten Fehler nie beobachten.

Wohin die Arbeit geht

Obwohl noch einige Dinge bezüglich der geometrischen Natur des Problems zu tun sind, untersuchen viele Forscher weiterhin, wie Ausreißer-Kanäle in verschiedenen Transformer-Köpfen verteilt sind und wie Hardware-Einschränkungen die Gruppierungen beeinflussen, die am billigsten sind: InnerQ falten die kanalweise Schlüssel-Normierung in die Schlüssel- und Abfrage-Gewichte während der Vorausfüllung ein. Daher entsteht keine zusätzliche Overhead bei der Laufzeit. Darüber hinaus speichert InnerQ Fenster von kürzlich generierten Tokenen und Aufmerksamkeitssenken-Tokenen in hoher Präzision. Indem es dies tut, eliminiert InnerQ die Gelegenheit, dass Ausreißer im Senken-Kanal benachbarte Kanäle kontaminieren.

Andere schlagen vor, dass anstelle des gesamten Caches nur genug Information gespeichert werden sollte, um den Schlüssel und/oder den Wert (en) auf Abruf zu rematerialisieren aus einer kleineren zwischengespeicherten Darstellung.

Schließlich ist es wichtig, daran zu denken, dass die Genauigkeit nicht der einzige Parameter ist, den die Quantisierung beeinflusst. Kürzlich veröffentlichte Forschung hat eine Verschlechterung der Ausrichtung aufgrund der Quantisierung von KV-Caches nachgewiesen. Darüber hinaus hat diese Forschung eine Verschlechterung der Ausrichtung auch in Produktions-vLLM-Dienstumgebungen mit FP8-Caches und einem trainingsfreien Wiederherstellungsprotokoll dokumentiert, das bis zu 97 Prozent dessen wiederherstellte, was an Ausrichtung verloren gegangen war. Wie dem auch sei, während eine Konfiguration ihre Benchmark-Ergebnisse behalten kann, bedeutet dies nicht unbedingt, dass sie alle anderen relevanten Parameter behält, die Ihnen wichtig sind.

Der allgemeine Grundsatz

Die Idee der Quantisierung wurde als “Präzisions-Budget” formuliert: Wie viele Bits kann ich mir leisten, zu opfern? Der KV-Cache zeigt, dass die nützlichere Frage die strukturelle ist. Präzision wird in Gruppen zugewiesen; die Gruppe ist die Einheit des Schadens, und die Achse, entlang derer Sie gruppieren, bestimmt, welche Elemente ihr Schicksal teilen. Die richtige Achse ist die, entlang derer Ihr Tensor konsumiert wird, d. h. die Art und Weise, wie Sie Ihren Tensor verwenden, und nicht, wie Ihr Tensor im Speicher gespeichert ist. Schlüssel werden durch ein Skalarprodukt gegen den Abfragevektor verwendet. Ein einzelner korrupter Kanal wird alle Scores vergiften. Werte werden durch ein sparsam gewichtetes Mittel über Token konsumiert. Daher wird ein einzelner korrupter Token einfach durch das Gewicht ausgeglichen.

Zwei Tensoren mit identischen Dimensionen und erzeugt durch zwei aufeinanderfolgende Schichten werden unterschiedlich behandelt. Es lohnt sich, für jede Aktivierung, die Sie komprimieren möchten, zu fragen: Welche Operation kontrahiert dies, und respektiert meine Gruppierung dies? Da ein einzelner korrupter Token einfach durch das Gewicht ausgeglichen wird. Zwei Tensoren mit identischen Dimensionen und erzeugt durch zwei aufeinanderfolgende Schichten werden unterschiedlich behandelt. Es lohnt sich, für jede Aktivierung, die Sie komprimieren möchten, zu fragen: Welche Operation kontrahiert dies, und respektiert meine Gruppierung dies? rage computation across tokens. Daher wird ein einzelner korrupter Token einfach durch das Gewicht ausgeglichen.