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 (`"flat"` | `"chunked"` | `"shard_major"`) 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), and validates `seq_len % num_shards == 0`
  up front for all modes. `"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"` uses a single
  layout-preserving `all_to_all` over a reshaped `(num_shards, chunk)` leading
  axis, so no relayout copies are inserted (also wired into the non-custom
  `attention="ulysses"` branch). 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,
  `[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.
- `"shard_major"` vs `"chunked"`: `"shard_major"` is the strict mode that
  requires the 4-D `q` / `out_num_shards` Pallas kernel path and raises
  `ValueError` when any gate precondition (`use_custom_kernel`,
  `S_local % block_q == 0`, `D == 128`, `heads_per_tile == 1`,
  `not transpose_out`, `attention_mask is None`) fails; `"chunked"` is the
  lenient mode that takes the 4-D kernel path when eligible and falls back to
  the chunked inverse `all_to_all` in `_ulysses_seq_to_heads` otherwise.
- Reachability: at 720p/81f with CP=4 the unpadded per-device sequence is
  18,900 tokens, which no launcher `block_q` divides (v6e `BQ=9472`, v7
  `BQ=6400` at this PR), so at this PR the launcher's `v6e`/`v7` profiles set
  `DEFAULT_ULYSSES_OUT_A2A=chunked`; once `feat/wan-fast-inference-optimizations`
  pads the sequence to 18,944 (`wan_seq_pad=lane`) and sets v7 `BQ=4736`
  (`9472 * 2 = 18,944` on v6e, `4736 * 4 = 18,944` on v7), it switches both
  profiles to `DEFAULT_ULYSSES_OUT_A2A=shard_major`.

Tests
- `ulysses_out_a2a_test.py` (8): flat vs chunked vs shard_major parity of the
  inverse exchange, the mode guard, the `seq_len % num_shards != 0` guard, the
  `save_residuals + shard-major` guard, the strict `"shard_major"` kernel-gate
  `ValueError` guard, the 4-D q ring kernel bit-exact vs 3-D (running in CI via
  `interpret=True`), and bit-exact wrapper output for the shard-major 1D Ulysses
  and 2D Ulysses+Ring paths vs the flat layout (spy-asserting
  `out_num_shards == U`).
- `wan_runtime_options_test.py` (12): the new switch's YAML default and enum
  validation.
- `end_to_end/tpu/run_wan_stack_tests.sh`: adds `ulysses_out_a2a_test.py`.

Author: Perseus14

end_to_end/tpu/run_wan_fast_inference.sh

@@ -68,7 +68,7 @@
 #                    v6e seed 777, 17.92 dB / 0.6556 on v6e seed 12345, and
 #                    17.69 dB / 0.6918 on tpu7x seed 12345 (visually clean).
 #                    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)
@@ -183,6 +183,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"
@@ -201,6 +202,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}';" \
@@ -221,6 +223,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
 
@@ -254,6 +257,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 devices remain. Defaults to CP=4 / DP=2 on an 8-device slice. On a
@@ -322,6 +326,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

@@ -49,5 +49,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

@@ -140,6 +140,13 @@ wan_cross_attn_kernel: "xla"
 # 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 (and opportunistically uses the Pallas 4D-q / out_num_shards path
+# when S_local % block_q == 0); "shard_major" strictly requires the Pallas
+# 4D-q / out_num_shards path and raises if any kernel precondition fails.
+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

@@ -137,6 +137,13 @@ wan_cross_attn_kernel: "xla"
 # 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 (and opportunistically uses the Pallas 4D-q / out_num_shards path
+# when S_local % block_q == 0); "shard_major" strictly requires the Pallas
+# 4D-q / out_num_shards path and raises if any kernel precondition fails.
+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

@@ -174,6 +174,13 @@ wan_cross_attn_kernel: "xla"
 # 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 (and opportunistically uses the Pallas 4D-q / out_num_shards path
+# when S_local % block_q == 0); "shard_major" strictly requires the Pallas
+# 4D-q / out_num_shards path and raises if any kernel precondition fails.
+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

@@ -119,6 +119,13 @@ wan_cross_attn_kernel: "xla"
 # 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 (and opportunistically uses the Pallas 4D-q / out_num_shards path
+# when S_local % block_q == 0); "shard_major" strictly requires the Pallas
+# 4D-q / out_num_shards path and raises if any kernel precondition fails.
+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

@@ -121,6 +121,13 @@ wan_cross_attn_kernel: "xla"
 # 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 (and opportunistically uses the Pallas 4D-q / out_num_shards path
+# when S_local % block_q == 0); "shard_major" strictly requires the Pallas
+# 4D-q / out_num_shards path and raises if any kernel precondition fails.
+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

@@ -121,6 +121,13 @@ wan_cross_attn_kernel: "xla"
 # 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 (and opportunistically uses the Pallas 4D-q / out_num_shards path
+# when S_local % block_q == 0); "shard_major" strictly requires the Pallas
+# 4D-q / out_num_shards path and raises if any kernel precondition fails.
+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.