[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,
         )