FlashAttention-3 📸

Блог и статья: tridao.me/blog/2024/flash3/
Код: github.com/Dao-AILab/flash-attention

1.5-2.0 раза быстрее чем FlashAttention 2 в FP16, почти в 3 раза быстрее в FP8

Flash 2 был оптимизирован под А100 и достигал 70% теоретических max flops, Flash 3 оптимизирован под H100 и достигает 75%. Для этого активно использовали библиотеку NVIDIA CUTLASS которая почти напрямую общается с железом и позволяет использовать хитрые фичи H100. Например использовать специальное железо (Tensor Memory Accelerator) для копирования тензоров между global и shared memory и Warpgroup Matrix Multiply-Accumulate (который я так и не понял что такое, но как-то связан с тензорными ядрами)

Сверху этого переписали код чтобы матричное умножение вычислялось (на tensor cores) параллельно с softmax (на multi-function unit) и сделали пару хаков чтобы FP8 давал более точные результаты

Выглядит очень классно, надо пробовать. Жаль что еще нету edge железа на Hopper