feat(attention): 2D Ulysses+Ring custom attention with fixed-m (pre-a2a norm reduction, unpadded ring K/V)

[Perseus14]

Summary

Adds fixed-m to the 2D Ulysses + Ring custom-kernel path (ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block), plus the supporting guards and GQA plumbing it needed. (The dot-product fallback layout fix and the tile_size_grid_search guards were split into feat/attn-dot-fallback-fixes below this PR; the XLA fused RMSNorm+RoPE producer was split into feat/fused-rmsnorm-rope-xla above it.)

What's in it

• Pre-all-to-all norm reduction: Q row norms, K max norms and the V max are reduced in one pmax over (ulysses, ring) on the pre-all-to-all shards, so they overlap the collective.
  • v_ok comes from that pmax, so it is identical on every rank.
  • all_fixed_global is mesh-uniform and drives the accumulate-vs-LSE lax.cond.
  • Blast radius: one ineligible head or Q-block on any device sends every rank to _lse_scan for that layer.
• Unpadded ring K/V: passed unpadded when the ring-shard KV length is 8-aligned (kv_pad_size=1). Q·log2(e) is folded in before the all-to-all.
• Virtual K-centering: works for R == 1 and R > 1. k_mean is computed as pmean over (ulysses, ring) and passed to the kernel; K is never rewritten.

  • Off by default on ring. Measured at the stack tip (Wan 2.2 720p/81f/40 steps, warm AOT, DVFS unpinned), it is never faster and it changes the bf16 output:

    TPU	Ring setup	Off	On	Delta
    v6e-8	R == 1	127.2 s	127.2 s	neutral
    v6e-8	R = 2	128.6 s	129.1 s	+0.4%
    tpu7x-8	U=2/R=2	107.66 s	108.04 s	+0.35%

  • use_k_centering=True forces it on.

• Other changes:
  • GQA plumbing through the custom Ulysses / ring paths, plus NotImplementedError on paths that don't support GQA.
  • ulysses_shards guards: non-ring Ulysses kernels reject a ulysses_shards they cannot honour; ring kernels warn once when U == CP makes the ring degenerate (R = 1) and suggest a valid U for the mesh and head count.

Performance

Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps (CP=4, DP=2, BQ=9472, BKV=4736, BKV_C=1024, DVFS unpinned) on TPU v6e-8, measured at this PR:

Config	Revision	Attention kernel	Generate	Denoise	Delta vs main
2D Ulysses+Ring (U=2, R=2)	main (f1d0fafb)	ulysses_ring_custom_fixed_m	155.0 s	152.5 s	—
2D Ulysses+Ring (U=2, R=2)	#478 (feat/ring-attention)	ulysses_ring_custom_fixed_m_per_q_block	138.8 s	136.6 s	−16.2 s generate / −15.9 s denoise (−10.4%)
1D Ulysses (U=4)	#478 (feat/ring-attention)	ulysses_custom_fixed_m_per_q_block	138.3 s	136.1 s	−8.8 s generate / −8.5 s denoise (−5.9%)

Tests

• New / updated: U=2/R=2 integration in ring_fixed_m_test.py (24 total) against a dense f32 reference (per_q_block × centering × GQA, ragged KV, sink head forcing _lse_scan, V outlier forcing v_ok=False, and B=2 _accumulate_scan vmap + d=64 k_mean head-dim padding), plus custom_splash_fixed_m_test.py (32), attention_config_guards_test.py (20), and custom_splash_unpadded_test.py (12, TPU-only).
• CI vs Local VM: 18 additional TPU kernel/multi-device tests skip under GITHUB_ACTIONS=true (36 of 110 cumulative); end_to_end/tpu/run_wan_stack_tests.sh runs all of them.
• Results (TPU v6e-8): 110 tests, all OK across the 6 cumulative stack test files (custom_splash_fixed_m 32, ring_fixed_m 24, dot_fallback_layout 6, tile_size_grid_search 16, attention_config_guards 20, custom_splash_unpadded 12).

Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491

[github-actions]

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

[gemini-code-assist]

Code Review

This pull request introduces optimized fused producers (fused_ln_adaln and fused_rmsnorm_rope) and implements exact Fixed-m support with Global Virtual K-Centering for Ulysses and Ring attention. It also adds support for Grouped Query Attention (GQA) across these attention kernels and includes comprehensive unit tests. The review feedback suggests moving an inline import of fused_rmsnorm_rope in attention_flax.py to the top of the file to avoid performance overhead in a hot path, and simplifying a double-negation conditional expression to improve code readability.

[eltsai]

for the e2e verification, did you check the output video quality as well?

[Perseus14]

Yes, I verified manually as well as both pixel-level fidelity (PSNR / SSIM against exact FP32 online softmax) and full VBench quality metrics on Wan 2.2 T2V-A14B (1280x720, 81 frames, 40 steps). I have updated the results in #491