[Bugfix][LoRA] Support ark_linear base layer in _get_lora_device (#47690)
Signed-off-by: AlejandroParedesLT <alejandroparedeslatorre@gmail.com>
Alejandro Paredes La Torre committed
Jul 11, 2026 at 17:13 UTC
9a48eef89a777fc7039d9765fb2f6fc618216948
2 files changed
+42
tests/lora/test_layers_utils.py
new
+39
@@ -0,0 +1,39 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+from torch import nn
7
+
8
+from vllm.lora.layers.utils import _get_lora_device
9
+
10
+pytestmark = pytest.mark.skip_global_cleanup
11
+
12
+
13
+def _param() -> nn.Parameter:
14
+ return nn.Parameter(torch.empty(1), requires_grad=False)
15
+
16
+
17
+def test_get_lora_device_unquantized():
18
+ base_layer = nn.Module()
19
+ base_layer.weight = _param()
20
+ assert _get_lora_device(base_layer) == base_layer.weight.device
21
+
22
+
23
+def test_get_lora_device_gptq_awq():
24
+ base_layer = nn.Module()
25
+ base_layer.qweight = _param()
26
+ assert _get_lora_device(base_layer) == base_layer.qweight.device
27
+
28
+
29
+def test_get_lora_device_ark_linear():
30
+ base_layer = nn.Module()
31
+ base_layer.ark_linear = nn.Module()
32
+ base_layer.ark_linear.qweight = _param()
33
+ assert _get_lora_device(base_layer) == base_layer.ark_linear.qweight.device
34
+
35
+
36
+def test_get_lora_device_unsupported_raises():
37
+ base_layer = nn.Module()
38
+ with pytest.raises(ValueError, match="Unsupported base layer"):
39
+ _get_lora_device(base_layer)
vllm/lora/layers/utils.py
+3
@@ -57,6 +57,9 @@ def _get_lora_device(base_layer: nn.Module) -> torch.device:
57
# GPTQ/AWQ
58
elif hasattr(base_layer, "qweight"):
59
return base_layer.qweight.device
60
+ # INC WNA16 (AutoRound)
61
+ elif hasattr(base_layer, "ark_linear"):
62
+ return base_layer.ark_linear.qweight.device
63
# MoE layer
64
elif hasattr(base_layer, "w2_weight"):
65
return base_layer.w2_weight.device