Technische Zusammenfassung: ProxyKV

Problemstellung

Die effiziente Inferenz mit langen Kontexten in Large Language Models (LLMs) wird derzeit durch den Speicherbedarf des Key-Value (KV) Caches limitiert. Mit zunehmender Sequenzlänge ($N$) übersteigt das lineare Wachstum des KV-Caches häufig die GPU-Speicherkapazität, was zu prohibitivem I/O-Overhead und Inferenz-Latenz führt.

Bestehende Methoden zur KV-Cache-Pruning stehen vor einem grundlegenden Zielkonflikt zwischen Scoring-Effizienz und Pruning-Präzision:

• Heuristische Methoden (z. B. H2O, SnapKV): Diese sind schnell und haben eine geringe Latenz, verlassen sich jedoch auf lokale Beobachtungsfenster und übersehen oft globale semantische Abhängigkeiten, die für komplexes Schlussfolgern erforderlich sind.
• Rekonstruktionsbasierte Methoden (z. B. KVZip): Diese erreichen eine hohe Präzision, indem sie einen sekundären, vom Lehrer erzwungenen Prefill-Durchlauf am Zielmodell durchführen, um den globalen Kontext zu rekonstruieren. Dies verursacht jedoch einen prohibitiven Prefill-Overhead, der die Latenzvorteile des Prunings zunichtemacht.

Die Kernherausforderung besteht darin, die Lücke zwischen latenzarmen Heuristiken und hochpräziser Rekonstruktion zu schließen, ohne die hohen Rechenkosten eines sekundären Prefills am Zielmodell in Kauf nehmen zu müssen.

Methodik

Die Autoren schlagen ProxyKV vor, ein plattformübergreifendes Proxy-Pruning-Framework, das die Aufgabe des Importance-Scoring vom Zielmodell auf einen leichten, asynchronen „Small-Model Proxy" auslagert.

Systemarchitektur

Das System arbeitet über eine Pipeline aus drei Komponenten:

1. Large-Model Target ($M_l$): Das primäre Modell, das für den kritischen Pfad des Prefills und die gesamte Decode-Phase verantwortlich ist.
2. Small-Model Proxy ($M_s$): Ein kleineres Modell derselben Familie (z. B. mit 1,5 Mrd. oder 4 Mrd. Parametern), das asynchron auf einem separaten Ausführungsstrom (oder CUDA-Stream) läuft. Es extrahiert cross-head-Aufmerksamkeitsmerkmale ($X$) aus dem Eingabekontext.
3. HybridAxialMapper ($\Phi_\theta$): Ein lernbares Modul, das die Merkmale des Proxies ($X$) auf zielgerichtete Importance-Scores ($\hat{Y}$) abbildet.

Arbeitsablauf:
Während der Inferenz wird der Eingabekontext parallel an beide Modelle gesendet. Der Proxy generiert Aufmerksamkeitsmerkmale, die vom HybridAxialMapper verarbeitet werden, um durch Top-K-Schwellenwertbildung eine binäre Pruning-Maske ($M$) zu erzeugen. Das Zielmodell führt dann die Dekodierung unter Verwendung nur des geprunten KV-Caches durch. Entscheidend ist, dass der KV-Cache des Proxies unmittelbar nach der Prefill-Phase freigegeben wird, wodurch der Speicheraufwand transient bleibt.

Schlüsselkomponenten

1. HybridAxialMapper-Architektur

Um die architektonische Lücke zwischen heterogenen Modellen (unterschiedliche Anzahl an Heads und Schichttiefen) zu überbrücken, setzt der Mapper eine dreistufige Pipeline ein, die die zeitliche Merkmalsextraktion von der cross-head-Ausrichtung entkoppelt:

• Stufe 1: Zeitliche Merkmalsextraktion: Ein 1D-Convolutional Stem projiziert rohe Proxy-Merkmale in einen modellagnostischen latenten Raum und normalisiert die Head-Konfiguration.
• Stufe 2: Time-Axis Context Encoder: Ein Transformer-Encoder verarbeitet die Sequenz mit sinusförmigen Positionsencodings, um den globalen Kontext zu erfassen, und nutzt einen gleitenden Fensteransatz für lange Sequenzen (bis zu 170.000 Token).
• Stufe 3: Head-Axis Cross-Attention: Lernbare Ziel-Queries aggregieren synthetische Head-Token vom Proxy in die spezifische Head-Konfiguration des Ziels. Diese Stufe ermöglicht es dem Mapper, auf beliebige Ziel-Head-Anzahlen zu skalieren, ohne das gesamte Netzwerk neu trainieren zu müssen.

2. Multi-Granularitäts-Hybrid-Loss

In Anbetracht dessen, dass die Einsatzmetrik die Qualität der Top-K-Behaltung und nicht die punktuelle Score-Rekonstruktion ist, entwerfen die Autoren eine zusammengesetzte Verlustfunktion ($L_{total}$), die aus fünf Termen besteht:

