Grouped Query Attention (GQA)

Grouped Query Attention (GQA) ist ein Attention-Design, bei dem sich mehrere Query-Heads einen kleineren Satz von Key- und Value-Heads teilen. Es bewahrt die verschiedenen Perspektiven von Multi-Head-Queries und reduziert gleichzeitig den für Keys und Values erforderlichen Speicher- und Datenversand. GQA ist besonders wertvoll bei der autoregressiven Inferenz, bei der ein Modell ein Token nach dem anderen generiert und wiederholt gespeicherte Attention-Zustände liest.

Wie GQA funktioniert#

Innerhalb eines attention mechanism wird jedes Token in drei Darstellungen projiziert:

• Query: Beschreibt die Informationen, die das aktuelle Token sucht.
• Key: Beschreibt, was jedes verfügbare Token darstellt.
• Value: Enthält die Informationen, die abgerufen werden, wenn eine Query zu einem Key passt.

In der konventionellen self-attention werden diese Projektionen in Heads unterteilt, damit das Modell verschiedene Beziehungen parallel lernen kann. GQA behält viele Query-Heads bei, weist jedoch Gruppen davon gemeinsamen Key- und Value-Heads zu. Beispielsweise kann eine Attention-Schicht acht Query-Heads und zwei Key-Value-Heads verwenden. Jeder Key-Value-Head bedient dann vier Query-Heads.

Die eigentliche Attention-Berechnung bleibt eine skalierte Dot-Product-Attention. Was sich ändert, ist die Head-Anordnung und die Menge der erzeugten, gespeicherten und gelesenen Key-Value-Daten. Frameworks wie die PyTorch scaled dot-product attention erfordern daher, dass die Anzahl der Query-Heads durch die Anzahl der Key-Value-Heads teilbar ist, wenn GQA aktiviert ist. (docs.pytorch.org)

GQA vs. Multi-Head- und Multi-Query-Attention#

GQA positioniert sich zwischen Multi-Head-Attention und Multi-Query-Attention:

• Multi-head attention: Jeder Query-Head hat seine eigenen Key- und Value-Heads. Dies bietet maximale Head-Unabhängigkeit, erzeugt jedoch den größten Key-Value-Speicherbedarf.
• Grouped Query Attention: Mehrere Query-Heads teilen sich jeden Key-Value-Head. Dies bringt die Repräsentationskapazität und die Speichereffizienz in ein ausgewogenes Verhältnis.
• Multi-Query-Attention: Alle Query-Heads teilen sich einen Key-Head und einen Value-Head. Dies minimiert die Key-Value-Speicherung, bietet jedoch eine geringere Key-Value-Vielfalt.

Wenn eine Schicht 32 Query-Heads hat, verwendet die Multi-Head-Attention möglicherweise auch 32 Key-Value-Heads, GQA eventuell 8 und die Multi-Query-Attention 1. GQA ist kein Ersatz für die umfassendere Transformer architecture; es ist eine mögliche Konfiguration innerhalb einer Transformer-Attention-Schicht. Die multi-head, multi-query, and grouped-query attention documentation von NVIDIA beschreibt diese Varianten primär darin, wie viele Key-Value-Heads die Query-Heads bedienen. (nvidia.github.io)

Warum GQA wichtig ist#

Bei der autoregressiven Generierung speichert ein Modell vorherige Keys und Values in einem KV cache. Bei der Standard-Multi-Head-Attention trägt jeder Attention-Head separate gecachte Keys und Values bei. Lange Sequenzen, große Batches und viele Modellschichten können daher beträchtlichen GPU-Speicher verbrauchen.

Da GQA weniger Key-Value-Heads verwendet, ist sein Cache proportional zur Reduzierung dieser Heads kleiner. Dies kann Folgendes verbessern:

• Speicherkapazität: Längere Prompts oder mehr gleichzeitige Anfragen passen in den verfügbaren Speicher.
• Generierungsdurchsatz: Für jedes generierte Token müssen weniger Key-Value-Daten gelesen werden.
• Inference latency: Ein reduzierter Speicherverkehr kann die Verarbeitungszeit von Token zu Token verkürzen.
• Context window-Skalierbarkeit: Das Verarbeiten langer Sequenzen wird praktikabler.

Diese Gewinne hängen von Hardware, Sequenzlänge, Batch-Größe und Kernel-Unterstützung ab. Produktions-Runtimes benötigen nach wie vor 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)

Praxisanwendungen#

Assistenten für langen Kontext: Ein Programmierassistent verarbeitet möglicherweise Tausende von Quellcode-Token, bevor er eine Antwort generiert. GQA reduziert den gecachten Key-Value-Zustand, der mit diesem Verlauf verbunden ist, sodass ein Server längere Dateien oder mehr gleichzeitige Benutzer verarbeiten kann, ohne dass der GPU-Speicher proportional wächst.

Multimodale Bild- und Video-Assistenten: Ein vision-language model kann Bild-Patches oder Videoframes als lange Token-Sequenzen darstellen. GQA reduziert den Attention-Cache-Druck, nachdem diese visuellen Token in den Sprachdecoder gelangt sind, wodurch möglicherweise mehr Speicher für zusätzliche Bilder, Frames oder Anfragen frei bleibt. Dies unterscheidet sich von einem Standard-Vision Transformer, bei dem die Attention alle Bild-Patches parallel verarbeiten kann, ohne autoreigenes KV-Caching.

Praktische Implementierung und Kompromisse#

PyTorch stellt GQA über enable_gqa=True bereit. Dieses dokumentierte CUDA-Beispiel verwendet acht Query-Heads und zwei geteilte Key-Value-Heads:

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 Query-Heads bei, während Keys und Values nur zwei Heads verwenden. Der PyTorch attention backend selector kann steuern, welcher unterstützte Kernel die Operation ausführt, während der PyTorch Transformer building-block guide einen breiteren Implementierungskontext bietet. (docs.pytorch.org)

GQA muss beim Entwerfen oder Anpassen eines Modells ausgewählt werden; es ist im Allgemeinen kein Inferenz-Schalter für einen bestehenden Checkpoint. Entwickler sollten die Framework-Unterstützung, die Head-Teilbarkeit, das numerische Verhalten, die Speichernutzung und die Aufgabenqualität überprüfen. Für die Bereitstellung in der Computer Vision helfen ergänzende Workflows wie die Ultralytics TensorRT integration und der Ultralytics benchmark mode dabei, zu messen, ob architektonische und Laufzeitoptimierungen signifikante Verbesserungen auf der Zielhardware bringen.