[codex] diffusion: enable group norm silu fuse by default (#23148)

python/sglang/jit_kernel/diffusion/group_norm_silu.py

@@ -0,0 +1,35 @@
+import torch
+from torch import nn
+
+
+def apply_group_norm_silu(
+    x: torch.Tensor,
+    norm: nn.Module,
+    activation: nn.Module,
+) -> torch.Tensor:
+    if (
+        x.is_cuda
+        and not torch.is_grad_enabled()
+        and not x.requires_grad
+        and isinstance(norm, nn.GroupNorm)
+        and isinstance(activation, nn.SiLU)
+        and not activation.inplace
+        and norm.affine
+        and norm.weight is not None
+        and norm.bias is not None
+    ):
+        from sglang.jit_kernel.diffusion.triton.group_norm_silu import (
+            triton_group_norm_silu,
+        )
+
+        return triton_group_norm_silu(
+            x,
+            norm.weight,
+            norm.bias,
+            num_groups=norm.num_groups,
+            eps=norm.eps,
+        )
+    return activation(norm(x))
+
+
+__all__ = ["apply_group_norm_silu"]

python/sglang/jit_kernel/diffusion/triton/group_norm_silu.py

@@ -240,6 +240,7 @@ def _can_use_triton_group_norm_silu(
 ) -> bool:
     return (
         x.is_cuda
+        and not torch.is_grad_enabled()
         and not x.requires_grad
         and x.dtype in _SUPPORTED_DTYPES
         and x.ndim in (2, 3, 4, 5)

python/sglang/jit_kernel/tests/diffusion/test_group_norm_silu.py

@@ -5,10 +5,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from sglang.jit_kernel.diffusion.group_norm_silu import apply_group_norm_silu
 from sglang.jit_kernel.diffusion.triton.group_norm_silu import triton_group_norm_silu
-from sglang.multimodal_gen.runtime.models.vaes.hunyuanvae import (
-    _apply_hunyuan_group_norm_silu,
-)
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=8, suite="stage-b-kernel-unit-1-gpu-large")
@@ -70,21 +68,18 @@ def test_triton_group_norm_silu(
 @torch.no_grad()
 @pytest.mark.parametrize("shape,num_groups", TEST_CASES[:2])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-def test_apply_hunyuan_group_norm_silu(
-    monkeypatch: pytest.MonkeyPatch,
+def test_apply_group_norm_silu(
     shape: tuple[int, ...],
     num_groups: int,
     dtype: torch.dtype,
 ) -> None:
-    monkeypatch.setenv("SGLANG_USE_CUDA_HUNYUANVIDEO_GROUP_NORM_SILU", "1")
-
     norm = nn.GroupNorm(num_groups, shape[1], eps=1e-5, affine=True).to(
         device=DEVICE, dtype=dtype
     )
     activation = nn.SiLU()
     hidden_states = torch.randn(shape, device=DEVICE, dtype=dtype)
 
-    actual = _apply_hunyuan_group_norm_silu(hidden_states, norm, activation)
+    actual = apply_group_norm_silu(hidden_states, norm, activation)
     expected = activation(norm(hidden_states))
 
     atol, rtol = _tol(dtype)