Как NVIDIA учит нейросети думать быстрее: оптимизация внимания для длинных контекстов

Инженеры NVIDIA опубликовали детальный разбор того, как архитектурные решения в механизме внимания трансформера влияют на скорость инференса. При контексте в 128 тысяч токенов внимание съедает до 85% времени — и это проблему нужно решать на уровне проектирования модели, а не только софтверной оптимизации.

Внимание — узкое место инференса. При коротком контексте оно отъедает 18% времени, при длинном — 85%. Исправлять это нужно не патчами, а на этапе проектирования архитектуры модели.

Проблема, о которой мало кто задумывается

Большие языковые модели стали работать с длинными контекстами — десятки и сотни тысяч токенов. Это хорошо для пользователя: можно закинуть в чат целый документ, кодовую базу или историю переписки. Но за этим удобством скрывается серьёзная техническая проблема.

Инженеры NVIDIA измерили, как распределяется время инференса у модели DeepSeek-R1 при разных длинах контекста. Результат наглядный:

• При 4K токенах механизм внимания (attention) забирает около 18% времени.
• При 32K токенах — уже заметно больше.
• При 128K токенах — 85% всего времени уходит именно на внимание.

Это значит, что все остальные компоненты модели — полносвязные слои, нормализация, активации — становятся почти бесплатными по сравнению с attention. И если вы хотите ускорить инференс длинных контекстов, оптимизировать нужно именно его.

Prefill и decode — это две разные задачи

Прежде чем углубляться в детали, важно понять базовое различие. Когда вы отправляете запрос языковой модели, происходит два этапа:

Prefill — модель «читает» весь ваш промпт целиком. Это параллельная обработка: все токены обрабатываются одновременно. Работа тяжёлая по вычислениям, но зато эффективно использует GPU — видеокарта занята матричными умножениями, и ей есть чем себя загрузить.

Decode — модель генерирует ответ по одному токену за шаг. Каждый новый токен нужно «сравнить» со всей историей разговора (KV-кэшем). Это совсем другая задача: вычислений мало, но нужно много раз прочитать данные из памяти. GPU в это время скорее простаивает, ожидая данные.

Именно поэтому параметры, которые ускоряют prefill, могут вообще не влиять на decode, и наоборот. NVIDIA подробно разобрала это на примере четырёх ключевых параметров.

Параметр первый: размер группы (group size)

Это, пожалуй, самый важный параметр для скорости генерации.

В стандартном механизме внимания (Multi-Head Attention, MHA) у каждого «головы» запроса (query head) есть своя собственная голова ключей и значений (KV head). Но это расточительно: KV-головы занимают много памяти, а при генерации каждого токена нужно перечитывать весь KV-кэш.

Grouped-Query Attention (GQA) решает эту проблему элегантно: несколько query-голов разделяют одну KV-голову. Число query-голов на одну KV-голову и называется размером группы (G). Если G = 1 — это классический MHA. Если G = 8 — восемь query-голов делят одну KV-голову. Если одна KV-голова на всех — это Multi-Query Attention (MQA).

Вот как это влияет на производительность:

При prefill — почти никак. Время остаётся практически неизменным при изменении G от 1 до 64. Почему? Потому что при обработке всего промпта параллельно GPU и так загружен вычислениями, и сокращение числа KV-голов не даёт видимого выигрыша.

При decode — критически важно. Удвоение G примерно удваивает «арифметическую интенсивность» операции (соотношение вычислений к объёму данных из памяти). Переход от G = 1 (MHA) к G = 8 даёт восьмикратный выигрыш. Причина проста: меньше KV-голов — меньше данных нужно прочитать из памяти на каждый генерируемый токен.

Это не маркетинговые цифры — инженеры NVIDIA подтвердили их экспериментально. На графике видно, что при генерации время падает примерно вдвое при каждом удвоении G. Правда, после G = 16 кривая начинает выравниваться — фиксированные накладные расходы начинают доминировать.

Вывод: выбирайте большой размер группы. Он не замедлит обработку промпта, но заметно ускорит генерацию.

Параметр второй: размерность головы (head dimension)

Размерность головы (Hsz) — это векторное пространство, в котором работает каждый «голова» внимания. Типичные значения: 64, 128, 256.

В отличие от размера группы, Hsz не влияет на арифметическую интенсивность — удвоение размерности удваивает и вычисления, и объём данных, и их соотношение остаётся тем же. Но на реальной производительности это всё равно сказывается, и причины — в особенностях работы GPU.

Видеокарта выполняет матричные умножения блоками (тайлами) фиксированного размера. Если ваша размерность не кратна этому размеру, часть тайла простаивает — но энергия и время на его обработку тратятся. Данные из памяти GPU тоже перемещаются блоками по 128 байт — и снова: если размерность не кратна 128, вы теряете эффективность.

