CUDA Graph：把一轮计算录下来

Graph 的契约：拓扑、参数与存储各管什么

这里讨论常规 CUDAGraph replay。CUDA 底层还有显式节点更新能力；Python 变量重新赋值不会替你更新捕获参数。

准备固定输入，在旁路 stream 上热身

import torch
static_x = torch.zeros(4, device="cuda")
def f(x):
    return x.square() + 1
warm = torch.cuda.Stream()
warm.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(warm):
    for _ in range(3):
        f(static_x)
torch.cuda.current_stream().wait_stream(warm)

zeros 先在当前 stream 上准备输入；wait_stream 建立先后关系。热身触发初始化与缓存，结束后主 stream 才进入捕获。示例要求 PyTorch + CUDA GPU。

捕获之后，把新值写进同一块存储

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_y = f(static_x)

static_x.fill_(3)       # 原地址，内容变为 3
g.replay()              # y = 3²+1 = 10
saved = static_y.clone() # 若要保留这一轮，先复制
static_x.fill_(4)
g.replay()              # static_y 被覆盖为 17

这段操作在同一 stream 按顺序进行。saved 独立保存 10，static_y 表示最近一次 replay 的输出。完整可下载示例还会与 eager 结果逐轮比较。

保存 Python 引用，并没有保存旧一轮的数值

第一步更新 A 的内容，B 尚未变化；第二步 replay 覆盖 B，旧引用也看到 17。只有独立复制的 saved 仍是 10。

一个容易漏掉的 bug：变量换了，graph 的指针没换

static_x = torch.zeros(4, device="cuda")
# ...捕获引用 static_x 原存储的 graph...
original_x = static_x          # 保留捕获用的存储
static_x = torch.full((4,), 9., device="cuda")
g.replay()                     # 仍读取 original_x 的 0
# 修复：original_x.copy_(static_x)，再 replay

变量名只是 Python 中的引用。Graph 记录的是原来那块存储的地址；释放原存储还可能引入更严重的生命周期错误。长期保留所有捕获用的输入与输出。

Python 分支在捕获时只走一次

# 不会在每次 replay 时重新运行这个 Python if
with torch.cuda.graph(g):
    if use_branch_a:
        static_y = branch_a(static_x)
    else:
        static_y = branch_b(static_x)

以后把 use_branch_a 改成 False，不会重建 graph。分支由 CPU 决定时，在 replay 前选择对应 graph；需要设备数据参与选择时，用支持捕获的设备算子表达。

CPU 计数器、print、普通 Python 列表更新等副作用同样不会随 GPU replay 重复执行。

捕获失败，沿着“要等 CPU 的哪一步”去找

torch.cuda.graph 会使用适合捕获的 side stream。不要把不支持的操作包进 try/except 后静默跳过；先保证 eager 路径和捕获路径计算相同。

KV 可以增长：稳定的是大块存储，变化的是索引

这个能力取决于 attention kernel 与服务引擎：若实现根据长度在 CPU 改变 launch 形状，仍可能要切换 bucket 或走 eager。

无效 padding 行不能误写真实请求的 KV；使用安全 slot、mask 和对应后端约定，并测试短请求结束后的复用。

形状桶的代价，要按请求分布求平均

教学 batch 序列 [1,3,5,8]；桶 [1,2,4,8] 对应 [1,4,8,8]。实际 token 17，计算槽位 21，padding 占 4/21≈19.0%。

如果新增 bucket=5，计算槽位降到 18，浪费率降到 1/18≈5.6%。代价是多捕获一张图、长期存储、管理和冷启动时间。

这里用每槽位相同成本估算；实际 attention 的长度、页表和 kernel 选择会让成本偏离线性。

Graph 省下的提交，要扣掉 padding 与拷贝

教学成本：eager 提交 6 μs + 有效计算 6 μs =12 μs；graph 提交 1 μs + bucket 计算 8 μs + copy/管理 1 μs =10 μs。净省 2 μs。

净收益只有 2 μs，捕获成本 600 μs 要多久回本？

沿用上一页：C=600 μs，te=12 μs，tg=10 μs。N=300 时正好打平；从 第 301 次 开始累计省时间。

若实际 B=2 却用了 bucket=8，额外计算可能让 tg≥te。这时重复再多也无法回本，应换小桶或保留 eager。

多一张 graph，可能少服务几条长请求

教学预算：剩余 8 GiB 可以给 KV / graph；每张 graph 独占 0.5 GiB，4 个 bucket 用 2 GiB，KV 剩 6 GiB。若每条请求 KV=0.25 GiB，可容纳 24 条。

增加到 8 个 bucket 后，graph 用 4 GiB，KV 只够 16 条。单轮更快不一定带来更高服务吞吐。

实际 pool 大小应测量；捕获的张量和 graph 生命周期会延长内存保留。不同图共享 pool 也有执行顺序与非并发约束。

换 stream 以后，“先写后读”需要显式依赖

# copy 与 replay 若不在同一个 stream，建立依赖
with torch.cuda.stream(copy_stream):
    static_x.copy_(new_x)
    ready.record()
with torch.cuda.stream(run_stream):
    run_stream.wait_event(ready)
    g.replay()
    done.record()

消费输出的 stream 同样要等待 done。下一轮不得提前覆盖 static_x / static_y；可以用 event 管理双缓冲或串行复用。跨 stream 使用的 tensor 还要保证分配器看到其使用寿命。

单 stream 顺序已经提供这些先后关系；增加 stream 必须有可重叠的工作才有收益。

把稳定的计算段捕获，动态调度留在边界

调度器决定活跃请求与 block 分配；计算段读取相同形状的 buffer；结果出来后再更新请求状态。这样每轮请求集合可以变化，而图中的算子布局保持稳定。

数据相关的动态 shape、跨进程操作或无法捕获的后端可以拆出去。拆成更多 graph 会增加 replay 与边界成本，必须重新测端到端收益。

检查 replay 是否真的覆盖了高频小算子：只捕获一个本来就很大的 GEMM，通常省不了多少主机关键路径。

打开 graph 后没有变快，用四组数据定位原因

正确性比较要覆盖多轮不同输入、不同 bucket、padding 和请求结束后的复用。只测第一次 replay，无法发现旧输入与输出覆盖问题。

两次 replay 后，a 与 b 分别保存什么？

先写出假设、单位与中间步骤，再展开核对。

少两个 bucket，节省了内存，增加了多少空算？

先写出假设、单位与中间步骤，再展开核对。

一个低频 bucket 是否值得提前捕获？

先写出假设、单位与中间步骤，再展开核对。

先在 CPU 复现存储别名，再用 GPU 验证真实 replay

python3 lab.py
# alias=17; clone=10; slots=21; break_even=301

CPU 实验模拟稳定 buffer、输出覆盖、形状桶与成本模型。下载同目录 gpu_example.py 后，在安装 PyTorch 与 CUDA 的机器运行 python3 gpu_example.py；逐轮断言 graph 与 eager 一致。

Python 3 标准库即可运行。模拟输出用于检验推导；GPU 性能需要在目标硬件上测量。