SGLang 里的 Breakable CUDA Graph：graph 构建快 5 倍，prefill 快 1.93 倍
SGLang 现在有两项 CUDA Graph 技术，做到了传统 CUDA Graph 在生产环境里做不到的事：构建 prefill graph 快 3.8–5.2 倍，prefill 本身相对 eager 最快 1.93 倍，而实现只用了 1/4 的代码量。Breakable CUDA Graph 由 SGLang 提出、命名并率先开源，现在已经是 SGLang prefill 路径上的默认方案。 1. Breakable CUDA Graph（BCG）解决了编译器切图方案启动慢、兼容性差的问题。 2. 在此之上，SGLang 用 token bucket 和 zero-length sentinel 让 prefill 能作为一整张 full CUDA Graph 被捕获。
背景
首先我们先要了解一下 CUDA Graph 的基本优化原理：一次推理并不是单个 kernel，而是由许多 GPU 操作组成的序列。CUDA Graph 会把 GPU workload 记录一次，之后一次性 replay 所有 kernel，省下 CPU 逐个 launch 小 kernel 的开销。
但在维护 SGLang 的过程中，我们发现把 CUDA Graph 用好远不止「全部操作录一遍」就完事了。有很多操作和 CUDA Graph 并不兼容。这些操作散落在 forward 中间，只要出现一个不兼容的，一大段 forward 就用不上 CUDA Graph。
因此，自然而然就会想到在这些不兼容操作前把 CUDA Graph 截断、之后再接上。下图对比同一段 forward 的三种执行方式：
BCG 的直觉很简单，工程实现并不简单。我们最初的答案是借助 torch.compile，下文称之为 tc_piecewise：先让编译器 trace 完整 forward，再在注册的切分点把图拆开。它能达成目标，但编译时间随模型复杂度增长，而且适配我们自己的定制优化 kernel 很别扭。这条路线的流程如下：
| Backend | 分段方式 | 边界怎么定 |
|---|---|---|
| full | 每个 shape 捕获一张完整 graph，无 eager region | 不需要边界 |
| breakable（BCG） | capture 过程中直接分段，segment 由 eager region 分隔 | @eager_on_graph 标记，capture 时插入 |
| tc_piecewise | torch.compile 先 trace 完整 forward，FX graph 在注册的切分点拆开 | 边界上的输入输出必须能被编译器表示 |
既然 capture 本身就是按顺序执行算子，断点为什么要靠编译器去找？顺着这个想法，我们做出了 Breakable CUDA Graph：在 capture 进行当中直接插入 eager 断点，不需要编译器理解 forward。同样的分段执行只用四分之一的代码量（521 行对 1771 行）；因为无需编译，构建 prefill graph 快 3.8–5.2 倍；只测 prefill，BCG 相对 eager 快 1.70 倍。BCG 最先落在 decode 路径上（#19102），随后扩展到 prefill（#22218），现在是 prefill 的默认方案。
接下来，我们会具体讲解 SGLang 工程上如何实现 BCG 和 Prefill CUDA Graph。
Breakable CUDA Graph
设计与机制
BCG 允许开发者用 @eager_on_graph 直接标记不兼容区域：capture
执行到被标记函数时，当前 graph segment 结束；函数以 eager 模式运行；随后开启新的
segment 继续 capture。forward 由此变成由 eager region 分隔的一系列 CUDA Graph
segment。
capture 阶段，被标记的函数会在两个 segment 之间运行一次，其返回的 tensor 被保留为持久化的 boundary buffer，device address 固定不变；接着后一个 segment 基于这个地址完成捕获。replay 时，eager function 正常执行并返回新的 tensor，BCG 再复制到保留的 buffer 中。下图按 capture 一次、replay 两次演示这一机制：
优势
启动更快。 对比传统基于编译器的 piecewise graph，准备阶段的主要开销来自编译而不是 capture：torch.compile 占 prefill graph 准备时间的 78–86%，并随模型复杂度增长。BCG 把这 78–86% 的占比完全消除，构建 prefill graph 的速度快 3.8–5.2 倍（每种配置捕获 42 个 shape，TP4，4×GB300）。
兼容范围更广。 SGLang 大量使用自定义 CUDA、Triton 和 JIT 编译 kernel 优化，它们并不是原生 PyTorch operator。为了让 torch.compile 识别这些 kernel，通常需要通过 torch.library 封装，并提供用于 trace 的 fake implementation。更重要的是，断点位置也会被限制：跨过边界的输入输出必须是编译器能表示的类型，碰上特殊的运行时状态或返回值，就得挪切分点，或者把 eager region 划大，只为给编译器一个能看懂的接口。这导致过去要让 CUDA Graph 正常工作，越来越像是在维护一个 torch.compile 集成项目。BCG 解除了 eager 断点处的这项限制：并且也不需要单独维护编译器。同样的分段执行，BCG 用 521 行实现，vs 旧方案 1771 行。
调试更容易。 CUDA Graph 一旦捕获，replay 的时候就是个黑盒：普通 Python
代码不会在内部执行，因此很难使用 print、assert 或逐步检查。BCG 自然保留了 eager
region，普通 Python 代码在每次 replay 时仍能执行。SGLang 还通过
--debug-cuda-graph 扩展了这一思路：这个选项会把整个 decode forward 包在一个
eager 断点中，模型以 eager
模式执行。这形成了一个很有用的调试边界：如果问题依然存在，原因很可能在模型或
runner 路径中；如果问题消失，则应优先怀疑 capture 本身。
BCG 也已用于 SGLang 的 diffusion stack（#27436）：单张 B200 上 512×512 的 Qwen-Image 端到端延迟从 6.48 秒降到 2.45 秒，Z-Image 从 1.231 秒降到 0.662 秒。
进一步：Prefill 的 full CUDA Graph
前面提到，BCG 靠断开 CUDA Graph 换来兼容性：不兼容算子留在图外，代价是每个断点处保留一小段 eager 执行。而相对应的，full CUDA Graph 把整个 forward 捕成一张图，没有任何 eager region，replay 时的发射次数也最少。这一节要把同样的捕法推到 prefill 上。
让 prefill 静态化
继承下来的技术难点。 相比于 decode AR step 固定的大小，prefill 的捕捉难度在于两个维度都会变：一次 prefill 的长度，以及这些 token 所属的 request 数。而 CUDA Graph 要求这两个维度都保持固定。加之，部分 attention backend 依赖运行时 metadata。因此，full CUDA Graph 一度很难应用于 prefill 阶段，这也是 prefill 路径采用 BCG 的主要原因之一。最近，#27988 重新设计了 request slot 和 attention metadata 的表示方式，让受支持的 attention backend 不必再放到 graph 外执行。
解决办法。 SGLang 做两件事。
- 用 token bucket 固定 token 维度：实时 batch pad 到最近的已捕获 token 数。
- 单独处理 request 维度：每张 graph 预留固定数量的 request slot，真实 request 占前面几个，剩下的重写成长度为 0 的 sentinel。
request 数超过预留的 slot 就回退到 eager。sentinel 和 attention metadata 都要在 graph 外、每次 replay 时重建，所以目前只有 FlashAttention（fa4）和 FlashInfer 这两个按这种方式准备 extend-mode metadata 的 backend 支持 prefill full capture。prefill full capture 目前仍是实验性功能，必须显式开启；engine 会提示 full 处于实验阶段，并建议生产 workload 使用 BCG 或 tc_piecewise。扩大 backend 支持范围、继续调整 bucket 和 slot 的选择，仍是后续工作。
如下图所示，图里展示了不同场景下 token 维和 request 维分别是怎么 pad 的：
Prefill benchmark
我们单独测量了 prefill：固定输入长度、只生成一个输出 token、每次处理一个 request，并在所有对比项中关闭 decode graph。测试使用 gpt-oss-120b（TP4，4×GB300），四条路径都能正常运行。结果显示，full capture 相比 eager 快 1.93 倍，BCG 快 1.70 倍，tc_piecewise 快 1.45 倍。因此 BCG 不只构建速度更快，replay 性能也比基于编译器的 backend 高 17%。差距来自每次 forward 的执行方式：BCG 直接 replay 已经记录好的 segment，tc_piecewise 每次都要重新调用 compiled callable，在执行自身捕获的 piece 之前，先支付 Torch Dynamo 的 guard check 和 dispatch 开销。在 GLM-5.2 上，只有 BCG 能完成 capture：tc_piecewise 无法 trace 它的 forward，full capture 也不支持它的 sparse attention；BCG 在该模型上相比 eager 快 1.60 倍。在 32 倍的 prompt 长度变化范围内，每条曲线都基本保持水平，这正是 launch 开销主导而非计算主导的典型特征：
CUDA Graph 的显存占用
前面讲的两件事都会影响显存：BCG 把 graph capture 切成多段，full capture 又把更大的 prefill shape 捕进来。CUDA Graph 显存是常驻的，replay 要求中间 tensor 的地址不变，capture 时分配下来就一直留到进程结束。所以有两个问题。
- BCG 切成多段之后，常驻显存不能跟着段数一起变多。
- 常驻的 graph 显存要能顶替最坏情况下的 activation 峰值，capture 就必须覆盖到最大的 prefill shape。
分段捕获内部的显存复用
设想一次 prefill forward 被 BCG 切成 N 段。如果每段的中间结果都各占一块显存并长期保留，graph 显存很容易成倍增长（N 段 × 捕获的 shape 数）。解决办法很简单：CUDA Graph 是按顺序 replay 的，同一时刻只有一段在执行，因此并不需要同时保住所有段的中间结果，后一段直接覆写前一段留下的数据，就能共享同一块显存。BCG 通过三种复用方式避免了这种增长。
- 不同 segment 共享一个显存池。 同一个 capture shape 的所有 segment 都使用同一个 CUDA Graph pool，让中间存储可以复用，而不是由每个 segment 分别长期占用。
- 在 Eager 断点处使用弱引用。 如果传入断点的 tensor storage 已由 graph pool 管理，就使用弱引用持有该 tensor，避免不必要的 Python reference 延长 tensor 生命周期。Tensor 弱引用技术来自 vLLM PR #9724；该 PR 引入这一技术，是为了让捕获的 graph 能够共享输出 buffer，而不是每张 graph 各自固定一块 buffer。
- 不同 capture size 共享一个输出 buffer。 所有 capture size 共用一块按最大尺寸分配的输出 buffer，每个 shape 只切出自己需要的行，避免为每个 shape 单独分配输出 buffer。
有了这些复用机制，即使 capture table 很大，显存开销仍然可控：在 GLM-5.2 上，为 78 层 MoE 捕获 42 个 shape，只增加了 2.4 GB graph 显存。
Capture 覆盖整个 chunked-prefill size
捕一个 shape，等于用常驻显存换掉它的瞬时峰值，这笔交换只对真正 replay 的 shape 成立。如果 capture 停在某个 size n，更大的 prefill size（>n）照旧回退 eager，峰值一点没少，小 shape 的常驻显存却已经花掉了，总显存比不开 graph 还高。所以重点永远是 capture 的上限，不是捕了多少个 shape。
chunked_prefill_size 是单次 prefill forward 的上限，一直捕到这个 size，最大的
prefill 也走 replay，最坏情况的 activation 峰值就被消除了。gpt-oss-120b 上从
0.56 GB 降到 0.001 GB，GLM-5.2 从 1.55 GB 降到 0.35 GB，总显存分别比不开 graph
的 baseline 低 0.51 GB 和 1.10 GB。如图所示，虽然从整体几百 GB
的显存占用来说，这个数字不算大，但是它意味着 CUDA Graph
带来的是节省而非成本。另一个好处是显存变得可预测：原来跟着 workload
波动的瞬时峰值，变成 capture 时定下的固定分配，engine 可以提前算进去，不用再为大
prefill 留 headroom：
峰值 0.56 GB
峰值 0.56 GB
峰值 0.001 GB
致谢
这项工作由 SGLang 团队与 Meta 团队合作完成。SGLang：Yuwei An*、Cheng Wan、Xiaoyu Zhang、Mick Qian、Baizhou Zhang、Yusheng Su、Ke Bao；Meta：Shiyang Chen*、Lianmin Zheng（* 同等贡献）。同时感谢 NVIDIA、AMD、Thinking Machines Lab 和 Meta PyTorch 团队在整个过程中给予的帮助。
延伸阅读
可视化教学由 Yichi Zhang 写作。