From 295784723a6b6c47e1ecc3c6acb212f4246a39e9 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Thu, 6 Aug 2026 19:57:34 +0800 Subject: [PATCH] [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 --- .../ops/diffusion/fused_gate_rmsnorm.py | 135 ++++++++++++++++++ .../runtime/models/dits/ideogram.py | 67 ++++++++- .../pipelines_core/stages/denoising.py | 53 ++++--- .../ops/diffusion/test_fused_gate_rmsnorm.py | 54 +++++++ 4 files changed, 282 insertions(+), 27 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/fused_gate_rmsnorm.py create mode 100644 test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py diff --git a/python/sglang/kernels/ops/diffusion/fused_gate_rmsnorm.py b/python/sglang/kernels/ops/diffusion/fused_gate_rmsnorm.py new file mode 100644 index 000000000..c8313cbf1 --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py b/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py index 2b13484a4..69666454d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py +++ b/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): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index bc1cb5997..edfddcd64 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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" diff --git a/test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py b/test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py new file mode 100644 index 000000000..5a300580b --- /dev/null +++ b/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__]))