Как работают CUDA graphs и когда их использовать?

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

CUDA graphs — это механизм записи последовательности операций на GPU (запуск ядер, копирование памяти) в единый исполняемый граф. Основная цель — устранить накладные расходы на каждый отдельный kernel launch (overhead в микросекундах). Для LLM inference с короткими запросами и фиксированной формой данных CUDA graphs дают ускорение 10–30%. Однако они плохо подходят для динамических размеров входов (разная длина последовательностей) и рекуррентных вычислений, так как требуют перестроения графа при каждом изменении формы.

1. Проблема: Overhead запуска CUDA-ядер

Каждый вызов CUDA kernel (например, операция умножения матриц или attention) требует:

• Переключения контекста на стороне CPU
• Подготовки параметров launch (grid, block)
• Отправки команд в GPU через драйвер

Хотя latency одного launch составляет единицы микросекунд (1–5 мкс), в современном LLM inference число kernel-запусков на один forward pass может достигать тысяч (например, для модели на 7B параметров с FlashAttention — около 1000–3000 запусков). Суммарный overhead становится значимым, особенно при batch=1, когда время вычислений ядер близко к overhead.

Пример:
При 2000 kernel launches overhead = 2000 × 3 мкс = 6 мс. Если чистый compute time forward pass 50 мс, то overhead добавляет 12% времени. При batch=1 это может быть до 30% overhead для коротких промптов (накладные расходы фиксированы, а compute мал).

2. Что такое CUDA Graph

CUDA graph — объект, который фиксирует последовательность операций (запуски ядер, копирования, события) и позволяет воспроизвести их одним вызовом cudaGraphLaunch. Запись происходит в два этапа:

1. Capture (захват) — запускаем интересующий фрагмент кода (например, forward pass модели) в специальном режиме. Драйвер записывает все операции.
2. Instantiation (создание экземпляра) — полученный граф оптимизируется, создается готовый к многократному запуску объект CUDAGraph.

Ключевое: CPU overhead переносится из фазы выполнения в фазу захвата. После захвата запуск графа требует минимальных действий CPU.

Пример использования в PyTorch

import torch

# Гипотетическая модель
model = torch.nn.Linear(1024, 1024).cuda()
input_tensor = torch.randn(1, 1024, device='cuda')

# Захват графа
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
    output = model(input_tensor)

# Первый запуск: CPU пересылает граф в GPU
graph.replay()

# Последующие запуски: overhead ~0
for _ in range(10):
    graph.replay()

Важно: между replay нужно обновлять входные буферы, если они меняются. Для этого используют отдельные буферы-заполнители (static input/output tensors).

3. Архитектура: Захват и воспроизведение

3.1. Механизм захвата

После вызова cudaStreamBeginCapture(capture_stream) все последующие операции на данном стриме записываются. Запись включает:

• Имена ядер, параметры launch (grid, block, аргументы)
• Вызовы cudaMemcpy, cudaMemset
• Обращения к выделенной памяти (адреса, размеры)

Ограничения

• Нельзя использовать динамическое выделение памяти внутри захвата (memset, cudaMalloc). Память должна быть фиксирована до захвата.
• Ветвления на GPU не допускаются — операции выполняются строго в порядке записи.

3.2. Воспроизведение

При cudaGraphLaunch(graph, stream) GPU получает предварительно распарсенный список команд. CPU не тратит время на подготовку каждого launch. Механизм называется single-kernel launch.

3.3. Граф и memory pool

Все тензоры, участвующие в графе, должны быть выделены до захвата. Для работы с переиспользованием памяти (например, torch.utils._cuda_pool) требуется аккуратное управление, чтобы адреса не менялись между replay.

4. Когда использовать CUDA Graphs

Ключевой признак overhead kernel launch составляет заметную долю от total inference time. Профилирование через ncu --set full покажет high "kernel launch overhead".

5. Когда не использовать

Пример плохого случая LLM генерация с autoregressive decoding — длина растет с каждым шагом. При каждом новом токене длина последовательности меняется => граф надо перестроить. Для такого лучше использовать graph caching: предзаписать графы для разных длин (например, 128, 256, 512) и переключать между ними.

6. Ускорение в LLM инференсе

Практические фреймворки (vLLM, TensorRT-LLM) активно используют CUDA graphs для ускорения. Типичная схема:

1. При запуске модели создается набор графов для разных длин последовательностей (block granularity).
2. Во время инференса выбирается ближайший граф по длине (либо pad).
3. Для коротких промптов (32–128 токенов) overhead практически исчезает.

Эксперименты

• LLaMA-7B, batch=1, prompt length=128: overhead ~50% от времени => после CUDA Graph acceleration ~20% относительно обычного запуска, общее ускорение 30–40% (за счет ускорения каждого kernel).
• GPT-175B с batch=1: overhead менее значим (compute доминирует), ускорение 5–10%.

7. Профилирование и отладка

Чтобы понять, нужны ли CUDA graphs, профилируйте:

sudo nsys profile -o profile_launch -t cuda,nvtx python inference.py

В отчёте ищите kernel launch overhead (раздел CUDA API). Если он >10% от общего времени — стоит внедрять.

Инструменты

• ncu (Nsight Compute): детальный анализ ядер
• nsys (Nsight Systems): визуализация временной шкалы
• PyTorch profiler: torch.profiler.profile с profile_memory=True

8. Связь с Agentic RAG

В архитектуре Agentic RAG агент может делать десятки последовательных вызовов LLM (tool calling, reasoning, retrieval summarization). Каждый вызов — короткий промпт (1024–2048 токенов). Если каждый вызов оптимизировать с помощью CUDA graphs, совокупная экономия latency может достигать 30–50%. Особенно это важно для real-time agent loops, где задержка в 100 мс критична.

Рекомендации для Agentic RAG

• Захватить граф для статической части модели (embedding, transformer layers) с фиксированной длиной.
• Для генерации (autoregressive) использовать caching графов по длинам.
• Если размер контекста меняется радикально (вызов разных инструментов), лучше захватить несколько графов под каждую типичную длину.

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

Задача Реализовать инференс небольшой LLM (GPT-2 124M) с CUDA Graph и без, замерить ускорение на коротких промптах.

Инструменты PyTorch 2.0+, CUDA 11+, Hugging Face Transformers, nsys.

Шаги:

1. Загрузить модель distilgpt2 на GPU.
2. Подготовить вход фиксированной длины (128 токенов) с torch.int64.
3. Написать функцию forward_without_graph(), запустить 100 раз, замерить среднее время.
4. Реализовать capture graph:

   with torch.no_grad():
       # Статический буфер
       input_ids = torch.zeros(1, 128, dtype=torch.long, device='cuda')
       graph = torch.cuda.CUDAGraph()
       with torch.cuda.graph(graph):
           output = model(input_ids).logits
   

5. В цикле (100 раз) копировать новые input_ids в статический буфер (buffer.copy_(new_inputs)), затем graph.replay().
6. Использовать torch.cuda.synchronize() перед замером.
7. Профилировать через nsys для наглядности overhead.
8. Повторить для разных длин (64, 256, 512) и сравнить.

Ожидаемый результат
Ускорение 15–30% для длины 64–128 токенов; для длины 1024+ ускорение падает до 5–10%. График зависимости ускорения от длины промпта.

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

Навигация

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