CUDA Graph 的执行机制与动态性

本文最后更新于 2026年9月13日

关于 CUDA Graph，一个常见的第一印象是“不可变”：输入形状不可变、内存地址不可变、控制流不可变。但如果只记住这些限制，就很容易产生一些看似合理的推论。例如，LLM 的 decode 阶段上下文长度不断增长，因此无法使用 CUDA Graph。

实际情况恰好相反，decode 是推理框架使用 CUDA Graph 的典型场景。

本文先从一次 kernel launch 的提交过程出发，解释 CUDA Graph 保存了什么，再讨论其动态性、LLM attention 的 capture，以及 PyTorch 如何管理 capture 期间产生的临时显存。

一次 kernel launch 提交了什么

CPU 启动 kernel 时，提交的是“执行这段 GPU 程序所需的工作描述”。理解这个过程，需要区分三类内容：

内容 例子 何时准备
kernel 代码 load、store、矩阵乘等机器指令 编译、加载时
一次执行的描述 程序入口、grid、block、参数、资源配置 launch 时
运行时数据 输入 tensor、权重、KV cache 执行前准备，执行中读取

例如：

1
scale<<<4, 256, 0, stream>>>(x_ptr, y_ptr, 1024, 2.0f);

这里需要指定运行哪个 kernel、启动多少线程，以及四个实参。对于 x_ptr，传入的是地址值；它指向的整个 tensor 不会随这次 launch 自动复制。kernel 代码通常也可以在加载后重复使用，无须每次重新传输。参见 CUDA Driver API。

下面以 NVIDIA GPU 上普通的 host launch 为例，分别讨论代码加载、参数准备和驱动提交。

代码加载与 kernel 启动

首先需要区分 kernel 的代码与某一次执行它的请求。

CUDA C++ 或 Triton 产生的 device 代码，最终需要转换为目标 GPU 能执行的机器指令。在 NVIDIA 平台上，PTX 是虚拟指令集，cubin 中包含面向具体架构的机器码，反汇编时通常称为 SASS。如果加载的是 PTX，驱动可以在加载过程中将其 JIT 编译为机器码；如果已经有适合当前架构的 cubin，则可以加载该二进制。

Driver API 用 module 和 function handle 表达这些对象。例如：

1
2
3
4
5
// 假设 CUDA context 已初始化，scale.cubin 中包含名为 scale 的 kernel。
CUmodule module;
CUfunction function;
cuModuleLoad(&module, "scale.cubin");
cuModuleGetFunction(&function, module, "scale");

这一步使驱动能够找到 kernel 的代码和相关 metadata。后续可以使用同一个 function handle 多次 launch，并不需要每次重新编译和传输整份 kernel 代码。lazy loading 可能将部分准备工作推迟到首次使用，但不改变代码加载与重复执行的区别。参见 CUDA Driver API 的 module 与 kernel execution。

参数如何传给 GPU

假设 scale 的签名是：

1
extern "C" __global__ void scale(const float* x, float* y, int n, float alpha);

在 Driver API 层，一次启动可以写成：

1
2
3
4
5
6
7
8
9
10
11
12
13
// d_x、d_y 是 CUdeviceptr 类型的 host 变量，保存已分配的 device buffer 地址。
// stream 是有效的 CUstream。
int n = 1024;
float alpha = 2.0f;
void* args[] = {&d_x, &d_y, &n, &alpha};

cuLaunchKernel(
    function,
    4, 1, 1,       // grid dimensions
    256, 1, 1,     // block dimensions
    0, stream,     // dynamic shared memory、stream
    args, nullptr
);

这里的 extern "C" 保证导出的 kernel 符号名为 scale，与前面的 cuModuleGetFunction 调用一致。cuLaunchKernel 在 CPU 上执行。args 中的 &d_x 告诉驱动从哪个 host 地址读取实参；读取到的 d_x 才是传给 kernel 的 device 地址值。参见 cuLaunchKernel 参数说明。

1
2
3
4
5
6
7
CPU 内存                           GPU

&d_x ──→ d_x = 0xA000
                  │
                  └──复制参数值──→ x = 0xA000
                                         │
                                         └──读取──→ device buffer 中的数据

图中的地址仅作示意。n 和 alpha 同样由 host 按值传入，因此 kernel 参数不必都是 device 指针。需要区分的是：传递指针值与搬运指针指向的数组，是两件不同的事。

从提交到执行

驱动将函数、参数和执行配置组织为 GPU 可以处理的工作描述，再通知 GPU 执行。这里需要区分启动计算的命令与计算本身的指令：

