Browse Summaries

Abstract

Diese Präsentation von Brett Barkley (UT Austin) untersucht die labelfreie Erkennung von Eingaben außerhalb des Trainingsbereichs (Out-of-Distribution, OOD) mittels eingefrorener (frozen) Repräsentationen von Basismodellen ohne nachgelagertes Fine-Tuning. Als Kernmetrik für die Systemzuverlässigkeit wird der End-to-End-Goodput definiert, welcher fehlerhafte oder unbegründete Modellausgaben (Halluzinationen) präventiv unterdrückt.

Im Rahmen einer Benchmark-Analyse über 59 Bild- und Sprachverarbeitungs-Tasks (unter Einsatz von Architekturen wie DINOv1 bis DINOv3 sowie Qwen-Modellen) wurden zwei Detektionsansätze auf den Aktivierungen der vorletzten Schicht evaluiert: ein komplexes Diffusionsmodell (Elucidated Diffusion) sowie die parametrische Mahalanobis-Distanz über multivariate Normalverteilungen. Die empirischen Resultate demonstrieren, dass mit zunehmender Modellskalierung die Leistungskluft zwischen hochkomplexen Algorithmen und einfacher Distanzmetrik kollabiert. Hochskalierte, eingefrorene Modelle besitzen inhärent suffiziente Repräsentationsräume zur OOD-Identifikation. Mittels einer schnellen OOD-Inferenz (5 ms Latenz via JAX) lässt sich eine echtzeitfähige Query-Filterung implementieren, welche signifikante Compute- und Token-Ressourcen insbesondere bei Chain-of-Thought-Inferenzketten einspart.

Wichtigste Punkte & Zeitstempel

• 0:00 Kontext und End-to-End-Goodput: Zuverlässige KI-Systeme erfordern End-to-End-Goodput, der nicht nur Token-Durchsatz misst, sondern sachlich fundierte Antworten sicherstellt und unbegründete Ausgaben bei epistemischer Unsicherheit unterbindet.
• 2:25 Problemdefinition der labelfreien OOD-Detektion: Das Betriebsszenario setzt Zugriff auf unbeschriftete In-Distribution-Daten (ID) voraus, verbietet jedoch jeglichen Zugriff auf Labels oder OOD-Beispieldaten während der Evaluierung.
• 3:00 Ineffizienz von aufgabenspezifischem Fine-Tuning: Konventionelle OOD-Methoden nutzen rechenintensives Parameter-Fine-Tuning (z. B. via LoRA oder Full-Parameter-Updates), was vortrainierte Repräsentationen kompromittiert; die Zielsetzung liegt stattdessen auf vollständig eingefrorenen Backbone-Modellen.
• 4:03 Methodischer Vergleich (Diffusion vs. Mahalanobis): Verglichen werden ein hochkomplexes generatives Diffusionsmodell auf den Vektoren der vorletzten Schicht und die klassische Mahalanobis-Distanz, die eine multivariate Normalverteilung an den Repräsentationsraum anpasst.
• 5:03 Skalierung und Inferenzgeschwindigkeit via JAX: Unter Verwendung des Bonsai-Frameworks in JAX trainieren die OOD-Diffusionsmodelle für Qwen-Backbones in etwa 2 Minuten. Die Detektionsinferenz benötigt 5 ms, verglichen mit etwa 100 ms für den Vorwärtslauf des 2-Milliarden-Parameter-Sprachmodells.
• 5:50 Benchmarking über multimodale Datensätze: Die Validierung erfolgte über 59 Vision- und NLP-Aufgabenpaarungen hinweg unter Nutzung moderner Open-Source-Implementierungen (u. a. DINO-Varianten und Qwen-Modelle).
• 6:16 Dominanz der Modellskalierung gegenüber Algorithmik: Entlang der AUROC-Metrik zeigt sich: Je größer das Modell und die Repräsentationskapazität (z. B. Übergang von DINOv1 zu DINOv3), desto irrelevanter wird die algorithmische Komplexität des Detektors; die 90 Jahre alte Mahalanobis-Distanz konvergiert zur Leistungsfähigkeit moderner Diffusionsmodelle.
• 7:37 Echtzeit-Abweisung von OOD-Prompts: Durch Definieren eines Schwellenwerts auf dem skalaren Anomalie-Score lassen sich OOD-Eingaben vor der autoregressiven Textgenerierung abfangen.
• 9:20 Token- und Latenzeinsparung bei Chain-of-Thought: Das präventive Abbrechen fehlschlagender Inferenzpfade verhindert Ressourcenverschwendung durch entgleiste Denkketten (Chain-of-Thought) in agentischen Systemen.

Diskussions-Highlights

• 10:51 Token-Level-Granularität vs. Query-Level-Aggregation: Auf Publikumsnachfrage bestätigt Barkley, dass das System auf der Gesamteingabe operiert, Repräsentationsabweichungen auf Sub-Token-Ebene jedoch messbar bleiben und sich zu einem robusten Gesamtanomaliewert aggregieren lassen.
• 11:47 Repräsentationsdimension und Modellarchitektur: Für das 2-Milliarden-Parameter-Modell Qwen-3 wird ein latenter Vektor der Dimension 1024 extrahiert, auf dem das Diffusionsmodell mit minimalem Parameter-Footprint aufsetzt.
• 12:00 Benchmark-Datensätze und Diffusionsalgorithmus: Als Text-Benchmark dient unter anderem CLINC150 (Banking-Queries vs. OOD-Domänen); algorithmisch basiert der komplexe Detektor auf dem Paradigma der Elucidated Diffusion (EDM).

Glossar

• Out-of-Distribution (OOD): Daten oder Eingaben, die statistisch signifikant von den Daten abweichen, mit denen ein maschinelles Lernmodell trainiert wurde.
• In-Distribution (ID): Daten, die der Wahrscheinlichkeitsverteilung des Trainingsdatensatzes entsprechen.
• Goodput: Ein aus der Netzwerktechnik stammender Begriff, der hier den Anteil der nützlichen, korrekten und semantisch gültigen Antworten am gesamten Systemdurchsatz beschreibt.
• Eingefrorenes Modell (Frozen Model): Ein neuronales Netz, dessen Parameter während weiterer Rechen- oder Analyseverfahren fixiert bleiben und nicht per Gradientenabstieg modifiziert werden.
• Vorletzte Schicht (Penultimate Layer Representation): Die Aktivierungsvektoren der Schicht direkt vor dem finalen Klassifikations- oder Ausgabekopf, welche hochgradig verdichtete semantische Merkmale enthalten.
• Mahalanobis-Distanz: Ein multivariates Distanzmaß, das den Abstand eines Punktes zu einer Verteilung unter Berücksichtigung der Kovarianzen zwischen den Merkmalen misst.
• Diffusionsmodell / Elucidated Diffusion: Eine Klasse generativer Modelle, die Daten durch schrittweises Entrauschen erzeugen; „Elucidated Diffusion“ (EDM) bezeichnet eine mathematisch fundierte Formulierung nach Karras et al. zur Stabilisierung des Trainings.
• AUROC (Area Under the Receiver Operating Characteristic Curve): Ein gütekritisches Evaluationsmaß für binäre Klassifikatoren. Ein Wert von 1,0 bedeutet eine fehlerfreie Unterscheidung zwischen In-Distribution und OOD.
• Chain-of-Thought (CoT): Eine Prompting- bzw. Inferenztechnik, bei der Sprachmodelle Zwischenschritte explizit generieren, um komplexe logische Probleme zu lösen, was den Token-Verbrauch drastisch erhöht.
• LoRA (Low-Rank Adaptation): Eine Methode des parameter-effizienten Fine-Tunings, bei der vortrainierte Gewichte fixiert und lediglich Matrizen niedrigen Rangs trainiert werden.
• CLINC150: Ein standardisierter Benchmark-Datensatz für Intent-Klassifikation und OOD-Erkennung, bestehend aus Abfragen aus dem Bankenbereich und expliziten Out-of-Scope-Eingaben.

Abstract

Dieser Vortrag von Liu Jiang (UC San Diego / Google) stellt FastFD vor, ein Open-Source-Framework zur disaggregierten Bereitstellung von Large Language Models (LLMs) durch die physische Entkopplung von Attention- und Feed-Forward-Netzwerk-Komponenten (FFN/MoE). In herkömmlichen, kollozierten Inferenzsystemen (wie vLLM) führt die Bindung von Attention- und FFN-Batchgrößen bei langen Kontexten zu massiver Unterauslastung: Da der wachsende KV-Cache die Batch-Größe der Attention limitiert, sinkt die Model FLOPs Utilization (MFU) des FFN auf einer NVIDIA GB200 von 28 % auf 0,5 %. FastFD löst diesen Engpass durch die Trennung in zustandsbehaftete Attention-Knoten und zustandslose MLP-Knoten. Durch schichtübergreifendes Microbatch-Overlapping, ein asymmetrisches M:N-Dispatch / N:M-Combine-Kommunikationsmuster, CPU-basiertes One-Step-Ahead-Scheduling und angepasste Fused Mega-Kernels wird die Rechen- und Transferzeit der FFN-Knoten vollständig hinter der Attention-Laufzeit verborgen. In Benchmarks mit FP8-quantisierten Modellen (MiniMax ~229B, Qwen ~235B) erzielt FastFD bei 8k- und 16k-Kontexten eine Durchsatzsteigerung von 1,35× bis 1,45× pro GPU gegenüber vLLM.

