Что такое CUDA graphs и как они ускоряют LLM инференс?

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

CUDA graphs — это механизм NVIDIA, который позволяет записать последовательность операций GPU (ядер, копий памяти) в граф и затем воспроизводить её одним вызовом. При LLM инференсе основной выигрыш (10–30% для коротких запросов) достигается за счёт устранения overhead|launch overhead|launch overhead|launch overhead|launch kernel overhead|launch overhead|launch overhead — накладных расходов на запуск каждого ядра CPU→GPU. Поскольку инференс LLM состоит из множества мелких ядер (операции attention, матричные умножения), CUDA graphs превращают сотни микросекундных вызовов в один, что особенно заметно на малых батчах и коротких последовательностях.

1. Термин: Kernel launch overhead (накладные расходы на запуск ядра)

Каждый вызов CUDA-ядра (например, torch.matmul или softmax) требует от CPU отправки команд GPU через драйвер. Этот процесс занимает порядка 1–10 микросекунд на ядро. При инференсе LLM с сотнями ядер на шаг (например, для модели 7B — около 300 ядер) суммарный overhead может достигать миллисекунд, что составляет значительную долю времени ответа, особенно при маленьких батчах (batch size = 1).

CUDA graphs решают эту проблему, записывая граф зависимостей между ядрами и воспроизводя его целиком, минуя повторные launch.

2. Как работают CUDA graphs

Процесс состоит из трёх этапов:

1. Захват (capture) — CPU запускает последовательность операций в специальном режиме. Вместо немедленного исполнения драйвер записывает граф: какие ядра, с какими параметрами, в каком порядке, какие зависимости по памяти.
2. Инстанцирование (instantiation) — из захваченного графа создаётся оптимизированный экземпляр, который можно многократно запускать.
3. Воспроизведение (replay) — единственный вызов cudaGraphLaunch запускает весь граф целиком. CPU не участвует в запуске каждого ядра.

import torch

# Пример захвата графа для простой операции
x = torch.randn(1, 512, device='cuda')
W = torch.randn(512, 1024, device='cuda')
b = torch.randn(1024, device='cuda')

# Захват
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    y = torch.mm(x, W)
    y = y + b
    y = torch.relu(y)

# Воспроизведение (можно менять входные данные через реплей-буферы)
g.replay()

3. Почему LLM инференс — идеальный кандидат для CUDA graphs

LLM (трансформеры) выполняют фиксированную последовательность операций на каждом шаге генерации:

• Self-attention: QKV проекции, attention scores, softmax, weighted sum.
• Feed-forward network: два линейных слоя с активацией.
• LayerNorm, residual connections.

Все эти операции имеют статический граф вычислений (при условии фиксированной длины последовательности и batch size). Это значит, что после первого шага можно захватить граф и на всех последующих шагах просто его воспроизводить, меняя только входные тензоры (через replay buffers).

4. Ускорение: 10–30% для коротких запросов

Формула оценки ускорения: Пусть T_kernel — суммарное время работы всех ядер, T_overhead — время на launch. Без графа: T_total = T_kernel + T_overhead. С графом: T_graph = T_kernel + T_overhead_capture (однократно) + T_replay (пренебрежимо мал). Ускорение ≈ T_overhead / (T_kernel + T_overhead). Для коротких запросов T_kernel мал, поэтому выигрыш велик.

5. Пример кода: инференс GPT-2 с CUDA graphs (упрощённо)

import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

model = GPT2LMHeadModel.from_pretrained('gpt2').cuda().eval()
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# Фиксируем длину и batch
batch_size = 1
seq_len = 64
input_ids = torch.randint(0, 50257, (batch_size, seq_len), device='cuda')

# Захват графа для одного forward pass
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    with torch.no_grad():
        logits = model(input_ids).logits

# Теперь можно менять input_ids через replay buffer (но размер должен совпадать)
# Для простоты используем тот же тензор
g.replay()

На практике в библиотеках (vLLM, TensorRT-LLM) графы захватываются для каждой комбинации (batch_size, seq_len) и переиспользуются при совпадении размеров.

6. Сравнение: обычный инференс vs CUDA graphs

7. Ограничения и подводные камни

• Фиксированные размеры: граф захватывается для конкретных batch_size и seq_len. Если они меняются, нужно захватывать новый граф (или использовать fallback).
• Overhead на захват: первый запуск с захватом может быть медленнее обычного (до 10 мс). Окупается только при многократном воспроизведении.
• Не все операции поддерживаются: например, операции с динамической памятью (torch.where с переменным числом элементов) могут не захватываться корректно.
• Совместимость с CUDA streams: графы работают в рамках одного стрима; при использовании нескольких стримов требуется осторожность.

8. Интеграция в LLM сервинг

Современные системы инференса LLM активно используют CUDA graphs:

• vLLM: автоматически захватывает графы для каждого размера блока (block size) и переиспользует их при continuous batching.
• TensorRT-LLM: строит графы на этапе компиляции модели, что даёт ещё больший выигрыш за счёт fusion ядер.
• PyTorch 2.0+: torch.compile может генерировать CUDA graphs для статических частей модели.

Пример из vLLM (упрощённо):

# vLLM создаёт граф для каждого слоя и шага генерации
graph = CUDAGraph()
with torch.cuda.graph(graph):
    hidden_states = self.attention(hidden_states)
    hidden_states = self.mlp(hidden_states)
# При декодировании вызывает graph.replay() с новыми данными

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

Задача: Сравнить latency инференса маленькой модели (например, GPT-2) с CUDA graphs и без.

Инструменты: Python, PyTorch, transformers, CUDA (GPU).

Шаги:

1. Загрузите GPT-2 и токенизатор.
2. Подготовьте входные тензоры фиксированной длины (например, 128 токенов).
3. Напишите функцию inference_without_graph, которая выполняет model.generate() (или просто forward pass) и замеряет время.
4. Напишите функцию inference_with_graph:
   • Захватите граф для одного forward pass (без generate, только логиты).
   • Используйте replay() для последующих вызовов.
   • Замерьте время 100 повторений.
5. Повторите для разных batch_size (1, 2, 4) и seq_len (64, 256, 1024).
6. Постройте таблицу ускорения.

Ожидаемый результат: Для batch_size=1, seq_len=64 ускорение ~20–30%; для больших размеров — меньше. Вы увидите, что overhead захвата окупается после нескольких десятков вызовов.

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

Навигация

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