feat(wan): fast serving with persistent AOT caching and tuned inference recipe

Delivers production-grade fast serving for Wan 2.2 T2V-A14B on Cloud TPU:
- Persistent AOT Compilation Caching with source hash and static graphdef metadata
- AOT signatures key non-static Python scalars by type, so guidance values share one executable
- Shared Wan 2.2 warmup for T2V and I2V: compile both experts without executing them, bounded in-flight steps
- Tuned Serving Recipe: Production launcher with platform-specific v6e and v7 profiles
- Clean Key-Centering & 4-Operand Kernel (mk, q, k, v) without VMEM register pressure
- Merged Pre-A2A Norm Reduction (pmax) removing ring all_gather relayout barrier
- Hoisted scalar-prefetch metadata in _lse_scan eliminating collective-permute stalls
- Lossless 6D unpatchify transpose optimization (collapsing p_t=1 to avoid 8D stride copies)

Author: Perseus14

end_to_end/tpu/run_wan_fast_inference.sh

@@ -101,10 +101,11 @@ svg_low_noise_density: -1.0
 # compute blocks and the sparse kernel for small ones. Must stay a dict so the
 # command line can override it with JSON.
 svg_flash_block_sizes: {}
-attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
+attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, tokamax_ring_custom, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block, ulysses_ring_custom_bidir
 
 use_base2_exp: True
 use_experimental_scheduler: True
+use_k_centering: "auto"  # auto: on for non-ring ulysses_custom_fixed_m* (free, virtual), off for ring paths (pmean + HBM copy)
 # 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.
@@ -309,6 +310,8 @@ dataset_config_name: ''
 jax_cache_dir: ''
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+aot_build_revision: ''
+wan_debug_cond_timers: False
 # Directory for memoized torch->flax converted weights ('' = disabled).
 converted_weights_dir: ''
 hf_data_dir: ''

src/maxdiffusion/aot_cache.py

@@ -262,6 +262,8 @@ dataset_config_name: ''
 jax_cache_dir: ''
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+aot_build_revision: ''
+wan_debug_cond_timers: False
 # Directory for memoized torch->flax converted weights ('' = disabled).
 converted_weights_dir: ''
 hf_data_dir: ''

src/maxdiffusion/configs/base_wan_14b.yml

@@ -127,8 +127,6 @@ use_base2_exp: True
 use_experimental_scheduler: True
 # auto: on for non-ring ulysses_custom_fixed_m* (virtual, no copy), off for ring
 # paths (materialized K - mean: an HBM copy, plus a pmean when R>1).
-# Note: the Wan pipelines do not pass this key to the attention layer yet, so
-# the layer's default ("auto") applies.
 use_k_centering: "auto"
 # For attention=ulysses_ring*, hidden Ulysses shard count; ring shards are context / this.
 ulysses_shards: -1
@@ -302,6 +300,8 @@ dataset_config_name: ''
 jax_cache_dir: ''
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+aot_build_revision: ''
+wan_debug_cond_timers: False
 # Directory for memoized torch->flax converted weights ('' = disabled).
 converted_weights_dir: ''
 hf_data_dir: ''

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -81,7 +81,7 @@ jit_initializers: True
 # Set true to load weights from pytorch
 from_pt: True
 split_head_dim: True
-attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
+attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block
 use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
@@ -253,6 +253,8 @@ dataset_config_name: ''
 jax_cache_dir: '.jax_cache'
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+aot_build_revision: ''
+wan_debug_cond_timers: False
 # Directory for memoized torch->flax converted weights ('' = disabled).
 converted_weights_dir: ''
 hf_data_dir: ''

src/maxdiffusion/configs/base_wan_27b.yml

@@ -83,7 +83,7 @@ jit_initializers: True
 # Set true to load weights from pytorch
 from_pt: True
 split_head_dim: True
-attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
+attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block
 use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
@@ -256,6 +256,8 @@ dataset_config_name: ''
 jax_cache_dir: ''
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+aot_build_revision: ''
+wan_debug_cond_timers: False
 # Directory for memoized torch->flax converted weights ('' = disabled).
 converted_weights_dir: ''
 hf_data_dir: ''

src/maxdiffusion/configs/base_wan_animate.yml

@@ -83,7 +83,7 @@ jit_initializers: True
 # Set true to load weights from pytorch
 from_pt: True
 split_head_dim: True
-attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
+attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block
 use_base2_exp: True
 use_experimental_scheduler: True
 # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this.
@@ -257,6 +257,8 @@ dataset_config_name: ''
 jax_cache_dir: ''
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+aot_build_revision: ''
+wan_debug_cond_timers: False
 # Directory for memoized torch->flax converted weights ('' = disabled).
 converted_weights_dir: ''
 hf_data_dir: ''