From 6e976f6d0025638b626347b20f3608bd3a1eacfb Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Fri, 28 Aug 2026 07:59:43 +0000 Subject: [PATCH 1/5] faster triton NA kernel for ltx-2.5 decoder --- .../layers/ltx2/na3d_mfma_flash.py | 275 ++++++++++++++++++ .../models/runner_models/ltx.py | 27 +- 2 files changed, 292 insertions(+), 10 deletions(-) create mode 100644 xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py diff --git a/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py b/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py new file mode 100644 index 000000000..5bba5a34b --- /dev/null +++ b/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Triton flash-style 3D neighbourhood attention for LTX-2.5 VAE decoder. + +Works on AMD (MFMA via ROCm) and NVIDIA (Tensor Cores via CUDA); Triton's JIT +handles both backends. Confirmed on gfx950 (MI350X), gfx942 (MI300X), and B200. + +Optimizations over the original na3d_eager_attn.py tiled SDPA implementation: + 1. BF16 tl.dot -> hardware matrix units (MFMA on AMD, TC on NVIDIA). + 2. BLOCK_KV = next_pow2(BLOCK_Q + KW - 1): single W-chunk per iteration. + 3. KT/KH/KW as tl.constexpr → static flat loop for Triton's software pipeliner. + 4. tl.exp -> tl.exp2 (folding log2e into the running max). + 5. Fused QKV GEMM in the processor: one (SEQ,C)×(C,3C) read instead of three. + 6. @triton.autotune over BLOCK_Q, num_stages, num_warps — keyed on (KT,KH,KW,W). + Empirical optimum for LTX-2.5 decode shapes: BLOCK_Q=16, num_stages=2, + num_warps=4 (confirmed by exhaustive search; see _MFMA_CONFIGS). + 7. int64 cast on pid_bnh to prevent signed-overflow when + pid_bnh × stride_bnh exceeds 2^31 (full non-tiled 12.5M-token sequence). + +AMDGCN_USE_BUFFER_OPS: leave unset (default) to support large tensors. Set to "1" +for tiled-only workloads (tensors < 2 GB) to use hardware-bounds-checked loads. +""" +from __future__ import annotations +import math +import os +import weakref +import triton +import triton.language as tl +import torch +import torch.nn.functional as F + +_LOG2E: float = math.log2(math.e) # ≈ 1.4426950408889634 + + +def _next_pow2(n: int) -> int: + p = 1 + while p < n: + p <<= 1 + return p + + +# --------------------------------------------------------------------------- +# Autotune configs +# --------------------------------------------------------------------------- +# For all LTX-2.5 kernel widths (KW=5,7,11): +# BLOCK_Q=16 -> BLOCK_KV = next_pow2(16+KW-1) = 32 (KW≤17) +# BLOCK_Q=32 -> BLOCK_KV = next_pow2(32+KW-1) = 64 (KW≤33) +# The W >= BLOCK_Q constraint is enforced by the pruner (det shapes with W=16 +# cannot use BLOCK_Q=32). +_MFMA_CONFIGS = [ + # Autotune sweeps BLOCK_Q, num_stages, and num_warps. + # For LTX-2.5 decode shapes the winner is BLOCK_Q=16, num_stages=2, num_warps=4. + triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 2}, num_warps=4), + triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 2}, num_warps=8), + triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 3}, num_warps=4), + triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 3}, num_warps=8), + triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 4}, num_warps=4), + triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 4}, num_warps=8), + # BLOCK_Q=32: halves program count but doubles BLOCK_KV (more masked compute). + triton.Config({'BLOCK_Q': 32, 'BLOCK_KV': 64, 'num_stages': 2}, num_warps=4), + triton.Config({'BLOCK_Q': 32, 'BLOCK_KV': 64, 'num_stages': 2}, num_warps=8), +] + + +def _prune_mfma_configs(configs, named_args, **kwargs): + """Drop BLOCK_Q > W configs (W >= BLOCK_Q required for same-(t,h) guarantee).""" + W = named_args['W'] + return [c for c in configs if c.kwargs['BLOCK_Q'] <= W] + + +# --------------------------------------------------------------------------- +# Inner kernel +# --------------------------------------------------------------------------- +@triton.autotune( + configs=_MFMA_CONFIGS, + key=['KT', 'KH', 'KW', 'W'], + prune_configs_by={'early_config_prune': _prune_mfma_configs}, +) +@triton.jit +def _na3d_mfma_fwd( + Q_ptr, K_ptr, V_ptr, Out_ptr, + stride_bnh, stride_seq, + T, H, W, SEQ, + HD: tl.constexpr, + KT: tl.constexpr, + KH: tl.constexpr, + KW: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_KV: tl.constexpr, +): + """Flash-attention inner kernel. Assumes W >= BLOCK_Q (same (t,h) row per block).""" + pid_q = tl.program_id(0) + pid_bnh = tl.program_id(1) + HW = H * W + + q_start = pid_q * BLOCK_Q + q_offs = tl.arange(0, BLOCK_Q) + q_idx = q_start + q_offs + q_mask = q_idx < SEQ + + q_t = q_idx // HW + q_h = (q_idx % HW) // W + q_w = q_idx % W + + q_t_ws = tl.minimum(tl.maximum(q_t - KT // 2, 0), T - KT) + q_h_ws = tl.minimum(tl.maximum(q_h - KH // 2, 0), H - KH) + q_w_ws = tl.minimum(tl.maximum(q_w - KW // 2, 0), W - KW) + + INF_I = 999999 + t_ws = tl.min(tl.where(q_mask, q_t_ws, INF_I)) + h_ws = tl.min(tl.where(q_mask, q_h_ws, INF_I)) + w_lo = tl.min(tl.where(q_mask, q_w_ws, INF_I)) + + hd_offs = tl.arange(0, HD) + kv_offs = tl.arange(0, BLOCK_KV) + # int64 cast: pid_bnh(int32) × stride_bnh can exceed 2^31 for large sequences. + base = pid_bnh.to(tl.int64) * stride_bnh + kv_w = w_lo + kv_offs + kv_ok = kv_w < W # W-boundary guard, constant across the KT×KH loop + + # Q loaded once into registers for the entire 121-iteration loop. + Q_tile = tl.load( + Q_ptr + base + q_idx[:, None] * stride_seq + hd_offs[None, :], + mask=q_mask[:, None], other=0.0, + ).to(tl.bfloat16) + + m_i = tl.full((BLOCK_Q,), -3e38, dtype=tl.float32) + l_i = tl.zeros((BLOCK_Q,), dtype=tl.float32) + acc = tl.zeros((BLOCK_Q, HD), dtype=tl.float32) + + # Flat KT×KH loop with static trip count — enables Triton's software pipeliner. + # valid_w, kv_ok, and q_mask are loop-invariant; the Triton compiler's LICM + # handles them without explicit pre-computation (which would raise VGPR pressure + # from 71 to ~87/lane, dropping occupancy from 7 to 6 waves/SIMD). + for kv_idx in range(KT * KH): + dt = kv_idx // KH + dh = kv_idx % KH + t_kv = t_ws + dt + h_kv = h_ws + dh + + row_base = t_kv * HW + h_kv * W + kv_flat = row_base + kv_w + + K_T = tl.load( + K_ptr + base + kv_flat[None, :] * stride_seq + hd_offs[:, None], + mask=kv_ok[None, :], other=0.0, + ).to(tl.bfloat16) + + V_tile = tl.load( + V_ptr + base + kv_flat[:, None] * stride_seq + hd_offs[None, :], + mask=kv_ok[:, None], other=0.0, + ).to(tl.bfloat16) + + # QK in log2 space (fold log2e into scale to use exp2 throughout). + scores = tl.dot(Q_tile, K_T, out_dtype=tl.float32) * 1.4426950408889634 + + valid_w = ( + (kv_w[None, :] >= q_w_ws[:, None]) & + (kv_w[None, :] < q_w_ws[:, None] + KW) + ) + scores = tl.where(valid_w & kv_ok[None, :] & q_mask[:, None], + scores, float('-inf')) + + m_new = tl.maximum(m_i, tl.max(scores, axis=1)) + exp_s = tl.exp2(scores - m_new[:, None]) + exp_diff = tl.exp2(m_i - m_new) + l_i = l_i * exp_diff + tl.sum(exp_s, axis=1) + acc = acc * exp_diff[:, None] + tl.dot( + exp_s.to(tl.bfloat16), V_tile, out_dtype=tl.float32 + ) + m_i = m_new + + safe_l = tl.where(l_i > 0, l_i, 1.0) + acc = acc / safe_l[:, None] + + tl.store( + Out_ptr + base + q_idx[:, None] * stride_seq + hd_offs[None, :], + acc.to(tl.bfloat16), + mask=q_mask[:, None], + ) + + +# --------------------------------------------------------------------------- +# Python launcher +# --------------------------------------------------------------------------- +def na3d_mfma_attn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + kernel_size: tuple[int, int, int]) -> torch.Tensor: + """BF16-MFMA 3D neighbourhood flash attention with autotuned tile sizes. + + Args: + q, k, v : (B, T, H, W, NH, HD=64) bfloat16. + kernel_size : (KT, KH, KW) neighbourhood window. + Returns: + Output (B, T, H, W, NH, HD) bfloat16. + """ + B, T, H, W, NH, HD = q.shape + assert HD == 64, "head_dim must be 64" + KT, KH, KW = kernel_size + SEQ = T * H * W + + # 2 GB buffer_load guard (only relevant when AMDGCN_USE_BUFFER_OPS=1). + tensor_bytes = B * NH * SEQ * HD * 2 + if os.environ.get("AMDGCN_USE_BUFFER_OPS") == "1" and tensor_bytes >= 2**31: + raise AssertionError( + f"Tensor {tensor_bytes/1e9:.1f} GB exceeds 2 GB AMDGCN buffer_load limit. " + "Unset AMDGCN_USE_BUFFER_OPS to use predicated flat loads with no size cap." + ) + + def _flat(t: torch.Tensor) -> torch.Tensor: + return t.permute(0, 4, 1, 2, 3, 5).reshape(B * NH, SEQ, HD).contiguous() + + q_f, k_f, v_f = _flat(q), _flat(k), _flat(v) + out_f = torch.empty_like(q_f) + + # Grid uses a lambda so Triton passes the autotuned BLOCK_Q via meta. + grid = lambda meta: (triton.cdiv(SEQ, meta['BLOCK_Q']), B * NH) + + _na3d_mfma_fwd[grid]( + q_f, k_f, v_f, out_f, + SEQ * HD, HD, + T, H, W, SEQ, + HD=HD, KT=KT, KH=KH, KW=KW, + ) + return out_f.reshape(B, NH, T, H, W, HD).permute(0, 2, 3, 4, 1, 5).contiguous() + + +# --------------------------------------------------------------------------- +# Attention processor +# --------------------------------------------------------------------------- +class LTX2VideoVaeMfmaAttnProcessor: + """Triton flash-NA3D attention for the LTX-2.5 diffusion decoder. + + Drop-in replacement for LTX2VideoVaeEagerSdpaAttnProcessor. Works on AMD + (MFMA) and NVIDIA (Tensor Cores) via Triton's JIT compilation. + Improvements over the SDPA fallback: + - Hardware matrix units via tl.dot (MFMA on ROCm, TC on CUDA). + - Autotuned BLOCK_Q, num_stages, num_warps per kernel/shape combination. + - Fused QKV GEMM: reads hidden_states once instead of 3×. + """ + + def __init__(self): + self._fused_qkv: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + + 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): + B, T, H, W, _ = hidden_states.shape + q, k, v = self._project_qkv_fused(attn, hidden_states) + out = na3d_mfma_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) diff --git a/xfuser/model_executor/models/runner_models/ltx.py b/xfuser/model_executor/models/runner_models/ltx.py index ab22127c0..01ec81f2d 100644 --- a/xfuser/model_executor/models/runner_models/ltx.py +++ b/xfuser/model_executor/models/runner_models/ltx.py @@ -544,17 +544,24 @@ 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 Triton flash-NA3D (AMD MFMA / NVIDIA TC). + try: + import triton # noqa: F401 — availability gate only + from xfuser.model_executor.layers.ltx2.na3d_mfma_flash import ( + LTX2VideoVaeMfmaAttnProcessor, + ) + diff_decoder.set_attn_processor(LTX2VideoVaeMfmaAttnProcessor()) + log("Diffusion decoder: using Triton flash-NA3D attention processor.") + except (ImportError, Exception): + # Triton unavailable — tiled PyTorch SDPA (works on CUDA, ROCm, CPU). + from xfuser.model_executor.layers.ltx2.na3d_eager_attn import ( + LTX2VideoVaeEagerSdpaAttnProcessor, + ) - diff_decoder.set_attn_processor(LTX2VideoVaeEagerSdpaAttnProcessor()) - log( - "Diffusion decoder: NATTEN unavailable; using tiled PyTorch SDPA fallback." - ) + diff_decoder.set_attn_processor(LTX2VideoVaeEagerSdpaAttnProcessor()) + log( + "Diffusion decoder: NATTEN and Triton unavailable; using tiled PyTorch SDPA fallback." + ) self.decode_pipe = LTX2VideoDiffusionDecodePipeline( diffusion_decoder=diff_decoder, scheduler=pipe.scheduler, From 67f971cbebf1e7db92b35cc84c1ccb0de4c70b9f Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 12:14:15 +0000 Subject: [PATCH 2/5] import na3d_flash from aiter --- .../layers/ltx2/na3d_mfma_flash.py | 237 +----------------- 1 file changed, 7 insertions(+), 230 deletions(-) diff --git a/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py b/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py index 5bba5a34b..193db1703 100644 --- a/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py +++ b/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py @@ -1,240 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -"""Triton flash-style 3D neighbourhood attention for LTX-2.5 VAE decoder. - -Works on AMD (MFMA via ROCm) and NVIDIA (Tensor Cores via CUDA); Triton's JIT -handles both backends. Confirmed on gfx950 (MI350X), gfx942 (MI300X), and B200. - -Optimizations over the original na3d_eager_attn.py tiled SDPA implementation: - 1. BF16 tl.dot -> hardware matrix units (MFMA on AMD, TC on NVIDIA). - 2. BLOCK_KV = next_pow2(BLOCK_Q + KW - 1): single W-chunk per iteration. - 3. KT/KH/KW as tl.constexpr → static flat loop for Triton's software pipeliner. - 4. tl.exp -> tl.exp2 (folding log2e into the running max). - 5. Fused QKV GEMM in the processor: one (SEQ,C)×(C,3C) read instead of three. - 6. @triton.autotune over BLOCK_Q, num_stages, num_warps — keyed on (KT,KH,KW,W). - Empirical optimum for LTX-2.5 decode shapes: BLOCK_Q=16, num_stages=2, - num_warps=4 (confirmed by exhaustive search; see _MFMA_CONFIGS). - 7. int64 cast on pid_bnh to prevent signed-overflow when - pid_bnh × stride_bnh exceeds 2^31 (full non-tiled 12.5M-token sequence). - -AMDGCN_USE_BUFFER_OPS: leave unset (default) to support large tensors. Set to "1" -for tiled-only workloads (tensors < 2 GB) to use hardware-bounds-checked loads. -""" +"""NA3D flash attention processor for the LTX-2.5 VAE diffusion decoder.""" from __future__ import annotations -import math -import os import weakref -import triton -import triton.language as tl import torch -import torch.nn.functional as F - -_LOG2E: float = math.log2(math.e) # ≈ 1.4426950408889634 - - -def _next_pow2(n: int) -> int: - p = 1 - while p < n: - p <<= 1 - return p - - -# --------------------------------------------------------------------------- -# Autotune configs -# --------------------------------------------------------------------------- -# For all LTX-2.5 kernel widths (KW=5,7,11): -# BLOCK_Q=16 -> BLOCK_KV = next_pow2(16+KW-1) = 32 (KW≤17) -# BLOCK_Q=32 -> BLOCK_KV = next_pow2(32+KW-1) = 64 (KW≤33) -# The W >= BLOCK_Q constraint is enforced by the pruner (det shapes with W=16 -# cannot use BLOCK_Q=32). -_MFMA_CONFIGS = [ - # Autotune sweeps BLOCK_Q, num_stages, and num_warps. - # For LTX-2.5 decode shapes the winner is BLOCK_Q=16, num_stages=2, num_warps=4. - triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 2}, num_warps=4), - triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 2}, num_warps=8), - triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 3}, num_warps=4), - triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 3}, num_warps=8), - triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 4}, num_warps=4), - triton.Config({'BLOCK_Q': 16, 'BLOCK_KV': 32, 'num_stages': 4}, num_warps=8), - # BLOCK_Q=32: halves program count but doubles BLOCK_KV (more masked compute). - triton.Config({'BLOCK_Q': 32, 'BLOCK_KV': 64, 'num_stages': 2}, num_warps=4), - triton.Config({'BLOCK_Q': 32, 'BLOCK_KV': 64, 'num_stages': 2}, num_warps=8), -] - - -def _prune_mfma_configs(configs, named_args, **kwargs): - """Drop BLOCK_Q > W configs (W >= BLOCK_Q required for same-(t,h) guarantee).""" - W = named_args['W'] - return [c for c in configs if c.kwargs['BLOCK_Q'] <= W] - - -# --------------------------------------------------------------------------- -# Inner kernel -# --------------------------------------------------------------------------- -@triton.autotune( - configs=_MFMA_CONFIGS, - key=['KT', 'KH', 'KW', 'W'], - prune_configs_by={'early_config_prune': _prune_mfma_configs}, -) -@triton.jit -def _na3d_mfma_fwd( - Q_ptr, K_ptr, V_ptr, Out_ptr, - stride_bnh, stride_seq, - T, H, W, SEQ, - HD: tl.constexpr, - KT: tl.constexpr, - KH: tl.constexpr, - KW: tl.constexpr, - BLOCK_Q: tl.constexpr, - BLOCK_KV: tl.constexpr, -): - """Flash-attention inner kernel. Assumes W >= BLOCK_Q (same (t,h) row per block).""" - pid_q = tl.program_id(0) - pid_bnh = tl.program_id(1) - HW = H * W - - q_start = pid_q * BLOCK_Q - q_offs = tl.arange(0, BLOCK_Q) - q_idx = q_start + q_offs - q_mask = q_idx < SEQ - - q_t = q_idx // HW - q_h = (q_idx % HW) // W - q_w = q_idx % W - - q_t_ws = tl.minimum(tl.maximum(q_t - KT // 2, 0), T - KT) - q_h_ws = tl.minimum(tl.maximum(q_h - KH // 2, 0), H - KH) - q_w_ws = tl.minimum(tl.maximum(q_w - KW // 2, 0), W - KW) - - INF_I = 999999 - t_ws = tl.min(tl.where(q_mask, q_t_ws, INF_I)) - h_ws = tl.min(tl.where(q_mask, q_h_ws, INF_I)) - w_lo = tl.min(tl.where(q_mask, q_w_ws, INF_I)) - - hd_offs = tl.arange(0, HD) - kv_offs = tl.arange(0, BLOCK_KV) - # int64 cast: pid_bnh(int32) × stride_bnh can exceed 2^31 for large sequences. - base = pid_bnh.to(tl.int64) * stride_bnh - kv_w = w_lo + kv_offs - kv_ok = kv_w < W # W-boundary guard, constant across the KT×KH loop - - # Q loaded once into registers for the entire 121-iteration loop. - Q_tile = tl.load( - Q_ptr + base + q_idx[:, None] * stride_seq + hd_offs[None, :], - mask=q_mask[:, None], other=0.0, - ).to(tl.bfloat16) - - m_i = tl.full((BLOCK_Q,), -3e38, dtype=tl.float32) - l_i = tl.zeros((BLOCK_Q,), dtype=tl.float32) - acc = tl.zeros((BLOCK_Q, HD), dtype=tl.float32) - - # Flat KT×KH loop with static trip count — enables Triton's software pipeliner. - # valid_w, kv_ok, and q_mask are loop-invariant; the Triton compiler's LICM - # handles them without explicit pre-computation (which would raise VGPR pressure - # from 71 to ~87/lane, dropping occupancy from 7 to 6 waves/SIMD). - for kv_idx in range(KT * KH): - dt = kv_idx // KH - dh = kv_idx % KH - t_kv = t_ws + dt - h_kv = h_ws + dh - - row_base = t_kv * HW + h_kv * W - kv_flat = row_base + kv_w - - K_T = tl.load( - K_ptr + base + kv_flat[None, :] * stride_seq + hd_offs[:, None], - mask=kv_ok[None, :], other=0.0, - ).to(tl.bfloat16) - - V_tile = tl.load( - V_ptr + base + kv_flat[:, None] * stride_seq + hd_offs[None, :], - mask=kv_ok[:, None], other=0.0, - ).to(tl.bfloat16) - - # QK in log2 space (fold log2e into scale to use exp2 throughout). - scores = tl.dot(Q_tile, K_T, out_dtype=tl.float32) * 1.4426950408889634 - - valid_w = ( - (kv_w[None, :] >= q_w_ws[:, None]) & - (kv_w[None, :] < q_w_ws[:, None] + KW) - ) - scores = tl.where(valid_w & kv_ok[None, :] & q_mask[:, None], - scores, float('-inf')) - - m_new = tl.maximum(m_i, tl.max(scores, axis=1)) - exp_s = tl.exp2(scores - m_new[:, None]) - exp_diff = tl.exp2(m_i - m_new) - l_i = l_i * exp_diff + tl.sum(exp_s, axis=1) - acc = acc * exp_diff[:, None] + tl.dot( - exp_s.to(tl.bfloat16), V_tile, out_dtype=tl.float32 - ) - m_i = m_new - - safe_l = tl.where(l_i > 0, l_i, 1.0) - acc = acc / safe_l[:, None] - - tl.store( - Out_ptr + base + q_idx[:, None] * stride_seq + hd_offs[None, :], - acc.to(tl.bfloat16), - mask=q_mask[:, None], - ) - - -# --------------------------------------------------------------------------- -# Python launcher -# --------------------------------------------------------------------------- -def na3d_mfma_attn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - kernel_size: tuple[int, int, int]) -> torch.Tensor: - """BF16-MFMA 3D neighbourhood flash attention with autotuned tile sizes. - - Args: - q, k, v : (B, T, H, W, NH, HD=64) bfloat16. - kernel_size : (KT, KH, KW) neighbourhood window. - Returns: - Output (B, T, H, W, NH, HD) bfloat16. - """ - B, T, H, W, NH, HD = q.shape - assert HD == 64, "head_dim must be 64" - KT, KH, KW = kernel_size - SEQ = T * H * W - - # 2 GB buffer_load guard (only relevant when AMDGCN_USE_BUFFER_OPS=1). - tensor_bytes = B * NH * SEQ * HD * 2 - if os.environ.get("AMDGCN_USE_BUFFER_OPS") == "1" and tensor_bytes >= 2**31: - raise AssertionError( - f"Tensor {tensor_bytes/1e9:.1f} GB exceeds 2 GB AMDGCN buffer_load limit. " - "Unset AMDGCN_USE_BUFFER_OPS to use predicated flat loads with no size cap." - ) - - def _flat(t: torch.Tensor) -> torch.Tensor: - return t.permute(0, 4, 1, 2, 3, 5).reshape(B * NH, SEQ, HD).contiguous() - - q_f, k_f, v_f = _flat(q), _flat(k), _flat(v) - out_f = torch.empty_like(q_f) - - # Grid uses a lambda so Triton passes the autotuned BLOCK_Q via meta. - grid = lambda meta: (triton.cdiv(SEQ, meta['BLOCK_Q']), B * NH) - - _na3d_mfma_fwd[grid]( - q_f, k_f, v_f, out_f, - SEQ * HD, HD, - T, H, W, SEQ, - HD=HD, KT=KT, KH=KH, KW=KW, - ) - return out_f.reshape(B, NH, T, H, W, HD).permute(0, 2, 3, 4, 1, 5).contiguous() -# --------------------------------------------------------------------------- -# Attention processor -# --------------------------------------------------------------------------- class LTX2VideoVaeMfmaAttnProcessor: - """Triton flash-NA3D attention for the LTX-2.5 diffusion decoder. + """Flash-NA3D attention processor for the LTX-2.5 diffusion decoder. - Drop-in replacement for LTX2VideoVaeEagerSdpaAttnProcessor. Works on AMD - (MFMA) and NVIDIA (Tensor Cores) via Triton's JIT compilation. - Improvements over the SDPA fallback: - - Hardware matrix units via tl.dot (MFMA on ROCm, TC on CUDA). - - Autotuned BLOCK_Q, num_stages, num_warps per kernel/shape combination. - - Fused QKV GEMM: reads hidden_states once instead of 3×. + Drop-in replacement for LTX2VideoVaeEagerSdpaAttnProcessor. Uses AITER's + BF16-MFMA flash kernel (gfx950 / MI350X, gfx942 / MI300X, B200) and + includes a fused QKV GEMM to reduce memory bandwidth. """ def __init__(self): @@ -270,6 +46,7 @@ def _project_qkv_fused(self, attn, hidden_states: torch.Tensor): def __call__(self, attn, hidden_states, block_mask=None): B, T, H, W, _ = hidden_states.shape q, k, v = self._project_qkv_fused(attn, hidden_states) - out = na3d_mfma_attn(q, k, v, kernel_size=tuple(attn.kernel_size)) + from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn + out = 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) From 5d5b4184e51cbf5473d22245f3b97fdd37bf5200 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 7 Sep 2026 05:49:24 +0000 Subject: [PATCH 3/5] LTX2VideoVaeMfmaAttnProcessor docstring clarification --- xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py b/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py index 193db1703..1c70bda2c 100644 --- a/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py +++ b/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py @@ -8,9 +8,9 @@ class LTX2VideoVaeMfmaAttnProcessor: """Flash-NA3D attention processor for the LTX-2.5 diffusion decoder. - Drop-in replacement for LTX2VideoVaeEagerSdpaAttnProcessor. Uses AITER's - BF16-MFMA flash kernel (gfx950 / MI350X, gfx942 / MI300X, B200) and - includes a fused QKV GEMM to reduce memory bandwidth. + Drop-in replacement for LTX2VideoVaeEagerSdpaAttnProcessor. Delegates to + AITER's BF16-MFMA flash kernel (gfx950 / MI350X, gfx942 / MI300X) with a + fused QKV GEMM to reduce memory bandwidth. """ def __init__(self): From ad20e4704abe72ed949fcd44e509b225e0bc2e7c Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 5 Oct 2026 09:55:14 +0000 Subject: [PATCH 4/5] style: ruff-format runner_models/ltx.py and fix E402 The PR lint job checks whole changed files, and ltx.py was not ruff-clean on main. Move the cache import into the top import block and apply ruff-format (v0.11.13, as pinned in .pre-commit-config.yaml). No behavior change (AST identical apart from the moved import). --- .../models/runner_models/ltx.py | 92 ++++++------------- 1 file changed, 29 insertions(+), 63 deletions(-) diff --git a/xfuser/model_executor/models/runner_models/ltx.py b/xfuser/model_executor/models/runner_models/ltx.py index b28c63878..b8eb0b8c0 100644 --- a/xfuser/model_executor/models/runner_models/ltx.py +++ b/xfuser/model_executor/models/runner_models/ltx.py @@ -9,6 +9,11 @@ from xfuser.core.utils.runner_utils import log from xfuser.core.utils.video_utils import encode_video_with_audio from xfuser.envs import PACKAGES_CHECKER +from xfuser.model_executor.cache import ( + DBCachePreset, + CacheDitAdapterConfig, + DBCacheSettings, +) from xfuser.model_executor.models.runner_models.base_model import ( DIFFUSERS_FROM_SOURCE, DefaultInputValues, @@ -37,11 +42,6 @@ "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, " "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." ) -from xfuser.model_executor.cache import ( - DBCachePreset, - CacheDitAdapterConfig, - DBCacheSettings, -) @register_model("dg845/LTX-2.3-Diffusers") @@ -117,9 +117,7 @@ def _load_model(self) -> DiffusionPipeline: xFuserLTX2VideoTransformer3DWrapper, ) - transformer = self.loader.load_transformer( - xFuserLTX2VideoTransformer3DWrapper - ) + transformer = self.loader.load_transformer(xFuserLTX2VideoTransformer3DWrapper) pipe = LTX2Pipeline.from_pretrained( pretrained_model_name_or_path=self.settings.model_name, @@ -146,9 +144,7 @@ def _load_model(self) -> DiffusionPipeline: subfolder="latent_upsampler", torch_dtype=torch.bfloat16, ) - upsample_pipe = LTX2LatentUpsamplePipeline( - vae=pipe.vae, latent_upsampler=latent_upsampler - ) + upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler) log("Enabling VAE tiling for LTX-2.3's full-resolution second-stage decode.") second_pipe.vae.enable_tiling() @@ -191,9 +187,7 @@ def _run_pipe(self, input_args: dict) -> DiffusionOutput: return_dict=False, ) - video_latent = self.upsample_pipe( - latents=video_latent, output_type="latent", return_dict=False - )[0] + video_latent = self.upsample_pipe(latents=video_latent, output_type="latent", return_dict=False)[0] self.second_pipe.transformer.enable_adapters() @@ -236,9 +230,7 @@ def _compile_model(self, input_args: dict) -> None: # two steps to warmup the torch compiler compile_args = copy.deepcopy(input_args) - compile_args["num_inference_steps"] = ( - 2 # Reduce steps for warmup # TODO: make this more generic - ) + compile_args["num_inference_steps"] = 2 # Reduce steps for warmup # TODO: make this more generic self._run_compile_warmup(compile_args) def save_output(self, output: DiffusionOutput) -> None: @@ -264,6 +256,8 @@ def _post_load_and_state_initialization(self, input_args: dict) -> None: super()._post_load_and_state_initialization(input_args) self.upsample_pipe.to(self.pipe.device) self.second_pipe.to(self.pipe.device) + + @register_model("Lightricks/LTX-2") @register_model("LTX-2") class xFuserLTX2VideoModel(xFuserModel): @@ -327,9 +321,7 @@ def _load_model(self) -> DiffusionPipeline: xFuserLTX2VideoTransformer3DWrapper, ) - transformer = self.loader.load_transformer( - xFuserLTX2VideoTransformer3DWrapper - ) + transformer = self.loader.load_transformer(xFuserLTX2VideoTransformer3DWrapper) pipe = LTX2Pipeline.from_pretrained( pretrained_model_name_or_path=self.settings.model_name, transformer=transformer, @@ -350,14 +342,10 @@ def _load_model(self) -> DiffusionPipeline: subfolder="latent_upsampler", torch_dtype=torch.bfloat16, ) - upsample_pipe = LTX2LatentUpsamplePipeline( - vae=pipe.vae, latent_upsampler=latent_upsampler - ) + upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler) - second_pipe.scheduler = ( - FlowMatchEulerDiscreteScheduler.from_config( # Scheduler for the 2nd stage - pipe.scheduler.config, use_dynamic_shifting=False, shift_terminal=None - ) + second_pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config( # Scheduler for the 2nd stage + pipe.scheduler.config, use_dynamic_shifting=False, shift_terminal=None ) self.second_pipe = second_pipe self.upsample_pipe = upsample_pipe @@ -382,9 +370,7 @@ def _run_pipe(self, input_args: dict) -> DiffusionOutput: generator=self._make_generator(input_args["seed"]), ) - video_latent = self.upsample_pipe( - latents=video_latent, output_type="latent", return_dict=False - )[0] + video_latent = self.upsample_pipe(latents=video_latent, output_type="latent", return_dict=False)[0] output = self.second_pipe( latents=video_latent, @@ -411,9 +397,7 @@ def _compile_model(self, input_args: dict) -> None: # two steps to warmup the torch compiler compile_args = copy.deepcopy(input_args) - compile_args["num_inference_steps"] = ( - 2 # Reduce steps for warmup # TODO: make this more generic - ) + compile_args["num_inference_steps"] = 2 # Reduce steps for warmup # TODO: make this more generic self._run_compile_warmup(compile_args) def save_output(self, output: DiffusionOutput) -> None: @@ -493,6 +477,7 @@ def _load_model(self) -> DiffusionPipeline: from xfuser.model_executor.models.transformers.transformer_ltx2 import ( xFuserLTX2VideoTransformer3DWrapper, ) + transformer_cls = xFuserLTX2VideoTransformer3DWrapper transformer = transformer_cls.from_pretrained( @@ -501,9 +486,7 @@ def _load_model(self) -> DiffusionPipeline: subfolder=self._TRANSFORMER_SUBFOLDER, ) - pipe_cls = ( - LTX2ImageToVideoPipeline if self.config.task == "i2v" else LTX2Pipeline - ) + pipe_cls = LTX2ImageToVideoPipeline if self.config.task == "i2v" else LTX2Pipeline pipe = pipe_cls.from_pretrained( self.settings.model_name, transformer=transformer, @@ -524,9 +507,7 @@ def _load_model(self) -> DiffusionPipeline: subfolder="latent_upsampler", torch_dtype=torch.bfloat16, ) - self.upsample_pipe = LTX2LatentUpsamplePipeline( - vae=pipe.vae, latent_upsampler=latent_upsampler - ) + self.upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler) # Diffusion decoder — replaces convolutional VAE decode for both distilled # and full model pipelines. @@ -557,6 +538,7 @@ def _load_model(self) -> DiffusionPipeline: from xfuser.model_executor.layers.ltx2.na3d_mfma_flash import ( LTX2VideoVaeMfmaAttnProcessor, ) + diff_decoder.set_attn_processor(LTX2VideoVaeMfmaAttnProcessor()) log("Diffusion decoder: using Triton flash-NA3D attention processor.") except (ImportError, Exception): @@ -566,9 +548,7 @@ def _load_model(self) -> DiffusionPipeline: ) diff_decoder.set_attn_processor(LTX2VideoVaeEagerSdpaAttnProcessor()) - log( - "Diffusion decoder: NATTEN and Triton unavailable; using tiled PyTorch SDPA fallback." - ) + log("Diffusion decoder: NATTEN and Triton unavailable; using tiled PyTorch SDPA fallback.") self.decode_pipe = LTX2VideoDiffusionDecodePipeline( diffusion_decoder=diff_decoder, scheduler=pipe.scheduler, @@ -600,21 +580,16 @@ def _validate_args(self, input_args: dict) -> None: if self.config.task == "i2v": images = input_args.get("input_images") or [] if len(images) != 1: - raise ValueError( - f"LTX-2.5 I2V requires exactly one input image, got {len(images)}." - ) + raise ValueError(f"LTX-2.5 I2V requires exactly one input image, got {len(images)}.") if self._DISTILLED: steps = input_args.get("num_inference_steps") if steps != 8: raise ValueError( - f"LTX-2.5 distilled uses a fixed 8-step schedule; " - f"num_inference_steps must be 8, got {steps}." + f"LTX-2.5 distilled uses a fixed 8-step schedule; num_inference_steps must be 8, got {steps}." ) guidance_scale = input_args.get("guidance_scale") if guidance_scale != 1.0: - log( - "Using guidance_scale=1.0. Other guidance scale values are not supported with this model." - ) + log("Using guidance_scale=1.0. Other guidance scale values are not supported with this model.") def _run_pipe(self, input_args: dict) -> DiffusionOutput: from diffusers.pipelines.ltx2.utils import ( @@ -651,9 +626,7 @@ def _run_pipe(self, input_args: dict) -> DiffusionOutput: **shared, ) - video_latent = self.upsample_pipe( - latents=video_latent, output_type="latent", return_dict=False - )[0] + video_latent = self.upsample_pipe(latents=video_latent, output_type="latent", return_dict=False)[0] # Stage-2: guidance disabled. stage2_shared = { @@ -688,9 +661,7 @@ def _run_pipe(self, input_args: dict) -> DiffusionOutput: shared["use_cross_timestep"] = True shared["negative_prompt"] = input_args["negative_prompt"] if self._SPATIO_TEMPORAL_GUIDANCE_BLOCKS is not None: - shared["spatio_temporal_guidance_blocks"] = ( - self._SPATIO_TEMPORAL_GUIDANCE_BLOCKS - ) + shared["spatio_temporal_guidance_blocks"] = self._SPATIO_TEMPORAL_GUIDANCE_BLOCKS video_latents, audio_latents = self.pipe( height=input_args["height"], width=input_args["width"], @@ -712,9 +683,7 @@ def _run_pipe(self, input_args: dict) -> DiffusionOutput: # Audio: pipeline returned raw audio latents when output_type="latent". # Decode via audio_vae -> mel -> vocoder -> waveform. - mel = self.pipe.audio_vae.decode( - audio_latents.to(self.pipe.audio_vae.dtype), return_dict=False - )[0] + mel = self.pipe.audio_vae.decode(audio_latents.to(self.pipe.audio_vae.dtype), return_dict=False)[0] audio_waveform = self.pipe.vocoder(mel) output = SimpleNamespace(frames=video_np, audio=audio_waveform) @@ -743,9 +712,7 @@ def _compile_model(self, input_args: dict) -> None: def _mark_cudagraph_step(module, args, kwargs): torch.compiler.cudagraph_mark_step_begin() - self.pipe.transformer.register_forward_pre_hook( - _mark_cudagraph_step, with_kwargs=True, prepend=True - ) + self.pipe.transformer.register_forward_pre_hook(_mark_cudagraph_step, with_kwargs=True, prepend=True) compile_args = copy.deepcopy(input_args) compile_args["num_inference_steps"] = 2 @@ -829,7 +796,6 @@ class xFuserLTX25FullVideoModel(_xFuserLTX25VideoModelBase): _TRANSFORMER_SUBFOLDER = "transformer_full" _DISTILLED = False # single-stage, no upsampler - # Full-model guidance parameters - LTX-2.4/2.5 params: # video: cfg_scale=3.0, stg_scale=1.0, rescale_scale=0.7, modality_scale=3.0, stg_blocks=[28] # audio: cfg_scale=7.0, stg_scale=1.0, rescale_scale=0.7, modality_scale=3.0, stg_blocks=[28] From 71e902bdffd6af3dc22fdedc7e874bbb5f339150 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 5 Oct 2026 09:57:37 +0000 Subject: [PATCH 5/5] LTX-2.5 decoder: harden AITER flash-NA3D fallback, add tests - Resolve the aiter na3d_flash import and check for a gfx942/gfx950 device when the processor is constructed, so the runner falls back to tiled SDPA instead of failing on the first decode call. - Route grids outside the kernel's asserted domain (bf16, W >= 16, KW <= 33, KW > 17 needs W >= 32, power-of-two head_dim) to the tiled SDPA processor; decoder tiling can produce remnant tiles as narrow as the kernel. - Runner: drop the triton import gate, catch only ImportError and RuntimeError, and log why the SDPA fallback was chosen. - Tests: CPU contract for the ImportError fallback; ROCm parity against the tiled SDPA processor for the LTX-2.5 kernel sizes, plus the narrow-tile path. --- .../layers/test_ltx2_na3d_flash.py | 67 ++++++++++++++++++ .../layers/ltx2/test_na3d_mfma_flash.py | 14 ++++ .../layers/ltx2/na3d_mfma_flash.py | 68 ++++++++++++++----- .../models/runner_models/ltx.py | 19 +++--- 4 files changed, 144 insertions(+), 24 deletions(-) create mode 100644 tests/integration/accelerator/model_executor/layers/test_ltx2_na3d_flash.py create mode 100644 tests/unit/model_executor/layers/ltx2/test_na3d_mfma_flash.py diff --git a/tests/integration/accelerator/model_executor/layers/test_ltx2_na3d_flash.py b/tests/integration/accelerator/model_executor/layers/test_ltx2_na3d_flash.py new file mode 100644 index 000000000..0b9848643 --- /dev/null +++ b/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) diff --git a/tests/unit/model_executor/layers/ltx2/test_na3d_mfma_flash.py b/tests/unit/model_executor/layers/ltx2/test_na3d_mfma_flash.py new file mode 100644 index 000000000..82a8e7117 --- /dev/null +++ b/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() diff --git a/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py b/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py index 1c70bda2c..fba42d8d8 100644 --- a/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py +++ b/xfuser/model_executor/layers/ltx2/na3d_mfma_flash.py @@ -1,34 +1,69 @@ # 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 BF16-MFMA flash kernel (gfx950 / MI350X, gfx942 / MI300X) with a - fused QKV GEMM to reduce memory bandwidth. + 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) + 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 = 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) + fused.bias = torch.nn.Parameter(b_cat, requires_grad=False) self._fused_qkv[attn] = fused return self._fused_qkv[attn] @@ -36,17 +71,18 @@ 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) + 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)) + 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) - from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn - out = na3d_flash_attn(q, k, v, kernel_size=tuple(attn.kernel_size)) + 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) diff --git a/xfuser/model_executor/models/runner_models/ltx.py b/xfuser/model_executor/models/runner_models/ltx.py index b8eb0b8c0..a5400fb7b 100644 --- a/xfuser/model_executor/models/runner_models/ltx.py +++ b/xfuser/model_executor/models/runner_models/ltx.py @@ -532,23 +532,26 @@ def _load_model(self) -> DiffusionPipeline: diff_decoder.set_attn_processor(LTX2VideoVaeNeighborhoodNattenProcessor()) log("Diffusion decoder: using NATTEN attention processor.") except (ImportError, RuntimeError, FileNotFoundError): - # NATTEN unavailable — try Triton flash-NA3D (AMD MFMA / NVIDIA TC). + # NATTEN unavailable — try AITER's Triton flash-NA3D kernel (ROCm gfx942/gfx950). try: - import triton # noqa: F401 — availability gate only from xfuser.model_executor.layers.ltx2.na3d_mfma_flash import ( LTX2VideoVaeMfmaAttnProcessor, ) - diff_decoder.set_attn_processor(LTX2VideoVaeMfmaAttnProcessor()) - log("Diffusion decoder: using Triton flash-NA3D attention processor.") - except (ImportError, Exception): - # Triton unavailable — tiled PyTorch SDPA (works on CUDA, ROCm, CPU). + 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, ) - diff_decoder.set_attn_processor(LTX2VideoVaeEagerSdpaAttnProcessor()) - log("Diffusion decoder: NATTEN and Triton unavailable; using tiled PyTorch SDPA fallback.") + 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,