FlashAttention und PagedAttention: zwei Lösungen für zwei verschiedene Engpässe

Wer sich mit LLMs beschäftigt, hört FlashAttention und PagedAttention oft in einem Atemzug, als wären sie Konkurrenten, und die Frage „Welches soll ich nehmen?“ taucht in ML-Communities ständig auf. Die Frage geht von einer falschen Annahme aus. Beide lösen zwei voneinander unabhängige Probleme in zwei verschiedenen Phasen im Leben eines Modells, und ein ernsthafter Serving-Stack nutzt oft beide gleichzeitig.

Warum Attention teuer ist

Self-Attention in einem Transformer berechnet Folgendes:

Attention(Q, K, V) = softmax(QKᵀ / √d) · V

Die Formel sieht harmlos aus, verbirgt aber zwei getrennte Engpässe, je nachdem, wann man sie ausführt.

1. Bei der Berechnung (Training und Prefill) wächst die Attention-Matrix mit O(N²) in der Sequenzlänge N. Ein Dokument mit 10.000 Tokens ergibt eine Matrix mit 100 Millionen Zellen, die immer wieder berechnet und im GPU-Speicher hin- und hergeschoben wird.
2. Bei der Generierung (Inferenz und Decoding) muss sich das Modell für jedes neue Token die Key- und Value-Vektoren aller vorherigen Tokens merken. Dieser Speicher ist der KV-Cache, und er wächst mit jedem Token und jeder zusätzlichen Nutzeranfrage.

Ein Restaurant ist ein gutes Bild dafür. Das erste Problem ist, wie effizient der Koch ein einzelnes Gericht zubereitet, also die Berechnung. Das zweite ist, wie gut das Restaurant sein Lager verwaltet, wenn Hunderte Bestellungen gleichzeitig laufen, also die Speicherverwaltung. FlashAttention beschleunigt die Berechnung, PagedAttention macht die Speichernutzung beim Serving effizient. Behalten Sie diese Trennung im Kopf, denn alles andere folgt daraus.

FlashAttention: schnellere Attention-Berechnung

Das Problem

Eine GPU hat zwei Arten von Speicher. HBM (High Bandwidth Memory) ist groß, Dutzende GB, aber vergleichsweise langsam. SRAM, der Speicher auf dem Chip, fasst nur wenige MB, ist dafür aber extrem schnell.

Eine Standardimplementierung baut die gesamte N-mal-N-Attention-Matrix auf und schreibt sie in den HBM, liest sie zurück, damit Softmax im SRAM laufen kann, schreibt das Ergebnis erneut und liest es dann noch einmal für die abschließende Multiplikation mit V. Dieses ständige Hin und Her zwischen langsamem und schnellem Speicher ist der eigentliche Engpass. Die Arithmetik selbst ist es nicht.

Stellen Sie sich einen Koch in einer kleinen, schnellen Küche (SRAM) vor, dessen Zutaten alle in einem riesigen Lagerhaus (HBM) einen Kilometer entfernt liegen. Standard-Attention heißt: Der Koch fährt hin, nimmt ein wenig mit, fährt zurück, kocht ein bisschen, fährt wieder hin, um Reste einzulagern, und das bei jedem Schritt. Die Fahrerei kostet mehr Zeit als das Kochen. FlashAttention ist eine einzige gut geplante Fahrt: eine passend große Ladung holen, alle Kochschritte in der Küche erledigen und erst ins Lager zurückkehren, wenn das Gericht fertig ist.

So funktioniert es

FlashAttention ändert nicht, was Attention berechnet, sondern wie. Die wichtigsten Techniken:

• Tiling teilt Q, K und V in kleine Blöcke auf, die in den SRAM passen.
• Online-Softmax erspart es, eine ganze Zeile im Speicher zu halten. Es führt ein laufendes Maximum und eine laufende Summe und aktualisiert beide mit jedem neuen Block, sodass die vollständige N-mal-N-Matrix nie materialisiert wird.
• Fusionierte Kernels führen die Matrixmultiplikation, die Softmax und die zweite Matrixmultiplikation in einem einzigen GPU-Kernel aus, ohne Zwischenschreibvorgänge in den HBM.
• Neuberechnung im Backward Pass bedeutet, dass die riesige Attention-Matrix nicht für die Gradienten gespeichert wird. Sie wird während der Backpropagation spontan neu aufgebaut, was etwas zusätzliche Rechenleistung kostet und dafür sehr viel Speicher spart.

