Advanced CUDA Graph Techniques in SGLang

发表时间: 2026-08 · Blog post by SGLang Team (lmsys.org)

原文: https://lmsys.org/blog/2026-08-17-advanced-cuda-graph

SGLang团队,2026年8月17日

速读

一句话结论
SGLang 重构了底层架构并提出了 Breakable CUDA Graph (BCG) 与预填充全图捕获技术,通过在运行时动态插入断点或填充请求槽位,彻底摆脱了编译器切图的限制,在消除内核启动开销的同时实现了最高 1.93 倍的推理加速。

要解决什么问题
现代大语言模型推理时,CPU 频繁向 GPU 下发大量细粒度算子会产生严重的内核启动开销(Kernel-launch overhead)。CUDA Graph 能通过录制和回放机制消除这一开销,但在实际推理引擎中面临着严峻的兼容性卡点:预填充(Prefill)阶段的注意力机制等操作高度依赖运行时元数据和动态控制流,导致整个前向传播无法被单一的图完全捕获。此前业界通常依赖 torch.compile 进行图追踪和切分(即 TC piecewise 路线),但这带来了极高的机制代价。首先是编译耗时极长,复杂模型构建图需数十至上百秒;其次,它要求所有自定义 CUDA 或 Triton 内核必须提供伪实现以满足编译器的类型推导;更致命的是,编译器的追踪能力强行限制了图边界的划分位置,使得图捕获逻辑与底层服务逻辑深度耦合。此外,如果图捕获未能覆盖最大的预填充长度,系统在处理长文本时仍会回退到动态执行,导致显存中同时存在常驻的图显存和极高的动态激活显存峰值,极大地限制了并发能力。

怎么做的
核心思路是将图捕获的执行状态与捕获策略解耦,并引入无需编译器的运行时断点机制。首先,架构被重构为 Runner 和 Backend 两层:Runner 负责管理捕获形状、静态输入缓冲区和批处理填充等独立于执行路径的状态;Backend 则决定具体的捕获策略。为了绕开动态算子无法入图的卡点,作者提出了 Breakable CUDA Graph (BCG) 方法。开发人员只需用 @eager_on_graph 装饰器标记不兼容的动态函数。在捕获阶段,当执行到该标记函数时,系统会主动关闭当前图片段的录制,以原生动态模式运行该函数,随后立即开启新的图片段录制。其关键设计在于跨边界的数据传递:动态函数返回的张量会被保留在一个持久化的边界缓冲区中,使其设备内存地址保持固定,后续的图片段直接基于该固定地址进行捕获。在回放时,系统按顺序交替执行图片段和动态函数,新产生的张量会被拷贝回该固定地址,从而保证数据流转正确。针对预填充阶段,作者还设计了全图捕获(Full CUDA Graph)策略。由于预填充的 Token 数和请求数都是动态的,该策略在 Token 维度向预设的桶大小进行填充,在请求维度则预留固定数量的请求槽位,将未使用的槽位改写为零长度的哨兵请求(Sentinel)。由于 FlashAttention 等后端按实际长度调度,零长度请求几乎不产生计算代价。在显存管理上,BCG 通过在所有图片段间共享同一个 CUDA Graph 内存池、在断点处使用弱引用避免延长张量生命周期,以及跨捕获大小共享单一最大输出缓冲区,防止了分段捕获导致的显存膨胀。只要将图捕获的上限覆盖到分块预填充(Chunked prefill)的最大尺寸,就能用固定的常驻图显存完全替代掉原本极高的动态激活显存峰值。