Wichtigste Punkte & Zeitstempel

• 00:11 FastFD-Systemarchitektur: Liu Jiang und Yichao (H-Lab, UC San Diego) präsentieren FastFD zur Trennung von Attention und FFN für hochskalierbare LLM-Inferenz-Engines.
• 00:51 Zustandsbehaftung vs. Zustandslosigkeit: Transformer-Schichten trennen sich fundamental in zustandsbehaftete Attention (Speicherbedarf skaliert über den KV-Cache mit der Kontextlänge) und zustandslose FFN/MLPs (rechengebunden, Effizienz hängt primär von Batch-Größe und Modellgewichten ab).
• 01:37 MFU-Zusammenbruch bei Kolokation: Bei geteilter Hardware (Time-Sharing wie in vLLM) erzwingt die Attention-Berechnung kleine Batch-Größen für FFNs; Simulationen auf NVIDIA GB200 belegen einen MFU-Absturz des FFN von 28 % auf 0,5 % bei steigender Kontextlänge.
• 02:52 Disaggregationsprinzip & M:N-Kommunikation: Attention-Server verwalten Router sowie KV-Cache; MLP-Server führen Experten-Berechnungen durch. Ein aggregiertes M:N-Dispatch- und N:M-Combine-Muster bündelt Requests mehrerer Attention-Server für hohe FFN-Batch-Größen.
• 04:17 Latenzüberhang beim Space-Sharing: Die Aufteilung auf getrennte GPU-Pools ohne Überlappung verdoppelt die Ausführungszeit gegenüber sequenziellem Time-Sharing auf Einzel-GPUs, was eine strikte Parallelisierung erzwingt.
• 05:32 Schichtübergreifendes Microbatch-Overlapping: Ein Dekodierungsschritt wird in Microbatches partitioniert; während der MLP-Server Schicht $k$ rechnet, führt der Attention-Server Schicht $k+1$ aus. Die FFN-Latenz und Kommunikation werden vollständig durch die Attention-Zeit maskiert.
• 06:35 One-Step-Ahead-Scheduling: Zur Vermeidung von Pipeline-Blasen zwischen generierten Tokens plant die CPU den Folge-Dekodierschritt asynchron voraus, während GPU/TPU den aktuellen Schritt verarbeitet.
• 07:07 Asymmetrischer Fused Mega-Kernel: Modifikation des DeepSeek-Mega-Kernels zur Fusion von Dispatch, Experten-GEMMs und Combine-Operationen für ungleiche Sender- (M) und Empfänger-Ränge (N).
• 08:09 Durchsatzgewinn von 1,35× bis 1,45×: Evaluierung auf FP8-quantisierten Modellen (~235B und ~229B Parameter) demonstriert bei 8k- und 16k-Workloads eine Durchsatzsteigerung von ca. 40 % gegenüber kolloziertem Serving.
• 08:54 Latenzsättigung im MiniMax-Benchmark: Mit steigendem Attention-zu-FFN-Verhältnis saturiert der Durchsatzgewinn bei flacher Inferenzlatenz, da die MLP-Ausführung vollständig hinter der dominanten Attention-Dauer verschwindet.
• 09:55 Hardware-Skalierung durch SRAM-Wachstum: Zukünftige Beschleunigerarchitekturen mit verdreifachtem On-Chip-SRAM werden FFN-Rechen- und Transferzeiten weiter minimieren, wodurch noch extremere Disaggregationsverhältnisse realisierbar werden.

Diskussions-Highlights

• 11:07 Einschränkung auf Decoding vs. Prefill: Das Evaluierungssetup misst ausschließlich den Steady-State-Decoding-Durchsatz. Eine Beschleunigung der Prefill-Phase ist gegenwärtig nicht implementiert, da der Codebasis Optimierungen wie CUDA-Piecewise-Graphs fehlen.
• 11:54 Physische Knotendisaggregation via NVL72: Die Disaggregation erfolgt auf Node-Ebene mit Expertenparallelismus 8 (EP8). Ein vollständiger Rechenknoten hostet Attention und KV-Cache, während ein separater Knoten als FFN/MoE-Pool dient; der Datenaustausch basiert auf High-Speed-NVLink-Interconnects.

Glossar

• FastFD: Open-Source-Serving-Architektur zur Entkopplung (Disaggregation) von Attention- und FFN-Berechnungen in Transformer-Inferenzsystemen.
• Attention-Mechanismus (zustandsbehaftet / stateful): Kernalgorithmus von Transformern zur Gewichtung von Token-Zusammenhängen; erfordert zur Vermeidung von Redundanz das Vorhalten historischer Zustände im Arbeitsspeicher.
• Feed-Forward-Netzwerk / FFN / MLP (zustandslos / stateless): Schicht aus dichten Matrixmultiplikationen, die Tokens isoliert transformiert und keinen Verlaufsspeicher über vergangene Tokens benötigt.
• KV-Cache (Key-Value Cache): Im GPU-VRAM allozierter Zwischenspeicher für Schlüssel- und Wertetensoren aller bisher generierten Tokens innerhalb einer Inferenzsequenz.
• Model FLOPs Utilization (MFU): Prozentualer Anteil der theoretisch maximalen Gleitkomma-Rechenleistung eines Hardware-Beschleunigers, der real für Modellberechnungen genutzt wird.
• vLLM: Industriestandard-Software zur LLM-Inferenz, die standardmäßig Attention und FFN auf denselben Beschleunigern ausführt (Kollokation).
• Disaggregation: Trennung traditionell integrierter Rechenschritte auf spezialisierte, physisch getrennte Hardware-Knoten.
• Microbatch Overlapping: Parallelisierungstechnik, bei der Anfragen in Mikro-Teilpakete zerlegt werden, um Ausführungsphasen unterschiedlicher Modellschichten hardwareübergreifend zeitlich zu verschränken.
• One-Step-Ahead Scheduling: Pipelining-Methode, bei der Host-Prozessoren (CPUs) Steuerbefehle für den Folgeschritt vorbereiten, während Beschleuniger noch den aktuellen Schritt ausführen, um Leerlaufzeiten zu eliminieren.
• Mega-Kernel: Hochoptimierter GPU-Kernel, der mehrere Einzelschritte (Datenversand, Matrixmultiplikation, Aggregation) ohne Zwischenspeichern im globalen Speicher direkt in Rechenregistern verschmilzt.
• M:N-Dispatch / N:M-Combine: Kommunikationsmuster zur asynchronen Datenverteilung von $M$ Attention-Knoten an $N$ Experten-Knoten und anschließender Rückführung.
• FP8 (8-Bit Floating Point): Numerisches Gleitkommaformat mit 8 Bit Datenbreite zur Halbierung des Speicherbedarfs und Erhöhung des Rechendurchsatzes gegenüber FP16/BF16.
• Expertenparallelismus (EP / Expert Parallelism, hier EP8): Verteilung der disjunkten neuronalen Netze ("Experten") eines Mixture-of-Experts-Modells (MoE) auf 8 getrennte GPUs oder Knoten.
• NVLink / NVL72: Proprietärer Hochgeschwindigkeits-Bus von NVIDIA zur direkten Kopplung von Beschleunigern mit extrem hoher Bandbreite (z. B. 72 GPUs in einer NVLink-Domain).
• Prefilling vs. Decoding: Die zwei Betriebsmodi der LLM-Inferenz; Prefilling verarbeitet den Eingabetext hochgradig parallel, während Decoding neue Tokens iterativ, Schritt für Schritt generiert.

Abstract

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.

Abstract

In diesem Vortrag analysiert ein Entwickler aus dem Flax-Team die Architekturüberholung von NNXv2, die als wesentliche Schnittstellenänderung mit Flax 0.13 veröffentlicht wird. Das Redesign adressiert Performance- und Komplexitätsgrenzen früherer Versionen durch die Einführung eines restriktiven Tree Mode (graph=False, graph_updates=False), der standardmäßig strikte PyTree-Hierarchien vorschreibt und das dynamische Anhängen neuer Modulattribute während der Laufzeit untersagt. Zu den funktionalen Erweiterungen gehören die Unterstützung mutierbarer Variablen innerhalb von Funktions-Closures, die explizite Zustandspartitionierung über die neue API unpack sowie die Minimierung von Python-Serialisierungs-Overheads via jit_partial. Dynamische Zwischenzustände und Metriken werden über das deklarative Muster sow und capture isoliert. Diese Maßnahmen reduzieren die interne Codepfad-Komplexität automatischer Transformationen signifikant und stellen die Vorwärtskompatibilität zur geplanten JAX-Initiative HiJAX sicher.

Wichtigste Punkte & Zeitstempel

