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

[jjuvonen-amd]

Summary

Adds an LTX-2.5 diffusion-decoder attention processor (LTX2VideoVaeMfmaAttnProcessor) that uses AITER's Triton 3D neighborhood flash attention (aiter.ops.triton.attention.na3d_flash, added in ROCm/aiter#5230, now merged), with a fused QKV projection.

The decoder now picks its attention processor in this order: NATTEN -> AITER flash-NA3D (ROCm gfx942/gfx950) -> tiled PyTorch SDPA.

• The processor checks for the AITER kernel and the GPU architecture when it is created. An AITER build without na3d_flash raises ImportError, and an unsupported device raises RuntimeError. The runner catches both and falls back to tiled SDPA, logging the reason.
• Calls with shapes the kernel doesn't support (non-bf16, W < 16, KW > 33, KW > 17 with W < 32, or a head_dim that isn't a power of two ≥ 16) go to the tiled SDPA processor. Decoder tiling can leave edge tiles as narrow as the kernel size, so this case does occur.
• runner_models/ltx.py is reformatted with ruff in a separate, style-only commit, because the lint job checks whole changed files and this file was not ruff-clean on main. It moves one import to fix E402; there is no behavior change.

Requires an AITER build that includes ROCm/aiter#5230. Without it, the decoder falls back to tiled SDPA as before.

Validation

Automated lint, formatting, CPU unit, and CPU integration checks run in PR CI. They do not cover accelerator behavior.

New tests:

• tests/unit/model_executor/layers/ltx2/test_na3d_mfma_flash.py: constructing the processor raises ImportError when AITER lacks na3d_flash. The runner's fallback depends on this.
• tests/integration/accelerator/model_executor/layers/test_ltx2_na3d_flash.py (rocm):
  • Compares the AITER output against the tiled SDPA processor for the LTX-2.5 window sizes (3,7,7), (3,5,5) and (11,11,11). Relative L2 error is about 1.3e-3; the threshold is 1e-2.
  • Checks that narrow grids (W < 16) take the SDPA path and match its output exactly.

Environment: 8x AMD Instinct MI350X (gfx950), ROCm 7.14, PyTorch 2.9.1, Triton 3.5.1, AITER main @ 7bfe279de (includes #5230), diffusers 0.41.0.dev0 @ 9f7aee482.

Suite	Command	Result
Lint	pre-commit run --from-ref <base> --to-ref HEAD	ruff and ruff-format passed
CPU unit	python -m pytest tests/unit -m "not accelerator"	896 passed, 18 deselected
CPU integration	python -m pytest tests/integration/cpu -m "not accelerator"	7 passed, 2 skipped, 3 deselected

The 2 CPU-integration skips are environmental: one because real GPUs are present, one because the installed diffusers is not a final release.

Accelerator tests

• NVIDIA tests passed
• AMD/ROCm tests passed
• Multi-GPU tests passed
• Not applicable (explain below)

Hardware, environment, command, and result:

• 1x MI350X: python -m pytest -m "accelerator and not nvidia and not multi_gpu" -> 94 passed, 2 skipped, 2 xfailed, 922 deselected
• 8x MI350X: python -m pytest -m "accelerator and multi_gpu and not nvidia" -> 4 passed, 1 skipped, 1015 deselected

End-to-end tests

• Appropriate model/pipeline E2E tests passed, including output-quality review
• Performance and memory were checked for kernel, distributed, or inference-path changes
• Not applicable (explain below)

Model, command, result, and any relevant performance or quality comparison:

Lightricks/LTX-2.5-Diffusers (distilled), i2v, 1536x1024, 121 frames, 8 steps, seed 42, MI350X:

# 1 GPU
torchrun --nproc_per_node=1 -m xfuser.runner --model Lightricks/LTX-2.5-Diffusers --task i2v \
  --input_images <image> --resize_input_images --height 1536 --width 1024 --num_frames 121 \
  --num_inference_steps 8 --seed 42 --use_torch_compile --attention_backend AITER \
  --num_iterations 2 --enable_tiling --prompt "<prompt>"
# 8 GPU: replace --enable_tiling with --ulysses_degree 8 --use_parallel_vae and use --nproc_per_node=8

The tiled SDPA baseline ran on the same branch with AITER's na3d_flash hidden, which forces the runner's fallback. Decode time is measured around decode_pipe(...) with a CUDA sync on each side, averaged over 2 iterations after warm-up (slowest rank on 8 GPUs).

Config	Tiled SDPA	AITER flash-NA3D	Speedup
1 GPU + tiling, decode	17,050 ms	9,145 ms	1.87x
1 GPU + tiling, end-to-end	23.74 s	15.87 s	1.50x
8 GPU (Ulysses 8 + parallel VAE), decode	4,116 ms	2,400 ms	1.71x
8 GPU, end-to-end	6.68 s	4.75 s	1.41x

Output quality: I visually compared frames from the 1-GPU and 8-GPU outputs against earlier reference videos. They show no tile seams or artifacts (videos attached below).

Tiled SDPA fallback

mi350_ltx_2_5_distilled_video_u8r1_tc_True_1536x1024_i2v_aiter-RMSNorm-pvae.mp4

With AITER na3d_flash kernel

ltx_2_5_distilled_video_u1r1_tc_True_1536x1024_i2v_0_triton_na_mi355x.mp4

Notes

• NVIDIA tests were not run because the AITER kernel is ROCm-only (gfx942/gfx950); on other devices the processor fails at construction and the existing NATTEN -> tiled SDPA path is unchanged.
• I checked performance but not peak memory. The fused QKV projection keeps a concatenated copy of each decoder attention block's Q/K/V weights alongside the originals. That's an estimated ~150 MB of bf16 weights for the default decoder config; I calculated this rather than measuring it.
• Validated on gfx950 only, but the kernel is validated to work on gfx942 (MI300X).
• The decode timings include the decoder's non-attention stages, so the speedup for attention alone is larger than the decode numbers above.

[avjves]

Hi, sorry the review took a while. This looks good to me! Can you fix the conflicts in ltx.py? :)

[jjuvonen-amd]

Hi, sorry the review took a while. This looks good to me! Can you fix the conflicts in ltx.py? :)

Thanks @avjves, the conflicts are now resolved and this merges cleanly.