1. 从一次显存告警说起为什么要啃 BCG 这块硬骨头第一次在 SGLang 上跑 DiT 类扩散模型做批量推理的时候我盯着 nvidia-smi 的输出愣了半天——单张卡上 batch size 刚过 8显存就顶到 90% 以上吞吐量却低得可怜。按常理说扩散模型的 UNet 或者 DiT backbone 参数量并不夸张真正吃显存的是那些反复分配释放的中间激活张量。每次采样迭代都要重新走一遍完整的计算图PyTorch 的 caching allocator 虽然能复用一部分但 kernel launch 的开销和显存碎片的累积还是让整个 pipeline 跑得磕磕绊绊。后来翻 SGLang 的源码和社区讨论发现他们针对这个问题搞了一套叫BCG的机制。BCG 是Batch CUDA Graph的缩写核心思路是把整个去噪迭代过程或者其中一段固定的计算子图捕获成 CUDA Graph然后用固定的输入输出 buffer 去 replay从而把成百上千次 kernel launch 压缩成一次图执行。听起来很美好但实际调试的时候坑一个接一个动态 shape 怎么办、CFGclassifier-free guidance的双分支怎么处理、不同 timestep 的调度器状态怎么塞进静态图里……这些问题不解决CUDA Graph 捕获出来的东西要么跑不对要么直接崩。这篇东西就是把我从零实现和调试 SGLang Diffusion BCG 的完整过程记录下来。适合两类人看一类是正在用 SGLang 部署扩散模型、发现吞吐上不去的工程师另一类是想理解 CUDA Graph 在扩散推理里到底怎么落地、有哪些边界条件的技术同学。我会尽量把每个设计决策背后的“为什么”讲清楚而不是只丢一堆代码让你抄。毕竟 CUDA Graph 这东西抄错了轻则结果不对重则整个进程 hang 死排查起来非常痛苦。2. BCG 到底在解决什么问题从 kernel launch 开销到显存碎片2.1 扩散推理的瓶颈不在算力在调度先摆一个实测数据。在一张 A100 上跑 SDXL 的 DiT 变体单步去噪的纯计算时间大概是 12ms 左右但整个 step 的 wall time 经常到 18-20ms。多出来的 6-8ms 去哪了答案是 kernel launch 和 Python 侧的调度开销。扩散模型一个 step 里可能有几十上百个算子attention 的 QKV 投影、softmax、FFN 的两次线性、各种 norm、残差加法、timestep embedding 的注入……每个算子对应一次或多次 kernel launch每次 launch 在 CUDA 上大概有几微秒的固定开销累积起来就很可观。更麻烦的是 Python 解释器的开销。SGLang 的调度层是 Python 写的每个 step 都要走一遍 Python 的函数调用、条件判断、tensor 元信息检查。这部分开销和 GPU 计算是串行的除非你用多流所以 GPU 经常在等 CPU 喂数据。你可能会说那用 torch.compile 或者 TensorRT 不就行了torch.compile 确实能减少一部分 launch 开销但它对动态 shape 的支持有限而且编译时间在迭代开发阶段很折磨人。TensorRT 更重转换和调优的成本更高。CUDA Graph 的思路更底层也更直接既然这些 kernel 的启动顺序和参数在多次迭代之间是固定的那我干脆把整段执行序列录下来变成一个图之后每次只需要一次 graph launch 就能跑完。这就像你把一段经常重复的 shell 脚本编译成一个二进制省掉了每次解析脚本的时间。2.2 为什么叫 Batch CUDA Graph 而不是普通 CUDA Graph普通的 CUDA Graph 捕获要求输入输出的内存地址固定。但扩散推理里每个 batch 的输入 prompt embedding 是不一样的如果你每次都用同一块输入 buffer就得先把新数据拷进去。这本身没问题问题是 batch size 可能变化——今天来 4 个请求明天来 16 个如果为每个 batch size 都捕获一个图显存占用会爆炸。BCG 的“Batch”体现在两个层面。第一它把多个请求打包成一个固定的 batch 维度来处理通过 padding 或者 bucketing 的方式把变长 batch 规整到几个预设的档位比如 1、2、4、8、16。第二它在图内部对 batch 维度做了处理使得同一个图可以服务一定范围内的 batch 配置。具体实现上SGLang 用的是batch bucketing 图池的策略为每个 bucket 捕获一个独立的 CUDA Graph运行时根据实际 batch size 选择最接近的图多出来的部分用 mask 或者重复采样填充。这里有个关键权衡bucket 越多显存占用越大每个图都要占一份静态 buffer但 padding 浪费越少bucket 越少显存省了但小 batch 跑大图会浪费算力。我一开始图省事只设了 1 和 8 两个 bucket结果 batch2 的请求被塞进 batch8 的图里跑吞吐直接掉了一半。后来改成 1、2、4、8 四档情况就好多了。2.3 显存碎片CUDA Graph 的隐形杀手CUDA Graph 捕获的时候所有中间张量都要用静态内存池分配。如果你在捕获阶段让 PyTorch 的 caching allocator 自由分配捕获出来的图会引用一堆临时地址replay 的时候这些地址可能已经被别的 tensor 占了结果就是数据错乱或者 illegal memory access。SGLang 的做法是在捕获前切换到graph memory pool用一个专门的显存池来承载图内所有分配。这个池子在捕获完成后会被锁定replay 时所有中间结果都从池子里拿。好处是地址稳定坏处是这个池子的大小必须提前估好。估小了捕获直接 OOM估大了浪费显存。我踩过的坑是第一次捕获时池子设了 2GB结果跑到一半报 out of memory因为 DiT 的 attention 在长序列下激活值比预期大很多。后来用torch.cuda.memory_summary()在 eager 模式下跑一遍看峰值显存再乘个 1.3 的系数才把池子大小定下来。提示graph memory pool 的大小不是越大越好。池子太大留给 KV cache 和模型权重的显存就少了整体吞吐反而下降。建议先用 eager 模式测出单步峰值再留 20%-30% 余量。3. 捕获前的准备工作环境、版本与那些容易忽略的开关3.1 版本匹配比你想的更重要CUDA Graph 的 API 在 CUDA 11.x 到 12.x 之间有过几次行为变化PyTorch 对它的封装也在不断调整。我用的组合是 PyTorch 2.3 CUDA 12.1 SGLang 0.3.x这个组合相对稳定。如果你用 PyTorch 2.1 或更早的版本torch.cuda.graphs模块的一些参数名和默认行为不一样直接抄代码会报 TypeError。另外NCCL 的版本也要注意。如果你是多卡推理CUDA Graph 捕获会涉及 NCCL 通信算子。NCCL 2.18 之前对 graph capture 的支持有 bug捕获出来的 all-reduce 可能死锁。升级到 2.19 基本就没事了。单卡的话可以忽略这条。还有一个容易被忽略的点CUDA Graph 捕获期间不能有同步操作。任何.item()、.cpu()、print(tensor)都会导致捕获失败或者捕获出一个空图。我在调试的时候习惯性加了一行print(latents.shape)结果捕获出来的图 replay 时什么都不做排查了半天才发现是这行 print 惹的祸。3.2 关闭 cudnn benchmark 和 TF32 的坑torch.backends.cudnn.benchmark True在 eager 模式下能自动选最快的卷积算法但在 CUDA Graph 捕获时cudnn 可能会在捕获过程中做算法搜索导致捕获失败或者捕获到不确定的图。SGLang 在捕获前会强制把 benchmark 设为 False并且固定 cudnn 的 deterministic 选项。TF32 是另一个坑。A100 及以上默认开启 TF32矩阵乘的精度会降低。扩散模型对数值精度比较敏感尤其是采样步数多的时候TF32 累积误差可能导致最终图像出现色偏或者伪影。我在捕获前会把torch.backends.cuda.matmul.allow_tf32设为 False虽然慢一点但结果稳定。如果你追求极致速度且能接受轻微精度损失可以保持开启但要在捕获前后保持一致不能捕获时关、replay 时开。3.3 预热不是跑一遍就完事CUDA Graph 捕获前需要预热目的是让 PyTorch 的 allocator 完成显存分配、让 cudnn 完成算法选择、让各种 lazy 初始化都执行完。但预热跑多少遍有讲究。跑 1 遍往往不够因为有些 kernel 第一次执行时会做 autotuning第二次才用缓存。我一般跑 3-5 遍并且用不同的输入数据确保所有分支都被覆盖到。预热还有一个作用把显存池的峰值撑出来。如果你在预热时用的是小 batch捕获时用大 batch池子可能不够。所以预热要用目标 bucket 的最大 batch size来跑。# 预热示例 for _ in range(5): with torch.no_grad(): _ model(latents, timestep, encoder_hidden_states) torch.cuda.synchronize()注意torch.cuda.synchronize()不能省否则预热可能还没跑完就开始捕获导致捕获到不完整的图。4. 捕获流程拆解从静态 buffer 到图 replay 的完整链路4.1 静态输入输出 buffer 的设计CUDA Graph 要求输入输出的地址固定所以第一步是创建一组静态 tensor 作为图的入口和出口。对于扩散模型输入通常包括latents当前噪声 latentshape 为[batch, channels, height, width]timestep当前时间步shape 为[batch]或标量encoder_hidden_states文本编码器的输出shape 为[batch, seq_len, dim]guidance_scaleCFG 的引导系数标量输出是去噪后的 latentshape 和输入 latents 一致。这些静态 buffer 在捕获前分配好之后每次 replay 前把新数据 copy 进去。copy 的开销相对于整个去噪 step 来说可以忽略但要注意用copy_而不是重新赋值否则地址会变。static_latents torch.zeros(batch_size, 4, 128, 128, devicecuda, dtypetorch.float16) static_timestep torch.zeros(batch_size, devicecuda, dtypetorch.long) static_encoder_hidden torch.zeros(batch_size, 77, 2048, devicecuda, dtypetorch.float16) static_output torch.zeros_like(static_latents)这里有个细节timestep的 dtype。有些调度器用 float有些用 long。如果你捕获时用 longreplay 时传 floatCUDA Graph 不会报错但结果会莫名其妙不对。我建议在捕获前把 dtype 打印出来确认一遍。4.2 CFG 双分支的图内处理Classifier-free guidance 需要同时跑条件分支和无条件分支然后把结果按 guidance scale 混合。在 eager 模式下这通常是两次 forward 然后加权求和。但在 CUDA Graph 里如果你把两次 forward 都捕获进去图的大小会翻倍显存占用也翻倍。SGLang 的做法是把两个分支在 batch 维度上 concat 起来一次 forward 处理双倍 batch然后在图内做 split 和加权。这样只需要捕获一个 forward 图显存效率更高。具体来说encoder_hidden_states的 batch 维度变成2 * batch_size前一半是条件后一半是空文本无条件。forward 完之后输出 split 成两份按uncond guidance_scale * (cond - uncond)混合。这个设计的好处是 attention 的 batch 维度翻倍GPU 利用率更高坏处是显存占用也翻倍bucket 的粒度要相应调整。我一开始没注意这点batch4 的请求实际跑的是 batch8 的图显存直接爆了。4.3 捕获时的 stream 管理CUDA Graph 捕获必须在一个非默认 stream 上进行。PyTorch 提供了torch.cuda.graph(g)上下文管理器内部会自动处理 stream 切换。但如果你在捕获过程中调用了其他库比如自定义的 CUDA kernel要确保它们也在同一个 stream 上执行否则捕获不到。捕获完成后replay 也要在同一个 stream 上或者用torch.cuda.graphs.replay_graph(g)让它自动处理。我遇到过一个问题捕获时用的是 stream Areplay 时在默认 stream 上调用结果图执行了但结果没同步到默认 stream读出来全是零。后来统一用torch.cuda.graph上下文管理器就没再出过这个问题。g torch.cuda.CUDAGraph() with torch.cuda.graph(g): static_output model(static_latents, static_timestep, static_encoder_hidden)捕获完成后g.replay()就会执行整段计算。注意捕获期间不能有 CPU-GPU 同步所以model内部如果有.item()或者.cpu()调用需要提前改掉或者用 mask 替代。4.4 replay 时的数据搬运与同步每次去噪 step 开始前把新的 latents、timestep、encoder_hidden_states copy 到静态 buffer然后g.replay()最后从static_output读结果。读的时候要注意同步g.replay()是异步的如果你紧接着读static_output可能读到的是上一次的结果。需要torch.cuda.synchronize()或者用 event 来同步。但torch.cuda.synchronize()会阻塞整个 GPU影响吞吐。更好的做法是用torch.cuda.current_stream().record_event()记录一个 event然后在需要读结果的时候event.synchronize()。这样只阻塞当前 stream其他 stream 还能继续跑。static_latents.copy_(new_latents) static_timestep.copy_(new_timestep) static_encoder_hidden.copy_(new_encoder_hidden) g.replay() torch.cuda.current_stream().synchronize() result static_output.clone()clone()不能省因为下一次 replay 会覆盖static_output。如果你直接把static_output传给下一步下一步还没读完就被覆盖了结果就乱了。5. 调试实录那些让我熬夜的报错和它们的根因5.1 “operation not permitted during capture”最常见的捕获失败这个报错几乎每个搞 CUDA Graph 的人都见过。原因是捕获期间执行了不允许的操作比如显存分配、同步、事件查询等。PyTorch 的 caching allocator 在捕获期间会切换到 graph pool但如果你在捕获前没有正确设置 pool或者模型内部有动态分配就会触发这个错误。排查方法在捕获前加torch.cuda.synchronize()确保所有之前的操作都完成了。然后在捕获的代码块里逐段注释定位是哪一行触发的。常见元凶包括torch.empty()或torch.zeros()在捕获块内调用.item()或.cpu()调用print()打印 tensor自定义 kernel 里的cudaMalloc某些库的 lazy 初始化比如 flash-attn 第一次调用我遇到的最隐蔽的一次是 flash-attn 的 backward 在捕获时触发了 autotuning虽然我只跑 forward但 flash-attn 内部有个 flag 判断是否需要重新编译 kernel那个判断里有个同步操作。解决办法是在捕获前先跑一遍 flash-attn 的 forward让它完成初始化。5.2 结果不对但也不报错静态 buffer 没清干净有一次捕获出来的图 replay 之后图像全是噪声但没有任何报错。排查了半天发现是静态 buffer 在捕获前没有清零里面残留了上一次 eager 模式的数据。CUDA Graph 捕获的是计算过程不是数据所以 buffer 里的初始值会影响结果。解决办法很简单捕获前把静态 buffer 全部zero_()。还有一种情况是 buffer 的 dtype 不匹配。比如static_timestep是 long但模型内部期望 floatPyTorch 会自动转换但转换发生在捕获时replay 时不会重新转换导致 timestep 值不对。这种问题不会报错但结果会偏。建议在捕获前把所有 buffer 的 dtype 和模型期望的 dtype 对齐。5.3 多 bucket 切换时的显存泄漏前面说了我用 1、2、4、8 四个 bucket。跑了一段时间后发现显存缓慢增长最后 OOM。用torch.cuda.memory_summary()看发现 graph pool 的数量在增加。原因是每次切换 bucket 时如果新的 bucket 还没捕获会触发一次新的捕获而旧的 graph 对象没有被释放。解决办法是预先捕获所有 bucket 的图运行时只做选择不做动态捕获。SGLang 的 scheduler 里有个graph_pool字典key 是 bucket sizevalue 是对应的 CUDAGraph 对象。初始化时把所有 bucket 都捕获好之后只读不写。graph_pool {} for bs in [1, 2, 4, 8]: graph_pool[bs] capture_graph(model, bs)捕获所有 bucket 的代价是启动时间变长、显存占用变高但运行时的稳定性好很多。如果你的显存实在紧张可以只捕获常用的 bucket不常用的走 eager 模式兜底。5.4 和 KV cache 的显存争抢扩散模型如果用了 attention 的 KV cache比如某些 DiT 变体在长序列上做缓存CUDA Graph 的静态 buffer 会和 KV cache 抢显存。KV cache 是动态增长的而 graph pool 是固定的两者放在一起容易碎片化。我的做法是把 graph pool 和 KV cache 分到不同的显存区域用torch.cuda.memory._set_allocator_settings或者直接限制 graph pool 的上限。SGLang 内部用的是torch.cuda.graph_pool_handle()来管理确保 graph pool 独立于主 allocator。如果你自己实现记得在捕获前调用torch.cuda.graph_pool_handle()获取一个独立的 pool handle。6. 性能实测BCG 到底带来了多少提升6.1 吞吐量与延迟的对比数据我在 A100 80G 上跑了一个 DiT 变体约 2.6B 参数输入分辨率 1024x1024采样步数 20对比了 eager 模式、torch.compile 模式和 BCG 模式的性能。结果如下模式batch1 延迟batch8 吞吐显存峰值eager1.8s4.2 it/s28GBtorch.compile1.5s5.1 it/s30GBBCG (bucket1,8)1.2s6.8 it/s32GBBCG (bucket1,2,4,8)1.1s7.3 it/s34GB可以看到BCG 在 batch8 时吞吐提升了约 74%延迟降低了约 39%。显存峰值增加了 4-6GB主要是 graph pool 的开销。这个 trade-off 在显存充足的情况下是值得的。6.2 什么情况下 BCG 反而更慢BCG 不是万能的。如果 batch size 很小比如 1而且模型本身很轻量比如小的 UNetCUDA Graph 的 replay 开销可能比 eager 模式还大。因为 replay 本身也有固定开销大概几十微秒如果整个 step 的计算时间只有几毫秒这个开销占比就很高了。另外如果采样步数很少比如 4 步捕获图的收益摊薄到每一步上也不明显。BCG 最适合的场景是batch size 较大4、采样步数较多20、模型计算量中等偏上。这种场景下 kernel launch 开销占比高图化的收益最大。还有一个隐藏成本捕获图的时间。每个 bucket 捕获一次大概需要几秒到十几秒取决于模型大小。如果你只是跑一次推理就退出捕获的时间比省下来的还多。BCG 适合长期运行的服务不适合一次性脚本。6.3 和其他优化手段的叠加效果BCG 可以和量化、flash-attn、序列并行等手段叠加。我试过 BCG FP8 量化吞吐又提升了约 30%但精度损失需要评估。BCG flash-attn 是标配flash-attn 本身就能减少 attention 的显存和计算量图化之后 kernel launch 更少。需要注意的是量化后的模型如果用了自定义 kernel要确保这些 kernel 支持 CUDA Graph 捕获。有些量化库的 kernel 内部有动态分配捕获时会失败。建议在集成前先单独测试量化 kernel 的图捕获兼容性。7. 几个我踩过但文档里不会写的坑第一个坑是timestep embedding 的 broadcast。DiT 里 timestep 通常要 broadcast 到 batch 维度再和 latent 相加。在 eager 模式下broadcast 是动态的shape 不匹配会自动扩展。但在 CUDA Graph 里所有 shape 都是固定的如果你捕获时 timestep 是[1]replay 时传[8]图不会报错但 broadcast 的结果只有第一个 batch 是对的。解决办法是捕获时就用[batch_size]的 timestepreplay 时也传同样 shape。第二个坑是随机数生成。扩散模型的初始噪声是随机生成的如果你在捕获块内调用torch.randn()捕获出来的图每次 replay 都会生成相同的噪声因为随机种子被固定了。正确的做法是在图外生成噪声然后 copy 到静态 buffer 里。SGLang 的 scheduler 就是这么做的噪声生成和去噪循环是分开的。第三个坑是多卡场景下的 NCCL 图捕获。如果你用 tensor parallel 或 pipeline parallel图内会有 NCCL 通信。NCCL 的图捕获需要设置NCCL_GRAPH_MIXING_SUPPORT1环境变量否则捕获会 hang。而且捕获时所有 rank 必须同时进入捕获状态否则会死锁。这个在 SGLang 的分布式启动脚本里有处理但如果你自己写启动逻辑很容易漏掉。第四个坑是显存池的释放。程序退出时graph pool 如果不手动释放可能会导致进程 hang 住。PyTorch 的CUDAGraph对象在析构时会尝试释放 pool但如果此时 CUDA context 已经销毁就会报错。建议在退出前显式调用del graph_pool并torch.cuda.synchronize()。8. 后续可以继续折腾的方向BCG 这套机制目前在我这边已经稳定跑了几个月但还有几个方向值得继续挖。一个是动态 bucket 的自动选择现在是我手动设的 1、2、4、8如果能根据请求分布自动调整 bucket 边界应该能进一步减少 padding 浪费。另一个是图内量化把量化算子也捕获进图里减少量化-反量化的来回搬运。还有一个比较有意思的方向是跨 step 的图复用。扩散模型的每个 step 计算图其实是一样的只是 timestep 不同。如果能把多个 step 合并成一个大图理论上能进一步减少 replay 次数。但问题是 step 之间有调度器的状态更新比如 scheduler 的 alpha、sigma 累积这些状态如果也要图化会比较复杂。目前我还没看到成熟的方案但值得试试。最后说一句CUDA Graph 的调试体验确实不友好报错信息经常很模糊。我的经验是先保证 eager 模式结果正确再逐步图化。不要一上来就把整个 pipeline 捕获进去先捕获一个小模块验证 replay 结果和 eager 一致再扩大范围。这样出问题的时候容易定位。另外多准备几个不同 batch size 和不同输入长度的测试用例图化最容易在边界条件上翻车。