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

Persistent per-shape AOT executable cache, converted-weights cache and a
tuned launcher for Wan inference on TPU.

User-visible changes
- AOT caching is opt-in. By default `generate_wan.run()` installs no cache
  directory and every call is plain jax.jit. When `aot_cache_dir` is set, the
  source revision must be reusable (`get_git_commit_hash()` or an explicit
  clean `aot_build_revision`, combined with the package content hash); a dirty
  or unversioned tree downgrades to an ephemeral temp dir rather than writing or
  reading stale persistent executables. `enable_zero_execution_warmup: True`
  (False in all 6 Wan ymls) uses an ephemeral temp dir that is removed in a
  `try/finally` block.
- `aot_cache_dir` must be a local filesystem path (`gs://` URLs raise
  `ValueError`), and `purge_compiled_cache` refuses symlink targets that resolve
  outside the cache root.
- `_prune_stale_generations` prunes `.aotx` files from older fingerprint
  generations beyond `_DEFAULT_KEEP_GENERATIONS = 3`.
- Converted-weights cache format is v3. Format v2 manifests whose
  `legacy_source_fingerprint` matches are migrated in place to v3 via
  `_migrate_manifest_header`; older or mismatched manifests are rejected and
  re-converted once. `try_load_converted_weights` validates leaf keys, shapes,
  and dtypes against `eval_shapes`.

AOT cache (aot_cache.py, generate_wan.py)
- Signatures key non-static Python scalars by type (jit traces them), so e.g.
  guidance values share one executable; statics are keyed by value; nnx
  GraphDef statics get a process-deterministic digest.
- Metadata fingerprint covers model/attention/tiles/mesh/VAE/dtype/remat/
  cache options, device kind, process count, jax/jaxlib/libtpu/flax/qwix/
  tokamax versions, matmul precision, PRNG impl, x64, threefry,
  LIBTPU_INIT_ARGS and XLA_FLAGS, plus the source revision.
- Multi-host: per-host `-p{idx}` .aotx files.
- `_align_inputs` raises on a leaf-count mismatch instead of truncating.
- The .aotx pickle envelope is loaded with an exact (module, name) allowlist.

Warmup & Denoise helpers (wan_denoise_utils.py)
- When aot_cache is in warmup mode, `compile_experts` compiles both experts'
  forward passes without executing them. Without an installed cache, warmup is
  an ordinary 2-step run.
- `InflightWindow` bounds queued denoise steps (`MAXD_QUEUE_MAX_INFLIGHT`,
  default 4, validated `>= 1`), shared by T2V and I2V.

Other
- Unpatchify: when p_t == 1, reshape/transpose through a 7D tensor instead of
  8D (same result; avoids an 8D strided copy).
- `use_k_centering` is plumbed from the Wan config to the attention layer;
  `use_k_centering`, `aot_build_revision`, and `enable_zero_execution_warmup`
  added to the Wan ymls.
- Dot-product attention: float32_qk_product uses preferred_element_type=f32
  on the QK einsum instead of upcasting Q/K (memory-efficient path unchanged).
- Video export is atomic (temp file + os.replace), process-0 only, and shared
  by `run()` and `inference_generate_video`; trainers print SSIM on process 0.
- Launcher (`end_to_end/tpu/run_wan_fast_inference.sh`): platform-detected v6e
  / v7 profiles, generic profile for everything else; derives `TPU_DEVICES`
  only for per-chip families (`v5litepod | v5e | v6e | v7* | tpu7*`) and exits
  with an error for other families unless `DP` and `CP` are set explicitly;
  DVFS pin is opt-in (`PIN_DVFS_P_STATE`).

Performance, measured with this PR at the top of the stack (without #488-#491;
Wan 2.2 T2V-A14B, 720p/81f/40 steps, CP=4, DP=2, warm AOT, launcher defaults):
  v6e-8:   generate 140.1s, denoise 137.0s
  tpu7x-8: generate 115.5s, denoise 112.8s (DVFS unpinned)
           generate 102.0s, denoise  99.4s (`PIN_DVFS_P_STATE=true`)

Tests
- aot_cache_test.py and converted_weights_cache_test.py: restricted-unpickler
  allowlist, dirty-tree ephemeral downgrade and symlink-escape purge guards,
  stale-generation pruning, float32_qk_product keeping bf16 QK operands,
  v2-to-v3 manifest migration and fail-closed mismatched/fingerprint-less
  manifests, leaf dtype validation, disk-space checks, warm start before the
  network, revision captured through HF symlinks, and distinct fingerprints for
  subfolders with identical index files.
- wan/wan_warmup_coverage_test.py (28): all tests run unconditionally in GitHub
  Actions; wan/wan_transformer_test.py adds `FlaxWanAttention` k-centering
  parity and dispatch tests.

Author: Perseus14

end_to_end/tpu/run_wan_fast_inference.sh

