[Feature][Wan2.2] Speed up diffusion model startup by multi-thread weight loading

[SamitHuang]

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

Purpose

The weight loading time for large diffusion model are large, ~3min for QwenImage, ~5min for Wan2.2-I2V 14B. This PR reduce weight loading time by loading safetensors shards in parallel with a thread pool instead of sequentially.

Helpful in:

• Reduce wait time for CI or benchmarking board
• Startup UX

API Changes

CLI: add --disable-multithread-weight-load to turn off, --num-weight-load-threads to set thread count

Use default 4 threads to reduce disk I/O contention on HDD/network storage

Test Plan

vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni 

vllm serve Qwen/Qwen-Image --omni --num-weight-load-threads 4

Test Result

on H800:

Wan2.2 I2V 14B

• Prev:

Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:19<00:00,  6.48s/it]
[Stage-0] INFO 02-26 03:47:27 [platform.py:77] Defaulting to diffusion attention backend FLASH_ATTN
Multi-thread loading safetensors shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 12/12 [02:05<00:00, 10.48s/it]
Multi-thread loading safetensors shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 12/12 [02:05<00:00, 10.46s/it]
[Stage-0] INFO 02-26 03:51:45 [diffusers_loader.py:277] Loading weights took 255.55 seconds
[Stage-0] INFO 02-26 03:51:45 [diffusion_model_runner.py:103] Model loading took 64.4626 GiB and 283.140117 seconds
[Stage-0] INFO 02-26 03:51:45 [diffusion_model_runner.py:108] Model runner: Model loaded successfully.

• Now:

Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:04<00:00,  1.53s/it]
[Stage-0] INFO 02-26 04:17:06 [platform.py:77] Defaulting to diffusion attention backend FLASH_ATTN
Multi-thread loading safetensors shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 12/12 [00:18<00:00,  1.51s/it]
Multi-thread loading safetensors shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 12/12 [00:19<00:00,  1.59s/it]
[Stage-0] INFO 02-26 04:17:53 [diffusers_loader.py:273] Loading weights took 44.93 seconds
[Stage-0] INFO 02-26 04:17:54 [diffusion_model_runner.py:103] Model loading took 64.4626 GiB and 55.996611 seconds
[Stage-0] INFO 02-26 04:17:54 [diffusion_model_runner.py:108] Model runner: Model loaded successfully.

Loading time cost: 283s -> 56s

QwenImage

Prev:

Loading safetensors checkpoint shards:   0% Completed | 0/9 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  11% Completed | 1/9 [00:17<02:18, 17.28s/it]
Loading safetensors checkpoint shards:  22% Completed | 2/9 [00:34<02:00, 17.15s/it]
Loading safetensors checkpoint shards:  33% Completed | 3/9 [00:51<01:42, 17.07s/it]
Loading safetensors checkpoint shards:  44% Completed | 4/9 [01:08<01:24, 16.96s/it]
Loading safetensors checkpoint shards:  56% Completed | 5/9 [01:09<00:44, 11.18s/it]
Loading safetensors checkpoint shards:  67% Completed | 6/9 [01:25<00:39, 13.12s/it]
Loading safetensors checkpoint shards:  78% Completed | 7/9 [01:42<00:28, 14.31s/it]
Loading safetensors checkpoint shards:  89% Completed | 8/9 [01:59<00:15, 15.10s/it]
Loading safetensors checkpoint shards: 100% Completed | 9/9 [02:16<00:00, 15.70s/it]
Loading safetensors checkpoint shards: 100% Completed | 9/9 [02:16<00:00, 15.17s/it]

[Stage-0] INFO 02-26 03:16:32 [diffusers_loader.py:227] Loading weights took 136.56 seconds
[Stage-0] INFO 02-26 03:16:33 [diffusion_model_runner.py:103] Model loading took 53.7462 GiB and 167.369310 seconds
[Stage-0] INFO 02-26 03:16:33 [diffusion_model_runner.py:108] Model runner: Model loaded successfully.

Now:

Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:06<00:00,  1.61s/it]
[Stage-0] INFO 02-26 03:37:10 [platform.py:77] Defaulting to diffusion attention backend FLASH_ATTN
Multi-thread loading safetensors shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 9/9 [00:13<00:00,  1.53s/it]
[Stage-0] INFO 02-26 03:37:27 [diffusers_loader.py:277] Loading weights took 15.34 seconds
[Stage-0] INFO 02-26 03:37:28 [diffusion_model_runner.py:103] Model loading took 53.7462 GiB and 27.148050 seconds
[Stage-0] INFO 02-26 03:37:28 [diffusion_model_runner.py:108] Model runner: Model loaded successfully.

Loading time cost: 168s -> 27s

---

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. Please providing the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please 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. Please run mkdocs serve to sync the documentation editions to ./docs.
• (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)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5ed55a2a5a

ℹ️ 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".

[chatgpt-codex-connector]

Wire new weight-load CLI flags into diffusion engine args

These new serve flags are parsed, but in the async serving path they are not propagated into diffusion engine_args (the default diffusion stage builder in vllm_omni/entrypoints/async_omni.py only forwards a fixed subset and omits enable_multithread_weight_load / num_weight_load_threads). Because _build_od_config only copies fields present in engine_args, --disable-multithread-weight-load (and custom thread counts) are silently ignored for default diffusion serving.

Useful? React with 👍 / 👎.

[gcanlin]

Really helpful feature. Could you please add it in docs?

[Isotr0py]

I think we can initialize DiffusersPipelineLoader with od_config, so that we don't need to access it from model:

def __init__(self, load_config: LoadConfig, od_config):
        self.load_config = load_config
        self.od_config = od_config

[SamitHuang]

good suggestion. updated

[SamitHuang]

Really helpful feature. Could you please add it in docs?

It's enabled by default. But it's ok to mention it in diffusion acceleration docs. thanks for reminding