CUDA Graph
把重复的工作一次提交
核心解释
CUDA Graph 把 kernel、数据拷贝等操作表示为节点,把它们之间的依赖表示为边。先定义并实例化图,后续再通过一次 launch 提交整组工作。许多准备工作可以提前完成,减少每次执行时的提交开销。[1]
当 GPU 很快完成一个 kernel,而 CPU 还没提交下一个时,执行时间线上会出现空隙。CUDA Graph 对这种重复运行、由许多短 kernel 组成的工作尤其有用。它本身不把这些 kernel 合成一个,也不减少矩阵乘法的运算量。[2]
Stream capture 记录操作与依赖,随后生成可执行图。普通 PyTorch 重放使用录制时的尺寸、布局和地址;输入数值可以改变。常见做法是把新数据写入保留的输入缓冲区,再 replay 并读取输出。[3]
训练可以捕获前向、反向和支持捕获的优化器步骤。推理则要处理不同 batch 和序列长度,常按尺寸分档,并在整图、分段图和普通执行之间选择。图的覆盖范围越大,需要一起管理的地址与依赖也越多。[4] [5]
是否值得使用,要看重复次数能否摊薄准备成本。设额外准备成本为 ,普通执行和重放的稳态耗时分别为 、。两者按相同工作范围计时,并包含各自需要的输入处理。重复 次时,可用下面的关系估算收益:
资料
CUDA Graphs 入门:看清 kernel 之间的空隙 ↗
从 Nsight 时间线看逐次启动、减少同步与整图提交的区别,再读最小 C++ 示例。
PyTorch:从捕获到训练重放 ↗
用执行时间线和训练例子建立直觉。先理解输入缓冲区与 replay,再查当前 API 的具体要求。
CUDA Programming Guide:图的执行与更新 ↗
按需查节点、依赖、stream capture、参数更新和条件节点。把 CUDA 原生能力与上层框架的接口区分开。
PyTorch:捕获约束与显存管理 ↗
从输入输出地址走到私有内存池、部分网络捕获与多 stream 协作。包含前向、反向和优化器的完整例子。
动态参数:更新节点,还是重新捕获? ↗
通过代码理解 stream capture 与显式节点 API 怎样组合,以及节点参数如何在下一次 launch 前更新。
CUDAGraph Trees:编译后的图怎样管理内存 ↗
重点读内存复用、不同 shape 的重新录制和跨迭代输出。理解 torch.compile 后,为什么仍会录制多张 CUDA Graph。
vLLM:Full、Piecewise 与运行时选择 ↗
沿 BatchDescriptor、dispatcher 和 wrapper 追踪捕获与重放。看 batch 组成、padding 和 attention 后端怎样影响图的选择。
NCCL:把通信捕获进图 ↗
看各 rank 如何一致地捕获并重放通信,再检查多 GPU、多 communicator 和混合执行的限制。
QA
CUDA Graph、kernel fusion 和 torch.compile 分别改变什么?
CUDA Graph 主要改变工作提交方式。Kernel fusion 将多个算子的计算合进更少的 kernel,还可能减少中间结果的访存。torch.compile 可以优化并生成 kernel,也可能接入 CUDA Graph。
编译成功并不表示整段代码已经被一张 CUDA Graph 捕获。判断是否使用了重放,要看对应后端的日志和执行时间线。
每次输入都不同,为什么还能重放同一张图?
普通 PyTorch 捕获记录了 kernel 使用的地址。重放前把新数据复制进保留的输入缓冲区,图就会在相同地址上读取新值。给 Python 变量换绑一个新 tensor,不会自动改掉图里的指针。
CUDA 原生 API 还可以按规则更新节点参数。那是对图的显式修改,与只替换输入缓冲区中的内容不同。
支持动态 shape 的编译,是否意味着一张 CUDA Graph 能覆盖所有尺寸?
一个编译结果可以处理多个 shape,但普通 CUDA Graph 重放仍使用录制时的尺寸、布局和地址。CUDAGraph Trees 会为不同 shape 录制相应的图。
推理系统常把请求整理到有限的尺寸档位,必要时 padding,再选择已有图。档位越多,捕获成本和显存占用通常也越高。
CUDA Graph 里可以有 if 和循环吗?
CUDA 原生条件节点可以让 GPU 选择分支或重复执行预先定义的子图。普通 Python 分支不会因为放进捕获区域,就自动变成每次重放时都会执行的控制逻辑。
PyTorch 2.12 已加入基于 GPU 条件的 torch.cond 捕获支持,发布说明限定了支持的后端。使用时要核对具体 PyTorch、CUDA 与编译后端,不能把原生条件节点能力直接套给所有调用路径。
为什么开启后显存增加,之前的输出还可能被覆盖?
图需要在重放期间保留使用的地址。PyTorch 用图私有内存池满足捕获中的分配,多张图可能保留多组内存。
手动捕获的输出缓冲区会被下一次重放写入。要长期保存某次结果,应在它被复用前复制到独立存储。CUDAGraph Trees 还会管理迭代与张量生命周期,跨迭代保存输出时要按其规则处理。
训练中的权重更新、Dropout 和反向传播能一起捕获吗?
只要相关操作支持捕获,就可以把前向、反向和优化器步骤放进图。权重数值可以在固定存储中更新。PyTorch 的示例演示了完整训练迭代。
支持的 CUDA 随机操作会管理 RNG 状态,重放并不要求每轮复用同一组 Dropout 随机数。动态 loss scaling、CPU 标量和自定义优化器逻辑则需要分别检查捕获支持。
一张卡能捕获成功,多卡训练为什么仍可能卡住?
NCCL 要求参与同一次通信的 rank 对是否捕获保持一致。重放包含该通信的图时,这些 rank 也必须使用对应捕获得到的图。
先检查通信顺序和各 rank 选择的执行路径,再检查进程内多 GPU、多 communicator 与混合执行的要求。只在单卡上验证,无法覆盖这些依赖。
推理为什么还要 Piecewise,直接 Full 不是更省开销吗?
Full 可以减少图之间返回主机的开销,但要求整段路径支持捕获。Piecewise 把可捕获的部分交给图,其余部分保留普通执行。
vLLM 会结合 batch 组成和 attention 后端能力选择模式。没有匹配的图时可以回退。评估时同时看请求延迟、捕获时间、padding 成本和显存占用。
捕获多个 stream,会不会把原来的并发变成串行?
图保留的是操作之间的依赖。正确表达 fork、join 和事件等待后,没有相互依赖的节点仍可被并发调度。实际重叠程度还受 GPU 资源限制。
PyTorch 允许多 stream 捕获,但各支路必须按捕获规则建立和汇合依赖。普通图共享私有内存池时,还要满足重放顺序与不可并发重放等内存使用要求。
怎样确认 CUDA Graph 真的缩短了训练步或请求耗时?
先在时间线上确认开销来自 CPU 提交或 kernel 之间的空隙。分别记录捕获与首次启动、稳态重放,以及包含输入复制和结果处理的端到端耗时。
CUDA launch 是异步的。测完整完成时间要等待对应工作结束,不能只计 API 返回时间。比较相同输入、精度和 batch,避免用每个 kernel 都同步的慢基线夸大收益。
真实请求分布下,应该保留哪些捕获尺寸?
待解
想结合命中率、padding 浪费、启动时间和显存预算,为一组实际推理请求选择尺寸档位。