Dieses Dokument fasst einen Google-Technologievortrag über die Optimierung von Large Language Models (LLMs) auf Tensor Processing Units (TPUs) zusammen, insbesondere im Hinblick auf den im April auf arXiv veröffentlichten „RPA“-Paper. Es beleuchtet TPU-Interna wie die Speicherhierarchie, das Tiling und die Matrix Multiply Units (MXU, 256x256 ab V5). Zu den Kernherausforderungen zählen unregelmäßige ("ragged") Eingaben aus Mixed-Batching und ineffiziente Scatter/Gather-Operationen bei PagedAttention. Die vorgestellte RPA-Architektur (Recurrent/Rapid-Paged Attention) löst dies durch die Fusion von KV-Cache-Aktualisierungen und asynchronen DMA-Pipelines, wodurch Bandbreitenauslastungen (MBU) von bis zu 86 % bei Decoder-Workloads und 73 % bei Prefill-Workloads erreicht werden. Abschließend werden zukünftige Optimierungen wie Sparse-Core-Auslagerung und gleichzeitige Mehrfachabfragen diskutiert.
Wichtigste Punkte & Zeitstempel
0:00 RPA-Paper auf TPUs: Vorstellung der ersten systematischen Analyse von TPU-Workloads für LLMs, publiziert im April auf arXiv durch Google-Softwareingenieur Jefen.
1:30 Herausforderungen bei PageAttention & Mixed Batching: PagedAttention erfordert dynamische Scatter/Gather-Operationen, die von TPU-Tensor-Cores unzureichend unterstützt werden; Mixed Batching erzeugt unregelmäßige ("ragged") Eingaben, was die statische Optimierung des XLA-Compilers erschwert.
3:08 TPU-Speicherhierarchie & Layout: Ein logischer Datenchunk entspricht 128 Werten zu 32 Bit (512 Bytes, die minimale DMA-Granularität eines Vektorworts). Physische Speicherlayouts müssen das hardwarenahe Tiling berücksichtigen.
4:42 MXU-Auslastungseinschränkungen: TPUs ab Version 5 verwenden eine 256x256-MXU (Matrix Multiply Unit). Bei einer Head-Dimension von $H_{dim} = 128$ wird die MXU-Auslastung physikalisch auf maximal 50 % limitiert.
5:56 Single-Threaded-Modell & Asynchrones DMA: TPU-Tensor-Cores arbeiten single-threaded und erzielen Parallelisierung durch asynchrone DMA-Befehle zur Überlappung von Prefetching und Berechnungen.
6:36 Fusionsdesign der RPA-Architektur: Integration der KV-Cache-Aktualisierung direkt in den Attention-Kernel, um die Latenz von Scatter/Gather-Operationen durch parallele DMA-Transfers vollständig zu kaschieren.
8:30 Decoder-Performance-Analyse: Decoder-Workloads erreichen bis zu 86 % Memory Bandwidth Utilization (MBU), erfordern jedoch Kontextlängen von 80K ($H_{dim}=128$) bzw. 4K ($H_{dim}=256$), um die Speicherbandbreite vollständig zu saturieren.
11:09 Prefill-Performance-Analyse: Prefill-Workloads erzielen 73 % MBU ohne Kausalmaske und 63 % mit Kausalmaske, wobei Vorverarbeitungs-Reshapes einen unvermeidbaren Overhead von 2 % bis 8 % verursachen.
16:06 Instruktionsbündelung & IDLE-Zustände: Die MXU ist durchgehend ausgelastet, während die Skalar-ALU (SAOU) bei Metadatenberechnungen untätig bleibt, was auf Einsparpotenziale durch Sparse-Core-Auslagerung hinweist.
17:34 Framework-Adoption: RPA dient als offizielles TPU-Backend für vLLM-TPU (seit Februar des Vorjahres), SGGLM (seit Oktober) und DeepMind-RL-Forschungsrepos.
18:19 Zukünftige Optimierungsansätze: Geplante Weiterentwicklungen umfassen die XLA-Unterstützung für unregelmäßige Formen, parallele Verarbeitung mehrerer Anfragen zur Bandbreitenauslastung bei kleineren Kontextlängen sowie Sparse-Core-Offloading.
Glossar
TPU (Tensor Processing Unit): Von Google entwickelte, anwendungsspezifische integrierte Schaltung (ASIC), die speziell für die Beschleunigung von Deep-Learning- und Machine-Learning-Workloads optimiert ist.
LLM (Large Language Model): Ein auf massiven Textdaten trainiertes Sprachmodell auf Basis von Transformern zur automatisierten Textgenerierung und -verarbeitung.
KV-Cache (Key-Value Cache): Ein Zwischenspeicher in Transformermodellen, der berechnete Schlüssel- (Key) und Wertvektoren (Value) früherer Tokens sichert, um redundante Berechnungen bei der Inferenz zu verhindern.
PagedAttention: Ein Speicherverwaltungssystem für den KV-Cache, das den Speicher in diskrete Blöcke (Pages) unterteilt, um Fragmentierung zu minimieren und dynamisches Caching zu ermöglichen.
XLA (Accelerated Linear Algebra): Ein domänenspezifischer Compiler von Google, der Rechengraphen linearer Algebra optimiert, um die Ausführungsgeschwindigkeit auf CPUs, GPUs und TPUs zu maximieren.
DMA (Direct Memory Access): Ein Hardware-Verfahren, bei dem Peripheriekomponenten oder Speicherbereiche Daten direkt unter Umgehung der Haupt-CPU austauschen, um den Datentransfer zu beschleunigen.
MXU (Matrix Multiply Unit): Die zentrale Recheneinheit eines TPU-Tensor-Cores, die für massiv parallele Matrix-Multiplikationen optimiert ist.
MBU (Memory Bandwidth Utilization): Eine prozentuale Kennzahl, die angibt, wie viel der maximal physikalisch verfügbaren Speicherbandbreite eines Prozessors während der Programmausführung genutzt wird.
HFU (Hardware Flops Utilization): Das Verhältnis der tatsächlich erreichten Gleitkommaoperationen pro Sekunde (FLOPS) zur theoretischen Höchstleistung der Hardware.
Scatter/Gather: Speicheroperationen, bei denen nicht zusammenhängende Daten aus verschiedenen Speicherbereichen gesammelt (Gather) oder verteilt (Scatter) werden, was auf bestimmten Hardware-Architekturen rechenintensiv ist.
Ragged Inputs: Unregelmäßige Datenstrukturen in Batches, bei denen Sequenzen unterschiedliche Längen aufweisen, was die hardwarenahe Vektorisierung und Kompilierung erschwert.
Causal Mask (Kausalmaske): Eine binäre Maske in Attention-Mechanismen, die verhindert, dass das Modell während der Inferenz oder des Trainings auf zukünftige Tokens zugreift.
Sparse Core: Ein spezialisierter Rechenkern auf TPU-Architekturen, der für spärliche (sparse) Berechnungen, Indizierungen und unregelmäßige Speicherzugriffe (wie Scatter/Gather) optimiert ist.
SAOU (Scalar ALU): Skalare Arithmetisch-Logische Einheit eines Prozessors, die im Gegensatz zu Vektoreinheiten einzelne, nicht-vektorisierte Werte verarbeitet.