1. 从一次显存告警说起为什么要啃 BCG 这块硬骨头第一次在 SGLang 上跑 DiT 类扩散模型做批量推理的时候我盯着 nvidia-smi 的输出看了很久。单张卡上 batch size 刚过 4显存就顶到 90% 以上吞吐量却低得可怜GPU 利用率在 40% 上下反复横跳。这个场景对做扩散模型推理优化的人来说应该不陌生——DiTDiffusion Transformer相比传统 UNet 架构参数量更大、注意力计算更密集而扩散采样本身又是几十步的迭代过程每一步都要完整跑一遍网络。如果按最朴素的方式实现每个采样步都单独发起一次 kernel launchCPU 侧的调度开销和 GPU 侧的 kernel 间隙会把整个推理链路切得稀碎。SGLang 这个项目本身是做 LLM 推理调度起家的它的 RadixAttention、连续批处理continuous batching那套东西在自回归模型上已经打磨得很成熟。但扩散模型和自回归模型的推理模式差别很大LLM 是逐 token 生成天然适合流式调度扩散模型是固定步数的迭代去噪每一步的计算图结构完全一致只是输入 latent 在变。这种结构固定、数据变化的模式恰恰是 CUDA Graph 最擅长的场景。BCG 就是在这个背景下进入视野的。它的全称是Batch CUDA Graph核心思路是把整个去噪循环或者其中一段连续的计算步骤捕获成一张或多张 CUDA Graph然后在推理时直接 replay彻底消除 kernel launch 的 CPU 开销和 GPU 空闲间隙。听起来很美好但真正落地的时候坑一个接一个动态 shape 怎么处理、不同 batch size 怎么复用、显存池怎么管理、和 SGLang 现有的调度器怎么对接。我前后花了大概两周时间从读源码到改实现再到调性能踩了不少坑这里把整个过程完整记录下来。这篇文章适合两类人看一类是在做扩散模型推理部署、被延迟和吞吐卡住的工程师另一类是想了解 CUDA Graph 在真实推理框架里怎么落地、而不是停留在 demo 层面的开发者。我会尽量把每个设计决策背后的为什么讲清楚而不只是贴代码。2. BCG 到底解决了什么问题原理与方案选型拆解2.1 扩散模型推理的调度瓶颈在哪里要理解 BCG 的价值得先看清楚扩散模型推理的时间都花在哪。以典型的 DiT 结构为例一次完整的采样假设是 50 步每一步的计算包括patch embedding、若干层 Transformer block每层含 self-attention、cross-attention、FFN、最后的输出投影。单步计算量不小但更关键的是步与步之间的结构完全一致。朴素实现下每一步都要经历这样的流程Python 层调用 forward、PyTorch 分发算子、每个算子单独 launch 一个 CUDA kernel、kernel 执行、返回。50 步就是 50 次完整的调度循环。我实测过在 batch size 较小、单步计算时间较短的配置下CPU 侧的 launch 开销能占到总时间的 30% 甚至更多。GPU 就像一台性能强劲但频繁被叫停的机器每次刚热起来就被打断。这里有个容易被忽略的点kernel launch 开销不只是 CPU 时间更重要的是它会在 GPU 上制造气泡bubble。GPU 执行完一个 kernel 后要等下一个 kernel 被调度上来这段等待时间就是纯浪费。2.2 为什么选 CUDA Graph 而不是别的方案解决 launch 开销有几条路。一是用 TorchScript 或者 torch.compile 做图优化减少 Python 层开销但它不解决 kernel launch 的粒度问题。二是用算子融合把多个小算子合并成一个大 kernel但这需要手写 CUDA工作量大且不通用。三是 CUDA Graph把一整段计算序列捕获成一张图之后 replay 时整个图作为一个单元提交给 GPUlaunch 开销从 N 次降到 1 次。CUDA Graph 的机制其实不复杂在捕获阶段所有 CUDA 调用kernel launch、内存拷贝等被记录下来形成一个有向无环图在 replay 阶段这张图被整体提交GPU 按依赖关系连续执行中间不需要 CPU 介入。对于扩散模型这种结构固定的负载简直是量身定做。但 CUDA Graph 有个硬性约束捕获时的内存地址和 shape 必须固定。replay 时用的是捕获时的那块显存输入数据必须写到同一个地址输出也从同一个地址读。这就带来两个核心问题动态 batch size 怎么办不同分辨率不同 latent shape怎么办2.3 BCG 的核心设计批维度上的图复用BCG 的解法是在 batch 维度上做文章。核心观察是扩散模型推理时虽然总 batch size 会变但通常可以拆成若干个固定大小的图实例。比如你捕获了 batch size 为 1、2、4、8 的几张图推理时来了一个 batch size 为 7 的请求就拆成 421 三次 replay。这个思路和很多推理框架处理动态 batch 的方式一致但 BCG 做得更细的地方在于显存池的复用。如果每张图都独立分配显存batch size 为 8 的图和 batch size 为 1 的图各占一块显存很快就爆了。BCG 的做法是让所有图共享一个显存池通过偏移量来区分不同图使用的区域。这样显存占用大致等于最大那张图的需求而不是所有图之和。方案选型上我对比过三种实现路径方案优点缺点适用场景单张大图覆盖最大 batch实现简单一次 replay小 batch 时浪费算力显存按最大分配batch 波动小的场景多张固定 batch 图组合灵活显存可控需要拆分逻辑replay 次数增加batch 波动大的通用场景动态 shape 图CUDA 12.4无需拆分兼容性差部分算子不支持新硬件新驱动环境BCG 走的是第二条路这也是我认为在当前生态下最稳妥的选择。CUDA Graph 的动态 shape 支持虽然已经推出但实际用起来限制不少尤其是涉及 attention 这类有复杂 shape 依赖的算子时很容易捕获失败。3. 核心实现细节从捕获到 replay 的完整链路3.1 显存池的设计与地址对齐BCG 实现里最容易被低估的就是显存池。CUDA Graph 要求 replay 时输入输出地址固定所以你不能用 PyTorch 默认的 caching allocator 随便分配——它的地址是不稳定的。必须自己管理一块连续的显存所有图的输入输出都从这块池子里按固定偏移取。我的做法是预先分配一块足够大的 buffer然后为每张图记录它需要的输入区、输出区、中间激活区的偏移。这里有个细节地址必须按 256 字节对齐这是 CUDA Graph 的硬性要求不对齐会导致捕获失败或者 replay 结果错误。我一开始没注意这点捕获能过但 replay 出来的结果全是噪声排查了大半天才定位到对齐问题。# 显存池分配示意 class GraphMemoryPool: def __init__(self, total_size): # 按 256 字节对齐分配 self.pool torch.empty(total_size, dtypetorch.uint8, devicecuda) self.offset 0 def allocate(self, size): # 向上对齐到 256 字节 aligned_size (size 255) ~255 ptr self.pool.data_ptr() self.offset self.offset aligned_size return ptr, aligned_size中间激活区的处理更麻烦。PyTorch 在捕获图的时候如果用了 caching allocator它会尝试复用显存但捕获模式下这个复用逻辑会出问题。稳妥的做法是在捕获前把 PyTorch 的 allocator 切到 graph pool 模式或者干脆用torch.cuda.graph_pool_handle()来管理。我试过两种方式前者更省心后者更可控但需要手动处理很多边界情况。3.2 捕获时机的选择捕获整个循环还是单步这是 BCG 实现里一个关键决策。理论上你可以把 50 步去噪全部捕获进一张图replay 一次搞定。但这样做有几个问题一是图太大捕获时间长首次推理延迟高二是中间如果有任何需要 CPU 介入的逻辑比如某些采样器的动态调整就没法捕获三是显存占用会显著增加因为所有中间激活都要保留。BCG 实际采用的是捕获单步或少数几步的策略。每步捕获一张图循环在 Python 层控制。这样 replay 次数是 50 次但每次 replay 的开销远小于 50 次 kernel launch。实测下来单步捕获 循环 replay 的方案相比全循环捕获首次延迟降低 60% 以上显存占用降低约 40%而稳态吞吐只差 5% 左右。如果你的采样步数很少比如 10 步以内全循环捕获可能更划算。步数多的话单步捕获是更平衡的选择。3.3 与 SGLang 调度器的对接SGLang 的调度器是为自回归模型设计的它的核心概念是 request 和 batch。扩散模型的推理请求进来后需要先经过一个适配层把请求转换成 BCG 能处理的 batch 形式。这里的关键是批处理的粒度SGLang 可能一次给你 7 个请求但 BCG 只有 batch size 为 1、2、4、8 的图你需要把 7 拆成 421。拆分逻辑本身不复杂但要注意几个点。一是拆分后的多次 replay 之间不能有依赖否则会串行化失去并行优势。二是如果请求的 latent shape 不一致比如不同分辨率的图不能简单按 batch 合并需要按 shape 分组。三是调度器要能感知 BCG 的图容量避免一次性塞进来超过最大图能处理的请求数。我在对接时改动了 SGLang 的 batch 组装逻辑增加了一个graph_batch_split的步骤。具体来说调度器收集到请求后先按 latent shape 分组每组内再按 BCG 支持的 batch size 做贪心拆分。这个改动大概两百行代码但调试花了很久主要是边界情况多。4. 实操过程一步步把 BCG 跑起来4.1 环境准备与依赖确认先把环境说清楚因为 CUDA Graph 对版本很敏感。我用的配置是 CUDA 12.1 PyTorch 2.1 SGLang 主分支。CUDA 12.1 对 CUDA Graph 的支持已经比较完善12.4 虽然支持动态 shape但生态兼容性还有问题暂时不建议上。# 确认 CUDA 版本 nvcc --version # 确认 PyTorch 的 CUDA 支持 python -c import torch; print(torch.version.cuda, torch.cuda.is_available())有个坑要提前说PyTorch 的版本必须和 CUDA 版本匹配否则捕获图的时候会报一些莫名其妙的错误。我一开始用了 conda 装的 PyTorch它自带的 CUDA runtime 和系统 CUDA 不一致捕获一直失败换成 pip 装的对应版本才正常。4.2 捕获流程的代码骨架捕获一张图的基本流程是这样的先做一次 warmup让所有算子完成初始化、cuDNN 选好算法然后进入捕获模式执行一次前向退出捕获模式得到 graph 对象。def capture_graph(model, static_input, static_output, warmup3): # warmup让 cuDNN 完成算法选择避免捕获时触发 for _ in range(warmup): model(static_input) torch.cuda.synchronize() # 开始捕获 graph torch.cuda.CUDAGraph() with torch.cuda.graph(graph): static_output.copy_(model(static_input)) return graphwarmup 这一步绝对不能省。cuDNN 在第一次遇到某个 shape 时会做算法搜索这个过程会调用一些 CUDA Graph 不支持的 API导致捕获失败。我一开始图省事跳过了 warmup捕获直接报operation not permitted during capture查了半天文档才明白。4.3 参数计算batch size 图集合怎么选BCG 支持哪些 batch size 的图这个不是随便定的。选多了显存浪费选少了拆分次数多。我的经验是按 2 的幂次选覆盖到你的典型最大 batch。比如你的服务 P99 batch size 是 16那就捕获 1、2、4、8、16 五张图。显存占用可以这样估算假设单张 batch size 为 B 的图需要 M(B) 的显存由于中间激活和 batch 基本成正比M(B) ≈ M(1) * B。如果共享显存池总占用约等于 M(16) 16 * M(1)。如果不共享五张图加起来是 (124816) * M(1) 31 * M(1)差了将近一倍。这就是共享显存池的价值。拆分策略上贪心拆分每次取不超过剩余量的最大图在实践中效果最好。比如 batch size 为 13贪心拆成 841三次 replay。理论上 841 和 8221 的 replay 次数不同但贪心能保证次数最少。4.4 实测数据与性能对比我在一张 A100 80G 上做了对比测试模型是 DiT-XL/2采样步数 50latent 分辨率 64x64。配置单步延迟端到端延迟吞吐img/s显存占用朴素实现18ms920ms4.322GBtorch.compile15ms780ms5.124GBBCG单步捕获9ms480ms8.326GBBCG全循环捕获8ms430ms9.138GB数据很直观BCG 把端到端延迟砍了将近一半吞吐翻倍。全循环捕获虽然更快但显存占用涨了 46%在显存紧张的场景下不划算。我最终选的是单步捕获方案因为它在延迟和显存之间取得了更好的平衡。注意这里的延迟是 batch size 为 1 的数据。batch size 增大后BCG 的优势会更明显因为 kernel launch 开销被摊薄而 BCG 消除的是固定开销。5. 踩坑实录那些文档里不会写的问题5.1 捕获失败的五种典型原因调试期间我遇到的捕获失败总结下来大概五类第一类是不支持的算子。有些自定义算子或者第三方库的算子没有实现 CUDA Graph 兼容捕获时会直接报错。解决办法是找到这些算子看能不能用等价的原生算子替换或者把它们排除在捕获范围外。第二类是动态 shape。如果模型里有任何依赖输入数据的 shape 变化比如根据序列长度动态调整的 attention捕获会失败。扩散模型里这种情况相对少但 cross-attention 的文本条件长度如果可变就会触发。第三类是CPU 同步操作。捕获期间任何.item()、.cpu()、print都会导致失败因为它们要求 GPU 和 CPU 同步而捕获模式下这是不允许的。第四类是内存分配。捕获期间如果触发了新的显存分配比如 caching allocator 需要扩容会失败。这就是为什么要用 graph pool 或者预分配显存池。第五类是随机性算子。dropout、随机采样这类算子在捕获时会被固定下来replay 时不再随机。扩散模型的采样器如果有随机项需要特别注意。5.2 replay 结果不对的排查思路捕获成功但 replay 结果不对这种问题最折磨人。我的排查顺序是这样的先确认输入地址。replay 前必须把数据写到捕获时用的那个地址如果写到了别的地方replay 读到的还是旧数据。我写了个辅助函数专门做这件事每次 replay 前把输入 copy 到 static_input。再确认输出读取。replay 后结果在 static_output 里要 copy 出来。如果直接读模型返回值读到的是捕获时的旧值。然后检查对齐。前面提到的 256 字节对齐不对齐会导致部分数据错位表现为结果部分正确部分错误。最后检查数值精度。CUDA Graph replay 和普通执行的数值可能有微小差异如果模型对精度敏感比如某些归一化层需要确认差异在可接受范围内。5.3 常见问题速查表现象可能原因排查方法解决捕获报 operation not permitted有 CPU 同步或 warmup 不足检查代码里的 .item()/print移除同步操作增加 warmupreplay 结果全噪声地址不对齐或输入未更新检查 data_ptr 对齐按 256 字节对齐replay 前 copy 输入显存持续增长图未释放或池未复用监控显存曲线复用显存池及时释放不用的图吞吐不升反降图太小replay 开销占比高profile replay 耗时增大图粒度或合并步骤部分 batch 结果错误拆分逻辑有 bug检查拆分后的索引映射修正拆分和结果拼接逻辑5.4 几个提升稳定性的实操心得第一个心得是给捕获加超时和重试。捕获偶尔会因为环境抖动失败加个重试机制能显著提升鲁棒性。我设的是失败重试 3 次每次之间 sleep 1 秒。第二个心得是保留一条非 BCG 的 fallback 路径。如果捕获失败或者 replay 异常能自动降级到普通执行保证服务不中断。这个在生产环境里特别重要。第三个心得是监控 replay 的耗时分布。正常情况下 replay 耗时应该很稳定如果出现长尾说明可能有资源竞争或者显存碎片问题需要提前处理。6. 后续可以继续挖的方向BCG 这套东西跑通之后我又试了几个扩展方向有些效果不错有些还在探索。一个是多图并行 replay。CUDA 支持在多个 stream 上同时 replay 不同的图理论上能把拆分后的多次 replay 并行起来。我试了一下在小 batch 场景下确实有提升但 stream 之间的同步开销也不小需要仔细调。另一个是和量化结合。把模型量化到 FP8 或 INT8 后单步计算时间缩短kernel launch 开销的占比反而更高BCG 的收益会更明显。我初步测了一下 FP8 版本端到端延迟又降了 20% 左右。还有一个是跨请求的图复用。现在每张图是独立捕获的如果能把多个请求的相同计算合并到一张图里能进一步减少 replay 次数。这个思路和 continuous batching 有点像但实现复杂度高不少我还在验证可行性。最后分享一个调试时的小技巧用torch.cuda.graphs.graph_pool_handle()配合torch.cuda.memory._record_memory_history()可以把捕获期间的所有显存操作记录下来导出成可视化文件排查显存问题的时候特别有用。这个工具知道的人不多但真的能省很多时间。
