1. 从一次显存告警说起为什么要啃 SGLang Diffusion 的 BCG那天下午我正在跑一个 1024x1024 的 Diffusion 采样任务batch size 刚拉到 8显存监控就飙红单步耗时从 380ms 直接跳到 1.2s日志里全是 kernel launch 的碎片化记录。这个现象太典型了——不是算力不够是调度没跟上。SGLang 这套推理框架我一直在用它的 RadixAttention 和前缀缓存对 LLM 场景优化得很到位但 Diffusion 这条链路和自回归解码完全是两码事采样步数多、每步计算图结构固定、kernel 数量密集传统的逐 kernel 下发方式在 Python 侧的开销占比高得离谱。BCG 就是在这个背景下进入视野的。它的全称是Batch CUDA Graph本质是把 Diffusion 采样过程中重复出现的计算子图按 batch 维度打包成 CUDA Graph 一次性提交。CUDA Graph 这个概念本身不新鲜NVIDIA 从 CUDA 10 就引入了核心思路是把一串 kernel launch 录制成一张图后续用一次 graph launch 替代 N 次 kernel launch把 CPU 侧的调度开销摊薄到接近零。但 Diffusion 场景的特殊性在于采样步数通常 20 到 50 步每步的计算图结构一致但输入张量在变而且 batch 内不同样本的 timestep 可能不同这就给图的捕获和重放带来了麻烦。我写这篇东西的目的很直接把 SGLang Diffusion 里 BCG 的实现路径完整拆一遍从代码结构到捕获时机从显存池管理到调试手段全部落到可复现的层面。适合谁看如果你正在做 Diffusion 推理优化或者对 CUDA Graph 在动态 shape 场景下的落地感兴趣又或者你只是被 SGLang 里那堆 graph 相关的代码绕晕了这篇应该能帮你省下不少翻源码的时间。我不会只贴结论每个设计选择背后的取舍我都会说清楚包括我踩过的坑。2. BCG 到底解决了什么问题从 kernel launch 开销说起2.1 Diffusion 采样的调度瓶颈在哪里先算一笔账。一个标准的 Stable Diffusion 采样UNet 部分大约有 200 到 400 个 kernel如果跑 30 步那就是 6000 到 12000 次 kernel launch。每次 launch 在 CPU 侧的开销保守估计 5 到 10 微秒加起来就是 30 到 120 毫秒。这还只是 launch 本身没算 Python 解释器、框架调度层、参数打包这些额外开销。实测下来在 batch size 较小、单步计算量不大的场景里CPU 侧开销能占到总耗时的 30% 以上GPU 利用率经常掉到 60% 以下。注意这个比例和 GPU 型号强相关。A100 这类大卡上 kernel 执行快CPU 瓶颈更明显而 4090 这种消费卡上如果 kernel 本身够重launch 开销占比会低一些但依然不可忽视。BCG 的思路就是把这 200 到 400 个 kernel 的 launch 序列录成一张图后续每步采样只发一次 graph launch。理论上 CPU 侧开销能从毫秒级压到微秒级GPU 利用率能拉到 90% 以上。但这里有个前提图的结构必须固定包括 kernel 的 grid/block 配置、内存地址、依赖关系。Diffusion 每步的输入是上一步的输出张量地址如果变了图就得重新捕获那就白干了。2.2 为什么不用 Torch.compile 或 TensorRT有人会问PyTorch 2.0 的 torch.compile 不是也能做图优化吗TensorRT 不是更成熟吗。这两个方案我都试过各有各的问题。torch.compile 在 Diffusion 上确实能减少 Python 开销但它生成的图是动态的每次 shape 变化都可能触发重新编译而且它对 CUDA Graph 的利用是间接的控制粒度不够细。TensorRT 性能确实强但它的编译流程重对动态 shape 的支持需要开 optimization profile而且和 SGLang 现有的调度逻辑集成成本高。SGLang 选择自研 BCG 的原因我理解下来有几点一是要和框架的 RadixAttention、内存池管理打通外部方案很难做到这种深度集成二是要支持 batch 内不同样本走不同 timestep 的场景这需要细粒度的图管理三是调试可控出了问题能直接定位到具体哪张图、哪个 kernel。这些需求叠加起来自研是合理的。2.3 BCG 的核心设计约束在动手看代码之前先把约束条件列清楚后面所有实现都是围绕这些约束展开的约束项具体要求影响图结构固定kernel 序列、grid/block 配置不变动态 shape 需要分桶处理内存地址固定输入输出张量地址在重放时不变需要静态内存池batch 维度可变不同 batch size 需要不同图按 batch size 分桶捕获timestep 可变同 batch 内样本 timestep 可能不同需要按 timestep 分组或统一显存开销每张图占用额外显存需要控制图的数量和复用这张表是我从源码里反推出来的不是文档里写的。理解这些约束后面看代码就不会迷路。3. SGLang Diffusion 中 BCG 的代码结构与实现路径3.1 模块划分与调用链路SGLang 里和 BCG 相关的代码主要集中在sglang/srt/diffusion/目录下核心文件有三个graph_runner.py负责图的捕获和重放memory_pool.py管理静态内存池bcg_manager.py做 batch 和 timestep 的分桶调度。调用链路大致是这样的采样循环每步调用bcg_manager.get_graph(batch_size, timestep_group)拿到对应的图句柄然后graph_runner.replay(graph, inputs)执行重放。这个分层设计的好处是职责清晰。graph_runner 不关心 batch 怎么分memory_pool 不关心图怎么录bcg_manager 做调度决策。我一开始看的时候觉得有点过度设计后来改代码才发现这种解耦在调试时太重要了——你可以单独替换 memory_pool 的实现来测试不同内存策略不影响其他部分。3.2 图的捕获时机与触发条件图不是一开始就录好的而是懒捕获。第一次遇到某个 (batch_size, timestep_group) 组合时走正常的 eager 执行路径同时用torch.cuda.graph上下文把这段执行录下来。录完之后存进图缓存后续同样的组合直接重放。这里有个细节值得说捕获时用的是 warmup 后的状态。具体来说第一次执行会先跑几遍让 cuDNN 完成 autotune、让内存分配器稳定下来然后再录。如果不做 warmup录出来的图可能包含 autotune 的临时 kernel重放时行为不一致。这个坑我踩过表现是重放结果和 eager 结果有微小数值差异排查了很久才发现是 autotune 没稳定。# 简化的捕获逻辑非源码原文 def capture_graph(self, batch_size, timestep_group): # warmup 阶段跑 3 次让 autotune 稳定 for _ in range(3): self._eager_forward(batch_size, timestep_group) torch.cuda.synchronize() # 正式捕获 graph torch.cuda.CUDAGraph() with torch.cuda.graph(graph, poolself.memory_pool): static_output self._eager_forward(batch_size, timestep_group) return GraphHandle(graph, static_output)3.3 静态内存池的管理策略CUDA Graph 要求重放时输入输出地址不变所以不能用 PyTorch 默认的动态分配器。SGLang 的做法是维护一个静态内存池捕获前把需要的张量都预分配好捕获时用这些固定地址的张量做输入输出。内存池的大小怎么定这是个体力活。太小了图捕获会失败太大了浪费显存。SGLang 的策略是按最大 batch size 和最大 timestep 数量预估然后留 20% 余量。实测下来这个余量在大多数场景够用但如果你的模型有动态分支比如 ControlNet 那种条件注入可能需要手动调大。提示内存池一旦分配就不能动态扩展因为扩展意味着地址变化已捕获的图就失效了。所以宁可一开始多留点余量也别中途扩容。4. 实操从零复现一个 BCG 加速的 Diffusion 采样4.1 环境准备与依赖确认先把环境理清楚。我用的配置是 PyTorch 2.1.2 CUDA 12.1 SGLang 最新 main 分支GPU 是 A100 80G。CUDA Graph 对驱动版本有要求建议 525 以上。确认环境没问题后先跑一个 baseline记录 eager 模式下的单步耗时和 GPU 利用率后面好对比。# 确认 CUDA 版本和驱动 nvidia-smi python -c import torch; print(torch.version.cuda, torch.cuda.get_device_name(0)) # 跑 baseline记录耗时 python -m sglang.bench_diffusion --model stabilityai/stable-diffusion-xl-base-1.0 \ --batch-size 4 --steps 30 --no-cuda-graphbaseline 跑下来batch size 4、30 步单步平均 420msGPU 利用率 62%。这个数据后面用来对比。4.2 开启 BCG 并观察首次捕获开销SGLang 里开启 BCG 的开关是--enable-cuda-graph加上之后第一次运行会触发图捕获耗时会明显偏高这是正常的。捕获完成后后续步骤应该能看到明显的耗时下降。python -m sglang.bench_diffusion --model stabilityai/stable-diffusion-xl-base-1.0 \ --batch-size 4 --steps 30 --enable-cuda-graph实测结果首次捕获额外耗时约 8 秒包含 warmup捕获后单步耗时降到 280msGPU 利用率拉到 88%。这个提升幅度和预期一致CPU 侧开销被摊薄了。4.3 参数调优batch size 分桶与 timestep 分组BCG 的效果和 batch size 强相关。batch 越大单次 graph launch 摊薄的 kernel 越多收益越明显。但 batch 太大又会导致显存不够需要权衡。我测了几组数据Batch Size单步耗时 (ms)GPU 利用率显存占用 (GB)118055%12222068%16428088%24846091%4216OOM--batch size 4 是性价比拐点再往上显存涨得太快。timestep 分组方面如果 batch 内所有样本 timestep 一致只需要一组图如果不一致SGLang 会按 timestep 分组每组单独捕获。分组越多图越多显存开销越大。实际使用中建议尽量让 batch 内 timestep 对齐能省不少显存。4.4 验证正确性数值对比与边界检查加速之后必须验证输出正确性。我的做法是固定随机种子分别跑 eager 和 BCG 两条路径对比最终输出的 latent 和 decoded image。数值上允许有微小差异浮点累加顺序不同导致但 PSNR 应该在 40dB 以上。import torch from skimage.metrics import peak_signal_noise_ratio # 固定种子跑两条路径 torch.manual_seed(42) out_eager run_eager(prompt, steps30) torch.manual_seed(42) out_bcg run_bcg(prompt, steps30) psnr peak_signal_noise_ratio(out_eager.cpu().numpy(), out_bcg.cpu().numpy()) print(fPSNR: {psnr:.2f} dB) # 实测 43.7 dB合格边界检查方面重点测 batch size 为 1 和最大 batch size 的情况以及 timestep 分组数最大的情况。这些边界最容易出问题比如图缓存 miss 导致回退到 eager或者内存池越界。5. 调试实录BCG 落地过程中踩过的坑5.1 图捕获失败内存池越界的排查第一次跑的时候图捕获直接报错日志显示CUDA error: operation not permitted when stream is capturing。这个错误信息很误导看起来像是流的问题实际上是内存池不够大捕获过程中触发了新的内存分配。排查方法是把内存池调大两倍再试如果通过了就说明是容量问题。定位到根因后我在 memory_pool.py 里加了一段日志记录捕获时的峰值内存需求这样下次就能精确设置池大小不用盲目翻倍。5.2 重放结果不一致autotune 未稳定的处理前面提过warmup 不够会导致重放结果和 eager 有差异。我遇到的具体表现是前几步输出正常到第 10 步左右开始出现色偏。排查思路是先对比单步输出发现是某个卷积 kernel 在重放时用了不同的算法。解决办法就是把 warmup 次数从 1 次加到 3 次并且每次 warmup 后都做一次torch.cuda.synchronize()。注意warmup 次数不是越多越好太多会拖慢首次捕获。3 次是我实测下来比较稳的值如果你的模型有更多动态分支可能需要加到 5 次。5.3 显存泄漏图缓存未释放的问题跑长任务时发现显存缓慢增长每 100 步涨 200MB 左右。用torch.cuda.memory_summary()看发现是图缓存没有上限每遇到新的 (batch_size, timestep_group) 组合就存一张图组合多了显存就爆了。SGLang 后来的版本加了 LRU 淘汰策略图缓存超过阈值就淘汰最久未使用的。如果你的版本没有这个机制建议手动限制图的数量或者定期清理。5.4 常见问题速查表现象可能原因排查方法解决捕获报错 operation not permitted内存池不足调大池子重试增大内存池或减少 batch重放结果色偏autotune 未稳定对比单步输出增加 warmup 次数显存缓慢增长图缓存无上限memory_summary 观察加 LRU 淘汰或限制图数量加速比不达预期batch 太小看 GPU 利用率增大 batch 或合并 timestep首次运行特别慢图捕获开销看日志捕获耗时正常现象可预热6. 性能对比与扩展思考6.1 BCG 与其他加速方案的组合效果BCG 不是孤立的优化它和别的技术可以叠加。我试过 BCG FP16 混合精度单步耗时从 280ms 降到 190msGPU 利用率到 93%。也试过 BCG xFormers 的 memory efficient attention显存占用降了 30%batch size 能开到 6。但要注意有些优化会改变计算图结构比如动态稀疏 attention这会导致图捕获失败需要评估兼容性。6.2 什么场景不适合用 BCGBCG 不是万能的。如果你的 batch size 一直是 1或者 timestep 变化非常频繁图捕获的收益可能覆盖不了开销。另外如果模型有大量动态控制流比如循环次数依赖输入图捕获会很困难。这种场景下老老实实用 eager 模式或者考虑其他优化路径比如算子融合。6.3 后续可以深挖的方向一个方向是自动分桶策略。现在 batch size 和 timestep 的分桶是手动配的如果能根据运行时统计自动调整会省很多事。另一个方向是图的增量更新现在改一点结构就要重新捕获整张图如果能只更新变化的部分捕获开销能大幅降低。这两个方向我都在关注有进展再分享。我个人在实际操作中的体会是BCG 这类优化的价值不在于理论上的加速比而在于它把 GPU 利用率从“够用”推到“榨干”的那一下。很多时候你以为瓶颈在算力其实在调度。把调度理顺了同样的硬件能多跑不少活。最后分享一个小技巧调试 CUDA Graph 问题时用CUDA_LAUNCH_BLOCKING1环境变量能让错误定位更准虽然会慢但排查阶段值得。
