[diffusion] Flatten Wan VAE RMSNorm row addressing (#35981)

2026-08-24 08:57:54 +08:00
parent b2eb0fa51e
commit e129fe21e5
4 changed files with 129 additions and 33 deletions
@@ -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` |
@@ -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))
    )
@@ -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()
@@ -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
 # ---------------------------------------------------------------------------