char-fork/llama.cpp

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

YiChen Lvmetal: per-op source split + parallel compile (#26561)b615f5b4b

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
7.1 KiB223 linesraw
1#include "common.h"
2
3template<typename T>
4kernel void kernel_soft_max(
5        constant ggml_metal_kargs_soft_max & args,
6        device const  char * src0,
7        device const  char * src1,
8        device const  char * src2,
9        device        char * dst,
10        threadgroup  float * buf [[threadgroup(0)]],
11        uint3 tgpig[[threadgroup_position_in_grid]],
12        uint3 tpitg[[thread_position_in_threadgroup]],
13        uint  sgitg[[simdgroup_index_in_threadgroup]],
14        uint  tiisg[[thread_index_in_simdgroup]],
15        uint3  tptg[[threads_per_threadgroup]]) {
16    const int32_t i03 = tgpig.z;
17    const int32_t i02 = tgpig.y;
18    const int32_t i01 = tgpig.x;
19
20    const int32_t i13 = i03%args.ne13;
21    const int32_t i12 = i02%args.ne12;
22    const int32_t i11 = i01;
23
24    device const float * psrc0 =                (device const float *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
25    device const     T * pmask = src1 != src0 ? (device const T *    ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
26    device const float * psrc2 = src2 != src0 ? (device const float *) (src2)                                                 : nullptr;
27    device       float * pdst  =                (device       float *) (dst  + i01*args.nb1  + i02*args.nb2  + i03*args.nb3);
28
29    float slope = 1.0f;
30
31    // ALiBi
32    if (args.max_bias > 0.0f) {
33        const int32_t h = i02;
34
35        const float base = h < args.n_head_log2 ? args.m0 : args.m1;
36        const int   exp  = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
37
38        slope = pow(base, exp);
39    }
40
41    // parallel max
42    float lmax = psrc2 ? psrc2[i02] : -INFINITY;
43
44    for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
45        lmax = MAX(lmax, psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f));
46    }
47
48    // find the max value in the block
49    float max_val = simd_max(lmax);
50    if (tptg.x > N_SIMDWIDTH) {
51        if (sgitg == 0) {
52            buf[tiisg] = -INFINITY;
53        }
54
55        threadgroup_barrier(mem_flags::mem_threadgroup);
56
57        if (tiisg == 0) {
58            buf[sgitg] = max_val;
59        }
60
61        threadgroup_barrier(mem_flags::mem_threadgroup);
62
63        max_val = buf[tiisg];
64        max_val = simd_max(max_val);
65    }
66
67    // parallel sum
68    float lsum = 0.0f;
69    for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
70        const float exp_psrc0 = exp((psrc0[i00]*args.scale + (pmask ? slope*pmask[i00] : 0.0f)) - max_val);
71        lsum += exp_psrc0;
72        pdst[i00] = exp_psrc0;
73    }
74
75    // This barrier fixes a failing test
76    // ref: https://github.com/ggml-org/ggml/pull/621#discussion_r1425156335
77    threadgroup_barrier(mem_flags::mem_none);
78
79    float sum = simd_sum(lsum);
80
81    if (tptg.x > N_SIMDWIDTH) {
82        if (sgitg == 0) {
83            buf[tiisg] = 0.0f;
84        }
85
86        threadgroup_barrier(mem_flags::mem_threadgroup);
87
88        if (tiisg == 0) {
89            buf[sgitg] = sum;
90        }
91
92        threadgroup_barrier(mem_flags::mem_threadgroup);
93
94        sum = buf[tiisg];
95        sum = simd_sum(sum);
96    }
97
98    if (psrc2) {
99        sum += exp(psrc2[i02] - max_val);
100    }
101
102    const float inv_sum = 1.0f/sum;
103
104    for (int i00 = tpitg.x; i00 < args.ne00; i00 += tptg.x) {
105        pdst[i00] *= inv_sum;
106    }
107}
108
109template<typename T>
110kernel void kernel_soft_max_4(
111        constant ggml_metal_kargs_soft_max & args,
112        device const  char * src0,
113        device const  char * src1,
114        device const  char * src2,
115        device        char * dst,
116        threadgroup  float * buf [[threadgroup(0)]],
117        uint3 tgpig[[threadgroup_position_in_grid]],
118        uint3 tpitg[[thread_position_in_threadgroup]],
119        uint  sgitg[[simdgroup_index_in_threadgroup]],
120        uint  tiisg[[thread_index_in_simdgroup]],
121        uint3  tptg[[threads_per_threadgroup]]) {
122    const int32_t i03 = tgpig.z;
123    const int32_t i02 = tgpig.y;
124    const int32_t i01 = tgpig.x;
125
126    const int32_t i13 = i03%args.ne13;
127    const int32_t i12 = i02%args.ne12;
128    const int32_t i11 = i01;
129
130    device const float4 * psrc4 =                (device const float4 *) (src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03);
131    device const      T * pmask = src1 != src0 ? (device const T *     ) (src1 + i11*args.nb11 + i12*args.nb12 + i13*args.nb13) : nullptr;
132    device const float *  psrc2 = src2 != src0 ? (device const float * ) (src2)                                                 : nullptr;
133    device       float4 * pdst4 =                (device       float4 *) (dst  + i01*args.nb1  + i02*args.nb2  + i03*args.nb3);
134
135    float slope = 1.0f;
136
137    if (args.max_bias > 0.0f) {
138        const int32_t h = i02;
139
140        const float base = h < args.n_head_log2 ? args.m0 : args.m1;
141        const int   exp  = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
142
143        slope = pow(base, exp);
144    }
145
146    // parallel max
147    float4 lmax4 = psrc2 ? psrc2[i02] : -INFINITY;
148
149    for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
150        lmax4 = fmax(lmax4, psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f)));
151    }
152
153    const float lmax = MAX(MAX(lmax4[0], lmax4[1]), MAX(lmax4[2], lmax4[3]));
154
155    float max_val = simd_max(lmax);
156    if (tptg.x > N_SIMDWIDTH) {
157        if (sgitg == 0) {
158            buf[tiisg] = -INFINITY;
159        }
160
161        threadgroup_barrier(mem_flags::mem_threadgroup);
162
163        if (tiisg == 0) {
164            buf[sgitg] = max_val;
165        }
166
167        threadgroup_barrier(mem_flags::mem_threadgroup);
168
169        max_val = buf[tiisg];
170        max_val = simd_max(max_val);
171    }
172
173    // parallel sum
174    float4 lsum4 = 0.0f;
175    for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
176        const float4 exp_psrc4 = exp((psrc4[i00]*args.scale + (float4)((pmask ? slope*pmask[i00] : 0.0f))) - max_val);
177        lsum4 += exp_psrc4;
178        pdst4[i00] = exp_psrc4;
179    }
180
181    const float lsum = lsum4[0] + lsum4[1] + lsum4[2] + lsum4[3];
182
183    // This barrier fixes a failing test
184    // ref: https://github.com/ggml-org/ggml/pull/621#discussion_r1425156335
185    threadgroup_barrier(mem_flags::mem_none);
186
187    float sum = simd_sum(lsum);
188
189    if (tptg.x > N_SIMDWIDTH) {
190        if (sgitg == 0) {
191            buf[tiisg] = 0.0f;
192        }
193
194        threadgroup_barrier(mem_flags::mem_threadgroup);
195
196        if (tiisg == 0) {
197            buf[sgitg] = sum;
198        }
199
200        threadgroup_barrier(mem_flags::mem_threadgroup);
201
202        sum = buf[tiisg];
203        sum = simd_sum(sum);
204    }
205
206    if (psrc2) {
207        sum += exp(psrc2[i02] - max_val);
208    }
209
210    const float inv_sum = 1.0f/sum;
211
212    for (int i00 = tpitg.x; i00 < args.ne00/4; i00 += tptg.x) {
213        pdst4[i00] *= inv_sum;
214    }
215}
216
217typedef decltype(kernel_soft_max<float>)    kernel_soft_max_t;
218typedef decltype(kernel_soft_max_4<float4>) kernel_soft_max_4_t;
219
220template [[host_name("kernel_soft_max_f16")]]   kernel kernel_soft_max_t   kernel_soft_max<half>;
221template [[host_name("kernel_soft_max_f32")]]   kernel kernel_soft_max_t   kernel_soft_max<float>;
222template [[host_name("kernel_soft_max_f16_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<half4>;
223template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4<float4>;