LTX-2.5 optimization: integrate faster AITER NA3D flash attention for diffusion decoder

tests/integration/accelerator/model_executor/layers/test_ltx2_na3d_flash.py

@@ -0,0 +1,67 @@
+import pytest
+import torch
+
+from xfuser.model_executor.layers.ltx2.na3d_eager_attn import LTX2VideoVaeEagerSdpaAttnProcessor
+from xfuser.model_executor.layers.ltx2.na3d_mfma_flash import LTX2VideoVaeMfmaAttnProcessor
+
+pytestmark = pytest.mark.rocm
+
+
+def _flash_processor() -> LTX2VideoVaeMfmaAttnProcessor:
+    if not torch.cuda.is_available() or torch.version.hip is None:
+        pytest.skip("AITER flash-NA3D requires a ROCm GPU")
+    try:
+        return LTX2VideoVaeMfmaAttnProcessor()
+    except (ImportError, RuntimeError) as exc:
+        pytest.skip(f"AITER flash-NA3D is unavailable: {exc}")
+
+
+def _attention(kernel_size: tuple[int, int, int]):
+    decoder = pytest.importorskip("diffusers.models.autoencoders.ltx2_diffusion_decoder")
+    torch.manual_seed(0)
+    return (
+        decoder.LTX2VideoVaeNeighborhoodAttention(dim=128, kernel_size=kernel_size, head_dim=64)
+        .to(device="cuda", dtype=torch.bfloat16)
+        .eval()
+    )
+
+
+def _fail_fallback(*args, **kwargs):
+    raise AssertionError("expected the AITER kernel path, not the SDPA fallback")
+
+
+@pytest.mark.parametrize(
+    ("kernel_size", "grid"),
+    [
+        ((3, 7, 7), (5, 16, 24)),
+        ((3, 5, 5), (4, 12, 16)),
+        ((11, 11, 11), (11, 12, 32)),
+    ],
+)
+@torch.no_grad()
+def test_flash_na3d_matches_tiled_sdpa(kernel_size, grid):
+    processor = _flash_processor()
+    processor._fallback = _fail_fallback
+    attn = _attention(kernel_size)
+    hidden_states = torch.randn(1, *grid, 128, device="cuda", dtype=torch.bfloat16)
+
+    output = processor(attn, hidden_states)
+    reference = LTX2VideoVaeEagerSdpaAttnProcessor()(attn, hidden_states)
+
+    assert output.shape == reference.shape
+    relative_l2 = torch.linalg.vector_norm((output - reference).float()) / torch.linalg.vector_norm(reference.float())
+    assert relative_l2 < 1e-2
+
+
+@torch.no_grad()
+def test_narrow_grid_uses_tiled_sdpa():
+    # Decoder tiling can leave remnant tiles narrower than the kernel's
+    # minimum width of 16; those must take the SDPA path instead of asserting.
+    processor = _flash_processor()
+    attn = _attention((3, 5, 5))
+    hidden_states = torch.randn(1, 4, 12, 8, 128, device="cuda", dtype=torch.bfloat16)
+
+    output = processor(attn, hidden_states)
+    reference = LTX2VideoVaeEagerSdpaAttnProcessor()(attn, hidden_states)
+
+    assert torch.equal(output, reference)

tests/unit/model_executor/layers/ltx2/test_na3d_mfma_flash.py

@@ -0,0 +1,14 @@
+import sys
+
+import pytest
+
+from xfuser.model_executor.layers.ltx2.na3d_mfma_flash import LTX2VideoVaeMfmaAttnProcessor
+
+
+def test_construction_raises_import_error_without_aiter_na3d_flash(monkeypatch):
+    # The LTX-2.5 runner falls back to tiled SDPA on ImportError, so an AITER
+    # build without na3d_flash must fail at construction, not on first decode.
+    monkeypatch.setitem(sys.modules, "aiter.ops.triton.attention.na3d_flash", None)
+
+    with pytest.raises(ImportError):
+        LTX2VideoVaeMfmaAttnProcessor()

xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+"""NA3D flash attention processor for the LTX-2.5 VAE diffusion decoder."""
+
+from __future__ import annotations
+
+import weakref
+
+import torch
+
+from xfuser.model_executor.layers.ltx2.na3d_eager_attn import (
+    LTX2VideoVaeEagerSdpaAttnProcessor,
+)
+
+# Architectures the AITER na3d_flash kernel accepts (asserted by its launcher).
+_SUPPORTED_ARCHS = ("gfx942", "gfx950")
+
+
+class LTX2VideoVaeMfmaAttnProcessor:
+    """Flash-NA3D attention processor for the LTX-2.5 diffusion decoder.
+
+    Drop-in replacement for LTX2VideoVaeEagerSdpaAttnProcessor. Delegates to
+    AITER's Triton flash-NA3D kernel (ROCm gfx950 / MI350X, gfx942 / MI300X)
+    with a fused QKV GEMM to reduce memory bandwidth. Grids outside the
+    kernel's supported domain (e.g. narrow decoder tiles) use the tiled SDPA
+    processor.
+
+    Construction raises ``ImportError`` when the installed AITER lacks
+    ``na3d_flash`` and ``RuntimeError`` on an unsupported device, so callers
+    can fall back to another processor.
+    """
+
+    def __init__(self):
+        from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn
+
+        if not torch.cuda.is_available():
+            raise RuntimeError("AITER flash-NA3D requires a ROCm device.")
+        arch = torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.split(":")[0]
+        if arch not in _SUPPORTED_ARCHS:
+            raise RuntimeError(f"AITER flash-NA3D supports {_SUPPORTED_ARCHS}; got {arch}.")
+        self._na3d_flash_attn = na3d_flash_attn
+        self._fallback = LTX2VideoVaeEagerSdpaAttnProcessor()
+        self._fused_qkv: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
+
+    @staticmethod
+    def _kernel_supports(attn, hidden_states: torch.Tensor) -> bool:
+        """Whether the grid is inside the domain the AITER launcher asserts."""
+        width = hidden_states.shape[3]
+        kernel_w = attn.kernel_size[2]
+        head_dim = attn.head_dim
+        return (
+            hidden_states.dtype == torch.bfloat16
+            and width >= 16
+            and kernel_w <= 33
+            and (kernel_w <= 17 or width >= 32)
+            and head_dim >= 16
+            and head_dim & (head_dim - 1) == 0
+        )
+
+    def _get_fused_qkv(self, attn) -> torch.nn.Linear:
+        if attn not in self._fused_qkv:
+            W_cat = torch.cat([attn.to_q.weight.data, attn.to_k.weight.data, attn.to_v.weight.data], dim=0)
+            b_cat = torch.cat([attn.to_q.bias.data, attn.to_k.bias.data, attn.to_v.bias.data], dim=0)
+            out_f, in_f = W_cat.shape
+            fused = torch.nn.Linear(in_f, out_f, bias=True, device=W_cat.device, dtype=W_cat.dtype)
+            fused.weight = torch.nn.Parameter(W_cat, requires_grad=False)
+            fused.bias = torch.nn.Parameter(b_cat, requires_grad=False)
+            self._fused_qkv[attn] = fused
+        return self._fused_qkv[attn]
+
+    def _project_qkv_fused(self, attn, hidden_states: torch.Tensor):
+        B, T, H, W, C = hidden_states.shape
+        shape = (B, T, H, W, attn.heads, attn.head_dim)
+        fused = self._get_fused_qkv(attn)
+        qkv = fused(hidden_states)
+        q_raw, k_raw, v_raw = qkv.chunk(3, dim=-1)
+        query = attn.norm_q(q_raw.view(shape))
+        key = attn.norm_k(k_raw.view(shape))
+        query = query * attn.scale
+        return attn.rope(query), attn.rope(key), v_raw.view(shape)
+
+    def __call__(self, attn, hidden_states, block_mask=None):
+        if not self._kernel_supports(attn, hidden_states):
+            return self._fallback(attn, hidden_states, block_mask)
+        B, T, H, W, _ = hidden_states.shape
+        q, k, v = self._project_qkv_fused(attn, hidden_states)
+        out = self._na3d_flash_attn(q, k, v, kernel_size=tuple(attn.kernel_size))
+        out = out.reshape(B, T, H, W, attn.heads * attn.head_dim)
+        return attn.to_out[0](out)

xfuser/model_executor/models/runner_models/ltx.py

@@ -542,15 +542,26 @@ def _load_model(self) -> DiffusionPipeline:
             diff_decoder.set_attn_processor(LTX2VideoVaeNeighborhoodNattenProcessor())
             log("Diffusion decoder: using NATTEN attention processor.")
         except (ImportError, RuntimeError, FileNotFoundError):
-            # NATTEN is not available
-            # Fall back to tiled PyTorch SDPA
-            # Works on CUDA, ROCm and CPU. Ported from LTX-2 EagerSdpaAttention.
-            from xfuser.model_executor.layers.ltx2.na3d_eager_attn import (
-                LTX2VideoVaeEagerSdpaAttnProcessor,
-            )
+            # NATTEN unavailable — try AITER's Triton flash-NA3D kernel (ROCm gfx942/gfx950).
+            try:
+                from xfuser.model_executor.layers.ltx2.na3d_mfma_flash import (
+                    LTX2VideoVaeMfmaAttnProcessor,
+                )
 
-            diff_decoder.set_attn_processor(LTX2VideoVaeEagerSdpaAttnProcessor())
-            log("Diffusion decoder: NATTEN unavailable; using tiled PyTorch SDPA fallback.")
+                processor = LTX2VideoVaeMfmaAttnProcessor()
+                log("Diffusion decoder: using AITER Triton flash-NA3D attention processor.")
+            except (ImportError, RuntimeError) as e:
+                # Tiled PyTorch SDPA (works on CUDA, ROCm, CPU).
+                from xfuser.model_executor.layers.ltx2.na3d_eager_attn import (
+                    LTX2VideoVaeEagerSdpaAttnProcessor,
+                )
+
+                processor = LTX2VideoVaeEagerSdpaAttnProcessor()
+                log(
+                    f"Diffusion decoder: NATTEN and AITER flash-NA3D unavailable ({e}); "
+                    "using tiled PyTorch SDPA fallback."
+                )
+            diff_decoder.set_attn_processor(processor)
         self.decode_pipe = LTX2VideoDiffusionDecodePipeline(
             diffusion_decoder=diff_decoder,
             scheduler=pipe.scheduler,