深入解析SGLang Diffusion的BCG实现与调试

在扩散模型的图像/视频生成任务中，多次重复的DiT forward调用会产生大量短耗时的GPU kernel，如果CPU kernel launch速度跟不上，就会导致GPU空闲间隙，显著降低整体利用率。本文将详细介绍SGLang框架中Diffusion任务的Breakable CUDA Graph（BCG）实现细节，以及在Prompt适配、Graph重放和性能调优中遇到的典型问题与解决方案。

一、Breakable CUDA Graph的实现逻辑

扩散模型的生成过程需要多次执行DiT前向传播，每个前向传播中包含投影、归一化、RoPE、残差连接、MLP等多个算子，这些算子会生成大量短耗时的kernel。以SANA1.5模型为例，在H200设备上的Eager模式下GPU利用率仅为28.5%，开启BCG后可以提升至93.3%，核心优化点就是减少CPU的kernel launch开销。

不过扩散模型中存在一些无法直接纳入CUDA Graph的操作，比如可变长度注意力（varlen attention）需要根据当前请求的mask动态生成cu_seqlens和indices，还有部分通信和动态分支逻辑需要在运行时处理。针对这个问题，SGLang在公共DiT注意力入口添加了@eager_on_graph装饰器：当捕获Graph到该位置时，先结束当前正在捕获的Graph，使用Eager模式执行注意力算子，之后再继续捕获后续的Graph segment。这样单个DiT forward就被拆分为多个独立的Graph片段，执行时按照“segment 0 → eager attention → segment 1 → …”的顺序重放，无需修改模型原有的forward代码结构。

为了匹配输入对应的Graph，系统会根据调用的kwargs生成唯一的signature：张量会记录形状和数据类型，Python的int、bool、string等常量会记录具体取值，list、tuple、dict等结构会递归处理，其他对象则按类型和标识区分。因此经过Prompt Padding后的输入，如果形状、数据类型以及其他参数一致，就可以复用同一个已捕获的Graph。

每次Graph重放前，需要将当前输入拷贝到捕获阶段分配的固定缓存中，再执行Graph。核心简化代码如下：

for buf, live in zip(entry.static_leaves, live_leaves):
    buf.copy_(live, non_blocking=True)
entry.graph.replay()
return clone(entry.output)

这里需要特别注意输出的拷贝操作：Graph的输出缓存会在下次重放时被覆盖，而CFG的正负分支可能复用同一块缓存。如果直接返回原始缓存，后一个分支执行完成后会覆盖前一个分支的结果，因此必须复制一份输出再返回。

目前BCG仅在warmup预热阶段完成捕获，服务启动后根据signature查找对应的Graph，未匹配到则回退到Eager模式。使用时需要注意，分辨率和视频帧数必须与warmup阶段一致，例如1024×1024的Graph不能用于1280×768的输入，17帧的视频模型Graph也无法适配121帧的请求。Prompt长度可以通过Padding处理，后续会详细说明适配策略。

二、不同扩散模型的Prompt Padding策略

通用Padding逻辑

以SANA模型为例，可以将19-token和47-token的Prompt都补全到64-token，新增位置的mask设为0，不参与注意力计算。这样两个输入的张量形状一致，可以复用同一个Graph。通用的Padding工具会在显式提供attention mask时自动补齐长度，默认的bucket为64、128、256、512、1024。如果输入长度超过最大bucket，则保留原长度，若该长度未经过warmup捕获，则回退到Eager模式。

Qwen-Image的特殊适配

Qwen-Image需要同步处理mask、文本RoPE缓存和txt_seq_lens字段，仅对embedding做Padding会导致这些字段长度不一致。其中txt_seq_lens需要设置为bucket值，避免该主机常量随Prompt长度变化，真实的文本范围由mask区分。例如将19个token补到64时，除了embedding外，还需要为64个位置准备对应的mask和文本位置编码，同时将txt_seq_lens设为64，而前19个位置标记为有效，后续补全的位置标记为0。

当更换Prompt后，Qwen-Image需要重新计算varlen attention的metadata。因为静态mask缓存的地址不变但内容已更新，如果仅使用data_ptr()作为缓存key，会复用旧的cu_seqlens数据。解决方案是使用DynamicVarlenMaskMeta工具，每次重放时根据当前mask生成metadata，第一个注意力层处理后复用该结果，下次重放时再更新。

Z-Image的Padding限制

Z-Image不能直接使用通用Padding逻辑。最初将不同Caption补到同一个bucket后，BCG和Eager模式输出的PSNR仅为21.24dB，生成的图片存在明显差异。逐层对比后发现，第一个差异出现在context_refiner.0，也就是第一个Caption自注意力层。

这是因为Z-Image本身会将Caption补到32的倍数，其中的learned pad token会作为寄存器参与注意力计算。如果继续补到更大的bucket，参与注意力的token数量会发生变化，导致生成结果不一致。修复方案是保留输入的原生Caption长度，相同长度的输入可以复用Graph，未经过warmup的长度则回退到Eager模式，修改后输出结果与Eager模式完全一致。

MiniMax-H3的Packed Sequence适配

MiniMax-H3将文本、视频/图像和音频数据整合在同一个packed sequence中，可以对Prompt embedding做Padding，但不能扩充主序列x，否则会改变sequence-parallel的行分区和GEMM形状，无法保证计算结果与Eager模式一致。因此除了文本bucket相同外，还要求主packed sequence属于同一个64行对齐组，跨组的请求回退到Eager模式。

固定文本长度的模型