Das Ergebnis ist exakt dasselbe wie bei Standard-Attention. FlashAttention ist keine Näherung, nur ein klügerer Weg zum selben Ziel.

Im Code

Auf unterstützten GPUs wie A100 und H100 wählt PyTorch von selbst einen Kernel im Stil von FlashAttention:

import torch
import torch.nn.functional as F

# PyTorch automatically dispatches to a FlashAttention-style
# kernel on supported GPUs (A100, H100, etc.)
q = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)  # (batch, heads, seq_len, head_dim)
k = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)


output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(output.shape)  # torch.Size([2, 8, 1024, 64])

Man kann auch direkt die spezialisierte Bibliothek flash-attn aufrufen:

from flash_attn import flash_attn_func
import torch

q = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)
k = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)

output = flash_attn_func(q, k, v, causal=True)

Wann es passt

• Große Transformer von Grund auf trainieren oder feinabstimmen.
• Lange Kontextfenster mit 8k, 32k oder 128k Tokens, bei denen das N²-Wachstum richtig wehtut.
• Jeder Fall, im Training oder in der Inferenz, in dem die reine Attention-Geschwindigkeit zählt.

PagedAttention: effizienter Speicher für den KV-Cache

Das Problem

Das ist ein reines Serving-Problem. Beim autoregressiven Decoding speichert das Modell die Key- und Value-Tensoren aller vorherigen Tokens zwischen, damit es sie nicht neu berechnen muss, und das ist der KV-Cache. Herkömmliche Serving-Systeme geraten hier in Schwierigkeiten:

• Die Länge der Ausgabe ist vorab unbekannt, also reservieren sie pro Anfrage einen großen zusammenhängenden Block, etwa für max_tokens = 2048, selbst wenn die Antwort am Ende 50 Tokens lang ist.
• Das führt zu starker interner Fragmentierung. Bis zu 60 bis 80 % des reservierten KV-Caches können leer bleiben und verschwendet werden.
• Nichts wird geteilt. Bei Beam Search oder parallelem Sampling, wo aus einem Prompt mehrere Ausgaben entstehen, bekommt jeder Zweig eine vollständige Kopie des KV-Caches, obwohl alle mit demselben Text beginnen.

Denken Sie an ein Mietshaus (GPU-Speicher) und Mieter (Anfragen). Beim alten Verfahren muss jeder neue Mieter vorsorglich ein ganzes Stockwerk reservieren. Die meisten nutzen zwei oder drei von zwanzig Zimmern, und niemand sonst darf die leeren anrühren. PagedAttention vergibt einzelne Zimmer, also Blöcke oder Seiten fester Größe, dort, wo Platz ist, und nur dann, wenn sie gebraucht werden. Ein Verzeichnis, die Block-Tabelle, hält fest, welche Zimmer wem gehören. Verwandte Mieter, etwa Beam-Search-Zweige mit gemeinsamem Prompt, können sich sogar Zimmer für ihre gemeinsamen Teile teilen. Die Idee ist direkt vom seitenbasierten virtuellen Speicher übernommen, den Betriebssysteme für Prozesse verwenden, daher der Name.

So funktioniert es

1. Der KV-Cache wird in Blöcke fester Größe zerlegt, zum Beispiel je 16 Tokens, statt in ein riesiges zusammenhängendes Array.
2. Jede Sequenz führt eine Block-Tabelle, die logische Blockindizes auf physische Speicherblöcke abbildet, wie eine Seitentabelle im Betriebssystem.
3. Blöcke werden bei Bedarf zugeteilt, erst wenn eine Sequenz so weit gewachsen ist, dass sie einen braucht, also gibt es keine verschwenderische Vorabreservierung.
4. Geteilt wird per Copy-on-Write. Bei Beam Search, parallelem Sampling oder einem gemeinsamen System-Prompt zeigen Sequenzen für ihr gemeinsames Präfix auf dieselben physischen Blöcke und kopieren erst, wenn ihr Inhalt auseinanderläuft.
5. Ein eigener CUDA-Kernel liest über die Block-Tabelle aus den verstreuten Blöcken. Weil die Blöcke nicht mehr zusammenhängend liegen, ist dieser spezialisierte Kernel die eigentliche PagedAttention-Operation.

