@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b69692a2d8
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 22:15 UTC

[Kernel] LoRA - Refactor sgmv kernels (#13110)

Varun Sundar Rabindranath committed Feb 20, 2025 at 17:58 UTC b69692a2d8804ee87efc43425e16f7f323a2f0a9
3 files changed +327 -129
vllm/lora/ops/triton_ops/kernel_utils.py new
+243
@@ -0,0 +1,243 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +Utilities for Punica kernel construction.
4 +"""
5 +import triton
6 +import triton.language as tl
7 +
8 +
9 +@triton.jit
10 +def mm_k(a_ptr, b_ptr, ak_stride, bk_stride, offset_k, K: tl.constexpr,
11 + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
12 + EVEN_K: tl.constexpr, SPLIT_K: tl.constexpr, CAST_TYPE: tl.constexpr,
13 + b_dtype: tl.constexpr):
14 + """
15 + Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of
16 + B (k x n), iterate, through the K dimension to compute the partial/complete
17 + matrix block product.
18 + If SPLIT_K == 1, the output m x n product is complete.
19 + If SPLIT_K > 1, the thread block computes partial outputs. The partial
20 + outputs are then atomically summed in the caller code.
21 + Args:
22 + a_ptr: Array of pointers, identifying rows of A
23 + b_ptr: Array of pointers, identifying columns of B
24 + ak_stride: K dimension stride of the A matrix
25 + bk_stride: K dimension stride of the B matrix
26 + K: Length of the K dimension
27 + BLOCK_M: M dimension of the output block m x n
28 + BLOCK_N: N dimension of the output block m x n
29 + BLOCK_K: K dimension atom
30 + EVEN_K: True if the blocks of A and B can be loaded without any
31 + masking.
32 + SPLIT_K: Parameter signifying parallelism in the K dimension.
33 + CAST_TYPE: if True, cast the values from the A matrix to the B
34 + matrix dtype.
35 + b_dtype: datatype of the B matrix
36 + """
37 + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
38 + for k in range(tl.cdiv(K, BLOCK_K * SPLIT_K)):
39 + if EVEN_K:
40 + tiled_a = tl.load(a_ptr)
41 + tiled_b = tl.load(b_ptr)
42 + else:
43 + tiled_a = tl.load(a_ptr,
44 + mask=offset_k[None, :]
45 + < K - k * (BLOCK_K * SPLIT_K),
46 + other=0)
47 + tiled_b = tl.load(b_ptr,
48 + mask=offset_k[:, None]
49 + < K - k * (BLOCK_K * SPLIT_K),
50 + other=0)
51 + if CAST_TYPE:
52 + tiled_a = tiled_a.to(b_dtype)
53 + accumulator += tl.dot(
54 + tiled_a,
55 + tiled_b,
56 + )
57 + a_ptr += BLOCK_K * SPLIT_K * ak_stride
58 + b_ptr += BLOCK_K * SPLIT_K * bk_stride
59 + return accumulator
60 +
61 +
62 +@triton.jit
63 +def do_expand_kernel(
64 + pid_n,
65 + lora_index,
66 + slice_id,
67 + input_ptr,
68 + lora_ptr,
69 + out_ptr,
70 + N,
71 + K,
72 + M_LEN,
73 + ram, # array identifying the rows of Input ptr to operate on
74 + slice_start_loc,
75 + # input ptr strides
76 + input_d0_stride,
77 + input_d1_stride,
78 + input_d2_stride,
79 + # lora ptr strides
80 + ls_d0_ptr,
81 + ls_d1_ptr,
82 + ls_d2_ptr,
83 + # out ptr strides
84 + output_d0_stride,
85 + output_d1_stride,
86 + # constants
87 + BLOCK_M: tl.constexpr,
88 + BLOCK_N: tl.constexpr,
89 + BLOCK_K: tl.constexpr,
90 + SAME_STRIDE: tl.constexpr,
91 + SLICE_NUM: tl.constexpr,
92 + EVEN_K: tl.constexpr,
93 + CAST_TYPE: tl.constexpr,
94 + ADD_INPUTS: tl.constexpr,
95 +):
96 + """
97 + Given an array of integers that identifies the rows of A, ram,
98 + a lora index that identifies which LoRA to use from lora_ptr, lora_index,
99 + a slice_id that identifies the input/output slice,
100 + compute the matrix product and store in the appropriate output location.
101 + Given that this is an expand kernel, we don't perform any split-K reduction
102 + as the K dimension is assumed to be small.
103 + """
104 +
105 + # ls_d*_ptr can be either an integer or a pointer
106 + if SAME_STRIDE:
107 + # integer
108 + cur_lora_d0_stride = ls_d0_ptr
109 + cur_lora_d1_stride = ls_d1_ptr
110 + cur_lora_d2_stride = ls_d2_ptr
111 + else:
112 + # pointer
113 + cur_lora_d0_stride = tl.load(ls_d0_ptr + slice_id)
114 + cur_lora_d1_stride = tl.load(ls_d1_ptr + slice_id)
115 + cur_lora_d2_stride = tl.load(ls_d2_ptr + slice_id)
116 +
117 + # Identify the input_ptr and lora_ptr from slice_id.
118 + if SLICE_NUM == 1:
119 + cur_input_ptr = input_ptr
120 + cur_lora_ptr = lora_ptr
121 + else:
122 + cur_input_ptr = input_ptr + slice_id * input_d0_stride
123 + cur_lora_ptr = tl.load(lora_ptr + slice_id).to(
124 + tl.pointer_type(out_ptr.dtype.element_ty))
125 +
126 + # Identify the column indices of B to process.
127 + offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
128 + rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N)
129 +
130 + # Identify A and B block pointers
131 + offset_k = tl.arange(0, BLOCK_K)
132 + a_ptr = (cur_input_ptr + ram[:, None] * input_d1_stride +
133 + offset_k[None, :] * input_d2_stride, )
134 + b_ptr = (cur_lora_ptr + cur_lora_d0_stride * lora_index +
135 + offset_k[:, None] * cur_lora_d2_stride +
136 + rbn[None, :] * cur_lora_d1_stride)
137 +
138 + # Compute the block matrix product.
139 + SPLIT_K = 1
140 + accumulator = mm_k(a_ptr, b_ptr, input_d2_stride, cur_lora_d2_stride,
141 + offset_k, K, BLOCK_M, BLOCK_N, BLOCK_K, EVEN_K, SPLIT_K,
142 + CAST_TYPE, cur_lora_ptr.dtype.element_ty)
143 +
144 + tiled_c = accumulator.to(cur_lora_ptr.dtype.element_ty)
145 + if SLICE_NUM == 1:
146 + cur_slice_start = slice_start_loc
147 + else:
148 + cur_slice_start = tl.load(slice_start_loc + slice_id)
149 +
150 + # Identify the C output pointers to store the results of the accumulator.
151 + offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + cur_slice_start
152 + offset_cm = tl.arange(0, BLOCK_M)
153 + c_ptr = (out_ptr + ram[:, None] * output_d0_stride +
154 + offset_cn[None, :] * output_d1_stride)
155 + c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :]
156 + < (cur_slice_start + N))
157 +
158 + if ADD_INPUTS:
159 + tiled_out = tl.load(c_ptr, mask=c_mask)
160 + tiled_c += tiled_out
161 + tl.store(c_ptr, tiled_c, mask=c_mask)
162 +
163 +
164 +@triton.jit
165 +def do_shrink_kernel(
166 + pid_n,
167 + pid_sk,
168 + slice_id,
169 + lora_index,
170 + input_ptr,
171 + lora_ptr,
172 + out_ptr,
173 + N,
174 + K,
175 + M_LEN,
176 + ram,
177 + # input strides
178 + input_d0_stride,
179 + input_d1_stride,
180 + # lora strides
181 + lora_d0_stride,
182 + lora_d1_stride,
183 + lora_d2_stride,
184 + # output strides
185 + output_d0_stride,
186 + output_d1_stride,
187 + output_d2_stride,
188 + scaling,
189 + BLOCK_M: tl.constexpr,
190 + BLOCK_N: tl.constexpr,
191 + BLOCK_K: tl.constexpr,
192 + EVEN_K: tl.constexpr,
193 + SPLIT_K: tl.constexpr,
194 + SLICE_NUM: tl.constexpr,
195 +):
196 + """
197 + Given an array of integers that identifies the rows of A, ram,
198 + a lora index that identifies which LoRA to use from lora_ptr, lora_index,
199 + a slice_id that identifies the input/output slice, compute the
200 + matrix product and store in the appropriate output location.
201 + """
202 +
203 + # Identify the lora_ptr from slice_id.
204 + if SLICE_NUM == 1:
205 + # current lora ptr
206 + cur_lora_ptr = lora_ptr
207 + else:
208 + # current lora ptr
209 + cur_lora_ptr = tl.load(lora_ptr + slice_id).to(
210 + tl.pointer_type(input_ptr.dtype.element_ty))
211 +
212 + # Identify the column indices of B to process.
213 + offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
214 + rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N)
215 +
216 + # Identify A and B block pointers
217 + offset_k = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K)
218 + a_ptr = (input_ptr + ram[:, None] * input_d0_stride +
219 + offset_k[None, :] * input_d1_stride)
220 + b_ptr = (cur_lora_ptr + lora_d0_stride * lora_index +
221 + rbn[None, :] * lora_d1_stride +
222 + offset_k[:, None] * lora_d2_stride)
223 +
224 + # Compute partial/complete block matrix product.
225 + accumulator = mm_k(a_ptr, b_ptr, input_d1_stride, lora_d2_stride, offset_k,
226 + K, BLOCK_M, BLOCK_N, BLOCK_K, EVEN_K, SPLIT_K, False,
227 + cur_lora_ptr.dtype.element_ty)
228 +
229 + # Identify the C output pointers to store the results of the accumulator.
230 + offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
231 + offset_cm = tl.arange(0, BLOCK_M)
232 + cur_out_ptr = (out_ptr if SLICE_NUM == 1 else out_ptr +
233 + slice_id * output_d0_stride)
234 + c_ptr = cur_out_ptr + ram[:, None] * output_d1_stride + offset_cn[
235 + None, :] * output_d2_stride
236 + c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < N)
237 +
238 + accumulator *= scaling
239 + # handles write-back with reduction-splitting
240 + if SPLIT_K == 1:
241 + tl.store(c_ptr, accumulator, mask=c_mask)
242 + else:
243 + tl.atomic_add(c_ptr, accumulator, mask=c_mask)
vllm/lora/ops/triton_ops/sgmv_expand.py
+44 -73
@@ -14,6 +14,7 @@ import triton.language as tl
14
15 from vllm.utils import direct_register_custom_op
16
17 +from .kernel_utils import do_expand_kernel
18 from .utils import _get_lora_b_ptr
19
20
@@ -63,86 +64,56 @@ def _sgmv_expand_kernel(
64 curr_N = N if SAME_STRIDE else tl.load(output_hs_ptr + slice_id)
65 pid_m = pid // cta_n_num
66 pid_n = pid % cta_n_num
67 +
68 M = tl.load(seq_lens + cur_batch)
67 - if pid_m * BLOCK_M > M:
69 + if pid_m * BLOCK_M >= M:
70 return
69 - if pid_n * BLOCK_N > curr_N:
71 + if pid_n * BLOCK_N >= curr_N:
72 return
73 lora_index = tl.load(lora_indices + cur_batch)
74 if lora_index == -1:
75 return
76
75 - cur_seq_start = tl.load(b_seq_start_loc + cur_batch)
76 - offset_m = tl.arange(0, BLOCK_M) + pid_m * BLOCK_M
77 - offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
78 - offset_k = tl.arange(0, BLOCK_K)
79 - ram = tl.max_contiguous(tl.multiple_of(offset_m % M, BLOCK_M), BLOCK_M)
80 - rbn = tl.max_contiguous(tl.multiple_of(offset_n % curr_N, BLOCK_N),
81 - BLOCK_N)
82 - # ls_d*_ptr can be either an integer or a pointer
83 - if SAME_STRIDE:
84 - # integer
85 - cur_lora_d0_stride = ls_d0_ptr
86 - cur_lora_d1_stride = ls_d1_ptr
87 - cur_lora_d2_stride = ls_d2_ptr
88 - else:
89 - # pointer
90 - cur_lora_d0_stride = tl.load(ls_d0_ptr + slice_id)
91 - cur_lora_d1_stride = tl.load(ls_d1_ptr + slice_id)
92 - cur_lora_d2_stride = tl.load(ls_d2_ptr + slice_id)
93 - if SLICE_NUM == 1:
94 - cur_input_ptr = input_ptr
95 - cur_lora_ptr = lora_ptr
96 -
97 - else:
98 - cur_input_ptr = input_ptr + slice_id * input_d0_stride
99 - cur_lora_ptr = tl.load(lora_ptr + slice_id).to(
100 - tl.pointer_type(out_ptr.dtype.element_ty))
101 -
102 - a_ptr = (cur_input_ptr + cur_seq_start * input_d1_stride +
103 - ram[:, None] * input_d1_stride +
104 - offset_k[None, :] * input_d2_stride, )
105 - b_ptr = (cur_lora_ptr + cur_lora_d0_stride * lora_index +
106 - offset_k[:, None] * cur_lora_d2_stride +
107 - rbn[None, :] * cur_lora_d1_stride)
108 - accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
109 - for k in range(tl.cdiv(K, BLOCK_K)):
110 - if EVEN_K:
111 - tiled_a = tl.load(a_ptr)
112 - tiled_b = tl.load(b_ptr)
113 - else:
114 - tiled_a = tl.load(a_ptr,
115 - mask=offset_k[None, :] < K - k * BLOCK_K,
116 - other=0)
117 - tiled_b = tl.load(b_ptr,
118 - mask=offset_k[:, None] < K - k * BLOCK_K,
119 - other=0)
120 - if CAST_TYPE:
121 - tiled_a = tiled_a.to(cur_lora_ptr.dtype.element_ty)
122 - accumulator += tl.dot(
123 - tiled_a,
124 - tiled_b,
125 - )
126 - a_ptr += BLOCK_K * input_d2_stride
127 - b_ptr += BLOCK_K * cur_lora_d2_stride
128 -
129 - tiled_c = accumulator.to(cur_lora_ptr.dtype.element_ty)
130 - if SLICE_NUM == 1:
131 - cur_slice_start = slice_start_loc
132 - else:
133 - cur_slice_start = tl.load(slice_start_loc + slice_id)
134 -
135 - offset_cm = cur_seq_start + tl.arange(0, BLOCK_M) + pid_m * BLOCK_M
136 - offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + cur_slice_start
137 - c_ptr = (out_ptr + offset_cm[:, None] * output_d0_stride +
138 - offset_cn[None, :] * output_d1_stride)
139 - M = tl.load(seq_lens + cur_batch)
140 - c_mask = (offset_cm[:, None] < (cur_seq_start + M)) & (
141 - offset_cn[None, :] < (cur_slice_start + curr_N))
142 - if ADD_INPUTS:
143 - tiled_out = tl.load(c_ptr, mask=c_mask)
144 - tiled_c += tiled_out
145 - tl.store(c_ptr, tiled_c, mask=c_mask)
77 + m_offset = tl.load(b_seq_start_loc + cur_batch)
78 +
79 + cta_m_len = min(BLOCK_M, M - (pid_m * BLOCK_M))
80 + cta_m_offset = m_offset + (pid_m * BLOCK_M)
81 + offset_m = tl.arange(0, BLOCK_M)
82 + ram = cta_m_offset + tl.max_contiguous(
83 + tl.multiple_of(offset_m % cta_m_len, BLOCK_M), BLOCK_M)
84 + do_expand_kernel(
85 + pid_n,
86 + lora_index,
87 + slice_id,
88 + input_ptr,
89 + lora_ptr,
90 + out_ptr,
91 + curr_N,
92 + K,
93 + cta_m_len,
94 + ram, # array identifying the rows of Input ptr to operate on
95 + slice_start_loc,
96 + # input ptr strides
97 + input_d0_stride,
98 + input_d1_stride,
99 + input_d2_stride,
100 + # lora ptr strides
101 + ls_d0_ptr,
102 + ls_d1_ptr,
103 + ls_d2_ptr,
104 + # out ptr strides
105 + output_d0_stride,
106 + output_d1_stride,
107 + # constants
108 + BLOCK_M,
109 + BLOCK_N,
110 + BLOCK_K,
111 + SAME_STRIDE,
112 + SLICE_NUM,
113 + EVEN_K,
114 + CAST_TYPE,
115 + ADD_INPUTS,
116 + )
117
118
119 @torch.inference_mode()
vllm/lora/ops/triton_ops/sgmv_shrink.py
+40 -56
@@ -14,6 +14,7 @@ import triton.language as tl
14
15 from vllm.utils import direct_register_custom_op
16
17 +from .kernel_utils import do_shrink_kernel
18 from .utils import _get_lora_a_ptr
19
20
@@ -62,67 +63,50 @@ def _sgmv_shrink_kernel(
63 pid_sk = pid_mix % SPLIT_K
64
65 M = tl.load(seq_lens + cur_batch)
65 - if pid_m * BLOCK_M > M:
66 + if pid_m * BLOCK_M >= M:
67 return
68 lora_index = tl.load(lora_indices + cur_batch)
69 if lora_index == -1:
70 return
70 - cur_seq_start = tl.load(b_seq_start_loc + cur_batch)
71 - offset_m = tl.arange(0, BLOCK_M) + pid_m * BLOCK_M
72 - offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
73 - offset_k = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K)
74 -
75 - ram = tl.max_contiguous(tl.multiple_of(offset_m % M, BLOCK_M), BLOCK_M)
76 - rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N)
77 - # input ptr
78 - a_ptr = (input_ptr + cur_seq_start * input_d0_stride +
79 - ram[:, None] * input_d0_stride +
80 - offset_k[None, :] * input_d1_stride)
71
82 - if SLICE_NUM == 1:
83 - # current lora ptr
84 - cur_lora_ptr = lora_ptr
85 - else:
86 - # current lora ptr
87 - cur_lora_ptr = tl.load(lora_ptr + slice_id).to(
88 - tl.pointer_type(input_ptr.dtype.element_ty))
89 -
90 - b_ptr = (cur_lora_ptr + lora_d0_stride * lora_index +
91 - rbn[None, :] * lora_d1_stride +
92 - offset_k[:, None] * lora_d2_stride)
93 -
94 - accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
95 - for k in range(0, tl.cdiv(K, BLOCK_K * SPLIT_K)):
96 - if EVEN_K:
97 - tiled_a = tl.load(a_ptr)
98 - tiled_b = tl.load(b_ptr)
99 - else:
100 - k_remaining = K - k * (BLOCK_K * SPLIT_K)
101 - tiled_a = tl.load(a_ptr,
102 - mask=offset_k[None, :] < k_remaining,
103 - other=0.0)
104 - tiled_b = tl.load(b_ptr,
105 - mask=offset_k[:, None] < k_remaining,
106 - other=0.0)
107 - accumulator += tl.dot(tiled_a, tiled_b)
108 -
109 - a_ptr += BLOCK_K * SPLIT_K * input_d1_stride
110 - b_ptr += BLOCK_K * SPLIT_K * lora_d2_stride
111 - offset_cm = cur_seq_start + tl.arange(0, BLOCK_M) + pid_m * BLOCK_M
112 -
113 - offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
114 - cur_out_ptr = (out_ptr if SLICE_NUM == 1 else out_ptr +
115 - slice_id * output_d0_stride)
116 - c_ptr = cur_out_ptr + offset_cm[:, None] * output_d1_stride + offset_cn[
117 - None, :] * output_d2_stride
118 - c_mask = (offset_cm[:, None] < (cur_seq_start + M)) & (offset_cn[None, :]
119 - < N)
120 - accumulator *= scaling
121 - # handles write-back with reduction-splitting
122 - if SPLIT_K == 1:
123 - tl.store(c_ptr, accumulator, mask=c_mask)
124 - else:
125 - tl.atomic_add(c_ptr, accumulator, mask=c_mask)
72 + m_offset = tl.load(b_seq_start_loc + cur_batch)
73 +
74 + cta_m_len = min(BLOCK_M, M - (pid_m * BLOCK_M))
75 + cta_m_offset = m_offset + (pid_m * BLOCK_M)
76 + offset_m = tl.arange(0, BLOCK_M)
77 + ram = cta_m_offset + tl.max_contiguous(
78 + tl.multiple_of(offset_m % cta_m_len, BLOCK_M), BLOCK_M)
79 +
80 + do_shrink_kernel(
81 + pid_n,
82 + pid_sk,
83 + slice_id,
84 + lora_index,
85 + input_ptr,
86 + lora_ptr,
87 + out_ptr,
88 + N,
89 + K,
90 + cta_m_len,
91 + ram,
92 + # input strides
93 + input_d0_stride,
94 + input_d1_stride,
95 + # lora strides
96 + lora_d0_stride,
97 + lora_d1_stride,
98 + lora_d2_stride,
99 + # output strides
100 + output_d0_stride,
101 + output_d1_stride,
102 + output_d2_stride,
103 + scaling,
104 + BLOCK_M,
105 + BLOCK_N,
106 + BLOCK_K,
107 + EVEN_K,
108 + SPLIT_K,
109 + SLICE_NUM)
110
111
112 @torch.inference_mode()