From 035486fce6c816d80abc4845fb68f342e97fd613 Mon Sep 17 00:00:00 2001
From: holykie <116775684+holykie@users.noreply.github.com>
Date: Tue, 15 Sep 2026 20:52:17 +0800
Subject: [PATCH] [Perf][Diffusion] Activate NPU/GPU matmul for LoRA delta
 computation when weights are on CPU

PR #7005 added compute_device parameter to _prepare_lora_delta() but the
caller in load_lora_into_module() passes None when weights are on CPU,
leaving the NPU/GPU matmul path never activated. This PR sets compute_dev
to get_local_device() when params are on CPU, with try/except fallback.

Tested on Ascend 910B3 with Wan2.2-T2V-A14B-Diffusers + Distill-LoRA:
- LoRA fusion: 46min -> 80s (35x speedup)
- Inference: unchanged (30.65s, 4 steps)

Signed-off-by: holykie <116775684+holykie@users.noreply.github.com>
---
 tests/diffusion/lora/test_loader.py | 79 +++++++++++++++++++++++++++++
 vllm_omni/diffusion/lora/loader.py  | 27 +++++++++-
 2 files changed, 104 insertions(+), 2 deletions(-)

diff --git a/tests/diffusion/lora/test_loader.py b/tests/diffusion/lora/test_loader.py
index 4fd8769372c..a10d609d825 100644
--- a/tests/diffusion/lora/test_loader.py
+++ b/tests/diffusion/lora/test_loader.py
@@ -18,6 +18,7 @@
     WanLoraLoaderMixin,
     _prepare_lora_delta,
     _remap_state_dict_keys,
+    _resolve_lora_compute_device,
 )
 
 HEAD_DIM = 24
@@ -521,3 +522,81 @@ def test_prepare_lora_delta_with_compute_device():
     assert delta is not None
     assert delta.shape == (HEAD_DIM, HEAD_DIM)
     assert len(used_keys) == 2
+
+
+def test_resolve_lora_compute_device_probes_unusable_accelerator(mocker: MockerFixture):
+    """cuda:N/npu:N may be reported without a physical device; the probe must reject it."""
+    mocker.patch(
+        "vllm_omni.diffusion.lora.loader.get_local_device",
+        return_value=torch.device("cuda:999"),
+    )
+    assert _resolve_lora_compute_device() is None
+
+
+def test_load_lora_into_module_cpu_weights_match_cpu_reference(mocker: MockerFixture):
+    """Regression test: params on CPU used to force compute_device=None and skip the
+    #7005 accelerator path. With an unusable accelerator, fusion must fall back to
+    CPU and still produce base + lora_B @ lora_A."""
+    pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
+    original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
+
+    lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
+    lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
+
+    mocker.patch(
+        "vllm_omni.diffusion.lora.loader.get_local_device",
+        return_value=torch.device("cuda:999"),
+    )
+
+    used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
+    assert len(used_keys) > 0
+    assert not (set(lora_state_dict.keys()) - used_keys)
+
+    for name, param in pipeline.transformer.named_parameters():
+        if not name.endswith(".weight"):
+            continue
+        base_key = name[: -len(".weight")]
+        key_a = f"{base_key}.lora_A.weight"
+        key_b = f"{base_key}.lora_B.weight"
+        if key_a in lora_state_dict and key_b in lora_state_dict:
+            expected = original_params[name] + (
+                lora_state_dict[key_b].to(param.dtype) @ lora_state_dict[key_a].to(param.dtype)
+            )
+            assert_close(param, expected)
+
+
+def test_load_lora_into_module_forwards_usable_accel_device(mocker: MockerFixture):
+    """A usable accelerator must be forwarded as compute_device for CPU params.
+    A spy forces the actual math back onto CPU so the test runs without a device."""
+    fake_dev = torch.device("cuda:0")
+    mocker.patch(
+        "vllm_omni.diffusion.lora.loader._resolve_lora_compute_device",
+        return_value=fake_dev,
+    )
+
+    seen_devices = []
+    original = _prepare_lora_delta
+
+    def spy(state_dict, base_key, *args, **kwargs):
+        seen_devices.append(kwargs.get("compute_device"))
+        kwargs["compute_device"] = None
+        return original(state_dict, base_key, *args, **kwargs)
+
+    mocker.patch(
+        "vllm_omni.diffusion.lora.loader._prepare_lora_delta",
+        side_effect=spy,
+    )
+
+    pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
+    original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
+
+    lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
+    lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
+
+    used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
+
+    # Every CPU param with a LoRA delta must have received the accelerator device.
+    assert seen_devices
+    assert all(d == fake_dev for d in seen_devices)
+    assert not (set(lora_state_dict.keys()) - used_keys)
+    assert_params_not_equal(pipeline.transformer, original_params)
diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py
index eb3f35188ea..8ed60ea066a 100644
--- a/vllm_omni/diffusion/lora/loader.py
+++ b/vllm_omni/diffusion/lora/loader.py
@@ -13,6 +13,7 @@
 from safetensors.torch import load_file
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline
 from vllm_omni.transformers_utils.repo_utils import hf_api
 
@@ -120,6 +121,22 @@ def _prepare_lora_delta(
     return delta, used_keys
 
 
+def _resolve_lora_compute_device() -> torch.device | None:
+    """Return the local accelerator device if actually usable, else None.
+
+    get_local_device() may return cuda:N/npu:N even when no physical device is
+    present, so probe with a tiny allocation before trusting the lookup.
+    """
+    try:
+        dev = get_local_device()
+        if dev.type != "cpu":
+            torch.empty(1, device=dev)
+            return dev
+    except Exception:
+        pass
+    return None
+
+
 def _load_lora_state_dict(
     pretrained_model_name_or_path: str,
     weights_name: str | None = None,
@@ -270,6 +287,9 @@ def load_lora_into_module(
             for param_name, weight_name, _ in module.stacked_params_mapping:
                 param_to_weight_names[param_name].append(weight_name)
 
+        # Resolve the accelerator fallback once, outside the parameter loop.
+        accel_dev = _resolve_lora_compute_device()
+
         for name, params in module.named_parameters(prefix):
             is_bias = False
             if name.endswith(".bias"):
@@ -280,7 +300,7 @@ def load_lora_into_module(
             else:
                 continue
 
-            compute_dev = params.device if params.device.type != "cpu" else None
+            compute_dev = params.device if params.device.type != "cpu" else accel_dev
             delta, used_keys = _prepare_lora_delta(
                 state_dict,
                 base_key,
@@ -324,6 +344,9 @@ def unload_module_lora(
             for param_name, weight_name, _ in module.stacked_params_mapping:
                 param_to_weight_names[param_name].append(weight_name)
 
+        # Resolve the accelerator fallback once, outside the parameter loop.
+        accel_dev = _resolve_lora_compute_device()
+
         for name, param in module.named_parameters(prefix):
             is_bias = False
             if name.endswith(".bias"):
@@ -334,7 +357,7 @@ def unload_module_lora(
             else:
                 continue
 
-            compute_dev = param.device if param.device.type != "cpu" else None
+            compute_dev = param.device if param.device.type != "cpu" else accel_dev
             delta, used_keys = _prepare_lora_delta(
                 state_dict,
                 base_key,