现代 CUDA Graph：从 Full Capture 到 Breakable CUDA Graph

CUDA Graph 最基础的理解是：

把一串 GPU kernel launch 录下来，以后不再让 CPU 一次次 launch，而是一次 replay。

这个理解没错，但对 2026 年的 LLM serving 已经不够了。

真正困难的部分不是 torch.cuda.CUDAGraph() 怎么写，而是：

• decode / prefill / speculative decoding 的 shape 都不一样；
• attention、MoE、collective、metadata preparation 中存在动态逻辑；
• 一个动态 op 不应该让整层模型全部退回 eager；
• graph capture 会长期占用显存；
• captured tensor 的地址必须稳定；
• graph 数量、bucket、padding、chunked prefill 又会反过来影响 scheduler。

因此，现代 CUDA Graph 更应该理解成一种 serving runtime 的静态化与分段执行机制，而不仅是一个 CUDA API 技巧。

本文主要基于 SGLang Team 2026 年的文章 Advanced CUDA Graph Techniques in SGLang 与 meta-pytorch/breakable-cuda-graphs。

---

1. CUDA Graph 到底解决什么问题？

假设一次 decode forward 包含很多 kernel：

RMSNorm
  ↓
QKV projection
  ↓
RoPE
  ↓
Attention
  ↓
Output projection
  ↓
RMSNorm
  ↓
Gate / Up projection
  ↓
SwiGLU
  ↓
Down projection
  ↓
...

GPU 真正做计算之前，CPU runtime 要不断执行：

CPU launch K1
CPU launch K2
CPU launch K3
...

对于大 GEMM，launch overhead 相对很小；但 decode 常常是很多很短的小 kernel，CPU dispatch gap 就会逐渐暴露出来。

CUDA Graph 把它变成：

capture once:
K1 → K2 → K3 → ... → Kn

replay:
CPU ─────── launch graph ───────► GPU executes whole recorded DAG

所以 CUDA Graph 的本质收益是：

$$T_{eager} = \sum_i (T_{launch,i} + T_{kernel,i})$$

变成近似：

$$T_{graph} = T_{graph\ replay} + \sum_i T_{kernel,i}$$

它主要消灭的是 host-side launch / dispatch overhead，并不会减少模型 FLOPs。

因此第一条判断规则非常重要：

CUDA Graph 对 launch-bound workload 最有效；如果已经是 compute-bound，大概率不会产生神奇加速。

---

2. 为什么“整个 forward 全 capture”在 serving 中很难？

传统 CUDA Graph 希望 replay 时 GPU 工作序列固定，而且 capture 期间使用的 device address 在 replay 时继续有效。

对于简单 workload：

fixed shape
fixed buffers
fixed control flow

非常自然。

但 serving 的 forward 实际上更像：

request batch
    ↓
runtime seq_len / batch_size / token_num
    ↓
attention metadata planning
    ↓
possibly collective / sparse routing / cache update
    ↓
model forward

问题在于其中一些步骤依赖实时状态：

• sequence length；
• request 数；
• KV Cache slot；
• attention metadata；
• MoE routing；
• DP / EP collective；
• prefix / hierarchical cache；
• LoRA configuration；
• speculative decoding state。

如果坚持：

“必须整个 forward 都能 capture，否则不用 CUDA Graph”

那么只要中间出现一个不兼容的 op，前后大量静态计算也会一起失去 CUDA Graph 收益。

这就是现代 CUDA Graph 设计的核心矛盾：

希望 capture 尽可能大
        ↕
serving runtime 天生动态

---

3. SGLang 的现代抽象：Runner 和 Backend 分离

SGLang 的一个很重要的重构，是把 CUDA Graph 分成两个问题。

Runner：决定“用什么输入状态运行”

Runner 管：

• capture shape；
• static input buffers；
• attention metadata；
• live batch 如何 pad 到 captured shape；
• decode / prefill / speculative decoding 的执行状态。

可以理解为：

Runner = shape contract + static buffers + runtime state adaptation

