feat(wan): fast serving with persistent AOT caching and tuned inference recipe by Perseus14 · Pull Request #479 · AI-Hypercomputer/maxdiffusion

Perseus14 · 2026-09-13T06:08:05Z

Summary

Fast serving for Wan on TPU: a persistent per-shape AOT executable cache, a converted-weights cache, and a tuned launcher (end_to_end/tpu/run_wan_fast_inference.sh).

User-visible changes

AOT caching is opt-in.
- By default generate_wan.run() is plain jax.jit.
- A persistent cache is used only when aot_cache_dir is set and the source revision is reusable. An explicit dirty:/unversioned: revision or enable_zero_execution_warmup: True (default False in all 6 Wan ymls) uses an ephemeral temp dir, torn down in a try/finally covering setup and install.
- aot_cache_dir must be a local path: gs:// raises ValueError regardless of revision.
Cache formats:
- .aotx format v2 (v1 files are recompiled once).
- Converted-weights cache format v3 (v1 caches are re-converted once with a warning and disk-space check; matching v2 manifests migrate header-in-place to v3).

What's in it

AOT signatures & metadata: keyed on scalar types, static values, a deterministic GraphDef digest, device kind, package versions, XLA/libtpu flags, PRNG/x64 settings, and a content hash of every non-test .py in maxdiffusion. Per-host .aotx files for multi-host.
.aotx loading: exact (module, name) pickle allowlist rejecting dotted names (defence in depth; directory write permissions are the real control).
Converted-weights fingerprint: repo id + HF snapshot revision (before resolving symlinks) + subfolder + index contents, fail-closed when missing, checked with local_files_only before network calls.
Warmup & Launcher: compile_experts compiles both experts without executing when AOT is installed; InflightWindow bounds queued steps; platform-detected v6e/v7/generic profiles with opt-in PIN_DVFS_P_STATE.

Performance

Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps (CP=4, DP=2), measured with this PR at top of stack:

TPU	Mode / DVFS	Load (`COMPILE_TE=true`)	Compile (`compile_experts`)	Generate	Denoise
v6e-8	Cold AOT (unpinned)	30.5 s	31.0 s	138.2 s	136.0 s
v6e-8	Warm AOT (unpinned)	34.7 s (vs ~200 s uncached HF)	3.1 s (vs 193.5 s cold JIT)	138.1 s	135.9 s
tpu7x-8	Warm AOT (unpinned)	—	—	115.5 s	112.8 s
tpu7x-8	Warm AOT (pinned)	—	—	102.0 s	99.4 s

Tests

aot_cache_test.py (41), converted_weights_cache_test.py (16), wan_transformer_test.py (15), wan_warmup_coverage_test.py (9 def test_; 14 with subtests).

