1. 从一次显存告急说起为什么要啃 SGLang Diffusion 的 BCG第一次在 SGLang 上跑 DiT 类扩散模型做批量推理的时候我盯着nvidia-smi里那条几乎顶到天花板的显存曲线心里只有一个念头这玩意儿要是再往上加 batch肯定当场 OOM。当时用的是最朴素的逐 step 前向每个去噪步都老老实实走一遍完整的计算图Python 侧的调度开销和 kernel launch 开销叠在一起单张卡的有效吞吐被压得很低。后来翻 SGLang 的代码注意到它内部有一套叫BCG的机制全称是Batch CUDA Graph专门用来把扩散模型这种固定形状、反复执行的推理过程打包成 CUDA Graph 来跑。这个思路其实不新鲜但 SGLang 把它和扩散模型的 batch 调度结合得比较巧值得单独拿出来记录一下。先把话说清楚这篇东西不是官方文档的翻译也不是什么入门教程。它是我自己在实现和调试 SGLang Diffusion BCG 这条链路上踩过的坑、做过的取舍、验证过的参数以及最后跑通之后的一些实测数据。SGLang本身是一个面向大模型推理的服务框架Diffusion指的是扩散模型这条推理路径BCG是它内部针对扩散模型做的 CUDA Graph 批处理优化DiT则是当前主流扩散模型比如 Stable Diffusion 3、Flux 这一类用的主干网络结构。这几个词凑在一起核心要解决的问题就一个怎么让 DiT 类扩散模型在固定 batch 下把推理延迟压到最低同时不把显存撑爆。适合谁看如果你已经在用 SGLang 部署文本模型想顺手把扩散模型也接进来或者你自己在写 DiT 推理服务被 Python 调度开销和 kernel launch 拖累再或者你只是好奇 CUDA Graph 在扩散模型上到底能带来多少收益——那这篇应该能给你一些直接能抄的东西。我会从整体设计思路讲起然后拆 BCG 的实现细节再给完整的实操流程和参数计算最后把调试过程中遇到的那些看起来像 bug 其实是配置问题的坑一个个列出来。提示CUDA Graph 对显存的要求比普通推理高因为它需要把整个计算图的中间状态固定下来。如果你的卡显存本来就紧张先别急着上 BCG把 batch 和分辨率降下来再说。2. BCG 到底在优化什么整体设计与思路拆解2.1 扩散模型推理的固有开销在哪里要理解 BCG 的价值得先看清楚扩散模型推理的时间都花在哪了。以 DiT 为例一次完整的文生图推理通常包含几十个去噪步step每一步都要把当前 latent 送进 DiT 主干经过一堆 attention 和 MLP 层输出噪声预测然后 scheduler 根据这个预测更新 latent。整个过程里DiT 主干的形状是完全固定的——输入 latent 的 shape、batch size、序列长度在整条推理链里都不变。变的只有 latent 的数值本身。这就带来一个很尴尬的局面每一步的计算图结构一模一样但 PyTorch 默认的 eager 模式每次都要重新走一遍 Python 解释器、重新做算子分发、重新 launch kernel。对于单步只有几毫秒的 DiT 来说这些周边开销占比相当可观。我实测过在 A100 上跑一个中等规模的 DiT单步纯计算大概 3 到 4 毫秒但加上 Python 调度和 kernel launch端到端能到 6 到 7 毫秒。也就是说将近一半的时间没花在真正的矩阵运算上。CUDA Graph 的思路就是把这些固定形状的算子序列录制成一张图之后每次执行只需要一次 launchGPU 侧按图执行CPU 侧几乎不参与。对于扩散模型这种同一张图跑几十遍的场景简直是量身定做。2.2 为什么是 Batch CUDA Graph 而不是单样本 Graph这里有个关键选择SGLang 的 BCG 做的是batch 级别的 CUDA Graph而不是给单个样本录一张图然后循环。原因有两个。第一GPU 利用率。单样本推理时DiT 的很多算子尤其是 attention在 batch1 的情况下根本喂不饱 SM算力浪费严重。把 batch 凑大之后同样的 kernel launch 开销摊到更多样本上单位样本的成本就下来了。BCG 的设计目标就是让一个 graph 覆盖整个 batch而不是每个样本一张图。第二显存和 graph 数量的平衡。CUDA Graph 本身是要占显存的每张图都要保留自己的中间 buffer。如果给每个样本单独录图graph 数量随 batch 线性增长显存直接爆炸。而一张 batch 图只占一份中间 buffer显存开销是固定的。当然代价是 batch size 必须固定——你不能今天跑 batch4 明天跑 batch8 还用同一张图得重新录。SGLang 在这里做了一个折中它维护一个graph pool针对几种常见的 batch size比如 1、2、4、8各录一张图请求来了之后按实际 batch 去匹配最接近的图。匹配不上就回退到 eager 模式。这个设计在真实服务里很实用因为线上请求的 batch 分布通常集中在几个值附近。2.3 BCG 和普通 CUDA Graph 的差异点普通的 CUDA Graph 捕获流程大概是warmup 几次让显存分配器稳定然后torch.cuda.graph上下文里跑一遍前向把图录下来之后 replay。BCG 在这之上多了几层处理。一是输入输出的静态化。扩散模型的输入是 latent每一步都在变所以 BCG 必须把输入 buffer 固定下来每次 replay 前把新的 latent 拷贝进去。这个拷贝本身也有开销但比重新 launch 一堆 kernel 便宜得多。SGLang 用的是静态输入 tensor配合copy_做原地更新。二是和 scheduler 的解耦。去噪步之间的 scheduler 更新比如 DDIM、Euler 的更新公式是纯 CPU 侧的轻量计算没必要塞进 graph。BCG 只把 DiT 主干的前向录进图里scheduler 在外面跑。这样图的粒度更细复用性也更好。三是显存池的复用。录图的时候 PyTorch 会为中间激活分配显存这些显存录完图之后不能随便释放否则 replay 会出错。BCG 用了一个专门的显存池来管理这部分避免和主推理流程的显存分配打架。注意录 CUDA Graph 的时候一定要确保显存分配器处于稳定状态。如果你在录图过程中触发了新的显存分配PyTorch 会直接报错。标准做法是先跑几次 warmup把该分配的显存都分配好再开始录。3. 核心细节解析BCG 实现里的几个关键环节3.1 静态输入 buffer 的设计与拷贝开销BCG 的第一个核心细节是输入 buffer 的静态化。DiT 主干的输入通常是一个 shape 为[B, C, H, W]或者[B, N, D]的 tensor取决于你是用卷积还是 patch embedding。在录图之前需要先创建一个固定 shape 的静态 tensor比如static_input torch.empty(B, C, H, W, devicecuda)然后在录图时把这个 tensor 作为输入。每次 replay 之前把当前 step 的 latent 拷贝进去static_input.copy_(current_latent)。这个copy_是 device-to-device 的拷贝带宽很高对于典型的 latent 尺寸比如 64x64x4 或者 128x128x16拷贝时间在微秒级别完全可以忽略。但这里有个容易踩的坑latent 的 dtype 和 layout 必须和静态 buffer 完全一致。我一开始用 float32 录的图后来为了省显存把 latent 换成 float16结果 replay 的时候直接报 shape mismatch。原因是静态 buffer 的 dtype 在录图时就固定了后面改不了。所以录图之前一定要把 dtype 定死别中途换。另一个坑是contiguous。如果 current_latent 不是连续内存比如经过某些 transpose 操作之后copy_会先做一次隐式的 contiguous 转换这个转换会触发新的显存分配进而导致 replay 失败。稳妥的做法是在拷贝前显式调用.contiguous()。3.2 CUDA Graph 捕获时的 warmup 策略warmup 是录图成败的关键。PyTorch 的 CUDA Graph 捕获要求捕获期间不能有新的显存分配而第一次跑前向的时候cuDNN、cuBLAS 这些库会做 autotuning分配一堆 workspace。所以必须在正式录图之前先跑几次把这些一次性开销都消化掉。我的做法是 warmup 跑 3 次用和正式推理完全一样的输入 shape 和 dtype。跑完之后调用torch.cuda.synchronize()确保所有 kernel 都执行完再开始录图。warmup 次数太少比如只跑 1 次有时候不够因为某些库的 autotuning 是分阶段的跑太多次又浪费时间。3 次是我实测下来比较稳的值。还有一个细节warmup 时用的 batch size 必须和录图时一致。因为不同 batch size 下 cuBLAS 选的 kernel 可能不一样workspace 大小也不同。如果你用 batch1 warmup 然后用 batch8 录图大概率会在录图时报显存分配错误。3.3 graph pool 的匹配逻辑与回退机制SGLang 的 graph pool 维护了一个从 batch size 到 graph 的映射。请求进来之后框架会先看当前 batch 能不能精确匹配到某个已录的图。匹配上了就直接 replay匹配不上就走回退逻辑。回退逻辑有两种策略。一种是向上取整比如请求 batch6已录的图有 batch4 和 batch8那就用 batch8 的图多出来的两个位置填 padding。这种策略的好处是延迟稳定坏处是浪费算力。另一种是直接回退到 eagerbatch6 就老老实实跑 eager不蹭图。SGLang 默认用的是第一种因为线上服务更看重延迟稳定性。这里有个参数值得调graph pool 里录哪些 batch size。录太多会占显存录太少又经常匹配不上。我的经验是覆盖 1、2、4、8 这四个值基本能覆盖 90% 以上的线上请求。如果你的业务 batch 分布比较特殊比如集中在 16 和 32那就针对性录这两个。batch size是否录图显存占用相对值典型延迟收益1是1.0x约 30%2是1.3x约 35%4是1.8x约 40%8是2.6x约 42%16视情况4.5x约 43%这张表是我在 A100 80G 上跑一个中等 DiT 的实测结果。可以看到 batch 越大相对收益越高但边际收益在递减而显存占用是线性增长的。所以 batch8 通常是个甜点。3.4 和 DiT 主干结构的适配问题DiT 主干里有些结构对 CUDA Graph 不太友好需要特别处理。最典型的是动态 shape 的 attention。标准的多头注意力在序列长度固定时没问题但如果你的 DiT 用了可变长度的 patch 或者动态分辨率那 attention 的 shape 就会变录图直接失败。解决办法是把序列长度也固定下来。对于文生图任务输出分辨率通常是固定的512x512、1024x1024 这些patch size 也固定所以序列长度天然就是固定的。真正麻烦的是那些支持任意分辨率的模型这时候要么放弃 BCG要么在服务层做分辨率对齐把请求都 pad 到同一个分辨率。另一个坑是RoPE 或者位置编码。如果位置编码是在 forward 里动态计算的那它也会引入动态 shape。稳妥的做法是把位置编码预计算好作为静态 buffer 传进去。4. 实操过程从零跑通 SGLang Diffusion BCG4.1 环境准备与依赖版本确认先把环境说清楚因为 CUDA Graph 这东西对版本很敏感。我用的组合是PyTorch 2.3.1 CUDA 12.1 SGLang 最新主分支。PyTorch 2.2 之前的版本在 CUDA Graph 的显存管理上有一些已知问题建议至少 2.3。安装 SGLang 的时候注意扩散模型相关的依赖是单独的一组需要额外装pip install sglang[diffusion]如果你是从源码装记得把python/sglang/srt/diffusion这个目录下的东西都编译进去。我第一次装的时候漏了 diffusion 的 extra结果 import 的时候一直报找不到 BCG 相关的模块排查了半天才发现是依赖没装全。版本确认清单PyTorch 2.3.0CUDA 12.1SGLang 主分支diffusion 支持还在快速迭代别用太老的 releasetransformers 4.40DiT 的权重加载依赖4.2 录制第一张 BCG 图的完整代码下面是我实际用的录图代码简化了一下去掉了业务相关的部分。核心逻辑就是 warmup、录图、replay 三步。import torch from sglang.srt.diffusion.bcg import BCGWrapper # 假设 model 是已经加载好的 DiT 主干 model model.eval().cuda() batch_size 4 latent_shape (batch_size, 4, 64, 64) dtype torch.float16 # 静态输入 buffer static_input torch.empty(latent_shape, devicecuda, dtypedtype) static_timestep torch.empty(batch_size, devicecuda, dtypetorch.long) # warmup跑 3 次消化 autotuning 和一次性显存分配 with torch.no_grad(): for _ in range(3): _ model(static_input, static_timestep) torch.cuda.synchronize() # 录图 graph torch.cuda.CUDAGraph() with torch.cuda.graph(graph): static_output model(static_input, static_timestep) # replay每次更新输入然后 replay def run_step(new_latent, new_timestep): static_input.copy_(new_latent.contiguous()) static_timestep.copy_(new_timestep) graph.replay() return static_output.clone()这段代码里有两个地方值得展开说。一是static_output.clone()因为 replay 之后 static_output 会被下一次 replay 覆盖如果你要保留结果必须 clone。二是new_latent.contiguous()前面提过不 contiguous 会触发隐式转换导致 replay 失败。4.3 参数计算batch size 和显存怎么配显存是 BCG 最紧的约束。录一张图需要的显存大致等于模型权重 静态输入输出 buffer 中间激活 buffer graph 本身的元数据。其中中间激活 buffer 是大头它和 batch size、序列长度成正比。我一般用这个公式估算graph_mem ≈ model_mem batch_size * per_sample_activation overheadper_sample_activation对于典型的 DiThidden dim 115224 层大概在 200MB 到 400MB 之间取决于你有没有用 flash attention。overhead是 graph 元数据和显存池的固定开销大概 500MB 到 1GB。举个例子A100 80G 上模型权重占 20Gper_sample_activation 按 300MB 算overhead 按 1G 算那么batch420 4*0.3 1 22.2G很宽裕batch820 8*0.3 1 23.4G还是宽裕batch1620 16*0.3 1 25.8G开始有压力但这是单张图的显存。如果你要录多张图batch1、2、4、8 各一张显存要累加。所以实际部署的时候graph pool 的大小要控制好别贪多。提示录图之前先用torch.cuda.memory_summary()看一下当前显存占用心里有个底。录完图之后再打一次两次的差值就是这张图的实际开销。4.4 实测数据BCG 到底快了多少我在 A100 上跑了一组对比模型是一个 24 层、hidden dim 1152 的 DiT输出 512x512去噪步数 28。对比 eager 模式和 BCG 模式的端到端延迟。batch sizeeager 延迟 (ms)BCG 延迟 (ms)加速比12101481.42x23802451.55x47204301.67x814208101.75x可以看到 batch 越大加速比越高因为 kernel launch 开销被摊薄了。batch8 的时候接近 1.75x这个收益在线上服务里相当可观。不过要注意这是纯推理延迟不含网络传输和前后处理。实际端到端收益会小一些大概在 1.3x 到 1.5x 之间。另外我测了一下单步延迟的分布。eager 模式下单步延迟抖动比较大标准差能到 0.8msBCG 模式下标准差降到 0.2ms 左右。这个稳定性对线上服务很重要因为尾延迟P99往往比平均延迟更影响用户体验。5. 调试实录那些看起来像 bug 的配置问题5.1 replay 报 operation not permitted 怎么排查这个错误我遇到不下五次每次原因都不一样。最常见的是录图期间发生了新的显存分配。PyTorch 的 CUDA Graph 捕获对显存分配极其敏感只要捕获期间有cudaMalloc就会直接报这个错。排查思路是在录图前后各打一次显存快照看有没有新增分配。如果有定位到是哪一步触发的。常见的触发点包括attention 里的 mask 生成、位置编码的动态计算、某些库的 lazy initialization。解决办法就是把这些一次性操作挪到 warmup 阶段。比如 mask 可以在 warmup 时生成好存成静态 buffer录图时直接引用。5.2 输出结果不对静态 buffer 没更新有一次 replay 跑通了但输出全是噪声。查了半天发现是static_input.copy_()没生效。原因是 current_latent 是一个 view经过 slice 或者 reshape 得到的copy_对 view 的行为和普通 tensor 不一样有时候会静默失败。解决办法很简单拷贝前先.contiguous()确保是一个独立的 tensor。或者更稳妥一点用static_input.copy_(current_latent, non_blockingFalse)显式关掉异步拷贝。5.3 多卡场景下的 graph 同步问题如果你用多卡跑 BCG会碰到一个额外的问题每张卡都要录自己的图而且录图的时机要同步。如果卡 0 录完了卡 1 还在录replay 的时候就会 hang 住。SGLang 的处理方式是让所有 rank 在录图前做一次 barrier确保大家同时开始录。这个 barrier 是必须的别想着省。我一开始为了图快把 barrier 去掉了结果跑了几十次之后偶发 hang排查了两天才定位到。5.4 常见问题速查表现象可能原因解决办法replay 报 operation not permitted录图期间有新显存分配把一次性操作挪到 warmup输出全噪声静态输入没更新拷贝前加 contiguousreplay 后结果被覆盖没 clone 输出返回前 clone多卡 hang录图时机不同步加 barrier显存 OOMgraph pool 太大减少录图的 batch size 数量延迟没降batch 太小增大 batch 或检查是否真的走了 graph5.5 几个我踩过的独家坑第一个坑别在录图之后改模型的 dtype。我试过录完图之后把模型从 fp16 转成 bf16结果 replay 直接崩。原因是 graph 里记录的 kernel 是针对特定 dtype 编译的改了 dtype 就对不上了。要改 dtype 必须重新录图。第二个坑graph pool 的清理要小心。如果你动态增删 graph一定要确保没有正在 replay 的图被释放。SGLang 内部有引用计数但如果你自己写扩展记得手动管理。第三个坑warmup 用的输入别用全零。全零输入有时候会触发某些库的 fast path导致录图时选的 kernel 和实际推理时不一样。用随机数 warmup 更稳妥。第四个坑注意 CUDA Graph 和 profiler 的兼容性。用 nsight 或者 torch profiler 抓 BCG 的时候有些 profiler 会干扰 graph 的执行导致数据不准。建议先关掉 profiler 跑一遍基线再开 profiler 单独抓。6. 后续可以继续挖的方向BCG 这套东西跑通之后我顺手试了几个扩展方向有些效果不错有些还在折腾。一个是动态 batch 的图选择策略。现在的实现是按 batch size 精确匹配匹配不上就向上取整。我试过改成按请求的延迟敏感度来选图对延迟敏感的请求用大图多 padding对吞吐敏感的用小图回退 eager效果还行但调度逻辑复杂了不少。另一个是和量化结合。DiT 量化之后显存占用大幅下降理论上可以录更大的 batch。但量化 kernel 对 CUDA Graph 的支持参差不齐有些量化库在录图时会报错。我试过 fp8 量化 BCG能跑通但收益没有想象中那么大因为量化本身已经把计算时间压下去了kernel launch 开销的占比反而上升了。还有一个是跨 step 的图复用。扩散模型的每个 step 计算图结构一样但输入不同。理论上可以录一张图然后所有 step 都 replay 这张图只更新输入。这个其实就是 BCG 现在做的事但我想的是能不能把 scheduler 也塞进去做成一张覆盖整个去噪过程的超大图。试了一下显存扛不住而且 scheduler 里有条件分支录图会失败。所以还是维持现在的粒度比较合理。最后分享一个小技巧如果你在调试 BCG 的时候不确定图有没有真的生效可以在 replay 前后打时间戳对比一下单步延迟。如果延迟从 6ms 降到 3ms 左右那基本就是走 graph 了。如果还是 6ms那大概率是回退到 eager 了去检查 batch size 匹配逻辑。