From e684f0cc948f4b9eac50b8e56775c09e75065bdd Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <792179245@qq.com> Date: Sat, 5 Sep 2026 13:31:31 +0000 Subject: [PATCH 1/2] [Diffusion] Keep the Wan VAE decoder channels_last and add a Triton NHWC nearest upsample Follow-up to #38020. The degenerate-stride layout issue found on the Qwen-Image VAE also affects the Wan 2.1 video VAE, where the decoder is a much larger share of short-step video pipelines, and it comes with a second problem: once the decoder does run channels_last, aten's upsample_nearest2d_nhwc kernel is several times slower than its NCHW sibling on the same bytes. Lossless (bit-exact, unconditional): - New kernel layout/nearest_upsample_nhwc_triton.py: integer-factor nearest upsample of a dense channels_last [N, C, H, W] tensor as a pure gather. For integer factors nearest and nearest-exact both map output index i to input index i // f, so the result is bitwise identical to nn.Upsample in either mode for every dtype. Registered as diffusion.nearest_upsample_nhwc with a can_use_ predicate. 0.061 ms vs 0.458 ms (aten NHWC) and 0.190 ms (aten NCHW) for [1, 192, 240, 416] bf16 on H200. - GatedChannelsLastUpsample runs it whenever the upsample input already has canonical NHWC strides (aten would run its NHWC kernel and return the same dense channels_last tensor), on the lossless path too. Quality-gated (extra-high / high): - resample_forward's upsample3d frame interleave (reshape / stack / reshape, which materialises NCDHW) writes the same values straight into a channels_last_3d buffer with one copy via _interleave_time_pairs. - The Wan installer now wraps WanUpsample with GatedChannelsLastUpsample and hands the gate to WanResample and WanAttentionBlock (residual operand order), so all 29 RMSNorm+SiLU sites hit the fused kernel instead of 25. The attention operand order is commutative but changes the reduction order of the eager norm that follows under fp32 autocast, which is how Wan pipelines decode (fp32 VAE weights, bf16 autocast); measured end-to-end as non-bit-exact, hence gated. Wan 2.1 VAE, 81 frames 480x832 on H200, torch.profiler GPU kernel time: bf16 weights: main lossless 1342 ms / high 1178 ms -> 1324 / 1112 ms. fp32 weights + bf16 autocast (the T2V pipeline setting): main 1633 / 1311 ms -> 1593 / 1215 ms. Lossless mp4 from TurboWan2.1-T2V-1.3B md5-identical to main; high vs main-high PSNR 36.3 dB mean / 35.0 dB min over 81 frames. Peak memory unchanged (4.577 GB). Qwen-Image high: 44.6 -> 42.9 ms. --- python/sglang/kernels/ops/diffusion/README.md | 9 +- .../sglang/kernels/ops/diffusion/__init__.py | 9 ++ .../layout/nearest_upsample_nhwc_triton.py | 149 ++++++++++++++++++ .../existing-fast-paths.md | 12 ++ .../runtime/models/vaes/wan_vae_cuda_opt.py | 27 +++- .../runtime/models/vaes/wanvae.py | 56 ++++++- .../kernels/ops/diffusion/test_layout.py | 44 ++++++ .../ops/diffusion/test_model_fast_paths.py | 87 ++++++++++ 8 files changed, 383 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 c28c6fb6bbdb..20a7770f6fa1 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 57de395deb83..cba573e8336b 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -398,6 +398,13 @@ _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 @@ "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 000000000000..81bd34a5628d --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/layout/nearest_upsample_nhwc_triton.py @@ -0,0 +1,149 @@ +# 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``. + +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 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: + 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: + f = float(s) + if f < 1.0 or f != int(f): + return None + out.append(int(f)) + return out[0], out[1] + + +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.""" + return ( + isinstance(x, torch.Tensor) + and x.is_cuda + and not torch.is_grad_enabled() + and not x.requires_grad + and mode in ("nearest", "nearest-exact") + and x.dim() == 4 + and x.numel() > 0 + and x.dtype in _SUPPORTED_DTYPES + # Dense NHWC. For C > 1 this fixes stride(C) == 1, which the kernel + # relies on by adding the channel index unscaled; for C == 1 that index + # is always 0, so no separate stride check is needed. + and x.is_contiguous(memory_format=torch.channels_last) + 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}") + if not ( + x.is_cuda + and x.dim() == 4 + and x.is_contiguous(memory_format=torch.channels_last) + ): + raise ValueError("nearest_upsample_nhwc needs a CUDA channels_last 4D tensor") + 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 f667a3699339..2382df6fd71c 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 fa9ecea36463..715e78dfaa89 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 @@ 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,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: and x.is_contiguous(memory_format=torch.channels_last) ): x = x.as_strided(x.shape, canonical) + # With canonical NHWC strides 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 x.stride() == canonical + 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 +122,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 +216,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 +231,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 c18bdf181c7f..c3865c5bbb26 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 @@ def forward(self, x): 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 ff8bdafb674f..4063d9caec33 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.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 @@ fused_pack_qkv, fused_pack_segmented_qkv, fused_scatter_to_padded, + nearest_upsample_nhwc, pack_qkv_destination_major, usp_merge_heads, ) @@ -657,3 +659,45 @@ def run(force_fallback, start): 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.is_contiguous(memory_format=torch.channels_last) + assert out.shape == ref.shape + 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 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") + with torch.enable_grad(): + assert not can_use_nearest_upsample_nhwc(x_cl, 2.0, "nearest") + with pytest.raises(ValueError): + nearest_upsample_nhwc(x_cl, 1.5) 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 238fde2e5c45..0b809c631bd1 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py @@ -120,6 +120,7 @@ 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 @@ -1060,6 +1061,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 From 8dc645408230144ec1b52ca757c0120de3814482 Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <792179245@qq.com> Date: Mon, 7 Sep 2026 13:22:31 +0000 Subject: [PATCH 2/2] [Diffusion] Harden nearest_upsample_nhwc input validation and layout contract Review follow-ups: - can_use_nearest_upsample_nhwc never raises: _integer_scale rejects non-numeric entries, NaN and +/-Inf (and bools) instead of throwing. - The predicate now requires C > 1 and canonical NHWC strides on every dim. That is exactly the condition under which aten's suggest_memory_format picks channels_last, so the Triton result is layout-identical as well as value-identical; with C == 1 the tensor is also NCHW-contiguous and aten returns e.g. strides (4, 4, 2, 1) for a [1, 1, 2, 2] output. - nearest_upsample_nhwc validates before launching and raises ValueError for requires_grad inputs (a direct call must not return a detached tensor), unsupported dtype/device/rank, and non-canonical layouts. - GatedChannelsLastUpsample relies on the predicate alone for the lossless dispatch instead of duplicating the stride check. - Tests assert out.stride() == ref.stride(), cover the malformed scale factors, the C == 1 case and the direct-call error paths. --- .../layout/nearest_upsample_nhwc_triton.py | 61 ++++++++++++++----- .../runtime/models/vaes/wan_vae_cuda_opt.py | 15 +++-- .../kernels/ops/diffusion/test_layout.py | 31 +++++++++- 3 files changed, 82 insertions(+), 25 deletions(-) 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 index 81bd34a5628d..e520b77d10f5 100644 --- a/python/sglang/kernels/ops/diffusion/layout/nearest_upsample_nhwc_triton.py +++ b/python/sglang/kernels/ops/diffusion/layout/nearest_upsample_nhwc_triton.py @@ -18,6 +18,15 @@ 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 @@ -26,6 +35,8 @@ from __future__ import annotations +import math + import torch import triton # type: ignore import triton.language as tl # type: ignore @@ -69,35 +80,46 @@ def _nearest_upsample_nhwc_kernel( 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 f < 1.0 or f != int(f): + 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.""" + 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 not x.requires_grad + 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 - # Dense NHWC. For C > 1 this fixes stride(C) == 1, which the kernel - # relies on by adding the channel index unscaled; for C == 1 that index - # is always 0, so no separate stride check is needed. - and x.is_contiguous(memory_format=torch.channels_last) + and _canonical_nhwc(x) and _integer_scale(scale_factor) is not None ) @@ -109,12 +131,23 @@ def nearest_upsample_nhwc(x: torch.Tensor, scale_factor) -> torch.Tensor: factors = _integer_scale(scale_factor) if factors is None: raise ValueError(f"scale_factor must be integer-valued, got {scale_factor}") - if not ( - x.is_cuda - and x.dim() == 4 - and x.is_contiguous(memory_format=torch.channels_last) - ): - raise ValueError("nearest_upsample_nhwc needs a CUDA channels_last 4D tensor") + # 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 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 715e78dfaa89..379d9aa43c3d 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 @@ -105,14 +105,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: and x.is_contiguous(memory_format=torch.channels_last) ): x = x.as_strided(x.shape, canonical) - # With canonical NHWC strides 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 x.stride() == canonical - and can_use_nearest_upsample_nhwc(x, up.scale_factor, up.mode) + # 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) diff --git a/test/registered/kernels/ops/diffusion/test_layout.py b/test/registered/kernels/ops/diffusion/test_layout.py index 4063d9caec33..ef11183ffad7 100644 --- a/test/registered/kernels/ops/diffusion/test_layout.py +++ b/test/registered/kernels/ops/diffusion/test_layout.py @@ -684,8 +684,8 @@ def test_nearest_upsample_nhwc_is_bit_exact(dtype, shape, scale, mode): 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.is_contiguous(memory_format=torch.channels_last) assert out.shape == ref.shape + assert out.stride() == ref.stride() # layout-identical, not just values assert torch.equal(out, ref) @@ -694,10 +694,35 @@ 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") - with torch.enable_grad(): - assert not can_use_nearest_upsample_nhwc(x_cl, 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)