char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

fairydreamingcuda : add support for f16->f16 GGML_OP_SET_ROWS (#25367)ed8c26150

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
15.0 KiB398 linesraw
1#include "set-rows.cuh"
2#include "cpy-utils.cuh"
3
4typedef void (*set_rows_kernel_t)(const char * src, char * dst);
5
6// Generic quantized set_rows kernel template
7template <typename idx_t, typename block_type, int qk, void (*quantize_func)(const float *, block_type *)>
8static __global__ void k_set_rows_quant(const float * __restrict__ src0,
9                                        const idx_t * __restrict__ src1,
10                                        block_type * __restrict__ dst,
11                                        const int64_t ne_total,
12                                        const int64_t ne10,
13                                        const int64_t ne11,
14                                        const int64_t ne12,
15                                        const int64_t ne13,
16                                        const int64_t s01,
17                                        const int64_t s02,
18                                        const int64_t s03,
19                                        const int64_t s10,
20                                        const int64_t s11,
21                                        const int64_t s12,
22                                        const int64_t s1,
23                                        const int64_t s2,
24                                        const int64_t s3,
25                                        const uint3   ne00,
26                                        const uint3   ne01,
27                                        const uint3   ne02,
28                                        const uint3   ne11_fd,
29                                        const uint3   ne12_fd) {
30    const int64_t i = int64_t(blockDim.x) * blockIdx.x + threadIdx.x;
31
32    if (i >= ne_total) {
33        return;
34    }
35
36    const int64_t i_base = i * qk;
37    uint32_t      tmp    = (uint32_t) i_base;
38    uint2         div_mod;
39
40    div_mod           = fast_div_modulo(tmp, ne00);
41    const int64_t i00 = div_mod.y;
42    tmp               = div_mod.x;
43
44    div_mod           = fast_div_modulo(tmp, ne01);
45    const int64_t i01 = div_mod.y;
46    tmp               = div_mod.x;
47
48    div_mod           = fast_div_modulo(tmp, ne02);
49    const int64_t i02 = div_mod.y;
50    const int64_t i03 = div_mod.x;
51
52    const int64_t i12 = fastmodulo((uint32_t) i03, ne12_fd);
53    const int64_t i11 = fastmodulo((uint32_t) i02, ne11_fd);
54    const int64_t i10 = i01;
55
56    ggml_cuda_pdl_sync();
57    const int64_t dst_row = *(src1 + i10*s10 + i11*s11 + i12*s12);
58
59    const float * src0_row = src0 + i01*s01 + i02*s02 + i03*s03;
60    block_type * dst_row_ptr = dst + (dst_row*s1 + i02*s2 + i03*s3) / sizeof(block_type);
61
62    const float * src_block = src0_row + i00;
63    block_type * dst_block = dst_row_ptr + i00 / qk;
64
65    quantize_func(src_block, dst_block);
66
67    GGML_UNUSED(ne10);
68    GGML_UNUSED(ne11);
69    GGML_UNUSED(ne12);
70    GGML_UNUSED(ne13);
71}
72
73// Template dispatch function for quantized set_rows
74template<typename idx_t, typename block_type, int qk, void (*quantize_func)(const float*, block_type*)>
75static void set_rows_cuda_quant(
76        const float * src0_d, const idx_t * src1_d, block_type * dst_d,
77        const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
78        const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13,
79        const size_t nb01, const size_t nb02, const size_t nb03,
80        const size_t nb10, const size_t nb11, const size_t nb12,
81        const size_t nb1, const size_t nb2, const size_t nb3,
82        cudaStream_t stream) {
83
84    GGML_ASSERT(ne00 % qk == 0);
85    const int64_t ne_total = (ne00 * ne01 * ne02 * ne03) / qk;
86    const int num_blocks = (ne_total + CUDA_SET_ROWS_BLOCK_SIZE - 1) / CUDA_SET_ROWS_BLOCK_SIZE;
87    const dim3 block_size(CUDA_SET_ROWS_BLOCK_SIZE);
88    const dim3 grid_size(num_blocks);
89
90    const int64_t s01 = nb01/sizeof(float);
91    const int64_t s02 = nb02/sizeof(float);
92    const int64_t s03 = nb03/sizeof(float);
93    const int64_t s10 = nb10/sizeof(idx_t);
94    const int64_t s11 = nb11/sizeof(idx_t);
95    const int64_t s12 = nb12/sizeof(idx_t);
96    const int64_t s1  = nb1;
97    const int64_t s2  = nb2;
98    const int64_t s3  = nb3;
99
100    if (ne_total > 0 && ne00 > 0 && ne01 > 0 && ne02 > 0 && ne11 > 0 && ne12 > 0) {
101        const uint3 ne00_fd = init_fastdiv_values((uint32_t) ne00);
102        const uint3 ne01_fd = init_fastdiv_values((uint32_t) ne01);
103        const uint3 ne02_fd = init_fastdiv_values((uint32_t) ne02);
104        const uint3 ne11_fd = init_fastdiv_values((uint32_t) ne11);
105        const uint3 ne12_fd = init_fastdiv_values((uint32_t) ne12);
106
107        k_set_rows_quant<idx_t, block_type, qk, quantize_func><<<grid_size, block_size, 0, stream>>>(
108            src0_d, src1_d, dst_d, ne_total, ne10, ne11, ne12, ne13, s01, s02, s03, s10, s11, s12, s1, s2, s3, ne00_fd,
109            ne01_fd, ne02_fd, ne11_fd, ne12_fd);
110    }
111}
112
113template <typename src_t, typename idx_t, typename dst_t>
114static __global__ void k_set_rows(const src_t * src0_ptr,
115                                  const idx_t * src1_ptr,
116                                  dst_t * dst_ptr,
117                                  const int64_t ne_total,
118                                  const int64_t ne10,
119                                  const int64_t ne11,
120                                  const int64_t ne12,
121                                  const int64_t ne13,
122                                  const int64_t s01,
123                                  const int64_t s02,
124                                  const int64_t s03,
125                                  const int64_t s10,
126                                  const int64_t s11,
127                                  const int64_t s12,
128                                  const int64_t s1,
129                                  const int64_t s2,
130                                  const int64_t s3,
131                                  const uint3   ne00,
132                                  const uint3   ne01,
133                                  const uint3   ne02,
134                                  const uint3   ne11_fd,
135                                  const uint3   ne12_fd) {
136    const src_t * GGML_CUDA_RESTRICT src0 = src0_ptr;
137    const idx_t * GGML_CUDA_RESTRICT src1 = src1_ptr;
138    dst_t       * GGML_CUDA_RESTRICT dst  = dst_ptr;
139    const int64_t i = int64_t(blockDim.x) * blockIdx.x + threadIdx.x;
140
141    if (i >= ne_total) {
142        return;
143    }
144
145    uint32_t tmp = (uint32_t) i;
146    uint2    div_mod;
147
148    div_mod           = fast_div_modulo(tmp, ne00);
149    const int64_t i00 = div_mod.y;
150    tmp               = div_mod.x;
151
152    div_mod           = fast_div_modulo(tmp, ne01);
153    const int64_t i01 = div_mod.y;
154    tmp               = div_mod.x;
155
156    div_mod           = fast_div_modulo(tmp, ne02);
157    const int64_t i02 = div_mod.y;
158    const int64_t i03 = div_mod.x;
159
160    const int64_t i12 = fastmodulo((uint32_t) i03, ne12_fd);
161    const int64_t i11 = fastmodulo((uint32_t) i02, ne11_fd);
162    const int64_t i10 = i01;
163
164    ggml_cuda_pdl_sync();
165    const int64_t dst_row = *(src1 + i10*s10 + i11*s11 + i12*s12);
166    ggml_cuda_pdl_lc();
167
168    const src_t * src0_row = src0 + i01*s01 + i02*s02 + i03*s03;
169    dst_t * dst_row_ptr    = dst + dst_row*s1 + i02*s2 + i03*s3;
170
171    dst_row_ptr[i00] = ggml_cuda_cast<dst_t>(src0_row[i00]);
172
173    GGML_UNUSED(ne10);
174    GGML_UNUSED(ne11);
175    GGML_UNUSED(ne12);
176    GGML_UNUSED(ne13);
177}
178
179template<typename src_t, typename idx_t, typename dst_t>
180static void set_rows_cuda(
181        const src_t * src0_d, const idx_t * src1_d, dst_t * dst_d,
182        const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
183        const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t ne13,
184        const size_t nb01, const size_t nb02, const size_t nb03,
185        const size_t nb10, const size_t nb11, const size_t nb12,
186        const size_t nb1, const size_t nb2, const size_t nb3,
187        cudaStream_t stream) {
188
189    const int64_t ne_total = ne00 * ne01 * ne02 * ne03;
190    const int num_blocks = (ne_total + CUDA_SET_ROWS_BLOCK_SIZE - 1) / CUDA_SET_ROWS_BLOCK_SIZE;
191    const dim3 block_size(CUDA_SET_ROWS_BLOCK_SIZE);
192    const dim3 grid_size(num_blocks);
193
194
195    const int64_t s01 = nb01/sizeof(src_t);
196    const int64_t s02 = nb02/sizeof(src_t);
197    const int64_t s03 = nb03/sizeof(src_t);
198    const int64_t s10 = nb10/sizeof(idx_t);
199    const int64_t s11 = nb11/sizeof(idx_t);
200    const int64_t s12 = nb12/sizeof(idx_t);
201    const int64_t s1  = nb1/sizeof(dst_t);
202    const int64_t s2  = nb2/sizeof(dst_t);
203    const int64_t s3  = nb3/sizeof(dst_t);
204
205    if (ne_total > 0 && ne00 > 0 && ne01 > 0 && ne02 > 0 && ne11 > 0 && ne12 > 0) {
206        const uint3 ne00_fd = init_fastdiv_values((uint32_t) ne00);
207        const uint3 ne01_fd = init_fastdiv_values((uint32_t) ne01);
208        const uint3 ne02_fd = init_fastdiv_values((uint32_t) ne02);
209        const uint3 ne11_fd = init_fastdiv_values((uint32_t) ne11);
210        const uint3 ne12_fd = init_fastdiv_values((uint32_t) ne12);
211
212        const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_size, block_size, 0, stream);
213        ggml_cuda_kernel_launch(k_set_rows<src_t, idx_t, dst_t>, launch_params,
214            src0_d, src1_d, dst_d, ne_total, ne10, ne11, ne12, ne13, s01,
215            s02, s03, s10, s11, s12, s1, s2, s3, ne00_fd, ne01_fd, ne02_fd,
216            ne11_fd, ne12_fd);
217    }
218}
219
220template<typename src_t, typename idx_t>
221static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
222    const src_t * src0_d = (const src_t *)src0->data;
223    const idx_t * src1_d = (const idx_t *)src1->data;
224
225    GGML_TENSOR_BINARY_OP_LOCALS
226
227    cudaStream_t stream = ctx.stream();
228
229
230    if (dst->type == GGML_TYPE_F32) {
231        set_rows_cuda(
232            src0_d, src1_d, (float*)dst->data,
233            ne00, ne01, ne02, ne03,
234            ne10, ne11, ne12, ne13,
235            nb01, nb02, nb03,
236            nb10, nb11, nb12,
237            nb1, nb2, nb3,
238            stream
239        );
240    } else if (dst->type == GGML_TYPE_F16) {
241        set_rows_cuda(
242            src0_d, src1_d, (half*)dst->data,
243            ne00, ne01, ne02, ne03,
244            ne10, ne11, ne12, ne13,
245            nb01, nb02, nb03,
246            nb10, nb11, nb12,
247            nb1, nb2, nb3,
248            stream
249        );
250    } else if (dst->type == GGML_TYPE_BF16) {
251        set_rows_cuda(
252            src0_d, src1_d, (nv_bfloat16*)dst->data,
253            ne00, ne01, ne02, ne03,
254            ne10, ne11, ne12, ne13,
255            nb01, nb02, nb03,
256            nb10, nb11, nb12,
257            nb1, nb2, nb3,
258            stream
259        );
260    } else if (dst->type == GGML_TYPE_Q4_0) {
261        set_rows_cuda_quant<idx_t, block_q4_0, QK4_0, quantize_f32_q4_0_block>(
262            src0_d, src1_d, (block_q4_0*)dst->data,
263            ne00, ne01, ne02, ne03,
264            ne10, ne11, ne12, ne13,
265            nb01, nb02, nb03,
266            nb10, nb11, nb12,
267            nb1, nb2, nb3,
268            stream
269        );
270    } else if (dst->type == GGML_TYPE_Q4_1) {
271        set_rows_cuda_quant<idx_t, block_q4_1, QK4_1, quantize_f32_q4_1_block>(
272            src0_d, src1_d, (block_q4_1*)dst->data,
273            ne00, ne01, ne02, ne03,
274            ne10, ne11, ne12, ne13,
275            nb01, nb02, nb03,
276            nb10, nb11, nb12,
277            nb1, nb2, nb3,
278            stream
279        );
280    } else if (dst->type == GGML_TYPE_Q5_0) {
281        set_rows_cuda_quant<idx_t, block_q5_0, QK5_0, quantize_f32_q5_0_block>(
282            src0_d, src1_d, (block_q5_0*)dst->data,
283            ne00, ne01, ne02, ne03,
284            ne10, ne11, ne12, ne13,
285            nb01, nb02, nb03,
286            nb10, nb11, nb12,
287            nb1, nb2, nb3,
288            stream
289        );
290    } else if (dst->type == GGML_TYPE_Q5_1) {
291        set_rows_cuda_quant<idx_t, block_q5_1, QK5_1, quantize_f32_q5_1_block>(
292            src0_d, src1_d, (block_q5_1*)dst->data,
293            ne00, ne01, ne02, ne03,
294            ne10, ne11, ne12, ne13,
295            nb01, nb02, nb03,
296            nb10, nb11, nb12,
297            nb1, nb2, nb3,
298            stream
299        );
300    } else if (dst->type == GGML_TYPE_Q8_0) {
301        set_rows_cuda_quant<idx_t, block_q8_0, QK8_0, quantize_f32_q8_0_block>(
302            src0_d, src1_d, (block_q8_0*)dst->data,
303            ne00, ne01, ne02, ne03,
304            ne10, ne11, ne12, ne13,
305            nb01, nb02, nb03,
306            nb10, nb11, nb12,
307            nb1, nb2, nb3,
308            stream
309        );
310    } else if (dst->type == GGML_TYPE_IQ4_NL) {
311        set_rows_cuda_quant<idx_t, block_iq4_nl, QK4_NL, quantize_f32_iq4_nl_block>(
312            src0_d, src1_d, (block_iq4_nl*)dst->data,
313            ne00, ne01, ne02, ne03,
314            ne10, ne11, ne12, ne13,
315            nb01, nb02, nb03,
316            nb10, nb11, nb12,
317            nb1, nb2, nb3,
318            stream
319        );
320    } else {
321        GGML_ABORT("unsupported type %s", ggml_type_name(dst->type));
322    }
323}
324
325template<>
326void set_rows_cuda<half, int32_t>(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
327    const half    * src0_d = (const half *)src0->data;
328    const int32_t * src1_d = (const int32_t *)src1->data;
329
330    GGML_TENSOR_BINARY_OP_LOCALS
331
332    cudaStream_t stream = ctx.stream();
333
334
335    if (dst->type == GGML_TYPE_F16) {
336        set_rows_cuda(
337            src0_d, src1_d, (half*)dst->data,
338            ne00, ne01, ne02, ne03,
339            ne10, ne11, ne12, ne13,
340            nb01, nb02, nb03,
341            nb10, nb11, nb12,
342            nb1, nb2, nb3,
343            stream
344        );
345    } else {
346        GGML_ABORT("unsupported type %s", ggml_type_name(dst->type));
347    }
348}
349
350template<>
351void set_rows_cuda<half, int64_t>(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
352    const half    * src0_d = (const half *)src0->data;
353    const int64_t * src1_d = (const int64_t *)src1->data;
354
355    GGML_TENSOR_BINARY_OP_LOCALS
356
357    cudaStream_t stream = ctx.stream();
358
359
360    if (dst->type == GGML_TYPE_F16) {
361        set_rows_cuda(
362            src0_d, src1_d, (half*)dst->data,
363            ne00, ne01, ne02, ne03,
364            ne10, ne11, ne12, ne13,
365            nb01, nb02, nb03,
366            nb10, nb11, nb12,
367            nb1, nb2, nb3,
368            stream
369        );
370    } else {
371        GGML_ABORT("unsupported type %s", ggml_type_name(dst->type));
372    }
373}
374
375
376void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
377    const ggml_tensor * src0 = dst->src[0];
378    const ggml_tensor * src1 = dst->src[1];
379
380    GGML_ASSERT(src0->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16));
381    GGML_ASSERT(src1->type == GGML_TYPE_I64 || src1->type == GGML_TYPE_I32);
382
383    if (src0->type == GGML_TYPE_F32) {
384        if (src1->type == GGML_TYPE_I64) {
385            set_rows_cuda<float, int64_t>(ctx, src0, src1, dst);
386        } else {
387            set_rows_cuda<float, int32_t>(ctx, src0, src1, dst);
388        }
389    } else if (src0->type == GGML_TYPE_F16) {
390        if (src1->type == GGML_TYPE_I64) {
391            set_rows_cuda<half, int64_t>(ctx, src0, src1, dst);
392        } else {
393            set_rows_cuda<half, int32_t>(ctx, src0, src1, dst);
394        }
395    } else {
396        GGML_ABORT("unsupported type %s", ggml_type_name(src0->type));
397    }
398}