文章摘要
本文介绍借助Agent辅助，针对SGLang框架完成Qwen-Image、FLUX.2、Wan、SANA Video等多个图像视频扩散模型的算子内核优化，相关代码已合入指定版本的SGLang主分支。文中梳理了扩散内核的架构设计、精度验证与质量分级机制，针对不同模型瓶颈完成融合、量化、数据搬运等优化，获得不同幅度的端到端性能提升，同时记录无效优化经验避免重复踩坑。

基于Agent的SGLang扩散模型Kernel优化实践：Qwen-Image、FLUX.2和视频模型

在近期的工作中，我们通过Agent辅助完成了多组扩散模型Kernel优化，并将其合入SGLang主分支。实际优化过程中我们发现一个常见误区：单个Kernel在微基准测试中表现优异，并不意味着接入完整模型后能带来同等收益。在扩散模型场景中，单次不起眼的`contiguous`或`cat`操作仅消耗几十微秒，但经过数十个Transformer Block、50次去噪步骤以及CFG双分支的累积，最终会占用相当可观的计算资源。

本文将完整记录基于Agent完成Qwen-Image、FLUX.2、Wan和SANA Video优化的全过程，包括如何从真实模型中定位性能瓶颈、Agent辅助编写Kernel的方法，以及优化过程中遇到的各类问题。本文基于2026年9月2日的SGLang主分支，Commit哈希为`99b910955377375a1385122680e0daefcf706f79`，该版本共注册43个扩散算子，对应47份实现，覆盖DiT Block、VAE以及序列并行的数据搬运逻辑。相关代码可参考：GitHub仓库。需要说明的是，后文出现的性能数据均来自不同模型和硬件环境，仅用于说明对应优化的有效性，不可直接累加对比。

SGLang扩散Kernel的现有架构

当前主分支的扩散Kernel目录按算子功能划分为多个子模块：

  • norm: 包含LayerNorm、RMSNorm、AdaLN以及残差归一化组合操作
  • modulate: 实现Scale/Shift变换、时间步嵌入以及逐Token调制逻辑
  • rope: 支持RoPE旋转位置编码、QK归一化以及QKV收尾处理
  • activation: 包含SiLU、SwiGLU激活函数以及后续量化逻辑
  • attention: 实现SANA-WM GDN、SANA Video线性注意力等专用注意力算子
  • layout: 包含Ulysses QKV重排、USP布局转换、变长序列打包以及Wan因果缓存
  • sites: 负责模型侧的算子挂载、回退逻辑以及质量控制
  • ext: 存放Rasterizer等需要单独编译的扩展算子

此外还有`python/sglang/kernels/kda_kernels`目录，该目录记录Kernel的来源，支持Triton和JIT CUDA等多种实现方式，当前已支持Qwen-Image的归一化算子、Cosmos3的因果Conv3D、LTX2的残差门控和QK归一化拆分RoPE等逻辑。

对外使用时统一通过以下入口导入：

from sglang.kernels.ops.diffusion import some_kernel

该模块采用延迟加载设计，未使用CuTe DSL、FlyDSL或特定扩展的模型在启动时不会加载对应依赖。这一细节看似与性能无关，但对推理框架的稳定性至关重要——我们曾遇到过Kernel本身没问题，但导入阶段导致不相关模型报错的情况。

在评估扩散Kernel是否可以合入时，我们通常按照以下顺序检查：

  1. 支持的GPU型号、数据类型、形状以及步幅条件
  2. 不支持的输入场景是否能回退到原始PyTorch实现
  3. 数值精度是否达到bit-exact，还是仅支持`quality=extra-high/high`等级别
  4. Kernel在真实模型中是否成功命中
  5. 最后才验证端到端性能是否有提升

这一顺序并非随意设定，尤其是第4点尤为重要：Triton或JIT CUDA编译失败后通常会静默回退到原始实现，如果仅通过程序运行状态和最终耗时判断，很容易误以为新Kernel已经生效，实际仍在使用旧版本实现。

扩散Kernel对数值精度的敏感性

数学等价不代表结果完全一致

