LLM Inside: выжимаем максимум из Decoder Attention на GPU

Меня зовут Андрей Шукшов. Я пишу YNMT в Яндекс R&D — это движок инференса, на котором работают почти все наши большие языковые модели (LLM). Бо́льшую часть времени я пытаюсь понять, почему некоторые вещи работают медленно и как сделать так, чтобы у них это получалось чуточку быстрее.

Если вы запускали локальную LLM, то, возможно, тоже удивлялись: почему железо, способное рендерить фотореалистичные миры в реальном времени, работает в темпе печатной машинки? В своей статье я попробую хотя бы отчасти ответить на этот вопрос. Под микроскопом посмотрим на механизм Attention в режиме генерации (декодирования) и, вооружившись лучшими современными практиками ускорения на GPU, объединим всю математику в один эффективный kernel, который выжмет максимум производительности из имеющегося у нас железа.

Почему GPU — грузовик, а не спорткар

Представьте, что CPU — это спорткар. У него несколько очень мощных и сложных ядер, спроектированных для минимизации задержки (latency). Иерархия кешей (L1/L2/L3), сложнейшая логика предсказания ветвлений, агрессивное внеочередное выполнение инструкций (out‑of‑order execution) — всё это нужно, чтобы справиться с задачей раньше, чем вы начнёте её считать. Но у нашей машинки небольшой багажник (количество ядер, регистров на каждое): на ней можно очень быстро привезти несколько вёдер песка на стройку. Можно, но зачем?

В свою очередь GPU — это грузовик. Он медленнее разгоняется и дольше едет, зато сразу привозит тонну песка, и вам не так уж важно, когда приедет второй рейс. Главная сила такой архитектуры — в пропускной способности (throughput). У GPU тысячи простых, но эффективных ядер. За счёт огромного параллелизма они компенсируют задержки доступа к DRAM.

Так как для работы с LLM нужно одновременно выполнять триллионы однотипных математических операций, GPU хорошо для этого подходит. Главное — постоянно держать видеокарту полностью загруженной.

Иллюзия многопоточности

CUDA — технология не новая. В 2007 году она должна была помочь CPU‑программистам, привыкшим писать многопоточный код, за разумное время портировать его на GPU и получить реальное ускорение.

Тут‑то и начинаются проблемы, ведь когда мы садимся писать код на CUDA, нам продают очень удобную ментальную модель: здесь можно запустить хоть миллион независимых потоков, и каждый будет выполнять свою работу. Но это всего лишь удобное упрощение, за которым скрывается железо со своими ограничениями.

Вы создаёте программу (kernel), которая описывает поведение одного конкретного потока. Но на самом деле GPU оперирует не отдельными потоками, а группами, например, по 32 потока, которые называются варпами (warp).

В отдельный момент весь варп выполняет одну команду аппаратного планировщика. У потоков в группе могут быть разные данные в регистрах, но команда (instruction pointer) одна на всех. С этим связана первая проблема — расхождение потоков (code divergence).

unsigned idx = threadIdx.x;
if(idx % 2 == 0)
	computeA();
else
	computeB();

Здесь мы наезжаем на классические проблемы SIMD (Single Instruction, Multiple Data). Потоки с чётными idx хотят пойти в ветку if, а с нечётными — в else. Однако варп так не умеет, поэтому железо сначала выполняет computeA() (при этом все нечётные потоки в варпе просто неактивны). Затем оно выполняет computeB() — в это время простаивают уже чётные потоки.

В итоге общее время выполнения команды равно сумме времени выполнения обеих веток. Только вот из‑за одного if мы теряем половину производительности. Но хорошая новость в том, что ветвления в LLM встречаются редко и не влияют на скорость (в gemm, которые занимают большую часть времени выполнения, всё линейно).

Streaming Multiprocessor

Теперь, когда мы вспомнили про варпы, посмотрим, где они живут. Основная рабочая единица GPU — это потоковый мультипроцессор (Streaming Multiprocessor, или SM). Упрощённо можно назвать его большим самостоятельным ядром GPU.

На борту современной видеокарты таких SM больше сотни, и именно между ними распределяется вся работа.

Внутри Streaming Multiprocessor, помимо прочего, находится несколько отдельных компонентов:

• CUDA Cores (ALU) — базовые арифметико‑логические устройства, рабочие лошадки для стандартных вычислений (сложение, умножение, битовые операции).

• Tensor Cores — специализированные блоки, заточенные под умножение небольших матриц с накоплением (MMA — Matrix Multiply‑Accumulate). Именно их производительность вы видите во всех маркетинговых материалах производителей GPU.