效果如何
实验在 4 张 GB300 GPU(TP4 并行)上进行,测试模型包括 gpt-oss-120b 和 GLM-5.2。对比基线包括代表原生动态执行的 Eager 路线,以及代表基于编译器切图路线的 TC piecewise 方法。在纯预填充任务(固定输入长度、单输出 Token、单请求)的量化测试中,BCG 展现了显著优势。在构建阶段,由于彻底移除了编译环节,BCG 的图构建速度比 TC piecewise 快 3.8 到 5.2 倍。在推理延迟方面,对于 gpt-oss-120b,全图捕获策略比 Eager 基线快 1.93 倍;BCG 比 Eager 快 1.70 倍,且由于回放时无需经过 Torch Dynamo 的守卫检查,BCG 的回放速度比 TC piecewise 还要快 17%。在 GLM-5.2 模型上,TC piecewise 因无法追踪前向传播而失效,全图捕获也因不支持其稀疏注意力而无法运行,而 BCG 依然有效并取得了比 Eager 快 1.60 倍的加速。在显存表现上,将捕获上限对齐分块预填充尺寸后,gpt-oss-120b 的动态显存峰值从 0.56 GB 骤降至 0.001 GB,总体显存占用甚至低于无图基线。该方法的局限性在于,全图捕获目前仍是实验性功能,仅支持 FlashAttention (FA4) 和 FlashInfer 这类能处理哨兵元数据的注意力后端;此外,虽然零长度请求槽位代价极低,但 Token 维度的填充仍会使稠密线性层产生真实的无效计算开销。

主要贡献

在现代大型语言模型(LLM)推理引擎中,CUDA Graph有望消除内核启动(kernel-launch)的开销。然而,要在真实的推理引擎中接近这一理论收益,需要尽可能多地将工作负载图化(graphing),同时不能牺牲兼容性、启动时间或内存。本文深入探讨了SGLang中CUDA Graph技术的演进与优化,其核心贡献如下:

  1. 架构重构:SGLang围绕通用的运行器/后端(runner/backend)接口重构了CUDA Graph的支持,使得不同的捕获策略能够在各种执行路径(如解码、预填充、推测解码)中复用。
  2. 首创Breakable CUDA Graph (BCG):SGLang首次提出、命名、实现并开源了Breakable CUDA Graph这种服务技术。该技术无需编译器即可通过显式插入急切中断(eager breaks)来实现图的分段捕获,大幅提升了启动速度和兼容性。
  3. 实现预填充的全CUDA Graph支持:SGLang率先使用FA4和FlashInfer注意力后端,通过请求填充(request padding)技术,为动态预填充工作负载实现了完整的CUDA Graph支持。
  4. CUDA Graph内存管理优化:深入探讨了CUDA Graph的内存管理机制,包括跨形状和图片段的内存重用,显著降低了驻留内存的倍增风险。
  5. 显著的性能提升:BCG目前是SGLang预填充的默认设置。与基于torch.compile的后端相比,BCG的代码量仅为约四分之一(521行 vs 1771行),预填充图的构建速度提高了3.8–5.2倍。在纯预填充测试中,BCG比急切执行(eager execution)快1.70倍,而全图捕获(full capture)则快1.93倍。

背景知识与挑战

推理步骤的复杂性。推理步骤并非单一内核,而是许多GPU操作的序列。在现代LLM服务引擎中,从CPU重复启动这些操作会引入明显的开销,特别是对于延迟敏感的工作负载。

CUDA Graph的作用与挑战。CUDA Graph通过记录一次GPU工作并以极低的启动开销重放来减少这种开销。然而,在现代推理引擎中有效应用CUDA Graph并不简单。图形设计必须适应不同的执行阶段,保持与复杂内核和依赖于运行时的行为的兼容性,并控制图形本身引入的捕获时间和内存开销。随着推理栈变得越来越复杂,适当的CUDA Graph集成变得越来越重要。

核心方法设计

SGLang中的CUDA Graph:运行器/后端的拆分与灵活组合

重构前的基础设施痛点。在重构之前,CUDA Graph的支持是围绕各个独立的执行路径发展起来的。解码(Decode)、预填充(Prefill)和推测解码(speculative decoding)各自拥有独立的CUDA Graph运行器(runner),在捕获形状、静态缓冲区、重放和图形配置方面存在重叠的逻辑。随着更多执行模式和捕获策略的加入,这种代码重复使得基础设施的重用变得更加困难,也使得与CUDA Graph相关的服务器参数变得越来越模糊。

