@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 340 lines 10.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Utilities for Punica kernel construction."""
4
5 from vllm.triton_utils import tl, triton
6
7
8 @triton.jit
9 def mm_k(
10 a_ptr: tl.tensor,
11 b_ptr: tl.tensor,
12 ak_stride: int,
13 bk_stride: int,
14 offset_k: tl.tensor,
15 K: tl.constexpr,
16 BLOCK_M: tl.constexpr,
17 BLOCK_N: tl.constexpr,
18 BLOCK_K: tl.constexpr,
19 EVEN_K: tl.constexpr,
20 SPLIT_K: tl.constexpr,
21 CAST_TYPE: tl.constexpr,
22 b_dtype: tl.constexpr,
23 USE_GDC: tl.constexpr,
24 base_k: int,
25 ):
26 """Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of
27 B (k x n), iterate, through the K dimension to compute the partial/complete
28 matrix block product.
29 If SPLIT_K == 1, the output m x n product is complete.
30 If SPLIT_K > 1, the thread block computes partial outputs. The partial
31 outputs are then atomically summed in the caller code.
32
33 Args:
34 a_ptr: Array of pointers, identifying rows of A
35 b_ptr: Array of pointers, identifying columns of B
36 ak_stride: K dimension stride of the A matrix
37 bk_stride: K dimension stride of the B matrix
38 offset_k: Offsets within the current BLOCK_K tile along K
39 K: Length of the K dimension
40 BLOCK_M: M dimension of the output block m x n
41 BLOCK_N: N dimension of the output block m x n
42 BLOCK_K: K dimension atom
43 EVEN_K: True if the blocks of A and B can be loaded without any
44 masking.
45 SPLIT_K: Parameter signifying parallelism in the K dimension.
46 CAST_TYPE: if True, cast the values from the A matrix to the B
47 matrix dtype.
48 b_dtype: datatype of the B matrix
49 USE_GDC: Whether to use PDL. True indicates use.
50 base_k: Base offset along K dimension for current SPLIT_K group
51
52 """
53 accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
54
55 # Step size along K for each iteration
56 STEP_K = BLOCK_K * SPLIT_K
57
58 # Total number of iterations (compile-time constant)
59 num_iters = tl.cdiv(K, STEP_K)
60
61 for k in range(num_iters):
62 # Current iteration's global K offset
63 iter_k = k * STEP_K + base_k
64
65 # Check if this iteration is completely valid (no masking needed)
66 block_end = iter_k + BLOCK_K
67
68 if EVEN_K:
69 # K is divisible by BLOCK_K, no masking ever needed
70 # pre-fetch lora weight
71 tiled_b = tl.load(b_ptr)
72 if USE_GDC:
73 tl.extra.cuda.gdc_wait()
74 tiled_a = tl.load(a_ptr)
75 if CAST_TYPE:
76 tiled_a = tiled_a.to(b_dtype)
77 accumulator += tl.dot(tiled_a, tiled_b)
78 else:
79 # Check if we need element-wise masking
80 if iter_k >= K:
81 # Entire block out of range, skip
82 pass
83 elif block_end <= K:
84 # Entire block in range, no masking needed (fast path)
85 tiled_b = tl.load(b_ptr)
86 if USE_GDC:
87 tl.extra.cuda.gdc_wait()
88 tiled_a = tl.load(a_ptr)
89 if CAST_TYPE:
90 tiled_a = tiled_a.to(b_dtype)
91 accumulator += tl.dot(tiled_a, tiled_b)
92 else:
93 # Partial block, need masking (only last iteration)
94 k_offsets = tl.arange(0, BLOCK_K)
95 mask = iter_k + k_offsets < K
96 tiled_b = tl.load(b_ptr, mask=mask[:, None], other=0.0)
97 if USE_GDC:
98 tl.extra.cuda.gdc_wait()
99 tiled_a = tl.load(a_ptr, mask=mask[None, :], other=0.0)
100 if CAST_TYPE:
101 tiled_a = tiled_a.to(b_dtype)
102 accumulator += tl.dot(tiled_a, tiled_b)
103
104 a_ptr += STEP_K * ak_stride
105 b_ptr += STEP_K * bk_stride
106
107 return accumulator
108
109
110 @triton.jit
111 def do_expand_kernel(
112 pid_n,
113 lora_index,
114 slice_id,
115 input_ptr,
116 lora_ptr,
117 out_ptr,
118 N,
119 K,
120 M_LEN,
121 ram, # array identifying the rows of Input ptr to operate on
122 slice_start_loc,
123 # input ptr strides
124 input_d0_stride,
125 input_d1_stride,
126 input_d2_stride,
127 # lora ptr strides
128 ls_d0_ptr,
129 ls_d1_ptr,
130 ls_d2_ptr,
131 # out ptr strides
132 output_d0_stride,
133 output_d1_stride,
134 # constants
135 BLOCK_M: tl.constexpr,
136 BLOCK_N: tl.constexpr,
137 BLOCK_K: tl.constexpr,
138 SAME_STRIDE: tl.constexpr,
139 SLICE_NUM: tl.constexpr,
140 EVEN_K: tl.constexpr,
141 CAST_TYPE: tl.constexpr,
142 ADD_INPUTS: tl.constexpr,
143 USE_GDC: tl.constexpr,
144 ):
145 """Given an array of integers that identifies the rows of A, ram,
146 a lora index that identifies which LoRA to use from lora_ptr, lora_index,
147 a slice_id that identifies the input/output slice,
148 compute the matrix product and store in the appropriate output location.
149 Given that this is an expand kernel, we don't perform any split-K reduction
150 as the K dimension is assumed to be small.
151 """
152 # ls_d*_ptr can be either an integer or a pointer
153 if SAME_STRIDE:
154 # integer
155 cur_lora_d0_stride = ls_d0_ptr
156 cur_lora_d1_stride = ls_d1_ptr
157 cur_lora_d2_stride = ls_d2_ptr
158 else:
159 # pointer
160 cur_lora_d0_stride = tl.load(ls_d0_ptr + slice_id)
161 cur_lora_d1_stride = tl.load(ls_d1_ptr + slice_id)
162 cur_lora_d2_stride = tl.load(ls_d2_ptr + slice_id)
163
164 # Identify the input_ptr and lora_ptr from slice_id.
165 if SLICE_NUM == 1:
166 cur_input_ptr = input_ptr
167 cur_lora_ptr = lora_ptr
168 else:
169 cur_input_ptr = input_ptr + slice_id * input_d0_stride
170 cur_lora_ptr = tl.load(lora_ptr + slice_id).to(
171 tl.pointer_type(out_ptr.dtype.element_ty)
172 )
173
174 # Identify the column indices of B to process.
175 offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
176 rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N)
177
178 # Identify A and B block pointers
179 offset_k = tl.arange(0, BLOCK_K)
180 a_ptr = (
181 cur_input_ptr
182 + ram[:, None] * input_d1_stride
183 + offset_k[None, :] * input_d2_stride
184 )
185 b_ptr = (
186 cur_lora_ptr
187 + cur_lora_d0_stride * lora_index
188 + offset_k[:, None] * cur_lora_d2_stride
189 + rbn[None, :] * cur_lora_d1_stride
190 )
191
192 # Compute the block matrix product.
193 SPLIT_K = 1
194
195 accumulator = mm_k(
196 a_ptr,
197 b_ptr,
198 input_d2_stride,
199 cur_lora_d2_stride,
200 offset_k,
201 K,
202 BLOCK_M,
203 BLOCK_N,
204 BLOCK_K,
205 EVEN_K,
206 SPLIT_K,
207 CAST_TYPE,
208 cur_lora_ptr.dtype.element_ty,
209 USE_GDC,
210 base_k=0,
211 )
212
213 tiled_c = accumulator.to(cur_lora_ptr.dtype.element_ty)
214 if SLICE_NUM == 1:
215 cur_slice_start = slice_start_loc
216 else:
217 cur_slice_start = tl.load(slice_start_loc + slice_id)
218
219 # Identify the C output pointers to store the results of the accumulator.
220 offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + cur_slice_start
221 offset_cm = tl.arange(0, BLOCK_M)
222 c_ptr = (
223 out_ptr
224 # int64 to keep the row offset from overflowing at long context lengths
225 + ram[:, None].to(tl.int64) * output_d0_stride
226 + offset_cn[None, :] * output_d1_stride
227 )
228 c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < (cur_slice_start + N))
229
230 if ADD_INPUTS:
231 tiled_out = tl.load(c_ptr, mask=c_mask)
232 tiled_c += tiled_out
233 tl.store(c_ptr, tiled_c, mask=c_mask)
234
235
236 @triton.jit
237 def do_shrink_kernel(
238 pid_n,
239 pid_sk,
240 slice_id,
241 lora_index,
242 input_ptr,
243 lora_ptr,
244 out_ptr,
245 N,
246 K,
247 M_LEN,
248 ram,
249 # input strides
250 input_d0_stride,
251 input_d1_stride,
252 # lora strides
253 lora_d0_stride,
254 lora_d1_stride,
255 lora_d2_stride,
256 # output strides
257 output_d0_stride,
258 output_d1_stride,
259 output_d2_stride,
260 scaling,
261 BLOCK_M: tl.constexpr,
262 BLOCK_N: tl.constexpr,
263 BLOCK_K: tl.constexpr,
264 EVEN_K: tl.constexpr,
265 SPLIT_K: tl.constexpr,
266 SLICE_NUM: tl.constexpr,
267 USE_GDC: tl.constexpr,
268 ):
269 """Given an array of integers that identifies the rows of A, ram,
270 a lora index that identifies which LoRA to use from lora_ptr, lora_index,
271 a slice_id that identifies the input/output slice, compute the
272 matrix product and store in the appropriate output location.
273 """
274 # Identify the lora_ptr from slice_id.
275 if SLICE_NUM == 1:
276 # current lora ptr
277 cur_lora_ptr = lora_ptr
278 else:
279 # current lora ptr
280 cur_lora_ptr = tl.load(lora_ptr + slice_id).to(
281 tl.pointer_type(input_ptr.dtype.element_ty)
282 )
283
284 # Identify the column indices of B to process.
285 offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
286 rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N)
287
288 # Identify A and B block pointers
289 offset_k = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K)
290 a_ptr = (
291 input_ptr
292 # int64 to keep the row offset from overflowing at long context lengths
293 + ram[:, None].to(tl.int64) * input_d0_stride
294 + offset_k[None, :] * input_d1_stride
295 )
296 b_ptr = (
297 cur_lora_ptr
298 + lora_d0_stride * lora_index
299 + rbn[None, :] * lora_d1_stride
300 + offset_k[:, None] * lora_d2_stride
301 )
302
303 # Compute partial/complete block matrix product.
304 accumulator = mm_k(
305 a_ptr,
306 b_ptr,
307 input_d1_stride,
308 lora_d2_stride,
309 offset_k,
310 K,
311 BLOCK_M,
312 BLOCK_N,
313 BLOCK_K,
314 EVEN_K,
315 SPLIT_K,
316 False,
317 cur_lora_ptr.dtype.element_ty,
318 False, # USE_GDC is always False in shrink kernel
319 base_k=pid_sk * BLOCK_K,
320 )
321 # GDC launch dependents hints the runtime system to launch dependent kernels.
322 if USE_GDC:
323 tl.extra.cuda.gdc_launch_dependents()
324 # Identify the C output pointers to store the results of the accumulator.
325 offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
326 offset_cm = tl.arange(0, BLOCK_M)
327 cur_out_ptr = out_ptr if SLICE_NUM == 1 else out_ptr + slice_id * output_d0_stride
328 c_ptr = (
329 cur_out_ptr
330 + ram[:, None] * output_d1_stride
331 + offset_cn[None, :] * output_d2_stride
332 )
333 c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < N)
334 accumulator *= scaling
335
336 # handles write-back with reduction-splitting
337 if SPLIT_K == 1:
338 tl.store(c_ptr, accumulator, mask=c_mask)
339 else:
340 tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed")