Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B
[Блогпост]

Прямой (да и обратный) проход через современную LLM, подразумевает запуск нескольких сотен кернелов (attention, mlp, нормализаций).

Команда из Стэнфорда обнаруживает, что скорость инференса маленьких LLM (1-3B параметров) упирается не в вычисления/память, а во время запуска кернелов. Эффективные движки для инференса (vLLM/SGLang) позволяют только на 50% использовать пропускную способность новых видеокарт (H100/GB200). Там некоторые операции уже слиты в один вызов кернела, но самих вызовов все еще остается много.

И авторы предлагают реализовать весь forward pass в виде одного Megakernel 😱!

Из нюансов реализации стоит выделить следующее:

1️⃣ Управление памятью. Так как за shared memory борются сразу несколько процессов, надо эффективно распределить ее и раздавать по запросу. Для это используется некий вариант paging.
2️⃣ Синхронизация. Теперь у нас много операций работающих асинхронно и требуется внутри кернела регулировать то, чтобы процесс не начал работать, пока не будут готовы все необходимые входы (т.е attention не запустился, пока не готовы Q, K, V).

В результате удается добиться ускорения на Llama-1B при инференсе с батчом 1:
🏃‍♂️ 2.5x против vLLM, 1.5x против SGLang на H100
🏃‍♂️ 3.5x против vLLM, 2.5x против SGLang на GB200

Утилизация ширины памяти для H100 достигает 78%.