| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from tests.lora.utils import ( |
| 8 | PunicaTensors, |
| 9 | assert_close, |
| 10 | generate_data, |
| 11 | generate_data_for_expand_nslices, |
| 12 | ) |
| 13 | from vllm.lora.ops.xpu_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink |
| 14 | from vllm.platforms import current_platform |
| 15 | |
| 16 | |
| 17 | def torch_bgmv_expand( |
| 18 | inputs: torch.Tensor, |
| 19 | lora_b_weights: torch.Tensor, |
| 20 | output_tensor: torch.Tensor, |
| 21 | lora_indices_tensor: torch.Tensor, |
| 22 | add_inputs: bool = True, |
| 23 | ): |
| 24 | selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype) |
| 25 | if len(selected_loras.shape) == 4: |
| 26 | selected_loras = selected_loras.squeeze(dim=1) |
| 27 | inputs = inputs.to(dtype=output_tensor.dtype) |
| 28 | outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras) |
| 29 | |
| 30 | limit = output_tensor.shape[0] |
| 31 | if outputs.shape[0] == 1 and output_tensor.shape[0] != 1: |
| 32 | limit = 1 |
| 33 | |
| 34 | # LoRA adapter and model may add different amounts of padding to output |
| 35 | common_len = min(outputs.shape[1], output_tensor.shape[1]) |
| 36 | |
| 37 | if add_inputs: |
| 38 | output_tensor[:, :common_len] += outputs[:limit, :common_len] |
| 39 | else: |
| 40 | output_tensor[:, :common_len] = outputs[:limit, :common_len] |
| 41 | |
| 42 | |
| 43 | def torch_bgmv_shrink( |
| 44 | inputs: torch.Tensor, |
| 45 | lora_b_weights: torch.Tensor, |
| 46 | output_tensor: torch.Tensor, |
| 47 | lora_indices_tensor: torch.Tensor, |
| 48 | scaling: float = 1.0, |
| 49 | ): |
| 50 | selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype) |
| 51 | if len(selected_loras.shape) == 4: |
| 52 | selected_loras = selected_loras.squeeze(dim=1) |
| 53 | inputs = inputs.to(dtype=output_tensor.dtype) |
| 54 | outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras) |
| 55 | |
| 56 | output_tensor[:, : outputs.shape[1]] = scaling * outputs[:] |
| 57 | |
| 58 | |
| 59 | def torch_bgmv_expand_slice( |
| 60 | inputs: torch.Tensor, |
| 61 | lora_b_weights: torch.Tensor, |
| 62 | output_tensor: torch.Tensor, |
| 63 | lora_indices_tensor: torch.Tensor, |
| 64 | slice_offset: int, |
| 65 | slice_size: int, |
| 66 | add_inputs: bool = True, |
| 67 | ): |
| 68 | selected_loras = lora_b_weights[lora_indices_tensor].to(dtype=output_tensor.dtype) |
| 69 | inputs = inputs.to(dtype=output_tensor.dtype) |
| 70 | if len(selected_loras.shape) == 4: |
| 71 | selected_loras = selected_loras.squeeze(dim=1) |
| 72 | outputs = torch.einsum("bi, boi -> bo", inputs, selected_loras) |
| 73 | |
| 74 | if add_inputs: |
| 75 | output_tensor[:, slice_offset : slice_offset + slice_size] += outputs[:] |
| 76 | else: |
| 77 | output_tensor[:, slice_offset : slice_offset + slice_size] = outputs[:] |
| 78 | |
| 79 | |
| 80 | def check_bgmv_shrink( |
| 81 | batches: int, |
| 82 | num_loras: int, |
| 83 | rank: int, |
| 84 | hidden_size: int, |
| 85 | dtype: torch.dtype, |
| 86 | device: str, |
| 87 | scaling: float, |
| 88 | ): |
| 89 | """Compare vllm.bgmv_shrink against a reference implementation.""" |
| 90 | seq_length = 1 |
| 91 | data: PunicaTensors = generate_data( |
| 92 | batches, |
| 93 | hidden_size, |
| 94 | num_loras, |
| 95 | rank, |
| 96 | seq_length, |
| 97 | dtype, |
| 98 | "shrink", |
| 99 | device, |
| 100 | ) |
| 101 | |
| 102 | bgmv_shrink( |
| 103 | data.inputs_tensor, |
| 104 | data.lora_weights, |
| 105 | data.our_out_tensor, |
| 106 | data.token_lora_mapping, |
| 107 | scaling, |
| 108 | ) |
| 109 | |
| 110 | torch_bgmv_shrink( |
| 111 | data.inputs_tensor, |
| 112 | data.lora_weights, |
| 113 | data.ref_out_tensor, |
| 114 | data.token_lora_mapping, |
| 115 | scaling, |
| 116 | ) |
| 117 | |
| 118 | data.ref_out_tensor = data.ref_out_tensor.to(torch.float32) |
| 119 | assert_close(data.our_out_tensor, data.ref_out_tensor) |
| 120 | |
| 121 | |
| 122 | def check_bgmv_expand( |
| 123 | batches: int, |
| 124 | num_loras: int, |
| 125 | rank: int, |
| 126 | hidden_size: int, |
| 127 | dtype: torch.dtype, |
| 128 | device: str, |
| 129 | add_inputs: bool, |
| 130 | ): |
| 131 | """Compare vllm.bgmv_expand against a reference implementation.""" |
| 132 | seq_length = 1 |
| 133 | data: PunicaTensors = generate_data( |
| 134 | batches, |
| 135 | hidden_size, |
| 136 | num_loras, |
| 137 | rank, |
| 138 | seq_length, |
| 139 | dtype, |
| 140 | "expand", |
| 141 | device, |
| 142 | ) |
| 143 | |
| 144 | bgmv_expand( |
| 145 | data.inputs_tensor, |
| 146 | data.lora_weights, |
| 147 | data.our_out_tensor, |
| 148 | data.token_lora_mapping, |
| 149 | add_inputs=add_inputs, |
| 150 | ) |
| 151 | torch_bgmv_expand( |
| 152 | data.inputs_tensor, |
| 153 | data.lora_weights, |
| 154 | data.ref_out_tensor, |
| 155 | data.token_lora_mapping, |
| 156 | add_inputs=add_inputs, |
| 157 | ) |
| 158 | assert_close(data.ref_out_tensor, data.our_out_tensor) |
| 159 | |
| 160 | |
| 161 | def check_bgmv_expand_slice( |
| 162 | batches: int, |
| 163 | num_loras: int, |
| 164 | rank: int, |
| 165 | hidden_size: int, |
| 166 | nslices: int, |
| 167 | dtype: torch.dtype, |
| 168 | device: str, |
| 169 | add_inputs: bool, |
| 170 | ): |
| 171 | """Compare vllm.bgmv_expand_slice against a reference implementation.""" |
| 172 | seq_length = 1 |
| 173 | data: PunicaTensors = generate_data_for_expand_nslices( |
| 174 | batches, |
| 175 | hidden_size, |
| 176 | num_loras, |
| 177 | rank, |
| 178 | seq_length, |
| 179 | dtype, |
| 180 | nslices, |
| 181 | device, |
| 182 | ) |
| 183 | |
| 184 | slice_offset = 0 |
| 185 | for index in range(nslices): |
| 186 | bgmv_expand_slice( |
| 187 | data.inputs_tensor, |
| 188 | data.lora_weights[index], |
| 189 | data.our_out_tensor, |
| 190 | data.token_lora_mapping, |
| 191 | slice_offset, |
| 192 | slice_size=hidden_size, |
| 193 | add_inputs=add_inputs, |
| 194 | ) |
| 195 | torch_bgmv_expand_slice( |
| 196 | data.inputs_tensor, |
| 197 | data.lora_weights[index], |
| 198 | data.ref_out_tensor, |
| 199 | data.token_lora_mapping, |
| 200 | slice_offset, |
| 201 | slice_size=hidden_size, |
| 202 | add_inputs=add_inputs, |
| 203 | ) |
| 204 | |
| 205 | slice_offset += hidden_size |
| 206 | assert_close(data.ref_out_tensor, data.our_out_tensor) |
| 207 | |
| 208 | |
| 209 | # General tests params that tests for variations in all dimensions |
| 210 | # except hidden_size. |
| 211 | test_params = { |
| 212 | "hidden_sizes": [2049], |
| 213 | "batches": [4], |
| 214 | "num_loras": [4], |
| 215 | "max_ranks": [32], |
| 216 | } |
| 217 | |
| 218 | DTYPES = [torch.float16, torch.bfloat16] |
| 219 | DEVICES = [f"xpu:{0}"] |
| 220 | SEED = [0] |
| 221 | |
| 222 | |
| 223 | @pytest.mark.parametrize("batches", test_params["batches"]) |
| 224 | @pytest.mark.parametrize("num_loras", test_params["num_loras"]) |
| 225 | @pytest.mark.parametrize("rank", test_params["max_ranks"]) |
| 226 | @pytest.mark.parametrize("hidden_size", test_params["hidden_sizes"]) |
| 227 | @pytest.mark.parametrize("dtype", DTYPES) |
| 228 | @pytest.mark.parametrize("device", DEVICES) |
| 229 | @pytest.mark.parametrize("seed", SEED) |
| 230 | @pytest.mark.parametrize("op_type", ["shrink", "expand"]) |
| 231 | @pytest.mark.skipif(not current_platform.is_xpu(), reason="skip for non xpu platform") |
| 232 | def test_bgmv( |
| 233 | batches: int, |
| 234 | num_loras: int, |
| 235 | rank: int, |
| 236 | hidden_size: int, |
| 237 | dtype: torch.dtype, |
| 238 | device: str, |
| 239 | seed: int, |
| 240 | op_type: str, |
| 241 | ): |
| 242 | if op_type == "shrink": |
| 243 | check_bgmv_shrink( |
| 244 | batches=batches, |
| 245 | num_loras=num_loras, |
| 246 | rank=rank, |
| 247 | hidden_size=hidden_size, |
| 248 | dtype=dtype, |
| 249 | device=device, |
| 250 | scaling=0.5, |
| 251 | ) |
| 252 | else: |
| 253 | check_bgmv_expand( |
| 254 | batches=batches, |
| 255 | num_loras=num_loras, |
| 256 | rank=rank, |
| 257 | hidden_size=hidden_size, |
| 258 | dtype=dtype, |
| 259 | device=device, |
| 260 | add_inputs=True, |
| 261 | ) |
| 262 | |
| 263 | |
| 264 | @pytest.mark.parametrize("batches", test_params["batches"]) |
| 265 | @pytest.mark.parametrize("num_loras", test_params["num_loras"]) |
| 266 | @pytest.mark.parametrize("rank", test_params["max_ranks"]) |
| 267 | @pytest.mark.parametrize("hidden_size", test_params["hidden_sizes"]) |
| 268 | @pytest.mark.parametrize("nslices", [2, 3]) |
| 269 | @pytest.mark.parametrize("dtype", DTYPES) |
| 270 | @pytest.mark.parametrize("device", DEVICES) |
| 271 | @pytest.mark.parametrize("seed", SEED) |
| 272 | @pytest.mark.skipif(not current_platform.is_xpu(), reason="skip for non xpu platform") |
| 273 | def test_bgmv_expand_nslices( |
| 274 | batches: int, |
| 275 | num_loras: int, |
| 276 | rank: int, |
| 277 | hidden_size: int, |
| 278 | nslices: int, |
| 279 | dtype: torch.dtype, |
| 280 | device: str, |
| 281 | seed: int, |
| 282 | ): |
| 283 | check_bgmv_expand_slice( |
| 284 | batches=batches, |
| 285 | num_loras=num_loras, |
| 286 | rank=rank, |
| 287 | hidden_size=hidden_size, |
| 288 | nslices=nslices, |
| 289 | dtype=dtype, |
| 290 | device=device, |
| 291 | add_inputs=True, |
| 292 | ) |