LTX-2/LTX-2.3、LongCat-Image和SANA-Video的文本编码已经是固定长度，分别为1024、512和默认300。不同长度的Prompt进入DiT时，文本形状已经一致，无需再次补齐。其中SANA-Video最初会经过通用Padding工具，将[1, 300, 2304]的输入继续补到512，导致每层交叉注意力多处理212个位置，同时多捕获了一组不必要的Graph。后续添加了专用处理逻辑，确认模型类型和输入长度为300时直接返回原输入，仅当用户修改max_sequence_length时才回到通用Padding流程，默认配置下无需额外设置--bcg-text-buckets 300。

三、开发过程中遇到的典型问题与解决方案

3.1 Z-Image Graph重放时的非法访存问题

BCG的注意力算子在Graph外执行，捕获阶段注意力输出的张量地址会被后续的Graph segment使用。重放阶段注意力会生成新的输出，需要将结果拷贝到捕获阶段的缓存中，再执行后续的segment。

早期实现中将该Eager输出转为弱引用（weak-ref），由于该张量是在Graph外分配的，没有强引用时可能被释放，导致后续segment读取该地址时出现非法访存。在Z-Image-Turbo上遇到的段错误就是这个原因，修复方案是将eager break output改回强引用，并添加CUDA回归测试。

Z-Image的RoPE和注意力metadata缓存也存在类似问题：这些缓存原本仅保存一份，捕获新的bucket时会替换旧值，但之前捕获的Graph仍然会读取旧张量的设备地址。解决方案是在活跃的Graph捕获期间保留这些缓存张量，避免在Graph仍需使用时被释放。

3.2 LTX-2的Warmup Shape与真实请求不一致问题

LTX-2使用通用视频warmup时，为了减少启动时间仅跑17帧，但真实请求通常为121帧，两者的形状不一致，无法复用Graph。同时warmup阶段会自动构造一张合成图像，LTX-2因此执行image-conditioned分支，该Graph无法用于纯文本到视频的请求。

另外还存在调用路径问题：LTX-2的two-stage denoise直接调用step.current_model(...)，没有经过通用DenoisingStage的BCG hook，因此即使模型被加入白名单，实际执行时也不会调用BCG runner。这些问题的修复可以参考对应PR。

如果启动日志中已经有captured标记，但真实请求未使用BCG，可以先检查模型调用是否经过runner，再对比warmup和真实请求的帧数、分辨率以及条件输入是否一致。

3.3 SANA的Signature生成开销问题

SANA接入BCG时，predict_noise每次forward都会构造一个包含72000个元素的嵌套None list，作为mask_strategy参数传给DiT。实际上没有任何DiT实现读取该参数，仅通过**kwargs接收了它。

在Eager模式下构造该列表约需要1.1ms，而BCG模式下还需要递归遍历整个列表生成signature，导致denoise耗时从0.67s增加到2.63s。该开销发生在Python侧，仅通过GPU kernel耗时难以发现，需要单独检查signature生成过程。

修复方案是删除该无用参数，修改后在H200、1024²配置下测试，Eager denoise耗时为699.2ms，BCG模式耗时降至408.1ms，端到端时间从0.821s降到0.608s。

3.4 GLM-Image精度对比中的随机性问题

GLM-Image在DiT前还有一段sampled AR prior过程，使用了do_sample=True参数。如果Eager和BCG模式分别生成的prior不同，即使DiT计算结果一致，最终生成的图片也可能不同。

修复方案是为prior绑定request seed，并额外执行一次same-prior重放：先保存采样结果，再让Eager和BCG使用同一份prior，这样两张图片的每个像素值都完全一致。需要注意检查DiT之前的所有随机过程，带有Prompt重写或随机条件输入的流水线也可以通过保存中间结果的方式进行精度对比。

四、BCG的性能测试效果

以下是部分B200设备上的测试数据，统计的是warmup之后的denoise延迟：

以SANA模型为例，5个profiled timesteps中，runtime launch次数从8412次减少到132次，Kernel之间的空闲间隙也明显减少。

BCG的收益与模型分辨率强相关：LTX-2在768×512×121、2×H200 CFG parallel配置下，GPU利用率从27.9%提升到96.2%；当分辨率换成1920×1088时，Eager模式的GPU利用率已经达到96.8%，BCG的端到端差异仅约0.2%。这种情况下launch开销占比已经很小，需要考虑GEMM、注意力或Kernel融合的进一步优化。

五、BCG的启动与日志检查

以SANA1.5为例，可以使用以下命令启动BCG模式：

python3 -m sglang.multimodal_gen.runtime.entrypoints.cli.main serve \
  --model-path Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers \
  --num-gpus 1 \
  --enable-breakable-cuda-graph \
  --warmup-resolutions 1024x1024 \
  --bcg-text-buckets 64 128 256 512 1024 \
  --enable-torch-compile false

需要注意的是，当前版本开启BCG后会跳过torch.compile，同时关闭Cache-DiT。在进行Eager和BCG的性能对比时，需要显式设置两边的compile和offload参数，因为--performance-mode speed的默认行为已经修改，直接使用旧命令可能导致测试配置不一致。

启动后可以先检查日志中是否包含[Diffusion BCG] captured标记，然后发送几个长短不同的Prompt，检查是否出现serving signature MISSED、capture failed或者Falling back to diffusers backend等日志。正常情况下，warmup和每次请求后的累计捕获数量应该保持一致，例如[5,5,5,5,5]。如果所有请求都使用Eager模式，该数量不会变化，因此需要同时查看miss日志。

相关代码位于diffusion BCG runner和同目录的模型专用Padding工具中，具体路径可参考项目仓库。