[he-yufeng]

Purpose

Fixes #8817.

DiffusionLoRAManager only homed the stacked LoRA buffers (lora_a_stacked / lora_b_stacked) to the execution device under DLO. Under every other offload strategy the buffers kept whatever device the base weights happened to sit on at creation time, and set_lora only copy_s into them. An adapter registered while the diffusion weights are offloaded (the standard RL weight-sync moment, where engines sleep to free memory) therefore stranded the buffers on CPU, and after wake DiffusionBaseLinearLayerWithLoRA.apply() matmulled device activations against host buffers. The same hole existed on the rank-growth path: _ensure_max_lora_rank re-runs create_lora_weights, which re-allocated on the offloaded weight device.

Sleep/offload never sees these sidecars (the offloader package has no notion of LoRA), so nothing ever re-homes them. The only safe placement is the compute device, under every strategy.

Fix

• DiffusionLoRAManager now always records the execution device as _resident_lora_device instead of gating it to OffloadStrategy.DISTRIBUTED_LAYER_WISE.
• The replacement path calls _set_diffusion_lora_buffer_device whenever the layer offers the hook. Layers without the hook keep the legacy placement (buffers follow the base weight device), so third-party layer types that predate the hook are unaffected.
• No change to DiffusionBaseLinearLayerWithLoRA: its create_lora_weights already re-homes via the recorded device, which now also covers the rank-growth reinit path, pinning the second exposure from the issue.

Tests

New test_lora_manager_homes_lora_buffers_to_compute_device_without_dlo in tests/diffusion/lora/test_lora_manager.py: a non-DLO pipeline, a recording layer, and a distinct compute device (meta). It asserts the replacement path homes the buffers and records the device that create_lora_weights consults on rank growth. It fails on current main and passes with this change. The existing DLO counterpart and the layer-level reallocation test keep passing unchanged.

Validation

• pytest tests/diffusion/lora/ -o addopts="": 73 passed; the 3 hunyuan_image3 failures and the 2 collection errors are pre-existing on main (byte-identical failure set with the fix stashed).
• New test fails without the source change (verified via git stash of the manager change) and passes with it.
• pre-commit run on the changed files: all hooks pass, including the locally enforced gates (check-spdx-header, mypy-3.10, CI-marks, forbidden imports, torch.cuda ratchet).
• Not run: GPU serving repro (no CUDA device on this machine). The device-stranding chain is exercised by the unit test through the same _replace_layers_with_lora and create_lora_weights paths the engine uses.

AI assistance: Kimi Code drafted the patch, test, and PR description. I reviewed the changes and verified the test reproduces the stranded-buffer behavior from #8817.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/diffusion/index.md, docs/design/module/diffusion/diffusion_model_integration.md.

Module owners: @david6666666 @Isotr0py @princepride

Routing: @david6666666 via module of the changed files, module named in the PR description; @Isotr0py via module of the changed files, module named in the PR description; @princepride via module of the changed files, module named in the PR description

@he-yufeng, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[Dong1017]

f155ab04: The decisive case is buffer reallocation while the base weights are offloaded. On H200 I forced the base layer onto CPU, grew the LoRA rank from 8 to 16, then moved only the base back to GPU and ran the production wrapped forward. Both recreated A/B buffers remained on CUDA, and the FP32 result matched base + (x @ A.T) @ B.T with zero observed error in this fixture.

This validates the ownership choice: the manager records the execution device, and create_lora_weights reapplies it after allocation. Placement must not follow the base weight's temporary offload location. The test exercises the #8817 reallocation mechanism without relying on the offloader to discover LoRA sidecars.

[he-yufeng]

Thanks for running the H200 pass, @Dong1017 — reallocation with the base offloaded is the case the fix is load-bearing for, and your matrix (rank grow with base on CPU, then base-only move back) exercises exactly the recreate path that used to leak buffers back to the host side. Glad the numbers line up.

[he-yufeng]

Rebased onto current main (da572cd, which includes the vLLM v0.32.0 alignment). The conflict was additive only: upstream's new max_num_batched_tokens validation stays, and the resident-device change now applies unconditionally on top of it. The premise is still live on main — _resident_lora_device remains gated on the DLO strategy there.

Re-verified locally after the rebase: the full tests/diffusion/lora/test_lora_manager.py file passes 56/56 (CPU, VLLM_TARGET_DEVICE=cpu), and the new test_lora_manager_homes_lora_buffers_to_compute_device_without_dlo still fails when run against the unpatched main-side manager, so the regression pin is intact.

[he-yufeng]

Post-rebase verification on 303f68d (CPU, torch 2.14.0, vLLM v0.32.0 source tree): tests/diffusion/lora/test_lora_manager.py is 56/56 green, including the non-DLO residency case. Against main's manager.py the new test_lora_manager_homes_lora_buffers_to_compute_device_without_dlo fails as expected, so the test still discriminates after the rebase.

[Gaohan123]

LGTM. And please submit a PR to branch v0.32.0. Thanks