#version 450
// Gated delta rule, SINGLE timestep (decode, seq_len==1). One invocation per (bh, v) state column,
// looping the k_dim in a local array. Matches gdn.rs::gated_delta_rule_recurrence for seq=1 and the
// CUDA gated_delta_rule_recurrence_kernel: state is [BH, K, V] with state[bh][j][v] at j*V + v.
// decay = exp(g[bh]); s[j] *= decay; kv_mem = sum_j s[j]*k[j]
// delta = (v[v] - kv_mem) * beta[bh]; s[j] += k[j]*delta; y[v] = sum_j s[j]*q[j]
// q is pre-scaled by the caller (1/sqrt(k_dim)), exactly as the CUDA recurrence wrapper does.
// MAX_K bounds the per-invocation k array; k_dim (head_k_dim) must be <= MAX_K (asserted host-side).
#define MAX_K 256u
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0) readonly buffer Q { float q[]; };
layout(set = 0, binding = 1) readonly buffer K { float k[]; };
layout(set = 0, binding = 2) readonly buffer V { float v[]; };
layout(set = 0, binding = 3) readonly buffer G { float g[]; };
layout(set = 0, binding = 4) readonly buffer Beta { float beta[]; };
layout(set = 0, binding = 5) buffer State { float state[]; };
layout(set = 0, binding = 6) writeonly buffer Out { float outp[]; };
layout(push_constant) uniform Pc { uint bh_count; uint k_dim; uint v_dim; };
void main() {
uint gid = gl_GlobalInvocationID.x;
uint total = bh_count * v_dim;
if (gid >= total) { return; }
uint bh = gid / v_dim;
uint v_idx = gid - bh * v_dim;
uint q_base = bh * k_dim; // q,k: [BH, K] (seq=1)
uint state_base = bh * k_dim * v_dim;
float decay = exp(g[bh]);
float beta_t = beta[bh];
float v_t = v[bh * v_dim + v_idx];
float s[MAX_K];
float kv_mem = 0.0;
for (uint j = 0u; j < k_dim; j++) {
float sj = state[state_base + j * v_dim + v_idx] * decay;
s[j] = sj;
kv_mem += sj * k[q_base + j];
}
float delta = (v_t - kv_mem) * beta_t;
float y_t = 0.0;
for (uint j = 0u; j < k_dim; j++) {
float sj = s[j] + k[q_base + j] * delta;
state[state_base + j * v_dim + v_idx] = sj;
y_t += sj * q[q_base + j];
}
outp[bh * v_dim + v_idx] = y_t;
}