FlashAttention: быстрее не за счёт математики, а за счёт памяти
FlashAttention считает точный attention, но почти не пишет большую матрицу N×N в медленную память GPU.
Зачем это нужно
Авторы утверждают, что в современных GPU вычисления стали быстрее памяти, и большинство операций трансформера упирается в обращения к памяти. Многие приближённые методы уменьшали FLOPs, но не давали реального ускорения, потому что игнорировали доступ к памяти.
Как это работает
Алгоритм разбивает матрицы K и V на блоки и загружает их в быструю on-chip SRAM. Внутри блока он проходит по блокам Q и записывает результат. Так большая матрица N×N не записывается в HBM. Результат точный, а не приближённый. Авторы доказывают, что число обращений к HBM меньше, чем у стандартного attention, и показывают нижнюю оценку для точных алгоритмов.
Пример
По данным статьи, FlashAttention до 3 раз быстрее стандартной реализации на длинах от 128 до 2K токенов и масштабируется до 64K. На вычислении attention для GPT-2 авторы получают ускорение в 7.6 раза. Block-sparse вариант в 2–4 раза быстрее самого FlashAttention.
Что запомнить
- Точный attention можно ускорить, сократив обращения к памяти.
- Матрица N×N не записывается в HBM, вычисления идут блоками в SRAM.
- Ускорение касается реального времени, а не только числа FLOPs.
Где читать дальше
- Dao et al., FlashAttention (arXiv:2205.14135), разделы 1 и 2