[bug]: Wan 2.2 A14B at 720p does not fit 24 GB (no working-memory estimate); SageAttention makes the overflow slower #9831
Description
Is there an existing issue for this problem?
- I have searched the existing issues
Install method
Manual (development checkout)
Operating system
Windows 10
GPU vendor
Nvidia (CUDA)
GPU model
RTX 4090
GPU VRAM
24 GB
Version number
main at 8111b50 (7.0.0-alpha.1). The SageAttention rows below use the stack #9657–#9665 at f58354f, which merges that main.
Browser
None (graphs sent through the queue API)
System Information
Python 3.12, torch 2.13.0+cu130. SageAttention rows: sageattention 2.2.0+cu130torch2.10.0andhigher.post6, triton-windows 3.7.1.post27.
What happened
Wan 2.2 T2V A14B (GGUF Q4_K_M, both experts) with the 4-step Lightning LoRA V1.1, at 1280×720 with 81 frames, does not fit a 24 GB card. How it fails depends on the settings:
enable_partial_loading |
device_working_mem_gb |
Attention | Whole generation |
|---|---|---|---|
| true (default) | 3 (default) | PyTorch SDPA | OOM in the first transformer block, also at 832×480 |
| false | 3 | PyTorch SDPA | 875 s; VRAM full, steps 177 / 111 / 215 / 192 s |
| true | 10 | PyTorch SDPA | 327 s; steps ~73 s |
| true | 10 | SageAttention (#9657 stack) | 687 s; steps 234 / 95 / 289 / 44 s |
| false | 3 | SageAttention | 1813 s; first step 1397 s |
At 832×480, with partial loading on and device_working_mem_gb: 10, both attention backends fit: 148 s with SDPA, 75 s with SageAttention.
What you expected to happen
With default settings, a 24 GB card either fits the generation or offloads enough model weights to fit, instead of running out of memory. SageAttention is never slower than PyTorch's attention.
How to reproduce the problem
- Install Wan 2.2 T2V A14B High Noise and Low Noise (GGUF Q4_K_M), the
Wan2.2-T2V-A14B-Diffuserscomponents, and the Lightning 4-step LoRAs (V1.1, high and low). - Keep the default
invokeai.yaml(partial loading on,device_working_mem_gb: 3). - Run
wan_model_loader→wan_lora_loader(high) →wan_lora_loader(low) →wan_video_denoise(1280×720, 81 frames, 4 steps, guidance 1.0) →wan_l2v. - Generation fails with
torch.OutOfMemoryErrorindiffusers/models/transformers/transformer_wan.py:502.
Additional context
Wan denoise has no working-memory estimate.
_get_wan_transformer_working_mem_bytes(invokeai/app/invocations/wan/wan_denoise.py) returnsNoneunlesswan_memory_optimizationis on. The expert is then loaded with the floor,device_working_mem_gb(3 GB), as its working memory.- At 720p with 81 frames the transformer sees 75,600 tokens, and the activations need far more than 3 GB.
- Partial loading on: the cache frees only 2,003 MB of the 10.8 GB UMT5 encoder before loading the 9.2 GB expert, so the first block runs out of memory with 21.3 GiB allocated.
- Partial loading off: the encoder is unloaded completely, but the card still overflows. Windows moves memory to shared system memory instead of raising OOM, so steps take two to three times as long.
- Workaround on 24 GB:
device_working_mem_gb: 10. The Z-Image, FLUX.2, Krea-2 and LTX-2 denoise nodes estimate their working memory; Wan does not.
SageAttention makes the overflow slower (the stack in #9657–#9665).
- It holds INT8/FP8 copies of Q, K and V for the duration of each call, about 1.2 GB more than SDPA at 75,600 tokens. When SDPA already fills the card, that pushes it over.
- Its fallback only catches
torch.OutOfMemoryError. With Windows' shared-memory fallback no OOM is raised, so the calls spill instead. - At 480p, where everything fits, SageAttention halves the time.
Possible fixes
- Wan: estimate the denoise working memory from the token count and pass it to
model_on_device(working_mem_bytes=...), as the nodes above do. - SageAttention: before a call, compare the size of the quantized copies with the free VRAM (
torch.cuda.mem_get_infoplus memory torch has reserved but not allocated). Take SDPA when they do not fit. That costs one driver call per eligible call, against kernels that take milliseconds.
Related: #9529 (Wan 2.2 A14B at 720p on a 24 GB RTX 5090 Laptop GPU), #9657 and #9665 (SageAttention).
Activity
docs(attention): add the Wan 2.2 A14B measurements
docs(attention): say that a generation which only just fits can get s…
feat(wan): run Wan 2.2 attention on SageAttention when enabled (#9665)