| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.lora.ops.torch_ops import ( |
| 10 | bgmv_expand, |
| 11 | bgmv_expand_slice, |
| 12 | bgmv_shrink, |
| 13 | sgmv_expand, |
| 14 | sgmv_expand_slice, |
| 15 | sgmv_shrink, |
| 16 | ) |
| 17 | |
| 18 | from .punica_base import PunicaWrapperBase |
| 19 | |
| 20 | |
| 21 | # The platforms that are compatible with the PyTorch-native implementation can |
| 22 | # inherit this class |
| 23 | class PunicaWrapperCPU(PunicaWrapperBase): |
| 24 | """PunicaWrapperCPU is designed to manage and provide metadata for the punica |
| 25 | kernel. The main function is to maintain the state information for |
| 26 | Multi-LoRA, and to provide the interface for the pytorch punica ops. |
| 27 | """ |
| 28 | |
| 29 | def __init__( |
| 30 | self, |
| 31 | max_num_batched_tokens: int, |
| 32 | max_batches: int, |
| 33 | device: torch.device | str, |
| 34 | **kwargs, |
| 35 | ): |
| 36 | PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device) |
| 37 | |
| 38 | def _shrink_prefill( |
| 39 | self, |
| 40 | y: torch.Tensor, |
| 41 | x: torch.Tensor, |
| 42 | w_t_all: torch.Tensor, |
| 43 | scale: float, |
| 44 | ): |
| 45 | # No LoRA request, so return directly |
| 46 | if self.no_lora: |
| 47 | return |
| 48 | sgmv_shrink( |
| 49 | x, |
| 50 | w_t_all, |
| 51 | y, |
| 52 | *self.prefill_metadata, |
| 53 | scale, |
| 54 | ) |
| 55 | |
| 56 | def _shrink_decode( |
| 57 | self, |
| 58 | y: torch.Tensor, |
| 59 | x: torch.Tensor, |
| 60 | w_t_all: torch.Tensor, |
| 61 | scale: float, |
| 62 | ): |
| 63 | bgmv_shrink(x, w_t_all, y, self.token_lora_indices, scale) |
| 64 | |
| 65 | def _expand_prefill( |
| 66 | self, |
| 67 | y: torch.Tensor, |
| 68 | x: torch.Tensor, |
| 69 | w_t_all: torch.Tensor, |
| 70 | add_inputs: bool, |
| 71 | ): |
| 72 | # No LoRA request, so return directly |
| 73 | if self.no_lora: |
| 74 | return |
| 75 | sgmv_expand( |
| 76 | x, |
| 77 | w_t_all, |
| 78 | y, |
| 79 | *self.prefill_metadata, |
| 80 | add_inputs, |
| 81 | ) |
| 82 | |
| 83 | def _expand_decode( |
| 84 | self, |
| 85 | y: torch.Tensor, |
| 86 | x: torch.Tensor, |
| 87 | w_t_all: torch.Tensor, |
| 88 | add_inputs: bool, |
| 89 | ): |
| 90 | bgmv_expand(x, w_t_all, y, self.token_lora_indices, add_inputs) |
| 91 | |
| 92 | def _expand_slice_prefill( |
| 93 | self, |
| 94 | y: torch.Tensor, |
| 95 | x: torch.Tensor, |
| 96 | w_t_all: torch.Tensor, |
| 97 | y_offset: int, |
| 98 | y_slice_size: int, |
| 99 | add_inputs: bool, |
| 100 | ): |
| 101 | # No LoRA request, so return directly |
| 102 | if self.no_lora: |
| 103 | return |
| 104 | sgmv_expand_slice( |
| 105 | x, |
| 106 | w_t_all, |
| 107 | y, |
| 108 | *self.prefill_metadata, |
| 109 | y_offset, |
| 110 | y_slice_size, |
| 111 | add_inputs, |
| 112 | ) |
| 113 | |
| 114 | def _expand_slice_decode( |
| 115 | self, |
| 116 | y: torch.Tensor, |
| 117 | x: torch.Tensor, |
| 118 | w_t_all: torch.Tensor, |
| 119 | y_offset: int, |
| 120 | y_slice_size: int, |
| 121 | add_inputs: bool, |
| 122 | ): |
| 123 | bgmv_expand_slice( |
| 124 | x, w_t_all, y, self.token_lora_indices, y_offset, y_slice_size, add_inputs |
| 125 | ) |
| 126 | |
| 127 | def _apply_expand( |
| 128 | self, |
| 129 | y: torch.Tensor, |
| 130 | x: torch.Tensor, |
| 131 | w_t_all: torch.Tensor, |
| 132 | y_offset: int, |
| 133 | y_slice_size: int, |
| 134 | add_inputs: bool = True, |
| 135 | ): |
| 136 | """Perform the ` y[:,y_offset:y_offset+y_slice_size]+=x@w_t_all` |
| 137 | computation, which is suitable for the |
| 138 | GEMM of lora'b. |
| 139 | """ |
| 140 | expand_slice_fun: Callable = ( |
| 141 | self._expand_slice_prefill if self.is_prefill else self._expand_slice_decode |
| 142 | ) |
| 143 | expand_slice_fun(y, x, w_t_all, y_offset, y_slice_size, add_inputs) |
| 144 | |
| 145 | def _apply_shrink( |
| 146 | self, y: torch.Tensor, x: torch.Tensor, w_t_all: torch.Tensor, scale: float |
| 147 | ): |
| 148 | """Perform the ` y+=x@w_t_all` computation, which is suitable for the |
| 149 | GEMM of lora'a. |
| 150 | When `is_prefill is` true, it indicates that it is currently the |
| 151 | prefill stage, and the `_shrink_prefill` function should be called. |
| 152 | Otherwise, it is the decode stage, and the _shrink_decode function |
| 153 | should be called. |
| 154 | """ |
| 155 | y_org = y |
| 156 | y = y.view(-1, y.shape[-1]) |
| 157 | shrink_fun: Callable = ( |
| 158 | self._shrink_prefill if self.is_prefill else self._shrink_decode |
| 159 | ) |
| 160 | shrink_fun(y, x, w_t_all, scale) |
| 161 | y = y.view_as(y_org) |
| 162 | |
| 163 | def add_shrink( |
| 164 | self, |
| 165 | y: tuple[torch.Tensor, ...] | torch.Tensor, |
| 166 | x: torch.Tensor, |
| 167 | lora_a_stacked: tuple[torch.Tensor, ...], |
| 168 | scale: float, |
| 169 | **kwargs: Any, |
| 170 | ): |
| 171 | """Performs GEMM for multiple slices of lora_a. |
| 172 | When `is_prefill is` true, it indicates that it is currently the |
| 173 | prefill stage, and the `_shrink_prefill` function should be called. |
| 174 | Otherwise, it is the decode stage, and the _shrink_decode function |
| 175 | should be called. |
| 176 | |
| 177 | Semantics: |
| 178 | for i in range(len(lora_a_stacked)): |
| 179 | y[i] += (x @ lora_a_stacked[i]) * scale |
| 180 | |
| 181 | Args: |
| 182 | y (Union[tuple[torch.Tensor, ...], torch.Tensor]): Output tensors |
| 183 | x (torch.Tensor): Input tensor |
| 184 | lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights |
| 185 | scale (float): Scaling factor for the operation |
| 186 | **kwargs: Unused; accepted for compatibility with the base class signature. |
| 187 | |
| 188 | """ |
| 189 | x = x.view(-1, x.shape[-1]) |
| 190 | # TODO fuse these kernels |
| 191 | for slice_idx in range(len(lora_a_stacked)): |
| 192 | self._apply_shrink(y[slice_idx], x, lora_a_stacked[slice_idx], scale) |
| 193 | |
| 194 | def add_expand( |
| 195 | self, |
| 196 | y: torch.Tensor, |
| 197 | x: tuple[torch.Tensor, ...] | torch.Tensor, |
| 198 | lora_b_stacked: tuple[torch.Tensor, ...], |
| 199 | output_slices: tuple[int, ...], |
| 200 | offset_start: int = 0, |
| 201 | add_inputs=True, |
| 202 | **kwargs: Any, |
| 203 | ) -> None: |
| 204 | """Performs GEMM for multiple slices of lora_b. |
| 205 | |
| 206 | Semantics: |
| 207 | for i in range(len(lora_b_stacked)): |
| 208 | slice = output_slices[i] |
| 209 | y[:, offset:offset+slice] += x[i] @ lora_b_stacked[i] |
| 210 | offset += slice |
| 211 | |
| 212 | Args: |
| 213 | y (torch.Tensor): Output tensor. |
| 214 | x (Union[tuple[torch.Tensor, ...], torch.Tensor]): Input tensors |
| 215 | lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight |
| 216 | output_slices (tuple[int, ...]): Every slice's size |
| 217 | offset_start (int): The starting position of y, defaults to 0. |
| 218 | add_inputs (bool): Defaults to True. |
| 219 | **kwargs: Unused; accepted for compatibility with the base class signature. |
| 220 | |
| 221 | """ |
| 222 | y_org = y |
| 223 | y = y.view(-1, y.shape[-1]) |
| 224 | offset_left = offset_start |
| 225 | for slice_idx in range(len(lora_b_stacked)): |
| 226 | self._apply_expand( |
| 227 | y, |
| 228 | x[slice_idx], |
| 229 | lora_b_stacked[slice_idx], |
| 230 | offset_left, |
| 231 | output_slices[slice_idx], |
| 232 | add_inputs=add_inputs, |
| 233 | ) |
| 234 | offset_left += output_slices[slice_idx] |
| 235 | y = y.view_as(y_org) |
| 236 | |
| 237 | def add_lora_embedding( |
| 238 | self, |
| 239 | y: torch.Tensor, |
| 240 | x: torch.Tensor, |
| 241 | lora_b_stacked: torch.Tensor, |
| 242 | add_inputs: bool = True, |
| 243 | **kwargs: Any, |
| 244 | ) -> None: |
| 245 | """Applies lora specifically for VocabParallelEmbeddingWithLoRA. |
| 246 | |
| 247 | Semantics: |
| 248 | y += x @ lora_b_stacked |
| 249 | |
| 250 | Args: |
| 251 | y (torch.Tensor): Output tensor. |
| 252 | x (torch.Tensor): Input tensor. |
| 253 | lora_b_stacked (torch.Tensor): lora_b's weights. |
| 254 | add_inputs (bool): Default to True. |
| 255 | **kwargs: Unused; accepted for compatibility with the base class signature. |
| 256 | |
| 257 | """ |
| 258 | # Embedding layer only need expand op |
| 259 | expand_fun: Callable = ( |
| 260 | self._expand_prefill if self.is_prefill else self._expand_decode |
| 261 | ) |
| 262 | expand_fun(y, x, lora_b_stacked, add_inputs) |
| 263 | |
| 264 | def add_lora_linear( |
| 265 | self, |
| 266 | y: torch.Tensor, |
| 267 | x: torch.Tensor, |
| 268 | lora_a_stacked: tuple[torch.Tensor, ...], |
| 269 | lora_b_stacked: tuple[torch.Tensor, ...], |
| 270 | scale: float, |
| 271 | output_slices: tuple[int, ...], |
| 272 | *, |
| 273 | buffer: tuple[torch.Tensor, ...] | None = None, |
| 274 | **kwargs: Any, |
| 275 | ) -> None: |
| 276 | """Applicable to linear-related lora. |
| 277 | |
| 278 | Semantics: |
| 279 | for i in range(len(lora_a_stacked)): |
| 280 | y[i] += ( |
| 281 | x[i].unsqueeze(0) |
| 282 | @ lora_a_stacked[indices[i], layer_idx, :, :] |
| 283 | @ lora_b_stacked[indices[i], layer_idx, :, :] |
| 284 | * scale |
| 285 | ).squeeze(0) |
| 286 | |
| 287 | Args: |
| 288 | y (torch.Tensor): Output tensor. Will be changed in-place. |
| 289 | x (torch.Tensor): Input tensor |
| 290 | lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weight. |
| 291 | lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight. |
| 292 | scale (float): Scaling factor. |
| 293 | output_slices (tuple[int, ...]): Every slice's size. |
| 294 | buffer (Optional[tuple[torch.Tensor, ...]]): Defaults to None. |
| 295 | **kwargs: Unused; accepted for compatibility with the base class signature. |
| 296 | |
| 297 | """ |
| 298 | assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices) |
| 299 | |
| 300 | if buffer is None: |
| 301 | r = lora_b_stacked[0].size(-1) |
| 302 | # We set the buffer to be float32 by default, consistent with the |
| 303 | # triton op |
| 304 | buffer = tuple( |
| 305 | torch.zeros((x.size(0), r), dtype=torch.float32, device=x.device) |
| 306 | for _ in range(len(output_slices)) |
| 307 | ) |
| 308 | self.add_shrink(buffer, x, lora_a_stacked, scale, **kwargs) |
| 309 | self.add_expand( |
| 310 | y, buffer, lora_b_stacked, output_slices, add_inputs=True, **kwargs |
| 311 | ) |
| 312 | |
| 313 | def add_lora_logits( |
| 314 | self, |
| 315 | y: torch.Tensor, |
| 316 | x: torch.Tensor, |
| 317 | lora_a_stacked: torch.Tensor, |
| 318 | lora_b_stacked: torch.Tensor, |
| 319 | scale, |
| 320 | *, |
| 321 | buffer: torch.Tensor | None = None, |
| 322 | **kwargs: Any, |
| 323 | ) -> None: |
| 324 | """Applies lora specifically for LogitsProcessorWithLoRA. |
| 325 | |
| 326 | Semantics: |
| 327 | buffer = (x @ lora_a_stacked) * scale |
| 328 | y += buffer @ lora_b_stacked |
| 329 | |
| 330 | Args: |
| 331 | y (torch.Tensor): Output tensor. |
| 332 | x (torch.Tensor): Input tensor. |
| 333 | lora_a_stacked (torch.Tensor): lora_a's weights. |
| 334 | lora_b_stacked (torch.Tensor):lora_b's weights. |
| 335 | scale (float): Scaling factor. |
| 336 | buffer (Optional[torch.Tensor]):Default to None. |
| 337 | **kwargs: Forwarded to `add_shrink` and `add_expand`. |
| 338 | |
| 339 | """ |
| 340 | y_org = y |
| 341 | y = y.view(-1, y.shape[-1]) |
| 342 | x = x.view(-1, x.shape[-1]) |
| 343 | r = lora_b_stacked.size(-1) |
| 344 | if buffer is None: |
| 345 | # We set the buffer to be float32 by default, consistent with the |
| 346 | # triton op |
| 347 | buffer = torch.zeros((x.size(0), r), dtype=torch.float32, device=x.device) |
| 348 | # LogitsProcessorWithLoRA always using bgmv. |
| 349 | bgmv_shrink(x, lora_a_stacked, buffer, self.sampler_indices, scale) |
| 350 | bgmv_expand(buffer, lora_b_stacked, y, self.sampler_indices, add_inputs=True) |
| 351 | y = y.view_as(y_org) |
| 352 | |
| 353 | def apply_lora_full_linear( |
| 354 | self, |
| 355 | y: torch.Tensor, |
| 356 | x: torch.Tensor, |
| 357 | weight_stacked: torch.Tensor, |
| 358 | bias_stacked: torch.Tensor, |
| 359 | module_enabled: torch.Tensor, |
| 360 | ) -> None: |
| 361 | indices = self.sampler_indices |
| 362 | adapter_y = torch.zeros( |
| 363 | (x.size(0), weight_stacked.size(-2)), |
| 364 | dtype=torch.float32, |
| 365 | device=x.device, |
| 366 | ) |
| 367 | bgmv_shrink(x, weight_stacked, adapter_y, indices, 1.0) |
| 368 | result = self._select_full_linear_output( |
| 369 | y, adapter_y, bias_stacked, module_enabled |
| 370 | ) |
| 371 | y.copy_(result) |