内容 作用
kernel 机器指令 描述线程执行哪些 load、store、算术和控制流操作
提交命令与工作描述符 指定运行哪段代码、采用什么执行配置、使用哪些参数与依赖 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
CPU / host
    框架调用或 CUDA C++ launch
          ↓
    Runtime / Driver：准备参数与工作描述
          ↓
    提交工作，通知 GPU
          │
          ▼
GPU
    命令前端取得工作，处理相关依赖
          ↓
    将 thread blocks 分派给具备资源的 SM
          ↓
    SM 内的 warp scheduler 发射就绪指令
          ↓
    执行 kernel 的 load / store / compute

host 指定 grid 和 block，具体的 block 分派和 warp 调度由 GPU 完成。因此，CPU 提交工作所需的时间，与 GPU 执行 kernel 的时间需要分别考察。

CUDA Graph 记录了什么

在普通的 eager 执行中，CPU 依次调用算子，经过框架的 dispatcher 和 CUDA runtime，将 kernel 提交到 GPU。对于计算时间很短的小算子，host 端的开销可能成为显著的性能瓶颈。

CUDA Graph 将一组 GPU 工作及其依赖关系预先定义为图，实例化后通过一次 graph launch 提交执行。stream capture 是构建这张图的一种方式，另外也可以通过显式的 Graph API 添加节点。CUDA 将这个过程划分为 definition、instantiation 和 execution 三个阶段，PyTorch 通常将最后一个阶段称为 replay。参见 CUDA Graphs 文档。

对于一段简单的逐元素计算，可以将其理解为：

1
2
3
4
5
6
7
8
eager：
    CPU 提交 add → CPU 提交 relu → CPU 提交 mul

capture 后：
    graph = [add → relu → mul]

replay：
    CPU 提交 graph

这里的三个 kernel 仍然是三个 kernel。CUDA Graph 本身不会将它们融合成一个 kernel，也不会自动消除中间结果的显存读写。kernel fusion 与 CUDA Graph 可以配合使用，但解决的是不同层面的问题。

API 层的图与执行时的表示

kernel 节点的参数结构 cudaKernelNodeParams 包含 func、gridDim、blockDim、sharedMemBytes，以及 kernelParams / extra 所描述的实参。节点间的依赖另外表达。参见 cudaKernelNodeParams。

以两个 kernel 为例，图在语义上可以表示为：

1
2
3
4
5
6
7
node A：scale(x_ptr, tmp_ptr, 1024, 2.0)
        grid=(4,1,1), block=(256,1,1)

node B：relu(tmp_ptr, y_ptr, 1024)
        grid=(4,1,1), block=(256,1,1)

dependency：A → B

cudaGraph_t 是图定义的句柄，对应上述节点、参数和依赖。调用 cudaGraphInstantiate 后，CUDA 校验图并完成执行所需的准备，返回一个可执行图实例的句柄 cudaGraphExec_t。后续 cudaGraphLaunch 接收的是这个实例，可以反复启动它：

1
2
3
4
5
6
// graph 是已经建好的 cudaGraph_t，stream 是执行所用的 CUDA stream。
cudaGraphExec_t executable;
cudaGraphInstantiate(&executable, graph, 0);

cudaGraphLaunch(executable, stream);
cudaGraphLaunch(executable, stream);  // 复用同一个实例，再执行一次
1
2
3
4
5
6
7
8
9
capture / 显式建图
    ↓
cudaGraph_t：图定义（节点、参数、依赖）
    ↓ instantiate
cudaGraphExec_t：可执行图实例，供后续 launch 使用
    ↓ upload（可以合并到首次 launch）
device 端可用的工作描述
    ↓ launch / replay
按依赖执行各节点

“不可变”的具体含义

指针值与指针指向的内容

考虑下面的 CUDA kernel：

1
2
3
4
5
6
extern "C" __global__ void scale(const float* x, float* y, int n, float alpha) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        y[i] = x[i] * alpha;
    }
}

假设 capture 时的调用是：

1
scale<<<4, 256, 0, stream>>>(x_ptr, y_ptr, 1024, 2.0f);

那么 graph 记录的是这次 launch：使用哪个 kernel、grid 和 block 的大小，以及 x_ptr、y_ptr、1024、2.0f 这些参数值。

每次 replay 时，kernel 都会从同一个 x_ptr 读取数据，向同一个 y_ptr 写入结果。但 x_ptr 指向的 1024 个浮点数可以改变。graph 保存了指针值，并没有保存这块显存内容的快照。

因此，修改输入的正确方式是在 replay 之前更新原有 buffer。下面假设 new_x 已准备完成，拷贝与 replay 在同一 stream 上提交：

1
2
static_x.copy_(new_x)
g.replay()

而下面的赋值只改变 Python 变量的绑定关系：

1
2
static_x = new_x
g.replay()

