Skip to content

Make AutoencoderKLWan stateless to fix torch.compile compatibility - #14771

Open
NotNANtoN wants to merge 1 commit into
huggingface:mainfrom
NotNANtoN:fix-wan-vae-stateless-compile
Open

NotNANtoN wants to merge 1 commit into
huggingface:mainfrom
NotNANtoN:fix-wan-vae-stateless-compile

Conversation

@NotNANtoN

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes #14770

This PR makes AutoencoderKLWan (used in Wan2.1 and Wan2.2 models) stateless during execution so that it compiles cleanly with torch.compile without graph breaks or issues from mutating internal module attributes.

Key Changes

  1. Stateless Feature Caching: Allocate feat_map = [None] * self._cached_conv_counts[...] locally inside _encode, _decode, tiled_encode, and tiled_decode and pass it down along with a local conv_idx = [0] list.
  2. Eliminate Mutable Defaults: Changed mutable default arguments from feat_idx=[0] to feat_idx=None in all residual / up / down / mid blocks.
  3. Backward Compatibility: clear_cache() is preserved as a method on AutoencoderKLWan so existing external scripts calling vae.clear_cache() do not break.

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 bitwise parity with eager baseline).
  • Unit Tests: All 32 unit tests in tests/models/autoencoders/test_models_autoencoder_wan.py pass cleanly.
  • Latency Benchmarks:
    • 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)

Self-Review Notes

  • Diff Scope: Confined strictly to src/diffusers/models/autoencoders/autoencoder_kl_wan.py.
  • Formatting: Verified with ruff check and ruff format.
  • Edge cases: Verified both standard encoding/decoding and tiled encoding/decoding (test_tiling_encode_shape, test_tiling_decode_shape).

Before submitting

Who can review?

@yiyixuxu @dg845 @asomoza

…#14770)

Allocate feat_map and conv_idx locally in _encode, _decode, tiled_encode, and tiled_decode instead of mutating module attributes, and replace mutable default argument feat_idx=[0] with feat_idx=None.
@github-actions github-actions Bot added fixes-issue models size/M PR with diff < 200 LOC and removed fixes-issue labels Sep 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

models size/M PR with diff < 200 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant