[WAN 2.2] 2D Ring Attention with Custom Kernel

[eltsai]

Summary

In this PR, we are adding support for

1. Ring attention with custom kernel (like Fix: Plumb vmem limit and bkv_compute_in to the custom kernel #416 we plumbed the vmem limits into the kernel)
2. 2D ring attention (ring + ulysses) with custom kernel
3. Bidirectional ring attention

(also formatting src/maxdiffusion/loaders/ltx2_lora_nnx_loader.py to remove formatting errors)

Results

Compared with previous best config

For dp2-cp4 on v7x-8, we see:

==================================================
  TIMING SUMMARY (2D Ring, U=2, R=2)
==================================================
  Load (checkpoint):     585.3s
  Compile:               144.3s
  Inference:             116.0s
  ────────────────────────────────────────
  Conditioning:            3.9s
    - VAE Encode:          0.0s
  Denoise Total:         109.8s
  VAE Decode:              2.2s
    - TPU Compute:         1.7s
    - Host Formatting:     0.5s
==================================================

which is a 11% speedup comparing best Ulysses config for the denoising step (9% e2e time):

==================================================
  TIMING SUMMARY (Ulysses)
==================================================
  Load (checkpoint):     141.2s
  Compile:               183.3s
  Inference:             126.9s
  ────────────────────────────────────────
  Conditioning:            2.2s
    - VAE Encode:          0.0s
  Denoise Total:         122.6s
  VAE Decode:              2.1s
    - TPU Compute:         1.7s
    - Host Formatting:     0.4s
==================================================

Best Configs

The 2D-ring (Ulysses×Ring) hybrid is optimal at cp4 and cp16, and ties pure Ulysses at cp8. Pure Ulysses never strictly wins.

• At CP4, 2D-ring (U=2, R=2, BQ=9472) (2.74 s/step) beats the best pure Ulysses config (U=4, BQ=8448) (3.06 s/step) by ~10%.
• At CP8, 2D-ring (U=4, R=2, BQ=9472) (1.595 sec/step) ties the best pure Ulysses config (U=8, BQ=8448) (1.590 sec/step).
• At CP16 a 2D ring wins (and is the only option, pure Ulysses U=16 impossible because num_head=40). The best is (U=8, R=2, BQ=9472).

Compared with ring attention

1. Using the custom kernel as base for ring attention gives a ~23% per-denoising-step speedup (cp=4: 4.85 → 3.74 s/step, BQ=4096).

2. A tile-size (BQ) search found the optimal ring tile is BQ=9472 for R<=8, and BQ=4096 for R>=16. It's worth ~19% speedup (dp1-cp8 U=4: 1.98→1.60 s/step at BQ 4096→9472).

For more details, see the internal doc

Note on bi-directional (wrap-free) ring attention.

Because of TPU topology, when R_degree >= 8 we start to see slow downs due to communication. I did some PoC on bi-directional ring, and see improvements for some cases when R_degree=16. I think it would be more beneficial when we scale up, so I am including it in this PR (despite that it doesn't give us optimal results on CP4, CP8 and CP16).

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[eltsai]

Pasting more results for CP4, CP8 and CP16:

dp1 means that the CFG is set as False (the stats for dp1-cp16 is basically what we will see on dp2-cp16 on v7x-32 with CFG enabled).

[Perseus14]

Thanks for the amazing work!

I have dropped a few comments for minor changes. PTAL!

Additionally in another PR, we should probably add a doc to the repo that talks about all the different attention kernels and where to use what.

[github-actions]

🤖 Hi @csgoogle, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @csgoogle, but I was unable to process your request. Please see the logs for more details.

[csgoogle]

nit: can we comment the best flashblock sizes here itself, it would be easy for anypne to look into