• 00:00:18 NNX-Zustandsmodell: NNX bildet veränderliche Zustände wie Modellparameter und Zähler über Variable-Abstraktionen ab, die als inspektierbare, Python-konforme PyTrees instanziiert werden.
• 00:02:46 Zustandsausbreitung in Transformationen: Die Zustandsfortführung erfolgt entweder manuell über nnx.State und nnx.update in reinen JAX-Pipelines oder transparent über NNX-spezifische Transformations-Wrapper (nnx.jit, nnx.grad).
• 00:04:16 Eliminierung dynamischer Graph-Updates: Das dynamische Erweitern von Modulen um neue Felder zur Laufzeit verursachte unverhältnismäßig hohe Implementierungs- und Optimierungskomplexität und wird vollständig entfernt.
• 00:05:28 Standardisierung auf Tree Mode: NNXv2 erzwingt standardmäßig strikte Baumstrukturen (graph=False); zyklische Graphstrukturen oder Parameter-Sharing müssen künftig explizit konfiguriert werden.
• 00:06:37 Ausrichtung an nativen JAX-Schnittstellen: Spezifische Abstraktionen wie state_axes oder diff_state entfallen zugunsten nativer JAX-Parameter wie in_axes; Transformationen geben mutierte Zustände intern vereinfacht als Hilfsausgaben aus.
• 00:08:21 Unterstützung von Variable-Closures: Variablen können künftig stabil als Closure-Argumente an transformierte Funktionen übergeben und innerhalb dieser modifiziert werden.
• 00:08:55 Zustandstrennung mit unpack: Die API unpack trennt Modellkomponenten vorab strikt in differenzierbare Parameter (params) und nicht-differenzierbare Zustände (nondiff).
• 00:09:07 Optimierung durch jit_partial: jit_partial flacht verschachtelte PyTree-Strukturen von Modell und Optimizer im Vorfeld ab, um den wiederkehrenden Python-Overhead pro Iterationsschritt vollständig zu beseitigen.
• 00:11:53 Zwischenwerterfassung via sow und capture: Die Erfassung von Zwischenaktivierungen und Metriken wird auf die APIs sow und capture umgestellt, welche Zwischenwerte in vorab instanziierten Variablen ablegen.
• 00:13:40 Konvergenz mit HiJAX und Release: NNXv2 harmonisiert die interne Architektur mit der kommenden JAX-Spezifikation HiJAX, um dedizierte NNX-Wrapper perspektivisch überflüssig zu machen; der Release erfolgt als Breaking Change in Flax 0.13.

Diskussions-Highlights

• 00:15:04 Anwendungsbereiche von sow und capture: Das Schnittstellenpaar dient der Extraktion von Zwischenmetriken oder ebenenspezifischen Zwischenverlusten (wie z. B. KL-Divergenzen), ohne Signaturen verschachtelter Schichten modifizieren zu müssen.
• 00:16:36 Erfassung rückwärtsgerichteter Signale via perturb: Über die komplementäre API perturb können in Verbindung mit capture auch Gradienten- und Backward-Metriken innerhalb des Modells isoliert und abgeführt werden.

Glossar

• NNX: Eine objektorientierte Schnittstellenerweiterung für das Framework Flax/JAX, die Python-Klassen mit veränderlichem Zustand (State) strukturiert abbildet.
• JAX: Eine von Google entwickelte Python-Bibliothek für Hochleistungs-Machine-Learning, die automatische Differenzierung mit Just-in-Time-Kompilierung über XLA verbindet.
• Flax: Eine modulare Bibliothek für neuronale Netzwerke, die auf JAX aufsetzt.
• PyTree: Eine hierarchische Datenstruktur in JAX (z. B. verschachtelte Dictionaries, Listen, Tupel), deren Blätter Tensoren bzw. Arrays darstellen.
• JIT (Just-in-Time Compilation): Das Übersetzen von Python- und JAX-Code in hochoptimierten Maschinencode für Beschleuniger (CPUs, GPUs, TPUs) zur Laufzeit mittels des XLA-Compilers.
• Tree Mode: Ein restriktiver Ausführungsmodus, der sicherstellt, dass Modellkomponenten rein hierarchische Baumstrukturen ohne zyklische Querverweise bilden.
• Closure: Ein Funktionsabschluss in der Programmierung; eine Funktion, die Variablen aus ihrer umgebenden Definitionsebene erfasst und mitführt.
• Auxiliary Output (Hilfsausgabe): Zusätzliche Ausgabewerte einer JAX-Funktion, die neben dem primären Rechenergebnis zurückgegeben werden (z. B. aktualisierte Statusvariablen).
• Flattening / Unflattening: Der Prozess des Zerlegens eines PyTrees in eine eindimensionale Blattliste (und Metadaten) sowie die anschließende Rekonstruktion für JAX-Operationen.
• sow / capture: Ein Entwurfsmuster in NNX, um an beliebigen Stellen im Modell Zwischenwerte zu deponieren (sow) und an übergeordneter Stelle gebündelt abzufangen (capture).
• KL-Loss (Kullback-Leibler-Divergenz): Ein mathematisches Distanzmaß zwischen zwei Wahrscheinlichkeitsverteilungen, das häufig als Regularisierungsterm in neuronalen Zwischenschichten eingesetzt wird.
• HiJAX: Eine geplante evolutionäre Architekturerweiterung des JAX-Kerns zur Vereinheitlichung und Vereinfachung höherwertiger Zustands- und Modellabstraktionen.
• perturb: Eine spezialisierte NNX-API zum Abgreifen oder Modifizieren von Gradientenflüssen während der automatischen Rückwärtspropagierung (Backward Pass).

Abstract

Diese Präsentation stellt TokuMX (auch als Tokamak X referenziert) vor, eine kuratierte Open-Source-Bibliothek für hochperformante, produktionsreife Custom-Kernels zur Beschleunigung von Training und Inferenz moderner KI-Modelle. Das Projekt adressiert die Herausforderung, State-of-the-Art-Architekturen optimal auf aktuellen Hardware-Beschleunigern abzubilden – darunter Google TPUs (Ironwood/TPU v5p, TPU 8/8i/80) sowie NVIDIA-GPUs (Hopper H100, Blackwell B200).

TokuMX kombiniert hardwarenahe Implementierungen mit robuster MLOps-Infrastruktur: Es umfasst automatisiertes Autotuning mit Caching optimaler Konfigurationen pro Tensorform, systematisches Performanz-Benchmarking via Xprof zur Regressionserkennung und kontinuierliche numerische Validierung gegen XLA als Referenz. Inhaltlich deckt der Kernel-Katalog moderne Architekturen wie lineare und hybride Attention (Mamba, Gated DeltaNet), Compressed Attention (DeepSeek, Kimi), Mixture-of-Experts-Primitiven (Token-Sortierung, Routing) sowie fusionierte Quantisierung ab. Die Bibliothek verbirgt Low-Level-JAX- und Palace-Code hinter nativen PyTorch-APIs und integriert sich in Serving- und Trainingsframeworks wie vLLM, MaxText und TuneX, ergänzt durch KI-gestützte Kernel-Entwicklungs-Workflows.

Wichtigste Punkte & Zeitstempel

• 0:00 TokuMX-Kernziel: Etablierung einer verlässlichen, kuratierten Bibliothek optimierter Custom-Kernels zur Maximierung der Trainings- und Inferenzleistung offener Modelle.
• 1:07 Hardware-Unterstützung: Primäre TPU-Ausrichtung auf Ironwood (TPU v5p) mit laufender Entwicklung für TPU 8 (8i für Inferenz, 80 für Training) sowie NVIDIA H100 und B200.
• 2:14 Automatisches Autotuning: Strukturierte Suche über Hyperparameter-Räume je nach Tensor-Inputformate mit Caching zur deterministischen Konfigurationsauswahl zur Laufzeit.
• 4:00 Benchmarking via Xprof: Niedrig-Overhead-Leistungsmessung mittels Xprof, automatisiertes Tracking nächtlicher Regressionen und Quellcode-Bisektion bei Kernel- oder Compiler-Änderungen.
• 5:04 Numerische Validierung: Laufender Abgleich von Kernel-Berechnungen gegen XLA als Referenz ("Golden Reference") mit engen Toleranzen zur Vermeidung trainingsdestabilisierender Drift.
• 6:12 Modellgetriebener Katalog: Laufende Erweiterung der Kernel-Unterstützung entlang moderner Open-Source-Modellarchitekturen statt isolierter Mikro-Benchmarks.
• 7:16 Lineare und Hybride Attention: Bereitstellung spezialisierter Kernels für neuartige Aufmerksamkeitsmechanismen wie Gated DeltaNet und Mamba-Varianten.
• 7:34 Compressed Attention & MoE: Kernel-Optimierungen für komprimierte Attention (DeepSeek, Kimi) sowie Beschleunigung kritischer Mixture-of-Experts-Pfade (Token-Sortierung, Expertendispatching).
• 8:25 Low-Level-APIs (Palace): Kontinuierliche Integration neuer Palace-Schnittstellen und Hardware-Features, einschließlich dedizierter Kollektivbeschleuniger auf TPU 8i.
• 10:15 Fused Quantization: Direktes Einbetten von Quantisierungsroutinen (z. B. blockweise Quantisierung) in die Attention-Kernels für verbesserte Speicher- und Durchsatzleistung.
• 11:31 PyTorch-Ergonomie: Kapselung der zugrundeliegenden JAX/Palace-Implementierungen in idiomatische PyTorch-APIs zur Vermeidung von Reibungsverlusten für Torch-Entwickler.
• 12:49 Framework-Integration: Konsolidierung fragmentierter Kernel-Implementierungen durch direkte Anbindung an High-Level-Systeme wie vLLM, MaxText und TuneX.
• 13:48 KI-unterstützte Kernel-Entwicklung: Symbiotischer Einsatz von LLM-Agenten zur automatisierten Code-Optimierung unter Nutzung manuell getunter Kernels als Referenzarchitektur.
• 15:18 Infrastruktur-Evolution: Geplante automatische Ableitung von Tuning-Profilen aus XLA-Traces und Open-Source-Portierung der bislang Google-internen Autotuning-Clusterarchitektur.

