LTX-2.5 optimizations: add parallel VAE decoding and AITER RMSNorm

[jjuvonen-amd]

Summary

LTX-2.5 diffusion VAE decoder is parallelized on N-GPU runs when --use_parallel_vae is enabled, and diffuser.models.normalization.RMSNorm layers are swapped to faster AITER's CK-Tile rms_norm when available.

Changes

1. Tile-parallel diffusion VAE decode (--use_parallel_vae)
New file: xfuser/model_executor/layers/ltx2/diffusion_decoder.py

xFuserLTX2VideoDiffusionDecoderWrapper subclasses LTX2VideoDiffusionDecoderModel and overrides tiled_decode. The decoder's existing tiled_decode already decomposes the work into an independent (temporal × height × width) tile grid where each tile runs forward_stage_4 + denoise (8 neighborhood-attention blocks) with no cross-tile dependency during compute. The wrapper distributes those tiles across the SP group via round-robin ownership:

• All ranks run forward_stages_1_to_3 on the full volume (cheap; avoids communicating the intermediate feature tensor)
• Each rank computes only its assigned tiles
• One all_reduce(SUM) of a [n_tiles, 3] shape-metadata tensor, then per-tile broadcast(src=owner) to deliver every tile to all ranks
• All ranks run the stock blend/assembly loop on the gathered tiles

For the default 1024×1536×121 config this produces 12 tiles, adequate for 2/4/8 GPUs. Noise determinism: the shipping 1-step x0 path uses per-tile seeding (base_seed XOR f(t,h,w)), which is rank-invariant without requiring non-owners to know tile shapes in advance. Multi-step paths draw a shared noise canvas identically on all ranks.

Activated by --use_parallel_vae. Automatically implies enable_tiling() on the decoder (required for tiled_decode dispatch). Falls back to the stock single-GPU tiled_decode when sp_world_size == 1 or the tile count is less than 2.

Note: PR #750's DistVAE tile-parallel mechanism does not cover LTX-2.5. The diffusion decoder uses pure nn.Linear + pixel-shuffle + 3D neighborhood-attention blocks, for which DistVAE has no adapter. This wrapper is the only viable multi-GPU path for this decoder.

2. AITER block RMSNorm acceleration

New file: xfuser/model_executor/layers/norms.py

diffusers.models.normalization.RMSNorm is a Python implementation that up-casts to float32 before normalizing; AITER's CK-Tile rms_norm kernel operates natively in bfloat16. norms.py provides:

• _aiter_rms_norm: torch.library.custom_op wrapper fixing a false mutation declaration in AITER's op schema (mutates_args=()) that would otherwise cause torch.compile to skip CUDA graph capture at every call site
• _AITERRMSNorm: nn.Module drop-in for diffusers.models.normalization.RMSNorm; handles both elementwise_affine=True and =False (ones buffer for the no-weight case)
• _replace_rms_norms_with_aiter(model): walks any model and replaces every diffusers.models.normalization.RMSNorm instance with _AITERRMSNorm, leaving torch.nn.RMSNorm (the C++ implementation, already faster than AITER) untouched
• Called in xFuserLTX2VideoTransformer3DWrapper.__init__ before model.to(device) so replacement buffers are placed correctly by the device move. LTX-2.5 has 384 such block norms (8 per block × 48 blocks), each an instance of the slow Python variant. (saves ~13ms/step)

The utility is model-agnostic. Any future model using diffusers.models.normalization.RMSNorm for block norms can call _replace_rms_norms_with_aiter directly.

3. File reorganization

LTX-2.5-specific layer files moved under xfuser/model_executor/layers/ltx2/:

• layers/ltx2_na3d_eager_attn.py renamed to layers/ltx2/na3d_eager_attn.py
• New: layers/ltx2/diffusion_decoder.py

Tests

LTX-2.5 Output quality remains similar after parallel-vae or AITER RMSNorm.

Original 1-GPU (MI350)

ltx_2_5_distilled_video_u1r1_tc_True_1536x1024_i2v_mi350_aiter-varlen-attn.mp4

8-GPU (MI350) parallel-VAE, AITER RMS-Norm

mi350_ltx_2_5_distilled_video_u8r1_tc_True_1536x1024_i2v_aiter-RMSNorm-pvae.mp4

[kTorp]

nitpick but is it supposed to say "Triton" here? Other comments mentions tiled PyTorch SDPA as the fallback

[jjuvonen-amd]

Good catch, I fixed the log message to use "tiled PyTorch SDPA".