[Feature]: Support VAE as a Separate Stage to Reduce GPU Memory Pressure in Diffusion Pipelines

[SamitHuang]

🚀 The feature, motivation and pitch

Problem

In vllm-omni's current diffusion pipeline architecture, the VAE (encoder/decoder) is loaded on the same GPU as the denoising model (DiT/UNet) within a single diffusion stage. VAE models can be very large — for example, QwenImage VAE and Wan VAE are 3D causal convolution networks with significant parameter counts and activation memory footprints. During the decode phase (latents → pixels), the VAE must coexist in GPU memory alongside the denoiser, leading to:

1. High peak GPU memory usage — the denoiser already occupies a large portion of VRAM; adding VAE on top can push close to or beyond device limits, especially for high-resolution / video generation.
2. Suboptimal memory utilization — VAE encode and decode are only needed at the beginning and end of the pipeline respectively; the denoiser does not need the VAE during its many iterative denoising steps. Despite this temporal non-overlap, both models occupy memory simultaneously.
3. OOM during VAE decode — several pipelines (e.g., Wan 2.2) already insert explicit torch.cuda.empty_cache() calls before VAE decode as a workaround for memory pressure, indicating this is a known pain point.

Current mitigations are insufficient

• vae_use_tiling / vae_use_slicing reduce activation memory during VAE forward passes but do not address the weight memory of the VAE itself.
• enable_cpu_offload (model-level offload) swaps DiT ↔ encoder modules but the current ModelLevelOffloadBackend does not include VAE in its sequential swap hooks — VAE is moved to GPU once and stays there.
• Distributed VAE (patch parallelism) spreads VAE computation across GPUs within the same diffusion worker group, but does not eliminate VAE weight memory from the denoiser's GPU.

Proposed solution: VAE as a separate omni stage

vllm-omni's stage system (StageConfig / OmniStage) already supports running different stages on different GPU devices via per-stage runtime.devices configuration and set_stage_devices(). We propose leveraging this existing multi-stage infrastructure to split VAE into its own stage, separate from the denoising stage:

┌─────────────┐      ┌──────────────────┐      ┌──────────────┐
│  VAE Encode │ ───► │   Denoising      │ ───► │  VAE Decode  │
│  (Stage A)  │      │   (Stage B)      │      │  (Stage C)   │
│  GPU 0      │      │   GPU 1,2,...    │      │  GPU 0       │
└─────────────┘      └──────────────────┘      └──────────────┘

• VAE Encode stage runs on a dedicated (possibly smaller / shared) GPU, encodes input images into latents, and passes them to the denoising stage.
• Denoising stage runs the iterative denoising loop on its own GPU(s) without loading VAE weights at all, freeing significant VRAM for larger models, higher resolution, or batch processing.
• VAE Decode stage receives final latents and decodes to pixel space. This could share the same GPU as VAE Encode (since they are temporally non-overlapping), or even share a GPU with the LLM stage.

Key benefits

1. Reduced per-GPU memory — The denoiser GPU no longer needs to hold VAE weights (often 100MB–2GB+ depending on precision and model).
2. Enable larger models / higher resolution — Freed VRAM can be used for longer context, higher resolution, or increased batch size on the denoiser.
3. Better multi-GPU utilization — VAE encode/decode can share a GPU with the LLM stage or other non-overlapping workloads, improving overall cluster utilization.
4. Consistent with existing architecture — Follows the same multi-stage pattern already used for LLM ↔ Diffusion separation.

Alternatives

1. Improve CPU offload to include VAE — Extend ModelLevelOffloadBackend to swap VAE ↔ DiT weights via CPU. Simpler but adds latency from CPU↔GPU transfers and doesn't improve GPU utilization across multiple devices.
2. Layer-wise offload for VAE — Offload individual VAE layers to CPU/disk. Very fine-grained but adds significant complexity and transfer overhead.
3. VAE quantization — Reduce VAE weight memory via INT8/INT4 quantization. Orthogonal optimization that can be combined with stage splitting, but may affect image quality.
4. Lazy VAE loading — Load VAE weights only when needed and unload after. Avoids holding weights but adds load/unload latency per request.

Additional context

Affected VAE implementations

VAE Class	Pipeline(s)	Notes
AutoencoderKLQwenImage	QwenImage, QwenImage Edit	3D causal conv VAE, large
DistributedAutoencoderKLWan	Wan 2.2 (t2v, i2v)	Video VAE, explicit OOM workarounds exist
AutoencoderKLFlux2	Flux2, Flux2 Klein
DistributedAutoencoderKL	SD3, Z-Image	Standard image VAE
AutoencoderKLConv3D	Hunyuan Image 3	3D conv VAE
AutoencoderKLLTX2Video + AutoencoderKLLTX2Audio	LTX2	Dual VAE (video + audio)
Bagel AutoEncoder	Bagel	Custom construction

Design considerations

• New stage type or config: Could introduce a StageType.VAE or reuse StageType.DIFFUSION with a sub-mode flag.