Glossar

• Autotuning: Ein systematischer, automatisierter Optimierungsprozess, der empirisch die performantesten Ausführungsparameter (z. B. Blockgrößen, Thread-Zuweisungen) eines Kernels für spezifische Tensor-Eingangsgrößen ermittelt.
• Collective Accelerator: Eine spezialisierte On-Chip-Hardwarekomponente (auf TPU 8i), die kollektive Kommunikationsoperationen (wie All-Reduce, All-Gather) parallel zu Matrixberechnungen ausführt.
• Compressed Attention: Verfahren zur Reduktion des Rechen- und Speicheraufwands im Aufmerksamkeitsmechanismus (z. B. durch Projektion oder Kompression des Key-Value-Caches), prominent genutzt in Modellen wie DeepSeek oder Kimi.
• Gated DeltaNet: Eine Architekturvariante linearer bzw. rekurrenter Aufmerksamkeitsmechanismen, die assoziative Speicherzustände mit adaptiven Gating-Mechanismen kombiniert.
• Ironwood: Codename für Googles TPU-v5p-Hardwaregeneration, ausgelegt für anspruchsvolle Large-Scale-Trainings- und Inferenz-Workloads.
• Kernel: Ein hardwarenah geschriebenes Rechenprogramm, das für die massiv-parallele Ausführung einzelner mathematischer Operationen (z. B. Matrixmultiplikation, Softmax) auf Beschleunigern (GPU/TPU) optimiert ist.
• Mamba: Eine Familie von State-Space-Modellen (SSMs), die lineare Skalierung bezüglich der Sequenzlänge bietet und als Alternative zu quadratisch skalierender Standard-Transformer-Attention dient.
• MaxText: Ein hochgradig skalierbares, in JAX/XLA geschriebenes Open-Source-LLM-Trainingsframework für Google Cloud TPUs und GPUs.
• MoE (Mixture of Experts): Modellarchitektur, die Tokens dynamisch an spezialisierte Sub-Netzwerke ("Experten") routet; erfordert optimierte Kernels für das Sortieren, Dispatching und Zusammenführen (Gather) der Tokens.
• NVIDIA B200 / H100: Hochleistungs-Grafikprozessoren (GPUs) basierend auf der Blackwell- (B200) bzw. Hopper-Architektur (H100) von NVIDIA.
• Palace: Eine hardwarenahe Programmierschnittstelle und domänenspezifische Sprache (DSL) innerhalb des JAX-Ökosystems zur direkten Programmierung von Google TPUs.
• PyTorch: Ein dominantes Open-Source-Machine-Learning-Framework; TokuMX stellt Wrapper bereit, um JAX-basierte Hardware-Kernels wie reguläre Torch-Operatoren anzusprechen.
• Quantized Attention: Fusionierte Ausführung von Aufmerksamkeitsberechnungen in reduzierten Zahlenformaten (z. B. INT8, FP8) oder blockweisen Skalierungen zur Reduktion von Speicherbandbreiten-Engpässen.
• TuneX: Ein Framework zur Feinabstimmung und Optimierung von Machine-Learning-Modellen auf verteilten Beschleuniger-Clustern.
• vLLM: Eine weitverbreitete Open-Source-Bibliothek zur hochdurchsatzstarken und latenzarmen Bereitstellung (Serving) von Large Language Models.
• XLA (Accelerated Linear Algebra): Ein domänenspezifischer Optimierungscompiler für Machine Learning, der Berechnungsgraphen analysiert, Operationen fusioniert und hier als mathematische Referenzlösung dient.
• Xprof: Ein hardwarenahes Profiling- und Analysetool zur detaillierten Erfassung von Laufzeiten, Pipeline-Blasen und Speicherengpässen auf TPUs und GPUs.

Abstract

Dieser Vortrag eines Google-Produktmanagers für KI-Infrastruktur analysiert das Profiling, Debugging und Optimieren von Machine-Learning-Workloads auf Google Tensor Processing Units (TPUs) mithilfe der Toolsuite XPROF. Nach einer Einführung in die TPU-Architektur – insbesondere die Balance zwischen Tensor Cores für dichte Matrixmultiplikationen, SparseCores für speichergebundene Operationen, High Bandwidth Memory (HBM) mit bis zu 35 TB/s und latenzarmen Inter-Chip Interconnects (ICI) – wird die Vermeidung von Ineffizienzen („Flying Blind“) adressiert. Anhand einer Fallstudie (Mistral 8x22B auf Trillium TPU v6e) wird demonstriert, wie die Vergrößerung der Batch-Größe von 1 auf 8 das Verhältnis von Rechenleistung zu Kommunikation optimiert: Die FLOP-Auslastung steigt von 27,6 % auf 60,74 %, der AllReduce-Overhead sinkt von 31 % auf unter 9 % und die Speichernutzung wird von 12 GB auf 29 GB (von 32 GB) skaliert. Abschließend stellt der Vortrag Erweiterungen von XPROF vor, darunter Low-Level-Kernel-Profiling (MLIR, LLO-Instruktionsbündel), ereignisbasierte Telemetrie mit Zugriff auf über 16.000 Hardware-Register (TPU v7 „Ironwood“), kontinuierliche 90-Sekunden-Snapshots via 2-GB-Ringpuffer sowie native Out-of-Tree-Integration in das PyTorch-Ökosystem via torch.profiler.

Wichtigste Punkte & Zeitstempel

• 0:00 Problemstellung bei mangelnder Hardware-Transparenz: Blindes Ausführen von ML-Workloads ohne Profiling führt zu verzögerten Forschungszyklen, explodierenden Cloud-Kosten und unidentifizierten Engpässen durch unausgelastete Beschleuniger.
• 0:01 TPU-Systemarchitektur und Rechenteilung: Maximale Auslastung basiert auf Systembalance. Die Rechenlast teilt sich auf Tensor Cores (dichte Matrixmultiplikation) und SparseCores (speicherintensive Einbettungs-Lookups) auf, gespeist durch HBM mit Bandbreiten von rund 35 TB/s und direkte ICI-Chip-Kopplung unter Umgehung der Host-CPU.
• 0:03 Datenfluss und Phasenunterschiede: Der Host-CPU-Orchestrator steuert Vorverarbeitung sowie XLA-Kompilierung via PCIe zur TPU. Training fokussiert Vorwärts-/Rückwärtspfade und Gradienten-Updates; Serving priorisiert Kompilierung, Quantisierung und dynamisches Batching.
• 0:04 Ursachen von Memory-Bound-Zuständen: Wenn Matrixmultiplikationseinheiten (MXUs) Daten schneller verarbeiten, als HBM sie über das Vektor-Memory (VM) in Register laden kann, verhungert die Recheneinheit und gerät in ein Speicher-Engpass-Szenario.
• 0:05 Fallstudie Mistral 8x22B (HLO-Op-Profil): Bei Ausführung auf Trillium (4K Sequenzlänge, Batch-Größe 1) blockiert die Kommunikationsoperation AllReduce 31 % der Laufzeit bei nur 27,6 % TPU-FLOP-Auslastung, da Chips auf Gradientenabgleiche warten.
• 0:06 Trace-Viewer-Analyse der Kollektivkommunikation: Eine Rückwärtsschicht benötigt 89 ms, wovon AllReduce 33 ms belegt. Die Datenübertragungsgröße entspricht mit 8×16K×16K exakt den Feed-Forward-Gewichten und ist unabhängig von der Batch-Größe.
• 0:07 Memory-Viewer und Speichergrenzen: Das Tool visualisiert dynamische Zuweisungen und den Speicher-Höchststand (High-Water Mark) zur OOM-Diagnose. Die Baseline nutzt lediglich 12 GB des verfügbaren 32-GB-Speichers aus.
• 0:08 Performance-Steigerung durch Batch-Skalierung: Das Anheben der Batch-Größe auf 8 verschiebt die Auslastung von 27,6 % auf 60,74 % FLOPs, drückt den AllReduce-Zeitanteil auf 6,7–8,7 %, steigert die Speichernutzung auf 29 GB/32 GB und überführt den Workload in den angestrebten Compute-Bound-Zustand.
• 0:09 Kernel-Profiling für Custom-Code: Zur Analyse spezialisierter Kernel (Pallas, Mosaic, Triton), die in Standard-Profilern als Blackbox erscheinen, bietet XPROF Tiefeninspektion von MLIR-Code über XLA-Debug-Flags (xla_enable_custom_call).
• 0:11 LLO-Instruktionsbündel und statische Analyse: XPROF visualisiert maschinennahe Instruktionsbündel pro Taktzyklus für MXU-Auslastung, Skalar-/Vektor-ALUs, Vektorfüllungen (HBM-zu-Register) sowie Cross-Core-Datenbewegungen (XLU).
• 0:12 Dynamische Hardware-Zähler (TPU v7 Ironwood): Statische Analysen übersehen Pipeline-Stalls; dynamisches Sampling von Hardware-Zählern auf Sub-Mikrosekunden-Ebene deckt Synchronisationspausen und die reale Auslastungslücke auf.
• 0:13 Hierarchie der Metriken und Rohdaten-Register: Hardware-Registermetriken von HBM und TPU bilden die absolute Entscheidungsgrundlage („Ground Truth“). XPROF stellt dafür über 16.000 unkorrumpierte Rohdaten-Hardware-Counter tabellarisch bereit.
• 0:14 Kontinuierliche Profiling-Snapshots: Ein stets aktiver 2-GB-Ringpuffer auf dem Host zeichnet die letzten 90 Sekunden Telemetrie bei minimalem Overhead auf und ermöglicht programmatische Snapshots direkt beim Auftreten sporadischer Regressionen.
• 0:15 Native PyTorch-Integration: Über Kineto bzw. Torch-TPU bindet sich XPROF nativ und out-of-tree ein, wodurch Entwickler ohne modifizierten Code direkt die Standard-API torch.profiler auf TPUs nutzen können.
• 0:17 Ressourcen und Profiling Summit: Bereitstellung von Repositories und Dokumentation über OpenXLA, JAX-Guides und den TPU Developer Hub sowie Ankündigung des jährlichen Profiling Summits (20.–21. August).

