Что такое FlashAttention и PagedAttention: ускорение инференса LLM в бэкендах

Современные большие языковые модели покоряют мир, но их эффективное развертывание упирается в фундаментальные ограничения — вычислительных возможностей систем далеко не всегда хватает при инференсе через ИИ-движки, а расход памяти при “сыром” развертывании может быть чудовищным. Долгое время эти проблемы были настоящим камнем преткновения, который делал инференс LLM менее доступным, чем он является на самом деле. Однако впоследствии индустрия все же смогла сдвинуть этот камень благодаря появлению двух революционных технологий — FlashAttention и PagedAttention. В этой статье мы расскажем вам, что из себя представляют эти передовые алгоритмы, зачем они нужны и как они смогли стать стандартом в области локального развертывания ИИ-моделей.

Как известно, трансформерная архитектура — это сердце любой LLM. И одним из важнейших компонентов этой архитектуры является так называемых механизм внимания, наделяющий языковые модели способностью анализировать контекст пользовательского запроса. Выход архитектуры трансформеров в 2017 году из-под крыла компании Google сделал возможным переход от рекуррентных нейронных сетей и сетей с долгой краткосрочной памятью к большим языковым моделям, а впоследствии и к агентским нейросетям. Несмотря на то, что трансформеры в свое время произвели революцию в мире ИИ, эта архитектура все-таки не идеальна, поскольку у нее есть две ключевые проблемы:

• Первая проблема — вычислительная сложность. Стандартный механизм внимания имеет квадратичную сложность O(n²) (время выполнения растёт пропорционально квадрату размера входных данных) относительно длины контекста. Это означает, что для обработки последовательности в 1000 токенов требуется порядка миллиона операций, а для 2000 токенов — уже четыре миллиона. Параллельно с вычислениями генерируется огромная промежуточная матрица внимания, которую необходимо хранить в высокоуровневой памяти GPU, что создает потребность в высоком объеме VRAM и пропускной способности, которые GPU далеко не всегда могут обеспечить в должной мере.
• Вторая проблема — управленческая. Для ускорения генерации токенов модель сохраняет ключи и значения (KV-кэш) уже обработанных токенов. При одновременной обработке множества запросов с переменной длиной контекста память под KV-кэш резервируется с запасом, что приводит к сильной фрагментации и неэффективному использованию дорогостоящей памяти GPU — память превращается в нечто вроде швейцарского сыра с множеством пустот, увеличивая задержку при передаче ИИ-данных.

Очевидно, что для масштабирования инференса требуются оптимизации на двух фронтах: алгоритма вычислений и системы управления памятью. В противном случае, большие языковые модели будут использовать ресурсы ИИ-ускорителей крайне неэффективно, требуя больше памяти, чем нужно для стабильной работы. 

Представленный в 2022 году алгоритм FlashAttention устраняет первую проблему. Вся его эффективность заключается в полном пересмотре механизма записи и чтения матрицы внимания: он просто перестает выполнять эти задачи в памяти GPU. Чтобы выполнить это, FlashAttention разбивает вычислительные операции на блоки. Входные матрицы Q, K и V разделяются на небольшие блоки, которые помещаются в сверхбыструю SRAM-память (кэш L3 видеокарт) на чипе GPU. Затем он применяет технику kernel fusion, при которой весь процесс вычисления внимания для блока выполняется за один проход непосредственно в SRAM, что устраняет необходимость промежуточного сохранения контекста в VRAM.

Благодаря этому, FlashAttention радикально сокращает количество операций с памятью, что приводит к ускорению вычислений в 2-3 раза на длинных контекстах, при этом значительно экономя память. Эти значительные преимущества при локальном запуске ИИ обеспечили FlashAttention прямой билет в интеграцию таких популярных движков ИИ-инференса коммерческого уровня, как vLLM и Triton Inference Server. Важно отметить, что все это происходит только в рамках одной последовательности контекста (запроса), а не длинных последовательностей, но и на них найдется управа в лице PagedAttention.

Технология PagedAttention, появившаяся в 2022 году, призвана решить вторую проблему при запуске LLM — неэффективное управление памятью KV-кэша. PagedAttention позволяет хранить KV-кэш для десятков и сотен запросов не в непрерывных, а в разрозненных блоках памяти GPU. При этом, процесс распределения контролируется не самим алгоритмом, а движком инференса ИИ, в результате чего память выделяется динамически, по мере необходимости, только под реально используемые токены. Это полностью устраняет проблему фрагментации и необходимость резервирования памяти под максимально возможную длину контекста для каждого запроса. 

