Grouped Query Attention (GQA)

Gruppierte Abfrageaufmerksamkeit (GQA) ist ein Aufmerksamkeitsentwurf, bei dem mehrere Abfrageköpfe eine kleinere Anzahl von Schlüssel- und Werteköpfen gemeinsam nutzen. Dabei bleiben die vielfältigen Perspektiven von Abfragen mit mehreren Köpfen erhalten, während der für Schlüssel und Werte erforderliche Speicherbedarf und Datentransfer sinken. GQA ist besonders bei der autoregressiven Inferenz nützlich, bei der ein Modell jeweils ein Token generiert und gespeicherte Aufmerksamkeitszustände wiederholt ausliest.

So funktioniert GQA#

Innerhalb eines Aufmerksamkeitsmechanismus wird jedes Token in drei Repräsentationen projiziert:

• Abfrage: Beschreibt die Information, nach der das aktuelle Token sucht.
• Schlüssel: Beschreibt, was jedes verfügbare Token repräsentiert.
• Wert: Enthält die abgerufene Information, wenn eine Abfrage zu einem Schlüssel passt.

Bei herkömmlicher Selbstaufmerksamkeit werden diese Projektionen in Köpfe aufgeteilt, damit das Modell verschiedene Beziehungen parallel lernen kann. GQA behält viele Abfrageköpfe bei, ordnet ihnen jedoch gemeinsam genutzte Schlüssel- und Werteköpfe zu. Eine Aufmerksamkeits-Schicht kann beispielsweise acht Abfrageköpfe und zwei Schlüssel-Werte-Köpfe verwenden. Jeder Schlüssel-Werte-Kopf bedient dann vier Abfrageköpfe.

Die Aufmerksamkeitsberechnung selbst bleibt eine skalierte Punktprodukt-Aufmerksamkeit. Geändert werden die Anordnung der Köpfe sowie die Menge der erzeugten, gespeicherten und ausgelesenen Schlüssel-Werte-Daten. Frameworks wie die skalierte Punktprodukt-Aufmerksamkeit von PyTorch erfordern daher, dass die Anzahl der Abfrageköpfe durch die Anzahl der Schlüssel-Werte-Köpfe teilbar ist, wenn GQA aktiviert ist. (docs.pytorch.org)

GQA im Vergleich zu Aufmerksamkeit mit mehreren Köpfen und mit mehreren Abfragen#

GQA liegt zwischen Aufmerksamkeit mit mehreren Köpfen und Aufmerksamkeit mit mehreren Abfragen:

• Aufmerksamkeit mit mehreren Köpfen: Jeder Abfragekopf hat eigene Schlüssel- und Werteköpfe. Das ermöglicht eine maximale Unabhängigkeit der Köpfe, führt jedoch zum größten Bedarf an Schlüssel-Werte-Speicher.
• Gruppierte Abfrageaufmerksamkeit: Mehrere Abfrageköpfe nutzen jeweils einen Schlüssel-Werte-Kopf gemeinsam. Dadurch werden Repräsentationskapazität und Speichereffizienz ausbalanciert.
• Aufmerksamkeit mit mehreren Abfragen: Alle Abfrageköpfe nutzen einen Schlüsselkopf und einen Wertekopf gemeinsam. Dadurch wird der Speicherbedarf für Schlüssel und Werte minimiert, allerdings steht eine geringere Vielfalt an Schlüsseln und Werten zur Verfügung.

Wenn eine Schicht 32 Abfrageköpfe hat, kann die Aufmerksamkeit mit mehreren Köpfen ebenfalls 32 Schlüssel-Werte-Köpfe verwenden, GQA möglicherweise 8 und die Aufmerksamkeit mit mehreren Abfragen 1. GQA ist kein Ersatz für die umfassendere Transformer-Architektur, sondern eine mögliche Konfiguration innerhalb einer Transformer-Aufmerksamkeitsschicht. Die Dokumentation von NVIDIA zu Aufmerksamkeit mit mehreren Köpfen, mehreren Abfragen und gruppierten Abfragen beschreibt diese Varianten als hauptsächlich dadurch verschieden, wie viele Schlüssel-Werte-Köpfe die Abfrageköpfe bedienen. (nvidia.github.io)

Warum GQA wichtig ist#