Backend：决定“怎么 capture”

Backend 管：

Full Graph

Breakable Graph

Torch-Compile Piecewise Graph

也就是说：

                 ┌──────────── Full CUDA Graph
Execution Runner ├──────────── Breakable CUDA Graph
                 └──────────── TC Piecewise CUDA Graph

这层抽象非常值得学。

现代 serving runtime 不应该让：

decode graph implementation
prefill graph implementation
speculative graph implementation

各自复制一套 shape、buffer、capture、replay 逻辑。

更好的方式是：

execution path 决定状态，capture backend 决定静态化策略。

---

4. 三种 CUDA Graph 策略

SGLang 当前可以从三个层次理解。

4.1 Full CUDA Graph

整个 forward 一次性 capture：

┌─────────────────────────────────────┐
│             CUDA Graph              │
│ Linear → Attn → Linear → MLP → ...  │
└─────────────────────────────────────┘

优点：

• replay launch 最少；
• host overhead 最低；
• 性能上限最高。

缺点：

• 所有路径都必须 graph-safe；
• shape / metadata 必须足够静态；
• prefill 很难做。

Decode 最适合，因为每个 request 一次通常只贡献一个 token，主要变化维度就是 batch size。

---

4.2 TC Piecewise CUDA Graph

一种早期方案是：

torch.compile
    ↓
trace full forward
    ↓
FX graph
    ↓
split at registered points
    ↓
compile + capture each piece

结果也是：

Graph A → eager / split → Graph B → Graph C

但核心问题是：compiler 必须先理解你的 forward。

对于大量 custom CUDA / Triton / JIT kernel，通常意味着还需要：

• custom op registration；
• fake implementation；
• tracing compatible inputs / outputs；
• compiler-friendly graph boundaries。

于是本来是“做 CUDA Graph”，最后容易变成“做 torch.compile integration”。

---

4.3 Breakable CUDA Graph（BCG）

BCG 的思路更直接：

不用 compiler 先理解整张图，而是在 capture 过程中显式告诉 runtime：这里断一下。

例如：

Graph Segment A
       ↓
   Eager Region
       ↓
Graph Segment B
       ↓
   Eager Region
       ↓
Graph Segment C

SGLang 中可以把某个 graph-incompatible region 标记成 eager break。

因此原先：

A → dynamic op → B

因为 dynamic op 不能 capture
=> A + dynamic op + B 全 eager

现在变成：

capture(A)
   ↓
eager(dynamic op)
   ↓
capture(B)

这就是 BCG 最重要的思想：

动态性不再污染整个 forward，而只污染真正动态的局部区域。

---

5. BCG 最关键的问题：Graph B 怎么读到 eager region 的新输出？

这是理解 Breakable CUDA Graph 最关键的一点。

CUDA Graph replay 要求地址稳定。

假设：

Graph A
  ↓
eager F(x)
  ↓
Graph B

capture Graph B 时，它看到的是一个 tensor 地址：

0xABC...

以后 replay Graph B 仍然会从这个地址读。

但 eager function 每次运行可能返回新的 Tensor：

replay #1: output at 0x111
replay #2: output at 0x222

那 Graph B 就坏了。

BCG 的解决思路是引入 persistent boundary buffer：

                 fixed device address
                       │
                       ▼
Graph A → eager F → [ bridge buffer ] → Graph B
                        0xABC...

Capture 时：

1. eager region 运行一次；
2. 为跨 break 的数据建立稳定 buffer；
3. Graph B 对着这个固定地址 capture。

Replay 时：

eager F produces fresh result
              ↓
copy / write back into bridge buffer
              ↓
Graph B replay reads same captured address

所以你可以把 BCG 理解为：

用固定地址的 bridge buffer，把动态 eager island 接回静态 CUDA Graph 世界。

这比单纯记 @eager_on_graph 有价值得多。

---

6. 一个最小的 Breakable CUDA Graph 心智模型

