[diffusion] Wan VAE RMSNorm+SiLU fusion behind quality=high (H200 FastWan2.2 e2e 9.611 -> 9.125 s) (#33546)

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

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

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Channels-last-3d Wan VAE RMSNorm(+SiLU) Triton kernel.
+
+Fuses the Wan VAE ``WanRMS_norm -> SiLU`` chain
+(``SiLU(F.normalize(x, dim=1) * scale * gamma + bias)`` on channel-first 5D
+activations) into one kernel for ``channels_last_3d`` tensors: one program
+reduces one (b, t, h, w) pixel's channel row with fully coalesced loads.
+
+Numerics contract: fp32 channel-norm statistics, every step materialized at
+the same dtype boundary as eager ``WanRMS_norm.forward`` (including the
+aten promotion to fp32 at ``* gamma`` for half-precision x with fp32 affine
+params -- the autocast case), SiLU in fp32. Bitwise equality with aten is
+still not guaranteed (different reduction and SiLU paths), so callers must
+keep this behind an opt-in gate. ``wan_rmsnorm_silu`` returns ``None`` for
+unsupported inputs (see ``can_use_wan_rmsnorm_silu``); callers must fall
+back to their reference path.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.srt.utils.custom_op import register_custom_op
+
+_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
+_MAX_CHANNELS = 1024
+
+
+@triton.jit
+def _wan_rmsnorm_silu_kernel(
+    x_ptr,
+    gamma_ptr,
+    bias_ptr,
+    out_ptr,
+    channels: tl.constexpr,
+    t_size,
+    h_size,
+    w_size,
+    x_stride_b,
+    x_stride_c,
+    x_stride_t,
+    x_stride_h,
+    x_stride_w,
+    out_stride_b,
+    out_stride_c,
+    out_stride_t,
+    out_stride_h,
+    out_stride_w,
+    rms_scale,
+    eps,
+    has_bias: tl.constexpr,
+    block_c: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, block_c)
+    mask = offsets < channels
+
+    w = row % w_size
+    tmp = row // w_size
+    h = tmp % h_size
+    tmp = tmp // h_size
+    t = tmp % t_size
+    b = tmp // t_size
+
+    x_base = b * x_stride_b + t * x_stride_t + h * x_stride_h + w * x_stride_w
+    out_base = b * out_stride_b + t * out_stride_t + h * out_stride_h + w * out_stride_w
+
+    x = tl.load(x_ptr + x_base + offsets * x_stride_c, mask=mask, other=0.0).to(
+        tl.float32
+    )
+    norm = tl.sqrt(tl.sum(x * x, axis=0))
+    inv_norm = 1.0 / tl.maximum(norm, eps)
+
+    # Eager op boundaries: normalize/*scale in x.dtype; *gamma/+bias in the
+    # promoted output dtype; SiLU in fp32, stored in the output dtype.
+    y = (x * inv_norm).to(x_ptr.dtype.element_ty)
+    gamma = tl.load(gamma_ptr + offsets, mask=mask, other=1.0)
+    y = (y * rms_scale).to(x_ptr.dtype.element_ty)
+    y = (y.to(tl.float32) * gamma.to(tl.float32)).to(out_ptr.dtype.element_ty)
+    if has_bias:
+        bias = tl.load(bias_ptr + offsets, mask=mask, other=0.0)
+        y = (y.to(tl.float32) + bias.to(tl.float32)).to(out_ptr.dtype.element_ty)
+    y = y.to(tl.float32)
+    y = y * tl.sigmoid(y)
+
+    tl.store(out_ptr + out_base + offsets * out_stride_c, y, mask=mask)
+
+
+def _fake_wan_rmsnorm_silu(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor,
+    rms_scale: float,
+    eps: float,
+    has_bias: bool,
+) -> torch.Tensor:
+    dtype = torch.promote_types(x.dtype, gamma.dtype)
+    return torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=dtype)
+
+
+@register_custom_op(
+    op_name="triton_wan_rmsnorm_silu_cuda",
+    fake_impl=_fake_wan_rmsnorm_silu,
+)
+def _triton_wan_rmsnorm_silu_cuda(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor,
+    rms_scale: float,
+    eps: float,
+    has_bias: bool,
+) -> torch.Tensor:
+    bsz, channels, t_size, h_size, w_size = x.shape
+    # Preserve the input strides so the VAE keeps its channels_last_3d layout.
+    dtype = torch.promote_types(x.dtype, gamma.dtype)
+    out = torch.empty_strided(x.shape, x.stride(), device=x.device, dtype=dtype)
+    block_c = triton.next_power_of_2(channels)
+    num_warps = 1 if block_c <= 64 else 4 if block_c <= 512 else 8
+
+    with torch.cuda.device(x.device):
+        _wan_rmsnorm_silu_kernel[(bsz * t_size * h_size * w_size,)](
+            x,
+            gamma,
+            bias,
+            out,
+            channels,
+            t_size,
+            h_size,
+            w_size,
+            *x.stride(),
+            *out.stride(),
+            rms_scale,
+            eps,
+            has_bias,
+            block_c,
+            num_warps=num_warps,
+        )
+    return out
+
+
+def _affine_supported(x: torch.Tensor, t: torch.Tensor) -> bool:
+    # Same dtype, or fp32 affine params on half-precision x (autocast case).
+    return (
+        t.is_cuda
+        and t.device == x.device
+        and (t.dtype == x.dtype or t.dtype == torch.float32)
+        and t.numel() == x.shape[1]
+    )
+
+
+def can_use_wan_rmsnorm_silu(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None,
+) -> bool:
+    return (
+        x.is_cuda
+        and not torch.is_grad_enabled()
+        and not x.requires_grad
+        and x.dtype in _SUPPORTED_DTYPES
+        and x.ndim == 5
+        and 0 < x.shape[1] <= _MAX_CHANNELS
+        and x.is_contiguous(memory_format=torch.channels_last_3d)
+        and _affine_supported(x, gamma)
+        and (bias is None or _affine_supported(x, bias))
+    )
+
+
+def wan_rmsnorm_silu(
+    x: torch.Tensor,
+    gamma: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    rms_scale: float | None = None,
+    eps: float = 1e-12,
+) -> torch.Tensor | None:
+    """Fused ``SiLU(F.normalize(x, dim=1) * rms_scale * gamma + bias)``.
+
+    Returns ``None`` when the input is unsupported; callers must fall back.
+    """
+    if not can_use_wan_rmsnorm_silu(x, gamma, bias):
+        return None
+
+    channels = x.shape[1]
+    gamma = gamma.reshape(channels).contiguous()
+    has_bias = bias is not None
+    bias = gamma if bias is None else bias.reshape(channels).contiguous()
+    if rms_scale is None:
+        rms_scale = channels**0.5
+    return _triton_wan_rmsnorm_silu_cuda(
+        x, gamma, bias, float(rms_scale), eps, has_bias
+    )
+
+
+__all__ = ["can_use_wan_rmsnorm_silu", "wan_rmsnorm_silu"]

