В эпоху, когда агентные системы и работа с длинным контекстом становятся стандартом, длина контекста моделей стремительно растет. Это приводит к тому, что операция внимания (attention) потребляет все большую долю времени инференса. Как показывают данные, доля времени, затрачиваемого на внимание, может значительно возрастать при увеличении контекста. В таких условиях то, как спроектирована архитектура модели — не просто как она реализована программно, а как она взаимодействует с аппаратной частью — становится решающим фактором производительности.
Ключевым подходом здесь является «кодизайн» (co-design) ИИ-моделей: проектирование архитектуры модели с учетом того, как именно GPU выполняет вычисления. В этой статье мы подробно разберем, как размер группы (group size), размерность заголовка (head dimension) и длина последовательности влияют на производительность плотного внимания. Мы также рассмотрим, как правильно настраивать параллелизм для NVIDIA GPU, чтобы избежать узких мест и максимизировать пропускную способность.
01Две стороны одной медали: Prefill и Decode
Чтобы понять, как оптимизировать модель, нужно сначала осознать фундаментальное различие между двумя фазами инференса: prefill (заполнение) и decode (генерация). Это две совершенно разные задачи с разными узкими местами.
Prefill обрабатывает весь входной промпт параллельно. Это порождает большие матричные умножения (GEMM), где размер матрицы M равен длине входной последовательности (ISL). Поскольку объем вычислений огромен, эта фаза является вычислительно ограниченной (compute-bound). GPU загружен полностью, ожидая завершения операций умножения и сложения.
Decode, напротив, генерирует один токен за раз. Размер матрицы M здесь мал. Основная нагрузка здесь ложится не на вычисления, а на чтение данных из высокоскоростной памяти (HBM). GPU постоянно ждет, пока данные из KV-кэша будут доставлены из памяти. Эта фаза является ограниченной пропускной способностью памяти (memory-bound).
Важно отметить, что такие техники, как спекулятивный декодинг (speculative decoding), могут увеличивать эффективный размер GEMM-M на этапе decode, сдвигая его ближе к вычислительной границе. Однако в стандартном сценарии различие между этими фазами критично для оптимизации.
Роль арифметической интенсивности и модели Roofline
Производительность GPU ограничена двумя факторами: пиковой вычислительной мощностью и пропускной способностью памяти. Переход между этими режимами определяется арифметической интенсивностью — отношением общего количества операций с плавающей точкой (FLOPs) к общему объему переданных байтов.
Существует так называемая «гребневая точка» (ridge point). Если арифметическая интенсивность модели выше этой точки, она находится в вычислительном режиме (compute-bound). Если ниже — в режиме, ограниченном памятью (memory-bound). Prefill, благодаря большим матрицам, находится далеко выше гребня и является вычислительно ограниченным. Decode, напротив, находится ниже гребня и сильно зависит от скорости памяти. Понимание этого позволяет нам целенаправленно менять параметры модели, чтобы сдвигать их в более выгодную зону.
02Параметр 1: Размер группы (Group Size, G)
Размер группы (G) определяет, сколько заголовков запросов (query heads) разделяют один заголовок ключа и значения (KV head). В Multi-Head Attention (MHA) G=1, в Grouped-Query Attention (GQA) G может быть различным, а в Multi-Query Attention (MQA) G равно общему числу заголовков запросов.
Влияние на Prefill
Для фазы prefill размер группы G практически не имеет значения. Поскольку prefill вычислительно ограничен, увеличение G не меняет существенно арифметическую интенсивность. Фактические замеры показывают, что изменение G от MHA до MQA меняет время выполнения prefill незначительно. Таким образом, prefill доминируется длиной последовательности (ISL), а не размером группы.
Влияние на Decode
Здесь ситуация кардинально меняется. Для decode арифметическая интенсивность пропорциональна размеру группы G. Увеличение G улучшает арифметическую интенсивность. Почему это важно? Потому что decode ограничен памятью. Увеличивая G, мы уменьшаем количество уника KV-заголовков, которые нужно читать из памяти для каждого токена. Это снижает нагрузку на память и улучшает использование вычислительных блоков GPU.
Замеры показывают, что время выполнения decode заметно сокращается при увеличении G. Например, переход от MHA к MQA дает колоссальный прирост скорости. Однако есть предел: при высоких значениях G кривая производительности начинает выравниваться. Это связано с тем, что накладные расходы на настройку ядра и постобработку становятся доминирующими по сравнению с самой операцией чтения. Тем не менее, тренд остается ясным: для decode нужно стремиться к высоким значениям G.
03Параметр 2: Размерность заголовка (Head Dimension, Hsz)
Размерность заголовка (Hsz) — это размер вектора, с которым работает каждый заголовок внимания. Типичные значения: 64, 128 или 256. В отличие от G, изменение Hsz не меняет арифметическую интенсивность, так как оно пропорционально увеличивает как количество операций (FLOPs), так и объем передаваемых данных (Bytes). Однако оно критически влияет на эффективность использования аппаратных ресурсов GPU.
Выравнивание и эффективность памяти
GPU NVIDIA обрабат
Источник: NVIDIA Developer ↗