[diffusion] Ideogram 4: fuse RMSNorm modulate/gate chains via the Z-Image Triton suite behind quality=high (H200 e2e -2.9%/-3.4%) (#33822)

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

python/sglang/kernels/ops/diffusion/fused_gate_rmsnorm.py

@@ -0,0 +1,135 @@
+"""Quality-gated fused RMSNorm modulate/gate sites (Z-Image Triton suite reuse).
+
+Adaln-style DiT blocks (Ideogram 4) spend four elementwise chains per block on
+modulate/gate around each RMSNorm: ``RMSNorm(x) * scale`` before
+attention/FFN and ``x + tanh(gate) * RMSNorm(out)`` after. The Z-Image
+bf16-native Triton kernels
+(:mod:`sglang.kernels.ops.diffusion.triton.zimage_native_norm`) fuse each
+chain into a single kernel (RMSNorm + tanh + mul + add in one pass).
+
+Z-Image mounts those kernels unconditionally because they reproduce its own
+native-bf16 reference RMSNorm. Ideogram's reference norm is ``F.rms_norm``
+(fp32 internal statistics), so the fused path is numerically close (the norm
+statistics round through bf16) but **not bit-exact**. Following the
+fused-linear-GELU precedent, sites are therefore mounted only for
+``quality="high"`` requests via :func:`mount_fused_gate_rmsnorm` /
+:func:`unmount_fused_gate_rmsnorm` at batch boundaries; the default
+``"lossless"`` path keeps the unmodified reference chain bit-for-bit.
+
+Mounting is all-or-nothing per transformer: if any marked site fails the
+static guards (non-bf16 norm weight, hidden size above the kernel limit, ...)
+every site on that transformer stays on the reference path.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Iterator
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+# Attributes of the site protocol (set by ``mark_fused_gate_rmsnorm_site``).
+_SITE_NORM_ATTRS = "_sgl_fused_gate_rmsnorm_norm_attrs"
+_SITE_ENABLED_ATTR = "_sgl_fused_gate_rmsnorm_enabled"
+
+# The Triton kernels mask a single block over the hidden dim.
+_MAX_HIDDEN_SIZE = 8192
+
+
+def fused_rmsnorm_scale(
+    x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, eps: float
+) -> torch.Tensor | None:
+    """``RMSNorm(x, weight, eps) * scale`` in one Triton kernel (or None)."""
+    from sglang.kernels.ops.diffusion.triton.zimage_native_norm import (
+        zimage_rmsnorm_scale,
+    )
+
+    return zimage_rmsnorm_scale(x, weight, scale, eps)
+
+
+def fused_rmsnorm_tanh_residual(
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    """``residual + tanh(gate) * RMSNorm(x, weight, eps)`` fused (or None)."""
+    from sglang.kernels.ops.diffusion.triton.zimage_native_norm import (
+        zimage_rmsnorm_tanh_residual,
+    )
+
+    return zimage_rmsnorm_tanh_residual(x, gate, residual, weight, eps)
+
+
+def _static_reject_reason(site: nn.Module) -> str | None:
+    """Why ``site`` may never use the fused kernels, or None if it may."""
+    try:
+        import triton  # type: ignore # noqa: F401
+    except ImportError:
+        return "triton unavailable"
+    for attr in getattr(site, _SITE_NORM_ATTRS, ()):
+        norm = getattr(site, attr, None)
+        weight = getattr(norm, "weight", None)
+        if weight is None or weight.dim() != 1:
+            return f"{attr}: missing or non-1D norm weight"
+        if weight.dtype != torch.bfloat16:
+            return f"{attr}: non-bf16 norm weight dtype {weight.dtype}"
+        if weight.numel() > _MAX_HIDDEN_SIZE:
+            return f"{attr}: hidden size {weight.numel()} above kernel limit"
+    return None
+
+
+def mark_fused_gate_rmsnorm_site(module: nn.Module, norm_attrs: tuple[str, ...]):
+    """Declare ``module`` as a fused RMSNorm modulate/gate site.
+
+    ``norm_attrs`` names the site's RMSNorm submodules (checked by the static
+    guards at mount time). The site starts unmounted
+    (``_sgl_fused_gate_rmsnorm_enabled = False``): the module's forward must
+    keep the reference path bit-exact until :func:`mount_fused_gate_rmsnorm`
+    enables it.
+    """
+    setattr(module, _SITE_NORM_ATTRS, tuple(norm_attrs))
+    setattr(module, _SITE_ENABLED_ATTR, False)
+
+
+def iter_fused_gate_rmsnorm_sites(root: nn.Module) -> Iterator[nn.Module]:
+    """Yield every marked site under ``root`` (including ``root``)."""
+    for module in root.modules():
+        if getattr(module, _SITE_NORM_ATTRS, None) is not None:
+            yield module
+
+
+def mount_fused_gate_rmsnorm(root: nn.Module) -> bool:
+    """Enable the fused kernels on every marked site under ``root``.
+
+    All-or-nothing: if any marked site fails the static guards, every site is
+    left (or reset) on the reference path and False is returned. Returns False
+    as well when ``root`` has no marked sites.
+    """
+    sites = list(iter_fused_gate_rmsnorm_sites(root))
+    if not sites:
+        return False
+    for site in sites:
+        reason = _static_reject_reason(site)
+        if reason is not None:
+            unmount_fused_gate_rmsnorm(root)
+            logger.info(
+                "fused gate RMSNorm: %s site failed static guards (%s); "
+                "keeping the whole model on the reference path",
+                type(site).__name__,
+                reason,
+            )
+            return False
+    for site in sites:
+        setattr(site, _SITE_ENABLED_ATTR, True)
+    return True
+
+
+def unmount_fused_gate_rmsnorm(root: nn.Module) -> None:
+    """Reset every marked site under ``root`` to the bit-exact reference path."""
+    for site in iter_fused_gate_rmsnorm_sites(root):
+        setattr(site, _SITE_ENABLED_ATTR, False)

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

@@ -7,6 +7,11 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from sglang.kernels.ops.diffusion.fused_gate_rmsnorm import (
+    fused_rmsnorm_scale,
+    fused_rmsnorm_tanh_residual,
+    mark_fused_gate_rmsnorm_site,
+)
 from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig
 from sglang.multimodal_gen.runtime.distributed import (
     divide,
@@ -294,6 +299,46 @@ class Ideogram4MLP(nn.Module):
         return self.w2(F.silu(self.w1(x)) * self.w3(x))
 
 
+def _norm_scale(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    norm: Ideogram4RMSNorm,
+    enable_fused: bool,
+) -> torch.Tensor:
+    """``RMSNorm(x) * (1 + scale)``, fused for ``quality="high"`` batches."""
+    if enable_fused:
+        y = fused_rmsnorm_scale(
+            x,
+            norm.weight.data.to(device=x.device, dtype=x.dtype).contiguous(),
+            1.0 + scale,
+            norm.eps,
+        )
+        if y is not None:
+            return y
+    return norm(x) * (1.0 + scale)
+
+
+def _gate_residual(
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    residual: torch.Tensor,
+    norm: Ideogram4RMSNorm,
+    enable_fused: bool,
+) -> torch.Tensor:
+    """``residual + tanh(gate) * RMSNorm(x)``, fused for ``quality="high"``."""
+    if enable_fused:
+        y = fused_rmsnorm_tanh_residual(
+            x,
+            gate,
+            residual,
+            norm.weight.data.to(device=x.device, dtype=x.dtype).contiguous(),
+            norm.eps,
+        )
+        if y is not None:
+            return y
+    return residual + torch.tanh(gate) * norm(x)
+
+
 class Ideogram4TransformerBlock(nn.Module):
     def __init__(
         self,
@@ -328,6 +373,13 @@ class Ideogram4TransformerBlock(nn.Module):
         self.ffn_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
         self.attention_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
         self.ffn_norm2 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
+        # quality="high" fusion sites: each RMSNorm modulate/gate chain
+        # collapses into one Triton kernel (Z-Image bf16-native suite). Off by
+        # default (bit-exact reference path); mounted per batch by the
+        # denoising stage.
+        mark_fused_gate_rmsnorm_site(
+            self, ("attention_norm1", "attention_norm2", "ffn_norm1", "ffn_norm2")
+        )
         self.adaln_modulation = _linear(
             adaln_dim,
             4 * hidden_size,
@@ -341,20 +393,21 @@ class Ideogram4TransformerBlock(nn.Module):
         scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaln_modulation(
             adaln_input
         ).chunk(4, dim=-1)
-        gate_msa = torch.tanh(gate_msa)
-        gate_mlp = torch.tanh(gate_mlp)
+        enable_fused = (
+            self._sgl_fused_gate_rmsnorm_enabled and not torch.compiler.is_compiling()
+        )
         attn_out = self.attention(
-            self.attention_norm1(x) * (1.0 + scale_msa),
+            _norm_scale(x, scale_msa, self.attention_norm1, enable_fused),
             cos=cos,
             sin=sin,
             attn_mask=attn_mask,
             attn_mask_meta=attn_mask_meta,
         )
-        x = x + gate_msa * self.attention_norm2(attn_out)
-        x = x + gate_mlp * self.ffn_norm2(
-            self.feed_forward(self.ffn_norm1(x) * (1.0 + scale_mlp))
+        x = _gate_residual(attn_out, gate_msa, x, self.attention_norm2, enable_fused)
+        ffn_out = self.feed_forward(
+            _norm_scale(x, scale_mlp, self.ffn_norm1, enable_fused)
         )
-        return x
+        return _gate_residual(ffn_out, gate_mlp, x, self.ffn_norm2, enable_fused)
 
 
 def _sinusoidal_embedding(t: torch.Tensor, dim: int, scale: float = 1e4):

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -20,6 +20,10 @@ from typing import Any
 import torch
 import torch.nn as nn
 
+from sglang.kernels.ops.diffusion.fused_gate_rmsnorm import (
+    mount_fused_gate_rmsnorm,
+    unmount_fused_gate_rmsnorm,
+)
 from sglang.kernels.ops.diffusion.fused_linear_gelu import (
     mount_fused_linear_gelu,
     unmount_fused_linear_gelu,
@@ -226,9 +230,9 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         # cache-dit state (for delayed mounting and idempotent control)
         self._cache_dit_enabled = False
         self._cached_num_steps = None
-        # fused linear+GELU state: whether the cublasLt-epilogue fusion is
-        # currently mounted on the transformers (quality="high" batches only).
-        self._fused_gelu_mounted = False
+        # quality="high" fusion state: whether the cublasLt linear+GELU and
+        # fused gate-RMSNorm sites are currently mounted on the transformers.
+        self._quality_fusions_mounted = False
         self._torch_compile_registry = CompiledModuleRegistry()
         # Breakable CUDA graph runners, one per transformer module (lazy).
         self._bcg_runners: dict[int, Any] = {}
@@ -450,38 +454,47 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         self, num_inference_steps: int | tuple[int, int], batch: Req
     ) -> None:
         """Apply request-dependent transformer acceleration in trace-safe order."""
-        self._maybe_toggle_fused_gelu(batch)
+        self._maybe_toggle_quality_fusions(batch)
         self._maybe_enable_cache_dit(num_inference_steps, batch)
         for transformer in filter(None, [self.transformer, self.transformer_2]):
             self._maybe_torch_compile(transformer)
 
-    def _maybe_toggle_fused_gelu(self, batch: Req) -> None:
-        """Mount/unmount the cublasLt linear+GELU fusion for this batch.
+    def _maybe_toggle_quality_fusions(self, batch: Req) -> None:
+        """Mount/unmount the ``quality="high"`` fusions for this batch.
 
-        The fused epilogue is numerically equivalent only at half-precision
-        rounding level (not bit-exact), so it is mounted for
-        ``quality="high"`` requests and unmounted otherwise -- the
-        ``"lossless"`` default runs the unmodified reference path bit-for-bit.
-        ``quality`` participates in the dynamic-batch signature, so a worker
-        batch is uniform in ``quality`` and this process-wide transition is
-        safe at the batch boundary. Mounting is all-or-nothing per
-        transformer (any ineligible marked site keeps the whole transformer
-        on the reference path); models without marked sites are no-ops.
+        The cublasLt linear+GELU epilogue and the fused gate-RMSNorm Triton
+        kernels are numerically equivalent only at half-precision rounding
+        level (not bit-exact), so they are mounted for ``quality="high"``
+        requests and unmounted otherwise -- the ``"lossless"`` default runs
+        the unmodified reference path bit-for-bit. ``quality`` participates
+        in the dynamic-batch signature, so a worker batch is uniform in
+        ``quality`` and this process-wide transition is safe at the batch
+        boundary. Mounting is all-or-nothing per transformer and per fusion
+        family (any ineligible marked site keeps the whole transformer on
+        that family's reference path); models without marked sites are
+        no-ops.
         """
         want = getattr(batch.sampling_params, "quality", "lossless") == "high"