• Shared Memory — крошечный, но быстрый кусочек памяти (128 КБ). Здесь можно хранить самые нужные инструменты и компоненты, чтобы не обращаться каждый раз к сравнительно медленной DRAM.

Чтобы выжать из GPU максимум, нужно хорошо понимать особенности архитектуры видеоускорителя: распараллеливать задачи, подбирать их размер, чтобы они правильно ложились на аппаратные блоки, избегать расхождения потоков в коде и активно использовать быструю Shared Memory. 

• GPU требует достаточного количества параллельных задач для загрузки всех потоковых мультипроцессоров.

• Каждая задача должна содержать достаточно вычислений для эффективной работы отдельного SM.

• Недостаточный параллелизм приводит к простою.

• Чтобы много считать на SM, нужно много грузить из VRAM.

А теперь посмотрим, как со всем этим взаимодействует LLM. 

Decoder Attention, когда сто вёрст не крюк

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

Модель отвечает не сразу целиком, а выдаёт текст слово за словом. Дело в том, как устроена генерация: грубо говоря, LLM смотрит на уже созданный текст и добавляет следующее слово так, чтобы текст был связным. 

Работа энкодера с точки зрения GPU в целом похожа на то, что происходит на обучении (умножение больших матриц, compute‑bound), и в сценариях вида «несколько тысяч токенов входа → несколько сотен токенов генерации» составляет относительно небольшой объём — доли или единицы секунд против десятков секунд на генерацию. Поэтому рассмотрим именно декодер. 

На схеме ниже он изображён справа. Каждый слой декодера состоит из нескольких операций. Из них две занимают основное время — матричное умножение внутри Feed‑forward network и Attention. В этой статье мы сосредоточимся на последнем (как более интересном с точки зрения возможных оптимизаций).

Во‑первых, это один из ключевых механизмов в LLM. Во‑вторых, на его работу уходит значительная часть времени всего декодера. Attention помогает выстроить связи между словами во входной последовательности. 

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

В основе реализации головы Attention лежит математическая формула, состоящая из трёх матриц.

Здесь всего три операции: умножение матриц Q на K, вычисление функции softmax от результата и умножение на матрицу V. Самый очевидный способ получить результат — сделать три последовательных шага, каждый из которых будет отдельным вызовом GPU‑ядра (kernel).

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

Если вернуться к аналогии со стройкой, наш грузовик едет на «склад» (DRAM), загружает в себя матрицы Q и K, выполняет умножение, а результат S… выгружает обратно на склад. Затем он едет на тот же склад, забирает матрицу S, считает softmax и результат P… снова выгружает на склад. GPU забирает P и V, умножает их и наконец‑то записывает финальный результат Y. При этом вся работа происходит на одном и том же производстве.

Таким образом, мы заставили наше железо трижды обращаться к самой медленной части системы, гоняя данные туда‑сюда. Получилось совсем неоптимально.

Камень преткновения

А почему нельзя объединить все три операции (Q @ K.T, softmax, умножение на V) в один кернел и сделать всё за один проход? Проблема в функции softmax.

Дело в том, что для эффективного вычисления GEMM0 нам по SM нужно распределить работу прямоугольными тайлами (выходными подматрицами ограниченного размера), а в Softmax каждый SM должен видеть одну строчку целиком. А это значит, что собрать обе операции внутри одного ядра (и не проиграть фатально по скорости) не получается.

Числитель — экспонента одного конкретного элемента, он считается относительно быстро и легко. А вот знаменатель — сумма экспонент по всем элементам вектора. Это число можно получить, только когда все потоки посчитают свои экспоненты и сложат результаты.

Особенность декодера

И вишенка на торте. Мы ведь говорим о декодере в режиме генерации текста — когда модель выдаёт ответ слово за словом. Матрица Q (Query) на самом деле представляет собой новый токен (слово), который модель только что сгенерировала и сейчас обрабатывает. Это всего лишь один вектор — одна строка чисел и забитое нулями пространство. А вот матрицы K (Key) и V (Value) — это совсем другая история. Это постоянно растущий кеш всех предыдущих токенов в последовательности. С каждым новым сгенерированным словом эти матрицы становятся всё больше и больше.

Поэтому первая операция — умножение Query на транспонированную Key — превращается из умножения матрицы на матрицу (GEMM) в умножение одного вектора на матрицу (GEMV).

С точки зрения GPU это худший из возможных сценариев. У операций типа GEMV низкая арифметическая плотность — на каждый байт данных, прочитанный из памяти, приходится очень мало вычислений. Такие задачи называют memory‑bound (ограниченные пропускной способностью памяти).

У нас есть изначально memory‑bound‑задача (где и так много времени тратится на чтение данных). И мы, из‑за природы функции softmax и её редукции, разбили вычисления на три ещё более неэффективных этапа, каждый из которых зависит от памяти.

