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 →
5.6 KiB179 linesraw
1#include "common.h"
2
3kernel void kernel_rwkv_wkv6_f32(
4    device const float * k,
5    device const float * v,
6    device const float * r,
7    device const float * tf,
8    device const float * td,
9    device const float * state_in,
10    device       float * dst,
11    constant    uint & B,
12    constant    uint & T,
13    constant    uint & C,
14    constant    uint & H,
15    uint3 tgpig[[threadgroup_position_in_grid]],
16    uint3 tpitg[[thread_position_in_threadgroup]],
17    uint3   ntg[[threads_per_threadgroup]])  {
18
19    const uint head_size = 64; // TODO: support head_size = 128
20    const uint batch_id = tgpig.x / H;
21    const uint head_id = tgpig.x % H;
22    const uint tid = tpitg.x;
23
24    if (batch_id >= B || head_id >= H) {
25        return;
26    }
27
28    const uint state_size = C * head_size;
29    const uint n_seq_tokens = T / B;
30
31    threadgroup float _k[head_size];
32    threadgroup float _r[head_size];
33    threadgroup float _tf[head_size];
34    threadgroup float _td[head_size];
35
36    float state[head_size];
37
38    for (uint i = 0; i < head_size; i++) {
39        state[i] = state_in[batch_id * state_size + head_id * head_size * head_size
40                          + i * head_size + tid];
41    }
42
43    threadgroup_barrier(mem_flags::mem_threadgroup);
44    _tf[tid] = tf[head_id * head_size + tid];
45    threadgroup_barrier(mem_flags::mem_threadgroup);
46
47    const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid;
48    const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid;
49
50    for (uint t = start_t; t < end_t; t += C) {
51        threadgroup_barrier(mem_flags::mem_threadgroup);
52        _k[tid] = k[t];
53        _r[tid] = r[t];
54        _td[tid] = td[t];
55        threadgroup_barrier(mem_flags::mem_threadgroup);
56
57        const float v_val = v[t];
58        float y = 0.0;
59
60        for (uint j = 0; j < head_size; j += 4) {
61            float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]);
62            float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]);
63            float4 tf_vec = float4(_tf[j], _tf[j+1], _tf[j+2], _tf[j+3]);
64            float4 td_vec = float4(_td[j], _td[j+1], _td[j+2], _td[j+3]);
65            float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
66
67            float4 kv = k_vec * v_val;
68
69            float4 temp = tf_vec * kv + s_vec;
70            y += dot(r_vec, temp);
71
72            s_vec = s_vec * td_vec + kv;
73            state[j]   = s_vec[0];
74            state[j+1] = s_vec[1];
75            state[j+2] = s_vec[2];
76            state[j+3] = s_vec[3];
77        }
78
79        dst[t] = y;
80    }
81
82    for (uint i = 0; i < head_size; i++) {
83        dst[T * C + batch_id * state_size + head_id * head_size * head_size
84            + i * head_size + tid] = state[i];
85    }
86}
87
88kernel void kernel_rwkv_wkv7_f32(
89    device const float * r,
90    device const float * w,
91    device const float * k,
92    device const float * v,
93    device const float * a,
94    device const float * b,
95    device const float * state_in,
96    device       float * dst,
97    constant    uint & B,
98    constant    uint & T,
99    constant    uint & C,
100    constant    uint & H,
101    uint3 tgpig[[threadgroup_position_in_grid]],
102    uint3 tpitg[[thread_position_in_threadgroup]],
103    uint3   ntg[[threads_per_threadgroup]])  {
104
105    const uint head_size = 64; // TODO: support head_size = 128
106    const uint batch_id = tgpig.x / H;
107    const uint head_id = tgpig.x % H;
108    const uint tid = tpitg.x;
109
110    if (batch_id >= B || head_id >= H) {
111        return;
112    }
113
114    const uint state_size = C * head_size;
115    const uint n_seq_tokens = T / B;
116
117    threadgroup float _r[head_size];
118    threadgroup float _w[head_size];
119    threadgroup float _k[head_size];
120    threadgroup float _a[head_size];
121    threadgroup float _b[head_size];
122
123    float state[head_size];
124
125    for (uint i = 0; i < head_size; i++) {
126        state[i] = state_in[batch_id * state_size + head_id * head_size * head_size
127                          + tid * head_size + i];
128    }
129
130    const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid;
131    const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid;
132
133    for (uint t = start_t; t < end_t; t += C) {
134        threadgroup_barrier(mem_flags::mem_threadgroup);
135        _r[tid] = r[t];
136        _w[tid] = w[t];
137        _k[tid] = k[t];
138        _a[tid] = a[t];
139        _b[tid] = b[t];
140        threadgroup_barrier(mem_flags::mem_threadgroup);
141
142        const float v_val = v[t];
143        float y = 0.0, sa = 0.0;
144
145        float4 sa_vec(0.0);
146
147        for (uint j = 0; j < head_size; j += 4) {
148            float4 a_vec = float4(_a[j], _a[j+1], _a[j+2], _a[j+3]);
149            float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
150            sa_vec += a_vec * s_vec;
151        }
152        sa = sa_vec[0] + sa_vec[1] + sa_vec[2] + sa_vec[3];
153
154        for (uint j = 0; j < head_size; j += 4) {
155            float4 r_vec = float4(_r[j], _r[j+1], _r[j+2], _r[j+3]);
156            float4 w_vec = float4(_w[j], _w[j+1], _w[j+2], _w[j+3]);
157            float4 k_vec = float4(_k[j], _k[j+1], _k[j+2], _k[j+3]);
158            float4 b_vec = float4(_b[j], _b[j+1], _b[j+2], _b[j+3]);
159            float4 s_vec = float4(state[j], state[j+1], state[j+2], state[j+3]);
160
161            float4 kv = k_vec * v_val;
162
163            s_vec = s_vec * w_vec + kv + sa * b_vec;
164            y += dot(s_vec, r_vec);
165
166            state[j]   = s_vec[0];
167            state[j+1] = s_vec[1];
168            state[j+2] = s_vec[2];
169            state[j+3] = s_vec[3];
170        }
171
172        dst[t] = y;
173    }
174
175    for (uint i = 0; i < head_size; i++) {
176        dst[T * C + batch_id * state_size + head_id * head_size * head_size
177            + tid * head_size + i] = state[i];
178    }
179}