Dieser Vortrag eines Google-Software-Ingenieurs analysiert das im April auf arXiv veröffentlichte Paper zu Recurrent/PagedAttention (RPA) auf Google Tensor Processing Units (TPUs) für Large Language Model (LLM) Workloads. Die Präsentation legt die internen Hardware- und Compiler-Mechanismen von TPUs offen – insbesondere die Speicherhierarchie, Tiling-Strukturen, Matrix Execution Units (MXU) und den XLA-Compiler-Stack.
Es wird aufgezeigt, wie traditionelle PagedAttention-Ansätze und Mixed-Batching-Verfahren (die Kombination aus Prefill- und Decode-Phasen) aufgrund von unregelmäßigen („ragged“) Datenmustern und dynamischen Scatter/Gather-Operationen fundamentale Engpässe auf den deterministischen, auf statische Formen optimierten Single-Thread-Tensor-Cores erzeugen. Zur Überwindung dieser Beschränkungen implementiert RPA eine Kernel-Fusion aus KV-Cache-Aktualisierung und FlashAttention, gekoppelt mit asynchronen DMAs für softwaregesteuertes Pipelining. Bei Decode-Workloads erzielt das Design eine Speicherauslastung (MBU) von bis zu 86 %, während bei Prefill-Workloads bis zu 73 % der theoretischen Spitzen-FLOPS (bzw. 63 % mit kausaler Maskierung) erreicht werden – vorausgesetzt, die Head-Dimension nutzt das 256×256-MXU-Layout vollständig aus. Abschließend werden bestehende Produktiv-Integrationen (vLLM, SGLang) sowie Optimierungsansätze wie SparseCore-Offloading und Multi-Request-Batching diskutiert.
Wichtigste Punkte & Zeitstempel
00:00:12 RPA-Paper & TPU-Architektur: Systematische Untersuchung von LLM-Workloads auf Google TPUs. Ziel ist die Offenlegung von TPU-Interna bezüglich Speicherhierarchie, Tiling-Verfahren, Packing, MXU-Dimensionierung und XLA-Compiler-Zyklen.
00:01:30 PagedAttention- und Mixed-Batch-Problematik: PagedAttention fragmentiert den KV-Cache und erfordert dynamisches Gather/Scatter, was auf TPU Tensor Cores architektonisch ineffizient ist. Mixed-Batching erzeugt dynamische, unregelmäßige Datenformen („ragged shapes“), die den statischen Optimierungsannahmen des XLA-Compilers widersprechen.
00:03:06 TPU-Speichergranularität & Tiling-Layout: Die minimale DMA-Granularität in XLA beträgt ein V-Word (128 Werte à 32-Bit bzw. 512 Bytes). Datenmatrizen werden in 2D-Kacheln (z. B. 4×128) partitioniert, wodurch das physische Speicherlayout nicht-linear zur logischen Repräsentation verläuft.
00:04:42 MXU-Dimensionierung & Head-Dim-Engpass: Die TPU v5 MXU basiert auf einem 256×256-Rechengitter, das die rechte Matrixseite (RHS) abbildet. Liegt die Dimension (z. B. head_dim = 128) unter 256, halbiert sich die Hardware-Auslastung der MXU zwangsläufig auf maximal 50 %.
00:05:49 Single-Threaded Execution & Asynchrone DMA: TPU Tensor Cores arbeiten strikt Single-Threaded ohne Hardware-Multithreading. Pipelining und Prefetching zwischen Rechenoperationen und Datentransfers müssen explizit über asynchrone DMA-Instruktionen gekoppelt werden.
00:06:36 Kernel-Fusion von KV-Cache & Attention: Um ungerichtete Scatter/Gather-Latenzen zu eliminieren, fusioniert RPA das KV-Cache-Update direkt in die Attention-Berechnung. Während Prefetch-Puffer Daten laden, wird das neu projizierte KV parallel über asynchrone DMAs in den Cache geschrieben, was die Latenz im rechengekoppelten Zustand vollständig verbirgt.
00:08:21 Decode-Performance & XLA-Preprocessing: Bei reinen Decode-Batches (128 Anfragen) erreicht RPA bis zu 86 % MBU. Ein vorgeschaltetes Reshaping erzwingt die kleinstmögliche XLA-Kachelung (512 Bytes), um ungerichtete Tensoren flexibel slicen zu können. Zur vollen Bandbreitensättigung sind 8.000 Kontext-Tokens bei head_dim = 128 bzw. 4.000 bei head_dim = 256 erforderlich.
00:11:09 Prefill-Durchsatz & Skalierungseffekte: Beim Prefill verdoppeln sich die Spitzen-FLOPS beim Wechsel von head_dim = 128 auf 256. Die Recheneffizienz erreicht bis zu 73 % ohne Maskierung und 63 % mit kausaler Maske. Das Tiling-Preprocessing verursacht dabei 2 % bis 8 % Laufzeit-Overhead.
00:13:36 Ablationsanalyse: Vergleichende Messungen zeigen, dass der KV-Cache-Schreibaufwand vollständig durch die Attention-Rechnung maskiert wird. Bei langen Sequenzen konvergiert die Gesamtlaufzeit erwartungsgemäß gegen die reine Rechenzeit (Compute-bound).
00:16:06 VLIW-Bundle-Auslastung & Skalar-Engpässe: Während die MXU-Instruktionen in FlashAttention voll ausgelastet sind, verbleibt die Skalareinheit (SAOU) ungenutzt. Für dynamische Scatter/Gather-Metadatenberechnungen wird Rechenzeit aufgewendet, weshalb eine Auslagerung auf den SparseCore vorgeschlagen wird.
00:17:38 Produktive Framework-Adoption: RPA dient als offizielles Backend für Inferenz-Frameworks wie vllm-tpu und SGLang sowie in DeepMinds Reinforcement-Learning-Repository Simpler.
00:18:14 Limitationen & Roadmap: Zukünftige Optimierungen umfassen native XLA-Unterstützung für unregelmäßige Shapes (Reduzierung des Preprocessing-Overheads), Multi-Request-Batching zur Sättigung kurzer Sequenzen unter 4K Tokens sowie das SparseCore-Offloading für Metadaten-Operationen.
Diskussions-Highlights
00:20:53 MBU versus HFU bei speichergebundenem Decode: Auf Nachfrage bezüglich der Hardware Flops Utilization (HFU) während 86 % MBU erklärt der Vortragende, dass HFU bei reinem Decode irrelevant ist, da der Prozess fundamental speicherbandbreitengebunden ist; entsprechende HFU-Zahlen wurden daher nicht isoliert quantifiziert.
00:23:12 Speicherübertragungsebenen & Steuerungshoheit: Unterscheidung zwischen Off-Chip-Transfers (HBM zu VMEM) und On-Chip-Transfers (VMEM zu Register). Asynchrone Off-Chip-DMAs unterliegen der expliziten Software-Pipeline-Steuerung des Entwicklers, während Registertransfers und temporäre Befehlsumordnungen vollautomatisch durch den XLA-Compiler erfolgen.
Glossar
Asynchronous DMA (Asynchroner direkter Speicherzugriff): Ein Hardwaremechanismus, der Datenübertragungen zwischen Speicherbereichen im Hintergrund durchführt, ohne dass die CPU oder der Rechenkern blockiert wird, wodurch Datentransfer und Berechnung überlappt werden können.
FlashAttention: Ein optimierter Algorithmus zur exakten Aufmerksamkeitsberechnung, der den Speicherzugriff minimiert, indem er Zwischenmatrizen kachelt und Lese-/Schreibzugriffe auf den langsamen Hauptspeicher drastisch reduziert.
HFU (Hardware FLOPS Utilization): Prozentualer Anteil der theoretisch maximal möglichen Gleitkommaoperationen pro Sekunde, den eine Hardware unter einer bestimmten Arbeitslast tatsächlich ausführt.
KV-Cache (Key-Value-Cache): Zwischenspeicher in Transformermodellen, der bereits berechnete Schlüssel- und Werte-Vektoren früherer Tokens ablegt, um redundante Neuberechnungen während der sequenziellen Generierung zu verhindern.
MBU (Memory Bandwidth Utilization): Die prozentuale Auslastung der physisch maximal verfügbaren Datentransferrate des Speichers (z. B. HBM).
Mixed Batching: Gleichzeitige Verarbeitung von Anfragen in der Prefill-Phase (Initialverarbeitung ganzer Prompt-Sequenzen) und der Decode-Phase (schrittweise Einzel-Token-Generierung) innerhalb desselben Batch-Durchlaufs.
MXU (Matrix Execution Unit): Die spezialisierte 2D-Systolic-Array-Hardwareeinheit auf Google TPUs, die für extrem schnelle Matrixmultiplikationen ausgelegt ist.
PagedAttention: Eine Speicherverwaltungsmethode für den KV-Cache, die dem virtuellen Paging in Betriebssystemen ähnelt; sie unterteilt Sequenzen in nicht-zusammenhängende physische Speicherblöcke, um interne Speicherfragmentierung zu vermeiden.
Ragged Tensors (Unregelmäßige Tensoren): Tensoren mit ungleichmäßigen oder dynamischen Dimensionen (z. B. Batches mit stark variierenden Sequenzlängen), die nicht in starre, dichte Matrizen passen.
RPA (Recurrent / PagedAttention auf TPU): Die vorgestellte Kernel-Implementierung von PagedAttention, die speziell auf die architektonischen Eigenheiten und Befehlssätze von Google TPUs abgestimmt ist.
Scatter / Gather (Streuen / Sammeln): Speicheroperationen, die Daten aus nicht zusammenhängenden Speicheradressen in einen zusammenhängenden Block aggregieren (Gather) bzw. Daten aus einem Block auf verstreute Speicherstellen verteilen (Scatter).
SparseCore: Ein spezialisierter Co-Prozessor auf TPU-Chips, der optimiert ist für unregelmäßige Speicherzugriffe, Einbettungsoperationen und spärliche (sparse) Datenstrukturen.
Tiling (Kachelung): Aufteilung großer Tensoren in kleinere, hardwarespezifisch dimensionierte Blöcke (Kacheln), um die lokalen On-Chip-Register und Speicherstrukturen optimal zu bedienen.
V-Word (Vector Word / Female Word): Hardwarespezifische Speichereinheit der TPU-Vektoreinheit; entspricht 128 Elementen à 32-Bit (512 Bytes) und bildet die minimale Adressierungs- und DMA-Granularität.
VLIW (Very Long Instruction Word): Eine Prozessorarchitektur, bei der mehrere unabhängige Operationen in einem einzigen langen Befehls-Bundle zusammengefasst und parallel auf verschiedenen Funktionseinheiten ausgeführt werden.
XLA (Accelerated Linear Algebra): Der domänenspezifische Graph- und Optimierungscompiler für ML-Frameworks (z. B. JAX, TensorFlow, PyTorch/XLA), der Rechenoperationen auf Beschleunigern wie TPUs optimiert und in Hardware-Instruktionen übersetzt.