• $L_{bin}$ (Multi-Ratio Binary): Überwacht die binäre Entscheidungsgrenze über mehrere Behaltungsverhältnisse hinweg und betont eine aggressive Kompression.
• $L_{mse}$ (Value-Weighted MSE): Stellt eine numerische Größenübereinstimmung für saliente Token sicher.
• $L_{fine}$ (Intra-Top-K Rank): Erzwingt eine relative Ordnung innerhalb der behaltenen Top-K-Menge.
• $L_{global}$ (Mass-Aware Global Rank): Optimiert die Trennung zwischen Top-K- und Nicht-Top-K-Token (Signal-Rausch-Grenze).
• $L_{cos}$ (Cosine Alignment): Richtet die globale Orientierung des vorhergesagten Score-Vektors auf das Ziel aus.

Dieses Loss-Design verschiebt das Ziel von einer starren Regression hin zur Konsistenz der relativen Rangfolge, was direkter mit der Pruning-Qualität übereinstimmt.

Hauptbeiträge

1. Asynchrones Proxy-Framework: Ein neuartiges Pruning-Framework, das den Scoring-Engpass aus dem kritischen Pfad des Ziels entfernt, indem es ihn auf ein asynchrones kleines Modell auslagert. Dies erzielt eine 3,21-fache Beschleunigung beim Prefilling auf Llama-3.1-8B (Dual-GPU) und hält Beschleunigungen bis zu 170.000 Token aufrecht.
2. HybridAxialMapper & Hybrid-Loss: Eine lernbare Architektur und Verlustfunktion, die strukturelle Lücken zwischen heterogenen Modellen erfolgreich überbrücken (z. B. 8B-Ziel vs. 1,5B-Proxy) und die Starrheit statischer Head-zu-Head-Ausrichtungen, wie sie in früheren Arbeiten verwendet wurden, überwinden.
3. Empirische Validierung: Umfassende Experimente mit den Familien Llama-3.1, Qwen-2.5 und Qwen-3 (7B–32B-Ziele) zeigen, dass ProxyKV ~98,7 % der mittleren Genauigkeit des KVZip-Orakels wiederherstellt und gleichzeitig die Latenz erheblich reduziert.

Ergebnisse

• Genauigkeit: ProxyKV entspricht der Leistung des KVZip-Orakels (der hochpräzisen Baseline) über die Benchmarks LongBench, SCBench und RULER hinweg. Es stellt ~98,7 % der mittleren Genauigkeit von KVZip wieder her und liegt bei den meisten Aufgaben innerhalb von 1–2 Punkten, selbst bei aggressiven Behaltungsverhältnissen ($\rho \le 0,2$).
• Effizienz:
  • Bei Llama-3.1-8B liefert ProxyKV eine 3,21-fache Beschleunigung der Prefill-Latenz im Vergleich zu KVZip (Dual-GPU) und ~1,5-fach auf einer einzigen geteilten GPU.
  • Der HybridAxialMapper selbst trägt nur 4,7 %–6,8 % zur Prefill-Wandzeit bei.
  • Die Methode skaliert effektiv auf 32B-Ziele (Qwen-3) mit einem Verhältnis von Ziel- zu Proxy-Größe von ca. 8:1.
• Robustheit: Das Framework zeigt eine robuste Zero-Shot-Verallgemeinerung auf zurückgehaltene Datensätze (z. B. SCBench RepoQA) und behält die Leistung bei komplexen Multi-Hop-Schlussfolgerungsaufgaben bei, bei denen heuristische Methoden (SnapKV) versagen.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass ProxyKV erfolgreich die „Scoring-Kosten–Genauigkeits-Lücke" überbrückt, die die Inferenz von LLMs mit langen Kontexten eingeschränkt hat. Durch die Nutzung von Intra-Familien-Aufmerksamkeitskorrelationen und einer lernbaren cross-axialen Abbildung erreicht es die hohe Präzision rekonstruktionsbasierter Methoden ohne den prohibitiven Prefill-Overhead.

Die Autoren heben hervor, dass statische Schicht-zu-Schicht-Ausrichtungen (wie sie bei IAM oder KVZip+IAM verwendet werden) zu einer bindenden Einschränkung werden, wenn Zielmodelle skalieren (z. B. auf 32B), während der lernbare HybridAxialMapper effektiv bleibt. Dies deutet darauf hin, dass ProxyKV einen skalierbaren Weg für eine effiziente Inferenz mit langen Kontexten bietet, der eine nahezu orakelgenaue Genauigkeit wiederherstellt und gleichzeitig die Latenz und den Speicherdruck während der Prefill-Phase erheblich reduziert. Die Methode wird als Deploy-Time-Lösung präsentiert, die nur einen leichten, Intra-Familien-Proxy und einen kleinen, trainierbaren Mapper erfordert, ohne dass eine vollständige Neukalibrierung des Modells oder komplexe Hardwareänderungen notwendig sind.