# PR #33886 完整报告 - 仓库：`sgl-project/sglang` - 标题：[diffusion] Z-Image bit-exact fused qk-norm (H200 Turbo 1024px e2e -6.4%) - 合并时间：2026-08-07 17:03 - 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/33886 --- # 执行摘要 - 一句话：Z-Image qk-norm 融合 Triton 内核，e2e 提速 6.4% - 推荐动作：建议精读。核心看点：(1) 如何用 Triton 逐位复刻 aten 归约顺序（vec8 lane 串行累加 + shfl-down butterfly）并做到 `torch.equal`；(2) 融合过程中吸收 `.contiguous()` 物化的设计，把 q/k 拷贝开销消掉；(3) 运行时自检从引入到移除的权衡——先以进程级安全网验证可行性，再收敛为单元测试契约。维护时重点关注 bit-exact 契约对 aten 实现细节的依赖程度，以及测试用例是否足以在 aten 变更时及时暴露回归。 # 功能与动机 PR #29742 为修复 Z-Image 精度，在主注意力路径（precomputed rope cache）传 `allow_inplace=False`，把 qk-norm 强制到 eager `ZImageRMSNorm` 链。PR body 的 torch-profiler 显示该链是 denoise 步骤最贵的开销：eager qk-norm 链占 7.9%、q/k `.contiguous()` 拷贝占 3.7%、独立 flashinfer rope 启动占 1.0%。目标是在不改变任何数值输出的前提下，用融合 kernel 吸收这部分开销。 # 实现拆解 1. 瓶颈定位与融合范围选择：PR #29742 使主注意力路径走 eager `ZImageRMSNorm` 链，profiler 定位 qk-norm 链与 q/k 物化为主要开销。本 PR 只对最热的 rope-cache 主路径做 fused 优化，其他 rope 分支（`freqs_cis`）、量化运行和 torch.compile 场景保持原 eager 回退。 2. 新增 bit-exact Triton kernel：在 `python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py` 实现 `_qk_rmsnorm_native_kernel`，逐位复刻 aten 的数值 DAG —— bf16 行 128 场景下 16 字节向量化加载、lane 内串行 fp32 累加 8 个连续平方（平方先按 bf16 舍入）、shfl-down butterfly 归约、fp32 opmath 加 eps 后再按 bf16 舍入。同时新增 `can_use_qk_rmsnorm_native` 与 `_qk_head_token_stride` 做前置布局检查（4D 视图、head 块连续、token 均匀 stride、bf16、head_dim=128、CUDA），不满足即返回 None 走回退。 3. 模型侧接线：`zimage.py` 新增 `zimage_native_qk_rmsnorm` 封装；`forward` 中只在 `qk_norm + rope_cos_sin_cache + 非 torch.compile + 开关开启` 时尝试 fused 路径，且把 q/k 的 contiguous 物化推迟到 kernel 内完成（v 仍保留拷贝）。fused 成功后复用 `apply_flashinfer_rope_qk_inplace`，positions 语义与 eager fallback 一致。 4. 安全网设计与收敛：提交 1 引入进程级 `torch.equal` 运行时自检（首次调用不匹配则永久禁用 fused 并告警）；提交 2 移除该全局状态，改为由单元测试固定 bit-exact 契约，避免隐藏运行时开关。 5. 测试与验证配套：新增 `test_zimage_qknorm_fusion.py`，覆盖 strided fused-qkv 视图上的 bit-exact（`torch.equal`）与不支持的 head_dim 回退；PR body 给出 9M+ 行跨尺度验证、E2E PNG md5 四臂全等，以及 H200 性能数据（e2e -6.4%）。 关键文件： - `python/sglang/multimodal_gen/runtime/models/dits/zimage.py`（模块 模型层；类别 source；类型 core-logic；符号 zimage_native_qk_rmsnorm）: Z-Image 主路径接线文件：新增 `zimage_native_qk_rmsnorm` 封装，并在 `forward` 的 rope-cache 分支引入 fused 路径与 eager fallback 的完整控制流，决定优化是否生效。 - `python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py`（模块 内核层；类别 source；类型 core-logic；符号 _qk_rmsnorm_native_kernel, _qk_head_token_stride, can_use_qk_rmsnorm_native, zimage_qk_rmsnorm_native）: PR 的核心：新增 bit-exact Triton kernel `_qk_rmsnorm_native_kernel`，逐位复刻 eager `ZImageRMSNorm` 的 aten 数值 DAG，并直接读取 strided fused-qkv 切片、写出 contiguous 结果，吸收 q/k 拷贝开销。 - `python/sglang/multimodal_gen/test/unit/test_zimage_qknorm_fusion.py`（模块 单元测试；类别 test；类型 test-coverage；符号 TestZImageQkNormFusion, _make_norm, test_bit_exact_on_strided_qkv_view, test_unsupported_head_dim_falls_back）: bit-exactness 的固定契约：用 `torch.equal` 校验 strided fused-qkv 视图上的融合结果与 eager `ZImageRMSNorm` 完全一致，并覆盖不支持 head_dim 时的回退行为。 关键符号：zimage_native_qk_rmsnorm, zimage_qk_rmsnorm_native, can_use_qk_rmsnorm_native, _qk_head_token_stride, _qk_rmsnorm_native_kernel, TestZImageQkNormFusion ## 关键源码片段 ### `python/sglang/multimodal_gen/runtime/models/dits/zimage.py` Z-Image 主路径接线文件：新增 `zimage_native_qk_rmsnorm` 封装，并在 `forward` 的 rope-cache 分支引入 fused 路径与 eager fallback 的完整控制流，决定优化是否生效。 ```python def zimage_native_qk_rmsnorm( q: torch.Tensor, k: torch.Tensor, norm_q: ZImageRMSNorm, norm_k: ZImageRMSNorm, head_dim: int, ) -> tuple[torch.Tensor, torch.Tensor] | None: # Fused per-head ZImageRMSNorm for q/k, bit-exact vs the eager fallback. # 用一个 Triton launch 直接读 strided fused-qkv 切片，返回 contiguous (q, k)， # 不支持时返回 None 交给调用方走 eager 回退。 from sglang.kernels.ops.diffusion.triton.zimage_native_norm import ( can_use_qk_rmsnorm_native, zimage_qk_rmsnorm_native, ) q_weight = norm_q.weight.data k_weight = norm_k.weight.data # 前置能力检查：任一输入不满足布局 / 精度条件就整体返回 None if not ( can_use_qk_rmsnorm_native(q, q_weight, head_dim) and can_use_qk_rmsnorm_native(k, k_weight, head_dim) ): return None q_out = zimage_qk_rmsnorm_native(q, q_weight, norm_q.variance_epsilon) k_out = zimage_qk_rmsnorm_native(k, k_weight, norm_k.variance_epsilon) if q_out is None or k_out is None: return None return q_out, k_out # ---- ZImageTransformerBlock.forward 中主路径接线（rope-cache 分支）---- # 只在最热路径（qk_norm + rope 缓存 + 非 torch.compile）尝试 fused， # 量化和 torch.compile 等场景永远走原 eager 链 try_native_qk_norm = ( self.qk_norm and rope_cos_sin_cache is not None and self.enable_zimage_qk_fusion and not torch.compiler.is_compiling() ) if self.use_fused_qkv: qkv, _ = self.to_qkv(hidden_states) q, k, v = qkv.split( [ self.local_num_heads * self.head_dim, self.local_num_kv_heads * self.head_dim, self.local_num_kv_heads * self.head_dim, ], dim=-1, ) # q/k 的 contiguous 物化推迟给 fused kernel；否则仍按原逻辑提前拷贝 if not try_native_qk_norm: q = q.contiguous() k = k.contiguous() v = v.contiguous() # ... if rope_cos_sin_cache is not None: if self.qk_norm: fused_qk = None if try_native_qk_norm: fused_qk = zimage_native_qk_rmsnorm( q, k, self.norm_q, self.norm_k, self.head_dim ) if fused_qk is not None: q, k = fused_qk # positions=None 时与 eager fallback 语义一致：按 seqlen arange 并跨 batch 重复 q, k = apply_flashinfer_rope_qk_inplace( q, k, rope_cos_sin_cache, head_size=self.head_dim, is_neox=False, positions=rope_positions, ) else: q = q.contiguous() k = k.contiguous() q, k = apply_qk_norm_with_optional_rope( q=q, k=k, q_norm=self.norm_q, k_norm=self.norm_k, head_dim=self.head_dim, cos_sin_cache=rope_cos_sin_cache, is_neox=False, positions=rope_positions, allow_inplace=False, # 保持 #29742 引入的原生 bf16 数值轨迹 ) ``` ### `python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py` PR 的核心：新增 bit-exact Triton kernel `_qk_rmsnorm_native_kernel`，逐位复刻 eager `ZImageRMSNorm` 的 aten 数值 DAG，并直接读取 strided fused-qkv 切片、写出 contiguous 结果，吸收 q/k 拷贝开销。 ```python @triton.jit def _qk_rmsnorm_native_kernel( y_ptr, x_ptr, weight_ptr, token_stride, nheads, n_rows, head_dim: tl.constexpr, eps: tl.constexpr, rows_per_prog: tl.constexpr, ): # Per-head RMSNorm replicating the eager ZImageRMSNorm kernel chain bit-for-bit. # 复刻要点：bf16 行 128 时 aten reduce_kernel 用 16 字节向量化加载， # 每个 lane 串行累加 8 个连续平方（平方先按 bf16 舍入），再做 shfl-down # butterfly 归约；所有中间结果在与 aten 相同的边界点按 bf16 舍入。 prog = tl.program_id(0) row_offs = tl.arange(0, rows_per_prog) rows = prog * rows_per_prog + row_offs row_mask = rows < n_rows # 把扁平行号 (token, head) 拆开，直接按 fused-qkv 的 stride 取地址 tokens = rows // nheads heads = rows % nheads base = x_ptr + tokens * token_stride + heads * head_dim offs = tl.arange(0, head_dim) x = tl.load(base[:, None] + offs[None, :], mask=row_mask[:, None], other=0.0) # aten::pow 的平方结果先舍入到 bf16，再进入 fp32 累加 sq = (x * x).to(tl.bfloat16) # 用有序 tl.split 链把 8 元素组按 lane 串行顺序拆开， # 保证 fp32 累加次序与 aten 完全一致（s0+...+s7） g = tl.reshape(sq, (rows_per_prog, head_dim // 8, 4, 2), can_reorder=False) p0, p1 = tl.split(g) p00, p01 = tl.split(tl.reshape(p0, (rows_per_prog, head_dim // 8, 2, 2), can_reorder=False)) p10, p11 = tl.split(tl.reshape(p1, (rows_per_prog, head_dim // 8, 2, 2), can_reorder=False)) s0, s4 = tl.split(tl.reshape(p00, (rows_per_prog, head_dim // 8, 1, 2), can_reorder=False)) s2, s6 = tl.split(tl.reshape(p01, (rows_per_prog, head_dim // 8, 1, 2), can_reorder=False)) s1, s5 = tl.split(tl.reshape(p10, (rows_per_prog, head_dim // 8, 1, 2), can_reorder=False)) s3, s7 = tl.split(tl.reshape(p11, (rows_per_prog, head_dim // 8, 1, 2), can_reorder=False)) acc = tl.reshape(s0, (rows_per_prog, head_dim // 8)).to(tl.float32) acc = acc + tl.reshape(s1, (rows_per_prog, head_dim // 8)).to(tl.float32) acc = acc + tl.reshape(s2, (rows_per_prog, head_dim // 8)).to(tl.float32) acc = acc + tl.reshape(s3, (rows_per_prog, head_dim // 8)).to(tl.float32) acc = acc + tl.reshape(s4, (rows_per_prog, head_dim // 8)).to(tl.float32) acc = acc + tl.reshape(s5, (rows_per_prog, head_dim // 8)).to(tl.float32) acc = acc + tl.reshape(s6, (rows_per_prog, head_dim // 8)).to(tl.float32) acc = acc + tl.reshape(s7, (rows_per_prog, head_dim // 8)).to(tl.float32) # shfl-down butterfly（half = 8/4/2/1 的两两有序折叠，等价 aten 的 warp 归约） acc = tl.sum(tl.reshape(acc, (rows_per_prog, 2, 8), can_reorder=False), axis=1) acc = tl.sum(tl.reshape(acc, (rows_per_prog, 2, 4), can_reorder=False), axis=1) acc = tl.sum(tl.reshape(acc, (rows_per_prog, 2, 2), can_reorder=False), axis=1) ssum = tl.sum(acc, axis=1) ms = (ssum / head_dim).to(tl.bfloat16) # aten 在 fp32 opmath 上加 eps，rsqrt 结果先舍入到 bf16 再参与后续乘法 rstd = tl.rsqrt((ms.to(tl.float32) + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16) weight = tl.load(weight_ptr + offs) # x * rstd 的乘积累积 bf16 舍入点与 eager mul 一致，再乘 weight y = ((x.to(tl.float32) * rstd.to(tl.float32)[:, None]).to(tl.bfloat16) * weight).to(tl.bfloat16) tl.store(y_ptr + rows[:, None] * head_dim + offs[None, :], y, mask=row_mask[:, None]) ``` # 评论区精华 该 PR 无公开 review 评论（review 评论为 0），唯一 issue 评论是作者 BBuf 贴的 CI 链接。最有价值的设计讨论体现在两次提交的演进中：提交 1 加入了进程级 `torch.equal` 运行时自检作为安全网（不匹配则永久禁用 fused 路径），提交 2（5c9e7df）将其移除，提交消息写道：“Drop the global self-check ... bit-exactness is pinned by the unit test instead of a hidden runtime switch”。这是一个值得注意的权衡：用显式单元测试契约替代进程内隐藏开关，消除运行时开销与全局副作用，代价是数值契约对 aten 实现变化的防护从运行时变成了 CI 时。 - 运行时 bit-exact 自检的去留 (design): 移除隐藏运行时开关与进程级全局状态，数值契约由 `test_zimage_qknorm_fusion.py` 固定；kernel 或 aten 行为变化时靠 CI 暴露。 # 风险与影响 - 风险： - 数值契约对 aten 实现的耦合：`_qk_rmsnorm_native_kernel` 精确复刻 aten `reduce_kernel` 在 bf16、head_dim=128 下的 16 字节向量化加载、lane 串行 fp32 累加与 shfl-down butterfly 顺序。若未来 aten dispatch、向量化宽度或 GPU 平台行为变化，bit-exact 可能被破坏；提交 2 移除运行时自检后，仅靠 `test_zimage_qknorm_fusion.py` 两个用例把关，覆盖规模有限（随机数据、非真实模型分布）。 - 适用范围窄但 fallback 完整：仅 CUDA + bf16 + head_dim=128 + 4D 且 head 块连续、token 均匀 stride 的输入走 fused；量化、torch.compile、`freqs_cis` 分支等仍走 eager。前置检查 `can_use_qk_rmsnorm_native` 返回 None 即回退，行为不变。 - 物化语义变化：q/k 不再提前 `contiguous()`，而是由 kernel 直接读 strided 切片并输出 contiguous 张量；v 仍保留拷贝。需关注后续 `apply_flashinfer_rope_qk_inplace` 与注意力后端对 q/k 布局的假设（本 PR 已在 E2E 验证）。 - 性能收益的普适性：-6.4% 是 H200 + 1024x1024 + Turbo 9 步 preset 的结果，其他分辨率、步数、GPU 型号的收益需要另行验证。 - 影响：影响范围集中于 `sglang/multimodal_gen` 的 Z-Image / Z-Image-Turbo 推理：默认 eager 配置下 rope-cache 主路径性能提升约 6%，且输出与之前逐位一致，用户无感知迁移成本。对量化、torch.compile、非常规布局用户完全透明（走原 eager 路径）。对团队而言，该 PR 提供了一种可复用的“bit-exact 融合 eager aten 链”模式：前置能力检查 + None 回退 + 测试固定契约，为后续吸收其他模型的规范化 kernel 链（如 MiniMax-H3、GLM-Image 的 RMSNorm 链）提供参考；同时新增一个需要与 aten 行为保持同步的 Triton kernel，带来长期维护成本。CI 新增一个 CUDA-dependent 单元测试（非 CUDA 环境 skip）。 - 风险标记：bit-exact 依赖 aten 数值实现 , 运行时自检已移除 , 仅 bf16/head_dim=128/CUDA 生效 , torch.compile/ 量化走回退路径 # 关联脉络 - PR #29742 #29742 修复 Z-Image 精度（标题未在上下文中提供）: PR body 明确引用：`allow_inplace=False` 强制 eager qk-norm 链，是本 PR 的性能瓶颈来源；本 PR 在其数值轨迹之上做 bit-exact 融合。 - PR #33923 [Diffusion] Route zimage and hunyuanvideo attention through USPAttention: 同样修改 `zimage.py` 的注意力路径，将 zimage 注意力迁入 USPAttention，与本 PR 同属 Z-Image 注意力链路演进线。 - PR #32667 [Diffusion] Add K/V-gather sequence parallel attention: 为 Diffusion 新增 K/V-gather 序列并行注意力并在 SP2 默认启用，同样涉及 `zimage.py` 注意力路径与 SP 布局，与本 PR 改动区域相邻。