@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 275 lines 8.38 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_shrink_kernel
13 from vllm.lora.ops.triton_ops.utils import (
14 _get_lora_a_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_shrink_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 scaling,
35 input_d0_stride,
36 input_d1_stride,
37 lora_d0_stride,
38 lora_d1_stride,
39 lora_d2_stride,
40 output_d0_stride,
41 output_d1_stride,
42 output_d2_stride,
43 BLOCK_M: tl.constexpr,
44 BLOCK_N: tl.constexpr,
45 BLOCK_K: tl.constexpr,
46 EVEN_K: tl.constexpr,
47 SPLIT_K: tl.constexpr,
48 GROUP_SIZE_M: tl.constexpr,
49 SLICE_NUM: tl.constexpr,
50 USE_GDC: tl.constexpr,
51 launch_pdl: tl.constexpr,
52 ):
53 cta_n_num = tl.cdiv(N, BLOCK_N)
54 cta_m_num = tl.cdiv(M, BLOCK_M)
55
56 pid_sk_m_n = tl.program_id(axis=0)
57 pid_sk = pid_sk_m_n % SPLIT_K
58
59 pid_m_n = pid_sk_m_n // SPLIT_K
60 num_pid_in_group = GROUP_SIZE_M * cta_n_num
61 group_id = pid_m_n // num_pid_in_group
62 first_pid_m = group_id * GROUP_SIZE_M
63 group_size_m = min(cta_m_num - first_pid_m, GROUP_SIZE_M)
64
65 # Column-major ordering within groups for better cache reuse
66 pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)
67 pid_n = (pid_m_n % num_pid_in_group) // group_size_m
68
69 slice_id = tl.program_id(axis=1)
70 lora_idx = tl.program_id(axis=2)
71
72 lora_id = tl.load(lora_ids + lora_idx)
73 if lora_id == -1:
74 # Early exit for the no-lora case.
75 return
76
77 lora_m_size = tl.load(num_tokens_per_lora + lora_idx)
78
79 cta_m_offset = pid_m * BLOCK_M
80 if cta_m_offset >= lora_m_size:
81 # Early exit CTA.
82 return
83
84 # num rows this CTA should process.
85 cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset)
86
87 # Identify all rows that this CTA should process.
88 lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx)
89 cta_lora_seq_indices = (
90 token_indices_sorted_by_lora_ids + lora_m_indices_start + cta_m_offset
91 )
92 # Load all relevant row indices.
93 offset_m = tl.arange(0, BLOCK_M) % cta_m_len
94 ram = tl.load(cta_lora_seq_indices + offset_m)
95
96 do_shrink_kernel(
97 pid_n,
98 pid_sk,
99 slice_id,
100 lora_id,
101 input_ptr,
102 lora_ptr,
103 out_ptr,
104 N,
105 K,
106 cta_m_len,
107 ram, # array identifying the rows of Input ptr to operate on
108 # input strides
109 input_d0_stride,
110 input_d1_stride,
111 # lora strides
112 lora_d0_stride,
113 lora_d1_stride,
114 lora_d2_stride,
115 # output strides
116 output_d0_stride,
117 output_d1_stride,
118 output_d2_stride,
119 scaling,
120 BLOCK_M,
121 BLOCK_N,
122 BLOCK_K,
123 EVEN_K,
124 SPLIT_K,
125 SLICE_NUM,
126 USE_GDC,
127 )
128
129
130 @torch.inference_mode()
131 def _lora_shrink(
132 inputs: torch.Tensor, # shape [num_tokens, hidden_size]
133 lora_a_weights: list[torch.Tensor], # shape [num_loras, lora_rank, hidden_size]
134 output_tensor: torch.Tensor, # shape [num_slices, num_tokens, lora_rank]
135 token_lora_mapping: torch.Tensor, # shape [num_tokens]
136 token_indices_sorted_by_lora_ids: torch.Tensor, # shape [num_tokens]
137 num_tokens_per_lora: torch.Tensor, # shape [max-loras + 1]
138 lora_token_start_loc: torch.Tensor, # shape [max-loras + 2]
139 lora_ids: torch.Tensor, # shape [max-loras + 1]
140 no_lora_flag_cpu: torch.Tensor, # shape [1]
141 num_active_loras: torch.Tensor, # CPU tensor [1], number of active LoRAs
142 scaling: float,
143 ) -> None:
144 """Args:
145 inputs (torch.Tensor): Input tensor
146 lora_a_weights (list[torch.Tensor]): LoRA weights
147 output_tensor (torch.Tensor): output tensor
148 token_lora_mapping (torch.Tensor): A tensor mapping each input token
149 to the lora-id related to that token. A value of -1 indicates that
150 LoRA doesn't apply to that token.
151 token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from
152 the A matrix grouped by LoRA IDs.
153 num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number
154 of tokens that are to be processed by LoRA ID lora_ids[i]
155 lora_token_start_loc (torch.Tensor): A cumulative sum of
156 num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that
157 lora_token_start_loc[i], along with num_tokens_per_lora[i]
158 identifies the region in token_indices_sorted_by_lora_ids that
159 LoRA lora_ids[i] should process.
160 lora_ids (torch.Tensor): LoRA ids to process.
161 no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates
162 if there are any requests that require LoRA.
163 num_active_loras (torch.Tensor): A CPU tensor of size 1, containing the
164 number of active LoRAs. Stored as a tensor (not int) so
165 torch.compile treats it as dynamic rather than a constant.
166 scaling (float): Scaling factor.
167
168 """
169 assert no_lora_flag_cpu.numel() == 1
170 if no_lora_flag_cpu.item():
171 # None of the inputs require LoRA.
172 return
173
174 assert inputs.dtype == lora_a_weights[0].dtype
175 assert inputs.dtype in [torch.float16, torch.bfloat16]
176 for weight in lora_a_weights:
177 assert weight.dtype in [torch.float16, torch.bfloat16]
178
179 assert inputs.size(1) == lora_a_weights[0].size(-1)
180 inputs = inputs.contiguous()
181 assert output_tensor.is_contiguous()
182
183 # metadata sanity check
184 M = inputs.size(0)
185 assert token_lora_mapping.size(0) == M
186 assert token_lora_mapping.size(0) == token_indices_sorted_by_lora_ids.size(0)
187 assert lora_ids.size(0) == num_tokens_per_lora.size(0)
188 assert lora_token_start_loc.size(0) == lora_ids.size(0) + 1
189
190 output_tensor.zero_()
191
192 (lora_ptr_tensor, lora_strides_d0, lora_strides_d1, lora_strides_d2) = (
193 _get_lora_a_ptr(lora_a_weights, inputs.device)
194 )
195 N, K = lora_a_weights[0].shape[-2:] # K=hidden_size,N=rank
196 NUM_SLICES = len(lora_a_weights)
197 MAX_LORAS = lora_ids.size(0)
198
199 # Triton kernel configs
200 kernel_config = get_lora_op_configs(
201 "shrink",
202 max_loras=MAX_LORAS,
203 batch=M,
204 hidden_size=K,
205 rank=N,
206 num_slices=NUM_SLICES,
207 )
208 BLOCK_M = kernel_config["block_m"]
209 BLOCK_N = kernel_config["block_n"]
210 BLOCK_K = kernel_config["block_k"]
211 SPLIT_K = kernel_config["split_k"]
212 NUM_WARPS = kernel_config["num_warps"]
213 NUM_STAGES = kernel_config["num_stages"]
214 NUM_CTAS = kernel_config["num_ctas"]
215 GROUP_SIZE_M = kernel_config.get("group_size_m", 8)
216 EVEN_K = K % (BLOCK_K * SPLIT_K) == 0 # type: ignore
217
218 # TODO (varun): This grid formulation maximizes parallelization at the
219 # cost of wasteful thread block launch when only few of the input tokens
220 # require LoRA. This might not be the best in all cases.
221 grid = (
222 SPLIT_K * triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),
223 NUM_SLICES,
224 num_active_loras.item(),
225 )
226
227 # PDL only works when dual-stream is being used.
228 use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM
229 _lora_shrink_kernel[grid](
230 inputs,
231 lora_ptr_tensor,
232 output_tensor,
233 M,
234 N,
235 K,
236 token_indices_sorted_by_lora_ids,
237 num_tokens_per_lora,
238 lora_token_start_loc,
239 lora_ids,
240 scaling,
241 inputs.stride(0),
242 inputs.stride(1),
243 lora_strides_d0,
244 lora_strides_d1,
245 lora_strides_d2,
246 output_tensor.stride(0),
247 output_tensor.stride(1),
248 output_tensor.stride(2),
249 BLOCK_M,
250 BLOCK_N,
251 BLOCK_K,
252 EVEN_K,
253 SPLIT_K,
254 GROUP_SIZE_M,
255 NUM_SLICES,
256 use_gdc,
257 num_warps=NUM_WARPS,
258 num_ctas=NUM_CTAS,
259 num_stages=NUM_STAGES,
260 launch_pdl=use_gdc,
261 )
262
263 return
264
265
266 try:
267 direct_register_custom_op(
268 op_name="lora_shrink",
269 op_func=_lora_shrink,
270 mutates_args=["output_tensor"],
271 )
272 lora_shrink = torch.ops.vllm.lora_shrink
273
274 except AttributeError:
275 lora_shrink = _lora_shrink