Что FlashAttention меняет в инференсе трансформеров
Автоматический перевод
Эта статья была автоматически переведена с оригинальной английской версии.
FlashAttention вычисляет точный аттеншн, не сохраняя полные матрицы оценок и вероятностей аттеншна в основной памяти GPU. Он обрабатывает небольшие блоки и объединяет их результаты с помощью онлайн-softmax. Это сокращает обмен данными с памятью и объём промежуточных данных, но сохраняет квадратичную зависимость числа арифметических операций плотного аттеншна от длины последовательности.
Снижение потребления памяти не делает рост вычислительных затрат плотного аттеншна линейным.
Разделяйте хранение и вычисления
Стандартный аттеншн с материализацией матриц вычисляет оценки для каждой видимой пары запрос/ключ, нормализует эти оценки и умножает их на значения. Квадратная матрица аттеншна без маски содержит N² элементов для N токенов.
Исходная статья о FlashAttention разбивает эти вычисления на блоки, которые помещаются в памяти на чипе. Алгоритм сохраняет обновляемую статистику нормализации и аккумуляторы выходных значений вместо записи всех вероятностей в HBM. При фиксированной размерности голов объём памяти для промежуточных результатов аттеншна растёт линейно с длиной последовательности. Алгоритм по-прежнему вычисляет плотные взаимодействия запрос/ключ, которые требует маска аттеншна.
Алгоритм вычисляет точный аттеншн; разный порядок операций с плавающей точкой может приводить к небольшим численным расхождениям. Это не приближённый метод разреженного аттеншна, который исключает отдельные пары токенов.
Реализация и нагрузка по-прежнему имеют значение
| Вопрос | Зачем это проверять |
|---|---|
| Поддерживает ли бэкенд эту GPU и эту точность? | Реализации ядер рассчитаны на определённое оборудование |
| Поддерживаются ли размерность голов, маска и структура кэша? | Неподдерживаемые формы могут привести к выбору другого бэкенда |
| Занимает ли аттеншн большую часть времени обработки запроса? | Большое ускорение ядра может мало повлиять на общее время |
| Нагрузка связана с префиллом или декодированием одного токена? | Объём параллельно выполняемой работы различается |
FlashAttention-2 дополнительно распараллеливает блоки последовательности запросов наряду с батчами и головами и сокращает другие накладные расходы выполнения. FlashAttention-3 использует возможности выполнения, характерные для Hopper. Названия версий описывают алгоритмы и реализации; они не определяют, какой бэкенд фактически выбрал ваш сервер.
Проверьте выбранный бэкенд в рантайме сервинга с зафиксированной версией. Измерьте латентность обработки промпта, пиковое потребление памяти и сквозной throughput при типичных длинах и размерах батча. Сравнивайте одинаковые маски и точность. Бенчмарк, который сообщает только TFLOPS ядра, не доказывает улучшения латентности токенов, которую наблюдает пользователь.
Инженерное руководство: FlashAttention рассматривает последовательные версии и их измерения с соответствующими ограничениями области применения.