[Perf][Model] Add optimized LTX-2.5 DiffVAE operators (#7308)

Signed-off-by: mglyn <1203789601@qq.com>

tests/diffusion/models/ltx2/test_ltx2_diffvae_attention.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Native FNA against an independent, test-only explicit neighborhood mask."""
+
+import pytest
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from vllm_omni.diffusion.models.ltx2.ops import resolve_ltx2_vae_operators
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
+
+
+def _selected_fna():
+    if not current_omni_platform.is_available():
+        pytest.skip("No accelerator is available")
+    operators = resolve_ltx2_vae_operators(current_omni_platform.get_torch_device())
+    if operators is None or operators.fna is None:
+        pytest.skip("LTX DiffVAE FNA is unavailable for this device or installation")
+    return operators.fna
+
+
+def _load_fna_ops():
+    return pytest.importorskip("vllm_omni.diffusion.models.ltx2.ops.fna", exc_type=ImportError)
+
+
+def _pack(x, tile):
+    b, t, h, w, heads, dim = x.shape
+    tt, th, tw = tile
+    nt, nh, nw = (t + tt - 1) // tt, (h + th - 1) // th, (w + tw - 1) // tw
+    padded = x.new_zeros((b, nt * tt, nh * th, nw * tw, heads, dim))
+    padded[:, :t, :h, :w] = x
+    # Outer and inner spatial tiles both flatten in W,H,T order.
+    return (
+        padded.view(b, nt, tt, nh, th, nw, tw, heads, dim)
+        .permute(0, 5, 3, 1, 6, 4, 2, 7, 8)
+        .reshape(b, nt * nh * nw * tt * th * tw, heads, dim)
+    )
+
+
+def _unpack(x, shape):
+    t, h, w = shape
+    b, _, heads, dim = x.shape
+    nt, nh, nw = ((n + 3) // 4 for n in shape)
+    return (
+        x.view(b, nw, nh, nt, 4, 4, 4, heads, dim)
+        .permute(0, 3, 6, 2, 5, 1, 4, 7, 8)
+        .reshape(b, nt * 4, nh * 4, nw * 4, heads, dim)[:, :t, :h, :w]
+    )
+
+
+@pytest.mark.parametrize("batch,shape", [(1, (11, 11, 11)), (2, (13, 12, 17))])
+@torch.inference_mode()
+def test_fna_matches_explicit_inward_shift_mask(batch, shape):
+    fna = _selected_fna()
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    q, k, v = [
+        torch.randn(batch, *shape, 4, 64, device="cuda", dtype=torch.bfloat16, generator=generator) for _ in range(3)
+    ]
+    q = q * 0.125
+    coordinates = torch.stack(
+        torch.meshgrid(*(torch.arange(n, device="cuda") for n in shape), indexing="ij"), dim=-1
+    ).reshape(-1, 3)
+    start = (coordinates - 5).clamp_min(0).minimum(torch.tensor(shape, device="cuda") - 11)
+    mask = ((coordinates[None] >= start[:, None]) & (coordinates[None] < start[:, None] + 11)).all(-1)
+    with sdpa_kernel(SDPBackend.MATH):
+        expected = (
+            torch.nn.functional.scaled_dot_product_attention(
+                *(x.flatten(1, 3).transpose(1, 2).float() for x in (q, k, v)), attn_mask=mask, scale=1.0
+            )
+            .transpose(1, 2)
+            .reshape(q.shape)
+        )
+    actual = _unpack(
+        fna(_pack(q, (4, 4, 4)), _pack(k, (4, 4, 8)), _pack(v, (4, 4, 8)), shape=shape, window=(11, 11, 11)),
+        shape,
+    ).float()
+    assert actual.isfinite().all()
+    assert float((actual - expected).norm() / expected.norm()) < 0.005
+    torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.025)
+
+
+def test_fna_rejects_cpu_without_fallback():
+    fna_ops = _load_fna_ops()
+    x = torch.zeros(1, 64, 4, 64, dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match="CUDA BF16"):
+        fna_ops.fna3d_tilelang(x, x, x, shape=(11, 11, 11), window=(11, 11, 11))

tests/diffusion/models/ltx2/test_ltx2_diffvae_ops.py

@@ -0,0 +1,409 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.ltx2.ops import qk_rms_norm as qk_rms_norm_ops
+from vllm_omni.diffusion.models.ltx2.ops import residual_adaln as residual_adaln_ops
+from vllm_omni.diffusion.models.ltx2.ops import resolve_ltx2_vae_operators
+from vllm_omni.diffusion.models.ltx2.ops import swiglu as swiglu_ops
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
+
+_DIM_SPLIT = (16, 24, 24)
+
+
+def _selected_operators():
+    if not current_omni_platform.is_available():
+        pytest.skip("No accelerator is available")
+    operators = resolve_ltx2_vae_operators(current_omni_platform.get_torch_device())
+    if operators is None:
+        pytest.skip("No LTX DiffVAE operators are available for this device")
+    return operators
+
+
+def _selected_fna():
+    operators = _selected_operators()
+    if operators.fna is None:
+        pytest.skip("LTX DiffVAE FNA is unavailable for this device or installation")
+    return operators.fna
+
+
+def _load_fna_ops():
+    return pytest.importorskip("vllm_omni.diffusion.models.ltx2.ops.fna", exc_type=ImportError)
+
+
+def test_diffvae_model_and_tests_do_not_import_tilelang_backend() -> None:
+    # Use a fresh interpreter: reloading the model in this process would leave
+    # already-imported distributed subclasses pointing at the old base class.
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import sys; "
+            "sys.modules['vllm_omni.diffusion.models.ltx2.ops.fna'] = None; "
+            "sys.modules['vllm.tilelang_utils'] = None; "
+            "from vllm_omni.diffusion.models.ltx2.vae.decoder import LTX2VideoDiffusionDecoderModel; "
+            "import tests.diffusion.models.ltx2.test_ltx2_diffvae_attention; "
+            "import tests.diffusion.models.ltx2.test_ltx2_diffvae_ops",
+        ],
+        capture_output=True,
+        text=True,
+        timeout=60,
+    )
+    assert result.returncode == 0, result.stdout + result.stderr
+
+
+def test_diffvae_fna_tile_map_covers_exact_neighborhood() -> None:
+    fna_ops = _load_fna_ops()
+    problem = fna_ops.Problem(13, 14, 17, (11, 11, 11))
+    indices, counts = fna_ops.build_kv_tile_map(problem)
+    q_grid_t, q_grid_h, _ = problem.q_grid
+    kv_grid_t, kv_grid_h, _ = problem.kv_grid
+
+    for q_block in range(indices.shape[0]):
+        selected = set(indices[q_block, : counts[q_block]].tolist())
+        q_outer_t = q_block % q_grid_t
+        q_outer_h = (q_block // q_grid_t) % q_grid_h
+        q_outer_w = q_block // (q_grid_t * q_grid_h)
+        for q_inner in range(fna_ops.Q_TOKENS):
+            query = (
+                q_outer_t * 4 + q_inner % 4,
+                q_outer_h * 4 + (q_inner // 4) % 4,
+                q_outer_w * 4 + q_inner // 16,
+            )
+            if any(index >= length for index, length in zip(query, problem.shape, strict=True)):
+                continue
+            starts = tuple(
+                fna_ops._window_start(index, window, length)
+                for index, window, length in zip(query, problem.window, problem.shape, strict=True)
+            )
+            for key_t in range(starts[0], starts[0] + problem.window[0]):
+                for key_h in range(starts[1], starts[1] + problem.window[1]):
+                    for key_w in range(starts[2], starts[2] + problem.window[2]):
+                        key_block = ((key_w // 8) * kv_grid_h + key_h // 4) * kv_grid_t + key_t // 4
+                        assert key_block in selected
+
+
+def test_diffvae_fna_metadata_uses_static_kernel_shapes() -> None:
+    fna_ops = _load_fna_ops()
+    for shape in ((11, 11, 11), (13, 14, 17), (66, 96, 128), (79, 192, 192)):
+        problem = fna_ops.Problem(*shape, (11, 11, 11))
+        _, t_patterns, h_patterns, w_patterns = fna_ops.build_kv_tile_metadata(problem)
+        assert {table.numel() for table in (t_patterns, h_patterns, w_patterns)} == {fna_ops.PATTERN_TABLE_SIZE}
+        assert {limit for _, _, limit in fna_ops.build_kv_tile_buckets(problem)} <= set(fna_ops.KV_BUCKET_LIMITS)
+        for ids, metadata, limit in fna_ops.build_kv_tile_buckets(problem):
+            assert metadata.shape == (ids.numel(), limit)
+
+
+def test_diffvae_fna_rejects_oversized_packed_tile_grid() -> None:
+    fna_ops = _load_fna_ops()
+    with pytest.raises(ValueError, match="15-bit metadata"):
+        fna_ops.Problem(128, 256, 256, (11, 11, 11))
+
+
+def test_diffvae_fna_propagates_launch_failures(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    fna_ops = _load_fna_ops()
+    fna = _selected_fna()
+    problem = fna_ops.Problem(11, 11, 11, (11, 11, 11))
+    query = torch.randn(1, problem.q_length, fna_ops.HEADS, 64, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(1, problem.kv_length, fna_ops.HEADS, 64, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    calls = 0
+
+    def fail_launch(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        raise RuntimeError("injected launch failure")
+
+    monkeypatch.setattr(fna_ops, "_fna3d_tilelang", fail_launch)
+    with torch.inference_mode():
+        for _ in range(2):
+            with pytest.raises(RuntimeError, match="injected launch failure"):
+                fna(
+                    query,
+                    key,
+                    value,
+                    shape=problem.shape,
+                    window=problem.window,
+                )
+    assert calls == 2
+
+
+@pytest.fixture
+def isolated_runtime_state():
+    targets = (
+        (qk_rms_norm_ops, "_FAILED_ROPE_KEYS"),
+        (qk_rms_norm_ops, "_VERIFIED_ROPE_KEYS"),
+        (swiglu_ops, "_FAILED_KEYS"),
+        (swiglu_ops, "_VERIFIED_KEYS"),
+        (residual_adaln_ops, "_FAILED_ADALN_KEYS"),
+        (residual_adaln_ops, "_VERIFIED_ADALN_KEYS"),
+        (residual_adaln_ops, "_FAILED_ADD_DEVICES"),
+        (residual_adaln_ops, "_VERIFIED_ADD_DEVICES"),
+    )
+    saved = [set(getattr(module, name)) for module, name in targets]
+    for module, name in targets:
+        getattr(module, name).clear()
+    try:
+        yield
+    finally:
+        for (module, name), state in zip(targets, saved, strict=True):
+            target = getattr(module, name)
+            target.clear()
+            target.update(state)
+
+
+@pytest.mark.parametrize(
+    ("dim", "rows"),
+    [
+        (256, 257),
+        (256, 16385),
+        (256, 131075),
+        (512, 65539),
+        (1024, 32771),
+        (2048, 16387),
+    ],
+)
+def test_diffvae_swiglu_tiled_is_bit_exact(dim: int, rows: int, isolated_runtime_state) -> None:
+    from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeSwiGLU
+
+    operators = _selected_operators()
+    generator = torch.Generator(device="cuda").manual_seed(20260829)
+    hidden_states = torch.randn((1, 1, 1, rows, dim), device="cuda", dtype=torch.bfloat16, generator=generator)
+    gate_weight = torch.randn((4 * dim, dim), device="cuda", dtype=torch.bfloat16, generator=generator)
+    up_weight = torch.randn_like(gate_weight)
+    down_weight = torch.randn((dim, 4 * dim), device="cuda", dtype=torch.bfloat16, generator=generator)
+    with torch.device("meta"):
+        reference = LTX2VideoVaeSwiGLU(dim, 4 * dim)
+    reference.w_gate.weight = torch.nn.Parameter(gate_weight, requires_grad=False)
+    reference.w_up.weight = torch.nn.Parameter(up_weight, requires_grad=False)
+    reference.w_down.weight = torch.nn.Parameter(down_weight, requires_grad=False)
+
+    with torch.inference_mode():
+        expected = reference(hidden_states)
+        actual = operators.swiglu(hidden_states, gate_weight, up_weight, down_weight)
+
+    assert actual is not None
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize(
+    ("dim", "expected_rows"),
+    [(256, 131072), (512, 65536), (1024, 32768), (2048, 16384)],
+)
+def test_diffvae_swiglu_tile_rows_keep_a_constant_hidden_budget(dim: int, expected_rows: int) -> None:
+    assert swiglu_ops._tile_rows(1_000_000, 4 * dim) == expected_rows
+
+
+def test_diffvae_swiglu_tiled_permanently_falls_back_after_failure(
+    monkeypatch: pytest.MonkeyPatch,
+    isolated_runtime_state,
+) -> None:
+    operators = _selected_operators()
+    hidden_states = torch.randn(1, 1, 1, 17, 256, device="cuda", dtype=torch.bfloat16)
+    gate_weight = torch.randn(1024, 256, device="cuda", dtype=torch.bfloat16)
+    up_weight = torch.randn_like(gate_weight)
+    down_weight = torch.randn(256, 1024, device="cuda", dtype=torch.bfloat16)
+    calls = 0
+
+    def fail_launch(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        raise RuntimeError("injected launch failure")
+
+    monkeypatch.setattr(swiglu_ops, "_launch", fail_launch)
+    with torch.inference_mode():
+        assert operators.swiglu(hidden_states, gate_weight, up_weight, down_weight) is None
+        assert operators.swiglu(hidden_states, gate_weight, up_weight, down_weight) is None
+
+    assert calls == 1
+
+
+@pytest.mark.parametrize("residual_count", [1, 2])
+def test_diffvae_residual_rms_norm_modulate_is_bit_exact(
+    residual_count: int,
+    isolated_runtime_state,
+) -> None:
+    operators = _selected_operators()
+    generator = torch.Generator(device="cuda").manual_seed(20260829)
+    shape = (2, 3, 7, 7, 256)
+    x = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator)
+    residual_a = torch.randn_like(x)
+    residual_b = torch.randn_like(x) if residual_count == 2 else None
+    norm_weight = torch.randn((256,), device="cuda", dtype=torch.bfloat16, generator=generator)
+    scale = torch.randn((2, 1, 1, 1, 256), device="cuda", dtype=torch.bfloat16, generator=generator)
+    shift = torch.randn_like(scale)
+    hidden_states = x + residual_a
+    if residual_b is not None:
+        hidden_states = hidden_states + residual_b
+    expected = torch.nn.functional.rms_norm(hidden_states, (256,), norm_weight, eps=1e-6) * (1 + scale) + shift
+
+    with torch.inference_mode():
+        actual = operators.residual_norm(
+            x,
+            residual_a,
+            residual_b,
+            norm_weight,
+            scale,
+            shift,
+            1e-6,
+        )
+
+    assert actual is not None
+    assert torch.equal(actual, expected)
+
+
+def test_diffvae_residual_add3_is_bit_exact(isolated_runtime_state) -> None:
+    operators = _selected_operators()
+    generator = torch.Generator(device="cuda").manual_seed(20260829)
+    tensors = [
+        torch.randn((2, 3, 7, 7, 256), device="cuda", dtype=torch.bfloat16, generator=generator) for _ in range(4)
+    ]
+    expected = tensors[0] + tensors[1]
+    expected = expected + tensors[2]
+    expected = expected + tensors[3]
+
+    with torch.inference_mode():
+        actual = operators.residual_add(*tensors)
+
+    assert actual is not None
+    assert torch.equal(actual, expected)
+
+
+def test_diffvae_residual_adaln_permanently_falls_back_after_failure(
+    monkeypatch: pytest.MonkeyPatch,
+    isolated_runtime_state,
+) -> None:
+    operators = _selected_operators()
+    shape = (1, 3, 7, 7, 256)
+    x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+    residual = torch.randn_like(x)
+    weight = torch.randn(256, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(1, 1, 1, 1, 256, device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn_like(scale)
+    calls = 0
+
+    def fail_launch(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        raise RuntimeError("injected launch failure")
+
+    monkeypatch.setattr(residual_adaln_ops, "_launch_adaln", fail_launch)
+    with torch.inference_mode():
+        for _ in range(2):
+            assert operators.residual_norm(x, residual, None, weight, scale, shift, 1e-6) is None
+
+    assert calls == 1
+
+
+def _tables(
+    frames: int,
+    height: int,
+    width: int,
+    dim_split: tuple[int, int, int] = _DIM_SPLIT,
+    base: float = 10000.0,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
+    tables = []
+    for length, dim in zip((frames, height, width), dim_split, strict=True):
+        exponents = torch.arange(0, dim, 2, dtype=torch.float64, device="cuda") / dim
+        inv_freqs = (1.0 / base**exponents).to(torch.float32)
+        positions = torch.arange(length, dtype=torch.float32, device="cuda")
+        angles = positions[:, None] * inv_freqs[None, :]
+        tables.append((angles.cos(), angles.sin()))
+    return tuple(tables)
+
+
+def _apply_rope_reference(
+    hidden_states: torch.Tensor,
+    tables: tuple[tuple[torch.Tensor, torch.Tensor], ...],
+    dim_split: tuple[int, int, int] = _DIM_SPLIT,
+) -> torch.Tensor:
+    outputs = []
+    offset = 0
+    for axis, (dim, (cos, sin)) in enumerate(zip(dim_split, tables, strict=True), 1):
+        chunk = hidden_states[..., offset : offset + dim]
+        pairs = chunk.reshape(*chunk.shape[:-1], dim // 2, 2)
+        even = pairs[..., 0].float()
+        odd = pairs[..., 1].float()
+        shape = [1, 1, 1, 1, 1, dim // 2]
+        shape[axis] = cos.shape[0]
+        cos = cos.reshape(shape)
+        sin = sin.reshape(shape)
+        rotated = torch.stack([even * cos - odd * sin, even * sin + odd * cos], dim=-1)
+        outputs.append(rotated.reshape(chunk.shape).to(hidden_states.dtype))
+        offset += dim
+    return torch.cat(outputs, dim=-1)
+
+
+@pytest.mark.parametrize(
+    ("frames", "height", "width", "heads"),
+    [(3, 7, 7, 1), (31, 136, 192, 4)],
+)
+def test_diffvae_qk_rms_norm_scale_rope_3d_is_bit_exact(
+    frames: int,
+    height: int,
+    width: int,
+    heads: int,
+    isolated_runtime_state,
+) -> None:
+    operators = _selected_operators()
+    generator = torch.Generator(device="cuda").manual_seed(20260829)
+    shape = (1, frames, height, width, heads, 64)
+    query = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator)
+    key = torch.randn_like(query)
+    query_weight = torch.randn((64,), device="cuda", dtype=torch.bfloat16, generator=generator)
+    key_weight = torch.randn((64,), device="cuda", dtype=torch.bfloat16, generator=generator)
+    tables = _tables(frames, height, width)
+    expected = (
+        _apply_rope_reference(torch.nn.functional.rms_norm(query, (64,), query_weight, 1e-6) * 0.125, tables),
+        _apply_rope_reference(torch.nn.functional.rms_norm(key, (64,), key_weight, 1e-6), tables),
+    )
+
+    with torch.inference_mode():
+        actual = operators.qk_norm_rope(
+            query,
+            key,
+            query_weight,
+            key_weight,
+            1e-6,
+            0.125,
+            _DIM_SPLIT,
+            10000.0,
+        )
+
+    assert actual is not None
+    assert torch.equal(actual[0], expected[0])
+    assert torch.equal(actual[1], expected[1])
+
+
+def test_diffvae_qk_rms_norm_scale_rope_3d_permanently_falls_back_after_failure(
+    monkeypatch: pytest.MonkeyPatch,
+    isolated_runtime_state,
+) -> None:
+    operators = _selected_operators()
+    query = torch.randn(1, 3, 7, 7, 1, 64, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn_like(query)
+    weight = torch.ones(64, device="cuda", dtype=torch.bfloat16)
+    calls = 0
+
+    def fail_launch(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        raise RuntimeError("injected launch failure")
+
+    monkeypatch.setattr(qk_rms_norm_ops, "_launch_combined", fail_launch)
+    with torch.inference_mode():
+        for _ in range(2):
+            assert operators.qk_norm_rope(query, key, weight, weight, 1e-6, 0.125, _DIM_SPLIT, 10000.0) is None
+
+    assert calls == 1

tests/diffusion/models/ltx2/test_ltx2_ops_platform.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+
+import pytest
+import torch
+from vllm.platforms.interface import DeviceCapability
+
+from vllm_omni.diffusion.models.ltx2 import ops as ltx2_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass(frozen=True)
+class _TensorProbe:
+    is_cuda: bool
+    device: torch.device
+
+
+@pytest.fixture(autouse=True)
+def clear_platform_cache():
+    ltx2_platform._cuda_compute_capability.cache_clear()
+    try:
+        yield
+    finally:
+        ltx2_platform._cuda_compute_capability.cache_clear()
+
+
+@pytest.mark.parametrize(
+    ("capability", "triton", "tilelang", "fusion_expected", "fna_expected"),
+    [
+        (80, True, True, False, False),
+        (90, True, True, True, True),
+        (90, True, False, True, False),
+        (90, False, True, False, False),
+        (100, True, True, False, False),
+        (103, True, True, False, False),
+        (110, True, True, False, False),
+        (120, True, True, False, False),
+        (121, True, True, False, False),
+    ],
+)
+def test_production_resolver_and_tensor_eligibility_agree(
+    monkeypatch, capability, triton, tilelang, fusion_expected, fna_expected
+):
+    monkeypatch.setattr(ltx2_platform, "HAS_TRITON", triton)
+    monkeypatch.setattr(ltx2_platform, "has_tilelang", lambda: tilelang)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: True)
+    monkeypatch.setattr(
+        ltx2_platform.current_omni_platform,
+        "get_device_capability",
+        lambda device_id=0: DeviceCapability(major=capability // 10, minor=capability % 10),
+    )
+    # Resolution must be lazy even when FNA is selected.
+    monkeypatch.setitem(sys.modules, "vllm_omni.diffusion.models.ltx2.ops.fna", None)
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False)
+    device = torch.device("cuda:0")
+    operators = ltx2_platform.resolve_ltx2_vae_operators(device)
+    assert (operators is not None) is fusion_expected
+    assert (operators is not None and operators.fna is not None) is fna_expected
+    if operators is not None:
+        from vllm_omni.diffusion.models.ltx2.ops import qk_rms_norm, residual_adaln, swiglu
+
+        assert operators.qk_norm_rope is qk_rms_norm.try_qk_rms_norm_scale_rope_3d_exact
+        assert operators.residual_norm is residual_adaln.try_residual_rms_norm_modulate_exact
+        assert operators.residual_add is residual_adaln.try_residual_add3_exact
+        assert operators.swiglu is swiglu.try_swiglu_tiled_exact
+
+    tensor = _TensorProbe(is_cuda=True, device=device)
+    with torch.inference_mode():
+        assert ltx2_platform.is_ltx2_fusion_eligible(tensor) is fusion_expected
+        assert ltx2_platform.is_ltx2_fna_eligible(tensor) is fna_expected
+
+
+@pytest.mark.parametrize(
+    ("device", "is_cuda", "available"),
+    [("cpu", True, True), ("cuda:0", False, True), ("cuda:0", True, False)],
+)
+def test_resolver_rejects_unsupported_platform_before_capability_lookup(monkeypatch, device, is_cuda, available):
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: is_cuda)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: available)
+
+    def unexpected_lookup(_index):
+        raise AssertionError("Unsupported platforms must not query CUDA capability")
+
+    monkeypatch.setattr(ltx2_platform, "_cuda_compute_capability", unexpected_lookup)
+    assert ltx2_platform.resolve_ltx2_vae_operators(torch.device(device)) is None
+
+
+@pytest.mark.parametrize(("compiling", "grad_enabled"), [(True, False), (False, True)])
+def test_runtime_guards_bypass_operator_resolution(monkeypatch, compiling, grad_enabled):
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: compiling)
+    monkeypatch.setattr(torch, "is_grad_enabled", lambda: grad_enabled)
+
+    def unexpected_resolve(_device):
+        raise AssertionError("Ineligible runtime must not select eager operators")
+
+    monkeypatch.setattr(ltx2_platform, "resolve_ltx2_vae_operators", unexpected_resolve)
+    tensor = _TensorProbe(is_cuda=True, device=torch.device("cuda:0"))
+    assert not ltx2_platform.is_ltx2_fusion_eligible(tensor)
+    assert not ltx2_platform.is_ltx2_fna_eligible(tensor)
+
+
+def test_device_capability_is_cached(monkeypatch):
+    calls = 0
+    monkeypatch.setattr(ltx2_platform, "HAS_TRITON", True)
+
+    def get_device_capability(device_id=0):
+        nonlocal calls
+        calls += 1
+        return DeviceCapability(major=9, minor=0)
+
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: True)
+    monkeypatch.setattr(ltx2_platform.current_omni_platform, "get_device_capability", get_device_capability)
+    assert ltx2_platform._cuda_compute_capability(0) == 90
+    assert ltx2_platform._cuda_compute_capability(0) == 90
+    assert calls == 1

tests/diffusion/models/ltx2/test_ltx2_vae.py

@@ -3,6 +3,8 @@
 
 """Unit tests for LTX video VAE tiling and distributed decode behavior."""
 
+from collections.abc import Callable
+from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any
 
@@ -35,6 +37,22 @@ def test_ltx_vocoder_deterministic_context(monkeypatch, initial_deterministic, f
     assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings
 
 
+@dataclass(frozen=True)
+class _OperatorSetProbe:
+    fna: Callable[..., torch.Tensor] | None
+
+
+@dataclass
+class _AttentionProbe:
+    kernel_size: tuple[int, int, int]
+    head_dim: int
+    to_out: tuple[torch.nn.Module, ...]
+    projected: torch.Tensor
+
+    def project_qkv(self, _hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self.projected, self.projected, self.projected
+
+
 def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
     import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
 
@@ -61,9 +79,164 @@ def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
 
 
 class TestLTXDiffusionDecoder:
+    def test_diffusion_decoder_reuses_diffusers_with_scoped_overrides(self, monkeypatch):
+        from diffusers.models.autoencoders.ltx2_diffusion_decoder import (
+            LTX2VideoDiffusionDecoderModel as DiffusersModel,
+        )
+
+        from vllm_omni.diffusion.models.ltx2.vae.decoder import (
+            LTX2VideoDiffusionDecoderModel,
+            LTX2VideoVaeDiffusionNABlock,
+            LTX2VideoVaeNeighborhoodAttention,
+            LTX2VideoVaeSwiGLU,
+        )
+
+        original_init = DiffusersModel.__init__
+        original_parameters = {}
+
+        def capture_parameters(self, *args, **kwargs):
+            original_init(self, *args, **kwargs)
+            original_parameters.update(self.named_parameters())
+
+        with torch.device("meta"):
+            reference = DiffusersModel()
+            monkeypatch.setattr(DiffusersModel, "__init__", capture_parameters)
+            model = LTX2VideoDiffusionDecoderModel()
+        assert isinstance(model, DiffusersModel)
+        assert all(parameter is original_parameters[name] for name, parameter in model.named_parameters())
+        assert {name: tensor.shape for name, tensor in model.state_dict().items()} == {
+            name: tensor.shape for name, tensor in reference.state_dict().items()
+        }
+
+        modules = tuple(model.decoder.modules())
+        expected_blocks = sum(model.config.decoder_stage_depths)
+        assert sum(isinstance(module, LTX2VideoVaeNeighborhoodAttention) for module in modules) == expected_blocks
+        assert sum(isinstance(module, LTX2VideoVaeSwiGLU) for module in modules) == expected_blocks
+        assert (
+            sum(isinstance(module, LTX2VideoVaeDiffusionNABlock) for module in modules)
+            == model.config.decoder_stage_depths[-1]
+        )
+
+        assert type(model.decoder).__module__.startswith("vllm_omni.")
+        assert model.tiled_decode.__func__ is DiffusersModel.tiled_decode
+        assert type(reference.decoder) is type(model.decoder).__bases__[0]
+
+    def test_natten_processor_routes_stage5_through_tilelang_fna(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.vae import decoder as diffvae_modules
+
+        processor = object.__new__(diffvae_modules.LTX2VideoVaeNeighborhoodNattenProcessor)
+        processor.backend = None
+        fallback_calls = []
+        processor._na3d = lambda *args, **kwargs: fallback_calls.append((args, kwargs))
+
+        hidden_states = torch.randn(1, 11, 11, 11, 256, dtype=torch.bfloat16)
+        projected = hidden_states.view(1, 11, 11, 11, 4, 64)
+
+        class FakeAttention:
+            kernel_size = (11, 11, 11)
+            head_dim = 64
+            to_out = (torch.nn.Identity(),)
+
+            @staticmethod
+            def project_qkv(_hidden_states):
+                return projected, projected, projected
+
+        def token_permute(value, _tile, *, flip_tiled_dims):
+            assert flip_tiled_dims
+            return value, value.shape, None
+
+        def token_unpermute(value, *, token_layout_shape, tile_shape, flip_tiled_dims):
+            assert token_layout_shape == projected.shape
+            assert tile_shape == (4, 4, 4)
+            assert flip_tiled_dims
+            return value
+
+        fast_path_calls = []
+
+        def fast_path(query, key, value, *, shape, window):
+            fast_path_calls.append((query, key, value, shape, window))
+            return query
+
+        monkeypatch.setattr(diffvae_modules, "is_ltx2_fna_eligible", lambda _tensor: True)
+        monkeypatch.setattr(
+            diffvae_modules, "resolve_ltx2_vae_operators", lambda _device: _OperatorSetProbe(fna=fast_path)
+        )
+        monkeypatch.setattr(
+            diffvae_modules,
+            "_load_natten_token_permutation",
+            lambda: (token_permute, token_unpermute),
+        )
+
+        output = processor(FakeAttention(), hidden_states)
+
+        assert torch.equal(output, hidden_states)
+        assert len(fast_path_calls) == 1
+        assert fast_path_calls[0][3:] == ((11, 11, 11), (11, 11, 11))
+        assert not fallback_calls
+
+    @pytest.mark.parametrize("reason", ["device", "head_dim", "channels", "dtype", "tilelang_missing", "kernel"])
+    def test_natten_processor_keeps_natten_when_fast_path_does_not_apply(self, monkeypatch, reason):
+        from vllm_omni.diffusion.models.ltx2.vae import decoder as ops
+
+        processor = object.__new__(ops.LTX2VideoVaeNeighborhoodNattenProcessor)
+        processor.backend = None
+        channels = 128 if reason == "channels" else 256
+        head_dim = 32 if reason == "head_dim" else 64
+        x = torch.zeros(1, 11, 11, 11, channels, dtype=torch.float32 if reason == "dtype" else torch.bfloat16)
+        projected = x.view(1, 11, 11, 11, channels // head_dim, head_dim)
+        attn = _AttentionProbe(
+            kernel_size=(3, 5, 5) if reason == "kernel" else (11, 11, 11),
+            head_dim=head_dim,
+            to_out=(torch.nn.Identity(),),
+            projected=projected,
+        )
+        calls = []
+
+        def natten(q, k, v, **kwargs):
+            calls.append(kwargs)
+            return q
+
+        def unexpected_native_load(*_args, **_kwargs):
+            raise AssertionError("Native FNA must not be loaded for an unsupported configuration")
+
+        processor._na3d = natten
+        monkeypatch.setattr(ops, "is_ltx2_fna_eligible", lambda _tensor: reason != "device")
+        monkeypatch.setattr(
+            ops,
+            "resolve_ltx2_vae_operators",
+            lambda _device: _OperatorSetProbe(fna=None if reason == "tilelang_missing" else unexpected_native_load),
+        )
+        assert torch.equal(processor(attn, x), x)
+        assert calls == [{"kernel_size": attn.kernel_size, "scale": 1.0, "backend": None}]
+
+    def test_diffusion_decoder_noneligible_block_uses_diffusers_forward(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.vae import decoder as diffvae_modules
+
+        block = object.__new__(diffvae_modules.LTX2VideoVaeDiffusionNABlock)
+        torch.nn.Module.__init__(block)
+        expected = torch.randn(1, 1, 1, 1, 1)
+        calls = []
+
+        def upstream_forward(_self, hidden_states, latent_context, modulation, block_mask):
+            assert block_mask is None
+            calls.append((hidden_states, latent_context, modulation))
+            return expected
+
+        monkeypatch.setattr(
+            diffvae_modules.diffusers_decoder.LTX2VideoVaeDiffusionNABlock,
+            "forward",
+            upstream_forward,
+        )
+        hidden_states = torch.randn(1, 1, 1, 1, 1)
+        latent_context = torch.randn_like(hidden_states)
+        modulation = (torch.randn_like(hidden_states),)
+
+        assert block.forward(hidden_states, latent_context, modulation) is expected
+        assert calls == [(hidden_states, latent_context, modulation)]
+
     @pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"])
     def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+        from vllm_omni.diffusion.models.ltx2.vae.distributed import (
             DistributedLTX2VideoDiffusionDecoderModel,
         )
 
@@ -75,7 +248,7 @@ class TestLTXDiffusionDecoder:
             model.set_parallel_size(2, mode=mode)
 
     def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+        from vllm_omni.diffusion.models.ltx2.vae.distributed import (
             DistributedLTX2VideoDiffusionDecoderModel,
         )
 
@@ -91,7 +264,7 @@ class TestLTXDiffusionDecoder:
         assert calls == [(4, "tile")]
 
     def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+        from vllm_omni.diffusion.models.ltx2.vae.decoder import (
             LTX2VideoDiffusionDecoder3d,
             LTX2VideoDiffusionDecoderModel,
         )
@@ -166,7 +339,7 @@ class TestLTXDiffusionDecoder:
             _ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5")
 
     def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+        from vllm_omni.diffusion.models.ltx2.vae.decoder import (
             convert_ltx25_native_diffusion_decoder_state_dict,
         )
 
@@ -213,7 +386,7 @@ class TestLTXDiffusionDecoder:
         assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted)
 
     def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+        from vllm_omni.diffusion.models.ltx2.vae.decoder import (
             convert_ltx25_native_diffusion_decoder_state_dict,
         )
 
@@ -411,8 +584,8 @@ class TestLTXDiffusionDecoder:
         assert output.output[1].numel() == 0
 
     def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch):
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+        from vllm_omni.diffusion.models.ltx2.vae.decoder import LTX2VideoDiffusionDecoderModel
+        from vllm_omni.diffusion.models.ltx2.vae.distributed import (
             DistributedLTX2VideoDiffusionDecoderModel,
         )
 
@@ -447,7 +620,7 @@ class TestLTXDiffusionDecoder:
     def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch):
         from diffusers.utils.torch_utils import randn_tensor
 
-        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+        from vllm_omni.diffusion.models.ltx2.vae.distributed import (
             DistributedLTX2VideoDiffusionDecoderModel,
         )
 

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -15,7 +15,6 @@ from typing import TYPE_CHECKING, Any
 
 import torch
 from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
-from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeNeighborhoodNattenProcessor
 from diffusers.pipelines.ltx2 import LTX2TextConnectors
 from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
 from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
@@ -36,12 +35,6 @@ from vllm_omni.transformers_utils.repo_utils import hf_api
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
-from .ltx2_diffusion_decoder import (
-    LTX25_NATIVE_ARTIFACT_REVISION,
-    LTX25_NATIVE_DIFFUSION_DECODER_FILENAME,
-    LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID,
-)
-from .ltx2_diffusion_decoder_distributed import DistributedLTX2VideoDiffusionDecoderModel
 from .ltx2_request import LTXCheckpointKind, validate_ltx_checkpoint
 from .ltx2_transformer import (
     LTX2VideoTransformer3DModel,
@@ -49,6 +42,13 @@ from .ltx2_transformer import (
     apply_split_rotary_emb,
     to_ltx_padding_mask,
 )
+from .vae.decoder import (
+    LTX25_NATIVE_ARTIFACT_REVISION,
+    LTX25_NATIVE_DIFFUSION_DECODER_FILENAME,
+    LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID,
+    LTX2VideoVaeNeighborhoodNattenProcessor,
+)
+from .vae.distributed import DistributedLTX2VideoDiffusionDecoderModel
 
 try:
     from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE

vllm_omni/diffusion/models/ltx2/ops/__init__.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Shared operator selection for LTX-2 eager kernels and their tests."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import cache
+
+import torch
+from vllm.triton_utils import HAS_TRITON
+from vllm.utils.import_utils import has_tilelang
+
+from vllm_omni.platforms import current_omni_platform
+
+# Pointwise fusions are launch/memory-traffic reductions whose first execution
+# for every runtime key is checked bit-for-bit; failures use the ordinary model
+# math. Both paths remain SM90-only until other architectures are validated.
+# The FNA schedule is performance-tuned for SM90 and has no silent fallback.
+_CUDA_COMPUTE_CAPABILITIES = frozenset({90})
+
+
+@dataclass(frozen=True)
+class LTX2VAEOperatorSet:
+    qk_norm_rope: Callable[..., tuple[torch.Tensor, torch.Tensor] | None]
+    residual_norm: Callable[..., torch.Tensor | None]
+    residual_add: Callable[..., torch.Tensor | None]
+    swiglu: Callable[..., torch.Tensor | None]
+    fna: Callable[..., torch.Tensor] | None
+
+
+def _fna3d_tilelang(*args, **kwargs):
+    # Selecting an operator set must not import the optional compiler backend.
+    from .fna import fna3d_tilelang
+
+    return fna3d_tilelang(*args, **kwargs)
+
+
+@cache
+def _operator_set(with_fna: bool) -> LTX2VAEOperatorSet:
+    from .qk_rms_norm import try_qk_rms_norm_scale_rope_3d_exact
+    from .residual_adaln import try_residual_add3_exact, try_residual_rms_norm_modulate_exact
+    from .swiglu import try_swiglu_tiled_exact
+
+    return LTX2VAEOperatorSet(
+        qk_norm_rope=try_qk_rms_norm_scale_rope_3d_exact,
+        residual_norm=try_residual_rms_norm_modulate_exact,
+        residual_add=try_residual_add3_exact,
+        swiglu=try_swiglu_tiled_exact,
+        fna=_fna3d_tilelang if with_fna else None,
+    )
+
+
+@cache
+def _cuda_compute_capability(device_index: int) -> int | None:
+    if not HAS_TRITON or not current_omni_platform.is_cuda() or not current_omni_platform.is_available():
+        return None
+    capability = current_omni_platform.get_device_capability(device_id=device_index)
+    return capability.to_int() if capability is not None else None
+
+
+def resolve_ltx2_vae_operators(device: torch.device) -> LTX2VAEOperatorSet | None:
+    """Select the production operators, or return None for unsupported devices."""
+    if device.type != "cuda" or not current_omni_platform.is_cuda() or not current_omni_platform.is_available():
+        return None
+    device_index = device.index if device.index is not None else torch.accelerator.current_device_index()
+    if _cuda_compute_capability(int(device_index)) not in _CUDA_COMPUTE_CAPABILITIES:
+        return None
+    return _operator_set(has_tilelang())
+
+
+def _eligible_operators(tensor: torch.Tensor) -> LTX2VAEOperatorSet | None:
+    if not tensor.is_cuda or torch.compiler.is_compiling() or torch.is_grad_enabled():
+        return None
+    return resolve_ltx2_vae_operators(tensor.device)
+
+
+def is_ltx2_fusion_eligible(tensor: torch.Tensor) -> bool:
+    """Allow self-verifying pointwise fusions on validated SM90 devices."""
+
+    return _eligible_operators(tensor) is not None
+
+
+def is_ltx2_fna_eligible(tensor: torch.Tensor) -> bool:
+    """Keep the strict, Hopper-tuned TileLang FNA schedule on SM90."""
+
+    operators = _eligible_operators(tensor)
+    return operators is not None and operators.fna is not None
+
+
+__all__ = ["LTX2VAEOperatorSet", "resolve_ltx2_vae_operators", "is_ltx2_fna_eligible", "is_ltx2_fusion_eligible"]

vllm_omni/diffusion/models/ltx2/ops/_numerics.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Numerical primitives shared by exact LTX-2 Triton kernels."""
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def round_bf16_to_fp32(value):
+    """RNE-round FP32 to BF16 precision while retaining an FP32 register."""
+
+    bits = value.to(tl.int32, bitcast=True)
+    rounding_bias = 0x7FFF + ((bits >> 16) & 1)
+    rounded_bits = (bits + rounding_bias) & -65536
+    return rounded_bits.to(tl.float32, bitcast=True)
+
+
+@triton.jit
+def add_rn_f32(x, y):
+    return tl.inline_asm_elementwise(
+        asm="add.rn.f32 $0, $1, $2;",
+        constraints="=f,f,f",
+        args=[x, y],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def mul_rn_f32(x, y):
+    return tl.inline_asm_elementwise(
+        asm="mul.rn.f32 $0, $1, $2;",
+        constraints="=f,f,f",
+        args=[x, y],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def fma_rn_f32(x, y, accumulator):
+    return tl.inline_asm_elementwise(
+        asm="fma.rn.f32 $0, $1, $2, $3;",
+        constraints="=f,f,f,f",
+        args=[x, y, accumulator],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def rsqrt_approx_f32(x):
+    return tl.inline_asm_elementwise(
+        asm="rsqrt.approx.f32 $0, $1;",
+        constraints="=f,f",
+        args=[x],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def shfl_down_f32(value, delta: tl.constexpr):
+    return tl.inline_asm_elementwise(
+        asm="shfl.sync.down.b32 $0, $1, $2, 0x1f, 0xffffffff;",
+        constraints="=f,f,n",
+        args=[value, delta],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+__all__ = [
+    "add_rn_f32",
+    "fma_rn_f32",
+    "mul_rn_f32",
+    "round_bf16_to_fp32",
+    "rsqrt_approx_f32",
+    "shfl_down_f32",
+]

vllm_omni/diffusion/models/ltx2/ops/fna.py

@@ -0,0 +1,525 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""TileLang FNA for the LTX-2.5 DiffVAE stage-5 schedule.
+
+This intentionally reuses NATTEN's token permutation.  It only replaces the
+middle BF16 attention kernel for Q tiles (4, 4, 4), KV tiles (4, 4, 8),
+head_dim=64, stride=dilation=1, and scale=1.0.
+"""
+
+import functools
+import math
+from collections import OrderedDict
+from dataclasses import dataclass
+
+import torch
+from vllm.tilelang_utils import T, tilelang
+
+Q_TILE = (4, 4, 4)
+KV_TILE = (4, 4, 8)
+Q_TOKENS = math.prod(Q_TILE)
+KV_TOKENS = math.prod(KV_TILE)
+HEAD_DIM = 64
+HEADS = 4
+LOG2_E = 1.4426950408889634
+KV_BUCKET_LIMITS = (32, 48, 60, 75)
+PATTERN_TABLE_SIZE = 16
+KV_TILE_ID_BITS = 15
+SHAPE_CACHE_SIZE = 16
+
+
+@dataclass(frozen=True)
+class Problem:
+    frames: int
+    height: int
+    width: int
+    window: tuple[int, int, int]
+
+    def __post_init__(self) -> None:
+        if min(self.frames, self.height, self.width) <= 0:
+            raise ValueError(f"invalid FNA problem: {self}")
+        if len(self.window) != 3 or any(x <= 0 for x in self.window):
+            raise ValueError(f"invalid FNA window: {self.window}")
+        if any(w > n for w, n in zip(self.window, self.shape)):
+            raise ValueError(f"window {self.window} exceeds shape {self.shape}")
+        if math.prod(self.kv_grid) > 1 << KV_TILE_ID_BITS:
+            raise ValueError(f"KV tile grid exceeds {KV_TILE_ID_BITS}-bit metadata encoding: {self.kv_grid}")
+
+    @property
+    def shape(self) -> tuple[int, int, int]:
+        return self.frames, self.height, self.width
+
+    @property
+    def q_padded(self) -> tuple[int, int, int]:
+        return tuple(_round_up(n, tile) for n, tile in zip(self.shape, Q_TILE))
+
+    @property
+    def kv_padded(self) -> tuple[int, int, int]:
+        return tuple(_round_up(n, tile) for n, tile in zip(self.shape, KV_TILE))
+
+    @property
+    def q_grid(self) -> tuple[int, int, int]:
+        return tuple(n // tile for n, tile in zip(self.q_padded, Q_TILE))
+
+    @property
+    def kv_grid(self) -> tuple[int, int, int]:
+        return tuple(n // tile for n, tile in zip(self.kv_padded, KV_TILE))
+
+    @property
+    def q_length(self) -> int:
+        return math.prod(self.q_padded)
+
+    @property
+    def kv_length(self) -> int:
+        return math.prod(self.kv_padded)
+
+
+def _round_up(value: int, multiple: int) -> int:
+    return (value + multiple - 1) // multiple * multiple
+
+
+def _window_start(index: int, window: int, length: int) -> int:
+    return min(max(index - window // 2, 0), length - window)
+
+
+@functools.lru_cache(maxsize=32)
+def _dimension_pattern_table(
+    length: int,
+    q_tile_size: int,
+    kv_tile_size: int,
+    window: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Map (Q tile, KV tile) pairs to a small table of packed 1D masks."""
+    q_tiles = _round_up(length, q_tile_size) // q_tile_size
+    kv_tiles = _round_up(length, kv_tile_size) // kv_tile_size
+    zero = (0,) * q_tile_size
+    patterns: list[tuple[int, ...]] = [zero]
+    pattern_ids = {zero: 0}
+    ids = torch.zeros((q_tiles, kv_tiles), dtype=torch.int64)
+
+    for q_tile in range(q_tiles):
+        for kv_tile in range(kv_tiles):
+            rows: list[int] = []
+            for q_inner in range(q_tile_size):
+                q = q_tile * q_tile_size + q_inner
+                bits = 0
+                if q < length:
+                    start = _window_start(q, window, length)
+                    for kv_inner in range(kv_tile_size):
+                        kv = kv_tile * kv_tile_size + kv_inner
+                        if kv < length and start <= kv < start + window:
+                            bits |= 1 << kv_inner
+                rows.append(bits)
+            signature = tuple(rows)
+            pattern_id = pattern_ids.get(signature)
+            if pattern_id is None:
+                pattern_id = len(patterns)
+                pattern_ids[signature] = pattern_id
+                patterns.append(signature)
+            ids[q_tile, kv_tile] = pattern_id
+
+    if len(patterns) > PATTERN_TABLE_SIZE:
+        raise AssertionError(f"4-bit pattern id overflow for length={length}: {len(patterns)}")
+    packed = []
+    for pattern in patterns:
+        value = sum(bits << (q * kv_tile_size) for q, bits in enumerate(pattern))
+        # Store uint32 bit patterns in an int32 tensor.  Arithmetic right shift
+        # followed by `& 1` still extracts the desired bit when bit 31 is set.
+        packed.append(value if value < (1 << 31) else value - (1 << 32))
+    # The four-bit pattern ids address at most 16 entries.  Pad to that fixed
+    # size so T/H/W-specific table lengths cannot specialize the GPU kernel.
+    packed.extend([0] * (PATTERN_TABLE_SIZE - len(packed)))
+    return ids, torch.tensor(packed, dtype=torch.int32)
+
+
+@functools.lru_cache(maxsize=SHAPE_CACHE_SIZE)
+def build_kv_tile_map(problem: Problem) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return packed KV tile ids and valid counts for each permuted Q tile."""
+    q_grid_t, q_grid_h, q_grid_w = problem.q_grid
+    kv_grid_t, kv_grid_h, _ = problem.kv_grid
+    rows: list[list[int]] = []
+
+    # Token permutation flattens outer tile modes as (W, H, T), with T
+    # changing fastest.  Keep KV tiles in the same increasing linear order.
+    for q_w_tile in range(q_grid_w):
+        for q_h_tile in range(q_grid_h):
+            for q_t_tile in range(q_grid_t):
+                ranges: list[range] = []
+                for q_tile_idx, q_tile_size, kv_tile_size, window, length in zip(
+                    (q_t_tile, q_h_tile, q_w_tile),
+                    Q_TILE,
+                    KV_TILE,
+                    problem.window,
+                    problem.shape,
+                ):
+                    first_q = q_tile_idx * q_tile_size
+                    last_q = min(first_q + q_tile_size, length) - 1
+                    first_k = _window_start(first_q, window, length)
+                    last_k = _window_start(last_q, window, length) + window - 1
+                    ranges.append(range(first_k // kv_tile_size, last_k // kv_tile_size + 1))
+
+                t_tiles, h_tiles, w_tiles = ranges
+                row = [
+                    (kv_w * kv_grid_h + kv_h) * kv_grid_t + kv_t
+                    for kv_w in w_tiles
+                    for kv_h in h_tiles
+                    for kv_t in t_tiles
+                ]
+                if not row:
+                    raise AssertionError("every Q tile must overlap at least one KV tile")
+                rows.append(row)
+
+    max_tiles = max(map(len, rows))
+    # Duplicate a valid tile in padding slots.  The kernel masks those slots
+    # with `counts`; avoiding -1 keeps every indirect global load in bounds.
+    indices = torch.tensor(
+        [row + [row[0]] * (max_tiles - len(row)) for row in rows],
+        dtype=torch.int32,
+    )
+    counts = torch.tensor([len(row) for row in rows], dtype=torch.int32)
+    return indices, counts
+
+
+@functools.lru_cache(maxsize=SHAPE_CACHE_SIZE)
+def build_kv_tile_metadata(
+    problem: Problem,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Pack KV tile id and three separable mask pattern ids into int32."""
+    indices, counts = build_kv_tile_map(problem)
+    q_grid_t, q_grid_h, _ = problem.q_grid
+    kv_grid_t, kv_grid_h, _ = problem.kv_grid
+    t_ids, t_patterns = _dimension_pattern_table(problem.frames, Q_TILE[0], KV_TILE[0], problem.window[0])
+    h_ids, h_patterns = _dimension_pattern_table(problem.height, Q_TILE[1], KV_TILE[1], problem.window[1])
+    w_ids, w_patterns = _dimension_pattern_table(problem.width, Q_TILE[2], KV_TILE[2], problem.window[2])
+
+    q_blocks = torch.arange(indices.shape[0], dtype=torch.int64)[:, None]
+    q_t = q_blocks % q_grid_t
+    q_h = (q_blocks // q_grid_t) % q_grid_h
+    q_w = q_blocks // (q_grid_t * q_grid_h)
+    kv = indices.to(torch.int64)
+    kv_t = kv % kv_grid_t
+    kv_h = (kv // kv_grid_t) % kv_grid_h
+    kv_w = kv // (kv_grid_t * kv_grid_h)
+    t_pattern_id = t_ids[q_t, kv_t]
+    h_pattern_id = h_ids[q_h, kv_h]
+    w_pattern_id = w_ids[q_w, kv_w]
+    metadata = kv | (t_pattern_id << 15) | (h_pattern_id << 19) | (w_pattern_id << 23)
+    active = torch.arange(indices.shape[1])[None, :] < counts[:, None]
+    metadata = torch.where(active, metadata, 0).to(torch.int32)
+    return metadata, t_patterns, h_patterns, w_patterns
+
+
+@functools.lru_cache(maxsize=SHAPE_CACHE_SIZE)
+def build_kv_tile_buckets(
+    problem: Problem,
+) -> tuple[tuple[torch.Tensor, torch.Tensor, int], ...]:
+    """Group Q tiles by an upper bound on their useful KV tile count."""
+    metadata, *_ = build_kv_tile_metadata(problem)
+    _, counts = build_kv_tile_map(problem)
+    maximum = int(counts.max())
+    if maximum > KV_BUCKET_LIMITS[-1]:
+        raise ValueError(
+            f"FNA neighborhood needs {maximum} KV tiles, exceeding the supported ceiling {KV_BUCKET_LIMITS[-1]}"
+        )
+    buckets = []
+    lower = 0
+    for limit in KV_BUCKET_LIMITS:
+        q_tile_ids = torch.nonzero((counts > lower) & (counts <= limit), as_tuple=False).flatten()
+        if q_tile_ids.numel():
+            # Small grids can have fewer columns than the first bucket limit.
+            # Zero metadata encodes an all-masked tile, preserving the fixed ABI.
+            bucket_metadata = torch.zeros((q_tile_ids.numel(), limit), dtype=torch.int32)
+            columns = min(limit, metadata.shape[1])
+            bucket_metadata[:, :columns] = metadata[q_tile_ids, :columns]
+            buckets.append(
+                (
+                    q_tile_ids.to(torch.int32),
+                    bucket_metadata,
+                    limit,
+                )
+            )
+        lower = limit
+        if limit >= maximum:
+            break
+    return tuple(buckets)
+
+
+@functools.lru_cache(maxsize=32)
+def compile_kernel(
+    max_kv_tiles: int,
+    *,
+    target: str = "auto",
+):
+    """Compile one T/H/W-dynamic kernel for a fixed KV-count ceiling."""
+    if max_kv_tiles not in KV_BUCKET_LIMITS:
+        raise ValueError(f"unsupported KV bucket ceiling: {max_kv_tiles}")
+
+    q_length = T.dynamic("q_length")
+    kv_length = T.dynamic("kv_length")
+    q_tile_count = T.dynamic("q_tile_count")
+
+    # Batch is fixed to one by the public guard.  Excluding it from the
+    # kernel ABI keeps every physical stride static even when token lengths
+    # vary, preserving vectorized global/shared-memory copies.
+    q_shape = [q_length, HEADS, HEAD_DIM]
+    kv_shape = [kv_length, HEADS, HEAD_DIM]
+    q_tile_id_shape = [q_tile_count]
+    map_shape = [q_tile_count, max_kv_tiles]
+    pattern_shape = [PATTERN_TABLE_SIZE]
+    dtype = T.bfloat16
+    accum_dtype = T.float32
+
+    @T.prim_func
+    def main(
+        query: T.Tensor(q_shape, dtype),
+        key: T.Tensor(kv_shape, dtype),
+        value: T.Tensor(kv_shape, dtype),
+        q_tile_ids: T.Tensor(q_tile_id_shape, T.int32),
+        kv_tile_metadata: T.Tensor(map_shape, T.int32),
+        t_mask_patterns: T.Tensor(pattern_shape, T.int32),
+        h_mask_patterns: T.Tensor(pattern_shape, T.int32),
+        w_mask_patterns: T.Tensor(pattern_shape, T.int32),
+        output: T.Tensor(q_shape, dtype),
+    ):
+        with T.Kernel(q_tile_count, HEADS, threads=128) as (bx, by):
+            q_shared = T.alloc_shared([Q_TOKENS, HEAD_DIM], dtype)
+            k_shared = T.alloc_shared([KV_TOKENS, HEAD_DIM], dtype)
+            v_shared = T.alloc_shared([KV_TOKENS, HEAD_DIM], dtype)
+            o_shared = T.alloc_shared([Q_TOKENS, HEAD_DIM], dtype)
+            scores = T.alloc_fragment([Q_TOKENS, KV_TOKENS], accum_dtype)
+            probs = T.alloc_fragment([Q_TOKENS, KV_TOKENS], dtype)
+            acc_o = T.alloc_fragment([Q_TOKENS, HEAD_DIM], accum_dtype)
+            row_max = T.alloc_fragment([Q_TOKENS], accum_dtype)
+            row_max_prev = T.alloc_fragment([Q_TOKENS], accum_dtype)
+            row_scale = T.alloc_fragment([Q_TOKENS], accum_dtype)
+            row_sum = T.alloc_fragment([Q_TOKENS], accum_dtype)
+            row_denom = T.alloc_fragment([Q_TOKENS], accum_dtype)
+
+            q_block = q_tile_ids[bx]
+            # CPU metadata only emits padded Q/KV tiles that are fully backed
+            # by the permuted tensors.  Expose those invariants so dynamic
+            # lengths do not add per-copy boundary predicates.
+            T.assume(q_block >= 0)
+            T.assume((q_block + 1) * Q_TOKENS <= q_length)
+            T.copy(
+                query[
+                    q_block * Q_TOKENS : (q_block + 1) * Q_TOKENS,
+                    by,
+                    :,
+                ],
+                q_shared,
+            )
+            T.fill(acc_o, 0)
+            T.fill(row_denom, 0)
+            T.fill(row_max, -T.infinity(accum_dtype))
+
+            for block_iter in T.Pipelined(max_kv_tiles, num_stages=1):
+                metadata = kv_tile_metadata[bx, block_iter]
+                kv_block = metadata & 0x7FFF
+                T.assume((kv_block + 1) * KV_TOKENS <= kv_length)
+                t_mask = t_mask_patterns[(metadata >> 15) & 0xF]
+                h_mask = h_mask_patterns[(metadata >> 19) & 0xF]
+                w_mask = w_mask_patterns[(metadata >> 23) & 0xF]
+                T.copy(
+                    key[
+                        kv_block * KV_TOKENS : (kv_block + 1) * KV_TOKENS,
+                        by,
+                        :,
+                    ],
+                    k_shared,
+                    prefer_instruction="cp_async",
+                )
+
+                for i, j in T.Parallel(Q_TOKENS, KV_TOKENS):
+                    q_t = i % Q_TILE[0]
+                    q_h = (i // Q_TILE[0]) % Q_TILE[1]
+                    q_w = i // (Q_TILE[0] * Q_TILE[1])
+                    k_t = j % KV_TILE[0]
+                    k_h = (j // KV_TILE[0]) % KV_TILE[1]
+                    k_w = j // (KV_TILE[0] * KV_TILE[1])
+                    valid = (
+                        (((t_mask >> (q_t * KV_TILE[0] + k_t)) & 1) != 0)
+                        & (((h_mask >> (q_h * KV_TILE[1] + k_h)) & 1) != 0)
+                        & (((w_mask >> (q_w * KV_TILE[2] + k_w)) & 1) != 0)
+                    )
+                    scores[i, j] = T.if_then_else(valid, 0, -T.infinity(accum_dtype))
+
+                T.gemm(
+                    q_shared,
+                    k_shared,
+                    scores,
+                    transpose_B=True,
+                    policy=T.GemmWarpPolicy.FullRow,
+                )
+                T.copy(row_max, row_max_prev)
+                T.fill(row_max, -T.infinity(accum_dtype))
+                T.reduce_max(scores, row_max, dim=1, clear=False)
+                # A KV tile is selected for the union of all 64 queries in
+                # this Q tile.  It can therefore be fully masked for one
+                # individual query.  Preserve that row's previous online
+                # softmax state instead of evaluating (-inf) - (-inf).
+                T.copy(row_max, row_sum)
+                for i in T.Parallel(Q_TOKENS):
+                    has_score = row_sum[i] > -T.infinity(accum_dtype)
+                    merged_max = T.max(row_sum[i], row_max_prev[i])
+                    row_max[i] = T.if_then_else(has_score, merged_max, row_max_prev[i])
+                    row_scale[i] = T.if_then_else(
+                        has_score,
+                        T.exp2((row_max_prev[i] - merged_max) * LOG2_E),
+                        1,
+                    )
+                    # Match Hopper NATTEN's two separately rounded FP32
+                    # products before exp2.
+                    row_sum[i] = T.if_then_else(
+                        has_score,
+                        merged_max * LOG2_E,
+                        -T.infinity(accum_dtype),
+                    )
+                for i, j in T.Parallel(Q_TOKENS, KV_TOKENS):
+                    scores[i, j] = T.if_then_else(
+                        row_sum[i] > -T.infinity(accum_dtype),
+                        T.exp2(scores[i, j] * LOG2_E - row_sum[i]),
+                        0,
+                    )
+                T.reduce_sum(scores, row_sum, dim=1)
+                for i in T.Parallel(Q_TOKENS):
+                    row_denom[i] = row_denom[i] * row_scale[i] + row_sum[i]
+                T.copy(scores, probs)
+                for i, j in T.Parallel(Q_TOKENS, HEAD_DIM):
+                    acc_o[i, j] *= row_scale[i]
+
+                T.copy(
+                    value[
+                        kv_block * KV_TOKENS : (kv_block + 1) * KV_TOKENS,
+                        by,
+                        :,
+                    ],
+                    v_shared,
+                    prefer_instruction="cp_async",
+                )
+                T.gemm(probs, v_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
+
+            # Match NATTEN's epilogue: compute one reciprocal per query row,
+            # then multiply every output accumulator by that shared alpha.
+            for i in T.Parallel(Q_TOKENS):
+                row_scale[i] = T.if_then_else(
+                    row_denom[i] > 0,
+                    T.call_extern("float32", "__frcp_rn", row_denom[i]),
+                    0,
+                )
+            for i, j in T.Parallel(Q_TOKENS, HEAD_DIM):
+                acc_o[i, j] *= row_scale[i]
+            T.copy(acc_o, o_shared)
+            T.copy(
+                o_shared,
+                output[
+                    q_block * Q_TOKENS : (q_block + 1) * Q_TOKENS,
+                    by,
+                    :,
+                ],
+            )
+
+    return tilelang.compile(
+        main,
+        out_idx=None,
+        target="auto" if target == "auto" else {"kind": "cuda", "arch": target},
+        pass_configs={
+            tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
+            # Keeping all work in the 128-thread consumer warpgroup avoids a
+            # separate producer warpgroup.  The generated 240-register
+            # consumer can then admit two CTAs per SM instead of one.
+            tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        },
+    )
+
+
+_DEVICE_MAP_CACHE: OrderedDict[tuple[Problem, torch.device], tuple[object, ...]] = OrderedDict()
+
+
+def _fna3d_tilelang(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    *,
+    shape: tuple[int, int, int],
+    window: tuple[int, int, int],
+) -> torch.Tensor:
+    problem = Problem(*shape, window=window)
+    if query.shape[1] != problem.q_length or key.shape[1] != problem.kv_length:
+        raise ValueError(f"permuted lengths do not match {problem}: Q={query.shape[1]}, KV={key.shape[1]}")
+    cache_key = (problem, query.device)
+    device_inputs = _DEVICE_MAP_CACHE.get(cache_key)
+    if device_inputs is None:
+        _, t_patterns, h_patterns, w_patterns = build_kv_tile_metadata(problem)
+        pattern_tables = tuple(
+            tensor.to(query.device, non_blocking=True) for tensor in (t_patterns, h_patterns, w_patterns)
+        )
+        buckets = tuple(
+            (
+                q_tile_ids.to(query.device, non_blocking=True),
+                metadata.to(query.device, non_blocking=True),
+                limit,
+            )
+            for q_tile_ids, metadata, limit in build_kv_tile_buckets(problem)
+        )
+        device_inputs = (pattern_tables, buckets)
+        _DEVICE_MAP_CACHE[cache_key] = device_inputs
+        if len(_DEVICE_MAP_CACHE) > SHAPE_CACHE_SIZE:
+            _DEVICE_MAP_CACHE.popitem(last=False)
+    else:
+        _DEVICE_MAP_CACHE.move_to_end(cache_key)
+
+    pattern_tables, buckets = device_inputs
+    output = torch.empty_like(query)
+    for q_tile_ids, metadata, limit in buckets:
+        from tilelang.backend.target import determine_target
+
+        target = str(determine_target("auto", return_object=True).attrs["arch"])
+        kernel = compile_kernel(limit, target=target)
+        for batch in range(query.shape[0]):
+            kernel(query[batch], key[batch], value[batch], q_tile_ids, metadata, *pattern_tables, output[batch])
+    return output
+
+
+def fna3d_tilelang(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    *,
+    shape: tuple[int, int, int],
+    window: tuple[int, int, int],
+) -> torch.Tensor:
+    """Run native stage-5 FNA, or raise for unsupported inputs/launch failures.
+
+    Q and KV use the documented token-tiled layout. Batch members share spatial
+    metadata; compilation is cached by target architecture and KV-count bucket.
+    """
+    if not (
+        query.is_cuda
+        and query.dtype is torch.bfloat16
+        and query.ndim == key.ndim == value.ndim == 4
+        and query.shape[0] > 0
+        and query.shape[-2:] == (HEADS, HEAD_DIM)
+        and query.numel() > 0
+        and query.is_contiguous()
+        and key.device == value.device == query.device
+        and key.dtype is value.dtype is query.dtype
+        and key.is_contiguous()
+        and value.is_contiguous()
+        and value.shape == key.shape
+        and key.shape[0] == query.shape[0]
+        and key.shape[-2:] == query.shape[-2:]
+        and len(shape) == 3
+        and tuple(window) == (11, 11, 11)
+    ):
+        raise ValueError(
+            "LTX DiffVAE FNA requires contiguous CUDA BF16 [B, tokens, 4, 64] Q/K/V and an 11x11x11 window."
+        )
+    if torch.is_grad_enabled() and any(x.requires_grad for x in (query, key, value)):
+        raise RuntimeError("LTX DiffVAE FNA is inference-only.")
+    with torch.accelerator.device_index(query.device.index):
+        return _fna3d_tilelang(query, key, value, shape=tuple(shape), window=tuple(window))
+
+
+__all__ = ["fna3d_tilelang"]

vllm_omni/diffusion/models/ltx2/ops/qk_rms_norm.py

@@ -0,0 +1,377 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Bit-exact multi-row Q/K RMSNorm for the LTX diffusion VAE decoder."""
+
+from __future__ import annotations
+
+import logging
+import math
+
+import torch
+import torch.nn.functional as F
+from vllm.triton_utils import tl, triton
+
+from . import is_ltx2_fusion_eligible
+from ._numerics import (
+    add_rn_f32,
+    fma_rn_f32,
+    mul_rn_f32,
+    round_bf16_to_fp32,
+    rsqrt_approx_f32,
+    shfl_down_f32,
+)
+from .rope_tables import get_rope_tables
+
+_HEAD_DIM = 64
+_ROWS_PER_PROGRAM = 2
+_VERIFY_ROWS = 256
+_FAILED_ROPE_KEYS: set[tuple[int | None, float, float, tuple[int, int, int], float]] = set()
+_VERIFIED_ROPE_KEYS: set[tuple[int | None, float, float, tuple[int, int, int], float]] = set()
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _shuffle_lane_zero(value):
+    return tl.inline_asm_elementwise(
+        asm="shfl.sync.idx.b32 $0, $1, 0, 0x1f, 0xffffffff;",
+        constraints="=f,f",
+        args=[value],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def _warp_sum(value):
+    value = add_rn_f32(value, shfl_down_f32(value, 16))
+    value = add_rn_f32(value, shfl_down_f32(value, 8))
+    value = add_rn_f32(value, shfl_down_f32(value, 4))
+    value = add_rn_f32(value, shfl_down_f32(value, 2))
+    value = add_rn_f32(value, shfl_down_f32(value, 1))
+    return _shuffle_lane_zero(value)
+
+
+@triton.jit
+def _qk_rms_norm_scale_rope_3d_kernel(
+    query_output_ptr,
+    key_output_ptr,
+    query_ptr,
+    key_ptr,
+    query_weight_ptr,
+    key_weight_ptr,
+    cos_t_ptr,
+    sin_t_ptr,
+    cos_h_ptr,
+    sin_h_ptr,
+    cos_w_ptr,
+    sin_w_ptr,
+    rows,
+    tokens_per_batch,
+    height,
+    width,
+    heads,
+    eps,
+    query_scale,
+    rows_per_program: tl.constexpr,
+    head_dim: tl.constexpr,
+    pairs_t: tl.constexpr,
+    pairs_h: tl.constexpr,
+):
+    threads = tl.arange(0, rows_per_program * 32)
+    local_row = threads // 32
+    lane = threads % 32
+    row = tl.program_id(0).to(tl.int64) * rows_per_program + local_row
+    valid_row = row < rows
+    reduction_valid = valid_row & (lane < 16)
+    reduction_base = row * head_dim + lane * 4
+
+    query_accumulator = tl.zeros((rows_per_program * 32,), dtype=tl.float32)
+    key_accumulator = tl.zeros((rows_per_program * 32,), dtype=tl.float32)
+    for element in tl.static_range(4):
+        query = tl.load(query_ptr + reduction_base + element, mask=reduction_valid, other=0.0).to(tl.float32)
+        key = tl.load(key_ptr + reduction_base + element, mask=reduction_valid, other=0.0).to(tl.float32)
+        query_accumulator = fma_rn_f32(query, query, query_accumulator)
+        key_accumulator = fma_rn_f32(key, key, key_accumulator)
+    query_rstd = rsqrt_approx_f32(_warp_sum(query_accumulator) / head_dim + eps)
+    key_rstd = rsqrt_approx_f32(_warp_sum(key_accumulator) / head_dim + eps)
+
+    tokens = (row // heads) % tokens_per_batch
+    pos_w = tokens % width
+    pos_h = (tokens // width) % height
+    pos_t = tokens // (height * width)
+    pairs_per_head: tl.constexpr = head_dim // 2
+    pairs_w: tl.constexpr = pairs_per_head - pairs_t - pairs_h
+    is_t = lane < pairs_t
+    is_h = (lane >= pairs_t) & (lane < pairs_t + pairs_h)
+    is_w = ~(is_t | is_h)
+    pair_t = lane
+    pair_h = lane - pairs_t
+    pair_w = lane - pairs_t - pairs_h
+    cos_t = tl.load(cos_t_ptr + pos_t * pairs_t + pair_t, mask=valid_row & is_t, other=0.0)
+    sin_t = tl.load(sin_t_ptr + pos_t * pairs_t + pair_t, mask=valid_row & is_t, other=0.0)
+    cos_h = tl.load(cos_h_ptr + pos_h * pairs_h + pair_h, mask=valid_row & is_h, other=0.0)
+    sin_h = tl.load(sin_h_ptr + pos_h * pairs_h + pair_h, mask=valid_row & is_h, other=0.0)
+    cos_w = tl.load(cos_w_ptr + pos_w * pairs_w + pair_w, mask=valid_row & is_w, other=0.0)
+    sin_w = tl.load(sin_w_ptr + pos_w * pairs_w + pair_w, mask=valid_row & is_w, other=0.0)
+    cos = cos_t + cos_h + cos_w
+    sin = sin_t + sin_h + sin_w
+
+    even_column = lane * 2
+    odd_column = even_column + 1
+    even_offset = row * head_dim + even_column
+    odd_offset = row * head_dim + odd_column
+    query_even = tl.load(query_ptr + even_offset, mask=valid_row).to(tl.float32)
+    query_odd = tl.load(query_ptr + odd_offset, mask=valid_row).to(tl.float32)
+    key_even = tl.load(key_ptr + even_offset, mask=valid_row).to(tl.float32)
+    key_odd = tl.load(key_ptr + odd_offset, mask=valid_row).to(tl.float32)
+    query_even_weight = tl.load(query_weight_ptr + even_column).to(tl.float32)
+    query_odd_weight = tl.load(query_weight_ptr + odd_column).to(tl.float32)
+    key_even_weight = tl.load(key_weight_ptr + even_column).to(tl.float32)
+    key_odd_weight = tl.load(key_weight_ptr + odd_column).to(tl.float32)
+
+    query_even = round_bf16_to_fp32(mul_rn_f32(query_even_weight, mul_rn_f32(query_rstd, query_even)))
+    query_odd = round_bf16_to_fp32(mul_rn_f32(query_odd_weight, mul_rn_f32(query_rstd, query_odd)))
+    query_even = round_bf16_to_fp32(mul_rn_f32(query_even, query_scale))
+    query_odd = round_bf16_to_fp32(mul_rn_f32(query_odd, query_scale))
+    key_even = round_bf16_to_fp32(mul_rn_f32(key_even_weight, mul_rn_f32(key_rstd, key_even)))
+    key_odd = round_bf16_to_fp32(mul_rn_f32(key_odd_weight, mul_rn_f32(key_rstd, key_odd)))
+
+    tl.store(
+        query_output_ptr + even_offset,
+        add_rn_f32(mul_rn_f32(query_even, cos), -mul_rn_f32(query_odd, sin)),
+        mask=valid_row,
+    )
+    tl.store(
+        query_output_ptr + odd_offset,
+        add_rn_f32(mul_rn_f32(query_even, sin), mul_rn_f32(query_odd, cos)),
+        mask=valid_row,
+    )
+    tl.store(
+        key_output_ptr + even_offset,
+        add_rn_f32(mul_rn_f32(key_even, cos), -mul_rn_f32(key_odd, sin)),
+        mask=valid_row,
+    )
+    tl.store(
+        key_output_ptr + odd_offset,
+        add_rn_f32(mul_rn_f32(key_even, sin), mul_rn_f32(key_odd, cos)),
+        mask=valid_row,
+    )
+
+
+def _weight_matches(value: torch.Tensor, weight: torch.Tensor) -> bool:
+    return (
+        weight.is_cuda
+        and weight.device == value.device
+        and weight.dtype is torch.bfloat16
+        and weight.shape == (_HEAD_DIM,)
+        and weight.is_contiguous()
+    )
+
+
+def _supported_inputs(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    query_weight: torch.Tensor,
+    key_weight: torch.Tensor,
+    eps: float,
+    query_scale: float,
+) -> bool:
+    return (
+        is_ltx2_fusion_eligible(query)
+        and query.dtype is torch.bfloat16
+        and query.ndim == 6
+        and query.shape[-1] == _HEAD_DIM
+        and query.numel() > 0
+        and query.is_contiguous()
+        and key.is_cuda
+        and key.device == query.device
+        and key.dtype is query.dtype
+        and key.shape == query.shape
+        and key.is_contiguous()
+        and _weight_matches(query, query_weight)
+        and _weight_matches(query, key_weight)
+        and math.isfinite(eps)
+        and eps > 0
+        and math.isfinite(query_scale)
+    )
+
+
+def _combined_supported_inputs(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    query_weight: torch.Tensor,
+    key_weight: torch.Tensor,
+    eps: float,
+    query_scale: float,
+    dim_split: tuple[int, int, int],
+    base: float,
+) -> bool:
+    return (
+        _supported_inputs(query, key, query_weight, key_weight, eps, query_scale)
+        and len(dim_split) == 3
+        and sum(dim_split) == _HEAD_DIM
+        and all(dim > 0 and dim % 2 == 0 for dim in dim_split)
+        and math.isfinite(base)
+        and base > 0
+    )
+
+
+def _launch_combined(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    query_weight: torch.Tensor,
+    key_weight: torch.Tensor,
+    eps: float,
+    query_scale: float,
+    dim_split: tuple[int, int, int],
+    base: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    query_output = torch.empty_like(query)
+    key_output = torch.empty_like(key)
+    rows = query.numel() // _HEAD_DIM
+    tables = get_rope_tables(query, dim_split, base)
+    cos_t, sin_t = tables[0]
+    cos_h, sin_h = tables[1]
+    cos_w, sin_w = tables[2]
+    with torch.accelerator.device_index(query.device.index):
+        _qk_rms_norm_scale_rope_3d_kernel[(triton.cdiv(rows, _ROWS_PER_PROGRAM),)](
+            query_output,
+            key_output,
+            query,
+            key,
+            query_weight,
+            key_weight,
+            cos_t,
+            sin_t,
+            cos_h,
+            sin_h,
+            cos_w,
+            sin_w,
+            rows,
+            query.shape[1] * query.shape[2] * query.shape[3],
+            query.shape[2],
+            query.shape[3],
+            query.shape[4],
+            eps,
+            query_scale,
+            rows_per_program=_ROWS_PER_PROGRAM,
+            head_dim=_HEAD_DIM,
+            pairs_t=dim_split[0] // 2,
+            pairs_h=dim_split[1] // 2,
+            num_warps=_ROWS_PER_PROGRAM,
+        )
+    return query_output, key_output
+
+
+def _verify_combined_prefix(
+    output: tuple[torch.Tensor, torch.Tensor],
+    query: torch.Tensor,
+    key: torch.Tensor,
+    query_weight: torch.Tensor,
+    key_weight: torch.Tensor,
+    eps: float,
+    query_scale: float,
+    dim_split: tuple[int, int, int],
+    base: float,
+) -> bool:
+    rows = min(query.numel() // _HEAD_DIM, _VERIFY_ROWS)
+    query_rows = query.reshape(-1, _HEAD_DIM)[:rows]
+    key_rows = key.reshape(-1, _HEAD_DIM)[:rows]
+    query_rows = F.rms_norm(query_rows, (_HEAD_DIM,), query_weight, eps) * query_scale
+    key_rows = F.rms_norm(key_rows, (_HEAD_DIM,), key_weight, eps)
+    tables = get_rope_tables(query, dim_split, base)
+    row_indices = torch.arange(rows, device=query.device)
+    tokens_per_batch = query.shape[1] * query.shape[2] * query.shape[3]
+    tokens = (row_indices // query.shape[4]) % tokens_per_batch
+    positions = (
+        tokens // (query.shape[2] * query.shape[3]),
+        (tokens // query.shape[3]) % query.shape[2],
+        tokens % query.shape[3],
+    )
+    cos = torch.cat([axis_tables[0][position] for axis_tables, position in zip(tables, positions, strict=True)], -1)
+    sin = torch.cat([axis_tables[1][position] for axis_tables, position in zip(tables, positions, strict=True)], -1)
+
+    def rotate(value: torch.Tensor) -> torch.Tensor:
+        pairs = value.reshape(rows, _HEAD_DIM // 2, 2)
+        even = pairs[..., 0].float()
+        odd = pairs[..., 1].float()
+        return (
+            torch.stack((even * cos - odd * sin, even * sin + odd * cos), -1).reshape(rows, _HEAD_DIM).to(value.dtype)
+        )
+
+    reference = rotate(query_rows), rotate(key_rows)
+    actual = tuple(value.reshape(-1, _HEAD_DIM)[:rows] for value in output)
+    return all(torch.equal(a, b) for a, b in zip(actual, reference, strict=True))
+
+
+def try_qk_rms_norm_scale_rope_3d_exact(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    query_weight: torch.Tensor,
+    key_weight: torch.Tensor,
+    eps: float,
+    query_scale: float,
+    dim_split: tuple[int, int, int],
+    base: float,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    """Normalize Q/K, scale Q, and apply paired 3D RoPE in one kernel."""
+
+    if not _combined_supported_inputs(
+        query,
+        key,
+        query_weight,
+        key_weight,
+        eps,
+        query_scale,
+        dim_split,
+        base,
+    ):
+        return None
+    runtime_key = (query.device.index, float(eps), float(query_scale), dim_split, float(base))
+    if runtime_key in _FAILED_ROPE_KEYS:
+        return None
+    try:
+        output = _launch_combined(
+            query,
+            key,
+            query_weight,
+            key_weight,
+            eps,
+            query_scale,
+            dim_split,
+            base,
+        )
+        if runtime_key not in _VERIFIED_ROPE_KEYS:
+            if not _verify_combined_prefix(
+                output,
+                query,
+                key,
+                query_weight,
+                key_weight,
+                eps,
+                query_scale,
+                dim_split,
+                base,
+            ):
+                _FAILED_ROPE_KEYS.add(runtime_key)
+                logger.warning(
+                    "Disabling LTX DiffVAE Q/K RMSNorm+3D RoPE fusion on %s after a bit-exactness mismatch",
+                    query.device,
+                )
+                return None
+            _VERIFIED_ROPE_KEYS.add(runtime_key)
+    except Exception as exc:  # noqa: BLE001 - fail closed after optimized-path failure
+        _FAILED_ROPE_KEYS.add(runtime_key)
+        logger.warning(
+            "Disabling LTX DiffVAE Q/K RMSNorm+3D RoPE fusion on %s after failure: %s",
+            query.device,
+            exc,
+        )
+        return None
+    return output
+
+
+__all__ = ["try_qk_rms_norm_scale_rope_3d_exact"]

vllm_omni/diffusion/models/ltx2/ops/residual_adaln.py

@@ -0,0 +1,489 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Exact LTX DiffVAE residual, RMSNorm, and AdaLN Triton fusions."""
+
+from __future__ import annotations
+
+import logging
+import math
+
+import torch
+import torch.nn.functional as F
+from vllm.triton_utils import tl, triton
+
+from . import is_ltx2_fusion_eligible
+from ._numerics import (
+    add_rn_f32,
+    fma_rn_f32,
+    mul_rn_f32,
+    round_bf16_to_fp32,
+    rsqrt_approx_f32,
+    shfl_down_f32,
+)
+
+_HIDDEN_SIZE = 256
+_POINTWISE_BLOCK = 1024
+_VERIFY_ROWS = 256
+_VERIFY_ELEMENTS = 1 << 20
+_FAILED_ADALN_KEYS: set[tuple[int | None, int, float]] = set()
+_VERIFIED_ADALN_KEYS: set[tuple[int | None, int, float]] = set()
+_FAILED_ADD_DEVICES: set[int | None] = set()
+_VERIFIED_ADD_DEVICES: set[int | None] = set()
+
+logger = logging.getLogger(__name__)
+
+
+@triton.jit
+def _two_warp_sum(values):
+    values = add_rn_f32(values, shfl_down_f32(values, 16))
+    values = add_rn_f32(values, shfl_down_f32(values, 8))
+    values = add_rn_f32(values, shfl_down_f32(values, 4))
+    values = add_rn_f32(values, shfl_down_f32(values, 2))
+    values = add_rn_f32(values, shfl_down_f32(values, 1))
+    threads = tl.arange(0, 64)
+    warp_0 = tl.sum(tl.where(threads == 0, values, 0.0))
+    warp_1 = tl.sum(tl.where(threads == 32, values, 0.0))
+    return add_rn_f32(warp_0, warp_1)
+
+
+@triton.jit
+def _load_ordered_residual(
+    x_ptr,
+    residual_a_ptr,
+    residual_b_ptr,
+    offset,
+    valid,
+    residual_count: tl.constexpr,
+):
+    value = add_rn_f32(
+        tl.load(x_ptr + offset, mask=valid, other=0.0).to(tl.float32),
+        tl.load(residual_a_ptr + offset, mask=valid, other=0.0).to(tl.float32),
+    )
+    value = round_bf16_to_fp32(value)
+    if residual_count == 2:
+        value = add_rn_f32(
+            value,
+            tl.load(residual_b_ptr + offset, mask=valid, other=0.0).to(tl.float32),
+        )
+        value = round_bf16_to_fp32(value)
+    return value
+
+
+@triton.jit
+def _store_modulated(
+    output_ptr,
+    norm_weight_ptr,
+    scale_ptr,
+    shift_ptr,
+    value,
+    offset,
+    column,
+    modulation_base,
+    reciprocal_rms,
+    valid,
+):
+    normalized = mul_rn_f32(value, reciprocal_rms)
+    norm_weight = tl.load(norm_weight_ptr + column, mask=valid, other=0.0).to(tl.float32)
+    normalized = round_bf16_to_fp32(mul_rn_f32(normalized, norm_weight))
+    scale = tl.load(scale_ptr + modulation_base + column, mask=valid, other=0.0).to(tl.float32)
+    shift = tl.load(shift_ptr + modulation_base + column, mask=valid, other=0.0).to(tl.float32)
+    one_plus_scale = round_bf16_to_fp32(add_rn_f32(1.0, scale))
+    product = round_bf16_to_fp32(mul_rn_f32(normalized, one_plus_scale))
+    tl.store(output_ptr + offset, add_rn_f32(product, shift), mask=valid)
+
+
+@triton.jit
+def _residual_rms_norm_modulate_kernel(
+    output_ptr,
+    x_ptr,
+    residual_a_ptr,
+    residual_b_ptr,
+    norm_weight_ptr,
+    scale_ptr,
+    shift_ptr,
+    rows_per_batch,
+    eps,
+    residual_count: tl.constexpr,
+    hidden_size: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    batch = row // rows_per_batch
+    threads = tl.arange(0, 64)
+    row_base = row * hidden_size
+
+    # PyTorch's BF16 RMSNorm kernel uses four contiguous values per thread for
+    # hidden=256. Keep the four ordered BF16 residual values in registers so
+    # the output pass does not read the large activations a second time.
+    column_0 = threads * 4
+    column_1 = column_0 + 1
+    column_2 = column_0 + 2
+    column_3 = column_0 + 3
+    valid_0 = column_0 < hidden_size
+    valid_1 = column_1 < hidden_size
+    valid_2 = column_2 < hidden_size
+    valid_3 = column_3 < hidden_size
+    offset_0 = row_base + column_0
+    offset_1 = row_base + column_1
+    offset_2 = row_base + column_2
+    offset_3 = row_base + column_3
+    value_0 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_0, valid_0, residual_count)
+    value_1 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_1, valid_1, residual_count)
+    value_2 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_2, valid_2, residual_count)
+    value_3 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_3, valid_3, residual_count)
+    accumulator = fma_rn_f32(value_0, value_0, tl.zeros((64,), dtype=tl.float32))
+    accumulator = fma_rn_f32(value_1, value_1, accumulator)
+    accumulator = fma_rn_f32(value_2, value_2, accumulator)
+    accumulator = fma_rn_f32(value_3, value_3, accumulator)
+
+    total = _two_warp_sum(accumulator)
+    reciprocal_rms = rsqrt_approx_f32(total / hidden_size + eps)
+    modulation_base = batch * hidden_size
+    _store_modulated(
+        output_ptr,
+        norm_weight_ptr,
+        scale_ptr,
+        shift_ptr,
+        value_0,
+        offset_0,
+        column_0,
+        modulation_base,
+        reciprocal_rms,
+        valid_0,
+    )
+    _store_modulated(
+        output_ptr,
+        norm_weight_ptr,
+        scale_ptr,
+        shift_ptr,
+        value_1,
+        offset_1,
+        column_1,
+        modulation_base,
+        reciprocal_rms,
+        valid_1,
+    )
+    _store_modulated(
+        output_ptr,
+        norm_weight_ptr,
+        scale_ptr,
+        shift_ptr,
+        value_2,
+        offset_2,
+        column_2,
+        modulation_base,
+        reciprocal_rms,
+        valid_2,
+    )
+    _store_modulated(
+        output_ptr,
+        norm_weight_ptr,
+        scale_ptr,
+        shift_ptr,
+        value_3,
+        offset_3,
+        column_3,
+        modulation_base,
+        reciprocal_rms,
+        valid_3,
+    )
+
+
+@triton.jit
+def _residual_add3_kernel(
+    output_ptr,
+    x_ptr,
+    residual_a_ptr,
+    residual_b_ptr,
+    residual_c_ptr,
+    elements,
+    block: tl.constexpr,
+):
+    offsets = tl.program_id(0).to(tl.int64) * block + tl.arange(0, block)
+    valid = offsets < elements
+    value = add_rn_f32(
+        tl.load(x_ptr + offsets, mask=valid, other=0.0).to(tl.float32),
+        tl.load(residual_a_ptr + offsets, mask=valid, other=0.0).to(tl.float32),
+    )
+    value = round_bf16_to_fp32(value)
+    value = add_rn_f32(
+        value,
+        tl.load(residual_b_ptr + offsets, mask=valid, other=0.0).to(tl.float32),
+    )
+    value = round_bf16_to_fp32(value)
+    value = add_rn_f32(
+        value,
+        tl.load(residual_c_ptr + offsets, mask=valid, other=0.0).to(tl.float32),
+    )
+    tl.store(output_ptr + offsets, value, mask=valid)
+
+
+def _same_activation(x: torch.Tensor, residual: torch.Tensor) -> bool:
+    return (
+        residual.dtype is torch.bfloat16
+        and residual.is_cuda
+        and residual.device == x.device
+        and residual.shape == x.shape
+        and residual.is_contiguous()
+    )
+
+
+def _modulation_matches(x: torch.Tensor, value: torch.Tensor) -> bool:
+    return (
+        value.dtype is torch.bfloat16
+        and value.is_cuda
+        and value.device == x.device
+        and value.shape[-1] == _HIDDEN_SIZE
+        and value.numel() == x.shape[0] * _HIDDEN_SIZE
+        and value.is_contiguous()
+    )
+
+
+def _norm_weight_matches(x: torch.Tensor, weight: torch.Tensor) -> bool:
+    return (
+        weight.dtype is torch.bfloat16
+        and weight.is_cuda
+        and weight.device == x.device
+        and weight.shape == (_HIDDEN_SIZE,)
+        and weight.is_contiguous()
+    )
+
+
+def _adaln_inputs_supported(
+    x: torch.Tensor,
+    residual_a: torch.Tensor,
+    residual_b: torch.Tensor | None,
+    norm_weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+) -> bool:
+    return (
+        is_ltx2_fusion_eligible(x)
+        and x.dtype is torch.bfloat16
+        and x.ndim >= 2
+        and x.shape[-1] == _HIDDEN_SIZE
+        and x.numel() > 0
+        and x.is_contiguous()
+        and _same_activation(x, residual_a)
+        and (residual_b is None or _same_activation(x, residual_b))
+        and _norm_weight_matches(x, norm_weight)
+        and _modulation_matches(x, scale)
+        and _modulation_matches(x, shift)
+        and math.isfinite(eps)
+        and eps > 0
+    )
+
+
+def _residual_reference(
+    x: torch.Tensor,
+    residual_a: torch.Tensor,
+    residual_b: torch.Tensor | None,
+) -> torch.Tensor:
+    hidden_states = x + residual_a
+    if residual_b is not None:
+        hidden_states = hidden_states + residual_b
+    return hidden_states
+
+
+def _adaln_reference(
+    x: torch.Tensor,
+    residual_a: torch.Tensor,
+    residual_b: torch.Tensor | None,
+    norm_weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    hidden_states = _residual_reference(x, residual_a, residual_b)
+    normalized = F.rms_norm(hidden_states, (_HIDDEN_SIZE,), weight=norm_weight, eps=eps)
+    return normalized * (1 + scale) + shift
+
+
+def _launch_adaln(
+    x: torch.Tensor,
+    residual_a: torch.Tensor,
+    residual_b: torch.Tensor | None,
+    norm_weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    output = torch.empty_like(x)
+    rows = x.numel() // _HIDDEN_SIZE
+    rows_per_batch = rows // x.shape[0]
+    with torch.accelerator.device_index(x.device.index):
+        _residual_rms_norm_modulate_kernel[(rows,)](
+            output,
+            x,
+            residual_a,
+            x if residual_b is None else residual_b,
+            norm_weight,
+            scale,
+            shift,
+            rows_per_batch,
+            eps,
+            residual_count=1 if residual_b is None else 2,
+            hidden_size=_HIDDEN_SIZE,
+            num_warps=2,
+        )
+    return output
+
+
+def _verify_adaln_prefix(
+    output: torch.Tensor,
+    x: torch.Tensor,
+    residual_a: torch.Tensor,
+    residual_b: torch.Tensor | None,
+    norm_weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+) -> bool:
+    rows_per_batch = x.numel() // (x.shape[0] * _HIDDEN_SIZE)
+    rows = min(rows_per_batch, _VERIFY_ROWS)
+    x_rows = x.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows]
+    residual_a_rows = residual_a.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows]
+    residual_b_rows = (
+        None if residual_b is None else residual_b.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows]
+    )
+    scale_row = scale.reshape(x.shape[0], 1, _HIDDEN_SIZE)[0]
+    shift_row = shift.reshape(x.shape[0], 1, _HIDDEN_SIZE)[0]
+    reference = _adaln_reference(
+        x_rows,
+        residual_a_rows,
+        residual_b_rows,
+        norm_weight,
+        scale_row,
+        shift_row,
+        eps,
+    )
+    actual = output.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows]
+    return torch.equal(actual, reference)
+
+
+def try_residual_rms_norm_modulate_exact(
+    x: torch.Tensor,
+    residual_a: torch.Tensor,
+    residual_b: torch.Tensor | None,
+    norm_weight: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    """Fuse one or two ordered BF16 residuals with RMSNorm and AdaLN."""
+
+    if not _adaln_inputs_supported(x, residual_a, residual_b, norm_weight, scale, shift, eps):
+        return None
+    residual_count = 1 if residual_b is None else 2
+    runtime_key = (x.device.index, residual_count, float(eps))
+    if runtime_key in _FAILED_ADALN_KEYS:
+        return None
+    try:
+        output = _launch_adaln(x, residual_a, residual_b, norm_weight, scale, shift, eps)
+        if runtime_key not in _VERIFIED_ADALN_KEYS:
+            if not _verify_adaln_prefix(
+                output,
+                x,
+                residual_a,
+                residual_b,
+                norm_weight,
+                scale,
+                shift,
+                eps,
+            ):
+                _FAILED_ADALN_KEYS.add(runtime_key)
+                logger.warning(
+                    "Disabling LTX DiffVAE residual-AdaLN fusion on %s after a bit-exactness mismatch",
+                    x.device,
+                )
+                return None
+            _VERIFIED_ADALN_KEYS.add(runtime_key)
+    except Exception as exc:  # noqa: BLE001 - fail closed after optimized-path failure
+        _FAILED_ADALN_KEYS.add(runtime_key)
+        logger.warning(
+            "Disabling LTX DiffVAE residual-AdaLN fusion on %s after failure: %s",
+            x.device,
+            exc,
+        )
+        return None
+    return output
+
+
+def _add_inputs_supported(x: torch.Tensor, residuals: tuple[torch.Tensor, ...]) -> bool:
+    return (
+        is_ltx2_fusion_eligible(x)
+        and x.dtype is torch.bfloat16
+        and x.numel() > 0
+        and x.is_contiguous()
+        and all(_same_activation(x, residual) for residual in residuals)
+    )
+
+
+def _launch_add3(
+    x: torch.Tensor,
+    residual_a: torch.Tensor,
+    residual_b: torch.Tensor,
+    residual_c: torch.Tensor,
+) -> torch.Tensor:
+    output = torch.empty_like(x)
+    with torch.accelerator.device_index(x.device.index):
+        _residual_add3_kernel[(triton.cdiv(x.numel(), _POINTWISE_BLOCK),)](
+            output,
+            x,
+            residual_a,
+            residual_b,
+            residual_c,
+            x.numel(),
+            block=_POINTWISE_BLOCK,
+            num_warps=4,
+        )
+    return output
+
+
+def try_residual_add3_exact(
+    x: torch.Tensor,
+    residual_a: torch.Tensor,
+    residual_b: torch.Tensor,
+    residual_c: torch.Tensor,
+) -> torch.Tensor | None:
+    """Fuse three ordered BF16 residual additions into one pointwise pass."""
+
+    residuals = (residual_a, residual_b, residual_c)
+    if not _add_inputs_supported(x, residuals) or x.device.index in _FAILED_ADD_DEVICES:
+        return None
+    try:
+        output = _launch_add3(x, *residuals)
+        if x.device.index not in _VERIFIED_ADD_DEVICES:
+            elements = min(x.numel(), _VERIFY_ELEMENTS)
+            reference = (
+                _residual_reference(
+                    x.reshape(-1)[:elements],
+                    residual_a.reshape(-1)[:elements],
+                    residual_b.reshape(-1)[:elements],
+                )
+                + residual_c.reshape(-1)[:elements]
+            )
+            if not torch.equal(output.reshape(-1)[:elements], reference):
+                _FAILED_ADD_DEVICES.add(x.device.index)
+                logger.warning(
+                    "Disabling LTX DiffVAE residual-add fusion on %s after a bit-exactness mismatch",
+                    x.device,
+                )
+                return None
+            _VERIFIED_ADD_DEVICES.add(x.device.index)
+    except Exception as exc:  # noqa: BLE001 - fail closed after optimized-path failure
+        _FAILED_ADD_DEVICES.add(x.device.index)
+        logger.warning(
+            "Disabling LTX DiffVAE residual-add fusion on %s after failure: %s",
+            x.device,
+            exc,
+        )
+        return None
+    return output
+
+
+__all__ = [
+    "try_residual_add3_exact",
+    "try_residual_rms_norm_modulate_exact",
+]

vllm_omni/diffusion/models/ltx2/ops/rope_tables.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Cached 3D RoPE tables for LTX DiffVAE kernels."""
+
+from __future__ import annotations
+
+import torch
+
+_TABLE_CACHE_MAX = 16
+_TABLE_CACHE: dict[
+    tuple[torch.device, int, int, int, tuple[int, int, int], float],
+    tuple[tuple[torch.Tensor, torch.Tensor], ...],
+] = {}
+
+
+def _axis_tables(
+    length: int,
+    dim: int,
+    device: torch.device,
+    base: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    exponents = torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim
+    inv_freqs = (1.0 / base**exponents).to(torch.float32)
+    positions = torch.arange(length, dtype=torch.float32, device=device)
+    angles = positions[:, None] * inv_freqs[None, :]
+    return angles.cos().contiguous(), angles.sin().contiguous()
+
+
+def get_rope_tables(
+    query: torch.Tensor,
+    dim_split: tuple[int, int, int],
+    base: float,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
+    num_frames, height, width = query.shape[1:4]
+    cache_key = (query.device, num_frames, height, width, dim_split, base)
+    cached = _TABLE_CACHE.get(cache_key)
+    if cached is not None:
+        return cached
+    tables = tuple(
+        _axis_tables(length, dim, query.device, base)
+        for length, dim in zip((num_frames, height, width), dim_split, strict=True)
+    )
+    if len(_TABLE_CACHE) >= _TABLE_CACHE_MAX:
+        _TABLE_CACHE.pop(next(iter(_TABLE_CACHE)))
+    _TABLE_CACHE[cache_key] = tables
+    return tables
+
+
+__all__ = ["get_rope_tables"]

vllm_omni/diffusion/models/ltx2/ops/swiglu.py

@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Bit-exact workspace-reusing SwiGLU for the LTX diffusion VAE decoder."""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+import torch.nn.functional as F
+from vllm.triton_utils import tl, triton
+
+from . import is_ltx2_fusion_eligible
+from ._numerics import mul_rn_f32, round_bf16_to_fp32
+
+_POINTWISE_BLOCK = 1024
+_BASE_TILE_ROWS = 16384
+_MAX_TILE_HIDDEN_ELEMENTS = _BASE_TILE_ROWS * 4 * 2048
+_FAILED_KEYS: set[tuple[int | None, int, int, int, int]] = set()
+_VERIFIED_KEYS: set[tuple[int | None, int, int, int, int]] = set()
+
+logger = logging.getLogger(__name__)
+
+
+def _tile_rows(rows: int, hidden: int) -> int:
+    """Bound projection storage while giving narrow MLPs larger GEMMs."""
+
+    budget_rows = _MAX_TILE_HIDDEN_ELEMENTS // hidden
+    aligned_rows = max(_BASE_TILE_ROWS, budget_rows // _BASE_TILE_ROWS * _BASE_TILE_ROWS)
+    return min(rows, aligned_rows)
+
+
+@triton.jit
+def _silu_mul_kernel(
+    gate_ptr,
+    up_ptr,
+    elements,
+    block: tl.constexpr,
+):
+    offsets = tl.program_id(0).to(tl.int64) * block + tl.arange(0, block)
+    valid = offsets < elements
+    gate = tl.load(gate_ptr + offsets, mask=valid).to(tl.float32)
+    up = tl.load(up_ptr + offsets, mask=valid).to(tl.float32)
+    # Eager rounds SiLU to BF16 before the separate BF16 multiply.
+    silu = round_bf16_to_fp32(mul_rn_f32(gate, tl.sigmoid(gate)))
+    tl.store(gate_ptr + offsets, mul_rn_f32(silu, up), mask=valid)
+
+
+def _weight_matches(
+    hidden_states: torch.Tensor,
+    weight: torch.Tensor,
+    shape: tuple[int, int],
+) -> bool:
+    return (
+        weight.is_cuda
+        and weight.device == hidden_states.device
+        and weight.dtype is torch.bfloat16
+        and weight.shape == shape
+        and weight.is_contiguous()
+    )
+
+
+def _supported_inputs(
+    hidden_states: torch.Tensor,
+    gate_weight: torch.Tensor,
+    up_weight: torch.Tensor,
+    down_weight: torch.Tensor,
+) -> bool:
+    if hidden_states.ndim < 3:
+        return False
+    dim = hidden_states.shape[-1]
+    hidden = gate_weight.shape[0] if gate_weight.ndim == 2 else 0
+    return (
+        is_ltx2_fusion_eligible(hidden_states)
+        and hidden_states.dtype is torch.bfloat16
+        and hidden_states.shape[0] == 1
+        and hidden_states.numel() > 0
+        and hidden_states.is_contiguous()
+        and 256 <= dim <= 2048
+        and hidden == 4 * dim
+        and _weight_matches(hidden_states, gate_weight, (hidden, dim))
+        and _weight_matches(hidden_states, up_weight, (hidden, dim))
+        and _weight_matches(hidden_states, down_weight, (dim, hidden))
+    )
+
+
+def _silu_mul_inplace(gate: torch.Tensor, up: torch.Tensor) -> None:
+    with torch.accelerator.device_index(gate.device.index):
+        _silu_mul_kernel[(triton.cdiv(gate.numel(), _POINTWISE_BLOCK),)](
+            gate,
+            up,
+            gate.numel(),
+            block=_POINTWISE_BLOCK,
+            num_warps=4,
+        )
+
+
+def _reference_tile(
+    hidden_states: torch.Tensor,
+    gate_weight: torch.Tensor,
+    up_weight: torch.Tensor,
+    down_weight: torch.Tensor,
+) -> torch.Tensor:
+    return F.linear(
+        F.silu(F.linear(hidden_states, gate_weight)) * F.linear(hidden_states, up_weight),
+        down_weight,
+    )
+
+
+def _launch(
+    hidden_states: torch.Tensor,
+    gate_weight: torch.Tensor,
+    up_weight: torch.Tensor,
+    down_weight: torch.Tensor,
+) -> torch.Tensor:
+    dim = hidden_states.shape[-1]
+    hidden = gate_weight.shape[0]
+    flat = hidden_states.reshape(-1, dim)
+    tile_rows = _tile_rows(flat.shape[0], hidden)
+    output = torch.empty_like(flat)
+    workspace_gate = torch.empty((tile_rows, hidden), device=flat.device, dtype=flat.dtype)
+    for start in range(0, flat.shape[0], tile_rows):
+        end = min(start + tile_rows, flat.shape[0])
+        tokens = end - start
+        tile = flat[start:end]
+        gate = workspace_gate[:tokens]
+        torch.mm(tile, gate_weight.t(), out=gate)
+        up = torch.mm(tile, up_weight.t())
+        _silu_mul_inplace(gate, up)
+        torch.mm(gate, down_weight.t(), out=output[start:end])
+    return output.reshape(hidden_states.shape)
+
+
+def _verify_edge_tiles(
+    output: torch.Tensor,
+    hidden_states: torch.Tensor,
+    gate_weight: torch.Tensor,
+    up_weight: torch.Tensor,
+    down_weight: torch.Tensor,
+) -> bool:
+    from diffusers.models.autoencoders.ltx2_diffusion_decoder import _SWIGLU_TILE_SIZE
+
+    dim = hidden_states.shape[-1]
+    flat = hidden_states.reshape(-1, dim)
+    actual = output.reshape(-1, dim)
+    tile_rows = _tile_rows(flat.shape[0], gate_weight.shape[0])
+    starts = {0, ((flat.shape[0] - 1) // tile_rows) * tile_rows}
+    for start in starts:
+        end = min(start + tile_rows, flat.shape[0])
+        # Match upstream's token tiling, not the optimized GEMM shape: changing
+        # the matmul row count can change its reduction order and BF16 output.
+        for reference_start in range(start, end, _SWIGLU_TILE_SIZE):
+            reference_end = min(reference_start + _SWIGLU_TILE_SIZE, end)
+            if not torch.equal(
+                actual[reference_start:reference_end],
+                _reference_tile(flat[reference_start:reference_end], gate_weight, up_weight, down_weight),
+            ):
+                return False
+    return True
+
+
+def try_swiglu_tiled_exact(
+    hidden_states: torch.Tensor,
+    gate_weight: torch.Tensor,
+    up_weight: torch.Tensor,
+    down_weight: torch.Tensor,
+) -> torch.Tensor | None:
+    """Run workspace-reusing tiled SwiGLU, or expose the original module."""
+
+    if not _supported_inputs(hidden_states, gate_weight, up_weight, down_weight):
+        return None
+    rows = hidden_states.numel() // hidden_states.shape[-1]
+    tile_rows = _tile_rows(rows, gate_weight.shape[0])
+    runtime_key = (
+        hidden_states.device.index,
+        hidden_states.shape[-1],
+        gate_weight.shape[0],
+        tile_rows,
+        (rows - 1) % tile_rows + 1,
+    )
+    if runtime_key in _FAILED_KEYS:
+        return None
+    try:
+        output = _launch(hidden_states, gate_weight, up_weight, down_weight)
+        if runtime_key not in _VERIFIED_KEYS:
+            if not _verify_edge_tiles(output, hidden_states, gate_weight, up_weight, down_weight):
+                _FAILED_KEYS.add(runtime_key)
+                logger.warning(
+                    "Disabling LTX DiffVAE tiled SwiGLU fusion on %s after a bit-exactness mismatch",
+                    hidden_states.device,
+                )
+                return None
+            _VERIFIED_KEYS.add(runtime_key)
+    except Exception as exc:  # noqa: BLE001 - fail closed after optimized-path failure
+        _FAILED_KEYS.add(runtime_key)
+        logger.warning(
+            "Disabling LTX DiffVAE tiled SwiGLU fusion on %s after failure: %s",
+            hidden_states.device,
+            exc,
+        )
+        return None
+    return output
+
+
+__all__ = ["try_swiglu_tiled_exact"]

vllm_omni/diffusion/models/ltx2/vae/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Ordinary LTX-2.5 DiffVAE, tiled decoding and acceleration operators."""

vllm_omni/diffusion/models/ltx2/vae/decoder.py

@@ -7,19 +7,24 @@
 
 """LTX-2.5 Native-weight adapter for Diffusers' diffusion VAE decoder."""
 
+import math
 from collections.abc import Mapping
+from functools import cache
 from typing import Any
 
 import torch
 from diffusers.models.autoencoders import (
     LTX2VideoDiffusionDecoderModel as DiffusersLTX2VideoDiffusionDecoderModel,
 )
+from diffusers.models.autoencoders import ltx2_diffusion_decoder as diffusers_decoder
 from diffusers.models.autoencoders.ltx2_diffusion_decoder import (
     LTX2VideoDiffusionDecoder3d as DiffusersLTX2VideoDiffusionDecoder3d,
 )
 from diffusers.models.autoencoders.vae import DecoderOutput
 from safetensors import safe_open
 
+from ..ops import is_ltx2_fna_eligible, is_ltx2_fusion_eligible, resolve_ltx2_vae_operators
+
 LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID = "Lightricks/LTX-2.5"
 LTX25_NATIVE_ARTIFACT_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
 LTX25_NATIVE_DIFFUSION_DECODER_FILENAME = "vae/ltx-2.5-video-vae-bf16.safetensors"
@@ -131,6 +136,177 @@ def load_ltx25_native_diffusion_decoder_state_dict(path: str) -> dict[str, torch
     return convert_ltx25_native_diffusion_decoder_state_dict(native_state_dict)
 
 
+@cache
+def _load_natten_token_permutation():
+    from kernels import get_kernel
+
+    natten = get_kernel("shi-labs/natten", version=1)
+    return natten.token_permute.token_permute_operation, natten.token_permute.token_unpermute_operation
+
+
+class LTX2VideoVaeNeighborhoodNattenProcessor(diffusers_decoder.LTX2VideoVaeNeighborhoodNattenProcessor):
+    """NATTEN dispatch with Omni FNA for eligible stage-5 grids."""
+
+    def __call__(
+        self, attn: "LTX2VideoVaeNeighborhoodAttention", hidden_states: torch.Tensor, block_mask=None
+    ) -> torch.Tensor:
+        batch_size, num_frames, height, width, channels = hidden_states.shape
+        query, key, value = attn.project_qkv(hidden_states)
+        # NATTEN's CUTLASS kernels silently produce wrong output for non-contiguous inputs.
+        query, key, value = query.contiguous(), key.contiguous(), value.contiguous()
+        operators = resolve_ltx2_vae_operators(hidden_states.device) if is_ltx2_fna_eligible(hidden_states) else None
+        # `scale=1.0`: the query is already scaled in `project_qkv`, as in the reference.
+        if (
+            tuple(attn.kernel_size) == (11, 11, 11)
+            and hidden_states.dtype == torch.bfloat16
+            and channels == 256
+            and attn.head_dim == 64
+            and operators is not None
+            and operators.fna is not None
+            # The native metadata encodes KV tile indices in 15 bits.
+            and math.prod((n + tile - 1) // tile for n, tile in zip((num_frames, height, width), (4, 4, 8))) <= 1 << 15
+        ):
+            token_permute, token_unpermute = _load_natten_token_permutation()
+            query_permuted, token_shape, _ = token_permute(query, (4, 4, 4), flip_tiled_dims=True)
+            key_permuted, _, _ = token_permute(key, (4, 4, 8), flip_tiled_dims=True)
+            value_permuted, _, _ = token_permute(value, (4, 4, 8), flip_tiled_dims=True)
+            output = operators.fna(
+                query_permuted,
+                key_permuted,
+                value_permuted,
+                shape=(num_frames, height, width),
+                window=tuple(attn.kernel_size),
+            )
+            hidden_states = token_unpermute(
+                output,
+                token_layout_shape=token_shape,
+                tile_shape=(4, 4, 4),
+                flip_tiled_dims=True,
+            )
+        else:
+            hidden_states = self._na3d(
+                query,
+                key,
+                value,
+                kernel_size=attn.kernel_size,
+                scale=1.0,
+                backend=self.backend,
+            )
+        hidden_states = hidden_states.reshape(batch_size, num_frames, height, width, channels)
+        return attn.to_out[0](hidden_states)
+
+
+class LTX2VideoVaeNeighborhoodAttention(diffusers_decoder.LTX2VideoVaeNeighborhoodAttention):
+    _available_processors = [
+        *diffusers_decoder.LTX2VideoVaeNeighborhoodAttention._available_processors,
+        LTX2VideoVaeNeighborhoodNattenProcessor,
+    ]
+
+    def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size, num_frames, height, width, _ = hidden_states.shape
+        shape = (batch_size, num_frames, height, width, self.heads, self.head_dim)
+        query = self.to_q(hidden_states).view(shape)
+        key = self.to_k(hidden_states).view(shape)
+        value = self.to_v(hidden_states).view(shape)
+
+        operators = resolve_ltx2_vae_operators(query.device) if is_ltx2_fusion_eligible(query) else None
+        optimized = (
+            None
+            if operators is None
+            else operators.qk_norm_rope(
+                query,
+                key,
+                self.norm_q.weight,
+                self.norm_k.weight,
+                self.norm_q.eps,
+                self.scale,
+                self.rope.rope_dim_split,
+                self.rope.base,
+            )
+        )
+        if optimized is not None:
+            query, key = optimized
+            return query, key, value
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+        query = query * self.scale
+        return self.rope(query), self.rope(key), value
+
+
+class LTX2VideoVaeSwiGLU(diffusers_decoder.LTX2VideoVaeSwiGLU):
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        operators = resolve_ltx2_vae_operators(hidden_states.device) if is_ltx2_fusion_eligible(hidden_states) else None
+        optimized = (
+            None
+            if operators is None
+            else operators.swiglu(hidden_states, self.w_gate.weight, self.w_up.weight, self.w_down.weight)
+        )
+        return super().forward(hidden_states) if optimized is None else optimized
+
+
+class LTX2VideoVaeDiffusionNABlock(diffusers_decoder.LTX2VideoVaeDiffusionNABlock):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        latent_context: torch.Tensor,
+        modulation: tuple[torch.Tensor, ...],
+        block_mask=None,
+    ) -> torch.Tensor:
+        if hidden_states.dtype is not torch.bfloat16 or not is_ltx2_fusion_eligible(hidden_states):
+            return super().forward(hidden_states, latent_context, modulation, block_mask)
+        operators = resolve_ltx2_vae_operators(hidden_states.device)
+        assert operators is not None
+
+        scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [
+            modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(self.num_mod_params)
+        ]
+        context_output = self.context_proj(latent_context)
+        attention_input = operators.residual_norm(
+            hidden_states,
+            context_output,
+            None,
+            self.norm1.weight,
+            scale_msa,
+            shift_msa,
+            self.norm1.eps,
+        )
+        if attention_input is None:
+            hidden_states = hidden_states + context_output
+            hidden_states = hidden_states + self.attn(
+                self.norm1(hidden_states) * (1 + scale_msa) + shift_msa, block_mask
+            )
+            hidden_states = hidden_states + self.mlp(self.norm2(hidden_states) * (1 + scale_mlp) + shift_mlp)
+            return hidden_states
+
+        attention_output = self.attn(attention_input, block_mask)
+        mlp_input = operators.residual_norm(
+            hidden_states,
+            context_output,
+            attention_output,
+            self.norm2.weight,
+            scale_mlp,
+            shift_mlp,
+            self.norm2.eps,
+        )
+        if mlp_input is None:
+            residual = hidden_states + context_output
+            residual = residual + attention_output
+            mlp_input = self.norm2(residual) * (1 + scale_mlp) + shift_mlp
+        mlp_output = self.mlp(mlp_input)
+        output = operators.residual_add(
+            hidden_states,
+            context_output,
+            attention_output,
+            mlp_output,
+        )
+        if output is not None:
+            return output
+        residual = hidden_states + context_output
+        residual = residual + attention_output
+        return residual + mlp_output
+
+
 class LTX2VideoDiffusionDecoder3d(DiffusersLTX2VideoDiffusionDecoder3d):
     """Diffusers decoder core with the short-clip NATTEN context fix."""
 
@@ -200,6 +376,15 @@ class LTX2VideoDiffusionDecoderModel(DiffusersLTX2VideoDiffusionDecoderModel):
         # to this behavior-only subclass preserves parameters and state-dict keys.
         self.decoder.__class__ = LTX2VideoDiffusionDecoder3d
         self.decoder.stage5_kernel = tuple(decoder_stage5_kernel)
+        overrides = {
+            diffusers_decoder.LTX2VideoVaeNeighborhoodAttention: LTX2VideoVaeNeighborhoodAttention,
+            diffusers_decoder.LTX2VideoVaeSwiGLU: LTX2VideoVaeSwiGLU,
+            diffusers_decoder.LTX2VideoVaeDiffusionNABlock: LTX2VideoVaeDiffusionNABlock,
+        }
+        # Preserve every module and parameter object; only execution methods change.
+        for module in self.decoder.modules():
+            if replacement := overrides.get(type(module)):
+                module.__class__ = replacement
 
     @classmethod
     def from_ltx25_native_checkpoint(

vllm_omni/diffusion/models/ltx2/vae/distributed.py

@@ -25,7 +25,7 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     TileTask,
 )
 
-from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
+from .decoder import LTX2VideoDiffusionDecoderModel
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name