独立的 meta-pytorch/breakable-cuda-graphs 已经把这种机制抽成了一个很小的 API。

概念代码可以写成：

import torch
from breakable_cuda_graphs import (
    CUDAGraphSequence,
    breakable_graph,
    no_graph,
)

@no_graph
def host_dependent_step(x, tmp):
    # 例如这里存在 .item() / runtime metadata / host logic
    if x.sum().item() > 0:
        tmp.copy_(x.relu())
    else:
        tmp.zero_()

static_x = torch.empty(4096, device="cuda")
static_tmp = torch.empty_like(static_x)
out = torch.empty_like(static_x)

def forward(x, tmp, y):
    tmp.copy_(x * 2)            # graph segment A
    host_dependent_step(tmp, tmp)  # eager break
    y.copy_(tmp + 1)            # graph segment B

# 先 warmup，再 capture
seq = CUDAGraphSequence()
with breakable_graph(seq):
    forward(static_x, static_tmp, out)

# replay 前只更新静态输入
static_x.copy_(new_input)
seq.replay()

真正需要记住的不是 API，而是三个 contract：

1. 静态地址 contract：captured graph 使用的 device buffer 地址不能随便变；
2. break boundary contract：eager region 和 graph segment 之间要通过稳定 buffer 传数据；
3. stream contract：进入 break 前，额外 side stream 必须正确 join 回 capture stream。

---

7. 为什么 BCG 比 compile piecewise 更“serving-native”？

二者最后都可能得到：

Graph A → Eager → Graph B → Eager → Graph C

区别在于 谁决定边界。

TC Piecewise

Python forward
    ↓
compiler trace
    ↓
FX graph
    ↓
compiler 能表示的 split point

BCG

Python forward
    ↓
developer / serving runtime 明确知道这里动态
    ↓
direct eager break during capture

因此 BCG 的 graph boundary 可以跟着真实 serving 语义走：

• dynamic attention planning；
• sparse attention indexer；
• runtime collective；
• cache bookkeeping；
• 特殊 Python-side logic。

而不需要为了让 compiler 能 trace，反过来修改 serving code 的结构。

这就是为什么它的工程价值很大：

BCG 把“哪些地方动态”从 compiler 推断问题，变成 runtime 显式 contract。

---

8. 为什么现代 CUDA Graph 一定要做 Shape Bucketing？

Graph capture 本质上还是静态 shape 世界。

例如 decode 真实 batch size：

1, 3, 5, 7, 12, 15, 21, ...

你不可能给每个值都 capture 一个 graph。

常见做法是建立 bucket：

captured batch sizes:
1, 2, 4, 8, 16, 32, 64

真实 batch=13：

13
 ↓
pad / map
 ↓
16 bucket
 ↓
replay graph(batch=16)

所以现代 CUDA Graph 的性能不只取决于 capture 本身，而取决于：

真实 workload distribution
        ↓
bucket policy
        ↓
padding waste
        ↓
graph hit rate
        ↓
最终收益

这也是为什么 CUDA Graph 已经变成 scheduler / runtime 设计的一部分。

---

9. Prefill 为什么比 Decode 难很多？

Decode 可以近似看成：

shape variable ≈ batch size

因为每个活跃 request 通常只 decode 一个 token。

Prefill 不一样。

它至少同时变化两个维度：

1. total token count
2. number of requests

例如：

Batch A:
4 requests × different prompt lengths
=> 2048 total tokens

Batch B:
17 requests × different prompt lengths
=> also 2048 total tokens

总 token 相同，但 request metadata 完全不同。

SGLang Full Prefill CUDA Graph 的思路是把这两个维度分别静态化。

---

10. Prefill 静态化：Token Bucket + Request Slots

10.1 Token 维度

总 token 数 pad 到最近的 captured bucket：

live tokens = 3650
       ↓
pad
       ↓
captured token bucket = 4096

这个 padding 是有真实代价的。

因为 padded token 会真的进入一部分 dense projection：

[X_real]
[X_pad ]
   ↓
Dense GEMM