Glossar

• AllReduce: Eine kollektive Kommunikationsoperation im verteilten Rechnen, die Daten (z. B. Gradienten) von allen beteiligten Chips sammelt, mathematisch aggregiert (z. B. summiert) und das Gesamtergebnis an alle Knoten zurückverteilt.
• Compute-Bound: Zustand, in dem die Ausführungsgeschwindigkeit eines Programms primär durch die Rechenleistung der Prozessorkerne (z. B. FLOP-Kapazität) und nicht durch Speicherbandbreite oder Netzwerktransfer limitiert wird.
• FLOPs / FLOP-Auslastung (Floating Point Operations per Second): Maß für die Anzahl der Gleitkommaoperationen pro Sekunde; die Auslastung beziffert den prozentualen Anteil der real erreichten Rechenleistung gegenüber dem theoretischen Hardware-Maximum.
• HBM (High Bandwidth Memory): Hochleistungsspeicherarchitektur mit gestapelten Speicherchips, die extrem breite Datenbusse und Durchsatzraten (im Vortrag ca. 35 TB/s) für Beschleuniger bereitstellt.
• High-Water Mark: Der exakte Spitzenwert des dynamisch allozierten Arbeitsspeichers während der gesamten Laufzeit eines Programms; kritisch für die Vorbeugung von Speicherüberläufen.
• HLO (High-Level Optimizer): Eine Zwischenrepräsentation (Intermediate Representation) im XLA-Compiler, die Rechenoperationen auf Graphebene abstrahiert und optimiert.
• ICI (Inter-Chip Interconnect): Ein proprietäres Hochgeschwindigkeitsnetzwerk von Google, das TPUs direkt chip-zu-chip ohne Umweg über Host-CPU oder standardmäßige Netzwerkkarten miteinander verbindet.
• Ironwood (TPU v7): Codename für die siebte Generation der Tensor Processing Units von Google.
• Kineto: Eine von Meta entwickelte Profiling-Bibliothek für PyTorch, die Hardware-Aktivitäten von Beschleunigern erfasst.
• LLO (Low-Level Optimizer): Compiler-Stufe unterhalb von HLO, die maschinennahe Instruktionsbündel direkt für die Funktionseinheiten der TPU generiert.
• Memory-Bound: Zustand, in dem Recheneinheiten auf Daten warten müssen, weil die Speicherbandbreite oder die Latenz des Speichertransfers die Verarbeitungsgeschwindigkeit begrenzt.
• Mistral 8x22B: Ein hochdimensionales Open-Weight-Sprachmodell, basierend auf einer Mixture-of-Experts-Architektur (MoE).
• MLIR (Multi-Level Intermediate Representation): Ein modulares Compiler-Framework zur Standardisierung und Optimierung von Zwischencode zwischen High-Level-Frameworks und Maschinencode.
• MXU (Matrix Multiply Unit): Dedizierte Hardware-Rechenwerke innerhalb des Tensor Cores einer TPU, die auf extrem parallele Matrixmultiplikationen spezialisiert sind.
• OOM (Out of Memory): Programmabbruch infolge der vollständigen Erschöpfung des physikalisch verfügbaren Arbeitsspeichers.
• Pallas / Mosaic / Triton: Spezialisierte Programmier-Frameworks und Domain Specific Languages (DSLs) zur manuellen Implementierung hocheffizienter GPU-/TPU-Rechenkerne.
• Performance Counter: Physikalische Zählerregister auf dem Chip, die spezifische Hardware-Ereignisse (z. B. Taktzyklen, Cache-Misses, Pipeline-Stalls) in Echtzeit zählen.
• SparseCore: Spezieller Rechenkern auf TPUs zur beschleunigten Abwicklung unregelmäßiger, spärlicher Speicherzugriffe (z. B. Embedding-Lookups), die klassische Matrix-Einheiten ineffizient belasten würden.
• Tensor Core: Hauptrechenkern einer TPU, spezialisiert auf dichte, hochgradig parallele Matrixberechnungen.
• Trillium (TPU v6e): Die sechste Generation der Google Cloud TPU, optimiert für Energieeffizienz und Skalierbarkeit bei KI-Modellen.
• XLA (Accelerated Linear Algebra): Ein domänenspezifischer Compiler für Machine Learning, der Frameworks wie JAX, PyTorch oder TensorFlow in optimierten Hardware-Binärcode überführt.
• XLU: Hardware-Funktionseinheit auf der TPU für Datenumordnung (Data Shuffling) und Kommunikation zwischen Prozessorkernen.
• XPROF: Die integrierte Performance-Profiling-Software-Suite von Google zur Diagnose und Optimierung von ML-Workloads auf TPUs und GPUs.

Zusammenfassung

Dieser Vortrag von David Masher (Cedars-Sinai Medical Center) erläutert die Integration professionellen medizinischen Fachwissens in Open-Source-Large-Language-Modelle (LLMs) durch das Training mit elektronischen Gesundheitsakten (EHR). Um die Leistung bei realen, komplexen klinischen Aufgaben zu verbessern, wurde das Modell Qwen 34B auf MDM-Abschnitten (Medical Decision Making) aus Notaufnahme-Notizen feinabgestimmt. Das Training erfolgte auf 8x A100 80GB GPUs über 172 Stunden. Die Ergebnisse zeigen eine signifikante Verbesserung der qualitativen ärztlichen Bewertung sowie eine Steigerung des F1-Scores bei ICD-10-Diagnosecodes um 16 %. Herausforderungen umfassen katastrophales Vergessen, den Verlust von Chain-of-Thought-Fähigkeiten und Rechenleistungsbeschränkungen, die durch Skalierung auf größere Modelle (14B/70B+) und Reinforcement Learning behoben werden sollen.

Wichtigste Punkte & Zeitstempel

• 0:00 Projektvorstellung: David Masher vom Cedars-Sinai Medical Center in Los Angeles präsentiert die Einbindung echten klinischen Wissens in Open-Source-LLMs für reale medizinische Anwendungen.
• 0:30 Limitationen von Benchmarks: Standard-LLMs bestehen medizinische Examina, scheitern jedoch an realen, longitudinalen und unvollständigen klinischen Abläufen, was durch Erweiterungen des Stanford-HELM-Frameworks belegt wird.
• 2:13 Das EHR-Datenproblem: Während das Internet Leitlinien und Lehrbücher liefert, bleibt die echte klinische Dokumentation in geschlossenen elektronischen Gesundheitsakten (EHR) verschlossen und erfordert den Zugriff im Krankenhaus.
• 3:38 MDM-Generierung: Extrahieren und Generieren von Medical Decision Making (MDM) Abschnitten aus Notaufnahmedaten ohne Informationsleck der vorläufigen Diagnose mithilfe des Modells Qwen 34B.
• 5:30 Hardware und Training: Training auf 8x A100 80GB GPUs, Hugging Face TRL (Transformer Reinforcement Learning), Cross-Entropy-Loss, 1 Epoche, 10% externen Daten zur Vermeidung von katastrophalem Vergessen; Gesamtdauer betrug 172 Stunden.
• 6:20 Qualitative & Quantitative Evaluation: Zwei Notärzte bewerteten das Modell positiv hinsichtlich Genauigkeit und Nützlichkeit; es gab eine 16%ige Steigerung des F1-Scores bei der Vorhersage von ICD-10-Codes.
• 9:00 Generalisierungsversuche: Bei der Erkennung seltener Ereignisse (In-Hospital-Herzstillstand) traten partieller Label-Kollaps und Überanpassung auf, die ein zusätzliches, vollständiges Fine-Tuning erforderten.
• 10:04 Verlust des Chain-of-Thought: Das Modell verlor nach dem Fine-Tuning seine inhärenten Reasoning-Fähigkeiten (leere Denk-Tags), da keine expliziten Reasoning-Spaces im Datensatz enthalten waren.
• 11:38 Skalierungspläne: Zukünftige Optimierungen umfassen das Training größerer Modelle (14 Milliarden bis 70 Milliarden+ Parameter via Google TPUs), präzisere RegEx-Segmentierung und erweiterte De-Identifizierung.
• 14:08 Q&A-Sitzung: Diskussion über implizites ärztliches Handeln (z.B. weggelassene Diagnosetests als diagnostisches Signal), Multi-Standort-Daten von Cedars-Sinai und zukünftige Reinforcement-Learning-Ansätze.

