Technische Zusammenfassung: RDKV – Rate-Distortion-Bit-Zuweisung für die gemeinsame Eviction und Quantisierung des KV-Caches

1. Problemstellung

Large Language Models (LLMs) stoßen während der Inferenz mit langen Kontexten auf einen kritischen Engpass: die Speicherbandbreite und -kapazität. Die Größe des Key-Value (KV)-Caches wächst linear mit der Sequenzlänge, was erfordert, dass bei jedem Decodierschritt Datenmengen von mehreren zehn Gigabyte aus dem externen High-Bandwidth Memory (HBM) in den On-Chip-Speicher neu eingelesen werden müssen. Dies macht die Inferenz speicherbandbreitenlimitiert statt rechenlimitiert.

Bestehende Kompressionsstrategien fallen im Allgemeinen in zwei isolierte Kategorien:

1. Eviction (Verdrängung): Bewertung von Tokens oder Kanälen und vollständiges Verwerfen der am schlechtesten bewerteten (binäres Behalten-oder-Wegwerfen). Dies lässt keine Präzisionsstufe für Tokens von moderater Wichtigkeit übrig.
2. Quantisierung: Beibehaltung aller Tokens, jedoch mit reduzierter Bit-Breite. Dies behält sogar vernachlässigbare Tokens bei und verschwendet Speicherplatz.

Während neuere Hybridansätze versuchen, diese zu kombinieren, verlassen sie sich oft auf heuristische Routing-Scores oder vorab festgelegte Verhältnisse, die durch Offline-Suche bestimmt wurden, und behandeln Eviction und Quantisierung nicht als ein einheitliches Optimierungsproblem.

2. Methodik: RDKV

Die Arbeit schlägt RDKV (Rate-Distortion KV-Cache-Kompression) vor, ein Framework, das die KV-Cache-Kompression als Rate-Distortion-Problem formuliert. In dieser Sichtweise sind Eviction und Quantisierung keine separaten Aktionen, sondern zwei Endpunkte einer einzigen Bit-Zuweisungskurve.

2.1 Theoretische Formulierung

RDKV leitet für jede Cache-Einheit (Token im V-Cache, Kanal im K-Cache) ein Verzerrungsgewicht her, basierend auf der Störung, die durch deren Entfernung oder Quantisierung bei der Attention-Berechnung verursacht wird.

• V-Cache (Tokens): Das Verzerrungsgewicht $w_t$ für ein Token ist der Total-Variations-Abstand der Attention-Verteilung nach der Eviction. Es wird als Summe der Attention-Scores berechnet, die das Token über alle Query-Positionen hinweg erhält: $w_t = \sum_{\tau} a_{\tau,t}$.
• K-Cache (Kanäle): Das Verzerrungsgewicht $w_c$ für einen Kanal ist die Spektralnorm der durch das Nullsetzen dieses Kanals verursachten Logit-Abweichung. Es wird als Produkt der Normen der entsprechenden Query- und Key-Spalten berechnet: $w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$.

Diese Gewichte dienen als Kostenkoeffizienten in einer Rate-Distortion-Zielfunktion:
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{unter der Bedingung} \quad \sum_{u} b_u \leq B$$
wobei $b_u$ die dem Einheit $u$ zugewiesene Bit-Breite ist und $B$ das gesamte Bit-Budget darstellt.

2.2 Bit-Zuweisungs-Pipeline

Die Zuweisung wird mittels Reverse Water-Filling (Lagrange-Relaxation) gelöst, angepasst für diskrete, hardwareunterstützte Bit-Breiten $\mathcal{B} = \{0, 2, 4, 8, 16\}$.

1. Stufe 1 (Gewichtung): Berechnung der Token-Gewichte ($w_t$) und Kanal-Gewichte ($w_c$) einmalig nach der Prefill-Phase unter Verwendung des vollen, unkomprimierten Caches.
2. Stufe 2 (Token-Zuweisung): Lösung eines Multiple-Choice-Knapsack-Problems (MCKP) für den V-Cache. Tokens erhalten Bit-Breiten aus $\{0, 2, 4, 8, 16\}$. Eine Bit-Breite von 0 entspricht einer Eviction.
3. Stufe 3 (Kanal-Zuweisung): Lösung des MCKP für den K-Cache, aber nur für die in Stufe 2 beibehaltenen Tokens. Kanäle erhalten Bit-Breiten unabhängig von der Token-V-Bit-Breite, was einem einzelnen Token gemischte Präzisions-Komponenten für K und V ermöglicht.

2.3 Effizienter gepackter Decode (TriZone)

Um tatsächliche Speicher- und Latenzeinsparungen zu realisieren, führt RDKV ein TriZone-Cache-Layout ein, das die Dekuantisierung in den Attention-Kernel integriert und so die Notwendigkeit vermeidet, Daten vor der Verarbeitung zu FP16 zu entpacken.

