Brunobkr/llama.cpp_AlgMor24_github

1#version 4502 3#extension GL_EXT_control_flow_attributes : require4 5#define BLOCK_SIZE 646layout(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_size42                          + 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_size80            + i * head_size + tid] = state[i];81    }82}83