From b712dd48feb8acb7c4214d646766c226a656d304 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Sat, 2 May 2026 20:55:51 +0800 Subject: [PATCH] [codex] diffusion: enable group norm silu fuse by default (#23148) --- .../jit_kernel/diffusion/group_norm_silu.py | 35 ++++++++++++++++++ .../diffusion/triton/group_norm_silu.py | 1 + .../tests/diffusion/test_group_norm_silu.py | 11 ++---- python/sglang/multimodal_gen/envs.py | 6 --- .../runtime/models/vaes/hunyuanvae.py | 37 +++++-------------- 5 files changed, 48 insertions(+), 42 deletions(-) create mode 100644 python/sglang/jit_kernel/diffusion/group_norm_silu.py diff --git a/python/sglang/jit_kernel/diffusion/group_norm_silu.py b/python/sglang/jit_kernel/diffusion/group_norm_silu.py new file mode 100644 index 000000000..67dbd892f --- /dev/null +++ b/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"] diff --git a/python/sglang/jit_kernel/diffusion/triton/group_norm_silu.py b/python/sglang/jit_kernel/diffusion/triton/group_norm_silu.py index dc614b594..83dc8e2f0 100644 --- a/python/sglang/jit_kernel/diffusion/triton/group_norm_silu.py +++ b/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) diff --git a/python/sglang/jit_kernel/tests/diffusion/test_group_norm_silu.py b/python/sglang/jit_kernel/tests/diffusion/test_group_norm_silu.py index 2da1cd6fa..c2bb57d0c 100644 --- a/python/sglang/jit_kernel/tests/diffusion/test_group_norm_silu.py +++ b/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) diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index 8b50b2107..0f110df3b 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -57,7 +57,6 @@ if TYPE_CHECKING: SGLANG_USE_RUNAI_MODEL_STREAMER: bool = True SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND: str | None = None SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D: bool = False - SGLANG_USE_CUDA_HUNYUANVIDEO_GROUP_NORM_SILU: bool = False SGLANG_USE_ROCM_VAE: bool = False SGLANG_USE_ROCM_CUDNN_BENCHMARK: bool = False @@ -251,11 +250,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D": _lazy_bool( "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D", "false" ), - # CUDA: enable the Triton GroupNorm+SiLU fast path in HunyuanVideo VAE - # residual blocks. - "SGLANG_USE_CUDA_HUNYUANVIDEO_GROUP_NORM_SILU": _lazy_bool( - "SGLANG_USE_CUDA_HUNYUANVIDEO_GROUP_NORM_SILU", "false" - ), # ================== cache-dit Env Vars ================== # Enable cache-dit acceleration for DiT inference "SGLANG_CACHE_DIT_ENABLED": _lazy_bool("SGLANG_CACHE_DIT_ENABLED"), diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py index 0c89d8887..3e3f62ecd 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py @@ -22,8 +22,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from sglang.jit_kernel.diffusion.triton.group_norm_silu import triton_group_norm_silu -from sglang.multimodal_gen import envs +from sglang.jit_kernel.diffusion.group_norm_silu import apply_group_norm_silu from sglang.multimodal_gen.configs.models.vaes import HunyuanVAEConfig from sglang.multimodal_gen.runtime.layers.activation import get_act_fn from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE @@ -46,26 +45,6 @@ def prepare_causal_attention_mask( return mask -def _apply_hunyuan_group_norm_silu( - hidden_states: torch.Tensor, - norm: nn.GroupNorm, - activation: nn.Module, -) -> torch.Tensor: - if ( - envs.SGLANG_USE_CUDA_HUNYUANVIDEO_GROUP_NORM_SILU - and isinstance(activation, nn.SiLU) - and norm.affine - ): - return triton_group_norm_silu( - hidden_states, - norm.weight, - norm.bias, - num_groups=norm.num_groups, - eps=norm.eps, - ) - return activation(norm(hidden_states)) - - class HunyuanVAEAttention(nn.Module): def __init__( @@ -280,12 +259,12 @@ class HunyuanVideoResnetBlockCausal3D(nn.Module): hidden_states = hidden_states.contiguous() residual = hidden_states - hidden_states = _apply_hunyuan_group_norm_silu( + hidden_states = apply_group_norm_silu( hidden_states, self.norm1, self.nonlinearity ) hidden_states = self.conv1(hidden_states) - hidden_states = _apply_hunyuan_group_norm_silu( + hidden_states = apply_group_norm_silu( hidden_states, self.norm2, self.nonlinearity ) hidden_states = self.dropout(hidden_states) @@ -655,8 +634,9 @@ class HunyuanVideoEncoder3D(nn.Module): assert self.mid_block is not None hidden_states = self.mid_block(hidden_states) - hidden_states = self.conv_norm_out(hidden_states) - hidden_states = self.conv_act(hidden_states) + hidden_states = apply_group_norm_silu( + hidden_states, self.conv_norm_out, self.conv_act + ) hidden_states = self.conv_out(hidden_states) return hidden_states @@ -777,8 +757,9 @@ class HunyuanVideoDecoder3D(nn.Module): hidden_states = up_block(hidden_states) # post-process - hidden_states = self.conv_norm_out(hidden_states) - hidden_states = self.conv_act(hidden_states) + hidden_states = apply_group_norm_silu( + hidden_states, self.conv_norm_out, self.conv_act + ) hidden_states = self.conv_out(hidden_states) return hidden_states