1. 为什么非要在SGLang里做Diffusion BCG先交代一下背景。SGLang这个框架早期主打LLM推理RadixAttention和连续批处理那套机制在长上下文、高并发场景确实好用。但很多人忽略了一点SGLang从2024年开始就在往多模态和扩散模型方向铺路ImageGen相关的接口和调度逻辑已经能支撑Diffusion Model的推理。而BCGBucketized Continuous Guidance出现在SGLang的Diffusion流程里本质上是为了解决扩散模型生成时引导强度不好精细控制的问题。先说BCG到底是什么。用过Stable Diffusion的朋友都知道CFGClassifier-Free Guidance它通过同时推理条件输入和无条件输入用两者的差异来控制生成结果对提示词的跟随程度。传统的CFG在一个推理步骤里做两次forward一次是有text condition的一次是没condition的然后按照guidance_scale这个标量把结果拉向条件方向。但CFG有个毛病guidance_scale是全局统一的从头到尾都用同一个强度。实际生成过程中早期去噪步骤需要强引导来确立构图后期步骤如果继续强引导就容易过曝、失真、色彩发闷。理想的做法是给不同timestep分配不同的引导强度比如前20%用高引导中段衰减后段再收一下。这就是Continuous Guidance的思路。BCG在此基础上做了bucket化处理把连续的引导曲线离散成若干区间同一个区间内的step共享一套引导参数张量并行下通信开销小、缓存友好、实现也简单得多。所以SGLang实现BCG的动机很清楚在保持diffusion推理吞吐优势的同时把CFG的控制粒度从整个生成过程细化到每个bucket区间让生成的画质和提示词跟随度都能上一个台阶。我在动手之前查了SGLang仓库发现它的Diffusion相关代码其实已经具备BCG的基础设施——比如BucketizedContinuousGuidance这个类在连续批处理的调度层占了一个坑但当时还是半成品状态很多分支没有实现。接下来的内容就是我把这块从占位符变成可用功能的过程包括实现思路、踩坑记录、调试链路的完整复盘。2. BCG的核心机制拆解guidance不是一路等强的节器是timestep本身。我画过一条衰减曲线横轴是去噪进度纵轴是引导强度然后按照bucket边界切成几段每一段内用常量或简单的线性插值。# 一个简化的bucket分配示意 # guidance_curve: 由用户指定的连续曲线比如从2.0衰减到1.2 # num_buckets: 例如8 # steps: 当前去噪总步数 def assign_bucket(timestep, steps, num_buckets): progress timestep / steps # 0~1早期progress小 bucket_id min(int(progress * num_buckets), num_buckets - 1) return bucket_id实际代码里SGLang不是这么简单的它要考虑cfg_scale在每个bucket的具体值、无条件token的batch组织方式以及在张量并行下怎么把这些tabulate好的参数同步到所有worker上。BCG和普通CFG最大的区别在这里普通CFG在SGLang的调度器看来就是一个batch里的两倍序列——条件prompt和无条件prompt一起进模型输出算diff后合并。但BCG不同同一个timestep内可能出现多个不同的cfg_scale值因为不同样本可能被分在不同的bucket。这就要求调度器在构造batch时不能把无条件token简单地拼在条件token后面而是得按bucket分组每个组用各自的guidance参数做后处理。这也是BCG实现里最容易出bug的地方我在下一节详细说。3. 张量并行下的bucket构造第一个真正的坎SGLang做diffusion推理时tp_size 1是标配场景模型太大单卡放不下。但这引入了一个问题张量并行时模型内部会切分attention和FFN的维度所有worker必须处理完全一致的batch结构。如果某个worker按bucketA处理样本另一个worker按bucketB处理输出拼接出来的结果一定是乱的。我开始实现时先选了最简单的方案一个batch内所有样本用同一个bucket参数。相当于把BCG退化成一个分段常数调度每次iteration只处理当前bucket段内的样本。这确实能跑通但有一个要命的瓶颈——如果某段bucket内的样本特别少而其他bucket样本很多整个batch就空转GPU利用率掉到惨不忍睹。后来我改成混合bucket方案也就是一个batch可以包含多个bucket的样本但这需要处理两个问题一是attention计算时的padding对齐二是所有TP worker上的bucket配置要一致。SGLang在LLM推理里有条经验就是凡是涉及调度决策的meta信息必须同步到所有TP worker。BCG的bucket指派也不例外。我在sglang的调度器里加了一个bucket plan的数据结构class BucketPlan: 记录当前batch内每个bucket的样本分布 def __init__(self, num_buckets, tp_size): self.num_buckets num_buckets self.tp_size tp_size self.bucket_record {} # seq_id - (bucket_id, guidance_scale) self.bucket_offset {} # local_batch内的起点 def compute_masked_attention(self): # 为不同bucket生成不同的attention mask避免串扰 pass这里有个很关键的细节不同bucket的样本虽然进了同一个batch但它们对应不同的guidance_scale所以在去噪步长和UNet输出处理上不能混用。我在model_worker端做后处理时用了torch的split按bucket边界把hidden state切回去再分别乘以各自的scale。guidance_scales [bucket_plan.bucket_record[sid].guidance_scale for sid in batch_ids] # hidden_states shape: [batch, seq_len, dim] # 按bucket分段处理 bucket_splits [bucket_plan.bucket_offset[i 1] - bucket_plan.bucket_offset[i] for i in range(num_buckets)] uncond_out, cond_out torch.split(hidden_states, bucket_splits, dim0) final_out uncond_out guidance_scales * (cond_out - uncond_out)这个split逻辑本身不难但我第一次跑就发现数值不对。排查到最后才发现问题不在split而在更前端的prompt拼接顺序——SGLang在处理无条件token时默认把它放在每个序列的最前面这和BCG的bucket编排冲突了。也就是说模型管线的input_embeds拼接逻辑需要把无条件token重新排布否则split出的位置跟你以为的bucket顺序完全对不上。这算是整个实现里第一个有含金量的坑。解决方法是在往tokenizer里塞prompt时不沿用默认的[uncond, cond]排列而是改成# 原始排列: [uncond_0, cond_0, uncond_1, cond_1, ...] # 改成: [uncond_0, uncond_1, ..., cond_0, cond_1, ...] # 这样split时按batch_size就能直接切出uncond和cond改完后张量并行的一致性也顺带解决了——因为所有worker都遵循同一套token排布规则。4. Debug完整链路从NotImplementedError到数值异常BCG实现过程中最耗时间的不是写代码是调试。我按时间顺序完整复盘一遍这个排查链路可以给以后做类似scheduler扩展的朋友参考。阶段一NotImplementedError占位符暴露的第一个设计缺陷我拿到SGLang源码后第一件事就在model_worker.py里搜BucketizedContinuousGuidance果不其然很多方法体就是一行raise NotImplementedError。这说明仓库里已经把接口占好了但逻辑还没补齐。我对照llama.cpp和Diffusers里CFG的实现思路先把这些占位方法补上最核心的是update_guidance_scales()和apply_guidance()。def update_guidance_scales(self, timestep, steps): # 按当前timestep动态更新bucket for seq_id in self.bucket_record: bucket_id assign_bucket(timestep, steps, self.num_buckets) self.bucket_record[seq_id].bucket_id bucket_id self.bucket_record[seq_id].guidance_scale self.bucket_cfg[bucket_id]这一步跑通不难很快就能消除NotImplementedError但只解决了能跑的问题距离跑对还差得远。阶段二全过程数值飘生成的图全糊第一版跑通后我拿SDXL测了一组固定prompt结果生成的图不仅糊而且颜色整体偏灰完全不像正常CFG的结果。我当时怀疑是guidance算错了方向——BCG的公式是uncond scale * (cond - uncond)如果把scale设成负数或者方向搞反就会往回拉。为了定位我在apply_guidance()前后各加了一次张量统计打印logging.info(pre-guidance norm: %.4f, hidden_states.norm(dim-1).mean()) logging.info(post-guidance norm: %.4f, guided.norm(dim-1).mean())然后拿一个简单case跑scale0.0时结果应等于uncond_outscale1.0时结果应等于cond_out。我分别测了scale0.0和scale1.0两个极端发现scale1.0的输出和纯condition跑出来的结果不一样这意味着split后的cond_out根本不是condition分支的输出。这下我把问题锁定到了token排列上。我在前面那节已经提过——SGLang默认把无条件token放在最前除非你显式修改input_embeds的顺序否则torch.split按batch_size切分时切出来的前半段是cond导致的计算完全错位。修复很简单改排序逻辑即可。数值上前后对比就很明显了上图是修复前的飘值下图是修复后稳定在预期范围内。阶段三随机性导致偶发崩溃报错在attention mask排序修好后大部分case能正常出图了但偶发崩溃报错信息指向attention的mask维度不匹配。我看了一下堆栈问题出在bucket_record的更新时机——我在model_worker的step内更新了bucket但调度器那边还在用旧bucket做padding/mask两者不同步。具体来说SGLang调度器都会在schedule()阶段就做好attention mask和位置编码。像我这样在worker内部改bucket等于已经发出去的batch在中间被换了个mask。正确的做法是把bucket更新放到调度器的schedule()里跟token batch同步构建。我重构了代码把BucketPlan的修改统一收敛到调度器侧def schedule_buckets(self, waiting_queue, running_batch): # 1. 给等待队列中所有seq分配bucket # 2. 对running batch中已经落桶的seq做动态迁移 # 3. 生成bucket plan并同步到tp workers return self.update_bucket_plan()这里动态迁移是个很有意思的点BCG允许同一个序列在生成过程中被分到不同bucketbucket切换时不需要打断量化只需在下一个schedule周期里换mask和guidance值。SGLang的连续批处理能轻松支持这种操作这也是BCG能落地在SGLang里的天然优势。阶段四精度衰减fp16下BCG数值和成熟库差1e-3大部分功能问题解决后我做了一次精度对比拿Diffusers官方的CFG结果当基准。同一prompt、同一seedSGLang的BCG输出在0到5步的中间latent上和基准有大约1e-3级别的差异。1e-3在图像像素域看不出什么但扩散模型是迭代过程误差会累积最后可能变成肉眼可见的色偏。到底是哪里引入的我先怀疑是guidance_scale在fp16下的精度损失毕竟计算cond - uncond时两个大数相减再乘以大scale精度会被放大。于是我把apply_guidance里的计算强制切成fp32torch.no_grad() def apply_guidance_bf16safe(self, hidden_states, scale): orig_dtype hidden_states.dtype hidden_states hidden_states.to(torch.float32) # 先算uncond和cond的分量再做线性组合 hidden_states uncond_out scale * (cond_out - uncond_out) return hidden_states.to(orig_dtype)切到fp32后误差缩小了一个数量级但还不够。继续深挖发现误差主要来源不是guidance计算而是预处理环节的position embedding——SGLang对timestep做embedding时在fp16下丢了一些精度。timestep embedding是整个UNet的条件输入提前精度损失会影响所有层。这个就没有简单修法了因为改timestep embedding的影响面太大。我采取的折中方案是在scheduler阶段提前把timestep转成fp32并缓存等输入embedding层时再转回fp16。虽然不能做到和Diffusers完全一致但误差从1e-3降到了2e-4左右肉眼已经分不出区别。5. 性能实测与边界case处理功能正确之后最关心的就是性能。BCG的价值在于能精细控制引导但如果性能比普通CFG掉太多就没人愿意用。我拿SGLang在单机8卡A100上跑了一组对比模型是SDXL分辨率1024x1024steps50。方案吞吐img/s显存占用GB备注普通CFG (无BCG)3.4262.1基准BCG 4 buckets3.1164.8约-9%吞吐可接受BCG 8 buckets2.8967.3约-15%吞吐控制更细BCG 16 buckets2.5472.6吞吐下降明显收益边际递减从这个数据能看出一个规律bucket数量从4涨到8时控制精度提升明显吞吐下降还好从8涨到16吞吐损失接近12%但生成质量的提升几乎看不出来。实际使用我建议默认设6到8个bucket就够除非你对引导曲线有特殊需求否则16个纯属浪费算力。BCG相对普通CFG还有另一个优势它能直接修改引导曲线的形状。我试过几种曲线线性衰减scale从2.0到1.2适合通用生成前高后平前20%步强引导定构图后80%弱引导保自然中部峰值中间阶段加强引导适合需要中段细节丰富的场景实测下来前高后平在中长steps设置下质量最好线条干净、色彩不过饱和。线性衰减最稳各种题材都不会翻车。中部峰值容易在某些prompt下出现结构性的扭曲慎用。边界case方面我整理了三个容易出问题的点边界case一bucket切换不能破坏KV Cache的连续性。SGLang的RadixAttention在LLM里主要服务prefill但diffusion场景下也有类似机制。bucket切换意味着attention mask变了如果KV Cache没跟着重新hash就会出现错位。我后来在bucket切换时强制触发一次cache reset代价是多做一次prefill但保证正确性。边界case二无条件token的长度不一致。有的prompt很长无条件token可能是空字符串长度只有1和条件token长度差异巨大。BCG在按bucket分割时如果无条件token长度不是统一值torch.split会直接炸。我在input_embeds拼接前统一做了padding把无条件token长度对齐到条件token的最大长度。边界case三不同TP大小下的bucket分布一致性。如果你在tp_size2跑通了BCG换到tp_size4发现推理结果不一样那八成是bucket plan没有广播给所有worker。SGLang在TP场景下有一个broadcast机制我在init_model_worker时把bucket配置表广播了一遍包括num_buckets、bucket_cfg数组确保所有worker的bucket边界完全一致。6. 实现BCG过程中学到的几件事整个调试过程下来有几个体会让我印象特别深写出来给后来者参考。第一SGLang这种框架里调度器是真正的核心。你可能会觉得BCG不过是在模型前向里加个参数改了也没什么了不起。但实际上SGLang的调度器掌控着token排布、padding方法、attention mask、Cache管理等所有环节BCG的任何一个改动如果不经过调度器就会寸步难行。我一开始图省事直接在worker内部处理bucket结果绕过了调度器对mask的管理崩溃报错一个接一个。最终方案就是把所有bucket决策全部收敛到调度器再由调度器统一生成执行plan。这大概也是SGLang设计者的本意——扩展新功能最好遵循框架的既有架构思路而不是另起炉灶。第二Diffusion的调试比LLM调试更依赖中间态的可视化。LLM输出是文本错了能直接看到tokenDiffusion的中间latent是张量光看数字很难判断错在哪。我后来加了一个debug钩子把每个step的latent dump成图片配合scale0.0和scale1.0两个极端case的对比图一眼就能看出guidance方向是否正确、token排列是否错位。这比看norm、看方差都直观得多。第三精度问题的定位要学会二分法。fp16下的误差来源很多tokenizer、embedding、attention、bucket切换、guidance计算都有可能。我先用最粗暴的方式锁定模块——把整个apply_guidance强制换成fp32误差变化判断这个模块是否主因再用同样的方式测embedding、测attention最后锁定了timestep embedding。这个方法听起来无聊但在多层管线debug时无比高效。第四如果只做一个优化我会优先考虑让bucket数量可配置且默认值保守。从我的实测来看4到8个bucket是甜区超过8个带来的质量提升微乎其微却会明显拖慢吞吐。跟业务方沟通时我一般会建议先用8个bucket跑一版看效果再决定要不要降到4个——留足调优空间别一上来就顶着上限跑。最后再分享一个小技巧用BCG时把第一个bucket的guidance设得稍微高一点比如1.8~2.2生成结果的初始构图会更稳定。这个区间在SDXL和SD3上都试过没有翻车风险。希望这篇记录能帮到正在研究SGLang diffusion路径的朋友BCG这套机制在扩散模型可控生成里还是挺有潜力的一个方向。
