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

> Explore LTX-2 attention backends like FlashAttention 3/4, PyTorch SDPA, and NATTEN. Learn to select the optimal backend for your hardware and accelerate your models.

- Repository: [Lightricks/LTX-2](https://github.com/Lightricks/LTX-2)
- Tags: deep-dive
- Published: 2026-08-18

---

**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](https://github.com/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

```python

# 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`](https://github.com/Lightricks/LTX-2/blob/main/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

```python
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.

```python
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 |

```python
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):**

```python
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):**

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

```

**macOS Apple Silicon:**

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

```

**All other platforms (Ada, Ampere, CPU):**

```python
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

```python
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

```python
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`](https://github.com/Lightricks/LTX-2/blob/main/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`](https://github.com/Lightricks/LTX-2/blob/main/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`](https://github.com/Lightricks/LTX-2/blob/main/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`](https://github.com/Lightricks/LTX-2/blob/main/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`](https://github.com/Lightricks/LTX-2/blob/main/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.