FastKV: Decoupling of Context Reduction and KV Cache Compression for Prefill-Decoding Acceleration

1. Problemstellung

Large Language Models (LLMs) mit langen Kontextfenstern (z. B. 128K Token oder mehr) stoßen bei der Inferenz auf zwei wesentliche Engpässe:

• Prefill-Phase: Die Berechnung der Aufmerksamkeitsmechanismen skaliert quadratisch mit der Eingabelänge. Lange Prompts führen zu enormen Rechenkosten und Latenzen.
• Decoding-Phase: Der Key-Value (KV) Cache wächst linear mit der Kontextlänge. Da bei jedem generierten Token auf diesen Cache zugegriffen werden muss, führt ein großer Cache zu hohem Speicherverbrauch und Bandbreitenengpässen, was den Durchsatz drastisch reduziert.

Herausforderung bestehender Ansätze:
Bisherige Methoden versuchen oft, diese Probleme zu lösen, leiden jedoch unter einem fundamentalen Trade-off:

• Decoding-only-Methoden (z. B. StreamingLLM, SnapKV) komprimieren den KV-Cache nachträglich, beschleunigen aber nicht den Prefill-Prozess.
• Prefill-aware-Methoden (z. B. GemFilter, PyramidInfer) reduzieren die zu verarbeitenden Token bereits während des Prefills. Dies beschleunigt zwar den Prefill, führt aber oft zu einer Kopplung (Coupling) zwischen der Reduktion der Rechenlast im Prefill und dem Budget für den KV-Cache im Decoding. Wenn zu aggressiv im Prefill geschnitten wird, gehen für tiefere Schichten wichtige Informationen verloren, was die Genauigkeit (Accuracy) stark beeinträchtigt.

2. Methodik: FastKV

FastKV ist ein Framework, das diese Kopplung aufbricht, indem es zwei Schlüsselbeobachtungen nutzt:

1. Schichtabhängige Kontextdynamik: In frühen Schichten des Modells ist der Fokus der Aufmerksamkeit sehr instabil; kritische Token ändern sich stark von Schicht zu Schicht. In späteren Schichten stabilisiert sich jedoch die Menge der wichtigen Token.
2. Sparsame Kontextnutzung: Während des Decodings konzentriert sich jede Schicht nur auf einen kleinen, spärlichen Teil der Token, auch wenn der gesamte Kontext im Prefill verarbeitet wurde.

Basierend darauf führt FastKV zwei Haupttechniken ein:

A. Token-Selective Propagation (TSP)

Anstatt den gesamten Kontext durch alle Schichten zu propagieren, führt FastKV eine zweistufige Prefill-Strategie ein:

• Phase 1 (Frühe Schichten bis zur TSP-Schicht): Das Modell verarbeitet den vollständigen Kontext. Dies stellt sicher, dass frühe Schichten, die unterschiedliche Token-Muster erkennen, Zugang zu allen Informationen haben.
• Phase 2 (Ab der TSP-Schicht): Eine dedizierte „Token-Selective Propagation"-Schicht (TSP-Layer) wählt basierend auf einer Salienz-Bewertung (Saliency Score) nur die informativsten Token aus. Nur diese Token (plus ein Fenster der neuesten Token) werden an die nachfolgenden Schichten weitergegeben.
• Wahl der TSP-Schicht: Die optimale Position der TSP-Schicht wird algorithmisch bestimmt, indem die Schicht gefunden wird, bei der die Ausgabe noch nahe am Vollkontext-Baseline liegt, aber eine signifikante Reduktion der nachfolgenden Berechnungen ermöglicht.

B. Entkoppelte KV-Cache-Reduktion (Decoupling)

Ein entscheidender Innovationsschritt ist die Trennung von zwei Parametern:

1. TSP-Rate: Bestimmt, wie viele Token nach der TSP-Schicht weiterverarbeitet werden (reduziert die Prefill-Rechenlast).
2. KV-Retention-Rate: Bestimmt, wie viele KV-Einträge pro Schicht für das Decoding gespeichert werden.

Im Gegensatz zu früheren Methoden, bei denen eine aggressive Prefill-Reduktion automatisch zu einem kleinen KV-Cache führt, erlaubt FastKV eine unabhängige Steuerung. Man kann z. B. einen konservativen TSP (hohe Genauigkeit im Prefill) mit einer aggressiven KV-Reduktion (hohe Geschwindigkeit im Decoding) kombinieren.

3. Schlüsselbeiträge

• Entkopplung von Prefill und Decoding: FastKV löst das Problem, dass die Reduktion der Rechenlast im Prefill nicht mehr zwingend die Genauigkeit im Decoding beeinträchtigt.
• Zweistufiger Prefill: Durch die Beibehaltung des Vollkontexts in frühen Schichten werden Informationsverluste vermieden, die bei früheren Ansätzen (wie GemFilter) zu Genauigkeitsverlusten führten.
• Flexibles Tuning: Die unabhängige Kontrolle von TSP-Rate und KV-Retention-Rate ermöglicht eine feine Abstimmung zwischen Latenz und Genauigkeit.

4. Ergebnisse

Die Evaluation erfolgte auf den Modellen LLaMA-3.1-8B-Instruct und Ministral-8B-Instruct mit Kontextlängen bis zu 128K Token.

• Genauigkeit (Accuracy):

  • FastKV erreicht auf dem LongBench-Benchmark eine Genauigkeit, die mit dem Vollkontext-Baseline vergleichbar ist (Verlust < 1%).
  • Im Vergleich dazu zeigen Methoden wie GemFilter oder PyramidInfer signifikante Genauigkeitsverluste (bis zu 11,58% bei GemFilter), da sie zu früh Token verwerfen.
  • Auf dem RULER- und Needle-in-a-Haystack-Benchmarks übertrifft FastKV alle anderen komprimierten Methoden und bleibt dem Vollkontext nahe.

• Geschwindigkeit (Latency):

  • Prefill: Bis zu 1,82-fache Beschleunigung im Vergleich zum Vollkontext.
  • Decoding: Bis zu 2,87-fache Beschleunigung im Vergleich zum Vollkontext.
  • FastKV ist die einzige Methode, die sowohl Prefill als auch Decoding signifikant beschleunigt, ohne dabei die Genauigkeit zu opfern.

5. Bedeutung und Fazit

FastKV stellt einen Paradigmenwechsel in der Optimierung von LLM-Inferenz für lange Kontexte dar. Indem es die Dynamik der Token-Wichtigkeit über die Schichten hinweg berücksichtigt und die Prefill-Berechnung von der KV-Cache-Speicherung entkoppelt, ermöglicht es eine skalierbare Inferenz.

• Praktische Relevanz: Die Methode macht lange Kontexte (128K+) in der Praxis effizient nutzbar, indem sie sowohl die Rechenkosten als auch den Speicherverbrauch senkt.
• Kompatibilität: FastKV ist orthogonal zu anderen Techniken wie Quantisierung oder Sparse-Attention-Kernen und kann mit diesen kombiniert werden, um die Effizienz weiter zu steigern.

Zusammenfassend beweist FastKV, dass eine intelligente, schichtabhängige Reduktion des Kontexts in Kombination mit einer flexiblen Cache-Verwaltung der Schlüssel ist, um den Trade-off zwischen Geschwindigkeit und Genauigkeit bei langen Sequenzen zu überwinden.