Optimizing Jagged Flash Attention with TLX: The Road Toward SOTA FA4 on Blackwell – PyTorch

Kỹ sư Meta đã xây dựng kernel Jagged Flash Attention (JFA) cho mô hình Generative Ads Model (GEM) sử dụng TLX, một tiện low-level Triton hỗ trợ cấp phát SMEM/TMEM và specialization warp trên GPU Blackwell B200. TLX kernel có khoảng 3.2K dòng code - ít hơn đáng kể so với ~10K dòng của FlashAttention-4 (FA4) viết bằng CuteDSL - nhưng vẫn vượt trội hiệu suất khoảng 13% forward và 50% backward trên các jagged production shape. Các cải tiến cấu trúc bao gồm warp specialization, persistent kernels, quản lý bộ nhớ rõ ràng, cùng các tối ưu hóa cụ thể như jagged tile scheduling với Cluster Launch Control, multi-stage dQ staging, early tensor-memory release, loop peeling, và 2-CTA collaborative MMA. Bài viết cũng giới thiệu các biến thể MXFP8 và block-sparse được xây dựng lại từ cùng một kernel skeleton.