#version 450

#extension GL_EXT_control_flow_attributes : require

layout(constant_id = 0) const uint S_V = 128;
layout(constant_id = 1) const uint KDA = 0;

layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;

layout(push_constant) uniform Parameters {
    uint H;
    uint n_tokens;
    uint n_seqs;
    uint s_off;
    uint sq1, sq2, sq3;
    uint sv1, sv2, sv3;
    uint sb1, sb2, sb3;
    uint neq1, rq3;
    float scale;
};

layout(binding = 0) readonly  buffer QBuf     { FLOAT_TYPE data_q[];     };
layout(binding = 1) readonly  buffer KBuf     { FLOAT_TYPE data_k[];     };
layout(binding = 2) readonly  buffer VBuf     { FLOAT_TYPE data_v[];     };
layout(binding = 3) readonly  buffer GBuf     { FLOAT_TYPE data_g[];     };
layout(binding = 4) readonly  buffer BetaBuf  { FLOAT_TYPE data_beta[];  };
layout(binding = 5) readonly  buffer StateBuf { FLOAT_TYPE data_state[]; };
layout(binding = 6)           buffer DstBuf   { FLOAT_TYPE data_dst[];   };

shared FLOAT_TYPE s_k[S_V];
shared FLOAT_TYPE s_q[S_V];
shared FLOAT_TYPE s_g[S_V]; // KDA only: cached exp(g[i])

void main() {
    const uint head_id = gl_WorkGroupID.x;
    const uint seq_id  = gl_WorkGroupID.y;
    const uint col     = gl_LocalInvocationID.x;

    const uint iq1 = head_id % neq1;
    const uint iq3 = seq_id / rq3;

    const uint state_size = S_V * S_V;
    const uint state_base = (seq_id * H + head_id) * state_size;

    FLOAT_TYPE state[S_V];
    [[unroll]] for (uint i = 0; i < S_V; i++) {
        state[i] = FLOAT_TYPE(data_state[state_base + col * S_V + i]);
    }

    uint attn_off = (seq_id * n_tokens * H + head_id) * S_V;

    for (uint t = 0; t < n_tokens; t++) {
        const uint q_off = iq3 * sq3 + t * sq2 + iq1 * sq1;
        const uint k_off = q_off;
        const uint v_off = seq_id * sv3 + t * sv2 + head_id * sv1;

        s_q[col] = FLOAT_TYPE(data_q[q_off + col]);
        s_k[col] = FLOAT_TYPE(data_k[k_off + col]);

        const uint gb_off = seq_id * sb3 + t * sb2 + head_id * sb1;

        if (KDA != 0) {
            const uint g_base = gb_off * S_V;
            s_g[col] = exp(FLOAT_TYPE(data_g[g_base + col]));
        }

        barrier();

        const FLOAT_TYPE v_val = FLOAT_TYPE(data_v[v_off + col]);
        const FLOAT_TYPE beta_val = FLOAT_TYPE(data_beta[gb_off]);

        if (KDA == 0) {
            const FLOAT_TYPE g_val = exp(FLOAT_TYPE(data_g[gb_off]));

            FLOAT_TYPE kv_col = 0.0;
            [[unroll]] for (uint i = 0; i < S_V; i += 4) {
                kv_col += dot(
                    vec4(state[i], state[i+1], state[i+2], state[i+3]),
                    vec4(s_k[i], s_k[i+1], s_k[i+2], s_k[i+3])
                );
            }

            FLOAT_TYPE delta_col = (v_val - g_val * kv_col) * beta_val;

            FLOAT_TYPE attn_col = 0.0;
            [[unroll]] for (uint i = 0; i < S_V; i += 4) {
                vec4 sv = vec4(state[i], state[i+1], state[i+2], state[i+3]);
                vec4 kv = vec4(s_k[i], s_k[i+1], s_k[i+2], s_k[i+3]);
                sv = g_val * sv + kv * delta_col;
                state[i] = sv.x; state[i+1] = sv.y; state[i+2] = sv.z; state[i+3] = sv.w;

                attn_col += dot(sv, vec4(s_q[i], s_q[i+1], s_q[i+2], s_q[i+3]));
            }

            data_dst[attn_off + col] = attn_col * scale;
        } else {
            FLOAT_TYPE kv_col = 0.0;
            [[unroll]] for (uint i = 0; i < S_V; i += 4) {
                vec4 gv = vec4(s_g[i], s_g[i+1], s_g[i+2], s_g[i+3]);
                vec4 sv = vec4(state[i], state[i+1], state[i+2], state[i+3]);
                vec4 kv = vec4(s_k[i], s_k[i+1], s_k[i+2], s_k[i+3]);
                kv_col += dot(gv * sv, kv);
            }

            FLOAT_TYPE delta_col = (v_val - kv_col) * beta_val;

            FLOAT_TYPE attn_col = 0.0;
            [[unroll]] for (uint i = 0; i < S_V; i += 4) {
                vec4 gv = vec4(s_g[i], s_g[i+1], s_g[i+2], s_g[i+3]);
                vec4 sv = vec4(state[i], state[i+1], state[i+2], state[i+3]);
                vec4 kv = vec4(s_k[i], s_k[i+1], s_k[i+2], s_k[i+3]);
                sv = gv * sv + kv * delta_col;
                state[i] = sv.x; state[i+1] = sv.y; state[i+2] = sv.z; state[i+3] = sv.w;

                attn_col += dot(sv, vec4(s_q[i], s_q[i+1], s_q[i+2], s_q[i+3]));
            }

            data_dst[attn_off + col] = attn_col * scale;
        }

        attn_off += S_V * H;
        barrier();
    }

    [[unroll]] for (uint i = 0; i < S_V; i++) {
        data_dst[s_off + state_base + col * S_V + i] = state[i];
    }
}