From 8d34d3b5deed95a3cdfa4ed3b58fdd86ddddca80 Mon Sep 17 00:00:00 2001 From: dengyunyang <584797741@qq.com> Date: Fri, 21 Aug 2026 00:29:21 +0800 Subject: [PATCH] [Perf][Hunyuan] Optimize for vae and patch_embed (#6306) Signed-off-by: dengyunyang <584797741@qq.com> --- .../model_executor/models/common/__init__.py | 0 .../test_fused_adaptive_group_norm_silu.py | 292 +++++++++++++++ .../common/test_fused_group_norm_silu.py | 338 ++++++++++++++++++ .../models/hunyuan_image3/__init__.py | 6 +- .../models/hunyuan_image3/autoencoder.py | 104 +----- .../hunyuan_image3_transformer.py | 138 +------ .../models/hunyuan_image3/layers/__init__.py | 22 ++ .../hunyuan_image3/layers/common/__init__.py | 11 + .../hunyuan_image3/layers/common/utils.py | 51 +++ .../hunyuan_image3/layers/native/__init__.py | 9 + .../layers/native/autoencoder_blocks.py | 44 +++ .../layers/native/transformer_blocks.py | 110 ++++++ .../hunyuan_image3/layers/nvidia/__init__.py | 9 + .../hunyuan_image3/layers/nvidia/_cudnn.py | 27 ++ .../layers/nvidia/autoencoder_blocks.py | 56 +++ .../layers/nvidia/transformer_blocks.py | 138 +++++++ .../models/common/ops/__init__.py | 17 + .../ops/fused_adaptive_group_norm_silu.py | 222 ++++++++++++ .../common/ops/fused_group_norm_silu.py | 207 +++++++++++ 19 files changed, 1583 insertions(+), 218 deletions(-) create mode 100644 tests/model_executor/models/common/__init__.py create mode 100644 tests/model_executor/models/common/test_fused_adaptive_group_norm_silu.py create mode 100644 tests/model_executor/models/common/test_fused_group_norm_silu.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/layers/__init__.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/layers/common/__init__.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/layers/common/utils.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/layers/native/__init__.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/layers/native/autoencoder_blocks.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/layers/native/transformer_blocks.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/__init__.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/_cudnn.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/autoencoder_blocks.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/transformer_blocks.py create mode 100644 vllm_omni/model_executor/models/common/ops/__init__.py create mode 100644 vllm_omni/model_executor/models/common/ops/fused_adaptive_group_norm_silu.py create mode 100644 vllm_omni/model_executor/models/common/ops/fused_group_norm_silu.py diff --git a/tests/model_executor/models/common/__init__.py b/tests/model_executor/models/common/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/common/test_fused_adaptive_group_norm_silu.py b/tests/model_executor/models/common/test_fused_adaptive_group_norm_silu.py new file mode 100644 index 000000000..e202321dc --- /dev/null +++ b/tests/model_executor/models/common/test_fused_adaptive_group_norm_silu.py @@ -0,0 +1,292 @@ +"""Unit tests for fused_adaptive_group_norm_silu operator. + +Tests numeric correctness against PyTorch native implementation: + F.silu(F.group_norm(x, num_groups, weight, bias, eps) * (1 + scale) + shift) +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.model_executor.models.common.ops import fused_adaptive_group_norm_silu + +# Skip tests if CUDA not available +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for Triton kernels"), +] + + +def _reference(x, weight, bias, scale, shift, num_groups, eps): + """AdaGN reference, computed in fp32. + + As with fused_group_norm_silu, a same-dtype reference is unusable: the + GroupNorm affine can cancel to near zero, and rounding its O(1) terms to + bf16 wipes out the remaining significant digits. + """ + B, C = x.shape[:2] + normed = F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps) + scale_b = scale.float().view(B, C, *([1] * (x.ndim - 2))) + shift_b = shift.float().view(B, C, *([1] * (x.ndim - 2))) + return F.silu(normed * (1 + scale_b) + shift_b) + + +def _tolerance(dtype): + if dtype == torch.float32: + return 1e-5, 1e-5 + if dtype == torch.float16: + return 2e-3, 2e-3 + return 1e-2, 1e-2 # bfloat16 + + +@pytest.mark.parametrize("batch_size", [1, 2]) +@pytest.mark.parametrize("channels", [32, 64, 128]) +@pytest.mark.parametrize("spatial_size", [(16, 16), (32, 32)]) +@pytest.mark.parametrize("num_groups", [8, 32]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_fused_adaptive_group_norm_silu_correctness(batch_size, channels, spatial_size, num_groups, dtype): + """Test numeric correctness against PyTorch native ops.""" + if channels % num_groups != 0: + pytest.skip(f"channels {channels} not divisible by num_groups {num_groups}") + + H, W = spatial_size + eps = 1e-6 + torch.manual_seed(42) + kw = dict(device="cuda", dtype=dtype) + + x = torch.randn(batch_size, channels, H, W, **kw) + weight = torch.randn(channels, **kw) + bias = torch.randn(channels, **kw) + scale = torch.randn(batch_size, channels, **kw) + shift = torch.randn(batch_size, channels, **kw) + + ref_out = _reference(x, weight, bias, scale, shift, num_groups, eps).to(dtype) + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, num_groups, eps) + + rtol, atol = _tolerance(dtype) + torch.testing.assert_close( + fused_out, + ref_out, + rtol=rtol, + atol=atol, + msg=f"Mismatch for batch={batch_size}, C={channels}, HW={spatial_size}, groups={num_groups}, dtype={dtype}", + ) + + +def test_broadcast_scale_shift_shapes(): + """scale/shift may arrive as (B, C) or already broadcast to (B, C, 1, 1).""" + B, C, H, W = 2, 64, 16, 16 + kw = dict(device="cuda", dtype=torch.float32) + x = torch.randn(B, C, H, W, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + flat = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + expanded = fused_adaptive_group_norm_silu(x, weight, bias, scale.view(B, C, 1, 1), shift.view(B, C, 1, 1), 32, 1e-6) + torch.testing.assert_close(flat, expanded, rtol=0, atol=0) + + +def test_output_dtype_matches_eager(): + """Output dtype must match what eager GroupNorm would produce.""" + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + B, C = 2, 64 + kw = dict(device="cuda", dtype=dtype) + x = torch.randn(B, C, 16, 16, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + eager_dtype = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)).dtype + out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + + assert out.dtype == eager_dtype, ( + f"Output dtype {out.dtype} != eager dtype {eager_dtype} for input dtype {dtype}" + ) + assert out.shape == x.shape + + +@pytest.mark.parametrize( + "autocast_dtype, input_dtype", + [ + (torch.float16, torch.float16), + (torch.bfloat16, torch.bfloat16), + (torch.float16, torch.float32), + (torch.bfloat16, torch.float32), + ], +) +def test_autocast_matches_eager(autocast_dtype, input_dtype): + """Inside autocast the fused op must stay a drop-in replacement. + + Also pins the fp32-activation corner so the inlined ``is_autocast_enabled + -> fp32`` branch cannot silently downcast an already-fp32 activation. + """ + B, C = 2, 64 + x = torch.randn(B, C, 16, 16, device="cuda", dtype=input_dtype) + weight = torch.randn(C, device="cuda", dtype=torch.float32) + bias = torch.randn(C, device="cuda", dtype=torch.float32) + scale = torch.randn(B, C, device="cuda", dtype=torch.float32) + shift = torch.randn(B, C, device="cuda", dtype=torch.float32) + + with torch.autocast("cuda", dtype=autocast_dtype): + eager_out = F.silu( + F.group_norm(x, 32, weight, bias, 1e-6) * (1 + scale.view(B, C, 1, 1)) + shift.view(B, C, 1, 1) + ) + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + + assert eager_out.dtype == torch.float32, "sanity check: eager group_norm is expected to return fp32 under autocast" + assert fused_out.dtype == eager_out.dtype, ( + f"Output dtype {fused_out.dtype} != eager dtype {eager_out.dtype} " + f"under autocast({autocast_dtype}) for input {input_dtype}" + ) + torch.testing.assert_close(fused_out, eager_out, rtol=1e-5, atol=1e-5) + + +def test_invalid_channels_not_divisible(): + """Test that assertion fires when channels not divisible by num_groups.""" + B, C = 1, 33 + kw = dict(device="cuda", dtype=torch.float32) + x = torch.randn(B, C, 16, 16, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + with pytest.raises(AssertionError, match="must be divisible by num_groups"): + fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("shape", [(2, 64, 3, 16, 16), (1, 32, 8, 8), (2, 64, 128)]) +def test_non_2d_spatial_matches_eager(dtype, shape): + """Any spatial rank must work; the wrapper collapses the spatial axes. + + GroupNorm reduces over the whole channel group and every spatial position, + so flattening is exact rather than an approximation. + """ + torch.manual_seed(0) + B, C = shape[:2] + kw = dict(device="cuda", dtype=dtype) + x = torch.randn(*shape, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + ref_out = _reference(x, weight, bias, scale, shift, 32, 1e-6).to(dtype) + + assert fused_out.shape == x.shape + rtol, atol = _tolerance(dtype) + torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("batch_size", [1, 2, 3]) +def test_scale_shift_from_chunk_are_not_contiguous(batch_size): + """Regression: scale/shift arrive as halves of a torch.chunk. + + ResBlock builds them as ``torch.chunk(emb_out, 2, dim=1)`` where emb_out is + (B, 2C, 1, 1). For B > 1 each half is a strided view, so the wrapper must + ``reshape`` (which copies) instead of ``view`` (which raises). + """ + C, H, W = 64, 16, 16 + kw = dict(device="cuda", dtype=torch.float32) + x = torch.randn(batch_size, C, H, W, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + + emb_out = torch.randn(batch_size, 2 * C, 1, 1, **kw) + scale, shift = torch.chunk(emb_out, 2, dim=1) + if batch_size > 1: + assert not scale.is_contiguous(), "test premise: chunk half is strided" + + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6) * (1.0 + scale) + shift) + + torch.testing.assert_close(fused_out, ref_out, rtol=1e-5, atol=1e-5) + + +def test_non_contiguous_input_matches_eager(): + """A channels-last / permuted activation must still be handled correctly.""" + B, C, H, W = 2, 64, 16, 16 + kw = dict(device="cuda", dtype=torch.float32) + x = torch.randn(B, H, W, C, **kw).permute(0, 3, 1, 2) + assert not x.is_contiguous(), "test premise: input is strided" + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6) * (1.0 + scale.view(B, C, 1, 1)) + shift.view(B, C, 1, 1)) + + torch.testing.assert_close(fused_out, ref_out, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize("spatial", [32, 64]) +def test_dit_sized_activation(spatial): + """Shapes the HunyuanImage3 DiT actually sees. + + The earlier kernel walked the spatial axis one position at a time with + ``spatial_size`` as a ``tl.constexpr``; at 64x64 that is 4096 serial steps + per pass. This test pins the realistic size so a regression to that layout + shows up as a timeout rather than silently costing throughput. + """ + B, C = 2, 256 + kw = dict(device="cuda", dtype=torch.bfloat16) + x = torch.randn(B, C, spatial, spatial, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-5) + ref_out = _reference(x, weight, bias, scale, shift, 32, 1e-5).to(torch.bfloat16) + + rtol, atol = _tolerance(torch.bfloat16) + torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_large_offset_small_variance(dtype): + """A large constant offset with a small variance must stay stable. + + Mirrors the fused_group_norm_silu case: ``x ~ 10000 ± 0.1`` has mean ~1e4 + and variance ~1e-2, so ``E[x^2] - E[x]^2`` cancels catastrophically in fp32 + (rounding the variance to ~8.0). The Welford reduction keeps it accurate. + fp32 tolerance is 0.1 (not 1e-6) because ``x``'s fp32 representation noise + (~0.001 ULP at 1e4) is amplified by ``rstd ~ 10`` to ~0.01-0.02 — the same + floor PyTorch's own reference hits; the old naive implementation would be + off by ~1.0 here. fp16/bf16 quantize the perturbation away, so for those + dtypes the meaningful assertion is just "no NaN". + """ + torch.manual_seed(0) + B, C, H, W = 2, 64, 16, 16 + num_groups = 32 + eps = 1e-6 + device = torch.device("cuda") + + x = torch.full((B, C, H, W), 10000.0, device=device, dtype=dtype) + x = x + torch.randn(B, C, H, W, device=device, dtype=dtype) * 0.1 + weight = torch.randn(C, device=device, dtype=dtype) + bias = torch.randn(C, device=device, dtype=dtype) + scale = torch.randn(B, C, device=device, dtype=dtype) + shift = torch.randn(B, C, device=device, dtype=dtype) + + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, num_groups, eps) + ref_out = _reference(x, weight, bias, scale, shift, num_groups, eps).to(dtype) + + assert not torch.isnan(fused_out).any(), "fused output must not contain NaN" + + if dtype == torch.float32: + # ~0.02 is the representation-noise floor at offset 1e4; old impl ~1.0. + torch.testing.assert_close(fused_out, ref_out, rtol=0.05, atol=0.1) + # fp16/bf16: perturbation is swallowed by the dtype — no meaningful close, + # the no-NaN assertion above is the regression guard. + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/model_executor/models/common/test_fused_group_norm_silu.py b/tests/model_executor/models/common/test_fused_group_norm_silu.py new file mode 100644 index 000000000..3e9d055af --- /dev/null +++ b/tests/model_executor/models/common/test_fused_group_norm_silu.py @@ -0,0 +1,338 @@ +"""Unit tests for fused_group_norm_silu operator. + +Tests numeric correctness against PyTorch native implementation: + F.silu(F.group_norm(x, num_groups, weight, bias, eps)) +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.model_executor.models.common.ops import fused_group_norm_silu + +# Skip tests if CUDA not available +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for Triton kernels"), +] + + +def _tolerance(dtype): + """Tolerance against an fp32 reference, rounded once to ``dtype``.""" + if dtype == torch.float32: + return 1e-5, 1e-6 + if dtype == torch.float16: + return 2e-3, 2e-3 + return 1e-2, 1e-2 # bfloat16 + + +@pytest.mark.parametrize("batch_size", [1, 2, 4]) +@pytest.mark.parametrize("channels", [32, 64, 128]) +@pytest.mark.parametrize("spatial_size", [(16, 16), (32, 32), (64, 64)]) +@pytest.mark.parametrize("num_groups", [8, 16, 32]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_fused_group_norm_silu_correctness(batch_size, channels, spatial_size, num_groups, dtype): + """Test numeric correctness against PyTorch native ops.""" + # Skip invalid configurations + if channels % num_groups != 0: + pytest.skip(f"channels {channels} not divisible by num_groups {num_groups}") + + H, W = spatial_size + device = torch.device("cuda") + eps = 1e-6 + + # Create inputs + torch.manual_seed(42) + x = torch.randn(batch_size, channels, H, W, device=device, dtype=dtype) + weight = torch.randn(channels, device=device, dtype=dtype) + bias = torch.randn(channels, device=device, dtype=dtype) + + # Reference implementation, computed in fp32. + # + # A same-dtype reference is NOT usable here: GroupNorm's affine step + # ``(x - mean) * rstd * weight + bias`` adds two O(1) terms that can cancel + # down to ~1e-3. Rounding those terms to bf16 (ULP ~8e-3 near 1.0) destroys + # every remaining significant digit -- measured cases flip sign, e.g. a true + # +0.006136 comes out as -0.000471. The low-precision "reference" is then + # far less accurate than the kernel it is supposed to validate: the fused + # output matches the correctly-rounded fp32 result for 99.997% of elements, + # a bf16 reference for only 65.2%. + ref_out = F.silu(F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps)).to(dtype) + + # Fused implementation + fused_out = fused_group_norm_silu(x, weight, bias, num_groups, eps) + + # Check correctness + # Tolerances are set from the measured worst case across this parameter + # sweep, leaving roughly 2x headroom for each dtype. + if dtype == torch.float32: + rtol, atol = 1e-5, 1e-6 + elif dtype == torch.float16: + rtol, atol = 2e-3, 2e-3 + else: # bfloat16 + rtol, atol = 1e-2, 1e-2 + + torch.testing.assert_close( + fused_out, + ref_out, + rtol=rtol, + atol=atol, + msg=f"Mismatch for batch={batch_size}, C={channels}, HW={spatial_size}, groups={num_groups}, dtype={dtype}", + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_fp32_accumulation_precision(dtype): + """Test that fp32 accumulation matches PyTorch's behavior. + + GroupNorm uses fp32 for mean/variance computation regardless of input dtype. + Our kernel should match this behavior. + """ + batch_size, channels, H, W = 2, 64, 32, 32 + num_groups = 32 + eps = 1e-6 + device = torch.device("cuda") + + # Create inputs with values that expose precision differences + torch.manual_seed(123) + x = torch.randn(batch_size, channels, H, W, device=device, dtype=dtype) * 100 + weight = torch.ones(channels, device=device, dtype=dtype) + bias = torch.zeros(channels, device=device, dtype=dtype) + + # Reference: fp32 accumulation, then a single round back to ``dtype``. + # (Same discipline as the main correctness test -- a same-dtype reference + # wouldn't isolate the accumulation path this test is named after.) + ref_out = F.silu(F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps)).to(dtype) + + # Fused + fused_out = fused_group_norm_silu(x, weight, bias, num_groups, eps) + + # Should match closely even with large values + if dtype == torch.float32: + rtol, atol = 1e-5, 1e-5 + else: + rtol, atol = 1e-2, 1e-2 + + torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol) + + +def test_edge_case_zero_input(): + """Test with zero input.""" + x = torch.zeros(1, 32, 16, 16, device="cuda") + weight = torch.ones(32, device="cuda") + bias = torch.zeros(32, device="cuda") + + ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)) + fused_out = fused_group_norm_silu(x, weight, bias, 32, 1e-6) + + torch.testing.assert_close(fused_out, ref_out, rtol=1e-5, atol=1e-6) + + +def test_edge_case_large_eps(): + """Test with large epsilon value.""" + x = torch.randn(1, 32, 16, 16, device="cuda") + weight = torch.randn(32, device="cuda") + bias = torch.randn(32, device="cuda") + eps = 1e-2 # Larger than typical + + ref_out = F.silu(F.group_norm(x, 32, weight, bias, eps)) + fused_out = fused_group_norm_silu(x, weight, bias, 32, eps) + + torch.testing.assert_close(fused_out, ref_out, rtol=1e-4, atol=1e-5) + + +def test_output_dtype_matches_eager(): + """Output dtype must match what eager GroupNorm+SiLU would produce. + + Outside autocast that means the input dtype; inside autocast it means fp32, + because group_norm is on autocast's fp32 cast policy. HunyuanImage3's VAE + runs fp32 weights under fp16 autocast, so getting this wrong would silently + downcast the post-norm activation. + """ + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + x = torch.randn(2, 64, 32, 32, device="cuda", dtype=dtype) + weight = torch.randn(64, device="cuda", dtype=dtype) + bias = torch.randn(64, device="cuda", dtype=dtype) + + eager_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)) + out = fused_group_norm_silu(x, weight, bias, 32, 1e-6) + + assert out.dtype == eager_out.dtype, ( + f"Output dtype {out.dtype} != eager dtype {eager_out.dtype} for input dtype {dtype}" + ) + assert out.shape == x.shape, f"Output shape {out.shape} != input shape {x.shape}" + + +@pytest.mark.parametrize( + "autocast_dtype, input_dtype", + [ + (torch.float16, torch.float16), + (torch.bfloat16, torch.bfloat16), + (torch.float16, torch.float32), + (torch.bfloat16, torch.float32), + ], +) +def test_autocast_matches_eager(autocast_dtype, input_dtype): + """Inside autocast the fused op must stay a drop-in replacement. + + Mirrors the HunyuanImage3 VAE setup: fp32 module weights, activations + arriving as fp16/bf16 from an upstream conv, all under autocast. Also pins + the fp32-activation corner (some VAE stages) so the inlined + ``is_autocast_enabled -> fp32`` branch cannot silently downcast it. + """ + x = torch.randn(2, 64, 32, 32, device="cuda", dtype=input_dtype) + weight = torch.randn(64, device="cuda", dtype=torch.float32) + bias = torch.randn(64, device="cuda", dtype=torch.float32) + + with torch.autocast("cuda", dtype=autocast_dtype): + eager_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)) + fused_out = fused_group_norm_silu(x, weight, bias, 32, 1e-6) + + assert eager_out.dtype == torch.float32, "sanity check: eager group_norm is expected to return fp32 under autocast" + assert fused_out.dtype == eager_out.dtype, ( + f"Output dtype {fused_out.dtype} != eager dtype {eager_out.dtype} " + f"under autocast({autocast_dtype}) for input {input_dtype}" + ) + torch.testing.assert_close(fused_out, eager_out, rtol=1e-5, atol=1e-5) + + +def test_invalid_channels_not_divisible(): + """Test that assertion fires when channels not divisible by num_groups.""" + x = torch.randn(1, 33, 16, 16, device="cuda") # 33 not divisible by 32 + weight = torch.randn(33, device="cuda") + bias = torch.randn(33, device="cuda") + + with pytest.raises(AssertionError, match="must be divisible by num_groups"): + fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6) + + +def test_invalid_weight_shape(): + """Test that assertion fires with wrong weight shape.""" + x = torch.randn(1, 64, 16, 16, device="cuda") + weight = torch.randn(32, device="cuda") # Wrong: should be 64 + bias = torch.randn(64, device="cuda") + + with pytest.raises(AssertionError, match="doesn't match channels"): + fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("shape", [(2, 64, 3, 16, 16), (1, 32, 1, 8, 8), (2, 64, 128)]) +def test_non_2d_spatial_matches_eager(dtype, shape): + """Spatial ranks other than 2 must work: the 3D VAE feeds (N, C, T, H, W). + + GroupNorm reduces over the whole channel group and every spatial position, + so flattening the spatial axes inside the wrapper has to be exact. + """ + torch.manual_seed(0) + C = shape[1] + x = torch.randn(*shape, device="cuda", dtype=dtype) + weight = torch.randn(C, device="cuda", dtype=dtype) + bias = torch.randn(C, device="cuda", dtype=dtype) + + fused_out = fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6) + ref_out = F.silu(F.group_norm(x.float(), 32, weight.float(), bias.float(), 1e-6)).to(dtype) + + assert fused_out.shape == x.shape + rtol, atol = _tolerance(dtype) + torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_permuted_input_matches_eager(dtype): + """A permuted (channel-stride-1) activation must still be correct. + + This is not a hypothetical layout: HunyuanImage3's ``UNetUp`` calls this op + on the output of ``rearrange(x, "b (h w) c -> b c h w")``. An earlier + version of the kernel indexed such inputs through their strides, which was + correct but uncoalesced enough to run at 0.44x of eager, so the wrapper now + materializes them first. + """ + torch.manual_seed(0) + B, C, H, W = 2, 64, 16, 16 + kw = dict(device="cuda", dtype=dtype) + x = torch.randn(B, H * W, C, **kw).permute(0, 2, 1).reshape(B, C, H, W) + assert not x.is_contiguous(), "test premise: input is strided" + + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + + fused_out = fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6) + ref_out = F.silu(F.group_norm(x.float(), 32, weight.float(), bias.float(), 1e-6)).to(dtype) + + rtol, atol = _tolerance(dtype) + torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("channels", [64, 128]) +def test_hunyuan_vae_config(batch_size, channels): + """Test with HunyuanImage3 VAE actual configuration. + + HunyuanImage3 ResnetBlock uses: + - num_groups=32 + - eps=1e-6 + - Typical spatial sizes: 16x16 to 128x128 (after tiling) + """ + spatial_sizes = [(16, 16), (32, 32), (64, 64), (128, 128)] + + for H, W in spatial_sizes: + x = torch.randn(batch_size, channels, H, W, device="cuda", dtype=torch.float32) + weight = torch.randn(channels, device="cuda") + bias = torch.randn(channels, device="cuda") + + ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)) + fused_out = fused_group_norm_silu(x, weight, bias, 32, 1e-6) + + torch.testing.assert_close( + fused_out, + ref_out, + rtol=1e-5, + atol=1e-6, + msg=f"HunyuanVAE config: batch={batch_size}, C={channels}, HW=({H},{W})", + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_large_offset_small_variance(dtype): + """A large constant offset with a small variance must stay stable. + + ``x ~ 10000 ± 0.1`` has mean ~1e4 and variance ~1e-2, so the naive + ``E[x^2] - E[x]^2`` subtracts two ~1e8 quantities and rounds the variance to + ~8.0 (the ULP of 1e8) — catastrophic cancellation, producing garbage + normalization. The Welford reduction keeps the variance accurate. + + Why fp32 tolerance is 0.1, not the usual 1e-6: at this offset, ``x`` is stored + in fp32 with ULP ~0.001, so ``x - mean`` carries ~0.001 of representation + noise, and ``rstd ~ 10`` amplifies it to ~0.01-0.02. Even PyTorch's own + GroupNorm reference hits this floor (the old naive implementation would be + off by ~1.0 here). fp16/bf16 quantize the 0.1 perturbation away entirely, so + for those dtypes the meaningful assertion is just "no NaN". + """ + torch.manual_seed(0) + B, C, H, W = 2, 64, 16, 16 + num_groups = 32 + eps = 1e-6 + device = torch.device("cuda") + + x = torch.full((B, C, H, W), 10000.0, device=device, dtype=dtype) + x = x + torch.randn(B, C, H, W, device=device, dtype=dtype) * 0.1 + weight = torch.randn(C, device=device, dtype=dtype) + bias = torch.randn(C, device=device, dtype=dtype) + + fused_out = fused_group_norm_silu(x, weight, bias, num_groups, eps) + # fp32 reference = PyTorch's Welford path, the ground truth. + ref_out = F.silu(F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps)).to(dtype) + + assert not torch.isnan(fused_out).any(), "fused output must not contain NaN" + + if dtype == torch.float32: + # ~0.02 is the representation-noise floor at offset 1e4; old impl ~1.0. + torch.testing.assert_close(fused_out, ref_out, rtol=0.05, atol=0.1) + # fp16/bf16: perturbation is swallowed by the dtype — no meaningful close, + # the no-NaN assertion above is the regression guard. + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/__init__.py index 625895d99..4c93e2397 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/__init__.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/__init__.py @@ -10,4 +10,8 @@ from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import ( HunyuanImage3Pipeline, ) -__all__ = ["HunyuanImage3Pipeline", "HunyuanImage3Model", "HunyuanImage3Text2ImagePipeline"] +__all__ = [ + "HunyuanImage3Pipeline", + "HunyuanImage3Model", + "HunyuanImage3Text2ImagePipeline", +] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/autoencoder.py b/vllm_omni/diffusion/models/hunyuan_image3/autoencoder.py index efba2f274..1ccbafa65 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/autoencoder.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/autoencoder.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import math from dataclasses import dataclass import numpy as np @@ -15,6 +14,10 @@ from diffusers.utils.torch_utils import randn_tensor from einops import rearrange from torch import Tensor, nn +# ResnetBlock is platform-dispatched in the layers package __init__: CUDA gets a +# fused GroupNorm+SiLU kernel, every other backend gets the plain PyTorch block. +from vllm_omni.diffusion.models.hunyuan_image3.layers import ResnetBlock + class DiagonalGaussianDistribution: def __init__(self, parameters: torch.Tensor, deterministic: bool = False): @@ -53,10 +56,6 @@ class DecoderOutput(BaseOutput): posterior: DiagonalGaussianDistribution | None = None -def swish(x: Tensor) -> Tensor: - return x * torch.sigmoid(x) - - def forward_with_checkpointing(module, *inputs, use_checkpointing=False): def create_custom_forward(module): def custom_forward(*inputs): @@ -70,46 +69,6 @@ def forward_with_checkpointing(module, *inputs, use_checkpointing=False): return module(*inputs) -class Conv3d(nn.Conv3d): - """ - Perform Conv3d on patches with numerical differences from nn.Conv3d within 1e-5. - Only symmetric padding is supported. - """ - - def forward(self, input): - B, C, T, H, W = input.shape - memory_count = (C * T * H * W) * 2 / 1024**3 - if memory_count > 2: - n_split = math.ceil(memory_count / 2) - assert n_split >= 2 - chunks = torch.chunk(input, chunks=n_split, dim=-3) - padded_chunks = [] - for i in range(len(chunks)): - if self.padding[0] > 0: - padded_chunk = F.pad( - chunks[i], - (0, 0, 0, 0, self.padding[0], self.padding[0]), - mode="constant" if self.padding_mode == "zeros" else self.padding_mode, - value=0, - ) - if i > 0: - padded_chunk[:, :, : self.padding[0]] = chunks[i - 1][:, :, -self.padding[0] :] - if i < len(chunks) - 1: - padded_chunk[:, :, -self.padding[0] :] = chunks[i + 1][:, :, : self.padding[0]] - else: - padded_chunk = chunks[i] - padded_chunks.append(padded_chunk) - padding_bak = self.padding - self.padding = (0, self.padding[1], self.padding[2]) - outputs = [] - for i in range(len(padded_chunks)): - outputs.append(super().forward(padded_chunks[i])) - self.padding = padding_bak - return torch.cat(outputs, dim=-3) - else: - return super().forward(input) - - class AttnBlock(nn.Module): """Attention with torch sdpa implementation.""" @@ -119,10 +78,12 @@ class AttnBlock(nn.Module): self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) - self.q = Conv3d(in_channels, in_channels, kernel_size=1) - self.k = Conv3d(in_channels, in_channels, kernel_size=1) - self.v = Conv3d(in_channels, in_channels, kernel_size=1) - self.proj_out = Conv3d(in_channels, in_channels, kernel_size=1) + # The memory concern is already handled via VAE tiling, + # so there's no need to chunk here — calling nn.Conv3d directly is fine. + self.q = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.k = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.v = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.proj_out = nn.Conv3d(in_channels, in_channels, kernel_size=1) def attention(self, h_: Tensor) -> Tensor: h_ = self.norm(h_) @@ -142,41 +103,12 @@ class AttnBlock(nn.Module): return x + self.proj_out(self.attention(x)) -class ResnetBlock(nn.Module): - def __init__(self, in_channels: int, out_channels: int): - super().__init__() - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - - self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) - self.conv1 = Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) - self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) - self.conv2 = Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) - if self.in_channels != self.out_channels: - self.nin_shortcut = Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) - - def forward(self, x): - h = x - h = self.norm1(h) - h = swish(h) - h = self.conv1(h) - - h = self.norm2(h) - h = swish(h) - h = self.conv2(h) - - if self.in_channels != self.out_channels: - x = self.nin_shortcut(x) - return x + h - - class DownsampleDCAE(nn.Module): def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True): super().__init__() factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 assert out_channels % factor == 0 - self.conv = Conv3d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) + self.conv = nn.Conv3d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) self.add_temporal_downsample = add_temporal_downsample self.group_size = factor * in_channels // out_channels @@ -196,7 +128,7 @@ class UpsampleDCAE(nn.Module): def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True): super().__init__() factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 - self.conv = Conv3d(in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) + self.conv = nn.Conv3d(in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) self.add_temporal_upsample = add_temporal_upsample self.repeats = factor * out_channels // in_channels @@ -233,7 +165,7 @@ class Encoder(nn.Module): self.num_res_blocks = num_res_blocks # downsampling - self.conv_in = Conv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) + self.conv_in = nn.Conv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) self.down = nn.ModuleList() block_in = block_out_channels[0] @@ -265,7 +197,7 @@ class Encoder(nn.Module): # end self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) - self.conv_out = Conv3d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + self.conv_out = nn.Conv3d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) self.gradient_checkpointing = False @@ -291,7 +223,7 @@ class Encoder(nn.Module): group_size = self.block_out_channels[-1] // (2 * self.z_channels) shortcut = rearrange(h, "b (c r) f h w -> b c r f h w", r=group_size).mean(dim=2) h = self.norm_out(h) - h = swish(h) + h = F.silu(h) h = self.conv_out(h) h += shortcut return h @@ -321,7 +253,7 @@ class Decoder(nn.Module): # z to block_in block_in = block_out_channels[0] - self.conv_in = Conv3d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + self.conv_in = nn.Conv3d(z_channels, block_in, kernel_size=3, stride=1, padding=1) # middle self.mid = nn.Module() @@ -351,7 +283,7 @@ class Decoder(nn.Module): # end self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) - self.conv_out = Conv3d(block_in, out_channels, kernel_size=3, stride=1, padding=1) + self.conv_out = nn.Conv3d(block_in, out_channels, kernel_size=3, stride=1, padding=1) self.gradient_checkpointing = False @@ -376,7 +308,7 @@ class Decoder(nn.Module): # end h = self.norm_out(h) - h = swish(h) + h = F.silu(h) h = self.conv_out(h) return h diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 45a94d76d..9328abe01 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -79,6 +79,13 @@ from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step from vllm_omni.diffusion.layers.fused_moe import FusedMoE from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding + +# ResBlock is platform-dispatched in the layers package __init__: CUDA gets +# fused GroupNorm+SiLU and AdaGN kernels, every other backend gets the plain +# PyTorch block. UNetDown and UNetUp below instantiate whichever one this +# resolves to. +from vllm_omni.diffusion.models.hunyuan_image3.layers import ResBlock +from vllm_omni.diffusion.models.hunyuan_image3.layers.common import conv_nd, normalization from vllm_omni.diffusion.utils.kv_utils import repeat_kv from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding from vllm_omni.platforms import current_omni_platform @@ -171,44 +178,6 @@ def real_batched_index_select(t, dim, idx): return torch.stack([torch.index_select(t[i], dim - 1, idx[i]) for i in range(len(t))]) -def conv_nd(dims, *args, **kwargs): # noqa: N802 - """ - Create a 1D, 2D, or 3D convolution module. - """ - if dims == 1: - return nn.Conv1d(*args, **kwargs) - elif dims == 2: - return nn.Conv2d(*args, **kwargs) - elif dims == 3: - return nn.Conv3d(*args, **kwargs) - raise ValueError(f"unsupported dimensions: {dims}") - - -def normalization(channels, **kwargs): - """ - Make a standard normalization layer. - :param channels: number of input channels. - :return: a nn.Module for normalization. - """ - return nn.GroupNorm(32, channels, **kwargs) - - -def linear(*args, **kwargs): - """ - Create a linear module. - """ - return nn.Linear(*args, **kwargs) - - -def zero_module(module): - """ - Zero out the parameters of a module and return it. - """ - for p in module.parameters(): - p.detach().zero_() - return module - - def _to_tuple(x, dim=2): if isinstance(x, int): return (x,) * dim @@ -3282,99 +3251,6 @@ class TimestepEmbedder(nn.Module): return t_emb -class ResBlock(nn.Module): - r""" - A residual block that can optionally change the number of channels. - - Args: - in_channels (`int`): - The number of input channels. - emb_channels (`int`): - The number of timestep embedding channels. - dropout (`float`): - The rate of dropout. - out_channels (`int`, *optional*): - If specified, the number of output channels. - use_conv (`bool`, *optional*): - If True and out_channels is specified, use a spatial convolution instead of a - smaller 1x1 convolution to change the channels in the skip connection. - dims (`int`, *optional*): - Determines if the signal is 1D, 2D, or 3D. - up (`bool`, *optional*): - If True, use this block for upsampling. - down (`bool`, *optional*): - If True, use this block for downsampling. - - """ - - def __init__( - self, - in_channels, - emb_channels, - out_channels=None, - dropout=0.0, - use_conv=False, - dims=2, - up=False, - down=False, - device=None, - dtype=None, - ) -> None: - factory_kwargs = {"dtype": dtype, "device": device} - super().__init__() - self.in_channels = in_channels - self.dropout = dropout - self.out_channels = out_channels or self.in_channels - self.use_conv = use_conv - - self.in_layers = nn.Sequential( - normalization(self.in_channels, **factory_kwargs), - nn.SiLU(), - conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs), # noqa: N802 - ) - - self.updown = up or down - self.h_upd = self.x_upd = nn.Identity() - - self.emb_layers = nn.Sequential(nn.SiLU(), linear(emb_channels, 2 * self.out_channels, **factory_kwargs)) - - self.out_layers = nn.Sequential( - normalization(self.out_channels, **factory_kwargs), - nn.SiLU(), - nn.Dropout(p=dropout), - zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, **factory_kwargs)), # noqa: N802 - ) - - if self.out_channels == self.in_channels: - self.skip_connection = nn.Identity() - elif use_conv: - self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs) # noqa: N802 - else: - self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 1, **factory_kwargs) # noqa: N802 - - def forward(self, x, emb) -> torch.Tensor: - if self.updown: - in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] - h = in_rest(x) - h = self.h_upd(h) - x = self.x_upd(x) - h = in_conv(h) - else: - h = self.in_layers(x) - - emb_out = self.emb_layers(emb) - while len(emb_out.shape) < len(h.shape): - emb_out = emb_out[..., None] - - # Adaptive Group Normalization - out_norm, out_rest = self.out_layers[0], self.out_layers[1:] - scale, shift = torch.chunk(emb_out, 2, dim=1) - h = out_norm(h) * (1.0 + scale) + shift - h = out_rest(h) - - return self.skip_connection(x) + h - - class UNetDown(nn.Module): def __init__( self, patch_size, in_channels, emb_channels, hidden_channels, out_channels, dropout=0.0, device=None, dtype=None diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/__init__.py new file mode 100644 index 000000000..39081df17 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/__init__.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Platform-dispatched layer implementations for HunyuanImage3. + +The two blocks are selected by platform: + +Layout:: + native/autoencoder_blocks.py default ResnetBlock + native/transformer_blocks.py default ResBlock + nvidia/autoencoder_blocks.py CUDA ResnetBlock (fused) + nvidia/transformer_blocks.py CUDA ResBlock (fused) +""" + +from vllm_omni.platforms import current_omni_platform + +if current_omni_platform.is_cuda(): + from vllm_omni.diffusion.models.hunyuan_image3.layers.nvidia import ResBlock, ResnetBlock +else: + from vllm_omni.diffusion.models.hunyuan_image3.layers.native import ResBlock, ResnetBlock + +__all__ = ["ResBlock", "ResnetBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/common/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/common/__init__.py new file mode 100644 index 000000000..5272e9d51 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/common/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .utils import ( + conv_nd, + linear, + normalization, + zero_module, +) + +__all__ = ["conv_nd", "linear", "normalization", "zero_module"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/common/utils.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/common/utils.py new file mode 100644 index 000000000..01dd1cc4e --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/common/utils.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Small platform-agnostic helpers for UNet-style blocks (adapted from diffusers). + +``conv_nd`` / ``normalization`` / ``linear`` / ``zero_module`` are used by the +``ResBlock`` / ``UNetDown`` / ``UNetUp`` implementations across HunyuanImage3's +platform-specific layer packages. They are plain ``nn.ConvNd`` / ``nn.GroupNorm`` +/ ``nn.Linear`` wrappers with no hardware-specific behaviour, so they live here +instead of being duplicated per backend. +""" + +from torch import nn + + +def conv_nd(dims, *args, **kwargs): # noqa: N802 + """ + Create a 1D, 2D, or 3D convolution module. + """ + if dims == 1: + return nn.Conv1d(*args, **kwargs) + elif dims == 2: + return nn.Conv2d(*args, **kwargs) + elif dims == 3: + return nn.Conv3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +def normalization(channels, **kwargs): + """ + Make a standard normalization layer. + :param channels: number of input channels. + :return: a nn.Module for normalization. + """ + return nn.GroupNorm(32, channels, **kwargs) + + +def linear(*args, **kwargs): + """ + Create a linear module. + """ + return nn.Linear(*args, **kwargs) + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/native/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/__init__.py new file mode 100644 index 000000000..55a695af7 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Native (plain PyTorch) layer implementations.""" + +from vllm_omni.diffusion.models.hunyuan_image3.layers.native.autoencoder_blocks import ResnetBlock +from vllm_omni.diffusion.models.hunyuan_image3.layers.native.transformer_blocks import ResBlock + +__all__ = ["ResBlock", "ResnetBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/native/autoencoder_blocks.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/autoencoder_blocks.py new file mode 100644 index 000000000..6e85a52f6 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/autoencoder_blocks.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""ResnetBlock for the HunyuanImage3 autoencoder — default fallback implementation. + +Uses native PyTorch ``nn.GroupNorm`` + ``F.silu``. Hardware-specific optimized +versions live in subdirectories (e.g., nvidia/) and are selected via __init__.py +dispatch. +""" + +import torch.nn.functional as F +from torch import nn + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + h = x + h = self.norm1(h) + h = F.silu(h) + h = self.conv1(h) + + h = self.norm2(h) + h = F.silu(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +__all__ = ["ResnetBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/native/transformer_blocks.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/transformer_blocks.py new file mode 100644 index 000000000..2322fa6e3 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/transformer_blocks.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""DiT ``ResBlock`` for HunyuanImage3 -- default implementation. + +Plain PyTorch: ``GroupNorm -> SiLU`` runs as two ops and the adaptive +GroupNorm as ``norm(h) * (1 + scale) + shift``. Hardware-specific variants +live in sibling packages (e.g. ``nvidia/``) and are picked by the dispatch in +``hunyuan_image3/__init__.py``. + +""" + +import torch +from torch import nn + +from vllm_omni.diffusion.models.hunyuan_image3.layers.common import conv_nd, linear, normalization, zero_module + + +class ResBlock(nn.Module): + r""" + A residual block that can optionally change the number of channels. + Args: + in_channels (`int`): + The number of input channels. + emb_channels (`int`): + The number of timestep embedding channels. + dropout (`float`): + The rate of dropout. + out_channels (`int`, *optional*): + If specified, the number of output channels. + use_conv (`bool`, *optional*): + If True and out_channels is specified, use a spatial convolution instead of a + smaller 1x1 convolution to change the channels in the skip connection. + dims (`int`, *optional*): + Determines if the signal is 1D, 2D, or 3D. + up (`bool`, *optional*): + If True, use this block for upsampling. + down (`bool`, *optional*): + If True, use this block for downsampling. + """ + + def __init__( + self, + in_channels, + emb_channels, + out_channels=None, + dropout=0.0, + use_conv=False, + dims=2, + up=False, + down=False, + device=None, + dtype=None, + ) -> None: + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.in_channels = in_channels + self.dropout = dropout + self.out_channels = out_channels or self.in_channels + self.use_conv = use_conv + + self.in_layers = nn.Sequential( + normalization(self.in_channels, **factory_kwargs), + nn.SiLU(), + conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs), # noqa: N802 + ) + + self.updown = up or down + self.h_upd = self.x_upd = nn.Identity() + + self.emb_layers = nn.Sequential(nn.SiLU(), linear(emb_channels, 2 * self.out_channels, **factory_kwargs)) + + self.out_layers = nn.Sequential( + normalization(self.out_channels, **factory_kwargs), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, **factory_kwargs)), # noqa: N802 + ) + + if self.out_channels == self.in_channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs) # noqa: N802 + else: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 1, **factory_kwargs) # noqa: N802 + + def forward(self, x, emb) -> torch.Tensor: + if self.updown: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + + emb_out = self.emb_layers(emb) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + + # Adaptive Group Normalization + out_norm, out_rest = self.out_layers[0], self.out_layers[1:] + scale, shift = torch.chunk(emb_out, 2, dim=1) + h = out_norm(h) * (1.0 + scale) + shift + h = out_rest(h) + + return self.skip_connection(x) + h + + +__all__ = ["ResBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/__init__.py new file mode 100644 index 000000000..78d3304a5 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""NVIDIA (CUDA) implementations.""" + +from .autoencoder_blocks import ResnetBlock # noqa: E402 +from .transformer_blocks import ResBlock # noqa: E402 + +__all__ = ["ResnetBlock", "ResBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/_cudnn.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/_cudnn.py new file mode 100644 index 000000000..4c541591d --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/_cudnn.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import contextmanager + +import torch + + +@contextmanager +def cudnn_settings(benchmark: bool, deterministic: bool): + """Temporarily set cuDNN backend flags, restoring the old values on exit. + + Used to scope ``torch.backends.cudnn.benchmark``/``deterministic`` to the + few blocks whose convolutions actually benefit from autotune, instead of + setting them globally at import time (which would silently change every + other model's conv behaviour in the same process). + """ + old_benchmark = torch.backends.cudnn.benchmark + old_deterministic = torch.backends.cudnn.deterministic + + torch.backends.cudnn.benchmark = benchmark + torch.backends.cudnn.deterministic = deterministic + try: + yield + finally: + torch.backends.cudnn.benchmark = old_benchmark + torch.backends.cudnn.deterministic = old_deterministic diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/autoencoder_blocks.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/autoencoder_blocks.py new file mode 100644 index 000000000..c6c68957e --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/autoencoder_blocks.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""ResnetBlock for the HunyuanImage3 autoencoder — NVIDIA CUDA + Triton implementation. + +Split out from autoencoder.py because its ``GroupNorm -> SiLU`` pairs are served +by :func:`fused_group_norm_silu`, a single Triton kernel that falls back to +native ``F.silu(F.group_norm(...))`` when Triton is unavailable. +""" + +from torch import nn + +from vllm_omni.diffusion.models.hunyuan_image3.layers.nvidia._cudnn import cudnn_settings +from vllm_omni.model_executor.models.common.ops import fused_group_norm_silu + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + # cuDNN autotune is scoped to this block rather than set at import time. + # These ResnetBlocks carry the bulk of the VAE's conv work: ~29 instances + # x 2 convs across every resolution (up to the 1024x1024 decode), versus + # the scattered convs left in autoencoder.py (AttnBlock's 1x1 at the + # lowest-res mid, down/up-samplers, conv_in/out) whose activations are + # negligible. A process-wide flag would instead silently change every + # other model's conv behaviour in the same process. + with cudnn_settings(benchmark=True, deterministic=True): + h = x + h = fused_group_norm_silu( + h, self.norm1.weight, self.norm1.bias, num_groups=self.norm1.num_groups, eps=self.norm1.eps + ) + h = self.conv1(h) + + h = fused_group_norm_silu( + h, self.norm2.weight, self.norm2.bias, num_groups=self.norm2.num_groups, eps=self.norm2.eps + ) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +__all__ = ["ResnetBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/transformer_blocks.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/transformer_blocks.py new file mode 100644 index 000000000..411dc6351 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/transformer_blocks.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""``ResBlock`` for HunyuanImage3 -- NVIDIA CUDA + Triton implementation. + +Two fusions over the default block: + +* ``in_layers``: ``GroupNorm -> SiLU`` becomes one :func:`fused_group_norm_silu` + kernel instead of two ops. +* ``out_layers``: ``GroupNorm -> adaptive modulation -> SiLU`` (i.e. + ``SiLU(norm(h) * (1 + scale) + shift)``) becomes one + :func:`fused_adaptive_group_norm_silu` kernel instead of four. +""" + +import torch +from torch import nn + +from vllm_omni.diffusion.models.hunyuan_image3.layers.common import conv_nd, normalization, zero_module +from vllm_omni.diffusion.models.hunyuan_image3.layers.nvidia._cudnn import cudnn_settings +from vllm_omni.model_executor.models.common.ops import ( + fused_adaptive_group_norm_silu, + fused_group_norm_silu, +) + + +class ResBlock(nn.Module): + r""" + A residual block that can optionally change the number of channels. + Args: + in_channels (`int`): + The number of input channels. + emb_channels (`int`): + The number of timestep embedding channels. + dropout (`float`): + The rate of dropout. + out_channels (`int`, *optional*): + If specified, the number of output channels. + use_conv (`bool`, *optional*): + If True and out_channels is specified, use a spatial convolution instead of a + smaller 1x1 convolution to change the channels in the skip connection. + dims (`int`, *optional*): + Determines if the signal is 1D, 2D, or 3D. + up (`bool`, *optional*): + If True, use this block for upsampling. + down (`bool`, *optional*): + If True, use this block for downsampling. + """ + + def __init__( + self, + in_channels, + emb_channels, + out_channels=None, + dropout=0.0, + use_conv=False, + dims=2, + up=False, + down=False, + device=None, + dtype=None, + ) -> None: + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.in_channels = in_channels + self.dropout = dropout + self.out_channels = out_channels or self.in_channels + self.use_conv = use_conv + + self.in_layers = nn.Sequential( + normalization(self.in_channels, **factory_kwargs), + nn.SiLU(), + conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs), # noqa: N802 + ) + + self.updown = up or down + self.h_upd = self.x_upd = nn.Identity() + + self.emb_layers = nn.Sequential(nn.SiLU(), nn.Linear(emb_channels, 2 * self.out_channels, **factory_kwargs)) + + self.out_layers = nn.Sequential( + normalization(self.out_channels, **factory_kwargs), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, **factory_kwargs)), # noqa: N802 + ) + + if self.out_channels == self.in_channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs) # noqa: N802 + else: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 1, **factory_kwargs) # noqa: N802 + + def forward(self, x, emb) -> torch.Tensor: + # cuDNN autotune is scoped to this block rather than set at import time. + # The benchmark=True win is concentrated in the convolutions right here: + # this ResBlock sits inside UNetDown/UNetUp (patch_embed / final_layer), + # the heaviest conv users in the DiT. A process-wide flag would instead + # silently change every other model's conv behaviour in the same process. + with cudnn_settings(benchmark=True, deterministic=True): + # ``in_layers``/``out_layers`` stay nn.Sequential and are indexed into + # rather than unpacked into named submodules -- that is what keeps the + # state_dict keys (``in_layers.0/.2``, ``out_layers.0/.3``) identical to + # the unfused block, so checkpoints load unchanged. + in_norm, in_conv = self.in_layers[0], self.in_layers[-1] + + # GroupNorm -> SiLU, one kernel instead of two. + h = fused_group_norm_silu(x, in_norm.weight, in_norm.bias, num_groups=in_norm.num_groups, eps=in_norm.eps) + if self.updown: + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + + emb_out = self.emb_layers(emb) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + + # Adaptive GroupNorm -> SiLU: GroupNorm -> mul -> add -> silu, one + # kernel instead of four. The ``nn.SiLU`` at ``out_layers[1]`` is + # skipped here but stays in the Sequential, so state_dict keys are + # unchanged. + out_norm, out_rest = self.out_layers[0], self.out_layers[2:] + scale, shift = torch.chunk(emb_out, 2, dim=1) + h = fused_adaptive_group_norm_silu( + h, + out_norm.weight, + out_norm.bias, + scale, + shift, + num_groups=out_norm.num_groups, + eps=out_norm.eps, + ) + h = out_rest(h) + + return self.skip_connection(x) + h + + +__all__ = ["ResBlock"] diff --git a/vllm_omni/model_executor/models/common/ops/__init__.py b/vllm_omni/model_executor/models/common/ops/__init__.py new file mode 100644 index 000000000..d9b522bdf --- /dev/null +++ b/vllm_omni/model_executor/models/common/ops/__init__.py @@ -0,0 +1,17 @@ +"""Common operators shared across models and hardware backends. + +This module contains fused operators that can be reused across multiple models +and compiled for different hardware targets (CUDA, ROCm, etc.). +""" + +from vllm_omni.model_executor.models.common.ops.fused_adaptive_group_norm_silu import ( + fused_adaptive_group_norm_silu, +) +from vllm_omni.model_executor.models.common.ops.fused_group_norm_silu import ( + fused_group_norm_silu, +) + +__all__ = [ + "fused_group_norm_silu", + "fused_adaptive_group_norm_silu", +] diff --git a/vllm_omni/model_executor/models/common/ops/fused_adaptive_group_norm_silu.py b/vllm_omni/model_executor/models/common/ops/fused_adaptive_group_norm_silu.py new file mode 100644 index 000000000..2755c3ece --- /dev/null +++ b/vllm_omni/model_executor/models/common/ops/fused_adaptive_group_norm_silu.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# ruff: noqa: N803 + +"""Fused Adaptive Group Normalization (AdaGN) + SiLU operator. + +This module implements the fused AdaGN+SiLU pattern commonly used in Diffusion +Transformer ResBlocks: + + output = SiLU(GroupNorm(x, weight, bias) * (1 + scale) + shift) + +where ``scale`` and ``shift`` are per-(batch, channel) conditioning signals +derived from the timestep embedding. Eager PyTorch spends four kernels and three +full-size intermediates on this; the fused version does it in one pass. + +Falls back to native PyTorch ops when Triton is unavailable. +""" + +import torch +import torch.nn.functional as F +from vllm.triton_utils import HAS_TRITON, tl, triton + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=1), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=2), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=4), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=6), + ], + key=["spatial_size", "C"], +) +@triton.jit +def _adaptive_group_norm_silu_kernel( + # Input/output pointers + x_ptr, + out_ptr, + # Affine parameters + weight_ptr, + bias_ptr, + # Conditioning signals, both (B, C) contiguous + scale_ptr, + shift_ptr, + # Shape info; x is contiguous (B, C, spatial_size) + C, + spatial_size, + num_groups: tl.constexpr, + eps: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + num_stages: tl.constexpr, +): + """ + Fused AdaGN + SiLU kernel: SiLU(GroupNorm(x) * (1 + scale) + shift). + Mean and variance use a parallel Welford reduction (block-level register + reduction merged with the Chan formula), which avoids the catastrophic + cancellation of E[x^2] - E[x]^2 on large-offset inputs. x is read once for + the statistics and once for normalization (2 reads total). + One program handles each (batch, group) pair. Channels are processed serially, + while the spatial axis is vectorized for diffusion workloads. + Moments are accumulated in fp32 to match PyTorch numerics. + """ + pid = tl.program_id(0) + + group_size = C // num_groups + n_idx = pid // num_groups + g_idx = pid % num_groups + + # === Pass 1: Welford reduction over the whole group (fp32) === + # Each block is reduced in registers (block mean, then centered M2), so the + # loaded values are reused instead of re-reading memory; blocks are merged + # with the Chan et al. parallel-Welford formula. This keeps both the mean and + # the variance accurate for inputs like ``10000 +- 0.1``, where the naive + # ``E[x^2] - E[x]^2`` cancels catastrophically, and reads x only once for the + # statistics (2 reads total including normalize). + n_total = tl.zeros([1], dtype=tl.float32) + mean_total = tl.zeros([1], dtype=tl.float32) + m2_total = tl.zeros([1], dtype=tl.float32) + + for c_offset in range(group_size): + c_idx = g_idx * group_size + c_offset + base = n_idx * C * spatial_size + c_idx * spatial_size + + for s_start in tl.range(0, spatial_size, BLOCK_SIZE, num_stages=num_stages): + offsets = s_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < spatial_size + + x_val = tl.load(x_ptr + base + offsets, mask=mask, other=0.0) + x_val = x_val.to(tl.float32) + + # block-level reduction in registers (no extra memory traffic) + n = tl.sum(tl.where(mask, 1.0, 0.0), axis=0) + bsum = tl.sum(x_val, axis=0) + bmean = bsum / n + bm2 = tl.sum(tl.where(mask, (x_val - bmean) * (x_val - bmean), 0.0), axis=0) + + # Chan et al. merge into the running (n, mean, m2) + delta = bmean - mean_total + new_n = n_total + n + mean_total = mean_total + delta * (n / new_n) + m2_total = m2_total + bm2 + delta * delta * (n_total * n / new_n) + n_total = new_n + + mean = mean_total + var = m2_total / n_total + rstd = 1.0 / tl.sqrt(var + eps) + + # === Pass 2: normalize, affine, then adaptive modulation === + for c_offset in range(group_size): + c_idx = g_idx * group_size + c_offset + base = n_idx * C * spatial_size + c_idx * spatial_size + + weight_val = tl.load(weight_ptr + c_idx).to(tl.float32) + bias_val = tl.load(bias_ptr + c_idx).to(tl.float32) + scale_val = tl.load(scale_ptr + n_idx * C + c_idx).to(tl.float32) + shift_val = tl.load(shift_ptr + n_idx * C + c_idx).to(tl.float32) + + for s_start in tl.range(0, spatial_size, BLOCK_SIZE, num_stages=num_stages): + offsets = s_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < spatial_size + + x_val = tl.load(x_ptr + base + offsets, mask=mask, other=0.0) + x_val = x_val.to(tl.float32) + + # GroupNorm: (x - mean) * rstd * weight + bias + norm_val = (x_val - mean) * rstd * weight_val + bias_val + + # AdaGN: norm * (1 + scale) + shift + out_val = norm_val * (1.0 + scale_val) + shift_val + + # SiLU: out * sigmoid(out) + out_val = out_val * tl.sigmoid(out_val) + + # ``tl.store`` casts to the output pointer's dtype, which the + # caller picked to match eager GroupNorm's autocast behaviour. + tl.store(out_ptr + base + offsets, out_val, mask=mask) + + +def fused_adaptive_group_norm_silu( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + num_groups: int, + eps: float = 1e-6, +) -> torch.Tensor: + """ + Fused Adaptive GroupNorm + SiLU. + + Computes SiLU(GroupNorm(x) * (1 + scale) + shift) in one fused operation, avoiding intermediate tensors. + + - x: (B, C, *spatial), any spatial rank ≥ 1. + - weight, bias: per-channel parameters, shape (C,). + - scale, shift: adaptive parameters with B * C elements, broadcastable per channel. + - num_groups: GroupNorm group count. + - eps: numerical stability epsilon. + + Returns a tensor with the same shape as x, matching eager F.group_norm dtype behavior. + + Spatial dimensions are flattened during computation and restored afterward. + GroupNorm statistics cover each channel group across all spatial positions, so the result is exact. + """ + assert x.ndim >= 3, f"Expected at least 3D input (B, C, *spatial), got {x.ndim}D" + B, C = x.shape[:2] + assert C % num_groups == 0, f"num_channels ({C}) must be divisible by num_groups ({num_groups})" + assert weight.ndim == 1 and weight.size(0) == C, f"Weight shape {tuple(weight.shape)} doesn't match channels {C}" + assert bias.ndim == 1 and bias.size(0) == C, f"Bias shape {tuple(bias.shape)} doesn't match channels {C}" + assert scale.numel() == B * C, f"scale has {scale.numel()} elements, expected B*C = {B * C}" + assert shift.numel() == B * C, f"shift has {shift.numel()} elements, expected B*C = {B * C}" + + # Fallback if Triton not available (NPU, CPU, ...) + if not HAS_TRITON: + broadcast = (B, C) + (1,) * (x.ndim - 2) + normed = F.group_norm(x, num_groups, weight, bias, eps) + return F.silu(normed * (1.0 + scale.reshape(broadcast)) + shift.reshape(broadcast)) + + orig_shape = x.shape + + # The kernel indexes x as a dense (B, C, spatial) block, so normalize the + # layout here. ``reshape``/``contiguous`` are no-ops for the common case of + # a contiguous activation coming out of a conv. + x_flat = x.contiguous().reshape(B, C, -1) + spatial_size = x_flat.size(2) + + # ``reshape`` rather than ``view``: scale/shift typically arrive as halves + # of a torch.chunk along dim 1, which is not contiguous when B > 1. + scale_2d = scale.reshape(B, C).contiguous() + shift_2d = shift.reshape(B, C).contiguous() + + # Allocate with the dtype eager GroupNorm would return, so the fused path + # stays a drop-in replacement inside autocast regions. + out_dtype = x_flat.dtype + if torch.is_autocast_enabled(x_flat.device.type): + out_dtype = torch.float32 + out_flat = torch.empty_like(x_flat, dtype=out_dtype) + + # Only B*num_groups programs are launched, which is well under the SM count + # for typical diffusion batches. A memory-bound kernel can still saturate + # HBM from few CTAs, but only with enough loads in flight, so widen the CTA + # for the large activations instead of leaving it at the 4-warp default. + + # One program per (batch, group) pair. + grid = (B * num_groups,) + + _adaptive_group_norm_silu_kernel[grid]( + x_flat, + out_flat, + weight, + bias, + scale_2d, + shift_2d, + C, + spatial_size, + num_groups=num_groups, + eps=eps, + ) + + return out_flat.reshape(orig_shape) + + +__all__ = ["fused_adaptive_group_norm_silu"] diff --git a/vllm_omni/model_executor/models/common/ops/fused_group_norm_silu.py b/vllm_omni/model_executor/models/common/ops/fused_group_norm_silu.py new file mode 100644 index 000000000..21fb245c6 --- /dev/null +++ b/vllm_omni/model_executor/models/common/ops/fused_group_norm_silu.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# ruff: noqa: N803 + +"""Fused GroupNorm + SiLU operator. + +This operator fuses GroupNorm followed by SiLU activation into a single kernel, +reducing memory traffic and kernel launch overhead. The implementation uses +Triton for CUDA/ROCm compatibility, and falls back to native PyTorch ops when +Triton is unavailable (NPU, CPU, ...), so callers never need a platform check. + +Measured against eager ``F.silu(F.group_norm(...))`` on one L20X, bf16, 32 +groups: 1.1-1.5x at the DiT ResBlock's activation sizes, where both paths are +dominated by launch overhead, and 2.2-2.9x at the VAE's decode-resolution +activations, where the saved memory traffic is what pays. +""" + +import torch +import torch.nn.functional as F +from vllm.triton_utils import HAS_TRITON, tl, triton + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=1), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=2), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=4), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=6), + ], + key=["spatial_size", "C"], +) +@triton.jit +def _group_norm_silu_kernel( + # Input/Output pointers + x_ptr, + out_ptr, + # Normalization parameters + weight_ptr, + bias_ptr, + # Shape info; x is contiguous (N, C, spatial_size) + C, + spatial_size, + num_groups: tl.constexpr, + eps: tl.constexpr, + # Block sizes + BLOCK_SIZE: tl.constexpr, + num_stages: tl.constexpr, +): + """ + Fused GroupNorm + SiLU kernel. + Computes SiLU(GroupNorm(x)) in one kernel, avoiding intermediate tensors. + Mean and variance use a parallel Welford reduction (block-level register + reduction merged with the Chan formula), which avoids the catastrophic + cancellation of E[x^2] - E[x]^2 on large-offset inputs. x is read once for + the statistics and once for normalization (2 reads total). + One program handles each (batch, group) pair. Channels are processed serially, + while spatial positions are vectorized for diffusion workloads. + Moments are accumulated in fp32 to match PyTorch numerics. + """ + pid = tl.program_id(0) + + group_size = C // num_groups + n_idx = pid // num_groups + g_idx = pid % num_groups + + # === Pass 1: Welford reduction over the whole group (fp32) === + # Each block is reduced in registers (block mean, then centered M2), so the + # loaded values are reused instead of re-reading memory; blocks are merged + # with the Chan et al. parallel-Welford formula. This keeps both the mean and + # the variance accurate for inputs like ``10000 +- 0.1``, where the naive + # ``E[x^2] - E[x]^2`` cancels catastrophically, and reads x only once for the + # statistics (2 reads total including normalize). + n_total = tl.zeros([1], dtype=tl.float32) + mean_total = tl.zeros([1], dtype=tl.float32) + m2_total = tl.zeros([1], dtype=tl.float32) + + for c_offset in range(group_size): + c_idx = g_idx * group_size + c_offset + base = n_idx * C * spatial_size + c_idx * spatial_size + + for s_start in tl.range(0, spatial_size, BLOCK_SIZE, num_stages=num_stages): + offsets = s_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < spatial_size + + x_val = tl.load(x_ptr + base + offsets, mask=mask, other=0.0) + x_val = x_val.to(tl.float32) + + # block-level reduction in registers (no extra memory traffic) + n = tl.sum(tl.where(mask, 1.0, 0.0), axis=0) + bsum = tl.sum(x_val, axis=0) + bmean = bsum / n + bm2 = tl.sum(tl.where(mask, (x_val - bmean) * (x_val - bmean), 0.0), axis=0) + + # Chan et al. merge into the running (n, mean, m2) + delta = bmean - mean_total + new_n = n_total + n + mean_total = mean_total + delta * (n / new_n) + m2_total = m2_total + bm2 + delta * delta * (n_total * n / new_n) + n_total = new_n + + mean = mean_total + var = m2_total / n_total + rstd = 1.0 / tl.sqrt(var + eps) + + # === Pass 2: normalize, apply affine, and SiLU === + for c_offset in range(group_size): + c_idx = g_idx * group_size + c_offset + base = n_idx * C * spatial_size + c_idx * spatial_size + + weight_val = tl.load(weight_ptr + c_idx).to(tl.float32) + bias_val = tl.load(bias_ptr + c_idx).to(tl.float32) + + for s_start in tl.range(0, spatial_size, BLOCK_SIZE, num_stages=num_stages): + offsets = s_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < spatial_size + + x_val = tl.load(x_ptr + base + offsets, mask=mask, other=0.0) + x_val = x_val.to(tl.float32) + + # Normalize and apply affine + norm_val = (x_val - mean) * rstd * weight_val + bias_val + + # Apply SiLU: x * sigmoid(x) + out_val = norm_val * tl.sigmoid(norm_val) + + # ``tl.store`` casts to the output pointer's dtype, which the + # caller picked to match eager GroupNorm's autocast behaviour. + tl.store(out_ptr + base + offsets, out_val, mask=mask) + + +def fused_group_norm_silu( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + num_groups: int = 32, + eps: float = 1e-6, +) -> torch.Tensor: + """ + Fused GroupNorm + SiLU activation. + Computes SiLU(GroupNorm(x, num_groups, weight, bias, eps)) in a single Triton kernel, + avoiding intermediate tensors and reducing memory traffic and launch overhead. + + - x: (N, C, *spatial), any spatial rank. + - weight, bias: per-channel parameters, shape (C,). + - num_groups: GroupNorm group count, default 32. + - eps: numerical stability epsilon, default 1e-6. + Uses fp32 accumulation for numeric alignment with PyTorch. + Returns the same shape and eager F.group_norm-compatible dtype. + Spatial dimensions are flattened during computation and restored afterward; + this is exact since GroupNorm reduces across each channel group and all spatial positions. + Non-contiguous inputs are materialized before launch. + """ + # Fallback if Triton not available (NPU, CPU, ...) + if not HAS_TRITON: + return F.silu(F.group_norm(x, num_groups, weight, bias, eps)) + + # Validate inputs + assert x.ndim >= 3, f"Expected at least 3D input (N, C, *spatial), got {x.ndim}D" + assert x.size(1) % num_groups == 0, f"Channels {x.size(1)} must be divisible by num_groups {num_groups}" + assert weight.ndim == 1 and weight.size(0) == x.size(1), ( + f"Weight shape {weight.shape} doesn't match channels {x.size(1)}" + ) + assert bias.ndim == 1 and bias.size(0) == x.size(1), f"Bias shape {bias.shape} doesn't match channels {x.size(1)}" + + # Collapse arbitrary spatial ranks into a single axis so one kernel serves + # both the 2D DiT blocks and the 3D VAE blocks. + # + # ``contiguous()`` is not just defensive. HunyuanImage3's UNetUp feeds this + # op straight out of ``rearrange(x, "b (h w) c -> b c h w")``, i.e. a + # permuted view whose *channel* stride is 1. Indexing that layout directly + # from the kernel makes every warp's spatial load stride by C elements, and + # the resulting uncoalesced traffic turned the op into 0.44x of eager at + # (1, 4096, 64, 64). One coalesced pre-pass costs far less than that, so + # normalize the layout here and let the kernel assume a dense block. + orig_shape = x.shape + B, C = orig_shape[0], orig_shape[1] + x_flat = x.contiguous().reshape(B, C, -1) + spatial_size = x_flat.size(2) + + # Allocate output with the dtype eager GroupNorm would return, so that the + # fused path stays a drop-in replacement inside autocast regions. + out_dtype = x_flat.dtype + if torch.is_autocast_enabled(x_flat.device.type): + out_dtype = torch.float32 + out_flat = torch.empty_like(x_flat, dtype=out_dtype) + + # Only B*num_groups programs are launched, which is well under the SM count + # for typical diffusion batches. A memory-bound kernel can still saturate + # HBM from few CTAs, but only with enough loads in flight, so widen the CTA + # for the large activations instead of leaving it at the 4-warp default. + + # One program per (batch, group) pair. + grid = (B * num_groups,) + + _group_norm_silu_kernel[grid]( + x_flat, + out_flat, + weight, + bias, + C, + spatial_size, + num_groups=num_groups, + eps=eps, + ) + + return out_flat.reshape(orig_shape)