-        if want == self._fused_gelu_mounted:
+        if want == self._quality_fusions_mounted:
             return
-        mounted = False
+        mounted_gelu = False
+        mounted_gate_norm = False
         for transformer in filter(None, [self.transformer, self.transformer_2]):
             if want:
-                mounted |= mount_fused_linear_gelu(transformer)
+                mounted_gelu |= mount_fused_linear_gelu(transformer)
+                mounted_gate_norm |= mount_fused_gate_rmsnorm(transformer)
             else:
                 unmount_fused_linear_gelu(transformer)
-        self._fused_gelu_mounted = want
-        if want and mounted:
+                unmount_fused_gate_rmsnorm(transformer)
+        self._quality_fusions_mounted = want
+        if want and mounted_gelu:
             logger.info(
                 "Mounted fused linear+GELU (cublasLt epilogue) for quality=high"
             )
+        if want and mounted_gate_norm:
+            logger.info(
+                "Mounted fused gate RMSNorm (Z-Image Triton suite) for quality=high"
+            )
 
     def _cache_dit_dual_model_name(self) -> str:
         return "wan2.2"

test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py

@@ -0,0 +1,54 @@
+"""Core checks for the quality-gated fused gate-RMSNorm (Z-Image suite reuse)."""
+
+import sys
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.kernels.ops.diffusion import fused_gate_rmsnorm as fgn
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+DIM, EPS = 4608, 1e-5  # Ideogram 4 hidden size / norm_eps
+
+
+class _Site(nn.Module):
+    def __init__(self, dtype=torch.bfloat16):
+        super().__init__()
+        self.norm = nn.RMSNorm(DIM, eps=EPS, device="cuda", dtype=dtype)
+        fgn.mark_fused_gate_rmsnorm_site(self, ("norm",))
+
+
+def test_fused_matches_ideogram_reference():
+    torch.manual_seed(0)
+    site = _Site()
+    w = site.norm.weight.data
+    x = torch.randn(1, 64, DIM, device="cuda", dtype=torch.bfloat16)
+    residual = torch.randn_like(x)
+    # adaln-style strided chunks, as produced by Ideogram's modulation .chunk()
+    mods = torch.randn(1, 1, 2 * DIM, device="cuda", dtype=torch.bfloat16)
+    scale, gate = mods.chunk(2, dim=-1)
+    assert fgn.mount_fused_gate_rmsnorm(site)
+    got_scale = fgn.fused_rmsnorm_scale(x, w, 1.0 + scale, EPS)
+    got_gate = fgn.fused_rmsnorm_tanh_residual(x, gate, residual, w, EPS)
+    ref_scale = F.rms_norm(x, (DIM,), w, EPS) * (1.0 + scale)
+    ref_gate = residual + torch.tanh(gate) * F.rms_norm(x, (DIM,), w, EPS)
+    # fused path uses bf16-native norm statistics: close, not bit-exact
+    torch.testing.assert_close(got_scale, ref_scale, atol=8e-2, rtol=4e-2)
+    torch.testing.assert_close(got_gate, ref_gate, atol=8e-2, rtol=4e-2)
+
+
+def test_mount_guards_all_or_nothing():
+    good, bad = _Site(), _Site(torch.float32)
+    assert not fgn.mount_fused_gate_rmsnorm(nn.ModuleList([good, bad]))
+    assert not good._sgl_fused_gate_rmsnorm_enabled
+    assert fgn.mount_fused_gate_rmsnorm(good)
+    fgn.unmount_fused_gate_rmsnorm(good)
+    assert not good._sgl_fused_gate_rmsnorm_enabled
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))