[Perf][Diffusion] Activate NPU/GPU matmul for LoRA delta computation when weights are on CPU (#7590)

Signed-off-by: holykie <116775684+holykie@users.noreply.github.com>

tests/diffusion/lora/test_loader.py

@@ -18,6 +18,7 @@ from vllm_omni.diffusion.lora.loader import (
     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)