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, 11:00 UTC
[Kernel] GGUF MoeVec kernel (#16780)
Signed-off-by: SzymonOzog <szymon.ozog@aleph-alpha.com> Signed-off-by: SzymonOzog <szymon.ozog@gmail.com> Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <2037008807@qq.com>
Szymon Ożóg committed
May 7, 2025 at 14:07 UTC
1a45a61387be295c969c620458b7ef4a13250ed2
8 files changed
+544
-16
csrc/ops.h
+4
@@ -178,6 +178,10 @@ torch::Tensor ggml_moe_a8(torch::Tensor X, torch::Tensor W,
178
torch::Tensor num_tokens_post_padded, int64_t type,
179
int64_t row, int64_t top_k, int64_t tokens);
180
181
+torch::Tensor ggml_moe_a8_vec(torch::Tensor X, torch::Tensor W,
182
+ torch::Tensor topk_ids, int64_t top_k,
183
+ int64_t type, int64_t row, int64_t tokens);
184
+
185
int64_t ggml_moe_get_block_size(int64_t type);
186
187
#ifndef USE_ROCM
csrc/quantization/gguf/gguf_kernel.cu
+137
@@ -13,6 +13,7 @@
13
#include "mmvq.cuh"
14
#include "mmq.cuh"
15
#include "moe.cuh"
16
+#include "moe_vec.cuh"
17
18
// Q8 gemv
19
template <typename scalar_t>
@@ -377,6 +378,142 @@ torch::Tensor ggml_moe_a8(torch::Tensor X, // input
378
return Y;
379
}
380
381
+torch::Tensor ggml_moe_a8_vec(torch::Tensor X, // input
382
+ torch::Tensor W, // expert weights
383
+ torch::Tensor topk_ids, int64_t top_k,
384
+ int64_t type, int64_t row, int64_t tokens) {
385
+ int col = X.sizes()[1];
386
+ const int padded = (col + 512 - 1) / 512 * 512;
387
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
388
+ auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
389
+ at::Tensor Y = torch::zeros({tokens * top_k, row}, options);
390
+ cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
391
+ options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
392
+ at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options);
393
+ VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_vec_a8", [&] {
394
+ quantize_row_q8_1_cuda<scalar_t>((scalar_t*)X.data_ptr(),
395
+ (void*)quant_X.data_ptr(), col, tokens,
396
+ stream);
397
+ switch (type) {
398
+ case 2:
399
+ moe_vec_q4_0_q8_1_cuda<scalar_t>(
400
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
401
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
402
+ col, row, quant_X.stride(0), stream);
403
+ break;
404
+ case 3:
405
+ moe_vec_q4_1_q8_1_cuda<scalar_t>(
406
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
407
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
408
+ col, row, quant_X.stride(0), stream);
409
+ break;
410
+ case 6:
411
+ moe_vec_q5_0_q8_1_cuda<scalar_t>(
412
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
413
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
414
+ col, row, quant_X.stride(0), stream);
415
+ break;
416
+ case 7:
417
+ moe_vec_q5_1_q8_1_cuda<scalar_t>(
418
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
419
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
420
+ col, row, quant_X.stride(0), stream);
421
+ break;
422
+ case 8:
423
+ moe_vec_q8_0_q8_1_cuda<scalar_t>(
424
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
425
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
426
+ col, row, quant_X.stride(0), stream);
427
+ break;
428
+ case 10:
429
+ moe_vec_q2_K_q8_1_cuda<scalar_t>(
430
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
431
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
432
+ col, row, quant_X.stride(0), stream);
433
+ break;
434
+ case 11:
435
+ moe_vec_q3_K_q8_1_cuda<scalar_t>(
436
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
437
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
438
+ col, row, quant_X.stride(0), stream);
439
+ break;
440
+ case 12:
441
+ moe_vec_q4_K_q8_1_cuda<scalar_t>(
442
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
443
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
444
+ col, row, quant_X.stride(0), stream);
445
+ break;
446
+ case 13:
447
+ moe_vec_q5_K_q8_1_cuda<scalar_t>(
448
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
449
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
450
+ col, row, quant_X.stride(0), stream);
451
+ break;
452
+ case 14:
453
+ moe_vec_q6_K_q8_1_cuda<scalar_t>(
454
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
455
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
456
+ col, row, quant_X.stride(0), stream);
457
+ break;
458
+ case 16:
459
+ moe_vec_iq2_xxs_q8_1_cuda<scalar_t>(
460
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
461
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
462
+ col, row, quant_X.stride(0), stream);
463
+ break;
464
+ case 17:
465
+ moe_vec_iq2_xs_q8_1_cuda<scalar_t>(
466
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
467
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
468
+ col, row, quant_X.stride(0), stream);
469
+ break;
470
+ case 18:
471
+ moe_vec_iq3_xxs_q8_1_cuda<scalar_t>(
472
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
473
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
474
+ col, row, quant_X.stride(0), stream);
475
+ break;
476
+ case 19:
477
+ moe_vec_iq1_s_q8_1_cuda<scalar_t>(
478
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
479
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
480
+ col, row, quant_X.stride(0), stream);
481
+ break;
482
+ case 20:
483
+ moe_vec_iq4_nl_q8_1_cuda<scalar_t>(
484
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
485
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
486
+ col, row, quant_X.stride(0), stream);
487
+ break;
488
+ case 21:
489
+ moe_vec_iq3_s_q8_1_cuda<scalar_t>(
490
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
491
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
492
+ col, row, quant_X.stride(0), stream);
493
+ break;
494
+ case 22:
495
+ moe_vec_iq2_s_q8_1_cuda<scalar_t>(
496
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
497
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
498
+ col, row, quant_X.stride(0), stream);
499
+ break;
500
+ case 23:
501
+ moe_vec_iq4_xs_q8_1_cuda<scalar_t>(
502
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
503
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
504
+ col, row, quant_X.stride(0), stream);
505
+ break;
506
+ case 29:
507
+ moe_vec_iq1_m_q8_1_cuda<scalar_t>(
508
+ (void*)W.data_ptr(), (void*)quant_X.data_ptr(),
509
+ (scalar_t*)Y.data_ptr(), (int*)topk_ids.data_ptr(), top_k, tokens,
510
+ col, row, quant_X.stride(0), stream);
511
+ break;
512
+ }
513
+ });
514
+ return Y;
515
+}
516
+
517
int64_t ggml_moe_get_block_size(int64_t type) {
518
switch (type) {
519
case 2:
csrc/quantization/gguf/moe_vec.cuh
new
+338
@@ -0,0 +1,338 @@
1
+// copied and adapted from
2
+// https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu
3
+template <typename scalar_t, int qk, int qi, typename block_q_t, int vdr,
4
+ vec_dot_q_cuda_t vec_dot_q_cuda>
5
+static __global__ void moe_vec_q(const void* __restrict__ vx,
6
+ const void* __restrict__ vy,
7
+ scalar_t* __restrict__ dst,
8
+ const int* topk_ids, const int topk,
9
+ const int ncols, const int nrows,
10
+ const int token_stride) {
11
+ const auto row = blockIdx.x * blockDim.y + threadIdx.y;
12
+
13
+ const auto token = blockIdx.z / topk;
14
+ const auto expert = (topk_ids)[blockIdx.z];
15
+
16
+ if (row >= nrows) {
17
+ return;
18
+ }
19
+
20
+ const int blocks_per_row = ncols / qk;
21
+ const int blocks_per_warp = vdr * WARP_SIZE / qi;
22
+
23
+ // partial sum for each thread
24
+ float tmp = 0.0f;
25
+
26
+ const block_q_t* x = ((const block_q_t*)vx) + expert * nrows * blocks_per_row;
27
+ const block_q8_1* y =
28
+ (const block_q8_1*)(((const int*)vy) + token * token_stride);
29
+
30
+ for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row;
31
+ i += blocks_per_warp) {
32
+ const int ibx = row * blocks_per_row + i; // x block index
33
+
34
+ const int iby = i * (qk / QK8_1); // y block index that aligns with ibx
35
+
36
+ const int iqs =
37
+ vdr *
38
+ (threadIdx.x %
39
+ (qi / vdr)); // x block quant index when casting the quants to int
40
+
41
+ tmp += vec_dot_q_cuda(&x[ibx], &y[iby], iqs);
42
+ }
43
+
44
+ // sum up partial sums and write back result
45
+#pragma unroll
46
+ for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) {
47
+ tmp += VLLM_SHFL_XOR_SYNC(tmp, mask);
48
+ }
49
+
50
+ if (threadIdx.x == 0) {
51
+ dst[blockIdx.z * nrows + row] = tmp;
52
+ }
53
+}
54
+
55
+template <typename scalar_t>
56
+static void moe_vec_q4_0_q8_1_cuda(const void* vx, const void* vy,
57
+ scalar_t* dst, const int* topk_ids,
58
+ const int top_k, const int tokens,
59
+ const int ncols, const int nrows,
60
+ const int token_stride,
61
+ cudaStream_t stream) {
62
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
63
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
64
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
65
+ moe_vec_q<scalar_t, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ,
66
+ vec_dot_q4_0_q8_1><<<block_nums, block_dims, 0, stream>>>(
67
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
68
+}
69
+
70
+template <typename scalar_t>
71
+static void moe_vec_q4_1_q8_1_cuda(const void* vx, const void* vy,
72
+ scalar_t* dst, const int* topk_ids,
73
+ const int top_k, const int tokens,
74
+ const int ncols, const int nrows,
75
+ const int token_stride,
76
+ cudaStream_t stream) {
77
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
78
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
79
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
80
+ moe_vec_q<scalar_t, QK4_0, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ,
81
+ vec_dot_q4_1_q8_1><<<block_nums, block_dims, 0, stream>>>(
82
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
83
+}
84
+
85
+template <typename scalar_t>
86
+static void moe_vec_q5_0_q8_1_cuda(const void* vx, const void* vy,
87
+ scalar_t* dst, const int* topk_ids,
88
+ const int top_k, const int tokens,
89
+ const int ncols, const int nrows,
90
+ const int token_stride,
91
+ cudaStream_t stream) {
92
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
93
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
94
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
95
+ moe_vec_q<scalar_t, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ,
96
+ vec_dot_q5_0_q8_1><<<block_nums, block_dims, 0, stream>>>(
97
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
98
+}
99
+
100
+template <typename scalar_t>
101
+static void moe_vec_q5_1_q8_1_cuda(const void* vx, const void* vy,
102
+ scalar_t* dst, const int* topk_ids,
103
+ const int top_k, const int tokens,
104
+ const int ncols, const int nrows,
105
+ const int token_stride,
106
+ cudaStream_t stream) {
107
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
108
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
109
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
110
+ moe_vec_q<scalar_t, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ,
111
+ vec_dot_q5_1_q8_1><<<block_nums, block_dims, 0, stream>>>(
112
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
113
+}
114
+
115
+template <typename scalar_t>
116
+static void moe_vec_q8_0_q8_1_cuda(const void* vx, const void* vy,
117
+ scalar_t* dst, const int* topk_ids,
118
+ const int top_k, const int tokens,
119
+ const int ncols, const int nrows,
120
+ const int token_stride,
121
+ cudaStream_t stream) {
122
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
123
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
124
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
125
+ moe_vec_q<scalar_t, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ,
126
+ vec_dot_q8_0_q8_1><<<block_nums, block_dims, 0, stream>>>(
127
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
128
+}
129
+
130
+template <typename scalar_t>
131
+static void moe_vec_q2_K_q8_1_cuda(const void* vx, const void* vy,
132
+ scalar_t* dst, const int* topk_ids,
133
+ const int top_k, const int tokens,
134
+ const int ncols, const int nrows,
135
+ const int token_stride,
136
+ cudaStream_t stream) {
137
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
138
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
139
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
140
+ moe_vec_q<scalar_t, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ,
141
+ vec_dot_q2_K_q8_1><<<block_nums, block_dims, 0, stream>>>(
142
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
143
+}
144
+
145
+template <typename scalar_t>
146
+static void moe_vec_q3_K_q8_1_cuda(const void* vx, const void* vy,
147
+ scalar_t* dst, const int* topk_ids,
148
+ const int top_k, const int tokens,
149
+ const int ncols, const int nrows,
150
+ const int token_stride,
151
+ cudaStream_t stream) {
152
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
153
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
154
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
155
+ moe_vec_q<scalar_t, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ,
156
+ vec_dot_q3_K_q8_1><<<block_nums, block_dims, 0, stream>>>(
157
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
158
+}
159
+
160
+template <typename scalar_t>
161
+static void moe_vec_q4_K_q8_1_cuda(const void* vx, const void* vy,
162
+ scalar_t* dst, const int* topk_ids,
163
+ const int top_k, const int tokens,
164
+ const int ncols, const int nrows,
165
+ const int token_stride,
166
+ cudaStream_t stream) {
167
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
168
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
169
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
170
+ moe_vec_q<scalar_t, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ,
171
+ vec_dot_q4_K_q8_1><<<block_nums, block_dims, 0, stream>>>(
172
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
173
+}
174
+
175
+template <typename scalar_t>
176
+static void moe_vec_q5_K_q8_1_cuda(const void* vx, const void* vy,
177
+ scalar_t* dst, const int* topk_ids,
178
+ const int top_k, const int tokens,
179
+ const int ncols, const int nrows,
180
+ const int token_stride,
181
+ cudaStream_t stream) {
182
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
183
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
184
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
185
+ moe_vec_q<scalar_t, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ,
186
+ vec_dot_q5_K_q8_1><<<block_nums, block_dims, 0, stream>>>(
187
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
188
+}
189
+
190
+template <typename scalar_t>
191
+static void moe_vec_q6_K_q8_1_cuda(const void* vx, const void* vy,
192
+ scalar_t* dst, const int* topk_ids,
193
+ const int top_k, const int tokens,
194
+ const int ncols, const int nrows,
195
+ const int token_stride,
196
+ cudaStream_t stream) {
197
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
198
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
199
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
200
+ moe_vec_q<scalar_t, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ,
201
+ vec_dot_q6_K_q8_1><<<block_nums, block_dims, 0, stream>>>(
202
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
203
+}
204
+
205
+template <typename scalar_t>
206
+static void moe_vec_iq2_xxs_q8_1_cuda(const void* vx, const void* vy,
207
+ scalar_t* dst, const int* topk_ids,
208
+ const int top_k, const int tokens,
209
+ const int ncols, const int nrows,
210
+ const int token_stride,
211
+ cudaStream_t stream) {
212
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
213
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
214
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
215
+ moe_vec_q<scalar_t, QK_K, QI2_XXS, block_iq2_xxs, 1, vec_dot_iq2_xxs_q8_1>
216
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k,
217
+ ncols, nrows, token_stride);
218
+}
219
+
220
+template <typename scalar_t>
221
+static void moe_vec_iq2_xs_q8_1_cuda(const void* vx, const void* vy,
222
+ scalar_t* dst, const int* topk_ids,
223
+ const int top_k, const int tokens,
224
+ const int ncols, const int nrows,
225
+ const int token_stride,
226
+ cudaStream_t stream) {
227
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
228
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
229
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
230
+ moe_vec_q<scalar_t, QK_K, QI2_XS, block_iq2_xs, 1, vec_dot_iq2_xs_q8_1>
231
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k,
232
+ ncols, nrows, token_stride);
233
+}
234
+
235
+template <typename scalar_t>
236
+static void moe_vec_iq2_s_q8_1_cuda(const void* vx, const void* vy,
237
+ scalar_t* dst, const int* topk_ids,
238
+ const int top_k, const int tokens,
239
+ const int ncols, const int nrows,
240
+ const int token_stride,
241
+ cudaStream_t stream) {
242
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
243
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
244
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
245
+ moe_vec_q<scalar_t, QK_K, QI2_S, block_iq2_s, 1, vec_dot_iq2_s_q8_1>
246
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k,
247
+ ncols, nrows, token_stride);
248
+}
249
+
250
+template <typename scalar_t>
251
+static void moe_vec_iq3_xxs_q8_1_cuda(const void* vx, const void* vy,
252
+ scalar_t* dst, const int* topk_ids,
253
+ const int top_k, const int tokens,
254
+ const int ncols, const int nrows,
255
+ const int token_stride,
256
+ cudaStream_t stream) {
257
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
258
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
259
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
260
+ moe_vec_q<scalar_t, QK_K, QI3_XXS, block_iq3_xxs, 1, vec_dot_iq3_xxs_q8_1>
261
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k,
262
+ ncols, nrows, token_stride);
263
+}
264
+
265
+template <typename scalar_t>
266
+static void moe_vec_iq1_s_q8_1_cuda(const void* vx, const void* vy,
267
+ scalar_t* dst, const int* topk_ids,
268
+ const int top_k, const int tokens,
269
+ const int ncols, const int nrows,
270
+ const int token_stride,
271
+ cudaStream_t stream) {
272
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
273
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
274
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
275
+ moe_vec_q<scalar_t, QK_K, QI1_S, block_iq1_s, 1, vec_dot_iq1_s_q8_1>
276
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k,
277
+ ncols, nrows, token_stride);
278
+}
279
+
280
+template <typename scalar_t>
281
+static void moe_vec_iq1_m_q8_1_cuda(const void* vx, const void* vy,
282
+ scalar_t* dst, const int* topk_ids,
283
+ const int top_k, const int tokens,
284
+ const int ncols, const int nrows,
285
+ const int token_stride,
286
+ cudaStream_t stream) {
287
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
288
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
289
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
290
+ moe_vec_q<scalar_t, QK_K, QI1_M, block_iq1_m, 1, vec_dot_iq1_m_q8_1>
291
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k,
292
+ ncols, nrows, token_stride);
293
+}
294
+
295
+template <typename scalar_t>
296
+static void moe_vec_iq4_nl_q8_1_cuda(const void* vx, const void* vy,
297
+ scalar_t* dst, const int* topk_ids,
298
+ const int top_k, const int tokens,
299
+ const int ncols, const int nrows,
300
+ const int token_stride,
301
+ cudaStream_t stream) {
302
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
303
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
304
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
305
+ moe_vec_q<scalar_t, QK4_NL, QI4_NL, block_iq4_nl, VDR_Q4_0_Q8_1_MMVQ,
306
+ vec_dot_iq4_nl_q8_1><<<block_nums, block_dims, 0, stream>>>(
307
+ vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
308
+}
309
+
310
+template <typename scalar_t>
311
+static void moe_vec_iq4_xs_q8_1_cuda(const void* vx, const void* vy,
312
+ scalar_t* dst, const int* topk_ids,
313
+ const int top_k, const int tokens,
314
+ const int ncols, const int nrows,
315
+ const int token_stride,
316
+ cudaStream_t stream) {
317
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
318
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
319
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
320
+ moe_vec_q<scalar_t, QK_K, QI4_XS, block_iq4_xs, 1, vec_dot_iq4_xs_q8_1>
321
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k,
322
+ ncols, nrows, token_stride);
323
+}
324
+
325
+template <typename scalar_t>
326
+static void moe_vec_iq3_s_q8_1_cuda(const void* vx, const void* vy,
327
+ scalar_t* dst, const int* topk_ids,
328
+ const int top_k, const int tokens,
329
+ const int ncols, const int nrows,
330
+ const int token_stride,
331
+ cudaStream_t stream) {
332
+ const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
333
+ const dim3 block_nums(block_num_y, 1, tokens * top_k);
334
+ const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
335
+ moe_vec_q<scalar_t, QK_K, QI3_XS, block_iq3_s, 1, vec_dot_iq3_s_q8_1>
336
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k,
337
+ ncols, nrows, token_stride);
338
+}
csrc/torch_bindings.cpp
+6
@@ -337,6 +337,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
337
"int type, SymInt row, SymInt top_k, SymInt tokens) -> Tensor");
338
ops.impl("ggml_moe_a8", torch::kCUDA, &ggml_moe_a8);
339
340
+ ops.def(
341
+ "ggml_moe_a8_vec(Tensor X, Tensor W, "
342
+ "Tensor topk_ids, int top_k, "
343
+ "int type, SymInt row, SymInt tokens) -> Tensor");
344
+ ops.impl("ggml_moe_a8_vec", torch::kCUDA, &ggml_moe_a8_vec);
345
+
346
ops.def("ggml_moe_get_block_size", &ggml_moe_get_block_size);
347
348
#ifndef USE_ROCM
tests/kernels/quantization/test_ggml.py
+6
@@ -36,3 +36,9 @@ def test_ggml_opcheck(quant_type):
36
opcheck(torch.ops._C.ggml_moe_a8,
37
(x, qweight, sorted_token_ids, expert_ids, num_tokens_post_padded,
38
quant_type, qweight.shape[0], 1, x.shape[0]))
39
+
40
+ topk_ids = torch.zeros((1, 1), device='cuda', dtype=torch.int32)
41
+
42
+ opcheck(
43
+ torch.ops._C.ggml_moe_a8_vec,
44
+ (x, qweight, topk_ids, 1, quant_type, qweight.shape[0], x.shape[0]))
tests/kernels/quantization/test_gguf.py
+5
-15
@@ -151,20 +151,7 @@ def test_mmq(num_tokens: int, hidden_size: int, dtype: torch.dtype,
151
@pytest.mark.parametrize("hidden_size", [512])
152
@pytest.mark.parametrize("top_k", [4, 8])
153
@pytest.mark.parametrize("dtype", DTYPES)
154
-@pytest.mark.parametrize(
155
- "quant_type",
156
- [
157
- # k-quants
158
- GGMLQuantizationType.Q2_K,
159
- GGMLQuantizationType.Q3_K,
160
- GGMLQuantizationType.Q4_K,
161
- GGMLQuantizationType.Q5_K,
162
- GGMLQuantizationType.Q6_K,
163
- # standard quants
164
- GGMLQuantizationType.Q4_0,
165
- GGMLQuantizationType.Q5_0,
166
- GGMLQuantizationType.Q8_0,
167
- ])
154
+@pytest.mark.parametrize("quant_type", QUANT_TYPES)
155
@torch.inference_mode()
156
def test_moe(num_tokens: int, hidden_size: int, dtype: torch.dtype,
157
quant_type: GGMLQuantizationType, top_k: int):
@@ -174,7 +161,10 @@ def test_moe(num_tokens: int, hidden_size: int, dtype: torch.dtype,
161
x = torch.rand((num_tokens, H), dtype=dtype, device="cuda")
162
163
topk_weights = torch.rand(num_tokens, top_k, device="cuda", dtype=dtype)
177
- topk_ids = torch.randint(0, E, (num_tokens, top_k), device="cuda")
164
+ topk_ids = torch.randint(0,
165
+ E, (num_tokens, top_k),
166
+ device="cuda",
167
+ dtype=torch.int32)
168
169
tensors = get_gguf_MoE_tensors(hidden_size, quant_type)
170
vllm/_custom_ops.py
+31
@@ -497,6 +497,24 @@ if hasattr(torch.ops._C, "ggml_dequantize"):
497
device=W.device)
498
499
500
+if hasattr(torch.ops._C, "ggml_moe_a8_vec"):
501
+
502
+ @register_fake("_C::ggml_moe_a8_vec")
503
+ def _ggml_moe_a8_vec_fake(
504
+ X: torch.Tensor,
505
+ W: torch.Tensor,
506
+ topk_ids: torch.Tensor,
507
+ top_k: int,
508
+ quant_type: int,
509
+ row: torch.SymInt,
510
+ tokens: torch.SymInt,
511
+ ) -> torch.Tensor:
512
+ tokens = X.size(0)
513
+ return torch.empty((tokens * top_k, row),
514
+ dtype=X.dtype,
515
+ device=W.device)
516
+
517
+
518
# cutlass
519
def cutlass_scaled_mm_supports_fp4(cuda_device_capability: int) -> bool:
520
return torch.ops._C.cutlass_scaled_mm_supports_fp4(cuda_device_capability)
@@ -1146,6 +1164,19 @@ def ggml_moe_a8(
1164
top_k, tokens)
1165
1166
1167
+def ggml_moe_a8_vec(
1168
+ X: torch.Tensor,
1169
+ W: torch.Tensor,
1170
+ topk_ids: torch.Tensor,
1171
+ top_k: int,
1172
+ quant_type: int,
1173
+ row: torch.SymInt,
1174
+ tokens: torch.SymInt,
1175
+) -> torch.Tensor:
1176
+ return torch.ops._C.ggml_moe_a8_vec(X, W, topk_ids, top_k, quant_type, row,
1177
+ tokens)
1178
+
1179
+
1180
def ggml_moe_get_block_size(quant_type: int) -> int:
1181
return torch.ops._C.ggml_moe_get_block_size(quant_type)
1182
vllm/model_executor/layers/quantization/gguf.py
+17
-1
@@ -145,7 +145,9 @@ def _fused_moe_gguf(
145
moe_align_block_size)
146
147
out_hidden_states = torch.empty_like(x)
148
- if qweight_type2 in MMQ_QUANT_TYPES and qweight_type in MMQ_QUANT_TYPES:
148
+ # unless we decent expert reuse we are better off running moe_vec kernel
149
+ if (qweight_type2 in MMQ_QUANT_TYPES and qweight_type in MMQ_QUANT_TYPES
150
+ and x.shape[0] > 64):
151
num_tokens, _ = x.shape
152
E, N, _ = w1.shape
153
top_k = topk_ids.shape[1]
@@ -163,6 +165,20 @@ def _fused_moe_gguf(
165
out = out.reshape(num_tokens, top_k, w2.shape[1]).mul_(
166
topk_weights.view(num_tokens, top_k, 1))
167
ops.moe_sum(out, out_hidden_states)
168
+ elif qweight_type2 in MMVQ_QUANT_TYPES and qweight_type in MMVQ_QUANT_TYPES:
169
+ num_tokens, _ = x.shape
170
+ E, N, _ = w1.shape
171
+ top_k = topk_ids.shape[1]
172
+
173
+ out = ops.ggml_moe_a8_vec(x, w1, topk_ids, top_k, qweight_type, N,
174
+ num_tokens)
175
+ out = act(out)
176
+
177
+ out = ops.ggml_moe_a8_vec(out, w2, topk_ids, 1, qweight_type2,
178
+ w2.shape[1], num_tokens * top_k)
179
+ out = out.reshape(num_tokens, top_k, w2.shape[1]).mul_(
180
+ topk_weights.view(num_tokens, top_k, 1))
181
+ ops.moe_sum(out, out_hidden_states)
182
else:
183
logger.warning_once("There is no support for fast MoE kernel "
184
"for current quantization method. "