Skip to content

Commit 269a067

Browse files
NexesenexNexesenex
committed
Reapply commit "CUDA: batched+noncont MMQ, refactor bs>1 MoE code (ggml-org#13199)" getrows
1 parent 4ff70ea commit 269a067

File tree

2 files changed

+53
-47
lines changed

2 files changed

+53
-47
lines changed

‎ggml/src/ggml-cuda/getrows.cu‎

Lines changed: 46 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,8 @@ static __global__ void k_get_rows(
3434
dfloat2 v;
3535
dequantize_kernel(src0_row, ib, iqs, v);
3636

37-
dst_row[iybs + iqs + 0] = v.x;
38-
dst_row[iybs + iqs + y_offset] = v.y;
37+
dst_row[iybs + iqs + 0] = float(v.x);
38+
dst_row[iybs + iqs + y_offset] = float(v.y);
3939
}
4040

4141
template<typename src0_t, typename dst_t>
@@ -62,7 +62,7 @@ static __global__ void k_get_rows_float(
6262
dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
6363
const src0_t * src0_row = (const src0_t *)((const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03);
6464

65-
dst_row[i00] = src0_row[i00];
65+
dst_row[i00] = float(src0_row[i00]);
6666
}
6767

6868
template<typename grad_t, typename dst_t>
@@ -88,74 +88,68 @@ static __global__ void k_get_rows_back_float(
8888
dst[dst_row*ncols + col] = sum;
8989
}
9090

91-
template<int qk, int qr, dequantize_kernel_t dq>
92-
static void get_rows_cuda(
93-
const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
94-
const void * src0_dd, const int32_t * src1_dd, float * dst_dd, cudaStream_t stream) {
95-
96-
GGML_TENSOR_BINARY_OP_LOCALS
97-
91+
template<int qk, int qr, dequantize_kernel_t dq, typename dst_t>
92+
static void get_rows_cuda_q(
93+
const void * src0_d, const int32_t * src1_d, dst_t * dst_d,
94+
const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03,
95+
const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12,
96+
const size_t nb1, const size_t nb2, const size_t nb3,
97+
cudaStream_t stream) {
9898
const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
9999
const int block_num_y = (ne00 + 2*CUDA_GET_ROWS_BLOCK_SIZE - 1) / (2*CUDA_GET_ROWS_BLOCK_SIZE);
100100
const dim3 block_nums(ne10, block_num_y, ne11*ne12);
101101

102102
// strides in elements
103-
//const size_t s0 = nb0 / ggml_element_size(dst);
104-
const size_t s1 = nb1 / ggml_element_size(dst);
105-
const size_t s2 = nb2 / ggml_element_size(dst);
106-
const size_t s3 = nb3 / ggml_element_size(dst);
103+
// const size_t s0 = nb0 / sizeof(dst_t);
104+
const size_t s1 = nb1 / sizeof(dst_t);
105+
const size_t s2 = nb2 / sizeof(dst_t);
106+
const size_t s3 = nb3 / sizeof(dst_t);
107107

108-
const size_t s10 = nb10 / ggml_element_size(src1);
109-
const size_t s11 = nb11 / ggml_element_size(src1);
110-
const size_t s12 = nb12 / ggml_element_size(src1);
111-
//const size_t s13 = nb13 / ggml_element_size(src1);
108+
const size_t s10 = nb10 / sizeof(int32_t);
109+
const size_t s11 = nb11 / sizeof(int32_t);
110+
const size_t s12 = nb12 / sizeof(int32_t);
111+
// const size_t s13 = nb13 / sizeof(int32_t);
112112

113113
GGML_ASSERT(ne00 % 2 == 0);
114114

115115
k_get_rows<qk, qr, dq><<<block_nums, block_dims, 0, stream>>>(
116-
src0_dd, src1_dd, dst_dd,
116+
src0_d, src1_d, dst_d,
117117
ne00, /*ne01, ne02, ne03,*/
118118
/*ne10, ne11,*/ ne12, /*ne13,*/
119119
/* s0,*/ s1, s2, s3,
120120
/* nb00,*/ nb01, nb02, nb03,
121121
s10, s11, s12/*, s13*/);
122-
123-
GGML_UNUSED(dst);
124122
}
125123

126-
template<typename src0_t>
124+
template<typename src0_t, typename dst_t>
127125
static void get_rows_cuda_float(
128-
const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
129-
const src0_t * src0_dd, const int32_t * src1_dd, float * dst_dd, cudaStream_t stream) {
130-
131-
GGML_TENSOR_BINARY_OP_LOCALS
132-
133-
GGML_ASSERT(ne13 == 1);
134-
126+
const src0_t * src0_d, const int32_t * src1_d, dst_t * dst_d,
127+
const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03,
128+
const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12,
129+
const size_t nb1, const size_t nb2, const size_t nb3,
130+
cudaStream_t stream) {
135131
const dim3 block_dims(CUDA_GET_ROWS_BLOCK_SIZE, 1, 1);
136132
const int block_num_y = (ne00 + CUDA_GET_ROWS_BLOCK_SIZE - 1) / CUDA_GET_ROWS_BLOCK_SIZE;
137133
const dim3 block_nums(ne10, block_num_y, ne11*ne12);
138134

139135
// strides in elements
140-
//const size_t s0 = nb0 / ggml_element_size(dst);
141-
const size_t s1 = nb1 / ggml_element_size(dst);
142-
const size_t s2 = nb2 / ggml_element_size(dst);
143-
const size_t s3 = nb3 / ggml_element_size(dst);
136+
// const size_t s0 = nb0 / sizeof(dst_t);
137+
const size_t s1 = nb1 / sizeof(dst_t);
138+
const size_t s2 = nb2 / sizeof(dst_t);
139+
const size_t s3 = nb3 / sizeof(dst_t);
144140

145-
const size_t s10 = nb10 / ggml_element_size(src1);
146-
const size_t s11 = nb11 / ggml_element_size(src1);
147-
const size_t s12 = nb12 / ggml_element_size(src1);
148-
//const size_t s13 = nb13 / ggml_element_size(src1);
141+
const size_t s10 = nb10 / sizeof(int32_t);
142+
const size_t s11 = nb11 / sizeof(int32_t);
143+
const size_t s12 = nb12 / sizeof(int32_t);
144+
// const size_t s13 = nb13 / sizeof(int32_t);
149145

150146
k_get_rows_float<<<block_nums, block_dims, 0, stream>>>(
151-
src0_dd, src1_dd, dst_dd,
147+
src0_d, src1_d, dst_d,
152148
ne00, /*ne01, ne02, ne03,*/
153149
/*ne10, ne11,*/ ne12, /*ne13,*/
154150
/* s0,*/ s1, s2, s3,
155151
/* nb00,*/ nb01, nb02, nb03,
156152
s10, s11, s12/*, s13*/);
157-
158-
GGML_UNUSED(dst);
159153
}
160154

161155
template <typename dst_t>
@@ -198,6 +192,10 @@ static void ggml_cuda_get_rows_switch_src0_type(
198192
get_rows_cuda_q<QK5_1, QR5_1, dequantize_q5_1>(src0_d, src1_d, dst_d,
199193
ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
200194
break;
195+
// case GGML_TYPE_Q6_0:
196+
// get_rows_cuda_q<QK6_0, QR6_0, dequantize_q6_0>(src0_d, src1_d, dst_d,
197+
// ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
198+
// break;
201199
case GGML_TYPE_Q8_0:
202200
get_rows_cuda_q<QK8_0, QR8_0, dequantize_q8_0>(src0_d, src1_d, dst_d,
203201
ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
@@ -242,20 +240,18 @@ void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
242240
const ggml_tensor * src0 = dst->src[0];
243241
const ggml_tensor * src1 = dst->src[1];
244242

245-
const void * src0_d = (const void *) src0->data;
246-
const int32_t * src1_d = (const int32_t *) src1->data;
247-
float * dst_d = (float *) dst->data;
248-
249243
cudaStream_t stream = ctx.stream();
250244

245+
GGML_TENSOR_BINARY_OP_LOCALS
246+
251247
GGML_ASSERT(src1->type == GGML_TYPE_I32);
252-
GGML_ASSERT(dst->type == GGML_TYPE_F32);
248+
GGML_ASSERT(ne13 == 1);
253249

254250
GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type));
255251
GGML_ASSERT(src1->nb[0] == ggml_type_size(src1->type));
256252
GGML_ASSERT(dst->nb[0] == ggml_type_size(dst->type));
257253

258-
switch (src0->type) {
254+
/* switch (src0->type) {
259255
case GGML_TYPE_F16:
260256
get_rows_cuda_float(src0, src1, dst, (const half *) src0_d, src1_d, dst_d, stream);
261257
break;
@@ -284,7 +280,10 @@ void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
284280
// TODO: k-quants
285281
GGML_ABORT("%s: unsupported type: %s\n", __func__, ggml_type_name(src0->type));
286282
break;
287-
}
283+
} */
284+
285+
get_rows_cuda(src0->data, src0->type, (const int32_t *) src1->data, dst->data, dst->type,
286+
ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
288287
}
289288

290289
void ggml_cuda_op_get_rows_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {

‎ggml/src/ggml-cuda/getrows.cuh‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,13 @@
33
#define CUDA_GET_ROWS_BLOCK_SIZE 256
44
#define CUDA_GET_ROWS_BACK_BLOCK_SIZE 256
55

6+
void get_rows_cuda(
7+
const void * src0_d, ggml_type src0_type, const int32_t * src1_d, void * dst_d, ggml_type dst_type,
8+
int64_t ne00, size_t nb01, size_t nb02, size_t nb03,
9+
int64_t ne10, int64_t ne11, int64_t ne12, size_t nb10, size_t nb11, size_t nb12,
10+
size_t nb1, size_t nb2, size_t nb3,
11+
cudaStream_t stream);
12+
613
void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
714

815
void ggml_cuda_op_get_rows_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

0 commit comments

Comments
 (0)