Im Code

Darauf basiert vLLM, eine der derzeit beliebtesten Serving-Engines für LLMs:

from vllm import LLM, SamplingParams

# vLLM automatically manages KV cache using PagedAttention under the hood
llm = LLM(model="meta-llama/Llama-2-7b-hf")
prompts = [
    "The future of AI is",
    "PagedAttention helps LLM serving by"
]

sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=100)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.outputs[0].text)

Als Nutzer sehen Sie PagedAttention nie direkt. Es arbeitet still im Hintergrund, damit der GPU-Speicher über alle gleichzeitigen Anfragen hinweg gut genutzt wird. Eine vereinfachte, konzeptionelle Sicht darauf, was eine Block-Tabelle enthalten könnte, sieht so aus, und das sind nicht die tatsächlichen Interna von vLLM:

# Simplified conceptual example (not actual vLLM internals)
block_table = {
    "request_1": [block_3, block_7, block_12],
    "request_2": [block_3, block_7, block_9],
    # request_1 & request_2 share blocks 3 & 7 -> same prompt prefix!
}

Wann es passt

• Produktives Serving mit vielen gleichzeitigen Nutzern.
• Chatbot-APIs, RAG-Pipelines und jedes System mit unvorhersehbarer Ausgabelänge.
• Beam Search oder paralleles Sampling, wo geteilter Speicher viel einspart.
• Jedes Setup, bei dem Sie den GPU-Durchsatz maximieren wollen, denn mehr Anfragen pro GPU bedeuten niedrigere Kosten pro Anfrage.

Beides zusammen

Die beiden lassen sich kombinieren, und in der Praxis geschieht das meist auch. PagedAttention ist die Verwaltungsschicht: Sie entscheidet, wo und wie KV-Cache-Blöcke gespeichert werden. FlashAttention oder fusionierte Kernels in seinem Stil bilden die Rechenschicht: Sie bestimmen, wie schnell das Produkt aus Q und transponiertem K sowie die Softmax laufen. vLLM verbindet PagedAttention für den Speicher mit optimierten, von FlashAttention inspirierten fusionierten Attention-Kernels für die Mathematik und erreicht so Speichereffizienz und Rechengeschwindigkeit zugleich.

Um beim Bild der Küche zu bleiben: PagedAttention wählt, aus welchem Regal im Lager die Zutaten kommen, und FlashAttention sorgt dafür, dass der Koch effizient arbeitet, wenn sie da sind.

Was brauchen Sie?

• Sie trainieren ein Modell oder stimmen es mit langen Sequenzen fein ab: FlashAttention.
• Sie stellen eine LLM-API für Hunderte oder Tausende gleichzeitige Nutzer bereit: PagedAttention, über vLLM oder etwas Ähnliches.
• Sie bauen einen Serving-Stack für die Produktion: beides. Die meisten modernen Frameworks, darunter vLLM, TensorRT-LLM und Hugging Face TGI, kombinieren intern bereits ähnliche Ideen.

Kurz gesagt

Das gemeinsame „Attention“ im Namen ist der einzige Grund, warum beide in einen Topf geworfen werden. Sie liegen auf verschiedenen Ebenen des Stacks. FlashAttention macht die Attention-Berechnung schneller und speicherschonender, indem es steuert, wie Daten zwischen HBM und SRAM wandern. PagedAttention macht die Speicherung des KV-Caches bei der Inferenz effizient und frei von Fragmentierung, indem es den seitenbasierten virtuellen Speicher von Betriebssystemen übernimmt.

Das eine optimiert die Berechnung, das andere die Speicherverwaltung beim Serving. Die meisten produktiven Serving-Systeme, vLLM eingeschlossen, nutzen heute beides, damit Modelle schnell laufen und günstig skalieren. Wenn also jemand fragt „PagedAttention oder FlashAttention, was ist besser?“, lautet die ehrliche Antwort: Es ist kein Entweder-oder, sondern ein Sowohl-als-auch.