Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
c9dc7b2 to
80719d9
Compare
80719d9 to
585247a
Compare
585247a to
2ee9c66
Compare
2ee9c66 to
8756432
Compare
8756432 to
2957831
Compare
8eb1b56 to
dc78d28
Compare
dc78d28 to
e98f895
Compare
e98f895 to
ba40a77
Compare
ba40a77 to
d7361d3
Compare
d7361d3 to
3540cb7
Compare
3540cb7 to
89af7b2
Compare
89af7b2 to
a1e652a
Compare
a1e652a to
838af4b
Compare
83f15bb to
5cb0364
Compare
syhuang22
left a comment
syhuang22
left a comment
There was a problem hiding this comment.
Tried this on v7x-8 too (same WAN 2.2 720p/81f/40 steps setup as my #478 comment):
- The new default (per_q_block with U=4 on CP=4) is really R=1, so there's no ring. It comes in at 115.8s vs the current default's 105.3s (+10%). Can we keep the v7 default and add the v6e recipe as a separate option?
- The -1.09% / bit-identical numbers look like v6e R=1 only. On v7 the output differs from main (the fused norm in #478), so mind scoping the description?
- The description mentions dual-transformer warmup, but I don't see warmup changes in this diff, and main already primes both transformers since #452. Am I missing something?
| # ulysses_custom_fixed_m: 3.5338 s/step | ||
| if [ -z "${ATTENTION:-}" ]; then | ||
| if [ "${FIXEDM:-1}" = "1" ]; then | ||
| ATTENTION="ulysses_ring_custom_fixed_m_per_q_block" |
There was a problem hiding this comment.
On v7 this default is ~10% slower than what we have now (115.8s vs 105.3s), and U=4 on CP=4 means R=1, so no ring at all. Could this be a v6e-only branch?
There was a problem hiding this comment.
Fixed in 09866a3e (run_wan_fast_inference.sh:L171-L200)! The script now branches on TPU_PROFILE:
v7defaults toulysses_ring_custom_fixed_mwithU=2(R=2onCP=4),BQ=6400,BKV=2048(achieving 114.8s E2E on v7x-8; 101.4s Denoise and 102.1s E2E with p_state=7).v6edefaults toulysses_custom_fixed_m_per_q_blockwithU=4(R=1onCP=4),BQ=9472,BKV=1024(achieving 138.7s E2E on v6e-8).
| export LIBTPU_INIT_ARGS="${LIBTPU_INIT_ARGS} ${EXTRA_LIBTPU:-}" | ||
| # Tuned collective/scheduler flag set for TPU. | ||
| # One line: libtpu stops parsing at a literal backslash. | ||
| DEFAULT_LIBTPU="--xla_tpu_spmd_rng_bit_generator_unsafe=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_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=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" |
There was a problem hiding this comment.
Three v7 flags got dropped here (async_collective_fusion, ..._fuse_all_gather, ..._multiple_steps). Intentional? Our v7 numbers were measured with them on.
There was a problem hiding this comment.
Restored in 09866a3e (run_wan_fast_inference.sh:L136)! On v6e (Ghostlite), --xla_tpu_enable_async_collective_fusion=true aborts backend initialization (FAILED_PRECONDITION: Continuation fusion for AllGather is enabled on platform ghostlite), which is why it was disabled for v6e. We split V6E_LIBTPU and V7_LIBTPU so v7 explicitly enables all three flags (--xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true) plus megacore fusion.
| # fast serving recipe (persistent AOT + relayout fusion + unpadded KV): 136.5s (3.4125 s/step) | ||
| # baseline denoise (unfused, default recipe): 138.0s (3.4500 s/step) | ||
| # unoptimized baseline (cold/AOT disabled): 139.8s (3.4950 s/step) | ||
| # Bit-identical output (MD5: 4c18d4ce340e0da3a1d4c617926ad525) applies to the tested 40-step generation. |
There was a problem hiding this comment.
This md5 only holds against your own baseline. Against main on v7 the output differs (fused RMSNorm in #478).
There was a problem hiding this comment.
With fused_rmsnorm_rope updated in #478 to match Flax's x * (rsqrt * scale) order, online attention is bit-identical to main on v7 again. Also updated the comments here (L152-L165) to explicitly document that fixed-m attention reorders float accumulation (giving +3.25 dB PSNR vs exact online softmax and VBench within 0.25%) rather than claiming bit-identical md5 against online softmax.
| detected_revision = commit_hash if commit_hash is not None else max_utils.get_git_commit_hash(check_dirty=True) | ||
| source_revision = _resolve_wan_aot_source_revision(config, detected_revision) | ||
| aot_cache_dir = getattr(config, "aot_cache_dir", "") | ||
| if aot_cache_dir and not _is_reusable_aot_revision(source_revision): |
There was a problem hiding this comment.
Any untracked file makes the tree dirty and turns persistent AOT off, and keying on the commit throws the cache away on every commit, even docs-only ones. Maybe hash the relevant source files instead?
There was a problem hiding this comment.
Great suggestion — implemented both improvements:
- Source content hashing (
generate_wan.py:L47-L93): Added_compute_wan_source_hash(), which hashes the sorted.pyfiles that actually affect Wan's lowered HLO (aot_cache.py,models/attention_flax.py,models/wan/,pipelines/wan/, andkernels/— 32 files, takes ~8.8 ms). Persistent AOT cache keys onsrc:<sha256[:16]>by default (unless overridden byaot_build_revision), so docs-only commits, changes to other models (flux,sdxl,ltx), or local edits to unrelated files never invalidate the AOT cache. - Untracked files ignored (
max_utils.py:L478): Passed--untracked-files=notogit status --porcelaininget_git_commit_hash(check_dirty=True)so untracked logs, videos, or scratch scripts never mark the working tree dirty.
| return compiled(leaves) | ||
| except Exception as e: # noqa: BLE001 - any failure means "use jit" | ||
| max_logging.log(f"[aot] {self.name}: compiled call failed ({e}); using jit") | ||
| max_logging.log(f"[aot] fast-path execution failed for {self.name}/{signature} ({e}); falling back to JIT") |
There was a problem hiding this comment.
With _align_inputs inside this try, a sharding mismatch now quietly falls back to JIT (a recompile mid-serving) with just a log line. I'd keep align errors loud.
There was a problem hiding this comment.
Moved flat = self._align_inputs(compiled, leaves) outside the try block (aot_cache.py:L236) and removed the blanket try/except inside _align_inputs (L259). Structural input pruning by XLA (len(flat_expected) != len(leaves)) still returns None cleanly, while any sharding mismatch or device_put error now raises loudly.
5cb0364 to
09e9a4a
Compare
526a9b5 to
905362a
Compare
09866a3 to
8631d1f
Compare
|
Thanks for testing PR #479 on v7x-8 @syhuang22! Updated the code and addressed your comments! |
8631d1f to
a909245
Compare
3a0c913 to
a415017
Compare
a415017 to
6b85dfc
Compare
6b85dfc to
7e38c04
Compare
7e38c04 to
1d65bfa
Compare
…ce 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 - 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)
2 participants
Summary
Stacked on top of #478 (
feat/ring-attention).Delivers production-grade fast serving for Wan 2.2 T2V-A14B on Cloud TPU:
compile: 76.6s cold → 8.3s warm ontpu7x-8; 67.6s cold → 8.65s–12.2s warm onv6e-8)._compute_wan_source_hash()) across 32 Wan/attention/kernel Python source files and includes--untracked-files=noinget_git_commit_hashso code edits automatically invalidate stale caches._align_inputsoutside the broadtry/exceptblock inaot_cache.pyso input sharding mismatches fail loudly instead of silently triggering recompilation.end_to_end/tpu/run_wan_fast_inference.sh):BQ=9472, BKV=1024, BKV_COMPUTE=512, BKV_COMPUTE_IN=512) and Cloud TPU v7 (tpu7x,BQ=6400, BKV=2048, BKV_COMPUTE=2048, BKV_COMPUTE_IN=2048), and tuned XLA scheduling flags.vn_local(_ring_fixed_m_norms_pre_a2a): Folds the K max-norm and 0-D scalar V max-norm (vn_local = (value.astype(jnp.float32) ** 2).max()) into the single(ulysses, ring)pmaxand removes the ring-axisall_gatherand per-headdynamic_sliceonvn_global, eliminating XLA relayout barriers inside the QKV projection fusion region and saving ~3.8s of compute fusion time._lse_scan): Replaces per-hop dynamic indexing (my_ring_index - hop) with hoistedmk_global_sq, allowing the ~190 MiB K/Vppermuteinsidelax.condto overlap kernel execution (collective-permute-done: 2405.4 ms → 47.3 ms; outputall-to-allBW: 173 GiB/s → 535 GiB/s), saving ~5.7s denoise time.kv_pad_size = 1): Eliminates explicitjnp.padHBM copies (37,800 -> 38,912), reducesppermuteICI payload by 1,112 tokens per shard (−2.86%), and slices the 19th KV block atslice_k_len = 936 <= 1024, skipping 384,000 inner VPU/MXU loop iterations across 40 steps (saving 0.4s denoise and 2.9s cold compile time with 100% bit-identical output).p_t == 1,transformer_wan.py): Collapses the unitary temporal patch axis8Dstride permutation(0, 7, 1, 4, 2, 5, 3, 6)6Dtranspose(0, 6, 1, 2, 4, 3, 5)), saving ~0.17s with 100% bit-exact output (max(abs(diff)) == 0.0).Verified End-to-End Latency Benchmarks (40 steps, 720p, 81 frames, identical configs)
Both platforms were benchmarked end-to-end on full 40-step Wan 2.2 T2V-A14B inference against
origin/main(1bc54811) under identical configurations (vae_decode_chunk=1, vae_spatial=8):1. Cloud TPU v6e-8 (Profile$U=4, R=1, B_Q=9472, B_{KV}=1024, B_{KV_COMPUTE}=512, CP=4, DP=2$ )
v6e,origin/main(1bc54811)a48e0815)138.36 s)8.65 sbest)2. Cloud TPU v7 (tpu7x-8) (Profile$U=2, R=2, B_Q=6400, B_{KV}=2048, B_{KV_COMPUTE}=2048, CP=4, DP=2$ )
v7,origin/main(1bc54811)a48e0815, Default)a48e0815,p_state 7)origin/main115.2 sbest)101.86 sbest)p_state 7p_state 7seed=12345)9fb82850...9530b119ab439714...9530b119ab439714...Note: v7x-8 denoise numbers verified across multiple warm & cold repetitions on two independent
tpu7x-8hosts (europe-west2-aandus-central1-c), both with default DVFS and withEXTRA_LIBTPU="--xla_tpu_dvfs_p_state=7".XProf Collective & Compute Profile Breakdown (tpu7x-8, 40 steps)
origin/mainpmax+ Hoisted Metadataa48e0815, Default /p_state 7)collective-permute-done(2-step)all-to-allBWp_state 7)Numerical Accuracy & Disclosures
ulysses_ring_customonline softmax without fixed-m, 40 steps, 81 frames,seed=12345):a48e0815) vs Exact: 20.84 dB PSNR (mean abs diff 10.52 — +3.25 dB closer to exact online softmax, 38% lower error)_lse_scan) evaluates per-hop fixed-m eligibility against the ring-wide max K norm (mk_global_sq) rather than the per-hop shard norm, the per-hop gate is strictly more conservative and numerically closer to exact online softmax (+3.25 dB PSNR).seed=12345.per_q_block=Trueremains supported and tested in unit tests, but requires a dedicated flash block size tile sweep before enabling in production on v7x-8.Verification
pytest src/maxdiffusion/tests/ring_fixed_m_test.py: 37 passed, 1 skipped in 261.62s ontpu7x-8(includingtest_gqa_ulysses_ring_custom_fixed_m_ragged).pytest src/maxdiffusion/tests/custom_splash_fixed_m_test.py: 23 passed.pytest src/maxdiffusion/tests/wan/wan_transformer_test.py: 58 passed, 1 skipped.pytest src/maxdiffusion/tests/attention_config_guards_test.py: 16 passed.origin/main(1bc54811).pyink --pyink-indentation=2 --line-length=125andruff checkclean.