diff --git a/python/sglang/kernels/ops/diffusion/triton/wan_rmsnorm_silu.py b/python/sglang/kernels/ops/diffusion/triton/wan_rmsnorm_silu.py new file mode 100644 index 000000000..275390cf7 --- /dev/null +++ b/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"] diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py b/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py index ee961ae78..ae6951d10 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/flux2_vae_cuda_opt.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py b/python/sglang/multimodal_gen/runtime/models/vaes/wan_vae_cuda_opt.py new file mode 100644 index 000000000..bd3934430 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 6fb73bb25..86613223f 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/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 diff --git a/test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py b/test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py new file mode 100644 index 000000000..b45ced7a5 --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py @@ -0,0 +1,70 @@ +"""Wan VAE decoder fast path: fused-kernel numerics and gate dispatch +(the lossless off-path must stay bit-exact).""" + +import sys + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from sglang.kernels.ops.diffusion.triton.wan_rmsnorm_silu import wan_rmsnorm_silu +from sglang.multimodal_gen.runtime.models.vaes.wan_vae_cuda_opt import ( + FusedWanRMSNormSiLU, + VaeFastPathGate, +) +from sglang.multimodal_gen.runtime.models.vaes.wanvae import WanRMS_norm +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _cl3d(shape, dtype): + return torch.randn(shape, device="cuda", dtype=dtype).contiguous( + memory_format=torch.channels_last_3d + ) + + +@torch.no_grad() +@pytest.mark.parametrize( + "x_dtype,affine_dtype,atol,rtol", + [ + (torch.float32, torch.float32, 1e-5, 1e-5), # FastWan2.2 fp32 decode + (torch.bfloat16, torch.float32, 1.5e-1, 3e-2), # Wan2.1 bf16 autocast + ], +) +def test_kernel_numerics(x_dtype, affine_dtype, atol, rtol) -> None: + torch.cuda.manual_seed(0) + x = _cl3d((1, 96, 3, 10, 14), x_dtype) + gamma = torch.randn((96, 1, 1, 1), device="cuda", dtype=affine_dtype) + for bias in (None, torch.randn_like(gamma)): + expected = F.silu( + F.normalize(x, dim=1) * 96**0.5 * gamma + (0 if bias is None else bias) + ) + actual = wan_rmsnorm_silu(x, gamma, bias) + assert actual is not None and actual.dtype == expected.dtype + assert actual.stride() == x.stride() + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + + +@torch.no_grad() +def test_fused_module_gate_dispatch() -> None: + # Gate off must stay bit-exact; gate on must route to the fused kernel. + torch.cuda.manual_seed(0) + norm = WanRMS_norm(96, images=False).to(device="cuda", dtype=torch.bfloat16) + norm.gamma.add_(torch.randn_like(norm.gamma)) + gate = VaeFastPathGate() + fused = FusedWanRMSNormSiLU(norm, gate) + # Parameter names must not change (weight transfer matches by name). + assert [n for n, _ in fused.named_parameters()] == ["gamma"] + x = _cl3d((1, 96, 3, 10, 14), torch.bfloat16) + assert torch.equal(fused(x), nn.SiLU()(norm(x))) + gate.enabled = True + expected = wan_rmsnorm_silu(x, norm.gamma, rms_scale=float(norm.scale)) + assert torch.equal(fused(x), expected) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))