Ulysses attention: repeat KV heads when there are fewer than context ranks

[parambole]

Summary

Allow Ulysses context parallelism (context_parallel_strategy: "ulysses") when num_kv_heads < ici_context_parallelism and ici_context_parallelism % num_kv_heads == 0 (such as Qwen3.5-397B-A17B with num_query_heads=32, num_kv_heads=2 at ici_context_parallelism=4 on 128 chips or ici_context_parallelism=8 on 256 chips) by repeating each KV head cp // num_kv_heads times in grouped order ([kv0, kv0, kv1, kv1] at cp=4; [kv0, kv0, kv0, kv0, kv1, kv1, kv1, kv1] at cp=8) before the Ulysses all-to-all.

Why

With context_parallel_strategy: "all_gather" and context_parallel_attention_load_balance: true on packed 65k sequences:

1. Each of the 15 full-attention layers executes DUAL_CHUNK_SWAP reorder_sequence collective-permutes on entry and exit (~200.3 ms/mb exposed across 15 layers) to keep GDN and MoE in contiguous sequence order.
2. On long packed segments, interior CP ranks (cp1..cp3) execute 136 fwd / 132 dkv Splash tiles (~1,162.5 ms/mb) vs ~80 tiles on cp0, causing fast chips (TPU:0, cp0 on a light row) to stall for ~749 ms/mb at the subsequent 256-way FSDP weight all-gather barriers (shared_expert wi_0/wi_1 in fwd, attention wq in bwd).

With context_parallel_strategy: "ulysses" and context_parallel_load_balance: false (context_parallel_attention_load_balance: false):

• Grouped KV-head replication (jnp.repeat along head axis 2):
  • At cp=4 (128 chips, 256 TensorCores), replication = 4 // 2 = 2: KV heads [kv0, kv1] repeat in place to [kv0, kv0, kv1, kv1]. After the 4-way all-to-all, CP rank r holds query heads [8r .. 8r+7] (8 Q heads) and KV head r // 2 (1 KV head) across the full 65,536-token sequence.
  • At cp=8 (256 chips, 512 TensorCores), replication = 8 // 2 = 4: KV heads [kv0, kv1] repeat in place to [kv0, kv0, kv0, kv0, kv1, kv1, kv1, kv1]. After the 8-way all-to-all, CP rank r holds query heads [4r .. 4r+3] (4 Q heads) and KV head r // 4 (1 KV head) across the full 65,536-token sequence.
  • Inside _ulysses_Custom's custom_vjp backward, dk and dv are scaled by 1 / replication before the reverse all-to-all, and autodiff through jnp.repeat sums the replication copies back onto each unique KV head, matching all_gather CP gradients.
• GatedDeltaNet (GDN) layers: Keep their existing contiguous sequence-parallel CP (context_parallel_load_balance=false), requiring zero sequence reordering between Attention, GDN, and MoE.

Full-Cell 60-Layer AOT HLO Verification (128-chip v7x topology fsdp=32, cp=4, ep=8)

Metric (per 4-layer scan body)	all_gather + attn_lb=true (#5758)	ulysses + lb=false (this PR)	Delta
Attention K3 load-balance collective-permutes (64 MiB each)	8 (200.3 ms/mb)	0	-100%
Attention K/V all-gathers (4-way, 64 MiB each)	4	0 (4-way all-to-all)	-100%
Backward dK/dV reduce-scatters (4-way, 32 MiB each)	2	0	-100%
Causal Splash tiles per Q head (cp0..cp3 on 61k segment)	80 (cp0) vs 136 (cp1..cp3)	~80 on all 4 CP ranks	0 CP rank skew
Backward wq 256-way AG compute instructions between start/done	18	79	4.4× overlap
Train-kernel device peak HBM	104.46 GiB	104.25 GiB	-0.21 GiB

Tests

• 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:
  • Helper unit tests for replication factor, grouped order ([kv0, kv0, kv1, kv1] vs tiled [kv0, kv1, kv0, kv1]), query/KV pairing at cp=4 and cp=8, and validator flags.
  • 4-device CPU numerical equivalence of AttentionOp.tpu_flash_attention (ulysses vs all_gather vs single-device reference): forward max |diff| = 0.0 (bit-exact), backward dQ/dK/dV max |diff| <= 1.66e-7.
  • Negative controls verifying that tiled head order (np.tile) or omitting the 1 / replication backward factor fail as expected.

[gemini-code-assist]

Code Review

This pull request introduces support for Ulysses context parallelism when there are fewer KV heads than context ranks (such as 2 KV heads with a context parallel size of 4 or 8, which is common in models like Qwen 3.5). This is achieved by repeating each KV head across the corresponding context ranks before the Ulysses all-to-all exchange. The changes span configuration validation, head sharding validation, and the attention operator, and are accompanied by comprehensive unit tests verifying correctness and gradient matching. I have no feedback to provide as there are no review comments to evaluate.