CODA: Wie Transformer-Blöcke als GEMM-Epilogue-Programme neu geschrieben werden
CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs
Wir stellen CODA vor, eine GPU-Kernel-Abstraktion, die viele Transformer-Operationen als GEMM-plus-epilogue-Programme ausdrückt. Dies ermöglicht es, Berechnungen wie Normalisierung und Aktivierungen auszuführen, während die GEMM-Ausgabetile noch auf dem Chip verbleiben, bevor sie in den Speicher geschrieben werden. Unser Ansatz kombiniert die Produktivität von Frameworks mit der Effizienz von Hardware und zeigt hohe Leistung bei sowohl menschlich als auch von LLMs verfassten Kernen.
Die Datenbewegung wird zu einem zunehmend wichtigen Flaschenhals in ansonsten hoch optimierten Trainingsstapeln, da große Zwischentensoren wiederholt durch den globalen Speicher bewegt werden, während kaum Arithmetik ausgeführt wird.
- rahen
Streng genommen ist das sehr domänenspezifisch und ermöglicht keine Performance, die Triton nicht bereits erreichen könnte (das Eliminieren von Global-Memory-Roundtrips durch Epilogue-Fusion ist nichts Neues). Die eigentliche Erkenntnis ist der Designwechsel für LLM-getriebene Codegenerierung gegenüber handgefertigten Kernen.
LLMs sind nach wie vor schlecht bei Low-Level-Hardware-Optimierungen, aber wirklich gut bei der High-Level-Zusammensetzung. Compiler-Abstraktionen mit einer eingeschränkten, komponierbaren API zu entwerfen, damit ein LLM leicht expertengeschriebene Blöcke zusammenfügen kann, ist ein kluger Schachzug. Ich vermute, dass dies sich schließlich zum Standard für Codegeneratoren entwickeln wird, wenn wir zu agentenbasierter Entwicklung übergehen.
- augment_me
TLDR:
Die Autoren erkennen, dass global zeilenweise abhängige Funktionen wie RMSNorm/LayerNorm eingebaute Skalierungsfaktoren haben, die in bestimmten Konfigurationen kommutativ sind, sodass sie nach einer nachfolgenden Projektion herausgezogen und teilweise auf Zeilentiles aggregiert werden können.
Also kann ((W1 @ gamma * globally_computed_scale) * W2 als (W1 @ gamma * W2) * globally_computed_scale geschrieben werden, solange wir nur zeilenweise Interaktionen für den Skalierungsfaktor haben.
Das wurde bisher meist nicht gemacht, weil von links nach rechts arbeitende Graph-Compiler wie torch.compile nicht davon ausgehen können, dass eine globale zeilenweise Reduktion zwischen GEMMs kommutativ ist.
- saagarjha
Leute, die nur CUTLASS GEMM Epilogue-Fusions geschrieben haben und jetzt ihren zweiten Kernel sehen: Das hier hat echt viel 'GEMM Epilogue Fusion'-Vibe.
- maxignol
« LLMs können erfolgreich CODA-Kerne verfassen »
Das könnte den Fortschritt in diesem Bereich dann beschleunigen.
- cold_harbor
Synthese allein ist der schwierige Teil. Mit Feedback aus der Ausführung — ausführen, profilieren, patchen — schließt sich die Lücke schnell. Es ist im Grunde ein RL-Problem im Verborgenen.