SGLang 深度解析 CUDA Graph 推理优化技术

SGLang 中的 CUDA Graph：Runner/Backend 分离与灵活组合

TC 分段式 CUDA Graph

可中断 CUDA Graph：无需编译器的 eager 中断

Prefill 完整 CUDA Graph

让 prefill 静态化

padding 的代价是什么？

CUDA Graph 的内存占用

分段捕获内的复用

通过 chunked-prefill 大小进行捕获

CUDA Graph 承诺消除 kernel-launch 开销，但在真实的推理引擎中要获得接近这一收益，需要在不牺牲兼容性、启动时间或内存的前提下将尽可能多的工作负载纳入图中。

在 SGLang 中，我们围绕通用的 runner/backend 接口重构了 CUDA Graph 支持，使不同的捕获策略可以在不同执行路径间复用。对于更复杂的 prefill 路径，SGLang 社区引入了可中断 CUDA Graph（Breakable CUDA Graph），并在使用 FA4 和 FlashInfer attention 后端的场景中率先实现了完整 CUDA Graph 支持。这两项技术最初都是在 SGLang 中作为开源服务技术开发的。我们还将深入探讨 CUDA Graph 的内存管理，包括跨 shape 和图段的内存复用，这正成为 SGLang 整体内存管理中越来越重要的一部分。

对于 prefill，可中断 CUDA Graph 现已是 SGLang 的默认方案。它以约四分之一的代码量（521 行对 1771 行）达到了与基于 torch.compile 的分段后端大致相同的分段执行效果，由于不涉及编译，构建 prefill 图的速度快了 3.8–5.2 倍，且对复杂功能具有更自然的覆盖范围。Prefill 完整 CUDA Graph 则更进一步，通过请求 padding 来捕获整个 forward，即使对于动态 prefill 工作负载也不例外。单独针对 prefill 测量，BCG 比 eager 执行快 1.70 倍，完整捕获则达到 1.93 倍。

推理步骤不是单个 kernel，而是一系列 GPU 操作。在现代 LLM 服务引擎中，反复从 CPU 启动这些操作会引入显著的开销，尤其对于延迟敏感的工作负载。CUDA Graph 通过一次性记录 GPU 工作并在后续重放时大幅降低启动开销来减少这种开销。

但在现代推理引擎中有效应用 CUDA Graph 并非易事。图的设计必须适配不同的执行阶段，与复杂的 kernel 和运行时相关行为保持兼容，并控制图本身带来的捕获时和内存开销。随着推理栈变得越来越复杂，适当的 CUDA Graph 集成变得越来越重要。

这篇文章将带你了解 CUDA Graph 支持在 SGLang 中是如何构建的，以及我们做了哪些改变：

• SGLang 中的 CUDA Graph：Runner/Backend 分离与灵活组合
• 可中断 CUDA Graph：无需编译器的 eager 中断
• Prefill 完整 CUDA Graph
• CUDA Graph 的内存占用

SGLang 中的 CUDA Graph：Runner/Backend 分离与灵活组合

在此重构之前，CUDA Graph 支持是围绕各个执行路径独立发展的。Decode、prefill 和投机解码各自都有独立的 CUDA Graph runner，捕获 shape、静态缓冲区、重放和图配置的逻辑相互重叠。随着更多执行模式和捕获策略的加入，这种重复使得基础设施复用变得困难，也使得 CUDA Graph 相关的服务端参数越来越模糊。

