minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-05-02 20:55:51 +08:00
committed by GitHub
parent 1360848ee1
commit b712dd48fe
5 changed files with 48 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/diffusion/group_norm_silu.py
+35
View File
@@ -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
+1
View File
@@ -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
+3 -8
View File
@@ -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)
python/sglang/multimodal_gen/envs.py
-6
View File
@@ -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"),
python/sglang/multimodal_gen/runtime/models/vaes/hunyuanvae.py
+9 -28
View File
@@ -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