python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py

@@ -1,36 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
-"""CUDA fast paths for the FLUX.2 VAE decoder (AutoencoderKLFlux2, diffusers
-``Decoder``) on image workloads.
+"""CUDA fast paths for the FLUX.2 VAE decoder (AutoencoderKLFlux2).
 
 All rewrites are mathematically exact re-associations of the original
-operators (weight folding is done lazily in fp32 on first fast-path use and
-written back to the model compute dtype). The wrappers are installed once at
-VAE load and stay in place; each forward dispatches on a shared
-request-scoped :class:`VaeFastPathGate`: requests with ``quality == "high"``
-run the fast paths, the ``"lossless"`` default runs the original module path
-bit-for-bit.
+operators. Wrappers are installed once at VAE load and dispatch on a
+request-scoped :class:`VaeFastPathGate` (published as
+``_sgl_vae_fast_path_gate``): ``quality == "high"`` runs the fast paths, the
+``"lossless"`` default runs the original module path bit-for-bit.
 
-- channels_last: run the decoder in channels_last so cuDNN convolutions run
-  natively in NHWC (removes the nchwToNhwc/nhwcToNchw transpose kernels
-  around every conv). The parameter layout is swapped at decode entry to
-  match the gate, so lossless decodes always run the NCHW baseline kernels
-  bit-for-bit. The mid-block attention needs a layout-safe forward because
-  diffusers' ``AttnProcessor2_0`` calls ``.view`` on the 4D activation,
-  which is illegal for channels_last tensors.
-- norm+SiLU: two-pass channels_last GroupNorm(+SiLU) Triton fusion (fp32
-  statistics) for the ResnetBlock2D norm1/norm2 + SiLU chains and the
-  decoder ``conv_norm_out``/``conv_act`` tail, which upcast to fp32 under
-  autocast and dominate the decode profile.
-- fused upsample: nearest-2x upsample + Conv2d(3x3, p1) ==
-  ConvTranspose2d(k4, s2, p1) with a lazily-summed kernel. Removes the 4x
-  upsampled intermediate materialization.
-- attention V/proj fold: fold the attention output projection into the V
-  projection of the single-head mid-block attention (softmax rows sum to 1,
-  so ``A @ (V W_v^T + b_v) W_o^T + b_o == A @ (V W_v'^T + b')``).
+- channels_last: run the decoder in NHWC so cuDNN convs skip the transpose
+  kernels; parameter layout is swapped at decode entry to match the gate.
+- norm+SiLU: two-pass channels_last GroupNorm(+SiLU) Triton fusion.
+- fused upsample: nearest-2x + Conv2d(3x3, p1) == ConvTranspose2d(k4, s2, p1).
+- attention V/proj fold: softmax rows sum to 1, so
+  ``A @ (V W_v^T + b_v) W_o^T + b_o == A @ (V W_v'^T + b')``.
 
