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, 06:00 UTC
[Kernel] GGUF MMVQ kernel for multiple input vectors (#18754)
Signed-off-by: SzymonOzog <szymon.ozog@gmail.com>
Szymon Ożóg committed
Jun 16, 2025 at 11:33 UTC
dec66d253b6e9c2745940bd4ccda786e7ac975d1
4 files changed
+95
-87
csrc/quantization/gguf/gguf_kernel.cu
+24
-23
@@ -92,111 +92,112 @@ torch::Tensor ggml_mul_mat_vec_a8(torch::Tensor W, // quant weight
92
torch::Tensor X, // input
93
int64_t type, int64_t row) {
94
int col = X.sizes()[1];
95
+ int vecs = X.sizes()[0];
96
const int padded = (col + 512 - 1) / 512 * 512;
97
const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
98
auto options = torch::TensorOptions().dtype(X.dtype()).device(W.device());
98
- at::Tensor Y = torch::empty({1, row}, options);
99
+ at::Tensor Y = torch::empty({vecs, row}, options);
100
cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
101
options = torch::TensorOptions().dtype(torch::kInt32).device(W.device());
101
- at::Tensor quant_X = torch::empty({1, padded / 32 * 9}, options);
102
+ at::Tensor quant_X = torch::empty({vecs, padded / 32 * 9}, options);
103
VLLM_DISPATCH_FLOATING_TYPES(X.scalar_type(), "ggml_mul_mat_vec_a8", [&] {
103
- quantize_row_q8_1_cuda<scalar_t>((scalar_t*)X.data_ptr(),
104
- (void*)quant_X.data_ptr(), col, 1, stream);
104
+ quantize_row_q8_1_cuda<scalar_t>(
105
+ (scalar_t*)X.data_ptr(), (void*)quant_X.data_ptr(), col, vecs, stream);
106
switch (type) {
107
case 2:
108
mul_mat_vec_q4_0_q8_1_cuda<scalar_t>(
109
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
109
- (scalar_t*)Y.data_ptr(), col, row, stream);
110
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
111
break;
112
case 3:
113
mul_mat_vec_q4_1_q8_1_cuda<scalar_t>(
114
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
114
- (scalar_t*)Y.data_ptr(), col, row, stream);
115
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
116
break;
117
case 6:
118
mul_mat_vec_q5_0_q8_1_cuda<scalar_t>(
119
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
119
- (scalar_t*)Y.data_ptr(), col, row, stream);
120
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
121
break;
122
case 7:
123
mul_mat_vec_q5_1_q8_1_cuda<scalar_t>(
124
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
124
- (scalar_t*)Y.data_ptr(), col, row, stream);
125
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
126
break;
127
case 8:
128
mul_mat_vec_q8_0_q8_1_cuda<scalar_t>(
129
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
129
- (scalar_t*)Y.data_ptr(), col, row, stream);
130
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
131
break;
132
case 10:
133
mul_mat_vec_q2_K_q8_1_cuda<scalar_t>(
134
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
134
- (scalar_t*)Y.data_ptr(), col, row, stream);
135
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
136
break;
137
case 11:
138
mul_mat_vec_q3_K_q8_1_cuda<scalar_t>(
139
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
139
- (scalar_t*)Y.data_ptr(), col, row, stream);
140
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
141
break;
142
case 12:
143
mul_mat_vec_q4_K_q8_1_cuda<scalar_t>(
144
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
144
- (scalar_t*)Y.data_ptr(), col, row, stream);
145
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
146
break;
147
case 13:
148
mul_mat_vec_q5_K_q8_1_cuda<scalar_t>(
149
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
149
- (scalar_t*)Y.data_ptr(), col, row, stream);
150
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
151
break;
152
case 14:
153
mul_mat_vec_q6_K_q8_1_cuda<scalar_t>(
154
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
154
- (scalar_t*)Y.data_ptr(), col, row, stream);
155
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
156
break;
157
case 16:
158
mul_mat_vec_iq2_xxs_q8_1_cuda<scalar_t>(
159
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
159
- (scalar_t*)Y.data_ptr(), col, row, stream);
160
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
161
break;
162
case 17:
163
mul_mat_vec_iq2_xs_q8_1_cuda<scalar_t>(
164
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
164
- (scalar_t*)Y.data_ptr(), col, row, stream);
165
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
166
break;
167
case 18:
168
mul_mat_vec_iq3_xxs_q8_1_cuda<scalar_t>(
169
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
169
- (scalar_t*)Y.data_ptr(), col, row, stream);
170
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
171
break;
172
case 19:
173
mul_mat_vec_iq1_s_q8_1_cuda<scalar_t>(
174
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
174
- (scalar_t*)Y.data_ptr(), col, row, stream);
175
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
176
break;
177
case 20:
178
mul_mat_vec_iq4_nl_q8_1_cuda<scalar_t>(
179
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
179
- (scalar_t*)Y.data_ptr(), col, row, stream);
180
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
181
break;
182
case 21:
183
mul_mat_vec_iq3_s_q8_1_cuda<scalar_t>(
184
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
184
- (scalar_t*)Y.data_ptr(), col, row, stream);
185
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
186
break;
187
case 22:
188
mul_mat_vec_iq2_s_q8_1_cuda<scalar_t>(
189
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
189
- (scalar_t*)Y.data_ptr(), col, row, stream);
190
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
191
break;
192
case 23:
193
mul_mat_vec_iq4_xs_q8_1_cuda<scalar_t>(
194
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
194
- (scalar_t*)Y.data_ptr(), col, row, stream);
195
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
196
break;
197
case 29:
198
mul_mat_vec_iq1_m_q8_1_cuda<scalar_t>(
199
(void*)W.data_ptr(), (void*)quant_X.data_ptr(),
199
- (scalar_t*)Y.data_ptr(), col, row, stream);
200
+ (scalar_t*)Y.data_ptr(), col, row, vecs, stream);
201
break;
202
}
203
});
csrc/quantization/gguf/mmvq.cuh
+65
-62
@@ -1,16 +1,19 @@
1
// copied and adapted from https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu
2
template <typename scalar_t, int qk, int qi, typename block_q_t, int vdr, vec_dot_q_cuda_t vec_dot_q_cuda>
3
-static __global__ void mul_mat_vec_q(const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, const int ncols, const int nrows) {
3
+static __global__ void mul_mat_vec_q(const void * __restrict__ vx, const void * __restrict__ vy, scalar_t * __restrict__ dst, const int ncols, const int nrows, const int nvecs) {
4
const auto row = blockIdx.x*blockDim.y + threadIdx.y;
5
+ const auto vec = blockIdx.y;
6
6
- if (row >= nrows) {
7
+ if (row >= nrows || vec >= nvecs) {
8
return;
9
}
10
11
const int blocks_per_row = ncols / qk;
12
const int blocks_per_warp = vdr * WARP_SIZE / qi;
13
+ const int nrows_y = (ncols + 512 - 1) / 512 * 512;
14
13
-// partial sum for each thread
15
+
16
+ // partial sum for each thread
17
float tmp = 0.0f;
18
19
const block_q_t * x = (const block_q_t *) vx;
@@ -19,7 +22,7 @@ static __global__ void mul_mat_vec_q(const void * __restrict__ vx, const void *
22
for (auto i = threadIdx.x / (qi/vdr); i < blocks_per_row; i += blocks_per_warp) {
23
const int ibx = row*blocks_per_row + i; // x block index
24
22
- const int iby = i * (qk/QK8_1); // y block index that aligns with ibx
25
+ const int iby = vec*(nrows_y/QK8_1) + i * (qk/QK8_1); // y block index that aligns with ibx
26
27
const int iqs = vdr * (threadIdx.x % (qi/vdr)); // x block quant index when casting the quants to int
28
@@ -33,177 +36,177 @@ static __global__ void mul_mat_vec_q(const void * __restrict__ vx, const void *
36
}
37
38
if (threadIdx.x == 0) {
36
- dst[row] = tmp;
39
+ dst[vec*nrows + row] = tmp;
40
}
41
}
42
43
template<typename scalar_t>
41
-static void mul_mat_vec_q4_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
44
+static void mul_mat_vec_q4_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
45
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
43
- const dim3 block_nums(block_num_y, 1, 1);
46
+ const dim3 block_nums(block_num_y, nvecs, 1);
47
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
48
mul_mat_vec_q<scalar_t, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1>
46
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
49
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
50
}
51
52
template<typename scalar_t>
50
-static void mul_mat_vec_q4_1_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
53
+static void mul_mat_vec_q4_1_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
54
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
52
- const dim3 block_nums(block_num_y, 1, 1);
55
+ const dim3 block_nums(block_num_y, nvecs, 1);
56
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
57
mul_mat_vec_q<scalar_t, QK4_0, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1>
55
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
58
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
59
}
60
61
template<typename scalar_t>
59
-static void mul_mat_vec_q5_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
62
+static void mul_mat_vec_q5_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
63
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
61
- const dim3 block_nums(block_num_y, 1, 1);
64
+ const dim3 block_nums(block_num_y, nvecs, 1);
65
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
66
mul_mat_vec_q<scalar_t, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1>
64
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
67
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
68
}
69
70
template<typename scalar_t>
68
-static void mul_mat_vec_q5_1_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
71
+static void mul_mat_vec_q5_1_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
72
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
70
- const dim3 block_nums(block_num_y, 1, 1);
73
+ const dim3 block_nums(block_num_y, nvecs, 1);
74
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
75
mul_mat_vec_q<scalar_t, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1>
73
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
76
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
77
}
78
79
template<typename scalar_t>
77
-static void mul_mat_vec_q8_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
80
+static void mul_mat_vec_q8_0_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
81
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
79
- const dim3 block_nums(block_num_y, 1, 1);
82
+ const dim3 block_nums(block_num_y, nvecs, 1);
83
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
84
mul_mat_vec_q<scalar_t, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1>
82
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
85
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
86
}
87
88
template<typename scalar_t>
86
-static void mul_mat_vec_q2_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
89
+static void mul_mat_vec_q2_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
90
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
88
- const dim3 block_nums(block_num_y, 1, 1);
91
+ const dim3 block_nums(block_num_y, nvecs, 1);
92
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
93
mul_mat_vec_q<scalar_t, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1>
91
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
94
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
95
}
96
97
template<typename scalar_t>
95
-static void mul_mat_vec_q3_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
98
+static void mul_mat_vec_q3_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
99
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
97
- const dim3 block_nums(block_num_y, 1, 1);
100
+ const dim3 block_nums(block_num_y, nvecs, 1);
101
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
102
mul_mat_vec_q<scalar_t, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1>
100
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
103
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
104
}
105
106
template<typename scalar_t>
104
-static void mul_mat_vec_q4_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
107
+static void mul_mat_vec_q4_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
108
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
106
- const dim3 block_nums(block_num_y, 1, 1);
109
+ const dim3 block_nums(block_num_y, nvecs, 1);
110
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
111
mul_mat_vec_q<scalar_t, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1>
109
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
112
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
113
}
114
115
template<typename scalar_t>
113
-static void mul_mat_vec_q5_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
116
+static void mul_mat_vec_q5_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
117
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
115
- const dim3 block_nums(block_num_y, 1, 1);
118
+ const dim3 block_nums(block_num_y, nvecs, 1);
119
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
120
mul_mat_vec_q<scalar_t, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1>
118
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
121
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
122
}
123
124
template<typename scalar_t>
122
-static void mul_mat_vec_q6_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
125
+static void mul_mat_vec_q6_K_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
126
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
124
- const dim3 block_nums(block_num_y, 1, 1);
127
+ const dim3 block_nums(block_num_y, nvecs, 1);
128
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
129
mul_mat_vec_q<scalar_t, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1>
127
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
130
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
131
}
132
133
template<typename scalar_t>
131
-static void mul_mat_vec_iq2_xxs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
134
+static void mul_mat_vec_iq2_xxs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
135
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
133
- const dim3 block_nums(block_num_y, 1, 1);
136
+ const dim3 block_nums(block_num_y, nvecs, 1);
137
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
138
mul_mat_vec_q<scalar_t, QK_K, QI2_XXS, block_iq2_xxs, 1, vec_dot_iq2_xxs_q8_1>
136
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
139
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
140
}
141
142
template<typename scalar_t>
140
-static void mul_mat_vec_iq2_xs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
143
+static void mul_mat_vec_iq2_xs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
144
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
142
- const dim3 block_nums(block_num_y, 1, 1);
145
+ const dim3 block_nums(block_num_y, nvecs, 1);
146
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
147
mul_mat_vec_q<scalar_t, QK_K, QI2_XS, block_iq2_xs, 1, vec_dot_iq2_xs_q8_1>
145
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
148
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
149
}
150
151
template<typename scalar_t>
149
-static void mul_mat_vec_iq2_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
152
+static void mul_mat_vec_iq2_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
153
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
151
- const dim3 block_nums(block_num_y, 1, 1);
154
+ const dim3 block_nums(block_num_y, nvecs, 1);
155
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
156
mul_mat_vec_q<scalar_t, QK_K, QI2_S, block_iq2_s, 1, vec_dot_iq2_s_q8_1>
154
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
157
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
158
}
159
160
template<typename scalar_t>
158
-static void mul_mat_vec_iq3_xxs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
161
+static void mul_mat_vec_iq3_xxs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
162
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
160
- const dim3 block_nums(block_num_y, 1, 1);
163
+ const dim3 block_nums(block_num_y, nvecs, 1);
164
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
165
mul_mat_vec_q<scalar_t, QK_K, QI3_XXS, block_iq3_xxs, 1, vec_dot_iq3_xxs_q8_1>
163
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
166
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
167
}
168
169
template<typename scalar_t>
167
-static void mul_mat_vec_iq1_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
170
+static void mul_mat_vec_iq1_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
171
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
169
- const dim3 block_nums(block_num_y, 1, 1);
172
+ const dim3 block_nums(block_num_y, nvecs, 1);
173
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
174
mul_mat_vec_q<scalar_t, QK_K, QI1_S, block_iq1_s, 1, vec_dot_iq1_s_q8_1>
172
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
175
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
176
}
177
178
template<typename scalar_t>
176
-static void mul_mat_vec_iq1_m_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
179
+static void mul_mat_vec_iq1_m_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
180
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
178
- const dim3 block_nums(block_num_y, 1, 1);
181
+ const dim3 block_nums(block_num_y, nvecs, 1);
182
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
183
mul_mat_vec_q<scalar_t, QK_K, QI1_M, block_iq1_m, 1, vec_dot_iq1_m_q8_1>
181
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
184
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
185
}
186
187
template<typename scalar_t>
185
-static void mul_mat_vec_iq4_nl_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
188
+static void mul_mat_vec_iq4_nl_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
189
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
187
- const dim3 block_nums(block_num_y, 1, 1);
190
+ const dim3 block_nums(block_num_y, nvecs, 1);
191
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
192
mul_mat_vec_q<scalar_t, QK4_NL, QI4_NL, block_iq4_nl, VDR_Q4_0_Q8_1_MMVQ, vec_dot_iq4_nl_q8_1>
190
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
193
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
194
}
195
196
template<typename scalar_t>
194
-static void mul_mat_vec_iq4_xs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
197
+static void mul_mat_vec_iq4_xs_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
198
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
196
- const dim3 block_nums(block_num_y, 1, 1);
199
+ const dim3 block_nums(block_num_y, nvecs, 1);
200
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
201
mul_mat_vec_q<scalar_t, QK_K, QI4_XS, block_iq4_xs, 1, vec_dot_iq4_xs_q8_1>
199
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
202
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
203
}
204
205
template<typename scalar_t>
203
-static void mul_mat_vec_iq3_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, cudaStream_t stream) {
206
+static void mul_mat_vec_iq3_s_q8_1_cuda(const void * vx, const void * vy, scalar_t * dst, const int ncols, const int nrows, const int nvecs, cudaStream_t stream) {
207
const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y;
205
- const dim3 block_nums(block_num_y, 1, 1);
208
+ const dim3 block_nums(block_num_y, nvecs, 1);
209
const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1);
210
mul_mat_vec_q<scalar_t, QK_K, QI3_XS, block_iq3_s, 1, vec_dot_iq3_s_q8_1>
208
- <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows);
211
+ <<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, ncols, nrows, nvecs);
212
}
vllm/_custom_ops.py
+1
-1
@@ -594,7 +594,7 @@ if hasattr(torch.ops._C, "ggml_dequantize"):
594
quant_type: int,
595
row: torch.SymInt,
596
) -> torch.Tensor:
597
- return torch.empty((1, row), dtype=X.dtype, device=W.device)
597
+ return torch.empty((X.shape[0], row), dtype=X.dtype, device=W.device)
598
599
@register_fake("_C::ggml_mul_mat_a8")
600
def _ggml_mul_mat_a8_fake(
vllm/model_executor/layers/quantization/gguf.py
+5
-1
@@ -99,6 +99,10 @@ MMQ_QUANT_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES
99
100
def _fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor,
101
qweight_type: int) -> torch.Tensor:
102
+ if qweight_type in IMATRIX_QUANT_TYPES:
103
+ mmvq_safe = 8 if qweight.shape[0] > 5120 else 16
104
+ else:
105
+ mmvq_safe = 2 if qweight.shape[0] > 5120 else 6
106
# HACK: when doing chunked prefill we don't generate output tokens
107
# so input to logits generator is empty which causes invalid parameter
108
if x.shape[0] == 0:
@@ -110,7 +114,7 @@ def _fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor,
114
if qweight_type in UNQUANTIZED_TYPES:
115
return x @ qweight.T
116
# enable MMVQ in contiguous batching with batch_size=1
113
- if x.shape[0] == 1 and qweight_type in MMVQ_QUANT_TYPES:
117
+ if x.shape[0] <= mmvq_safe and qweight_type in MMVQ_QUANT_TYPES:
118
y = ops.ggml_mul_mat_vec_a8(qweight, x, qweight_type, qweight.shape[0])
119
# Use MMQ Kernel if it's available (standard + k-quants)
120
elif qweight_type in MMQ_QUANT_TYPES: