【模型推理】浅谈CUDA Graph(WIP)

Mon, 20 Apr 2026 04:27:43 +0000

CUDA Graph 算是在推理优化里被提到最多的特性之一了。原理倒是不复杂：把原本一次次 launch 的 kernel 打包成一张静态图，一次性扔给 GPU，省掉中间的调度开销。

但这个"打包"具体怎么做？PyTorch 里怎么用？以及最重要的是——在 H200 上到底能快多少？

这篇文章就是把这几个问题摸清楚的过程记录。

1. LLM 推理在忙什么

Graph 能减少 kernel launch 的开销，但为什么 LLM 推理特别需要这个？先简单说下背景。

老生常谈的问题，LLM 推理分两个阶段：

Prefill

输入是用户的一整段 prompt
一次性做完整的 self-attention 计算，然后讲产生的KVCache写入显存或者其他存储设备，从而在Decode阶段使用。
特点是计算密集，kernel消耗的时间比较多，GPU 利用率相对高，同时对内存IO要要求比较低。总之就compute-bound，算力是决定他的速度的关键。

Decode

自回归过程，每次只生成一个 token
要用到之前计算的 KV Cache，只做必要的更新，对于单个请求推理的时候我们输入model的Tensor维度为[1, D], 其中D是Embedding的维度。
特点是memory-bound, 单个请求计算量少，但是需要load KVCache巨大，而且推理的时间越长，KVCache load的压力就越大。

对于这个推理过程，我们可以认为CPU step2step的将模型每一层计算所需要的kernel相关数据和元数据准备好，调用launch通知GPU执行，然后异步的准备下一次计算所需要的信息的过程。

此时因为不同Prefill和Decode不同特点，CPU与GPU直接的协作和等待有不同的表现

Prefill 由于Kernel计算时间比较长，因此CPU侧的开销其实基本可以被GPU的运算Overlap掉，因此关注点核心应该是GPU Kernel如何跑更快。

Decode阶段，其实就是两个地方了，首先Kernel计算量比较小，因此开销时间很少，这种情况就出现了GPU等CPU等情况，这也是就是为什么我们常说Decode阶段需要开CUDA Graph的，原因，CUDA Graph的目标就是消除GPU等CPU的问题，让GPU计算更高效。

如今其实也有主流的推理框架也支持Prefill阶段的CUDA Graph，这个主要是因为在推理过程中模型有大量的小kernel，比如说Norm这类的Kernel，这种其实也类似于Decode阶段的样子，GPU也需要等待CPU，因此CUDA Graph也是可以计算的，至于细节后面我会再写一篇专门介绍Prefill Piecewise CUDA Graph的一些细节。

2. Pytorch在推理的时候CPU纠结在做什么？

因为像Pytorch这种框架一般对于一个算子支持很多Kernel的实现，因此为了支持这么多Kernel的实现就回有大量的dispatch的代码用来运行时自动路由到不同的kernel上，同时也会有很多准备工作，所有工作完成之后就调用cudaLaunchKernel交给CUDA的runtime层做一些准备工作并发送给GPU。在GPU运行的时候其实也会有一些操作，比如说GPU侧做Kernel环境的准备，在执行后也会有一些清理状态的后处理。总结来说就如图，一个Kernel从Pytorch发起执行到做完运行有这几个步骤：

CPU:
1. Pytorch C++ Dispatch
2. Kernel Launcher Prepare
3. cudaLaunchKernel
GPU:
1. Command Processor
2. Kernel Execute
3. Post-Process

2.1. Pytorch Dispatch & Kernel Prepare

Pytorch Dispatch: 这个其实是所有支持多种Kernel+使用Python实现执行前端框架的通病。

CUDA on Code Now

【模型推理】浅谈CUDA Graph(WIP)

1. LLM 推理在忙什么

2. Pytorch在推理的时候CPU纠结在做什么？

2.1. Pytorch Dispatch & Kernel Prepare