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 的思路就是把这些固定形状的算子序列录制成一张图，之后每次执行只需要一次 launch，GPU 侧按图执行，CPU 侧几乎不参与。对于扩散模型这种"同一张图跑几十遍"的场景，简直是量身定做。
2.2 为什么是 Batch CUDA Graph 而不是单样本 Graph
这里有个关键选择：SGLang 的 BCG 做的是batch 级别的 CUDA Graph，而不是给单个样本录一张图然后循环。原因有两个。
第一，GPU 利用率。单样本推理时，DiT 的很多算子（尤其是 attention）在 batch=1 的情况下根本喂不饱 SM，算力浪费严重。把 batch 凑大之后，同样的 kernel launch 开销摊到更多样本上，单位样本的成本就下来了。BCG 的设计目标就是让一个 graph 覆盖整个 batch，而不是每个样本一张图。
第二，显存和 graph 数量的平衡。CUDA Graph 本身是要占显存的，每张图都要保留自己的中间 buffer。如果给每个样本单独录图，graph 数量随 batch 线性增长，显存直接爆炸。而一张 batch 图只占一份中间 buffer，显存开销是固定的。当然代价是 batch size 必须固定——你不能今天跑 batch=4 明天跑 batch=8 还用同一张图，得重新录。
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, device='cuda')，然后在录图时把这个 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 大小也不同。如果你用 batch=1 warmup 然后用 batch=8 录图，大概率会在录图时报显存分配错误。
3.3 graph pool 的匹配逻辑与回退机制
SGLang 的 graph pool 维护了一个从 batch size 到 graph 的映射。请求进来之后，框架会先看当前 batch 能不能精确匹配到某个已录的图。匹配上了就直接 replay，匹配不上就走回退逻辑。
回退逻辑有两种策略。一种是向上取整：比如请求 batch=6，已录的图有 batch=4 和 batch=8，那就用 batch=8 的图，多出来的两个位置填 padding。这种策略的好处是延迟稳定，坏处是浪费算力。另一种是直接回退到 eager：batch=6 就老老实实跑 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 越大，相对收益越高，但边际收益在递减，而显存占用是线性增长的。所以 batch=8 通常是个甜点。
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 BCG
4.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.0
- CUDA >= 12.1
- SGLang 主分支（diffusion 支持还在快速迭代，别用太老的 release）
- transformers >= 4.40（DiT 的权重加载依赖）
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, device='cuda', dtype=dtype) static_timestep = torch.empty(batch_size, device='cuda', dtype=torch.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对于典型的 DiT（hidden dim 1152，24 层）大概在 200MB 到 400MB 之间，取决于你有没有用 flash attention。overhead是 graph 元数据和显存池的固定开销，大概 500MB 到 1GB。
举个例子，A100 80G 上，模型权重占 20G，per_sample_activation 按 300MB 算，overhead 按 1G 算，那么：
- batch=4：20 + 4*0.3 + 1 = 22.2G，很宽裕
- batch=8：20 + 8*0.3 + 1 = 23.4G，还是宽裕
- batch=16：20 + 16*0.3 + 1 = 25.8G，开始有压力
但这是单张图的显存。如果你要录多张图（batch=1、2、4、8 各一张），显存要累加。所以实际部署的时候，graph pool 的大小要控制好，别贪多。
提示：录图之前先用
torch.cuda.memory_summary()看一下当前显存占用，心里有个底。录完图之后再打一次，两次的差值就是这张图的实际开销。
4.4 实测数据：BCG 到底快了多少
我在 A100 上跑了一组对比，模型是一个 24 层、hidden dim 1152 的 DiT，输出 512x512，去噪步数 28。对比 eager 模式和 BCG 模式的端到端延迟。
| batch size | eager 延迟 (ms) | BCG 延迟 (ms) | 加速比 |
|---|---|---|---|
| 1 | 210 | 148 | 1.42x |
| 2 | 380 | 245 | 1.55x |
| 4 | 720 | 430 | 1.67x |
| 8 | 1420 | 810 | 1.75x |
可以看到 batch 越大加速比越高，因为 kernel launch 开销被摊薄了。batch=8 的时候接近 1.75x，这个收益在线上服务里相当可观。不过要注意，这是纯推理延迟，不含网络传输和前后处理。实际端到端收益会小一些，大概在 1.3x 到 1.5x 之间。
另外我测了一下单步延迟的分布。eager 模式下单步延迟抖动比较大，标准差能到 0.8ms；BCG 模式下标准差降到 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_blocking=False)，显式关掉异步拷贝。
5.3 多卡场景下的 graph 同步问题
如果你用多卡跑 BCG，会碰到一个额外的问题：每张卡都要录自己的图，而且录图的时机要同步。如果卡 0 录完了卡 1 还在录，replay 的时候就会 hang 住。
SGLang 的处理方式是让所有 rank 在录图前做一次 barrier，确保大家同时开始录。这个 barrier 是必须的，别想着省。我一开始为了图快把 barrier 去掉了，结果跑了几十次之后偶发 hang，排查了两天才定位到。
5.4 常见问题速查表
| 现象 | 可能原因 | 解决办法 |
|---|---|---|
| replay 报 operation not permitted | 录图期间有新显存分配 | 把一次性操作挪到 warmup |
| 输出全噪声 | 静态输入没更新 | 拷贝前加 contiguous |
| replay 后结果被覆盖 | 没 clone 输出 | 返回前 clone |
| 多卡 hang | 录图时机不同步 | 加 barrier |
| 显存 OOM | graph pool 太大 | 减少录图的 batch size 数量 |
| 延迟没降 | batch 太小 | 增大 batch 或检查是否真的走了 graph |
5.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 匹配逻辑。