Glossar

• LLM (Large Language Model): Ein auf massiven Textmengen trainiertes KI-Modell zur Sprachverarbeitung und Texterzeugung.
• EHR (Electronic Health Record): Elektronische Gesundheitsakte, die die vollständige und vertrauliche Patientenakte eines Krankenhauses digital speichert.
• MDM (Medical Decision Making): Ein standardisierter Abschnitt in medizinischen Notizen, der Patientensymptome, wahrscheinlichste Diagnosen und Behandlungspläne zusammenfasst.
• Qwen 34B: Ein Open-Source-LLM der Qwen-Familie mit 34 Milliarden Parametern.
• GPU (Graphics Processing Unit): Spezialisierte Hardware (hier Nvidia A100), die für das rechenintensive Training von KI-Modellen optimiert ist.
• Cross-Entropy-Loss: Eine Standard-Verlustfunktion in der statistischen Modellierung zur Messung der Differenz zwischen vorhergesagten und tatsächlichen Wahrscheinlichkeiten.
• F1-Score: Eine statistische Metrik zur Messung der Genauigkeit eines Modells, die Präzision (Genauigkeit) und Recall (Vollständigkeit) kombiniert.
• ICD-10: Die zehnte Revision der Internationalen Statistischen Klassifikation der Krankheiten und verwandter Gesundheitsprobleme, genutzt für standardisierte Diagnosecodes.
• Chain-of-Thought (CoT): Eine Prompting- oder Trainingsmethode, bei der das Modell seine logischen Zwischenschritte vor dem finalen Ergebnis ausgibt.
• RegEx (Regular Expressions): Formale Suchmuster zur automatisierten Textanalyse und -extraktion in Datensätzen.
• De-Identifizierung: Das Entfernen personenbezogener und sensibler Daten (PHI) aus medizinischen Datensätzen zum Datenschutz.

Zusammenfassung

Diese Präsentation erläutert die Integration des oneAPI-Backends in XLA, um JAX-Workloads auf Intel-GPUs auszuführen, am Beispiel der Battlemage-Architektur (Arc Pro B70). Es werden die Hardware-Komponenten, das oneAPI-Software-Stack (oneMKL, oneDNN, oneCCL), die XLA-Kompilierungspipeline (inklusive oneDNN-Pfad, MLIR-Codegen und experimentellem Triton-Support), die Runtime-Ausführung über Level Zero und SYCL sowie die Multi-GPU-Skalierung via oneCCL detailliert beleuchtet. Zudem werden der PJRT-basierte JAX-Plugin-Status und die Roadmap für Upstream-Optimierungen vorgestellt.

Wichtigste Punkte & Zeitstempel

• 0:09 Intel Battlemage-Architektur: Vorstellung der Arc Pro B70 Workstation-GPU basierend auf der Xe2-Architektur mit 32 GB GDDR6 VRAM, Xe-Cores (je 8 Vektor- und Matrix-Engines) und 256 KB L1/lokalem Speicher pro Core.
• 0:26 oneAPI-Framework: Einbindung des plattformübergreifenden Standards mit Kernbibliotheken wie oneMKL (Mathematik), oneDNN (Deep Learning) und oneCCL (kollektive Kommunikation) für heterogene Hardware.
• 0:40 XLA-Kompilierung für GEMM und Convolutions: Umleitung von HLO-Dot- und Convolutions-Operationen über den oneDNN-Pfad in einzeln gefuste SYCL-Kernels via JIT-Kompilierung, wodurch Overhead eliminiert wird.
• 6:26 MLIR-Codegen & SPIR-V: Verarbeitung von Elementwise- und Reduktions-Operationen über MLIR-Transformatoren, Zielselektion zu SPIR-V-Intrinsics und JIT-Übersetzung in native ISA durch den Intel Graphics Compiler via Level Zero.
• 8:32 Triton-Codegen: Dreiphasiger Kompilierungspfad (Triton IR, TTIG-Dialekt, Triton-GenOps) mit integriertem Autotuner für das Benchmarken und Optimieren von GEMM-Tiling-Konfigurationen.
• 10:57 Runtime-Infrastruktur: Initialisierung des SYCL-Executors und Laden von SPIR-V-Binärdateien in Level-Zero-Module zur asynchronen Ausführung auf GPU-Streams.
• 12:09 Multi-GPU-Skalierung: Nutzung von oneCCL zur Bereitstellung hardwarebeschleunigter Kommunikationsprimitiven (All-Reduce, Broadcast) für verteiltes Training.
• 13:27 JAX-Integration: Bereitstellung eines PJRT-basierten oneAPI-Plugins mit drei Wheels (jaxlib, jax-oneapi-plugin, jax-oneapi-pjrt) für validierte Modelle wie Gemma 2, SDXL und GPT-2.
• 14:14 Roadmap & Upstream-Status: Laufende Upstream-Integration von SPIR-V-Codegen und PJRT; geplante Features umfassen PTI-Profiling, Triton-Fusionskernel, SYCL-Graph-Command-Buffer und Unterstützung zukünftiger Rechenzentrum-GPUs wie Crescent Island.

Glossar

• XLA (Accelerated Linear Algebra): Ein Domain-spezifischer Optimierungscompiler für lineare Algebra, der Berechnungs-Graphen aus Frameworks wie TensorFlow oder JAX in maschinennahen Code übersetzt.
• JAX: Ein Python-Framework für Hochleistungs-Numerik und maschinelles Lernen, das automatische Differenzierung und JIT-Kompilierung unterstützt.
• Battlemage / Arc Pro B70: Intels diskrete Grafikprozessor-Generation (Xe2-Architektur) für professionelle Workstations, optimiert für lokale KI-Inferenz und Grafikberechnungen.
• oneAPI: Ein offener, plattformübergreifender Programmierstandard von Intel zur einheitlichen Softwareentwicklung über CPUs, GPUs und FPGAs hinweg.
• oneMKL / oneDNN / oneCCL: Intel-Bibliotheken für optimierte mathematische Routinen (Math Kernel Library), Deep-Neural-Network-Primitive bzw. kollektive Multiknoten-Kommunikation.
• SYCL: Eine C++-basierte Programmierschnittstelle für heterogenes Computing, die die Entwicklung von Code für verschiedene Hardware-Beschleuniger ermöglicht.
• MLIR (Multi-Level Intermediate Representation): Eine modulare Compiler-Infrastruktur, die strukturierte Zwischenrepräsentationen auf verschiedenen Abstraktionsebenen ermöglicht.
• SPIR-V: Ein offener Standard-Zwischencode (Bytecode) für parallele Shader und Kernel, der von Khronos-Standards und Intel Level Zero genutzt wird.
• Level Zero: Eine Low-Level-Hardwareschnittstelle von Intel zur direkten, effizienten Steuerung von GPU-Ressourcen, Queues und Speicher-Managements ohne Treiber-Overhead.
• PJRT: Eine einheitliche Laufzeitschnittstelle (Plugin Device Runtime) in JAX, die die Anbindung verschiedener Hardware-Beschleuniger standardisiert.
• Triton: Ein von OpenAI entwickelter Compiler und eine Programmiersprache zur automatischen Erstellung hochleistungsfähiger GPU-Kernel.

Zusammenfassung

Dieser Vortrag von Brett Barkley (Doktorand an der UT Austin) präsentiert einen skalierbaren, label-freien Ansatz zur Out-of-Distribution (OOD)-Erkennung mittels vortrainierter Repräsentationen eingefrorener Modelle (wie DinoV3 und Qwen-2B) ohne Fine-Tuning. Das Verfahren vergleicht einen rechenintensiven Ansatz (Diffusionsmodelle, trainiert auf den Repräsentationen der vorletzten Schicht) mit einem simplen, klassischen Ansatz (dem 90 Jahre alten Mahalanobis-Abstand) und erreicht Inferenzzeiten von nur 5 Millisekunden. Evaluengeleitet über 59 Vision- und Sprach-Task-Pairings, zeigt die Arbeit, dass die Modellskalierung weitaus wichtiger ist als die Komplexität des Erkennungsalgorithmus, was eine effiziente Online-OOD-Abfrageerkennung zur Verhinderung von Halluzinationen und Token-Verschwendung (z. B. bei Chain-of-Thought-Prozessen) in modernen LLMs ermöglicht.