Математическая магия

Чтобы исправить положение, применим Online Softmax — алгоритм, который позволяет вычислять softmax итеративно. 

Представьте, что вы ищете самого высокого человека в огромной толпе. Обычный softmax работает так: сначала собирает всех на стадионе, измеряет рост каждого, а потом сравнивает. Это требует полного прохода по всем данным, перед тем как вы получите результат. Online Softmax предлагает другой подход: «Держите в голове рост первого человека. Встретили кого‑то выше? Отлично, теперь он самый высокий. Встретили кого‑то ниже? Игнорируйте».

В терминах softmax, мы делаем то же самое, но одновременно поддерживаем два состояния для каждой строки вычислений. Первое — текущий максимум (m): самое большое значение, которое мы видели до сих пор. Второе — сумма экспонент (z): сумма, нормализованная относительно этого максимума. Когда появляется новый элемент, x, мы просто обновляем эти два значения по хитрой формуле. 

Это позволяет обрабатывать матрицу по частям, не дожидаясь её полного вычисления.

Собираем всё в один CUDA‑кернел

Теперь, когда у нас есть Online Softmax, можно упаковать все три операции в один большой кернел. Это называется Fused Kernel.

Вместо того чтобы трижды отправлять грузовик на склад (в медленную DRAM), мы заставим его сделать всего одну поездку и организуем работу так, чтобы вся основная возня происходила уже на стройплощадке.

Thread Block берёт небольшой кусочек (tile) матриц K и V из медленной глобальной памяти (DRAM) и переносит их в Shared Memory. Используя только данные из Shared Memory, потоки вычисляют соответствующую часть Q @ K.T. Прямо здесь, не покидая SM, мы на лету обновляем значения Online Softmax — текущий максимум и сумму экспонент. Результат softmax сразу же домножается на кусочек матрицы V, который уже лежит в Shared Memory. Процесс повторяется для следующего блока матриц K и V, пока мы не пройдём их целиком.

Все эти действия происходят последовательно в одном потоке вычислений, не покидая быструю Shared Memory.

Почему не Tensor Cores

Казалось бы, для умножения матриц логично задействовать тензорные ядра — они же для этого и созданы. Но тут начинаются нюансы. Tensor Cores заточены под работу с блоками матриц, скажем, 16 × 16 элементов, а в режиме декодера Q у нас всего один вектор. Получается, что матрица в умножении почти целиком состоит из нулей, так что обрабатывать её в Tensor Core — всё равно что стрелять из пушки по воробьям. Утилизация будет ничтожной.

Поэтому лучше написать Fused Kernel с использованием обычных CUDA Cores (ALU). Их много, они более универсальны и хорошо подходят для задачи с низкой арифметической плотностью, какой и является GEMV.

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

Важная ремарка: на самом деле в ситуациях с GQA и спекулятивным декодированием тензорные ядра могут быть уже жизненно необходимы. Подробнее об этом вы можете прочитать в техническом отчёте Alice AI.

Анализ производительности

Первые же замеры показывают солидное ускорение: в зависимости от параметров задачи Fused Kernel работает в 1,5–1,7 раза быстрее, чем набор из трёх стандартных вызовов.

Победа? Вроде бы да, но быстрее относительно чего? Может быть, это эффект низкой базы?

Для memory‑bound‑задач вроде нашей есть удобный эталон — утилизация пропускной способности памяти (Memory Bandwidth Utilization). Чтобы рассчитать этот показатель, берём общий объём данных, который прочитал и записал кернел (ReadTraffic + WriteTraffic), и делим на время его работы. Важно понимать, что речь идёт об объёме трафика, который реально прокачивается через шину DRAM, то есть зачастую в этом контексте: ReadTraffic + WriteTraffic = InputSize + OutputSize.

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

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

Результаты тестов

Посмотрим на финальные цифры. Вне зависимости от размера батча или длины последовательности наш Fused Kernel стабильно показывает утилизацию шины в 85–91%.

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

Результат 90% означает, что наш алгоритм практически полностью насыщает шину памяти. Мы упёрлись в физический предел оборудования.

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

Первое: разберитесь, как работает железо. Это нужно, чтобы утилизировать его по максимуму и не оставлять простаивающим.

Второе: поймите, что происходит с вашими данными и какую задачу вы решаете. Тогда сможете применить математические трюки вроде Online Softmax и получить дополнительное ускорение.

Третье: оглядитесь вокруг. Может быть, кто‑то уже решил вашу задачу :)

Заглядывайте под капот, не бойтесь математики, и ваше железо ответит вам взаимностью.

Спасибо, что дочитали, буду рад ответить на вопросы в комментариях!