Fix the Wan VAE attention block, which never attends over spatial positions

[boshjerns]

Summary

WanAttentionBlock (the self-attention in the mid-block of the Wan VAE encoder and decoder) builds q, k and v in the PyTorch SDPA layout (batch, heads, seq, dim) and passes them to jax.nn.dot_product_attention, which takes (batch, seq, heads, dim). JAX therefore sees a sequence of length 1 with HW heads. Softmax over a single key is 1, so the attention output is exactly v, and the block computes proj(v) + x instead of attending over the HW positions of the frame.

This PR reshapes q, k and v to (N*T, H*W, 1, C) (one head over the spatial positions), which is what the PyTorch AutoencoderKLWan computes. With the fix, the MaxDiffusion VAE matches diffusers to within float32 rounding on real Wan2.1 weights (max abs diff 2e-5 or less).

AutoencoderKLWan is the VAE used by the Wan 2.1 and Wan 2.2 A14B pipelines, so this affects every encode and decode there. autoencoder_kl_wan_2p2.py has the same block; it is not imported anywhere at the moment, and I fixed it too so the two copies stay consistent.

Issue

In WanAttentionBlock.__call__:

qkv = qkv.reshape(batch_size * time, 1, -1, channels * 3)
qkv = jnp.transpose(qkv, (0, 1, 3, 2))
q, k, v = jnp.split(qkv, 3, axis=-2)
q = jnp.transpose(q, (0, 1, 3, 2))  # (N*T, 1, H*W, C)
...
x = jax.nn.dot_product_attention(q, k, v)  # reads axis 1 as seq (length 1) and axis 2 as heads (H*W)

Minimal check on CPU with random weights:

import jax, jax.numpy as jnp
from flax import nnx
from maxdiffusion.models.wan.autoencoder_kl_wan import WanAttentionBlock

dim = 32
blk = WanAttentionBlock(dim=dim, rngs=nnx.Rngs(0))
x = jax.random.normal(jax.random.key(1), (1, 2, 4, 6, dim))
out, _, _ = blk(x)

# the same block with the attention replaced by its value input
h = blk.norm(x.reshape(2, 4, 6, dim))
v = jnp.split(blk.to_qkv(h), 3, axis=-1)[2]
no_attn = blk.proj(v).reshape(x.shape) + x
print("max |out - (proj(v) + x)| =", float(jnp.max(jnp.abs(out - no_attn))))

main prints 0.0, this PR prints 2.41.

The VAE weights are trained with the attention in place, so this changes real outputs. I compared the MaxDiffusion Wan2.1 VAE (Wan-AI/Wan2.1-T2V-1.3B-Diffusers, loaded with load_wan_vae) against diffusers AutoencoderKLWan with the same weights, both in float32 on CPU, on the first 9 frames of the public hiker.mp4 test video at 256x448 (script below):

	main	this PR
encoder: latent mean, relative L2 vs diffusers	8.2e-3	1.3e-6
encoder: latent mean, max abs diff	0.16	2.2e-5
decoder (diffusers latents in): PSNR vs diffusers	47.6 dB	131.6 dB
decoder: max abs diff (output range [-1, 1])	0.34	1.3e-5
round trip PSNR vs the input video (diffusers: 37.31 dB)	37.07 dB	37.31 dB

