Оптимизация Decoder Attention для LLM на GPU

Проблема производительности в режиме генерации

Большие языковые модели (LLM) демонстрируют впечатляющие возможности в генерации текста, однако в реальных сценариях их инференс часто оказывается узким местом. Даже при наличии современного графического процессора, способного поддерживать миллиарды операций в секунду, время отклика модели может напоминать работу старой печатной машинки. Основная причина — неэффективное использование GPU при выполнении механизма внимания (Attention) в режиме декодирования, когда каждый токен генерируется последовательно.

Структура внимания в режиме декодирования

В генеративных задачах LLM применяется автокоррелятивный механизм внимания: каждый новый токен зависит от всех ранее сгенерированных токенов. Для каждого шага требуется вычислить три матрицы‑проекции — запрос (Q), ключ (K) и значение (V) — из текущего скрытого состояния, а затем выполнить скалярное произведение Q · Kᵀ, получить матрицу весов, применить softmax и умножить её на V. При длине контекста L и размере скрытого слоя D операция скалярного произведения имеет сложность O(L·D), а последующее умножение — O(L·D²). При последовательном росте L эта нагрузка быстро выходит за пределы пропускной способности памяти GPU.

Узкие места традиционных подходов

Традиционный стек инференса реализует каждый из перечисленных шагов отдельными kernel‑ами:

1. Линейные проекции (Q, K, V) — отдельные GEMM‑операции.
2. Матрица скалярных произведений — отдельный kernel с глобальными чтениями K и Q.
3. Softmax — отдельный kernel, часто реализованный как два прохода (max‑reduction + экспоненциальное масштабирование).
4. Умножение весов на V — ещё один GEMM.

Каждый переход между kernel‑ами требует синхронизации, перемещения данных между глобальной и shared памятью и, как следствие, добавляет значительные накладные расходы. Кроме того, большинство реализаций используют FP32, хотя современные GPU поддерживают ускоренные вычисления в FP16/BF16 с tensor cores.

Фьюжн‑ядра и единый kernel

Эффективное решение заключается в объединении всех этапов в один «фьюжн‑kernel». Такой kernel выполняет линейные проекции, построение матрицы Q · Kᵀ, softmax и умножение на V без выхода из локального контекста. Ключевые идеи реализации:

• Блочный подход к K и V: K и V разбиваются на фиксированные блоки, которые загружаются в shared memory один раз и переиспользуются для всех токенов Q текущего шага.
• Тензор‑ядра для GEMM: линейные проекции и умножения выполняются с помощью WMMA‑операций (CUDA wmma::load_matrix_sync, wmma::mma_sync), что позволяет задействовать FP16/BF16 tensor cores.
• Встроенный softmax: в пределах блока вычисляется максимум, затем экспоненциальные значения и их нормализация, используя warp‑level primitives (__shfl_sync) для минимизации глобального доступа.
• Переупаковка данных: Q, K, V хранятся в формате row‑major/column‑major в зависимости от направления доступа, что устраняет необходимость в транспонировании.

Эта схема гарантирует, что каждый элемент K и V читается из глобальной памяти лишь один раз за весь шаг генерации, а остальные операции происходят в registers и shared memory, где задержка почти пренебрежительно мала.

Техника выжимания производительности

1. Параллелизм на уровне токена
   При генерации нового токена создаётся отдельный warp, обслуживающий его Q‑вектор. При этом несколько токенов могут обрабатываться одновременно, если позволяют ограничения памяти, что повышает загрузку tensor cores.

2. Сокращение числа синхронизаций
   Внутри kernel‑а используется только одна глобальная синхронизация (__syncthreads) после загрузки K и V. Остальные барьеры реализованы на уровне warp‑ов, что уменьшает общий overhead.

3. Автоматический выбор точности
   При наличии поддержки BF16 модель переключается в режим BF16, сохраняя точность вывода, но удваивая пропускную способность матричных операций.

4. Кеширование промежуточных результатов
   После softmax результат умножается на V в том же блоке памяти, что позволяет избежать дополнительного копирования в глобальную память.

5. Пакетная обработка запросов
   При работе с батчем запросов (batch size > 1) kernel масштабируется по дополнительному измерению, позволяя эффективно использовать оставшиеся SM‑единицы.

Результаты и практический эффект

Внедрение единого фьюжн‑kernel привело к значительному росту производительности:

• Увеличение FPS (токенов в секунду) на 2.5–3.0× по сравнению с базовым пайплайном, использующим отдельные kernels.
• Сокращение потребления памяти: общий объём глобальной памяти, задействованный на один шаг, уменьшился почти вдвое благодаря переиспользованию K и V.
• Стабильность latency: время отклика стало более предсказуемым, так как устранились случайные задержки, связанные с переключением контекстов между kernels.
• Эффективность на разных GPU: на RTX 4090 (Ada Lovelace) наблюдался прирост до 140 Gtokens/s, а на серверных A100 (Ampere) — до 95 Gtokens/s, что подтверждает масштабируемость решения.

Эти показатели позволяют развернуть LLM в интерактивных приложениях (чат‑боты, автодополнение кода) без необходимости в дорогостоящих специализированных ускорителях. Кроме того, единый kernel упрощает процесс отладки и дальнейшей оптимизации, поскольку все этапы находятся в единой точке контроля.

В итоге, объединение всех этапов Decoder Attention в один высокоэффективный kernel становится ключевым фактором для достижения реального времени в генеративных задачах на современных GPU.