MHA vs. MQA vs. GQA: wie Attention den KV-Speicherbedarf verändert

Automatische Übersetzung

Dieser Artikel wurde automatisch aus der englischen Originalversion übersetzt.

MHA weist jedem Query-Head einen eigenen Key- und Value-Head zu. MQA nutzt einen gemeinsamen Key/Value-Head für alle Query-Heads. GQA nutzt einen gemeinsamen Key/Value-Head innerhalb jeder Gruppe von Query-Heads. Weniger KV-Heads verkleinern den konventionellen Attention-Cache und reduzieren die Datenmenge, die Decode daraus liest.

Prüfe beim Abschätzen der Cache-Größe num_key_value_heads. Allein anhand der Anzahl der Query-Heads kannst du die Cache-Größe deutlich überschätzen.

Die gemeinsame Nutzung der Heads direkt vergleichen

Nimm acht Query-Heads sowie identische Head-Dimensionen, Layer-Anzahlen, Sequenzlängen und Cache-Präzision an:

AttentionQuery-HeadsKV-HeadsCache relativ zu MHA
Multi-Head Attention, MHA88100%
Multi-Query Attention, MQA8112.5%
Grouped-Query Attention, GQA8225%

Ein KV-Head umfasst einen Key-Head und einen Value-Head. Die Query-Projektionen bleiben getrennt. Shazeers MQA-Paper und das GQA-Paper von Ainslie et al. definieren diese Änderungen.

Bei Llama 3 70B teilen sich 64 Query-Heads acht KV-Heads: Das ergibt achtmal weniger reine KV-Daten als bei einem ansonsten identischen Modell mit 64 KV-Heads. Llama 3.1 405B hat 128 Query-Heads und acht KV-Heads, was im selben Vergleich eine Reduktion um den Faktor sechzehn ergibt. Metas Architekturtabelle liefert die Zahlen.

Was die Reduktion nicht verspricht

Ein achtmal kleinerer KV cache bedeutet weder achtmal weniger GPU-Gesamtspeicher noch achtmal schnellere Generierung. Weights, Activations, Kernel-Ausführung und Kommunikation bleiben bestehen. Verteilte Implementierungen können KV-Heads außerdem über Partitionen hinweg replizieren.

Das GQA-Paper fand bei den getesteten Modellen der T5-Familie einen brauchbaren Kompromiss zwischen Qualität und Geschwindigkeit. Einen trainierten MHA-Checkpoint auf GQA umzustellen, ist eine Architekturänderung, die Anpassung und Evaluation erfordert; es ist keine Einstellung zur Speicherzuweisung.

Prüfe die Attention-Konfiguration des Checkpoints, verwende die Anzahl der KV-Heads in der Cache-Formel und miss den Server bei den benötigten Kontextlängen und der benötigten Anzahl gleichzeitiger Anfragen. Berücksichtige beim Vergleich von Checkpoints auch die Qualität bei deinen Aufgaben.

Engineering-Leitfaden: GQA und MQA zeigt die Muster der gemeinsamen Head-Nutzung.