FlashAttention и PagedAttention: два решения для двух разных узких мест

Кто хоть немного работает с LLM, слышит о FlashAttention и PagedAttention так, будто это соперники, а вопрос «что из них выбрать?» постоянно всплывает в ML-сообществах. Сама постановка неверна. Они решают две не связанные между собой задачи на двух разных этапах жизни модели, и серьёзный стек для сервинга часто использует обе одновременно.

Почему внимание обходится дорого

Self-attention в трансформере вычисляет вот что:

Attention(Q, K, V) = softmax(QKᵀ / √d) · V

Формула выглядит безобидно, но за ней скрываются два отдельных узких места, и какое из них важнее, зависит от того, когда вы её считаете.

1. При вычислении (обучение и prefill) матрица внимания растёт как O(N²) от длины последовательности N. Документ из 10 000 токенов даёт матрицу на 100 миллионов ячеек, которую снова и снова вычисляют и гоняют по памяти GPU.
2. При генерации (инференс и декодирование) для каждого нового токена модель должна помнить векторы Key и Value всех предыдущих токенов. Это хранилище называется KV-кэшем, и оно растёт с каждым токеном и каждым новым запросом пользователя.

Удобная аналогия — ресторан. Первая проблема — насколько эффективно повар готовит одно блюдо, то есть вычисления. Вторая — насколько хорошо ресторан распоряжается складом, когда в работе сотни заказов, то есть управление памятью. FlashAttention ускоряет вычисления, а PagedAttention делает расход памяти при сервинге эффективным. Держите это разделение в голове, всё остальное вытекает из него.

FlashAttention: быстрее считать внимание

Проблема

У GPU два вида памяти. HBM (high bandwidth memory) большая, десятки гигабайт, но сравнительно медленная. SRAM, память на кристалле, вмещает всего несколько мегабайт, зато очень быстрая.

Стандартная реализация строит всю матрицу внимания N на N и записывает её в HBM, читает обратно, чтобы посчитать softmax в SRAM, снова записывает результат, а затем ещё раз читает для финального умножения на V. Настоящее узкое место — вся эта беготня между медленной и быстрой памятью. Сама арифметика им не является.

Представьте повара на маленькой, но быстрой кухне (SRAM), у которого все продукты лежат на огромном складе (HBM) в полутора километрах. Стандартное внимание — это когда повар едет на склад, берёт немного, возвращается, немного готовит, снова едет отвезти остатки и так на каждом шаге. Дорога отнимает больше времени, чем готовка. FlashAttention — это одна хорошо спланированная поездка: привезти партию подходящего размера, сделать на кухне все шаги готовки и вернуться на склад, только когда блюдо готово.

Как это работает

FlashAttention не меняет того, что вычисляет внимание. Он меняет то, как это делается. Основные приёмы:

• Тайлинг разбивает Q, K и V на небольшие блоки, которые помещаются в SRAM.
• Онлайн-softmax избавляет от необходимости держать в памяти целую строку. Он хранит текущий максимум и текущую сумму и обновляет их с приходом каждого нового блока, так что полная матрица N на N никогда не материализуется.
• Слитые ядра (fused kernels) выполняют матричное умножение, softmax и второе умножение внутри одного ядра GPU, без промежуточных записей в HBM.
• Пересчёт в обратном проходе означает, что огромная матрица внимания не хранится ради градиентов. Её заново строят на лету во время обратного распространения, тратя немного лишних вычислений, чтобы сильно сэкономить память.

Результат точно такой же, как у стандартного внимания. FlashAttention — не приближение, а просто более умный способ прийти к тому же ответу.

В коде

На поддерживаемых GPU, таких как A100 и H100, PyTorch сам переключается на ядро в стиле FlashAttention:

import torch
import torch.nn.functional as F

# PyTorch automatically dispatches to a FlashAttention-style
# kernel on supported GPUs (A100, H100, etc.)
q = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)  # (batch, heads, seq_len, head_dim)
k = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)


output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(output.shape)  # torch.Size([2, 8, 1024, 64])

Можно также напрямую вызвать специализированную библиотеку flash-attn:

from flash_attn import flash_attn_func
import torch

q = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)
k = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)

output = flash_attn_func(q, k, v, causal=True)

Когда он нужен

• Обучение больших трансформеров с нуля или их дообучение.
• Длинные контекстные окна на 8k, 32k или 128k токенов, где рост N² действительно бьёт больно.
• Любой случай, при обучении или инференсе, где важна чистая скорость внимания.

PagedAttention: эффективная память для KV-кэша

Проблема

Это чисто проблема сервинга. При авторегрессивном декодировании модель кэширует тензоры Key и Value всех предыдущих токенов, чтобы не пересчитывать их, — это и есть KV-кэш. Традиционные системы сервинга упираются здесь в трудности:

• Длина ответа заранее неизвестна, поэтому они выделяют на каждый запрос один большой непрерывный кусок, например под max_tokens = 2048, даже если ответ в итоге займёт 50 токенов.
• Это вызывает сильную внутреннюю фрагментацию. До 60–80% выделенного KV-кэша может простаивать пустым и пропадать зря.
• Ничего не разделяется. При beam search или параллельном сэмплировании, когда из одного промпта растут несколько ответов, каждая ветка получает полную копию KV-кэша, хотя все они начинаются с одного и того же текста.

Представьте многоквартирный дом (память GPU) и жильцов (запросы). При старом подходе каждый новый жилец заранее бронирует целый этаж, на всякий случай. Большинство пользуется двумя-тремя комнатами из двадцати, а пустые больше никому трогать нельзя. PagedAttention выдаёт отдельные комнаты, блоки или страницы фиксированного размера, там, где есть место, и только тогда, когда они нужны. Справочник, таблица блоков, записывает, какие комнаты кому принадлежат. Родственные жильцы, например ветки beam search с общим промптом, могут даже делить комнаты под общие части. Идея прямо позаимствована у страничной виртуальной памяти, которую операционные системы используют для процессов, отсюда и название.

Как это работает

1. KV-кэш режется на блоки фиксированного размера, скажем по 16 токенов, вместо одного гигантского непрерывного массива.
2. Каждая последовательность ведёт таблицу блоков, которая сопоставляет логические индексы блоков физическим блокам памяти, как таблица страниц в ОС.
3. Блоки выделяются по требованию, только когда последовательность выросла настолько, что ей нужен новый, поэтому расточительного предварительного выделения нет.
4. Разделение работает через copy-on-write. При beam search, параллельном сэмплировании или общем системном промпте последовательности указывают на одни и те же физические блоки для общего префикса и копируют их, только когда содержимое расходится.
5. Специальное ядро CUDA читает разбросанные блоки через таблицу блоков. Поскольку блоки больше не лежат подряд, именно это специализированное ядро и есть собственно операция PagedAttention.

В коде

На этом работает vLLM, один из самых популярных на сегодня движков для сервинга LLM:

from vllm import LLM, SamplingParams

# vLLM automatically manages KV cache using PagedAttention under the hood
llm = LLM(model="meta-llama/Llama-2-7b-hf")
prompts = [
    "The future of AI is",
    "PagedAttention helps LLM serving by"
]

sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=100)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.outputs[0].text)

Как пользователь вы никогда не видите PagedAttention напрямую. Он тихо работает в фоне, чтобы память GPU хорошо использовалась для всех параллельных запросов. Упрощённый концептуальный взгляд на то, что может хранить таблица блоков, выглядит так, и это не реальное внутреннее устройство vLLM:

# Simplified conceptual example (not actual vLLM internals)
block_table = {
    "request_1": [block_3, block_7, block_12],
    "request_2": [block_3, block_7, block_9],
    # request_1 & request_2 share blocks 3 & 7 -> same prompt prefix!
}

Когда он нужен

• Продакшен-сервинг для множества одновременных пользователей.
• API чат-ботов, RAG-пайплайны и любые системы с непредсказуемой длиной ответов.
• Beam search или параллельное сэмплирование, где общая память даёт большую экономию.
• Любая конфигурация, где нужно выжать максимум пропускной способности из GPU: больше запросов на GPU — ниже стоимость одного запроса.

Вместе

Эти два подхода действительно сочетаются, и на практике обычно так и бывает. PagedAttention — слой управления: он решает, где и как хранятся блоки KV-кэша. FlashAttention или слитые ядра в его духе — вычислительный слой: он определяет, насколько быстро работают произведение Q на транспонированную K и softmax. vLLM сочетает PagedAttention для памяти с оптимизированными слитыми ядрами внимания, вдохновлёнными FlashAttention, для математики, и получает сразу и эффективность памяти, и скорость вычислений.

Если вернуться к кухне, PagedAttention выбирает, с какой полки склада брать продукты, а FlashAttention следит, чтобы повар эффективно готовил, когда они приехали.

Что нужно именно вам?

• Обучаете модель или дообучаете её на длинных последовательностях: FlashAttention.
• Отдаёте API LLM сотням или тысячам одновременных пользователей: PagedAttention, через vLLM или что-то подобное.
• Строите стек сервинга продакшен-уровня: оба. Большинство современных фреймворков, включая vLLM, TensorRT-LLM и Hugging Face TGI, уже сочетают похожие идеи внутри.

Коротко

Общее слово «Attention» в названиях — единственная причина, по которой их валят в одну кучу. Они находятся на разных слоях стека. FlashAttention делает вычисление внимания быстрее и экономнее по памяти, управляя тем, как данные перемещаются между HBM и SRAM. PagedAttention делает хранение KV-кэша при инференсе эффективным и свободным от фрагментации, заимствуя у операционных систем страничную виртуальную память.

Один оптимизирует вычисления, другой — управление памятью при сервинге. Большинство продакшен-систем сервинга сегодня, включая vLLM, используют оба, чтобы модели работали быстро и масштабировались дёшево. Так что когда кто-то спрашивает «PagedAttention или FlashAttention, что лучше?», честный ответ такой: здесь не «или», а «и».