graph 不会重新查找 static_x 这个名字，也不会因此改用 new_x 的地址。如果原来的输入 tensor 失去引用，其内存还可能被 allocator 回收复用。因此，capture 之外分配的输入、权重等 buffer，需要由调用方保证生命周期。

标量也可以变成运行时数据

上面的 alpha 是一个按值传入的标量，capture 时已经固定为 2.0f。即使随后修改 host 端的 alpha 变量，graph 中的参数也不会随之变化。

如果需要每次 replay 使用不同的缩放系数，可以改为：

1
2
3
4
5
6
7
8
__global__ void scale_dynamic(
    const float* x, float* y, int n, const float* alpha_ptr
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        y[i] = x[i] * alpha_ptr[0];
    }
}

此时固定的是 alpha_ptr，缩放系数则由 GPU 在执行时从显存读取。只要在 replay 之前更新该地址的内容，并保证更新与 replay 之间的执行顺序，系数就可以变化。

同一个数值，作为 host 标量传入，或者作为 device buffer 的内容传入，对 CUDA Graph 来说具有不同的含义。后面讨论的 sequence length、block table 和条件判断，都与这个区别有关。

固定 shape 与可变的有效长度

继续使用前面的 scale。假设输入 tensor 的 shape 是 [1024]，host 根据元素数量计算 grid，再启动 kernel：

1
2
3
int n = 1024;                    // 来自输入 tensor 的元素数量
int blocks = (n + 255) / 256;     // 每个线程处理一个元素
scale<<<blocks, 256, 0, stream>>>(x_ptr, y_ptr, n, 2.0f);

capture 得到的是 grid=4、n=1024 的这次 launch。之后 replay 不会重新执行上面计算 n 和 blocks 的 host 代码。

如果输入改为 [2048]，普通执行会重新计算，得到 grid=8、n=2048。但原 graph 仍然启动 4 个 block，并向 kernel 传入 1024：它只处理前 1024 个元素。即使原来的输入输出 buffer 足够大，这张图也不会自动扩大计算范围。

输入 shape 普通执行时计算出的 launch 原 graph 的 replay
[1024] 4 个 block，n=1024 4 个 block，n=1024
[2048] 8 个 block，n=2048 仍为 4 个 block，n=1024

这就是 shape 限制的具体来源：由 shape 决定的 launch 配置和参数，在 capture 时已经确定。 对普通 PyTorch CUDA Graph，输入需要保持 capture 时的 size 和 layout，参见 PyTorch 的约束说明。

padding 与 kernel 内部的分支

CUDA Graph 并不要求 kernel 内部每次执行完全相同的指令路径。

通过 padding，可以将不同长度的输入放入相同 shape 的 buffer。为便于展示，假设 buffer 容量固定为 8，两次输入的有效长度分别为 5 和 7：

1
2
3
4
5
6
7
下标：          0  1  2  3  4  5  6  7

第一次输入：   [1, 2, 3, 4, 5, 0, 0, 0]   有效长度 = 5
               └──有效数据──┘ └padding┘

第二次输入：   [1, 2, 3, 4, 5, 6, 7, 0]   有效长度 = 7
               └─────有效数据─────┘  padding

本例约定 0 表示 padding，有效数据均非零。kernel 读取输入值，遇到 padding 时直接返回：

1
2
3
4
5
6
7
8
9
10
11
12
13
__global__ void scale_padded(
    const int* x, int* y, int alpha
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int value = x[i];
    if (value == 0) {
        return;  // padding 位置不参与计算
    }
    y[i] = value * alpha;
}

// 仅为演示：固定启动 8 个线程，每个线程对应一个元素。
scale_padded<<<1, 8, 0, stream>>>(x_ptr, y_ptr, 2);
第一次 replay 第二次 replay
输入输出 tensor 的 shape [8] [8]
启动的线程数 8 8
x[5] 的值 0（padding） 6（有效数据）
下标为 5 的线程 读到 0，直接返回 读到 6，计算 y[5] = 6 × 2

第一次 replay 中，i = 5 的线程读到 0 后直接返回；第二次 replay 中，i = 5 的线程读到 6 并执行乘法。分支条件由 GPU 执行时读取的输入值决定，无须修改 graph。

调用方只使用有效数据对应的输出，padding 位置不会被此 kernel 更新。

为什么 decode 可以使用 CUDA Graph

首先需要明确 LLM 如何组织计算。用 T 表示本轮输入的总 token 数，H 表示 hidden size，hidden states 通常表示为 [T, H]：每一行对应一个 token 的隐藏向量。