所以 token padding 是昂贵的 padding。

---

10.2 Request 维度

Graph 预留固定数量 request slots：

Captured slots = 16
Live requests   = 11

布局可以看成：

slot 0  real request
slot 1  real request
...
slot 10 real request
slot 11 zero-length sentinel
slot 12 zero-length sentinel
...
slot 15 zero-length sentinel

空 slot 设置成：

seq_len    = 0
extend_len = 0

FlashAttention variable-length scheduler 会根据真实序列长度产生工作量，所以 zero-length request 基本不会产生真正 attention 计算。

因此：

request-slot padding 很便宜，token padding 很贵。

这是设计 prefill bucket 时很关键的非对称性。

---

11. Full Prefill Graph 的工程本质

你可以把 full prefill capture 看成强行构造一个静态 execution contract：

(token_bucket, fixed_request_slots)
                  ↓
      fixed-shape forward contract
                  ↓
           full graph replay

但 runtime 仍然需要在 replay 前更新：

• live request slots；
• zero-length sentinels；
• attention metadata；
• true token count。

因此 Full Graph 不代表：

“整个 serving 系统什么都不用动态做了。”

而是：

把动态 metadata preparation 推到 replay 之前，使真正 GPU forward 足够静态。

这是一个更准确的理解。

---

12. Full / Breakable / Piecewise 应该怎么选？

可以按这个顺序判断。

场景	推荐思路
Decode，shape 容易 bucket	Full CUDA Graph 优先
Prefill，大部分静态但存在少数动态 op	Breakable CUDA Graph
Prefill，backend 支持完整静态 metadata contract	尝试 Full CUDA Graph
已经依赖 torch.compile，而且 tracing 稳定	TC piecewise 仍可用
Custom CUDA / Triton / Python runtime logic 很多	BCG 通常更自然
调试 graph correctness	先全 eager / debug graph path

核心不是“Full 一定最好”，而是：

$$\text{收益} \approx \text{减少的 launch overhead} - \text{padding / eager break / bookkeeping 成本}$$

---

13. BCG 的性能为什么可能接近 Full Graph？

假设原始 forward 有 200 个 kernel，但只有 attention metadata planning 需要 eager。

Eager：

200 次左右 host launch / dispatch

BCG：

Graph A replay
    ↓
small eager island
    ↓
Graph B replay

虽然不是一次 graph launch，但仍然可能把 200 次 launch 压缩成：

2 次 graph replay + 少量 eager launches

所以只要 eager island 足够小：

BCG ≈ Full Graph performance

而工程兼容性大幅增加。

SGLang 文章在特定 prefill benchmark 中报告：

Eager        1.00×
TC piecewise 1.45×
BCG          1.70×
Full Graph   1.93×

这些不是通用常数，但很好地体现了层级：

capture coverage 越高
host dispatch 越少
理论上越接近 full graph 上限

---

14. 为什么 BCG 的启动更快？

TC piecewise 大致需要：

trace
 ↓
compile
 ↓
split
 ↓
compile pieces
 ↓
capture

BCG：

warmup
 ↓
one capture pass
 ↓
遇到 explicit break 就分段

因此它不需要先完成 compiler compilation。

SGLang 在其测试中观察到 BCG 构建 prefill graphs 比 compile-based piecewise 快约 3.8–5.2×。

对于 serving 系统，这个指标很重要，因为 graph 不只是 steady-state throughput 问题，还影响：

• cold start；
• CI；
• model startup；
• capture table 初始化；
• 大模型部署弹性。

---

15. CUDA Graph 的隐藏成本：显存常驻

CUDA Graph capture 的一个常见误区是只看 latency，不看 memory。

Graph capture 时使用的中间存储往往需要长期保持可 replay：

Eager activation:
request arrives → allocate → compute → release

Graph memory:
capture → resident → resident → resident

所以如果：

42 shapes
× 多个 graph segments
× 每个 segment 独立 allocator pool

显存很容易爆炸。

现代实现必须认真做 memory reuse。

