NERDMANSakana AI and NVIDIA Introduce TwELL with CUDA Kernels for 20.5% Inference and 21.9% Training Speedup in LLMs
NVIDIA und Sakana AI stellen Twell vor: CUDA-Kernel für 20,5 % schnelleres Inferenz und 21,9 % schnelleres Training von LLMs
Die Skalierung großer Sprachmodelle (LLMs) ist teuer. Jeder Token, der während der Inferenz verarbeitet wird, und jeder Gradient, der während des Trainings berechnet wird, durchläuft die Feedforward-Schichten, die für über zwei Drittel der Modellparameter und mehr als 80 % der gesamten FLOPs in größeren Modellen verantwortlich sind. Ein Forscherteam von Sakana AI und NVIDIA hat an einer neuen Forschungsmethode gearbeitet, direkt auf diesen Engpass zielt – nicht durch eine Änderung der Architektur, sondern durch eine deutlich günstigere Berechnung innerhalb der Feedforward-Schichten.
Die Lösung: Twell – Ein neuer CUDA-Kernel für effizientere Feedforward-Schichten
Das Team stellt Twell vor, einen hochoptimierten CUDA-Kernel, der speziell für die linearen Projektionen in den dichten Feedforward-Blöcken von Transformer-Modellen entwickelt wurde. Der Kern von Twell ist ein neuer, hochparalleler Algorithmus, der die Berechnung der SiLU/Gated Linear Unit (GLU)-Aktivierungsfunktion mit der anschließenden linearen Projektion fusioniert. Diese Fusion eliminiert den teuren Lese-/Schreibzugriff auf den globalen Speicher für die Zwischenaktivierungen, was den Engpass bei der Speicherbandbreite erheblich reduziert.
Messbare Leistungssteigerungen in der Praxis
Die Benchmark-Ergebnisse sind beeindruckend: Inferenz: Twell beschleunigt die Vorhersage um 20,5 % im Vergleich zu den optimierten Baselines in PyTorch (compiled) und der state-of-the-art `vLLM`-Bibliothek. Training: Die Vorwärts- und Rückwärtsdurchläufe während des Trainings werden um 21,9 % beschleunigt.
Diese Geschwindigkeitsgewinne wurden auf dem Llama 3.1B-Modell mit 8 Milliarden Parametern validiert und sind direkt auf andere gängige Modelle wie Mistral und Gemma übertragbar.
**Wie funktioniert es?Die Magie von Twell liegt in seiner Fähigkeit, die Berechnung der GLU-Aktivierung und der linearen Projektion in einem einzigen, zusammengeführten Kernel durchzuführen. Normalerweise würde dies zwei separate Operationen erfordern: 1. Berechnung der GLU-Aktivierung (z.B. `silu(x) * v`), die ein großes Zwischen tensor erzeugt. 2. Eine lineare Projektion (Matrizen multiplikation) dieses großen Tensors.
Durch das separate Schreiben und Lesen des großen Zwischen tensors in den globalen Speicher entsteht der Bandbreiten engpass. Twell fusioniert beide Schritte, berechnet das Endergebnis direkt, ohne das Zwischen tensor materialisieren zu müssen. Dies reduziert den Speicherdruck erheblich und macht den Prozess rechenbound statt speicherbound.
Einordnung und Bedeutung
Twell ist ein hervorragendes Beispiel für spezialisierte Kernel-Optimierung, die bestehende Architekturen effizienter macht, ohne sie zu ändern. Es adressiert ein klares Leistungsproblem in modernen LLMs. Vergleich mit anderen Ansätzen: Während Projekte wie Groq2 die Feedforward-Schichten durch Mixture of Experts (MoE) ersetzen wollen, damit die Architektur verändern, optimiert Twell die Berechnung innerhalb der Standard-FFN-Blöcke. Es ist somit ein komplementärer, nicht ein konkurrierender Ansatz. Für wen ist es relevant? Jeder, der große Transformer-Modelle trainiert oder deployed – von Forschungslaboren bis hin zu Unternehmen, betreibt Inference-Dienstleistungen. Die Leistungsgewinne bedeuten direkte Kosteneinsparungen bei den riesigen GPU-Rechenkosten. * Zukunft: Die Arbeit unterstreicht, Bedeutung der Hardware-nahen Optimierung im KI-Bereich. Solche Kernel könnten zukünftig direkt in Deep Learning-Frameworks wie PyTorch integriert werden und so die Effizienz standardmäßig für alle Nutzer verbessern.
Urteil: Ein kleiner Kernel mit großer Wirkung
Twell ist ein spezifisches, aber äußerst effektives Tool. Es zeigt, dass selbst in hochoptimierten Stacks wie PyTorch mit `torch.compile` und `vLLM` noch erhebliche Leistungsreserven durch maßgeschneiderte, hardwarebewusste Programmierung gehoben werden können. Für Teams, an der Grenze der Effizienz bei Training und Inferenz von LLMs arbeiten, ist Twell ein sofortiges und praktisches Upgrade. Die Kooperation zwischen einem innovativen KI-Lab (Sakana AI) und einem Hardware-Giganten (NVIDIA) hat hier ein Benchmark-Resultat hervorgebracht, von dem die gesamte Community profitieren kann.
