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 →
9.8 KiB308 linesraw
1#include "common.h"
2
3// F == 1 : norm (no fuse)
4// F == 2 : norm + mul
5// F == 3 : norm + mul + add
6template <typename T, short F>
7kernel void kernel_norm_fuse_impl(
8        constant ggml_metal_kargs_norm & args,
9        device const char * src0,
10        device const char * src1_0,
11        device const char * src1_1,
12        device       char * dst,
13        threadgroup float * shmem_f32 [[threadgroup(0)]],
14        uint3   tgpig[[threadgroup_position_in_grid]],
15        ushort3 tpitg[[thread_position_in_threadgroup]],
16        ushort  sgitg[[simdgroup_index_in_threadgroup]],
17        ushort  tiisg[[thread_index_in_simdgroup]],
18        ushort3   ntg[[threads_per_threadgroup]]) {
19    if (sgitg == 0) {
20        shmem_f32[tiisg] = 0.0f;
21    }
22
23    const int i01 = tgpig.x;
24    const int i02 = tgpig.y;
25    const int i03 = tgpig.z;
26
27    device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
28
29    device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
30    device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);
31
32    T sumft(0.0f);
33
34    float sumf = 0.0f;
35
36    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
37        sumft += x[i00];
38    }
39    sumf = dot(sumft, T(1.0f));
40    sumf = simd_sum(sumf);
41
42    threadgroup_barrier(mem_flags::mem_threadgroup);
43
44    if (tiisg == 0) {
45        shmem_f32[sgitg] = sumf;
46    }
47
48    threadgroup_barrier(mem_flags::mem_threadgroup);
49
50    sumf = shmem_f32[tiisg];
51    sumf = simd_sum(sumf);
52
53    const float mean = sumf/args.ne00;
54
55    device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
56
57    sumf = 0.0f;
58    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
59        y[i00] = x[i00] - mean;
60        sumf += dot(y[i00], y[i00]);
61    }
62    sumf = simd_sum(sumf);
63
64    threadgroup_barrier(mem_flags::mem_threadgroup);
65
66    if (tiisg == 0) {
67        shmem_f32[sgitg] = sumf;
68    }
69
70    threadgroup_barrier(mem_flags::mem_threadgroup);
71
72    sumf = shmem_f32[tiisg];
73    sumf = simd_sum(sumf);
74
75    const float variance = sumf/args.ne00;
76
77    const float scale = 1.0f/sqrt(variance + args.eps);
78    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
79        if (F == 1) {
80            y[i00] = (y[i00]*scale);
81        }
82        if (F == 2) {
83            y[i00] = (y[i00]*scale)*f0[i00];
84        }
85        if (F == 3) {
86            y[i00] = (y[i00]*scale)*f0[i00] + f1[i00];
87        }
88    }
89}
90
91typedef decltype(kernel_norm_fuse_impl<float4, 1>) kernel_norm_fuse_t;
92
93template [[host_name("kernel_norm_f32")]]         kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 1>;
94template [[host_name("kernel_norm_mul_f32")]]     kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 2>;
95template [[host_name("kernel_norm_mul_add_f32")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float, 3>;
96
97template [[host_name("kernel_norm_f32_4")]]         kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 1>;
98template [[host_name("kernel_norm_mul_f32_4")]]     kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 2>;
99template [[host_name("kernel_norm_mul_add_f32_4")]] kernel kernel_norm_fuse_t kernel_norm_fuse_impl<float4, 3>;
100
101// F == 1 : rms_norm (no fuse)
102// F == 2 : rms_norm + mul
103// F == 3 : rms_norm + mul + add
104template <typename T, short F>
105kernel void kernel_rms_norm_fuse_impl(
106        constant ggml_metal_kargs_norm & args,
107        device const char * src0,
108        device const char * src1_0,
109        device const char * src1_1,
110        device       char * dst,
111        threadgroup float * shmem_f32 [[threadgroup(0)]],
112        uint3   tgpig[[threadgroup_position_in_grid]],
113        ushort3 tpitg[[thread_position_in_threadgroup]],
114        ushort  sgitg[[simdgroup_index_in_threadgroup]],
115        ushort  tiisg[[thread_index_in_simdgroup]],
116        ushort3   ntg[[threads_per_threadgroup]]) {
117    if (sgitg == 0) {
118        shmem_f32[tiisg] = 0.0f;
119    }
120
121    const int i01 = tgpig.x;
122    const int i02 = tgpig.y;
123    const int i03 = tgpig.z;
124
125    device const T * x = (device const T *) (src0 + i03*args.nbf3[0] + i02*args.nbf2[0] + i01*args.nbf1[0]);
126
127    device const T * f0 = (device const T *) (src1_0 + (i03%args.nef3[1])*args.nbf3[1] + (i02%args.nef2[1])*args.nbf2[1] + (i01%args.nef1[1])*args.nbf1[1]);
128    device const T * f1 = (device const T *) (src1_1 + (i03%args.nef3[2])*args.nbf3[2] + (i02%args.nef2[2])*args.nbf2[2] + (i01%args.nef1[2])*args.nbf1[2]);
129
130    float sumf = 0.0f;
131
132    // parallel sum
133    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
134        sumf += dot(x[i00], x[i00]);
135    }
136    sumf = simd_sum(sumf);
137
138    threadgroup_barrier(mem_flags::mem_threadgroup);
139
140    if (tiisg == 0) {
141        shmem_f32[sgitg] = sumf;
142    }
143
144    threadgroup_barrier(mem_flags::mem_threadgroup);
145
146    sumf = shmem_f32[tiisg];
147    sumf = simd_sum(sumf);
148
149    const float mean  = sumf/args.ne00;
150    const float scale = 1.0f/sqrt(mean + args.eps);
151
152    device T * y = (device T *) (dst + i03*args.nb3 + i02*args.nb2 + i01*args.nb1);
153    for (int i00 = tpitg.x; i00 < args.ne00_t; i00 += ntg.x) {
154        if (F == 1) {
155            y[i00] = (x[i00]*scale);
156        }
157        if (F == 2) {
158            y[i00] = (x[i00]*scale)*f0[i00];
159        }
160        if (F == 3) {
161            y[i00] = (x[i00]*scale)*f0[i00] + f1[i00];
162        }
163    }
164}
165
166typedef decltype(kernel_rms_norm_fuse_impl<float4, 1>) kernel_rms_norm_fuse_t;
167
168template [[host_name("kernel_rms_norm_f32")]]         kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 1>;
169template [[host_name("kernel_rms_norm_mul_f32")]]     kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 2>;
170template [[host_name("kernel_rms_norm_mul_add_f32")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float, 3>;
171
172template [[host_name("kernel_rms_norm_f32_4")]]         kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 1>;
173template [[host_name("kernel_rms_norm_mul_f32_4")]]     kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 2>;
174template [[host_name("kernel_rms_norm_mul_add_f32_4")]] kernel kernel_rms_norm_fuse_t kernel_rms_norm_fuse_impl<float4, 3>;
175
176template <typename T0, typename T>
177kernel void kernel_l2_norm_impl(
178        constant ggml_metal_kargs_l2_norm & args,
179        device const char * src0,
180        device       char * dst,
181        threadgroup float * shmem_f32 [[threadgroup(0)]],
182        uint3   tgpig[[threadgroup_position_in_grid]],
183        ushort3 tpitg[[thread_position_in_threadgroup]],
184        ushort  sgitg[[simdgroup_index_in_threadgroup]],
185        ushort  tiisg[[thread_index_in_simdgroup]],
186        ushort3   ntg[[threads_per_threadgroup]]) {
187    const int i03 = tgpig.z;
188    const int i02 = tgpig.y;
189    const int i01 = tgpig.x;
190
191    if (sgitg == 0) {
192        shmem_f32[tiisg] = 0.0f;
193    }
194
195    device const T0 * x = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
196    device       T  * y = (device       T  *) (dst  + i03*args.nb3  + i02*args.nb2  + i01*args.nb1);
197
198    float sumf = 0.0f;
199
200    // parallel sum
201    for (int i00 = tpitg.x; i00 < args.ne00; i00 += ntg.x) {
202        sumf += dot(x[i00], x[i00]);
203    }
204    sumf = simd_sum(sumf);
205
206    threadgroup_barrier(mem_flags::mem_threadgroup);
207
208    if (tiisg == 0) {
209        shmem_f32[sgitg] = sumf;
210    }
211
212    threadgroup_barrier(mem_flags::mem_threadgroup);
213
214    sumf = shmem_f32[tiisg];
215    sumf = simd_sum(sumf);
216
217    const float scale = 1.0f/max(sqrt(sumf), args.eps);
218
219    for (int i00 = tpitg.x; i00 < args.ne00; i00 += ntg.x) {
220        y[i00] = x[i00] * scale;
221    }
222}
223
224typedef decltype(kernel_l2_norm_impl<float, float>) kernel_l2_norm_t;
225
226template [[host_name("kernel_l2_norm_f32_f32")]]   kernel kernel_l2_norm_t kernel_l2_norm_impl<float,  float>;
227template [[host_name("kernel_l2_norm_f32_f32_4")]] kernel kernel_l2_norm_t kernel_l2_norm_impl<float4, float4>;
228
229kernel void kernel_group_norm_f32(
230        constant ggml_metal_kargs_group_norm & args,
231        device const float * src0,
232        device       float * dst,
233        threadgroup float  * buf [[threadgroup(0)]],
234        uint tgpig[[threadgroup_position_in_grid]],
235        uint tpitg[[thread_position_in_threadgroup]],
236        uint sgitg[[simdgroup_index_in_threadgroup]],
237        uint tiisg[[thread_index_in_simdgroup]],
238        uint   ntg[[threads_per_threadgroup]]) {
239    const int64_t ne = args.ne00*args.ne01*args.ne02;
240    const int64_t gs = args.ne00*args.ne01*((args.ne02 + args.ngrp - 1) / args.ngrp);
241
242    int start = tgpig * gs;
243    int end   = start + gs;
244
245    start += tpitg;
246
247    if (end >= ne) {
248        end = ne;
249    }
250
251    float tmp = 0.0f; // partial sum for thread in warp
252
253    for (int j = start; j < end; j += ntg) {
254        tmp += src0[j];
255    }
256
257    threadgroup_barrier(mem_flags::mem_threadgroup);
258    tmp = simd_sum(tmp);
259    if (ntg > N_SIMDWIDTH) {
260        if (sgitg == 0) {
261            buf[tiisg] = 0.0f;
262        }
263
264        threadgroup_barrier(mem_flags::mem_threadgroup);
265
266        if (tiisg == 0) {
267            buf[sgitg] = tmp;
268        }
269
270        threadgroup_barrier(mem_flags::mem_threadgroup);
271
272        tmp = buf[tiisg];
273        tmp = simd_sum(tmp);
274    }
275
276    const float mean = tmp / gs;
277    tmp = 0.0f;
278
279    for (int j = start; j < end; j += ntg) {
280        float xi = src0[j] - mean;
281        dst[j] = xi;
282        tmp += xi * xi;
283    }
284
285    tmp = simd_sum(tmp);
286    if (ntg > N_SIMDWIDTH) {
287        if (sgitg == 0) {
288            buf[tiisg] = 0.0f;
289        }
290
291        threadgroup_barrier(mem_flags::mem_threadgroup);
292
293        if (tiisg == 0) {
294            buf[sgitg] = tmp;
295        }
296
297        threadgroup_barrier(mem_flags::mem_threadgroup);
298
299        tmp = buf[tiisg];
300        tmp = simd_sum(tmp);
301    }
302
303    const float variance = tmp / gs;
304    const float scale = 1.0f/sqrt(variance + args.eps);
305    for (int j = start; j < end; j += ntg) {
306        dst[j] *= scale;
307    }
308}