-Install is all-or-nothing and fail-closed: without Triton, or with any
-attention block lacking the layout-safe rewrite, no wrapper is installed and
-every request runs the unmodified decoder.
+Install is all-or-nothing and fail-closed.
 """
 
 from types import MethodType
@@ -55,12 +39,8 @@ except ImportError:  # pragma: no cover
 
 
 class VaeFastPathGate:
-    """Mutable fast-path flag shared by every wrapper of one VAE.
-
-    Published on the VAE as ``_sgl_vae_fast_path_gate``; ``DecodingStage``
-    enables it for the duration of a decode when the request's ``quality``
-    sampling param is ``"high"``.
-    """
+    """Mutable fast-path flag shared by every wrapper of one VAE; enabled by
+    ``DecodingStage`` while decoding a ``quality == "high"`` request."""
 
     __slots__ = ("enabled",)
 
@@ -77,13 +57,9 @@ GATE_ATTR = "_sgl_vae_fast_path_gate"
 
 
 class FusedGroupNormSiLU(nn.Module):
-    """GroupNorm + SiLU fused with the two-pass channels_last Triton kernel.
-
-    fp32 statistics and affine/SiLU application, output in the input dtype.
-    Falls back to the original module chain (norm + F.silu, bit-identical to
-    the original norm + nn.SiLU pair) for unsupported inputs and whenever
-    the fast-path gate is disabled.
-    """
+    """GroupNorm + SiLU fused with the two-pass channels_last Triton kernel;
+    falls back to the original op chain (bit-identical to norm + ``nn.SiLU``)
+    for unsupported inputs and whenever the gate is off."""
 
     def __init__(self, norm: nn.GroupNorm, gate: VaeFastPathGate) -> None:
         super().__init__()
@@ -181,12 +157,8 @@ def _fold_upsample2x_conv2d_weight(conv: nn.Conv2d) -> torch.Tensor:
 
 class FusedUpsample2xConv2d(nn.Module):
     """ConvTranspose2d(k4, s2, p1) equivalent of diffusers Upsample2D
