[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)