---

16. SGLang 的三层 Graph Memory Reuse

16.1 Segments 共享 Graph Pool

不要：

Segment A → pool A
Segment B → pool B
Segment C → pool C

而是：

Segment A ┐
Segment B ├── shared CUDA Graph pool
Segment C ┘

这样不同 segment 的中间存储可以复用。

16.2 Eager Break 使用 weak reference 控制生命周期

如果 graph pool 已经拥有某块 storage，就不要因为 Python tensor reference 又把生命周期无意义拉长。

因此对可复用 graph-owned storage，可以使用 weak-reference 思路减少额外 pinning。

16.3 不同 Capture Shape 共享最大 Output Buffer

不要：

shape 8  → output buffer 8
shape 16 → output buffer 16
shape 32 → output buffer 32

而是：

max output buffer 32
├── [:8]
├── [:16]
└── [:32]

这跟你写 kernel workspace 时的 max-buffer reuse 是同一个思想。

---

17. 哪块显存不能随便复用？

就是前面说的 跨 eager break boundary buffer。

因为下一个 graph segment capture 的是它的地址：

Graph B captured input pointer = P

所以它必须：

persistent
fixed address
updated in place

这块 buffer 的生命周期本质上是 graph topology 的一部分。

---

18. 为什么 capture ceiling 要覆盖 chunked-prefill size？

这是文章里一个很容易被忽略但非常好的系统点。

假设：

chunked_prefill_size = 8192

但你只 capture 到：

4096

那么 8192-token prefill 还是 eager：

小 shape:
resident graph memory

最大 shape:
仍然出现 eager transient activation peak

于是你同时付了：

Graph resident memory
+
Worst eager activation peak

这是最差情况。

如果 capture ceiling 一直到 8192：

largest legal prefill
      ↓
也走 graph
      ↓
worst transient activation peak 消失 / 显著下降

因此决定显存的关键不只是：

“capture 了多少个 shape？”

更重要的是：

“最大的 capture shape 到哪里？”

这和 allocator 里的 high-water mark 思维很类似。

---

19. 现代 CUDA Graph 的 Debug 方法

Full CUDA Graph 的一个缺点是 replay 时 Python 不会重新执行 graph 内部逻辑。

所以：

print(...)
assert ...
breakpoint()

往往看不到预期行为。

BCG 天然保留 eager island，因此更容易定位问题。

更进一步，可以让整个 forward 都作为 eager break：

Graph Runner / static buffer / metadata path
                保留
                   ↓
Model forward completely eager

这样可以做二分：

问题 eager 也存在
=> model / runner / metadata bug

问题只有 capture 后存在
=> CUDA Graph capture / lifetime / stream / address bug

这是非常实用的 CUDA Graph debug 心法。

---

20. CUDA Graph 的真正“现代”之处是什么？

我认为可以总结成五点。

20.1 从“全图或不用”变成“最大化 capture coverage”

Old:
Can whole forward be graphed?

Modern:
Which regions are stable enough to graph?

20.2 Graph boundary 变成 serving contract

Dynamic op 不再被视为失败，而是明确的：

static region
→ dynamic island
→ static region

20.3 Shape bucket 与 scheduler 联合设计

CUDA Graph 不是 executor 的局部开关，而会影响：

• batching；
• chunk size；
• token bucket；
• batch bucket；
• graph hit rate；
• padding waste。

20.4 Memory pool 是 Graph 设计的一部分

Graph 性能优化必须同时考虑：

latency
startup
resident memory
activation peak

20.5 Debug path 必须与 optimized path 共用 runner

否则你 debug 的不是实际 serving execution contract。

---

21. 和你熟悉的 Prefill / Decode 联系起来

你可以把 CUDA Graph 放进这张图里：

                        LLM Serving
                            │
              ┌─────────────┴─────────────┐
              │                           │
           Prefill                     Decode
              │                           │
       compute-heavy               launch / memory-heavy
              │                           │
       dynamic token count           one token / request
       dynamic req count             batch-size buckets
              │                           │
    BCG / Full Prefill CG          Full CUDA Graph
              │                           │
 token bucket + req slots          high replay hit rate