Wichtigste Punkte & Zeitstempel

• 0:00 Einführung & Autorenschaft: Brett Barkley (UT Austin) stellt seine Forschung zur label-freien OOD-Erkennung vor, entstanden in Zusammenarbeit mit seinem Betreuer David Friedlander Kyle, Preston Culbertson (Cornell) sowie Unterstützung durch das JAX- und Bonsai-Team von Google und die National Science Foundation.
• 1:48 Kernproblem (Wissen von Modellen): Die zentrale Forschungsfrage lautet: „Wissen Modelle, was sie nicht wissen?“ LLMs liefern oft ungeerdete, hochpräzise klingende Antworten („Halluzinationen“), was bei kritischen Domänen (z. B. Medizin) erhebliche Sicherheitsrisiken birgt.
• 2:26 Betriebsszenario (Label-freie Anomaly Detection): Das System operiert unter realistischen Bedingungen: Zugriff auf ungelabelte In-Distribution-Daten, ohne jeglichen Zugriff auf Out-of-Distribution-Daten oder vordefinierte Labels.
• 3:15 Verzicht auf Fine-Tuning: Anstatt teure Fine-Tuning-Verfahren (wie LoRA oder Vollparameter-Anpassung) an eingefrorenen Transformern durchzuführen, wird untersucht, ob Basismodelle Anomalien bereits inhrent in ihren Repräsentationen erkennen können.
• 4:03 Detektor-Architekturen: Es werden zwei Varianten verglichen: ein komplexes System mittels Diffusionsmodell ("elucidated diffusion") auf Basis der vorletzten Schicht und ein extrem simpler Ansatz über den Mahalanobis-Abstand (Multidimensionale Gauß-Verteilung).
• 5:02 Skalierung mit JAX: Mithilfe von JAX und Open-Source-Implementierungen des Bonsai-Teams konnten Diffusionsmodelle extrem schnell trainiert werden (z. B. Qwen in 2 Minuten); die Inferenz dauert 5 Millisekunden im Vergleich zu 100 ms für das LLM selbst.
• 5:52 Benchmark-Umfang: Die Evaluierung erfolgte breit gefächert über 59 Vision- und Sprachtask-Pairings unter Verwendung von Modellen wie DinoV3 und Qwen-3 (bzw. Qwen-2B).
• 6:16 Skalierung schlägt Algorithmendesign: Der Leistungsvergleich zeigt, dass mit steigender Modellskalierung (von DinoV1 zu DinoV3) die Unterschiede zwischen simplen und komplexen Detektoren verschwinden und beide exzellente Trennergebnisse (Area Under ROC Curve) erzielen.
• 7:37 Online-OOD-Abfrageerkennung: Praktischer Einsatz als Schwellenwert-Filter vor der LLM-Generierung: Liegt der skalare Anomaliescore über der definierten Grenze, wird die Generierung abgebrochen, um Token-Verschwendung bei langen Chain-of-Thought-Prozessen (z. B. Waymos internem "Antigravity") zu verhindern.
• 10:27 Q&A & Technische Details: In der Fragerunde wird geklärt, dass das Verfahren auf Repräsentationsräumen (z. B. 1024-dimensionalen Vektoren bei Qwen-2B) operiert, Token-Muster bewertet und Benchmarks wie Clink 150 (Bankabfragen) nutzt.

Glossar

• Out-of-Distribution (OOD): Daten, Anfragen oder Muster, die sich außerhalb der Verteilung befinden, auf die ein Modell trainiert wurde (unbekannte oder neuartige Eingaben).
• In-Distribution (ID): Daten, die aus derselben statistischen Verteilung stammen wie der ursprüngliche Trainingsdatensatz.
• Fine-Tuning: Das nachträgliche Anpassen der Gewichte eines vortrainierten maschinellen Lernmodells anhand eines spezifischen, zielgerichteten Datensatzes.
• LoRA (Low-Rank Adaptation): Eine ressourceneffiziente Methode zum Fine-Tuning großer Sprachmodelle, bei der die originalen Gewichte eingefroren und stattdessen kleine, zusätzliche Matrizen trainiert werden.
• Transformer: Eine neuronale Netzwerkarchitektur, die auf Selbstaufmerksamkeitsmechanismen (Self-Attention) basiert und als Fundament moderner Sprach- und Bildmodelle dient.
• Mahalanobis-Abstand: Ein statistisches Maß zur Bestimmung des Abstands zwischen einem Datenpunkt und einer Verteilung, das im Gegensatz zum euklidischen Abstand Korrelationen zwischen den Variablen berücksichtigt.
• Diffusionsmodell: Eine generative KI-Architektur, die Daten durch schrittweises Entrauschen generiert; hier zweckentfremdet zur Extraktion eines skalaren Anomaliescores.
• JAX: Ein Hochleistungs-Framework für numerische Berechnungen und maschinelles Lernen in Python, optimiert für automatische Differentiation und Hardware-Beschleunigung (GPU/TPU).
• DinoV3: Ein fortschrittliches, selbstüberwachtes Bildverarbeitungsmodell (Vision Transformer) von Meta.
• Qwen: Eine Familie großer Sprachmodelle (Large Language Models), entwickelt von der Alibaba Group.
• ROC Curve (Receiver Operating Characteristic): Ein grafisches Diagramm zur Bewertung der Güte von Klassifikationsmodellen; die Area Under the Curve (AUC) misst die Fähigkeit zur perfekten Trennung.
• Chain-of-Thought (CoT): Eine Prompting-Technik, bei der ein Sprachmodell schrittweise Argumentationsschritte erarbeitet, um komplexe Probleme zu lösen, was jedoch den Token-Verbrauch drastisch erhöht.
• Token: Die kleinste von einem Sprachmodell verarbeitete Texteinheit, die ganzen Wörtern oder Wortteilen (Subwörtern) entsprechen kann.

Zusammenfassung

Dieser Vortrag von Tomas behandelt Mosaic, einen Compiler, der als Bindeglied zwischen dem Programmiermodell Palace und zielspezifischen Backends (wie Tensor Cores und Sparse Cores) für TPUs (Tensor Processing Units) dient. Mosaic verarbeitet MLIR (Multi-Level Intermediate Representation) mit Kern- und TPU-spezifischen Dialekten und ist primär für die Verwaltung, Transformation und Optimierung von Vektorkonstrukten (Vector Layouts) verantwortlich. Der Vortrag erläutert die vier Komponenten von Vektorlayouts (Bitbreite, Tiling, Offsets mit Replikation und implizite Dimensionen), die Handhabung von Layout-Konvertierungen (Relayouts) sowie die drei zentralen Compiler-Phasen: Infer Vector Layout, Einfügen von Relayouts und Apply Vector Layout (Unrolling). Abschließend werden Unterschiede zu Sparse-Core-Architekturen sowie die Verarbeitung gepackter Sub-32-Bit-Datentypen (z. B. FP8, BF16) erörtert.

Wichtigste Punkte & Zeitstempel

• 0:16 Mosaic-Rolle im Compiler-Stack: Mosaic agiert als Zwischenschicht-Compiler zwischen Palace und TPU-Backends, verarbeitet MLIR-Eingaben und führt hardwarenahe Transformationen für Tensordaten durch.
• 1:16 TPU-Vektorregister (VREGs): TPU-Vektorregister sind zweidimensionale Strukturen aus 8 Sublanes und 128 Lanes; Tensoren müssen für Operationen zwingend in diese Register geladen werden.
• 3:43 Vektor-Layouts (Vector Layouts): Layouts definieren, wie logische N-dimensionale Tensordaten in physischen Registern repräsentiert werden, gesteuert über vier Kernkomponenten.
• 4:39 Tiling-Mechanismus: Tiling bestimmt die Aufteilung von Tensoren in Register-Slices (z. B. 8x128 oder 4x256), bedingt durch Speicherstrukturen von XLA und zur Minimierung von Padding-Overhead.
• 6:46 Offsets und Replikation: Offsets verwalten Padding-Daten aus Slicing-Operationen oder nicht ausgerichteten Speicherzugriffen; Replikations-Flags ermöglichen identische Werte über Dimensionen hinweg für effiziente Reduktionen.
• 9:41 Implizite Dimensionen: Historisch für komplexe Operationen (Reduktion gefolgt von Reshape) eingeführt, erlauben implizite Dimensionen das Vermeidung kostspieliger Transpositionen zwischen Sublanes und Lanes.
• 11:51 Relayout-Operationen: Da binäre Operationen identische Layouts erfordern, fügt Mosaic automatische Konvertierungen (Relayouts) ein, deren Aufwand optimiert werden muss.
• 13:48 Die drei Hauptphasen des Compilers: Der Optimierungsprozess gliedert sich in Infer Vector Layout (globale/lokale Layout-Zuweisung), das Einfügen von Relayouts und Apply Vector Layout (Unrolling der Operationen für das Ziel-Backend).
• 16:36 Sparse Core (Sparsore) Unterstützung: Sparse Cores nutzen kleinere Registerstrukturen (z. B. 8x32 oder 16x32 Bit) und verlangen separate Layout-Logiken für nicht-zusammenhängend gepackte Datenelemente.
• 19:04 C++-Nutzung und Layout-Heuristiken: Fortgeschrittene Anwender nutzen Mosaic direkt via C++, um Python-Schichten zu umgehen; bestehende Layout-Heuristiken sind teilweise lokal und werden kontinuierlich verbessert.
• 21:41 Gepackte Datentypen (Packed Types): Sub-32-Bit-Typen wie FP8, BF16 oder INT16 packen mehrere Elemente in ein Vektorregister, wodurch die logische Registerkapazität entsprechend multipliziert wird.