-    (nearest-2x interpolate + Conv2d(3x3, p1)).
-
-    nearest 2x upsampling is pure pixel replication (no arithmetic), so the
-    fusion only re-associates the conv taps; the kernel is summed lazily in
-    fp32 on first fast-path use and written back to the conv dtype. With the
-    fast-path gate disabled the original Upsample2D runs bit-for-bit.
+    (nearest-2x interpolate + Conv2d(3x3, p1)); the kernel is summed lazily
+    in fp32 on first use. Gate off runs the original Upsample2D bit-for-bit.
     """
 
     def __init__(self, upsample: nn.Module, gate: VaeFastPathGate) -> None:

python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CUDA fast path for the Wan VAE decoder (AutoencoderKLWan).
+
+Fuses every decoder ``WanRMS_norm -> SiLU`` chain into one Triton kernel on
+the channels_last_3d layout. Wrappers are installed once at VAE load and
+dispatch on a request-scoped :class:`VaeFastPathGate` (published as
+``_sgl_vae_fast_path_gate``): ``quality == "high"`` runs the fused kernel
+(not bitwise-identical to aten, hence gated), the ``"lossless"`` default
+runs the original module path bit-for-bit. Install is all-or-nothing and
+fail-closed.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.models.vaes.flux2_vae_cuda_opt import (
+    GATE_ATTR,
+    VaeFastPathGate,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+try:
+    from sglang.kernels.ops.diffusion.triton.wan_rmsnorm_silu import wan_rmsnorm_silu
+
+    _HAS_TRITON = True
+except ImportError:  # pragma: no cover
+    _HAS_TRITON = False
+
+
+class FusedWanRMSNormSiLU(nn.Module):
+    """``WanRMS_norm`` + SiLU fused via the channels_last_3d Triton kernel;
+    falls back to the original op chain (bit-identical to norm + ``nn.SiLU``)
+    for unsupported inputs and whenever the gate is off. Steps aside under
+    ``torch.compile``, where Inductor already fuses this chain."""
+
+    def __init__(self, norm: nn.Module, gate: VaeFastPathGate) -> None:
+        super().__init__()
+        # Keep the norm's parameters registered directly on the wrapper so
+        # parameter names stay `...norm1.gamma` (weight transfer and
+        # state_dict load match by name).
+        self.gamma = norm.gamma
+        self.bias = norm.bias
+        self.scale = float(norm.scale)
+        self._sgl_gate = gate
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._sgl_gate.enabled and not torch.compiler.is_compiling():
+            bias = self.bias if isinstance(self.bias, torch.Tensor) else None
+            y = wan_rmsnorm_silu(x, self.gamma, bias, rms_scale=self.scale)
+            if y is not None:
+                return y
+        # WanRMS_norm.forward (channel-first) + SiLU, same ops in the same
+        # order, so the off-path stays bit-identical.
+        return F.silu(F.normalize(x, dim=1) * self.scale * self.gamma + self.bias)
+
+
+def _is_plain_silu(act: object) -> bool:
+    return isinstance(act, nn.SiLU) and not act.inplace
+
+
+def _norm_fusable(norm: object, wan_rms_norm_cls: type) -> bool:
+    return (
+        type(norm) is wan_rms_norm_cls
+        and getattr(norm, "channel_first", False)
+        and isinstance(getattr(norm, "gamma", None), torch.Tensor)
+    )
+
+
+def _install_norm_silu(decoder: nn.Module, gate: VaeFastPathGate) -> int | None:
+    """Wrap every decoder ``WanRMS_norm -> SiLU`` chain; ``None`` = fail closed."""
+    from sglang.multimodal_gen.runtime.models.vaes.wanvae import (
+        WanResidualBlock,
+        WanRMS_norm,
+    )
+
+    res_blocks = [m for m in decoder.modules() if isinstance(m, WanResidualBlock)]
+    eligible = [
+        m
+        for m in res_blocks
+        if type(m) is WanResidualBlock
+        and _is_plain_silu(m.nonlinearity)
+        and _norm_fusable(m.norm1, WanRMS_norm)
+        and _norm_fusable(m.norm2, WanRMS_norm)
+    ]
+    if len(eligible) != len(res_blocks):
+        logger.warning(
+            "Wan VAE: %d/%d residual blocks non-standard; skipping fast path.",
+            len(res_blocks) - len(eligible),
+            len(res_blocks),
+        )
+        return None
+    if not (
+        _norm_fusable(getattr(decoder, "norm_out", None), WanRMS_norm)
+        and _is_plain_silu(getattr(decoder, "nonlinearity", None))
+    ):
+        logger.warning("Wan VAE: non-standard output head; skipping fast path.")
+        return None
+
+    count = 0
+    for m in eligible:
+        m.norm1 = FusedWanRMSNormSiLU(m.norm1, gate)
+        m.norm2 = FusedWanRMSNormSiLU(m.norm2, gate)
+        m.nonlinearity = nn.Identity()
+        count += 2
+    decoder.norm_out = FusedWanRMSNormSiLU(decoder.norm_out, gate)
+    decoder.nonlinearity = nn.Identity()
+    return count + 1
+
+
+def maybe_optimize_wan_vae(vae: nn.Module) -> nn.Module:
+    """Install the quality-gated CUDA Wan VAE decoder fast path."""
+    from sglang.multimodal_gen.runtime.models.vaes.wanvae import (
+        AutoencoderKLWan,
+        WanDecoder3d,
+    )
+
+    if not isinstance(vae, AutoencoderKLWan):
+        return vae
+    decoder = getattr(vae, "decoder", None)
+    if type(decoder) is not WanDecoder3d:
+        return vae
+    if decoder.use_parallel_decode and decoder.world_size > 1:
+        logger.info("Wan VAE: spatial-parallel decode; skipping fast path.")
+        return vae
+    if not _HAS_TRITON:
+        logger.warning("Wan VAE: Triton unavailable; skipping fast path.")
+        return vae
+
+    gate = VaeFastPathGate()
+    n_norm = _install_norm_silu(decoder, gate)
+    if n_norm is None:
+        return vae
+    setattr(vae, GATE_ATTR, gate)
+    logger.info(
+        "Wan VAE: installed quality-gated fast path (%d RMSNorm+SiLU fusions).",
+        n_norm,
+    )
+    return vae

python/sglang/multimodal_gen/runtime/platforms/cuda.py

@@ -539,22 +539,25 @@ class CudaPlatformBase(Platform):
 
     @classmethod
     def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module:
-        """Install the quality-gated FLUX.2 VAE decoder fast paths.
+        """Install the quality-gated FLUX.2 / Wan VAE decoder fast paths.
 
         Requests with quality == "high" run the fast paths; the "lossless"
         default runs the original module path bit-for-bit. See
-        flux2_vae_cuda_opt for details.
+        flux2_vae_cuda_opt and wan_vae_cuda_opt for details.
         """
         try:
             from sglang.multimodal_gen.runtime.models.vaes.flux2_vae_cuda_opt import (
                 maybe_optimize_flux2_vae,
             )
+            from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import (
+                maybe_optimize_wan_vae,
+            )
 
             vae = maybe_optimize_flux2_vae(vae)
+            vae = maybe_optimize_wan_vae(vae)
         except Exception:
             logger.warning(
-                "Failed to apply CUDA FLUX.2 VAE optimizations; "
-                "using the unmodified VAE.",
+                "Failed to apply CUDA VAE optimizations; using the unmodified VAE.",
                 exc_info=True,
             )
         return vae