Metas TLX-basierte Jagged Flash Attention schlägt FA4 auf Blackwell mit dreimal weniger Code

Warum die Attention in GEM neu geschrieben werden musste

Metas Werbemodelle, darunter das Generative Ads Model (GEM) und die Kunlun-Architektur, berechnen Attention über jagged User-Sequenzen, also über Eingaben variabler Länge. Die Autoren geben an, dass das Auffüllen (Padding) solcher Sequenzen auf eine feste Länge bis zu 50 Prozent der Rechenleistung verschwenden kann. GEM packt die Sequenzen stattdessen lückenlos hintereinander und hält ihre Grenzen in einem Offsets-Tensor fest. Jagged Flash Attention wendet den FlashAttention-Algorithmus direkt auf diese gepackten Q/K/V-Tensoren an, ohne aufgefüllte Tokens zu materialisieren.

Dem Beitrag zufolge ist die Attention der langsamste einzelne Kernel in GEM. Maximale Leistung auf Blackwell erforderte bislang handgeschriebenes CuteDSL oder CUDA, was nach Angaben der Autoren die Entwicklung verlangsamt und die Erweiterung auf neue Varianten wie Sliding Window oder Block-Sparse Attention erschwert. Ausgangspunkt war ein ausgereifter Triton-basierter JFA-Kernel, der Datenbewegung und Scheduling auf dem Chip dem Compiler überließ. Dadurch standen die Tensor Cores still, sobald die einzige Warp Group keine Matrixmultiplikation anstieß.

Was TLX ändert

TLX ergänzt Tritons kachelbasiertes Programmiermodell um explizite, hardwarebewusste Kontrolle. Es ermöglicht die explizite Belegung von Shared Memory und Tensor Memory, Warp Specialization über asynchrone Tasks, Barrieren, asynchrones TMA und MMA sowie Cluster Launch Control (CLC).

Mit dieser Kontrolle teilte das Team jede CTA in rollenspezialisierte Warps auf: für TMA-Loads, Tensor-Core-Matmuls, Softmax- und Korrekturrechnungen sowie das Schreiben im Epilog, im Backward-Pass zusätzlich ein eigener Warp für die dQ-Reduktion. Die Puffer auf dem Chip werden von Hand zugewiesen, auch die Pipeline-Tiefe wird gewählt, etwa Triple-Buffering für K/V. Außerdem überlagern die Entwickler Tensor-Memory-Puffer mit nicht überlappenden Lebensdauern per Aliasing und setzen explizite Producer/Consumer-Barrieren ein. Beide Pässe sind persistent: Je eine CTA pro SM durchläuft die Kacheln in einer Schleife. Nach Angaben der Autoren lässt sich nichts davon in reinem Triton ausdrücken.

Beim Entwicklungsaufwand umfasst der TLX-Kernel laut Beitrag etwa 3,2K Zeilen Code auf Triton-Ebene und damit rund dreimal weniger als die etwa 10K Zeilen der CuteDSL-Kernel von FA4.

Zusätzliche Optimierungen

Das Team fand seine Optimierungen mit NVIDIA Nsight Compute, ptxas-Dump-Logs und TritonBench. Einige stammen aus FA4, andere leiteten die Entwickler aus dem Profiling der Jagged-Broadcast-Q-Last ab. Dabei wird ein dichtes Q von allen Sequenzen im Batch gemeinsam genutzt, und sein Gradient dQ muss über den gesamten Batch summiert werden. Aufgeführt sind folgende Optimierungen:

• Scheduling jagged Kacheln über SMs hinweg: Software-Lastverteilung und Cluster Launch Control halten die SMs trotz stark schwankender Längen ausgelastet.
• Mehrstufiges dQ-Staging: Eine doppelt gepufferte Shared-Memory-Pipeline verbirgt die umkämpfte dQ-Reduce-Add-Operation, die die Autoren als Engpass Nummer eins im Backward-Pass bezeichnen.
• Frühe Freigabe von Tensor Memory: Wird der dQ-Puffer vor den letzten Schreibvorgängen freigegeben, kann der MMA-Warp die dQ-Matmul der nächsten Kachel früher starten.
• Loop Peeling: Die Aufteilung der KV-Schleife in einen verzweigungsfreien Hauptdurchlauf und einen kleinen maskierten Rest beseitigt den Masken-Overhead pro Iteration und die dadurch verursachten Register Spills.
• 2-CTA-kollaboratives MMA: Zwei CTAs arbeiten an einer Matmul zusammen, um die Tensor-Core-Auslastung im Backward-Pass zu steigern; dieser Ansatz stammt aus FA4.

Für den Forward-Pass sortiert das Team die Kacheln auf dem Host absteigend nach KV-Workload und verteilt sie im Zickzackmuster an die SMs. Laut Beitrag brachte das bei vernachlässigbaren CPU-Kosten rund 20 Prozent mehr Leistung im Forward-Kernel.

Was das für die Kernel-Entwicklung bedeutet

Die Autoren argumentieren, dass der Verbleib im High-Level-Modell von Triton den Kernel für Modellierungsingenieure zugänglich macht, nicht nur für Kernel-Spezialisten: Sie können ihn lesen, erweitern und fusionieren. Alle Benchmarks liefen in bfloat16 auf der B200 gegen FA4, das der Beitrag als aktuellen Stand der Technik auf Blackwell bezeichnet. Der Code ist im GitHub-Repository ads_model_kernel_library von Meta unter tlx_jfa verfügbar.