From e129fe21e54beb0594c34826b58f1e316f43d81c Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Mon, 24 Aug 2026 08:57:54 +0800 Subject: [PATCH] [diffusion] Flatten Wan VAE RMSNorm row addressing (#35981) --- python/sglang/kernels/ops/diffusion/README.md | 2 +- .../diffusion/norm/wan_rmsnorm_silu_triton.py | 40 ++----- .../diffusion/bench_wan_rmsnorm_silu.py | 107 ++++++++++++++++++ .../kernels/ops/diffusion/test_norm.py | 13 +++ 4 files changed, 129 insertions(+), 33 deletions(-) create mode 100644 test/registered/kernels/benchmark/diffusion/bench_wan_rmsnorm_silu.py diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index 985972af4..3dd406bde 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/python/sglang/kernels/ops/diffusion/README.md @@ -104,7 +104,7 @@ Several norms look interchangeable and are not. Start here. |---|---|---|---| | `triton_group_norm_silu` / `apply_group_norm_silu` | Triton | close | NCHW-contiguous, any channels-per-group, always applies SiLU | | `group_norm_silu_4d` / `group_norm_silu_rows` | Triton | close | **channels_last only**; power-of-two `C <= 2048`; optional SiLU. This is what lets a VAE decoder run channels_last end-to-end with no `nchwToNhwc` | -| `wan_rmsnorm_silu` | Triton | close | `channels_last_3d` 5D, Wan VAE channel-first RMSNorm + SiLU | +| `wan_rmsnorm_silu` | Triton | close | dense `channels_last_3d` 5D (`stride(C) == 1`), Wan VAE channel-first RMSNorm + SiLU | | `rmsnorm_scale` / `rmsnorm_tanh_residual` | Triton | bf16-native statistics | Z-Image (matches its own reference exactly), Ideogram 4 (gated) | | `zimage_qk_rmsnorm_native` | Triton | bit-exact | Z-Image per-head QK RMSNorm | | `fused_qk_head_layernorm` | Triton | bit-exact | per-head LN on q/k, `dim_head % 4 == 0`, `<= 128` | diff --git a/python/sglang/kernels/ops/diffusion/norm/wan_rmsnorm_silu_triton.py b/python/sglang/kernels/ops/diffusion/norm/wan_rmsnorm_silu_triton.py index 7cd1a455e..d7ea71a73 100644 --- a/python/sglang/kernels/ops/diffusion/norm/wan_rmsnorm_silu_triton.py +++ b/python/sglang/kernels/ops/diffusion/norm/wan_rmsnorm_silu_triton.py @@ -34,19 +34,6 @@ def _wan_rmsnorm_silu_kernel( 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, @@ -55,20 +42,11 @@ def _wan_rmsnorm_silu_kernel( row = tl.program_id(0).to(tl.int64) offsets = tl.arange(0, block_c) mask = offsets < channels + # Dense channels-last-3d stores each pixel as one contiguous channel row. + # Address it directly instead of recovering b/t/h/w with integer div/mod. + row_offsets = row * channels + offsets - 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 - ) + x = tl.load(x_ptr + row_offsets, 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) @@ -84,7 +62,7 @@ def _wan_rmsnorm_silu_kernel( y = y.to(tl.float32) y = y * tl.sigmoid(y) - tl.store(out_ptr + out_base + offsets * out_stride_c, y, mask=mask) + tl.store(out_ptr + row_offsets, y, mask=mask) def _fake_wan_rmsnorm_silu( @@ -125,11 +103,6 @@ def _triton_wan_rmsnorm_silu_cuda( bias, out, channels, - t_size, - h_size, - w_size, - *x.stride(), - *out.stride(), rms_scale, eps, has_bias, @@ -163,6 +136,9 @@ def can_use_wan_rmsnorm_silu( and x.numel() > 0 and 0 < x.shape[1] <= _MAX_CHANNELS and x.is_contiguous(memory_format=torch.channels_last_3d) + # Size-one channel tensors can satisfy the memory-format predicate + # while retaining channel-first strides, so require dense rows too. + and x.stride(1) == 1 and _affine_supported(x, gamma) and (bias is None or _affine_supported(x, bias)) ) diff --git a/test/registered/kernels/benchmark/diffusion/bench_wan_rmsnorm_silu.py b/test/registered/kernels/benchmark/diffusion/bench_wan_rmsnorm_silu.py new file mode 100644 index 000000000..6b1d71178 --- /dev/null +++ b/test/registered/kernels/benchmark/diffusion/bench_wan_rmsnorm_silu.py @@ -0,0 +1,107 @@ +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.diffusion import wan_rmsnorm_silu +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=20, + stage="base-b-kernel-benchmark", + runner_config="1-gpu-large", + disabled="standalone benchmark", +) + +DEVICE = "cuda" + + +@dataclass(frozen=True) +class Case: + name: str + shape: tuple[int, int, int, int, int] + x_dtype: torch.dtype + affine_dtype: torch.dtype + atol: float + rtol: float + + +CASES = [ + Case( + "fastwan21_c96_t4_h480_w832", + (1, 96, 4, 480, 832), + torch.bfloat16, + torch.float32, + 1.5e-1, + 3e-2, + ), + Case( + "fastwan22_c256_t4_h384_w576", + (1, 256, 4, 384, 576), + torch.float32, + torch.float32, + 1e-5, + 1e-5, + ), +] +CASE_BY_NAME = {case.name: case for case in CASES} +CASE_NAMES = list(CASE_BY_NAME) + + +@torch.no_grad() +def native_wan_rmsnorm_silu( + x: torch.Tensor, gamma: torch.Tensor, bias: torch.Tensor +) -> torch.Tensor: + return F.silu(F.normalize(x, dim=1) * x.shape[1] ** 0.5 * gamma + bias) + + +@torch.no_grad() +def sglang_wan_rmsnorm_silu( + x: torch.Tensor, gamma: torch.Tensor, bias: torch.Tensor +) -> torch.Tensor: + return wan_rmsnorm_silu(x, gamma, bias) + + +def make_inputs(case: Case) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + generator = torch.Generator(device=DEVICE) + generator.manual_seed(case.shape[1] * 1009 + case.shape[-1]) + x = torch.randn( + case.shape, + device=DEVICE, + dtype=case.x_dtype, + generator=generator, + ).contiguous(memory_format=torch.channels_last_3d) + gamma = torch.randn( + (case.shape[1], 1, 1, 1), + device=DEVICE, + dtype=case.affine_dtype, + generator=generator, + ) + bias = torch.randn_like(gamma) + return x, gamma, bias + + +@marker.parametrize("case_name", CASE_NAMES) +@marker.benchmark("provider", ["torch", "sglang"]) +def benchmark(case_name: str, provider: str) -> marker.BenchResult: + case = CASE_BY_NAME[case_name] + x, gamma, bias = make_inputs(case) + expected = native_wan_rmsnorm_silu(x, gamma, bias) + actual = sglang_wan_rmsnorm_silu(x, gamma, bias) + torch.testing.assert_close(actual, expected, atol=case.atol, rtol=case.rtol) + assert actual.stride() == x.stride() + + fn = native_wan_rmsnorm_silu if provider == "torch" else sglang_wan_rmsnorm_silu + return marker.do_bench( + fn, + input_args=(x, gamma, bias), + use_cuda_graph=False, + replay_iters=100, + memory_args=(x, gamma, bias), + memory_output="out", + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/kernels/ops/diffusion/test_norm.py b/test/registered/kernels/ops/diffusion/test_norm.py index d2db4cf7d..261db1ae0 100644 --- a/test/registered/kernels/ops/diffusion/test_norm.py +++ b/test/registered/kernels/ops/diffusion/test_norm.py @@ -282,6 +282,19 @@ def test_wan_rmsnorm_silu_rejects_empty_input(): wan_rmsnorm_silu(x, gamma) +@torch.no_grad() +def test_wan_rmsnorm_silu_rejects_non_dense_channel_rows(): + # C=1 can report channels_last_3d compatibility while retaining NCTHW + # strides. The flat-row kernel requires stride(C)=1 and must fall back. + x = torch.randn(1, 1, 2, 3, 4, device=DEVICE, dtype=torch.bfloat16) + assert x.is_contiguous(memory_format=torch.channels_last_3d) + assert x.stride(1) != 1 + gamma = torch.ones(1, 1, 1, 1, device=DEVICE, dtype=torch.float32) + assert not can_use_wan_rmsnorm_silu(x, gamma, None) + with pytest.raises(ValueError): + wan_rmsnorm_silu(x, gamma) + + # --------------------------------------------------------------------------- # CuTe-DSL fused (residual +) norm + scale/shift # ---------------------------------------------------------------------------