Что такое kernel fusion и как он применяется в LLM serving?

Краткий тезис

Kernel fusion — это техника оптимизации, при которой несколько последовательных вычислительных операций (ядер, kernels) объединяются в одно ядро, выполняемое на GPU. В контексте LLM serving это позволяет значительно снизить накладные расходы на запуск ядер (kernel launch overhead) и уменьшить объём передаваемых данных между глобальной памятью и регистрами (memory traffic). Ключевые примеры применения — FlashAttention (слияние операций attention), fused LayerNorm + activation и fused QKV projection. Без kernel fusion современные LLM не могли бы достичь требуемой скорости инференса.

1. Что такое kernel fusion (слияние ядер)

Kernel в контексте GPU — это функция, выполняемая на устройстве (GPU) над большим массивом данных. Обычно одна операция (например, умножение матриц, ReLU, LayerNorm) реализуется отдельным ядром. Kernel fusion — объединение нескольких таких ядер в одно, чтобы:

• Уменьшить количество запусков ядер (kernel launches) — каждый запуск требует передачи команд через драйвер и планировщик, что занимает микросекунды, но при большом числе вызовов становится заметным.
• Сократить чтение/запись промежуточных результатов в глобальную память GPU (HBM). Вместо того чтобы записывать результат первого ядра в HBM, а затем читать его вторым ядром, fused kernel держит данные в регистрах или shared memory.

Термин memory-bound (ограниченный по памяти) — операция, скорость которой упирается в пропускную способность памяти, а не в вычислительную мощность. Большинство операций в LLM serving (например, LayerNorm, softmax, element-wise) являются memory-bound, поэтому fusion особенно эффективно.

2. Проблема: накладные расходы и memory traffic

При последовательном выполнении нескольких операций (например, x = matmul(Q, K) → x = softmax(x) → x = matmul(x, V)) каждое ядро:

• Запускается отдельно (kernel launch overhead: ~5-20 мкс на ядро).
• Читает входные данные из HBM.
• Записывает результат в HBM.
• Следующее ядро читает этот результат снова.

Для attention в LLM это особенно критично: матрицы Q, K, V имеют размер [batch, heads, seq_len, d_head], и промежуточные результаты (например, Q @ K^T) могут быть огромными (seq_len²). Без fusion мы бы записывали и читали гигабайты данных.

Roofline model — модель, показывающая, что производительность операции ограничена либо пропускной способностью памяти (memory-bound), либо вычислительной мощностью (compute-bound). Для memory-bound операций ускорение достигается уменьшением объёма передаваемых данных, что и делает fusion.

3. Примеры kernel fusion в LLM serving

3.1 FlashAttention (fused attention)

FlashAttention — самый известный пример. Он объединяет все шаги attention (QK^T, softmax, dropout, PV) в одно ядро, используя tiling (разбиение на блоки) и online softmax (вычисление softmax без записи всей матрицы). Результат: ускорение в 2-4 раза и снижение использования памяти с O(N²) до O(N).

3.2 Fused LayerNorm + activation

После LayerNorm часто идёт нелинейная активация (например, SiLU, GELU). Вместо двух ядер (LayerNorm → активация) пишут одно, которое сразу применяет активацию к нормализованным значениям, не сохраняя промежуточный результат.

3.3 Fused QKV projection

В трансформерах три линейных слоя (Q, K, V) часто объединяются в один большой слой (QKV projection). Но даже внутри одного слоя можно выполнить умножение на веса и последующее разделение на три части в одном ядре, избегая лишних копий.

3.4 Fused softmax + reduction

В некоторых реализациях softmax объединяют с последующим reduction (например, суммированием для loss). Это уменьшает количество проходов по данным.

3.5 Fused RMSNorm + RoPE

В современных LLM (LLaMA, Mistral) используется RMSNorm и RoPE (Rotary Position Embedding). Их можно выполнить в одном ядре, применив позиционное кодирование сразу после нормализации.

4. Как реализуется kernel fusion

4.1 Ручное написание CUDA-ядер

Опытные инженеры пишут fused kernels на CUDA C++, используя shared memory и регистры. Это даёт максимальный контроль, но требует глубоких знаний архитектуры GPU.

4.2 Triton (OpenAI)

Triton — язык программирования для написания GPU-ядер на Python. Он автоматически управляет shared memory и tiling, позволяя быстро прототипировать fused kernels. Пример fused add + ReLU:

import triton
import triton.language as tl

@triton.jit
def fused_add_relu_kernel(x_ptr, y_ptr, output_ptr, n, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = tl.maximum(x + y, 0)  # fused add + relu
    tl.store(output_ptr + offsets, output, mask=mask)

4.3 Компиляторы (TVM, TensorRT, XLA)

Библиотеки вроде TensorRT и XLA автоматически обнаруживают последовательности операций и сливают их при компиляции графа. Это удобно, но не всегда даёт оптимальный результат.

5. Влияние на производительность

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

6. Инструменты и библиотеки

• FlashAttention (https://github.com/Dao-AILab/flash-attention) — fused attention ядра.
• vLLM — использует PagedAttention (вариант fusion) и fused kernels для быстрого serving.
• TensorRT-LLM — компилирует модели с автоматическим fusion.
• xFormers (Meta) — набор оптимизированных строительных блоков, включая fused attention.
• Triton — для написания собственных fused kernels.

7. Связь с архитектурой Agentic RAG

Хотя вопрос напрямую про LLM serving, в Agentic RAG агенты часто вызывают LLM многократно (планирование, генерация, проверка). Каждый вызов требует быстрого инференса. Kernel fusion — одна из ключевых оптимизаций, позволяющая снизить latency до приемлемого уровня для интерактивных агентов. Без неё агенты были бы слишком медленными.

Пет-проект для закрепления

Задача Написать fused kernel для операции LayerNorm + SiLU и сравнить производительность с двумя отдельными ядрами.

Инструменты Python, PyTorch, Triton.

Шаги:

1. Реализовать обычный LayerNorm и SiLU через PyTorch (два отдельных вызова).
2. Написать fused kernel на Triton, который принимает входной тензор, вычисляет нормализацию (среднее, дисперсия) и сразу применяет SiLU.
3. Замерить latency для тензоров размера [batch, seq_len, hidden_dim] (например, 4×1024×4096) на GPU.
4. Построить график ускорения в зависимости от размера.

Ожидаемый результат Fused kernel должен быть в 1.5-2 раза быстрее за счёт уменьшения числа запусков и чтений/записей из HBM.

Связь с другими вопросами

Навигация

• Предыдущий: 308
• Следующий: 310
• Индекс: 00. Индекс разборов