重构 [#23906] 将这些职责分离为两层。Runner 管理捕获和重放所需的执行特定状态：捕获的 shape、静态输入缓冲区、attention 元数据，以及将 live batch padding 到捕获 shape 的操作。Backend 则决定如何捕获该执行——是作为一个完整图、分段的可中断段，还是编译器生成的分片。

由于 runner 只依赖于通用的 backend 接口，每个执行路径可以独立选择其捕获策略。Prefill 和 decode 有独立的 runner，投机解码在此基础上增加更多：EAGLE draft、draft-extend 和 frozen-KV MTP draft 步骤各自由基于 decode runner 构建的独立 runner，而 target verify 本身就是 decode runner，对每个请求捕获超过一个 token。

Runner 为捕获和重放准备每个执行路径，而 backend 则决定如何将 forward 转化为可重放的图：作为一个完整图、在捕获期间分段，或在捕获前先 tracing 并分割。

完整后端（Full backend）为每个选定的 shape 捕获一个 torch.cuda.CUDAGraph，没有 eager 区域，且在三种后端中重放时启动次数最少。这对 decode 来说很自然：每个请求贡献一个 token，因此主要的 shape 变量是 batch size，可以由一组捕获的 batch-size bucket 覆盖。Prefill 在更多维度上存在变化，因此更难处理；我们将在独立章节中讨论。

可中断 CUDA Graph（BCG）捕获图安全的区域，同时允许选定的操作在图段之间 eager 运行。不兼容的操作可以用 @eager_on_graph 标记；捕获在标记函数之前停止，并在之后恢复，产生由 eager 区域分隔的一系列 CUDA Graph 段。

与基于编译器的分段捕获不同，这些中断是在捕获期间直接插入的，而不是通过先 tracing 完整模型来发现的。下一节我们将讨论这个机制以及 SGLang 为何转向这个设计。

TC 分段式 CUDA Graph

第三种后端通过编译器实现类似的分段。torch.compile 以 fullgraph=True tracing forward，在注册的分割点处拆分生成的 FX 图，每个分片独立编译和捕获。这是 SGLang 对部分 CUDA Graph 捕获的第一个解决方案，目前仍在 breakable 捕获尚未通过验证的平台上使用。

可中断 CUDA Graph：无需编译器的 eager 中断

CUDA Graph 传统上要求被捕获的区域完全图兼容。但在实践中，现代推理工作负载包含无法直接捕获的操作。Prefill attention 是一个常见的例子：某些 attention 后端依赖于运行时元数据和 host 端准备。一个不兼容的操作因此可能阻止 CUDA Graph 覆盖 forward 的更大一部分。

我们引入了可中断 CUDA Graph（BCG）来使捕获更加灵活。该机制和 @eager_on_graph 装饰器最初作为 CUDA Graph 调试模式的一部分登陆于 [#19102]，随后在 [#22218] 中被构建为 prefill 的可中断分段后端。BCG 不再要求整个 forward 都图兼容，而是允许选定的操作 eager 运行，同时捕获周围的图兼容区域。从高层次看，forward 变成了一系列由显式 eager 中断连接的 CUDA Graph 段。

CUDA Graph 在重放遵循固定的 GPU 操作序列且无需 host 参与时效果最佳。然而，真实的推理 forward 包含不适合这种模型的操作：attention 后端可能需要根据实时序列长度做规划，collectives 可能涉及运行时协调，而服务功能可能动态更新状态。

每当这样一个操作出现就放弃 CUDA Graph 将使 forward 的很大一部分未被捕获。BCG 则允许开发者直接用 @eager_on_graph 标记不兼容的区域。在捕获期间，当执行到达被标记的函数时，当前图段关闭，函数 eager 运行，之后在新的段中恢复捕获。

在重放时，记录的图段和 eager 函数按相同顺序运行。跨越 eager 中断的张量由前面的捕获段创建，并注册为持久边界缓冲区，因此其设备地址保持固定。随后的捕获段针对同一地址进行捕获。在重放期间，eager 函数因此将其新计算的结果写回这个边界缓冲区，而不是返回新分配的张量，使下一个段能够从其最初捕获时使用的地址读取更新后的值。BCG 从不检查或 tracing eager 区域内的操作：它们只需要正确执行即可。

从功能角度看，BCG 和早期基于 torch-compile 的分段后端产生相同类型的可重放结构：由 eager 区域分隔的 CUDA Graph 段。关键区别在于这种结构的构建方式。TC 分段先让编译器理解完整的 forward，然后再拆分生成的图。BCG 则在捕获进行时直接放置分割点。

更快的启动。 对于基于编译器的分段图，编译——而非捕获——主导了设置过程：torch.compile 占用了准备 prefill 图所花费时间的 78–86%，且随模型复杂度增长，在 235B MoE 上达到 90 秒，在 GLM-5.2 上达到 158 秒。BCG 完全省略了这个阶段，在单次捕获过程中达到分段执行。

编译开销在日常开发中也同样明显。在当时的 CI 配置中，编译往往在各测试运行之间重复进行，使得 CUDA Graph 测试明显变慢。更完善的缓存可以缓解这一问题，但从捕获路径中移除编译器也消除了开发循环中这一额外的复杂性来源。

更广泛的兼容性。SGLang 严重依赖自定义的 CUDA、Triton 和 JIT 编译内核，这些并非 PyTorch 原生算子。为了使这些内核对 torch.compile 可见，我们经常需要通过 torch.library 包装它们，并为追踪提供假实现。这在内核堆栈中引入了大量针对编译器的脚手架代码。

更重要的是，编译器还限制了图边界可以放置的位置。穿越注册算子边界的输入和输出必须能够被编译器表示。当自然边界涉及更专业的运行时状态或返回类型时，我们有时不得不寻找不同的切割点，或扩大 eager 区域以暴露编译器可以处理的接口。随着服务堆栈的扩展，编译器边界越来越影响那些原本与编译无关的代码结构。

BCG 在 eager 断点处消除了这一约束：图系统不需要理解标记函数的实现方式或追踪其内部逻辑，这使得图边界可以遵循服务逻辑而非编译器的追踪和类型要求。由于 CUDA Graph 必须与 DP attention、MoE all-to-all 后端、LoRA、PD disaggregation、分层缓存、确定性推理以及其它快速演进的功能共存，使 CUDA Graph 工作的难度日益感觉像是一个 torch.compile 集成项目。新内核往往意味着自定义算子注册和假实现，而新功能可能迫使我们仅仅为了满足编译器而移动图边界。有了 BCG，不兼容的区域可以保持为普通的 eager 执行，大大减少了这种针对编译器的工程开销。

天然可调试。捕获的 CUDA Graph 以不透明单元回放：普通的 Python 不会在其中执行，这使得打印、断言和逐步检查变得困难。BCG 自然地在 eager 区域保留了普通 Python 在每次回放时仍然运行的能力。

SGLang 通过 --debug-cuda-graph [#19102] 扩展了这一思路，它实际上将整个前向传播包装在一个 eager 断点中。模型随后以 eager 方式执行，同时仍然经过 CUDA Graph 运行器、静态缓冲区、回放路径和元数据准备。这提供了一个有用的调试边界：如果问题依然存在，可能在模型或运行器路径中；如果消失了，则捕获本身成为主要怀疑对象。

BCG 也已被 SGLang 的扩散堆栈采用 [#27436]。扩散在去噪过程中反复执行相同的 DiT 前向传播，当这些前向传播包含许多小型、launch-bound 内核时，CUDA Graph 尤其有用。

捕获实际服务的形状。分辨率、视频帧数、提示词条件长度、CFG 模式以及所选的 transformer 都会影响捕获签名。我们对实际服务的形状进行预热，对于未见的签名则回退到 eager 执行。

在动态操作周围断开。诸如动态 attention 和运行时相关元数据准备等操作保持 eager，而 BCG 捕获它们周围的稳定计算，无需 torch.compile 理解完整的 DiT 前向传播。

利用重复的去噪结构。扩散在去噪步骤中反复执行相同的 DiT 结构。BCG 捕获稳定区域一次，然后在整个去噪循环中回放它们，而动态区域保持 eager。

在执行受 launch 限制时这尤其有效。例如，预热后，Qwen-Image 在单卡 B200 上以 512×512 分辨率从 6.48 秒改善到 2.45 秒端到端延迟，Z-Image 从 1.231 秒改善到 0.662 秒。

预热后的端到端延迟。每对柱状图使用相同的模型工作负载和随机种子。

更广泛的教训是，BCG 消除了 launch 开销；它不会减少模型 FLOPs 或使计算受限的内核更便宜。它的优势在暴露的 launch 间隙占执行时间比重较大时最大。

Prefill 的完整 CUDA Graph

由于每个请求贡献一个 token，完整 CUDA Graph 对于 decode 来说很简单：主要变化维度是 batch size。Prefill 更难，因为批在两个维度上同时变化——token 总数和这些 token 所属的请求数——而捕获的图要求这两个维度都保持固定。再加上依赖运行时元数据的 attention 后端，这使得完整 CUDA Graph 难以应用于 prefill，也是我们在这里采用 Breakable CUDA Graph 的主要原因之一。

最近，我们找到了使 prefill 执行对完整 CUDA Graph 足够静态的方法 [#27988]，包括重构请求槽和 attention 元数据的表示方式，使得支持的后端不再需要留在图外。

使 prefill 静态化

SGLang 用 token bucket 修复 token 维度。活跃批次被填充到最近的捕获 token 数量，类似于 decode 将 batch size 填充到捕获的 bucket。

请求维度单独处理。每个捕获的图预留固定数量的请求槽。活跃请求占据前面的槽；未使用的槽被重写为零长度哨兵，零序列和扩展长度以及偏移量停在真实 token 之后。如果批次包含的请求多于图的槽数，则回退到 eager 执行。

在回放时，token 被填充到捕获的 bucket，而未使用的请求槽用零长度哨兵填充。

哨兵元数据必须在每次回放时重写，因为捕获的图仍然读取整个请求表。Attention 元数据同样在图外为填充的批次重建，然后回放。因此，完整的 prefill 捕获目前要求 attention 后端支持这种元数据准备方式，包括 FlashAttention 和 FlashInfer。

填充的代价是什么？

两种形式的填充代价截然不同。

填充的 token 是真实的工作。它们成为捕获批次中的实际行，因此作为相同 GEMM 的一部分通过密集投影。SGLang 单独携带真实的 token 数量，允许 MoE 路由、attention 和线性 attention 内核跳过大部分填充区域，但密集计算仍需为这些额外行付出代价。

空的请求槽则便宜得多。在 FlashAttention 的变长调度器中，工作来自每个序列的实际长度，而非为每个请求槽分配固定量的计算。因此零长度请求基本上不贡献 attention 工作；它主要增加元数据和少量调度开销。

这种不对称性很重要：token 填充是昂贵的维度，而请求槽填充相对便宜。

完整的 prefill 捕获仍然是实验性功能。它必须被显式启用——引擎警告 full 是实验性的，并指向 breakable 或 tc_piecewise 用于生产工作负载——而且它目前主要在 FlashAttention（fa4）和 FlashInfer 后端上工作，因为它们以捕获路径需要的方式构建 extend-mode 元数据。扩大后端支持以及调整 bucket 和槽的选择仍在进行中。

有了三种捕获 prefill 的方式和 eager 基线，剩余问题是每个方式在回放时的代价。在仅测量 prefill 的条件下——固定输入长度、单输出 token、一次一个请求、各分支中禁用 decode 图——在 gpt-oss-120b（TP4，4×GB300）上，四条路径都能运行：完整捕获比 eager 快 1.93 倍，BCG 1.70 倍，TC piecewise 1.45 倍，因此 BCG 在回放时也比基于编译器的后端快 17%，不仅在构建时如此。在 GLM-5.2 上只有 BCG 能够捕获——TC piecewise 无法追踪前向传播，完整捕获对其稀疏 attention 没有路径——比 eager 快 1.60 倍。在 32 倍的提示词长度范围内每条曲线都是平的，这是 launch 开销而非计算的特征。

gpt-oss-120b 上的仅 prefill 延迟，四种后端都能运行。

CUDA Graph 的内存占用

内存带来两个独立挑战：防止分段捕获增加常驻内存，以及捕获足够远以使常驻图内存实际取代最差的 eager 激活峰值。

分段捕获内的复用

分段后端很容易成倍增加图内存：每个捕获的形状包含多个图段，每段都有必须在回放时保持有效的中间结果。BCG 通过三种形式的复用避免这种成倍增长。

多个分段共享同一份显存池。每个已捕获形状的各个分段都使用同一个 CUDA Graph 显存池，允许中间存储被复用，而不必为每个分段单独 pinning（固定）内存。

在 eager break 处使用弱引用。当图池已经拥有某个张量的存储时，传入 break 的张量以弱引用方式持有，避免不必要的 Python 引用导致张量生命周期被延长。这一张量弱引用技术来自 vLLM [#9724]，该 PR 引入此技术使得捕获的图可以共享输出缓冲区，而不必各自 pinning 一份。

多个捕获形状共享一个输出缓冲区。各个捕获形状共用一个最大尺寸的输出缓冲区，按每个形状所需的行数进行切片，而不是为每个形状单独分配一个输出缓冲区。

有一个值无法用这种方式处理：跨 eager break 传递数据的张量。下一个图分段按照其地址进行捕获，因此该缓冲区必须保持存活状态，并在每次 replay 时原地更新。

借助这些复用机制，即使是一个大型捕获表也依然相当克制：GLM-5.2 上 78 层 MoE 的 42 个形状仅增加了 2.4 GB 的图显存。

通过 chunked-prefill size 进行捕获

CUDA Graphs 改变了 prefill 内存使用的形态。图显存在捕获期间分配，并在服务器整个生命周期内一直保留，属于常驻内存。Eager 激活是瞬态的：每次 prefill 分配工作内存，而最大支持尺寸决定峰值。

捕获一个 prefill 形状会将大部分瞬态工作集移入图的常驻显存池。但这只对实际 replay 图的形状有帮助。如果捕获阶梯止步于最大 prefill 尺寸以下，则最大的 prefill 仍会回退到 eager 执行并保持原有的激活峰值——同时服务器还要为下方所有常驻图付出代价。

这使得捕获上限比捕获形状的数量更为重要。由于 chunked_prefill_size 限制了单个 prefill 的最大前向大小，捕获超过该尺寸可以消除最差的 eager 激活峰值。

无图常驻基线以上的 prefill 内存，在恰好为 chunked-prefill size 的一次 prefill 后测量。

低于 chunk 值的上限略微高于无图基线：它们增加了常驻图，而激活峰值保持不变。一旦上限达到 chunk size，最大的 prefill 终于 replay 一个图，该峰值骤然塌缩——在 gpt-oss-120b 上几乎归零（0.56 GB 降至 0.001 GB），在 GLM-5.2 上从 1.55 GB 降至 0.35 GB，因为其稀疏注意力索引器仍在 break 处 eager 运行。

捕获超过 chunked-prefill size 带来了两重收益：

更低的总内存。激活峰值不再按请求计费，总内存降至无图基线以下——gpt-oss-120b 降低 0.51 GB，GLM-5.2 降低 1.10 GB。相对于数百 GB 的总占用空间来看幅度不大，但这是收益而非代价。

可预测的内存使用。工作负载依赖的激活峰值变为在捕获时确定的固定分配。引擎可以提前核算该内存，而不必为仅在大 prefill 期间出现的瞬态峰值预留余量。

本工作是 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 团队一路上的帮助。

（* 同等贡献）