Исследование показывает, что оптимальные значения — 128 или 256. Размерность 64 формально платит за тайл шириной 128 (то есть тратит ресурсы впустую). Размерность 512 и выше упирается в лимит тензорной памяти (TMEM) на чипе.

Есть ещё один нюанс: при prefill операция softmax не зависит от Hsz (она работает с матрицей скоров запросов к ключам, а не с размерностью головы). Поэтому увеличение Hsz сдвигает баланс в пользу матричных умножений, которые эффективно выполняются на тензорных ядрах. По сути, более широкая голова «прячет» softmax за математикой.

Вывод: используйте Hsz = 128 или 256. Это не хак, а инженерный компромисс между вычислительной эффективностью и ограничениями железа.

Параметр третий: длина последовательности

Здесь всё интуитивно, но важны масштабы.

Prefill растёт квадратично с длиной промпта. Каждый токен «смотрит» на каждый другой — это ISL² операций. Удвоение промпта — примерно четырёхкратный рост времени обработки. При коротких промптах рост ниже квадратичного, потому что фиксированные накладные расходы занимают заметную долю, но при длинных последовательностях квадратичная зависимость проявляется чётко.

Decode растёт линейно с длиной KV-кэша. Каждый шаг генерации читает весь накопленный кэш — но генерирует только один токен. Удвоение контекста — удвоение времени на каждый шаг.

Асимметрия важна: если вы можете каким-то образом сократить «эффективную» длину KV-кэша, выигрыш будет значительным. NVIDIA предлагает несколько подходов:

• Сжатие KV-кэша — хранить не все ключи и значения, а только важные.
• Разреженное или скользящее окно внимания — токен «смотрит» не на всю историю, а только на ближайшее окно.
• Гибридные архитектуры — например, модель Nemotron 3 от NVIDIA, где не все слои несут полный глобальный KV-кэш: одни слои работают с полным контекстом, другие — только с локальным окном.

Вывод: длину контекста задаёт задача, но можно сократить «стоимость» длинного контекста архитектурными средствами.

Параметр четвёртый: параллелизм и распределение по GPU

Когда модель слишком большая для одной видеокарты, её раскидывают по нескольким GPU с помощью тензорного параллелизма (Tensor Parallelism, TP). Каждая карта получает часть «голов» внимания.

Но есть подводный камень. В GQA несколько query-голов делят одну KV-голову — и эта связка должна целиком находиться на одном GPU. Если TP превышает число KV-голов (KH), одна KV-голова вынужденно дублируется на нескольких картах. Это пустая трата памяти и пропускной способности шины — данные копируются, но не дают нового вычислительного пара.

Пример: у Nemotron 3 всего две KV-головы. С TP = 2 каждая карта получает по одной KV-голове — ок. Но уже TP = 4 приведёт к дублированию. Как масштабироваться дальше?

Инженеры NVIDIA реализовали в фреймворке TensorRT-LLM альтернативные стратегии:

• Attention Data Parallelism (ADP) — запросы (а не головы) распределяются по картам.
• KV Parallelism (KVP) — длинный KV-кэш шардируется между картами.
• Wide EP — комбинация ADP для внимания и Expert Parallelism для полносвязных слоёв.
• Helix Parallelism — комбинация KVP для внимания и Expert Parallelism для MoE-слоёв.

Вывод: при проектировании модели учитывайте, сколько KV-голов она будет иметь. От этого зависит, как эффективно масштабировать инференс на кластере GPU.

Чек-лист для разработчиков

NVIDIA свела всё к четырём простым правилам:

1. Большой размер группы (G) — ускоряет генерацию, не влияет на обработку промпта. Спекулятивная генерация (speculative decoding) дополнительно повышает эффективность при заданном G.

2. Размерность головы 128 или 256 — совпадает с размером тайлов GPU и блоками передачи данных из памяти. Большая размерность также помогает скрыть накладные расходы softmax при prefill.

3. Сокращайте эффективный KV-кэш — через сжатие, разреженное внимание или гибридные архитектуры.

4. Согласуйте параллелизм с числом KV-голов — TP не должен превышать KH. Для моделей с малым числом KV-голов используйте ADP, KVP и их комбинации.

Насколько это важно для обычного пользователя?

Прямо сейчас — опосредованно. Эти рекомендации ориентированы на разработчиков и архитекторов моделей, а не на конечных пользователей. Но последствия ощутимы: модели, спроектированные с учётом этих принципов, будут быстрее отвечать при длинных контекстах, потреблять меньше памяти и лучше масштабироваться на серверном железе.

По сути, NVIDIA продвигает идею co-design — совместного проектирования модели и аппаратной платформы. Вместо того чтобы с