diffusion: add HunyuanVideo GroupNorm+SiLU fast path (#22814)

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

@@ -0,0 +1,349 @@
+import math
+
+import torch
+import torch.nn.functional as F
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.srt.utils.custom_op import register_custom_op
+
+_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
+_LARGE_GROUP_THRESHOLD = 1 << 18
+_BLOCK_SIZE = 4096
+_BLOCKS_PER_PROGRAM = 2
+_CHUNK_SIZE = _BLOCK_SIZE * _BLOCKS_PER_PROGRAM
+
+
+@triton.jit
+def _group_norm_silu_contiguous_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    channels,
+    spatial_size,
+    channels_per_group,
+    group_size,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    group_id = tl.program_id(0).to(tl.int64)
+    batch_id = tl.program_id(1).to(tl.int64)
+
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    sum_val = tl.zeros((), dtype=tl.float32)
+    sum_sq = tl.zeros((), dtype=tl.float32)
+    for off in range(0, group_size, BLOCK_SIZE):
+        idx = off + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        sum_val += tl.sum(x, axis=0)
+        sum_sq += tl.sum(x * x, axis=0)
+
+    inv_group = 1.0 / group_size
+    mean = sum_val * inv_group
+    var = sum_sq * inv_group - mean * mean
+    rstd = tl.rsqrt(var + eps)
+
+    weight_group_offset = group_id * channels_per_group
+    for off in range(0, group_size, BLOCK_SIZE):
+        idx = off + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        channel_offsets = weight_group_offset + idx // spatial_size
+        weight = tl.load(weight_ptr + channel_offsets, mask=mask, other=1.0).to(
+            tl.float32
+        )
+        bias = tl.load(bias_ptr + channel_offsets, mask=mask, other=0.0).to(tl.float32)
+        y = (x - mean) * rstd
+        y = y * weight + bias
+        y = y * tl.sigmoid(y)
+        tl.store(output_ptr + group_base + idx, y, mask=mask)
+
+
+@triton.jit
+def _group_norm_stats_kernel(
+    input_ptr,
+    partial_sum_ptr,
+    partial_sq_ptr,
+    channels,
+    spatial_size,
+    num_groups,
+    channels_per_group,
+    group_size,
+    chunks_per_row,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCKS_PER_PROGRAM: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    chunk_id = tl.program_id(1).to(tl.int64)
+
+    batch_id = row // num_groups
+    group_id = row - batch_id * num_groups
+    chunk_start = chunk_id * BLOCK_SIZE * BLOCKS_PER_PROGRAM
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+
+    sum_val = tl.zeros((), dtype=tl.float32)
+    sum_sq = tl.zeros((), dtype=tl.float32)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    for block_id in range(BLOCKS_PER_PROGRAM):
+        idx = chunk_start + block_id * BLOCK_SIZE + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        sum_val += tl.sum(x, axis=0)
+        sum_sq += tl.sum(x * x, axis=0)
+
+    partial_index = row * chunks_per_row + chunk_id
+    tl.store(partial_sum_ptr + partial_index, sum_val)
+    tl.store(partial_sq_ptr + partial_index, sum_sq)
+
+
+@triton.jit
+def _group_norm_finalize_stats_kernel(
+    partial_sum_ptr,
+    partial_sq_ptr,
+    stats_ptr,
+    chunks_per_row,
+    group_size,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    sum_val = tl.zeros((), dtype=tl.float32)
+    sum_sq = tl.zeros((), dtype=tl.float32)
+    base = row * chunks_per_row
+    for off in range(0, chunks_per_row, BLOCK_SIZE):
+        idx = off + offsets
+        mask = idx < chunks_per_row
+        sum_val += tl.sum(
+            tl.load(partial_sum_ptr + base + idx, mask=mask, other=0.0), axis=0
+        )
+        sum_sq += tl.sum(
+            tl.load(partial_sq_ptr + base + idx, mask=mask, other=0.0), axis=0
+        )
+
+    inv_group = 1.0 / group_size
+    mean = sum_val * inv_group
+    var = sum_sq * inv_group - mean * mean
+    rstd = tl.rsqrt(var + eps)
+    tl.store(stats_ptr + row * 2, mean)
+    tl.store(stats_ptr + row * 2 + 1, rstd)
+
+
+@triton.jit
+def _group_norm_apply_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    stats_ptr,
+    channels,
+    spatial_size,
+    num_groups,
+    channels_per_group,
+    group_size,
+    chunks_per_row,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCKS_PER_PROGRAM: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    chunk_id = tl.program_id(1).to(tl.int64)
+
+    batch_id = row // num_groups
+    group_id = row - batch_id * num_groups
+    chunk_start = chunk_id * BLOCK_SIZE * BLOCKS_PER_PROGRAM
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+    weight_group_offset = group_id * channels_per_group
+
+    mean = tl.load(stats_ptr + row * 2)
+    rstd = tl.load(stats_ptr + row * 2 + 1)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    for block_id in range(BLOCKS_PER_PROGRAM):
+        idx = chunk_start + block_id * BLOCK_SIZE + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        channel_offsets = weight_group_offset + idx // spatial_size
+        weight = tl.load(weight_ptr + channel_offsets, mask=mask, other=1.0).to(
+            tl.float32
+        )
+        bias = tl.load(bias_ptr + channel_offsets, mask=mask, other=0.0).to(tl.float32)
+        y = (x - mean) * rstd
+        y = y * weight + bias
+        y = y * tl.sigmoid(y)
+        tl.store(output_ptr + group_base + idx, y, mask=mask)
+
+
+def _group_norm_silu_native(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+) -> torch.Tensor:
+    return F.silu(F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps))
+
+
+def _can_use_triton_group_norm_silu(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+) -> bool:
+    return (
+        x.is_cuda
+        and not x.requires_grad
+        and x.dtype in _SUPPORTED_DTYPES
+        and x.ndim in (2, 3, 4, 5)
+        and x.shape[1] % num_groups == 0
+        and weight.is_cuda
+        and bias.is_cuda
+        and weight.dtype == x.dtype
+        and bias.dtype == x.dtype
+        and weight.ndim == 1
+        and bias.ndim == 1
+        and weight.shape == bias.shape == (x.shape[1],)
+    )
+
+
+def _launch_one_pass(
+    x_contiguous: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+) -> torch.Tensor:
+    batch_size, channels = x_contiguous.shape[:2]
+    spatial_size = math.prod(x_contiguous.shape[2:]) if x_contiguous.ndim > 2 else 1
+    channels_per_group = channels // num_groups
+    group_size = channels_per_group * spatial_size
+
+    x_flat = x_contiguous.reshape(batch_size, channels, spatial_size, 1)
+    y_flat = torch.empty_like(x_flat)
+    block_size = min(4096, triton.next_power_of_2(max(1, min(group_size, 4096))))
+
+    _group_norm_silu_contiguous_kernel[(num_groups, batch_size)](
+        x_flat,
+        weight,
+        bias,
+        y_flat,
+        channels,
+        spatial_size,
+        channels_per_group,
+        group_size,
+        eps,
+        BLOCK_SIZE=block_size,
+    )
+    return y_flat.reshape_as(x_contiguous)
+
+
+def _launch_chunked(
+    x_contiguous: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+) -> torch.Tensor:
+    batch_size, channels = x_contiguous.shape[:2]
+    spatial_size = math.prod(x_contiguous.shape[2:]) if x_contiguous.ndim > 2 else 1
+    channels_per_group = channels // num_groups
+    group_size = channels_per_group * spatial_size
+    rows = batch_size * num_groups
+    chunks_per_row = triton.cdiv(group_size, _CHUNK_SIZE)
+
+    x_flat = x_contiguous.reshape(-1)
+    y = torch.empty_like(x_contiguous)
+    y_flat = y.reshape(-1)
+    partial_sum = torch.empty(
+        (rows, chunks_per_row), device=x_contiguous.device, dtype=torch.float32
+    )
+    partial_sq = torch.empty_like(partial_sum)
+    stats = torch.empty((rows, 2), device=x_contiguous.device, dtype=torch.float32)
+
+    _group_norm_stats_kernel[(rows, chunks_per_row)](
+        x_flat,
+        partial_sum,
+        partial_sq,
+        channels,
+        spatial_size,
+        num_groups,
+        channels_per_group,
+        group_size,
+        chunks_per_row,
+        BLOCK_SIZE=_BLOCK_SIZE,
+        BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
+        num_warps=8,
+        num_stages=3,
+    )
+
+    reduce_block = min(1024, triton.next_power_of_2(max(1, chunks_per_row)))
+    _group_norm_finalize_stats_kernel[(rows,)](
+        partial_sum,
+        partial_sq,
+        stats,
+        chunks_per_row,
+        group_size,
+        eps,
+        BLOCK_SIZE=reduce_block,
+        num_warps=4,
+        num_stages=2,
+    )
+
+    _group_norm_apply_kernel[(rows, chunks_per_row)](
+        x_flat,
+        weight,
+        bias,
+        y_flat,
+        stats,
+        channels,
+        spatial_size,
+        num_groups,
+        channels_per_group,
+        group_size,
+        chunks_per_row,
+        BLOCK_SIZE=_BLOCK_SIZE,
+        BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
+        num_warps=8,
+        num_stages=3,
+    )
+    return y
+
+
+@register_custom_op(op_name="triton_group_norm_silu_cuda", out_shape="x")
+def _triton_group_norm_silu_cuda(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    if not _can_use_triton_group_norm_silu(x, weight, bias, num_groups):
+        return _group_norm_silu_native(x, weight, bias, num_groups, eps)
+
+    x_contiguous = x.contiguous()
+    spatial_size = math.prod(x_contiguous.shape[2:]) if x_contiguous.ndim > 2 else 1
+    channels_per_group = x_contiguous.shape[1] // num_groups
+    group_size = channels_per_group * spatial_size
+
+    with torch.cuda.device(x.device):
+        if group_size >= _LARGE_GROUP_THRESHOLD:
+            return _launch_chunked(x_contiguous, weight, bias, num_groups, eps)
+        return _launch_one_pass(x_contiguous, weight, bias, num_groups, eps)
+
+
+def triton_group_norm_silu(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    return _triton_group_norm_silu_cuda(x, weight, bias, num_groups, eps)
+
+
+__all__ = ["triton_group_norm_silu"]

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

@@ -0,0 +1,109 @@
+import sys
+
+import pytest
+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.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")
+register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
+
+DEVICE = "cuda"
+DTYPES = [torch.float16, torch.bfloat16, torch.float32]
+TEST_CASES = [
+    pytest.param((2, 64, 32, 32), 32, id="image_2d"),
+    pytest.param((1, 64, 4, 16, 16), 32, id="video_3d"),
+    pytest.param((4, 128), 32, id="token_2d"),
+]
+LARGE_TILE_CASE = ((1, 128, 20, 256, 256), 32)
+
+
+def _tol(dtype: torch.dtype) -> tuple[float, float]:
+    if dtype == torch.float32:
+        return 1e-5, 1e-5
+    if dtype == torch.bfloat16:
+        return 7e-2, 2e-2
+    return 3e-3, 3e-3
+
+
+@pytest.fixture(autouse=True)
+def cuda_setup():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    torch.cuda.manual_seed(0)
+
+
+def _reference(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    return F.silu(F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps))
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("shape,num_groups", TEST_CASES)
+@pytest.mark.parametrize("dtype", DTYPES)
+def test_triton_group_norm_silu(
+    shape: tuple[int, ...], num_groups: int, dtype: torch.dtype
+) -> None:
+    channels = shape[1]
+    x = torch.randn(shape, device=DEVICE, dtype=dtype)
+    weight = torch.randn(channels, device=DEVICE, dtype=dtype)
+    bias = torch.randn(channels, device=DEVICE, dtype=dtype)
+
+    actual = triton_group_norm_silu(x, weight, bias, num_groups=num_groups)
+    expected = _reference(x, weight, bias, num_groups)
+
+    atol, rtol = _tol(dtype)
+    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+
+
+@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,
+    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)
+    expected = activation(norm(hidden_states))
+
+    atol, rtol = _tol(dtype)
+    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+
+
+@torch.no_grad()
+def test_triton_group_norm_silu_large_tile_bf16() -> None:
+    shape, num_groups = LARGE_TILE_CASE
+    x = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16)
+    weight = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
+    bias = torch.randn(shape[1], device=DEVICE, dtype=torch.bfloat16)
+
+    actual = triton_group_norm_silu(x, weight, bias, num_groups=num_groups)
+    expected = _reference(x, weight, bias, num_groups)
+
+    atol, rtol = _tol(torch.bfloat16)
+    torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

python/sglang/multimodal_gen/envs.py

@@ -57,6 +57,7 @@ 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
 
@@ -250,6 +251,11 @@ 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

@@ -22,6 +22,8 @@ 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.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
@@ -44,6 +46,26 @@ 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__(
@@ -258,12 +280,14 @@ class HunyuanVideoResnetBlockCausal3D(nn.Module):
         hidden_states = hidden_states.contiguous()
         residual = hidden_states
 
-        hidden_states = self.norm1(hidden_states)
-        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = _apply_hunyuan_group_norm_silu(
+            hidden_states, self.norm1, self.nonlinearity
+        )
         hidden_states = self.conv1(hidden_states)
 
-        hidden_states = self.norm2(hidden_states)
-        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = _apply_hunyuan_group_norm_silu(
+            hidden_states, self.norm2, self.nonlinearity
+        )
         hidden_states = self.dropout(hidden_states)
         hidden_states = self.conv2(hidden_states)