feat(wan): fast inference optimizations (fused cross-attention, token patch embed, shard-major A2A, lane padding)

[Perseus14]

Summary

Four Wan 2.2 inference optimizations. Each is a config switch that defaults off in the yml; the launcher's v6e and v7 profiles turn them on.

Switch	What it does
wan_cross_attn_kernel="pallas"	Text cross-attention kernel keeping one head block's 512-token KV resident in VMEM with an exact single-pass softmax (no online rescaling). Falls back to XLA (logging once under separate warn keys for input vs VMEM) when shapes/sharding don't fit, the TPU isn't validated (v6e/tpu7x), or estimate_vmem_bytes (~42 MiB at Wan shard shape, validated at KV=512) exceeds the 64 MiB budget.
wan_patch_embed_mode="tokens"	Patch embedding as patchify + matmul on the token-sharded layout, avoiding the activation all-gather.
wan_ulysses_out_a2a="shard_major" (alias chunked)	One layout-preserving all-to-all over a (num_shards, chunk) axis, avoiding XLA relayout copies. With the custom splash kernel, writes shard-major output directly on both v6e (1D Ulysses) and tpu7x (2D Ulysses+Ring).
wan_seq_pad="lane"	Pads video tokens to a per-shard multiple of 128 (75,600 → 75,776). Pad keys are masked in custom Ulysses/ring kernels and pad query rows are excluded from the fixed-m Q-norm bound. Paths that cannot mask padding raise NotImplementedError.

Other changes:

• All switches plumb through attention_config_entries(config); pyconfig no longer sets Wan options for non-Wan models.
• The launcher's generic profile no longer sets ici_ag_pipelining / ag_backward_pipelining.

Performance

Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps, launcher defaults (fused RoPE on), warm AOT, DVFS unpinned:

TPU	Denoise	Generate
v6e-8	125.9 s	128.05 s
tpu7x-8	105.1 s	107.6 s

• tpu7x-8 with DVFS pinned (PIN_DVFS_P_STATE=true, fused RoPE off): 93.6 s denoise / 96.1 s generate.
• Output videos were visually inspected on both platforms: clean, sharp, no artifacts.

Tests

cross_attention_pallas_test.py (39), wan_seq_pad_test.py (17, including U=2×R=2 aligned real_len + chunked shard-major layout), wan_patch_embed_test.py (7), ulysses_out_a2a_test.py (5), wan_runtime_options_test.py (12).

• Pallas-active dispatch tests skip on unvalidated TPUs (vmem_budget_is_validated), while test_falls_back_to_xla_on_unvalidated_tpu_kind covers the fallback.
• CI vs Local VM (end_to_end/tpu/run_wan_stack_tests.sh):
  • Full local runner on TPU v6e-8: 350 passed, 1 skipped (off-TPU fallback only) in 13m 44s (824.7 s).
  • With GITHUB_ACTIONS=true (all 17 stack test files): 231 passed, 120 skipped in ~2m 16s.
  • Emulated unvalidated TPU (device_kind="TPU v4", rope_accum_is_measured=False): 0 failures across all 17 stack test files.

Stack: #477 → #478 → #479 → #488 → #491

[github-actions]

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 performance optimizations for Wan models on TPU, including a fused short-KV cross-attention Pallas kernel, token-sharded patch embedding, inverse Ulysses all-to-all optimizations, and token sequence padding. A critical issue was identified in the new Pallas cross-attention kernel where the BlockSpec index mapping fails to scale grid indices by their block sizes, which would result in incorrect memory slicing and wrong outputs.

[eltsai]

thanks for these stacks of PRs @Perseus14 ! I wonder if we can do an experiment on videos with lengths of 2.5 sec, 5 sec (already covered), 7.5 sec, and 10 sec to see if we would see any regression on other seq length?

[Perseus14]

Thanks @eltsai — ran the sweep at 41 / 81 / 121 / 161 frames (2.5 / 5 / 7.5 / 10 s @ 16 fps), 720p, 40 steps, seed 12345, same prompt, DP=2×CP=4, DVFS unpinned, on tpu7x-8 and v6e-8.

A = main (10af306), B = this stack with the #491 switches off (= #477→#488), C = this stack with the launcher defaults. Denoise seconds for 40 steps (load/compile/VAE excluded):

frames	tpu7x A	tpu7x B	tpu7x C	C vs A	v6e A	v6e B	v6e C	C vs A
41	56.5	43.0	41.5	1.36×	58.8	53.0	49.6	1.19×
81	123.0	109.1	105.0	1.17×	144.6	132.8	125.9	1.15×
121	234.2	216.2	212.7	1.10×	290.8*	308.1	313.0 / 254.6*	1.14×*
161	385.0	362.1	356.5	1.08×	454.7	430.4	419.2	1.08×

No regression vs main at any length; the gain narrows as attention dominates the step. The #491 switches themselves (B→C) are +1.6–3.9 % on tpu7x and +2.7–6.9 % on v6e, except v6e @ 121 f (−1.6 %).

Note that all runs use the launcher's 81 f-tuned flash block sizes (v6e: BQ 9472 / BKV 1024; tpu7x: BQ 4736 / BKV 2048). We did not do a block-size sweep per length: the sequence length (and therefore the per-shard length and its alignment to BQ) differs for each frame count, so the optimal block sizes need to be found separately for each frame choice. The non-81 f numbers above are therefore conservative for all three configurations.

* v6e @ 121 f is the one length where the launcher's v6e profile (Ulysses U=4, BQ 9472) is the wrong pick: the per-shard length (27,900 tokens) is not block-aligned, and the 2D-ring recipe that is also in this stack (ATTENTION=ulysses_ring_custom_fixed_m ULYSSES_SHARDS=2 BQ=4736 BKV=2048) is 19 % faster there (254.6 s; main with its own ring recipe: 290.8 s, with its Ulysses kernel 657.6 s). At 41 / 81 / 161 f Ulysses remains the better v6e choice by 1–37 %.

All runs were bit-reproducible across cold/warm AOT (identical md5s), and the outputs are visually clean on both platforms.