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, що краще?», чесна відповідь така: тут не «чи», а «і».