常见 Transformer 中的 Linear、RMSNorm 和 MLP 都是 token level 的运算：对每个 token 的隐藏向量独立应用相同的变换。 例如，Linear 的矩阵乘 Y = X @ W 可以按行理解为：

1
2
3
4
5
6
X：[T, H]           W：[H, D]          Y：[T, D]

token 0 的向量  ────── @ W ──────→ token 0 的输出
token 1 的向量  ────── @ W ──────→ token 1 的输出
token 2 的向量  ────── @ W ──────→ token 2 的输出
      ...                              ...

这些 token 共享权重 W，但每一行的结果不依赖其他行。实现时通常将所有行合并为一次矩阵乘，以提高 GPU 利用率。RMSNorm 则沿每行的 hidden dimension 归一化，MLP 也分别变换每个 token 的向量。因此，对这些算子而言，T 个 token 来自一条请求还是多条请求，不影响逐 token 的计算方式。

attention 负责 token 之间的信息交互：每个当前 token 的 query，需要读取该请求中允许访问的 token 的 K/V。因此，attention 除了当前输入，还需要知道请求边界和历史 KV 长度。

在使用 KV cache 的普通单 token decode 中，每条请求每轮只输入一个新 token。Linear、RMSNorm、MLP 等算子只处理本轮的 T 个 token；attention 通过 KV cache 读取历史信息，无须重新计算历史 token 的这些逐 token 变换。只要本轮 token 数 T 固定，当前 hidden states 就保持为 [T, H]，即使历史 KV 长度仍在增长。

例如，连续两轮 decode 的状态可以是：

1
2
3
                    第 t 轮              第 t+1 轮
hidden states       [4, H]               [4, H]
sequence lengths    [128, 256, 80, 512]  [129, 257, 81, 513]

sequence lengths tensor 的 shape 始终为 [4]，变化的是其中的四个整数。

对于采用 paged KV cache 的实现，KV cache 可以来自预先分配的物理 block 池。逻辑序列访问哪些 block，由 block table 描述。跨越一个 block 边界时，框架可以分配新的逻辑映射，更新表项，而不移动整个 KV cache。

attention kernel 启动时，传入的是 sequence lengths buffer 的指针，而不是其中每条请求的长度值。capture 固定了这个指针；每轮 decode 只需更新 buffer 中的长度，例如从 [128, 256, 80, 512] 更新为 [129, 257, 81, 513]。kernel 在 GPU 上执行时，再从同一个地址读取本轮的长度，确定需要访问多少历史 KV。

sequence lengths 的内容改变，不需要更新 kernel 的启动参数。 在 kernel、grid、block 等其余启动配置也保持不变的前提下，不同 sequence length 就可以复用同一张 CUDA Graph。

block table 和 slot mapping 也可以采用相同的方式：保持 buffer 地址不变，通过更新表项指定读取哪些 KV block，以及将新 token 的 K/V 写入哪里。

下面是这一类实现的数据流示意，具体字段和布局由 backend 决定：

1
2
3
4
5
固定地址的 hidden states ───────────────┐
固定地址的 KV cache block 池 ───────────┤
固定地址、内容可更新的 sequence lengths ─┼─→ attention kernel
固定地址、内容可更新的 block table ─────┤
固定地址、内容可更新的 slot mapping ────┘

这并不意味着所有 attention kernel 都天然支持这种用法。如果某个实现根据 host 端的最大序列长度选择 split 数量、分配 workspace，或者改变 launch grid，那么仅更新 device metadata 就不够了。backend 仍需要提供与 CUDA Graph 兼容的调度方式。

再考虑 T 的变化。continuous batching 中，请求会加入或退出，每轮输入的 token 数也会随之变化。常见做法是预先为若干 T capture graph，例如：

1
2
3
capture sizes：1, 2, 4, 8, 16, 32
实际 T：     13
选择 graph： T = 16

框架将输入填充到 16，并通过有效长度、mask 或安全的占位 metadata 处理多出的条目。尤其是 KV cache 写入，padding 请求必须具有正确的屏蔽语义。

因此，decode 的动态性通常由两部分共同承担：graph 外部选择适合当前 T 的图，graph 内部通过 device metadata 处理变化的序列状态。比如vLLM 的 batch descriptor 和 CUDA Graph dispatcher 就是在这一层管理 capture 与 replay，参见其 CUDA Graphs 设计文档。

prefill attention 的限制来自哪里

普通 decode 中，每条请求本轮只有一个 query token；prefill 中，每条请求本轮的 query token 数也会变化。即使总 token 数 T 相同，attention 的 launch grid 仍可能不同。

例如，两次 prefill 都有两条请求、总共 256 个 token，但请求长度分别是 [128, 128] 和 [192, 64]。Linear、RMSNorm、MLP 的输入都可以保持为 [256, H]；attention 则需要按请求处理各自的 query。

