[holykie]

Purpose

PR #7005 added a compute_device parameter to _prepare_lora_delta() to allow
offloading the LoRA delta matmul (delta = lora_B @ lora_A) to a non-CPU
device. However, the calling code in load_lora_into_module() sets
compute_dev = params.device if params.device.type != "cpu" else None,
which passes None when weights are on CPU — the common case during
pre-sharding LoRA fusion (PR #6948). As a result, the NPU/GPU matmul path
is never activated.

Root Cause

load_lora_into_module():
    compute_dev = params.device if params.device.type != "cpu" else None
    #                                              ↑ None when params on CPU
    #                                                → _prepare_lora_delta() skips
    #                                                  NPU/GPU branch → CPU matmul

Solution

Extract a _resolve_lora_compute_device() helper that returns the local accelerator device only when it is actually usable, else None. Since get_local_device() may return cuda:N / npu:N even without a physical device, the helper probes with a tiny allocation before trusting the lookup:

def _resolve_lora_compute_device() -> torch.device | None:
    try:
        dev = get_local_device()
        if dev.type != "cpu":
            torch.empty(1, device=dev)
            return dev
    except Exception:
        pass
    return None

It is called once per load_lora_into_module / unload_module_lora call, outside the parameter loop, and its result is used as the compute device for CPU-resident weights:
accel_dev = _resolve_lora_compute_device()

for name, params in module.named_parameters(prefix):
compute_dev = params.device if params.device.type != "cpu" else accel_dev
delta, used_keys = _prepare_lora_delta(..., compute_device=compute_dev)
This activates PR #7005's existing compute_device infrastructure: when the probe succeeds, a/b are moved to the accelerator and the delta matmul runs there; when it fails, compute_dev stays None and computation remains on CPU, so CPU-only hosts never attempt an H2D transfer.

Test Plan

Environment:

• Host: Kunpeng-920 5250 (192 cores), 2TB RAM
• NPU: Ascend 910B3 × 8 (64GB HBM each)
• Model: Wan2.2-T2V-A14B-Diffusers (118GB) + Wan2.2-Distill-Loras rank64 4-step
• Config: HSDP 4-card, USP=4
• Docker image: quay.io/ascend/vllm-omni:v0.28.0 (with [Bugfix][Diffusion] Fuse distilled LoRA weights before HSDP sharding #6948 + [Perf][Diffusion] Optimize HSDP startup via Rank-0 shared weight loading and accelerated LoRA delta computation #7005 applied)
• OMP_NUM_THREADS=1 (Docker default)

Test:

• 4-card HSDP with --lora-backend distill + dual LoRA files
• Compare LoRA fusion time before and after fix
• Verify inference output quality unchanged

vLLM Version: 0.28.0

vLLM-Omni Commit: eb11446b7 (with #6948 + #7005 manually applied)

Test Result

LoRA Fusion Time

Metric	Before (CPU matmul)	After (NPU matmul)	Speedup
Fusion per transformer	~22 min	~40s	33x
Total fusion (2 transformers)	~46 min	~80s	35x
Total startup time	~65 min	~5 min	13x

Inference Performance (unchanged)

Metric	Before	After
Inference time (4 steps)	31.32s	30.65s
Per-step latency	7.8s	7.7s
Output video quality	Correct	Correct

Logs

INFO Fusing distilled LoRA weights ... into Wan22Pipeline
INFO 800 lora keys loaded into WanTransformer3DModel.  (1st transformer, 39s)
INFO 800 lora keys loaded into WanTransformer3DModel.  (2nd transformer, 80s total)
INFO Application startup complete.
INFO [OmniTiming] total=30.65s
INFO | denoise_step_latency_ms  |  7,661.482 |
INFO | num_inference_steps      |      4.000 |

Unit Tests

pytest -o addopts="" tests/diffusion/lora/test_loader.py tests/diffusion/model_loader/test_diffusers_loader.py -v

(Run on local environment with the fix applied; all existing tests pass
without modification — the try/except fallback ensures CPU environments
are unaffected.)

[vllm-omni-review-bot]

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

Module owners: @RuixiangMa @wtomin @david6666666

Routing: @RuixiangMa via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @wtomin via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @david6666666 via module of the changed files, module named in the PR description

@holykie, 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.

[holykie]

This PR activates the compute_device infrastructure from #7005 by setting compute_dev = get_local_device() when weights are on CPU (the common case during pre-sharding LoRA fusion from #6948). Tested on Ascend 910B3: LoRA fusion time 46min → 80s (35x speedup). Pre-commit, build (3.11/3.12), and DCO all pass. Ready for review.

[SamitHuang]

This change does activate the #7005 compute_device path for the pre-sharding CPU-weight case, and the Wan2.2 HSDP numbers (46 min → 80 s fusion) are enough to justify the direction. It is not merge-ready as written: the try/except does not implement a CPU fallback, get_local_device() is not availability-checked, and there is no regression test for the bug this PR claims to fix.

Please extract a single device-resolution helper, fail over only when the accelerator is actually usable (and around H2D/matmul, not the lookup), and add a CPU-module unit test that mocks the accelerator device and checks fused weights match a CPU B @ A reference.

[SamitHuang]

get_local_device() almost never raises: CUDA/NPU get_torch_device() just returns cuda:{LOCAL_RANK} / npu:{LOCAL_RANK} even when no accelerator is present (vllm_omni/platforms/cuda/platform.py, vllm_omni/platforms/npu/platform.py). This except Exception: pass therefore does not restore CPU matmul. _prepare_lora_delta() then does .to(compute_device) with no fallback, so CPU-only runs and pytest.mark.cpu tests in tests/diffusion/lora/test_loader.py can crash on H2D. Guard with device availability (and wrap the copy/matmul, not only the lookup).

[holykie]

Addressed: Availability probe:** Extracted a single _resolve_lora_compute_device() helper that probes the device with torch.empty(1, device=dev) before use. An unusable device resolves to None and delta computation stays on CPU

[SamitHuang]

No regression test covers the bug this PR describes: params.device.type == "cpu" used to force compute_device=None and skip the #7005 path. Existing test_prepare_lora_delta_with_compute_device only passes torch.device("cpu") into _prepare_lora_delta directly. Add a load_lora_into_module test on a CPU module that patches get_local_device to a fake accelerator device and asserts the fused weight equals a CPU lora_B @ lora_A reference (and still works when the helper returns CPU / is unavailable).

[holykie]

Addressed: test_load_lora_into_module_cpu_weights_match_cpu_reference: patches get_local_device to a fake accelerator (cuda:999) on a CPU module and asserts fused weights equal the CPU reference base + lora_B @ lora_A. All 20 tests pass on a no-GPU/NPU container.

[SamitHuang]

This is a copy of the load-path block, and both sit inside the per-parameter loop (800 keys in the Wan log). Resolve the compute device once per load_lora_into_module / unload_module_lora call via a shared helper (availability check + fallback). That keeps load/unload identical and avoids repeating a broken try/except in two places.

[holykie]

Addressed: load_lora_into_module / unload_module_lora call, outside the parameter loop; both paths share it, no duplicated blocks.

[SamitHuang]

The shared _resolve_lora_compute_device() helper plus the torch.empty(1, device=dev) probe address the earlier availability/fallback issue, and resolving once outside the parameter loop in both load and unload is the right structure.

The new tests still do not lock in the bug this PR exists to fix. Both cases patch get_local_device to cuda:999 so _resolve_lora_compute_device() returns None and fusion stays on CPU. Reverting compute_dev to None when params.device is CPU would still pass CI. Please add a test that a usable accelerator device is forwarded into _prepare_lora_delta for CPU-resident weights (spy _prepare_lora_delta and force CPU compute inside the spy so the file can stay pytest.mark.cpu). Also update the PR body: it still documents the old per-parameter try/except around get_local_device().

[SamitHuang]

This patches get_local_device to cuda:999, so _resolve_lora_compute_device() returns None and load_lora_into_module never takes the #7005 accelerator path. The assertion only checks CPU base + B @ A after fallback — the same result as the pre-fix compute_dev = None behavior. Add a second test that patches _resolve_lora_compute_device to a non-CPU device, spies _prepare_lora_delta, asserts compute_device is that device for CPU params, and sets compute_device=None inside the spy so the numerical check can stay on CPU.

[holykie]

Addressed: Added test_load_lora_into_module_forwards_usable_accel_device :

• Patches _resolve_lora_compute_device to return a non-CPU device, bypassing the probe to simulate a usable accelerator.
• Spies _prepare_lora_delta and asserts every CPU param receives that device as compute_device — this fails if the assignment regresses back to else None.
• The spy forces compute_device=None for the actual math, so the numerical check still runs on CPU without a physical device.
  21 tests pass on a no-GPU/NPU container.

[SamitHuang]

This is the line that actually activates NPU/GPU matmul for CPU weights, and nothing in the new tests observes it. A regression that restores else None would still pass test_load_lora_into_module_cpu_weights_match_cpu_reference. Cover this assignment (usable accel_dev forwarded as compute_device) rather than only the unusable-device fallback.

[holykie]

Addressed: Added test_load_lora_into_module_forwards_usable_accel_device :
Patches _resolve_lora_compute_device to return a non-CPU device, bypassing the probe to simulate a usable accelerator.
Spies _prepare_lora_delta and asserts every CPU param receives that device as compute_device — this fails if the assignment regresses back to else None.
The spy forces compute_device=None for the actual math, so the numerical check still runs on CPU without a physical device.
21 tests pass on a no-GPU/NPU container.

[SamitHuang]

Previous blocking items are addressed on 035486f.

_resolve_lora_compute_device() probes with torch.empty(1, device=dev) and is resolved once per load/unload. test_load_lora_into_module_forwards_usable_accel_device spies _prepare_lora_delta and asserts CPU params receive the patched accelerator as compute_device, so restoring else None would fail CI. Fallback numerical fusion is covered by test_load_lora_into_module_cpu_weights_match_cpu_reference. Wan2.2 HSDP numbers (46 min → 80 s fusion) remain the evidence for the perf claim.

Please update the PR description: the Solution section still shows the old per-parameter try/except around get_local_device(), which is not what this diff does. Unload forwarding is still untested; that is the same assignment as load and is not blocking.

[SamitHuang]

LGTM