@@ -37,5 +37,10 @@ TESTS=(
   "$T/custom_splash_unpadded_test.py"
   # XLA-fused RMSNorm+RoPE producer (feat/fused-rmsnorm-rope-xla)
   "$T/fused_producers_test.py"
+  # Wan fast serving / AOT cache (feat/wan-fast-serving)
+  "$T/aot_cache_test.py"
+  "$T/converted_weights_cache_test.py"
+  "$T/wan/wan_transformer_test.py"
+  "$T/wan/wan_warmup_coverage_test.py"
 )
 PYTHONPATH="src${PYTHONPATH:+:$PYTHONPATH}" exec "${PYTHON:-python3}" -m pytest -q -rs "${TESTS[@]}" "$@"

end_to_end/tpu/run_wan_stack_tests.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* (virtual, no copy), off for ring paths (virtual via k_mean, plus pmean when R>1)
 # 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,12 @@ dataset_config_name: ''
 jax_cache_dir: ''
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+# Zero-execution warmup (compile the transformer passes without executing them) is always
+# on when aot_cache_dir is set. When aot_cache_dir is '', True installs an ephemeral AOT
+# cache for zero-execution warmup; False keeps plain jax.jit (warmup executes the model).
+enable_zero_execution_warmup: False
+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

@@ -102,6 +102,7 @@ attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cud
 
 use_base2_exp: True
 use_experimental_scheduler: True
+use_k_centering: "auto"
 # 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.
@@ -262,6 +263,12 @@ dataset_config_name: ''
 jax_cache_dir: ''
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+# Zero-execution warmup (compile the transformer passes without executing them) is always
+# on when aot_cache_dir is set. When aot_cache_dir is '', True installs an ephemeral AOT
+# cache for zero-execution warmup; False keeps plain jax.jit (warmup executes the model).
+enable_zero_execution_warmup: False
+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

@@ -129,8 +129,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 (virtual K-centering via k_mean, plus a pmean across ring shards 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
@@ -304,6 +302,12 @@ dataset_config_name: ''
 jax_cache_dir: ''
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+# Zero-execution warmup (compile the transformer passes without executing them) is always
+# on when aot_cache_dir is set. When aot_cache_dir is '', True installs an ephemeral AOT
+# cache for zero-execution warmup; False keeps plain jax.jit (warmup executes the model).
+enable_zero_execution_warmup: False
+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,9 +81,10 @@ 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
+use_k_centering: "auto"
 # 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.
@@ -253,6 +254,12 @@ dataset_config_name: ''
 jax_cache_dir: '.jax_cache'
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+# Zero-execution warmup (compile the transformer passes without executing them) is always
+# on when aot_cache_dir is set. When aot_cache_dir is '', True installs an ephemeral AOT
+# cache for zero-execution warmup; False keeps plain jax.jit (warmup executes the model).
+enable_zero_execution_warmup: False
+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,9 +83,10 @@ 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
+use_k_centering: "auto"
 # 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.
@@ -256,6 +257,12 @@ dataset_config_name: ''
 jax_cache_dir: ''
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+# Zero-execution warmup (compile the transformer passes without executing them) is always
+# on when aot_cache_dir is set. When aot_cache_dir is '', True installs an ephemeral AOT
+# cache for zero-execution warmup; False keeps plain jax.jit (warmup executes the model).
+enable_zero_execution_warmup: False
+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,9 +83,10 @@ 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
+use_k_centering: "auto"
 # 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.
@@ -257,6 +258,12 @@ dataset_config_name: ''
 jax_cache_dir: ''
 # Directory for per-shape AOT serialized executables ('' = disabled).
 aot_cache_dir: ''
+# Zero-execution warmup (compile the transformer passes without executing them) is always
+# on when aot_cache_dir is set. When aot_cache_dir is '', True installs an ephemeral AOT
+# cache for zero-execution warmup; False keeps plain jax.jit (warmup executes the model).
+enable_zero_execution_warmup: False
+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_i2v_14b.yml

@@ -19,7 +19,6 @@
 import time
 import os
 import subprocess
-import uuid
 from maxdiffusion.checkpointing.ltx2_checkpointer import LTX2Checkpointer
 from maxdiffusion import aot_cache, pyconfig, max_logging, max_utils
 from absl import app
@@ -149,7 +148,7 @@ def _canonical_aot_value(value):
 
 def _non_reusable_aot_revision():
   """Returns a unique identity so unversioned source can never hit old HLO."""
-  return f"unversioned:{uuid.uuid4().hex}"
+  return aot_cache.non_reusable_aot_revision()
 
 
 def _resolve_ltx2_aot_source_revision(config, commit_hash=None):
@@ -161,9 +160,7 @@ def _resolve_ltx2_aot_source_revision(config, commit_hash=None):
 
 
 def _is_reusable_aot_revision(source_revision) -> bool:
-  if source_revision is None or not str(source_revision).strip():
-    return False
-  return not str(source_revision).startswith(("dirty:", "unversioned:"))
+  return aot_cache.is_reusable_aot_revision(source_revision)
 
 
 def ltx2_aot_metadata(config, pipeline, source_revision=None):