LTX-2 Attention Backends: FlashAttention 3/4, PyTorch SDPA, and NATTEN Explained

LTX-2 supports FlashAttention 3 and 4, PyTorch SDPA (with CUDNN/FLASH/EFFICIENT variants), MPS-SDPA for Apple Silicon, and NATTEN for 3-D neighborhood attention, with automatic hardware-aware selection and manual override options.

This guide covers every attention backend available in the Lightricks/LTX-2 repository, how the library automatically selects the optimal kernel for your hardware, and how to manually configure backends when needed.

FlashAttention 3 and 4 Support

LTX-2 integrates the latest FlashAttention implementations for maximum throughput on NVIDIA Hopper and Blackwell GPUs.

FlashAttention 3

Availability: CUDA-only, requires flash-attn-3 package (flash_attn_interface)

Best for: NVIDIA H100 (sm_90) GPUs running unmasked attention

# Auto-selected on Hopper when package is installed

from ltx_core.model.transformer.attention import Attention, AttentionOps, AttentionFunction

ops = AttentionOps(
    attention_function=AttentionFunction.FLASH_ATTENTION_3.to_callable(),
)
attn = Attention(query_dim=512, heads=8, ops=ops)

As implemented in packages/ltx-core/src/ltx_core/model/transformer/attention.py (lines 50-60), this backend delivers peak memory efficiency but does not support attention masks.

FlashAttention 4

Availability: CUDA-only, requires flash-attn-4 package (flash_attn.cute.flash_attn_func)

Best for: NVIDIA B200 (sm_100) GPUs and forward-looking deployments

from ltx_core.model.transformer.attention import Attention, AttentionOps, AttentionFunction

ops = AttentionOps(
    attention_function=AttentionFunction.FLASH_ATTENTION_4.to_callable(),
)
attn = Attention(query_dim=512, heads=8, ops=ops)

The FlashAttention4 class (lines 70-80) represents the bleeding-edge kernel. When available on Blackwell hardware, LTX-2 automatically prefers this over FlashAttention 3.

PyTorch SDPA Backend Variants

PyTorch's Scaled Dot-Product Attention (SDPA) serves as the default and most portable backend, with multiple kernel variants.

Variant	Hardware	Mask Support	Selection Criteria
SDPA CUDNN	CUDA only	Yes	torch.backends.cuda.cudnn_attention for compatible shapes
SDPA FLASH	CUDA only	No	Built-in torch flash kernel when mask not required
SDPA EFFICIENT	CUDA only	Yes	torch.backends.cuda.efficient_attention when supported
SDPA MATH	CPU/CUDA/MPS	Yes	Pure-Python fallback, always works

The _resolve_sdpa_variant function (lines 104-128) dispatches to the fastest available variant based on tensor shapes, whether a mask is present, and PyTorch's internal capability checks.

SDPA FLASH vs. FlashAttention 3/4

• SDPA FLASH: PyTorch's built-in kernel, always available on CUDA, no extra package needed
• FlashAttention 3/4: Standalone packages with optimized cuDNN integration, higher peak performance

Both require unmasked attention; when masks are needed, LTX-2 automatically falls back to SDPA CUDNN or EFFICIENT.

MPS-SDPA for Apple Silicon

macOS deployments receive hardware-optimized attention through the MPS-SDPA backend.

Availability: macOS Apple Silicon, requires mps-sdpa package

Advantage: Fused MPSGraph kernel with no score-matrix materialization, reducing memory pressure on unified memory architectures.

from ltx_core.model.transformer.attention import Attention, AttentionOps, MaskedAttentionFunction

ops = AttentionOps(
    masked_attention_function=MaskedAttentionFunction.MPS_SDPA.to_callable(),
)
attn = Attention(query_dim=256, heads=8, ops=ops)

The MPSSdpaAttention class (lines 100-118) wraps the mps-sdpa callable. If the package is unavailable, the automatic selector falls back to standard SDPA.

NATTEN for 3-D Neighborhood Attention

Video VAE workloads use Neighborhood Attention to restrict attention to local spatiotemporal patches.

Backend Options

Implementation	Package Required	Performance
NATTEN	natten (optional extra)	Fastest, optimized CUDA kernels
Triton fallback	None	Good, no installation needed
Eager SDPA	None	Slowest, reference implementation

from ltx_core.model.video_vae.transformer.attention import NeighborhoodAttention3D

# Automatic: NATTEN if installed, else Triton/SDPA fallback

na = NeighborhoodAttention3D(dim=256, kernel_size=(3, 3, 3))

# Manual override to Triton (no natten dependency)

from ltx_kernels.vae.na_attn_dsl import TritonNAAttention
na.attention_function = TritonNAAttention()

The NattenAttention class (lines 53-83) handles the 3-D windowing logic, with automatic fallback and clear warnings via warn_no_natten when the optional dependency is missing.

Automatic Backend Selection

LTX-2 caches the optimal backend choice per-process using _select_primary_attention and _select_masked_attention (cached with @functools.cache).