Covers unpickler allowlist, run() AOT gating and teardown (including install failure), gs:// rejection, float32_qk_product keeping bf16 QK operands, fail-closed manifests, disk-space checks, revision/subfolder fingerprinting, and real TPU executable reload.
CI vs Local VM: 4 heavy wan_transformer tests skip in GITHUB_ACTIONS=true (40 of 199 cumulative def test_, or 204 counting subtests); all 57 AOT/converted-weights and 9 warmup tests (14 with subtests) run in CI.
Results (TPU v6e-8): 199 def test_ passed (204 counting the 5 warmup subtests; 81 def test_ / 86 with subtests in this PR's 4 test files).

Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491

github-actions · 2026-09-13T06:08:25Z

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

gemini-code-assist

Code Review

This pull request introduces several optimizations and features for Wan model inference, including a fast-path signature cache in the AOT cache, support for fused LayerNorm and AdaLN kernels, and updated configuration options. The review feedback highlights a potential concurrency issue in the signature cache that could lead to race conditions, suggests moving inline imports of the fused kernel to the top of the file to reduce overhead in hot paths, and recommends avoiding os.path.join on GCS URIs to ensure cross-platform safety.

eltsai · 2026-09-23T23:27:38Z

I was trying to reproduce the results, and I got:

Configuration	Total inference	Denoise (40 steps)	VAE decode	Warm compile
Baseline — `origin/main@1bc54811`	110.8 s	107.1 s	1.5 s	19.3 s
+ #477 + #478 + #479	101.2 s	98.7 s	0.7 s	7.5 s
+ #488 (full stack)	99.4 s	96.8 s	0.7 s	7.6 s

Deltas vs baseline

Configuration	Total inference	Denoise	VAE decode
#477 + #478 + #479	−9.6 s (−8.7%)	−8.4 s (−7.8%)	−0.8 s (−53%)
+ #488	−11.4 s (−10.4%)	−10.3 s (−9.6%)	−0.8 s (−53%)

this is on v7x-8. The stats are slightly different, but I think the perf improvement is bigger than 3.8%

mbohlool

LGTM after addressing comments.

…ce 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.

Perseus14 requested a review from entrpn as a code owner September 13, 2026 06:08

gemini-code-assist Bot reviewed Sep 13, 2026

View reviewed changes

Perseus14 force-pushed the feat/wan-fast-serving branch from c9dc7b2 to 80719d9 Compare September 13, 2026 06:16

Perseus14 force-pushed the feat/wan-fast-serving branch from 80719d9 to 585247a Compare September 13, 2026 06:25

Perseus14 force-pushed the feat/wan-fast-serving branch from 585247a to 2ee9c66 Compare September 13, 2026 06:30

Perseus14 force-pushed the feat/wan-fast-serving branch from 2ee9c66 to 8756432 Compare September 13, 2026 06:49

Perseus14 force-pushed the feat/wan-fast-serving branch from 8756432 to 2957831 Compare September 13, 2026 07:15

Perseus14 force-pushed the feat/wan-fast-serving branch 2 times, most recently from 8eb1b56 to dc78d28 Compare September 13, 2026 15:29

Perseus14 force-pushed the feat/wan-fast-serving branch from dc78d28 to e98f895 Compare September 13, 2026 15:40

Perseus14 force-pushed the feat/wan-fast-serving branch from e98f895 to ba40a77 Compare September 13, 2026 16:21

Perseus14 force-pushed the feat/wan-fast-serving branch from ba40a77 to d7361d3 Compare September 13, 2026 16:38

Perseus14 force-pushed the feat/wan-fast-serving branch from d7361d3 to 3540cb7 Compare September 13, 2026 18:29

Perseus14 force-pushed the feat/wan-fast-serving branch from 3540cb7 to 89af7b2 Compare September 13, 2026 18:53

Perseus14 force-pushed the feat/wan-fast-serving branch from 89af7b2 to a1e652a Compare September 13, 2026 19:09

Perseus14 force-pushed the feat/wan-fast-serving branch from a1e652a to 838af4b Compare September 13, 2026 19:49

Perseus14 force-pushed the feat/wan-fast-serving branch from 8631d1f to a909245 Compare September 17, 2026 12:31

Perseus14 force-pushed the feat/wan-fast-serving branch 2 times, most recently from 3a0c913 to a415017 Compare September 17, 2026 14:34

syhuang22 previously approved these changes Sep 17, 2026

View reviewed changes

github-actions Bot added the pull ready label Sep 17, 2026

Perseus14 force-pushed the feat/wan-fast-serving branch from a415017 to 6b85dfc Compare September 17, 2026 18:41

Perseus14 force-pushed the feat/wan-fast-serving branch from 6b85dfc to 7e38c04 Compare September 17, 2026 18:51

Perseus14 force-pushed the feat/wan-fast-serving branch from 7e38c04 to 1d65bfa Compare September 17, 2026 19:07

Perseus14 self-assigned this Sep 17, 2026

Perseus14 requested a review from eltsai September 17, 2026 19:09

Perseus14 mentioned this pull request Sep 22, 2026

feat(wan): fused RMSNorm+RoPE Pallas producer; Wan switches in config, not env #488

Open

eltsai reviewed Sep 23, 2026

View reviewed changes

Comment thread src/maxdiffusion/generate_wan.py

eltsai previously approved these changes Oct 2, 2026

View reviewed changes

eltsai previously approved these changes Oct 5, 2026

View reviewed changes

mbohlool previously approved these changes Oct 19, 2026

View reviewed changes

Navigation Menu

Repository navigation

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

feat(wan): fast serving with persistent AOT caching and tuned inference recipe#479
Perseus14 wants to merge 1 commit into
feat/fused-rmsnorm-rope-xlafrom
feat/wan-fast-serving

Perseus14 commented Sep 13, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 13, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

eltsai commented Sep 23, 2026

Uh oh!

Uh oh!

mbohlool left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

Perseus14 commented Sep 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

User-visible changes

What's in it

Performance

Tests

Uh oh!

github-actions Bot commented Sep 13, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

eltsai commented Sep 23, 2026

Deltas vs baseline

Uh oh!

Uh oh!

mbohlool left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Perseus14 commented Sep 13, 2026 •

edited

Loading