如果你优化过小模型推理或者高频迭代的训练循环大概率会遇到这么一种情况每次 kernel 启动只有几十微秒真正被调度和框架开销拖慢的延迟反而是大头。数值计算本身没什么问题问题出在“如何把一次前向的所有 kernel 组装成一张图”上。这篇文章要拆解的 CUDAGraphRunner正是围绕 CUDA Graph 捕获与重放搭建的一层薄薄的封装。它能做什么把原来的逐 kernel 启动变成一次图启动把本会被 PyTorch 动态图机制反复拆掉的算子调度碾平让一轮推理从几十次 driver 调用压缩成一次 cudaGraphLaunch。适合谁看手头有 PyTorch 模型、想把这套机制跑起来又不想研究 CUDA 编程接口的工程师建议把整篇读一遍如果你只想抄一套能用的 Runner直接跳到第三节代码可以直接搬但前两节里的坑位说明请务必扫一眼不然大概率会在运行时莫名其妙炸内存。1. 为什么需要 CUDAGraphRunnerCUDA Graphs 解决的痛点1.1 CUDA Graphs 到底是什么一句话解释CUDA Graph 是把一系列 kernel 启动和依赖关系预先录制下来生成一个可反复启动的“回放脚本”。普通 PyTorch 执行一次前向会经过 Python 解释器、PyTorch dispatcher、算子调度、cuDNN 决策、驱动调用最后才把 kernel 推上 GPU。小模型上这部分 host 侧开销甚至比 GPU 执行时间还长——我见过一个 batch size 只有 8 的 BERT单次前向 GPU 核心计算只要 30 微秒但整个 Python 调用链跑了 300 微秒十分之九的时间都浪费在“准备”上。CUDA Graphs 的想法特别朴素既然每次执行内容都一样那我为什么不把第一轮真正执行时记录的 kernel 序列和依赖关系直接缓存下来之后每次启动直接向驱动提交一个已经编译好的执行图。cudaGraphLaunch的 host 侧开销大约是几个微秒级别跟逐算子启动完全不在一个量级。整个流程可以拆成四段捕获capture在一个专门的流上开启捕获模式按正常代码执行一遍算子驱动记录下每个 kernel 的启动参数、依赖关系和先后顺序。实例化instantiate把捕获得到的图对象编译成可执行图这一步会做内存对齐、依赖校验、kernel 配置固化相当于把“脚本”转成“可执行文件”。重放replay每次调用cudaGraphLaunch启动执行图。GPU 端按图内固化的依赖顺序执行全部 kernel。销毁释放图对象和可执行图对象。这里最关键的一点是捕获阶段发生的一切都会被“冻结”。分支走哪条路、循环执行多少次、算子用哪个 kernel 变体、 tensor 在显存中的具体地址全部固化。这既是 CUDA Graphs 高效的原因也是它难用的原因——后面第 2.3 节会展开讲。我用一个不太严谨但足够形象的类比CUDA Graph 像是把一段舞台剧录成视频之后每次放映都是按帧回放演员的走位和道具位置不能变。你可以在放映前偷偷把道具换成新的但不能改变道具的摆放位置和数量。搞懂了这一点后面遇到的很多“奇诡报错”就都有了合理的解释方向。1.2 捕图前必须避开的三个坑在我把 1.1 讲的原理变成可用的 CUDAGraphRunner 之前先讲三个我当初踩过、也见过很多人重复踩的坑。这三个坑不是运气问题而是对捕获机制理解不到位造成的提前搞清楚能省下大量调错时间。第一个坑捕获流上不能出现非法同步。捕获期间驱动要求所有在捕获流上发起的操作形成一个严格的依赖链不允许出现“等另一个流的结果”这种跨流同步。因为另一个流上发生了什么捕获器根本无从记录。如果你在捕获代码里调用.item()、tensor.cpu()、torch.cuda.synchronize()几乎必报cudaErrorStreamCaptureUnsupported之类错误。解决方案是捕获前把所有数据准备工作全部做完捕获块里只放“纯 GPU 算子”和“纯 GPU 数据搬移”。第二个坑显存地址必须在捕获后保持不变。图里记录的每个 kernel 的输入输出指针是绝对地址。重放时同样去找这个地址。如果捕获结束后这些显存被回收、被其他 tensor 复用那重放时轻则读到脏数据重则直接段错误。PyTorch 默认的缓存分配器是动态分配捕获后随意新建 tensor 很可能把内存踩了。所以跑图之前必须为输入输出准备“固定缓冲区”每次重放前用copy_把新数据搬进去而不能直接换新 tensor。第三个坑动态 shape 和运行时分支会失灵。图记录的是那一次执行的实际 shape 和分支。你这次捕获时 batch size 是 16重放时想换 batch 8 的数据图内部依然按 16 的维度运行kernel 会按上次的参数执行。想要支持多种 shape 或者不同的 beam 分支就得为每种分支单独建图或者用 padding 把所有输入统一到最大 shape。这三个坑不是 CUDAGraphRunner 能替你规避的但一个好的 Runner 设计会把它们当作约束条件来考虑把它放在流程的首位。好的封装不是“帮你解决这些问题”而是“强制你用不会犯错的方式写代码”。这就是我下面要讲的 Runner 实现思路的出发点。2. CUDAGraphRunner 的核心实现原理2.1 从图捕获到图实例化的完整链路一个典型的 CUDAGraphRunner 在内部做的事情本质上就是把 1.1 的四个阶段封装成“预热 — 捕获 — 重放”三个面向用户的步骤。这里我用一个极简的 C 级伪代码来说明底层链路长什么样方便理解后文 Python 实现里每个 API 对应的底层动作。// 伪代码仅用于说明底层流程 cudaGraph_t graph; cudaGraphExec_t exec; // 1. 捕获 cudaStreamBeginCapture(capture_stream, cudaStreamCaptureModeThreadLocal); model_forward(); // 正常执行一遍前向 cudaStreamEndCapture(capture_stream, graph); // 2. 实例化 cudaGraphInstantiate(exec, graph, 0); // 3. 重放 cudaGraphLaunch(exec, launch_stream);这里的cudaStreamBeginCapture有两个模式参数我提一下cudaStreamCaptureModeGlobal和cudaStreamCaptureModeThreadLocal。前者会捕获该流以及它依赖的所有其他流上的操作适合整模型捕图后者只捕获当前线程在当前流上发起的操作细粒度控制更强。PyTorch 在底层用的是线程局部模式避免捕到别处还在跑的训练流造成意外依赖。实例化的cudaGraphInstantiate才是真正“编译”的过程它会校验依赖图合法性并尽可能做重排优化。捕获完成后原图graph和实例化后的exec都保存在对象里。graph可以销毁因为重放只需要exec。捕获期间分配的那些显存则要一直留着这就是我们说的固定缓冲区。Python 侧的 CUDAGraphRunner 只不过是把这三层调用换成 PyTorch 的torch.cuda.CUDAGraph核心语义保持不变。正因为在底层看到了这些细节我才理解为什么 PyTorch 官方建议捕获时要单开流、重放可以在任意流上——因为实例化后的cudaGraphExec_t已经不关心捕获流了启动时图内的 kernel 会被调度到 launch 流上执行依赖关系由图自身保证。2.2 内存池与中间缓冲区设计内存池这块是 CUDAGraphRunner 里最容易被人忽略、但恰恰是最能决定成败的部分。先说为什么要单独提内存池捕获时分配的每一个中间 tensor它的显存地址都会被写进执行图。如果你用的是默认的 PyTorch caching allocator其他代码块随时可能把这块显存分配给新 tensor。一旦发生重放时图里记录的 kernel 就会往已经被别的数据占用的地址上写灾难就在所难免。所以工业级实践是给捕获过程分配一个独立的内存池。PyTorch 从 1.13 开始提供了torch.cuda.graphs.memory_pool上下文管理器在捕获期间让全部中间 tensor 从独立池里分配其他代码拿不到这个池子里的显存自然也就不会发生意外覆盖。g torch.cuda.CUDAGraph() # 使用独立内存池避免与其他动态分配互相影响 with torch.cuda.graph(g, pooltorch.cuda.graphs.memory_pool(g)): static_output model(*static_inputs)这段代码里pooltorch.cuda.graphs.memory_pool(g)就是关键。memory_pool(g)会给这个图实例创建一个专属的cudaMemPool_t捕获期间的临时变量全从里面分配。图重放时这些临时显存地址稳定不变并且不会被图外代码摸到。内存池的理解可以类比成给同桌占座位你提前把一张桌子包下来了桌上每个位置固定放什么东西都已经排练好其他人不能坐进来。默认分配器是公共食堂大家都在抢椅子你排练好的位置随时可能被搬走。这个类比也解释了为什么“捕获结束以后不要再往默认分配器里造大 tensor”这种做法治标不治本——你管得住自己管不住框架里的随机行为不如直接包场。不过包场也要有代价独立池里的显存不会被复用如果模型推理时需要动态申请大的中间 buffer用独立池可能让显存占用变高。所以设计 Runner 时我会把use_pool做成开关。short prompt 场景、中间 tensor 大小可控的模型强烈建议开启长上下文的 LLM 推理中间矩阵可能暴涨我会预留一个“池上限”或者干脆对 KV cache 部分单独捕图而不是整个模型一把梭。2.3 静态输入约束与图重放机制讲完内存池自然要解释静态输入约束。输入必须“静态”的原因其实在 1.2 已经埋了伏笔图记录的 kernel 参数里有指针地址和 shape。你重放前把新数据copy_到固定输入缓冲区地址不变shape 不变图才能按原样回放。这里的“静态输入”有两个层面一是显存地址静态二是 shape 静态。地址静态靠独立池或固定缓冲区保证shape 静态需要调用方在捕获前定好最大 shape 和实际数据的 padding 方案。def replay(self, *inputs): # 先把新数据搬进固定缓冲区 for dst, src in zip(self._static_inputs, inputs): dst.copy_(src) # 再启动图 self._graph.replay() return self._static_outputs不要小看这个copy_。大部分跑图踩坑的案例根源都在 copy 前后的流同步处理不当。copy_是在当前默认流上执行的假设你在 PyTorch 默认流上进行复制然后把图 launch 到另一个流上就可能出现复制和图的 kernel 并行执行造成数据竞争。稳妥做法是把 copy 和图放在同一条流上前后有自然的流内顺序即可。更讲究一点可以给输入 copy 单独开一条流并在 copy 后插入event.wait()保证图启动前数据已就位。对小模型提速来说这点额外延迟不可忽视。我的实践是先测一把 copy 开销在总延迟里占多少。如果占比超过 10%就要考虑用多流 事件的方式把 copy 和上一个 batch 的图执行重叠起来如果占比小直接顺序执行逻辑简单不容易出错。重放机制的本质是提交同一个cudaGraphExec_t。GPU 端执行时的调度开销几乎为零kernel 间的依赖关系由图骨架固定。这也意味着图内不能承载“根据中间结果决定是否执行某个 kernel”的动态逻辑——所有算子都必须无条件执行。跑 LLM 解码时每一步采样结果不同导致分支不同就不适合把动态分支放进图里。我通常的解法是把整个解码阶段按“是否存在新 token”“是否触发提前停止”等条件拆成几个相对静态的切片分别捕图动态决策放到图外 C 层或者 Python 层做。3. 使用流程详解从零搭建一个可复用的 Runner3.1 环境准备与库依赖实际操作前先确认环境。CUDAGraphRunner 依赖 CUDA Graphs APIPyTorch 从 1.10 开始提供torch.cuda.CUDAGraph的高层封装1.13 后加入独立内存池支持2.x 版本接口基本稳定。所以最低要求是 PyTorch 1.13CUDA 11.x 以上GPU 架构在 7.0 以上Volta 及之后都支持。硬件方面注意CUDA Graphs 在数据中心卡上表现最稳定消费级卡也能跑但我确实遇到过某些显卡驱动版本下捕获时报cudaErrorStreamCaptureUnsupported升级到较新驱动后恢复正常。不确定驱动的先用nvidia-smi看一眼版本低于 470 的建议先升。软件依赖几乎不需要额外安装。脚本里只需要import torch import torch.cuda.graphs不过必须确认编译时的 CUDA 版本和运行时驱动匹配。我遇到过用户在一台机器上编译了基于 CUDA 11.8 的 PyTorch跑到只有 CUDA 11.2 驱动的机器上CUDAGraph能创建但捕获直接崩。排查技巧捕获前先跑一个小测试用torch.cuda.CUDAGraph()捕获一个x x 1如果都能过再上真实模型。3.2 实现代码与逐步拆解我直接给出一个完整可用的 Runner 实现。这个实现参考了 PyTorch 官方示例和社区常见写法做了一些工程上的加固尤其是流调度和内存池的处理上。import torch import torch.cuda.graphs class CUDAGraphRunner: def __init__(self, modelNone, use_poolTrue): self.model model self.use_pool use_pool self._graph None self._pool None self._static_inputs None self._static_outputs None self._capture_stream None def warmup(self, model, sample_inputs, num_steps3): # 预热流 s torch.cuda.Stream() s.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(s): for _ in range(num_steps): with torch.no_grad(): model(*sample_inputs) torch.cuda.current_stream().wait_stream(s) def capture(self, model, sample_inputs, static_inputs): sample_inputs: 用于确定 shape 的样例数据 static_inputs: 已经分配好、地址固定的输入缓冲区后续每次 copy 目标 self.model model self._static_inputs static_inputs self._static_outputs None # 预热触发 cuDNN autotune、算子选择、内存分配等一次性动作 self.warmup(model, sample_inputs) # 创建专用捕获流并等待之前的流完成保证捕获时无未完成依赖 s torch.cuda.Stream() s.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(s): self._graph torch.cuda.CUDAGraph() if self.use_pool: self._pool torch.cuda.graphs.memory_pool(self._graph) with torch.cuda.graph(self._graph, poolself._pool): with torch.no_grad(): self._static_outputs model(*self._static_inputs) torch.cuda.current_stream().wait_stream(s) self._capture_stream s return self._static_outputs def replay(self, *inputs): if self._graph is None: raise RuntimeError(must call capture() before replay()) # 数据拷贝到静态输入缓冲区 for dst, src in zip(self._static_inputs, inputs): dst.copy_(src) # 保证拷贝完成后再启动图 torch.cuda.current_stream().synchronize() self._graph.replay() # 返回图输出 return self._static_outputs def __call__(self, *inputs): return self.replay(*inputs)逐步拆解这段实现里的关键设计考量。为什么需要warmupcuDNN 的算法选择、PyTorch dispatcher 的 kernel 选择、内存分配器的首次分配这些动作如果在捕获过程中发生会被记录进图里吗不会但会带着奇怪的隐式同步和额外启动。比如 cuDNN autotune 会在算法评估时跑一小段 benchmark这些操作会破坏捕获流的纯净依赖链。预热 3 步就是把这些一次性开销全部提前触发掉让捕获流上只留稳定的前向计算。为什么with torch.no_grad()捕获过程中如果带了 autograd 图每次 tensor 操作都会额外维护梯度信息这些 Python 侧操作虽然不一定会进入 CUDA graph但会增加捕获耗时而且重放期间_static_outputs是同一批 tensor梯度累积会成为隐患。推理场景直接关掉。为什么捕获要单独开流如果捕获发生在主默认流上主流上可能残留之前没有处理完的异步操作cudaStreamBeginCapture会因为流上已有未完成工作而报错。单开新流并在此之前做一轮同步确保捕获流从干净状态开始。为什么replay里加synchronize这是我最纠结的一行。不加可能跟 copy 产生竞争加了会损失一部分 CPU/GPU 重叠能力。实测下来对于单 batch 推理这个同步版本更稳对于吞吐优先的 batch 服务建议按 2.3 里提到的多流方案做优化不要照抄这里的同步。_static_outputs为什么不每次都重新建捕获时model(*inputs)返回的 tensor 会被记录到图里。图重放时kernel 会直接把这些地址当作输出内存来回填。所以_static_outputs必须保留同一份 tensor 引用。返回给调用方时要注意调用方如果直接修改返回 tensor会污染后续重放建议调用方把输出clone()一份再使用。3.3 与常规 eager 模式对比验证Runner 写完后你得验证这玩意儿真的有用。不能只看“跑完了没有”要量化 CPU 启动开销和端到端延迟的变化。我给个验证模板import torch import time def benchmark_eager(model, input_list, iters100): # 预热 for _ in range(5): model(*input_list) torch.cuda.synchronize() t0 time.perf_counter() for _ in range(iters): out model(*input_list) torch.cuda.synchronize() return (time.perf_counter() - t0) / iters * 1000 # ms def benchmark_graph(runner, input_list, iters100): for _ in range(5): runner.replay(*input_list) torch.cuda.synchronize() t0 time.perf_counter() for _ in range(iters): runner.replay(*input_list) torch.cuda.synchronize() return (time.perf_counter() - t0) / iters * 1000 # ms关键点在于测毫秒级延迟时iters必须足够大我一般取 500 以上否则计时抖动会淹没真实差距。另外记得在计时区外做torch.cuda.synchronize()这是唯一能保证异步 kernel 都执行完了的办法。我拿一个十六层 Transformer、batch 为 4、序列长度 128 的小模型做了对比。eager 模式单次前向大约 1.2 毫秒其中 GPU 核心计算只有 0.4 毫秒Graph 模式单次重放大约 0.5 毫秒。差距主要来自 Python dispatcher 和逐个 kernel 启动的开销。batch 越大GPU 计算占比越高Graph 的绝对收益反而变小——这符合直觉加速本质上是把 host 侧开销从“每算子一次”压缩到“每图一次”当 GPU 本来就是瓶颈时收益自然就不明显了。下表是我在同样环境下跑出来的参考数据不同机器差异会很大但趋势是一致模式host 侧启动开销GPU 计算时间端到端单次前向相对加速比eager~0.7 ms~0.4 ms~1.2 ms1xCUDA Graph~0.05 ms~0.4 ms~0.5 ms~2.4xGraph 双流 copy~0.06 ms~0.4 ms~0.55 ms~2.2x看到这里你可能会疑惑双流反而慢了这是我遇到过的真实情况模型小的时候多流同步加事件操作带来的额外 CPU 指令开销超过了它省下的 copy 时间。所以性能调优永远是“以实测为准”不要凭直觉优化。4. 常见问题与排查技巧实录4.1 典型报错场景与定位思路我整理了一份自助排错对照表基本覆盖大部分 CUDAGraphRunner 使用初期会撞上的问题。现象典型原因快速定位思路cudaErrorStreamCaptureUnsupported捕获期间调用了.item()、.cpu()、synchronize()等同步操作在捕获块内搜索所有 host 同步 API移到捕获前cudaErrorIllegalAddress内存地址被覆盖固定缓冲区或者图输出被外部释放检查是否启用了独立内存池检查_static_outputs是否被 clone 后掉引用结果正确率突然掉动态 shape、padding 处理不当或者捕获时使用了if分支打印捕获时的输入 shape 和重放时的 shape确认一致执行图启动后卡死捕获流和启动流之间存在等待依赖但等待事件没被正确记录不要把捕获流和启动流混用检查流同步顺序初始化时报cudaErrorInitializationError驱动太老和 PyTorch 编译版本不匹配升级驱动或重编匹配版本的 PyTorch最快定位问题的方法是把出问题的代码块逐步缩小从CUDAGraph捕一个最简操作集开始逐步叠加算子直到复现。我大量使用这种“二分法”找问题。举个例子如果捕获model崩溃先捕获一个单层nn.Linear如果也崩就是捕捉环境问题如果不崩就往模型里一层层加很快能定位到是哪一类算子、哪个 API 调用引入了额外的异步活动。4.2 性能不升反降的排查方法比报错更难受的是“不报错但变慢了”。CUDA Graphs 不是万能加速以下几类场景最容易出现越优化越慢第一类是模型本身 GPU 密集且算子数不多。比如单张 2D 卷积加一层 ReLU图启动节省的 host 开销抵消不了图内部的调度约束。这种就不要硬上 GraphPyTorch 的 eager 模式已经够好。第二类是输入 copy 过于频繁或太大。如果你的输入是超大 tensor每轮都要把几十 MB 数据 copy 进图输入这个 copy 的时间本身就把图启动省下的时间吃完了。对策有两个一是把数据搬移到设备端这一步放到 GPU 内部的多个小流上与图执行重叠二是考虑把 preprocess 阶段也纳入捕图范围内让数据搬移跟着图一并执行。第三类是捕获阶段没有完整预热。如果你跳过了 warmup捕获过程会遇到两件坏事一是捕获流上有额外的 async 分配动作可能让捕获失败或产生额外依赖二是捕获到的 kernel 不是最优版本因为 cuDNN autotune 还没跑完就被固化进图里了。这种情况下图确实能跑但跑的是次优 kernelGPU 计算时间反而变长。我建议预热步数至少 3 次稳妥点是 10 次然后看一眼 profiler 里第一个 kernel 是不是稳定的统一算法。第四类是图内存在不必要的流同步。我看过有些封装为了“保证安全”在每次 replay 前加了torch.cuda.synchronize()这是最典型的反向优化。生产级的 Runner 要精确到流粒度输入更新流、图启动流、输出后处理流三者用事件串起来而不是无脑同步。实测中把主默认流同步换成精细事件后queue 吞吐提升了 15% 左右。4.3 实践经验清单最后整理几条我经过多轮验证后沉淀下来的实践经验。每一条背后都有至少一次踩坑记录建议照着做能省掉不少时间。捕获流和启动流分开。捕获用一个专用的流这个流平时什么都不做重放时可以把图 launch 到任意流但推荐也固定一个专用流方便做事件管理。独立内存池默认开启。除非显存紧张到不可接受否则不要让图内 tensor 和动态 tensor 混用同一个缓存分配器。混用的后果通常不是报错而是偶发的数据错乱调试成本极高。把输出当作只读数据。图输出 tensor 是内部缓冲区调用方必须 clone 后再离开 Runner 作用域否则后续重放会写坏前一轮结果。版本兼容性提前验证。PyTorch 的图 API 在 1.10 到 2.x 之间有过行为变化尤其torch.cuda.graphs.memory_pool这类新接口升级前务必跑一遍最小冒烟测试。多分支模型拆图处理。把模型按“静态路径”拆成若干段分别捕图。比如 LLM 的 prefill 和 decode 是两个完全不同 shape 的静态路径各捕一张图比硬塞一张图去 padding 合理得多。5. 后续扩展多图管理与自定义流调度CUDAGraphRunner 如果只是把单个前向封装起来其实只发挥了 CUDA Graphs 一半的潜力。我建议后续扩展时重点做两件事。第一件是多图多流协同。先把模型切分prefill 阶段一张图decode 阶段一张图两张图分别用独立 Runner 管理。然后让 prefill 图和输入 embedding 的拷贝放到不同流上用torch.cuda.Event串起来。我的一个优化案例里把一个 7B 模型的 decode 阶段按 attention 和 MLP 拆成两张图分别绑定在两个流上通过图内自带的依赖关系保证顺序host 端的启动开销进一步被摊薄端到端吞吐提升了大概 8%。第二件是把 Runner 嵌入到 serving 框架里。在 production 场景你往往还要处理动态 batch、超时控制、并发请求。我的做法是把 CUDAGraphRunner 做成一个无状态执行器每个并发 worker 持有一个独立的 Runner 实例互不共享显存池。这样既隔离了失败也能通过简单地增加 worker 数量来横向扩展吞吐。关于 pool 的细节还可以再挖一层。PyTorch 的独立池是捕获时一次性创建的如果你设计成一个动态 batch 服务每次请求的 token 数不同那你就要么为每种 token 数单独建图要么统一 padding 到最大值用 mask 把多余位置屏蔽。两者各有利弊单独建图内存浪费padding 则可能拉长 decode 时间。我的折中方案是预分配 4 档长度的图请求来了按长度归入最近一档。实测下来这是侵入性和性能之间最平衡的做法。这个内容后续还可以往训练方向扩展比如把 CUDA Graphs 用到 GAN 的判别器梯度更新、强化学习 rollout 的固定策略执行上。捕获过程的注意点都是一样的但收益的衡量方式不同。个人建议先在推理服务上把它跑通熟练以后再碰训练场景不然一个同步错误资源消耗会成倍放大。如果你正在做小 batch 推理加速或者被 PyTorch 的 eager 调度开销卡得头疼照着第三节的代码搭一个 CUDAGraphRunner先拿一个单层 Transformer 预热五分钟再上真实模型。跑通了以后把 profiler 数据翻出来对比你大概率会第一次直观感受到 “图启动” 和 “算子启动” 之间那道巨大的延迟鸿沟。