Hardware-Aware Priority Order

NVIDIA Hopper (sm_9, H100):

if major == 9:
    if flash_attn_interface is not None:   # FA3

        return FlashAttention3()
    if flash_attn_4_func is not None:      # FA4

        return FlashAttention4()

NVIDIA Blackwell (sm_10, B200):

if major == 10 and flash_attn_4_func is not None:
    return FlashAttention4()

macOS Apple Silicon:

if _on_macos():
    return MPSSdpaAttention() if _mps_sdpa_available() else _sdpa_full_priority()

All other platforms (Ada, Ampere, CPU):

return _sdpa_full_priority()  # PyTorch's internal dispatcher

Masked Attention Path

For attention requiring additive masks:

• macOS: MPS-SDPA if available
• All platforms: Full-priority SDPA dispatcher (CUDNN → EFFICIENT → MATH)

FlashAttention kernels are excluded from masked paths by design.

Manual Backend Configuration

Override automatic selection using the AttentionFunction and MaskedAttentionFunction enums.

Force Specific Unmasked Backend

from ltx_core.model.transformer.attention import (
    Attention,
    AttentionOps,
    AttentionFunction,
    MaskedAttentionFunction,
)

ops = AttentionOps(
    attention_function=AttentionFunction.FLASH_ATTENTION_4.to_callable(),
    masked_attention_function=MaskedAttentionFunction.AUTOMATIC.to_callable(),
)
attn = Attention(query_dim=256, heads=8, ops=ops)

Force PyTorch SDPA Everywhere

ops = AttentionOps(
    attention_function=AttentionFunction.PYTORCH.to_callable(),
    masked_attention_function=MaskedAttentionFunction.PYTORCH.to_callable(),
)

Enum members raise RuntimeError immediately if the requested backend is unavailable, preventing silent degradation.

Choosing the Best Backend for Your Hardware

Hardware	Recommended Unmasked	Recommended Masked	Installation Notes
NVIDIA H100	FlashAttention 3	SDPA CUDNN	pip install flash-attn-3
NVIDIA B200	FlashAttention 4	SDPA CUDNN	pip install flash-attn-4
NVIDIA RTX 4090/3090	SDPA FLASH	SDPA EFFICIENT	None (built into PyTorch)
Apple M1/M2/M3	MPS-SDPA	MPS-SDPA	pip install mps-sdpa
CPU-only	SDPA MATH	SDPA MATH	None

For reproducibility across different machines, pin your backend explicitly. For maximum performance, rely on AttentionFunction.AUTOMATIC and MaskedAttentionFunction.AUTOMATIC.

Key Source Files

Path	Purpose
packages/ltx-core/src/ltx_core/model/transformer/attention.py	Core backends, enums, and selection logic (lines 50-128)
packages/ltx-core/src/ltx_core/model/video_vae/transformer/attention.py	3-D Neighborhood Attention implementation
packages/ltx-core/src/ltx_core/model/transformer/ops.py	Low-level attention callable interfaces
packages/ltx-kernels/src/ltx_kernels/vae/na_attn_dsl.py	Triton fallback for neighborhood attention

Summary

• LTX-2 supports nine distinct attention backends spanning FlashAttention 3/4, four PyTorch SDPA variants, MPS-SDPA, and NATTEN with fallbacks
• Automatic selection in attention.py (lines 84-106) optimizes for GPU architecture and mask requirements without user intervention
• Manual override via AttentionFunction and MaskedAttentionFunction enums enables reproducibility and debugging
• FlashAttention 3/4 and NATTEN require optional packages; the library degrades gracefully to PyTorch-native implementations when unavailable
• MPS-SDPA delivers optimized inference on Apple Silicon, while SDPA MATH provides cross-platform compatibility

Frequently Asked Questions

How do I check which attention backend LTX-2 is actually using?

Inspect the AttentionOps object after construction or enable PyTorch logging. The automatic selector caches its choice, so the same backend persists for the process lifetime. For explicit verification, force a specific backend and catch RuntimeError if unavailable.

Why doesn't FlashAttention work with attention masks?

FlashAttention 3 and 4 optimize memory access by fusing the attention computation into a single kernel that doesn't materialize the full score matrix. This design inherently prevents applying additive masks. LTX-2 automatically routes masked calls to SDPA variants that support masking.

Can I use LTX-2 on AMD GPUs or Intel Arc?

The PyTorch SDPA backend runs on any hardware PyTorch supports, but optimized kernels (FlashAttention, CUDNN) are CUDA-only. Intel and AMD GPUs use the SDPA MATH fallback by default. Community Triton kernels may offer acceleration on alternative hardware.

What's the performance difference between NATTEN and its Triton fallback?

NATTEN's optimized CUDA kernels typically achieve 2-3× higher throughput than the Triton fallback on video VAE workloads. However, the Triton implementation requires no extra dependencies and maintains correctness, making it suitable for development or when compilation fails.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}