| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TYPE_CHECKING |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.utils.torch_utils import async_tensor_h2d |
| 9 | |
| 10 | if TYPE_CHECKING: |
| 11 | # avoid circuit import |
| 12 | from vllm.lora.layers import LoRAMapping |
| 13 | |
| 14 | |
| 15 | def compute_meta( |
| 16 | token_lora_tensor: torch.Tensor, |
| 17 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int, bool]: |
| 18 | """Get the information required for the sgmv kernel. With the features: |
| 19 | 1. If consecutive requests in the batch use the same LoRA, this function |
| 20 | will combine them into a single request, improving sgmv kernel inference |
| 21 | performance. |
| 22 | 2. At the beginning of each prefill stage inference, recalculations are |
| 23 | needed based on the input, but only once. |
| 24 | """ |
| 25 | lora_indices_tensor, seq_length_tensor = torch.unique_consecutive( |
| 26 | token_lora_tensor, return_counts=True |
| 27 | ) |
| 28 | cum_result = torch.cumsum(seq_length_tensor, dim=0) |
| 29 | b_seq_start_tensor = torch.zeros_like(seq_length_tensor) |
| 30 | b_seq_start_tensor[1:].copy_(cum_result[:-1]) |
| 31 | max_length = seq_length_tensor.max().item() |
| 32 | token_nums = seq_length_tensor.sum().item() |
| 33 | batch_size = lora_indices_tensor.size(0) |
| 34 | no_lora = False |
| 35 | # -1 means no lora should be applied. Use `no_lora` to determine whether |
| 36 | # the current step requires LoRA. If LoRA is not needed, the prefill stage |
| 37 | # does not need to launch the triton kernel, which can improve performance |
| 38 | if batch_size == 1 and lora_indices_tensor == -1: |
| 39 | no_lora = True |
| 40 | return ( |
| 41 | b_seq_start_tensor, |
| 42 | seq_length_tensor, |
| 43 | lora_indices_tensor, |
| 44 | batch_size, |
| 45 | max_length, |
| 46 | token_nums, |
| 47 | no_lora, |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | # TODO see if this can be vectorized |
| 52 | def convert_mapping( |
| 53 | mapping: "LoRAMapping", |
| 54 | lora_index_to_id: list[int | None], |
| 55 | max_loras: int, |
| 56 | vocab_size: int, |
| 57 | device: torch.device, |
| 58 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[int]]: |
| 59 | """Converts LoRAMapping to index tensors. |
| 60 | |
| 61 | Args: |
| 62 | mapping: LoRAMapping mapping rows in a batch to LoRA ids. |
| 63 | lora_index_to_id: List mapping LoRA ids to LoRA indices. |
| 64 | max_loras: Maximum number of LoRAs. |
| 65 | vocab_size: Model vocab size. |
| 66 | device: Device the returned index tensors are created on. |
| 67 | |
| 68 | Returns: |
| 69 | A tuple of tensors: |
| 70 | base_indices: Tensor of shape [batch_size] mapping batch rows to |
| 71 | LoRA indices. |
| 72 | sampler_indices: Tensor of shape [batch_size] mapping requests to |
| 73 | LoRA indices for sampler. For generation, this will be the |
| 74 | same as base_indices. For prefill, this will map requests |
| 75 | to LoRA indices. |
| 76 | sampler_indices_padded: Tensor of shape [batch_size] mapping |
| 77 | requests to LoRA indices for sampler with padding. |
| 78 | Same as sampler_indices, but -1 is replaced with |
| 79 | max_loras. |
| 80 | embeddings_indices: Tensor of shape [batch_size] mapping |
| 81 | requests to the row offset of their LoRA.lora_a embeddings. |
| 82 | indices_len: List of lengths of the above tensors. It contains |
| 83 | (base_indices, sampler_indices, sampler_indices_padded, |
| 84 | embeddings_indices). |
| 85 | |
| 86 | """ |
| 87 | index_mapping_indices: list[int] = list(mapping.index_mapping).copy() |
| 88 | embedding_indices = index_mapping_indices.copy() |
| 89 | lora_indices = index_mapping_indices.copy() |
| 90 | |
| 91 | # Build a reverse lookup (LoRA id -> index) once instead of repeatedly |
| 92 | # calling list.index(), which is an O(num_loras) linear scan performed for |
| 93 | # every prompt token and every batch token on each step. |
| 94 | lora_id_to_index = { |
| 95 | lora_id: index |
| 96 | for index, lora_id in enumerate(lora_index_to_id) |
| 97 | if lora_id is not None |
| 98 | } |
| 99 | |
| 100 | prompt_mapping: list[int] = [ |
| 101 | lora_id_to_index[x] if x > 0 else -1 for x in mapping.prompt_mapping |
| 102 | ] |
| 103 | lora_idx = None |
| 104 | for i in range(len(index_mapping_indices)): |
| 105 | lora_idx = ( |
| 106 | lora_id_to_index[index_mapping_indices[i]] |
| 107 | if index_mapping_indices[i] > 0 |
| 108 | else -1 |
| 109 | ) |
| 110 | embedding_indices[i] = lora_idx if index_mapping_indices[i] > 0 else 0 |
| 111 | lora_indices[i] = lora_idx |
| 112 | |
| 113 | indices_list: list[list[int] | torch.Tensor] = [ |
| 114 | index_mapping_indices, |
| 115 | lora_indices, |
| 116 | embedding_indices, |
| 117 | ] |
| 118 | |
| 119 | indices = async_tensor_h2d(indices_list, dtype=torch.long, device=device) |
| 120 | prompt_mapping_tensor = async_tensor_h2d( |
| 121 | prompt_mapping, dtype=torch.long, device=device |
| 122 | ) |
| 123 | embeddings_indices = indices[2] * vocab_size |
| 124 | base_indices = indices[1] |
| 125 | sampler_indices = prompt_mapping_tensor |
| 126 | sampler_indices_padded = sampler_indices.clone() |
| 127 | sampler_indices_padded = torch.where( |
| 128 | sampler_indices_padded == -1, max_loras - 1, sampler_indices_padded |
| 129 | ) |
| 130 | sampler_indices_padded = torch.arange( |
| 131 | 0, len(sampler_indices_padded), device=device, dtype=torch.long |
| 132 | ) + (sampler_indices_padded * len(sampler_indices_padded)) |
| 133 | |
| 134 | # Contain length of indices tensors. Used to index into each tensor. |
| 135 | indices_len = [ |
| 136 | base_indices.shape[-1], |
| 137 | sampler_indices.shape[-1], |
| 138 | sampler_indices_padded.shape[-1], |
| 139 | embeddings_indices.shape[-1], |
| 140 | ] |
| 141 | |
| 142 | return ( |
| 143 | base_indices, |
| 144 | sampler_indices, |
| 145 | sampler_indices_padded, |
| 146 | embeddings_indices, |
| 147 | indices_len, |
| 148 | ) |