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 producer, guard and layout fixes it needed.

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:
  • XLA fused RMSNorm+RoPE producer (fused_producers.py), bit-exact with nnx.RMSNorm. It is jax.jit-wrapped (heads/dim/eps static) so eager callers also get fusion — this is what fixed the wan_vace_transformer_test OOM on the v4-8 CI runner (the test calls a full 75,600-token block eagerly; unfused, the producer materialised ~1.5 GB of extra intermediates). Inlined, i.e. a no-op, under the pipeline's outer jit.
  • GQA plumbing, plus NotImplementedError on paths that don't support GQA.
  • ulysses_shards guards, dot-fallback split_head_dim reshape fix, and tile-search hardening.

Tests

• New: U=2/R=2 integration 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 attention_config_guards_test (20), custom_splash_unpadded_test (12), dot_fallback_layout_test (5), fused_producers_test (6), tile_size_grid_search_test (16).
• CI vs Local VM: 18 additional TPU kernel/multi-device tests skip under GITHUB_ACTIONS=true (36 of 115 cumulative); end_to_end/tpu/run_wan_stack_tests.sh runs all of them.
• Results (TPU v6e-8): 115 passed (run_wan_stack_tests.sh, 284.2 s).

Stack: #477 → #478 → #479 → #488 → #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