#version 450
// Gated delta-rule recurrence (decode/sequential path), mirroring cuda/gdn.cu
// gated_delta_rule_recurrence_kernel. V-tiled: one workgroup per (v_tile, bh),
// one invocation per V-column. Each invocation keeps its state column s[k] in a
// private register array and walks the sequence once.
//
// q,k: [BH, S, K] v: [BH, S, V] g,beta: [BH, S]
// state: [BH, K, V] (in/out, written in place) output: [BH, S, V]
//
// k_dim is a runtime push constant; MAX_K bounds the private array (matches the
// CUDA fallback's MAX_K=256, covering every shipping GDN config: K=64/128).
//
// Inactive lanes (v_idx >= v_dim) do NOT early-return: they still hit every
// barrier() and help with the cooperative k/q loads, so control flow stays
// uniform across the workgroup (a hard requirement for barrier()). Their per-
// column compute and the final state write are guarded by `lane_on`.
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
const uint MAX_K = 256u;
const uint BV = 64u;
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 seq_len; uint k_dim; uint v_dim; };
shared float k_buf[MAX_K];
shared float q_buf[MAX_K];
void main() {
uint v_tile = gl_WorkGroupID.x;
uint bh = gl_WorkGroupID.y;
uint tid = gl_LocalInvocationID.x;
uint v_idx = v_tile * BV + tid;
bool lane_on = v_idx < v_dim;
uint qk_bh = bh * seq_len * k_dim;
uint v_bh = bh * seq_len * v_dim;
uint gb_bh = bh * seq_len;
uint state_bh = bh * k_dim * v_dim;
uint out_bh = bh * seq_len * v_dim;
float s[MAX_K];
if (lane_on) {
for (uint j = 0u; j < k_dim; j++) {
s[j] = state[state_bh + j * v_dim + v_idx];
}
}
for (uint t = 0u; t < seq_len; t++) {
for (uint j = tid; j < k_dim; j += BV) {
k_buf[j] = k[qk_bh + t * k_dim + j];
}
barrier();
float delta = 0.0;
if (lane_on) {
float decay = exp(g[gb_bh + t]);
float beta_t = beta[gb_bh + t];
float v_t = v[v_bh + t * v_dim + v_idx];
float kv_mem = 0.0;
for (uint j = 0u; j < k_dim; j++) {
s[j] *= decay;
kv_mem = fma(s[j], k_buf[j], kv_mem);
}
delta = (v_t - kv_mem) * beta_t;
}
for (uint j = tid; j < k_dim; j += BV) {
q_buf[j] = q[qk_bh + t * k_dim + j];
}
barrier();
if (lane_on) {
float y_t = 0.0;
for (uint j = 0u; j < k_dim; j++) {
s[j] = fma(k_buf[j], delta, s[j]);
y_t = fma(s[j], q_buf[j], y_t);
}
outp[out_bh + t * v_dim + v_idx] = y_t;
}
barrier();
}
if (lane_on) {
for (uint j = 0u; j < k_dim; j++) {
state[state_bh + j * v_dim + v_idx] = s[j];
}
}
}