SparDA: Sparse entkoppelte Aufmerksamkeit für effiziente LLM-Inferenz mit langem Kontext

Zusammenfassung

Sparse Attention reduziert Rechenaufwand und Speicherbandbreite bei der Inferenz von LLMs mit langen Kontexten. Es bleiben jedoch zwei zentrale Herausforderungen bestehen: (1) Die KV-Cache-Kapazität wächst weiterhin mit der Sequenzlänge, und das Auslagern auf die CPU führt zu einem PCIe-Transfer-Engpass; (2) der Sparse-Selektionsschritt selbst behält eine O(T²)-Komplexität und kann bei langen Kontexten den dominierenden Kostenfaktor der Attention darstellen. Wir schlagen SparDA vor, eine entkoppelte Sparse-Attention-Architektur, die eine vierte schichtweise Projektion, den Forecast, neben Query, Key und Value einführt. Der Forecast prognostiziert die KV-Blöcke, die von der nächsten Schicht benötigt werden, und ermöglicht so eine vorausschauende Selektion, die das CPU-zu-GPU-Prefetch mit der Ausführung der aktuellen Schicht überlappt. Da der Forecast von der Attention-Query entkoppelt ist, verwendet unsere GQA-Implementierung einen Forecast-Kopf pro GQA-Gruppe, was den Selektionsaufwand im Vergleich zum ursprünglichen Multi-Head-Selektor reduziert. SparDA fügt weniger als 0,5 % Parameter hinzu und trainiert nur die Forecast-Projektionen, indem es die Aufmerksamkeitsverteilung des ursprünglichen Selektors nachahmt. Auf zwei sparse-vortrainierten 8B-Modellen erreicht SparDA eine vergleichbare oder leicht verbesserte Genauigkeit und erzielt bis zu 1,25-fachen Prefill-Speedup sowie 1,7-fachen Decode-Speedup gegenüber der Sparse-Attention-Offload-Baseline. Durch die Ermöglichung größerer nutzbarer Batch-Größen auf einer einzelnen GPU erreicht SparDA zudem einen bis zu 5,3-fach höheren Decode-Durchsatz als die Sparse-Baseline ohne Offload. Unser Quellcode ist verfügbar unter https://github.com/NVlabs/SparDA.

English

Sparse attention reduces compute and memory bandwidth for long-context LLM inference. However, two key challenges remain: (1) KV cache capacity still grows with sequence length, and offloading to CPU memory introduces a PCIe transfer bottleneck; (2) the sparse selection step itself retains O(T^2) complexity and can dominate attention cost at long contexts. We propose SparDA, a decoupled sparse attention architecture that introduces a fourth per-layer projection, the Forecast, alongside Query, Key, and Value. The Forecast predicts the KV blocks needed by the next layer, enabling lookahead selection that overlaps CPU-to-GPU prefetch with current-layer execution. Because Forecast is decoupled from the attention query, our GQA implementation uses one Forecast head per GQA group, reducing selection overhead versus the original multi-head selector. SparDA adds <0.5% parameters and trains only the Forecast projections by matching the original selector's attention distribution. On two sparse-pretrained 8B models, SparDA matches or slightly improves accuracy and delivers up to 1.25times prefill speedup and 1.7times decode speedup over the sparse-attention offload baseline. By enabling larger feasible batch sizes on a single GPU, SparDA further reaches up to 5.3times higher decode throughput than the non-offload sparse baseline. Our source code is available at https://github.com/NVlabs/SparDA.