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, 19:45 UTC
[Kernel] GGUF MoE kernel (#14613)
Signed-off-by: SzymonOzog <szymon.ozog@aleph-alpha.com>
Szymon Ożóg committed
Mar 12, 2025 at 04:33 UTC
e22ee1e7a269853e422a83e8e646fae22c414215
8 files changed
+1070
-25
csrc/ops.h
+8
@@ -151,6 +151,14 @@ torch::Tensor ggml_mul_mat_vec_a8(torch::Tensor W, torch::Tensor X,
151
torch::Tensor ggml_mul_mat_a8(torch::Tensor W, torch::Tensor X, int64_t type,
152
int64_t row);
153
154
+torch::Tensor ggml_moe_a8(torch::Tensor X, torch::Tensor W,
155
+ torch::Tensor sorted_token_ids,
156
+ torch::Tensor expert_ids,
157
+ torch::Tensor num_tokens_post_padded, int64_t type,
158
+ int64_t row, int64_t top_k, int64_t tokens);
159
+
160
+int64_t ggml_moe_get_block_size(int64_t type);
161
+
162
#ifndef USE_ROCM
163
void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
164
torch::Tensor const& B, torch::Tensor const& A_sf,
csrc/quantization/gguf/gguf_kernel.cu
+138
-4
@@ -12,6 +12,7 @@
12
#include "dequantize.cuh"
13
#include "mmvq.cuh"
14
#include "mmq.cuh"
15
+#include "moe.cuh"
16
17
// Q8 gemv
18
template <typename scalar_t>
@@ -59,10 +60,14 @@ static void quantize_row_q8_1_cuda(const scalar_t* x, void* vy, const int kx,
60
const int64_t kx_padded = (kx + 512 - 1) / 512 * 512;
61
const int block_num_x =
62
(kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE;
62
- const dim3 num_blocks(block_num_x, ky, 1);
63
- const dim3 block_size(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1);
64
- quantize_q8_1<scalar_t>
65
- <<<num_blocks, block_size, 0, stream>>>(x, vy, kx, kx_padded);
63
+ constexpr int MAX_BLOCK_SIZE = 65535;
64
+ for (int off = 0; off < ky; off += MAX_BLOCK_SIZE) {
65
+ const int num_blocks_y = std::min(ky, off + MAX_BLOCK_SIZE) - off;
66
+ const dim3 num_blocks(block_num_x, num_blocks_y, 1);
67
+ const dim3 block_size(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1);
68
+ quantize_q8_1<<<num_blocks, block_size, 0, stream>>>(
69
+ &x[off * kx], (int32_t*)vy + off * (kx_padded / 32 * 9), kx, kx_padded);
70
+ }
71
}
72
73
torch::Tensor ggml_dequantize(torch::Tensor W, // quant weight
@@ -263,3 +268,132 @@ torch::Tensor ggml_mul_mat_a8(torch::Tensor W, // quant weight
268
});
269
return Y;
270
}
271
+
272
+torch::Tensor ggml_moe_a8(torch::Tensor X, // input
273
+ torch::Tensor W, // expert weights
274
+ torch::Tensor sorted_token_ids,
275
+ torch::Tensor expert_ids,
276
+ torch::Tensor num_tokens_post_padded, int64_t type,
277
+ int64_t row, int64_t top_k, int64_t tokens) {
278
+ int col = X.sizes()[1];
279
+ int padded = (col + 512 - 1) / 512 * 512;
280
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
281
+ auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
282
+ at::Tensor Y = torch::empty({tokens * top_k, row}, options);
283
+ cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
284
+ options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
285
+ at::Tensor quant_X = torch::empty({tokens, padded / 32 * 9}, options);
286
+ VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_moe_a8", [&] {
287
+ quantize_row_q8_1_cuda((scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(),
288
+ col, tokens, stream);
289
+ switch (type) {
290
+ case 2:
291
+ ggml_moe_q4_0_q8_1_cuda(
292
+ (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
293
+ (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
294
+ (int*)expert_ids.data_ptr(),
295
+ (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
296
+ tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
297
+ break;
298
+ case 3:
299
+ ggml_moe_q4_1_q8_1_cuda(
300
+ (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
301
+ (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
302
+ (int*)expert_ids.data_ptr(),
303
+ (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
304
+ tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
305
+ break;
306
+ case 6:
307
+ ggml_moe_q5_0_q8_1_cuda(
308
+ (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
309
+ (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
310
+ (int*)expert_ids.data_ptr(),
311
+ (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
312
+ tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
313
+ break;
314
+ case 7:
315
+ ggml_moe_q5_1_q8_1_cuda(
316
+ (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
317
+ (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
318
+ (int*)expert_ids.data_ptr(),
319
+ (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
320
+ tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
321
+ break;
322
+ case 8:
323
+ ggml_moe_q8_0_q8_1_cuda(
324
+ (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
325
+ (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
326
+ (int*)expert_ids.data_ptr(),
327
+ (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
328
+ tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
329
+ break;
330
+ case 10:
331
+ ggml_moe_q2_K_q8_1_cuda(
332
+ (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
333
+ (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
334
+ (int*)expert_ids.data_ptr(),
335
+ (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
336
+ tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
337
+ break;
338
+ case 11:
339
+ ggml_moe_q3_K_q8_1_cuda(
340
+ (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
341
+ (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
342
+ (int*)expert_ids.data_ptr(),
343
+ (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
344
+ tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
345
+ break;
346
+ case 12:
347
+ ggml_moe_q4_K_q8_1_cuda(
348
+ (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
349
+ (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
350
+ (int*)expert_ids.data_ptr(),
351
+ (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
352
+ tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
353
+ break;
354
+ case 13:
355
+ ggml_moe_q5_K_q8_1_cuda(
356
+ (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
357
+ (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
358
+ (int*)expert_ids.data_ptr(),
359
+ (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
360
+ tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
361
+ break;
362
+ case 14:
363
+ ggml_moe_q6_K_q8_1_cuda(
364
+ (void*)quant_X.data_ptr(), (void*)W.data_ptr(),
365
+ (scalar_t*)Y.data_ptr(), (int*)sorted_token_ids.data_ptr(),
366
+ (int*)expert_ids.data_ptr(),
367
+ (int*)num_tokens_post_padded.data_ptr(), W.stride(0), col, row,
368
+ tokens, padded, row, top_k, sorted_token_ids.sizes()[0], stream);
369
+ break;
370
+ }
371
+ });
372
+ return Y;
373
+}
374
+
375
+int64_t ggml_moe_get_block_size(int64_t type) {
376
+ switch (type) {
377
+ case 2:
378
+ return MMQ_X_Q4_0;
379
+ case 3:
380
+ return MMQ_X_Q4_1;
381
+ case 6:
382
+ return MMQ_X_Q5_0;
383
+ case 7:
384
+ return MMQ_X_Q5_1;
385
+ case 8:
386
+ return MMQ_X_Q8_0;
387
+ case 10:
388
+ return MMQ_X_Q2_K;
389
+ case 11:
390
+ return MMQ_X_Q3_K;
391
+ case 12:
392
+ return MMQ_X_Q4_K;
393
+ case 13:
394
+ return MMQ_X_Q5_K;
395
+ case 14:
396
+ return MMQ_X_Q6_K;
397
+ }
398
+ return 0;
399
+}
csrc/quantization/gguf/moe.cuh
new
+739
@@ -0,0 +1,739 @@
1
+#include <cstdint>
2
+
3
+/* Adapted from ./csrc/quantization/gguf/mmq.cuh
4
+ based on ./vllm/model_executor/layers/fused_moe/fused_moe.py */
5
+template <typename scalar_t, int qk, int qr, int qi, bool need_sum,
6
+ typename block_q_t, int mmq_x, int mmq_y, int nwarps,
7
+ allocate_tiles_cuda_t allocate_tiles, load_tiles_cuda_t load_tiles,
8
+ int vdr, vec_dot_q_mul_mat_cuda_t vec_dot>
9
+static __device__ __forceinline__ void moe_q(
10
+ const void* __restrict__ vx, const void* __restrict__ vy,
11
+ scalar_t* __restrict__ dst, const int* __restrict__ sorted_token_ids,
12
+ const int* __restrict__ expert_ids,
13
+ const int* __restrict__ num_tokens_post_padded, const int exp_stride,
14
+ const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y,
15
+ const int nrows_dst, const int top_k) {
16
+ const int blocks_per_row_x = ncols_x / qk;
17
+ const int blocks_per_col_y = nrows_y / QK8_1;
18
+ const int blocks_per_warp = WARP_SIZE_GGUF / qi;
19
+
20
+ const int ncols_dst = ncols_y * top_k;
21
+
22
+ const int row_dst_0 = blockIdx.x * mmq_y;
23
+ const int& row_x_0 = row_dst_0;
24
+
25
+ const int col_dst_0 = blockIdx.y * mmq_x;
26
+
27
+ int token_offs[mmq_x / nwarps];
28
+ for (int i = 0; i < mmq_x; i += nwarps) {
29
+ token_offs[i / nwarps] = sorted_token_ids[col_dst_0 + threadIdx.y + i];
30
+ }
31
+
32
+ const int exp_idx = expert_ids[blockIdx.y];
33
+ if (exp_idx > 255 || exp_idx < 0) return;
34
+ if (blockIdx.y * mmq_x > num_tokens_post_padded[0]) return;
35
+
36
+ const block_q_t* x = (const block_q_t*)((char*)vx + exp_idx * exp_stride);
37
+ const block_q8_1* y = (const block_q8_1*)(vy);
38
+
39
+ int* tile_x_ql = nullptr;
40
+ half2* tile_x_dm = nullptr;
41
+ int* tile_x_qh = nullptr;
42
+ int* tile_x_sc = nullptr;
43
+
44
+ allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc);
45
+
46
+ __shared__ int tile_y_qs[mmq_x * WARP_SIZE_GGUF];
47
+ __shared__ half2 tile_y_ds[mmq_x * WARP_SIZE_GGUF / QI8_1];
48
+
49
+ float sum[mmq_y / WARP_SIZE_GGUF][mmq_x / nwarps] = {{0.0f}};
50
+
51
+ for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) {
52
+ load_tiles(x + row_x_0 * blocks_per_row_x + ib0, tile_x_ql, tile_x_dm,
53
+ tile_x_qh, tile_x_sc, threadIdx.y, nrows_x - row_x_0 - 1,
54
+ threadIdx.x, blocks_per_row_x);
55
+
56
+ const int n_per_r = ((qk * blocks_per_warp) / qr);
57
+#pragma unroll
58
+ for (int ir = 0; ir < qr && ib0 * qk + ir * n_per_r < ncols_x; ++ir) {
59
+ const int kqs = ir * WARP_SIZE_GGUF + threadIdx.x;
60
+ const int kbxd = kqs / QI8_1;
61
+
62
+#pragma unroll
63
+ for (int i = 0; i < mmq_x; i += nwarps) {
64
+ const int col_y_eff = token_offs[i / nwarps] / top_k;
65
+ const int block_x = ib0 * (qk / QK8_1) + kbxd;
66
+ if (col_y_eff < ncols_y && block_x < blocks_per_col_y) {
67
+ const block_q8_1* by0 = &y[col_y_eff * blocks_per_col_y + block_x];
68
+ const int index_y =
69
+ (threadIdx.y + i) * WARP_SIZE_GGUF + kqs % WARP_SIZE_GGUF;
70
+ tile_y_qs[index_y] =
71
+ get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1);
72
+ }
73
+ }
74
+
75
+ if (threadIdx.x < n_per_r / QK8_1) {
76
+ const int kby = threadIdx.x % (WARP_SIZE_GGUF / QI8_1);
77
+ const int col_y_eff = token_offs[threadIdx.y] / top_k;
78
+ const int block_x =
79
+ ib0 * (qk / QK8_1) + ir * (WARP_SIZE_GGUF / QI8_1) + kby;
80
+
81
+ if (col_y_eff < ncols_y && block_x < blocks_per_col_y) {
82
+ const half2* dsi_src = &y[col_y_eff * blocks_per_col_y + block_x].ds;
83
+ half2* dsi_dst =
84
+ &tile_y_ds[threadIdx.y * (WARP_SIZE_GGUF / QI8_1) + kby];
85
+
86
+ if (need_sum) {
87
+ *dsi_dst = *dsi_src;
88
+ } else {
89
+ float* dfi_dst = (float*)dsi_dst;
90
+ *dfi_dst = __low2float(*dsi_src);
91
+ }
92
+ }
93
+ }
94
+ __syncthreads();
95
+
96
+ // #pragma unroll // unrolling this loop causes too much register pressure
97
+ for (int k = ir * WARP_SIZE_GGUF / qr; k < (ir + 1) * WARP_SIZE_GGUF / qr;
98
+ k += vdr) {
99
+#pragma unroll
100
+ for (int j = 0; j < mmq_x; j += nwarps) {
101
+#pragma unroll
102
+ for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) {
103
+ sum[i / WARP_SIZE_GGUF][j / nwarps] +=
104
+ vec_dot(tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs,
105
+ tile_y_ds, threadIdx.x + i, threadIdx.y + j, k);
106
+ }
107
+ }
108
+ }
109
+ __syncthreads();
110
+ }
111
+ }
112
+
113
+#pragma unroll
114
+ for (int j = 0; j < mmq_x; j += nwarps) {
115
+ const int col_dst = token_offs[j / nwarps];
116
+ if (col_dst >= ncols_dst) {
117
+ return;
118
+ }
119
+
120
+#pragma unroll
121
+ for (int i = 0; i < mmq_y; i += WARP_SIZE_GGUF) {
122
+ const int row_dst = row_dst_0 + threadIdx.x + i;
123
+ if (row_dst >= nrows_dst) {
124
+ continue;
125
+ }
126
+ dst[col_dst * nrows_dst + row_dst] = sum[i / WARP_SIZE_GGUF][j / nwarps];
127
+ }
128
+ }
129
+}
130
+
131
+#if defined(USE_ROCM)
132
+ #define MMQ_X_Q4_0 64
133
+ #define MMQ_Y_Q4_0 128
134
+ #define NWARPS_Q4_0 8
135
+#else
136
+ #define MMQ_X_Q4_0 4
137
+ #define MMQ_Y_Q4_0 32
138
+ #define NWARPS_Q4_0 4
139
+#endif
140
+
141
+template <typename scalar_t, bool need_check>
142
+static __global__ void
143
+#if defined(USE_ROCM)
144
+__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_0, 2)
145
+#endif
146
+ moe_q4_0(const void* __restrict__ vx, const void* __restrict__ vy,
147
+ scalar_t* __restrict__ dst, const int* sorted_token_ids,
148
+ const int* expert_ids, const int* num_tokens_post_padded,
149
+ const int exp_stride, const int ncols_x, const int nrows_x,
150
+ const int ncols_y, const int nrows_y, const int nrows_dst,
151
+ const int top_k) {
152
+ const int mmq_x = MMQ_X_Q4_0;
153
+ const int mmq_y = MMQ_Y_Q4_0;
154
+ const int nwarps = NWARPS_Q4_0;
155
+
156
+ moe_q<scalar_t, QK4_0, QR4_0, QI4_0, true, block_q4_0, mmq_x, mmq_y, nwarps,
157
+ allocate_tiles_q4_0<mmq_y>, load_tiles_q4_0<mmq_y, nwarps, need_check>,
158
+ VDR_Q4_0_Q8_1_MMQ, vec_dot_q4_0_q8_1_mul_mat>(
159
+ vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
160
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
161
+}
162
+
163
+template <typename scalar_t>
164
+static void ggml_moe_q4_0_q8_1_cuda(
165
+ const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids,
166
+ const int* expert_ids, const int* num_tokens_post_padded,
167
+ const int exp_stride, const int ncols_x, const int nrows_x,
168
+ const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k,
169
+ const int tokens_post_padded, cudaStream_t stream) {
170
+ int mmq_x = MMQ_X_Q4_0;
171
+ int mmq_y = MMQ_Y_Q4_0;
172
+ int nwarps = NWARPS_Q4_0;
173
+
174
+ const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
175
+ const int block_num_y = (tokens_post_padded) / mmq_x;
176
+ const dim3 block_nums(block_num_x, block_num_y, 1);
177
+ const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
178
+
179
+ if (nrows_x % mmq_y == 0) {
180
+ constexpr bool need_check = false;
181
+ moe_q4_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
182
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
183
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
184
+ } else {
185
+ constexpr bool need_check = true;
186
+ moe_q4_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
187
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
188
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
189
+ }
190
+}
191
+
192
+#if defined(USE_ROCM)
193
+ #define MMQ_X_Q4_1 64
194
+ #define MMQ_Y_Q4_1 128
195
+ #define NWARPS_Q4_1 8
196
+#else
197
+ #define MMQ_X_Q4_1 4
198
+ #define MMQ_Y_Q4_1 32
199
+ #define NWARPS_Q4_1 4
200
+#endif
201
+
202
+template <typename scalar_t, bool need_check>
203
+static __global__ void
204
+#if defined(USE_ROCM)
205
+__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_1, 2)
206
+#endif
207
+ moe_q4_1(const void* __restrict__ vx, const void* __restrict__ vy,
208
+ scalar_t* __restrict__ dst, const int* sorted_token_ids,
209
+ const int* expert_ids, const int* num_tokens_post_padded,
210
+ const int exp_stride, const int ncols_x, const int nrows_x,
211
+ const int ncols_y, const int nrows_y, const int nrows_dst,
212
+ const int top_k) {
213
+ const int mmq_x = MMQ_X_Q4_1;
214
+ const int mmq_y = MMQ_Y_Q4_1;
215
+ const int nwarps = NWARPS_Q4_1;
216
+
217
+ moe_q<scalar_t, QK4_1, QR4_1, QI4_1, true, block_q4_1, mmq_x, mmq_y, nwarps,
218
+ allocate_tiles_q4_1<mmq_y>, load_tiles_q4_1<mmq_y, nwarps, need_check>,
219
+ VDR_Q4_1_Q8_1_MMQ, vec_dot_q4_1_q8_1_mul_mat>(
220
+ vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
221
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
222
+}
223
+
224
+template <typename scalar_t>
225
+static void ggml_moe_q4_1_q8_1_cuda(
226
+ const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids,
227
+ const int* expert_ids, const int* num_tokens_post_padded,
228
+ const int exp_stride, const int ncols_x, const int nrows_x,
229
+ const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k,
230
+ const int tokens_post_padded, cudaStream_t stream) {
231
+ int mmq_x = MMQ_X_Q4_1;
232
+ int mmq_y = MMQ_Y_Q4_1;
233
+ int nwarps = NWARPS_Q4_1;
234
+
235
+ const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
236
+ const int block_num_y = (tokens_post_padded) / mmq_x;
237
+ const dim3 block_nums(block_num_x, block_num_y, 1);
238
+ const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
239
+
240
+ if (nrows_x % mmq_y == 0) {
241
+ constexpr bool need_check = false;
242
+ moe_q4_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
243
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
244
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
245
+ } else {
246
+ constexpr bool need_check = true;
247
+ moe_q4_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
248
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
249
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
250
+ }
251
+}
252
+
253
+#if defined(USE_ROCM)
254
+ #define MMQ_X_Q5_0 64
255
+ #define MMQ_Y_Q5_0 128
256
+ #define NWARPS_Q5_0 8
257
+#else
258
+ #define MMQ_X_Q5_0 4
259
+ #define MMQ_Y_Q5_0 32
260
+ #define NWARPS_Q5_0 4
261
+#endif
262
+
263
+template <typename scalar_t, bool need_check>
264
+static __global__ void
265
+#if defined(USE_ROCM)
266
+__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_0, 2)
267
+#endif
268
+ moe_q5_0(const void* __restrict__ vx, const void* __restrict__ vy,
269
+ scalar_t* __restrict__ dst, const int* sorted_token_ids,
270
+ const int* expert_ids, const int* num_tokens_post_padded,
271
+ const int exp_stride, const int ncols_x, const int nrows_x,
272
+ const int ncols_y, const int nrows_y, const int nrows_dst,
273
+ const int top_k) {
274
+ const int mmq_x = MMQ_X_Q5_0;
275
+ const int mmq_y = MMQ_Y_Q5_0;
276
+ const int nwarps = NWARPS_Q5_0;
277
+
278
+ moe_q<scalar_t, QK5_0, QR5_0, QI5_0, false, block_q5_0, mmq_x, mmq_y, nwarps,
279
+ allocate_tiles_q5_0<mmq_y>, load_tiles_q5_0<mmq_y, nwarps, need_check>,
280
+ VDR_Q5_0_Q8_1_MMQ, vec_dot_q5_0_q8_1_mul_mat>(
281
+ vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
282
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
283
+}
284
+
285
+template <typename scalar_t>
286
+static void ggml_moe_q5_0_q8_1_cuda(
287
+ const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids,
288
+ const int* expert_ids, const int* num_tokens_post_padded,
289
+ const int exp_stride, const int ncols_x, const int nrows_x,
290
+ const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k,
291
+ const int tokens_post_padded, cudaStream_t stream) {
292
+ const int mmq_x = MMQ_X_Q5_0;
293
+ const int mmq_y = MMQ_Y_Q5_0;
294
+ const int nwarps = NWARPS_Q5_0;
295
+
296
+ const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
297
+ const int block_num_y = (tokens_post_padded) / mmq_x;
298
+ const dim3 block_nums(block_num_x, block_num_y, 1);
299
+ const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
300
+
301
+ if (nrows_x % mmq_y == 0) {
302
+ constexpr bool need_check = false;
303
+ moe_q5_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
304
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
305
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
306
+ } else {
307
+ constexpr bool need_check = true;
308
+ moe_q5_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
309
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
310
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
311
+ }
312
+}
313
+
314
+#if defined(USE_ROCM)
315
+ #define MMQ_X_Q5_1 64
316
+ #define MMQ_Y_Q5_1 128
317
+ #define NWARPS_Q5_1 8
318
+#else
319
+ #define MMQ_X_Q5_1 4
320
+ #define MMQ_Y_Q5_1 32
321
+ #define NWARPS_Q5_1 4
322
+#endif
323
+
324
+template <typename scalar_t, bool need_check>
325
+static __global__ void
326
+#if defined(USE_ROCM)
327
+__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_1, 2)
328
+#endif
329
+ moe_q5_1(const void* __restrict__ vx, const void* __restrict__ vy,
330
+ scalar_t* __restrict__ dst, const int* sorted_token_ids,
331
+ const int* expert_ids, const int* num_tokens_post_padded,
332
+ const int exp_stride, const int ncols_x, const int nrows_x,
333
+ const int ncols_y, const int nrows_y, const int nrows_dst,
334
+ const int top_k) {
335
+ const int mmq_x = MMQ_X_Q5_1;
336
+ const int mmq_y = MMQ_Y_Q5_1;
337
+ const int nwarps = NWARPS_Q5_1;
338
+
339
+ moe_q<scalar_t, QK5_1, QR5_1, QI5_1, true, block_q5_1, mmq_x, mmq_y, nwarps,
340
+ allocate_tiles_q5_1<mmq_y>, load_tiles_q5_1<mmq_y, nwarps, need_check>,
341
+ VDR_Q5_1_Q8_1_MMQ, vec_dot_q5_1_q8_1_mul_mat>(
342
+ vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
343
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
344
+}
345
+
346
+template <typename scalar_t>
347
+static void ggml_moe_q5_1_q8_1_cuda(
348
+ const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids,
349
+ const int* expert_ids, const int* num_tokens_post_padded,
350
+ const int exp_stride, const int ncols_x, const int nrows_x,
351
+ const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k,
352
+ const int tokens_post_padded, cudaStream_t stream) {
353
+ const int mmq_x = MMQ_X_Q5_1;
354
+ const int mmq_y = MMQ_Y_Q5_1;
355
+ const int nwarps = NWARPS_Q5_1;
356
+
357
+ const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
358
+ const int block_num_y = (tokens_post_padded) / mmq_x;
359
+ const dim3 block_nums(block_num_x, block_num_y, 1);
360
+ const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
361
+
362
+ if (nrows_x % mmq_y == 0) {
363
+ constexpr bool need_check = false;
364
+ moe_q5_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
365
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
366
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
367
+ } else {
368
+ constexpr bool need_check = true;
369
+ moe_q5_1<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
370
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
371
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
372
+ }
373
+}
374
+
375
+#if defined(USE_ROCM)
376
+ #define MMQ_X_Q8_0 64
377
+ #define MMQ_Y_Q8_0 128
378
+ #define NWARPS_Q8_0 8
379
+#else
380
+ #define MMQ_X_Q8_0 4
381
+ #define MMQ_Y_Q8_0 32
382
+ #define NWARPS_Q8_0 4
383
+#endif
384
+
385
+template <typename scalar_t, bool need_check>
386
+static __global__ void
387
+#if defined(USE_ROCM)
388
+__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q8_0, 2)
389
+#endif
390
+ moe_q8_0(const void* __restrict__ vx, const void* __restrict__ vy,
391
+ scalar_t* __restrict__ dst, const int* sorted_token_ids,
392
+ const int* expert_ids, const int* num_tokens_post_padded,
393
+ const int exp_stride, const int ncols_x, const int nrows_x,
394
+ const int ncols_y, const int nrows_y, const int nrows_dst,
395
+ const int top_k) {
396
+ const int mmq_x = MMQ_X_Q8_0;
397
+ const int mmq_y = MMQ_Y_Q8_0;
398
+ const int nwarps = NWARPS_Q8_0;
399
+
400
+ moe_q<scalar_t, QK8_0, QR8_0, QI8_0, false, block_q8_0, mmq_x, mmq_y, nwarps,
401
+ allocate_tiles_q8_0<mmq_y>, load_tiles_q8_0<mmq_y, nwarps, need_check>,
402
+ VDR_Q8_0_Q8_1_MMQ, vec_dot_q8_0_q8_1_mul_mat>(
403
+ vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
404
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
405
+}
406
+
407
+template <typename scalar_t>
408
+static void ggml_moe_q8_0_q8_1_cuda(
409
+ const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids,
410
+ const int* expert_ids, const int* num_tokens_post_padded,
411
+ const int exp_stride, const int ncols_x, const int nrows_x,
412
+ const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k,
413
+ const int tokens_post_padded, cudaStream_t stream) {
414
+ const int mmq_x = MMQ_X_Q8_0;
415
+ const int mmq_y = MMQ_Y_Q8_0;
416
+ const int nwarps = NWARPS_Q8_0;
417
+
418
+ const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
419
+ const int block_num_y = (tokens_post_padded) / mmq_x;
420
+ const dim3 block_nums(block_num_x, block_num_y, 1);
421
+ const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
422
+
423
+ if (nrows_x % mmq_y == 0) {
424
+ constexpr bool need_check = false;
425
+ moe_q8_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
426
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
427
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
428
+ } else {
429
+ constexpr bool need_check = true;
430
+ moe_q8_0<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
431
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
432
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
433
+ }
434
+}
435
+
436
+#if defined(USE_ROCM)
437
+ #define MMQ_X_Q2_K 64
438
+ #define MMQ_Y_Q2_K 128
439
+ #define NWARPS_Q2_K 8
440
+#else
441
+ #define MMQ_X_Q2_K 4
442
+ #define MMQ_Y_Q2_K 32
443
+ #define NWARPS_Q2_K 4
444
+#endif
445
+
446
+template <typename scalar_t, bool need_check>
447
+static __global__ void
448
+#if defined(USE_ROCM)
449
+__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q2_K, 2)
450
+#endif
451
+ moe_q2_K(const void* __restrict__ vx, const void* __restrict__ vy,
452
+ scalar_t* __restrict__ dst, const int* sorted_token_ids,
453
+ const int* expert_ids, const int* num_tokens_post_padded,
454
+ const int exp_stride, const int ncols_x, const int nrows_x,
455
+ const int ncols_y, const int nrows_y, const int nrows_dst,
456
+ const int top_k) {
457
+ const int mmq_x = MMQ_X_Q2_K;
458
+ const int mmq_y = MMQ_Y_Q2_K;
459
+ const int nwarps = NWARPS_Q2_K;
460
+
461
+ moe_q<scalar_t, QK_K, QR2_K, QI2_K, false, block_q2_K, mmq_x, mmq_y, nwarps,
462
+ allocate_tiles_q2_K<mmq_y>, load_tiles_q2_K<mmq_y, nwarps, need_check>,
463
+ VDR_Q2_K_Q8_1_MMQ, vec_dot_q2_K_q8_1_mul_mat>(
464
+ vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
465
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
466
+}
467
+
468
+template <typename scalar_t>
469
+static void ggml_moe_q2_K_q8_1_cuda(
470
+ const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids,
471
+ const int* expert_ids, const int* num_tokens_post_padded,
472
+ const int exp_stride, const int ncols_x, const int nrows_x,
473
+ const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k,
474
+ const int tokens_post_padded, cudaStream_t stream) {
475
+ const int mmq_x = MMQ_X_Q2_K;
476
+ const int mmq_y = MMQ_Y_Q2_K;
477
+ const int nwarps = NWARPS_Q2_K;
478
+
479
+ const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
480
+ const int block_num_y = (tokens_post_padded) / mmq_x;
481
+ const dim3 block_nums(block_num_x, block_num_y, 1);
482
+ const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
483
+
484
+ if (nrows_x % mmq_y == 0) {
485
+ constexpr bool need_check = false;
486
+ moe_q2_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
487
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
488
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
489
+ } else {
490
+ constexpr bool need_check = true;
491
+ moe_q2_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
492
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
493
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
494
+ }
495
+}
496
+
497
+#if defined(USE_ROCM)
498
+ #define MMQ_X_Q3_K 64
499
+ #define MMQ_Y_Q3_K 128
500
+ #define NWARPS_Q3_K 8
501
+#else
502
+ #define MMQ_X_Q3_K 4
503
+ #define MMQ_Y_Q3_K 32
504
+ #define NWARPS_Q3_K 4
505
+#endif
506
+
507
+template <typename scalar_t, bool need_check>
508
+static __global__ void
509
+#if defined(USE_ROCM)
510
+__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q3_K, 2)
511
+#endif
512
+ moe_q3_K(const void* __restrict__ vx, const void* __restrict__ vy,
513
+ scalar_t* __restrict__ dst, const int* sorted_token_ids,
514
+ const int* expert_ids, const int* num_tokens_post_padded,
515
+ const int exp_stride, const int ncols_x, const int nrows_x,
516
+ const int ncols_y, const int nrows_y, const int nrows_dst,
517
+ const int top_k) {
518
+
519
+ const int mmq_x = MMQ_X_Q3_K;
520
+ const int mmq_y = MMQ_Y_Q3_K;
521
+ const int nwarps = NWARPS_Q3_K;
522
+
523
+ moe_q<scalar_t, QK_K, QR3_K, QI3_K, false, block_q3_K, mmq_x, mmq_y, nwarps,
524
+ allocate_tiles_q3_K<mmq_y>, load_tiles_q3_K<mmq_y, nwarps, need_check>,
525
+ VDR_Q3_K_Q8_1_MMQ, vec_dot_q3_K_q8_1_mul_mat>(
526
+ vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
527
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
528
+}
529
+template <typename scalar_t>
530
+static void ggml_moe_q3_K_q8_1_cuda(
531
+ const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids,
532
+ const int* expert_ids, const int* num_tokens_post_padded,
533
+ const int exp_stride, const int ncols_x, const int nrows_x,
534
+ const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k,
535
+ const int tokens_post_padded, cudaStream_t stream) {
536
+ const int mmq_x = MMQ_X_Q3_K;
537
+ const int mmq_y = MMQ_Y_Q3_K;
538
+ const int nwarps = NWARPS_Q3_K;
539
+
540
+ const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
541
+ const int block_num_y = (tokens_post_padded) / mmq_x;
542
+ const dim3 block_nums(block_num_x, block_num_y, 1);
543
+ const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
544
+
545
+ if (nrows_x % mmq_y == 0) {
546
+ constexpr bool need_check = false;
547
+ moe_q3_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
548
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
549
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
550
+ } else {
551
+ constexpr bool need_check = true;
552
+ moe_q3_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
553
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
554
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
555
+ }
556
+}
557
+
558
+#if defined(USE_ROCM)
559
+ #define MMQ_X_Q4_K 64
560
+ #define MMQ_Y_Q4_K 128
561
+ #define NWARPS_Q4_K 8
562
+#else
563
+ #define MMQ_X_Q4_K 4
564
+ #define MMQ_Y_Q4_K 32
565
+ #define NWARPS_Q4_K 4
566
+#endif
567
+
568
+template <typename scalar_t, bool need_check>
569
+static __global__ void
570
+#if defined(USE_ROCM)
571
+__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q4_K, 2)
572
+#endif
573
+ moe_q4_K(const void* __restrict__ vx, const void* __restrict__ vy,
574
+ scalar_t* __restrict__ dst, const int* sorted_token_ids,
575
+ const int* expert_ids, const int* num_tokens_post_padded,
576
+ const int exp_stride, const int ncols_x, const int nrows_x,
577
+ const int ncols_y, const int nrows_y, const int nrows_dst,
578
+ const int top_k) {
579
+ const int mmq_x = MMQ_X_Q4_K;
580
+ const int mmq_y = MMQ_Y_Q4_K;
581
+ const int nwarps = NWARPS_Q4_K;
582
+
583
+ moe_q<scalar_t, QK_K, QR4_K, QI4_K, true, block_q4_K, mmq_x, mmq_y, nwarps,
584
+ allocate_tiles_q4_K<mmq_y>, load_tiles_q4_K<mmq_y, nwarps, need_check>,
585
+ VDR_Q4_K_Q8_1_MMQ, vec_dot_q4_K_q8_1_mul_mat>(
586
+ vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
587
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
588
+}
589
+
590
+template <typename scalar_t>
591
+static void ggml_moe_q4_K_q8_1_cuda(
592
+ const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids,
593
+ const int* expert_ids, const int* num_tokens_post_padded,
594
+ const int exp_stride, const int ncols_x, const int nrows_x,
595
+ const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k,
596
+ const int tokens_post_padded, cudaStream_t stream) {
597
+ const int mmq_x = MMQ_X_Q4_K;
598
+ const int mmq_y = MMQ_Y_Q4_K;
599
+ const int nwarps = NWARPS_Q4_K;
600
+
601
+ const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
602
+ const int block_num_y = (tokens_post_padded) / mmq_x;
603
+ const dim3 block_nums(block_num_x, block_num_y, 1);
604
+ const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
605
+
606
+ if (nrows_x % mmq_y == 0) {
607
+ constexpr bool need_check = false;
608
+ moe_q4_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
609
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
610
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
611
+ } else {
612
+ constexpr bool need_check = true;
613
+ moe_q4_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
614
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
615
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
616
+ }
617
+}
618
+
619
+#if defined(USE_ROCM)
620
+ #define MMQ_X_Q5_K 64
621
+ #define MMQ_Y_Q5_K 128
622
+ #define NWARPS_Q5_K 8
623
+#else
624
+ #define MMQ_X_Q5_K 4
625
+ #define MMQ_Y_Q5_K 32
626
+ #define NWARPS_Q5_K 4
627
+#endif
628
+
629
+template <typename scalar_t, bool need_check>
630
+static __global__ void
631
+#if defined(USE_ROCM)
632
+__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q5_K, 2)
633
+#endif
634
+ moe_q5_K(const void* __restrict__ vx, const void* __restrict__ vy,
635
+ scalar_t* __restrict__ dst, const int* sorted_token_ids,
636
+ const int* expert_ids, const int* num_tokens_post_padded,
637
+ const int exp_stride, const int ncols_x, const int nrows_x,
638
+ const int ncols_y, const int nrows_y, const int nrows_dst,
639
+ const int top_k) {
640
+ const int mmq_x = MMQ_X_Q5_K;
641
+ const int mmq_y = MMQ_Y_Q5_K;
642
+ const int nwarps = NWARPS_Q5_K;
643
+
644
+ moe_q<scalar_t, QK_K, QR5_K, QI5_K, true, block_q5_K, mmq_x, mmq_y, nwarps,
645
+ allocate_tiles_q5_K<mmq_y>, load_tiles_q5_K<mmq_y, nwarps, need_check>,
646
+ VDR_Q5_K_Q8_1_MMQ, vec_dot_q5_K_q8_1_mul_mat>(
647
+ vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
648
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
649
+}
650
+
651
+template <typename scalar_t>
652
+static void ggml_moe_q5_K_q8_1_cuda(
653
+ const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids,
654
+ const int* expert_ids, const int* num_tokens_post_padded,
655
+ const int exp_stride, const int ncols_x, const int nrows_x,
656
+ const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k,
657
+ const int tokens_post_padded, cudaStream_t stream) {
658
+ const int mmq_x = MMQ_X_Q5_K;
659
+ const int mmq_y = MMQ_Y_Q5_K;
660
+ const int nwarps = NWARPS_Q5_K;
661
+
662
+ const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
663
+ const int block_num_y = (tokens_post_padded) / mmq_x;
664
+ const dim3 block_nums(block_num_x, block_num_y, 1);
665
+ const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
666
+
667
+ if (nrows_x % mmq_y == 0) {
668
+ constexpr bool need_check = false;
669
+ moe_q5_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
670
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
671
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
672
+ } else {
673
+ constexpr bool need_check = true;
674
+ moe_q5_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
675
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
676
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
677
+ }
678
+}
679
+
680
+#if defined(USE_ROCM)
681
+ #define MMQ_X_Q6_K 64
682
+ #define MMQ_Y_Q6_K 128
683
+ #define NWARPS_Q6_K 8
684
+#else
685
+ #define MMQ_X_Q6_K 4
686
+ #define MMQ_Y_Q6_K 32
687
+ #define NWARPS_Q6_K 4
688
+#endif
689
+
690
+template <typename scalar_t, bool need_check>
691
+static __global__ void
692
+#if defined(USE_ROCM)
693
+__launch_bounds__(WARP_SIZE_GGUF* NWARPS_Q6_K, 2)
694
+#endif
695
+ moe_q6_K(const void* __restrict__ vx, const void* __restrict__ vy,
696
+ scalar_t* __restrict__ dst, const int* sorted_token_ids,
697
+ const int* expert_ids, const int* num_tokens_post_padded,
698
+ const int exp_stride, const int ncols_x, const int nrows_x,
699
+ const int ncols_y, const int nrows_y, const int nrows_dst,
700
+ const int top_k) {
701
+ const int mmq_x = MMQ_X_Q6_K;
702
+ const int mmq_y = MMQ_Y_Q6_K;
703
+ const int nwarps = NWARPS_Q6_K;
704
+
705
+ moe_q<scalar_t, QK_K, QR6_K, QI6_K, false, block_q6_K, mmq_x, mmq_y, nwarps,
706
+ allocate_tiles_q6_K<mmq_y>, load_tiles_q6_K<mmq_y, nwarps, need_check>,
707
+ VDR_Q6_K_Q8_1_MMQ, vec_dot_q6_K_q8_1_mul_mat>(
708
+ vx, vy, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
709
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
710
+}
711
+
712
+template <typename scalar_t>
713
+static void ggml_moe_q6_K_q8_1_cuda(
714
+ const void* inp, const void* w, scalar_t* dst, const int* sorted_token_ids,
715
+ const int* expert_ids, const int* num_tokens_post_padded,
716
+ const int exp_stride, const int ncols_x, const int nrows_x,
717
+ const int ncols_y, const int nrows_y, const int nrows_dst, const int top_k,
718
+ const int tokens_post_padded, cudaStream_t stream) {
719
+ const int mmq_x = MMQ_X_Q6_K;
720
+ const int mmq_y = MMQ_Y_Q6_K;
721
+ const int nwarps = NWARPS_Q6_K;
722
+
723
+ const int block_num_x = (nrows_x + mmq_y - 1) / mmq_y;
724
+ const int block_num_y = (tokens_post_padded) / mmq_x;
725
+ const dim3 block_nums(block_num_x, block_num_y, 1);
726
+ const dim3 block_dims(WARP_SIZE_GGUF, nwarps, 1);
727
+
728
+ if (nrows_x % mmq_y == 0) {
729
+ constexpr bool need_check = false;
730
+ moe_q6_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
731
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
732
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
733
+ } else {
734
+ constexpr bool need_check = true;
735
+ moe_q6_K<scalar_t, need_check><<<block_nums, block_dims, 0, stream>>>(
736
+ w, inp, dst, sorted_token_ids, expert_ids, num_tokens_post_padded,
737
+ exp_stride, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst, top_k);
738
+ }
739
+}
csrc/torch_bindings.cpp
+10
@@ -305,6 +305,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
305
"ggml_mul_mat_a8(Tensor W, Tensor X, int type, SymInt row) -> Tensor");
306
ops.impl("ggml_mul_mat_a8", torch::kCUDA, &ggml_mul_mat_a8);
307
308
+ // moe kernel for GGML.
309
+ ops.def(
310
+ "ggml_moe_a8(Tensor X, Tensor W, "
311
+ "Tensor sorted_token_ids, Tensor expert_ids, Tensor "
312
+ "num_tokens_post_padded, "
313
+ "int type, SymInt row, SymInt top_k, SymInt tokens) -> Tensor");
314
+ ops.impl("ggml_moe_a8", torch::kCUDA, &ggml_moe_a8);
315
+
316
+ ops.def("ggml_moe_get_block_size", &ggml_moe_get_block_size);
317
+
318
#ifndef USE_ROCM
319
// fp8_marlin Optimized Quantized GEMM for FP8 weight-only.
320
ops.def(
tests/kernels/test_ggml.py
+13
@@ -22,3 +22,16 @@ def test_ggml_opcheck(quant_type):
22
(qweight, x, quant_type, qweight.shape[0]))
23
opcheck(torch.ops._C.ggml_mul_mat_vec_a8,
24
(qweight, x, quant_type, qweight.shape[0]))
25
+
26
+ shape = [256, 1024, 336]
27
+ qweight = torch.randint(0, 100, shape, device='cuda', dtype=torch.uint8)
28
+ x = torch.rand((1, 1024), device='cuda', dtype=torch.float16)
29
+ sorted_token_ids = torch.arange(776, device='cuda')
30
+ expert_ids = torch.randint(0, 256, (194, ), device='cuda')
31
+ num_tokens_post_padded = torch.tensor([1],
32
+ dtype=torch.int64,
33
+ device='cuda')
34
+
35
+ opcheck(torch.ops._C.ggml_moe_a8,
36
+ (x, qweight, sorted_token_ids, expert_ids, num_tokens_post_padded,
37
+ quant_type, qweight.shape[0], 1, x.shape[0]))
tests/kernels/test_gguf.py
+64
@@ -8,9 +8,13 @@ from gguf import GGMLQuantizationType, GGUFReader, ReaderTensor, dequantize
8
from huggingface_hub import snapshot_download
9
10
import vllm._custom_ops as ops
11
+from vllm.model_executor.layers.activation import SiluAndMul
12
+from vllm.model_executor.layers.fused_moe import fused_experts
13
+from vllm.model_executor.layers.quantization.gguf import _fused_moe_gguf
14
from vllm.platforms import current_platform
15
16
GGUF_SAMPLE = snapshot_download("Isotr0py/test-gguf-sample")
17
+GGUF_SAMPLE_MOE = snapshot_download("SzymonOzog/test-gguf-moe-sample")
18
19
20
def get_gguf_sample_tensors(
@@ -22,6 +26,15 @@ def get_gguf_sample_tensors(
26
return GGUFReader(sample_file).tensors
27
28
29
+def get_gguf_MoE_tensors(
30
+ hidden_size: int,
31
+ quant_type: GGMLQuantizationType) -> list[ReaderTensor]:
32
+ sample_dir = GGUF_SAMPLE_MOE
33
+ filename = f"Quant_{quant_type.name}_{hidden_size}.gguf"
34
+ sample_file = Path(sample_dir) / filename
35
+ return GGUFReader(sample_file).tensors
36
+
37
+
38
DTYPES = [torch.half, torch.bfloat16, torch.float32]
39
# Hidden_size for testing, must match the sample file in HF repo,
40
# we have `hidden_size = 256, 1024` for test in HF repo currently.
@@ -132,3 +145,54 @@ def test_mmq(num_tokens: int, hidden_size: int, dtype: torch.dtype,
145
ref_output,
146
atol=atols[dtype],
147
rtol=rtols[dtype])
148
+
149
+
150
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
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
+ ])
168
+@torch.inference_mode()
169
+def test_moe(num_tokens: int, hidden_size: int, dtype: torch.dtype,
170
+ quant_type: GGMLQuantizationType, top_k: int):
171
+ current_platform.seed_everything(0)
172
+ H, E = 1024, 256
173
+
174
+ x = torch.rand((num_tokens, H), dtype=dtype, device="cuda")
175
+
176
+ 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")
178
+
179
+ tensors = get_gguf_MoE_tensors(hidden_size, quant_type)
180
+
181
+ w13 = tensors[0]
182
+ w2 = tensors[1]
183
+
184
+ w13_dequant = torch.tensor(dequantize(w13.data, quant_type),
185
+ device="cuda").to(dtype)
186
+
187
+ w2_dequant = torch.tensor(dequantize(w2.data, quant_type),
188
+ device="cuda").to(dtype)
189
+ act = SiluAndMul()
190
+
191
+ output = _fused_moe_gguf(x, torch.tensor(w13.data, device="cuda"),
192
+ torch.tensor(w2.data,
193
+ device="cuda"), topk_weights,
194
+ topk_ids, quant_type, quant_type, act)
195
+
196
+ ref_output = fused_experts(x, w13_dequant, w2_dequant, topk_weights,
197
+ topk_ids).reshape(output.shape)
198
+ torch.testing.assert_close(output, ref_output, atol=1, rtol=1e-1)
vllm/_custom_ops.py
+37
@@ -448,6 +448,23 @@ if hasattr(torch.ops._C, "ggml_dequantize"):
448
batch = X.size(0)
449
return torch.empty((batch, row), dtype=X.dtype, device=W.device)
450
451
+ @register_fake("_C::ggml_moe_a8")
452
+ def _ggml_moe_a8_fake(
453
+ X: torch.Tensor,
454
+ W: torch.Tensor,
455
+ sorted_token_ids: torch.Tensor,
456
+ expert_ids: torch.Tensor,
457
+ num_tokens_post_padded: torch.Tensor,
458
+ quant_type: int,
459
+ row: torch.SymInt,
460
+ top_k: torch.SymInt,
461
+ tokens: torch.SymInt,
462
+ ) -> torch.Tensor:
463
+ tokens = X.size(0)
464
+ return torch.empty((tokens * top_k, row),
465
+ dtype=torch.float16,
466
+ device=W.device)
467
+
468
469
# cutlass
470
def cutlass_scaled_fp4_mm(a: torch.Tensor, b: torch.Tensor,
@@ -1034,6 +1051,26 @@ def ggml_mul_mat_a8(
1051
return torch.ops._C.ggml_mul_mat_a8(W, X, quant_type, row)
1052
1053
1054
+def ggml_moe_a8(
1055
+ X: torch.Tensor,
1056
+ W: torch.Tensor,
1057
+ sorted_token_ids: torch.Tensor,
1058
+ expert_ids: torch.Tensor,
1059
+ num_tokens_post_padded: torch.Tensor,
1060
+ quant_type: int,
1061
+ row: int,
1062
+ top_k: int,
1063
+ tokens: int,
1064
+) -> torch.Tensor:
1065
+ return torch.ops._C.ggml_moe_a8(X, W, sorted_token_ids, expert_ids,
1066
+ num_tokens_post_padded, quant_type, row,
1067
+ top_k, tokens)
1068
+
1069
+
1070
+def ggml_moe_get_block_size(quant_type: int) -> int:
1071
+ return torch.ops._C.ggml_moe_get_block_size(quant_type)
1072
+
1073
+
1074
# mamba
1075
def causal_conv1d_fwd(x: torch.Tensor, weight: torch.Tensor,
1076
bias_: Optional[torch.Tensor],
vllm/model_executor/layers/quantization/gguf.py
+61
-21
@@ -8,7 +8,9 @@ from gguf import GGMLQuantizationType as WeightType
8
from torch.nn.parameter import Parameter, UninitializedParameter
9
10
from vllm import _custom_ops as ops
11
+from vllm.logger import init_logger
12
from vllm.model_executor.layers.activation import SiluAndMul
13
+from vllm.model_executor.layers.fused_moe.fused_moe import moe_align_block_size
14
from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,
15
FusedMoEMethodBase)
16
from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
@@ -18,6 +20,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
20
VocabParallelEmbedding)
21
from vllm.model_executor.utils import set_weight_attrs
22
23
+logger = init_logger(__name__)
24
+
25
26
class GGUFConfig(QuantizationConfig):
27
"""Config class for GGUF."""
@@ -119,6 +123,59 @@ def _fuse_mul_mat(x: torch.Tensor, qweight: torch.Tensor,
123
return y
124
125
126
+def _fused_moe_gguf(
127
+ x: torch.Tensor,
128
+ w1: torch.Tensor,
129
+ w2: torch.Tensor,
130
+ topk_weights: torch.Tensor,
131
+ topk_ids: torch.Tensor,
132
+ qweight_type: int,
133
+ qweight_type2: int,
134
+ act,
135
+) -> torch.Tensor:
136
+ out_hidden_states = torch.empty_like(x)
137
+ if qweight_type2 in MMQ_QUANT_TYPES and qweight_type in MMQ_QUANT_TYPES:
138
+ num_tokens, _ = x.shape
139
+ E, N, _ = w1.shape
140
+ top_k = topk_ids.shape[1]
141
+ BLOCK_SIZE = ops.ggml_moe_get_block_size(qweight_type)
142
+
143
+ sorted_token_ids, expert_ids, num_tokens_post_padded = \
144
+ moe_align_block_size(topk_ids, BLOCK_SIZE, E)
145
+ out = ops.ggml_moe_a8(x, w1, sorted_token_ids, expert_ids,
146
+ num_tokens_post_padded, qweight_type, N, top_k,
147
+ num_tokens)
148
+ out = act(out)
149
+ out = ops.ggml_moe_a8(out, w2, sorted_token_ids, expert_ids,
150
+ num_tokens_post_padded, qweight_type2,
151
+ w2.shape[1], 1, num_tokens * top_k)
152
+ out = out.reshape(num_tokens, top_k, w2.shape[1]).mul_(
153
+ topk_weights.view(num_tokens, top_k, 1))
154
+ ops.moe_sum(out, out_hidden_states)
155
+ else:
156
+ logger.warning_once("There is no support for fast MoE kernel "
157
+ "for current quantization method. "
158
+ "Falling back to slow implementation. ")
159
+ for tok, (w, idx) in enumerate(zip(topk_weights, topk_ids)):
160
+ inp = x[tok].reshape((1, ) + x.shape[1:])
161
+ current_hidden_state = None
162
+ for ww, ii in zip(w, idx):
163
+ expert_up = w1[ii]
164
+
165
+ out = _fuse_mul_mat(inp, expert_up, qweight_type)
166
+ out = act(out)
167
+
168
+ expert_down = w2[ii]
169
+ current_state = _fuse_mul_mat(out, expert_down,
170
+ qweight_type2).mul_(ww)
171
+ if current_hidden_state is None:
172
+ current_hidden_state = current_state
173
+ else:
174
+ current_hidden_state.add_(current_state)
175
+ out_hidden_states[tok] = current_hidden_state
176
+ return out_hidden_states
177
+
178
+
179
class GGUFLinearMethod(LinearMethodBase):
180
"""Linear method for GGUF.
181
@@ -285,27 +342,10 @@ class GGUFMoEMethod(FusedMoEMethodBase):
342
custom_routing_function=custom_routing_function,
343
scoring_func=scoring_func,
344
e_score_correction_bias=e_score_correction_bias)
288
- final_hidden_states = torch.empty_like(x)
289
- for tok, (w, idx) in enumerate(zip(topk_weights, topk_ids)):
290
- inp = x[tok].reshape((1, ) + x.shape[1:])
291
- current_hidden_state = None
292
- for ww, ii in zip(w, idx):
293
- expert_up = layer.w13_qweight[ii]
294
-
295
- out = _fuse_mul_mat(inp, expert_up,
296
- layer.w13_qweight_type.weight_type)
297
- out = self.act(out)
298
-
299
- expert_down = layer.w2_qweight[ii]
300
- current_state = _fuse_mul_mat(
301
- out, expert_down,
302
- layer.w2_qweight_type.weight_type).mul_(ww)
303
- if current_hidden_state is None:
304
- current_hidden_state = current_state
305
- else:
306
- current_hidden_state.add_(current_state)
307
- final_hidden_states[tok] = current_hidden_state
308
- return final_hidden_states
345
+ return _fused_moe_gguf(x, layer.w13_qweight, layer.w2_qweight,
346
+ topk_weights, topk_ids,
347
+ layer.w13_qweight_type.weight_type,
348
+ layer.w2_qweight_type.weight_type, self.act)
349
350
351
class GGUFEmbeddingMethod(GGUFLinearMethod):