Вышел 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