@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 1,494 lines 49.4 KB
Raw
1 /*
2 Adapted from https://github.com/turboderp/exllamav2 and
3 https://github.com/qwopqwop200/GPTQ-for-LLaMa
4 */
5
6 #include <cstdint>
7 #include <cstdio>
8
9 #include "../../torch_utils.h"
10 #include <torch/csrc/stable/ops.h>
11 #include <cuda_runtime.h>
12 #include <cuda_fp16.h>
13
14 #include "compat.cuh"
15 #include "matrix_view.cuh"
16 #include "qdq_2.cuh"
17 #include "qdq_3.cuh"
18 #include "qdq_4.cuh"
19 #include "qdq_8.cuh"
20
21 namespace vllm {
22 namespace gptq {
23
24 #define BLOCK_KN_SIZE 128
25 #define BLOCK_M_SIZE_MAX 8
26 #define MAX_GROUPS_IN_BLOCK (BLOCK_KN_SIZE / 32)
27 #define MAX_Q_GEMM_ROWS 50
28 #define MAX_Q_GEMM_ROWS_8BIT 24
29 #define MAX_ALT_GEMM_ROWS 8
30 #define THREADS_X 32
31 #define THREADS_Y 32
32 #define DIVIDE(x, size) (((x) + (size) - 1) / (size))
33
34 #if defined(USE_ROCM)
35 #include <hipblas/hipblas.h>
36 __host__ __forceinline__ hipblasStatus_t __compat_hipblasHgemm(
37 hipblasHandle_t handle, hipblasOperation_t transA,
38 hipblasOperation_t transB, int m, int n, int k, const half* alpha,
39 const half* AP, int lda, const half* BP, int ldb, const half* beta,
40 half* CP, int ldc) {
41 return hipblasHgemm(handle, transA, transB, m, n, k,
42 reinterpret_cast<const hipblasHalf*>(alpha),
43 reinterpret_cast<const hipblasHalf*>(AP), lda,
44 reinterpret_cast<const hipblasHalf*>(BP), ldb,
45 reinterpret_cast<const hipblasHalf*>(beta),
46 reinterpret_cast<hipblasHalf*>(CP), ldc);
47 }
48 #define hipblasHgemm __compat_hipblasHgemm
49
50 // Previous version of PyTorch were converting to rocBLAS instead of hipBLAS.
51 #define rocblas_operation_none HIPBLAS_OP_N
52 #define rocblas_hgemm __compat_hipblasHgemm
53 #endif
54
55 __forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
56 half2 result = {};
57 const half2* a2_ptr = (const half2*)a_ptr;
58 #pragma unroll
59 for (int i = 0; i < 4; i++) result = __hfma2(dq[i], *a2_ptr++, result);
60 return __half2float(__low2half(result)) + __half2float(__high2half(result));
61 }
62
63 __forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr,
64 const float g_result,
65 const float qs_f) {
66 half2 result = {};
67 const half2* a2_ptr = (const half2*)a_ptr;
68 #pragma unroll
69 for (int i = 0; i < 4; i++) result = __hfma2(dq[i], *a2_ptr++, result);
70 float result_f =
71 __half2float(__low2half(result)) + __half2float(__high2half(result));
72 return fma(result_f, qs_f, g_result);
73 }
74
75 __forceinline__ __device__ half dot22_8_h(half2 (&dq)[4], const half* a_ptr,
76 const half g_result,
77 const half qs_h) {
78 // Use FP32 accumulator to avoid potential overflow since unscaled weights are
79 // in the range -128..127
80
81 float result = {};
82 #pragma unroll
83 for (int i = 0; i < 4; i++) {
84 half2 w01 = dq[i];
85 float w0 = __low2float(w01);
86 float w1 = __high2float(w01);
87 float x0 = __half2float(*a_ptr++);
88 float x1 = __half2float(*a_ptr++);
89 result = fma(w0, x0, result);
90 result = fma(w1, x1, result);
91 }
92 float qs = __half2float(qs_h);
93 result *= qs;
94 half result_h = __float2half_rn(result);
95 return __hadd(result_h, g_result);
96 }
97
98 __forceinline__ __device__ half dot22_16_h(half2 (&dq)[8], const half* a_ptr,
99 const half g_result,
100 const half qs_h) {
101 half2 result = {};
102 const half2* a2_ptr = (const half2*)a_ptr;
103 #pragma unroll
104 for (int i = 0; i < 8; i++) result = __hfma2(dq[i], *a2_ptr++, result);
105 half result_h = __hadd(__low2half(result), __high2half(result));
106 return __hfma(result_h, qs_h, g_result);
107 }
108
109 __forceinline__ __device__ half dot22_32_h(half2 (&dq)[16], const half* a_ptr,
110 const half g_result,
111 const half qs_h) {
112 half2 result = {};
113 const half2* a2_ptr = (const half2*)a_ptr;
114 #pragma unroll
115 for (int i = 0; i < 16; i += 1) result = __hfma2(dq[i], *a2_ptr++, result);
116 half result_h = __hadd(__low2half(result), __high2half(result));
117 return __hfma(result_h, qs_h, g_result);
118 }
119
120 typedef void (*fp_gemm_half_q_half_gptq_kernel)(const half*, const uint32_t*,
121 const uint32_t*, const half*,
122 half*, const int, const int,
123 const int, const int,
124 const bool);
125
126 template <bool first_block, int m_count>
127 __global__ void gemm_half_q_half_gptq_4bit_kernel(
128 const half* __restrict__ a, const uint32_t* __restrict__ b_q_weight,
129 const uint32_t* __restrict__ b_gptq_qzeros,
130 const half* __restrict__ b_gptq_scales, half* __restrict__ c,
131 const int size_m, const int size_n, const int size_k, const int groups,
132 const bool use_v2_format) {
133 MatrixView_half a_(a, size_m, size_k);
134 MatrixView_half_rw c_(c, size_m, size_n);
135 MatrixView_q4_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
136 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
137
138 // GPTQv2 and GPTQv1 handles zero points differently
139 int zero_offset = use_v2_format ? 0 : 1;
140
141 auto t = threadIdx.x;
142
143 // Block
144 auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
145 auto offset_m = blockIdx.y * m_count;
146 auto offset_k = blockIdx.z * BLOCK_KN_SIZE;
147
148 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
149
150 int n = offset_n + t * 4;
151
152 // Preload block_a
153 __shared__ half block_a[m_count][BLOCK_KN_SIZE];
154
155 if (offset_k + t < end_k) {
156 for (int m = 0; m < m_count; ++m) {
157 const half* a_ptr = a_.item_ptr(offset_m + m, 0);
158 half* block_a_ptr = block_a[m];
159
160 block_a_ptr[t] = a_ptr[offset_k + t];
161 }
162 }
163
164 // Zero output
165 if (n >= size_n) return;
166
167 __syncthreads();
168
169 // Find initial group
170 int groupsize = size_k / groups;
171 int group = offset_k / groupsize;
172 int nextgroup = offset_k + groupsize;
173
174 // a, b offset
175 int qk = offset_k / (32 / 4);
176
177 const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
178 const half* a_ptr = &block_a[0][0];
179 int a_stride = BLOCK_KN_SIZE;
180
181 // Initial group
182 int zeros[4];
183 float scales[4];
184 half2 z1z16[4][2];
185 half2 y1y16[4][2];
186 b_gptq_qzeros_.item4(zeros, group, n);
187 b_gptq_scales_.item4_f(scales, group, n);
188 dequant_4bit_8_prep_zero(zeros[0] + zero_offset, z1z16[0], y1y16[0]);
189 dequant_4bit_8_prep_zero(zeros[1] + zero_offset, z1z16[1], y1y16[1]);
190 dequant_4bit_8_prep_zero(zeros[2] + zero_offset, z1z16[2], y1y16[2]);
191 dequant_4bit_8_prep_zero(zeros[3] + zero_offset, z1z16[3], y1y16[3]);
192
193 // Column result
194 float block_c[m_count][4] = {};
195
196 // Dequantize and multiply
197 int k = offset_k;
198 while (k < end_k) {
199 if (k == nextgroup) {
200 group++;
201 nextgroup += groupsize;
202 b_gptq_qzeros_.item4(zeros, group, n);
203 b_gptq_scales_.item4_f(scales, group, n);
204 dequant_4bit_8_prep_zero(zeros[0] + zero_offset, z1z16[0], y1y16[0]);
205 dequant_4bit_8_prep_zero(zeros[1] + zero_offset, z1z16[1], y1y16[1]);
206 dequant_4bit_8_prep_zero(zeros[2] + zero_offset, z1z16[2], y1y16[2]);
207 dequant_4bit_8_prep_zero(zeros[3] + zero_offset, z1z16[3], y1y16[3]);
208 }
209
210 #pragma unroll
211 for (int j = 0; j < 4; j++) {
212 const int4* b_ptr4 = (int4*)b_ptr;
213 int4 load_int4 = *b_ptr4;
214
215 half2 dq[4][4];
216 dequant_4bit_8_gptq(load_int4.x, dq[0], z1z16[0], y1y16[0], size_n,
217 false);
218 dequant_4bit_8_gptq(load_int4.y, dq[1], z1z16[1], y1y16[1], size_n,
219 false);
220 dequant_4bit_8_gptq(load_int4.z, dq[2], z1z16[2], y1y16[2], size_n,
221 false);
222 dequant_4bit_8_gptq(load_int4.w, dq[3], z1z16[3], y1y16[3], size_n,
223 false);
224
225 #pragma unroll
226 for (int m = 0; m < m_count; m++) {
227 block_c[m][0] = fma(dot22_8_f(dq[0], a_ptr + m * a_stride), scales[0],
228 block_c[m][0]);
229 block_c[m][1] = fma(dot22_8_f(dq[1], a_ptr + m * a_stride), scales[1],
230 block_c[m][1]);
231 block_c[m][2] = fma(dot22_8_f(dq[2], a_ptr + m * a_stride), scales[2],
232 block_c[m][2]);
233 block_c[m][3] = fma(dot22_8_f(dq[3], a_ptr + m * a_stride), scales[3],
234 block_c[m][3]);
235 }
236
237 b_ptr += size_n;
238 a_ptr += 8;
239 }
240
241 k += 32;
242 }
243
244 for (int m = 0; m < m_count; m++) {
245 half2* out = (half2*)c_.item_ptr(offset_m + m, n);
246 half2 result01 = __halves2half2(__float2half_rn(block_c[m][0]),
247 __float2half_rn(block_c[m][1]));
248 half2 result23 = __halves2half2(__float2half_rn(block_c[m][2]),
249 __float2half_rn(block_c[m][3]));
250 atomicAdd(out, result01);
251 atomicAdd(out + 1, result23);
252 }
253 }
254
255 template <bool first_block, int m_count>
256 __global__ void gemm_half_q_half_gptq_2bit_kernel(
257 const half* __restrict__ a, const uint32_t* __restrict__ b_q_weight,
258 const uint32_t* __restrict__ b_gptq_qzeros,
259 const half* __restrict__ b_gptq_scales, half* __restrict__ c,
260 const int size_m, const int size_n, const int size_k, const int groups,
261 const bool use_v2_format) {
262 MatrixView_half a_(a, size_m, size_k);
263 MatrixView_half_rw c_(c, size_m, size_n);
264 MatrixView_q2_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
265 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
266
267 // GPTQv2 and GPTQv1 handles zero points differently
268 int zero_offset = use_v2_format ? 0 : 1;
269
270 auto t = threadIdx.x;
271
272 // Block
273 auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
274 auto offset_m = blockIdx.y * m_count;
275 auto offset_k = blockIdx.z * BLOCK_KN_SIZE;
276
277 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
278
279 int n = offset_n + t * 4;
280
281 // Preload block_a
282 __shared__ half block_a[m_count][BLOCK_KN_SIZE];
283
284 if (offset_k + t < end_k) {
285 for (int m = 0; m < m_count; ++m) {
286 const half* a_ptr = a_.item_ptr(offset_m + m, 0);
287 half* block_a_ptr = block_a[m];
288
289 block_a_ptr[t] = a_ptr[offset_k + t];
290 }
291 }
292
293 // Zero output
294 if (n >= size_n) return;
295
296 __syncthreads();
297
298 // Find initial group
299 int groupsize = size_k / groups;
300 int group = offset_k / groupsize;
301 int nextgroup = offset_k + groupsize;
302
303 // a, b offset
304 int qk = offset_k / (32 / 2);
305
306 const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
307 const half* a_ptr = &block_a[0][0];
308 int a_stride = BLOCK_KN_SIZE;
309
310 // Initial group
311 int zeros[4];
312 half scales[4];
313 b_gptq_qzeros_.item4(zeros, group, n);
314 b_gptq_scales_.item4(scales, group, n);
315 // Column result
316 half block_c[m_count][4] = {};
317
318 // Dequantize and multiply
319 int k = offset_k;
320 while (k < end_k) {
321 if (k == nextgroup) {
322 group++;
323 nextgroup += groupsize;
324 b_gptq_qzeros_.item4(zeros, group, n);
325 b_gptq_scales_.item4(scales, group, n);
326 }
327
328 #pragma unroll
329 for (int j = 0; j < 1; j++) {
330 const int4* b_ptr4 = (int4*)b_ptr;
331 int4 load_int4 = *b_ptr4;
332
333 half2 dq[4][8];
334 dequant_2bit_16(load_int4.x, dq[0], size_n, zeros[0] + zero_offset);
335 dequant_2bit_16(load_int4.y, dq[1], size_n, zeros[1] + zero_offset);
336 dequant_2bit_16(load_int4.z, dq[2], size_n, zeros[2] + zero_offset);
337 dequant_2bit_16(load_int4.w, dq[3], size_n, zeros[3] + zero_offset);
338
339 #pragma unroll
340 for (int m = 0; m < m_count; m++) {
341 block_c[m][0] =
342 dot22_16_h(dq[0], a_ptr + m * a_stride, block_c[m][0], scales[0]);
343 block_c[m][1] =
344 dot22_16_h(dq[1], a_ptr + m * a_stride, block_c[m][1], scales[1]);
345 block_c[m][2] =
346 dot22_16_h(dq[2], a_ptr + m * a_stride, block_c[m][2], scales[2]);
347 block_c[m][3] =
348 dot22_16_h(dq[3], a_ptr + m * a_stride, block_c[m][3], scales[3]);
349 }
350
351 b_ptr += size_n;
352 a_ptr += 16;
353 }
354
355 k += 16;
356 }
357
358 for (int m = 0; m < m_count; m++) {
359 half2* out = (half2*)c_.item_ptr(offset_m + m, n);
360 half2 result01 = __halves2half2(block_c[m][0], block_c[m][1]);
361 half2 result23 = __halves2half2(block_c[m][2], block_c[m][3]);
362 atomicAdd(out, result01);
363 atomicAdd(out + 1, result23);
364 }
365 }
366
367 template <bool first_block, int m_count>
368 __global__ void gemm_half_q_half_gptq_3bit_kernel(
369 const half* __restrict__ a, const uint32_t* __restrict__ b_q_weight,
370 const uint32_t* __restrict__ b_gptq_qzeros,
371 const half* __restrict__ b_gptq_scales, half* __restrict__ c,
372 const int size_m, const int size_n, const int size_k, const int groups,
373 const bool use_v2_format) {
374 MatrixView_half a_(a, size_m, size_k);
375 MatrixView_half_rw c_(c, size_m, size_n);
376 MatrixView_q3_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
377 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
378
379 // GPTQv2 and GPTQv1 handles zero points differently
380 int zero_offset = use_v2_format ? 0 : 1;
381
382 auto t = threadIdx.x;
383
384 // Block
385 auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
386 auto offset_m = blockIdx.y * m_count;
387 auto offset_k = blockIdx.z * BLOCK_KN_SIZE;
388
389 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
390
391 int n = offset_n + t * 4;
392
393 // Preload block_a
394 __shared__ half block_a[m_count][BLOCK_KN_SIZE];
395
396 if (offset_k + t < end_k) {
397 for (int m = 0; m < m_count; ++m) {
398 const half* a_ptr = a_.item_ptr(offset_m + m, 0);
399 half* block_a_ptr = block_a[m];
400
401 block_a_ptr[t] = a_ptr[offset_k + t];
402 }
403 }
404
405 // Zero output
406 if (n >= size_n) return;
407
408 __syncthreads();
409
410 // Find initial group
411 int groupsize = size_k / groups;
412 int group = offset_k / groupsize;
413 int nextgroup = offset_k + groupsize;
414
415 // a, b offset
416 int qk = offset_k / 32 * 3;
417
418 const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
419 const half* a_ptr = &block_a[0][0];
420 int a_stride = BLOCK_KN_SIZE;
421
422 // Initial group
423 int zeros[4];
424 half scales[4];
425 b_gptq_qzeros_.item4(zeros, group, n);
426 b_gptq_scales_.item4(scales, group, n);
427 // Column result
428 half block_c[m_count][4] = {};
429
430 // Dequantize and multiply
431 int k = offset_k;
432 while (k < end_k) {
433 if (k == nextgroup) {
434 group++;
435 nextgroup += groupsize;
436 b_gptq_qzeros_.item4(zeros, group, n);
437 b_gptq_scales_.item4(scales, group, n);
438 }
439
440 #pragma unroll
441 for (int j = 0; j < 1; j++) {
442 int4 load_int4[3];
443 load_int4[0] = *((int4*)b_ptr);
444 b_ptr += size_n;
445 load_int4[1] = *((int4*)b_ptr);
446 b_ptr += size_n;
447 load_int4[2] = *((int4*)b_ptr);
448 b_ptr += size_n;
449
450 half2 dq[4][16];
451 dequant_3bit_32(load_int4[0].x, load_int4[1].x, load_int4[2].x, dq[0],
452 size_n, zeros[0] + zero_offset);
453 dequant_3bit_32(load_int4[0].y, load_int4[1].y, load_int4[2].y, dq[1],
454 size_n, zeros[1] + zero_offset);
455 dequant_3bit_32(load_int4[0].z, load_int4[1].z, load_int4[2].z, dq[2],
456 size_n, zeros[2] + zero_offset);
457 dequant_3bit_32(load_int4[0].w, load_int4[1].w, load_int4[2].w, dq[3],
458 size_n, zeros[3] + zero_offset);
459
460 #pragma unroll
461 for (int m = 0; m < m_count; m++) {
462 block_c[m][0] =
463 dot22_32_h(dq[0], a_ptr + m * a_stride, block_c[m][0], scales[0]);
464 block_c[m][1] =
465 dot22_32_h(dq[1], a_ptr + m * a_stride, block_c[m][1], scales[1]);
466 block_c[m][2] =
467 dot22_32_h(dq[2], a_ptr + m * a_stride, block_c[m][2], scales[2]);
468 block_c[m][3] =
469 dot22_32_h(dq[3], a_ptr + m * a_stride, block_c[m][3], scales[3]);
470 }
471 a_ptr += 32;
472 }
473
474 k += 32;
475 }
476
477 for (int m = 0; m < m_count; m++) {
478 half2* out = (half2*)c_.item_ptr(offset_m + m, n);
479 half2 result01 = __halves2half2(block_c[m][0], block_c[m][1]);
480 half2 result23 = __halves2half2(block_c[m][2], block_c[m][3]);
481 atomicAdd(out, result01);
482 atomicAdd(out + 1, result23);
483 }
484 }
485
486 template <bool first_block, int m_count>
487 __global__ void gemm_half_q_half_gptq_8bit_kernel(
488 const half* __restrict__ a, const uint32_t* __restrict__ b_q_weight,
489 const uint32_t* __restrict__ b_gptq_qzeros,
490 const half* __restrict__ b_gptq_scales, half* __restrict__ c,
491 const int size_m, const int size_n, const int size_k, const int groups,
492 const bool use_v2_format) {
493 MatrixView_half a_(a, size_m, size_k);
494 MatrixView_half_rw c_(c, size_m, size_n);
495 MatrixView_q8_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
496 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
497
498 // GPTQv2 and GPTQv1 handles zero points differently
499 int zero_offset = use_v2_format ? 0 : 1;
500
501 auto t = threadIdx.x;
502
503 // Block
504 auto offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
505 auto offset_m = blockIdx.y * m_count;
506 auto offset_k = blockIdx.z * BLOCK_KN_SIZE;
507
508 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
509
510 int n = offset_n + t * 4;
511
512 // Preload block_a
513 __shared__ half block_a[m_count][BLOCK_KN_SIZE];
514
515 if (offset_k + t < end_k) {
516 for (int m = 0; m < m_count; ++m) {
517 const half* a_ptr = a_.item_ptr(offset_m + m, 0);
518 half* block_a_ptr = block_a[m];
519
520 block_a_ptr[t] = a_ptr[offset_k + t];
521 }
522 }
523
524 // Zero output
525 if (n >= size_n) return;
526
527 __syncthreads();
528
529 // Find initial group
530 int groupsize = size_k / groups;
531 int group = offset_k / groupsize;
532 int nextgroup = offset_k + groupsize;
533
534 // a, b offset
535 int qk = offset_k / (32 / 8);
536
537 const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
538 const half* a_ptr = &block_a[0][0];
539 int a_stride = BLOCK_KN_SIZE;
540
541 // Initial group
542 int zeros[4];
543 half scales[4];
544 b_gptq_qzeros_.item4(zeros, group, n);
545 b_gptq_scales_.item4(scales, group, n);
546 // Column result
547 half block_c[m_count][4] = {};
548
549 // Dequantize and multiply
550 int k = offset_k;
551 while (k < end_k) {
552 if (k == nextgroup) {
553 group++;
554 nextgroup += groupsize;
555 b_gptq_qzeros_.item4(zeros, group, n);
556 b_gptq_scales_.item4(scales, group, n);
557 }
558
559 #pragma unroll
560 for (int j = 0; j < 4; j++) {
561 int4 load_int4[2];
562 load_int4[0] = *((int4*)b_ptr);
563 b_ptr += size_n;
564 load_int4[1] = *((int4*)b_ptr);
565 b_ptr += size_n;
566
567 half2 dq[4][4];
568 dequant_8bit_8(load_int4[0].x, load_int4[1].x, dq[0], size_n,
569 zeros[0] + zero_offset);
570 dequant_8bit_8(load_int4[0].y, load_int4[1].y, dq[1], size_n,
571 zeros[1] + zero_offset);
572 dequant_8bit_8(load_int4[0].z, load_int4[1].z, dq[2], size_n,
573 zeros[2] + zero_offset);
574 dequant_8bit_8(load_int4[0].w, load_int4[1].w, dq[3], size_n,
575 zeros[3] + zero_offset);
576
577 for (int m = 0; m < m_count; m++) {
578 block_c[m][0] =
579 dot22_8_h(dq[0], a_ptr + m * a_stride, block_c[m][0], scales[0]);
580 block_c[m][1] =
581 dot22_8_h(dq[1], a_ptr + m * a_stride, block_c[m][1], scales[1]);
582 block_c[m][2] =
583 dot22_8_h(dq[2], a_ptr + m * a_stride, block_c[m][2], scales[2]);
584 block_c[m][3] =
585 dot22_8_h(dq[3], a_ptr + m * a_stride, block_c[m][3], scales[3]);
586 }
587 a_ptr += 8;
588 }
589 k += 32;
590 }
591
592 for (int m = 0; m < m_count; m++) {
593 half2* out = (half2*)c_.item_ptr(offset_m + m, n);
594 half2 result01 = __halves2half2(block_c[m][0], block_c[m][1]);
595 half2 result23 = __halves2half2(block_c[m][2], block_c[m][3]);
596 atomicAdd(out, result01);
597 atomicAdd(out + 1, result23);
598 }
599 }
600
601 fp_gemm_half_q_half_gptq_kernel pick_gemm_half_q_half_gptq_kernel(
602 bool first_block, const int m_count, const int bit) {
603 #define SELECT_KERNEL(M_COUNT) \
604 if (m_count == M_COUNT) { \
605 if (bit == 2) return gemm_half_q_half_gptq_2bit_kernel<true, M_COUNT>; \
606 if (bit == 3) return gemm_half_q_half_gptq_3bit_kernel<true, M_COUNT>; \
607 if (bit == 4) return gemm_half_q_half_gptq_4bit_kernel<true, M_COUNT>; \
608 if (bit == 8) return gemm_half_q_half_gptq_8bit_kernel<true, M_COUNT>; \
609 }
610 #if BLOCK_M_SIZE_MAX >= 1
611 SELECT_KERNEL(1);
612 #endif
613 #if BLOCK_M_SIZE_MAX >= 2
614 SELECT_KERNEL(2);
615 #endif
616 #if BLOCK_M_SIZE_MAX >= 3
617 SELECT_KERNEL(3);
618 #endif
619 #if BLOCK_M_SIZE_MAX >= 4
620 SELECT_KERNEL(4);
621 #endif
622 #if BLOCK_M_SIZE_MAX >= 5
623 SELECT_KERNEL(5);
624 #endif
625 #if BLOCK_M_SIZE_MAX >= 6
626 SELECT_KERNEL(6);
627 #endif
628 #if BLOCK_M_SIZE_MAX >= 7
629 SELECT_KERNEL(7);
630 #endif
631 #if BLOCK_M_SIZE_MAX >= 8
632 SELECT_KERNEL(8);
633 #endif
634 return NULL;
635 }
636
637 void gemm_half_q_half_cuda_part(const half* a, const uint32_t* b_q_weight,
638 const uint32_t* b_gptq_qzeros,
639 const half* b_gptq_scales, half* c, int size_m,
640 int size_n, int size_k, int m_count, int groups,
641 bool use_v2_format, int bit) {
642 dim3 blockDim, gridDim;
643 blockDim.x = BLOCK_KN_SIZE;
644 blockDim.y = 1;
645 blockDim.z = 1;
646 gridDim.x = DIVIDE(size_n, BLOCK_KN_SIZE * 4);
647 gridDim.y = DIVIDE(size_m, m_count);
648 gridDim.z = DIVIDE(size_k, BLOCK_KN_SIZE);
649
650 fp_gemm_half_q_half_gptq_kernel kernel =
651 pick_gemm_half_q_half_gptq_kernel(true, m_count, bit);
652
653 const cudaStream_t stream = get_current_cuda_stream();
654 kernel<<<gridDim, blockDim, 0, stream>>>(a, b_q_weight, b_gptq_qzeros,
655 b_gptq_scales, c, size_m, size_n,
656 size_k, groups, use_v2_format);
657 }
658
659 __global__ void reconstruct_exllama_8bit_kernel(
660 const uint32_t* __restrict__ b_q_weight,
661 const uint32_t* __restrict__ b_gptq_qzeros,
662 const half* __restrict__ b_gptq_scales, const int size_k, const int size_n,
663 const int groups, const bool use_v2_format, half* __restrict__ b) {
664 MatrixView_half_rw b_(b, size_k, size_n);
665 MatrixView_q8_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
666 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
667
668 // GPTQv2 and GPTQv1 handles zero points differently
669 int zero_offset = use_v2_format ? 0 : 1;
670
671 auto offset_k = BLOCK_KN_SIZE * blockIdx.y;
672 auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
673
674 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
675
676 auto t = threadIdx.x;
677
678 // Column
679 int n = offset_n + t * 4;
680 if (n >= size_n) return;
681
682 // Find initial group
683 int groupsize = size_k / groups;
684 int group = offset_k / groupsize;
685 int nextgroup = offset_k + groupsize;
686
687 // b offset
688 int qk = offset_k / (32 / 8);
689
690 const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
691
692 // Initial zeros/scale
693 int zeros[4];
694 half2 scales[4];
695 b_gptq_qzeros_.item4(zeros, group, n);
696 b_gptq_scales_.item4_h2(scales, group, n);
697
698 int k = offset_k;
699 int lk = 0;
700
701 while (k < end_k) {
702 if (k == nextgroup) {
703 group++;
704 nextgroup += groupsize;
705 b_gptq_qzeros_.item4(zeros, group, n);
706 b_gptq_scales_.item4_h2(scales, group, n);
707 }
708
709 for (int p = 0; p < 4; p++) {
710 int4 load_int4[2];
711 load_int4[0] = *((int4*)b_ptr);
712 b_ptr += size_n;
713 load_int4[1] = *((int4*)b_ptr);
714 b_ptr += size_n;
715
716 half2 dq[4][4];
717 dequant_8bit_8(load_int4[0].x, load_int4[1].x, dq[0], size_n,
718 zeros[0] + zero_offset);
719 dequant_8bit_8(load_int4[0].y, load_int4[1].y, dq[1], size_n,
720 zeros[1] + zero_offset);
721 dequant_8bit_8(load_int4[0].z, load_int4[1].z, dq[2], size_n,
722 zeros[2] + zero_offset);
723 dequant_8bit_8(load_int4[0].w, load_int4[1].w, dq[3], size_n,
724 zeros[3] + zero_offset);
725
726 for (int j = 0; j < 4; j++) {
727 for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
728 b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]),
729 __low2half(dq[2][j]), __low2half(dq[3][j]));
730 b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
731 __high2half(dq[1][j]), __high2half(dq[2][j]),
732 __high2half(dq[3][j]));
733 }
734 }
735 k += 32;
736 }
737 }
738
739 __global__ void reconstruct_exllama_4bit_kernel(
740 const uint32_t* __restrict__ b_q_weight,
741 const uint32_t* __restrict__ b_gptq_qzeros,
742 const half* __restrict__ b_gptq_scales, const int size_k, const int size_n,
743 const int groups, const bool use_v2_format, half* __restrict__ b) {
744 MatrixView_half_rw b_(b, size_k, size_n);
745 MatrixView_q4_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
746 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
747
748 // GPTQv2 and GPTQv1 handles zero points differently
749 int zero_offset = use_v2_format ? 0 : 1;
750
751 auto offset_k = BLOCK_KN_SIZE * blockIdx.y;
752 auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
753
754 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
755
756 auto t = threadIdx.x;
757
758 // Column
759 int n = offset_n + t * 4;
760 if (n >= size_n) return;
761
762 // Find initial group
763 int groupsize = size_k / groups;
764 int group = offset_k / groupsize;
765 int nextgroup = offset_k + groupsize;
766
767 // b offset
768 int qk = offset_k / (32 / 4);
769
770 const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
771
772 // Initial zeros/scale
773 int zeros[4];
774 half2 scales[4];
775 half2 z1z16[4][2];
776 half2 y1y16[4][2];
777 b_gptq_qzeros_.item4(zeros, group, n);
778 b_gptq_scales_.item4_h2(scales, group, n);
779 dequant_4bit_8_prep_zero(zeros[0] + zero_offset, z1z16[0], y1y16[0]);
780 dequant_4bit_8_prep_zero(zeros[1] + zero_offset, z1z16[1], y1y16[1]);
781 dequant_4bit_8_prep_zero(zeros[2] + zero_offset, z1z16[2], y1y16[2]);
782 dequant_4bit_8_prep_zero(zeros[3] + zero_offset, z1z16[3], y1y16[3]);
783
784 int k = offset_k;
785 int lk = 0;
786
787 while (k < end_k) {
788 if (k == nextgroup) {
789 group++;
790 nextgroup += groupsize;
791 b_gptq_qzeros_.item4(zeros, group, n);
792 b_gptq_scales_.item4_h2(scales, group, n);
793 dequant_4bit_8_prep_zero(zeros[0] + zero_offset, z1z16[0], y1y16[0]);
794 dequant_4bit_8_prep_zero(zeros[1] + zero_offset, z1z16[1], y1y16[1]);
795 dequant_4bit_8_prep_zero(zeros[2] + zero_offset, z1z16[2], y1y16[2]);
796 dequant_4bit_8_prep_zero(zeros[3] + zero_offset, z1z16[3], y1y16[3]);
797 }
798
799 for (int p = 0; p < 4; p++) {
800 half2 dq[4][4];
801 const int4* b_ptr4 = (int4*)b_ptr;
802 int4 load_int4 = *b_ptr4;
803
804 dequant_4bit_8_gptq(load_int4.x, dq[0], z1z16[0], y1y16[0], size_n,
805 false);
806 dequant_4bit_8_gptq(load_int4.y, dq[1], z1z16[1], y1y16[1], size_n,
807 false);
808 dequant_4bit_8_gptq(load_int4.z, dq[2], z1z16[2], y1y16[2], size_n,
809 false);
810 dequant_4bit_8_gptq(load_int4.w, dq[3], z1z16[3], y1y16[3], size_n,
811 false);
812
813 b_ptr += size_n;
814 for (int j = 0; j < 4; j++) {
815 for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
816 b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]),
817 __low2half(dq[2][j]), __low2half(dq[3][j]));
818 b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
819 __high2half(dq[1][j]), __high2half(dq[2][j]),
820 __high2half(dq[3][j]));
821 }
822 }
823 k += 32;
824 }
825 }
826
827 __global__ void reconstruct_exllama_3bit_kernel(
828 const uint32_t* __restrict__ b_q_weight,
829 const uint32_t* __restrict__ b_gptq_qzeros,
830 const half* __restrict__ b_gptq_scales, const int size_k, const int size_n,
831 const int groups, const bool use_v2_format, half* __restrict__ b) {
832 MatrixView_half_rw b_(b, size_k, size_n);
833 MatrixView_q3_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
834 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
835
836 // GPTQv2 and GPTQv1 handles zero points differently
837 int zero_offset = use_v2_format ? 0 : 1;
838
839 auto offset_k = BLOCK_KN_SIZE * blockIdx.y;
840 auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
841
842 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
843
844 auto t = threadIdx.x;
845
846 // Column
847 int n = offset_n + t * 4;
848 if (n >= size_n) return;
849
850 // Find initial group
851 int groupsize = size_k / groups;
852 int group = offset_k / groupsize;
853 int nextgroup = offset_k + groupsize;
854
855 // b offset
856 int qk = offset_k / 32 * 3;
857
858 const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
859
860 // Initial zeros/scale
861 int zeros[4];
862 half2 scales[4];
863 b_gptq_qzeros_.item4(zeros, group, n);
864 b_gptq_scales_.item4_h2(scales, group, n);
865
866 int k = offset_k;
867 int lk = 0;
868
869 while (k < end_k) {
870 if (k == nextgroup) {
871 group++;
872 nextgroup += groupsize;
873 b_gptq_qzeros_.item4(zeros, group, n);
874 b_gptq_scales_.item4_h2(scales, group, n);
875 }
876
877 for (int p = 0; p < 1; p++) {
878 int4 load_int4[3];
879 load_int4[0] = *((int4*)b_ptr);
880 b_ptr += size_n;
881 load_int4[1] = *((int4*)b_ptr);
882 b_ptr += size_n;
883 load_int4[2] = *((int4*)b_ptr);
884 b_ptr += size_n;
885
886 half2 dq[4][16];
887 dequant_3bit_32(load_int4[0].x, load_int4[1].x, load_int4[2].x, dq[0],
888 size_n, zeros[0] + zero_offset);
889 dequant_3bit_32(load_int4[0].y, load_int4[1].y, load_int4[2].y, dq[1],
890 size_n, zeros[1] + zero_offset);
891 dequant_3bit_32(load_int4[0].z, load_int4[1].z, load_int4[2].z, dq[2],
892 size_n, zeros[2] + zero_offset);
893 dequant_3bit_32(load_int4[0].w, load_int4[1].w, load_int4[2].w, dq[3],
894 size_n, zeros[3] + zero_offset);
895
896 for (int j = 0; j < 16; j++) {
897 for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
898 b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]),
899 __low2half(dq[2][j]), __low2half(dq[3][j]));
900 b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
901 __high2half(dq[1][j]), __high2half(dq[2][j]),
902 __high2half(dq[3][j]));
903 }
904 }
905 k += 32;
906 }
907 }
908
909 __global__ void reconstruct_exllama_2bit_kernel(
910 const uint32_t* __restrict__ b_q_weight,
911 const uint32_t* __restrict__ b_gptq_qzeros,
912 const half* __restrict__ b_gptq_scales, const int size_k, const int size_n,
913 const int groups, const bool use_v2_format, half* __restrict__ b) {
914 MatrixView_half_rw b_(b, size_k, size_n);
915 MatrixView_q2_row b_gptq_qzeros_(b_gptq_qzeros, groups, size_n);
916 MatrixView_half b_gptq_scales_(b_gptq_scales, groups, size_n);
917
918 // GPTQv2 and GPTQv1 handles zero points differently
919 int zero_offset = use_v2_format ? 0 : 1;
920
921 auto offset_k = BLOCK_KN_SIZE * blockIdx.y;
922 auto offset_n = BLOCK_KN_SIZE * blockIdx.x * 4;
923
924 int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
925
926 auto t = threadIdx.x;
927
928 // Column
929 int n = offset_n + t * 4;
930 if (n >= size_n) return;
931
932 // Find initial group
933 int groupsize = size_k / groups;
934 int group = offset_k / groupsize;
935 int nextgroup = offset_k + groupsize;
936
937 // b offset
938 int qk = offset_k / (32 / 2);
939
940 const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
941
942 // Initial zeros/scale
943 int zeros[4];
944 half2 scales[4];
945 b_gptq_qzeros_.item4(zeros, group, n);
946 b_gptq_scales_.item4_h2(scales, group, n);
947
948 int k = offset_k;
949 int lk = 0;
950
951 while (k < end_k) {
952 if (k == nextgroup) {
953 group++;
954 nextgroup += groupsize;
955 b_gptq_qzeros_.item4(zeros, group, n);
956 b_gptq_scales_.item4_h2(scales, group, n);
957 }
958
959 for (int p = 0; p < 2; p++) {
960 const int4* b_ptr4 = (int4*)b_ptr;
961 int4 load_int4 = *b_ptr4;
962
963 half2 dq[4][8];
964 dequant_2bit_16(load_int4.x, dq[0], size_n, zeros[0] + zero_offset);
965 dequant_2bit_16(load_int4.y, dq[1], size_n, zeros[1] + zero_offset);
966 dequant_2bit_16(load_int4.z, dq[2], size_n, zeros[2] + zero_offset);
967 dequant_2bit_16(load_int4.w, dq[3], size_n, zeros[3] + zero_offset);
968
969 b_ptr += size_n;
970 for (int j = 0; j < 8; j++) {
971 for (int v = 0; v < 4; v++) dq[v][j] = __hmul2(scales[v], dq[v][j]);
972 b_.set4(offset_k + lk++, n, __low2half(dq[0][j]), __low2half(dq[1][j]),
973 __low2half(dq[2][j]), __low2half(dq[3][j]));
974 b_.set4(offset_k + lk++, n, __high2half(dq[0][j]),
975 __high2half(dq[1][j]), __high2half(dq[2][j]),
976 __high2half(dq[3][j]));
977 }
978 }
979 k += 32;
980 }
981 }
982
983 void reconstruct_exllama(const uint32_t* b_q_weight,
984 const uint32_t* b_gptq_qzeros,
985 const half* b_gptq_scales, half* out, int height,
986 int width, int groups, bool use_v2_format, int bit) {
987 dim3 blockDim, gridDim;
988 blockDim.x = BLOCK_KN_SIZE;
989 blockDim.y = 1;
990 gridDim.y = DIVIDE(height, BLOCK_KN_SIZE);
991 gridDim.x = DIVIDE(width, BLOCK_KN_SIZE);
992
993 auto reconstruct_exllama_kernel = reconstruct_exllama_4bit_kernel;
994 if (bit == 2) {
995 reconstruct_exllama_kernel = reconstruct_exllama_2bit_kernel;
996 } else if (bit == 3) {
997 reconstruct_exllama_kernel = reconstruct_exllama_3bit_kernel;
998 } else if (bit == 8) {
999 reconstruct_exllama_kernel = reconstruct_exllama_8bit_kernel;
1000 }
1001
1002 const cudaStream_t stream = get_current_cuda_stream();
1003 reconstruct_exllama_kernel<<<gridDim, blockDim, 0, stream>>>(
1004 b_q_weight, b_gptq_qzeros, b_gptq_scales, height, width, groups,
1005 use_v2_format, out);
1006 }
1007
1008 __global__ void gemm_half_q_half_alt_4bit_kernel(
1009 const half2* __restrict__ vec, const uint32_t* __restrict__ mat,
1010 half* __restrict__ mul, const half* __restrict__ scales,
1011 const uint32_t* __restrict__ zeros, int batch, int height, int width,
1012 int groups, bool use_v2_format) {
1013 int zero_width = width / 8;
1014 int vec_height = height * 4;
1015 const int blockwidth2 = BLOCK_KN_SIZE / 2;
1016 auto b = blockIdx.y * BLOCK_M_SIZE_MAX;
1017 int b_end = min(BLOCK_M_SIZE_MAX, batch - b);
1018 auto h = BLOCK_KN_SIZE * blockIdx.z / 8;
1019 int h_end = min(BLOCK_KN_SIZE / 8, height - h) * 4;
1020 auto w = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1021
1022 // GPTQv2 and GPTQv1 handles zero points differently
1023 int zero_offset = use_v2_format ? 0 : 1;
1024
1025 __shared__ half2 blockvec[BLOCK_M_SIZE_MAX][blockwidth2];
1026 if (threadIdx.x < h_end) {
1027 for (int m = 0; m < b_end; ++m) {
1028 blockvec[m][threadIdx.x] =
1029 vec[(m + b) * vec_height + blockIdx.z * BLOCK_KN_SIZE / 2 +
1030 threadIdx.x];
1031 }
1032 }
1033
1034 __shared__ half2 deq2[256][8];
1035 auto val = threadIdx.x / 8;
1036 auto off = threadIdx.x % 8;
1037 for (; val < 256; val += BLOCK_KN_SIZE / 8) {
1038 deq2[val][off] =
1039 __halves2half2(__int2half_rn(val & 0xF), __int2half_rn(val >> 4));
1040 }
1041
1042 __syncthreads();
1043
1044 int i = width * h + w;
1045 int g_h = h * 8;
1046 int group_size = vec_height * 2 / groups;
1047 int k = 0;
1048 int z_w = w / 8;
1049 int z_mod = (w % 8) * 4;
1050 half2 res2;
1051 half res[BLOCK_M_SIZE_MAX] = {};
1052
1053 unsigned int tmp;
1054 while (k < h_end) {
1055 tmp = mat[i];
1056 half2 scales_tmp[4];
1057 half2 zeros_tmp[4];
1058 for (int tmp_k = 0; tmp_k < 4; tmp_k++) {
1059 int g = (g_h + (k + tmp_k) * 2) / group_size;
1060 int g2 = (g_h + (k + tmp_k) * 2 + 1) / group_size;
1061 half scale_f = scales[g * width + w];
1062 half scale_f2 = scales[g2 * width + w];
1063 half2 scale = __halves2half2(scale_f, scale_f2);
1064 half2 zero = __halves2half2(
1065 __hmul(scale_f,
1066 __int2half_rn(-((zeros[g * zero_width + z_w] >> z_mod) & 0xF) -
1067 zero_offset)),
1068 __hmul(
1069 scale_f2,
1070 __int2half_rn(-((zeros[g2 * zero_width + z_w] >> z_mod) & 0xF) -
1071 zero_offset)));
1072 scales_tmp[tmp_k] = scale;
1073 zeros_tmp[tmp_k] = zero;
1074 }
1075 for (int m = 0; m < b_end; m++) {
1076 #ifndef USE_ROCM
1077 res2 = {};
1078 #else
1079 res2.x = __half_as_ushort(__float2half(0));
1080 res2.y = __half_as_ushort(__float2half(0));
1081 #endif
1082 res2 = __hfma2(
1083 __hfma2(deq2[(tmp >> 0) & 0xff][off], scales_tmp[0], zeros_tmp[0]),
1084 blockvec[m][k + 0], res2);
1085 res2 = __hfma2(
1086 __hfma2(deq2[(tmp >> 8) & 0xff][off], scales_tmp[1], zeros_tmp[1]),
1087 blockvec[m][k + 1], res2);
1088 res2 = __hfma2(
1089 __hfma2(deq2[(tmp >> 16) & 0xff][off], scales_tmp[2], zeros_tmp[2]),
1090 blockvec[m][k + 2], res2);
1091 res2 = __hfma2(
1092 __hfma2(deq2[(tmp >> 24) & 0xff][off], scales_tmp[3], zeros_tmp[3]),
1093 blockvec[m][k + 3], res2);
1094 #ifndef USE_ROCM
1095 res[m] = __hadd(res[m], __hadd(res2.x, res2.y));
1096 #else
1097 res[m] = __hadd(
1098 res[m], __hadd(__ushort_as_half(res2.x), __ushort_as_half(res2.y)));
1099 #endif
1100 }
1101 i += width;
1102 k += 4;
1103 }
1104 for (int m = 0; m < b_end; m++) {
1105 atomicAdd(&mul[(b + m) * width + w], res[m]);
1106 }
1107 }
1108
1109 __global__ void gemm_half_q_half_alt_8bit_kernel(
1110 const half2* __restrict__ vec, const uint32_t* __restrict__ mat,
1111 half* __restrict__ mul, const half* __restrict__ scales,
1112 const uint32_t* __restrict__ zeros, int batch, int height, int width,
1113 int groups, bool use_v2_format) {
1114 int zero_width = width / 4;
1115 int vec_height = height * 2;
1116 const int blockwidth2 = BLOCK_KN_SIZE / 2;
1117 auto b = blockIdx.y * BLOCK_M_SIZE_MAX;
1118 int b_end = min(BLOCK_M_SIZE_MAX, batch - b);
1119 auto h = BLOCK_KN_SIZE * blockIdx.z / 4;
1120 int h_end = min(BLOCK_KN_SIZE / 4, height - h) * 2;
1121 auto w = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1122
1123 // GPTQv2 and GPTQv1 handles zero points differently
1124 int zero_offset = use_v2_format ? 0 : 1;
1125
1126 __shared__ half2 blockvec[BLOCK_M_SIZE_MAX][blockwidth2];
1127 if (threadIdx.x < h_end) {
1128 for (int m = 0; m < b_end; ++m) {
1129 blockvec[m][threadIdx.x] =
1130 vec[(m + b) * vec_height + blockIdx.z * BLOCK_KN_SIZE / 2 +
1131 threadIdx.x];
1132 }
1133 }
1134
1135 __syncthreads();
1136
1137 int i = width * h + w;
1138 int g_h = h * 4;
1139 int group_size = vec_height * 2 / groups;
1140 int k = 0;
1141 int z_w = w / 4;
1142 int z_mod = (w % 4) * 8;
1143 half2 res2;
1144 half res[BLOCK_M_SIZE_MAX] = {};
1145
1146 unsigned int tmp;
1147 while (k < h_end) {
1148 tmp = mat[i];
1149 half2 scales_tmp[2];
1150 half2 zeros_tmp[2];
1151 for (int tmp_k = 0; tmp_k < 2; tmp_k++) {
1152 int g = (g_h + (k + tmp_k) * 2) / group_size;
1153 int g2 = (g_h + (k + tmp_k) * 2 + 1) / group_size;
1154 half scale_f = scales[g * width + w];
1155 half scale_f2 = scales[g2 * width + w];
1156 half2 scale = __halves2half2(scale_f, scale_f2);
1157 half2 zero = __halves2half2(
1158 __hmul(scale_f, __int2half_rn(
1159 -((zeros[g * zero_width + z_w] >> z_mod) & 0xff) -
1160 zero_offset)),
1161 __hmul(
1162 scale_f2,
1163 __int2half_rn(-((zeros[g2 * zero_width + z_w] >> z_mod) & 0xff) -
1164 zero_offset)));
1165 scales_tmp[tmp_k] = scale;
1166 zeros_tmp[tmp_k] = zero;
1167 }
1168 for (int m = 0; m < b_end; m++) {
1169 #ifndef USE_ROCM
1170 res2 = {};
1171 #else
1172 res2.x = __half_as_ushort(__float2half(0));
1173 res2.y = __half_as_ushort(__float2half(0));
1174 #endif
1175 half2 v12 = __halves2half2(__int2half_rn(tmp & 0xFF),
1176 __int2half_rn((tmp >> 8) & 0xFF));
1177 res2 = __hfma2(__hfma2(v12, scales_tmp[0], zeros_tmp[0]),
1178 blockvec[m][k + 0], res2);
1179 half2 v34 = __halves2half2(__int2half_rn((tmp >> 16) & 0xFF),
1180 __int2half_rn((tmp >> 24) & 0xFF));
1181 res2 = __hfma2(__hfma2(v34, scales_tmp[1], zeros_tmp[1]),
1182 blockvec[m][k + 1], res2);
1183 #ifndef USE_ROCM
1184 res[m] = __hadd(res[m], __hadd(res2.x, res2.y));
1185 #else
1186 res[m] = __hadd(
1187 res[m], __hadd(__ushort_as_half(res2.x), __ushort_as_half(res2.y)));
1188 #endif
1189 }
1190 i += width;
1191 k += 2;
1192 }
1193 for (int m = 0; m < b_end; m++) {
1194 atomicAdd(&mul[(b + m) * width + w], res[m]);
1195 }
1196 }
1197
1198 void gemm_half_q_half_alt(const half* a, const uint32_t* b_q_weight,
1199 const uint32_t* b_gptq_qzeros,
1200 const half* b_gptq_scales, half* c, int size_m,
1201 int size_n, int size_k, int groups,
1202 bool use_v2_format, int bit) {
1203 dim3 blockDim, gridDim;
1204 blockDim.x = BLOCK_KN_SIZE;
1205 blockDim.y = 1;
1206 blockDim.z = 1;
1207 gridDim.x = DIVIDE(size_n, BLOCK_KN_SIZE);
1208 gridDim.y = DIVIDE(size_m, BLOCK_M_SIZE_MAX);
1209 gridDim.z = DIVIDE(size_k, BLOCK_KN_SIZE);
1210
1211 auto kernel = gemm_half_q_half_alt_4bit_kernel;
1212 if (bit == 8) {
1213 kernel = gemm_half_q_half_alt_8bit_kernel;
1214 }
1215
1216 const cudaStream_t stream = get_current_cuda_stream();
1217 kernel<<<gridDim, blockDim, 0, stream>>>(
1218 (const half2*)a, b_q_weight, c, b_gptq_scales, b_gptq_qzeros, size_m,
1219 size_k / 32 * bit, size_n, groups, use_v2_format);
1220 }
1221
1222 template <class T, int bit>
1223 __global__ void reconstruct_gptq_kernel(
1224 const uint32_t* __restrict__ w, const half* __restrict__ w_scales,
1225 const uint32_t* __restrict__ w_zeros, const int height, const int width,
1226 const int num_groups, const bool use_v2_format, half* __restrict__ out) {
1227 // Start of block
1228
1229 auto column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1230 auto row = blockIdx.y * 32 / bit;
1231 if (column >= width) return;
1232
1233 // Views
1234
1235 MatrixView_half_rw out_(out, height, width);
1236 MatrixView_half w_scales_(w_scales, num_groups, width);
1237 T w_zeros_(w_zeros, num_groups, width);
1238
1239 // GPTQv2 and GPTQv1 handles zero points differently
1240 int zero_offset = use_v2_format ? 0 : 1;
1241
1242 uint32_t w_read = w[blockIdx.y * width + column];
1243 half* out_ptr = out_.item_ptr(row, column);
1244
1245 #pragma unroll
1246 for (int s = 0; s < 32; s += bit) {
1247 int group = (row + s / bit) / (height / num_groups);
1248 half w_scale = w_scales_.item(group, column);
1249 uint32_t w_zero = w_zeros_.item(group, column) + zero_offset;
1250 half w_item =
1251 __hmul(__int2half_rn((int)((w_read >> s) & ((1 << bit) - 1)) - w_zero),
1252 w_scale);
1253 *out_ptr = w_item;
1254 out_ptr += out_.width;
1255 }
1256 }
1257
1258 __global__ void reconstruct_gptq_3bit_kernel(
1259 const uint32_t* __restrict__ w, const half* __restrict__ w_scales,
1260 const uint32_t* __restrict__ w_zeros, const int height, const int width,
1261 const int num_groups, const bool use_v2_format, half* __restrict__ out) {
1262 // Start of block
1263 auto column = BLOCK_KN_SIZE * blockIdx.x + threadIdx.x;
1264 auto row = blockIdx.y * 32;
1265 if (column >= width) return;
1266
1267 // Views
1268
1269 MatrixView_half_rw out_(out, height, width);
1270 MatrixView_half w_scales_(w_scales, num_groups, width);
1271 MatrixView_q3_row w_zeros_(w_zeros, num_groups, width);
1272
1273 // GPTQv2 and GPTQv1 handles zero points differently
1274 int zero_offset = use_v2_format ? 0 : 1;
1275
1276 uint32_t w1 = w[(blockIdx.y * 3) * width + column];
1277 uint32_t w2 = w[(blockIdx.y * 3 + 1) * width + column];
1278 uint32_t w3 = w[(blockIdx.y * 3 + 2) * width + column];
1279 half* out_ptr = out_.item_ptr(row, column);
1280
1281 #pragma unroll
1282 for (int i = 0; i < 32; i += 1) {
1283 int group = (row + i) / (height / num_groups);
1284 half w_scale = w_scales_.item(group, column);
1285 uint32_t w_zero = w_zeros_.item(group, column) + zero_offset;
1286 int w_item;
1287 if (i == 10) {
1288 w_item = (w1 >> 30) | ((w2 << 2) & 0x4);
1289 } else if (i == 21) {
1290 w_item = (w2 >> 31) | ((w3 << 1) & 0x6);
1291 } else if (i < 10) {
1292 w_item = ((w1 >> (i * 3)) & 0x7);
1293 } else if (i < 21) {
1294 w_item = ((w2 >> (i * 3 - 32)) & 0x7);
1295 } else {
1296 w_item = ((w3 >> (i * 3 - 64)) & 0x7);
1297 }
1298 *out_ptr = __hmul(__int2half_rn(w_item - w_zero), w_scale);
1299 out_ptr += out_.width;
1300 }
1301 }
1302
1303 void reconstruct_gptq(const uint32_t* b_q_weight, const uint32_t* b_gptq_qzeros,
1304 const half* b_gptq_scales, half* out, int height,
1305 int width, int groups, bool use_v2_format, int bit) {
1306 dim3 blockDim, gridDim;
1307 blockDim.x = BLOCK_KN_SIZE;
1308 blockDim.y = 1;
1309 gridDim.y = DIVIDE(height, 32 / bit);
1310 gridDim.x = DIVIDE(width, BLOCK_KN_SIZE);
1311
1312 auto kernel = reconstruct_gptq_kernel<MatrixView_q4_row, 4>;
1313 if (bit == 2) {
1314 kernel = reconstruct_gptq_kernel<MatrixView_q2_row, 2>;
1315 } else if (bit == 8) {
1316 kernel = reconstruct_gptq_kernel<MatrixView_q8_row, 8>;
1317 } else if (bit == 3) {
1318 kernel = reconstruct_gptq_3bit_kernel;
1319 gridDim.y = DIVIDE(height, 32);
1320 }
1321
1322 const cudaStream_t stream = get_current_cuda_stream();
1323 kernel<<<gridDim, blockDim, 0, stream>>>(b_q_weight, b_gptq_scales,
1324 b_gptq_qzeros, height, width, groups,
1325 use_v2_format, out);
1326 }
1327
1328 void gemm_half_q_half_cuda(cublasHandle_t cublas_handle, const half* a,
1329 const uint32_t* b_q_weight,
1330 const uint32_t* b_gptq_qzeros,
1331 const half* b_gptq_scales, half* c, half* temp_dq,
1332 int size_m, int size_n, int size_k, int groups,
1333 bool use_exllama, bool use_v2_format, int bit) {
1334 bool use_reconstruct;
1335 if (use_exllama) {
1336 use_reconstruct = ((bit == 8 && size_m > MAX_Q_GEMM_ROWS_8BIT) ||
1337 (bit != 8 && size_m > MAX_Q_GEMM_ROWS));
1338 } else {
1339 // The 2/3-bit kernels are somehow slower than dequant + gemm baseline, so
1340 // we disabled them for now.
1341 use_reconstruct = (bit < 4 || size_m > MAX_ALT_GEMM_ROWS);
1342 }
1343 if (use_reconstruct) {
1344 // Reconstruct FP16 matrix, then cuBLAS
1345 if (use_exllama) {
1346 reconstruct_exllama(b_q_weight, b_gptq_qzeros, b_gptq_scales, temp_dq,
1347 size_k, size_n, groups, use_v2_format, bit);
1348 } else {
1349 reconstruct_gptq(b_q_weight, b_gptq_qzeros, b_gptq_scales, temp_dq,
1350 size_k, size_n, groups, use_v2_format, bit);
1351 }
1352
1353 const half alpha = __float2half(1.0f);
1354 const half beta = __float2half(0.0f);
1355 cublasHgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, size_n, size_m, size_k,
1356 &alpha, temp_dq, size_n, a, size_k, &beta, c, size_n);
1357 } else if (use_exllama) {
1358 // Quantized matmul
1359 int max_chunks = size_m / BLOCK_M_SIZE_MAX;
1360 int last_chunk = max_chunks * BLOCK_M_SIZE_MAX;
1361 int last_chunk_size = size_m - last_chunk;
1362
1363 if (max_chunks) {
1364 gemm_half_q_half_cuda_part(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c,
1365 last_chunk, size_n, size_k, BLOCK_M_SIZE_MAX,
1366 groups, use_v2_format, bit);
1367 }
1368
1369 if (last_chunk_size) {
1370 gemm_half_q_half_cuda_part(
1371 a + last_chunk * size_k, b_q_weight, b_gptq_qzeros, b_gptq_scales,
1372 c + last_chunk * size_n, last_chunk_size, size_n, size_k,
1373 last_chunk_size, groups, use_v2_format, bit);
1374 }
1375 } else {
1376 gemm_half_q_half_alt(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, c, size_m,
1377 size_n, size_k, groups, use_v2_format, bit);
1378 }
1379 }
1380
1381 __global__ void shuffle_4bit_kernel(uint32_t* __restrict__ b_q_weight,
1382 const int size_k, const int size_n) {
1383 auto n = blockIdx.x * THREADS_X + threadIdx.x;
1384 if (n >= size_n) return;
1385 int k = 0;
1386 uint32_t* b_ptr = b_q_weight + n;
1387 while (k < size_k) {
1388 shuffle_4bit_8(b_ptr, size_n);
1389 b_ptr += 1 * size_n;
1390 k += 8;
1391 }
1392 }
1393
1394 __global__ void shuffle_8bit_kernel(uint32_t* __restrict__ b_q_weight,
1395 const int size_k, const int size_n) {
1396 auto n = blockIdx.x * THREADS_X + threadIdx.x;
1397 if (n >= size_n) return;
1398 int k = 0;
1399 uint32_t* b_ptr = b_q_weight + n;
1400 while (k < size_k) {
1401 shuffle_8bit_4(b_ptr, size_n);
1402 b_ptr += 1 * size_n;
1403 k += 4;
1404 }
1405 }
1406
1407 __global__ void shuffle_2bit_kernel(uint32_t* __restrict__ b_q_weight,
1408 const int size_k, const int size_n) {
1409 auto n = blockIdx.x * THREADS_X + threadIdx.x;
1410 if (n >= size_n) return;
1411 int k = 0;
1412 uint32_t* b_ptr = b_q_weight + n;
1413 while (k < size_k) {
1414 shuffle_2bit_16(b_ptr, size_n);
1415 b_ptr += 1 * size_n;
1416 k += 16;
1417 }
1418 }
1419
1420 __global__ void shuffle_3bit_kernel(uint32_t* __restrict__ b_q_weight,
1421 const int size_k, const int size_n) {
1422 auto n = blockIdx.x * THREADS_X + threadIdx.x;
1423 if (n >= size_n) return;
1424 int k = 0;
1425 uint32_t* b_ptr = b_q_weight + n;
1426 while (k < size_k) {
1427 shuffle_3bit_32(b_ptr, size_n);
1428 b_ptr += 3 * size_n;
1429 k += 32;
1430 }
1431 }
1432
1433 void shuffle_exllama_weight(uint32_t* q_weight, int height, int width,
1434 int bit) {
1435 dim3 blockDim, gridDim;
1436 blockDim.x = THREADS_X;
1437 blockDim.y = 1;
1438 gridDim.x = DIVIDE(width, THREADS_X);
1439 gridDim.y = 1;
1440 auto shuffle_kernel = shuffle_4bit_kernel;
1441 if (bit == 2) {
1442 shuffle_kernel = shuffle_2bit_kernel;
1443 } else if (bit == 3) {
1444 shuffle_kernel = shuffle_3bit_kernel;
1445 } else if (bit == 8) {
1446 shuffle_kernel = shuffle_8bit_kernel;
1447 }
1448 const cudaStream_t stream = get_current_cuda_stream();
1449 shuffle_kernel<<<gridDim, blockDim, 0, stream>>>(q_weight, height, width);
1450 }
1451
1452 } // namespace gptq
1453 } // namespace vllm
1454
1455 torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
1456 torch::stable::Tensor b_q_weight,
1457 torch::stable::Tensor b_gptq_qzeros,
1458 torch::stable::Tensor b_gptq_scales,
1459 bool use_exllama, bool use_v2_format,
1460 int64_t bit) {
1461 const torch::stable::accelerator::DeviceGuard device_guard(
1462 a.get_device_index());
1463 const int64_t num_groups = b_gptq_qzeros.size(0);
1464 STD_TORCH_CHECK(num_groups > 0, "GPTQ requires at least one scale group");
1465 STD_TORCH_CHECK(b_gptq_scales.size(0) == num_groups,
1466 "GPTQ scales and zero-points must have the same number of "
1467 "groups");
1468 STD_TORCH_CHECK(a.size(1) % num_groups == 0, "Input size K = ", a.size(1),
1469 " is not divisible by the number of groups = ", num_groups);
1470 auto c = torch::stable::new_zeros(a, {a.size(0), b_q_weight.size(1)});
1471 auto temp_dq =
1472 torch::stable::empty({b_q_weight.size(0) * 32 / bit, b_q_weight.size(1)},
1473 a.scalar_type(), std::nullopt, a.device());
1474
1475 vllm::gptq::gemm_half_q_half_cuda(
1476 get_current_cuda_blas_handle(), (const half*)a.data_ptr(),
1477 (const uint32_t*)b_q_weight.data_ptr(),
1478 (const uint32_t*)b_gptq_qzeros.data_ptr(),
1479 (const half*)b_gptq_scales.data_ptr(), (half*)c.data_ptr(),
1480 (half*)temp_dq.data_ptr(),
1481 c.size(0), // m
1482 c.size(1), // n
1483 a.size(1), // k
1484 num_groups, use_exllama, use_v2_format, bit);
1485 return c;
1486 }
1487
1488 void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit) {
1489 const torch::stable::accelerator::DeviceGuard device_guard(
1490 q_weight.get_device_index());
1491 vllm::gptq::shuffle_exllama_weight((uint32_t*)q_weight.data_ptr(),
1492 q_weight.size(0) * 32 / bit,
1493 q_weight.size(1), bit);
1494 }