wan: causal 3D VAE with feat-cache streaming (2.1 + 2.2)

[jamesbrink]

Part of the Wan video family (tracking: #741). Stack layer 3.

Scope (wan/03-vae)

New crates/mold-inference/src/wan/model/vae.rs: Wan's causal 3D VAE with feat-cache temporal streaming — the property that makes 121-frame 720p decode fit 24 GB (O(1) memory in T). Implementation strategy: CausalConv3d as Conv2d-slices summed per output frame (proven by Ltx2VideoCausalConv3d; candle has no cudnn conv3d and doesn't need one).

2.1 VAE (all 1.3B/14B/A14B models; z=16, stride 4×8×8, dim 96, dim_mult [1,2,4,4], 2 res blocks, temporal downsample [F,T,T]):

• CausalConv3d with left-only temporal pad; when cache_x supplied, prepend cached tail and reduce pad
• Feat-cache protocol: CACHE_T=2 trailing frames per conv, flat cache indexed by conv-visit order (feat_idx counter); encode chunks 1/4/4/4 pixel frames; decode 1 latent frame per iteration; 'Rep' sentinel path in upsample3d for the first chunk; downsample3d caches the raw last frame. Wrong cache order/chunking ⇒ temporal seams every 4 frames.
• RMS_norm (per-channel, gamma [dim,1,1]), single-head per-frame 2D attention mid-block, encoder/decoder per ComfyUI key layout (decoder.upsamples.{0..14} flat Sequential — see report §1.3 for exact keys)
• Per-channel latent mean/std vectors in code (16 values each; encode (µ−mean)/std, decode input z·std+mean, output clamp [−1,1]); encode returns distribution mean (no sampling)
• Spatial tiled-decode fallback reusing vae_tiling.rs arithmetic (ComfyUI: 256 px tiles, ¼ overlap) for OOM headroom

2.2 VAE (TI2V-5B only; z=48, stride 4×16×16):

• Pixel-space 2× patchify (3→12 ch) before conv_in, unpatchify after; encoder dim 160, decoder dim 256 (asymmetric); AvgDown3D avg-pool shortcuts / DupUp3D repeat shortcuts, which drop the first factor_t−1 output frames when first_chunk; nested downsamples.{s}.downsamples.{j} key layout; 48-value mean/std tables copied verbatim

Acceptance

• Unit test: streaming (chunked, feat-cached) encode/decode ≡ full-clip forward on a small random tensor, both VAEs
• Golden test vs diffusers AutoencoderKLWan on a fixed image → encode → decode round trip
• clippy -D warnings, workspace tests; CHANGELOG entry

Key facts: tmp/wan-research/report-official-wan.md §3, report-comfyui.md §1.3/§4, report-diffusers.md §5.

[jamesbrink]

Merged to main as 1e792b9 (PR #755); post-merge CI green. Both VAE generations with the feat-cache streaming protocol, mutation-tested chunk-plan invariance, layouts verified against real checkpoint headers; carried every real-weights UAT decode (9.6-27.3s for 33-53f).