@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 294 lines 8.66 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 import torch
10
11 from vllm import envs
12 from vllm.lora.ops.triton_ops.kernel_utils import do_expand_kernel
13 from vllm.lora.ops.triton_ops.utils import (
14 _get_lora_b_ptr,
15 get_lora_op_configs,
16 supports_pdl,
17 )
18 from vllm.triton_utils import tl, triton
19 from vllm.utils.torch_utils import direct_register_custom_op
20
21
22 @triton.jit
23 def _lora_expand_kernel(
24 input_ptr,
25 lora_ptr,
26 out_ptr,
27 M,
28 N,
29 K,
30 token_indices_sorted_by_lora_ids,
31 num_tokens_per_lora,
32 lora_token_start_loc,
33 lora_ids,
34 slice_start_loc,
35 input_d0_stride,
36 input_d1_stride,
37 input_d2_stride, # 1
38 ls_d0_ptr,
39 ls_d1_ptr,
40 ls_d2_ptr, # 1
41 output_d0_stride,
42 output_d1_stride, # 1
43 output_hs_ptr,
44 BLOCK_M: tl.constexpr,
45 BLOCK_N: tl.constexpr,
46 BLOCK_K: tl.constexpr,
47 EVEN_K: tl.constexpr,
48 ADD_INPUTS: tl.constexpr,
49 CAST_TYPE: tl.constexpr,
50 SLICE_NUM: tl.constexpr,
51 SAME_STRIDE: tl.constexpr,
52 USE_GDC: tl.constexpr,
53 launch_pdl: tl.constexpr,
54 ):
55 cta_n_num = tl.cdiv(N, BLOCK_N)
56 cta_m_num = tl.cdiv(M, BLOCK_M)
57
58 pid_mn = tl.program_id(axis=0)
59 pid_m = pid_mn % cta_m_num
60 pid_n = (pid_mn // cta_m_num) % cta_n_num
61
62 slice_id = tl.program_id(axis=1)
63 lora_idx = tl.program_id(axis=2)
64
65 lora_id = tl.load(lora_ids + lora_idx)
66 if lora_id == -1:
67 # Early exit for the no-lora case.
68 return
69
70 lora_m_size = tl.load(num_tokens_per_lora + lora_idx)
71
72 cta_m_offset = pid_m * BLOCK_M
73 if cta_m_offset >= lora_m_size:
74 # Early exit CTA.
75 return
76
77 # When the output dimensions of each slice are the same,cur_n=N, otherwise
78 # cur_n=tl.load(output_hs_ptr + slice_id), this situation exists in GQA's
79 # qkv linear.
80 curr_N = N if SAME_STRIDE else tl.load(output_hs_ptr + slice_id)
81 if pid_n * BLOCK_N >= curr_N:
82 # Early exit CTA.
83 return
84
85 # num rows this CTA should process.
86 cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset)
87
88 # Identify all rows that this CTA should process.
89 lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx)
90 cta_lora_seq_indices = (
91 token_indices_sorted_by_lora_ids + lora_m_indices_start + cta_m_offset
92 )
93
94 # Load all relevant row indices.
95 offset_m = tl.arange(0, BLOCK_M) % cta_m_len
96 ram = tl.load(cta_lora_seq_indices + offset_m)
97
98 do_expand_kernel(
99 pid_n,
100 lora_id,
101 slice_id,
102 input_ptr,
103 lora_ptr,
104 out_ptr,
105 curr_N,
106 K,
107 cta_m_len,
108 ram, # array identifying the rows of Input ptr to operate on
109 slice_start_loc,
110 # input ptr strides
111 input_d0_stride,
112 input_d1_stride,
113 input_d2_stride,
114 # lora ptr strides
115 ls_d0_ptr,
116 ls_d1_ptr,
117 ls_d2_ptr,
118 # out ptr strides
119 output_d0_stride,
120 output_d1_stride,
121 # constants
122 BLOCK_M,
123 BLOCK_N,
124 BLOCK_K,
125 SAME_STRIDE,
126 SLICE_NUM,
127 EVEN_K,
128 CAST_TYPE,
129 ADD_INPUTS,
130 USE_GDC,
131 )
132
133
134 @torch.inference_mode()
135 def _lora_expand(
136 inputs: torch.Tensor, # shape [num_slices, num_tokens, lora_rank]
137 lora_b_weights: list[torch.Tensor], # shape [num_lora, hidden_size, lora_rank]
138 output_tensor: torch.Tensor, # shape [num_tokens, hidden_size * num_slices]
139 token_lora_mapping: torch.Tensor, # shape [num_tokens]
140 token_indices_sorted_by_lora_ids: torch.Tensor, # shape [num_tokens]
141 num_tokens_per_lora: torch.Tensor, # shape [max-loras + 1]
142 lora_token_start_loc: torch.Tensor, # shape [max-loras + 2]
143 lora_ids: torch.Tensor, # shape [max-loras + 1]
144 no_lora_flag_cpu: torch.Tensor, # shape [1]
145 num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs
146 offset_start: int = 0,
147 add_inputs: bool = False,
148 ) -> None:
149 """Args:
150 inputs (torch.Tensor): input tensor
151 lora_b_weights (list[torch.Tensor]): lora'b weight
152 output_tensor (torch.Tensor): output tensor
153 token_lora_mapping (torch.Tensor): A tensor mapping each input token
154 to the lora-id related to that token. A value of -1 indicates that
155 LoRA doesn't apply to that token.
156 token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from
157 the A matrix grouped by LoRA IDs.
158 num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number
159 of tokens that are to be processed by LoRA ID lora_ids[i]
160 lora_token_start_loc (torch.Tensor): A cumulative sum of
161 num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that
162 lora_token_start_loc[i], along with num_tokens_per_lora[i]
163 identifies the region in token_indices_sorted_by_lora_ids that
164 LoRA lora_ids[i] should process.
165 lora_ids (torch.Tensor): LoRA ids to process.
166 no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates
167 if there are any requests that require LoRA.
168 offset_start (int, optional): Offset start for output_tensor.
169 Defaults to 0.
170 add_inputs (bool, optional): Whether to add the input tensor to the
171 output tensor. Defaults to False.
172
173 """
174 assert no_lora_flag_cpu.numel() == 1
175 if no_lora_flag_cpu.item():
176 # None of the inputs require LoRA.
177 return
178
179 assert inputs.dtype in [torch.float16, torch.bfloat16, torch.float32]
180 for weight in lora_b_weights:
181 assert weight.dtype in [torch.float16, torch.bfloat16]
182
183 assert inputs.size(0) == len(lora_b_weights)
184 assert output_tensor.is_contiguous()
185
186 # metadata sanity check.
187 M = inputs.size(1)
188 assert token_lora_mapping.size(0) == M
189 assert token_lora_mapping.size(0) == token_indices_sorted_by_lora_ids.size(0)
190 assert lora_ids.size(0) == num_tokens_per_lora.size(0)
191 assert lora_token_start_loc.size(0) == lora_ids.size(0) + 1
192
193 (
194 slice_start_tensor,
195 lora_ptr_tensor,
196 lora_strides_d0_tensor,
197 lora_strides_d1_tensor,
198 lora_strides_d2_tensor,
199 hidden_sizes_tensor,
200 same_stride,
201 MAX_N,
202 ) = _get_lora_b_ptr(lora_b_weights, offset_start, inputs.device)
203
204 K = lora_b_weights[0].shape[-1] # K= rank
205 ADD_INPUTS = add_inputs
206 MAX_LORAS = lora_ids.size(0)
207 CAST_TYPE = False
208 NUM_SLICES = len(lora_b_weights)
209
210 # Triton kernel configs.
211 kernel_config = get_lora_op_configs(
212 op_type="expand",
213 max_loras=MAX_LORAS,
214 batch=M,
215 hidden_size=MAX_N,
216 rank=K,
217 num_slices=NUM_SLICES,
218 add_inputs=add_inputs,
219 )
220 BLOCK_M = kernel_config["block_m"]
221 BLOCK_N = kernel_config["block_n"]
222 BLOCK_K = kernel_config["block_k"]
223 NUM_WARPS = kernel_config["num_warps"]
224 NUM_CTAS = kernel_config["num_ctas"]
225 NUM_STAGES = kernel_config["num_stages"]
226
227 EVEN_K = K % BLOCK_K == 0 # type: ignore
228
229 if inputs.dtype == torch.float32 and lora_b_weights[0].dtype in [
230 torch.float16,
231 torch.bfloat16,
232 ]:
233 CAST_TYPE = True
234
235 # TODO (varun): This grid formulation maximizes parallelization at the
236 # cost of wasteful thread block launch when only a few input tokens require
237 # LoRA. This might not be the best in all cases.
238 grid = (
239 triton.cdiv(M, BLOCK_M) * triton.cdiv(MAX_N, BLOCK_N),
240 NUM_SLICES,
241 num_active_loras.item(),
242 )
243
244 # PDL only works when dual-stream is being used.
245 use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM
246 _lora_expand_kernel[grid](
247 inputs,
248 lora_ptr_tensor,
249 output_tensor,
250 M,
251 MAX_N,
252 K,
253 token_indices_sorted_by_lora_ids,
254 num_tokens_per_lora,
255 lora_token_start_loc,
256 lora_ids,
257 slice_start_tensor,
258 inputs.stride(0),
259 inputs.stride(1),
260 inputs.stride(2),
261 lora_strides_d0_tensor,
262 lora_strides_d1_tensor,
263 lora_strides_d2_tensor,
264 output_tensor.stride(0),
265 output_tensor.stride(1),
266 hidden_sizes_tensor,
267 BLOCK_M,
268 BLOCK_N,
269 BLOCK_K,
270 EVEN_K,
271 ADD_INPUTS,
272 CAST_TYPE,
273 NUM_SLICES,
274 same_stride,
275 use_gdc,
276 num_warps=NUM_WARPS,
277 num_ctas=NUM_CTAS,
278 num_stages=NUM_STAGES,
279 launch_pdl=use_gdc,
280 )
281
282 return
283
284
285 try:
286 direct_register_custom_op(
287 op_name="lora_expand",
288 op_func=_lora_expand,
289 mutates_args=["output_tensor"],
290 )
291 lora_expand = torch.ops.vllm.lora_expand
292
293 except AttributeError:
294 lora_expand = _lora_expand