Während der autoregressiven Generierung speichert ein Modell vorherige Schlüssel und Werte in einem KV-Cache. Bei standardmäßiger Aufmerksamkeit mit mehreren Köpfen trägt jeder Aufmerksamkeitskopf separate zwischengespeicherte Schlüssel und Werte bei. Lange Sequenzen, große Stapel und zahlreiche Modellschichten können daher beträchtlichen GPU-Speicher belegen.

Da GQA weniger Schlüssel-Werte-Köpfe verwendet, ist sein Cache proportional zur Verringerung dieser Kopfanzahl kleiner. Das kann Folgendes verbessern:

• Speicherkapazität: Längere Eingabeaufforderungen oder mehr gleichzeitige Anfragen können in den verfügbaren Speicher passen.
• Generierungsdurchsatz: Für jedes generierte Token müssen weniger Schlüssel-Werte-Daten gelesen werden.
• Inferenzlatenz: Weniger Speicherverkehr kann die Verarbeitungszeit von Token zu Token verkürzen.
• Skalierbarkeit des Kontextfensters: Die Verarbeitung langer Sequenzen wird praktikabler.

Diese Vorteile hängen von der Hardware, der Sequenzlänge, der Stapelgröße und der Kernel-Unterstützung ab. Produktionslaufzeitumgebungen benötigen weiterhin eine effiziente Cache-Zuweisung. Das TensorRT-LLM-KV-Cache-System von NVIDIA kombiniert beispielsweise GQA-Unterstützung mit blockbasiertem Caching, Wiederverwendung und Auslagerung. (nvidia.github.io)

Anwendungen in der Praxis#

Assistenten mit langem Kontext: Ein Programmierassistent kann Tausende Quellcode-Tokens verarbeiten, bevor er eine Antwort generiert. GQA reduziert den mit diesem Verlauf verbundenen zwischengespeicherten Schlüssel-Werte-Zustand, sodass ein Server längere Dateien oder mehr gleichzeitige Benutzer verarbeiten kann, ohne den GPU-Speicher proportional zu erhöhen.

Multimodale Bild- und Videoassistenten: Ein Bild-Sprach-Modell kann Bildbereiche oder Videoframes als lange Token-Sequenzen darstellen. GQA verringert den Druck auf den Aufmerksamkeits-Cache, nachdem diese visuellen Tokens in den Sprach-Decoder gelangen, und schafft so möglicherweise mehr Speicher für zusätzliche Bilder, Frames oder Anfragen. Das unterscheidet sich von einem standardmäßigen visuellen Transformer-Modell, bei dem die Aufmerksamkeit möglicherweise alle Bildbereiche parallel verarbeitet, ohne einen autoregressiven KV-Cache zu verwenden.

Praktische Implementierung und Abwägungen#

PyTorch stellt GQA über enable_gqa=True bereit. Dieses dokumentierte CUDA-Beispiel verwendet acht Abfrageköpfe und zwei gemeinsam genutzte Schlüssel-Werte-Köpfe:

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

assert torch.cuda.is_available(), "A CUDA device is required."

query = torch.randn(1, 8, 16, 32, device="cuda")
key = torch.randn(1, 2, 16, 32, device="cuda")
value = torch.randn(1, 2, 16, 32, device="cuda")

# Four query heads share each key-value head.
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(
        query,
        key,
        value,
        enable_gqa=True,
    )

print(output.shape)

Die Ausgabe behält acht Abfrageköpfe bei, während Schlüssel und Werte nur zwei Köpfe verwenden. Der Auswahlmechanismus für Aufmerksamkeits-Backends in PyTorch kann steuern, welcher unterstützte Kernel den Vorgang ausführt, während der Leitfaden von PyTorch zu Transformer-Bausteinen einen umfassenderen Implementierungskontext bietet. (docs.pytorch.org)

GQA muss beim Entwurf oder bei der Anpassung eines Modells ausgewählt werden; für einen vorhandenen Checkpoint ist es im Allgemeinen kein Schalter, der zur Inferenzzeit umgelegt werden kann. Entwickler sollten die Framework-Unterstützung, die Teilbarkeit der Kopfanzahl, das numerische Verhalten, den Speicherbedarf und die Aufgabenqualität überprüfen. Für die Bereitstellung in der Computer Vision helfen ergänzende Arbeitsabläufe wie die TensorRT-Integration von Ultralytics und der Benchmark-Modus von Ultralytics dabei zu messen, ob architektonische und Laufzeitoptimierungen auf der Zielhardware tatsächlich relevante Verbesserungen bewirken.