KV-Cache in LLMs: Verständnis und Implementierung von Grund auf
Der KV-Cache ist eine Technik für effiziente LLM-Inferenz, die Zwischenvektoren für Schlüssel und Werte speichert, um sie wiederzuverwenden und redundante Berechnungen zu vermeiden. Der Artikel erklärt das Konzept und bietet eine Implementierung von Grund auf in PyTorch, die zeigt, wie man einen Multi-Head-Attention-Mechanismus modifiziert, um Schlüssel und Werte während der Textgenerierung zu cachen.
Der KV-Cache speichert Key- (K) und Value- (V) Vektoren aus früheren Attention-Berechnungen, um sie bei der Generierung nachfolgender Tokens wiederzuverwenden, was die Inferenz beschleunigt, indem Neuberechnungen vermieden werden. Ohne einen KV-Cache berechnet jeder Generierungsschritt die Keys und Values für alle vorherigen Tokens neu; mit dem Cache werden nur die Vektoren des neuen Tokens berechnet und an den Cache angehängt. Der Artikel stellt eine Code-Implementierung basierend auf einem GPT-ähnlichen Modell aus dem Buch des Autors vor. Die wichtigsten Änderungen sind: Hinzufügen von Cache-Puffern (cache_k und cache_v) in der MultiHeadAttention-Klasse, Ändern der Forward-Methode, um den Cache bedingt zu verwenden, Hinzufügen einer Reset-Methode und Weiterleiten des use_cache-Flags durch das Modell. Bei der Generierung verarbeitet das Modell, wenn use_cache True ist, nach dem initialen Prompt nur das neue Token, während es ohne Cache bei jedem Schritt die vollständige Sequenz verarbeitet. Ein einfacher Leistungsvergleich zeigt, dass der KV-Cache die Generierungsgeschwindigkeit für das getestete Beispiel etwa verdoppelt.
Quelle: Sebastian Raschka —
Original