DuoAttention: Effiziente Inferenz von Long-Context LLM mit Abruf- und Streaming-Köpfen

DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads

October 14, 2024

Autoren: Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, Song Han

cs.AI

Zusammenfassung

Die Bereitstellung von langen Kontexten großer Sprachmodelle (LLMs) ist unerlässlich, birgt jedoch erhebliche Rechen- und Speicherherausforderungen. Das Zwischenspeichern aller Schlüssel- und Wertezustände (KV) über alle Aufmerksamkeitsköpfe hinweg verbraucht erheblichen Speicherplatz. Bestehende KV-Zwischenspeicherverfahren beschädigen entweder die Fähigkeiten der LLMs im langen Kontext oder bieten nur begrenzte Effizienzverbesserungen. In diesem Papier identifizieren wir, dass nur ein Bruchteil der Aufmerksamkeitsköpfe, auch als Abrufköpfe bezeichnet, für die Verarbeitung langer Kontexte entscheidend ist und eine vollständige Aufmerksamkeit über alle Token hinweg erfordert. Im Gegensatz dazu benötigen alle anderen Köpfe, die sich hauptsächlich auf aktuelle Token und Aufmerksamkeitssenken konzentrieren - als Streaming-Köpfe bezeichnet - keine vollständige Aufmerksamkeit. Basierend auf dieser Erkenntnis stellen wir DuoAttention vor, ein Framework, das nur einen vollständigen KV-Zwischenspeicher für Abrufköpfe verwendet, während für Streaming-Köpfe ein leichtgewichtiger, konstanter KV-Zwischenspeicher verwendet wird. Dies reduziert sowohl den Dekodier- als auch den Vorabfüllspeicher und die Latenz der LLMs, ohne die Fähigkeit im langen Kontext zu beeinträchtigen. DuoAttention verwendet einen leichten, optimierungsbasierten Algorithmus mit synthetischen Daten, um Abrufköpfe genau zu identifizieren. Unsere Methode reduziert den Speicherbedarf für Inferenzen im langen Kontext um bis zu 2,55-fach für MHA- und 1,67-fach für GQA-Modelle, beschleunigt das Dekodieren um bis zu 2,18-fach und 1,50-fach und beschleunigt das Vorabfüllen um bis zu 1,73-fach und 1,63-fach für MHA- und GQA-Modelle, jeweils mit minimalen Genauigkeitsverlusten im Vergleich zur vollständigen Aufmerksamkeit. Besonders in Kombination mit Quantisierung ermöglicht DuoAttention das Dekodieren von Llama-3-8B mit einer Kontextlänge von 3,3 Millionen auf einer einzelnen A100-GPU. Der Code ist verfügbar unter https://github.com/mit-han-lab/duo-attention.

English

Deploying long-context large language models (LLMs) is essential but poses significant computational and memory challenges. Caching all Key and Value (KV) states across all attention heads consumes substantial memory. Existing KV cache pruning methods either damage the long-context capabilities of LLMs or offer only limited efficiency improvements. In this paper, we identify that only a fraction of attention heads, a.k.a, Retrieval Heads, are critical for processing long contexts and require full attention across all tokens. In contrast, all other heads, which primarily focus on recent tokens and attention sinks--referred to as Streaming Heads--do not require full attention. Based on this insight, we introduce DuoAttention, a framework that only applies a full KV cache to retrieval heads while using a light-weight, constant-length KV cache for streaming heads, which reduces both LLM's decoding and pre-filling memory and latency without compromising its long-context abilities. DuoAttention uses a lightweight, optimization-based algorithm with synthetic data to identify retrieval heads accurately. Our method significantly reduces long-context inference memory by up to 2.55x for MHA and 1.67x for GQA models while speeding up decoding by up to 2.18x and 1.50x and accelerating pre-filling by up to 1.73x and 1.63x for MHA and GQA models, respectively, with minimal accuracy loss compared to full attention. Notably, combined with quantization, DuoAttention enables Llama-3-8B decoding with 3.3 million context length on a single A100 GPU. Code is provided in https://github.com/mit-han-lab/duo-attention.