以 FlashAttention 2 的普通 forward 路径为例，host 按最长 query 长度计算分块数，再为每条请求、每个 head 启动对应数量的 block。假设选用的 kernel 每个 block 处理 128 个 query token，则有：

第一次 prefill 第二次 prefill
两条请求的 query 长度 [128, 128] [192, 64]
总 token 数 T 256 256
最长 query 长度 128 192
每条请求、每个 head 启动的 block 数 1 2

这里的分块数按 ceil(最长 query 长度 / 128) 计算；128 是本例选取的 query 分块大小。对应逻辑见 FlashAttention 的 grid 计算，最长 query 长度由 host API 传入。

如果按第一组输入 capture，图中为每条请求记录的只有一个 block。第二组输入中，第一条请求需要两个 block；直接 replay 原图，就缺少处理后 64 个 query token 的 block。更新 device buffer 中的请求长度和边界，不会改变已经捕获的 grid。

这就是仅按 T 复用 prefill graph 可能失败的具体原因。要让这两组输入复用同一张图，可以提前按足够大的 query 长度上限建图，让 kernel 从 device metadata 读取实际长度并跳过多余分块，同时保证其他启动参数和 buffer 地址也兼容。

如果 attention 路径仍需要 CPU 每轮重新计算 launch，框架可以采用 piecewise CUDA Graph：只捕获 attention 前后的计算，attention 每轮正常启动。

动态控制流与 torch.cond

前面讨论的是同一个 kernel 内部的动态性。如果数据决定的是执行哪一组 kernel，就涉及 graph 层面的控制流。

例如：

1
2
3
4
if pred.item():
    y = x.sin()
else:
    y = x.cos()

假设 pred 是 CUDA tensor，这里的 .item() 需要将结果交给 CPU，Python 才能决定下一步调用哪个算子。这种依赖 GPU 结果的 host 分支不能直接放入普通 stream capture。

如果 pred 是 Python bool，capture 时可以完成分支判断，但只会记录当时实际调用的分支。之后修改这个 bool，replay 也不会重新执行 Python 的 if。

CUDA 的 conditional nodes 提供了另一种表达方式：预先构建分支或循环体，由 device 端的条件决定执行哪部分图。条件可以通过 kernel 中的 cudaGraphSetConditional 设置，判断过程不需要将控制权交回 CPU。参见 NVIDIA 的 Conditional Nodes 介绍。

PyTorch 中对应的高层表达是 torch.cond。这里使用 torch.compile，是为了把 torch.cond 表达的条件分支转换成 CUDA Graph 的条件节点。

下面将分支前的 x * 2、条件分支和分支后的 z * 3 写入同一个 forward，再对整个 forward 使用 cudagraphs backend。capture 的范围是整个 forward，包括分支前后的计算。 backend 在首次几次调用中自动完成 warmup 和 capture，后续调用则 replay 整张图：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import torch


def positive_branch(h):
    return h + 1


def negative_branch(h):
    return h - 1


def forward(pred, x):
    h = x * 2
    z = torch.cond(pred, positive_branch, negative_branch, (h,))
    return z * 3


compiled_forward = torch.compile(forward, backend="cudagraphs", fullgraph=True)

with torch.inference_mode():
    x = torch.ones(8, device="cuda")
    pred = torch.tensor(True, device="cuda")

    # backend 在这些调用中完成整个 forward 的 warmup、capture 和首次 replay。
    for _ in range(3):
        compiled_forward(pred, x)

    pred.fill_(True)
    y = compiled_forward(pred, x)   # replay，结果为 8 个 9.0

    pred.fill_(False)
    y = compiled_forward(pred, x)   # replay 同一张图，结果为 8 个 3.0

对应的整张 CUDA Graph 可以表示为：

1
2
3
4
5
6
                    forward 的 CUDA Graph
┌────────────────────────────────────────────────────┐
│                  ┌─ true：z = h + 1 ──┐             │
│ h = x * 2 → 条件 ┤                    ├→ y = z * 3  │
│                  └─ false：z = h - 1 ─┘             │
└────────────────────────────────────────────────────┘

两次调用更新的是同一个 device bool tensor pred 的内容。compiled_forward(pred, x) 内部通过一次 graph launch 执行分支前的计算、选中的分支和分支后的计算。示例没有显式的 with torch.cuda.graph(g):，是因为整张图的 capture 和 replay 都由 backend 管理。

普通 CUDA Graph capture 本身不要求 torch.compile。在本文实测的 PyTorch 2.13.0+cu130 中，仅将 torch.cond 放入 with torch.cuda.graph(g):，仍会在读取条件时触发 capture 错误；提前 warmup 两个分支也不能改变这条执行路径。上面的 backend 用法参见 Data Dependent Control Flow。

