[Perf][Hunyuan] Optimize for vae and patch_embed (#6306)

Signed-off-by: dengyunyang <584797741@qq.com>

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"])

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"])

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",
+]

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
 

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

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"]

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"]

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

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"]

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"]

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"]

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"]

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

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"]

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"]

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",
+]

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"]

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)