Ulysses attention: layout pin, saved exchanged Q/K/V, packed Q/K/V all-to-all (on #5779)

[khatwanimohit]

Three follow-ups to #5779 (Ulysses + KV-head replication), one commit each. All three come from diffing the compiled 60-layer Qwen3.5-397B trainer HLO (128 v7x / tpu7x-256 AOT, tunix#2884 recipe flags + the Ulysses keys below).

Commit	Change	Effect (60-layer AOT, per microbatch)
bf0abf8cd	Pin the attention block's input/output (and cotangent) layout under Ulysses	Removes the 8 transposed 512 MiB MoE-token copies per block ahead of the EP reduce-scatter; copy/transpose traffic 685 → 631 GiB
3d5b76d47	ulysses_qkv remat key: save the exchanged Q/K/V	With offload: remat-replayed all-to-alls 30 → 15, all exchanges stay async, device peak 103.35 → 99.14 GiB (+4.69 GiB host)
59904c5ab	One packed all-to-all for Q/K/V (fwd, remat and bwd)	Ulysses all-to-alls 183 → 93; bit-identical to separate exchanges

Full comparison (device peak is AOT, which overestimates runtime HBM):

Variant	Device peak GiB	All-to-alls async / sync	Remat replays	Copy GiB	Transposed MoE RS
#5721 baseline (all_gather CP + attention LB)	105.16	–	–	618	0
+ #5779 Ulysses	103.43	183 / 0	60	685	8 per block
+ layout pin + packed exchange	103.35	93 / 0	30	631	0
+ ulysses_qkv=device	112.01	48 / 30	15	645	0
+ ulysses_qkv=offload	99.14 (+4.69 host)	80 / 0	15	633	0

ulysses_qkv=device is not recommended: it costs +8.6 GiB and, with no replayed work left to overlap them with, XLA runs both backward all-to-alls synchronously.

Flags to enable

MaxText:

context_parallel_strategy=ulysses
context_parallel_load_balance=false
context_parallel_attention_load_balance=false
ulysses_qkv=offload

libtpu (trainer LIBTPU_INIT_ARGS, e.g. TRAINER_LIBTPU_INIT_ARGS in the tunix recipe; no recipe sets it today, and without it every Ulysses all-to-all is synchronous):

--xla_tpu_enable_async_all_to_all=true

The layout pin and the packed exchange need no flag; they apply whenever context_parallel_strategy=ulysses. ulysses_qkv defaults to remat, so nothing changes for other configs.

Testing

• pytest tests/unit/ulysses_kv_head_replication_test.py tests/unit/ulysses_attention_test.py tests/unit/configs_value_test.py tests/unit/attention_test.py -k "ulysses or kv_head or replication" (CPU): 35 passed, 55 subtests passed, 9 TPU-only skipped. The attention-op check runs the packed path against all_gather and a dense reference (max normalized error ≤ 3.4e-7, positive controls still fail as intended). The HLO count now expects 2 forward all-to-alls (packed Q/K/V, output) instead of 4.
• Packed vs separate exchanges on a CPU mesh: loss and dQ/dK/dV diff 0.0.
• AOT numbers above are from compiles of [Qwen3.5-397B MLPerf] Consolidated top-ROI trainer & GDN/MoE optimizations on atwigg/mlperf #5721 + Ulysses attention: repeat KV heads when there are fewer than context ranks #5779 + these commits with the tunix#2884 rendered flags. Not yet measured on TPU.

[gemini-code-assist]

Code Review

This pull request introduces optimizations for Ulysses context-parallel attention in MaxText. Key changes include implementing a packed all-to-all collective for Q, K, and V tensors to reduce communication overhead, adding a custom VJP to pin tensor layouts to row-major to prevent inefficient layout propagation in MoE layers, and updating the configuration to support rematerialization policies for the new packed QKV exchange. Unit tests were updated to reflect the reduced number of all-to-all collectives. There are no review comments to address.