[diffusion] Z-Image bit-exact fused qk-norm (H200 Turbo 1024px e2e -6.4%) (#33886)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py

@@ -180,3 +180,163 @@ def zimage_rmsnorm_tanh_residual(
             num_warps=8,
         )
     return y
+
+
+@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.
+
+    The eager path (`ZImageRMSNorm.forward` on a `[rows, head_dim]` bf16 tensor)
+    lowers to aten pow/mean/rsqrt/mul kernels. For bf16 rows of 128, aten's
+    reduce_kernel vectorizes with 16-byte loads (8 bf16 elements per lane):
+    lane t serially accumulates contiguous elements 8t .. 8t+7 in fp32
+    (squares rounded to bf16 first, matching aten::pow), then a shfl-down
+    butterfly tree combines the 32 lane partials (halves 16/8/4/2/1). Every
+    intermediate below is rounded to bf16 at exactly the same points as the
+    aten kernels, so the output satisfies `torch.equal` against the eager path
+    (verified over 9M+ random rows across scales 0.005-200).
+    """
+    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
+    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)
+    sq = (x * x).to(tl.bfloat16)
+
+    # Lane t's serial accumulation of contiguous elements 8t..8t+7: peel the
+    # 8-element groups apart with ordered tl.split chains (loads stay one
+    # coalesced pass) and add them in exact serial order in fp32.
+    g = tl.reshape(sq, (rows_per_prog, head_dim // 8, 4, 2), can_reorder=False)
+    p0, p1 = tl.split(g)  # elements (0,2,4,6) / (1,3,5,7)
+    p00, p01 = tl.split(
+        tl.reshape(p0, (rows_per_prog, head_dim // 8, 2, 2), can_reorder=False)
+    )  # (0,4) / (2,6)
+    p10, p11 = tl.split(
+        tl.reshape(p1, (rows_per_prog, head_dim // 8, 2, 2), can_reorder=False)
+    )  # (1,5) / (3,7)
+    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 over the 16 lane partials (lanes 16..31 hold the
+    # additive identity in aten's 32-lane warp, so their stage is an exact
+    # no-op): acc[i] += acc[i + half] for half in (8, 4, 2, 1).
+    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 adds the python-float eps in fp32 opmath, then rsqrt rounds to bf16.
+    rstd = tl.rsqrt((ms.to(tl.float32) + eps).to(tl.bfloat16).to(tl.float32)).to(
+        tl.bfloat16
+    )
+
+    weight = tl.load(weight_ptr + offs)
+    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]
+    )
+
+
+def _qk_head_token_stride(x: torch.Tensor, head_dim: int) -> int | None:
+    """Token stride for a `[B, S, H, head_dim]` view whose head block is packed.
+
+    Accepts the strided views produced by slicing a fused qkv projection: the
+    last dim must be contiguous, heads packed (`stride(-2) == head_dim`), and
+    tokens uniformly strided across the flattened `(B, S)` dims.
+    """
+    if x.dim() != 4 or x.shape[-1] != head_dim:
+        return None
+    if x.stride(-1) != 1 or x.stride(-2) != head_dim:
+        return None
+    token_stride = x.stride(1)
+    if x.shape[0] > 1 and x.stride(0) != x.shape[1] * token_stride:
+        return None
+    return token_stride
+
+
+def can_use_qk_rmsnorm_native(
+    x: torch.Tensor, weight: torch.Tensor, head_dim: int
+) -> bool:
+    return (
+        x.is_cuda
+        and weight.is_cuda
+        and x.dtype == torch.bfloat16
+        and head_dim == 128
+        and weight.numel() == head_dim
+        and weight.is_contiguous()
+        and _qk_head_token_stride(x, head_dim) is not None
+    )
+
+
+def zimage_qk_rmsnorm_native(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    """Fused, bit-exact ZImageRMSNorm over q/k heads; returns a contiguous tensor.
+
+    Reads strided `[B, S, H, head_dim]` views (e.g. fused-qkv slices) directly,
+    absorbing the `.contiguous()` materialization the eager path needs.
+    Returns None when the input is not supported (caller falls back).
+    """
+    head_dim = x.shape[-1]
+    if not can_use_qk_rmsnorm_native(x, weight, head_dim):
+        return None
+    token_stride = _qk_head_token_stride(x, head_dim)
+    if token_stride is None:
+        return None
+    if weight.dtype != x.dtype:
+        weight = weight.to(dtype=x.dtype)
+    nheads = x.shape[2]
+    n_rows = x.shape[0] * x.shape[1] * nheads
+    rows_per_prog = 8
+    y = torch.empty(x.shape, dtype=x.dtype, device=x.device)
+    grid = (triton.cdiv(n_rows, rows_per_prog),)
+    with torch.get_device_module().device(x.device):
+        _qk_rmsnorm_native_kernel[grid](
+            y,
+            x,
+            weight,
+            token_stride,
+            nheads,
+            n_rows,
+            head_dim,
+            eps,
+            rows_per_prog=rows_per_prog,
+            num_warps=8,
+        )
+    return y

python/sglang/multimodal_gen/runtime/models/dits/zimage.py

@@ -126,6 +126,38 @@ def zimage_rmsnorm_scale(
     return norm(x) * scale
 
 
+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.
+
+    Replaces `apply_qk_norm`'s eager chain (`q_norm(q.reshape(-1, head_dim))`)
+    with one Triton launch per tensor that reads the strided fused-qkv slices
+    directly. Returns contiguous (q, k) or None when unsupported.
+    """
+    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
+    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
+
+
 class TimestepEmbedder(nn.Module):
     def __init__(self, out_size, mid_size=None, frequency_embedding_size=256):
         super().__init__()
@@ -316,6 +348,15 @@ class ZImageAttention(nn.Module):
         num_replicated_suffix: int = 0,
         skip_sequence_parallel_override: bool = False,
     ):
