CUDA Graph：捕获 / 回放 / 失败降级

谁该读这一篇？ 想搞清"CUDA Graph 究竟省了什么 / 什么时候不应该开"的工程师。 前置阅读： 03-code-walkthrough/04-model-runner.md §9。 耗时： 30 分钟 学完能： 1. 解释 CUDA Graph 的核心收益（消除 kernel launch 开销）； 2. 说出 SGLang 在哪一步捕获 / 回放； 3. 知道什么变量会导致 graph fallback； 4. 调整 --cuda-graph-bs、--cuda-graph-max-bs 配合业务； 5. 排查 "Capture cuda graph 卡 60s+" 类启动问题。

1. 一个 kernel launch 也是不便宜的

PyTorch 跑一次 forward 要发起几百到几千次 CUDA kernel launch，每次大约 5-10 μs。一次 decode 的"算"时间可能只有 1.5 ms，但这些 launch 累计起来就有 0.5-1 ms，不算小数。

CUDA Graph：把"这一串 launch 的拓扑"录下来，下次直接提交整张图，省掉 host ↔ device 的多次同步。省下来的就是上面那 0.5-1 ms。

2. SGLang 怎么用 CUDA Graph

源码：cuda_graph_runner.py（1435 行）。

2.1 启动时捕获

ModelRunner.init_cuda_graphs → CudaGraphRunner.capture（cuda_graph_runner.py:813）

def capture(self):
    bs_to_capture = get_batch_sizes_to_capture(self.model_runner)
    # 通常: [1, 2, 4, 8, 16, 32, 64, 128, 256] 或自定义
    for bs in bs_to_capture:
        self.capture_one_batch_size(bs)

capture_one_batch_size（917）：

def capture_one_batch_size(self, bs):
    dummy_batch = self.make_dummy_decode_batch(bs)
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph, pool=self.global_pool):
        out = self.model_runner.model.forward(dummy_batch.input_ids, ..., dummy_batch)
    self.graph_pool[bs] = (graph, out)

pool=self.global_pool 让多个 graph 共享一块 memory pool，避免显存放大。

启动日志里 "Capture cuda graph begin/end" 就是这一步，耗时 5-30s。

2.2 运行时回放

forward_decode → cuda_graph_runner.replay（cuda_graph_runner.py:1260）

def replay(self, forward_batch):
    bs = forward_batch.batch_size
    graph, out_buf = self.graph_pool[bs]
    self.replay_prepare(forward_batch)         # 把输入塞进 input buffer
    graph.replay()
    return out_buf                              # output 在固定 buffer 里

replay_prepare 做两件事：

copy forward_batch.input_ids 等到捕获时分配的 input buffer。
copy 当前 batch 的 attention metadata（page table 等）到捕获时分配的元数据 buffer。

关键不变量：所有"形状可变"的张量都必须用捕获时分配的 buffer，运行时只 copy 内容、不重新分配。这是 CUDA Graph 的硬约束（捕获时记录的是地址）。

3. 什么 batch size 被捕获

get_batch_sizes_to_capture（cuda_graph_runner.py:483）：

def get_batch_sizes_to_capture(model_runner, num_tokens_per_bs=1):
    # 默认指数序列 [1, 2, 4, 8, 16, ..., max_running_requests]
    # 可以用 --cuda-graph-bs 强制
    ...

可调参数：

参数	含义
`--cuda-graph-bs`	显式列出要捕获的 batch size，如 `"1,4,16,64,128"`
`--cuda-graph-max-bs`	最大 bs（默认与 max_running_requests 相同）
`--disable-cuda-graph`	完全关掉
`--disable-cuda-graph-padding`	关掉 batch size padding

Padding 机制

如果 batch size 是 5，没有 5 的 graph，会 padding 到下一个有的（8）：

real batch_size = 5
padded to 8
graph_replay(bs=8)
output[:5] 取前 5 个有效

代价：算了 3 个 dummy 请求。但比 fallback 到 eager 还划算。关 padding 后小 batch 走 eager 路径。

4. Fallback 触发条件

CUDA Graph 不能用时回退到 eager forward。触发条件（cuda_graph_runner.py:replay）：

情况	原因
batch 含 prefill 请求	Graph 不支持变长输入
batch_size > max_bs	没捕获过
使用了 LoRA 但 adapter 没 captured	LoRA 路径未支持 graph
受约束解码切换 grammar	grammar 元数据不在 graph 里
Speculative draft tree 形状改变	tree attention 不固定

可观察 metrics：

sglang:cuda_graph_replay_count_total{...}
sglang:cuda_graph_fallback_count_total{...}

fallback_count / total 持续 > 5% 说明 graph 用得不充分，要排查为什么。

5. Memory Pool 与显存

捕获每个 graph 都需要一份"中间张量"显存。 SGLang 让多个 graph 共享同一个 memory pool：

self.global_pool = torch.cuda.graph_pool_handle()

效果：N 个 graph 不会让显存翻 N 倍，只是共享池增大一些。但仍有开销，启动日志里 "Memory pool end" 之后会看到显存被多吃几 GB。

6. CUDA Graph + Overlap 调度

Overlap 模式下 graph 在 forward_stream 上 replay；准备工作（replay_prepare）在 schedule_stream 上做。两个 stream 同时跑，graph replay 直接拿"上一轮 prepare 好的 buffer"。

实现复杂，但收益是 TPOT 再降 10-15%。详见 cuda_graph_runner.py:replay_prepare（1174）。

7. 与 torch.compile 的关系

# 启动时
def set_torch_compile_config():
    torch._dynamo.config.cache_size_limit = ...
    torch._inductor.config.coordinate_descent_tuning = True

