Вышел Flash Attention 4!
(На прошлой неделе)

На этот раз специально оптимизировали под новую архитектуру Blackwell (B200 и GB200), в отличие от FA3 для Hopper. На всяких 5090 дополнительного прироста нет.

Для BF16 даёт до 1.3x ускорения относительно cuDNN 9.13 и до 2.7x по сравнению с Triton! Выжимает солидные 1.6 PFLOPs/s (71% от теоретического пика B200). При этом в новых версиях cuDNN часть этих оптимизаций уже тоже появилась.

Среди ключевых трюков — software-эмуляция exp, conditional softmax rescaling, а в backward — использование tensor memory* и 2-CTA MMA**, что заметно снижает давление на shared memory. К этому бонусом идёт то, что весь kernel код теперь написан на Python (CuTe-DSL) без жестких C++ темплейтов, поэтому компиляция стала в 20-30 раз быстрее.

* Tensor memory — это новый сверхбыстрый on-chip буфер Blackwell рядом с tensor cores, куда можно складывать промежуточные результаты и меньше дёргать shared memory.
** 2-CTA MMA — режим, где одну matmul считает пара CTA (группа потоков), а не одна, что позволяет брать более крупные тайлы и заметно уменьшать shared-memory traffic → backward работает эффективнее.

Для задротов, очень рекомендую вот этот разбор на YouTube о том, как FA4 работает.

Пейпер
Код

@ai_newz