Runner与Backend的职责分离。重构(https://github.com/sgl-project/sglang/pull/23906)将这些职责分离为两层。运行器(runner)负责管理捕获和重放所需的特定于执行的状态:包括捕获的形状、静态输入缓冲区、注意力元数据,以及将活动批次(live batches)填充到捕获的形状中。后端(backend)则决定如何捕获该执行,无论是作为一个完整的图(full graph)、一系列可中断的片段(breakable segments),还是编译器生成的片段(compiler-generated pieces)。

各执行路径的独立策略选择。因为运行器仅依赖于一个通用的后端接口,所以每个执行路径可以独立选择其捕获策略。预填充和解码拥有各自独立的运行器,而推测解码则增加了更多:EAGLE草稿(draft)、草稿扩展(draft-extend)和冻结KV(frozen-KV)MTP草稿步骤各自获得了建立在解码运行器之上的专属运行器,而目标验证(target verify)本身就是解码运行器,每次请求捕获超过一个token。
图1 运行器为捕获和重放准备每个执行路径,而后端决定如何将前向传播转化为可重放的图:作为完整图、在捕获期间分段,或在捕获前追踪并分割。

全图(Full CUDA Graph)后端设计。全图后端为每个选定的形状捕获一个torch.cuda.CUDAGraph,没有急切执行(eager)区域,并且在三种后端中具有最少的重放时间启动次数。这对于解码来说非常自然:每个请求贡献一个token,因此主要的形状变量是批处理大小(batch size),这可以通过一组捕获的批处理大小桶(buckets)来覆盖。预填充在更多维度上变化,因此更难实现;这将在其专门的章节中讨论。

可中断图(Breakable CUDA Graph, BCG)机制。Breakable CUDA Graph (BCG) 捕获图安全的区域,同时允许选定的操作在图片段之间急切地(eagerly)运行。不兼容的操作可以被标记为@eager_on_graph;捕获在到达标记的函数之前停止,并在之后恢复,从而产生一系列由急切区域分隔的CUDA Graph片段。与基于编译器的分段捕获不同,这些中断是在捕获期间直接插入的,而不是通过首先追踪整个模型来发现的。

TC分段(TC piecewise CUDA Graph)机制。第三种后端通过编译器实现了类似的分段。torch.compile通过fullgraph=True追踪前向传播,生成的FX图在注册的分割点(split points)处被分割,每个片段被独立编译和捕获。这是SGLang对部分CUDA Graph捕获的最初解决方案,并且仍然在那些尚未验证可中断捕获的平台上发布使用。

Breakable CUDA Graph:无需编译器的急切中断

传统CUDA Graph的限制与实际冲突。CUDA Graph传统上要求被捕获的区域必须完全兼容图执行。在实践中,现代推理工作负载包含无法直接捕获的操作。预填充注意力(Prefill attention)就是一个常见的例子:一些注意力后端依赖于运行时元数据和主机端的准备工作。因此,单个不兼容的操作可能会阻止CUDA Graph覆盖前向传播中大得多的部分。

BCG的引入与核心思想。SGLang引入并首次开源了Breakable CUDA Graph (BCG),以使捕获更加灵活。BCG不是用编译器追踪整个前向传播,而是在捕获进行时插入显式的急切中断(eager breaks):选定的操作急切地运行,而捕获则在它们周围的图兼容区域恢复。从高层次来看,前向传播变成了一系列由显式急切中断连接的CUDA Graph片段。

BCG的开源历史与首创性。BCG起源于SGLang。其开发记录在公开的、按时间顺序排列的源码历史中:

  • 2026年2月21日:SGLang发布了第一个BCG实现。 https://github.com/sgl-project/sglang/pull/19102的初始提交已经包含了BreakableCUDAGraph、一个急切中断装饰器,以及结束和恢复CUDA Graph捕获的运行时机制。该PR于4月11日合并。
  • 2026年4月:SGLang将BCG扩展到预填充。 https://github.com/sgl-project/sglang/pull/22218明确基于#19102,将BCG构建为一个无需编译器的可中断分段后端用于预填充,并于4月24日合并。BCG随后成为SGLang默认的预填充CUDA Graph策略。
    这里的原创性声明特指这种用于开源LLM服务的运行时可中断捕获/重放机制:SGLang在它被其他地方采用之前提出、命名、实现并开源了该机制。这并不是声明图形分段或急切回退的更广泛概念没有现有技术。

BCG的设计与机制细节。当重放遵循固定的GPU操作序列且没有主机参与时,CUDA Graph效果最好。然而,真实的推理前向传播包含不自然适合该模型操作:注意力后端可能会根据实时序列长度进行计划,集合通信(collectives)可能涉及运行时协调,服务功能可能会动态更新状态。

处理不兼容操作的方法。如果每当出现这样一个操作就放弃CUDA Graph,将会导致前向传播的很大一部分未被捕获。相反,BCG允许开发者直接使用@eager_on_graph标记不兼容的区域。在捕获期间,当执行到达标记的函数时,当前的图片段被关闭,该函数急切地运行,然后捕获在一个新的片段中恢复。

重放时的内存与地址处理。在重放时,记录的图片段和急切函数以相同的顺序运行。在捕获期间,标记的函数在图片段之间运行一次,它返回的张量被保留为一个持久的边界缓冲区(persistent boundary buffer),因此其设备地址保持固定;随后的片段是针对该地址捕获的。在每次重放时,急切函数正常运行并返回一个新的张量,BCG将其复制到保留的缓冲区中,因此下一个片段从最初捕获时使用的地址读取更新后的值。BCG从不检查或追踪急切区域内的操作:它们只需要正确执行即可。

BCG与TC分段后端的对比。从功能角度来看,BCG和早期的基于torch-compile的分段后端产生了相同的可重放结构:由急切区域分隔的CUDA Graph片段。关键区别在于这种结构是如何构建的。TC分段首先要求编译器理解完整的前向传播,然后分割生成的图。BCG则在捕获发生时直接放置分割点。

消除编译开销带来更快的启动。对于基于编译器的分段图,编译(而非捕获)主导了设置过程:torch.compile占用了准备预填充图所花费时间的78–86%,并且它随着模型复杂度的增加而增长,在235B MoE上达到90秒,在GLM-5.2上达到158秒。BCG完全消除了该阶段,在单次捕获通道中就达到了分段执行。
图2 构建预填充CUDA Graph的时间,42个捕获形状,4×GB300上的TP4配置。

开发体验的改善。编译开销在日常开发中也很明显。在当时的CI设置中,编译经常在测试运行中重复,使得CUDA Graph测试明显变慢。更好的缓存可以缓解这个问题,但从捕获路径中移除编译器也从开发循环中移除了这个额外的复杂性来源。

摆脱编译器限制实现更广的兼容性。SGLang严重依赖定制的CUDA、Triton和JIT编译的内核,这些不是原生的PyTorch算子。为了让torch.compile看到这些内核,通常必须通过torch.library包装它们并提供用于追踪的伪实现。这在整个内核栈中引入了特定于编译器的脚手架。

图边界放置的自由度。更重要的是,编译器还限制了图边界可以放置的位置。跨越注册算子边界的输入和输出必须是编译器可表示的。当自然边界涉及更专门的运行时状态或返回类型时,有时必须寻找不同的切入点或扩大急切区域,仅仅是为了暴露一个编译器可以处理的接口。随着服务栈的增长,编译器边界越来越影响原本与编译无关的代码结构。

降低特定于编译器的工程开销。BCG在急切中断处消除了这种限制:图系统不需要理解标记的函数是如何实现的,也不需要追踪其内部,从而允许图边界遵循服务逻辑,而不是编译器的追踪和类型要求。由于CUDA Graph必须与DP注意力、MoE all-to-all后端、LoRA、PD解耦、分层缓存、确定性推理以及其他快速演进的功能共存,让CUDA Graph工作越来越感觉像是一个torch.compile集成项目。新内核通常意味着自定义算子注册和伪实现,而新功能可能会迫使移动图边界仅仅是为了满足编译器。有了BCG,不兼容的区域可以保持普通的急切执行,大大减少了这种特定于编译器的工程开销。

天然可调试性。捕获的CUDA Graph作为一个不透明的单元重放:普通的Python不会在其中执行,这使得打印、断言和逐步检查变得困难。BCG自然地留下了急切区域,在每次重放时普通的Python仍然在其中运行。

调试边界的扩展。SGLang通过--debug-cuda-graph [https://github.com/sgl-project/sglang/pull/19102] 扩展了这个想法,它有效地将整个前向传播包装在一个急切中断中。然后模型急切地执行,同时仍然通过CUDA Graph运行器、静态缓冲区、重放路径和元数据准备。这提供了一个有用的调试边界:如果问题仍然存在,很可能在模型或运行器路径中;如果它消失了,捕获本身就成为首要怀疑对象。

BCG在扩散模型中的应用。BCG也被SGLang的扩散栈(diffusion stack)采用 [https://github.com/sgl-project/sglang/pull/27436]。扩散在去噪过程中重复执行相同的DiT前向传播,当这些前向传播包含许多受限于启动(launch-bound)的小内核时,CUDA Graph特别有用。

扩散模型的具体策略。

  • 捕获真实的服务形状。 分辨率、视频帧数、提示条件长度、CFG模式和选定的Transformer都会影响捕获签名。预热实际提供的形状,并对未见过的签名回退到急切执行。
  • 在动态操作周围中断。 动态注意力(dynamic attention)和依赖于运行时的元数据准备等操作保持急切执行,而BCG捕获它们周围稳定的计算,无需torch.compile理解完整的DiT前向。
  • 利用重复的去噪结构。 扩散在去噪步骤中重复执行相同的DiT结构。BCG捕获稳定的区域一次,并在整个去噪循环中重放它们,而动态区域保持急切执行。

性能提升说明。这在执行受限于启动时特别有效。例如,预热后,单张B200上512×512的Qwen-Image端到端延迟从6.48秒改善到2.45秒,Z-Image从1.231秒改善到0.662秒。
图3 预热后的端到端延迟。每对柱状图使用相同的模型工作负载和种子。

核心优势总结。更广泛的教训是,BCG消除了启动开销;它不减少模型FLOPs或使受限于计算(compute-bound)的内核变得更便宜。当暴露的启动间隙占执行时间的有意义的比例时,它的优势最大。

预填充的全CUDA Graph(Full CUDA Graph for Prefill)

预填充面临的二维变量挑战。全图CUDA Graph对于解码来说很简单,因为每个请求贡献一个token:主要变化的维度是批处理大小。预填充更难,因为批次同时在两个维度上变化——总token数和这些token所属的请求数——而捕获的图要求两者都保持固定。再加上依赖于运行时元数据的注意力后端,这使得全图CUDA Graph难以应用于预填充,这也是在那里采用Breakable CUDA Graph的主要原因之一。

使预填充静态化的突破。最近,SGLang找到了使预填充执行足够静态以用于全图CUDA Graph的方法 [https://github.com/sgl-project/sglang/pull/27988],包括重构请求槽(request slots)和注意力元数据的表示方式,以便受支持的注意力后端不再必须留在图外。

固定Token与请求维度。SGLang使用token桶(token buckets)固定token维度。一个活动批次被填充到最接近的捕获token数,就像解码将批处理大小填充到捕获的桶一样。请求维度是单独处理的。每个捕获的图保留固定数量的请求槽。活动请求占据前几个槽;未使用的槽被重写为零长度的哨兵(zero-length sentinels),其序列长度和扩展长度为零,且偏移量停放在真实token之后。如果一个批次包含的请求多于图的槽数,它将回退到急切执行。
图4 在重放时,token被填充到捕获的桶中,而未使用的请求槽被填充为零长度的哨兵。

重放时的元数据重建要求。哨兵元数据必须在每次重放时重写,因为捕获的图仍会读取整个请求表。同样,在重放之前,注意力元数据也会在图外为填充的批次重建。因此,目前的全面预填充捕获需要支持这种元数据准备风格的注意力后端,包括FlashAttention和FlashInfer。

填充成本的非对称性分析。这两种形式的填充具有非常不同的成本。填充的token是真实的工作。它们成为捕获批次中的实际行,因此作为相同GEMM的一部分通过密集投影。SGLang单独携带真实的token计数,允许MoE路由、注意力和线性注意力内核跳过大部分填充区域,但密集计算仍需为这些额外的行付出代价。

空请求槽的低成本。空请求槽要便宜得多。在FlashAttention的可变长度调度器中,工作量是从每个序列的实际长度推导出来的,而不是为每个请求槽分配固定数量的计算。因此,零长度请求基本上不贡献注意力工作;它主要增加元数据和少量的调度开销。这种不对称性很重要:token填充是昂贵的维度,而请求槽填充相对便宜。

全图预填充的实验性状态。全图预填充捕获仍然是一个实验性功能。它必须被显式启用——引擎会警告full是实验性的,并指出生产工作负载应使用breakable或tc_piecewise——并且它目前主要在FlashAttention (fa4) 和FlashInfer后端上工作,因为这些后端是以捕获路径需要的方式构建扩展模式元数据的。扩大后端支持以及调整桶和槽的选择仍在未来计划中。

CUDA Graph的内存占用

内存面临的双重挑战。内存带来了两个独立的挑战:保持分段捕获不使驻留内存成倍增加,以及捕获足够深以使驻留图内存真正取代最坏的急切激活峰值。

避免分段捕获中的内存乘法效应。分段后端很容易使图内存成倍增加:每个捕获的形状包含多个图片段,每个片段都有必须保持有效以供重放的中间变量。BCG通过三种形式的重用避免了这种乘法。

  • 跨片段共享单一内存池。 捕获形状的每个片段使用相同的CUDA Graph池,允许中间存储被重用,而不是为每个片段单独固定(pinned)。
  • 急切中断处的弱引用。 传入中断的张量在图池已经拥有其存储时被弱引用持有,避免了延长张量生命周期的不必要的Python引用。这种张量弱引用技术来自vLLM [https://github.com/vllm-project/vllm/pull/9724],引入它是为了捕获的图可以共享输出缓冲区,而不是每个图固定自己的缓冲区。
  • 跨捕获大小共享一个输出缓冲区。 捕获大小共享一个最大尺寸的输出缓冲区,将其切片为每个形状所需的行,而不是每个形状分配一个输出缓冲区。

跨边界数据的特殊处理。有一个值不能这样处理:跨越急切中断携带数据的张量。下一个图片段是针对其地址捕获的,因此该缓冲区必须保持活动状态,并在每次重放时就地更新。借助这些重用机制,即使是大型捕获表也保持适度:在GLM-5.2上,跨越78层MoE的42个形状增加了2.4 GB的图内存。

预填充大小对内存的影响。CUDA Graphs改变了预填充内存使用的形状。图内存是驻留的(resident):它在捕获期间分配,并在服务器的生命周期内保留。急切激活是瞬态的:每次预填充都会分配工作内存,最大的受支持预填充决定了峰值。

捕获天花板的重要性。捕获一个预填充形状将大部分瞬态工作集转移到图的驻留内存池中。但这只对实际重放图的形状有帮助。如果捕获阶梯停在最大预填充大小以下,最大的预填充仍然会回退到急切执行并保留原始激活峰值——同时服务器还要为它下面的所有驻留图付费。这使得捕获天花板(capture ceiling)比捕获形状的数量更重要。由于chunked_prefill_size限制了最大单次预填充前向传播,通过捕获达到该大小可以消除最坏的急切激活峰值。
图6 高于无图驻留基线的预填充内存,在刚好等于分块预填充大小的一次预填充后测量。

达到分块大小的内存收益。低于分块大小(chunk size)的天花板略微高于无图基线:它们增加了驻留图,而激活峰值保持在原来的位置。一旦天花板达到分块大小,最大的预填充终于重放了一个图,那个峰值就会崩溃——在gpt-oss-120b上基本降为零(0.56 GB到0.001 GB),在GLM-5.2上从1.55 GB降至0.35 GB,其稀疏注意力索引器仍在中断处急切运行。

收益总结。通过捕获达到分块预填充大小买到了两件事:

  • 更低的总内存。 激活峰值不再按请求支付,总内存降至无图基线以下——在gpt-oss-120b上低0.51 GB,在GLM-5.2上低1.10 GB。与几百GB的占用空间相比,这显得很适度,但它是节省而不是成本。
  • 可预测的内存使用。 依赖于工作负载的激活尖峰变成了在捕获时建立的固定分配。引擎可以预先核算该内存,而不是为仅在大型预填充期间出现的瞬态峰值保留余量。

实验环境配置

  • 硬件配置:

    • 主要测试平台:4×GB300,配置为TP4(张量并行度为4)。
    • 扩散模型测试平台:单张 NVIDIA B200 GPU。
  • 软件配置:

    • 推理引擎:SGLang,结合Meta PyTorch。
    • 依赖库:torch.compile,FlashAttention (fa4),FlashInfer。
  • 模型架构:

    • LLM测试:gpt-oss-120b,GLM-5.2(78层MoE,包含稀疏注意力),235B MoE模型。
    • 扩散模型测试:Qwen-Image(512×512分辨率),Z-Image。

实验评估结果

  1. 预填充图构建时间评估

    • 实验内容:在4×GB300 (TP4) 平台上,测量构建42个捕获形状的预填充CUDA Graph所需的时间。
    • 实验结果:基于编译器的torch.compile在235B MoE上耗时90秒,在GLM-5.2上耗时158秒。而BCG完全移除了编译阶段,构建速度提高了3.8–5.2倍。
    • 分析结论:torch.compile占据了图准备时间的78%–86%,BCG显著降低了启动开销。(见图2)
  2. 扩散模型端到端延迟评估

    • 实验内容:在单张B200 GPU上测量扩散模型在预热后的端到端延迟。
    • 实验结果:512×512的Qwen-Image延迟从6.48秒降至2.45秒;Z-Image从1.231秒降至0.662秒。
    • 分析结论:BCG在受限于启动(launch-bound)的小内核连续调用场景(如扩散去噪循环)中效果显著。(见图3)
  3. 预填充延迟基准测试

    • 实验内容:在gpt-oss-120b (TP4, 4×GB300) 上测试纯预填充延迟(固定输入长度,生成1个token,一次一个请求,禁用解码图)。
    • 实验结果:全图捕获(Full capture)比急切执行快1.93倍;BCG快1.70倍;TC分段快1.45倍。在GLM-5.2上,仅BCG能够成功捕获,速度比急切执行快1.60倍。
    • 分析结论:BCG不仅在构建时更快,在重放时也比基于编译器的后端快17%,因为它避免了Torch Dynamo的守卫检查(guard checks)和调度开销。所有曲线在32倍的提示长度范围内保持平坦,证实了这是启动开销而非计算瓶颈。(见图5)
    图5 gpt-oss-120b上的纯预填充延迟,其中所有四种后端都在运行。
  4. 内存占用评估

    • 实验内容:测量在刚好等于分块预填充大小(chunked-prefill size)进行一次预填充后,超出无图驻留基线的预填充内存。
    • 实验结果:捕获直到分块预填充大小消除了最坏的急切激活峰值。在gpt-oss-120b上,峰值从0.56 GB降至0.001 GB;在GLM-5.2上,峰值从1.55 GB降至0.35 GB。总内存均低于无图基线(gpt-oss-120b低0.51 GB,GLM-5.2低1.10 GB)。
    • 分析结论:通过将瞬态激活峰值转化为固定的驻留内存分配,不仅降低了整体内存占用,还使得内存使用具有高度的可预测性。(见图6)

结论与展望

本文系统性地展示了SGLang中高级CUDA Graph技术的演进。通过将运行器与后端分离,SGLang实现了高度灵活的图捕获策略。其首创的Breakable CUDA Graph (BCG) 摆脱了对编译器的依赖,通过显式插入急切中断,在显著缩短启动时间、增强系统兼容性的同时,提供了出色的可调试性。结合请求填充技术,SGLang进一步实现了预填充阶段的全图捕获。此外,通过共享内存池和弱引用技术,成功控制了分段图带来的内存膨胀,并通过覆盖最大分块预填充大小,将不可预测的瞬态内存峰值转化为更低且确定的常驻内存。

未来,SGLang团队计划进一步扩大全图预填充捕获的注意力后端支持范围,并继续优化桶和槽的分配策略。同时,BCG机制的独立库化将促进该技术在更广泛的学术和工业推理及训练引擎中的应用。

补充细节:独立代码库

独立库的发布。可中断捕获机制本身不仅与SGLang绑定,因此在与Meta PyTorch团队的合作下,将其提取为一个独立的库:https://github.com/meta-pytorch/breakable-cuda-graphs。

API设计。它在一个小型的PyTorch风格的API背后暴露了相同的捕获模型:一个类似于torch.cuda.graph的breakable_graph(...)上下文管理器,一个标记函数在图片段之间急切运行的@no_graph装饰器,以及一个按顺序重放捕获片段和急切区域的CUDAGraphSequence。如果在上下文中没有发生急切中断,该块被捕获为单个图片段——标准的CUDA Graph模式作为平凡情况出现。

应用前景。如果系统需要部分CUDA Graph捕获——在另一个推理引擎、训练循环或研究原型中——可以直接在学术或工业工作中使用该库,而无需采用SGLang的运行器栈。SGLang将继续开发该机制,并期望独立库和树内(in-tree)实现随着时间的推移而融合。