| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | from transformers import PreTrainedConfig |
| 7 | |
| 8 | from vllm import envs |
| 9 | from vllm.config import get_current_vllm_config |
| 10 | from vllm.config.lora import LoRAConfig |
| 11 | from vllm.distributed.utils import divide |
| 12 | from vllm.forward_context import ( |
| 13 | ForwardContext, |
| 14 | get_forward_context, |
| 15 | is_forward_context_available, |
| 16 | ) |
| 17 | from vllm.model_executor.layers.linear import ( |
| 18 | ColumnParallelLinear, |
| 19 | LinearBase, |
| 20 | QuantizeMethodBase, |
| 21 | ReplicatedLinear, |
| 22 | RowParallelLinear, |
| 23 | ) |
| 24 | from vllm.platforms import current_platform |
| 25 | from vllm.utils.multi_stream_utils import maybe_execute_in_parallel |
| 26 | from vllm.utils.torch_utils import direct_register_custom_op |
| 27 | |
| 28 | from .base import BaseLayerWithLoRA |
| 29 | from .utils import _get_lora_aux_cuda_stream, _get_lora_device |
| 30 | |
| 31 | if envs.VLLM_LORA_ENABLE_DUAL_STREAM: |
| 32 | |
| 33 | def lora_linear_async( |
| 34 | layer_name: str, |
| 35 | output_size: int, |
| 36 | x: torch.Tensor, |
| 37 | bias: torch.Tensor | None = None, |
| 38 | ) -> torch.Tensor: |
| 39 | forward_context: ForwardContext = get_forward_context() |
| 40 | self = forward_context.no_compile_layers[layer_name] |
| 41 | return self._apply_async_impl(x, bias) |
| 42 | |
| 43 | def lora_linear_async_fake( |
| 44 | layer_name: str, |
| 45 | output_size: int, |
| 46 | x: torch.Tensor, |
| 47 | bias: torch.Tensor | None = None, |
| 48 | ) -> torch.Tensor: |
| 49 | # The real function reshapes output back to the original 3D shape |
| 50 | # when the input has an extra batch dimension (transformers backend). |
| 51 | if x.ndim == 3: |
| 52 | return torch.empty( |
| 53 | (x.size(0), x.size(1), output_size), |
| 54 | device=x.device, |
| 55 | dtype=x.dtype, |
| 56 | ) |
| 57 | return torch.empty( |
| 58 | (x.size(0), output_size), |
| 59 | device=x.device, |
| 60 | dtype=x.dtype, |
| 61 | ) |
| 62 | |
| 63 | direct_register_custom_op( |
| 64 | op_name="lora_linear_async", |
| 65 | op_func=lora_linear_async, |
| 66 | fake_impl=lora_linear_async_fake, |
| 67 | ) |
| 68 | |
| 69 | |
| 70 | class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): |
| 71 | # The adapter branch consumes the original activation, independently of |
| 72 | # whether the wrapped base layer can consume a pre-quantized activation. |
| 73 | requires_unquantized_input = True |
| 74 | |
| 75 | def __init__(self, base_layer: LinearBase): |
| 76 | super().__init__() |
| 77 | |
| 78 | self._enable_aux_cuda_stream = envs.VLLM_LORA_ENABLE_DUAL_STREAM |
| 79 | self.base_layer = base_layer |
| 80 | self.input_size = self.base_layer.input_size |
| 81 | # Ensure tp_size and tp_rank consistency with the base_layer. |
| 82 | self.tp_size = self.base_layer.tp_size |
| 83 | self.tp_rank = self.base_layer.tp_rank |
| 84 | self.device = _get_lora_device(self.base_layer) |
| 85 | self._init_lora_stream_context() |
| 86 | self.output_slices: tuple[int, ...] |
| 87 | self.output_size: int |
| 88 | self.n_slices: int |
| 89 | |
| 90 | def _init_lora_stream_context(self) -> None: |
| 91 | if not self._enable_aux_cuda_stream: |
| 92 | return |
| 93 | vllm_config = get_current_vllm_config() |
| 94 | self._lora_stream = _get_lora_aux_cuda_stream() |
| 95 | assert current_platform.is_cuda_alike() |
| 96 | self._events = [torch.cuda.Event(), torch.cuda.Event()] |
| 97 | # lora_linear avoids prefix conflicts with the base layer |
| 98 | self.layer_name = self.base_layer.prefix + ".lora_linear_async" |
| 99 | compilation_config = vllm_config.compilation_config |
| 100 | if self.layer_name in compilation_config.static_forward_context: |
| 101 | raise ValueError("Duplicate layer name: {}".format(self.layer_name)) |
| 102 | compilation_config.static_forward_context[self.layer_name] = self |
| 103 | |
| 104 | def create_lora_weights( |
| 105 | self, |
| 106 | max_loras: int, |
| 107 | lora_config: LoRAConfig, |
| 108 | model_config: PreTrainedConfig | None = None, |
| 109 | ) -> None: |
| 110 | self.lora_config = lora_config |
| 111 | if isinstance(self.base_layer, ReplicatedLinear): |
| 112 | lora_a_out_size = lora_config.max_lora_rank |
| 113 | lora_b_out_size = self.output_size |
| 114 | |
| 115 | elif isinstance(self.base_layer, ColumnParallelLinear): |
| 116 | lora_a_out_size = ( |
| 117 | lora_config.max_lora_rank |
| 118 | if not lora_config.fully_sharded_loras |
| 119 | else divide(lora_config.max_lora_rank, self.tp_size) |
| 120 | ) |
| 121 | lora_b_out_size = self.output_size |
| 122 | |
| 123 | elif isinstance(self.base_layer, RowParallelLinear): |
| 124 | lora_a_out_size = lora_config.max_lora_rank |
| 125 | lora_b_out_size = ( |
| 126 | self.output_size |
| 127 | if not lora_config.fully_sharded_loras |
| 128 | else divide(self.output_size, self.tp_size) |
| 129 | ) |
| 130 | else: |
| 131 | raise NotImplementedError |
| 132 | |
| 133 | self.lora_a_stacked = tuple( |
| 134 | torch.zeros( |
| 135 | max_loras, |
| 136 | 1, |
| 137 | lora_a_out_size, |
| 138 | self.input_size, |
| 139 | dtype=lora_config.lora_dtype, |
| 140 | device=self.device, |
| 141 | ) |
| 142 | for _ in range(self.n_slices) |
| 143 | ) |
| 144 | self.lora_b_stacked = tuple( |
| 145 | torch.zeros( |
| 146 | max_loras, |
| 147 | 1, |
| 148 | lora_b_out_size, |
| 149 | lora_config.max_lora_rank, |
| 150 | dtype=lora_config.lora_dtype, |
| 151 | device=self.device, |
| 152 | ) |
| 153 | for _ in range(self.n_slices) |
| 154 | ) |
| 155 | self.output_slices = (self.lora_b_stacked[0].shape[2],) |
| 156 | |
| 157 | def reset_lora(self, index: int): |
| 158 | for s_index in range(self.n_slices): |
| 159 | self.lora_a_stacked[s_index][index] = 0 |
| 160 | self.lora_b_stacked[s_index][index] = 0 |
| 161 | |
| 162 | def set_lora( |
| 163 | self, |
| 164 | index: int, |
| 165 | lora_a: torch.Tensor | list[torch.Tensor], |
| 166 | lora_b: torch.Tensor | list[torch.Tensor], |
| 167 | ): |
| 168 | # Except for QKVParallelLinearWithLoRA and |
| 169 | # MergedColumnParallelLinearWithLoRA, all other linear LoRA layers |
| 170 | # store weights in a tuple of size 1. These two layers will |
| 171 | # override this function. |
| 172 | assert isinstance(lora_a, torch.Tensor) |
| 173 | assert isinstance(lora_b, torch.Tensor) |
| 174 | assert ( |
| 175 | len(self.lora_a_stacked) == len(self.lora_b_stacked) == self.n_slices == 1 |
| 176 | ) |
| 177 | |
| 178 | self.reset_lora(index) |
| 179 | if self.tp_size > 1: |
| 180 | lora_a = self.slice_lora_a(lora_a) |
| 181 | lora_b = self.slice_lora_b(lora_b) |
| 182 | |
| 183 | self.lora_a_stacked[0][index, 0, : lora_a.shape[0], : lora_a.shape[1]].copy_( |
| 184 | lora_a, non_blocking=True |
| 185 | ) |
| 186 | self.lora_b_stacked[0][index, 0, : lora_b.shape[0], : lora_b.shape[1]].copy_( |
| 187 | lora_b, non_blocking=True |
| 188 | ) |
| 189 | |
| 190 | def _get_quant_method(self) -> QuantizeMethodBase: |
| 191 | quant_method = self.base_layer.quant_method |
| 192 | if quant_method is None: |
| 193 | raise RuntimeError( |
| 194 | f"{type(self.base_layer).__name__} must define quant_method for LoRA." |
| 195 | ) |
| 196 | return quant_method |
| 197 | |
| 198 | def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: |
| 199 | # is_forward_context_available for tower modules |
| 200 | if self._enable_aux_cuda_stream and is_forward_context_available(): |
| 201 | output_size = sum(self.output_slices) |
| 202 | return torch.ops.vllm.lora_linear_async( |
| 203 | self.layer_name, output_size, x, bias |
| 204 | ) |
| 205 | else: |
| 206 | return self._apply_sync(x, bias) |
| 207 | |
| 208 | def _apply_sync( |
| 209 | self, x: torch.Tensor, bias: torch.Tensor | None = None |
| 210 | ) -> torch.Tensor: |
| 211 | output = self._get_quant_method().apply(self.base_layer, x, bias) |
| 212 | return self._apply_lora_to_output(x, output) |
| 213 | |
| 214 | def _apply_base_forward(self, x: torch.Tensor) -> torch.Tensor: |
| 215 | base_output = self.base_layer(x) |
| 216 | output = base_output[0] if isinstance(base_output, tuple) else base_output |
| 217 | return self._apply_lora_to_output(x, output) |
| 218 | |
| 219 | def _apply_lora_to_output( |
| 220 | self, x: torch.Tensor, output: torch.Tensor |
| 221 | ) -> torch.Tensor: |
| 222 | original_shape = output.shape if output.ndim == 3 else None |
| 223 | |
| 224 | # In transformers backend, x and output have extra batch dimension like |
| 225 | # (1, seq_len, hidden_dim), while punica expects (seq_len, hidden_dim), |
| 226 | # therefore we need to flatten the batch dimensions. |
| 227 | if x.ndim == 3 and output.ndim == 3: |
| 228 | output = output.flatten(0, 1) |
| 229 | x = x.flatten(0, 1) |
| 230 | |
| 231 | lora_output: torch.Tensor | None = self.punica_wrapper.add_lora_linear( |
| 232 | output, x, self.lora_a_stacked, self.lora_b_stacked, 1.0, self.output_slices |
| 233 | ) |
| 234 | if not current_platform.can_update_inplace(): |
| 235 | output = lora_output |
| 236 | |
| 237 | # Reshape the flattened output back to its original shape, |
| 238 | # as some MM encoders cannot handle flattened inputs. |
| 239 | if original_shape is not None: |
| 240 | output = output.reshape(original_shape) |
| 241 | |
| 242 | return output |
| 243 | |
| 244 | def _apply_async_impl( |
| 245 | self, x: torch.Tensor, bias: torch.Tensor | None = None |
| 246 | ) -> torch.Tensor: |
| 247 | """Forward pass with base linear and LoRA on separate CUDA streams |
| 248 | for overlap, using maybe_execute_in_parallel. |
| 249 | Base layer runs on default stream; LoRA runs on aux stream. |
| 250 | """ |
| 251 | assert envs.VLLM_LORA_ENABLE_DUAL_STREAM |
| 252 | assert x.ndim in (2, 3) |
| 253 | num_tokens = x.size(0) if x.ndim == 2 else x.size(1) |
| 254 | output_size = sum(self.output_slices) |
| 255 | |
| 256 | def base_fn() -> torch.Tensor: |
| 257 | return self._get_quant_method().apply(self.base_layer, x, bias) |
| 258 | |
| 259 | def lora_fn() -> torch.Tensor: |
| 260 | # Must be zeros, not empty: _lora_expand_kernel exits early (without |
| 261 | # writing) when lora_id == -1 (no active LoRA). If uninitialized, |
| 262 | # output.add_(lora_result) below would corrupt the base output. |
| 263 | lora_output = torch.zeros( |
| 264 | (num_tokens, output_size), |
| 265 | device=self.device, |
| 266 | dtype=x.dtype, |
| 267 | ) |
| 268 | |
| 269 | # Flatten the batch dimension for the transformers backend |
| 270 | # (which uses shape (1, seq_len, hidden)), matching _apply_sync. |
| 271 | x_2d = x.flatten(0, 1) if x.ndim == 3 else x |
| 272 | self.punica_wrapper.add_lora_linear( |
| 273 | lora_output, |
| 274 | x_2d, |
| 275 | self.lora_a_stacked, |
| 276 | self.lora_b_stacked, |
| 277 | 1.0, |
| 278 | self.output_slices, |
| 279 | add_inputs=False, |
| 280 | ) |
| 281 | return lora_output |
| 282 | |
| 283 | output, lora_result = maybe_execute_in_parallel( |
| 284 | base_fn, |
| 285 | lora_fn, |
| 286 | self._events[0], |
| 287 | self._events[1], |
| 288 | self._lora_stream, |
| 289 | ) |
| 290 | |
| 291 | original_shape = output.shape if output.ndim == 3 else None |
| 292 | |
| 293 | # In transformers backend, x and output have extra batch dimension like |
| 294 | # (1, seq_len, hidden_dim), while punica expects (seq_len, hidden_dim), |
| 295 | # therefore we need to flatten the batch dimensions. |
| 296 | if x.ndim == 3 and output.ndim == 3: |
| 297 | output = output.flatten(0, 1) |
| 298 | x = x.flatten(0, 1) |
| 299 | |
| 300 | output.add_(lora_result) |
| 301 | |
| 302 | # Reshape the flattened output back to its original shape, |
| 303 | # as some MM encoders cannot handle flattened inputs. |
| 304 | if original_shape is not None: |
| 305 | output = output.reshape(original_shape) |
| 306 | |
| 307 | return output |
| 308 | |
| 309 | @property |
| 310 | def weight(self) -> torch.Tensor: |
| 311 | # unquantizedLinear |
| 312 | if hasattr(self.base_layer, "weight"): |
| 313 | return self.base_layer.weight |
| 314 | # Compressed Tensor |
| 315 | elif hasattr(self.base_layer, "weight_packed"): |
| 316 | return self.base_layer.weight_packed |
| 317 | # GPTQ/AWQ |
| 318 | elif hasattr(self.base_layer, "qweight"): |
| 319 | return self.base_layer.qweight |
| 320 | # marlin |
| 321 | elif hasattr(self.base_layer, "B"): |
| 322 | return self.base_layer.B |
| 323 | else: |
| 324 | raise ValueError(f"Unsupported base layer: {self.base_layer}") |
| 325 | |
| 326 | @property |
| 327 | def bias(self) -> torch.Tensor | None: |
| 328 | if hasattr(self.base_layer, "bias"): |
| 329 | return self.base_layer.bias |
| 330 | else: |
| 331 | return None |