Cut WAN family time-to-ready 220-304s -> 33-51s: AOT executable cache, zero-exec warmup, converted-weights cache

[syhuang22]

Cut WAN family time-to-ready: AOT executable cache, zero-exec warmup, converted-weights cache

Serving-oriented startup optimizations for the whole WAN family (2.1 / 2.2 T2V, 2.2 I2V, VACE). A process restart previously re-paid tracing, XLA compilation (or persistent-cache lookup), a full torch->flax weight conversion, and two real denoise steps of warmup. All four are either memoized on disk or eliminated.

Mechanisms

1. Per-shape AOT executable cache (aot_cache.py): cached_jit replaces jax.jit at definition sites (identical behavior until install()). After warmup the compiled executables are serialized (jax.experimental.serialize_executable); later processes deserialize on background threads — no trace, no lowering, no cache lookup. One file per input signature, so resolutions never collide. Wired into the transformer forward passes (base, CFG, cfg-cache, VACE) and VAE encode/decode.
2. Zero-exec warmup (aot_cache.warmup_mode()): compilation only needs avals. During warmup, wrapped calls lower+compile (or reuse the deserialized executable) and return sharded zeros, so downstream executables (scheduler, VAE) still warm faithfully — without paying real 14B denoise steps.
3. Converted-weights cache (wan_utils.py, converted_weights_dir): the torch->flax conversion (per-tensor transpose + scan-stack + dtype cast) is a pure function of the checkpoint; it is now saved once as per-tensor .npy + manifest and re-loaded with parallel reads. Self-validating against eval shapes and the cast policy (mismatch = silent reconvert).

Both caches are off by default (aot_cache_dir: '', converted_weights_dir: '').

Results (v7x-8, 720p/81f, 2D-ring recipe, same warm JAX persistent cache on both sides; only the two cache knobs differ)

Pipeline	Load (off -> hit)	Compile window (off -> hit)	Time-to-ready (off -> hit)
WAN 2.1 T2V 14B	24.2s -> 25.0s	195.5s -> 8.3s (23.5x)	219.7s -> 33.3s
WAN 2.2 T2V A14B	29.6s -> 28.1s	55.9s -> 8.7s (6.4x)	85.5s -> 36.8s
WAN 2.2 I2V A14B	29.2s -> 27.9s	151.7s -> 16.3s (9.3x)	180.9s -> 44.2s
WAN 2.1 VACE 14B	29.7s -> 29.6s	274.6s -> 21.5s (12.8x)	304.3s -> 51.1s

First run per (model, shape) builds the caches and lands in between (compile 11-33s).

Highlights: VAE decode cold compile 144.9s -> 0.2s (2.1), I2V image VAE encode 89.5s -> 1.0s, VACE (unscanned, so all 40 blocks were traced every process) 274.6s -> 21.5s.

Load is flat in this A/B because both sides already include this PR's loader changes (fused cast, chunked staged put, threaded 2.1 component load) and the remaining critical path is the CPU text-encoder load; measured against the previous loader the transformer host path alone is 26.2s -> 8.5s, verified bitwise-equal on all 1095 params.

Correctness

• Denoise step time identical across all variants (5.9s +-0.4 run noise); videos generated with all mechanisms on validated frame-exact (h264, 1280x720, 81 frames, zero decode errors on full decode) for all four pipelines.
• Loader-path equivalence verified bitwise (1095/1095 params) between the fused-cast path and the historical two-pass path.
• The AOT cache is keyed by a config fingerprint (model, attention, flash block sizes, mesh, dtypes, scan, jax version) plus a per-call input signature; unknown signatures fall back to plain jit and are serialized after the next warmup. Executables deserialize with mesh-ordered execution_devices.
• Unit tests: aot_cache_test.py (9, incl. cross-process signature determinism and in-trace inlining) and converted_weights_cache_test.py (5, incl. bf16 bit-view round-trip and cache invalidation).

Example

end_to_end/tpu/run_wan_fast_inference.sh — end-to-end serving example (tuned v7 flag set + both caches): ./run_wan_fast_inference.sh 22 40 "a cat surfing".

[eltsai]

LGTM! I was able to reproduce the stats on v7x