[feature] cpu offloading support for diffusion

[LawJarp-A]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

FIX #412

Implementation

Uses a mutual-exclusion swap pattern via PyTorch forward pre-hooks:

• DiT (transformer) and encoders alternate GPU access
• Before encoder forward: DiT → CPU, encoder → GPU
• Before DiT forward: encoders → CPU, DiT → GPU

Key features:

• Single config flag: dit_cpu_offload=True
• Zero pipeline code changes required - hooks handle device placement automatically
• Pinned CPU memory for faster PCIe transfers (pin_cpu_memory=True by default)
• Compatible with TeaCache and Cache-DiT

Files changed:

• vllm_omni/diffusion/offload.py - New file with SequentialOffloader class
• vllm_omni/diffusion/data.py - Simplified config to single dit_cpu_offload flag
• vllm_omni/diffusion/worker/gpu_worker.py - Hook application after model load

TODO

• Try to fix some model loading is very slow when enabling offloading
• add tests for this
• add doc

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[ZJY0516]

Could you please test the cache functionality(cache-dit and tea-cache) when offload enabled?

[LawJarp-A]

Could you please test the cache functionality(cache-dit and tea-cache) when offload enabled?

Tested it @ZJY0516
Is this a agreeable approach for you rather than adding it as a hook like TeaCache?

Category	Config	Time (s)	Speedup vs Baseline
No Cache	no_cache_no_offload	4.06	baseline
No Cache	no_cache_with_offload	14.20	0.29×
Cache-DiT	cache_dit_no_offload	2.51	1.62×
Cache-DiT	cache_dit_with_offload	10.70	1.33× vs offload
TeaCache	teacache_no_offload	3.22	1.26×
TeaCache	teacache_with_offload	11.63	1.22× vs offload

[ZJY0516]

I was wondering what if we want to overlap computation and offload?

[SamitHuang]

seems it's a synchronized and blocking transfer from CPU to GPU. Will it be more efficient to use async transfers with CUDA streams?

[hsliuustc0106]

please check whether this blog will help https://zhuanlan.zhihu.com/p/1986157623695922659

[ZJY0516]

Given the high demand for offloading, I think we can leave efficient overlap in following pr. @LawJarp-A @SamitHuang @hsliuustc0106

[ZJY0516]

I didn't see any evidence that text encoder is offloaded when I run it locally.

from vllm_omni import Omni

if __name__ == "__main__":

    m = Omni(model="Qwen/Qwen-Image",text_encoder_cpu_offload=True)

    outputs = m.generate(
        "a photo of a cat sitting on a laptop keyboard",
        height=1024,
        width=1024,
        num_inference_steps=50,
        num_outputs_per_prompt=1,
    )

[bwyangseek]

That makes sense. I’ve previously worked on overlapping GPU compute with CPU communication using a dual-thread, multi-stream, and thread-switching approach. A similar micro-batch + async model could help overlap offloading with compute instead of serializing steps.
Happy to explore this in a follow-up PR once the current one lands.
Could you please clarify if handling this overlap falls under this PR, the roadmap, or should be a separate task? @LawJarp-A @SamitHuang @hsliuustc0106

[ZJY0516]

It should be a separate task. @bwyangseek

[ZJY0516]

The main problem is model loading is very slow when enabling offloading

[SamitHuang]

thanks for the nice work.
what kind of models are tested? are t2i, i2i and t2v all covered?

[SamitHuang]

What is the reason of latency? is it proportional the model size?

[ZJY0516]

I suspect the root cause is CPU-side model initialization time, which seems proportional to model size(not sure). While Z-Image incurs only a small delay, Qwen-Image's much larger size results in over a minute of initialization latency.

[yuanheng-zhao]

A note for this part:

If DiT modules are initialized on GPU first (not the current case), here it comes two copies of parameters happening on CPU: GPU params -> CPU params -> CPU pinned params. Creation of empty tensor with pin_memory =True and then in-place copy from existing param data will be a method to ensure a single-copy path, for the case the dit modules are not residing on cpu at some moment.

For now, it does not make a diff, since the usage in gpu worker ensures the model loader to load dit modules on cpu directly when dit_cpu_offload is enabled:

class GPUWorker:
    def init_device_and_model(self) -> None:
        ...
        load_device = "cpu" if self.od_config.dit_cpu_offload else str(self.device)
        ...
                    self.pipeline = model_loader.load_model(
                        od_config=self.od_config,
                        load_device=load_device,
                    )

cc @ZJY0516

[ZJY0516]

Thanks for the reminder. Initializing on the CPU is currently a necessary workaround due to our architecture, which requires all components to be initialized simultaneously. This can lead to significant slowness in some cases. To address this, we plan to redesign the architecture to allow per-component GPU initialization, followed by selective offloading to the CPU

[hsliuustc0106]

Please attach your design doc using this template in your RFC :)

[ZJY0516]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: dd7c1c1b4c

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[hsliuustc0106]

change to enable-cpu-offload

[bwyangseek]

It should be a separate task. @bwyangseek

Thank you for your kind reply！I have proposed a PR #798 based on the previous discussions to implement overlapping, with improvements and optimizations. The focus is on the DiT computation, where instead of loading the entire DiT model onto the GPU as in #497, the DiT is split into block-level execution. During the computation of one block, the weights for the next block are asynchronously prefetched using independent CUDA streams to achieve compute/copy overlap, thus improving GPU utilization rather than just reducing peak memory usage. Additionally, I plan to inherit part of the ideas from #497 and modify the SequentialOffloader such that after the encoders finish computation, DiT is not fully loaded onto the GPU but performs an overlap forward.

Currently, the code has a preliminary framework, which has been submitted in PR #798, but I am still refining the logic, code, and conducting tests. I would appreciate your feedback on this approach. Thank you! @ZJY0516