feat(wan): fast inference optimizations (fused cross-attention, token patch embed, shard-major A2A, lane padding)

end_to_end/tpu/run_wan_fast_inference.sh

@@ -63,6 +63,8 @@
 #                    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_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)
 set -euo pipefail
@@ -132,8 +134,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:-"--xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false"}
+# 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); it must stay
 # false on v7, where libtpu fails backend init with it enabled ("Continuation
 # fusion for AllGather ... not supported ... other than Viperlite").
@@ -144,13 +148,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_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false --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:   not re-measured at this PR
+# Reference benchmarks (Wan 2.2 T2V-A14B, 720p/81f/40-step, CP=4, DP=2, full 5-PR
+# stack, these launcher defaults, warm AOT cache, DVFS unpinned):
+#   v6e-8:   128.05s generate (denoise 125.9s)
+#   tpu7x-8: 107.6s  generate (denoise 105.1s)
+# tpu7x-8 with PIN_DVFS_P_STATE=true: 96.1s generate (denoise 93.6s), measured
+# with the fused RoPE kernel off (pinned with it on was not measured).
 case "$TPU_PROFILE" in
   v6e)
     PLATFORM_LIBTPU="$V6E_LIBTPU"
@@ -167,12 +172,16 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_CHUNK=1
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=true
+    DEFAULT_CROSS_ATTN_KERNEL=pallas
+    DEFAULT_PATCH_EMBED_MODE=tokens
+    DEFAULT_ULYSSES_OUT_A2A=chunked
+    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
@@ -183,6 +192,10 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_CHUNK=1
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=true
+    DEFAULT_CROSS_ATTN_KERNEL=pallas
+    DEFAULT_PATCH_EMBED_MODE=tokens
+    DEFAULT_ULYSSES_OUT_A2A=chunked
+    DEFAULT_SEQ_PAD=lane
     ;;
   *)
     echo "== warning: unrecognised or non-v6e/v7 accelerator '${ACCEL_TYPE:-unknown}';" \
@@ -201,6 +214,10 @@ case "$TPU_PROFILE" in
     DEFAULT_VAE_CHUNK=1
     DEFAULT_VAE_SPATIAL=8
     DEFAULT_FUSED_ROPE=false
+    DEFAULT_CROSS_ATTN_KERNEL=xla
+    DEFAULT_PATCH_EMBED_MODE=tokens
+    DEFAULT_ULYSSES_OUT_A2A=flat
+    DEFAULT_SEQ_PAD=off
     ;;
 esac
 
@@ -232,6 +249,10 @@ USE_FUSED_ROPE_KERNEL=${USE_FUSED_ROPE_KERNEL:-$DEFAULT_FUSED_ROPE}
 FUSED_ROPE_ARGS=()
 [ -n "${FUSED_ROPE_BLOCK_S:-}" ] && FUSED_ROPE_ARGS+=("fused_rope_block_s=$FUSED_ROPE_BLOCK_S")
 [ -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}
+WAN_SEQ_PAD=${WAN_SEQ_PAD:-$DEFAULT_SEQ_PAD}
 
 # 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
@@ -294,6 +315,10 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \
   use_kv_cache=true use_base2_exp=true use_experimental_scheduler=true \
   use_fused_rope_kernel="$USE_FUSED_ROPE_KERNEL" \
   ${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" \
+  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

@@ -42,5 +42,11 @@ TESTS=(
   # Fused RMSNorm+RoPE producer, Wan switches in config (feat/wan-custom-kernels)
   "$T/fused_rmsnorm_rope_pallas_test.py"
   "$T/wan_runtime_options_test.py"
+  # Cross-attention kernel, token patch embed, shard-major A2A, lane padding
+  # (feat/wan-fast-inference-optimizations)
+  "$T/cross_attention_pallas_test.py"
+  "$T/ulysses_out_a2a_test.py"
+  "$T/wan_patch_embed_test.py"
+  "$T/wan_seq_pad_test.py"
 )
 PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@"

src/maxdiffusion/__init__.py

@@ -16,6 +16,7 @@
 
 __version__ = "0.22.0.dev0"
 
+
 from typing import TYPE_CHECKING
 
 from .utils import (

src/maxdiffusion/configs/base_wan_14b.yml

@@ -133,6 +133,27 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
+wan_cross_attn_kernel: "xla"
+# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
+wan_cross_attn_cpu_interpret: False
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# 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"
+# 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

@@ -130,6 +130,27 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
+wan_cross_attn_kernel: "xla"
+# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
+wan_cross_attn_cpu_interpret: False
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# 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"
+# 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

@@ -162,6 +162,27 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
+wan_cross_attn_kernel: "xla"
+# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
+wan_cross_attn_cpu_interpret: False
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# 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"
+# 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

@@ -112,6 +112,27 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
+wan_cross_attn_kernel: "xla"
+# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
+wan_cross_attn_cpu_interpret: False
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# 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"
+# 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

@@ -114,6 +114,27 @@ wan_cfg_before_unpatchify: True
 wan_cross_attn_prescale_kv: False
 # Fused RoPE accumulation: "auto" (measured per-platform default), "dtype", "f32".
 wan_rope_accum: "auto"
+# Text cross-attention backend where it would otherwise fall back to XLA dot-product
+# (KV shorter than flash_min_seq_length): "xla", or "pallas" for the fused kernel that
+# keeps the KV in VMEM (within bf16 rounding of "xla", not bit-identical).
+wan_cross_attn_kernel: "xla"
+# When wan_cross_attn_kernel is "pallas", allow CPU execution via Pallas interpret mode (tests only).
+wan_cross_attn_cpu_interpret: False
+# Patch embedding: "conv" (nnx.Conv; its channel-sharded output is resharded by
+# 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"
+# 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.