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

🚀 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.