[Diffusion] Keep the Wan VAE decoder channels_last and add a Triton NHWC nearest upsample (#38182)

Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

python/sglang/kernels/ops/diffusion/README.md

@@ -161,13 +161,20 @@ tensor copy per residual site.
 `usp_merge_heads`, `pack_qkv_destination_major`, `fused_pack_qkv`,
 `fused_pack_segmented_qkv`, `fused_scatter_to_padded`,
 `fused_causal_conv3d_cat_pad_cuda`,
-`cat_pad_channels_last_3d`, `dup_up3d_add`, `fused_temb_table_slices`,
+`cat_pad_channels_last_3d`, `dup_up3d_add`, `nearest_upsample_nhwc`,
+`fused_temb_table_slices`,
 and `ltx2_ada_values9` are bit-exact data movement or same-order arithmetic.
 `fused_layernorm_modulate_fp8_quant_raw` folds FLUX.2 LayerNorm, adaLN
 modulation, and static FP8 quantization. `try_flux2_token_cat_fp8` and
 `try_flux2_token_cat_nvfp4` fuse branch concatenation directly into the
 quantized representation selected by the FLUX.2 checkpoint path.
 
+`nearest_upsample_nhwc` replaces `nn.Upsample(nearest / nearest-exact,
+integer factor)` on a dense channels_last input with a Triton gather: same
+values and layout as aten, but aten's own NHWC nearest kernel is several times
+slower than its NCHW sibling, which is what the Wan-family VAE decoders hit
+once they run channels_last end-to-end.
+
 `fused_temb_table_slices` is worth knowing about: the eager
 `(table + temb.float()).chunk(6, dim=2)` materializes ~8 GB of fp32 at
 704p/121f *and* hands six strided slices downstream, whose `.contiguous()`

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

@@ -398,6 +398,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
         _CUDA,
         "Wan causal VAE main + DupUp3D(src).",
     ),