• Inter-stage data transfer: Latent tensors are relatively small compared to pixel-space tensors, making inter-stage transfer efficient.

• Pipeline refactoring: Pipelines would need to be split into encode(), denoise(), decode() callable segments rather than monolithic __call__().

• Backward compatibility: When no separate VAE stage is configured, behavior should remain identical to today (VAE loaded within the diffusion stage).

• Shared GPU with LLM: Since LLM inference and VAE encode/decode are temporally non-overlapping in most omni workflows, they could share the same GPU device to minimize total GPU count.

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[hsliuustc0106]

@erfgss the metrics in your PR should be aligned with the 3 fundamental stages in diffusion: 1) vae encode 2) denoising 3) vae decode #1977

[hsliuustc0106]

@yuanheng-zhao @asukaqaq-s PTAL

[asukaqaq-s]

I think this setup is worthwhile. For a single request, the peak memory should be around 4.46 GB for VAE, while LightTAE reaches about 16.95 GB in parallel mode.

Other potential benefits I see:

1. After splitting by stage, the diffusion engine only needs to focus on denoise-step execution, without handling pre/post-processing. That should make it easier to introduce continuous batching in a cleaner way and improve GPU utilization.
2. It should also help with the streaming / real-time diffusion work the community is discussing. I also tested a 14B streaming model on Hopper, and for a single request the gain from SP=4 to SP=8 was relatively small (~10%), so some GPUs could potentially be reassigned to VAE.

I previously worked on the step/stage pipeline interface in #1368, and I think it could be combined with this direction. I do have a few questions:

1. I do not think the diffusion engine/runner needs very large changes. The bigger question is how to design the scheduler/engine and runner layer for the VAE encoder/decoder. If we split by stage, it seems more reasonable to place them into separate stages, similar to Qwen-Omni.
2. How should we design the interface for data transfer and communication across stages, especially when some prepare/process logic is needed after receiving stage outputs (for example, DiT doing scheduler preparation after receiving latents)?
3. Should we also consider splitting the Qwen-Image text encoder (7B model)? For TI2V tasks, how should the text encoder and VAE encoder be placed?

I am very interested in this area and would be happy to discuss.

[yuanheng-zhao]

Agree that this feature is useful. It was actually explored and applied by some video-gen serving systems from recent years.

While the diagram in Proposed solution section shows VAE encode and decode in two different stages, I think a single VAE stage handling both encode and decode is better, for its weight sharing and reusability.

┌─────────────┐      ┌──────────────────┐      ┌──────────────┐
│  VAE Encode │ ───► │   Denoising      │ ───► │  VAE Decode  │
│  (Stage A)  │      │   (Stage B)      │      │  (Stage A)   │
│  GPU 0      │      │   GPU 1,2,...    │      │  GPU 0       │
└─────────────┘      └──────────────────┘      └──────────────┘

I'd like to participate and help as well.

[yuanheng-zhao]

I think this setup is worthwhile. For a single request, the peak memory should be around 4.46 GB for VAE, while LightTAE reaches about 16.95 GB in parallel mode.

Other potential benefits I see:

1. After splitting by stage, the diffusion engine only needs to focus on denoise-step execution, without handling pre/post-processing. That should make it easier to introduce continuous batching in a cleaner way and improve GPU utilization.
2. It should also help with the streaming / real-time diffusion work the community is discussing. I also tested a 14B streaming model on Hopper, and for a single request the gain from SP=4 to SP=8 was relatively small (~10%), so some GPUs could potentially be reassigned to VAE.

I previously worked on the step/stage pipeline interface in #1368, and I think it could be combined with this direction. I do have a few questions:

1. I do not think the diffusion engine/runner needs very large changes. The bigger question is how to design the scheduler/engine and runner layer for the VAE encoder/decoder. If we split by stage, it seems more reasonable to place them into separate stages, similar to Qwen-Omni.
2. How should we design the interface for data transfer and communication across stages, especially when some prepare/process logic is needed after receiving stage outputs (for example, DiT doing scheduler preparation after receiving latents)?
3. Should we also consider splitting the Qwen-Image text encoder (7B model)? For TI2V tasks, how should the text encoder and VAE encoder be placed?

I am very interested in this area and would be happy to discuss.

Shall we create and apply a VAE model runner analogous to DiffusionModelRunner, which should be simpler and cleaner compared with that for diffusion stage.

2. We might want to split current processing and handling in diffusion stage to be handled in VAE stage as well as custom_process_input_func of diffusion stage
3. I think this is model/task-specific. An optional text encoder could be provided in the newly-added (todo) stage but this might make confusion; Or, we might want to put text encoder to another separate stage later?

cc @asukaqaq-s @SamitHuang

[princepride]

@yuanheng-zhao @asukaqaq-s hi, can you help implement this feature, thank you 🙏

[asukaqaq-s]

@yuanheng-zhao @asukaqaq-s hi, can you help implement this feature, thank you 🙏

Honored to help. I’d be happy to work on this.
maybe we can discuss the task split first and divide the work accordingly.