torch.cond 保留了两个分支的结构，使运行时选择可以在 device 端完成。示例中的两个分支都产生相同 shape、dtype 和 layout 的输出，这也使后续计算能够使用一致的 buffer 约定。

它并不意味着任意 Python 控制流或任意动态输出形状都能进入 CUDA Graph。更准确的理解是：分支结构提前建立，每次 replay 可以选择不同的执行路径。

capture 期间的执行语义

“执行”在这里需要区分 host 与 device。

在 with torch.cuda.graph(g): 内，Python 代码仍然执行，算子的 host 端实现也会被调用。框架需要计算 metadata、选择 kernel、分配输出 tensor，并调用 CUDA API，这样才能知道应当向图中添加哪些节点。

但提交到正在 capture 的 stream 中的 GPU 工作不会在此时执行，而是被记录为 graph 节点。capture_end 得到图，后续 graph launch 才会执行其中的工作。这是 CUDA stream capture 的定义。

下面的例子可以同时观察这两个层面的行为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import torch


static_x = torch.zeros(4, device="cuda")
static_y = torch.full_like(static_x, -1)

# 在 side stream 上 warmup，使算子完成必要的初始化。
warmup_stream = torch.cuda.Stream()
warmup_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warmup_stream):
    for _ in range(3):
        torch.add(static_x, 1, out=static_y)
torch.cuda.current_stream().wait_stream(warmup_stream)

# 恢复一个便于观察的初值。
static_y.fill_(-1)
torch.cuda.synchronize()

g = torch.cuda.CUDAGraph()
host_calls = 0

with torch.cuda.graph(g):
    host_calls += 1
    torch.add(static_x, 1, out=static_y)

# 以下读取均位于 capture 之外。
print(host_calls, static_y.tolist())
# 预期：1 [-1.0, -1.0, -1.0, -1.0]

static_x.fill_(3)
g.replay()
print(host_calls, static_y.tolist())
# 预期：1 [4.0, 4.0, 4.0, 4.0]

static_x.fill_(5)
g.replay()
print(host_calls, static_y.tolist())
# 预期：1 [6.0, 6.0, 6.0, 6.0]

capture 之后，host_calls 已经增加，但 static_y 仍然是事先写入的 -1。replay 之后，GPU 输出发生变化，Python 计数器却不会继续增加。

这里使用原地写入，是为了保留一个确定的输出初值。如果在 capture 内使用 static_y = static_x + 1，PyTorch 同样会创建输出 tensor，但在首次 replay 前，不应把它的内容当作这次加法的结果。

一些封装会在 capture 前运行 warmup，在 capture 后立即 replay，因此观察外层函数时，可能会看到输出已经完成计算。那是封装额外执行的结果，不能据此认为 stream capture 本身执行了被记录的 kernel。

为什么同步和读取结果会失败

capture 期间，GPU 工作尚未被提交执行，自然不能要求 CPU 等待其中某个结果完成。常见的不兼容操作包括：

  • 对正在 capture 的 stream 调用 cudaStreamSynchronize 或查询执行状态。
  • 对 captured event 调用 cudaEventSynchronize 或查询完成状态。
  • 调用覆盖当前 capture 的 device 级同步，例如 torch.cuda.synchronize()。
  • 通过 CUDA tensor 的 .item() 或需要同步的拷贝，将数据读回 CPU 并参与后续计算。

这些限制针对的是 capture 中的执行状态及其依赖，并不是所有与 stream 或 event 有关的 API 都被禁止。例如，符合 capture 规则的 cudaEventRecord 和 cudaStreamWaitEvent 可以用来建立图内依赖。CUDA 文档对不允许的操作和跨 stream 依赖作了分别说明。

这也解释了为什么“异步 API”不一定就能 capture：如果 host 后续逻辑必须读取拷贝结果才能确定 launch，仍然没有消除数据依赖。是否兼容，需要检查完整调用路径。

CUDA Graph 与 stream

stream 用来组织异步提交的 GPU 工作。在通常的执行语义下，同一 stream 中的操作按提交顺序执行；不同 stream 之间需要通过 event 等机制建立依赖。CUDA Graph 则将这些操作及其依赖保存下来，供后续重复执行。

stream 在 capture 和 replay 中承担不同的作用：

阶段 stream 的作用
capture 将 stream 内的先后顺序、stream 之间的等待关系转化为图中的依赖
replay 确定整张 graph 与图外其他工作的执行顺序

多个 stream 如何进入同一张图

如果在一个 stream 中依次 capture 四个 kernel，得到的依赖就是一条链。要表达可以并行的分支，可以在 capture 中使用多个 stream。