+    (
+        "diffusion.nearest_upsample_nhwc",
+        KernelBackend.TRITON,
+        "layout.nearest_upsample_nhwc_triton:nearest_upsample_nhwc",
+        _CUDA,
+        "Wan-family VAE channels_last integer-factor nearest upsample.",
+    ),
     (
         "diffusion.flux2_token_cat_nvfp4",
         KernelBackend.JIT,
@@ -542,6 +549,8 @@ _EXPORTS: dict[str, str] = {
     "fused_scatter_to_padded": "layout.varlen_pack_pad_triton",
     "cat_pad_channels_last_3d": "layout.wan_causal_cache_triton",
     "dup_up3d_add": "layout.wan_causal_cache_triton",
+    "nearest_upsample_nhwc": "layout.nearest_upsample_nhwc_triton",
+    "can_use_nearest_upsample_nhwc": "layout.nearest_upsample_nhwc_triton",
     "try_flux2_token_cat_nvfp4": "layout.flux2_token_cat_nvfp4_jit",
     # Fusion-site policy: quality gate, first-sight verification, mount
     "BitExactFusionGate": "sites.bitexact_gate",

python/sglang/kernels/ops/diffusion/layout/nearest_upsample_nhwc_triton.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Bit-exact channels_last nearest upsample for the Wan-family VAE decoders.
+
+``nn.Upsample(scale_factor=2, mode="nearest-exact")`` on a channels_last
+(NHWC) input dispatches to aten's ``upsample_nearest2d_nhwc_out_frame``. On an
+H200 with a ``[1, 192, 240, 416]`` bf16 input that kernel takes 0.458 ms
+against 0.190 ms for aten's own NCHW kernel on the same bytes; this gather
+takes 0.061 ms. The Wan / Qwen-Image VAE decoders hit the aten NHWC kernel
+once per up block per chunk as soon as they run channels_last end-to-end.
+
+Numerical contract: bit-exact vs ``nn.Upsample`` in ``nearest`` and
+``nearest-exact`` mode for integer scale factors. Write an output index as
+``i = k * f + r`` with ``0 <= r <= f - 1``. ``nearest`` reads ``floor(i / f)
+= k``; ``nearest-exact`` reads ``floor((i + 0.5) / f) = floor(k + (r + 0.5) /
+f) = k`` because ``(r + 0.5) / f < 1``. Both therefore read input ``i // f``,
+so the op is a pure gather that never touches a value, and the result is
+bitwise identical for any dtype the predicate admits (bf16 / fp16 / fp32, the
+ones the tests cover). The kernel walks the output in its NHWC memory order,
+so the stores and the gathered loads are both contiguous along ``C``.
+
+Layout contract: the output is dense channels_last, which is what aten returns
+exactly when ``suggest_memory_format()`` says channels_last. That requires
+``C > 1`` (with ``C == 1`` the tensor is also NCHW-contiguous and aten picks the
+NCHW kernel, returning e.g. strides ``(4, 4, 2, 1)`` for a ``[1, 1, 2, 2]``
+output) and canonical NHWC strides ``(H*W*C, 1, W*C, C)`` on every dim,
+including size-1 dims (aten's stride test does not skip them, unlike
+``is_contiguous``). The predicate enforces both, so a call it admits is
+value- and layout-identical to ``nn.Upsample``.
+
+Verified (``torch.equal`` vs ``F.interpolate``): ``[1, 192, 240, 416]``,
+``[4, 192, 120, 208]``, ``[1, 96, 480, 832]``, ``[1, 3, 5, 7]`` at factor 2
+and ``[2, 3, 5, 7]`` at factor ``(3, 2)`` for all three dtypes; end-to-end
+inside the Wan 2.1 (81 frames, 480x832) and Qwen-Image (1024x1024) decoders.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+_MAX_INT32 = 2**31 - 1
+_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
+
+
+@triton.jit
+def _nearest_upsample_nhwc_kernel(
+    x_ptr,
+    out_ptr,
+    total,
+    C,
+    out_h,
+    out_w,
+    fh,
+    fw,
+    sxn,
+    sxh,
+    sxw,
+    IDX64: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    if IDX64:
+        offs = pid.to(tl.int64) * BLOCK + tl.arange(0, BLOCK).to(tl.int64)
+    else:
+        offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < total
+    # Output is dense NHWC: offs = ((n * out_h + h) * out_w + w) * C + c.
+    c = offs % C
+    t = offs // C
+    w = t % out_w
+    t = t // out_w
+    h = t % out_h
+    n = t // out_h
+    src = n * sxn + (h // fh) * sxh + (w // fw) * sxw + c
+    vals = tl.load(x_ptr + src, mask=mask)
+    tl.store(out_ptr + offs, vals, mask=mask)
+
+
+def _integer_scale(scale) -> tuple[int, int] | None:
+    """``(fh, fw)`` when ``scale`` is a finite integer-valued factor (scalar or
+    pair) of at least 1; ``None`` for anything else, never an exception."""
+    if isinstance(scale, bool):
+        return None
+    if isinstance(scale, (int, float)):
+        scale = (scale, scale)
+    if not isinstance(scale, (tuple, list)) or len(scale) != 2:
+        return None
+    out = []
+    for s in scale:
+        if isinstance(s, bool) or not isinstance(s, (int, float)):
+            return None
+        f = float(s)
+        if not math.isfinite(f) or f < 1.0 or f != int(f):
+            return None
+        out.append(int(f))
+    return out[0], out[1]
+
+
+def _canonical_nhwc(x: torch.Tensor) -> bool:
+    """Dense channels_last with ``C > 1``: the exact condition under which
+    aten's nearest upsample runs its NHWC kernel and returns a dense
+    channels_last tensor (see the module docstring)."""
+    _, c, h, w = x.shape
+    return c > 1 and x.stride() == (h * w * c, 1, w * c, c)
+
+
+def can_use_nearest_upsample_nhwc(x: torch.Tensor, scale_factor, mode: str) -> bool:
+    """True when ``F.interpolate(x, scale_factor=..., mode=...)`` is a plain
+    integer-factor gather on a dense channels_last 4D tensor whose result is
+    value- and layout-identical to aten's. Never raises."""
+    return (
+        isinstance(x, torch.Tensor)
+        and x.is_cuda
+        and not (torch.is_grad_enabled() and x.requires_grad)
+        and mode in ("nearest", "nearest-exact")
+        and x.dim() == 4
+        and x.numel() > 0
+        and x.dtype in _SUPPORTED_DTYPES
+        and _canonical_nhwc(x)
+        and _integer_scale(scale_factor) is not None
+    )
+
+
+def nearest_upsample_nhwc(x: torch.Tensor, scale_factor) -> torch.Tensor:
+    """Integer-factor nearest upsample of a channels_last ``[N, C, H, W]``
+    tensor, returned dense channels_last. Bit-exact vs ``nn.Upsample`` in
+    ``nearest`` and ``nearest-exact`` modes; raises on unsupported input."""
+    factors = _integer_scale(scale_factor)
+    if factors is None:
+        raise ValueError(f"scale_factor must be integer-valued, got {scale_factor}")
+    # Re-check everything the predicate checks: a direct call must fail loudly
+    # rather than return a detached tensor (autograd) or a differently laid-out
+    # one (see the layout contract in the module docstring).
+    if torch.is_grad_enabled() and x.requires_grad:
+        raise ValueError(
+            "nearest_upsample_nhwc is inference-only (input requires grad)"
+        )
+    if not (x.is_cuda and x.dim() == 4 and x.dtype in _SUPPORTED_DTYPES):
+        raise ValueError(
+            "nearest_upsample_nhwc needs a CUDA 4D bf16/fp16/fp32 tensor, got "
+            f"{x.device.type} {x.dim()}D {x.dtype}"
+        )
+    if not _canonical_nhwc(x):
+        raise ValueError(
+            "nearest_upsample_nhwc needs dense channels_last strides with C > 1, "
+            f"got shape {tuple(x.shape)} strides {tuple(x.stride())}"
+        )
+    fh, fw = factors
+    n, c, h, w = x.shape
+    out_h, out_w = h * fh, w * fw
+    out = torch.empty(
+        (n, c, out_h, out_w),
+        device=x.device,
+        dtype=x.dtype,
+        memory_format=torch.channels_last,
+    )
+    total = out.numel()
+    if total == 0:
+        return out
+    sxn, _, sxh, sxw = x.stride()
+    BLOCK = 1024
+    grid = (triton.cdiv(total, BLOCK),)
+    with torch.get_device_module().device(x.device):
+        _nearest_upsample_nhwc_kernel[grid](
+            x,
+            out,
+            total,
+            c,
+            out_h,
+            out_w,
+            fh,
+            fw,
+            sxn,
+            sxh,
+            sxw,
+            IDX64=total >= _MAX_INT32 or x.numel() >= _MAX_INT32,
+            BLOCK=BLOCK,
+        )
+    return out

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md

@@ -219,6 +219,18 @@ framework-specific optimization workflow.
 - Numerical contract: these are bit-exact data-movement / same-order-add
   replacements and run independently of the request-gated Wan RMSNorm+SiLU
   path. Unsupported layouts or padding fall back to the aten chain.
+- `nearest_upsample_nhwc` (`layout/nearest_upsample_nhwc_triton.py`) is the
+  bit-exact Triton gather behind `GatedChannelsLastUpsample` in
+  `wan_vae_cuda_opt.py`; it runs on the lossless path whenever the 2D
+  upsample input already has canonical NHWC strides (aten's own NHWC nearest
+  kernel is several times slower), while the stride canonicalisation, the
+  channels_last_3d `upsample3d` frame interleave (`_interleave_time_pairs` in
+  `wanvae.py`) and the attention residual operand order (`identity + x`, so
+  the sum keeps channels_last_3d) stay behind the quality gate: the first two
+  change which cuDNN conv2d algorithm runs, and the layout change alters the
+  reduction order of the eager RMSNorm that follows under fp32 autocast
+  (Wan pipelines decode with fp32 VAE weights under bf16 autocast, where this
+  was measured non-bit-exact even though the add itself is commutative).
 - The Qwen-Image VAE (`autoencoder_kl_qwenimage.py`) uses the same
   `cat_pad_channels_last_3d` + compact-cache helper for every causal conv
   slot; single-frame image decodes keep the compact cache at the reference

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

@@ -24,7 +24,9 @@ logger = init_logger(__name__)
 
 try:
     from sglang.kernels.ops.diffusion import (
+        can_use_nearest_upsample_nhwc,
         can_use_wan_rmsnorm_silu,
+        nearest_upsample_nhwc,
         wan_rmsnorm_silu,
     )
 
@@ -74,6 +76,11 @@ class GatedChannelsLastUpsample(nn.Module):
     NHWC conv2d is not guaranteed to pick the same cuDNN algorithm as the NCHW
     one, so the stride canonicalisation stays behind the quality gate; with
     the gate off the layout is left exactly as the eager module sees it.
+
+    The upsample itself then runs the Triton ``nearest_upsample_nhwc`` gather
+    (bit-exact vs ``nn.Upsample`` for integer factors) instead of aten's
+    ``upsample_nearest2d_nhwc`` kernel, which is several times slower than
+    its NCHW sibling on the same bytes.
     """
 
     def __init__(self, upsample: nn.Upsample, gate: VaeFastPathGate) -> None:
@@ -98,6 +105,15 @@ class GatedChannelsLastUpsample(nn.Module):
             and x.is_contiguous(memory_format=torch.channels_last)
         ):
             x = x.as_strided(x.shape, canonical)
+        # The predicate admits exactly the inputs on which aten itself would
+        # run its NHWC kernel and return a dense channels_last tensor, so the
+        # Triton gather is a layout- and value-identical replacement: it runs
+        # on the lossless path too (the gate only controls the stride
+        # canonicalisation above).
+        if up.size is None and can_use_nearest_upsample_nhwc(
+            x, up.scale_factor, up.mode
+        ):
+            return nearest_upsample_nhwc(x, up.scale_factor)
         return up(x)
 
 
@@ -105,8 +121,9 @@ def _install_module_gates(
     decoder: nn.Module, gate: VaeFastPathGate, module_classes: tuple[type, ...]
 ) -> int:
     """Hand the gate to layout-sensitive modules that consult ``_sgl_gate`` in
-    their own forward (the attention block's residual operand order, so the
-    sum keeps the channels_last_3d layout)."""
+    their own forward: the ``Resample`` (channels_last_3d ``upsample3d`` frame
+    interleave, see ``resample_forward``) and the attention block (residual
+    operand order, so the sum keeps the channels_last_3d layout)."""
     count = 0
     for m in decoder.modules():
         if type(m) in module_classes:
@@ -198,9 +215,12 @@ 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,
+        WanAttentionBlock,
         WanDecoder3d,
+        WanResample,
         WanResidualBlock,
         WanRMS_norm,
+        WanUpsample,
     )
 
     if not isinstance(vae, AutoencoderKLWan):
@@ -210,6 +230,8 @@ def maybe_optimize_wan_vae(vae: nn.Module) -> nn.Module:
         decoder_cls=WanDecoder3d,
         residual_block_cls=WanResidualBlock,
         rms_norm_cls=WanRMS_norm,
+        upsample_cls=WanUpsample,
+        gated_module_classes=(WanResample, WanAttentionBlock),
         label="Wan VAE",
     )
 

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

@@ -354,6 +354,45 @@ class WanUpsample(nn.Upsample):
         return super().forward(x.float()).type_as(x)
 
 
+def _interleave_time_pairs(self, x, b, c, t, h, w):
+    """``time_conv`` doubles the channels; split them into two frame halves
+    and interleave along time: ``[B, 2C, T, H, W] -> [B, C, 2T, H, W]``.
+
+    The eager ``reshape / stack / reshape`` materialises the result in NCDHW,
+    which sends the following 2D upsample and conv2d down their NCHW paths
+    and hands every up block an NCDHW tensor (the fused RMSNorm+SiLU then
+    falls back and the residual adds run strided). With the decode-scoped
+    fast-path gate on and a channels_last_3d input, write the same values
+    straight into a channels_last_3d buffer with one copy instead. Values are
+    identical either way; the layout change is gated because the NHWC conv2d
+    it enables need not pick the same cuDNN algorithm as the NCHW one.
+    """
+    gate = getattr(self, "_sgl_gate", None)
+    if (
+        gate is not None
+        and gate.enabled
+        and not torch.compiler.is_compiling()
+        # Dense channels_last_3d (2C >= 2 channels, so this fixes stride(C) == 1
+        # and the channel split below is a view).
+        and x.is_contiguous(memory_format=torch.channels_last_3d)
+    ):
+        out = torch.empty(
+            (b, c, t * 2, h, w),
+            device=x.device,
+            dtype=x.dtype,
+            memory_format=torch.channels_last_3d,
+        )
+        # out viewed as [B, C, T, 2, H, W] receives x viewed as [B, 2, C, T, H, W]
+        # with the pair axis moved next to time.
+        out.view(b, c, t, 2, h, w).copy_(
+            x.view(b, 2, c, t, h, w).permute(0, 2, 3, 1, 4, 5)
+        )
+        return out
+    x = x.reshape(b, 2, c, t, h, w)
+    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+    return x.reshape(b, c, t * 2, h, w)
+
+
 def resample_forward(self, x):
     b, c, t, h, w = x.size()
     first_frame = is_first_frame.get()
@@ -370,17 +409,12 @@ def resample_forward(self, x):
             else:
                 x = _run_cached_causal_conv(self.time_conv, x, _feat_cache, idx)
                 _feat_idx += 1
-
-                x = x.reshape(b, 2, c, t, h, w)
-                x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
-                x = x.reshape(b, c, t * 2, h, w)
+                x = _interleave_time_pairs(self, x, b, c, t, h, w)
             feat_cache.set(_feat_cache)
             feat_idx.set(_feat_idx)
         elif not first_frame and hasattr(self, "time_conv"):
             x = self.time_conv(x)
-            x = x.reshape(b, 2, c, t, h, w)
-            x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
-            x = x.reshape(b, c, t * 2, h, w)
+            x = _interleave_time_pairs(self, x, b, c, t, h, w)
     t = x.shape[2]
     x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
     x = self.resample(x)
@@ -476,6 +510,14 @@ def attention_block_forward(self, x):
     x = x.view(batch_size, num_frames, channels, height, width)
     x = x.permute(0, 2, 1, 3, 4)
 
+    gate = getattr(self, "_sgl_gate", None)
+    if gate is not None and gate.enabled and not torch.compiler.is_compiling():
+        # ``identity`` carries the decoder's channels_last_3d layout; putting
+        # it first lets the sum inherit that layout so the next block's fused
+        # RMSNorm+SiLU applies. The add itself is commutative, but the layout
+        # changes the reduction order of the *eager* norm that follows under
+        # fp32 autocast, so this is quality-gated rather than lossless.
+        return identity + x
     return x + identity