char-fork/llama.cpp

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

Georgi Gerganovmetal : single-source fusion table + fusion debug rework (#28164)a2878d30d

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
9.2 KiB264 linesraw
1#include "common.h"
2
3constant short FC_gated_delta_net_ne20 [[function_constant(FC_GATED_DELTA_NET + 0)]];
4constant short FC_gated_delta_net_ne30 [[function_constant(FC_GATED_DELTA_NET + 1)]];
5constant short FC_gated_delta_net_K    [[function_constant(FC_GATED_DELTA_NET + 2)]];
6
7#if 1
8template<short NSG>
9kernel void kernel_gated_delta_net_impl(
10        constant ggml_metal_kargs_gated_delta_net & args,
11        device const char * q,
12        device const char * k,
13        device const char * v,
14        device const char * g,
15        device const char * b,
16        device const char * s,
17        device       char * dst,
18        device       char * dst_fuse,
19        uint3 tgpig[[threadgroup_position_in_grid]],
20        uint3 tpitg[[thread_position_in_threadgroup]],
21        uint3   ntg[[threads_per_threadgroup]])  {
22#define S_v FC_gated_delta_net_ne20
23#define G   FC_gated_delta_net_ne30
24#define K   FC_gated_delta_net_K
25
26    const uint tx = tpitg.x;
27    const uint ty = tpitg.y;
28
29    const uint i23 = tgpig.z; // B (n_seqs)
30    const uint i21 = tgpig.y; // H (head)
31    const uint i20 = tgpig.x*NSG + ty; // row within S_v
32
33    const uint i01 = i21 % args.ne01;
34    const uint i11 = i21 % args.ne11;
35
36    const float scale = 1.0f / sqrt((float)S_v);
37
38    // input state layout [S_v, S_v, H, n_seqs] (s0 only): per-seq stride is H*D.
39    // state is stored transposed: M[i20][is] = S[is][i20], so row i20 is contiguous
40    const uint state_in_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
41    device const float * s_ptr = (device const float *) (s) + state_in_base;
42
43    float ls[NSG];
44
45    FOR_UNROLL (short j = 0; j < NSG; j++) {
46        const short is = tx*NSG + j;
47        ls[j] = s_ptr[is];
48    }
49
50    device float * dst_attn = (device float *) (dst) + (i23*args.ne22*args.ne21 + i21)*S_v + i20;
51
52    device const float * q_ptr = (device const float *) (q + i23*args.nb03 + i01*args.nb01);
53    device const float * k_ptr = (device const float *) (k + i23*args.nb13 + i11*args.nb11);
54    device const float * v_ptr = (device const float *) (v + i23*args.nb23 + i21*args.nb21);
55
56    device const float * b_ptr = (device const float *) (b) + (i23*args.ne22*args.ne21 + i21);
57    device const float * g_ptr = (device const float *) (g) + (i23*args.ne22*args.ne21 + i21)*G;
58
59    // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
60    // When n_tokens < K, only slots 0..n_tokens-1 are written; older slots are caller-owned.
61
62    // output state base offset: after attention scores
63    const uint attn_size = args.ne22 * args.ne21 * S_v * args.ne23;
64    // output state per-slot size: S_v * S_v * H * n_seqs
65    const uint state_size_per_snap = S_v * S_v * args.ne21 * args.ne23;
66    // per-(seq,head) offset within a slot
67    const uint state_out_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
68
69    // when fused with the cache cpy, write the snapshots straight into the cache buffer using
70    // the slot stride; otherwise append them after the attn scores (nb_out == 0)
71    const bool fused = args.nb_out > 0;
72    const device float * state_out = fused ? (device float *)dst_fuse : (device float *)dst + attn_size;
73    const uint slot_stride = fused ? (uint)args.nb_out : state_size_per_snap;
74
75    for (short t = 0; t < args.ne22; t++) {
76        float s_k = 0.0f;
77
78        if (G == 1) {
79            const float g_exp = exp(g_ptr[0]);
80
81            FOR_UNROLL (short j = 0; j < NSG; j++) {
82                const short is = tx*NSG + j;
83                ls[j] *= g_exp;
84
85                s_k += ls[j]*k_ptr[is];
86            }
87        } else {
88            // KDA
89            FOR_UNROLL (short j = 0; j < NSG; j++) {
90                const short is = tx*NSG + j;
91                ls[j] *= exp(g_ptr[is]);
92
93                s_k += ls[j]*k_ptr[is];
94            }
95        }
96
97        s_k = simd_sum(s_k);
98
99        const float d = (v_ptr[i20] - s_k)*b_ptr[0];
100
101        float y = 0.0f;
102
103        FOR_UNROLL (short j = 0; j < NSG; j++) {
104            const short is = tx*NSG + j;
105            ls[j] += k_ptr[is]*d;
106
107            y += ls[j]*q_ptr[is];
108        }
109
110        y = simd_sum(y);
111
112        if (tx == 0) {
113            dst_attn[t*args.ne21*S_v] = y*scale;
114        }
115
116        q_ptr += args.ns02;
117        k_ptr += args.ns12;
118        v_ptr += args.ns22;
119
120        b_ptr += args.ne21;
121        g_ptr += args.ne21*G;
122
123        if (K > 1) {
124            const int target_slot = (int)args.ne22 - 1 - (int)t;
125            if (target_slot >= 0 && target_slot < (int)K) {
126                device float * dst_state = (device float *)state_out + (uint)target_slot * slot_stride + state_out_base;
127                FOR_UNROLL (short j = 0; j < NSG; j++) {
128                    const short is = tx*NSG + j;
129                    dst_state[is] = ls[j];
130                }
131            }
132        }
133    }
134
135    if (K == 1) {
136        device float * dst_state = (device float *)state_out + state_out_base;
137        FOR_UNROLL (short j = 0; j < NSG; j++) {
138            const short is = tx*NSG + j;
139            dst_state[is] = ls[j];
140        }
141    }
142
143#undef S_v
144#undef G
145#undef K
146}
147
148typedef decltype(kernel_gated_delta_net_impl<4>) kernel_gated_delta_net_t;
149
150template [[host_name("kernel_gated_delta_net_f32_1")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<1>;
151template [[host_name("kernel_gated_delta_net_f32_2")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<2>;
152template [[host_name("kernel_gated_delta_net_f32_4")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<4>;
153
154#else
155// a simplified version of the above
156// no performance improvement, so keep the above version for now
157
158template<typename T, short NSG>
159kernel void kernel_gated_delta_net_impl(
160        constant ggml_metal_kargs_gated_delta_net & args,
161        device const char * q,
162        device const char * k,
163        device const char * v,
164        device const char * g,
165        device const char * b,
166        device const char * s,
167        device       char * dst,
168        device       char * dst_fuse,
169        uint3 tgpig[[threadgroup_position_in_grid]],
170        uint3 tpitg[[thread_position_in_threadgroup]],
171        uint3   ntg[[threads_per_threadgroup]])  {
172#define S_v FC_gated_delta_net_ne20
173#define G   FC_gated_delta_net_ne30
174
175    const uint tx = tpitg.x;
176    const uint ty = tpitg.y;
177
178    const uint i23 = tgpig.z; // B
179    const uint i21 = tgpig.y; // H
180    const uint i20 = tgpig.x*NSG + ty;
181
182    const uint i01 = i21 % args.ne01;
183    const uint i11 = i21 % args.ne11;
184
185    const float scale = 1.0f / sqrt((float)S_v);
186
187    device const float * s_ptr = (device const float *) (s) + (i23*args.ne21 + i21)*S_v*S_v + i20;
188
189    float lsf[NSG];
190
191    FOR_UNROLL (short j = 0; j < NSG; j++) {
192        const short is = tx*NSG + j;
193        lsf[j] = s_ptr[is*S_v];
194    }
195
196    thread T * ls = (thread T *) (lsf);
197
198    device float * dst_attn = (device float *) (dst) + (i23*args.ne22*args.ne21 + i21)*S_v + i20;
199
200    device const float * q_ptr = (device const float *) (q + i23*args.nb03 + i01*args.nb01);
201    device const float * k_ptr = (device const float *) (k + i23*args.nb13 + i11*args.nb11);
202    device const float * v_ptr = (device const float *) (v + i23*args.nb23 + i21*args.nb21);
203
204    device const float * b_ptr  = (device const float *) (b) + (i23*args.ne22*args.ne21 + i21);
205    device const float * g_ptr  = (device const float *) (g) + (i23*args.ne22*args.ne21 + i21)*G;
206
207    for (short t = 0; t < args.ne22; t++) {
208        device const T * qt_ptr = (device const T *) (q_ptr);
209        device const T * kt_ptr = (device const T *) (k_ptr);
210        device const T * gt_ptr = (device const T *) (g_ptr);
211
212        if (G == 1) {
213            *ls *= exp(g_ptr[0]);
214        } else {
215            // KDA
216            *ls *= exp(gt_ptr[tx]);
217        }
218
219        const float s_k = simd_sum(dot(*ls, kt_ptr[tx]));
220
221        const float d = (v_ptr[i20] - s_k)*b_ptr[0];
222
223        *ls += kt_ptr[tx]*d;
224
225        const float y = simd_sum(dot(*ls, qt_ptr[tx]));
226
227        if (tx == 0) {
228            *dst_attn = y*scale;
229        }
230
231        q_ptr += args.ns02;
232        k_ptr += args.ns12;
233        v_ptr += args.ns22;
234
235        b_ptr += args.ne21;
236        g_ptr += args.ne21*G;
237
238        dst_attn += args.ne21*S_v;
239    }
240
241    // when fused with the cache cpy, write the snapshots straight into the cache buffer using
242    // the slot stride; otherwise append them after the attn scores (nb_out == 0)
243    const bool fused = args.nb_out > 0;
244    const device float * state_out = fused ? (device float *)dst_fuse : (device float *)dst + args.ne23*args.ne22*args.ne21*S_v;
245    const uint slot_stride = fused ? (uint)args.nb_out : S_v*S_v;
246
247    device float * dst_state  = (device float *)state_out + (i23*args.ne21 + i21)*slot_stride + i20;
248    device T     * dstt_state = (device T     *) (dst_state);
249
250    FOR_UNROLL (short j = 0; j < NSG; j++) {
251        const short is = tx*NSG + j;
252        dst_state[is*S_v] = lsf[j];
253    }
254
255#undef S_v
256#undef G
257}
258
259typedef decltype(kernel_gated_delta_net_impl<float4, 4>) kernel_gated_delta_net_t;
260
261template [[host_name("kernel_gated_delta_net_f32_1")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float,  1>;
262template [[host_name("kernel_gated_delta_net_f32_2")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float2, 2>;
263template [[host_name("kernel_gated_delta_net_f32_4")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl<float4, 4>;
264#endif