扩散模型会重复执行数十次去噪步骤，每一步的输出都会作为下一步的输入，因此单次BF16舍入位置的微小变化，都可能在后续步骤中不断累积误差。

以常见的AdaLN计算为例：

hidden_states = layer_norm(hidden_states)
hidden_states = hidden_states * (1 + scale) + shift

最直接的融合方式是将LayerNorm和后续的乘加操作合并为FP32计算，最后再写回BF16格式，从数学角度看这一实现没有问题，甚至精度看起来更高。但PyTorch原生实现会在LayerNorm输出和后续乘加的特定位置执行BF16舍入，如果融合Kernel没有保留这些边界细节，最终结果将无法达到bit-exact。

我们曾在ERNIE-Image上测试过普通的FP32单通道归一化融合，单个Kernel的输出结果看似与参考实现接近，但在完成50次去噪步骤后，最终生成的图片与参考版本的PSNR仅为18.83dB。这一结果直接推动我们重写了bit-exact的LayerNorm实现路径。

当前SGLang中的JIT CUDA实现会尽可能复现PyTorch 2.11的`vectorized_layer_norm_kernel`执行顺序，包括128线程Welford计算、固定的Shuffle Tree、`div.rn`以及对应的RSQ路径，AdaLN后续的每一步也会在与PyTorch完全相同的位置转回BF16格式。相关实现可参考：layernorm_modulate_triton.py

首次真实请求的额外验证

即使单元测试通过，在实际模型运行时仍会执行一次`torch.equal`检查。这是因为Torch版本、CUDA驱动和编译器版本都可能影响底层实现，静态测试无法覆盖所有运行环境。

实际逻辑可以简化为：

if compiling_or_capturing() and not verified:
    return reference_path(x)
candidate = fused_path(x)
reference = reference_path(x)
return gate.accept_or_fallback(candidate, reference)

首次Eager模式请求会同时运行融合Kernel和参考实现，如果`torch.equal`验证通过，后续请求将继续使用融合Kernel；如果结果不一致，该模块或该形状签名将永久回退到参考实现。在CUDA Graph捕获期间无法插入主机同步操作，因此未经过验证的路径会直接使用原始实现。对应代码可参考：bitexact_gate.py

质量分级机制

部分融合操作确实会改变半精度运算顺序，但经过模型级验证后仍可安全使用。当前SGLang扩散模块提供三个质量级别：

  • lossless: 基于参考路径的bit-exact替换，为默认级别
  • extra-high: 在lossless基础上启用经过质量验证的DiT/VAE Kernel融合
  • high: 继续启用模型自主管理的Cache-DiT或低精度解码等优化

非bit-exact融合并非每层单独开关，SGLang会在Batch边界检查整个Transformer模块，只要其中一个位点不满足条件，整组融合都不会挂载，避免同一请求中部分模块使用新路径、部分模块使用旧路径的情况。

需要注意区分量化Checkpoint自身的实现逻辑：当用户加载FP8或NVFP4 Checkpoint后，`quality=lossless`的参考实现是该Checkpoint原生的量化路径，而非将其恢复为BF16模型。例如Qwen-Image的Bias Absorption会调整BF16舍入点，目前仅在验证过的SM103架构上启用，不能称为Pixel-Exact。

此外，`quality=extra-high/high`的DiT融合目前无法与Breakable CUDA Graph直接混用，因为BCG预热阶段捕获的是lossless分支，请求到来后挂载新位点会导致Graph Replay跳过新逻辑。仅修改VAE解码的质量路径不受此限制，因为BCG不会捕获VAE模块。

Qwen-Image FP8和NVFP4优化

合并QKV投影时的数值问题

Qwen-Image原生的FP8注意力模块中，图像和文本分支共有6个Q/K/V投影操作，每个投影前还需要执行一次FP8量化。GEMM本身已经比较高效，但Profile数据显示，GEMM前后的量化操作、QK RMSNorm、RoPE和`cat`操作占用了大量时间。

最初我们尝试将QKV权重统一到三份权重中的最大Scale，这样合并操作最简单，性能提升也明显，但最终生成图片的质量未达标，SSIM仅为0.883575，LPIPS为0.075531。最终我们采用保留每个分片自身FP8 Scale的方案，将6次GEMM和6次激活量化缩减为2次GEMM和2次量化，中间通过Strided View拆分结果，避免额外的Copy操作。

