[MAX] Add Wan VAE and refactor autoencoder module

[jglee-sqbits]

Stacked PRs:

• [MAX] Add Wan video generation examples and comparison benchmark #19
• [MAX] Add Wan I2V diffusion pipeline #18
• [MAX] Add Wan T2V diffusion pipeline with MoE support #17
• [MAX] Add Wan transformer model with block-level compilation #16
• ->[MAX] Add Wan VAE and refactor autoencoder module #15
• [MAX] Add UMT5 text encoder for Wan diffusion #14
• [MAX] Add UniPC multistep scheduler for Wan diffusion #13

---

[MAX] Add Wan VAE and refactor autoencoder module

Summary

Add a Wan VAE (3D causal video autoencoder) and restructure the autoencoder module to separate Module V2 and V3 implementations.

Description

Wan VAE

• Implements the Wan 3D causal VAE with temporal caching for chunked encode/decode
• Encoder: processes video in temporal chunks (first frame + subsequent chunks) with cached convolution state to maintain temporal consistency
• Decoder: same chunked approach with 3 specialized graphs (post-quant conv, first frame, subsequent frames)
• Uses symbolic spatial dims for resolution flexibility
• Adds 3D convolution support via cuDNN (conv.mojo) with depth-tiled execution for large volumes

Autoencoder restructuring

• Moves existing Module V3 (Flux) autoencoder files to autoencoders_modulev3/
• The autoencoders/ directory now contains Module V2 graph-based implementations (Wan VAE, Qwen Image VAE)
• Updates flux1_modulev3 and flux2_modulev3 import paths accordingly

This follows the same pattern as modular#6278 which established the V2/V3 split.

Dependencies

Should be merged before modular#6301 (transformer) and modular#6302 (pipeline-t2v), which import from autoencoders.

Checklist

• PR is small and focused
• I ran ./bazelw run format to format my changes

Assisted-by: Claude Code

Assisted-by: Claude Code

[gemini-code-assist]

Code Review

This pull request introduces support for 3D convolutions using cuDNN and adds a new Wan VAE autoencoder implementation. The changes include new cuDNN descriptor APIs, depth-tiled convolution for large tensors, and a new architecture module for Wan. My feedback highlights that the hardcoded total_cache_slots calculation in vae.py is fragile and should be dynamic, and that the comptime if in conv.mojo is redundant and can be simplified to a standard if statement.

[gemini-code-assist]

The total_cache_slots property calculation is hardcoded and fragile. If the number of blocks or cache slots changes in the future, this will lead to silent errors. Consider calculating this dynamically based on the actual number of modules in down_blocks and mid_block.

[gemini-code-assist]

The use of comptime if inside the conv_gpu function for algorithm selection is redundant here, as filter_is_fcrs is already a compile-time constant. This can be simplified to a standard if statement for better readability without impacting performance.