From f651b487644409447f76d68fcce8667099f46df2 Mon Sep 17 00:00:00 2001 From: pengdurice Date: Tue, 2 Jun 2026 02:28:57 -0700 Subject: [PATCH] Apply apply_group_norm_silu to LTX-2 latent upsampler (#26045) Signed-off-by: pengdurice Co-authored-by: github-actions[bot] --- .../diffusion/bench_group_norm_silu.py | 36 ++- .../models/upsampler/latent_upsampler.py | 14 +- .../test_latent_upsampler_group_norm_silu.py | 306 ++++++++++++++++++ 3 files changed, 348 insertions(+), 8 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py diff --git a/python/sglang/jit_kernel/benchmark/diffusion/bench_group_norm_silu.py b/python/sglang/jit_kernel/benchmark/diffusion/bench_group_norm_silu.py index 4a6e0f4c6..86971b525 100644 --- a/python/sglang/jit_kernel/benchmark/diffusion/bench_group_norm_silu.py +++ b/python/sglang/jit_kernel/benchmark/diffusion/bench_group_norm_silu.py @@ -38,8 +38,28 @@ CASES = [ Case("video_3d_small", (1, 64, 4, 16, 16), 32), Case("threshold_3d", (1, 128, 1, 256, 256), 32), Case("hunyuan_video_large", (1, 128, 20, 256, 256), 32), + # LTX-2 latent upsampler (`LatentUpsampler` + `ResBlock`) operates on + # `[B, mid_channels=512, F, H, W]` tensors with num_groups=32. The + # `small` and `pre_720p` cases stay in the default set; the larger + # `post_720p` case is opt-in via LARGE_CASES below. + Case("ltx2_upsampler_small", (1, 512, 8, 45, 80), 32), + Case("ltx2_upsampler_pre_720p", (1, 512, 16, 90, 160), 32), ] -CASE_BY_NAME = {case.name: case for case in CASES} + +# Cases too large to fit comfortably alongside the native-path intermediates +# on consumer GPUs (e.g. 24 GB L4). Opt in with `--cases large` (large only), +# `--cases all-large` (default + large), or by name. +# +# `ltx2_upsampler_post_720p` is ~471M bf16 elements (~940 MB tensor) and the +# eager `silu(group_norm(x))` reference materializes mean / variance / +# normalized / silu intermediates -- working set lands around 5 GB. On +# H100 / H200 this is fine and surfaces the asymptotic ~14x kernel speedup; +# on a 24 GB GPU it can OOM, so it's gated out of `--cases all`. +LARGE_CASES = [ + Case("ltx2_upsampler_post_720p", (1, 512, 16, 180, 320), 32), +] + +CASE_BY_NAME = {case.name: case for case in CASES + LARGE_CASES} def dtype_from_name(name: str) -> torch.dtype: @@ -70,6 +90,10 @@ def parse_dtypes(text: str) -> list[torch.dtype]: def parse_cases(text: str) -> list[Case]: if text == "all": return CASES + if text == "large": + return LARGE_CASES + if text == "all-large": + return CASES + LARGE_CASES names = [item.strip() for item in text.split(",") if item.strip()] missing = sorted(set(names) - CASE_BY_NAME.keys()) if missing: @@ -239,7 +263,15 @@ def main() -> None: parser = argparse.ArgumentParser( description="Benchmark fused GroupNorm+SiLU against PyTorch GroupNorm+SiLU." ) - parser.add_argument("--cases", default="all") + parser.add_argument( + "--cases", + default="all", + help=( + "Comma-separated case names, or one of: 'all' (default-sized " + "cases only), 'large' (high-memory cases only -- requires " + "H100/H200-class GPU), 'all-large' (both). See CASES + LARGE_CASES." + ), + ) parser.add_argument("--dtypes", default="bf16,fp16") parser.add_argument("--rounds", type=int, default=3) parser.add_argument("--warmup", type=int, default=25) diff --git a/python/sglang/multimodal_gen/runtime/models/upsampler/latent_upsampler.py b/python/sglang/multimodal_gen/runtime/models/upsampler/latent_upsampler.py index 35573b478..579e6807e 100644 --- a/python/sglang/multimodal_gen/runtime/models/upsampler/latent_upsampler.py +++ b/python/sglang/multimodal_gen/runtime/models/upsampler/latent_upsampler.py @@ -8,6 +8,7 @@ import torch import torch.nn.functional as F from einops import rearrange +from sglang.jit_kernel.diffusion.group_norm_silu import apply_group_norm_silu from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( LayerwiseOffloadableModuleMixin, ) @@ -110,8 +111,11 @@ class ResBlock(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: residual = x x = self.conv1(x) - x = self.norm1(x) - x = self.activation(x) + # Fused GroupNorm + SiLU on the first norm of the block. The second + # norm (line below) is followed by `silu(norm + residual)`, which the + # current `apply_group_norm_silu` helper does not cover -- left on + # the eager path until a `group_norm_add_silu` helper exists. + x = apply_group_norm_silu(x, self.norm1, self.activation) x = self.conv2(x) x = self.norm2(x) x = self.activation(x + residual) @@ -242,8 +246,7 @@ class LatentUpsampler(torch.nn.Module, LayerwiseOffloadableModuleMixin): if self.dims == 2: x = rearrange(latent, "b c f h w -> (b f) c h w") x = self.initial_conv(x) - x = self.initial_norm(x) - x = self.initial_activation(x) + x = apply_group_norm_silu(x, self.initial_norm, self.initial_activation) for block in self.res_blocks: x = block(x) x = self.upsampler(x) @@ -253,8 +256,7 @@ class LatentUpsampler(torch.nn.Module, LayerwiseOffloadableModuleMixin): x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) else: x = self.initial_conv(latent) - x = self.initial_norm(x) - x = self.initial_activation(x) + x = apply_group_norm_silu(x, self.initial_norm, self.initial_activation) for block in self.res_blocks: x = block(x) diff --git a/python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py b/python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py new file mode 100644 index 000000000..649d3d983 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py @@ -0,0 +1,306 @@ +from __future__ import annotations + +from unittest.mock import patch + +import pytest +import torch + +import sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler as lu_mod +from sglang.jit_kernel.diffusion.group_norm_silu import apply_group_norm_silu +from sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler import ( + LatentUpsampler, + ResBlock, + SpatialRationalResampler, +) + + +def _resblock_eager_reference(block: ResBlock, x: torch.Tensor) -> torch.Tensor: + residual = x + x = block.activation(block.norm1(block.conv1(x))) # fused site + x = block.norm2(block.conv2(x)) + return block.activation(x + residual) + + +def _latent_upsampler_eager_reference( + upsampler: LatentUpsampler, latent: torch.Tensor +) -> torch.Tensor: + from einops import rearrange + + b, _, f, _, _ = latent.shape + if upsampler.dims == 2: + x = rearrange(latent, "b c f h w -> (b f) c h w") + x = upsampler.initial_activation( + upsampler.initial_norm(upsampler.initial_conv(x)) + ) + for block in upsampler.res_blocks: + x = _resblock_eager_reference(block, x) + x = upsampler.upsampler(x) + for block in upsampler.post_upsample_res_blocks: + x = _resblock_eager_reference(block, x) + x = upsampler.final_conv(x) + return rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) + + x = upsampler.initial_activation( + upsampler.initial_norm(upsampler.initial_conv(latent)) + ) + for block in upsampler.res_blocks: + x = _resblock_eager_reference(block, x) + if upsampler.temporal_upsample: + x = upsampler.upsampler(x)[:, :, 1:, :, :] + elif isinstance(upsampler.upsampler, SpatialRationalResampler): + x = upsampler.upsampler(x) + else: + x = rearrange(x, "b c f h w -> (b f) c h w") + x = upsampler.upsampler(x) + x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) + for block in upsampler.post_upsample_res_blocks: + x = _resblock_eager_reference(block, x) + return upsampler.final_conv(x) + + +@pytest.mark.parametrize( + "batch,channels,dims,spatial", + [ + (1, 64, 2, (16, 16)), + (2, 64, 2, (8, 24)), + (1, 128, 3, (2, 8, 8)), + ], +) +def test_resblock_forward_parity(batch, channels, dims, spatial): + torch.manual_seed(0) + block = ResBlock(channels=channels, dims=dims).eval() + torch.manual_seed(1) + x = torch.randn(batch, channels, *spatial, dtype=torch.float32) + + with torch.no_grad(): + out = block(x) + ref = _resblock_eager_reference(block, x) + + torch.testing.assert_close(out, ref, atol=0.0, rtol=0.0) + + +@pytest.mark.parametrize( + "dims,latent_shape,mid_channels,num_blocks_per_stage,rational_resampler", + [ + (2, (1, 32, 2, 16, 16), 64, 2, False), + (3, (1, 32, 2, 16, 16), 64, 2, False), + (3, (1, 32, 2, 16, 16), 64, 2, True), + ], +) +def test_latent_upsampler_forward_parity( + dims, latent_shape, mid_channels, num_blocks_per_stage, rational_resampler +): + torch.manual_seed(2) + upsampler = LatentUpsampler( + in_channels=latent_shape[1], + mid_channels=mid_channels, + num_blocks_per_stage=num_blocks_per_stage, + dims=dims, + spatial_upsample=True, + temporal_upsample=False, + spatial_scale=2.0, + rational_resampler=rational_resampler, + ).eval() + torch.manual_seed(3) + latent = torch.randn(*latent_shape, dtype=torch.float32) + + with torch.no_grad(): + out = upsampler(latent) + ref = _latent_upsampler_eager_reference(upsampler, latent) + + torch.testing.assert_close(out, ref, atol=0.0, rtol=0.0) + + +def test_resblock_fuses_exactly_one_site(): + torch.manual_seed(4) + block = ResBlock(channels=64, dims=2).eval() + x = torch.randn(1, 64, 16, 16, dtype=torch.float32) + + with patch.object( + lu_mod, "apply_group_norm_silu", wraps=lu_mod.apply_group_norm_silu + ) as spy: + with torch.no_grad(): + block(x) + assert spy.call_count == 1 + + +@pytest.mark.parametrize( + "dims,num_blocks_per_stage,rational_resampler,expected_calls", + [ + (2, 2, False, 1 + 2 * 2), + (2, 4, False, 1 + 4 * 2), + (3, 2, False, 1 + 2 * 2), + (3, 2, True, 1 + 2 * 2), + ], +) +def test_latent_upsampler_fuses_expected_sites( + dims, num_blocks_per_stage, rational_resampler, expected_calls +): + torch.manual_seed(5) + upsampler = LatentUpsampler( + in_channels=32, + mid_channels=64, + num_blocks_per_stage=num_blocks_per_stage, + dims=dims, + spatial_upsample=True, + temporal_upsample=False, + spatial_scale=2.0, + rational_resampler=rational_resampler, + ).eval() + latent = torch.randn(1, 32, 2, 16, 16, dtype=torch.float32) + + with patch.object( + lu_mod, "apply_group_norm_silu", wraps=lu_mod.apply_group_norm_silu + ) as spy: + with torch.no_grad(): + upsampler(latent) + assert spy.call_count == expected_calls + + +# CUDA Triton fast path ------------------------------------------------------- + +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Triton fused group_norm_silu requires CUDA", +) + +# bf16 keeps kernel-level tolerance because its fp32-equivalent exponent range +# absorbs multi-layer conv drift; fp16 needs a looser tolerance on the e2e +# upsampler test where 8+ downstream convs amplify fused-vs-eager rounding. +_RESBLOCK_TOL = {torch.bfloat16: (7e-2, 2e-2), torch.float16: (3e-3, 3e-3)} +_UPSAMPLER_TOL = {torch.bfloat16: (7e-2, 2e-2), torch.float16: (2e-2, 1e-1)} + + +@requires_cuda +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + "batch,channels,dims,spatial", + [ + (1, 64, 2, (16, 16)), + (1, 128, 3, (2, 8, 8)), + ], +) +def test_resblock_forward_parity_cuda(dtype, batch, channels, dims, spatial): + torch.manual_seed(0) + device = torch.device("cuda") + block = ResBlock(channels=channels, dims=dims).to(device=device, dtype=dtype).eval() + torch.manual_seed(1) + x = torch.randn(batch, channels, *spatial, device=device, dtype=dtype) + + with torch.no_grad(): + out = block(x) + ref = _resblock_eager_reference(block, x) + + atol, rtol = _RESBLOCK_TOL[dtype] + torch.testing.assert_close(out, ref, atol=atol, rtol=rtol) + + +@requires_cuda +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + "dims,latent_shape,mid_channels,num_blocks_per_stage,rational_resampler", + [ + (2, (1, 32, 2, 16, 16), 64, 2, False), + (3, (1, 32, 2, 16, 16), 64, 2, False), + (3, (1, 32, 2, 16, 16), 64, 2, True), + ], +) +def test_latent_upsampler_forward_parity_cuda( + dtype, dims, latent_shape, mid_channels, num_blocks_per_stage, rational_resampler +): + torch.manual_seed(2) + device = torch.device("cuda") + upsampler = ( + LatentUpsampler( + in_channels=latent_shape[1], + mid_channels=mid_channels, + num_blocks_per_stage=num_blocks_per_stage, + dims=dims, + spatial_upsample=True, + temporal_upsample=False, + spatial_scale=2.0, + rational_resampler=rational_resampler, + ) + .to(device=device, dtype=dtype) + .eval() + ) + torch.manual_seed(3) + latent = torch.randn(*latent_shape, device=device, dtype=dtype) + + with torch.no_grad(): + out = upsampler(latent) + ref = _latent_upsampler_eager_reference(upsampler, latent) + + atol, rtol = _UPSAMPLER_TOL[dtype] + torch.testing.assert_close(out, ref, atol=atol, rtol=rtol) + + +@requires_cuda +def test_resblock_actually_uses_triton_kernel_cuda(): + from sglang.jit_kernel.diffusion.triton import group_norm_silu as triton_mod + + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.bfloat16 + block = ResBlock(channels=64, dims=2).to(device=device, dtype=dtype).eval() + x = torch.randn(1, 64, 16, 16, device=device, dtype=dtype) + + with patch.object( + triton_mod, + "triton_group_norm_silu", + wraps=triton_mod.triton_group_norm_silu, + ) as spy: + with torch.no_grad(): + block(x) + assert spy.call_count >= 1 + + +@requires_cuda +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + "dims,latent_shape,mid_channels", + [ + (2, (1, 32, 2, 16, 16), 64), + (3, (1, 32, 2, 16, 16), 64), + ], +) +def test_initial_groupnorm_silu_parity_cuda_local( + dtype, dims, latent_shape, mid_channels +): + # Sharp parity at the fused initial-norm boundary, before downstream convs + # can amplify drift; uses kernel-level tolerance instead of the looser e2e. + torch.manual_seed(2) + device = torch.device("cuda") + upsampler = ( + LatentUpsampler( + in_channels=latent_shape[1], + mid_channels=mid_channels, + num_blocks_per_stage=2, + dims=dims, + spatial_upsample=True, + temporal_upsample=False, + spatial_scale=2.0, + rational_resampler=False, + ) + .to(device=device, dtype=dtype) + .eval() + ) + torch.manual_seed(3) + latent = torch.randn(*latent_shape, device=device, dtype=dtype) + + with torch.no_grad(): + if dims == 2: + from einops import rearrange + + b, _, f, _, _ = latent.shape + x_in = rearrange(latent, "b c f h w -> (b f) c h w") + else: + x_in = latent + x_after_conv = upsampler.initial_conv(x_in) + out_fused = apply_group_norm_silu( + x_after_conv, upsampler.initial_norm, upsampler.initial_activation + ) + out_eager = upsampler.initial_activation(upsampler.initial_norm(x_after_conv)) + + atol, rtol = _RESBLOCK_TOL[dtype] + torch.testing.assert_close(out_fused, out_eager, atol=atol, rtol=rtol)