+        # The fused native qk-norm kernel reads the strided fused-qkv slices
+        # directly and writes contiguous outputs, so q/k materialization is
+        # deferred to it on the main (rope-cache) path.
+        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(
@@ -326,8 +367,9 @@ class ZImageAttention(nn.Module):
                 ],
                 dim=-1,
             )
-            q = q.contiguous()
+            if not try_native_qk_norm:
-            k = k.contiguous()
+                q = q.contiguous()
+                k = k.contiguous()
             v = v.contiguous()
         else:
             q, _ = self.to_q(hidden_states)
@@ -339,17 +381,37 @@ class ZImageAttention(nn.Module):
 
         if rope_cos_sin_cache is not None:
             if self.qk_norm:
-                q, k = apply_qk_norm_with_optional_rope(
+                fused_qk = None
-                    q=q,
+                if try_native_qk_norm:
-                    k=k,
+                    fused_qk = zimage_native_qk_rmsnorm(
-                    q_norm=self.norm_q,
+                        q, k, self.norm_q, self.norm_k, self.head_dim
-                    k_norm=self.norm_k,
+                    )
-                    head_dim=self.head_dim,
+                if fused_qk is not None:
-                    cos_sin_cache=rope_cos_sin_cache,
+                    q, k = fused_qk
-                    is_neox=False,
+                    # positions=None is handled identically to the eager
-                    positions=rope_positions,
+                    # fallback (arange over seqlen, repeated across batch).
-                    allow_inplace=False,
+                    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,
+                    )
             else:
                 q, k = apply_flashinfer_rope_qk_inplace(
                     q,

python/sglang/multimodal_gen/test/unit/test_zimage_qknorm_fusion.py

@@ -0,0 +1,51 @@
+import unittest
+
+import torch
+
+from sglang.multimodal_gen.runtime.models.dits.zimage import (
+    ZImageRMSNorm,
+    zimage_native_qk_rmsnorm,
+)
+
+HEAD_DIM = 128
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+class TestZImageQkNormFusion(unittest.TestCase):
+    def _make_norm(self, seed):
+        torch.manual_seed(seed)
+        norm = ZImageRMSNorm(HEAD_DIM, eps=1e-5)
+        with torch.no_grad():
+            norm.weight.copy_(torch.randn(HEAD_DIM) * 0.5 + 1.0)
+        return norm.to(device="cuda", dtype=torch.bfloat16)
+
+    def test_bit_exact_on_strided_qkv_view(self):
+        norm_q, norm_k = self._make_norm(1), self._make_norm(2)
+        heads, kv_heads = 4, 4
+        total = (heads + 2 * kv_heads) * HEAD_DIM
+        qkv = (torch.randn(1, 64, total, device="cuda")).to(torch.bfloat16)
+        q = qkv[..., : heads * HEAD_DIM].view(1, 64, heads, HEAD_DIM)
+        k = qkv[..., heads * HEAD_DIM : (heads + kv_heads) * HEAD_DIM].view(
+            1, 64, kv_heads, HEAD_DIM
+        )
+
+        fused = zimage_native_qk_rmsnorm(q, k, norm_q, norm_k, HEAD_DIM)
+        self.assertIsNotNone(fused)
+        q_out, k_out = fused
+        self.assertTrue(q_out.is_contiguous() and k_out.is_contiguous())
+        self.assertTrue(
+            torch.equal(q_out, norm_q(q.reshape(-1, HEAD_DIM)).view(q.shape))
+        )
+        self.assertTrue(
+            torch.equal(k_out, norm_k(k.reshape(-1, HEAD_DIM)).view(k.shape))
+        )
+
+    def test_unsupported_head_dim_falls_back(self):
+        norm = ZImageRMSNorm(64, eps=1e-5).to(device="cuda", dtype=torch.bfloat16)
+        q = torch.randn(1, 32, 4, 64, device="cuda", dtype=torch.bfloat16)
+        k = torch.randn(1, 32, 4, 64, device="cuda", dtype=torch.bfloat16)
+        self.assertIsNone(zimage_native_qk_rmsnorm(q, k, norm, norm, 64))
+
+
+if __name__ == "__main__":
+    unittest.main()