char-fork/llama.cpp

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

Aman Guptaggml: add SWIGLU_CLAMP (#27930)0190529ec

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
24.7 KiB721 linesraw
1#include "unary.cuh"
2#include "convert.cuh"
3
4static __device__ __forceinline__ float op_abs(float x) {
5    return fabsf(x);
6}
7
8static __device__ __forceinline__ float op_sgn(float x) {
9    return (x > 0.f ? 1.f : ((x < 0.f ? -1.f : 0.f)));
10}
11
12static __device__ __forceinline__ float op_neg(float x) {
13    return -x;
14}
15
16static __device__ __forceinline__ float op_step(float x) {
17    return x > 0.0f;
18}
19
20static __device__ __forceinline__ float op_gelu(float x) {
21    return ggml_cuda_op_gelu_single(x);
22}
23
24static __device__ __forceinline__ float op_gelu_erf(float x) {
25    const float SQRT_2_INV = 0.70710678118654752440084436210484f;
26
27    return 0.5f*x*(1.0f + erff(x*SQRT_2_INV));
28}
29
30static __device__ __forceinline__ float op_gelu_quick(float x) {
31    const float GELU_QUICK_COEF = -1.702f;
32
33    return x * (1.0f / (1.0f + expf(GELU_QUICK_COEF * x)));
34}
35
36static __device__ __forceinline__ float op_silu(float x) {
37    return ggml_cuda_op_silu_single(x);
38}
39
40static __device__ __forceinline__ float op_tanh(float x) {
41    return tanhf(x);
42}
43
44static __device__ __forceinline__ float op_relu(float x) {
45    return fmaxf(x, 0);
46}
47
48static __device__ __forceinline__ float op_sigmoid(float x) {
49    return 1.0f / (1.0f + expf(-x));
50}
51
52static __device__ __forceinline__ float op_hardsigmoid(float x) {
53    return fminf(1.0f, fmaxf(0.0f, (x + 3.0f) / 6.0f));
54}
55
56static __device__ __forceinline__ float op_hardswish(float x) {
57    return x * fminf(1.0f, fmaxf(0.0f, (x + 3.0f) / 6.0f));
58}
59
60static __device__ __forceinline__ float op_exp(float x) {
61    return expf(x);
62}
63
64static __device__ __forceinline__ float op_sqr(float x) {
65    return x * x;
66}
67
68static __device__ __forceinline__ float op_relu_sqr(float x) {
69    const float r = fmaxf(x, 0.0f);
70    return r * r;
71}
72
73static __device__ __forceinline__ float op_sqrt(float x) {
74    return sqrtf(x);
75}
76
77static __device__ __forceinline__ float op_sin(float x) {
78    return sinf(x);
79}
80
81static __device__ __forceinline__ float op_cos(float x) {
82    return cosf(x);
83}
84
85static __device__ __forceinline__ float op_log(float x) {
86    return logf(x);
87}
88
89static __device__ __forceinline__ float op_expm1(float x) {
90    return expm1f(x);
91}
92
93static __device__ __forceinline__ float op_softplus(float x) {
94    return (x > 20.0f) ? x : logf(1.0f + expf(x));
95}
96
97static __device__ __forceinline__ float op_elu(float x) {
98    return (x > 0.f) ? x : expm1f(x);
99}
100
101static __device__ __forceinline__ float op_floor(float x) {
102    return floorf(x);
103}
104
105static __device__ __forceinline__ float op_ceil(float x) {
106    return ceilf(x);
107}
108
109static __device__ __forceinline__ float op_round(float x) {
110    return round(x);
111}
112
113static __device__ __forceinline__ float op_trunc(float x) {
114    return trunc(x);
115}
116
117template <float (*op)(float), typename T>
118static __global__ void unary_op_kernel(const T * x, T * dst, const int k) {
119    ggml_cuda_pdl_lc();
120    const int i = blockDim.x*blockIdx.x + threadIdx.x;
121
122    if (i >= k) {
123        return;
124    }
125
126    ggml_cuda_pdl_sync();
127    dst[i] = (T)op((float)x[i]);
128}
129
130template <float (*op)(float), typename T>
131static void unary_cuda(const T * x, T * dst, const int k, cudaStream_t stream) {
132    const int num_blocks = (k + CUDA_NEG_BLOCK_SIZE - 1) / CUDA_NEG_BLOCK_SIZE;
133    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params((dim3)num_blocks, CUDA_NEG_BLOCK_SIZE, 0, stream);
134    ggml_cuda_kernel_launch(unary_op_kernel<op, T>, launch_params, x, dst, k);
135}
136
137template <float (*op)(float)>
138void ggml_cuda_op_unary(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
139    const ggml_tensor * src0 = dst->src[0];
140    const void * src0_d = src0->data;
141    void * dst_d = dst->data;
142    cudaStream_t stream = ctx.stream();
143
144    GGML_ASSERT(ggml_is_contiguous(src0));
145
146    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
147    GGML_ASSERT( dst->type == GGML_TYPE_F32 ||  dst->type == GGML_TYPE_F16);
148    GGML_ASSERT(src0->type == dst->type);
149
150    if (src0->type == GGML_TYPE_F16) {
151        unary_cuda<op>((const half *)src0_d, (half *)dst_d, ggml_nelements(src0), stream);
152    } else {
153        unary_cuda<op>((const float *)src0_d, (float *)dst_d, ggml_nelements(src0), stream);
154    }
155}
156
157void ggml_cuda_op_abs(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
158    ggml_cuda_op_unary<op_abs>(ctx, dst);
159}
160
161void ggml_cuda_op_sgn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
162    ggml_cuda_op_unary<op_sgn>(ctx, dst);
163}
164
165void ggml_cuda_op_neg(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
166    ggml_cuda_op_unary<op_neg>(ctx, dst);
167}
168
169void ggml_cuda_op_step(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
170    ggml_cuda_op_unary<op_step>(ctx, dst);
171}
172
173void ggml_cuda_op_gelu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
174    ggml_cuda_op_unary<op_gelu>(ctx, dst);
175}
176
177void ggml_cuda_op_gelu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
178    ggml_cuda_op_unary<op_gelu_erf>(ctx, dst);
179}
180
181void ggml_cuda_op_gelu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
182    ggml_cuda_op_unary<op_gelu_quick>(ctx, dst);
183}
184
185void ggml_cuda_op_silu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
186    ggml_cuda_op_unary<op_silu>(ctx, dst);
187}
188
189void ggml_cuda_op_tanh(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
190    ggml_cuda_op_unary<op_tanh>(ctx, dst);
191}
192
193void ggml_cuda_op_relu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
194    ggml_cuda_op_unary<op_relu>(ctx, dst);
195}
196
197void ggml_cuda_op_sigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
198    ggml_cuda_op_unary<op_sigmoid>(ctx, dst);
199}
200
201void ggml_cuda_op_hardsigmoid(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
202    ggml_cuda_op_unary<op_hardsigmoid>(ctx, dst);
203}
204
205void ggml_cuda_op_hardswish(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
206    ggml_cuda_op_unary<op_hardswish>(ctx, dst);
207}
208
209void ggml_cuda_op_exp(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
210    ggml_cuda_op_unary<op_exp>(ctx, dst);
211}
212
213void ggml_cuda_op_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
214    ggml_cuda_op_unary<op_sqr>(ctx, dst);
215}
216
217void ggml_cuda_op_sqrt(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
218    ggml_cuda_op_unary<op_sqrt>(ctx, dst);
219}
220
221void ggml_cuda_op_sin(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
222    ggml_cuda_op_unary<op_sin>(ctx, dst);
223}
224
225void ggml_cuda_op_cos(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
226    ggml_cuda_op_unary<op_cos>(ctx, dst);
227}
228
229void ggml_cuda_op_log(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
230    ggml_cuda_op_unary<op_log>(ctx, dst);
231}
232
233void ggml_cuda_op_elu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
234    ggml_cuda_op_unary<op_elu>(ctx, dst);
235}
236
237void ggml_cuda_op_floor(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
238    ggml_cuda_op_unary<op_floor>(ctx, dst);
239}
240
241void ggml_cuda_op_ceil(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
242    ggml_cuda_op_unary<op_ceil>(ctx, dst);
243}
244
245void ggml_cuda_op_round(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
246    ggml_cuda_op_unary<op_round>(ctx, dst);
247}
248
249void ggml_cuda_op_trunc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
250    ggml_cuda_op_unary<op_trunc>(ctx, dst);
251}
252
253void ggml_cuda_op_expm1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
254    ggml_cuda_op_unary<op_expm1>(ctx, dst);
255}
256
257void ggml_cuda_op_softplus(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
258    ggml_cuda_op_unary<op_softplus>(ctx, dst);
259}
260/* gated ops */
261
262template <float (*op)(float), typename T>
263static __global__ void unary_gated_op_kernel(const T * x, const T * g, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1) {
264    ggml_cuda_pdl_lc();
265    const int64_t i = int64_t(blockDim.x)*blockIdx.x + threadIdx.x;
266
267    if (i >= k) {
268        return;
269    }
270
271    // perform base op and multiply with gate (either offset in same tensor or a separate one)
272    const int64_t j0 = (i / n) * o0 + (i % n);
273    const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
274
275    ggml_cuda_pdl_sync();
276    dst[i] = (T)(op((float)x[j0]) * (float)g[j1]);
277}
278
279template <float (*op)(float), typename T>
280static void unary_gated_cuda(const T * x, const T * g, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, cudaStream_t stream) {
281    const int64_t num_blocks = (k + CUDA_GLU_BLOCK_SIZE - 1) / CUDA_GLU_BLOCK_SIZE;
282    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params((dim3)num_blocks, CUDA_GLU_BLOCK_SIZE, 0, stream);
283    ggml_cuda_kernel_launch(unary_gated_op_kernel<op, T>, launch_params, x, g, dst, k, n, o0, o1);
284}
285
286template <float (*op)(float)>
287void ggml_cuda_op_unary_gated(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
288    const ggml_tensor * src0 = dst->src[0];
289    const ggml_tensor * src1 = dst->src[1];
290    void * src0_d = src0->data;
291    void * src1_d = src1 ? src1->data : src0->data;
292    const int64_t src0_o = src0->nb[1];
293    const int64_t src1_o = src1 ? src1->nb[1] : src0->nb[1];
294    void * dst_d = dst->data;
295    const int64_t nc = src1 ? src0->ne[0] : src0->ne[0] / 2;
296    cudaStream_t stream = ctx.stream();
297
298    GGML_ASSERT(ggml_is_contiguous_1(src0));
299    GGML_ASSERT(src0->nb[0] == ggml_element_size(src0));
300    GGML_ASSERT(ggml_is_contiguous(dst));
301
302    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
303    GGML_ASSERT( dst->type == GGML_TYPE_F32 ||  dst->type == GGML_TYPE_F16);
304    GGML_ASSERT(src0->type == dst->type);
305    GGML_ASSERT(dst->ne[0] == nc);
306    GGML_ASSERT(ggml_nrows(dst) == ggml_nrows(src0));
307
308    if (src1) {
309        GGML_ASSERT(ggml_is_contiguous_1(src1));
310        GGML_ASSERT(src1->nb[0] == ggml_element_size(src1));
311        GGML_ASSERT(src1->ne[0] == nc);
312        GGML_ASSERT(src0->type == src1->type);
313    }
314
315    const int32_t swapped = ((const int32_t *) dst->op_params)[1];
316
317    if (src0->type == GGML_TYPE_F16) {
318        half * src0_p = (half *) src0_d;
319        half * src1_p = (half *) src1_d;
320
321        if (!src1) {
322            src0_p += swapped ? nc : 0;
323            src1_p += swapped ? 0 : nc;
324        }
325
326        unary_gated_cuda<op>(src0_p, src1_p, (half *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(half), src1_o / sizeof(half), stream);
327    } else {
328        float * src0_p = (float *) src0_d;
329        float * src1_p = (float *) src1_d;
330
331        if (!src1) {
332            src0_p += swapped ? nc : 0;
333            src1_p += swapped ? 0 : nc;
334        }
335
336        unary_gated_cuda<op>(src0_p, src1_p, (float *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), stream);
337    }
338}
339
340void ggml_cuda_op_reglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
341    ggml_cuda_op_unary_gated<op_relu>(ctx, dst);
342}
343
344void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
345    ggml_cuda_op_unary_gated<op_gelu>(ctx, dst);
346}
347
348void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
349    ggml_cuda_op_unary_gated<op_silu>(ctx, dst);
350}
351
352void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
353    ggml_cuda_op_unary_gated<op_gelu_erf>(ctx, dst);
354}
355
356void ggml_cuda_op_geglu_quick(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
357    ggml_cuda_op_unary_gated<op_gelu_quick>(ctx, dst);
358}
359
360// swiglu_oai
361
362template <typename T>
363static __global__ void swiglu_oai_kernel(const T * x, const T * g, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, float alpha, float limit) {
364    const int64_t i = int64_t(blockDim.x)*blockIdx.x + threadIdx.x;
365
366    if (i >= k) {
367        return;
368    }
369
370    // perform base op and multiply with gate (either offset in same tensor or a separate one)
371    const int64_t j0 = (i / n) * o0 + (i % n);
372    const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
373
374    float xi = x[j0];
375    float gi = g[j1];
376
377    dst[i] = ggml_cuda_op_swiglu_oai_single(xi, gi, alpha, limit);
378}
379
380template <typename T>
381static void swiglu_oai_cuda(const T * x, const T * g, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, const float alpha, const float limit, cudaStream_t stream) {
382    const int64_t num_blocks = (k + CUDA_GLU_BLOCK_SIZE - 1) / CUDA_GLU_BLOCK_SIZE;
383    swiglu_oai_kernel<<<num_blocks, CUDA_GLU_BLOCK_SIZE, 0, stream>>>(x, g, dst, k, n, o0, o1, alpha, limit);
384}
385
386void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
387    const ggml_tensor * src0 = dst->src[0];
388    const ggml_tensor * src1 = dst->src[1];
389    void * src0_d = src0->data;
390    void * src1_d = src1 ? src1->data : src0->data;
391    const int64_t src0_o = src0->nb[1];
392    const int64_t src1_o = src1 ? src1->nb[1] : src0->nb[1];
393    void * dst_d = dst->data;
394    const int64_t nc = src1 ? src0->ne[0] : src0->ne[0] / 2;
395    cudaStream_t stream = ctx.stream();
396
397    GGML_ASSERT(ggml_is_contiguous_1(src0));
398    GGML_ASSERT(src0->nb[0] == ggml_element_size(src0));
399    GGML_ASSERT(ggml_is_contiguous(dst));
400
401    GGML_ASSERT(src0->type == GGML_TYPE_F32);
402    GGML_ASSERT( dst->type == GGML_TYPE_F32);
403    GGML_ASSERT(src0->type == dst->type);
404    GGML_ASSERT(dst->ne[0] == nc);
405    GGML_ASSERT(ggml_nrows(dst) == ggml_nrows(src0));
406
407    if (src1) {
408        GGML_ASSERT(ggml_is_contiguous_1(src1));
409        GGML_ASSERT(src1->nb[0] == ggml_element_size(src1));
410        GGML_ASSERT(src1->ne[0] == nc);
411        GGML_ASSERT(src0->type == src1->type);
412    }
413
414    //const int32_t swapped = ((const int32_t *) dst->op_params)[1];
415    const int32_t swapped = ggml_get_op_params_i32(dst, 1);
416    const float alpha = ggml_get_op_params_f32(dst, 2);
417    const float limit = ggml_get_op_params_f32(dst, 3);
418
419    float * src0_p = (float *) src0_d;
420    float * src1_p = (float *) src1_d;
421
422    if (!src1) {
423        src0_p += swapped ? nc : 0;
424        src1_p += swapped ? 0 : nc;
425    }
426
427    swiglu_oai_cuda(src0_p, src1_p, (float *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), alpha, limit, stream);
428}
429
430// swiglu_clamp
431
432template <typename T>
433static __global__ void swiglu_clamp_kernel(const T * gate, const T * up, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, float limit) {
434    const int64_t i = int64_t(blockDim.x)*blockIdx.x + threadIdx.x;
435
436    if (i >= k) {
437        return;
438    }
439
440    const int64_t j0 = (i / n) * o0 + (i % n);
441    const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);
442
443    dst[i] = (T) ggml_cuda_op_swiglu_clamp_single((float) gate[j0], (float) up[j1], limit);
444}
445
446template <typename T>
447static void swiglu_clamp_cuda(const T * gate, const T * up, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, const float limit, cudaStream_t stream) {
448    const int64_t num_blocks = (k + CUDA_GLU_BLOCK_SIZE - 1) / CUDA_GLU_BLOCK_SIZE;
449    swiglu_clamp_kernel<<<num_blocks, CUDA_GLU_BLOCK_SIZE, 0, stream>>>(gate, up, dst, k, n, o0, o1, limit);
450}
451
452void ggml_cuda_op_swiglu_clamp(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
453    const ggml_tensor * src0 = dst->src[0];
454    const ggml_tensor * src1 = dst->src[1];
455    void * src0_d = src0->data;
456    void * src1_d = src1 ? src1->data : src0->data;
457    const int64_t src0_o = src0->nb[1];
458    const int64_t src1_o = src1 ? src1->nb[1] : src0->nb[1];
459    void * dst_d = dst->data;
460    const int64_t nc = src1 ? src0->ne[0] : src0->ne[0] / 2;
461    cudaStream_t stream = ctx.stream();
462
463    GGML_ASSERT(ggml_is_contiguous_1(src0));
464    GGML_ASSERT(src0->nb[0] == ggml_element_size(src0));
465    GGML_ASSERT(ggml_is_contiguous(dst));
466
467    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
468    GGML_ASSERT(src0->type == dst->type);
469    GGML_ASSERT(dst->ne[0] == nc);
470    GGML_ASSERT(ggml_nrows(dst) == ggml_nrows(src0));
471
472    if (src1) {
473        GGML_ASSERT(ggml_is_contiguous_1(src1));
474        GGML_ASSERT(src1->nb[0] == ggml_element_size(src1));
475        GGML_ASSERT(src1->ne[0] == nc);
476        GGML_ASSERT(src0->type == src1->type);
477    }
478
479    const int32_t swapped = ggml_get_op_params_i32(dst, 1);
480    const float limit = ggml_get_op_params_f32(dst, 3);
481
482    if (src0->type == GGML_TYPE_F16) {
483        half * src0_p = (half *) src0_d;
484        half * src1_p = (half *) src1_d;
485
486        if (!src1) {
487            src0_p += swapped ? nc : 0;
488            src1_p += swapped ? 0 : nc;
489        }
490
491        swiglu_clamp_cuda(src0_p, src1_p, (half *) dst_d, ggml_nelements(dst), nc, src0_o / sizeof(half), src1_o / sizeof(half), limit, stream);
492    } else {
493        float * src0_p = (float *) src0_d;
494        float * src1_p = (float *) src1_d;
495
496        if (!src1) {
497            src0_p += swapped ? nc : 0;
498            src1_p += swapped ? 0 : nc;
499        }
500
501        swiglu_clamp_cuda(src0_p, src1_p, (float *) dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), limit, stream);
502    }
503}
504
505/* CUDA kernel + launcher for xIELU */
506
507template <typename T>
508static __global__ void xielu_kernel(const T * x, T * dst, const int k, float alpha_n, float alpha_p, float beta, float eps) {
509    const int i = blockDim.x*blockIdx.x + threadIdx.x;
510
511    if (i >= k) {
512        return;
513    }
514
515    const float xi = ggml_cuda_cast<float>(x[i]);
516
517    const float gate_pos = (xi > 0.0f);
518    const float y_pos = alpha_p * xi * xi + beta * xi;
519    const float min_v_eps = fminf(xi, eps);
520    const float y_neg = (expm1f(min_v_eps) - xi) * alpha_n + beta * xi;
521    const float out = gate_pos * y_pos + (1.0f - gate_pos) * y_neg;
522
523    dst[i] = ggml_cuda_cast<T>(out);
524}
525
526template <typename T>
527static void xielu_cuda(const T * x, T * dst, const int k, float alpha_n, float alpha_p, float beta, float eps, cudaStream_t stream) {
528    const int num_blocks = (k + CUDA_XIELU_BLOCK_SIZE) / CUDA_XIELU_BLOCK_SIZE;
529    xielu_kernel<<<num_blocks, CUDA_XIELU_BLOCK_SIZE, 0, stream>>>(x, dst, k, alpha_n, alpha_p, beta, eps);
530}
531
532void ggml_cuda_op_xielu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
533    const ggml_tensor * src0 = dst->src[0];
534    const void * src0_d = src0->data;
535    void * dst_d = dst->data;
536    cudaStream_t stream = ctx.stream();
537
538    GGML_ASSERT(ggml_is_contiguous(src0));
539
540    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
541    GGML_ASSERT( dst->type == GGML_TYPE_F32 ||  dst->type == GGML_TYPE_F16);
542    GGML_ASSERT(src0->type == dst->type);
543
544    const float alpha_n = ggml_get_op_params_f32(dst, 1);
545    const float alpha_p = ggml_get_op_params_f32(dst, 2);
546    const float beta    = ggml_get_op_params_f32(dst, 3);
547    const float eps     = ggml_get_op_params_f32(dst, 4);
548
549    if (src0->type == GGML_TYPE_F16) {
550        xielu_cuda((const half *)src0_d, (half *)dst_d, ggml_nelements(src0), alpha_n, alpha_p, beta, eps, stream);
551    } else {
552        xielu_cuda((const float *)src0_d, (float *)dst_d, ggml_nelements(src0), alpha_n, alpha_p, beta, eps, stream);
553    }
554}
555
556
557
558/* silu_back */
559
560static __device__ __forceinline__ float op_silu_back(float grad, float x) {
561    const float s = 1.0f / (1.0f + expf(-x));
562    return grad * s * (1.0f + x * (1.0f - s));
563}
564
565template <class T>
566static __global__ void silu_back_kernel(const T * grad, const T * xf, T * dst, const int k) {
567    const int i = blockDim.x*blockIdx.x + threadIdx.x;
568
569    if (i >= k) {
570        return;
571    }
572
573    dst[i] = (T)op_silu_back((float)grad[i], (float)xf[i]);
574}
575
576template <class T>
577static void silu_back_cuda(const T * grad, const T * x, T * dst, const int k, cudaStream_t stream) {
578    const int num_blocks = (k + CUDA_SILU_BACK_BLOCK_SIZE - 1) / CUDA_SILU_BLOCK_SIZE;
579    silu_back_kernel<<<num_blocks, CUDA_SILU_BACK_BLOCK_SIZE, 0, stream>>>(grad, x, dst, k);
580}
581
582void ggml_cuda_op_silu_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
583    const ggml_tensor * src0 = dst->src[0]; // input from forward pass
584    const ggml_tensor * src1 = dst->src[1]; // grads of forward pass output
585
586    const float * src0_d = (const float *) src0->data;
587    const float * src1_d = (const float *) src1->data;
588    float       * dst_d  = (float       *) dst->data;
589
590    cudaStream_t stream = ctx.stream();
591
592    GGML_ASSERT(ggml_is_contiguous(src0));
593
594    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
595    GGML_ASSERT( dst->type == GGML_TYPE_F32 ||  dst->type == GGML_TYPE_F16);
596    GGML_ASSERT(src0->type == dst->type);
597
598    if (src0->type == GGML_TYPE_F16) {
599        silu_back_cuda((const half *)src0_d, (const half *)src1_d, (half *)dst_d, ggml_nelements(src0), stream);
600    } else {
601        silu_back_cuda((const float*)src0_d, (const float*)src1_d, (float *)dst_d, ggml_nelements(src0), stream);
602    }
603}
604
605/* leaky relu */
606
607static __device__ __forceinline__ float op_leaky_relu(float x, const float negative_slope) {
608    return fmaxf(x, 0) + fminf(x, 0.0f) * negative_slope;
609}
610
611template <class T>
612static __global__ void leaky_relu_kernel(const T * x, T * dst, const int k, const float negative_slope) {
613    const int i  = blockDim.x*blockIdx.x + threadIdx.x;
614
615    if (i >= k) {
616        return;
617    }
618
619    dst[i] = (T)op_leaky_relu((float)x[i], negative_slope);
620}
621
622template <class T>
623static void leaky_relu_cuda(const T * x, T * dst, const int k, const float negative_slope, cudaStream_t stream) {
624    const int num_blocks = (k + CUDA_RELU_BLOCK_SIZE - 1) / CUDA_RELU_BLOCK_SIZE;
625    leaky_relu_kernel<<<num_blocks, CUDA_RELU_BLOCK_SIZE, 0, stream>>>(x, dst, k, negative_slope);
626}
627
628void ggml_cuda_op_leaky_relu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
629    const ggml_tensor * src0 = dst->src[0];
630    const void * src0_d = src0->data;
631    void * dst_d = dst->data;
632    cudaStream_t stream = ctx.stream();
633
634    GGML_ASSERT(ggml_is_contiguous(src0));
635
636    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
637    GGML_ASSERT( dst->type == GGML_TYPE_F32 ||  dst->type == GGML_TYPE_F16);
638    GGML_ASSERT(src0->type == dst->type);
639
640    float negative_slope;
641    memcpy(&negative_slope, dst->op_params, sizeof(float));
642
643    if (src0->type == GGML_TYPE_F16) {
644        leaky_relu_cuda((const half *)src0_d, (half *)dst_d, ggml_nelements(src0), negative_slope, stream);
645    } else {
646        leaky_relu_cuda((const float *)src0_d, (float *)dst_d, ggml_nelements(src0), negative_slope, stream);
647    }
648}
649
650/* fused unary + mul */
651
652template <float (*op)(float)>
653static void ggml_cuda_op_unary_mul_impl(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node) {
654    // unary_node: UNARY op applied to unary_node->src[0]
655    // mul_node:   MUL(a, b) where one of a/b is unary_node
656    // Output goes to mul_node->data
657
658    const ggml_tensor * unary_src = unary_node->src[0];  // input to the unary op
659    const ggml_tensor * other_src = (mul_node->src[0] == unary_node) ? mul_node->src[1] : mul_node->src[0];
660
661    GGML_ASSERT(ggml_is_contiguous_1(unary_src));
662    GGML_ASSERT(unary_src->nb[0] == ggml_element_size(unary_src));
663    GGML_ASSERT(ggml_is_contiguous_1(other_src));
664    GGML_ASSERT(other_src->nb[0] == ggml_element_size(other_src));
665    GGML_ASSERT(ggml_are_same_shape(unary_src, other_src));
666
667    GGML_ASSERT(unary_src->type == GGML_TYPE_F32 || unary_src->type == GGML_TYPE_F16);
668    GGML_ASSERT(unary_src->type == other_src->type);
669    GGML_ASSERT(unary_src->type == mul_node->type);
670
671    cudaStream_t stream = ctx.stream();
672
673    const int64_t k  = ggml_nelements(mul_node);
674    const int64_t nc = unary_src->ne[0];
675    const int64_t unary_stride = unary_src->nb[1];
676    const int64_t other_stride = other_src->nb[1];
677
678    if (unary_src->type == GGML_TYPE_F16) {
679        unary_gated_cuda<op>((const half *) unary_src->data, (const half *) other_src->data,
680                             (half *) mul_node->data, k, nc,
681                             unary_stride / sizeof(half), other_stride / sizeof(half), stream);
682    } else {
683        unary_gated_cuda<op>((const float *) unary_src->data, (const float *) other_src->data,
684                             (float *) mul_node->data, k, nc,
685                             unary_stride / sizeof(float), other_stride / sizeof(float), stream);
686    }
687}
688
689void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node) {
690    switch (ggml_get_unary_op(unary_node)) {
691        case GGML_UNARY_OP_SILU:
692            ggml_cuda_op_unary_mul_impl<op_silu>(ctx, unary_node, mul_node);
693            break;
694        case GGML_UNARY_OP_SIGMOID:
695            ggml_cuda_op_unary_mul_impl<op_sigmoid>(ctx, unary_node, mul_node);
696            break;
697        case GGML_UNARY_OP_SOFTPLUS:
698            ggml_cuda_op_unary_mul_impl<op_softplus>(ctx, unary_node, mul_node);
699            break;
700        default:
701            GGML_ABORT("Unsupported unary op for fused unary+mul");
702    }
703}
704
705/* fused relu + sqr */
706
707void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node) {
708    const ggml_tensor * src = relu_node->src[0];
709    cudaStream_t stream = ctx.stream();
710
711    GGML_ASSERT(ggml_is_contiguous(src));
712    GGML_ASSERT(src->type == GGML_TYPE_F32 || src->type == GGML_TYPE_F16);
713    GGML_ASSERT(src->type == sqr_node->type);
714
715    const int k = ggml_nelements(src);
716    if (src->type == GGML_TYPE_F16) {
717        unary_cuda<op_relu_sqr>((const half *)src->data, (half *)sqr_node->data, k, stream);
718    } else {
719        unary_cuda<op_relu_sqr>((const float *)src->data, (float *)sqr_node->data, k, stream);
720    }
721}