FlashAttention — алгоритм вычисления внимания, который уменьшает затраты памяти и ускоряет работу трансформеров за счёт блочного вычисления.
Определение
FlashAttention — это оптимизированный алгоритм вычисления self-attention, разработанный для уменьшения использования памяти и ускорения инференса и обучения трансформеров. В традиционном attention вычисления производятся на полных матрицах Q, K, V и промежуточных тензорах, что требует хранения больших массивов в GPU-памяти. FlashAttention избегает этих затрат, вычисляя attention по блокам и не создавая промежуточные матрицы полностью.
Алгоритм стал стандартом в современных LLM благодаря способности ускорять обучение и генерацию при больших длинах контекста.
Как работает
FlashAttention использует чередование вычислений и чтения данных таким образом, чтобы избежать хранения полной матрицы внимания в GPU-памяти. Ключевые элементы алгоритма:
- блочная обработка — Q, K и V разбиваются на маленькие блоки, которые помещаются в SRAM (быструю память GPU);
- streaming softmax — softmax вычисляется инкрементально без хранения полного attention-score тензора;
- минимизация чтения/записи — данные читаются из HBM по одному разу, что уменьшает пропускную нагрузку;
- отсутствие промежуточных матриц — не создаются QKᵀ и softmax(QKᵀ) в полном размере;
- аккумуляция по блокам — результат attention формируется частями, без переполнения памяти.
Благодаря этому FlashAttention достигает почти оптимального использования GPU-памяти и снижает вычислительные задержки даже для больших контекстов (до десятков тысяч токенов).
Где применяется
- LLM с контекстами 4k–128k токенов.
- Трансформеры для обучения больших датасетов.
- Генерация в реальном времени.
- Мультимодальные модели с тяжёлым attention.
- Архитектуры, требующие минимального использования GPU-памяти.
- Модели, оптимизированные для inference-нагрузки.
Практические примеры использования
При обучении больших трансформеров FlashAttention снижает потребление памяти в разы, позволяя увеличивать batch size или длину контекста без увеличения VRAM.
В LLM-инференсе FlashAttention ускоряет генерацию, так как attention-операции выполняются быстрее и требуют меньше чтения данных из HBM.
В мультимодальных моделях с тяжёлыми attention-блоками алгоритм уменьшает задержку и делает вычисления предсказуемыми при больших токенных последовательностях.
В продуктах реального времени FlashAttention используется как основной оптимизационный слой, позволяя моделям работать быстрее без изменения архитектуры.
Преимущества и ограничения
- Плюс: значительное снижение памяти.
- Плюс: ускорение внимания в обучении и инференсе.
- Плюс: не требует изменения архитектуры трансформера.
- Плюс: оптимально использует SRAM и снижает нагрузку на HBM.
- Минус: сложность реализации на уровне ядра GPU.
- Минус: поддержка зависит от фреймворков и версий CUDA.
- Минус: требует переработки kernel-операций для модифицированного внимания.
- Минус: некоторые формы линейного внимания несовместимы с классическим FlashAttention.
Связанные термины
- Self-attention
- Long-context architectures
- FlashAttention-2
- Memory optimization
- KV-cache
- Tensor cores
- Block-sparse attention