Glossar

• Mosaic: Ein spezialisierter Compiler im TPU-Software-Stack zur Steuerung von Speicher- und Register-Layouts zwischen High-Level-APIs und Hardware-Backends.
• Palace: Ein High-Level-Programmiermodell für TPUs (oft in Python oder C++), das MLIR-Code für Mosaic generiert.
• MLIR (Multi-Level Intermediate Representation): Ein modulares Compiler-Infrastruktur-Framework, das als Eingabeformat für Mosaic dient und domänenspezifische Dialekte verwendet.
• TPU (Tensor Processing Unit): Ein von Google entwickelter, kundenspezifischer Prozessor (ASIC), der speziell für die Beschleunigung von Deep-Learning- und Tensor-Berechnungen optimiert ist.
• VREG (Vector Register): Ein Hardware-Vektorregister auf der TPU, das als 2D-Matrix (z. B. 8 Sublanes mal 128 Lanes) organisiert ist.
• Tiling: Die Zerlegung großer, mehrdimensionaler Tensoren in kleinere Blöcke (Tiles), um sie effizient in Hardware-Registern zu verarbeiten.
• Padding / Offsets: Unnütze oder mit Platzhalterdaten gefüllte Speicherbereiche in Registern, die durch Slicing oder unaligned Memory Loads entstehen.
• Unrolling: Ein Compiler-Optimierungsschritt, bei dem abstrakte Operationen oder Schleifen in explizite, sequenzielle Maschineninstruktionen aufgelöst werden.
• Tensor Core / Sparsore: Spezifische Hardware-Ausführungseinheiten (Backends) der TPU für dichte bzw. dünnbesetzte (sparse) Matrixoperationen.
• FP8 / BF16 / INT16: Numerische Datenformate mit reduzierter Bitbreite (8-Bit-Gleitkomma, 16-Bit-Bfloat, 16-Bit-Ganzzahl) zur Steigerung von Rechen- und Speichereffizienz.
• Sublanes und Lanes: Die physischen Dimensionen der internen TPU-Vektorregisterstruktur, die die parallele Datenverarbeitung steuern.

Zusammenfassung

Diese Präsentation stellt UniDiffuser vor, ein auf der ICML veröffentlichtes Framework für multimodale Argumentation und Generierung, entwickelt an der University of Illinois Urbana-Champaign. UniDiffuser überwindet die Fragmentierung heutiger Systeme – bei denen getrennte Vision-Language-Modelle (VLMs) für das Verständnis und Diffusionsmodelle für die Generierung eingesetzt werden –, indem es ein einziges diskretes Diffusions-Framework nutzt. Durch den Einsatz von diskretem Flow Matching und effizienten LoRA-Adaptern vereint das Modell multimodales Verständnis, Bildgenerierung und intelligentes Editieren in einer einzigen probabilistischen Schnittstelle. Das System erreicht eine State-of-the-Art-Leistung bei geringem Rechenaufwand, kürzeren Inferenzschritten und hoher TPU-Kompatibilität.

Wichtigste Punkte & Zeitstempel

• 0:21 Fragmentierte Multimodale Systeme: Heutige KI-Architekturen leiden unter einer Trennung zwischen sprachbasierten Verständnismodellen und bildgenerierenden Diffusionsmodellen oder erfordern ineffiziente hybride Ansätze mit uneinheitlichen Trainingszielen.
• 0:49 UniDiffuser-Ziel: Entwicklung einer einzigen probabilistischen Schnittstelle, die visuelle Fragebeantwortung (VQA), Text-zu-Bild-Generierung und instruktionsgesteuertes Editieren bei hoher Inferenz-Effizienz vereint.
• 1:52 Vorteile der Diskreten Diffusion: Im Gegensatz zu autoregressiven Modellen mit wachsendem KV-Cache und tausenden Dekodierungsschritten ermöglicht die diskrete Diffusion das parallele Entrauschen von 1024-Token-Rastern in nur 20 Schritten.
• 3:05 Drei-Stufen-Trainingspipeline: Das Modell repurportiert vortrainierte VLMs als diskrete Flow-Matching-Modelle. Von den 4 Milliarden Parametern werden lediglich 900 Millionen über LoRA-Adapter trainiert (Stufe 1: Verständnis mit KL-Divergenz, Stufe 2: Bildgenerierung, Stufe 3: Editieren mittels M ref DPO).
• 4:57 Benchmark-Überlegenheit: UniDiffuser erzielt wettbewerbsfähige oder überlegene Ergebnisse im Vergleich zu wesentlich größeren Modellen (wie Bagel mit 7B und When Image mit 27B) über sechs Benchmarks hinweg bei gesteigertem Inferenz-Durchsatz.
• 5:38 Zero-Shot-Fähigkeiten: Das Framework demonstriert generalisierte Null-Schuss-Fähigkeiten bei OCR, mathematischem Reasoning aus Bildern, In-Painting, referenzbasiertem Editieren und kompositorischen Multi-Objekt-Szenen.
• 6:42 Open Source & JAX-Implementierung: Das System wird in PyTorch und JAX bereitgestellt. Die Inferenzzeit beträgt circa 12 Sekunden auf einer H200-GPU und etwa 9 Sekunden auf einer TPU V4.
• 7:33 Kommerzielle Serving-Optimierungen: Geplante Verbesserungen für den Produktivbetrieb umfassen intelligentes Anfrage-Batching basierend auf Maskenmustern sowie maßgeschneiderte Kernel für den konfidenzbasierten Rekonstruktionsschritt.
• 9:09 Architektur-Einordnung: UniDiffuser basiert auf einer Diffusions-LLM-Architektur, bei der im Gegensatz zur kontinuierlichen Diffusion (wie bei Flux) diskrete Token-Repräsentationen verwendet werden.

Glossar

• Vision-Language-Modell (VLM): Ein künstliches neuronales Netz, das darauf trainiert ist, sowohl visuelle Inhalte (Bilder) als auch Text zu verarbeiten und logisch miteinander zu verknüpfen.
• Autoregressiv: Eine Generierungsmethode, bei der Text oder Daten sequenziell – Token für Token auf Basis der vorherigen Tokens – vorhergesagt werden.
• KV-Cache (Key-Value Cache): Ein temporärer Speicher in Transformer-Modellen, der berechnete Schlüssel- und Wert-Vektoren zwischenspeichert, um redundante Berechnungen bei der sequenziellen Generierung zu vermeiden.
• Diskrete Diffusion: Ein generatives Diffusionsmodell, das auf diskreten Einheiten (wie Vokabular-Tokens oder maskierten Rasterpunkten) anstelle von kontinuierlichen mathematischen Räumen operiert.
• Flow Matching: Ein generatives Modellierungsverfahren, das kontinuierliche Wahrscheinlichkeitsströme lernt, um Daten effizient von einer Ausgangsverteilung in eine Zielverteilung zu transformieren.
• LoRA (Low-Rank Adaptation): Eine Methode zur effizienten Feinabstimmung großer KI-Modelle, bei der nur ein minimaler Bruchteil der Modellparameter über zusätzliche Low-Rank-Matrizen modifiziert wird.
• KL-Divergenz (Kullback-Leibler-Divergenz): Ein statistisches Maß, das quantifiziert, wie stark sich eine Wahrscheinlichkeitsverteilung von einer anderen Referenzverteilung unterscheidet.
• M ref DPO (Multimodal Reference Direct Preference Optimization): Ein Optimierungsverfahren zur Ausrichtung (Alignment) von KI-Modellen an menschlichen Präferenzen unter Einbezug eines Referenzbildes.
• TPU (Tensor Processing Unit): Ein von Google entwickelter, spezialisierter Mikroprozessor (ASIC), der für die beschleunigte Berechnung von Deep-Learning-Algorithmen optimiert ist.
• Inferenz: Die Ausführungsphase eines bereits trainierten KI-Modells, in der neue Daten verarbeitet werden, um Vorhersagen oder Ergebnisse zu generieren.
• Kontinuierliche Diffusion: Ein Diffusionsverfahren, das in einem stetigen, kontinuierlichen Raum (wie bei klassischen Bildgeneratoren) anstelle von diskreten Tokens arbeitet.