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

Third of four pieces split out of the Wan fast-inference optimizations PR.
The switch defaults off ("flat") in the Wan YAMLs and in wan_runtime_options;
it is a plain config key plumbed to FlaxWanAttention / NNXAttentionOp /
AttentionOp and down to the Ulysses kernels through the dispatch context.

- `_ulysses_seq_to_heads` is the inverse Ulysses exchange (attention output,
  head-sharded -> sequence-sharded). "flat" is the previous
  `all_to_all(split_axis=seq, concat_axis=1)`; XLA reorders the whole output
  so the sequence axis is major before the collective and back again after,
  four relayout copies per layer at the 720p shard shape (~1.5 ms on v6e,
  ~0.7 ms on tpu7x). "chunked" / "shard_major" is a single layout-preserving
  all_to_all over a reshaped (num_shards, chunk) leading axis (not chunked
  compute/comm overlap), so no relayout copies are inserted. Same data
  movement, same values.
- With the custom splash kernel (and S_local % block_q == 0, D == 128,
  heads_per_tile == 1) the forward Q exchange also stays shard-major (4-D q,
  `[B, U, H/U, S_local, D]`) 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 and a 4-D shard-major output for `_splash_attention_forward` and
  `_splash_attention_forward_ring`; ring_attention_kernel.py threads
  `out_num_shards` through `_custom_ring_attention_forward` /
  `make_custom_ring_attention` and sizes the o/l/m accumulators accordingly
  (`save_residuals` with `out_num_shards > 1` raises NotImplementedError:
  the lse contract stays `(num_q_heads, orig_q_seq_len)`).
  `_compute_fixed_m_metadata` accepts the 5-D shard-major query.
- Reachability: 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 (v6e BQ=9472, v7 BQ=6400
  at this PR), so on this PR alone the launcher's `chunked` only changes the
  inverse all-to-all; the 4-D q / `out_num_shards` kernel path engages once
  feat/wan-fast-inference-optimizations pads the sequence to 18,944
  (`wan_seq_pad=lane`) and sets BQ=4736 (4736 x 4 = 18,944). The rough
  1-1.5 s per 40-step 720p run measured for the kernel path on a pre-review
  revision was therefore taken with that padding active (DVFS state
  unrecorded; not re-measured since).
- run_wan_fast_inference.sh: the v6e and v7 profiles set
  WAN_ULYSSES_OUT_A2A=chunked (overridable); the generic profile keeps flat.

Tests
- ulysses_out_a2a_test.py (6): flat vs chunked vs shard_major parity of the
  inverse exchange, the mode guard, the save_residuals + shard-major guard,
  the 4-D q ring kernel bit-exact vs 3-D (interpret mode off TPU, compiled on
  TPU), and bit-exact wrapper output for the shard-major 1D Ulysses and 2D
  Ulysses+Ring paths vs the flat layout -- both wrapper tests spy on
  make_splash_mha / make_custom_ring_attention and assert out_num_shards == U,
  i.e. that the Pallas shard-major path was actually taken.
- wan_runtime_options_test.py (12): the new switch's YAML default and enum
  validation.
- CI: the three kernel bit-exact tests skip in GitHub Actions.
- end_to_end/tpu/run_wan_stack_tests.sh: adds ulysses_out_a2a_test.py.
Verified (this head): TPU v6e-8 (jax 0.11.2), every file of
end_to_end/tpu/run_wan_stack_tests.sh at this head run directly with unittest:
338 tests, all OK, 1 skipped (cross_attention_pallas_test: off-TPU fallback only);
ulysses_out_a2a_test.py 6 (all 4 multi-device tests ran), wan_runtime_options_test.py 12.

Author: Perseus14

end_to_end/tpu/run_wan_fast_inference.sh

@@ -63,7 +63,7 @@
 #                    against the XLA-path video, visually clean); on tpu7x the
 #                    output was bit-identical in that measurement.
 #                    USE_FUSED_ROPE_KERNEL=false restores the XLA producer.
-#   WAN_CROSS_ATTN_KERNEL / WAN_PATCH_EMBED_MODE
+#   WAN_CROSS_ATTN_KERNEL / WAN_PATCH_EMBED_MODE / WAN_ULYSSES_OUT_A2A
 #                    override Wan fast-inference kernel/layout switches
 #   DP / CP / PER_DEVICE_BATCH / SEED
 #                    override mesh parallelism, per-device batch, or RNG seed (default 12345)
@@ -171,6 +171,7 @@ case "$TPU_PROFILE" in
     DEFAULT_FUSED_ROPE=true
     DEFAULT_CROSS_ATTN_KERNEL=pallas
     DEFAULT_PATCH_EMBED_MODE=tokens
+    DEFAULT_ULYSSES_OUT_A2A=chunked
     ;;
   v7)
     PLATFORM_LIBTPU="$V7_LIBTPU"
@@ -189,6 +190,7 @@ case "$TPU_PROFILE" in
     DEFAULT_FUSED_ROPE=true
     DEFAULT_CROSS_ATTN_KERNEL=pallas
     DEFAULT_PATCH_EMBED_MODE=tokens
+    DEFAULT_ULYSSES_OUT_A2A=chunked
     ;;
   *)
     echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \
@@ -209,6 +211,7 @@ case "$TPU_PROFILE" in
     DEFAULT_FUSED_ROPE=false
     DEFAULT_CROSS_ATTN_KERNEL=xla
     DEFAULT_PATCH_EMBED_MODE=tokens
+    DEFAULT_ULYSSES_OUT_A2A=flat
     ;;
 esac
 
@@ -242,6 +245,7 @@ FUSED_ROPE_ARGS=()
 [ -n "${FUSED_ROPE_HEAD_BLOCK:-}" ] && FUSED_ROPE_ARGS+=("fused_rope_head_block=$FUSED_ROPE_HEAD_BLOCK")
 WAN_CROSS_ATTN_KERNEL=${WAN_CROSS_ATTN_KERNEL:-$DEFAULT_CROSS_ATTN_KERNEL}
 WAN_PATCH_EMBED_MODE=${WAN_PATCH_EMBED_MODE:-$DEFAULT_PATCH_EMBED_MODE}
+WAN_ULYSSES_OUT_A2A=${WAN_ULYSSES_OUT_A2A:-$DEFAULT_ULYSSES_OUT_A2A}
 
 # Mesh: context parallelism carries the Ulysses shards, data parallelism takes
 # whatever chips remain. Defaults to CP=4 / DP=2 on an 8-chip slice. On a slice
@@ -306,6 +310,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \
   ${FUSED_ROPE_ARGS[@]+"${FUSED_ROPE_ARGS[@]}"} \
   wan_cross_attn_kernel="$WAN_CROSS_ATTN_KERNEL" \
   wan_patch_embed_mode="$WAN_PATCH_EMBED_MODE" \
+  wan_ulysses_out_a2a="$WAN_ULYSSES_OUT_A2A" \
   fps=16 ${GUIDANCE_ARGS[@]+"${GUIDANCE_ARGS[@]}"} \
   seed="${SEED:-12345}" \
   flash_block_sizes="$FLASH_BLOCK_SIZES" \

end_to_end/tpu/run_wan_stack_tests.sh

@@ -48,5 +48,7 @@ TESTS=(
   "$T/cross_attention_pallas_test.py"
   # Token-sharded patch embedding (feat/wan-token-patch-embed)
   "$T/wan_patch_embed_test.py"
+  # Shard-major Ulysses output A2A (feat/ulysses-shard-major-a2a)
+  "$T/ulysses_out_a2a_test.py"
 )
 PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@"

src/maxdiffusion/configs/base_wan_14b.yml

@@ -143,6 +143,11 @@ wan_cross_attn_cpu_interpret: False
 # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
 # replicated kernel, no activation collective; equal up to accumulation order).
 wan_patch_embed_mode: "conv"
+# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded):
+# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit
+# leading sequence-chunk axis so XLA does not relayout the output around the
+# collective (used by the custom Ulysses / Ulysses-ring kernels only).
+wan_ulysses_out_a2a: "flat"
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -140,6 +140,11 @@ wan_cross_attn_cpu_interpret: False
 # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
 # replicated kernel, no activation collective; equal up to accumulation order).
 wan_patch_embed_mode: "conv"
+# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded):
+# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit
+# leading sequence-chunk axis so XLA does not relayout the output around the
+# collective (used by the custom Ulysses / Ulysses-ring kernels only).
+wan_ulysses_out_a2a: "flat"
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.

src/maxdiffusion/configs/base_wan_27b.yml

@@ -172,6 +172,11 @@ wan_cross_attn_cpu_interpret: False
 # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
 # replicated kernel, no activation collective; equal up to accumulation order).
 wan_patch_embed_mode: "conv"
+# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded):
+# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit
+# leading sequence-chunk axis so XLA does not relayout the output around the
+# collective (used by the custom Ulysses / Ulysses-ring kernels only).
+wan_ulysses_out_a2a: "flat"
 # For attention=ulysses_ring*, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.

src/maxdiffusion/configs/base_wan_animate.yml

@@ -122,6 +122,11 @@ wan_cross_attn_cpu_interpret: False
 # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
 # replicated kernel, no activation collective; equal up to accumulation order).
 wan_patch_embed_mode: "conv"
+# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded):
+# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit
+# leading sequence-chunk axis so XLA does not relayout the output around the
+# collective (used by the custom Ulysses / Ulysses-ring kernels only).
+wan_ulysses_out_a2a: "flat"
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -124,6 +124,11 @@ wan_cross_attn_cpu_interpret: False
 # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
 # replicated kernel, no activation collective; equal up to accumulation order).
 wan_patch_embed_mode: "conv"
+# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded):
+# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit
+# leading sequence-chunk axis so XLA does not relayout the output around the
+# collective (used by the custom Ulysses / Ulysses-ring kernels only).
+wan_ulysses_out_a2a: "flat"
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -124,6 +124,11 @@ wan_cross_attn_cpu_interpret: False
 # token, an all-gather on tpu7x) or "tokens" (token-sharded patches times the
 # replicated kernel, no activation collective; equal up to accumulation order).
 wan_patch_embed_mode: "conv"
+# Inverse Ulysses all-to-all (attention output, head-sharded -> sequence-sharded):
+# "flat" splits the sequence axis directly; "chunked" exchanges over an explicit
+# leading sequence-chunk axis so XLA does not relayout the output around the
+# collective (used by the custom Ulysses / Ulysses-ring kernels only).
+wan_ulysses_out_a2a: "flat"
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
 # Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder.