From 98ea89f57d5cc2f365821b7d8c188fbc121e2f2a Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Sun, 25 Oct 2026 20:23:34 +0000 Subject: [PATCH] perf(wan): batch-sharded cross-attn KV cache and batch-aware VAE mesh Two memory fixes for Wan 2.2 T2V at per-device batch >= 1, motivated by a B=1..16 batch-scaling sweep at 720p/81f/40 steps on v6e-8 and tpu7x-8. 1. Cross-attention KV cache (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. It is now one jitted executable that pins the batch axis to ('data', 'fsdp') on input and output and scans the layers one at a time. The cross-attn consumer already re-constrains K/V to that layout, so nothing is resharded at use time; each device just 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 case the pipelines already tolerate for the CFG cache) the batch constraint is skipped and the cache stays replicated as before, rather than raising inside the jit. Sharding the 512-token text axis on 'context' as well was tried and rejected: XLA hoists the per-layer all-gather out of the layer scan, so the forward program's reservation grew by 5.8 GiB (net loss on v6e-8). 2. VAE mesh (wan_pipeline.py), T2V pipelines only With vae_spatial=-1 (auto) the VAE mesh is now ('redundant', 'vae_spatial') = (N, 1) when the global batch is a multiple of the device count N, i.e. one video per device with no cross-chip collectives, and (1, N) otherwise (unchanged). Mixed meshes were measured and are not a win: at B=2 on v6e-8 a (2, 4) mesh needed a 21.2 GiB decode program vs 15.5 GiB for (1, 8). An explicit vae_spatial still overrides. The batch-aware choice is gated by a new batch_parallel_vae argument on _create_common_components, enabled only by wan_pipeline_2_1 / 2_2. The VAE hard-constrains activations to P('redundant', ..., 'vae_spatial', ...), so a (N, 1) mesh requires every tensor through the VAE to have a batch divisible by N. That holds for T2V (only the final full-batch decode) but not for VACE (reference images encoded one at a time, [1, 1, H, W, C]) or the i2v / animate conditioning encodes; those keep the (1, N) mesh. Tests: new tests/wan/wan_vae_mesh_test.py (CPU, 8 host devices) covers the _auto_vae_spatial table for B in {0..16} gated and ungated, a guard that only the T2V pipelines set the gate, and _batch_sharding / _shard_kv_cache_on_batch layouts including the non-divisible fallback. Measured on main (DP=2, CP=4 unless noted, use_kv_cache=true): tpu7x-8 B=2 peak HBM 27.49 -> 16.65 GiB, denoise 270.5 -> 269.8 s tpu7x-8 B=8 peak HBM 68.12 -> 24.32 GiB, VAE decode 41.6 -> 6.3 s, denoise 897.8 -> 893.4 s v6e-8 B=8 (DP=1, FSDP=2, CP=4, vae_decode_chunk=1) main OOMs in compute_kv_cache (6.25 GiB alloc, 4.99 GiB free); with this change it completes: peak 17.41 GiB, denoise 1302.6 s, VAE decode 4.8 s. --- .../wan/transformers/transformer_wan.py | 106 ++++++++++++-- .../pipelines/wan/wan_pipeline.py | 45 +++++- .../pipelines/wan/wan_pipeline_2_1.py | 3 + .../pipelines/wan/wan_pipeline_2_2.py | 3 + .../tests/wan/wan_vae_mesh_test.py | 132 ++++++++++++++++++ 5 files changed, 278 insertions(+), 11 deletions(-) create mode 100644 src/maxdiffusion/tests/wan/wan_vae_mesh_test.py diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py index f147e4414..31a2d7f90 100644 --- a/src/maxdiffusion/models/wan/transformers/transformer_wan.py +++ b/src/maxdiffusion/models/wan/transformers/transformer_wan.py @@ -16,6 +16,7 @@ from typing import Tuple, Optional, Dict, Union, Any import contextlib +import functools import math import jax import jax.numpy as jnp @@ -38,6 +39,90 @@ from ...gradient_checkpoint import GradientCheckpointType BlockSizes = common_types.BlockSizes +BATCH = common_types.BATCH + + +def _batch_sharding(mesh, ndim: int, batch_axis: int, batch_size: Optional[int] = None): + """NamedSharding with the activation-batch mesh axes on ``batch_axis``, rest replicated. + + Returns None (caller skips the constraint, array stays replicated) when there + is no mesh, when the ambient ``axis_rules`` context does not map BATCH to any + mesh axis, or when ``batch_size`` is not divisible by the product of the + batch mesh axes -- e.g. ``data_parallelism > 1`` with ``per_device_batch_size + = 1`` and CFG off, which the pipelines already tolerate for the CFG cache. + ``with_sharding_constraint`` would otherwise raise inside the jit. + """ + if mesh is None: + return None + batch_mesh_axes = nn.logical_to_mesh_axes((BATCH,))[0] + if batch_mesh_axes is None: + return None + if batch_size is not None: + axes = (batch_mesh_axes,) if isinstance(batch_mesh_axes, str) else tuple(batch_mesh_axes) + num_shards = math.prod(mesh.shape[a] for a in axes) + if batch_size % num_shards != 0: + return None + spec = [None] * ndim + spec[batch_axis] = batch_mesh_axes + return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec)) + + +def _shard_kv_cache_on_batch(kv_cache, mesh): + """Pins every K/V cache array to the activation batch sharding. + + Cache arrays are either [B, S, H*D] (one layer) or [L, B, S, H*D] (stacked), + so the batch axis is the first non-layer axis. The cross-attention consumer + re-constrains K/V to (BATCH, CROSS_ATTN_HEAD, None, D_KV), i.e. batch on + ('data', 'fsdp') and everything else replicated, so this is exactly the + layout it needs — no resharding at use time, and each device only holds its + own B_local slice instead of the full CFG-doubled batch. + """ + if mesh is None: + return kv_cache + + def _constrain(x): + batch_axis = 1 if x.ndim == 4 else 0 + sharding = _batch_sharding(mesh, x.ndim, batch_axis, batch_size=x.shape[batch_axis]) + return x if sharding is None else jax.lax.with_sharding_constraint(x, sharding) + + return jax.tree.map(_constrain, kv_cache) + + +@functools.partial(jax.jit, static_argnames=("num_layers", "mesh")) +def _scanned_compute_kv(blocks_graphdef, blocks_state, encoder_hidden_states, encoder_attention_mask, num_layers, mesh): + """Cross-attention K/V cache for all layers as ONE sharded executable. + + Replaces the eager ``nnx.vmap`` over the stacked blocks. Eager vmap + materialises the K/V projections, the FP32 RMSNorm upcast and the bf16 + downcast for all ``num_layers`` at once *and* on every device (the inputs + are replicated), which is what blew the 32 GiB parts at B=8. + + Here the batch axis is sharded on ('data', 'fsdp') before the projections, + and ``nnx.scan`` processes one layer per iteration so the FP32 temporaries + are bounded to a single layer and freed before the next. The stacked + [L, B, S, H*D] outputs stay batch-sharded. + """ + blocks = nnx.merge(blocks_graphdef, blocks_state) + batch_size = encoder_hidden_states.shape[0] + in_sharding = _batch_sharding(mesh, encoder_hidden_states.ndim, 0, batch_size=batch_size) + if in_sharding is not None: + encoder_hidden_states = jax.lax.with_sharding_constraint(encoder_hidden_states, in_sharding) + if encoder_attention_mask is not None: + encoder_attention_mask = jax.lax.with_sharding_constraint( + encoder_attention_mask, _batch_sharding(mesh, encoder_attention_mask.ndim, 0, batch_size=batch_size) + ) + + def scan_fn(carry, block): + kv = block.compute_kv(encoder_hidden_states, encoder_attention_mask) + return carry, _shard_kv_cache_on_batch(kv, mesh) + + _, kv_cache = nnx.scan( + scan_fn, + length=num_layers, + in_axes=(nnx.Carry, 0), + out_axes=(nnx.Carry, 0), + )(None, blocks) + return _shard_kv_cache_on_batch(kv_cache, mesh) def get_frequencies(max_seq_len: int, theta: int, attention_head_dim: int): @@ -588,6 +673,7 @@ def __init__( out_channels = out_channels or in_channels self.num_layers = num_layers self.scan_layers = scan_layers + self.mesh = mesh self.enable_jax_named_scopes = enable_jax_named_scopes attention_config = { "use_base2_exp": False, @@ -741,16 +827,17 @@ def compute_kv_cache( encoder_attention_mask = jnp.concatenate([encoder_attention_mask, text_mask], axis=1) if self.scan_layers: - - @nnx.vmap( - in_axes=(0, None, None), - out_axes=0, - transform_metadata={nnx.PARTITION_NAME: "layers_per_stage"}, + # Jitted + scanned + batch-sharded; see _scanned_compute_kv for why the + # previous eager nnx.vmap was a 6.8 GiB-per-video HBM hit on every device. + blocks_graphdef, blocks_state = nnx.split(self.blocks) + kv_cache = _scanned_compute_kv( + blocks_graphdef, + blocks_state, + encoder_hidden_states, + encoder_attention_mask, + num_layers=self.num_layers, + mesh=self.mesh, ) - def _compute_kv(block, enc_states, enc_mask): - return block.compute_kv(enc_states, enc_mask) - - kv_cache = _compute_kv(self.blocks, encoder_hidden_states, encoder_attention_mask) else: kv_cache_list = [] for block in self.blocks: @@ -761,6 +848,7 @@ def _compute_kv(block, enc_states, enc_mask): k_list = [d[k][0] for d in kv_cache_list] v_list = [d[k][1] for d in kv_cache_list] kv_cache[k] = (jnp.stack(k_list, axis=0), jnp.stack(v_list, axis=0)) + kv_cache = _shard_kv_cache_on_batch(kv_cache, self.mesh) return kv_cache, encoder_attention_mask diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline.py b/src/maxdiffusion/pipelines/wan/wan_pipeline.py index 4006e2fd1..55e616302 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline.py @@ -994,6 +994,45 @@ def _decode_latents_to_video(self, latents: jax.Array, trace: Optional[dict] = N video = np.asarray(video) return video + @staticmethod + def _auto_vae_spatial(config, total_devices: int, batch_parallel_vae: bool = False) -> int: + """Picks the `vae_spatial` mesh extent when the config leaves it at -1 (auto). + + The VAE mesh is ('redundant', 'vae_spatial') with redundant * vae_spatial == + total_devices. Batch is sharded on 'redundant' and width / conv channels on + 'vae_spatial'. Spatial sharding is what lets a single video fit, but with + vae_spatial == total_devices the batch axis is fully replicated, so at + B == total_devices every chip holds all B videos' activations and the + causal conv stack pays a ~10x slowdown (41.6s -> 6.3s at B=8 on tpu7x-8) + and OOMs on 32 GiB parts. + + Measured on v6e-8 / tpu7x-8 (720p, 81 frames): + * B % total_devices == 0: pure data parallelism (vae_spatial == 1) is + strictly better -- 1 video per device, zero cross-chip collectives. + * 1 < B < total_devices: mixed meshes are NOT a win. At B=2 on v6e-8 + the (2, 4) mesh needed a 21.2 GiB decode program vs 15.5 GiB for + (1, 8), even though per-device activation volume is identical. + So this only switches to vae_spatial == 1 when the batch is a multiple of + the device count and keeps full spatial sharding otherwise. Setting + `vae_spatial` explicitly in the config still overrides this. + + `batch_parallel_vae` gates the batch-aware choice. The VAE hard-constrains + every activation to P('redundant', ..., 'vae_spatial', ...), so a + vae_spatial == 1 mesh requires *every* tensor that goes through the VAE to + have a batch divisible by total_devices. That holds for the T2V pipelines + (the only VAE call is the final decode of the full batch), but not for + pipelines that also encode conditioning through the VAE at a different + batch -- VACE encodes reference images one at a time ([1, 1, H, W, C]), + animate / i2v encode conditioning clips. Those pipelines leave the gate + off and keep the full-spatial (1, N) mesh. + """ + if not batch_parallel_vae: + return total_devices + batch_size = int(getattr(config, "global_batch_size_to_train_on", 1) or 1) + if batch_size >= total_devices and batch_size % total_devices == 0: + return 1 + return total_devices + @classmethod def _create_common_components( cls, @@ -1002,6 +1041,7 @@ def _create_common_components( load_text_encoder=True, load_scheduler=True, i2v=False, + batch_parallel_vae=False, ): devices_array = max_utils.create_device_mesh(config) mesh = Mesh(devices_array, config.mesh_axes) @@ -1010,7 +1050,7 @@ def _create_common_components( total_devices = math.prod(devices_array.shape) if vae_spatial == -1: - vae_spatial = total_devices + vae_spatial = cls._auto_vae_spatial(config, total_devices, batch_parallel_vae=batch_parallel_vae) assert ( total_devices % vae_spatial == 0 @@ -1021,7 +1061,8 @@ def _create_common_components( vae_mesh = Mesh(vae_devices_array, ("redundant", "vae_spatial")) max_logging.log( - f"Created VAE specific mesh with axes ('redundant', 'vae_spatial') to support spatial sharding of {vae_spatial}." + f"Created VAE specific mesh with axes ('redundant', 'vae_spatial') = " + f"({total_devices // vae_spatial}, {vae_spatial}) to support spatial sharding of {vae_spatial}." ) # logical axis rules for VAE encoding/decoding diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py index 0cd0d7f7b..f80ca786a 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py @@ -68,6 +68,9 @@ def _load_and_init( load_vae=load_vae, load_text_encoder=load_text_encoder, load_scheduler=load_scheduler, + # T2V only runs the VAE on the final full-batch decode, so the + # batch-aware (N, 1) mesh is safe here; see _auto_vae_spatial. + batch_parallel_vae=True, ) transformer = None try: diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py index f55bd79e5..109549dde 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py @@ -79,6 +79,9 @@ def _load_and_init( load_vae=load_vae, load_text_encoder=load_text_encoder, load_scheduler=load_scheduler, + # T2V only runs the VAE on the final full-batch decode, so the + # batch-aware (N, 1) mesh is safe here; see _auto_vae_spatial. + batch_parallel_vae=True, ) low_noise_transformer, high_noise_transformer = None, None transformer_executor = concurrent.futures.ThreadPoolExecutor(max_workers=2) diff --git a/src/maxdiffusion/tests/wan/wan_vae_mesh_test.py b/src/maxdiffusion/tests/wan/wan_vae_mesh_test.py new file mode 100644 index 000000000..2e2f63ed9 --- /dev/null +++ b/src/maxdiffusion/tests/wan/wan_vae_mesh_test.py @@ -0,0 +1,132 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +CPU checks for the batch-aware VAE mesh selection and the batch-sharded +cross-attention KV cache helpers. +""" + +import os +import types + +os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=4") + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from flax.linen import partitioning as nn_partitioning +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P + +from maxdiffusion.pipelines.wan.wan_pipeline import WanPipeline +from maxdiffusion.models.wan.transformers import transformer_wan + + +def _cfg(batch_size): + return types.SimpleNamespace(global_batch_size_to_train_on=batch_size) + + +@pytest.mark.parametrize( + "batch_size, total_devices, expected", + [ + (0, 8, 8), + (1, 8, 8), + (2, 8, 8), + (3, 8, 8), + (4, 8, 8), + (8, 8, 1), + (12, 8, 8), + (16, 8, 1), + (4, 4, 1), + (1, 1, 1), + ], +) +def test_auto_vae_spatial_batch_parallel(batch_size, total_devices, expected): + """With the gate on, B % N == 0 picks pure data parallelism, otherwise full spatial.""" + assert WanPipeline._auto_vae_spatial(_cfg(batch_size), total_devices, batch_parallel_vae=True) == expected + + +@pytest.mark.parametrize("batch_size", [0, 1, 2, 8, 16]) +def test_auto_vae_spatial_default_is_full_spatial(batch_size): + """Without the gate (i2v / vace / animate) auto always resolves to the old value.""" + assert WanPipeline._auto_vae_spatial(_cfg(batch_size), 8) == 8 + assert WanPipeline._auto_vae_spatial(_cfg(batch_size), 8, batch_parallel_vae=False) == 8 + + +def test_only_t2v_pipelines_enable_batch_parallel_vae(): + """Guard against the gate leaking into pipelines that VAE-encode at batch != global batch.""" + pipeline_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "pipelines", "wan") + enabled = set() + for name in sorted(os.listdir(pipeline_dir)): + if not name.endswith(".py"): + continue + with open(os.path.join(pipeline_dir, name)) as f: + if "batch_parallel_vae=True" in f.read(): + enabled.add(name) + assert enabled == {"wan_pipeline_2_1.py", "wan_pipeline_2_2.py"} + + +@pytest.fixture +def mesh_and_rules(): + all_devices = jax.devices() + if len(all_devices) < 2: + pytest.skip("Requires at least 2 JAX devices to test multi-shard batch sharding") + usable = (len(all_devices) // 2) * 2 + devices = np.array(all_devices[:usable]).reshape(2, 1, -1, 1) + mesh = Mesh(devices, ("data", "fsdp", "context", "tensor")) + rules = (("activation_batch", ("data", "fsdp")),) + return mesh, rules + + +def test_batch_sharding_divisible(mesh_and_rules): + mesh, rules = mesh_and_rules + with mesh, nn_partitioning.axis_rules(rules): + sharding = transformer_wan._batch_sharding(mesh, ndim=3, batch_axis=0, batch_size=4) + assert sharding is not None + assert sharding.spec == P(("data", "fsdp"), None, None) + + +def test_batch_sharding_not_divisible_returns_none(mesh_and_rules): + """B=3 on data*fsdp=2 must skip the constraint instead of raising inside the jit.""" + mesh, rules = mesh_and_rules + with mesh, nn_partitioning.axis_rules(rules): + assert transformer_wan._batch_sharding(mesh, ndim=3, batch_axis=0, batch_size=3) is None + + +def test_batch_sharding_without_rules_or_mesh(mesh_and_rules): + mesh, _ = mesh_and_rules + assert transformer_wan._batch_sharding(None, ndim=3, batch_axis=0) is None + with mesh, nn_partitioning.axis_rules(()): + assert transformer_wan._batch_sharding(mesh, ndim=3, batch_axis=0) is None + + +def test_shard_kv_cache_on_batch_layouts(mesh_and_rules): + """[B,S,D] shards on axis 0, stacked [L,B,S,D] on axis 1; odd batches stay replicated.""" + mesh, rules = mesh_and_rules + cache = { + "text": (jnp.zeros((4, 16, 8)), jnp.zeros((2, 4, 16, 8))), + "odd": (jnp.zeros((3, 16, 8)), jnp.zeros((2, 3, 16, 8))), + } + with mesh, nn_partitioning.axis_rules(rules): + out = jax.jit(lambda c: transformer_wan._shard_kv_cache_on_batch(c, mesh))(cache) + k, v = out["text"] + # JAX canonicalises the spec of the returned array (size-1 'fsdp' is dropped), + # so compare semantically rather than by literal PartitionSpec. + assert k.sharding.is_equivalent_to(NamedSharding(mesh, P(("data", "fsdp"), None, None)), k.ndim) + assert v.sharding.is_equivalent_to(NamedSharding(mesh, P(None, ("data", "fsdp"), None, None)), v.ndim) + assert k.addressable_shards[0].data.shape[0] == 2 + assert v.addressable_shards[0].data.shape[1] == 2 + k_odd, v_odd = out["odd"] + assert k_odd.shape == (3, 16, 8) and v_odd.shape == (2, 3, 16, 8) + assert k_odd.sharding.is_equivalent_to(NamedSharding(mesh, P()), k_odd.ndim)