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

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.