huggingface / huggingface/diffusers

[Bug / Optimization] Make AutoencoderKLWan stateless to fix torch.compile compatibility and enable 1.45x speedup

Open
#14,770 0 comments 0 reactions 0 assignees View on GitHub
bug models
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

### Description

`AutoencoderKLWan` (used in Wan2.1 and Wan2.2 models) mutates internal instance attributes during its forward pass in `_encode`, `_decode`, `tiled_encode`, and `tiled_decode`:
- `self._enc_feat_map`
- `self._enc_conv_idx`
- `self._feat_map`
- `self._conv_idx`

Furthermore, the intermediate blocks (`WanResample`, `WanResidualBlock`, `WanMidBlock`, `WanResidualDownBlock`, `WanEncoder3d`, `WanResidualUpBlock`, `WanUpBlock`, `WanDecoder3d`) use a mutable default argument `feat_idx=[0]`.

Because module attributes are repeatedly mutated during the forward pass and `self.clear_cache()` reallocates lists on `self`, compiling `AutoencoderKLWan` with `torch.compile` causes graph breaks and prevents efficient compilation.

### Reproduction

```python
import torch
from diffusers import AutoencoderKLWan

vae = AutoencoderKLWan.from_pretrained(
"Wan-AI/Wan2.1-T2V-1.3B", subfolder="vae", torch_dtype=torch.bfloat16
).cuda()
vae = torch.compile(vae)

# Input shape: (B, C, T, H, W) e.g. 5 frames
x = torch.randn(1, 3, 5, 480, 640, dtype=torch.bfloat16, device="cuda")
with torch.no_grad():
# Mutates self._enc_feat_map and self._enc_conv_idx during forward pass
latents = vae.encode(x).latent_dist.sample()
```

### Proposed Solution

Make `AutoencoderKLWan` stateless during execution:
1. Allocate `feat_map = [None] * self._cached_conv_counts["encoder"]` (and `decoder`) locally inside `_encode`, `_decode`, `tiled_encode`, and `tiled_decode`, passing the local list and a local `conv_idx = [0]` to child modules.
2. Change mutable default arguments from `feat_idx=[0]` to `feat_idx=None` (initializing to `[0]` inside the method when `None`).
3. Preserve `clear_cache()` as a method for backward compatibility.

### Verification & Performance

Tested on NVIDIA RTX 4090 (PyTorch 2.11 / CUDA 12.8, BF16, 5 frames 480x640):
- **Numerical Parity**: `max_abs_diff = 0.0` (exact match to uncompiled eager baseline).
- **Unit Tests**: All 32 unit tests in `tests/models/autoencoders/test_models_autoencoder_wan.py` pass cleanly.
- **Speedup**:
- Eager uncompiled: `84.43 ms`
- `torch.compile(mode="default")`: `61.10 ms` (**1.38x speedup**)
- `torch.compile(mode="max-autotune-no-cudagraphs")`: `58.38 ms` (**1.45x speedup**)

We have a clean PR ready with these changes if the maintainers would like us to open it!

Contributor guide

Open the contributing guide

Research direction

Start with AutoencoderKLWan's _encode, _decode, tiled_encode, and tiled_decode methods, then inspect the listed Wan intermediate blocks and tests/models/autoencoders/test_models_autoencoder_wan.py. Run the 32 existing tests first; done means forward execution no longer mutates module state, the tests retain numerical parity, and torch.compile compatibility is verified.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.