perf(wan): pad sequence to 128-lane multiple for self-attention (wan_seq_pad="lane")

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_ULYSSES_OUT_A2A
+#   WAN_CROSS_ATTN_KERNEL / WAN_PATCH_EMBED_MODE / WAN_ULYSSES_OUT_A2A / WAN_SEQ_PAD
 #                    override Wan fast-inference kernel/layout switches
 #   DP / CP / PER_DEVICE_BATCH / SEED
 #                    override mesh parallelism, per-device batch, or RNG seed (default 12345)
@@ -146,8 +146,10 @@ if [ -z "${TPU_PROFILE:-}" ]; then
 fi
 
 # Keep LIBTPU flags single-line (literal backslashes truncate libtpu flag parsing).
-COMMON_LIBTPU=${COMMON_LIBTPU:-"--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_memory_bound_loop_optimizer_options=enabled:true --xla_tpu_enable_dot_strength_reduction=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_tpu_enable_ici_ag_pipelining=true --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true"}
-V6E_LIBTPU=${V6E_LIBTPU:-""}
+# ici_ag_pipelining and ag_backward_pipelining are configured per-profile below
+# (disabled on v6e, enabled on v7).
+COMMON_LIBTPU=${COMMON_LIBTPU:-"--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_memory_bound_loop_optimizer_options=enabled:true --xla_tpu_enable_dot_strength_reduction=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true"}
+V6E_LIBTPU=${V6E_LIBTPU:-"--xla_tpu_enable_ici_ag_pipelining=false --xla_tpu_enable_ag_backward_pipelining=false"}
 # fuse_all_gather is false in every profile (see COMMON_LIBTPU); libtpu fails
 # backend init with it enabled on any non-v5e platform, including v6e and v7
 # ("Continuation fusion for AllGather ... not supported ... other than Viperlite").
@@ -158,13 +160,14 @@ V7_DVFS_FLAG=""
 if [ "$PIN_DVFS_P_STATE" = "true" ]; then
   V7_DVFS_FLAG=" --xla_tpu_dvfs_p_state=7"
 fi
-V7_LIBTPU=${V7_LIBTPU:-"--xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true${V7_DVFS_FLAG}"}
+V7_LIBTPU=${V7_LIBTPU:-"--xla_tpu_enable_ici_ag_pipelining=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true --xla_tpu_enable_offloading_copy_to_sparsecore=false${V7_DVFS_FLAG}"}
 
-# Reference benchmarks (Wan 2.2 T2V-A14B, 720p/81f/40-step, CP=4, DP=2, warm AOT
-# cache, these launcher defaults, measured at this PR without the later PRs):
-#   tpu7x-8: 115.5s generate (denoise 112.8s) DVFS unpinned;
-#            102.0s generate (denoise  99.4s) with PIN_DVFS_P_STATE=true
-#   v6e-8:   140.1s generate (denoise 137.0s)
+# Reference benchmarks (Wan 2.2 T2V-A14B, 720p/81f/40-step, CP=4, DP=2, full
+# stack, these launcher defaults, warm AOT cache, DVFS unpinned):
+#   v6e-8:   126.6s - 127.5s generate (denoise 124.5s - 125.0s)
+#   tpu7x-8: 107.8s generate (denoise 104.9s; 109.6s / 106.6s with fused RoPE off)
+# tpu7x-8 with PIN_DVFS_P_STATE=true: 94.6s generate (denoise 91.7s) with the
+# fused RoPE kernel on, and 96.4s generate (denoise 93.4s) with it off.
 case "$TPU_PROFILE" in
   v6e)
     PLATFORM_LIBTPU="$V6E_LIBTPU"
@@ -183,13 +186,14 @@ case "$TPU_PROFILE" in
     DEFAULT_FUSED_ROPE=true
     DEFAULT_CROSS_ATTN_KERNEL=pallas
     DEFAULT_PATCH_EMBED_MODE=tokens
-    DEFAULT_ULYSSES_OUT_A2A=chunked
+    DEFAULT_ULYSSES_OUT_A2A=shard_major
+    DEFAULT_SEQ_PAD=lane
     ;;
   v7)
     PLATFORM_LIBTPU="$V7_LIBTPU"
     DEFAULT_ATTENTION="ulysses_ring_custom_fixed_m"
     DEFAULT_U=2
-    DEFAULT_BQ=6400
+    DEFAULT_BQ=4736
     DEFAULT_BKV=2048
     DEFAULT_BKV_COMPUTE=2048
     DEFAULT_BKV_COMPUTE_IN=2048
@@ -202,7 +206,8 @@ case "$TPU_PROFILE" in
     DEFAULT_FUSED_ROPE=true
     DEFAULT_CROSS_ATTN_KERNEL=pallas
     DEFAULT_PATCH_EMBED_MODE=tokens
