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

- Fused short KV text cross-attention Pallas kernel (wan_cross_attn_kernel='pallas'):
  Implements fused online softmax attention keeping 512-token text KV in VMEM,
  bypassing XLA dot-product fallback and saving 5.6s on v6e-8 and 3.4s on tpu7x-8.
- Token-sharded patch embedding (wan_patch_embed_mode='tokens'):
  Replaces Conv3d patch embedding with token-sharded patch projection matmul,
  eliminating activation all-gather collectives on tpu7x-8 (-0.5s) and v6e-8 (-0.2s).
- Chunked inverse Ulysses all-to-all (wan_ulysses_out_a2a='chunked'):
  Splits sequence-to-heads A2A over leading chunk axis to prevent XLA from
  generating costly layout conversions around the collective on v6e-8 (-1.8s).
- Lane sequence padding (wan_seq_pad='lane'):
  Pads video tokens to 128-token multiple per shard (75,600 -> 75,776), enabling
  zero-copy Ulysses reshapes with TokenPadding masking in attention kernels (-1.1s v6e, -0.2s v7).
- Tuned platform recipes & compiler flags in run_wan_fast_inference.sh:
  v6e-8: --xla_tpu_enable_ici_ag_pipelining=false --xla_tpu_enable_ag_backward_pipelining=false.
  tpu7x-8: --xla_tpu_enable_offloading_copy_to_sparsecore=false, BQ=4736.
- Benchmark results (Wan 2.2 T2V-14B, 720p/81f/40-step, CP=4, DP=2):
  v6e-8:   132.0s -> 122.2s denoise (-7.4%) / 124.3s e2e
  tpu7x-8:  97.6s ->  93.2s denoise (-4.5%) /  95.7s e2e

Author: Perseus14

end_to_end/tpu/run_wan_fast_inference.sh

@@ -78,15 +78,15 @@ 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"}
+V6E_LIBTPU=${V6E_LIBTPU:-"--xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=false --xla_tpu_enable_ici_ag_pipelining=false --xla_tpu_enable_ag_backward_pipelining=false"}
 # fuse_all_gather must stay false on v7: with true, libtpu fails backend init on
 # tpu7x ("Continuation fusion for AllGather ... not supported ... other than
 # Viperlite"). The tpu7x reference numbers below were measured with false.
-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 --xla_tpu_dvfs_p_state=7"}
+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 --xla_tpu_dvfs_p_state=7 --xla_tpu_enable_offloading_copy_to_sparsecore=false"}
 
 # Reference benchmarks (Wan 2.2 T2V-A14B, 720p/81f/40-step, CP=4, DP=2):
-#   v6e-8:   131.7s e2e (denoise 129.6s)
-#   tpu7x-8:  99.4s e2e (denoise  96.7s)
+#   v6e-8:   124.3s e2e (denoise 122.2s)
+#   tpu7x-8:  95.7s e2e (denoise  93.2s)
 case "$TPU_PROFILE" in
   v6e)
     PLATFORM_LIBTPU="$V6E_LIBTPU"
@@ -102,12 +102,13 @@ case "$TPU_PROFILE" in
     DEFAULT_BATCHED_TE=true
     DEFAULT_VAE_CHUNK=1
     DEFAULT_VAE_SPATIAL=8
+    DEFAULT_ULYSSES_OUT_A2A=chunked
     ;;
   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
@@ -117,6 +118,7 @@ case "$TPU_PROFILE" in
     DEFAULT_BATCHED_TE=true
     DEFAULT_VAE_CHUNK=1
     DEFAULT_VAE_SPATIAL=8
+    DEFAULT_ULYSSES_OUT_A2A=flat
     ;;
   *)
     echo "== warning: unrecognised accelerator '${ACCEL_TYPE:-unknown}';" \
@@ -134,6 +136,7 @@ case "$TPU_PROFILE" in
     DEFAULT_BATCHED_TE=true
     DEFAULT_VAE_CHUNK=1
     DEFAULT_VAE_SPATIAL=8
+    DEFAULT_ULYSSES_OUT_A2A=flat
     ;;
 esac
 
