[Test][LoRA] Use lightweight CPU reference and skip heavy cleanup in punica ops tests (#47534)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Chaojun Zhang committed
Jul 6, 2026 at 11:29 UTC
6569df6a3ecf2f59313a40aaf001e1b208fcd2f7
1 file changed
+79
-49
tests/lora/test_punica_ops.py
+79
-49
@@ -5,7 +5,6 @@ from threading import Lock
5
import pytest
6
import torch
7
8
-import vllm.lora.ops.torch_ops as torch_ops
8
import vllm.lora.ops.triton_ops as triton_ops
9
from vllm.lora.ops.triton_ops import LoRAKernelMeta
10
from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
@@ -22,6 +21,59 @@ def reset_device(reset_default_device):
21
pass
22
23
24
+@pytest.fixture(autouse=True)
25
+def cleanup_fixture():
26
+ """Override conftest's cleanup_fixture— not needed for punica tests."""
27
+ yield
28
+
29
+
30
+@pytest.fixture(autouse=True)
31
+def dynamo_reset():
32
+ """Override conftest's dynamo_reset — not needed for punica tests."""
33
+ yield
34
+
35
+
36
+def _cpu_bgmv_shrink(
37
+ inputs, lora_weight, output, seq_len_tensor, lora_indices, scaling=1.0
38
+):
39
+ """Memory-efficient shrink reference: per-LoRA matmul loop on CPU.
40
+ output[mask] = scaling * inputs[mask] @ weight.T"""
41
+ exploded = torch.repeat_interleave(lora_indices, seq_len_tensor)
42
+ for lid in exploded.unique():
43
+ if lid < 0:
44
+ continue
45
+ mask = exploded == lid
46
+ inp = inputs[mask].to(output.dtype)
47
+ w = lora_weight[lid].to(output.dtype)
48
+ output[mask] = scaling * (inp @ w.T)
49
+
50
+
51
+def _cpu_bgmv_expand(
52
+ inputs,
53
+ lora_weight,
54
+ output,
55
+ seq_len_tensor,
56
+ lora_indices,
57
+ offset=0,
58
+ add_inputs=False,
59
+):
60
+ """Memory-efficient expand reference: per-LoRA matmul loop on CPU.
61
+ output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T"""
62
+ exploded = torch.repeat_interleave(lora_indices, seq_len_tensor)
63
+ for lid in exploded.unique():
64
+ if lid < 0:
65
+ continue
66
+ mask = exploded == lid
67
+ inp = inputs[mask].to(output.dtype)
68
+ w = lora_weight[lid].to(output.dtype)
69
+ n = w.shape[0]
70
+ result = inp @ w.T
71
+ if add_inputs:
72
+ output[mask, offset : offset + n] += result
73
+ else:
74
+ output[mask, offset : offset + n] = result
75
+
76
+
77
# Utility shrink and expand operations used as reference implementations.
78
def sgmv_shrink_for_nslices(
79
nslices: int,
@@ -36,22 +88,21 @@ def sgmv_shrink_for_nslices(
88
num_tokens: int,
89
scaling: float,
90
):
39
- """
40
- Wrapper around torch_ops.sgmv_shrink that handles any nslices.
41
- """
91
+ """CPU reference for sgmv_shrink using per-LoRA matmul loop."""
92
+ inp_cpu = inputs_tensor.cpu()
93
+ seq_cpu = seq_len_tensor.cpu()
94
+ idx_cpu = prompt_lora_mapping.cpu()
95
+ out_cpu = out_tensor.cpu()
96
for index in range(nslices):
43
- torch_ops.sgmv_shrink(
44
- inputs_tensor,
45
- lora_weights_lst[index],
46
- out_tensor[index],
47
- b_seq_start_loc,
48
- seq_len_tensor,
49
- prompt_lora_mapping,
50
- batches,
51
- max_seq_length,
52
- num_tokens,
53
- scaling,
97
+ _cpu_bgmv_shrink(
98
+ inp_cpu,
99
+ lora_weights_lst[index].cpu(),
100
+ out_cpu[index],
101
+ seq_cpu,
102
+ idx_cpu,
103
+ scaling=scaling,
104
)
105
+ out_tensor.copy_(out_cpu)
106
107
108
def sgmv_expand_for_nslices(
@@ -68,42 +119,21 @@ def sgmv_expand_for_nslices(
119
num_tokens: int,
120
add_inputs: bool,
121
) -> None:
71
- """
72
- Wrapper around torch_ops.sgmv_expand that handles any nslices.
73
- """
74
- if nslices == 1:
75
- # Verify the torch's sgmv_expand op
76
- torch_ops.sgmv_expand(
77
- inputs_tensor[0],
78
- lora_weights_lst[0],
79
- out_tensor,
80
- b_seq_start_loc,
81
- seq_len_tensor,
82
- prompt_lora_mapping,
83
- batches,
84
- max_seq_length,
85
- num_tokens,
122
+ """CPU reference for sgmv_expand using per-LoRA matmul loop."""
123
+ seq_cpu = seq_len_tensor.cpu()
124
+ idx_cpu = prompt_lora_mapping.cpu()
125
+ out_cpu = out_tensor.cpu()
126
+ for index in range(nslices):
127
+ _cpu_bgmv_expand(
128
+ inputs_tensor[index].cpu(),
129
+ lora_weights_lst[index].cpu(),
130
+ out_cpu,
131
+ seq_cpu,
132
+ idx_cpu,
133
+ offset=hidden_size * index,
134
add_inputs=add_inputs,
135
)
88
- else:
89
- slice_offset = 0
90
- for index in range(nslices):
91
- lora_weights = lora_weights_lst[index]
92
- torch_ops.sgmv_expand_slice(
93
- inputs_tensor[index],
94
- lora_weights,
95
- out_tensor,
96
- b_seq_start_loc,
97
- seq_len_tensor,
98
- prompt_lora_mapping,
99
- batches,
100
- max_seq_length,
101
- num_tokens,
102
- slice_offset,
103
- hidden_size,
104
- add_inputs=add_inputs,
105
- )
106
- slice_offset += hidden_size
136
+ out_tensor.copy_(out_cpu)
137
138
139
_dict_lock = Lock()