这里不能简单说：

“Prefill compute-bound，所以 CUDA Graph 没意义。”

大型 GEMM 本身可能 compute-bound，但整个 prefill forward 中仍然可能夹着很多短 kernel、metadata、routing、normalization、activation 与 custom ops。

SGLang 的 prefill benchmark 在 prompt length 跨 32× 范围时 CUDA Graph 加速曲线仍比较平，正说明其收益里存在明显 launch-overhead 成分。

---

22. 面试时怎么回答“讲一下现代 CUDA Graph”

可以用下面这段结构。

CUDA Graph 的基本作用是把一串 kernel launch capture 成可 replay 的 GPU DAG，减少 CPU launch overhead。传统 full graph 要求 shape、地址和控制流都比较静态，但现代 LLM serving 很多路径是动态的，例如 prefill attention metadata、MoE 或 runtime cache state，所以现在更重要的是 partial graph capture。

SGLang 的一个做法是 Breakable CUDA Graph：把 graph-safe 的区域 capture，遇到显式标记的动态 op 就结束当前 graph segment，运行 eager op，然后继续 capture 下一个 segment。跨 break 的 tensor 需要有固定地址的 persistent boundary buffer，replay 时 eager region 把新结果写回这个 buffer，后面的 graph 才能继续使用 capture 时的地址。

另外实际 serving 还要做 shape bucketing。Decode 主要 bucket batch size；Prefill 同时有 token count 和 request count 两个维度，可以用 token bucket 加固定 request slots，把空 slot 写成 zero-length sentinel。Token padding 会产生真实 dense compute，所以贵；request slot padding 基本只有 metadata overhead，所以便宜。

工程上还要关注 graph memory pool reuse、跨 shape output buffer reuse、capture ceiling 是否覆盖 chunked-prefill size，以及 replay hit rate。CUDA Graph 的收益主要来自消除 launch gap，所以 launch-bound workload 收益最大，compute-bound kernel 本身不会因为 graph 而变快。

如果这段能顺下来，已经不是“知道 torch.cuda.CUDAGraph API”，而是真正理解 CUDA Graph 在现代 serving runtime 里的位置。

---

23. 一张最终心智图

                         Runtime Request
                               │
                               ▼
                      ┌─────────────────┐
                      │     Runner      │
                      │ shape / metadata│
                      │ static buffers  │
                      └────────┬────────┘
                               │
               choose captured shape / bucket
                               │
                               ▼
                 ┌─────────────────────────┐
                 │     Capture Backend     │
                 └────────────┬────────────┘
                              │
             ┌────────────────┼────────────────┐
             │                │                │
             ▼                ▼                ▼
         Full Graph          BCG         TC Piecewise
             │                │                │
             │        Graph A │                │
             │                ▼                │
             │            Eager op             │
             │                │                │
             │      persistent bridge buf      │
             │                │                │
             │                ▼                │
             │            Graph B              │
             └────────────────┼────────────────┘
                              │
                              ▼
                           Replay
                              │
                              ▼
                 fewer CPU launch gaps

现代 CUDA Graph 的核心不是：

“怎么把所有代码塞进一张 Graph？”

而是：

“怎么建立稳定的 shape、地址、memory 与 dynamic-boundary contract，让尽可能大的 GPU 工作集可以低成本 replay？”

这才是 CUDA Graph 从一个 CUDA API 技巧走向 serving runtime 基础设施的关键一步。

---

References

• SGLang Team, Advanced CUDA Graph Techniques in SGLang, 2026-08-17.
• Meta PyTorch, breakable-cuda-graphs.
• SGLang PR #19102: Introduce CUDA graph debug mode with breakable CUDA graph.
• SGLang PR #22218: Breakable Piecewise CUDA Graph.
• SGLang PR #27988: Full CUDA Graph Support for Prefill.