@@ -161,6 +164,7 @@ COMPILE_TE=${COMPILE_TE:-$DEFAULT_COMPILE_TE}
 USE_BATCHED_TE=${USE_BATCHED_TE:-$DEFAULT_BATCHED_TE}
 VAE_SPATIAL=${VAE_SPATIAL:-$DEFAULT_VAE_SPATIAL}
 VAE_DECODE_CHUNK=${VAE_DECODE_CHUNK:-$DEFAULT_VAE_CHUNK}
+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
@@ -222,6 +226,10 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \
   text_encoder_dtype=bfloat16 compile_text_encoder="$COMPILE_TE" use_batched_text_encoder="$USE_BATCHED_TE" \
   use_kv_cache=true use_base2_exp=true use_experimental_scheduler=true \
   use_fused_rope_kernel=true \
+  wan_cross_attn_kernel="${WAN_CROSS_ATTN_KERNEL:-pallas}" \
+  wan_patch_embed_mode="${WAN_PATCH_EMBED_MODE:-tokens}" \
+  wan_ulysses_out_a2a="${WAN_ULYSSES_OUT_A2A}" \
+  wan_seq_pad="${WAN_SEQ_PAD:-lane}" \
   fps=16 "${GUIDANCE_ARGS[@]}" \
   seed="${SEED:-12345}" \
   flash_block_sizes="$FLASH_BLOCK_SIZES" \

src/maxdiffusion/configs/base_wan_14b.yml

@@ -133,6 +133,23 @@ 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"
+# 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, sequence -> heads): "flat" splits
+# the sequence axis directly; "chunked" exchanges over an explicit leading chunk
+# axis so XLA does not relayout the output around the collective. Same values.
+wan_ulysses_out_a2a: "flat"
+# Self-attention token padding: "off", or "lane" to pad each context shard of the
+# video token sequence to a multiple of 128 tokens (per ring segment on the ring
+# path). 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,23 @@ 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"
+# 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, sequence -> heads): "flat" splits
+# the sequence axis directly; "chunked" exchanges over an explicit leading chunk
+# axis so XLA does not relayout the output around the collective. Same values.
+wan_ulysses_out_a2a: "flat"
+# Self-attention token padding: "off", or "lane" to pad each context shard of the
+# video token sequence to a multiple of 128 tokens (per ring segment on the ring
+# path). 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

@@ -155,6 +155,23 @@ 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"
+# 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, sequence -> heads): "flat" splits
+# the sequence axis directly; "chunked" exchanges over an explicit leading chunk
+# axis so XLA does not relayout the output around the collective. Same values.
+wan_ulysses_out_a2a: "flat"
+# Self-attention token padding: "off", or "lane" to pad each context shard of the
+# video token sequence to a multiple of 128 tokens (per ring segment on the ring
+# path). 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,23 @@ 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"
+# 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, sequence -> heads): "flat" splits
+# the sequence axis directly; "chunked" exchanges over an explicit leading chunk
+# axis so XLA does not relayout the output around the collective. Same values.
+wan_ulysses_out_a2a: "flat"
+# Self-attention token padding: "off", or "lane" to pad each context shard of the
+# video token sequence to a multiple of 128 tokens (per ring segment on the ring
+# path). 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,23 @@ 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"
+# 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, sequence -> heads): "flat" splits
+# the sequence axis directly; "chunked" exchanges over an explicit leading chunk
+# axis so XLA does not relayout the output around the collective. Same values.
+wan_ulysses_out_a2a: "flat"
+# Self-attention token padding: "off", or "lane" to pad each context shard of the
+# video token sequence to a multiple of 128 tokens (per ring segment on the ring
+# path). 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

@@ -114,6 +114,23 @@ 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"
+# 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, sequence -> heads): "flat" splits
+# the sequence axis directly; "chunked" exchanges over an explicit leading chunk
+# axis so XLA does not relayout the output around the collective. Same values.
+wan_ulysses_out_a2a: "flat"
+# Self-attention token padding: "off", or "lane" to pad each context shard of the
+# video token sequence to a multiple of 128 tokens (per ring segment on the ring
+# path). 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.