perf(attention): shard-major Ulysses output all-to-all (wan_ulysses_out_a2a="chunked"/"shard_major")

[Perseus14]

Stacked on feat/wan-token-patch-embed. Third of four pieces split out of #491. Adds wan_ulysses_out_a2a: "chunked" | "shard_major" (default "flat"): _ulysses_seq_to_heads performs the inverse Ulysses exchange as a single layout-preserving all_to_all over a reshaped (num_shards, chunk) leading axis, so XLA stops inserting the four relayout copies per layer that the flat split_axis=seq form costs at the 720p shard shape (~1.5 ms on v6e, ~0.7 ms on tpu7x). Same data movement, same values. With the custom splash kernel (S_local % block_q == 0, D == 128, heads_per_tile == 1) the forward Q exchange also stays shard-major (4-D q) and the kernel writes the shard-major output directly (out_num_shards), on both the 1D Ulysses path (v6e) and the 2D Ulysses+Ring path (tpu7x): custom_splash_attention.py gains a 4-D q BlockSpec / 4-D output, ring_attention_kernel.py threads out_num_shards through the ring accumulators (save_residuals=True with out_num_shards > 1 raises NotImplementedError), and _compute_fixed_m_metadata accepts the 5-D shard-major query.

Reachability note: The Pallas shard-major Q/out path requires S_local % block_q == 0. At 720p/81f with CP=4 the per-device sequence is 18,900 tokens, which no launcher block_q divides (BQ=9472 on v6e, BQ=6400 on v7 at this PR), so on this PR alone the launcher's chunked setting only changes the inverse all-to-all; the 4-D q / out_num_shards Pallas kernel path engages once #491 (feat/wan-fast-inference-optimizations) pads the sequence to 18,944 (wan_seq_pad=lane) and sets BQ=4736 (4736 × 4 = 18,944).

E2E latency (TPU v6e-8, Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps, CP=4, DP=2, wan_ulysses_out_a2a="chunked", DVFS unpinned): 128.9 s generate / 126.7 s denoise (29.6 s compile), down from 130.6 s generate / 128.5 s denoise on #510 (−1.7 s generate / −1.8 s denoise end-to-end from the layout-preserving inverse Ulysses all-to-all alone).

Tests: ulysses_out_a2a_test.py (6, new; flat vs chunked vs shard_major parity, mode guard, save_residuals + shard-major guard, bit-exact 4-D ring kernel output, and bit-exact wrapper output on the 1D and 2D paths asserting via spy that out_num_shards == U), wan_runtime_options_test.py (+enum validation).

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 the wan_ulysses_out_a2a configuration option (supporting "flat", "chunked", and "shard_major" modes) to optimize the inverse Ulysses exchange and avoid expensive post-kernel relayout copies on TPUs. This is achieved by updating the custom splash and ring attention kernels to support 4D shard-major query inputs and output sharding, alongside adding a new test suite to verify correctness. The review feedback highlights a critical bug in _ulysses_ring_custom_attention where use_shard_major fails to check for transpose_out, potentially causing layout mismatches. Additionally, the reviewer suggests using the more idiomatic jnp.square instead of ** 2 in JAX operations and wrapping a long line in ring_attention_kernel.py to comply with PEP 8 line length guidelines.