下面用四个操作构造一个例子。所有 buffer 都提前分配，并在 warmup、capture 和 replay 期间保留引用：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import torch


static_x = torch.zeros(8, device="cuda")
z = torch.empty_like(static_x)
a = torch.empty_like(static_x)
b = torch.empty_like(static_x)
static_y = torch.empty_like(static_x)
s0 = torch.cuda.Stream()
s1 = torch.cuda.Stream()


def workload():  # 在 s0 为当前 stream 时调用。
    torch.mul(static_x, 2, out=z)       # K0

    s1.wait_stream(s0)
    torch.add(z, 1, out=a)             # K1，在 s0 上
    with torch.cuda.stream(s1):
        torch.mul(z, 3, out=b)         # K2，在 s1 上

    s0.wait_stream(s1)
    torch.add(a, b, out=static_y)       # K3，在 s0 上


s0.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s0):
    for _ in range(3):
        workload()
torch.cuda.current_stream().wait_stream(s0)

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g, stream=s0):
    workload()

第一次 wait_stream 使 K2 依赖 K0，第二次使 K3 依赖 K2。结合 s0 自身的顺序，最终得到：

1
2
3
              ┌── K1：a = z + 1 ──┐
K0：z = x * 2 ┤                    ├── K3：y = a + b
              └── K2：b = z * 3 ──┘

K1 和 K2 都读取 z，但写入不同的输出，因此可以并行；K3 必须等待两者完成。是否实际重叠执行，还取决于 GPU 资源和调度。

这里的 wait_stream 通过 event 表达依赖，capture 时会将依赖加入图中。它不要求 CPU 等待 kernel 执行完成。参与同一次 capture 的其他 stream，需要在结束前通过依赖汇合回起始 stream；仅切换当前 stream，并不会自动建立这种关系。参见 PyTorch 多 stream capture 示例。

PyTorch 如何处理 allocator

当我们计算一个巨大的神经网络时，除了权重是固定的、输入是给定的，很多中间隐藏层的输出都是临时创建、计算、销毁的。以不记录梯度的 forward 为例：

1
2
3
4
5
6
@torch.no_grad()
def forward(x, w1, w2):
    h = x @ w1
    h = torch.relu(h)
    y = h @ w2
    return y

在这个函数返回时，局部变量 h 不再被 Python 持有。但 CUDA Graph 中的后续 replay 仍然会向中间结果对应的地址写入数据，再从这些地址读取。显然，不能直接将这些地址交给无关 tensor 使用。

另一方面，如果每个临时 tensor 都必须永久保留独占显存，深层网络的内存开销又会非常大。解决这个问题，需要区分 tensor 对象的生命周期、allocator 中 block 的复用，以及底层显存的释放。

缓存分配与 stream 顺序

先看普通 eager 执行。PyTorch 的 native caching allocator 会缓存已经申请的显存块，tensor 释放时，通常只是将对应 block 归还缓存，并不立即调用 cudaFree。

对于同一 stream 上的使用，allocator 可以利用顺序关系复用内存，不必每次都等 GPU 完成。

假设某段代码的内存使用关系如下：

1
2
3
4
kernel A：写入临时 buffer p
kernel B：读取 p，写入另一个 buffer
          p 对应的临时 tensor 释放
kernel C：将 p 复用为自己的输出

即使 CPU 提交 kernel C 时，GPU 尚未执行完 kernel A，也不妨碍这种复用。只要 A、B、C 按上述依赖顺序执行，C 对 p 的覆盖就发生在 B 读取之后。

跨 stream 则不能仅凭 host 调用顺序判断安全性。allocator 通常需要结合 record_stream 和 event 等机制，确认其他 stream 不再使用该 block。相关实现可以参考 CUDACachingAllocator.cpp 开头关于 stream 与内存复用的说明。

graph-private memory pool

capture 期间，PyTorch 的 native caching allocator 会将相应分配导向 graph-private memory pool。临时 tensor 仍然可以创建和释放，满足依赖条件的 block 也仍然可以在池内复用；需要额外保证的是，这些地址在后续 replay 时仍由该 graph 使用。

可以将前面的例子延伸为：

1
2
3
4
5
6
7
8
capture 期间记录：
    A 写 p → B 读 p → C 写 p

第一次 replay：
    A 写 p → B 读 p → C 写 p

第二次 replay：
    A 写 p → B 读 p → C 写 p

p 可以先后服务于不同临时结果。每次 replay 都保留相同的依赖关系，因此也保留了复用所需的先后顺序。

replay 时不会重新进入 Python forward，也不会重新执行其中的 tensor 分配逻辑。kernel 直接使用 capture 时确定的地址。这意味着 allocator 不需要在每次 replay 中重新构造上述分配与释放过程。

