1. 从一次显存告警说起：BCG到底在解决什么问题
第一次在SGLang里跑Diffusion推理的时候，我盯着显存监控看了很久。模型加载完之后，显存占用比预期高出一大截，batch稍微开大一点就直接OOM。当时第一反应是模型权重太大，但算了一下参数量，按FP16算下来不该这么夸张。后来把profiler打开，才发现真正吃显存的不是权重本身，而是推理过程中反复分配释放的中间激活张量，以及CUDA Graph捕获时留下的那份静态内存池。
这就是BCG要处理的核心矛盾。BCG是Bucketed CUDA Graph的缩写，直译过来就是"分桶的CUDA Graph"。要理解它，得先理解CUDA Graph本身在Diffusion推理里为什么重要，以及它为什么又会带来新的问题。
Diffusion模型，尤其是DiT架构这一类，推理过程是一个典型的多步迭代：每一步都要跑一遍完整的去噪网络，步数通常在20到50步之间。如果每一步都走常规的kernel launch流程，CPU端要发出成百上千个kernel调用指令，每个调用都有固定的launch overhead。步数一多，这个开销累积起来非常可观，GPU经常处于"等CPU发指令"的空转状态。CUDA Graph的思路是把一整段固定的kernel调用序列录制成一张图，之后每次执行只需要一次launch，把CPU端的开销压到最低。
但问题来了。Diffusion推理里，不同batch size、不同序列长度、不同分辨率对应的张量形状是不一样的。CUDA Graph要求录制时的形状和执行时的形状完全一致，形状一变，图就失效了。如果为每一种可能的形状都录一张图，显存会被这些静态图撑爆。BCG就是在这个背景下出现的：它把不同形状的请求按某种规则归到有限的几个"桶"里，每个桶对应一张录好的CUDA Graph，请求进来先判断落到哪个桶，然后复用那张图。
我这次调试的目标，就是搞清楚SGLang Diffusion里BCG这套机制具体是怎么落地的，桶是怎么划分的，图是怎么捕获和复用的，以及在实际跑的时候会遇到哪些坑。
2. SGLang Diffusion里BCG的代码落点与调用链路
2.1 从入口到图捕获的完整路径
要调试一个机制，第一步永远是找到它在代码里的位置。SGLang Diffusion的推理入口通常从pipeline的forward开始，但BCG的介入点不在最外层，而是在模型执行器的层面。我沿着调用栈往下追，大致路径是这样的：请求进来之后先经过调度器，调度器决定这一批请求怎么组batch，然后交给模型runner，runner在真正执行去噪循环之前，会先判断当前这批请求的形状能不能命中已有的CUDA Graph。
这里有个关键设计：BCG不是在去噪循环的每一步都做判断，而是在整个循环开始之前就确定好要用哪张图。因为去噪循环内部的张量形状在步与步之间是稳定的，只有batch维度和序列维度会随请求变化。所以BCG的桶划分主要围绕这两个维度展开。
我在代码里找到的桶管理逻辑大致是这样的结构：有一个bucket manager负责维护所有已捕获的图，每个图对应一个key，key由batch size和序列长度（或者更准确地说，是token数量）组合而成。请求进来时，manager会根据当前请求的实际形状去查表，如果命中就直接复用，没命中就判断是否需要新捕获一张图。
2.2 桶的key是怎么算出来的
这一步是理解BCG的关键。桶的key不是简单的(batch_size, seq_len)二元组，因为那样组合太多，桶会爆炸。实际实现里通常会对形状做一定的归并。比如batch size可能按2的幂次分档，1、2、4、8这样；序列长度则可能按一个固定的步长向上取整，比如每64个token一档。
我实测下来，这种归并策略的效果取决于你的实际请求分布。如果你的请求batch size集中在某几个值附近，归并带来的浪费就很小；如果请求形状非常分散，那要么桶数量多导致显存吃紧，要么归并粒度粗导致计算浪费。这里没有一个万能参数，得根据自己的负载来调。
有个细节值得注意：归并的时候是向上取整还是向下取整，直接影响正确性。向上取整意味着实际张量比图里录的要小，多出来的部分需要padding或者mask掉；向下取整则可能导致张量比图大，直接报错。SGLang的实现里我看到的做法是向上取整，然后在执行时对超出部分做处理。这个处理逻辑如果写得不严谨，很容易出现结果不对但又不报错的情况，非常隐蔽。
2.3 图捕获时机的选择
CUDA Graph的捕获本身是有成本的，而且捕获期间不能有其他CUDA操作干扰。所以BCG不会在服务启动时就把所有可能的桶都录好，那样启动太慢。常见的做法是懒捕获：第一个命中某个桶的请求触发这张图的录制，录制完成后再执行。后续命中同一个桶的请求就直接复用。
懒捕获带来的一个副作用是首请求延迟。第一个落到某个桶的请求，要额外承担一次图捕获的时间。如果桶划分得比较细，服务刚起来那段时间会频繁触发捕获，延迟抖动会比较明显。我在调试时观察到，冷启动阶段的前几十个请求延迟明显高于稳态，就是这个原因。
提示：如果你对首请求延迟敏感，可以考虑在服务启动后用一个预热脚本，把预期的几种典型形状各跑一遍，提前把图录好。这样线上流量进来时基本都是命中已有图的状态。
3. 调试BCG时我踩过的几个坑
3.1 形状不匹配导致的静默错误
前面提到归并时向上取整需要处理padding。我遇到过一次结果异常：同样的prompt，单独跑没问题，一旦和另一个请求组batch跑，输出就偏了。排查了很久才定位到，是BCG复用的图里序列长度比实际请求大，多出来的位置没有被正确mask，注意力计算时把padding位置的无效值也算进去了。
这类问题的麻烦之处在于它不报错，只是结果悄悄变差。我的排查方法是：固定一个prompt，分别用batch size 1和batch size 2跑，对比输出。如果batch 2的结果和batch 1不一致，基本可以确定是padding或mask的问题。进一步可以用一个全零的padding输入去验证，看输出是否受影响。
修复思路是在图执行前，把实际有效长度之外的位置显式置为注意力mask能识别的值，通常是负无穷或者一个极大的负数，确保softmax之后这些位置的权重趋近于零。
3.2 显存池碎片化
CUDA Graph捕获时会占用一块静态显存作为图的私有内存池。BCG维护多张图，就意味着多块这样的内存池。如果桶数量多，这些池子加起来会占掉相当一部分显存，而且它们是常驻的，不会因为请求结束而释放。
我一开始没意识到这一点，把桶的粒度设得很细，结果跑了一段时间后显存占用居高不下，新的请求进不来。后来把桶数量压下来，显存才恢复正常。这里的经验是：桶的数量和显存占用是直接的正相关，划分桶的时候要把这部分开销算进去。
一个实用的估算方法是：先跑一张图，看它占多少显存，然后乘以你预期的桶数量，再加上模型权重和激活的常规占用，看是否在预算内。如果超了，就得考虑合并桶或者降低归并精度。
3.3 图捕获期间的同步问题
CUDA Graph捕获要求捕获期间不能有跨流的同步操作，也不能有动态内存分配。Diffusion推理里有些操作，比如某些自定义kernel或者动态shape的算子，天然不满足这个条件。我在调试时遇到过捕获直接失败的情况，报错信息指向某个算子不支持图捕获。
解决办法通常有两个：一是把不支持的算子排除在图之外，让它走常规执行路径，图只覆盖支持的部分；二是找这个算子的替代实现。SGLang的实现里对哪些算子能进图、哪些不能进图应该是有白名单或黑名单的，调试时可以先确认当前用的算子是否在支持列表里。
注意：图捕获失败有时候不会给出很明确的错误信息，可能只是捕获后执行结果不对。遇到这种情况，可以先关掉BCG，用常规路径跑一遍确认结果正确，再开BCG对比，这样能快速判断问题是否出在图捕获上。
4. BCG和DiT架构配合时的特殊考量
4.1 DiT的序列长度特性
DiT架构和传统的UNet架构在形状特性上有个明显区别：DiT把图像切成patch之后当成序列来处理，序列长度和图像分辨率、patch大小直接相关。分辨率一变，序列长度就变，而且变化幅度可能很大。这对BCG的桶划分提出了更高要求。
举个例子，512x512的图像，patch size为2，序列长度大概是(512/2)^2加上一些条件token，量级在几万。如果换成1024x1024，序列长度直接翻四倍。这种跨度下，如果桶的归并粒度太粗，小分辨率请求会被归到大分辨率的桶里，浪费大量计算；如果太细，桶数量又会失控。
我的做法是针对常用的几档分辨率分别设桶，而不是用一个统一的归并策略。因为实际业务里分辨率往往是固定的几档，不会连续变化，按档设桶反而更高效。
4.2 单流DiT和多流DiT的差异
单流DiT和多流DiT在BCG上的表现也不一样。单流结构里，所有token走同一条计算路径，形状规整，图捕获相对简单。多流结构里，不同模态或不同分支的token可能走不同的路径，形状更复杂，图捕获时需要处理的分支更多。
我在调试多流DiT时发现，如果不同分支的序列长度不一致，BCG的桶key需要把每个分支的长度都考虑进去，否则会出现某个分支形状不匹配的问题。这种情况下，桶的key会变成一个多维向量，管理复杂度上升。一个简化思路是把所有分支的长度归并到同一个档位，牺牲一点灵活性换取管理的简洁。
4.3 和注意力机制的交互
DiT里的注意力计算，尤其是长序列下的注意力，是显存和计算的大头。BCG把注意力算子录进图里之后，注意力相关的中间张量也会被静态分配。这部分显存如果没算准，很容易成为压垮显存的最后一根稻草。
我实测下来，注意力部分的显存占用和序列长度大致是平方关系（如果是全注意力）或者线性关系（如果是某种稀疏或分块注意力）。做桶划分和显存预算时，这一项要单独拎出来算，不能和普通线性层的占用混在一起估。
5. 一套可复现的BCG调试流程
5.1 环境准备与基线确认
调试之前，先把基线跑通。关掉BCG，用常规执行路径跑一组固定输入，记录输出结果、延迟、显存占用。这组数据是后面所有对比的参照。基线一定要在干净环境下跑，避免其他进程干扰。
我习惯用固定的随机种子，确保每次输入一致。Diffusion模型本身有随机性，如果不固定种子，输出对比就没有意义。种子固定之后，同样的输入应该得到完全一样的输出，这是判断BCG是否引入偏差的前提。
5.2 开启BCG并观察桶命中情况
开启BCG后，第一件事是看桶的命中率。如果大部分请求都命中已有桶，说明桶划分合理；如果频繁触发新桶捕获，说明划分策略和实际负载不匹配。SGLang里应该有相关的日志或统计指标可以看，没有的话可以自己在bucket manager里加计数。
我一般会构造几组不同形状的请求，覆盖预期的形状范围，然后看桶的数量增长曲线。理想情况下，桶数量应该在跑了一小段时间后趋于稳定，不再增长。如果一直增长，说明归并策略有问题，需要调整。
5.3 逐桶验证输出正确性
桶命中正常之后，要逐个验证每个桶的输出是否正确。方法是：对每个桶，用一组能命中它的输入，分别跑BCG和常规路径，对比输出。如果所有桶的输出都一致，说明BCG没有引入偏差。
这一步不能省。我见过太多情况是BCG跑起来没问题，但某些边界形状下结果悄悄偏了。逐桶验证能把这些边界情况揪出来。验证时要注意覆盖桶的边界，比如刚好等于桶容量和略小于桶容量的输入，这两种情况下的padding处理逻辑可能不同。
5.4 压力测试与显存监控
正确性确认之后，做压力测试。逐步加大请求并发，观察显存占用、延迟、吞吐的变化。重点关注两个点：一是显存是否在某个并发下突然跳升（可能是触发了新桶捕获），二是延迟是否稳定（如果抖动大，可能是桶命中率下降）。
压力测试要跑足够长时间，让各种形状的请求都有机会出现。短时间的测试可能覆盖不到所有桶，漏掉潜在问题。
6. 几个容易被忽略的调优细节
6.1 桶的预热策略
前面提过预热，这里展开说一下具体怎么做。预热的核心是让服务在正式接流量之前，把预期的桶都录好。预热的输入不需要真实数据，随机张量就行，只要形状对。预热脚本可以按预期的形状分布，每个形状跑几次，确保对应的桶都被触发。
预热的顺序也有讲究。如果桶之间有依赖（比如大桶的图捕获会占用显存，影响小桶的捕获），预热顺序不对可能导致某些桶捕获失败。我的做法是从小桶到大桶依次预热，先占小的，再占大的，避免大桶把显存吃光导致小桶录不进去。
6.2 动态调整桶的粒度
桶的粒度不是一成不变的。服务刚上线时，请求分布还不明确，可以先设粗一点，桶少一些，保证稳定。等积累了一段时间的请求形状数据，再根据实际分布调整粒度，把高频形状单独设桶，低频形状归并。
这个调整过程最好是动态的，但实现起来复杂。折中方案是定期离线分析请求日志，手动调整配置后重启。虽然不够优雅，但胜在可控。
6.3 和调度器的配合
BCG的桶命中率和调度器的组batch策略直接相关。如果调度器把形状差异很大的请求组到一个batch里，这个batch的形状就会被归并到较大的桶，造成浪费。反过来，如果调度器倾向于把形状相近的请求组在一起，桶命中率和计算效率都会更好。
所以调BCG的时候，不能只看BCG本身，还要看调度器怎么组batch。有些实现里调度器会考虑形状相似度，有些则只看请求到达时间。如果是后者，可以考虑在调度层面加一点形状感知的逻辑，让组出来的batch形状更集中。
6.4 监控指标的建立
最后说监控。BCG相关的关键指标有几个：桶命中率、桶数量、图捕获次数、图捕获耗时、各桶的显存占用。这些指标要持续监控，才能及时发现异常。
我一般会在服务里暴露这些指标，接到监控系统上。桶命中率下降、桶数量异常增长、图捕获耗时突增，这些都是需要关注的信号。尤其是桶数量，如果它一直涨，说明有形状在持续触发新桶，可能是某个上游逻辑在生成异常形状的请求。
7. 关于BCG值不值得上的个人判断
聊了这么多实现和调试细节，最后说点实在的：BCG不是所有场景都值得上。它的收益主要来自减少kernel launch开销，如果你的模型推理本身计算量很大，launch开销占比很小，那BCG带来的提升可能有限，反而增加了复杂度和显存占用。
我自己的判断标准是看两个数：一是单步去噪的kernel数量，如果超过几百个，BCG的收益会比较明显；二是请求的形状分布，如果形状很集中，桶好划分，BCG上起来就顺；如果形状很分散，桶管理成本高，可能得不偿失。
另外，BCG和具体的硬件、驱动版本也有关系。不同环境下图捕获的行为可能有差异，上线前一定要在目标环境里充分测试。我在调试过程中就遇到过某个算子在特定驱动版本下无法进图的情况，换了个版本就好了。这类问题没有通用解，只能靠实测。
整体来说，BCG是一套用复杂度换性能的机制，用得好能显著降低延迟抖动、提升吞吐，用不好就是给自己挖坑。建议先在测试环境把整套流程跑通，把桶划分、预热、监控都搭起来，确认收益确实存在，再考虑上生产。
