SGLang 扩散模型融合算子包sglang.kernels.ops.diffusion完全指南DiT / VAE 位精确内核与质量门控机制【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglang导读本文档面向在 SGLang 中接入扩散模型multimodal-generation服务、或需要为 DiT Transformer 块、VAE 编码器/解码器编写高性能融合算子的开发者。文章以python/sglang/kernels/ops/diffusion/README.md为骨架结合 包入口 的注册表与导出表、位精确门控、质量门控 以及各算子域的 Triton / CuTe-DSL / FlyDSL / JIT CUDA 实现深入讲解该包的设计哲学、算子选择矩阵、数值契约与接入新算子的完整流程。读完本文你将掌握为什么扩散算子必须区分位精确与近似两条门控路径、如何在众多看起来可互换的归一化算子中正确选型、以及如何把一个新的融合内核接入 SGLang 的统一内核注册表并保证服务质量不劣化。定位一个按模型定制的融合算子包而非通用算子库与 SGLang 面向 LLM 的算子分组不同sglang.kernels.ops.diffusion中几乎没有任何通用算子。每一个内核都对应某个具体模型中的某一条具体 eager 算子链例如 ERNIE-Image 的 adaLN 链、Wan VAE 的 channel-first RMSNormSiLU、FLUX.2 的 QK RMSNormRoPEQKV 打包、LingBot Video 的分组受限 MoE 路由等其价值不仅来自带宽优化更来自它精确复现了哪一条舍入边界。这一点在扩散模型场景下格外关键多步去噪multi-step denoising会把单步的舍入差异逐级放大为肉眼可见的质量损失。因此本文档给出一个核心判断标准——差不多close与位精确bit-exact是两个不同的产品二者的挂载门控策略完全不同契约类型判定方式挂载策略位精确bit-exacttorch.equal与 eager 链逐位一致无条件挂载unconditionally mounted无需质量门控近似close数值接近但不逐位一致仅质量门控只对qualityextra-high与qualityhigh的请求在 batch 边界、按整个 Transformer 全有或全无地挂载一个关键反例一个看起来无害的普通 fp32 单遍 norm 融合实际上并不安全——在 ERNIE-Image 上它把 50 步去噪轨迹推到了 PSNR 18.83 dB这正是该包重写为位精确实现的直接动因见 rmsnorm_scale_shift_bitexact.py 的模块文档。另外模型/checkpoint 原生的model/checkpoint-native选择——例如与所选模型或部署路径绑定的宽松契约内核、稀疏算子、FP8/NVFP4 量化产出器——不属于请求级quality档位管制的范围quality 档位既不会选择也不会禁用这些独立的模型级选择。导入面只从包门面导入绝不从子模块导入包的对外契约非常明确from sglang.kernels.ops.diffusion import fused_rmsnorm_scale_shift_bitexact必须从包facade导入绝不能从子模块导入。内部文件布局norm、rope等子包随时可能调整而门面接口是稳定的。唯一的例外是刻意只测单一后端的测试。这种懒解析lazy resolutionPEP 562设计在 包入口 中体现得十分彻底各后端拥有互不相交的重依赖Triton、CUTLASS/CuTe-DSL、ROCm 上的 FlyDSL若在包导入时急切地 re-export会把所有后端变成所有平台上的硬性 import-time 依赖因此入口模块维护一张_EXPORTS字典符号 → 所属子模块通过模块级__getattr__在首次属性访问时才import_module并缓存到globals()后续查找完全绕过__getattr__同时维护_SPECS元数据元组通过register_kernel(KernelSpec(...))把每个算子注册进 SGLang 的统一内核注册表。注册只登记元数据——不导入 torch、不触发任何后端、不触发 JIT 构建。同一算子携带多个后端如scale_residual_norm_scale_shift同时有 Triton、CuTe-DSL、FlyDSL 三种实现时调用方通过select_kernel按名字选用。以入口文件中的注册表为例一个算子条目包含op名、后端类型KernelBackend.TRITON / JIT / KDA / CUTE_DSL / FLYDSL / AOT、目标函数路径相对或全限定、能力约束CUDA / HIP / SM100如 Qwen-Image 的 QKV epilogue 明确要求CapabilityRequirement.cuda(min_sm(10, 0))、以及一句话描述。目录布局算子域、编译器后缀、以及不是内核的目录普通实现按**算子域operator domain**划分子包编译器/后端则体现在文件名的后缀_triton、_jit、_cutedsl、_flydsl以及表示数值语义的_bitexact。带 Kernel Design AgentsKDA来源的实现位于sglang.kernels.kda_kernels门面仍是其唯一受支持的运行时导入面。norm/ RMSNorm / LayerNorm / GroupNorm 及其融合 epilogue modulate/ adaLN modulate、gating、timestep 条件化 rope/ 旋转位置编码及融合进其中的 QK-norm 链 activation/ SiLU / GLU / GELU 融合 attention/ 稀疏线性注意力、gated delta-net routing/ 扩散模型 MoE 路由与专家选择 layout/ 纯数据搬运USP/Ulysses relayout、varlen pack、causal pad common/ 数值原语、平台谓词、非 Triton 回退 sites/ 请求作用域的挂载策略 —— 不是内核见下文 ext/ JIT C/CUDA 扩展Hunyuan3D raster/inpaint—— 不是内核 ../../kda_kernels/ Agent 生成的实现及其 JIT CUDA 源码实际目录与 README 完全一致例如 norm 下有rmsnorm_scale_shift_bitexact.py、scale_residual_norm_cutedsl.py、group_norm_silu_triton.py等sites 下有bitexact_gate.py、quality_gate.py及各模型站点文件。数值契约与质量策略位精确、质量门控与非内核边界位精确内核 → 无条件挂载位精确内核逐位复现 eager 链的每一个 aten 舍入边界有时精确到归约树reduction tree本身。例如sglang.kernels.kda_kernels/layernorm_modulate_triton.py复现 torch 2.11 的vectorized_layer_norm_kernel128 线程 Welford、_rcp4保护的倒数、shfl.down折叠顺序、div.rnMUFU.RSQnorm/rmsnorm_scale_shift_bitexact.py复现 flashinfer 的 CuTe-DSLRMSNormKernel的 fragment 顺序与shfl.bfly折叠。即使如此这些内核仍会在首次见到真实输入时用sites/bitexact_gate.py对实时 eager 链做一次自校验一旦不匹配就永久回退——因为它们复现的 dispatch 本身可能随平台变化。BitExactFusionGate 提供两种校验模式once-for-all默认第一次torch.equal通过后融合路径对之后所有调用永久生效GLM / Ernie 使用per-signature每个不同的签名sig独立校验FLUX / Sana 使用以匹配按形状分派的 aten LayerNorm dispatch。can_attempt_once()明确禁止在torch.compile追踪或 CUDA graph 捕获期间做首次校验会执行 eager 参考链 host 同步而accept_or_fallback()在输出与参考不一致时打印告警、永久disable()并返回 eager 参考输出保证正确性优先。它还提供了flashinfer_rmsnorm_diagnostic_hint()回调用于在精确度失配时输出 FlashInfer 归一化后端诊断检测_USE_CUDA_NORM标志、FLASHINFER_USE_CUDA_NORM环境变量及 flashinfer 相关包版本。非位精确内核 → 质量门控近似内核只能通过sites/quality_gate.py中的QualityGatedFusion挂载到被标记的nn.Module站点上且仅对qualityextra-high和qualityhigh的请求生效在 batch 边界、按每个 Transformer 全有或全无all-or-nothingextra-high只追加这些请求门控的 DiT/VAE 融合high是累积的还可能启用模型自有的近似路径例如 Cache-DiT 或低精度 decode。QualityGatedFusion通过mark/mount/unmount/is_enabled维护站点的 marker 属性与 enabled 标志enabled 是一个普通模块属性以便编译后的模型 forward 读取时不依赖该辅助对象mount还支持reject_reason静态守卫回调——任何站点未通过守卫时整个模型保持参考路径。一个质量门控的实际例子SANA-Video 线性注意力SANA-Video 的质量门控线性注意力站点第一个 GEMM 保持 BF16 输入但要求 FP32 累加/输出第二个 GEMM 用 FP32 运行而默认路径在两次 GEMM 之前就把 Q/K/V 提升到 FP32。这个半精度输入 FP32 输出的混合契约正是通过站点文件sites/sana_video_linear_attention_site.py实现的。什么是不是内核sites 与 extsites/重写nn.Module树mark / mount / unmount不是算子它是唯一允许在函数内部、懒加载地引用multimodal_gen类型的地方因为检查模型模块就是它的全部工作ext/构建没有后端维度、也没有数值契约的 C/CUDA 扩展如 Hunyuan3D rasterizer它们共享本包的构建机制但刻意独立成目录。入口点协议predicate kernel 配对每个公开内核都是谓词 内核配对if can_use_op(...): out op(...) else: out reference chain关键约束内核在遇到不支持的输入时直接 raise绝不返回None——一个静默的None太容易被调用方漏掉检查而失败模式将是一张看起来错误的图像而不是一个异常。这一点在group_limited_topk的入口函数中可见can_use_group_limited_topk失败时直接raise ValueError(...)列出全部前置条件非空连续 CUDA float32[tokens, experts]张量、每组至少两个 2 的幂专家、1 n_group、0 topk_group n_group、top_k不超过所选组容量。同时公开内核用register_custom_op注册并配套fake_implmeta/fake 实现确保在 torch.compile / 元设备模式下可追踪见rmsnorm_scale_shift_bitexact.py与group_limited_topk_triton.py。选择矩阵从看起来可互换的算子中正确选型README 明确警告好几个归一化算子看起来可互换实际并不是。从这里开始选型。Norm scale/shiftadaLNEntry pointBackendContractApplies tofuse_scale_shift_kernelTritonclose连续 BLCscalar/row/token 调制加 causal-video[B, F, 1, C]使用静态封顶的 2 的幂 tile 避免请求期 autotuningfused_rmsnorm_scale_shift_bitexactTriton对 flashinfer CuTe RMSNorm aten modulate 位精确bf16、连续行、H 64 * threads_per_rowfused_scale_residual_rmsnorm_scale_shift_bitexactTriton位精确含前置 residual-gate add同上fused_layernorm_modulateTriton对 atenvectorized_layer_norm位精确bf16、N % 4 0、16B 对齐fused_norm_scale_shift/fused_scale_residual_norm_scale_shiftCuTe-DSLfp32 统计量、closefp16/bf16/fp32、LN 或 RMS、多种广播模式flydsl_norm_scale_shift/flydsl_fused_residual_norm_scale_shiftFlyDSLclose仅 ROCm gfx950try_fused_scale_residual_norm_scale_shift_nvfp4JIT CUDA匹配所选 NVFP4 产出器契约Qwen residual LayerNorm/调制 FC1 NVFP4 量化fuse_layernorm_scale_shift_gate_select01_kernelTritonclose每个 token 在两行调制之间选择Qwen-Imagenorm_infer/rms_norm_fnTritontorch/NPU/MPS 回退close通用入口上面都不适用时用它以fused_rmsnorm_scale_shift_bitexact为例其数值契约在源码文档中逐条写死rmsnorm_scale_shift_bitexact.pyRMSNorm.forward_cuda分派到 flashinfer CuTe-DSLRMSNormKernel连续 bf16 行、H 64 * threads_per_rowH≤3072 时 32否则 64时每个线程tx拥有列{8*TPR*b 8*tx v}fragment 按 v 最快、再 b 排序平方在 fp32 中分别舍入、不使用 FMA用不透明mul.rn.f32阻止编译器收缩成 FMA然后对 64 个 fragment 值做有序的顺序 fadd 链无 reassocwarp 归约用shfl.bfly偏移 1,2,4,8,16——即相邻对折叠树rstd rsqrt.approx.f32(sum_sq / H eps)输出y (bf16)(float(x) * rstd * (w 0.0))只做一次最终舍入aten modulate 链在每步之后舍入到 bf16round(1 scale)、round(y * that)、round(prod shift)残差变体重现 eager 对round(gate * update)、round(residual that)再把舍入结果送入同一忠实 norm。作者还特别提示了一个易踩的坑折叠阶段num_warps必须与被复现内核的 warps-per-row 一致否则会触发病态的 Triton 布局转换实测 25µs → 500µs。Norm 变体Entry pointBackendContractApplies totriton_group_norm_silu/apply_group_norm_siluTritoncloseNCHW 连续、任意 channels-per-group、总是施加 SiLUgroup_norm_silu_4d/group_norm_silu_rowsTritonclose仅 channels_last2 的幂C 2048可选 SiLU。这让 VAE decoder 可以端到端保持 channels_last无需nchwToNhwcwan_rmsnorm_siluTritonclose稠密channels_last_3d5Dstride(C) 1、Wan VAE channel-first RMSNorm SiLUrmsnorm_scale/rmsnorm_tanh_residualTritonbf16 原生统计量Z-Image与其自身参考精确一致、Ideogram 4门控zimage_qk_rmsnorm_nativeTritonbit-exactZ-Image 每头 QK RMSNormfused_qk_head_layernormTritonbit-exact每头 LN on q/k、dim_head % 4 0、 128triton_one_pass_rms_normTritonclose独立 RMSNorm单遍残差门控Residual gatingEntry pointBackendContractApplies toresidual_gate_addKDAJIT CUDA对residual update * gate位精确连续张量或转置稠密[B, tokens, hidden]残差/输出 连续 update row-broadcast gateSANA-Video转置稠密路径用共享内存 tile 以逻辑行主序读取 update同时保持残差读取与输出写在其[B, hidden, tokens]底层布局上合并coalesced。README 特别警告不要仅仅为了走普通路径而插入.contiguous()——那会在每个残差站点恢复一整次张量拷贝。RoPE / QK-normEntry pointBackendContractfused_inplace_qknorm_ropeJIT CUDA相对拆分基线只多一步 bf16 舍入round_norm_before_ropeTrue时精确支持 compact 与全宽 NeoX/interleaved cachefused_qknorm_rope_pack_kvJIT CUDA同上额外打包前缀 K/Vtry_fused_flux2_qkv_epilogueKDAJIT CUDA对所选 BF16 链位精确FLUX.2 QK RMSNorm RoPE 联合 QKV 打包try_fused_qwen_qkv_epilogueJIT CUDA对所选 BF16 链位精确Qwen-Image QK RMSNorm RoPE 联合 QKV 写入SM100fused_rope_rotate_half_bitexactTriton位精确仅逐元素fused_interleaved_rope_fp64JIT CUDA对 SANA-Video 配对的 fp64 RoPE 位精确fused_inplace_helios_qk_ropeJIT CUDA对 Helios 转置频率布局的配对就地 RoPE 位精确ltx2_qknorm_split_rope_cudaKDAJIT CUDAclose在 B200 上验证fused_ltx25_decoder_ropeJIT CUDA由缓存的 compact 轴表配对 3D RoPE位精确apply_rotary_embeddingTriton回退close通用入口hunyuan_qkv_rope_packTriton位精确单遍打包 QKV 并施加 RoPEMoE 路由Entry pointBackendContractApplies togroup_limited_topkTriton所选专家 id 集合与受守卫的 CUDAtorch.topk(..., sortedFalse)链一致输出顺序未指定LingBot Video 默认开启的 sigmoidbias 分组受限路由连续 fp32[tokens, experts]、每组至少两个 2 的幂专家group_limited_topk_triton.py 的实现细节值得展开参考的 LingBot Video 路由用一个小内核链完成分组受限 top-k每组 top-2 与求和、组 top-k、scatter_进零掩码、expand/reshape广播、masked_fill为-inf、最终专家 top-k。在 launch-bound 的单 GPU 上这条链纯粹是开销——每个中间张量都很小整个计算受带宽与 launch 限制。融合内核每个 token 一个 program一次性加载自己的 score 行、在寄存器中归约组内和、用-inf掩蔽未选组、写出 top-k 专家 id并在 128 专家 / 4 组 / 选 2 组 / top-8 的生产配置上保持所选集合一致。注意其对重复最大值的处理掩掉所有等于m1的值会丢掉组内第二大的项组内出现重复最大值时因此实现用tl.min(tl.where(g m1, group_e, BLOCK_EPG))只移除恰一份第一个最大值。数据搬运与量化布局产出器以下算子要么是位精确的数据搬运要么是保持运算次序的算术same-order arithmeticusp_merge_heads、pack_qkv_destination_major、fused_pack_qkv、fused_pack_segmented_qkv、fused_scatter_to_padded、fused_causal_conv3d_cat_pad_cuda、cat_pad_channels_last_3d、dup_up3d_add、fused_temb_table_slices、ltx2_ada_values9。此外fused_layernorm_modulate_fp8_quant_raw把 FLUX.2 的 LayerNorm、adaLN 调制与静态 FP8 量化折叠进单内核try_flux2_token_cat_fp8与try_flux2_token_cat_nvfp4把分支拼接直接融合进 FLUX.2 checkpoint 路径所选的量化表示。fused_temb_table_slices尤其值得关注eager 版本(table temb.float()).chunk(6, dim2)在 704p/121f 时物化约 8GB 的 fp32并把六个带步长的切片交给下游它们的.contiguous()调用又把每个切片各复制一遍——融合后这些中间拷贝全部消失。在真实模型中的接线方式位精确内核在python/sglang/multimodal_gen/runtime/models/dits/下的各 DiT 模型中使用。以 ERNIE-Image 为例ernie_image.py模块顶部为每条融合路径各建一个BitExactFusionGatefused-norm、fused gated-norm、fused RoPE、fused QKNormRoPE、fused GELU-mul正向传播中调用fused_rmsnorm_scale_shift_bitexact等入口通过门的accept_or_fallback在首次校验通过后启用、失败则永久回退到 eager。类似的接线也出现在flux.py、flux_2.py、glm_image.py、qwen_image.py、sana.py等文件中与 README 中 dispatch 可能在调用方脚下变化 的警告相互印证。添加一个新内核六步流程README 给出了官方接入流程结合源码可完整还原每一步放对位置普通实现放进对应算子域子包norm/、rope/、modulate/等KDA 工作流生成的实现放进sglang.kernels.kda_kernels连同其源码修订信息与任何 JIT CUDA 源文件。登记在 包入口 的_EXPORTS符号 → 子模块映射供 PEP 562 懒导入与_SPECSKernelSpec注册元数据中各登记一条。规范上_EXPORTS按域 → 模块 → 符号排序新公开内核只能出现在这里。给出can_use_*谓词不支持的输入直接 raise不返回None。在模块 docstring 中声明数值契约包括在哪些形状上验证过例如rmsnorm_scale_shift_bitexact声明在(1,4216,4096)/(1,4096,4096)/(2,1140,4096)/(1,128,2048)bf16 上做过torch.equal验证。若非位精确走sites/门控必须同时挂载到extra-high与high绝不能在默认的lossless路径上生效。测试算子域测试放进test/registered/kernels/ops/diffusion/当前包含test_norm.py、test_rope.py、test_modulate.py、test_layout.py、test_routing.py、test_sites.py、test_model_fast_paths.py等 19 个测试文件模型接线测试进test_model_fast_paths.py。另有test/registered/e2e/diffusion/下的端到端用例如 test_diffusion_unit.py以run_diffusion_suite(unit)方式组织 1-GPU / 2-GPU / B200 / BCG 等 CI 泳道。小结sglang.kernels.ops.diffusion是 SGLang 中一个纪律性极强的融合算子包它用位精确 → 无条件挂载、近似 → 质量门控的双轨契约把扩散模型多步去噪对舍入误差的敏感性变成可工程化的准则用predicate kernel协议和 PEP 562 懒门面隔离了 Triton / CuTe-DSL / FlyDSL / JIT CUDA 多种后端的重依赖再用sites/把改模型模块这一非算子职责隔离在算子域之外。无论是为新的 DiT/VAE 模型接入融合算子还是理解 SGLang 扩散服务的数值质量保障机制本文的选择矩阵与接入流程都可以直接作为工作起点。【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