PyTorch 源码中的 Note [Interaction with CUDA graph capture] 专门说明了这一点：private pool 内仍然允许复用，但 graph 存活期间会保留该池所需的内存高水位。参见 allocator 源码。

由此也可以解释，为什么 capture 后显存占用可能高于普通 eager：普通路径中的空闲 block 可以被其他工作复用，而 graph 的 private pool 需要为未来 replay 保留地址。这个占用并不等于“所有层的临时 tensor 大小相加”，还取决于生命周期重叠、workspace、池内复用和碎片。

输出 tensor 也具有同样的地址约束。如果需要保留某次输出用于之后的计算，应在它被后续 replay 覆盖之前复制：

1
2
3
g.replay()
saved_output = static_output.clone()
g.replay()  # static_output 的内容会被覆盖

这里假定 clone 与 replay 具有正确的 stream 顺序。仅写 saved_output = static_output 不会保存结果的副本。

private pool 的生命周期还与 graph 和其中仍然存活的 tensor 有关。销毁 graph 并不意味着所有输出引用同时失效，因此也不保证显存立刻全部归还给 CUDA。具体生命周期见 Graph memory management。

capture 期间的底层显存分配

private pool 并不意味着所有内存都必须在 capture 前精确分配完毕。需要区分 PyTorch 的 tensor 分配、底层 cudaMalloc，以及 CUDA Graph 自身的内存节点。

对于 native allocator，已有可复用 block 时，tensor 分配主要是 host 端的缓存管理。如果确实需要新的底层分配，当前源码中的 cudaMallocMaybeCapturing 会在 capture 状态下临时使用 cudaStreamCaptureModeRelaxed，再进行分配。

这种分配在 capture 阶段取得真实地址，随后由 private pool 保留，并不会在每次 replay 时重新申请一份显存。对应实现同样位于 CUDACachingAllocator.cpp。

这与 cudaMallocAsync 被捕获为 graph 的 allocation node 是不同的机制，后者由 CUDA 提供图内分配与释放语义，参见 Graph Memory Nodes。不能将 native caching allocator 的工作方式直接套用到所有 allocator backend。

这里的 relaxed mode 也不表示 kernel 会在 capture 时执行，更不会让依赖 GPU 结果的 host 同步变得有效。它只是允许框架在自行保证地址生命周期的前提下完成这类分配。

无法查询 event 时如何回收

剩下的问题是，allocator 在 eager 下可能需要查询 event，capture 时却不能照常依赖这些执行状态。

PyTorch 对此有专门的 capture 路径。native allocator 在 capture 期间跳过常规的 process_events，对需要跨 stream 跟踪的部分回收操作作延后处理。同一 stream 中可以根据依赖顺序证明安全的复用，仍然能够进行。

较新的实现还提供 graph_capture_record_stream_reuse 选项，尝试利用图的拓扑关系判断部分跨 stream 内存何时可以安全复用。这里使用的是图中已经建立的依赖，而不是查询尚未执行的 kernel 是否完成。对应逻辑可以从 allocator 源码中的 process_events、deferred_blocks 和该配置项继续阅读。

因此，capture 中的内存管理并不是暂停所有分配和释放，而是将“可以根据依赖关系确定的复用”与“需要运行时完成状态才能确定的回收”分别处理。

多张 graph 能否共享内存池

如果每个 T 对应的 graph 都使用完全独立的 private pool，推理服务会为许多不会同时执行的 graph 重复保留内存。PyTorch 允许通过 pool 参数让不同 capture 共享池，例如：

1
2
3
4
5
6
7
8
9
# 假设输入已经准备好，两个 workload 都已完成 warmup。
g1 = torch.cuda.CUDAGraph()
g2 = torch.cuda.CUDAGraph()

with torch.cuda.graph(g1):
    out1 = workload1(static_input1)

with torch.cuda.graph(g2, pool=g1.pool()):
    out2 = workload2(static_input2)

共享池只是允许复用，不会自动判断调用方的数据依赖是否安全。对于前后相接的图段，需要维持与 capture 兼容的执行顺序和 tensor 生命周期，不能让后一个图段覆盖前一个图段尚未消费的数据。

对于相互独立、只会择一执行的不同 T 对应的 graph，可以在不并发执行且旧输出不再需要时复用池内空间。PyTorch 文档也列举了推理服务中多个 graph 共享池的用法，参见 Sharing memory across captures。

从这个角度看，capture 为临时内存确定了一套可以重复使用的地址安排。Python 局部变量可以结束生命周期，但图所依赖的地址、执行顺序和内存所有权必须继续成立。