[diffusion] Clean up kernels and shared fast paths (#34085)

test/registered/kernels/ops/diffusion/test_ernie_residual_gate_add.py

@@ -1,35 +0,0 @@
-"""ERNIE residual-gate fast path must stay bit-exact vs the eager pair."""
-
-import sys
-
-import pytest
-import torch
-
-from sglang.multimodal_gen.runtime.models.dits.ernie_image import (
-    _ernie_residual_gate_add,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-
-
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
-def test_residual_gate_add_is_bit_exact(dtype):
-    # Real ERNIE-Image shapes: hidden 4096, 1024^2 image tokens + text tokens.
-    # fp32 exercises the eager fallback (fast path is half-dtype only).
-    torch.manual_seed(0)
-    residual = torch.randn(1, 4216, 4096, device="cuda", dtype=dtype)
-    update = torch.randn_like(residual)
-    gate = torch.randn(1, 1, 4096, device="cuda", dtype=dtype)
-    out = _ernie_residual_gate_add(residual, update, gate)
-    assert torch.equal(out, residual + gate * update)
-
-    # Full-shape gate takes the same kernel path and must stay exact too.
-    gate_full = gate.expand_as(residual).contiguous()
-    out_full = _ernie_residual_gate_add(residual, update, gate_full)
-    assert torch.equal(out_full, residual + gate_full * update)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_flux2_vae_fastpath.py

@@ -32,6 +32,11 @@ def test_flux2_vae_fastpath():
         gn_kernel.group_norm_silu_4d(x.contiguous(), gn.weight, gn.bias, 32, 1e-6)
         is None
     )
+    assert gn_kernel.group_norm_silu_4d(x, gn.weight.cpu(), gn.bias, 32, 1e-6) is None
+    assert (
+        gn_kernel.group_norm_silu_4d(x[..., :0, :], gn.weight, gn.bias, 32, 1e-6)
+        is None
+    )
 
     gate.enabled = True
     fast = fused_gn(x)

test/registered/kernels/ops/diffusion/test_fused_gate_rmsnorm.py

@@ -1,4 +1,4 @@
-"""Core checks for the quality-gated fused gate-RMSNorm (Z-Image suite reuse)."""
+"""Core checks for the quality-gated fused gate-RMSNorm path."""
 
 import sys
 
@@ -44,10 +44,10 @@ def test_fused_matches_ideogram_reference():
 def test_mount_guards_all_or_nothing():
     good, bad = _Site(), _Site(torch.float32)
     assert not fgn.mount_fused_gate_rmsnorm(nn.ModuleList([good, bad]))
-    assert not good._sgl_fused_gate_rmsnorm_enabled
+    assert not fgn.fused_gate_rmsnorm_active(good)
     assert fgn.mount_fused_gate_rmsnorm(good)
     fgn.unmount_fused_gate_rmsnorm(good)
-    assert not good._sgl_fused_gate_rmsnorm_enabled
+    assert not fgn.fused_gate_rmsnorm_active(good)
 
 
 if __name__ == "__main__":

test/registered/kernels/ops/diffusion/test_fused_linear_gelu.py

@@ -21,7 +21,7 @@ class _Site(nn.Module):
         gelu.mark_fused_gelu_site(self, "proj")
 
     def forward(self, x):
-        if self._sgl_fused_gelu_enabled and gelu.can_fuse_linear_gelu(self.proj, x):
+        if gelu.fused_gelu_active(self) and gelu.can_fuse_linear_gelu(self.proj, x):
             return gelu.fused_linear_gelu_tanh(x, self.proj.weight, self.proj.bias)
         return F.gelu(self.proj(x), approximate="tanh")
 
@@ -59,7 +59,7 @@ def test_mount_guards_and_lossless_path():
     good, bad = _Site(), _Site(torch.float32)
     model = nn.ModuleList([good, bad])
     assert not gelu.mount_fused_linear_gelu(model)
-    assert not good._sgl_fused_gelu_enabled
+    assert not gelu.fused_gelu_active(good)
 
     x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
     ref = good(x)
@@ -72,5 +72,15 @@ def test_mount_guards_and_lossless_path():
     assert not gelu.can_fuse_linear_gelu(good.proj, x.float())
 
 
+@torch.no_grad()
+def test_mounted_site_torch_compile_fullgraph():
+    site = _Site()
+    x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
+    assert gelu.mount_fused_linear_gelu(site)
+    expected = site(x)
+    actual = torch.compile(site, fullgraph=True)(x)
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/kernels/ops/diffusion/test_fused_ln_modulate.py

@@ -50,6 +50,32 @@ def test_fused_ln_modulate_guards_and_mount_protocol():
     assert not mount_fused_ln_modulate(nn.Module())  # no marked sites
 
 
+@torch.no_grad()
+def test_mounted_ln_modulate_site_torch_compile_fullgraph():
+    class Site(nn.Module):
+        def __init__(self):
+            super().__init__()
+            mark_fused_ln_modulate_site(self)
+
+        def forward(self, x, scale, shift):
+            if fused_ln_modulate_active(self) and can_fuse_ln_modulate(x, scale, shift):
+                return fused_ln_modulate(x, scale, shift, eps=1e-6)
+            return (
+                nn.functional.layer_norm(x, (x.shape[-1],), eps=1e-6)
+                * (1 + scale[:, None])
+                + shift[:, None]
+            )
+
+    site = Site()
+    assert mount_fused_ln_modulate(site)
+    x = torch.randn(1, 64, 128, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(1, 128, device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn_like(scale)
+    expected = site(x, scale, shift)
+    actual = torch.compile(site, fullgraph=True)(x, scale, shift)
+    torch.testing.assert_close(actual, expected, atol=0.0625, rtol=0.05)
+
+
 if __name__ == "__main__":
     import sys
 

test/registered/kernels/ops/diffusion/test_modulate_scale_shift.py

@@ -3,6 +3,7 @@ import torch
 
 from sglang.kernels.ops.diffusion.modulate_scale_shift import (
     can_use_modulate_scale_shift_cuda,
+    modulate_scale_shift,
     modulate_scale_shift_cuda,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
@@ -48,6 +49,7 @@ def test_modulate_scale_shift_guards_reject_fp32():
     x = torch.randn((1, 64, 64), device="cuda", dtype=torch.float32)
     row = torch.randn((1, 64), device="cuda", dtype=torch.float32)
     assert not can_use_modulate_scale_shift_cuda(x, row, row)
+    assert torch.equal(modulate_scale_shift(x, row, row), _eager(x, row, row))
 
 
 if __name__ == "__main__":

test/registered/kernels/ops/diffusion/test_native_bf16_rmsnorm.py

@@ -1,9 +1,9 @@
 import pytest
 import torch
 
-from sglang.kernels.ops.diffusion.triton.zimage_native_norm import (
-    zimage_rmsnorm_scale,
-    zimage_rmsnorm_tanh_residual,
+from sglang.kernels.ops.diffusion.triton.native_bf16_rmsnorm import (
+    rmsnorm_scale,
+    rmsnorm_tanh_residual,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
 
@@ -21,26 +21,28 @@ def _native_bf16_rmsnorm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
     return ((x * rstd).to(torch.bfloat16) * weight).to(torch.bfloat16)
 
 
-def test_zimage_native_norm_rejects_cpu_inputs():
+def test_native_bf16_rmsnorm_rejects_unsupported_inputs():
     x = torch.randn(2, 3, 16, dtype=torch.bfloat16)
     weight = torch.randn(16, dtype=torch.bfloat16)
     modulation = torch.randn(2, 1, 16, dtype=torch.bfloat16)
     residual = torch.randn_like(x)
 
-    assert zimage_rmsnorm_scale(x, weight, modulation, EPS) is None
-    assert zimage_rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
+    assert rmsnorm_scale(x, weight, modulation, EPS) is None
+    assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
+    assert rmsnorm_scale(x, weight[:-1], modulation, EPS) is None
+    assert rmsnorm_tanh_residual(x, modulation, residual[..., :-1], weight, EPS) is None
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
-def test_zimage_rmsnorm_scale_matches_native_bf16(shape):
+def test_rmsnorm_scale_matches_native_bf16(shape):
     torch.manual_seed(0)
     batch, _, dim = shape
     x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
     weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16)
     scale = torch.randn(batch, 1, dim, device="cuda", dtype=torch.bfloat16)
 
-    actual = zimage_rmsnorm_scale(x, weight, scale, EPS)
+    actual = rmsnorm_scale(x, weight, scale, EPS)
     expected = (_native_bf16_rmsnorm(x, weight) * scale).to(torch.bfloat16)
 
     assert actual is not None
@@ -49,7 +51,7 @@ def test_zimage_rmsnorm_scale_matches_native_bf16(shape):
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.parametrize("shape", [(1, 32, 2560), (2, 17, 256)])
-def test_zimage_rmsnorm_tanh_residual_matches_native_bf16(shape):
+def test_rmsnorm_tanh_residual_matches_native_bf16(shape):
     torch.manual_seed(0)
     batch, _, dim = shape
     x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
@@ -57,7 +59,7 @@ def test_zimage_rmsnorm_tanh_residual_matches_native_bf16(shape):
     residual = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
     weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16)
 
-    actual = zimage_rmsnorm_tanh_residual(x, gate, residual, weight, EPS)
+    actual = rmsnorm_tanh_residual(x, gate, residual, weight, EPS)
     norm = _native_bf16_rmsnorm(x, weight)
     gated = (torch.tanh(gate.float()).to(torch.bfloat16) * norm).to(torch.bfloat16)
     expected = (residual + gated).to(torch.bfloat16)
@@ -68,15 +70,15 @@ def test_zimage_rmsnorm_tanh_residual_matches_native_bf16(shape):
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def test_zimage_native_norm_rejects_hidden_size_above_limit():
+def test_native_bf16_rmsnorm_rejects_hidden_size_above_limit():
     dim = 8448
     x = torch.empty(1, 1, dim, device="cuda", dtype=torch.bfloat16)
     weight = torch.empty(dim, device="cuda", dtype=torch.bfloat16)
     modulation = torch.empty(1, 1, dim, device="cuda", dtype=torch.bfloat16)
     residual = torch.empty_like(x)
 
-    assert zimage_rmsnorm_scale(x, weight, modulation, EPS) is None
-    assert zimage_rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
+    assert rmsnorm_scale(x, weight, modulation, EPS) is None
+    assert rmsnorm_tanh_residual(x, modulation, residual, weight, EPS) is None
 
 
 if __name__ == "__main__":

test/registered/kernels/ops/diffusion/test_qknorm_rope.py

@@ -84,6 +84,17 @@ def fused_qknorm_rope(
     )
 
 
+def test_qknorm_rope_rejects_unsupported_dtypes() -> None:
+    from sglang.kernels.ops.diffusion.qknorm_rope import (
+        can_use_fused_inplace_qknorm_rope,
+    )
+
+    assert not can_use_fused_inplace_qknorm_rope(128, 128, False, torch.float32)
+    assert not can_use_fused_inplace_qknorm_rope(
+        128, 128, False, torch.bfloat16, torch.float64
+    )
+
+
 BS_LIST = [2**n for n in range(13)]
 BS_LIST += [x + 1 for x in BS_LIST]
 BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 129, 257, 2049, 4097])
@@ -205,5 +216,28 @@ def test_qknorm_rope_preserves_split_bf16_rounding() -> None:
     assert torch.equal(k_ref, k_fused)
 
 
+def test_qknorm_rope_accepts_empty_token_dimension() -> None:
+    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
+
+    num_heads, head_dim = 8, 128
+    q = torch.empty(0, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
+    k = torch.empty_like(q)
+    weight = torch.ones(head_dim, device=DEVICE, dtype=DTYPE)
+    cache = create_cos_sin_cache(head_dim, 1)
+    positions = torch.empty(0, device=DEVICE, dtype=torch.int64)
+
+    fused_inplace_qknorm_rope(
+        q,
+        k,
+        weight,
+        weight,
+        cache,
+        positions,
+        is_neox=False,
+        rope_dim=head_dim,
+    )
+    assert q.numel() == k.numel() == 0
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_quality_gate.py

@@ -0,0 +1,47 @@
+import sys
+
+import pytest
+import torch.nn as nn
+
+from sglang.kernels.ops.diffusion.quality_gate import QualityGatedFusion
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def test_quality_gate_mounts_and_unmounts_all_sites():
+    fusion = QualityGatedFusion(
+        name="test fusion",
+        marker_attr="_test_fusion_site",
+        enabled_attr="_test_fusion_enabled",
+    )
+    root = nn.ModuleList([nn.Module(), nn.Module()])
+    for index, site in enumerate(root):
+        fusion.mark(site, index)
+
+    assert [fusion.metadata(site) for site in root] == [0, 1]
+    assert fusion.mount(root)
+    assert all(fusion.is_enabled(site) for site in root)
+    fusion.unmount(root)
+    assert not any(fusion.is_enabled(site) for site in root)
+
+
+def test_quality_gate_rejection_is_all_or_nothing():
+    fusion = QualityGatedFusion(
+        name="test fusion",
+        marker_attr="_test_fusion_site",
+        enabled_attr="_test_fusion_enabled",
+    )
+    root = nn.ModuleList([nn.Module(), nn.Module()])
+    for index, site in enumerate(root):
+        fusion.mark(site, index)
+
+    assert not fusion.mount(
+        root, reject_reason=lambda site: "rejected" if fusion.metadata(site) else None
+    )
+    assert not any(fusion.is_enabled(site) for site in root)
+    assert not fusion.mount(nn.Module())
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_residual_gate_add.py

@@ -5,6 +5,7 @@ import torch
 
 from sglang.kernels.ops.diffusion.residual_gate_add import (
     can_use_residual_gate_add_cuda,
+    residual_gate_add,
     residual_gate_add_cuda,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
@@ -25,6 +26,8 @@ CASES = [
     ((1, 4608, 3072), (1, 1, 3072)),
     # FLUX.2-dev (D=6144) joint sequence.
     ((1, 4608, 6144), (1, 1, 6144)),
+    # ERNIE-4.5-VL 1024^2 image tokens plus text tokens.
+    ((1, 4216, 4096), (1, 1, 4096)),
 ]
 
 
@@ -55,6 +58,7 @@ def test_residual_gate_add_matches_torch(residual_shape, gate_shape):
     out = residual_gate_add_cuda(residual, update, gate)
     ref = residual + update * gate
     _assert_matches_torch(out, ref)
+    assert torch.equal(residual_gate_add(residual, update, gate), ref)
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@@ -79,6 +83,13 @@ def test_can_use_residual_gate_add_cuda_rejects_unsupported_inputs():
     assert not can_use_residual_gate_add_cuda(residual, update.float(), gate)
     assert not can_use_residual_gate_add_cuda(residual, update[:, ::2], gate)
     assert not can_use_residual_gate_add_cuda(residual, update, gate[:, :, ::2])
+    empty_residual = residual[:, :0]
+    empty_update = update[:, :0]
+    assert not can_use_residual_gate_add_cuda(empty_residual, empty_update, gate)
+    assert torch.equal(
+        residual_gate_add(empty_residual, empty_update, gate),
+        empty_residual + empty_update * gate,
+    )
 
     # Only [1, ..., 1, D] row-broadcast gates are supported; a batched
     # [B>1, 1, D] gate is not row-broadcast here and must fall back.
@@ -88,6 +99,10 @@ def test_can_use_residual_gate_add_cuda_rejects_unsupported_inputs():
     assert not can_use_residual_gate_add_cuda(
         batched_residual, batched_update, batched_gate
     )
+    assert torch.equal(
+        residual_gate_add(batched_residual, batched_update, batched_gate),
+        batched_residual + batched_update * batched_gate,
+    )
 
 
 def test_residual_gate_add_custom_op_torch_compile_fullgraph():
@@ -96,7 +111,7 @@ def test_residual_gate_add_custom_op_torch_compile_fullgraph():
     gate = torch.randn((1, 1, 128), device="cuda", dtype=torch.bfloat16)
 
     def fn(residual, update, gate):
-        return residual_gate_add_cuda(residual, update, gate)
+        return residual_gate_add(residual, update, gate)
 
     compiled = torch.compile(fn, fullgraph=True)
     out = compiled(residual, update, gate)

test/registered/kernels/ops/diffusion/test_scale_shift.py

@@ -0,0 +1,40 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion.triton.scale_shift import (
+    try_fused_scaled_residual_add_exact,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_scaled_residual_add_is_bit_exact(dtype):
+    torch.manual_seed(0)
+    residual = torch.randn(2, 17, 64, device="cuda", dtype=torch.float32)
+    x = torch.randn(2, 17, 64, device="cuda", dtype=dtype)
+    scale = torch.randn(64, device="cuda", dtype=torch.float32)
+
+    actual = try_fused_scaled_residual_add_exact(residual, x, scale)
+    expected = residual + x * scale
+    assert actual is not None
+    assert torch.equal(actual, expected)
+
+
+@torch.no_grad()
+def test_scaled_residual_add_rejects_unsupported_inputs():
+    residual = torch.empty(2, 3, 8, device="cuda", dtype=torch.float32)
+    x = torch.empty_like(residual)
+    scale = torch.empty(8, device="cuda", dtype=torch.float32)
+
+    assert try_fused_scaled_residual_add_exact(residual, x, scale) is None
+    assert try_fused_scaled_residual_add_exact(residual, x.half(), scale[:-1]) is None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_timestep_embedding.py

@@ -134,12 +134,12 @@ def test_timestep_embedding_perf():
         end = torch.cuda.Event(enable_timing=True)
 
         for _ in range(warmup_times):
-            output_fn = kernel_fn(*args, **kwargs)
+            kernel_fn(*args, **kwargs)
         torch.cuda.synchronize()
 
         start.record()
         for _ in range(repeat_times):
-            output_fn = kernel_fn(*args, **kwargs)
+            kernel_fn(*args, **kwargs)
         end.record()
         end.synchronize()
         return start.elapsed_time(end) / repeat_times

test/registered/kernels/ops/diffusion/test_ulysses_qkv.py

@@ -0,0 +1,54 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion.triton.ulysses_qkv import (
+    pack_qkv_destination_major,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_pack_qkv_destination_major_is_bit_exact(dtype):
+    torch.manual_seed(0)
+    rows, world_size, global_heads, head_size = 17, 4, 12, 64
+    q, k, v = (
+        torch.randn(rows, global_heads, head_size, device="cuda", dtype=dtype)
+        for _ in range(3)
+    )
+
+    local_heads = global_heads // world_size
+    expected = torch.empty(
+        world_size,
+        rows,
+        local_heads,
+        3 * head_size,
+        device="cuda",
+        dtype=dtype,
+    )
+    for index, tensor in enumerate((q, k, v)):
+        shards = tensor.view(rows, world_size, local_heads, head_size).permute(
+            1, 0, 2, 3
+        )
+        expected[..., index * head_size : (index + 1) * head_size].copy_(shards)
+
+    actual = pack_qkv_destination_major(q, k, v, world_size)
+    assert torch.equal(actual, expected)
+
+
+def test_pack_qkv_destination_major_validates_inputs():
+    q = torch.empty(2, 4, 8, device="cuda", dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match="same 3D shape"):
+        pack_qkv_destination_major(q, q[:, :-1], q, 2)
+    with pytest.raises(ValueError, match="divide global_heads"):
+        pack_qkv_destination_major(q, q, q, 3)
+    with pytest.raises(ValueError, match="expected shape"):
+        pack_qkv_destination_major(q, q, q, 2, out=torch.empty_like(q))
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernels/ops/diffusion/test_varlen_pack_pad.py

@@ -55,7 +55,10 @@ def _build_mask(bs, s_txt, s_img, valid_txt_lens):
 
 def _ref_pack(q, k, v, indices):
     bs, seq = q.shape[:2]
-    flat = lambda t: t.reshape(bs * seq, *t.shape[2:])
+
+    def flat(t):
+        return t.reshape(bs * seq, *t.shape[2:])
+
     return (
         flat(q).index_select(0, indices),
         flat(k).index_select(0, indices),
@@ -64,7 +67,6 @@ def _ref_pack(q, k, v, indices):
 
 
 def _ref_scatter(out_unpad, indices, bs, seq):
-    n_valid = indices.shape[0]
     _, num_heads, head_dim = out_unpad.shape
     flat = torch.zeros(
         bs * seq, num_heads, head_dim, dtype=out_unpad.dtype, device=DEVICE

test/registered/kernels/ops/diffusion/test_wan_vae_fastpath.py

@@ -66,5 +66,14 @@ def test_fused_module_gate_dispatch() -> None:
     assert torch.equal(fused(x), expected)
 
 
+@torch.no_grad()
+def test_kernel_rejects_empty_input() -> None:
+    x = torch.empty(1, 96, 0, 2, 2, device="cuda", dtype=torch.bfloat16).to(
+        memory_format=torch.channels_last_3d
+    )
+    gamma = torch.ones(96, 1, 1, 1, device="cuda", dtype=torch.bfloat16)
+    assert wan_rmsnorm_silu(x, gamma) is None
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))