SGLang 支持 torch.compile 编译模型（compilation/）。两者关系：

torch.compile 在 op 层做 fusion，产出更少的 kernel。
CUDA Graph 在更上层把"一串 kernel launch" 录下来。
配合用：先 compile（少 kernel）再 capture（去 launch 开销）。

启动：

--enable-torch-compile           # 开 torch.compile
--torch-compile-max-bs 128       # compile 的 batch size 上限

8. 启动慢排查

"Capture cuda graph"卡 60s+：

减 --cuda-graph-bs 列表（默认 9 个 bs，减到 5 个）。
关 --enable-torch-compile（compile 也会慢）。
模型超大（70B+）单 bs capture 就要 10s，多 bs 累计很慢。
调试期间直接 --disable-cuda-graph 跳过。

9. 生产部署建议

场景	配置
主流 7B-70B 模型	默认即可
模型超大 + 启动时间敏感	自定义 `--cuda-graph-bs "1,32,128"`（少 bs）
LoRA 多 adapter	关 graph（fallback 太多）
Disaggregated decoder	默认开
Debug 模式	`--disable-cuda-graph` 看 eager 行为

10. 小结

CUDA Graph 录一串 kernel launch，下次 replay 省 launch 开销，TPOT 降 ~30%。
SGLang 启动时给一组常见 batch size 各捕获一份；运行时 batch_size + 形状匹配才能 replay。
Padding 机制：batch_size 不匹配时填到下一档。
Prefill / LoRA / spec draft 形状变化触发 fallback。
关掉 graph 启动快、运行慢；生产建议保留默认开。

11. 自检

CUDA Graph 省的是什么开销？

答案

**Kernel launch 开销**——每次 `kernel<<>>(...)` 调用 host 都要做：(a) 参数打包；(b) 写入 GPU command queue；(c) 设置 stream 同步。每次 launch 平均 5-10 μs。一次 LLM forward 有几百到几千个 kernel（每层 attention + MLP + layernorm + ...），累计 launch 开销 0.5-1 ms。 Graph 录制把"这串 launch 的拓扑"录下来，replay 时一次 `graph.replay()` 把整个序列直接提交，省掉 host ↔ device 的多次同步。省的是 host-side overhead，不是 GPU 算力；所以单 token decode（GPU 算力本就少）收益最大。

为什么 prefill 不能用 graph？

答案

CUDA Graph 录制时把张量地址和形状都"烙进"图里，replay 时形状必须完全一致。 Prefill 每请求的 prompt 长度不同 → input_ids shape 每次变 → 录的 graph 跑下次新 prompt 直接形状不匹配崩。不是绝对不能用：如果限定 prefill chunk size 固定（chunked prefill 模式下每 chunk 都是 4096 token），形状一致可以 capture。SGLang 的 `piecewise_cuda_graph_runner.py` 就是干这事，对长 prompt 业务把 prefill 也吃进 graph 加速。

Padding 和 fallback 哪个代价大？

答案

**通常 fallback 代价大**。 - Padding 到下一档（如 real bs=5 → graph bs=8）：算了 3 个 dummy 请求的算力（GPU 浪费 ~37%），但仍享 graph 的 launch 省开销。整体时延仅比 graph bs=5 高一点点，比 eager 快得多。 - Fallback 到 eager：完全失去 graph 的 launch 省 → TPOT 增 30-50%（取决于模型层数）。规律：除非 batch size 落在 graph 列表的"大洞"里（如 bs=200 但 graph 只有 [128, 256]，padding 到 256 浪费 ~22%），padding 都比 fallback 划算。关闭 padding（`--disable-cuda-graph-padding`）时小 batch 走 eager 路径，建议保留 padding 默认开。

Graph 和 torch.compile 重复吗？

答案

**不重复，互补**。 torch.compile 在算子层做 fusion + 代码生成：把 elementwise + matmul + layernorm 一系列小 op 融合成少数几个大 kernel（Inductor 产出）。**产出更少 kernel**。 CUDA Graph 在更上层把"一串 kernel launch" 录下来。**让剩下的 launch 不花 CPU 时间**。两者顺序：先 compile（少 kernel）再 capture（去 launch 开销），叠加收益 TPOT 再降 5-10%。 SGLang 启动 `--enable-torch-compile --cuda-graph-max-bs 128` 同时开。代价：torch.compile 本身首次 capture 慢（30-60s），总启动时间长。

启动 capture 卡 90s，怎么诊断 + 缓解？

答案

诊断（按概率）： (a) **graph batch size 列表太长**：默认 9 个（1,2,4,...,256），70B 模型每个 5-10s，累计 1.5 分钟。 (b) **同时开 torch.compile**：每个 bs 多花 5-20s。 (c) **TP 大** + NCCL 慢：每个 graph 录制都要走 AllReduce，节点间 NCCL 慢拖累。 (d) **模型超大**：DeepSeek-V3 671B 每 bs capture 几十秒正常。缓解： - `--cuda-graph-bs "1,32,128"` 减档（牺牲少数 batch size 的 padding 表现）。 - `--cuda-graph-max-bs 128` 上限切低。 - 关 torch.compile（debug 阶段）。 - 把模型 ramdisk + NCCL 调好。 - 接受 90s：用 K8s `startupProbe initialDelaySeconds=120s`，扩容预热池避免突发流量等待。

12. 下一步

03-code-walkthrough/04-model-runner.md — Graph 的上层用户。
01-flashinfer.md — Attention backend 的 graph 兼容性。
源码：cuda_graph_runner.py、piecewise_cuda_graph_runner.py。

上游源码：sglang/python/sglang/srt/model_executor/cuda_graph_runner.py。