Можно было бы подумать, что алгоритмы FlashAttention и PagedAttention конкурируют друг с другом, ведь они выполняют практически одну и ту же задачу, хоть и разными способами. На деле все наоборот: FlashAttention и PagedAttention идеально дополняют друг друга, решая разные сложности инференса моделей-трансформеров. Именно их комбинация в современных движках локального вывода ИИ дает максимальный прирост производительности. 

Это работает так: FlashAttention эффективно вычисляет механизм внимания для одного запроса, минимизируя затраты времени и памяти на операциях, а PagedAttention эффективно хранит и управляет данными KV-кэша для тысяч таких запросов, максимально используя доступный объем памяти GPU. Идеальным примером их синергии является движок vLLM. Под капотом он использует PagedAttention для управления виртуальной памятью KV-кэша всех поступающих запросов. Когда же требуется непосредственно вычислить шаг внимания для конкретного запроса, он задействует оптимизированные ядра FlashAttention, что делает каждый отдельный вычислительный акт максимально быстрым и экономичным.

На сегодняшний день FlashAttention и PagedAttention стали стандартом де-факто для высокопроизводительного инференса. Они интегрированы во все ведущие фреймворки и платформы локального запуска ИИ. vLLM изначально построен вокруг связки PagedAttention и FlashAttention. PyTorch и компилятор Triton предоставляют нативную поддержку ядер FlashAttention для реализации пользовательских моделей. Промышленные движки, такие как TensorRT-LLM от NVIDIA и Text Generation Inference (TGI) от Hugging Face, также активно интегрируют обе технологии для достижения рекордных показателей скорости и пропускной способности.

Комбинированный эффект от их использования впечатляет. FlashAttention сам по себе ускоряет вычисления внимания на длинных контекстах в 2-3 раза. PagedAttention изолированно может увеличить количество обрабатываемых запросов в секунду также в несколько раз, устранив потери на фрагментацию. Вместе они обеспечивают высокую пропускную способность и низкую задержку обращения к памяти даже в условиях интенсивной нагрузки с длинными контекстами и большим количеством пользователей. Это напрямую влияет на экономическую эффективность использования железа для инференса: один сервер может обслуживать больше запросов, снижая стоимость локального вывода на один запрос. В противном случае, каждый сервер будет тратить неприлично много ресурсов GPU даже для одного запроса, что свело бы на нет все преимущества работы с локальными LLM. Именно по этой причине FlashAttention и PagedAttention применяются только в индустриальных движках инференса — в пользовательских движках, вроде Ollama, LM Studio и llama.cpp, подразумевается работа только с одним юзером, поэтому алгоритмы оптимизации в них просто не нужны.

Технологиям FlashAttention и PagedAttention с самого их зарождения был предписан большой успех, но у этих алгоритмов есть ряд ограничений. Оба решения заточены исключительно под работу с GPU и в особенности с экосистемой CUDA, что создает зависимость от конкретного типа железа. Для пользователей open-source сообщества, экспериментирующих с инференсом на CPU, это стало большой проблемой, ведь они буквально лишены возмозможности оптимизировать инференс из-за философии FlashAttention и PagedAttention. 

Тем не менее, история развития этих алгоритмов все еще продолжается, ведь разработчики постоянно расширяют их функциональные возможности и оптимизируют производительность. Уже выпущены более эффективные FlashAttention-2 и FlashAttention-3, где были внедрены такие улучшения, как, оптимизированная параллелизация и работа с разделяемой памятью, новые архитектурные улучшения, поддержка новых функций и многое другое. Кроме того, начиная с FlashAttention-2 технология получила полную совместимость центральными процессорами Intel и AMD, что открыло дорогу к производительному инференсу на мощных CPU.

PagedAttention тоже не уступает своему старшему брату, получая новые обновления до более совершенных версий, совместимых с новыми поколениями GPU. В vLLM уже доступны PagedAttention V1 и V2, которые являются отдельными алгоритмами, предназначенными для выполнения разных типов задач. V1 используется при длине последовательности до 8192 токенов, оптимизируя параллелизацию большого количества однотипных запросов. V2 задействуется для более длинных контекстов и в сценариях, где V1 может столкнуться с нехваткой разделяемой памяти. В Roadmap движка vLLM также упоминалось, что технология PagedAttention получит грандиозное обновление в виде поддержки реализации управления KV-кэшем на центральных процессорах, что наконец-то позволит энтузиастам оптимизировать инференс ИИ на CPU.

Если вы хотите стать настоящим ИИ-гиком, то понимание современного стека технологий инференса LLM и, в особенности, алгоритмов PagedAttention и FlashAttention, заложит мощный фундамент для получения новых знаний, ведь именно симбиоз этих систем лежит в основе производительности и рентабельности современных сервисов на базе искусственного интеллекта. А следующим шагом познания инференса LLM станет оптимизация самих весов моделей, о чем мы расскажем вам в нашей следующей статье.