@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 536 lines 22.5 KB
Raw
1 /*
2 Adapted from https://github.com/mit-han-lab/llm-awq
3 @article{lin2023awq,
4 title={AWQ: Activation-aware Weight Quantization for LLM Compression and
5 Acceleration}, author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang,
6 Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023}
7 }
8 */
9
10 #include <torch/csrc/stable/tensor.h>
11 #include <torch/csrc/stable/ops.h>
12 #include "libtorch_stable/torch_utils.h"
13
14 #include "libtorch_stable/quantization/awq/dequantize.cuh"
15
16 #include <cuda_fp16.h>
17
18 namespace vllm {
19 namespace awq {
20
21 template <int N>
22 __global__ void __launch_bounds__(64)
23 gemm_forward_4bit_cuda_m16nXk32(int G, int split_k_iters,
24 half* __restrict__ A, int* __restrict__ B,
25 half* __restrict__ scaling_factors,
26 int* __restrict__ zeros, int M, int IC,
27 int OC, half* __restrict__ C) {
28 // Only support matrix n = 64 or 128
29 assert(N == 64 || N == 128);
30 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750
31 assert(false);
32 #else
33 static constexpr uint32_t ZERO = 0x0;
34 float C_warp[32];
35 __shared__ half A_shared[16 * (32 + 8)];
36 __shared__ half B_shared[32 * (N + 8)];
37
38 int j_factors1 = ((OC + N - 1) / N);
39 int blockIdx_y = blockIdx.x % ((M + 16 - 1) / 16 * j_factors1);
40 int blockIdx_z = blockIdx.x / ((M + 16 - 1) / 16 * j_factors1);
41
42 half A_shared_warp[8];
43 half B_shared_warp[N / 4];
44 for (int j_0_4_init = 0; j_0_4_init < N / 32; ++j_0_4_init) {
45 for (int i = 0; i < 8; ++i) {
46 C_warp[(j_0_4_init * 8) + i] = 0.0;
47 }
48 }
49
50 static constexpr int row_stride_warp = 32 * 8 / 32;
51 static constexpr int row_stride = 2 * 32 * 8 / N;
52 // TODO: Haotian: blockIdx_y / j_factors1 in A loading to support bsz > 16
53 bool ld_A_flag =
54 (blockIdx_y / j_factors1 * 16 + threadIdx.y * row_stride_warp +
55 threadIdx.x * 8 / 32) < M; // threadIdx.y is warp_id
56 // bool wb_C_flag = (threadIdx.x / 4) < M;
57
58 half* A_ptr =
59 A +
60 (((int)blockIdx_y) / j_factors1 * 16 +
61 (((int)threadIdx.y) * row_stride_warp) + ((int)threadIdx.x) / (32 / 8)) *
62 IC +
63 (((int)threadIdx.x) % (32 / 8)) * 8;
64
65 int* B_ptr = B + ((int)threadIdx.y) * (OC / 8) * (256 / N) +
66 (((int)threadIdx.x) / (N / 8)) * (OC / 8) +
67 (((int)blockIdx_y) % j_factors1) * (N / 8) +
68 (((int)threadIdx.x) % (N / 8)) * 1;
69 // Why * 1 in the above line?
70
71 half* A_shared_ptr = A_shared +
72 ((int)threadIdx.y) * row_stride_warp * (32 + 8) +
73 (((int)threadIdx.x) / (32 / 8)) * (32 + 8) +
74 (((int)threadIdx.x) % (32 / 8)) * 8;
75
76 half* B_shared_ptr = B_shared +
77 ((int)threadIdx.y) * (row_stride / 2) * (N + 8) +
78 (((int)threadIdx.x) / (N / 8)) * (N + 8) +
79 (((int)threadIdx.x) % (N / 8)) * 8;
80
81 int* zeros_ptr = zeros + (((int)blockIdx_y) % j_factors1) * (N / 8) +
82 ((int)threadIdx.x) % (N / 8);
83
84 half* scaling_factors_ptr = scaling_factors +
85 (((int)blockIdx_y) % j_factors1) * N +
86 (((int)threadIdx.x) % (N / 8)) * 8;
87
88 half* C_ptr =
89 C +
90 static_cast<long long>(blockIdx_z) * M * OC // blockIdz.x -> split_k dim
91 + (((int)blockIdx_y) % j_factors1) * N + ((int)threadIdx.y) * (N / 2) +
92 (((int)threadIdx.x) % 4) * 2;
93
94 // preload s.f. and zeros
95 int k_bound = (IC / 32 + split_k_iters - 1) / split_k_iters;
96 if ((k_bound - 1) * split_k_iters * 32 + blockIdx_z * 32 >= IC) k_bound -= 1;
97 for (int _k_0_0 = 0; _k_0_0 < k_bound; ++_k_0_0) {
98 int k_0_0 = _k_0_0 * split_k_iters + blockIdx_z;
99 __syncthreads();
100 // TODO: Haotian: blockIdx_y / j_factors1 in A loading to support bsz > 16
101 if (ld_A_flag) {
102 *(uint4*)(A_shared_ptr) = *(uint4*)(A_ptr + (k_0_0 * 32));
103 } else {
104 *(uint4*)(A_shared_ptr) = make_uint4(0, 0, 0, 0);
105 }
106
107 // for (int ax0_ax1_fused_0 = 0; ax0_ax1_fused_0 < 2; ++ax0_ax1_fused_0) {
108 uint32_t zeros_loaded = *(uint32_t*)(zeros_ptr + k_0_0 * 32 / G * (OC / 8));
109 uint4 B_loaded_zero = dequantize_s4_to_fp16x2(zeros_loaded);
110 uint4 B_loaded_scale =
111 *(uint4*)(scaling_factors_ptr + k_0_0 * 32 / G * (OC));
112 /*
113 if (blockIdx_z == 0 && blockIdx_y == 0 && k_0_0 == 0 && threadIdx.x == 0 &&
114 threadIdx.y == 0){ printf("%x %x %x %x %x %x %x %x\n", B_loaded_scale.x,
115 B_loaded_scale.y, B_loaded_scale.z, B_loaded_scale.w, B_loaded_zero.x,
116 B_loaded_zero.y, B_loaded_zero.z, B_loaded_zero.w);
117 }
118 */
119 // uint4 B_loaded_scale = make_uint4(0, 0, 0, 0);
120 int* B_ptr_local = B_ptr + k_0_0 * 32 * (OC / 8);
121
122 for (int ax0_ax1_fused_0 = 0; ax0_ax1_fused_0 < N / 16; ++ax0_ax1_fused_0) {
123 // B: 32 x 136 (128+8) float16
124 // each warp: 32 x 4
125 // each thr: read 32 bit -> convert to 8xFP16 (a UINT4) -> scale and minus
126 // zero -> WB UINT4
127 // *(uint4*)(B_shared + ((((ax0_ax1_fused_0 * 544) + (((int)threadIdx.y) *
128 // 272)) + ((((int)threadIdx.x) >> 4) * 136)) + ((((int)threadIdx.x) & 15)
129 // * 8))) = *(uint4*)(B + ((((((k_0_0 * 163840) + (ax0_ax1_fused_0 *
130 // 20480)) + (((int)threadIdx.y) * 10240)) + ((((int)threadIdx.x) >> 4) *
131 // 5120)) + (((int)blockIdx_y) * 128)) + ((((int)threadIdx.x) & 15) *
132 // 8))); row stride in shared memory: (NWARPS * 32 * 8 / cta_N)
133 uint32_t B_loaded =
134 *(uint32_t*)(B_ptr_local + ax0_ax1_fused_0 * row_stride * (OC / 8));
135 uint4 B_loaded_fp16 = dequantize_s4_to_fp16x2(B_loaded);
136
137 // - zero and * scale
138 // TODO (Haotian): can save 4 assembly instructions if sormulate as deq =
139 // q * scale - zero * scale.
140 asm volatile("sub.f16x2 %0, %1, %2;\n"
141 : "=r"(B_loaded_fp16.x)
142 : "r"(B_loaded_fp16.x), "r"(B_loaded_zero.x));
143 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
144 : "=r"(B_loaded_fp16.x)
145 : "r"(B_loaded_fp16.x), "r"(B_loaded_scale.x), "r"(ZERO));
146 asm volatile("sub.f16x2 %0, %1, %2;\n"
147 : "=r"(B_loaded_fp16.y)
148 : "r"(B_loaded_fp16.y), "r"(B_loaded_zero.y));
149 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
150 : "=r"(B_loaded_fp16.y)
151 : "r"(B_loaded_fp16.y), "r"(B_loaded_scale.y), "r"(ZERO));
152 asm volatile("sub.f16x2 %0, %1, %2;\n"
153 : "=r"(B_loaded_fp16.z)
154 : "r"(B_loaded_fp16.z), "r"(B_loaded_zero.z));
155 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
156 : "=r"(B_loaded_fp16.z)
157 : "r"(B_loaded_fp16.z), "r"(B_loaded_scale.z), "r"(ZERO));
158 asm volatile("sub.f16x2 %0, %1, %2;\n"
159 : "=r"(B_loaded_fp16.w)
160 : "r"(B_loaded_fp16.w), "r"(B_loaded_zero.w));
161 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
162 : "=r"(B_loaded_fp16.w)
163 : "r"(B_loaded_fp16.w), "r"(B_loaded_scale.w), "r"(ZERO));
164 /*
165 if (ax0_ax1_fused_0 == 0 && blockIdx_z == 0 && blockIdx_y == 0 && k_0_0 ==
166 0 && threadIdx.x == 17 && threadIdx.y == 0){ printf("[x] %X %X %X %X\n",
167 B_loaded_fp16.x, B_loaded_fp16.y, B_loaded_fp16.z, B_loaded_fp16.w);
168 }
169 */
170
171 // write back
172 *(uint4*)(B_shared_ptr + ax0_ax1_fused_0 * row_stride * (N + 8)) =
173 B_loaded_fp16;
174 }
175 __syncthreads();
176
177 for (int k_0_1 = 0; k_0_1 < 2; ++k_0_1) {
178 {
179 unsigned int addr;
180 __asm__ __volatile__(
181 "{ .reg .u64 addr; cvta.to.shared.u64 addr, %1; cvt.u32.u64 %0, "
182 "addr; }\n"
183 : "=r"(addr)
184 : "l"((void*)((&(A_shared[(k_0_1 * 16)])) +
185 (((((int)threadIdx.x) & 15) * 40) +
186 ((((int)threadIdx.x) >> 4) * 8)))));
187
188 __asm__ __volatile__(
189 "ldmatrix.sync.aligned.m8n8.x4.shared.b16"
190 "{%0, %1, %2, %3}, [%4];\n"
191 : "=r"(((unsigned*)(A_shared_warp + 0))[0]),
192 "=r"(((unsigned*)(A_shared_warp + 0))[1]),
193 "=r"(((unsigned*)(A_shared_warp + 0))[2]),
194 "=r"(((unsigned*)(A_shared_warp + 0))[3])
195 : "r"(addr));
196 }
197
198 for (int ax1_0 = 0; ax1_0 < N / 32; ++ax1_0) {
199 {
200 unsigned int addr;
201 __asm__ __volatile__(
202 "{ .reg .u64 addr; cvta.to.shared.u64 addr, %1; cvt.u32.u64 %0, "
203 "addr; }\n"
204 : "=r"(addr)
205 : "l"((void*)((&(B_shared[(((k_0_1 * (N * 16 + 128)) +
206 (((int)threadIdx.y) * (N / 2))) +
207 (ax1_0 * 16))])) +
208 (((((int)threadIdx.x) & 15) * (N + 8)) +
209 ((((int)threadIdx.x) >> 4) * 8)))));
210 __asm__ __volatile__(
211 "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16"
212 "{%0, %1, %2, %3}, [%4];\n"
213 : "=r"(((unsigned*)(B_shared_warp + (ax1_0 * 8)))[0]),
214 "=r"(((unsigned*)(B_shared_warp + (ax1_0 * 8)))[1]),
215 "=r"(((unsigned*)(B_shared_warp + (ax1_0 * 8)))[2]),
216 "=r"(((unsigned*)(B_shared_warp + (ax1_0 * 8)))[3])
217 : "r"(addr));
218 }
219 }
220 for (int j_0_4 = 0; j_0_4 < N / 32; ++j_0_4) {
221 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
222 {
223 __asm__ __volatile__(
224 "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32"
225 "{%0, %1, %2, %3}, {%4, %5}, {%6}, {%7, %8, %9, %10};\n"
226 : "=f"(((float*)(C_warp + (j_0_4 * 8)))[0]),
227 "=f"(((float*)(C_warp + (j_0_4 * 8)))[1]),
228 "=f"(((float*)(C_warp + (j_0_4 * 8)))[2]),
229 "=f"(((float*)(C_warp + (j_0_4 * 8)))[3])
230 : "r"(((unsigned*)(A_shared_warp + 0))[0]),
231 "r"(((unsigned*)(A_shared_warp + 0))[1]),
232 "r"(((unsigned*)(B_shared_warp + (j_0_4 * 8)))[0]),
233 "f"(((float*)(C_warp + (j_0_4 * 8)))[0]),
234 "f"(((float*)(C_warp + (j_0_4 * 8)))[1]),
235 "f"(((float*)(C_warp + (j_0_4 * 8)))[2]),
236 "f"(((float*)(C_warp + (j_0_4 * 8)))[3]));
237 }
238
239 {
240 __asm__ __volatile__(
241 "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32"
242 "{%0, %1, %2, %3}, {%4, %5}, {%6}, {%7, %8, %9, %10};\n"
243 : "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[0]),
244 "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[1]),
245 "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[2]),
246 "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[3])
247 : "r"(((unsigned*)(A_shared_warp + 0))[0]),
248 "r"(((unsigned*)(A_shared_warp + 0))[1]),
249 "r"(((unsigned*)(B_shared_warp + ((j_0_4 * 8) + 4)))[0]),
250 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[0]),
251 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[1]),
252 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[2]),
253 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[3]));
254 }
255
256 {
257 __asm__ __volatile__(
258 "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32"
259 "{%0, %1, %2, %3}, {%4, %5}, {%6}, {%7, %8, %9, %10};\n"
260 : "=f"(((float*)(C_warp + (j_0_4 * 8)))[0]),
261 "=f"(((float*)(C_warp + (j_0_4 * 8)))[1]),
262 "=f"(((float*)(C_warp + (j_0_4 * 8)))[2]),
263 "=f"(((float*)(C_warp + (j_0_4 * 8)))[3])
264 : "r"(((unsigned*)(A_shared_warp + 0))[2]),
265 "r"(((unsigned*)(A_shared_warp + 0))[3]),
266 "r"(((unsigned*)(B_shared_warp + (j_0_4 * 8)))[1]),
267 "f"(((float*)(C_warp + (j_0_4 * 8)))[0]),
268 "f"(((float*)(C_warp + (j_0_4 * 8)))[1]),
269 "f"(((float*)(C_warp + (j_0_4 * 8)))[2]),
270 "f"(((float*)(C_warp + (j_0_4 * 8)))[3]));
271 }
272
273 {
274 __asm__ __volatile__(
275 "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32"
276 "{%0, %1, %2, %3}, {%4, %5}, {%6}, {%7, %8, %9, %10};\n"
277 : "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[0]),
278 "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[1]),
279 "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[2]),
280 "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[3])
281 : "r"(((unsigned*)(A_shared_warp + 0))[2]),
282 "r"(((unsigned*)(A_shared_warp + 0))[3]),
283 "r"(((unsigned*)(B_shared_warp + ((j_0_4 * 8) + 4)))[1]),
284 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[0]),
285 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[1]),
286 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[2]),
287 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[3]));
288 }
289 #else
290 {
291 __asm__ __volatile__(
292 "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32"
293 "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, "
294 "%13};\n"
295 : "=f"(((float*)(C_warp + (j_0_4 * 8)))[0]),
296 "=f"(((float*)(C_warp + (j_0_4 * 8)))[1]),
297 "=f"(((float*)(C_warp + (j_0_4 * 8)))[2]),
298 "=f"(((float*)(C_warp + (j_0_4 * 8)))[3])
299 : "r"(((unsigned*)(A_shared_warp + 0))[0]),
300 "r"(((unsigned*)(A_shared_warp + 0))[1]),
301 "r"(((unsigned*)(A_shared_warp + 0))[2]),
302 "r"(((unsigned*)(A_shared_warp + 0))[3]),
303 "r"(((unsigned*)(B_shared_warp + (j_0_4 * 8)))[0]),
304 "r"(((unsigned*)(B_shared_warp + (j_0_4 * 8)))[1]),
305 "f"(((float*)(C_warp + (j_0_4 * 8)))[0]),
306 "f"(((float*)(C_warp + (j_0_4 * 8)))[1]),
307 "f"(((float*)(C_warp + (j_0_4 * 8)))[2]),
308 "f"(((float*)(C_warp + (j_0_4 * 8)))[3]));
309 }
310
311 {
312 __asm__ __volatile__(
313 "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32"
314 "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, "
315 "%13};\n"
316 : "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[0]),
317 "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[1]),
318 "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[2]),
319 "=f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[3])
320 : "r"(((unsigned*)(A_shared_warp + 0))[0]),
321 "r"(((unsigned*)(A_shared_warp + 0))[1]),
322 "r"(((unsigned*)(A_shared_warp + 0))[2]),
323 "r"(((unsigned*)(A_shared_warp + 0))[3]),
324 "r"(((unsigned*)(B_shared_warp + ((j_0_4 * 8) + 4)))[0]),
325 "r"(((unsigned*)(B_shared_warp + ((j_0_4 * 8) + 4)))[1]),
326 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[0]),
327 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[1]),
328 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[2]),
329 "f"(((float*)(C_warp + ((j_0_4 * 8) + 4)))[3]));
330 }
331
332 #endif
333 }
334 }
335 }
336
337 // TODO: Shang: Hoist loop invariance.
338 for (int ax1_0_1 = 0; ax1_0_1 < (N / 32); ++ax1_0_1) {
339 for (int local_id = 0; local_id < 8; ++local_id) {
340 int row_offset = (((int)blockIdx_y) / j_factors1) * 16 +
341 ((int)threadIdx.x) / 4 + (local_id % 4) / 2 * 8;
342 if (row_offset < M) {
343 *(C_ptr + ax1_0_1 * 16 + row_offset * OC + (local_id / 4) * 8 +
344 local_id % 2) = __float2half(C_warp[(ax1_0_1 * 8) + local_id]);
345 }
346 }
347 }
348 #endif
349 }
350
351 __global__ void __launch_bounds__(64)
352 dequantize_weights(int* __restrict__ B, half* __restrict__ scaling_factors,
353 int* __restrict__ zeros, half* __restrict__ C, int G) {
354 static constexpr uint32_t ZERO = 0x0;
355 half B_shared[32 * (128 + 8)];
356
357 half* B_shared_ptr2 = B_shared;
358
359 int N = blockDim.x * gridDim.x; // 2
360 int col = (blockIdx.x * blockDim.x + threadIdx.x);
361 int row = blockIdx.y * blockDim.y + threadIdx.y;
362 int index1 = 8 * col + 8 * row * N;
363 half* C_ptr2 = C + index1;
364
365 int index2 = col + row * N;
366 int* B_ptr2 = B + index2;
367
368 int index3 = col + (int)(row / G) * N;
369 int* zeros_ptr2 = zeros + index3;
370 int index4 = 8 * col + (int)(row / G) * N * 8;
371 half* scaling_factors_ptr2 = scaling_factors + index4;
372
373 uint32_t zeros_loaded = *(uint32_t*)(zeros_ptr2);
374 uint4 B_loaded_zero = dequantize_s4_to_fp16x2(zeros_loaded);
375 uint4 B_loaded_scale = *(uint4*)(scaling_factors_ptr2);
376
377 uint32_t B_loaded = *(uint32_t*)B_ptr2;
378 uint4 B_loaded_fp16 = dequantize_s4_to_fp16x2(B_loaded);
379 asm volatile("sub.f16x2 %0, %1, %2;\n"
380 : "=r"(B_loaded_fp16.x)
381 : "r"(B_loaded_fp16.x), "r"(B_loaded_zero.x));
382 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
383 : "=r"(B_loaded_fp16.x)
384 : "r"(B_loaded_fp16.x), "r"(B_loaded_scale.x), "r"(ZERO));
385 asm volatile("sub.f16x2 %0, %1, %2;\n"
386 : "=r"(B_loaded_fp16.y)
387 : "r"(B_loaded_fp16.y), "r"(B_loaded_zero.y));
388 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
389 : "=r"(B_loaded_fp16.y)
390 : "r"(B_loaded_fp16.y), "r"(B_loaded_scale.y), "r"(ZERO));
391 asm volatile("sub.f16x2 %0, %1, %2;\n"
392 : "=r"(B_loaded_fp16.z)
393 : "r"(B_loaded_fp16.z), "r"(B_loaded_zero.z));
394 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
395 : "=r"(B_loaded_fp16.z)
396 : "r"(B_loaded_fp16.z), "r"(B_loaded_scale.z), "r"(ZERO));
397 asm volatile("sub.f16x2 %0, %1, %2;\n"
398 : "=r"(B_loaded_fp16.w)
399 : "r"(B_loaded_fp16.w), "r"(B_loaded_zero.w));
400 asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
401 : "=r"(B_loaded_fp16.w)
402 : "r"(B_loaded_fp16.w), "r"(B_loaded_scale.w), "r"(ZERO));
403
404 *(uint4*)B_shared_ptr2 = B_loaded_fp16;
405
406 for (int i = 0; i < 8; ++i) {
407 *(C_ptr2 + i) = B_shared[i];
408 }
409 }
410
411 } // namespace awq
412 } // namespace vllm
413
414 torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
415 torch::stable::Tensor _scaling_factors,
416 torch::stable::Tensor _zeros,
417 int64_t split_k_iters, int64_t thx,
418 int64_t thy) {
419 int in_c = _kernel.size(0);
420 int qout_c = _kernel.size(1);
421 int out_c = qout_c * 8;
422 int G = in_c / _scaling_factors.size(0);
423
424 int x_thread = thx;
425 int y_thread = thy;
426
427 int x_blocks = 1;
428 int y_blocks = 1;
429 if (thx == 0) {
430 x_thread = qout_c;
431 }
432 if (thy == 0) {
433 y_thread = in_c;
434 }
435 if (thx == 0 && thy == 0) {
436 x_thread = 8;
437 y_thread = 8;
438 x_blocks = (int)(qout_c / 8);
439 y_blocks = (int)(in_c / 8);
440 }
441
442 const torch::stable::accelerator::DeviceGuard device_guard(
443 _scaling_factors.get_device_index());
444
445 auto _de_kernel =
446 torch::stable::empty({in_c, out_c}, _scaling_factors.scalar_type(),
447 std::nullopt, _scaling_factors.device());
448
449 auto kernel = reinterpret_cast<int*>(_kernel.mutable_data_ptr<int>());
450 auto de_kernel = reinterpret_cast<half*>(
451 _de_kernel.mutable_data_ptr<torch::headeronly::Half>());
452 auto scaling_factors = reinterpret_cast<half*>(
453 _scaling_factors.mutable_data_ptr<torch::headeronly::Half>());
454 auto zeros = reinterpret_cast<int*>(_zeros.mutable_data_ptr<int>());
455
456 dim3 num_blocks(x_blocks, y_blocks);
457 dim3 threads_per_block(x_thread, y_thread);
458
459 const cudaStream_t stream = get_current_cuda_stream();
460 vllm::awq::dequantize_weights<<<num_blocks, threads_per_block, 0, stream>>>(
461 kernel, scaling_factors, zeros, de_kernel, G);
462
463 return _de_kernel;
464 }
465
466 // in_feats: M, IC [float16]
467 // kernel: IC, OC // 8 [int32] -> cast to IC, OC [uint4b]
468 // scaling_factors: IC // G, OC [float16]
469 // zeros: IC // G, OC // 8 [int32] -> cast to IC // G, OC [uint4b]
470 // assume that batch_size < 16 for now
471
472 torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
473 torch::stable::Tensor _kernel,
474 torch::stable::Tensor _scaling_factors,
475 torch::stable::Tensor _zeros,
476 int64_t split_k_iters) {
477 int num_in_feats = _in_feats.size(0);
478 int num_in_channels = _in_feats.size(1);
479 const torch::stable::accelerator::DeviceGuard device_guard(
480 _in_feats.get_device_index());
481
482 auto _out_feats = torch::stable::empty(
483 {split_k_iters, num_in_feats, _kernel.size(1) * 8},
484 _in_feats.scalar_type(), std::nullopt, _in_feats.device());
485 int num_out_feats = _out_feats.size(-2);
486 int num_out_channels = _out_feats.size(-1);
487
488 auto in_feats = reinterpret_cast<half*>(
489 _in_feats.mutable_data_ptr<torch::headeronly::Half>());
490 auto kernel = reinterpret_cast<int*>(_kernel.mutable_data_ptr<int>());
491 auto out_feats = reinterpret_cast<half*>(
492 _out_feats.mutable_data_ptr<torch::headeronly::Half>());
493 auto scaling_factors = reinterpret_cast<half*>(
494 _scaling_factors.mutable_data_ptr<torch::headeronly::Half>());
495 auto zeros = reinterpret_cast<int*>(_zeros.mutable_data_ptr<int>());
496 int group_size = num_in_channels / _scaling_factors.size(0);
497
498 if (num_out_channels % 64 != 0)
499 throw std::invalid_argument("OC is not multiple of cta_N = 64");
500 if (num_out_channels % 8 != 0)
501 throw std::invalid_argument("OC is not multiple of pack_num = 8");
502 if (group_size % 32 != 0)
503 throw std::invalid_argument("Group size should be a multiple of 32");
504 if (num_out_channels % group_size != 0)
505 throw std::invalid_argument("OC is not multiple of Group size");
506
507 const cudaStream_t stream = get_current_cuda_stream();
508 if (num_out_channels % 128 == 0) {
509 int j_factors1 = num_out_channels / 128 / 1;
510 dim3 num_blocks((num_out_feats + 16 - 1) / 16 * j_factors1 * split_k_iters);
511 // threadIdx.x: 32
512 // threadIdx.y: i_factors[2] * j_factors[2]
513 dim3 threads_per_block(32, 2);
514 vllm::awq::gemm_forward_4bit_cuda_m16nXk32<128>
515 <<<num_blocks, threads_per_block, 0, stream>>>(
516 group_size, split_k_iters, in_feats, kernel, scaling_factors, zeros,
517 num_in_feats, num_in_channels, num_out_channels, out_feats);
518 } else if (num_out_channels % 64 == 0) {
519 int j_factors1 = num_out_channels / 64 / 1;
520 dim3 num_blocks(1 * (num_out_feats + 16 - 1) / 16 * j_factors1 *
521 split_k_iters);
522
523 // threadIdx.x: 32
524 // threadIdx.y: i_factors[2] * j_factors[2]
525 dim3 threads_per_block(32, 2);
526 vllm::awq::gemm_forward_4bit_cuda_m16nXk32<64>
527 <<<num_blocks, threads_per_block, 0, stream>>>(
528 group_size, split_k_iters, in_feats, kernel, scaling_factors, zeros,
529 num_in_feats, num_in_channels, num_out_channels, out_feats);
530 }
531 // Pass the reduction dim via a named, lifetime-stable local rather than a
532 // bare `0`.
533 constexpr int64_t sum_dim = 0;
534 return torch::stable::sum(
535 _out_feats, torch::headeronly::IntHeaderOnlyArrayRef(&sum_dim, 1));
536 }