我们还在GEMM后添加了一个Blackwell JIT Kernel，一次性完成以下操作：

  1. Q/K RMSNorm
  2. Interleaved RoPE旋转位置编码
  3. V分支复制
  4. 文本和图像Token写入联合QKV缓冲区

优化前这些步骤会生成多个中间张量并触发多次Kernel Launch，融合后GEMM输出无需先写入临时张量再被下一个Kernel读取。相关PR可参考：#37123

在单张GB300显卡、1024x1024分辨率、两次去噪步骤的测试环境中，Profile数据显示：

  • 静态FP8量化次数从1320次降至840次
  • QK Norm + RoPE操作从240次降至0次
  • QKV Cat操作从360次降至0次
  • 新增120次融合QKV Epilogue Kernel
  • 总Kernel Launch次数从4989次降至3309次
  • Profile窗口内GPU耗时从105.874ms降至98.955ms

端到端测试配置为单张GB300、1024x1024分辨率、50个去噪步骤、CFG系数4，去噪步骤耗时从310.168ms降至245.247ms，请求总时间从15.635秒降至12.391秒，提升20.75%，最终生成的PNG图片逐像素一致。

将FP8量化融合到前序算子

完成QKV优化后，追踪数据中仍存在大量单独的FP8量化Kernel，这些Kernel的输入通常刚被LayerNorm或残差归一化完整遍历过一次，随后又需要从显存重新读取以执行量化。

我们在PR #37156中实现了两组Producer融合：

  • LayerNorm/AdaLN同时输出BF16和静态E4M3格式
  • Gated Residual + LayerNorm/AdaLN同时输出BF16和静态E4M3格式

BF16结果继续用于残差等分支，FP8结果直接送入下一个GEMM。这里有一个容易被忽略的细节：Scale计算必须使用`div.full.f32`，CUDA中的快速Reciprocal操作对于某些真实Checkpoint的Scale值会相差1个FP32 ULP，刚好可能跨越E4M3的舍入中点，最终导致量化结果相差1个bit。

单独测试该PR时，GB300端到端性能提升21.13%，与前面的QKV投影优化组合后，相对于原始主分支提升约33.9%——这一结果并非两个PR的百分比直接相加，而是组合路径经过完整测试后的实际效果。相关PR可参考：#37156

NVFP4路径也实现了类似优化，PR #37129将以下链路合并为单个Kernel：

residual = residual + gate * (attention + bias)
hidden_states = layer_norm(residual)
hidden_states = hidden_states * (1 + scale) + shift
hidden_states_fp4, block_scale = fp4_quantize(hidden_states)

最终Kernel直接写入E2M1打包值和128x4 Swizzled E4M3 Scale，对应链路的GPU耗时下降23.6%，但完整请求仅提升1.19%。这一结果符合预期：目标链路在整个Pipeline中的占比本身不高，即使局部耗时下降23.6%，折算到端到端性能提升也较为有限。相关PR可参考：#37129

CFG调制缓存

Qwen-Image的条件和非条件分支是串行执行的，两个分支共享时间步调节点，原生实现会在两个分支中重复计算图像/文本调制投影。

我们在PR #37090中添加了请求内的一次性缓存，第一个CFG分支写入结果，第二个分支使用后立即清空。缓存会检查输入Identity/版本、数据类型和设备，训练、编译、BCG以及捕获场景仍使用原始路径。

两次去噪步骤的Profile数据显示，量化GEMM次数减少224次，普通`addmm`操作减少16次，相当于每个去噪步骤减少120次投影计算。GB300 NVFP4端到端性能提升5.72%，输出结果保持bit-exact。相关PR可参考：#37090

Bias Absorption并非无损优化

Qwen-Image输出投影后存在单独的Bias Add操作，我们在PR #37116中将该Bias延后，在下一次残差/归一化操作中统一处理，GB300 NVFP4端到端性能提升6.07%。

