[Perseus14]

Two memory fixes for Wan 2.2 T2V at per-device batch >= 1, from a B=1..16 batch-scaling sweep at 720p / 81 frames / 40 steps on v6e-8 and tpu7x-8.

1. Cross-attention KV cache is batch-sharded and scanned (transformer_wan.py)

compute_kv_cache ran an eager nnx.vmap over all 40 layers on replicated prompt embeds, so every device materialised the full CFG-doubled cache plus the FP32 RMSNorm temporaries for every layer at once (~6.8 GiB per video per device at 720p). It is now a single jitted executable that pins the batch axis to ('data', 'fsdp') on input and output and nnx.scans the layers one at a time. The cross-attention consumer already re-constrains K/V to (BATCH, CROSS_ATTN_HEAD, None, D_KV), so nothing is resharded at use time — each device simply holds its own B_local slice.

If the batch is not divisible by data * fsdp (e.g. data_parallelism > 1 with per_device_batch_size = 1 and CFG off — the same case the pipelines already tolerate for the CFG cache), the batch constraint is skipped and the cache stays replicated, as before, instead of raising inside the jit.

2. vae_spatial=-1 picks a batch-aware VAE mesh for T2V (wan_pipeline.py)

The VAE mesh is ('redundant', 'vae_spatial'); batch is sharded on redundant, width / conv channels on vae_spatial. With the old auto value (vae_spatial = N) the batch axis is fully replicated, so at B = N every chip holds all N videos' activations and the causal conv stack pays ~10x (and OOMs on 32 GiB parts). Auto now resolves to (N, 1) — one video per device, no cross-chip collectives — when the global batch is a multiple of the device count, and to the unchanged (1, N) otherwise. Mixed meshes were measured and are not a win (at B=2 on v6e-8, (2, 4) needed a 21.2 GiB decode program vs 15.5 GiB for (1, 8)). An explicit vae_spatial still overrides.

Scope: the batch-aware choice is gated by a new batch_parallel_vae argument on _create_common_components, enabled only by the T2V pipelines (wan_pipeline_2_1, wan_pipeline_2_2). The VAE hard-constrains every activation to P('redundant', …, 'vae_spatial', …), so a (N, 1) mesh requires every tensor that goes through the VAE to have a batch divisible by N. That holds for T2V (the only VAE call is the final full-batch decode) but not for VACE (reference images are encoded one at a time, [1, 1, H, W, C]) or the i2v / animate conditioning encodes. Those pipelines keep the previous (1, N) mesh, so nothing changes for them.

Measurements (on main; DP=2, CP=4, use_kv_cache=true unless noted)

Platform	B		Compile	Denoise	VAE decode (TPU)	Peak HBM / device
tpu7x-8	2	main	229.9s	270.5s	2.2s	27.49 GiB
tpu7x-8	2	this PR	193.0s	269.8s	2.2s	16.65 GiB
tpu7x-8	8	main	265.8s	897.8s	41.6s	68.12 GiB
tpu7x-8	8	this PR	233.7s	893.4s	6.3s	24.32 GiB
v6e-8 (DP=1, FSDP=2, CP=4, vae_decode_chunk=1)	8	main	—	—	—	OOM in compute_kv_cache (6.25 GiB alloc, 4.99 GiB free)
v6e-8 (same)	8	this PR	192.2s	1302.6s	4.8s	17.41 GiB — completes

Denoise time is unchanged within noise; the wins are −64% peak HBM and 6.6x faster VAE decode at B=8 on tpu7x-8, and B=8 with the KV cache becoming possible at all on v6e-8.

Validation

• CPU checks on an 8-device mesh (DP=2, CP=4): scanned cache output is P(None, 'data') with B_local shards and matches a per-layer eager reference to 1e-5.
• New tests/wan/wan_vae_mesh_test.py (CPU, 8 host devices): _auto_vae_spatial table for B in {0,1,2,3,4,8,12,16} with the gate on and off, a guard that only the T2V pipelines set batch_parallel_vae=True, and _batch_sharding / _shard_kv_cache_on_batch layout checks including the non-divisible-batch fallback.
• wan_kv_cache_test.py is TPU-only (attention=flash); not runnable on CPU before or after this change.
• pyink clean. Note: the Linter workflow is currently red on main for two unrelated files; style: fix pyink formatting on main (2 files) #503 fixes that.

[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 optimizes memory usage and performance for the Wan model by replacing eager nnx.vmap with a jitted, scanned, and batch-sharded implementation (_scanned_compute_kv) for cross-attention KV cache computation. It also introduces dynamic vae_spatial mesh extent selection in the VAE pipeline to prevent OOMs and improve decoding speed. A critical issue was identified in the JIT compilation of _scanned_compute_kv, where blocks_graphdef must be marked as a static argument to avoid a JAX TypeError since it is not a valid JAX PyTree.