char-fork/llama.cpp

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

Pranav Uttarkarvulkan backend ops: implemented GATED_LINEAR_ATTN (#25601)f26efa02a

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.5 KiB82 linesraw
1#version 450
2
3#extension GL_EXT_control_flow_attributes : require
4
5#define BLOCK_SIZE 64
6layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
7
8layout(push_constant) uniform Parameters {
9    uint B;
10    uint T;
11    uint C;
12    uint H;
13    float scale;
14};
15
16layout(binding = 0) readonly buffer KBuf { A_TYPE k[]; };
17layout(binding = 1) readonly buffer VBuf { A_TYPE v[]; };
18layout(binding = 2) readonly buffer QBuf { A_TYPE q[]; };
19layout(binding = 3) readonly buffer GBuf { A_TYPE g[]; };
20layout(binding = 4) readonly buffer StateBuf { A_TYPE state_in[]; };
21layout(binding = 5) buffer DstBuf { A_TYPE dst[]; };
22
23shared A_TYPE _k[BLOCK_SIZE], _q[BLOCK_SIZE], _g[BLOCK_SIZE];
24
25void main() {
26    const uint head_size = BLOCK_SIZE;
27    const uint batch_id = gl_WorkGroupID.x / H;
28    const uint head_id = gl_WorkGroupID.x % H;
29    const uint tid = gl_LocalInvocationID.x;
30
31    const uint state_size = C * head_size;
32    const uint n_seq_tokens = T / B;
33
34    if (batch_id >= B || head_id >= H) {
35        return;
36    }
37
38    // state[i] holds column tid of this head's state matrix: S[i][tid]
39    A_TYPE state[BLOCK_SIZE];
40    [[unroll]] for (uint i = 0; i < head_size; i++) {
41        state[i] = state_in[batch_id * state_size + head_id * head_size * head_size
42                          + i * head_size + tid];
43    }
44
45    const uint start_t = batch_id * n_seq_tokens * C + head_id * head_size + tid;
46    const uint end_t = (batch_id + 1) * n_seq_tokens * C + head_id * head_size + tid;
47
48    for (uint t = start_t; t < end_t; t += C) {
49        barrier();
50        _k[tid] = k[t];
51        _q[tid] = q[t];
52        _g[tid] = g[t];
53        barrier();
54
55        const A_TYPE v_val = v[t];
56        A_TYPE y = 0.0;
57
58        [[unroll]] for (uint i = 0; i < head_size; i += 4) {
59            vec4 k_vec = vec4(_k[i], _k[i+1], _k[i+2], _k[i+3]);
60            vec4 q_vec = vec4(_q[i], _q[i+1], _q[i+2], _q[i+3]);
61            vec4 g_vec = vec4(_g[i], _g[i+1], _g[i+2], _g[i+3]);
62            vec4 s_vec = vec4(state[i], state[i+1], state[i+2], state[i+3]);
63
64            vec4 kv = k_vec * v_val;
65
66            s_vec = s_vec * g_vec + kv;
67            y += dot(q_vec, s_vec);
68
69            state[i]   = s_vec.x;
70            state[i+1] = s_vec.y;
71            state[i+2] = s_vec.z;
72            state[i+3] = s_vec.w;
73        }
74
75        dst[t] = y * scale;
76    }
77
78    [[unroll]] for (uint i = 0; i < head_size; i++) {
79        dst[T * C + batch_id * state_size + head_id * head_size * head_size
80            + i * head_size + tid] = state[i];
81    }
82}