perf(wan): batch-sharded cross-attn KV cache and batch-aware VAE mesh

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
 

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

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:

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)

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)