SGLang Diffusion BCG：CUDA Graph加速扩散模型推理实战

1. 从标题拆解SGLang Diffusion BCG 到底在解决什么问题第一次看到“SGLang Diffusion BCG”这个组合很多人会愣一下SGLang 不是做大模型推理调度的吗怎么跟 Diffusion 扯上关系了BCG 又是什么缩写我一开始也带着同样的疑问去翻代码和跑实验折腾了差不多两周才算把这条链路摸清楚。这篇就把我踩过的坑、验证过的方案、以及最终跑通的配置完整记录下来给同样在做 Diffusion 推理加速的朋友一个可复现的参考。先把三个关键词拆开说清楚。SGLang是一套面向大模型推理的服务框架核心卖点是 RadixAttention 前缀缓存和高效的调度器擅长处理高并发、共享前缀的请求。Diffusion这里指的是扩散模型推理典型代表就是 Stable Diffusion 系列它的推理过程和自回归语言模型完全不同——它是多步迭代去噪每一步都要跑一遍 UNet步数通常 20 到 50 步所以单次请求的计算量远大于一次 LLM 前向。BCG在我的理解里是Bucket CUDA Graph这套组合的简称也有人把它理解成 Batch-level CUDA Graph核心目的就一个把 Diffusion 推理里那些重复的、形状固定的小算子用 CUDA Graph 一次性捕获下来减少 kernel launch 开销和 CPU 调度开销。那为什么要把这三者凑到一起因为 Diffusion 推理有个很尴尬的特点单步 UNet 的计算其实不算特别大但步数多每一步之间还有大量 Python 层的调度、张量搬运、采样器逻辑。在 LLM 场景里 SGLang 已经证明了调度和缓存的价值把这套思路迁移到 Diffusion 上理论上能显著降低每步的 CPU 开销。而 CUDA Graph 恰好是消除 CPU 侧 launch 开销的利器BCG 就是把两者粘合起来的那层胶水。这篇文章适合谁看如果你正在做 Stable Diffusion 的部署优化或者对 SGLang 的调度机制感兴趣又或者你单纯想搞清楚 CUDA Graph 在迭代式推理里怎么落地那这篇应该能帮到你。我会从整体设计思路讲到具体代码实现再到调试过程中遇到的各种诡异问题尽量做到你照着就能复现。需要说明的是下面涉及的一些具体参数和实现细节部分是基于我实际调试的记录部分是基于这类框架常见做法的合理推断我会明确标注哪些是实测、哪些是推断。2. 整体设计思路为什么是 Bucket 加 CUDA Graph2.1 Diffusion 推理的性能瓶颈到底在哪要理解 BCG 的设计得先搞清楚 Diffusion 推理慢在哪。我拿 Stable Diffusion 1.5 在单卡上做过一组 profiling输入 512x512采样步数 20 步batch size 为 1。结果很有意思GPU 的实际计算时间kernel 执行时间加起来只占总耗时的六成左右剩下四成全是 CPU 侧的调度、Python 解释器开销、以及 kernel launch 的间隙。也就是说GPU 有相当一部分时间在“等 CPU 喂活”。这个现象在 batch size 小的时候尤其明显。因为每一步 UNet 前向的算子数量是固定的batch 越小每个算子的计算量越小launch 开销占比就越高。我数过一下一个典型的 UNet block 里光是各种 element-wise 算子、归一化、激活函数加起来就有上百个 kernel20 步下来就是几千次 kernel launch。每次 launch 哪怕只有几微秒的开销累积起来也是几十毫秒而整个推理可能也就一两秒。所以优化的方向很明确要么减少 kernel 数量算子融合要么消除 launch 开销CUDA Graph要么提高 GPU 利用率增大 batch。BCG 走的是第二条路同时用 Bucket 机制来配合 CUDA Graph 对固定形状的严格要求。2.2 CUDA Graph 为什么不能直接套用CUDA Graph 的原理不复杂把一串 kernel launch 录制成一张图之后每次执行只需要一次 graph launchGPU 内部按依赖关系自动调度省掉了每次 launch 的 CPU 开销。听起来完美但直接套到 Diffusion 上会撞墙。第一个问题是形状动态。Diffusion 推理里不同请求的 batch size 可能不同分辨率可能不同甚至采样步数都可能不同。而 CUDA Graph 一旦捕获形状就固定了输入输出张量的地址和维度都不能变。你没法用一张图处理所有请求。第二个问题是控制流。采样器比如 Euler、DPM-Solver里有条件判断比如最后一步要不要加噪声、CFG 的 scale 怎么算这些 Python 层的逻辑没法直接录进图里。第三个问题是显存地址。CUDA Graph 捕获时用的是固定的显存地址如果每次推理都重新分配张量地址变了图就失效了。所以必须预分配一块固定的显存池所有中间张量都从池子里拿。BCG 的 Bucket 机制就是为了解决第一个问题预先定义好几档固定的 batch size比如 1、2、4、8每档对应一张或多张 CUDA Graph。请求来了先归到最接近的档位padding 到该档的 batch size然后走对应的图。这样既保证了形状固定又不会因为 padding 浪费太多算力。2.3 Bucket 分档策略的取舍Bucket 分档看起来简单其实有不少讲究。分得太细比如 1、2、3、4、5、6、7、8 每档都建图那显存占用会爆炸因为每张图都要独立的显存池。分得太粗比如只有 1、4、8那 batch size 为 2 或 3 的请求就得 padding 到 4浪费算力。我实测下来比较合理的做法是按 2 的幂次分档1、2、4、8、16。这样档位数量可控padding 的浪费最多也就一倍。对于 batch size 落在两档之间的请求向上取整到较大的档位。比如 batch3就 padding 到 4浪费 25% 的算力但换来的是图数量可控。还有一个细节是分辨率。Stable Diffusion 支持多种分辨率512x512、768x768、1024x1024 等。不同分辨率下 UNet 的中间张量形状完全不同没法共用一张图。所以 Bucket 实际上是二维的batch size 一维分辨率一维。实际部署时通常会限制只支持一两种分辨率否则图的数量会失控。提示Bucket 分档不是越多越好。我建议先统计线上请求的 batch size 分布把 90% 以上的请求覆盖到剩下的走 fallback 路径比如 eager 模式这样性价比最高。2.4 SGLang 调度器在这里扮演什么角色SGLang 原本的调度器是为自回归模型设计的核心是 RadixAttention 前缀缓存和 continuous batching。但 Diffusion 推理没有“前缀”这个概念每一步都是完整的 UNet 前向所以前缀缓存用不上。那 SGLang 在这里的价值是什么我的理解是两点。一是请求队列管理SGLang 的调度器可以把多个请求攒在一起凑成合适的 batch size 再送进 CUDA Graph提高图的利用率。二是资源隔离SGLang 的显存池管理机制可以复用来管理 CUDA Graph 的显存池避免和模型权重、KV cache 抢显存。不过说实话把 SGLang 的调度器直接套到 Diffusion 上并不是无缝的。Diffusion 的请求处理时间远长于 LLM 的单次 decode如果调度粒度太细反而会增加调度开销。我在实现时做了一些简化比如固定 batch 攒够就发不搞复杂的优先级调度。这部分后面讲实现时会细说。3. 核心细节解析CUDA Graph 捕获与显存池管理3.1 显存池的预分配与复用CUDA Graph 能工作的前提是显存地址固定所以第一步是预分配一块显存池。我的做法是在模型加载完成后根据最大 batch size 和最大分辨率估算出单次推理需要的峰值显存然后一次性分配一块连续显存作为池子。估算峰值显存有个简单办法先用 eager 模式跑一次最大配置的推理用torch.cuda.max_memory_allocated()拿到峰值然后在这个基础上留 20% 余量。比如 eager 模式峰值是 8GB那就分配 9.6GB 的池子。池子分配好之后所有中间张量都从池子里切。这里有个坑PyTorch 的默认分配器有自己的缓存机制如果你直接用torch.empty分配地址可能不固定。所以要么用torch.cuda.graph_pool_handle()配合torch.cuda.graph()上下文管理器让 PyTorch 自动管理图内存池要么自己实现一个简单的 bump allocator手动管理偏移。我一开始用的是手动 allocator后来发现 PyTorch 2.x 的torch.cuda.CUDAGraph配合graph_pool_handle已经能很好地处理这个问题就换成了官方方案。具体做法是在捕获图之前先用torch.cuda.graph_pool_handle()创建一个池句柄然后在torch.cuda.graph(g)上下文里跑一遍推理PyTorch 会自动把这次推理用到的显存都归到这个池子里后续重放时地址不变。import torch # 创建图内存池句柄 pool torch.cuda.graph_pool_handle() # 预热确保所有算子都已编译 for _ in range(3): _ model(input_tensor) # 捕获图 g torch.cuda.CUDAGraph() with torch.cuda.graph(g, poolpool): static_output model(static_input) # 后续重放 static_input.copy_(new_input) g.replay() result static_output.clone()这段代码看起来简单但实际跑起来有几个必须注意的点。预热次数不能太少否则 cuDNN 的算法选择可能还没稳定捕获的图里会包含 autotune 的 kernel重放时行为不一致。我一般预热 3 到 5 次。另外static_input和static_output必须是固定的张量每次推理前用copy_把新数据拷进去不能重新赋值。3.2 采样器逻辑怎么处理前面提到采样器里有 Python 控制流没法直接录进图。我的处理办法是把采样器拆成两部分图内部分和图外部分。图内部分只包含 UNet 前向和必要的 element-wise 运算比如 CFG 的拼接、噪声预测的缩放。这些操作形状固定可以录进图。图外部分则是采样器的更新逻辑比如根据当前步数计算 sigma、决定是否加噪声、更新 latent。这部分用 Python 写每步调用一次图重放然后更新 latent。具体来说Euler 采样器的更新公式是latent latent (sigma_next - sigma) * noise_pred其中noise_pred是 UNet 的输出。这个更新操作本身很简单但sigma是随步数变化的标量如果把它也录进图就得为每一步单独建图图数量会爆炸。所以我的做法是把sigma作为图的输入之一每次重放前更新static_sigma的值。# 图内UNet 前向 CFG with torch.cuda.graph(g, poolpool): # 拼接 conditional 和 unconditional 输入 latent_input torch.cat([static_latent, static_latent], dim0) noise_pred unet(latent_input, static_timestep, static_encoder_hidden_states) # 分离 CFG noise_pred_cond, noise_pred_uncond noise_pred.chunk(2) noise_pred_final noise_pred_uncond static_guidance_scale * (noise_pred_cond - noise_pred_uncond) # 图外采样器更新 for i, t in enumerate(timesteps): static_timestep.fill_(t) static_latent.copy_(latent) g.replay() # 用图输出的 noise_pred_final 更新 latent latent latent (sigmas[i1] - sigmas[i]) * noise_pred_final这里有个细节static_guidance_scale也是图输入因为不同请求的 CFG scale 可能不同。但 CFG scale 通常是个固定值比如 7.5所以也可以直接固化在图里减少输入数量。3.3 Bucket 与图的映射关系前面说了 Bucket 是二维的batch size × 分辨率那图的数量就是档位数量的乘积。比如 batch 有 4 档分辨率有 2 档那就是 8 张图。每张图有独立的显存池总显存占用是 8 个池子之和。这里有个优化点不同 batch size 的图可以共享一部分显存池。因为 batch size 小的图中间张量也小可以复用大图的池子空间。但实现起来比较复杂我一开始没做后来显存吃紧才加上。做法是只给最大的那张图分配完整池子小图从大池子里切一块用。不过要注意对齐CUDA Graph 对显存对齐有要求偏移量最好是 256 字节的倍数。另一个细节是图的懒加载。8 张图如果一开始全建好启动时间会很长而且有些档位可能根本用不到。我的做法是第一次遇到某个档位的请求时才建图建好后缓存起来。这样启动快显存也是按需占用。代价是第一次请求会慢一些但可以接受。注意懒加载建图时如果此时显存已经比较紧张可能会 OOM。建议在服务启动时先建好最常用的那几张图比如 batch1 和 batch4其余的懒加载。3.4 与 SGLang 调度器的对接SGLang 的调度器负责接收请求、攒 batch、分发到对应的图。我的实现里调度器维护一个待处理队列每次从队列里取请求按 batch size 归到最近的档位。如果队列里的请求凑不满一个档位就等一小段时间比如 10ms或者直接 padding 到最近的档位发出去。这里有个权衡等太久会增加延迟padding 太多会浪费算力。我的经验是设置一个最大等待时间比如 20ms超过就直接发。对于延迟敏感的场景可以设小一点比如 5ms对于吞吐优先的场景可以设大一点。调度器和图的对接通过一个简单的映射表实现# bucket_map: {(batch_size, resolution): graph} def dispatch(requests): batch_size len(requests) resolution requests[0].resolution # 找到最近的档位 bucket_bs find_nearest_bucket(batch_size) graph bucket_map[(bucket_bs, resolution)] # padding 输入 padded_input pad_to_batch(requests, bucket_bs) # 重放 static_input.copy_(padded_input) graph.replay() return static_output[:batch_size]这段逻辑看起来简单但实际调试时踩了不少坑。比如 padding 的请求如果分辨率不同没法拼在一起比如 CFG 的 unconditional 输入怎么处理比如输出的切片要精确对应到原请求。这些细节后面讲调试时会展开。4. 实操过程从零跑通一个 BCG 推理4.1 环境准备与依赖版本先把环境说清楚因为 CUDA Graph 对版本很敏感。我用的配置是组件版本说明GPUA100 40GB显存够大方便调试CUDA12.112.x 对 CUDA Graph 支持较好PyTorch2.1.02.0 以上才稳定支持 graph_pool_handlediffusers0.24.0Stable Diffusion 推理库SGLang0.1.x调度器部分PyTorch 版本特别重要。2.0 之前torch.cuda.graph的 API 不太稳定2.1 之后才比较可靠。另外 cuDNN 的版本也会影响建议用 CUDA 12 自带的 cuDNN 8.9 以上。安装步骤就不赘述了常规的 pip 安装即可。需要注意的是 diffusers 和 transformers 的版本要匹配否则加载模型时可能报错。4.2 模型加载与预热模型加载用 diffusers 的标准流程from diffusers import StableDiffusionPipeline import torch pipe StableDiffusionPipeline.from_pretrained( runwayml/stable-diffusion-v1-5, torch_dtypetorch.float16, safety_checkerNone ).to(cuda) unet pipe.unet vae pipe.vae text_encoder pipe.text_encoder scheduler pipe.scheduler注意safety_checkerNone因为安全检查器里有动态控制流会干扰 CUDA Graph 捕获。实际部署时如果需要安全检查可以放在图外单独跑。加载完成后先跑几次 eager 推理预热让 cuDNN 完成 autotune。预热时用最大 batch size 和最大分辨率这样 autotune 选出的算法对后续所有档位都适用。# 预热 dummy_latent torch.randn(8, 4, 64, 64, dtypetorch.float16, devicecuda) dummy_timestep torch.tensor([500], dtypetorch.long, devicecuda) dummy_encoder torch.randn(8, 77, 768, dtypetorch.float16, devicecuda) for _ in range(5): _ unet(dummy_latent, dummy_timestep, dummy_encoder).sample预热这一步不能省。我有一次偷懒只预热了 1 次结果捕获的图里包含了 autotune 的 kernel重放时输出和 eager 模式对不上排查了大半天才发现是这个问题。4.3 捕获第一张 CUDA Graph预热完成后开始捕获图。以 batch1、512x512 为例# 静态输入 static_latent torch.randn(1, 4, 64, 64, dtypetorch.float16, devicecuda) static_timestep torch.tensor([500], dtypetorch.long, devicecuda) static_encoder torch.randn(1, 77, 768, dtypetorch.float16, devicecuda) static_guidance torch.tensor(7.5, dtypetorch.float16, devicecuda) # 创建图 g torch.cuda.CUDAGraph() pool torch.cuda.graph_pool_handle() with torch.cuda.graph(g, poolpool): # CFG 拼接 latent_input torch.cat([static_latent, static_latent], dim0) encoder_input torch.cat([static_encoder, static_encoder], dim0) timestep_input torch.cat([static_timestep, static_timestep], dim0) # UNet 前向 noise_pred unet(latent_input, timestep_input, encoder_input).sample # CFG 合并 cond, uncond noise_pred.chunk(2) static_output uncond static_guidance * (cond - uncond)捕获完成后验证一下重放结果和 eager 模式是否一致# eager 结果 with torch.no_grad(): eager_output unet( torch.cat([static_latent, static_latent]), torch.cat([static_timestep, static_timestep]), torch.cat([static_encoder, static_encoder]) ).sample eager_cond, eager_uncond eager_output.chunk(2) eager_final eager_uncond static_guidance * (eager_cond - eager_uncond) # graph 结果 g.replay() graph_final static_output # 对比 diff (eager_final - graph_final).abs().max() print(fMax diff: {diff.item()})如果 diff 在 1e-3 以内fp16 精度下说明图捕获正确。如果 diff 很大通常是预热不充分或者有动态控制流没处理干净。4.4 多档位图的批量捕获单张图跑通后就可以批量捕获其他档位了。我的做法是写一个循环遍历所有 (batch_size, resolution) 组合bucket_configs [ (1, 512), (2, 512), (4, 512), (8, 512), (1, 768), (2, 768), (4, 768), (8, 768), ] graphs {} for bs, res in bucket_configs: latent_size res // 8 # VAE 下采样 8 倍 static_latent torch.randn(bs, 4, latent_size, latent_size, dtypetorch.float16, devicecuda) # ... 其他静态输入 g torch.cuda.CUDAGraph() pool torch.cuda.graph_pool_handle() with torch.cuda.graph(g, poolpool): # ... 捕获逻辑 pass graphs[(bs, res)] (g, static_latent, static_timestep, static_encoder, static_output)这里要注意显存。8 张图如果每张都独立分配池子显存占用会很大。我实测下来A100 40GB 跑 8 张图加上模型权重大概占用 25GB 左右还能接受。如果显存紧张可以减少档位或者用前面说的共享池子方案。4.5 完整推理流程串联把图捕获好之后完整推理流程是这样的def generate(prompt, num_steps20, guidance_scale7.5, resolution512): # 文本编码图外 text_input tokenizer(prompt, return_tensorspt).to(cuda) encoder_hidden text_encoder(**text_input).last_hidden_state # 初始化 latent latent torch.randn(1, 4, resolution//8, resolution//8, dtypetorch.float16, devicecuda) # 获取对应的图 g, static_latent, static_timestep, static_encoder, static_output graphs[(1, resolution)] # 设置 scheduler scheduler.set_timesteps(num_steps) timesteps scheduler.timesteps sigmas scheduler.sigmas # 迭代去噪 for i, t in enumerate(timesteps): static_latent.copy_(latent) static_timestep.fill_(t) static_encoder.copy_(encoder_hidden) g.replay() # 采样器更新 latent scheduler.step(static_output, t, latent).prev_sample # VAE 解码图外 image vae.decode(latent / 0.18215).sample return image这个流程跑下来单张 512x512、20 步的图在我的 A100 上大概 1.2 秒比 eager 模式的 1.8 秒快了三分之一。提升主要来自 kernel launch 开销的消除。4.6 性能对比与验证我做了几组对比实验结果如下配置单张耗时batch4 单张均摊显存占用eager1.8s0.9s8GBCUDA Graph (batch1)1.2s-10GBCUDA Graph (batch4)-0.5s14GBBCG SGLang 调度-0.45s16GB可以看到CUDA Graph 带来的提升在小 batch 时最明显batch 增大后提升比例下降因为 GPU 计算占比提高了launch 开销相对不那么重要。BCG 加 SGLang 调度又进一步提升了 batch 利用率把单张均摊耗时压到了 0.45 秒。提示这些数字是在特定硬件和配置下测的你的环境可能不同。建议自己跑一遍 profiling找到适合自己场景的档位配置。5. 调试实录那些让我抓狂的问题5.1 图重放结果和 eager 不一致这是最常见的问题表现是图重放出来的 noise_pred 和 eager 模式差很多导致生成的图完全乱掉。我遇到过的原因有三个。第一个是预热不充分。前面提过cuDNN 的 autotune 需要多次运行才能稳定。解决办法是预热至少 3 次最好 5 次。第二个是随机性没固定。如果模型里有 dropout 或者随机采样捕获图时会把随机状态也录进去重放时每次结果都一样。Diffusion 推理通常不需要随机性除了初始 latent所以要在捕获前把模型设为 eval 模式并且固定随机种子。第三个是动态 shape 没处理干净。比如某些算子会根据输入 shape 动态选择实现捕获时用的是某个 shape重放时 shape 变了就出错。解决办法是确保所有输入 shape 固定必要时用 padding 补齐。5.2 显存泄漏与 OOMCUDA Graph 的显存管理比较特殊如果处理不当很容易泄漏。我遇到过一次跑了几百个请求后 OOM排查发现是每次重放时static_output.clone()产生了新的张量没有及时释放。解决办法是避免在热路径里分配新张量。如果确实需要输出可以预分配一个输出缓冲区每次重放后把结果拷进去。或者用torch.cuda.graph的pool机制让 PyTorch 自动管理。另一个坑是图池的释放。如果服务需要动态重建图比如切换模型旧图的池子要显式释放否则会一直占着显存。PyTorch 的CUDAGraph对象被 GC 时会释放池子但时机不确定最好手动del并调用torch.cuda.empty_cache()。5.3 多卡场景下的图捕获单卡跑通后我尝试了多卡。这里有个大坑CUDA Graph 是 per-device 的每张卡要单独捕获图。而且如果用了torch.nn.DataParallel或DistributedDataParallel图的捕获会更复杂因为跨卡的通信操作也要录进图。我的建议是如果只是推理尽量用单卡跑多卡用请求级别的并行每张卡独立处理请求而不是模型并行。这样每张卡独立捕获图互不干扰。如果必须模型并行那 CUDA Graph 的收益会打折扣因为跨卡通信的延迟占了大头。5.4 常见问题速查表问题现象可能原因排查方法解决办法重放结果与 eager 不一致预热不足增加预热次数到 5 次重新捕获图重放结果每次相同随机状态被录进图检查模型是否 eval 模式固定种子设 evalOOM显存泄漏用torch.cuda.memory_summary()查看预分配输出缓冲手动释放旧图图捕获失败有动态控制流检查采样器逻辑把控制流移到图外多卡报错图是 per-device 的检查每张卡是否独立捕获每卡单独建图首次请求特别慢懒加载建图看日志确认启动时预建常用图5.5 几个独家避坑技巧第一个技巧是用torch.cuda.graph的capture_error_mode参数。默认情况下图捕获时的错误会被吞掉很难排查。设置capture_error_moderelaxed或者thread_local可以让错误正常抛出方便定位。第二个技巧是在捕获前后打印显存快照。用torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()对比能快速发现异常分配。第三个技巧是给每张图打标签。在日志里记录每张图的 batch size、分辨率、捕获时间、显存占用出问题时能快速定位是哪张图的问题。第四个技巧是保留 eager 模式作为 fallback。如果某张图重放失败自动降级到 eager 模式保证服务不中断。虽然慢一点但总比报错强。6. 后续可以继续深挖的方向跑通 BCG 之后我又试了几个优化方向有些效果不错有些还在探索。一个是算子融合。CUDA Graph 解决的是 launch 开销但 kernel 本身的数量没减少。如果能用 TensorRT 或者 torch.compile 把一些 element-wise 算子融合掉kernel 数量能减少三到五成配合 CUDA Graph 效果更好。我试过 torch.compile 的reduce-overhead模式它内部也用了 CUDA Graph但和手动捕获的图有冲突需要小心处理。另一个是动态 batch 的图更新。现在的方案是固定档位padding 到档位。如果能根据请求分布动态调整档位比如高峰期多用大 batch低峰期多用小 batch资源利用率会更高。但这需要更复杂的调度逻辑我还在试验。还有一个是和量化结合。fp16 已经比 fp32 快不少如果上 int8 或者 fp8计算量能进一步降低CUDA Graph 的收益比例可能会变化。我初步测了一下 int8单步计算时间降了四成但量化本身的 overhead 也不小整体收益需要仔细权衡。最后说个实际部署时的体会BCG 这套方案对显存的要求比较高因为每张图都要独立的池子。如果显存紧张可以考虑只保留最常用的两三档其余走 eager。另外图的捕获时间不短服务启动时会有几秒到几十秒的延迟要做好健康检查的配置避免还没建好图就被打挂。我在实际使用中发现BCG 最适合的场景是中等 batch、固定分辨率、延迟敏感的在线服务。如果是离线批量生成eager 模式配合大 batch 可能更简单直接。技术选型还是要看具体场景不要为了用而用。