-    DEFAULT_ULYSSES_OUT_A2A=chunked
+    DEFAULT_ULYSSES_OUT_A2A=shard_major
+    DEFAULT_SEQ_PAD=lane
     ;;
   *)
     echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \
@@ -224,6 +229,7 @@ case "$TPU_PROFILE" in
     DEFAULT_CROSS_ATTN_KERNEL=xla
     DEFAULT_PATCH_EMBED_MODE=tokens
     DEFAULT_ULYSSES_OUT_A2A=flat
+    DEFAULT_SEQ_PAD=off
     ;;
 esac
 
@@ -258,6 +264,7 @@ FUSED_ROPE_ARGS=()
 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}
+WAN_SEQ_PAD=${WAN_SEQ_PAD:-$DEFAULT_SEQ_PAD}
 
 # 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
@@ -327,6 +334,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \
   wan_cross_attn_kernel="$WAN_CROSS_ATTN_KERNEL" \
   wan_patch_embed_mode="$WAN_PATCH_EMBED_MODE" \
   wan_ulysses_out_a2a="$WAN_ULYSSES_OUT_A2A" \
+  wan_seq_pad="$WAN_SEQ_PAD" \
   fps=16 ${GUIDANCE_ARGS[@]+"${GUIDANCE_ARGS[@]}"} \
   seed="${SEED:-12345}" \
   flash_block_sizes="$FLASH_BLOCK_SIZES" \

end_to_end/tpu/run_wan_stack_tests.sh

@@ -51,5 +51,7 @@ TESTS=(
   "$T/wan_patch_embed_test.py"
   # Shard-major Ulysses output A2A (feat/ulysses-shard-major-a2a)
   "$T/ulysses_out_a2a_test.py"
+  # Lane sequence padding (feat/wan-fast-inference-optimizations)
+  "$T/wan_seq_pad_test.py"
 )
 PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@"

src/maxdiffusion/configs/base_wan_14b.yml

@@ -147,6 +147,12 @@ wan_patch_embed_mode: "conv"
 # 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"
+# Self-attention token padding: "off", or "lane" to append pad tokens at the end
+# of the video token sequence (or at the end of each ring segment on the ring
+# path) so every context shard has a multiple of 128 tokens. Aligned shards make
+# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys.
+# Only the custom Ulysses / Ulysses-ring kernels use it.
+wan_seq_pad: "off"
 # 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

@@ -144,6 +144,12 @@ wan_patch_embed_mode: "conv"
 # 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"
+# Self-attention token padding: "off", or "lane" to append pad tokens at the end
+# of the video token sequence (or at the end of each ring segment on the ring
+# path) so every context shard has a multiple of 128 tokens. Aligned shards make
+# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys.
+# Only the custom Ulysses / Ulysses-ring kernels use it.
+wan_seq_pad: "off"
 # 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

@@ -181,6 +181,12 @@ wan_patch_embed_mode: "conv"
 # 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"
+# Self-attention token padding: "off", or "lane" to append pad tokens at the end
+# of the video token sequence (or at the end of each ring segment on the ring
+# path) so every context shard has a multiple of 128 tokens. Aligned shards make
+# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys.
+# Only the custom Ulysses / Ulysses-ring kernels use it.
+wan_seq_pad: "off"
 # 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

@@ -126,6 +126,12 @@ wan_patch_embed_mode: "conv"
 # 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"
+# Self-attention token padding: "off", or "lane" to append pad tokens at the end
+# of the video token sequence (or at the end of each ring segment on the ring
+# path) so every context shard has a multiple of 128 tokens. Aligned shards make
+# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys.
+# Only the custom Ulysses / Ulysses-ring kernels use it.
+wan_seq_pad: "off"
 # 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

@@ -128,6 +128,12 @@ wan_patch_embed_mode: "conv"
 # 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"
+# Self-attention token padding: "off", or "lane" to append pad tokens at the end
+# of the video token sequence (or at the end of each ring segment on the ring
+# path) so every context shard has a multiple of 128 tokens. Aligned shards make
+# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys.
+# Only the custom Ulysses / Ulysses-ring kernels use it.
+wan_seq_pad: "off"
 # 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

@@ -128,6 +128,12 @@ wan_patch_embed_mode: "conv"
 # 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"
+# Self-attention token padding: "off", or "lane" to append pad tokens at the end
+# of the video token sequence (or at the end of each ring segment on the ring
+# path) so every context shard has a multiple of 128 tokens. Aligned shards make
+# the Ulysses all-to-all reshapes free; the custom kernels mask the pad keys.
+# Only the custom Ulysses / Ulysses-ring kernels use it.
+wan_seq_pad: "off"
 # 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.