CUDA GraphCUDA Graph
CUDA 图 (CUDA Graph) 是一个由核函数 (kernel) 启动操作以及其他工作组成的图结构,主机端(Host)可以将其一次性整体提交给设备端(Device)执行。
CUDA 图的核心应用场景是减少由主机端在短时间内进行识别、配置和提交大量核函数所带来的开销 (overhead)。 每次启动核函数通常需要微秒($\mu s$)级别的耗时,因此如果需要在几毫秒内启动数百个核函数,这种开销就会变得非常明显。这在 低延迟大语言模型 (LLM) 推理 中是一个非常普遍的现象。
CUDA 图通常通过 CUDA 运行时 (CUDA Runtime) 中的流捕获 API (Stream Capture API) 来创建。它允许将单个 CUDA 流(Stream)上发生的所有操作捕获下来,并在随后进行重放,示例如下:
// capture
cudaStreamBeginCapture(stream);
kernelGemm<<<{32, 20},64,19200,stream>>>(a, b, c);
kernelEpilogue<<<{256,2},{8,32},0,stream>>>(c, c);
cudaStreamEndCapture(stream, &graph);
// launch
cudaGraphInstantiate(&graphExec, graph, flags);
cudaGraphLaunch(graphExec, stream);
NVIDIA 在 此处 对 CUDA Runtime 中的 CUDA Graphs 接口进行了详细记录。
该 API 也被 PyTorch 进行了封装(例如通过 torch.cuda.graph 上下文管理器),这也是在神经网络训练和推理中通常捕获 CUDA 图的方式。
以下是一个 CUDA 图的示例,它是在 B200 GPU 上执行 torch.Linear 层时捕获到的:
┌─────────────────────────────────────────────────────────────────────────┐
│ ┌───────────────────────────────────────────────────────────────────┐ │
│ │ NODE 0: KERNEL │ │
│ ├───────────────────────────────────────────────────────────────────┤ │
│ │ ID: 0 (topoId: 1) │ │
│ │ Kernel: cutlass3x_sm100_simt_sgemm_f32_f32_f32_f32_f32_ │ │
│ │ 64x32x16_1x1x1_3_tnn_align1_bias_f32_relu │ │
│ │ <<<{32,20},64,19200>>> │ │
│ │ Node handle: 0x0000564604539520 │ │
│ │ Func handle: 0x0000564603AFCC00 │ │
│ └───────────────────────────────────────────────────────────────────┘ │
│ │ │
│ │ │
│ ▼ │
│ ┌───────────────────────────────────────────────────────────────────┐ │
│ │ NODE 1: KERNEL │ │
│ ├───────────────────────────────────────────────────────────────────┤ │
│ │ ID: 1 (topoId: 0) │ │
│ │ Kernel: _ZN8cublasLt8epilogue4impl12globalKernelILi8E... │ │
│ │ <<<{256,2},{8,32},0>>> │ │
│ │ Node handle: 0x0000564604539C88 │ │
│ │ Func handle: 0x00005646044770F0 │ │
│ └───────────────────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────────────┘
请注意,核函数是通过指针来识别的(例如 0x564603AFCC00),输入和输出同样由指针定义。 这些以及其他对设备资源的引用,导致 CUDA 图无法被直接序列化,同时也使其失去了可移植性——除非通过完整的主机与设备内存检查点设置及恢复技术 来实现。
英文原文(Modal GPU Glossary)
A CUDA Graph is a graph of kernel launches and other work that can be submitted by the host to the device all at once.
The primary use case for CUDA Graphs is reducing overhead from host identification, configuration, and submission of large numbers of kernels in short periods. Each launch takes on the order of microseconds, so if hundreds of kernels need to be launched in milliseconds, this overhead can be very noticeable. This is commonly the case for low-latency LLM inference.
CUDA Graphs are most commonly created via the stream capture API in the CUDA Runtime, which allows all of the operations that occur on a single CUDA stream to be captured and then later replayed, like
// capture
cudaStreamBeginCapture(stream);
kernelGemm<<<{32, 20},64,19200,stream>>>(a, b, c);
kernelEpilogue<<<{256,2},{8,32},0,stream>>>(c, c);
cudaStreamEndCapture(stream, &graph);
// launch
cudaGraphInstantiate(&graphExec, graph, flags);
cudaGraphLaunch(graphExec, stream);
The CUDA Runtime interface to CUDA Graphs is documented by NVIDIA here.
This API is wrapped by PyTorch, e.g. via the torch.cuda.graph context manager, which is how CUDA Graphs are generally captured for neural network training and inference.
Below is a sample CUDA Graph, captured from a B200 GPU executing a torch.Linear layer:
┌─────────────────────────────────────────────────────────────────────────┐
│ ┌───────────────────────────────────────────────────────────────────┐ │
│ │ NODE 0: KERNEL │ │
│ ├───────────────────────────────────────────────────────────────────┤ │
│ │ ID: 0 (topoId: 1) │ │
│ │ Kernel: cutlass3x_sm100_simt_sgemm_f32_f32_f32_f32_f32_ │ │
│ │ 64x32x16_1x1x1_3_tnn_align1_bias_f32_relu │ │
│ │ <<<{32,20},64,19200>>> │ │
│ │ Node handle: 0x0000564604539520 │ │
│ │ Func handle: 0x0000564603AFCC00 │ │
│ └───────────────────────────────────────────────────────────────────┘ │
│ │ │
│ │ │
│ ▼ │
│ ┌───────────────────────────────────────────────────────────────────┐ │
│ │ NODE 1: KERNEL │ │
│ ├───────────────────────────────────────────────────────────────────┤ │
│ │ ID: 1 (topoId: 0) │ │
│ │ Kernel: _ZN8cublasLt8epilogue4impl12globalKernelILi8E... │ │
│ │ <<<{256,2},{8,32},0>>> │ │
│ │ Node handle: 0x0000564604539C88 │ │
│ │ Func handle: 0x00005646044770F0 │ │
│ └───────────────────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────────────┘
Notice that kernels are identified by pointers, e.g. 0x564603AFCC00. Inputs and outputs are also defined by pointers. These and other references to device resources prevent serialization of CUDA Graphs and make them non-portable, outside of fully checkpointing and then restoring the host and device memory.
相关词条
本词条改编自 Modal GPU Glossary(CC BY 4.0)· 中文翻译 miter6/gpu-glossary-zh,MAE 整理排版。