diff --git a/python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py b/python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py index 82eb4595a..c605cee05 100644 --- a/python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index f558a7c15..3d6184b9f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/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() - k = k.contiguous() + if not try_native_qk_norm: + 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( - 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, - ) + 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 is handled identically to the eager + # fallback (arange over seqlen, repeated across 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, + ) else: q, k = apply_flashinfer_rope_qk_inplace( q, diff --git a/python/sglang/multimodal_gen/test/unit/test_zimage_qknorm_fusion.py b/python/sglang/multimodal_gen/test/unit/test_zimage_qknorm_fusion.py new file mode 100644 index 000000000..6acfe0ae6 --- /dev/null +++ b/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()