该优化改变了BF16舍入位置，最终验证结果为SSIM 0.950545、PSNR 29.262dB、LPIPS 0.042544，因此当前代码仅在SM103架构上启用。我们认为必须明确标注数值口径，避免将“视觉效果正常”误写为“结果完全一致”。相关PR可参考：#37116

FLUX.2量化Kernel优化

FP8 Producer、QKV和Token Cat优化

FLUX.2的ModelOpt FP8路径与Qwen-Image类似，GEMM之间穿插了大量量化和数据整理操作。我们在PR #37162中完成了以下优化：

  1. 为80处Norm/Activation Producer增加FP8输出
  2. 合并QKV投影，保留原有Scale
  3. 将QK Norm、RoPE和联合QKV打包合并为单个Epilogue Kernel
  4. 将注意力/MLP分支的Token Cat和FP8量化合并

相关PR可参考：#37162

在GB300、1024x1024分辨率、5个活跃时间步的测试环境中，Profile数据显示：

  • 静态FP8量化次数从1196次降至199次
  • Cat操作从836次降至386次
  • 旧的QK Norm/RoPE操作从96次降至0次
  • 新增48次联合Epilogue Kernel和309次Token-Cat + Quant操作
  • 总Launch次数从7335次降至5878次，减少19.9%
  • GPU Kernel耗时从799.555ms降至782.259ms，仅下降2.16%

这一结果符合预期：模型中仍存在大量GEMM和注意力操作，减少Kernel Launch次数并不等同于同等比例的GPU耗时下降。在50步常驻基准测试中，端到端耗时从6847.800ms降至6628.402ms，提升3.20%，运行时峰值保留内存减少438MB，输出结果保持Pixel-Exact。

NVFP4的三组优化

针对FLUX.2 NVFP4路径，我们后续又完成了三组优化：

第一组是门控残差归一化，将当前块延后的残差更新合并到下一次归一化/调制操作中，相关Kernel耗时下降21.85%，端到端提升0.71%，结果保持bit-exact。PR链接：#37112

第二组是Token Cat + NVFP4量化，FLUX.2单流块会将6144维注意力输出和18432维MLP输出拼接为24576维张量，再量化后送入输出投影。新的SM103 Kernel直接从两个输入读取数据，写入E2M1打包值和128x4 Swizzled Scale，无需先生成BF16拼接张量。目标链路GPU耗时下降58.63%，端到端提升2.39%，最终打包字节完全一致。PR链接：#37141

第三组是FC1 + SwiGLU + FC2输入量化，该路径复用了现有的CuTe DSL NVFP4 GEMM Epilogue，不再生成BF16的SwiGLU中间结果。端到端提升3.32%，但并非bit-exact，因此当前主分支将其放在`quality=extra-high/high`级别。质量测试结果为SSIM 0.956636、PSNR 28.908dB、LPIPS 0.027606。需要注意的是，该优化导致峰值显存增加1386MB，原因是回退布局需要额外保存一份权重。在GB300显存充足的场景下可以换取3.32%的性能提升，但如果部署更关注并发容量，则未必划算。PR链接：#37096

Wan视频模型的数据搬运优化

8GB临时张量问题

Wan2.2-TI2V的逐Token调制原生实现包含以下代码：

(table.unsqueeze(0) + temb.float()).chunk(6, dim=2)

在704p分辨率、121帧的输入场景中，该代码会先生成约8GB的FP32临时张量，随后返回6个步幅较大的视图。下游使用这些视图时还需要执行`contiguous`操作，在5步追踪测试中，这部分逻辑加上后续的具体化操作占每个步骤约14%的耗时。

我们在PR #34584中使用Triton Kernel直接写入`(6, B, S, D)`布局，输入仅读取一次，6份结果均为连续存储，后续的`contiguous`操作自然消失。此外我们还将每个块重复构造的RoPE缓存移至外部。

该优化未改变Float32加法顺序，最终结果保持bit-exact。50步测试中，H100的去噪耗时从56.53秒降至49.10秒，提升13.1%；H200从54.54秒降至47.66秒，提升12.6%。该测试中Eager模式甚至比编译模式更快，因为真正的瓶颈是巨大的中间张量和数据搬运，而非Python本身的开销。相关PR可参考：#34584

