Хабр разобрал механизмы внимания в современных LLM
На Хабре вышла статья о механизмах внимания в современных LLM. Автор объясняет, как MHA, MQA, GQA и MLA влияют на размер KV cache и скорость декода.
- В статье разобраны четыре уровня системы: архитектура голов, паттерн внимания, kernel и позиционное кодирование
- MQA и GQA сокращают KV cache, но не уменьшают число query-голов и квадратичную сложность prefill
- При 32 слоях, 128K токенов, 8 KV-головах и head dimension 128 кэш в BF16 занимает около 16 GiB
На Хабре опубликована статья «Механизмы внимания в современных LLM». Автор разбирает, из чего складывается эффективность attention в больших языковых моделях и почему одна «новая формула» редко даёт выигрыш.
Четыре уровня системы
В разговорах об attention часто смешивают архитектуру, паттерн связей, позиционное кодирование и GPU kernel. Автор разделяет их: архитектура голов (MHA, MQA, GQA, MLA) определяет, как формируются и хранятся Q, K и V; паттерн внимания (full, sliding window, block sparse, learned top k) — какие прошлые токены доступны текущему; kernel (FlashAttention, cuDNN SDPA, SageAttention) — как та же математика выполняется на GPU; позиции (RoPE, ALiBi, NoPE, iRoPE) — как в Q и K попадает порядок токенов.
PagedAttention стоит рядом, но отвечает за другое: это менеджмент страниц KV cache в serving engine. Он уменьшает фрагментацию и позволяет делить блоки кэша между запросами, не меняя формулу attention.
Стоимость длинного контекста
В causal self attention каждый токен строит query и сравнивает его с keys всех доступных предыдущих токенов. Для последовательности длины n полное внимание создаёт n на n попарных взаимодействий, объём вычислений префилла растёт как O(n²). При авторегрессионном декоде новый токен сравнивается со всем накопленным контекстом, поэтому работа на шаг растёт примерно линейно с длиной истории.
Сам attention на декоде часто упирается не в FLOPS, а в чтение KV cache из памяти. Автор приводит оценку: при 32 слоях, 128K токенов, 8 KV-головах, head dimension 128 и BF16 кэш занимает около 16 GiB. Если заменить 8 KV-голов на 32, получится около 64 GiB. Поэтому число query heads само по себе почти ничего не говорит о стоимости serving — важен параметр num_key_value_heads.
MHA, MQA и GQA
Классический Multi Head Attention (MHA) даёт каждой голове собственные Q, K и V проекции. Это выразительная схема, но каждая KV-голова добавляет данные в кэш, что повышает memory bandwidth при декоде.
Multi Query Attention (MQA) оставляет много query heads, но заставляет их разделять одну пару K и V. Кэш сокращается примерно во столько раз, во сколько число query heads больше числа KV heads. Декод становится дешевле по bandwidth, но общая KV-проекция может ухудшать качество.
Grouped Query Attention (GQA) — компромисс: query heads делятся на группы, каждая группа использует свою KV-голову. Например, 32 query heads и 8 KV heads означают четыре query heads на одну пару K и V. GQA даёт качество, близкое к MHA, сохраняя большую часть выигрыша MQA по кэшу, поэтому стал распространённым выбором для современных LLM.
Автор подчёркивает: MQA и GQA уменьшают число KV-голов, но не число query-голов. На prefill каждая query-head всё ещё вычисляет attention scores для доступных позиций, и квадратичная зависимость full attention от длины сохраняется. Главный выигрыш — меньший KV cache и меньший объём данных, который читается из памяти GPU во время декода.
Для профиТехнические детали: архитектура, цифры, ссылки
Статья описывает четыре уровня системы attention:
- Архитектура голов — MHA, MQA, GQA, MLA. Влияет на качество, размер KV cache и bandwidth декода.
- Паттерн внимания — full, sliding window, block sparse, learned top k. Влияет на стоимость префилла и длинного декода.
- Kernel — FlashAttention, cuDNN SDPA, SageAttention. Влияет на фактическую скорость, HBM traffic и точность формата.
- Позиции — RoPE, ALiBi, NoPE, iRoPE. Влияет на экстраполяцию и поведение на длинном контексте.
Оценка KV cache для одного запроса: первый множитель 2 — отдельные K и V. При 32 слоях, 128K токенов, 8 KV-головах, head dimension 128 и BF16 кэш занимает около 16 GiB; при 32 KV-головах — около 64 GiB.
Автор отмечает, что расчёты даны для BF16 (2 байта на элемент) как базовая оценка, но современный инференс может использовать FP8, FP4, INT8, INT4 и смешанные форматы. Поэтому цифры стоит читать как оценку порядка величины, а не как фиксированную конфигурацию.
Вопросы и ответы
- Чем GQA отличается от MHA?
- В MHA у каждой головы свои Q, K и V. В GQA query heads делятся на группы, и каждая группа использует свою KV-голову. Это сокращает KV cache при качестве, близком к MHA.
- Почему декод упирается в память, а не в вычисления?
- На декоде новый токен сравнивается со всем накопленным контекстом, и основное время уходит на чтение KV cache из памяти GPU, а не на FLOPS.
- Сокращают ли MQA и GQA квадратичную сложность attention?
- Нет. Они уменьшают число KV-голов и размер кэша, но на prefill каждая query-head всё ещё вычисляет attention scores, и зависимость O(n²) для full attention сохраняется.
