Что такое FlashAttention и какой выигрыш по памяти он даёт?

FlashAttention — это оптимизированный алгоритм вычисления механизма внимания (attention) в трансформерах, который значительно снижает потребление памяти и ускоряет вычисления по сравнению с классическим подходом.

Основная идея FlashAttention — вычислять attention без явного хранения больших матриц промежуточных результатов (например, матрицы весов внимания QK^T), используя блоковую обработку и эффективное использование кэш-памяти GPU. Это позволяет уменьшить пиковое потребление памяти примерно в 3-4 раза, что особенно важно при обучении больших языковых моделей (LLM) с длинными последовательностями.

Примерно, если классический attention требует памяти O(N^2) для последовательности длины N, FlashAttention снижает эту нагрузку за счёт вычислений на лету и оптимизаций, что позволяет запускать модели с большими контекстами на тех же GPU.

Таким образом, FlashAttention — это способ сделать attention более эффективным по памяти и скорости, что критично для масштабных моделей.