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.
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.