@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 337 lines 11.3 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5
6 from vllm.triton_utils import tl, triton
7
8 AWQ_TRITON_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
9
10
11 @triton.jit
12 def awq_dequantize_kernel(
13 qweight_ptr, # quantized matrix
14 scales_ptr, # scales, per group
15 zeros_ptr, # zeros, per group
16 group_size, # Should always be one of the supported group sizes
17 result_ptr, # Output matrix
18 num_cols, # input num cols in qweight
19 num_rows, # input num rows in qweight
20 BLOCK_SIZE_X: tl.constexpr,
21 BLOCK_SIZE_Y: tl.constexpr,
22 ):
23 # Set up the pids.
24 pid_x = tl.program_id(axis=0)
25 pid_y = tl.program_id(axis=1)
26
27 # Compute offsets and masks for qweight_ptr.
28 offsets_y = pid_y * BLOCK_SIZE_Y + tl.arange(0, BLOCK_SIZE_Y)
29 offsets_x = pid_x * BLOCK_SIZE_X + tl.arange(0, BLOCK_SIZE_X)
30 offsets = num_cols * offsets_y[:, None] + offsets_x[None, :]
31
32 masks_y = offsets_y < num_rows
33 masks_x = offsets_x < num_cols
34
35 masks = masks_y[:, None] & masks_x[None, :]
36
37 # Compute offsets and masks for result output ptr.
38 result_offsets_y = pid_y * BLOCK_SIZE_Y + tl.arange(0, BLOCK_SIZE_Y)
39 result_offsets_x = pid_x * BLOCK_SIZE_X * 8 + tl.arange(0, BLOCK_SIZE_X * 8)
40 result_offsets = (
41 8 * num_cols * result_offsets_y[:, None] + result_offsets_x[None, :]
42 )
43
44 result_masks_y = result_offsets_y < num_rows
45 result_masks_x = result_offsets_x < num_cols * 8
46 result_masks = result_masks_y[:, None] & result_masks_x[None, :]
47
48 # Load the weights.
49 iweights = tl.load(qweight_ptr + offsets, masks, 0.0)
50 iweights = tl.interleave(iweights, iweights)
51 iweights = tl.interleave(iweights, iweights)
52 iweights = tl.interleave(iweights, iweights)
53
54 # Create reverse AWQ order as tensor: [0, 4, 1, 5, 2, 6, 3, 7]
55 # that will map given indices to the correct order.
56 reverse_awq_order_tensor = (
57 (tl.arange(0, 2) * 4)[None, :] + tl.arange(0, 4)[:, None]
58 ).reshape(8)
59
60 # Use this to compute a set of shifts that can be used to unpack and
61 # reorder the values in iweights and zeros.
62 shifts = reverse_awq_order_tensor * 4
63 shifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_Y * BLOCK_SIZE_X, 8))
64 shifts = tl.reshape(shifts, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8))
65
66 # Unpack and reorder: shift out the correct 4-bit value and mask.
67 iweights = (iweights >> shifts) & 0xF
68
69 # Compute zero offsets and masks.
70 zero_offsets_y = pid_y * BLOCK_SIZE_Y // group_size + tl.arange(0, 1)
71 zero_offsets_x = pid_x * BLOCK_SIZE_X + tl.arange(0, BLOCK_SIZE_X)
72 zero_offsets = num_cols * zero_offsets_y[:, None] + zero_offsets_x[None, :]
73
74 zero_masks_y = zero_offsets_y < num_rows // group_size
75 zero_masks_x = zero_offsets_x < num_cols
76 zero_masks = zero_masks_y[:, None] & zero_masks_x[None, :]
77
78 # Load the zeros.
79 zeros = tl.load(zeros_ptr + zero_offsets, zero_masks, 0.0)
80 zeros = tl.interleave(zeros, zeros)
81 zeros = tl.interleave(zeros, zeros)
82 zeros = tl.interleave(zeros, zeros)
83 zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8))
84
85 # Unpack and reorder: shift out the correct 4-bit value and mask.
86 zeros = (zeros >> shifts) & 0xF
87
88 # Compute scale offsets and masks.
89 scale_offsets_y = pid_y * BLOCK_SIZE_Y // group_size + tl.arange(0, 1)
90 scale_offsets_x = pid_x * BLOCK_SIZE_X * 8 + tl.arange(0, BLOCK_SIZE_X * 8)
91 scale_offsets = num_cols * 8 * scale_offsets_y[:, None] + scale_offsets_x[None, :]
92 scale_masks_y = scale_offsets_y < num_rows // group_size
93 scale_masks_x = scale_offsets_x < num_cols * 8
94 scale_masks = scale_masks_y[:, None] & scale_masks_x[None, :]
95
96 # Load the scales.
97 scales = tl.load(scales_ptr + scale_offsets, scale_masks, 0.0)
98 scales = tl.broadcast_to(scales, (BLOCK_SIZE_Y, BLOCK_SIZE_X * 8))
99
100 # Dequantize.
101 iweights = (iweights - zeros) * scales
102 iweights = iweights.to(result_ptr.type.element_ty)
103
104 # Finally, store.
105 tl.store(result_ptr + result_offsets, iweights, result_masks)
106
107
108 @triton.jit
109 def awq_gemm_kernel(
110 a_ptr,
111 b_ptr,
112 c_ptr,
113 zeros_ptr,
114 scales_ptr,
115 M,
116 N,
117 K,
118 group_size,
119 BLOCK_SIZE_M: tl.constexpr,
120 BLOCK_SIZE_N: tl.constexpr,
121 BLOCK_SIZE_K: tl.constexpr,
122 SPLIT_K: tl.constexpr,
123 ):
124 pid = tl.program_id(axis=0)
125 pid_z = tl.program_id(1)
126
127 # NOTE: This doesn't work in TRITON_INTERPRET=1 mode. Use below instead.
128 # num_pid_n = (N + BLOCK_SIZE_N - 1) // BLOCK_SIZE_N
129 num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
130
131 pid_m = pid // num_pid_n
132 pid_n = pid % num_pid_n
133
134 accumulator_dtype = c_ptr.type.element_ty
135
136 # NOTE: This doesn't work in TRITON_INTERPRET=1 mode. Use below instead.
137 # accumulator = tl.arange(0, BLOCK_SIZE_N)
138 # accumulator = tl.broadcast_to(accumulator[None, :],
139 # (BLOCK_SIZE_M, BLOCK_SIZE_N))
140 # accumulator = accumulator & 0x0
141 # accumulator = accumulator.to(accumulator_dtype)
142 accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype)
143
144 # Create reverse AWQ order as tensor: [0, 4, 1, 5, 2, 6, 3, 7]
145 # that will map given indices to the correct order.
146 reverse_awq_order_tensor = (
147 (tl.arange(0, 2) * 4)[None, :] + tl.arange(0, 4)[:, None]
148 ).reshape(8)
149
150 # Create the necessary shifts to use to unpack.
151 shifts = reverse_awq_order_tensor * 4
152 shifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_K * (BLOCK_SIZE_N // 8), 8))
153 shifts = tl.reshape(shifts, (BLOCK_SIZE_K, BLOCK_SIZE_N))
154
155 # Offsets and masks.
156 offsets_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
157 masks_am = offsets_am < M
158
159 offsets_bn = pid_n * (BLOCK_SIZE_N // 8) + tl.arange(0, BLOCK_SIZE_N // 8)
160 masks_bn = offsets_bn < N // 8
161
162 offsets_zn = pid_n * (BLOCK_SIZE_N // 8) + tl.arange(0, BLOCK_SIZE_N // 8)
163 masks_zn = offsets_zn < N // 8
164
165 offsets_sn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
166 masks_sn = offsets_sn < N
167
168 offsets_k = pid_z * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
169 offsets_a = K * offsets_am[:, None] + offsets_k[None, :]
170 offsets_b = (N // 8) * offsets_k[:, None] + offsets_bn[None, :]
171
172 a_ptrs = a_ptr + offsets_a
173 b_ptrs = b_ptr + offsets_b
174
175 # NOTE: Use this in TRITON_INTERPRET=1 mode instead of tl.cdiv
176 # block_offset = BLOCK_SIZE_K * SPLIT_K
177 # for k in range(0, (K + block_offset - 1) // (block_offset)):
178 for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)):
179 masks_k = offsets_k < K
180 masks_a = masks_am[:, None] & masks_k[None, :]
181 a = tl.load(a_ptrs, mask=masks_a, other=0.0)
182
183 masks_b = masks_k[:, None] & masks_bn[None, :]
184 b = tl.load(b_ptrs, mask=masks_b, other=0.0)
185 b = tl.interleave(b, b)
186 b = tl.interleave(b, b)
187 b = tl.interleave(b, b)
188
189 # Dequantize b.
190 offsets_szk = (
191 BLOCK_SIZE_K * SPLIT_K * k + pid_z * BLOCK_SIZE_K
192 ) // group_size + tl.arange(0, 1)
193 offsets_z = (N // 8) * offsets_szk[:, None] + offsets_zn[None, :]
194 masks_zk = offsets_szk < K // group_size
195 masks_z = masks_zk[:, None] & masks_zn[None, :]
196 zeros_ptrs = zeros_ptr + offsets_z
197 zeros = tl.load(zeros_ptrs, mask=masks_z, other=0.0)
198 zeros = tl.interleave(zeros, zeros)
199 zeros = tl.interleave(zeros, zeros)
200 zeros = tl.interleave(zeros, zeros)
201 zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N))
202
203 offsets_s = N * offsets_szk[:, None] + offsets_sn[None, :]
204 masks_sk = offsets_szk < K // group_size
205 masks_s = masks_sk[:, None] & masks_sn[None, :]
206 scales_ptrs = scales_ptr + offsets_s
207 scales = tl.load(scales_ptrs, mask=masks_s, other=0.0)
208 scales = tl.broadcast_to(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N))
209
210 b = (b >> shifts) & 0xF
211 zeros = (zeros >> shifts) & 0xF
212 b = (b - zeros) * scales
213 b = b.to(c_ptr.type.element_ty)
214
215 # Accumulate results.
216 accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
217
218 offsets_k += BLOCK_SIZE_K * SPLIT_K
219 a_ptrs += BLOCK_SIZE_K * SPLIT_K
220 b_ptrs += BLOCK_SIZE_K * SPLIT_K * (N // 8)
221
222 c = accumulator.to(c_ptr.type.element_ty)
223 offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
224 offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
225 c_ptrs = c_ptr + pid_z * N * M + N * offs_cm[:, None] + offs_cn[None, :]
226 c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
227 tl.store(c_ptrs, c, mask=c_mask)
228
229
230 # qweights - [K , M // 8], int32
231 # scales - [K // G, M ], float16
232 # zeros - [K // G, M // 8], int32
233 def awq_dequantize_triton(
234 qweight: torch.Tensor,
235 scales: torch.Tensor,
236 zeros: torch.Tensor,
237 block_size_x: int = 32,
238 block_size_y: int = 32,
239 ) -> torch.Tensor:
240 K = qweight.shape[0]
241 M = scales.shape[1]
242 group_size = qweight.shape[0] // scales.shape[0]
243
244 assert K > 0 and M > 0
245 assert scales.shape[0] == K // group_size and scales.shape[1] == M
246 assert zeros.shape[0] == K // group_size and zeros.shape[1] == M // 8
247 assert group_size <= K
248 assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K
249
250 # Result tensor:
251 # number of rows = same as input tensor
252 # number of cols = 8 x input tensor num cols
253 result = torch.empty(
254 qweight.shape[0],
255 qweight.shape[1] * 8,
256 device=qweight.device,
257 dtype=scales.dtype,
258 )
259
260 Y = qweight.shape[0] # num rows
261 X = qweight.shape[1] # num cols
262
263 grid = lambda META: (
264 triton.cdiv(X, META["BLOCK_SIZE_X"]),
265 triton.cdiv(Y, META["BLOCK_SIZE_Y"]),
266 )
267 awq_dequantize_kernel[grid](
268 qweight,
269 scales,
270 zeros,
271 group_size,
272 result,
273 X,
274 Y,
275 BLOCK_SIZE_X=block_size_x,
276 BLOCK_SIZE_Y=block_size_y,
277 )
278
279 return result
280
281
282 # input - [M, K]
283 # qweight - [K, N // 8]
284 # qzeros - [K // G, N // 8]
285 # scales - [K // G, N]
286 # split_k_iters - parallelism along K-dimension, int, power of 2.
287 def awq_gemm_triton(
288 input: torch.Tensor,
289 qweight: torch.Tensor,
290 scales: torch.Tensor,
291 qzeros: torch.Tensor,
292 split_k_iters: int,
293 block_size_m: int = 32,
294 block_size_n: int = 32,
295 block_size_k: int = 32,
296 ) -> torch.Tensor:
297 M, K = input.shape
298 N = qweight.shape[1] * 8
299 group_size = qweight.shape[0] // qzeros.shape[0]
300
301 assert N > 0 and K > 0 and M > 0
302 assert qweight.shape[0] == K and qweight.shape[1] == N // 8
303 assert qzeros.shape[0] == K // group_size and qzeros.shape[1] == N // 8
304 assert scales.shape[0] == K // group_size and scales.shape[1] == N
305 assert split_k_iters & (split_k_iters - 1) == 0 and split_k_iters != 0
306 assert split_k_iters <= 32
307 assert group_size <= K
308 assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K
309
310 grid = lambda META: (
311 triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
312 split_k_iters,
313 )
314
315 result = torch.zeros((split_k_iters, M, N), dtype=scales.dtype, device=input.device)
316
317 # A = input, B = qweight, C = result
318 # A = M x K, B = K x N, C = M x N
319 awq_gemm_kernel[grid](
320 input,
321 qweight,
322 result,
323 qzeros,
324 scales,
325 M,
326 N,
327 K,
328 group_size,
329 BLOCK_SIZE_M=block_size_m,
330 BLOCK_SIZE_N=block_size_n,
331 BLOCK_SIZE_K=block_size_k,
332 SPLIT_K=split_k_iters,
333 )
334
335 result = result.sum(0)
336
337 return result