From db1de6ff4cd631f7b679003d69ccc4154542a737 Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <158081477+Dayuxiaoshui@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:31:35 +0800 Subject: [PATCH] [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 | 9 +- .../sglang/kernels/ops/diffusion/__init__.py | 9 + .../layout/nearest_upsample_nhwc_triton.py | 182 ++++++++++++++++++ .../existing-fast-paths.md | 12 ++ .../runtime/models/vaes/wan_vae_cuda_opt.py | 26 ++- .../runtime/models/vaes/wanvae.py | 56 +++++- .../kernels/ops/diffusion/test_layout.py | 69 +++++++ .../ops/diffusion/test_model_fast_paths.py | 87 +++++++++ 8 files changed, 440 insertions(+), 10 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/layout/nearest_upsample_nhwc_triton.py diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index 7543fb83f..8bc44770b 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/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()` diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 8761786c7..38668208e 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/layout/nearest_upsample_nhwc_triton.py b/python/sglang/kernels/ops/diffusion/layout/nearest_upsample_nhwc_triton.py new file mode 100644 index 000000000..e520b77d1 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md index f667a3699..2382df6fd 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md +++ b/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 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 index fa9ecea36..379d9aa43 100644 --- 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 @@ -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", ) diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py index c18bdf181..c3865c5bb 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py +++ b/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 diff --git a/test/registered/kernels/ops/diffusion/test_layout.py b/test/registered/kernels/ops/diffusion/test_layout.py index ff8bdafb6..ef11183ff 100644 --- a/test/registered/kernels/ops/diffusion/test_layout.py +++ b/test/registered/kernels/ops/diffusion/test_layout.py @@ -21,6 +21,7 @@ from sglang.kernels.jit.utils import get_ci_test_range from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func from sglang.kernels.ops.diffusion import ( build_inv_indices, + can_use_nearest_upsample_nhwc, can_use_usp_merge_heads, cat_pad_channels_last_3d, dup_up3d_add, @@ -33,6 +34,7 @@ from sglang.kernels.ops.diffusion import ( fused_pack_qkv, fused_pack_segmented_qkv, fused_scatter_to_padded, + nearest_upsample_nhwc, pack_qkv_destination_major, usp_merge_heads, ) @@ -657,3 +659,70 @@ def test_wan_cached_conv_chunk_loop_bitwise(pads_temporal_only): if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) + + +# ------------------------------------------------------------------------- +# Wan-family VAE channels_last nearest upsample (pure gather, bit-exact) +# ------------------------------------------------------------------------- + + +@torch.no_grad() +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize( + "shape,scale", [((1, 192, 30, 52), 2), ((4, 96, 15, 26), 2), ((2, 3, 5, 7), (3, 2))] +) +@pytest.mark.parametrize("mode", ["nearest", "nearest-exact"]) +def test_nearest_upsample_nhwc_is_bit_exact(dtype, shape, scale, mode): + x = torch.randn(shape, device="cuda", dtype=dtype).contiguous( + memory_format=torch.channels_last + ) + sf = ( + (float(scale), float(scale)) + if isinstance(scale, int) + else tuple(float(v) for v in scale) + ) + assert can_use_nearest_upsample_nhwc(x, sf, mode) + ref = F.interpolate(x, scale_factor=sf, mode=mode) + out = nearest_upsample_nhwc(x, sf) + assert out.shape == ref.shape + assert out.stride() == ref.stride() # layout-identical, not just values + assert torch.equal(out, ref) + + +@torch.no_grad() +def test_nearest_upsample_nhwc_rejects_unsupported_inputs(): + x = torch.randn(2, 8, 6, 6, device="cuda", dtype=torch.bfloat16) + assert not can_use_nearest_upsample_nhwc(x, 2.0, "nearest-exact") # NCHW + x_cl = x.contiguous(memory_format=torch.channels_last) + assert can_use_nearest_upsample_nhwc(x_cl, 2.0, "nearest-exact") + assert not can_use_nearest_upsample_nhwc(x_cl, 1.5, "nearest-exact") + assert not can_use_nearest_upsample_nhwc(x_cl, 2.0, "bilinear") + assert not can_use_nearest_upsample_nhwc(x_cl[:, :, :0], 2.0, "nearest") + # Malformed scale factors must yield False, never raise. + for bad in ( + (None, 2), + float("nan"), + float("inf"), + (2, float("-inf")), + "2", + True, + (2,), + ): + assert not can_use_nearest_upsample_nhwc(x_cl, bad, "nearest") + # C == 1 is also NCHW-contiguous: aten picks its NCHW kernel and returns a + # differently laid-out tensor, so the predicate must reject it. + x1 = torch.randn(1, 1, 1, 1, device="cuda", dtype=torch.bfloat16) + assert x1.is_contiguous(memory_format=torch.channels_last) + assert not can_use_nearest_upsample_nhwc(x1, 2.0, "nearest-exact") + with pytest.raises(ValueError): + nearest_upsample_nhwc(x1, 2.0) + # Direct calls validate too: no silent autograd drop, no layout surprise. + with pytest.raises(ValueError): + nearest_upsample_nhwc(x_cl, 1.5) + with pytest.raises(ValueError): + nearest_upsample_nhwc(x, 2.0) # NCHW + with torch.enable_grad(): + xg = x_cl.clone().requires_grad_(True) + assert not can_use_nearest_upsample_nhwc(xg, 2.0, "nearest") + with pytest.raises(ValueError): + nearest_upsample_nhwc(xg, 2.0) diff --git a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py index 811fdc9eb..8eb00e3d1 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py @@ -123,6 +123,7 @@ from sglang.multimodal_gen.runtime.models.vaes import ( from sglang.multimodal_gen.runtime.models.vaes import flux2_vae_cuda_opt as vae_opt from sglang.multimodal_gen.runtime.models.vaes import ( wan_vae_cuda_opt, + wanvae, ) from sglang.multimodal_gen.runtime.models.vaes.autoencoder import AutoencoderKL from sglang.multimodal_gen.runtime.models.vaes.fast_path_gate import use_vae_fast_path @@ -1073,6 +1074,92 @@ def test_wan_vae_rejects_empty_input() -> None: assert not can_use_wan_rmsnorm_silu(x, gamma, None) +@torch.no_grad() +def test_wan_vae_time_interleave_matches_stack_and_keeps_layout() -> None: + # time_conv output [B, 2C, T, H, W] -> interleaved [B, C, 2T, H, W]. + b, c, t, h, w = 1, 8, 3, 6, 10 + x = _wan_cl3d((b, 2 * c, t, h, w), torch.bfloat16) + ref = torch.stack( + (x.reshape(b, 2, c, t, h, w)[:, 0], x.reshape(b, 2, c, t, h, w)[:, 1]), 3 + ).reshape(b, c, 2 * t, h, w) + + class _Holder: # stands in for the WanResample instance + pass + + holder = _Holder() + off = wanvae._interleave_time_pairs(holder, x, b, c, t, h, w) + assert torch.equal(off, ref) and off.is_contiguous() # eager NCDHW path + holder._sgl_gate = VaeFastPathGate() + holder._sgl_gate.enabled = True + on = wanvae._interleave_time_pairs(holder, x, b, c, t, h, w) + assert torch.equal(on, ref) + assert on.is_contiguous(memory_format=torch.channels_last_3d) + + +@torch.no_grad() +def test_wan_vae_upsample_wrapper_dispatch() -> None: + gate = VaeFastPathGate() + up = GatedChannelsLastUpsample( + wanvae.WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), gate + ) + aten = nn.Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact") + # Canonical NHWC (multi-frame chunk): the Triton gather replaces aten's + # NHWC kernel on both paths, same values and layout. + x = torch.randn(4, 8, 6, 6, device="cuda", dtype=torch.bfloat16).contiguous( + memory_format=torch.channels_last + ) + out = up(x) + assert torch.equal(out, aten(x)) + assert out.is_contiguous(memory_format=torch.channels_last) + # Degenerate batch stride (single frame): gate off keeps aten's NCHW + # result, gate on canonicalises and stays channels_last. + x1 = ( + _wan_cl3d((1, 8, 1, 6, 6), torch.bfloat16) + .permute(0, 2, 1, 3, 4) + .reshape(1, 8, 6, 6) + ) + ref1 = aten(x1) + off = up(x1) + assert torch.equal(off, ref1) and off.is_contiguous() + gate.enabled = True + on = up(x1) + assert torch.equal(on, ref1) + assert on.is_contiguous(memory_format=torch.channels_last) + # NCHW input is untouched on either path. + xn = torch.randn(2, 8, 6, 6, device="cuda", dtype=torch.bfloat16) + assert torch.equal(up(xn), aten(xn)) and up(xn).is_contiguous() + + +@torch.no_grad() +def test_wan_vae_decoder_install_wires_resample_gate() -> None: + torch.manual_seed(0) + dec = wanvae.WanDecoder3d( + dim=16, z_dim=4, dim_mult=[1, 1], num_res_blocks=1, temperal_upsample=[True] + ).to("cuda", torch.bfloat16) + keys = set(dec.state_dict()) + gate = VaeFastPathGate() + n_norm = wan_vae_cuda_opt._install_norm_silu( + dec, + gate, + residual_block_cls=wanvae.WanResidualBlock, + rms_norm_cls=WanRMS_norm, + label="test", + ) + n_up = wan_vae_cuda_opt._install_channels_last_upsample( + dec, gate, wanvae.WanUpsample + ) + n_gated = wan_vae_cuda_opt._install_module_gates( + dec, gate, (wanvae.WanResample, wanvae.WanAttentionBlock) + ) + # one Resample plus the mid block's attention + assert n_norm == 13 and n_up == 1 and n_gated == 2 + assert set(dec.state_dict()) == keys + resample = next(m for m in dec.modules() if type(m) is wanvae.WanResample) + assert resample._sgl_gate is gate + attn = next(m for m in dec.modules() if type(m) is wanvae.WanAttentionBlock) + assert attn._sgl_gate is gate + + # ------------------------------------------------------------------------- # Qwen-Image VAE (Wan 2.1 VAE) -- lossless causal-conv data movement and the # quality-gated RMSNorm+SiLU / channels_last upsample fast path