Wan因果VAE解码优化

在短步数视频模型中，VAE解码通常是性能瓶颈。LongLive2在H200上的完整Pipeline耗时4.513秒，其中解码占2.802秒，比例高达62%。进一步分析Profile数据，真正的Conv操作仅约1.26秒，周围的Clone、Cat、Pad、布局转换、`repeat_interleave`和Shortcut Add操作占用了大量时间。

原生的因果Conv输入处理依次执行：

cache = hidden_states[:, :, -CACHE_T:].clone()
hidden_states = torch.cat([old_cache, hidden_states], dim=2)
hidden_states = F.pad(hidden_states, padding)
hidden_states = hidden_states.contiguous(memory_format=torch.channels_last_3d)

我们在PR #34125中新增了两个Kernel：

  • cat_pad_channels_last_3d: 一次性完成缓存+隐藏状态拼接、Padding、NDHWC格式写出，并顺便生成下一个Chunk使用的紧凑缓存
  • dup_up3d_add: 将`repeat_interleave + permute().contiguous() + add`合并为单次Gather + Add操作

LongLive2无损解码窗口的Profile数据显示：

  • `aten::copy_`从513ms降至151ms
  • `fill_`从60ms降至0ms
  • `cat`从39ms降至6ms
  • `DupUp3D + Shortcut Add`从175ms降至50ms

新的Cat-Pad Kernel占用155ms，按照完整事件名称统计，总GPU工作量减少约470ms，与解码Wall Time减少479ms基本吻合。最终解码耗时从2.802秒降至2.323秒，提升17.1%；完整Pipeline耗时从4.513秒降至4.031秒，提升10.7%；峰值显存从49.6GB降至46.1GB，无损路径的帧流MD5与主分支一致。相关PR可参考：#34125

Ulysses QKV仅执行一次All-to-All

序列并行中还有一个典型问题：原生的Ulysses注意力会分别对Q、K、V执行一次All-to-All通信，每次通信前后还需要布局转换。继续优化单个布局转换Kernel虽然有收益，但三次Collective操作本身仍存在。

我们在PR #33667中将QKV写入Destination-Major打包缓冲区，仅执行一次All-to-All通信。为避免每个步骤重复申请缓冲区，我们还添加了可复用的暂存缓冲区。

H200x4的微基准测试显示性能提升1.28-1.31倍，打包Kernel提升约2.7倍。在Wan2.2 480p模型中，去噪步骤提升3.4%，端到端提升2.1%，输出逐帧一致。相关PR可参考：#33667

SANA Video线性注意力优化

SANA Video的线性注意力包含两次GEMM操作，原生实现会先将输入转换为FP32格式，第一段GEMM使用FP32输入和输出，随后执行第二段GEMM。

我们在PR #35728中将第一段GEMM修改为BF16输入、FP32累加和FP32输出，第二段GEMM仍保留FP32精度。微基准测试从0.7352ms降至0.3898ms，提升1.89倍。

该改动会影响累积精度，因此仅在`quality=extra-high/high`级别启用。在B300、832x480分辨率、81帧、50个去噪步骤的测试中，端到端耗时从53.402秒降至49.542秒，提升7.23%；去噪阶段耗时提升7.58%。81帧的SSIM均值/最小值为0.95467/0.94467，PSNR均值/最小值为31.63/30.19dB。相关PR可参考：#35728

未带来端到端收益的优化实验

在Kernel优化过程中，我们认为失败的结果也应该记录在PR中，避免后续重复尝试相同的方向，尤其是Agent辅助优化时，历史记录可以节省Token消耗。

第一个失败实验是前文提到的Qwen-Image QKV通用Scale方案，该方案性能表现良好，但图片质量明显下降，因此最终未合入主分支。

第二个实验是Qwen-Image Final Adaptive LayerNorm，融合Kernel在4608行时从99.360us降至29.632us，微基准测试提升3.353倍，但接入完整模型后，每个步骤仅节省约102us，端到端性能提升落在噪声区间内，甚至有一组测试慢了约0.4%。该Kernel最终复用了通用bit-exact实现，但PR中未宣称模型加速。PR链接：#37144