I also saw it end to end, with a separate JAX implementation of the same decoder that I could switch between the two attention behaviours (so these numbers are not from this PR's code). For a Wan2.2 I2V-A14B clip (832x464, 81 frames) generated on a TPU v6e and compared with the PyTorch pipeline on a GPU from the same initial noise, fixing only the decoder attention raised the mean PSNR vs the GPU output from 25.65 to 25.96 dB, and frame 0 from 36.4 to 38.7 dB.

Changes

• models/wan/autoencoder_kl_wan.py: WanAttentionBlock reshapes qkv to (N*T, H*W, 1, 3C), splits on the channel axis and passes the result to jax.nn.dot_product_attention directly. The four transposes go away.
• models/wan/autoencoder_kl_wan_2p2.py: the same fix in its copy of WanAttentionBlock.
• tests/wan/wan_vae_test.py: adds TorchWanAttentionBlock, a PyTorch reference written the same way as the diffusers block (like the existing TorchWanRMS_norm), and test_wan_attention_matches_pytorch, which copies its weights into both JAX blocks and compares outputs on a small random input (atol/rtol 1e-5). On main both sub-tests fail (max abs diff 1.6).

The block now does the work it is meant to do, so it costs more than before. With the default vae_encode_chunk: 4 and vae_decode_chunk: 1 the mid-block sees one latent frame per call. At 720p that is 90x160 = 14,400 positions, so the XLA attention path holds at least a 14,400 x 14,400 float32 score matrix (about 0.8 GB) per latent frame in the call; vae_encode_chunk: -1 puts all latent frames in one call and scales this up. With vae_spatial sharding, XLA gathers qkv and every device computes the full attention (I checked the compiled HLO on 8 host CPU devices: one all-gather, the same count as on main, and the sharded output matches the unsharded one to 5e-7). Constraining q to be sharded along the sequence would split that work; I kept this change minimal. In the separate implementation above, the same jax.nn.dot_product_attention call added about 0.02 s to a whole 81-frame 480p decode on a TPU v6e.

Testing

• python -m pytest src/maxdiffusion/tests/wan/wan_vae_test.py -k attention (CPU, jax 0.11.2, flax 0.12.10, torch 2.10.0): 2 passed with this PR. On main, test_wan_attention_matches_pytorch fails for both blocks.
• python -m pytest src/maxdiffusion/tests/wan/wan_vae_test.py on CPU, without test_load_checkpoint (which downloads the 14B VAE): 12 passed, including the full-size encode and decode tests.
• pyink --check --pyink-indentation=2 --line-length=125 on the three files: unchanged.
• pylint --fail-under=7 src/maxdiffusion/: passes, no new messages in the touched files.
• ruff check .: passes.

Comparison script against diffusers (CPU, float32)

import os

import flax
import jax
import jax.numpy as jnp
import numpy as np
import torch
from diffusers import AutoencoderKLWan as HFAutoencoderKLWan
from flax import nnx
from flax.linen import partitioning as nn_partitioning
from jax.sharding import Mesh

import maxdiffusion
from maxdiffusion import pyconfig
from maxdiffusion.max_utils import create_device_mesh
from maxdiffusion.models.wan.autoencoder_kl_wan import AutoencoderKLWan, AutoencoderKLWanCache
from maxdiffusion.models.wan.wan_utils import load_wan_vae
from maxdiffusion.utils import load_video
from maxdiffusion.video_processor import VideoProcessor

flax.config.update("flax_always_shard_variable", False)
REPO = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"


def psnr(a, b):  # both in [-1, 1]
  return 10 * np.log10(1.0 / np.mean(((a - b) / 2.0) ** 2))


frames = load_video("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/hiker.mp4")[:9]
video = np.asarray(VideoProcessor(vae_scale_factor=8).preprocess_video(frames, height=256, width=448), np.float32)

hf = HFAutoencoderKLWan.from_pretrained(REPO, subfolder="vae", torch_dtype=torch.float32).eval()
with torch.no_grad():
  hf_mean = hf.encode(torch.from_numpy(video)).latent_dist.mean
  hf_dec = hf.decode(hf_mean).sample.numpy()
hf_mean = hf_mean.numpy()

pyconfig.initialize([None, os.path.join(os.path.dirname(maxdiffusion.__file__), "configs", "base_wan_1_3b.yml")], unittest=True)
config = pyconfig.config
mesh = Mesh(create_device_mesh(config).flatten().reshape(-1, 1), ("redundant", "vae_spatial"))
with mesh, nn_partitioning.axis_rules(config.vae_logical_axis_rules):
  vae = AutoencoderKLWan.from_config(REPO, subfolder="vae", rngs=nnx.Rngs(0), mesh=mesh)
  cache = AutoencoderKLWanCache(vae)
  graphdef, state = nnx.split(vae)
  params = load_wan_vae(REPO, state.to_pure_dict(), "cpu")
  flat = dict(nnx.to_flat_state(state))
  for path, val in flax.traverse_util.flatten_dict(params).items():
    flat[path][...] = jnp.asarray(val, jnp.float32)
  vae = nnx.merge(graphdef, nnx.from_flat_state(flat))
  md_mean = np.asarray(vae.encode(jnp.asarray(video), feat_cache=cache).latent_dist.mean).transpose(0, 4, 1, 2, 3)
  md_dec = np.asarray(vae.decode(jnp.asarray(hf_mean.transpose(0, 2, 3, 4, 1)), feat_cache=cache).sample)
  md_dec = md_dec.transpose(0, 4, 1, 2, 3)

print("latent rel L2:", np.linalg.norm(md_mean - hf_mean) / np.linalg.norm(hf_mean))
print("decode PSNR vs diffusers:", psnr(md_dec, hf_dec))

[gemini-code-assist]

Code Review

This pull request refactors the spatial self-attention block in both autoencoder_kl_wan.py and autoencoder_kl_wan_2p2.py to simplify the reshaping and splitting of the QKV projections before passing them to jax.nn.dot_product_attention. Additionally, a new test test_wan_attention_matches_pytorch has been added to wan_vae_test.py along with a PyTorch reference implementation TorchWanAttentionBlock to verify that the JAX attention blocks produce equivalent outputs to PyTorch. There are no review comments to address, and I have no feedback to provide.