main
py 552 lines 17.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Based on:
4 Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
5 Punica: Multi-Tenant LoRA Serving.
6 https://arxiv.org/abs/2310.18547
7 """
8
9 from abc import ABC, abstractmethod
10 from typing import TYPE_CHECKING
11
12 import torch
13
14 from .utils import compute_meta, convert_mapping
15
16 if TYPE_CHECKING:
17 # avoid circuit import
18 from vllm.lora.layers import LoRAMapping
19
20
21 class PunicaWrapperABC(ABC):
22 """PunicaWrapper ABC."""
23
24 @abstractmethod
25 def update_metadata(
26 self,
27 mapping: "LoRAMapping",
28 lora_index_to_id: list[int | None],
29 max_loras: int,
30 vocab_size: int,
31 **kwargs,
32 ) -> None:
33 """Update the lora-related metadata"""
34 raise NotImplementedError
35
36 @abstractmethod
37 def add_shrink(
38 self,
39 y: tuple[torch.Tensor, ...] | torch.Tensor,
40 x: torch.Tensor,
41 lora_a_stacked: tuple[torch.Tensor, ...],
42 scale: float,
43 **kwargs,
44 ) -> torch.Tensor | None:
45 """Performs GEMM for multiple slices of lora_a."""
46 raise NotImplementedError
47
48 @abstractmethod
49 def add_expand(
50 self,
51 y: torch.Tensor,
52 x: tuple[torch.Tensor, ...] | torch.Tensor,
53 lora_b_stacked: tuple[torch.Tensor, ...],
54 output_slices: tuple[int, ...],
55 offset_start: int = 0,
56 add_inputs=True,
57 **kwargs,
58 ) -> torch.Tensor | None:
59 """Performs GEMM for multiple slices of lora_b."""
60 raise NotImplementedError
61
62 @abstractmethod
63 def add_lora_embedding(
64 self,
65 y: torch.Tensor,
66 x: torch.Tensor,
67 lora_b_stacked: torch.Tensor,
68 add_inputs: bool = True,
69 **kwargs,
70 ) -> torch.Tensor | None:
71 """Applies lora specifically for VocabParallelEmbeddingWithLoRA,
72 and this layer only requires the expand operation.
73 """
74 raise NotImplementedError
75
76 @abstractmethod
77 def add_lora_linear(
78 self,
79 y: torch.Tensor,
80 x: torch.Tensor,
81 lora_a_stacked: tuple[torch.Tensor, ...],
82 lora_b_stacked: tuple[torch.Tensor, ...],
83 scale: float,
84 output_slices: tuple[int, ...],
85 *,
86 buffer: tuple[torch.Tensor, ...] | None = None,
87 **kwargs,
88 ) -> torch.Tensor | None:
89 """Applicable to linear-related lora."""
90 raise NotImplementedError
91
92 @abstractmethod
93 def add_lora_logits(
94 self,
95 y: torch.Tensor,
96 x: torch.Tensor,
97 lora_a_stacked: torch.Tensor,
98 lora_b_stacked: torch.Tensor,
99 scale,
100 *,
101 buffer: torch.Tensor | None = None,
102 **kwargs,
103 ) -> torch.Tensor | None:
104 """Applies lora specifically for LogitsProcessorWithLoRA."""
105 raise NotImplementedError
106
107 @abstractmethod
108 def apply_lora_full_linear(
109 self,
110 y: torch.Tensor,
111 x: torch.Tensor,
112 weight_stacked: torch.Tensor,
113 bias_stacked: torch.Tensor,
114 module_enabled: torch.Tensor,
115 ) -> torch.Tensor | None:
116 """Apply request-routed full linear weights to selected rows."""
117 raise NotImplementedError
118
119
120 class PunicaWrapperBase(PunicaWrapperABC):
121 """PunicaWrapperBase is designed to manage and provide metadata for the punica
122 kernel. The main function is to maintain the state information for
123 Multi-LoRA, and to provide the interface for the punica.
124 """
125
126 def __init__(
127 self,
128 max_num_batched_tokens: int,
129 max_batches: int,
130 device: torch.device | str,
131 **kwargs,
132 ):
133 self._token_lora_indices = torch.empty(
134 max_num_batched_tokens, dtype=torch.long, device=device
135 )
136 self._sampler_indices = torch.empty(
137 max_num_batched_tokens, dtype=torch.long, device=device
138 )
139 self._sampler_indices_padded = torch.empty(
140 max_num_batched_tokens, dtype=torch.long, device=device
141 )
142 self._embeddings_indices = torch.empty(
143 max_num_batched_tokens, dtype=torch.long, device=device
144 )
145
146 # 4 is the number of indices tensors.
147 # base_indices, sampler_indices, sampler_indices_padded,
148 # embeddings_indices
149 self.indices_len: list[int | None] = [None] * 4
150 # these attributes are the information required for sgmv kernel
151 self._seq_start_locs = torch.empty(max_batches, dtype=torch.long, device=device)
152 self._seq_lengths = torch.empty(max_batches, dtype=torch.long, device=device)
153 self._lora_indices_per_batch = torch.empty(
154 max_batches, dtype=torch.long, device=device
155 )
156 self.device: torch.device = device
157 self.max_length: int = 0
158 self.token_nums: int = 0
159 self.batch_size: int = -1
160 self.is_prefill = False
161 self.no_lora = False
162
163 def _update_base_metadata(
164 self,
165 mapping: "LoRAMapping",
166 lora_index_to_id: list[int | None],
167 max_loras: int,
168 vocab_size: int,
169 ):
170 (
171 base_indices,
172 sampler_indices,
173 sampler_indices_padded,
174 embeddings_indices,
175 indices_len,
176 ) = convert_mapping(
177 mapping,
178 lora_index_to_id,
179 max_loras,
180 vocab_size,
181 self.device,
182 )
183 self._token_lora_indices[: base_indices.shape[0]].copy_(base_indices)
184 self._sampler_indices[: sampler_indices.shape[0]].copy_(sampler_indices)
185 self._sampler_indices_padded[: sampler_indices_padded.shape[0]].copy_(
186 sampler_indices_padded
187 )
188 self._embeddings_indices[: embeddings_indices.shape[0]].copy_(
189 embeddings_indices
190 )
191
192 self.indices_len[:] = indices_len
193
194 def _update_prefill_metadata(self, token_lora_tensor: torch.Tensor) -> None:
195 (
196 b_seq_start_tensor,
197 seq_length_tensor,
198 lora_indices_tensor,
199 batch_size,
200 max_length,
201 token_nums,
202 no_lora,
203 ) = compute_meta(token_lora_tensor)
204
205 self._seq_start_locs[: b_seq_start_tensor.shape[0]].copy_(b_seq_start_tensor)
206 self._seq_lengths[: seq_length_tensor.shape[0]].copy_(seq_length_tensor)
207 self._lora_indices_per_batch[: lora_indices_tensor.shape[0]].copy_(
208 lora_indices_tensor
209 )
210 self.batch_size = batch_size
211 self.max_length = max_length
212 self.token_nums = token_nums
213 self.no_lora = no_lora
214
215 def _select_full_linear_output(
216 self,
217 y: torch.Tensor,
218 adapter_y: torch.Tensor,
219 bias_stacked: torch.Tensor,
220 module_enabled: torch.Tensor,
221 ) -> torch.Tensor:
222 indices = self.sampler_indices
223 assert indices.size(0) == y.size(0), (
224 "Full linear rows do not match LoRA request mapping"
225 )
226 safe_indices = indices.clamp_min(0)
227 use_full = (indices >= 0) & module_enabled[safe_indices]
228 adapter_y = adapter_y.to(y.dtype) + bias_stacked[safe_indices].to(y.dtype)
229 return torch.where(use_full.unsqueeze(-1), adapter_y, y)
230
231 @property
232 def prefill_metadata(
233 self,
234 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int]:
235 """This property provides a convenient way to access the necessary
236 metadata for prefill-related kernel computations.
237 1. seq_start_locs: Tensor of sequence start positions.
238 2. seq_lengths: Tensor of sequence lengths.
239 3. lora_indices_per_batch: Tensor of lora indices, and an index of
240 -1 means no lora should be applied.
241 4. batch_size: Batch size after clustering identical lora indices.
242 5. max_length: The maximum sequence length in the batch.
243 6. token_nums: The token numbers in the batch.
244 """
245 return (
246 self._seq_start_locs[: self.batch_size],
247 self._seq_lengths[: self.batch_size],
248 self._lora_indices_per_batch[: self.batch_size],
249 self.batch_size,
250 self.max_length,
251 self.token_nums,
252 )
253
254 @property
255 def token_lora_indices(self) -> torch.Tensor:
256 """This property provides the lora indices corresponding to each token
257 in the batch. An index of -1 means no lora should be applied.
258 """
259 token_lora_len = self.indices_len[0]
260 return self._token_lora_indices[:token_lora_len]
261
262 @property
263 def sampler_indices(self) -> torch.Tensor:
264 """This property is used to access the lora indices specifically for
265 LogitsProcessorWithLoRA.
266 """
267 sampler_indices_len = self.indices_len[1]
268 return self._sampler_indices[:sampler_indices_len]
269
270 @property
271 def sampler_indices_padded(self) -> torch.Tensor:
272 """This property provides access to padded sampler indices."""
273 indices_padded_len = self.indices_len[2]
274 return self._sampler_indices_padded[:indices_padded_len]
275
276 def update_metadata(
277 self,
278 mapping: "LoRAMapping",
279 lora_index_to_id: list[int | None],
280 max_loras: int,
281 vocab_size: int,
282 **kwargs,
283 ):
284 self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size)
285
286 if mapping.is_prefill:
287 # Update metadata required for prefill-related operators.
288 self._update_prefill_metadata(self.token_lora_indices)
289 self.is_prefill = True
290 else:
291 self.is_prefill = False
292
293 @abstractmethod
294 def add_shrink(
295 self,
296 y: tuple[torch.Tensor, ...] | torch.Tensor,
297 x: torch.Tensor,
298 lora_a_stacked: tuple[torch.Tensor, ...],
299 scale: float,
300 **kwargs,
301 ) -> torch.Tensor | None:
302 """Performs GEMM for multiple slices of lora_a.
303
304 Semantics:
305 for i in range(len(lora_a_stacked)):
306 y[i] += (x @ lora_a_stacked[i]) * scale
307
308 Args:
309 y (Union[tuple[torch.Tensor, ...], torch.Tensor]): Output tensors
310 x (torch.Tensor): Input tensor
311 lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights
312 scale (float): Scaling factor for the operation
313
314 """
315 # TODO: implement it based on torch ops
316 raise NotImplementedError
317
318 @abstractmethod
319 def add_expand(
320 self,
321 y: torch.Tensor,
322 x: tuple[torch.Tensor, ...] | torch.Tensor,
323 lora_b_stacked: tuple[torch.Tensor, ...],
324 output_slices: tuple[int, ...],
325 offset_start: int = 0,
326 add_inputs=True,
327 **kwargs,
328 ) -> torch.Tensor | None:
329 """Performs GEMM for multiple slices of lora_b.
330
331 Semantics:
332 offset = offset_start
333 for i in range(len(lora_b_stacked)):
334 slice = output_slices[i]
335 y[:, offset:offset+slice] += x[i] @ lora_b_stacked[i]
336 offset += slice
337
338 Args:
339 y (torch.Tensor): Output tensor.
340 x (Union[tuple[torch.Tensor, ...], torch.Tensor]): Input tensors
341 lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight
342 output_slices (tuple[int, ...]): Every slice's size
343 offset_start (int): The starting position of y, defaults to 0
344 add_inputs (bool): Defaults to True.
345
346 """
347 # TODO: implement it based on torch ops
348 raise NotImplementedError
349
350 @abstractmethod
351 def add_lora_embedding(
352 self,
353 y: torch.Tensor,
354 x: torch.Tensor,
355 lora_b_stacked: torch.Tensor,
356 add_inputs: bool = True,
357 **kwargs,
358 ) -> torch.Tensor | None:
359 """Applies lora specifically for VocabParallelEmbeddingWithLoRA.
360 and this layer only requires the expand operation.
361 Semantics:
362 y += x @ lora_b_stacked
363
364 Args:
365 y (torch.Tensor): Output tensor.
366 x (torch.Tensor): Input tensor.
367 lora_b_stacked (torch.Tensor): lora_b's weights.
368 add_inputs (bool): Default to True.
369
370 """
371 # TODO: implement it based on torch ops
372 raise NotImplementedError
373
374 @abstractmethod
375 def add_lora_linear(
376 self,
377 y: torch.Tensor,
378 x: torch.Tensor,
379 lora_a_stacked: tuple[torch.Tensor, ...],
380 lora_b_stacked: tuple[torch.Tensor, ...],
381 scale: float,
382 output_slices: tuple[int, ...],
383 *,
384 buffer: tuple[torch.Tensor, ...] | None = None,
385 **kwargs,
386 ) -> torch.Tensor | None:
387 """Applicable to linear-related lora.
388
389 Semantics:
390 for i in range(len(lora_a_stacked)):
391 y[i] += (
392 x[i].unsqueeze(0)
393 @ lora_a_stacked[indices[i], layer_idx, :, :]
394 @ lora_b_stacked[indices[i], layer_idx, :, :]
395 * scale
396 ).squeeze(0)
397
398 Args:
399 y (torch.Tensor): Output tensor. Will be changed in-place.
400 x (torch.Tensor): Input tensor
401 lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weight.
402 lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight.
403 scale (float): Scaling factor.
404 output_slices (tuple[int, ...]): Every slice's size.
405 buffer (Optional[tuple[torch.Tensor, ...]]): Defaults to None.
406
407 """
408 # TODO: implement it based on torch ops
409 raise NotImplementedError
410
411 @abstractmethod
412 def add_lora_logits(
413 self,
414 y: torch.Tensor,
415 x: torch.Tensor,
416 lora_a_stacked: torch.Tensor,
417 lora_b_stacked: torch.Tensor,
418 scale,
419 *,
420 buffer: torch.Tensor | None = None,
421 **kwargs,
422 ) -> torch.Tensor | None:
423 """Applies lora specifically for LogitsProcessorWithLoRA.
424
425 Semantics:
426 buffer = (x @ lora_a_stacked) * scale
427 y += buffer @ lora_b_stacked
428
429 Args:
430 y (torch.Tensor): Output tensor.
431 x (torch.Tensor): Input tensor.
432 lora_a_stacked (torch.Tensor): lora_a's weights.
433 lora_b_stacked (torch.Tensor):lora_b's weights.
434 scale (float): Scaling factor.
435 buffer (Optional[torch.Tensor]):Default to None.
436
437 """
438 # TODO: implement it based on torch ops
439 raise NotImplementedError
440
441 def moe_lora_align_block_size(
442 self,
443 topk_ids: torch.Tensor,
444 num_tokens: int,
445 block_size: int,
446 num_experts: int,
447 max_loras: int,
448 adapter_enabled: torch.Tensor,
449 expert_map: torch.Tensor | None = None,
450 pad_sorted_ids: bool = False,
451 naive_block_assignment: bool = False,
452 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
453 """Aligns tokens and experts into block-sized chunks for LoRA-based
454 mixture-of-experts (MoE) execution.
455 """
456 # TODO: implement it based on torch ops
457 raise NotImplementedError
458
459 def add_lora_fused_moe(
460 self,
461 y: torch.Tensor,
462 x: torch.Tensor,
463 lora_a_stacked: tuple[torch.Tensor, ...],
464 lora_b_stacked: tuple[torch.Tensor, ...],
465 topk_weights: torch.Tensor,
466 sorted_token_ids: torch.Tensor | None,
467 expert_ids: torch.Tensor,
468 num_tokens_post_padded: torch.Tensor | None,
469 max_lora_rank: int,
470 top_k_num: int,
471 shrink_config,
472 expand_config,
473 adapter_enabled: torch.Tensor,
474 mul_routed_weight=False,
475 fully_sharded: bool = False,
476 offset: int = 0,
477 token_lora_mapping: torch.Tensor | None = None,
478 ):
479 """Performs a fused forward computation for LoRA of
480 Mixture-of-Experts (MoE) layer.
481 """
482 # TODO: implement it based on torch ops
483 raise NotImplementedError
484
485 def add_lora_w13(
486 self,
487 y: torch.Tensor,
488 x: torch.Tensor,
489 lora_a_stacked: tuple[torch.Tensor, ...],
490 lora_b_stacked: tuple[torch.Tensor, ...],
491 topk_ids: torch.Tensor,
492 topk_weights: torch.Tensor,
493 expert_map: torch.Tensor | None,
494 w1: torch.Tensor,
495 w2: torch.Tensor,
496 num_tokens: int,
497 top_k_num: int,
498 max_loras: int,
499 adapter_enabled: torch.Tensor,
500 local_num_experts: int,
501 top_k: int,
502 num_slices: int,
503 fully_sharded: bool,
504 use_tuned_config: bool,
505 add_inputs: bool = True,
506 token_lora_mapping: torch.Tensor | None = None,
507 ) -> tuple[
508 torch.Tensor | None,
509 torch.Tensor | None,
510 torch.Tensor | None,
511 torch.Tensor | None,
512 ]:
513 """Apply w13 LoRA to y (intermediate_cache1) in-place before activation.
514
515 When `token_lora_mapping` is provided it overrides the punica_wrapper's
516 global mapping — used by EP+LoRA to pass the per-rank-local mapping
517 after all-to-all dispatch.
518
519 Returns (sorted_token_ids_lora, expert_ids_lora,
520 num_tokens_post_padded_lora, token_lora_mapping)
521 for reuse by add_lora_w2.
522 """
523 raise NotImplementedError
524
525 def add_lora_w2(
526 self,
527 y: torch.Tensor,
528 x: torch.Tensor,
529 lora_a_stacked: tuple[torch.Tensor, ...],
530 lora_b_stacked: tuple[torch.Tensor, ...],
531 topk_weights: torch.Tensor,
532 sorted_token_ids_lora: torch.Tensor | None,
533 expert_ids_lora: torch.Tensor | None,
534 num_tokens_post_padded_lora: torch.Tensor | None,
535 token_lora_mapping: torch.Tensor | None,
536 num_tokens: int,
537 w1: torch.Tensor,
538 w2: torch.Tensor,
539 top_k_num: int,
540 max_loras: int,
541 adapter_enabled: torch.Tensor,
542 top_k: int,
543 fully_sharded: bool,
544 tp_rank: int,
545 use_tuned_config: bool,
546 add_inputs: bool = True,
547 ) -> None:
548 """Apply w2 LoRA to y (intermediate_cache3) in-place before moe_sum.
549
550 Reuses routing tensors returned by add_lora_w13.
551 """
552 raise NotImplementedError