FlashAttention

FlashAttention — алгоритм, который вычисляет точное внимание трансформера кратно быстрее за счёт умного обращения с памятью GPU. Стандартное внимание материализует полную матрицу оценок N×N в медленной памяти HBM, и на длинных последовательностях узкое место — не арифметика, а перегонка этой матрицы туда-обратно. FlashAttention (Tri Dao и соавт., 2022) разбивает вычисление на плитки: блоки запросов и ключей остаются в быстрой набортной SRAM, softmax и матричные умножения сливаются в один проход, а при обратном распространении небольшие фрагменты пересчитываются вместо хранения. Итог — тот же математический результат при кратном ускорении и памяти, растущей линейно, а не квадратично с длиной последовательности; это один из ключей к сегодняшним контекстным окнам в 100K+ токенов. Преемники (FlashAttention-2 и 3) сильнее выжимают новое железо. Для разработчиков это инфраструктура, а не ручка API: используемые вами фреймворки обучения и серверы инференса включают его по умолчанию. Ключевое слово в этом описании — точное. FlashAttention не приближение, покупающее скорость ценой точности: это перестройка того же вычисления с учётом ввода-вывода, дающая математически идентичный наивной реализации результат внимания. Меняется лишь то, где живут промежуточные значения: разбивая запросы, ключи и значения на блоки, помещающиеся в набортной SRAM, и пересчитывая несколько дешёвых величин на обратном проходе вместо их хранения, алгоритм избегает многократной записи и чтения большой промежуточной матрицы оценок в память HBM — а на длинных последовательностях узкое место именно этот трафик памяти, а не арифметика. Поскольку экономия касается и обучения, и инференса, она проявляется как меньшее число GPU-часов на обоих концах жизненного цикла модели — снижение себестоимости, которое просачивается вниз по стеку в цену за токен, а не продаётся как отдельная функция. Стоит проговорить и другое: FlashAttention — не архитектура модели. Это оптимизация уровня вычислительного ядра, применимая к любому стандартному трансформеру; поэтому она подключается к готовым моделям без переобучения и стала стандартным компонентом продакшн-движков инференса и массовых фреймворков обучения, а не преимуществом, которым владеет один вендор. Для разработчиков последствия косвенные, но реальные: длинноконтекстные функции доступны по цене отчасти благодаря ей, и если вы размещаете модель сами, проверка того, что ваш серверный стек действительно использует ядро FlashAttention, а не откатился молча на наивный путь из-за неподдерживаемого поколения GPU или типа данных, — самая дешёвая из доступных проверок производительности, потому что откат функционально корректен и именно поэтому легко остаётся незамеченным. И ещё одна оговорка: квадратичная вычислительная стоимость внимания никуда не исчезает — дешевеет только трафик памяти. Длинный контекст остаётся дорогим; FlashAttention переводит его из невозможного в управляемое, но не делает бесплатным. Различие существенно, когда вы планируете бюджет цены и задержки для промптов на сотню тысяч токенов.