SANA的LayerNorm + Modulate也遇到过类似情况：GPU Kernel耗时确实减少，但Eager模式本身受限于CPU Launch开销，新Kernel在Python侧的调用成本抵消了性能收益，最终仅在BCG路径启用，Eager模式继续使用原始实现。PR链接：#34015

还有一种更直接的情况：Profile中完全没有新Kernel的名称，这通常意味着Shape Guard未通过、后端编译失败或进入了编译/捕获回退逻辑。遇到这种情况时，应优先解决Kernel命中问题，继续运行基准测试没有意义。

Agent在本次优化中的作用

在这批Kernel优化中，Agent的参与度非常高，我们使用过GPT-5.5、GPT-5.6、Claude Opus 4.8和Fable 5等模型。根据开发记录粗略统计，约95%的代码是在Agent帮助下直接或间接完成的。

这里的帮助不仅限于生成Kernel代码，还包括读取完整Profile数据、对比主分支和候选实现的Kernel调用次数、回到模型代码中寻找可以消除的中间张量。实现完成后，Agent还会协助补充Shape Guard、回退逻辑、单元测试、模型基准测试和PR Review。我们的主要工作是确定优化的模型路径、检查数据、判断数值和质量是否可接受，然后决定哪些实现可以进入主分支，以及清理和重构大量AI生成的代码。

SGLang扩散Kernel的验证流程

最后我们总结一下目前常用的基于Agent的验证流程：

  1. 首先确认Native Backend和Kernel确实命中，启动日志中不应出现Diffusers回退信息，torch profiler中需找到新Kernel的名称和正确的调用次数。
  2. 然后进行Kernel正确性验证：bit-exact路径使用`torch.equal`，输入覆盖连续存储、步幅、不同数据类型、边界形状以及回退场景；允许误差的Kernel使用`rtol/atol`参数验证。
  3. 模型测试固定Prompt、Seed、分辨率、去噪步骤、CFG系数和Checkpoint：bit-exact路径直接保存输出MD5；非bit-exact路径运行多组Prompt的SSIM、PSNR和LPIPS指标；视频模型除了计算平均值，还需关注最差帧和时间方向上的稳定性。
  4. 性能测试分为微基准测试、去噪步骤耗时和常驻端到端测试：微基准测试用于解释具体节省了哪些读写或Launch操作，真正决定是否接入的是后两项测试。冷启动耗时会单独统计，不会混入常驻请求。
  5. 计时前需要进行预热，边界处执行CUDA Synchronize。主分支和PR测试最好在同一台机器上交错执行，例如`main, PR, PR, main`，以减少GPU频率和机器负载波动带来的影响。

SGLang现在为新增扩散Kernel约定了统一的入口规范：Stable Facade、`KernelSpec`、公开的`can_use`接口、数值说明、Kernel测试和模型接线，非bit-exact实现还需放入质量门控。完整说明可参考：README.md

总结

通过这段时间的优化工作，我们对扩散模型Kernel优化的最大感受是：不能仅关注GEMM、注意力等核心操作。Qwen-Image和FLUX.2中存在大量量化、归一化、RoPE和Cat操作，Wan视频模型中还有更复杂的布局转换、因果缓存和VAE数据搬运，以及一些跨越模块边界的融合机会。这些操作单独看开销不大，但放在完整的去噪循环中累积后，会对整体性能产生显著影响。

此外，微基准测试性能优异并不代表模型端到端性能一定提升，Qwen Final AdaLN就是典型的反例。每个PR最好都将Kernel调用次数、GPU耗时、去噪步骤耗时和端到端结果一并展示，尽可能完整地呈现优化效果。

数值精度问题也不能等到最后才处理，是bit-exact、质量门控还是属于量化Checkpoint自身的数值路径，应该在编写Kernel之前就明确确定。边界条件明确后，该融合的地方大胆融合，不满足条件的场景则回退到原始实现。

最后需要说明的是，SGLang扩散Kernel的全部工作是由SGLang团队和社区共同完成的，本文仅记录了我们借助Agent添加的部分快速路径优化，不能代表SGLang扩散Kernel的全部实现和贡献。

以上内容不代表本平台立场，仅供读者参考