• Zone A: Gepackter alter Cache-Stream, der beibehaltene K-Zeilen und V-Zeilen mit Bit-Breiten $\{2, 4, 8\}$ enthält. Diese sind nach Bit-Breite gruppiert, um einheitliche Dekuantisierungspfade zu ermöglichen.
• Zone B: FP16-V-Zeilen für Tokens, die eine 16-Bit-Präzision zugewiesen erhielten.
• Zone C: FP16-K- und V-Zeilen für neu generierte Tokens.

Der Attention-Kernel berechnet Logits und Ausgaben durch Kombination dieser drei Zonen, wobei die Dekuantisierung direkt in die Matrixmultiplikation integriert ist und zusätzlichen Speicherbandbreiten-Overhead eliminiert.

3. Hauptbeiträge

1. Rate-Distortion-Formulierung: Die Arbeit stellt die KV-Cache-Kompression als einheitliches Rate-Distortion-Problem dar und beweist, dass die optimale Lösung inhärent Einheiten mit Null-Rate (Eviction) und Einheiten mit endlicher Rate (Quantisierung) mischt. Dies eliminiert die Notwendigkeit für heuristische Verhältnisse oder Offline-Suchen.
2. Diskrete Zuweisung via Reverse Water-Filling: Die Autoren instanziieren die Zuweisung als diskretes Knapsack-Problem, das durch Lagrange-Relaxation lösbar ist, und weisen optimale Bit-Breiten aus $\{0, 2, 4, 8, 16\}$ basierend auf Verzerrungsgewichten zu, die aus Attention-Dynamiken abgeleitet werden.
3. TriZone gepacktes Layout: Eine hardwarebewusste Implementierung, die die Dekuantisierung in den Attention-Kernel integriert und die theoretischen Bit-Breiten-Einsparungen in tatsächliche Reduktionen des Spitzenspeichers und der Decodier-Latenz umwandelt.

4. Experimentelle Ergebnisse

Experimente wurden an LLaMA-3.1-8B-Instruct, Mistral-7B, Qwen3-4B, LLaMA-2-13B und Qwen2.5-72B über LongBench, RULER und InfiniteBench durchgeführt.

• Genauigkeit: RDKV übertrifft die beste bewertete Baseline (Ada-SnapKV) im Durchschnitt um 9,1 %. Auf LongBench stellt es 97,81 % der Genauigkeit des Voll-Caches wieder her, wobei nur 2,48 % des Caches beibehalten werden.
• Leistung bei langen Kontexten:
  • RULER: Bei einem Kontext von 128K erreicht RDKV eine Genauigkeit von 66,95 % im Vergleich zu 58,43 % für die stärkste Baseline (AdaKV), eine Lücke von 8,5 Punkten.
  • InfiniteBench: RDKV belegt unter den Kompressionsmethoden den ersten Platz mit einer durchschnittlichen Punktzahl von 39,46 und übertrifft AdaKV um 1,40 Punkte.
  • Needle-in-a-Haystack: RDKV behält über alle Tiefen hinweg eine nahezu einheitliche Abrufgenauigkeit (durchschnittlich 0,99) bei, wohingegen Eviction-only-Baselines bei mittleren Tiefen signifikant versagen.
• Effizienz:
  • Latenz: Mit einem 128K-Kontext erreicht RDKV eine 4,5-fache Decodier-Geschwindigkeitssteigerung im Vergleich zur Decodierung mit Voll-Cache FlashAttention-2.
  • Speicher: Es erreicht eine 1,9-fache Reduktion des Spitzenspeichers bei 128K Kontext und führt erfolgreich 256K Kontext auf einer einzelnen A100 64GB GPU aus, während FullKV an Speicherplatzmangel (OOM) scheitert.

5. Bedeutung und Behauptungen

Die Arbeit behauptet, dass die gemeinsame Behandlung von Eviction und Quantisierung innerhalb eines einzigen Optimierungsrahmens überlegen ist gegenüber gestuften oder verhältnisbasierten Ansätzen. Durch die direkte Herleitung von Zuweisungsscores aus der bei der Attention-Berechnung verursachten Verzerrung vermeidet RDKV die „binäre Falle" von Eviction-only-Methoden (die moderate Tokens verwerfen) und die „Verschwendung" von Quantisierung-only-Methoden (die vernachlässigbare Tokens beibehalten).

Die Autoren betonen, dass ihre Methode:

• Kein erneutes Training des Modells erfordert.
• Architekturagnostisch ist und über verschiedene Modellgrößen und -familien hinweg konsistent funktioniert.
• Einen pareto-optimalen Kompromiss zwischen Latenz und Genauigkeit bietet und bei langen Kontextlängen Baselines in beiden Metriken strikt übertrifft.
• Zeigt, dass das „kontinuierliche Optimum" der Rate-Distortion-Theorie, wenn es für Hardware diskretisiert wird, natürlich eine Strategie mit gemischten Bit-Breiten liefert, die aktuelle State-of-the-Art-Hybriden übertrifft.

Die Arbeit schließt, dass dieser gemeinsame Ansatz die Genauigkeit verbessert, ohne die Effizienz zu opfern, und damit Inferenz mit langen Kontexten auf ressourcenbeschränkter Hardware ermöglicht.