// metal/kernels/gdn_prompt.metal - a GDN layer's recurrence for a PROMPT chunk of T tokens, in three passes: the work
// that does not touch the state runs parallel over the tokens; only the state update is a scan, and the columns of a
// head's state are independent in it (column j needs k, q, v_j, the gate and beta), so a lane owns one column.
//
//   gdn_pre    per (token, head): conv (4 taps; history oldest first) -> SiLU; q, k L2 norm * 1/sqrt(128) (the decode
//              step's arithmetic); kq = k . q per key head; gate = exp(softplus(alpha + dt) * a), beta = sigmoid(beta);
//              the chunk's last 3 conv inputs -> hist_out
//   gdn_scan   per (value head, 32 columns x CB simdgroups): the state column in registers (128 floats a lane);
//              per token a = S^T k, b = S^T q (old state), delta = (v - g a) beta, o = (g b + delta kq) / sqrt(128)
//              (= S_new^T q / sqrt(128)), S = g S + k delta^T; q and k staged per TB tokens in threadgroup memory
//   gdn_post   per (token, head): y = o * rsqrt(mean(o^2) + eps) * gamma * sigmoid(z) -> fp16 (the out GEMM's operand)
//              and, with `f32`, fp32
//
// State layout as Strata's: S[(i * h_v + h) * 128 + j].  Value head h reads q/k head h % h_k.
// The kernels are new entries (nothing of gdn.metal changes).

struct GdnPArgs {
    uint n_tok;
    uint h_k;        // 16
    uint h_v;        // 48
    uint channels;   // 10240
    float eps;       // 1e-6
    uint f32;        // gdn_post: also write fp32 y
};

static inline float gp_conv(device const float* qkv, device const float* hist, device const float* cw, int ch, int t,
                            int channels) {
    const float4 w = *(device const float4*)(cw + ch * 4);
    float s = 0.0f;
#pragma clang loop unroll(full)
    for (int k = 0; k < 4; ++k) {
        const int src = t - 3 + k;
        const float x = src >= 0 ? qkv[src * channels + ch] : hist[ch * 3 + (src + 3)];
        s = fma(x, w[k], s);
    }
    s = s + 0.0f;
    return s / (1.0f + exp(-s));
}

// grid ((h_k + h_v) / 8, T) = (8, T) or (h_k + h_v / 2, T) = (40, T), 256 threads: when grid_x == (h_k + h_v) / 8,
// each simdgroup (32 lanes) owns one key head (x < h_k / 8: q, k and kq in registers) or one value head (else: vc and
// gb), with zero threadgroup memory and zero barriers; grid_x == h_k + h_v / 2 keeps the legacy grid.
kernel void gdn_pre(device const float* qkv [[buffer(0)]], device const float* hist [[buffer(1)]],
                    device const float* conv_w [[buffer(2)]], device const float* alpha [[buffer(3)]],
                    device const float* beta [[buffer(4)]], device const float* dt [[buffer(5)]],
                    device const float* ssm_a [[buffer(6)]], device float* qn [[buffer(7)]], device float* kn [[buffer(8)]],
                    device float* vc [[buffer(9)]], device float2* gb [[buffer(10)]], device float* kq [[buffer(11)]],
                    device float* hist_out [[buffer(12)]], constant GdnPArgs& a [[buffer(13)]],
                    uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]],
                    ushort tid [[thread_index_in_threadgroup]], ushort sg [[simdgroup_index_in_threadgroup]],
                    ushort lane [[thread_index_in_simdgroup]]) {
    constexpr int S = 128;
    const int x = int(tg.x), t = int(tg.y), hk_n = int(a.h_k), hv = int(a.h_v), C = int(a.channels), T = int(a.n_tok);
    const bool packed = int(tgs.x) * 8 == hk_n + hv;
    // alpha and beta as the halves of one [T][2 h_v] GEMM (beta == alpha + h_v), else each [T][h_v]
    const int ab = beta == alpha + hv ? 2 * hv : hv;
    const int k_tgs = packed ? hk_n / 8 : hk_n;
    if (x < k_tgs) {
        const int kh = packed ? x * 8 + int(sg) : (sg == 0 ? x : hk_n);
        if (kh >= hk_n) return;
        float cv_q[4], cv_k[4];
        if (t >= 3) {
            device const float* pq = qkv + (t - 3) * C + kh * S + int(lane);
            device const float* pk = qkv + (t - 3) * C + (hk_n + kh) * S + int(lane);
#pragma clang loop unroll(full)
            for (int m = 0; m < 4; ++m) {
                const int d = int(lane) + m * 32;
                const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d;
                const float4 wq = *(device const float4*)(conv_w + ch_q * 4);
                const float4 wk = *(device const float4*)(conv_w + ch_k * 4);
                device const float* pqm = pq + m * 32;
                device const float* pkm = pk + m * 32;
                float sq = fma(pqm[0], wq.x, 0.0f);
                sq = fma(pqm[C], wq.y, sq);
                sq = fma(pqm[2 * C], wq.z, sq);
                sq = fma(pqm[3 * C], wq.w, sq) + 0.0f;
                cv_q[m] = sq / (1.0f + exp(-sq));
                float sk = fma(pkm[0], wk.x, 0.0f);
                sk = fma(pkm[C], wk.y, sk);
                sk = fma(pkm[2 * C], wk.z, sk);
                sk = fma(pkm[3 * C], wk.w, sk) + 0.0f;
                cv_k[m] = sk / (1.0f + exp(-sk));
            }
            if (t == T - 1)
                for (int m = 0; m < 4; ++m) {
                    const int d = int(lane) + m * 32;
                    const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d;
                    for (int k = 0; k < 3; ++k) {
                        hist_out[ch_q * 3 + k] = pq[m * 32 + (k + 1) * C];
                        hist_out[ch_k * 3 + k] = pk[m * 32 + (k + 1) * C];
                    }
                }
        } else {
#pragma clang loop unroll(full)
            for (int m = 0; m < 4; ++m) {
                const int d = int(lane) + m * 32;
                const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d;
                cv_q[m] = gp_conv(qkv, hist, conv_w, ch_q, t, C);
                cv_k[m] = gp_conv(qkv, hist, conv_w, ch_k, t, C);
            }
            if (t == T - 1)
                for (int m = 0; m < 4; ++m) {
                    const int d = int(lane) + m * 32;
                    const int ch_q = kh * S + d, ch_k = (hk_n + kh) * S + d;
                    for (int k = 0; k < 3; ++k) {
                        const int src = T - 3 + k;
                        hist_out[ch_q * 3 + k] = src >= 0 ? qkv[src * C + ch_q] : hist[ch_q * 3 + (src + 3)];
                        hist_out[ch_k * 3 + k] = src >= 0 ? qkv[src * C + ch_k] : hist[ch_k * 3 + (src + 3)];
                    }
                }
        }
        const float4 sq = simd_sum(float4(cv_q[0] * cv_q[0], cv_q[1] * cv_q[1], cv_q[2] * cv_q[2], cv_q[3] * cv_q[3]));
        const float4 sk = simd_sum(float4(cv_k[0] * cv_k[0], cv_k[1] * cv_k[1], cv_k[2] * cv_k[2], cv_k[3] * cv_k[3]));
        const float tot_q = (sq.x + sq.y) + (sq.z + sq.w), tot_k = (sk.x + sk.y) + (sk.z + sk.w);
        const float rs_q = rsqrt(tot_q / float(S) + a.eps / float(S));
        const float rs_k = rsqrt(tot_k / float(S) + a.eps / float(S));
        const float scale = 1.0f / sqrt(float(S));
        const int base = (t * hk_n + kh) * S + int(lane);
        float4 prod;
#pragma clang loop unroll(full)
        for (int m = 0; m < 4; ++m) {
            const float nq = (rs_q * cv_q[m]) * scale, nk = (rs_k * cv_k[m]) * scale;
            qn[base + m * 32] = nq;
            kn[base + m * 32] = nk;
            prod[m] = nk * nq;
        }
        const float4 pv = simd_sum(prod);
        if (lane == 0) kq[t * hk_n + kh] = (pv.x + pv.y) + (pv.z + pv.w);
    } else if (packed) {
        const int h = (x - k_tgs) * 8 + int(sg);
        if (t >= 3) {
            const int d = int(lane) * 4, ch = (2 * hk_n + h) * S + d;
            device const float4* p = (device const float4*)(qkv + (t - 3) * C + ch);
            device const float4* w = (device const float4*)(conv_w + ch * 4);
            const float4 w0 = w[0], w1 = w[1], w2 = w[2], w3 = w[3];
            float4 cv = fma(p[0], float4(w0.x, w1.x, w2.x, w3.x), float4(0.0f));
            cv = fma(p[C / 4], float4(w0.y, w1.y, w2.y, w3.y), cv);
            cv = fma(p[C / 2], float4(w0.z, w1.z, w2.z, w3.z), cv);
            cv = fma(p[3 * C / 4], float4(w0.w, w1.w, w2.w, w3.w), cv);
            cv = cv + 0.0f;
            cv = cv / (1.0f + exp(-cv));
            *((device float4*)(vc + (t * hv + h) * S + d)) = cv;
            if (t == T - 1)
                for (int k = 0; k < 3; ++k) {
                    const float4 v = p[(k + 1) * C / 4];
                    for (int m = 0; m < 4; ++m) hist_out[(ch + m) * 3 + k] = v[m];
                }
        } else {
#pragma clang loop unroll(full)
            for (int m = 0; m < 4; ++m) {
                const int d = int(lane) + m * 32;
                const int ch = 2 * hk_n * S + h * S + d;
                const float cv = gp_conv(qkv, hist, conv_w, ch, t, C);
                if (t == T - 1)
                    for (int k = 0; k < 3; ++k) {
                        const int src = T - 3 + k;
                        hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)];
                    }
                vc[(t * hv + h) * S + d] = cv;
            }
        }
        if (lane == 0) {
            const float v = alpha[t * ab + h] + dt[h];
            const float e = exp(v), u = 1.0f + e;
            const float sp = v > 20.0f ? v : (u == 1.0f ? e : log(u) * e / (u - 1.0f));
            gb[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h])));
        }
    } else {
        const int h = 2 * (x - k_tgs) + (tid >> 7), d = tid & 127;
        const int ch = 2 * hk_n * S + h * S + d;
        const float cv = gp_conv(qkv, hist, conv_w, ch, t, C);
        if (t == T - 1)
            for (int k = 0; k < 3; ++k) {
                const int src = T - 3 + k;
                hist_out[ch * 3 + k] = src >= 0 ? qkv[src * C + ch] : hist[ch * 3 + (src + 3)];
            }
        vc[(t * hv + h) * S + d] = cv;
        if (d == 0) {
            const float v = alpha[t * ab + h] + dt[h];
            const float e = exp(v), u = 1.0f + e;
            const float sp = v > 20.0f ? v : (u == 1.0f ? e : log(u) * e / (u - 1.0f));
            gb[t * hv + h] = float2(exp(sp * ssm_a[h]), 1.0f / (1.0f + exp(-beta[t * ab + h])));
        }
    }
}

// grid (h_v * 4 / CB), CB * 32 threads; TB tokens of q and k staged at a time
template <int CB, int TB>
kernel void gdn_scan_t(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]],
                       device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]],
                       device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]],
                       device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]],
                       uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
                       ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
    constexpr int S = 128, NB = 4 / CB;
    threadgroup float4 qs[TB][32], ks[TB][32];
    const int h = int(tg.x) / NB, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k);
    const int j = ((int(tg.x) % NB) * CB + sg) * 32 + lane;
    const int T = int(a.n_tok);
    const float scale = 1.0f / sqrt(float(S));
    float st[S];
#pragma clang loop unroll(full)
    for (int i = 0; i < S; ++i) st[i] = state[(i * hv + h) * S + j];
    for (int t0 = 0; t0 < T; t0 += TB) {
        const int nb = min(TB, T - t0);
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (int idx = tid; idx < nb * 32; idx += CB * 32) {
            const int tt = idx >> 5, w = idx & 31;
            qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w];
            ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (int tt = 0; tt < nb; ++tt) {
            const int t = t0 + tt;
            const float2 g_b = gb[t * hv + h];
            const float kqv = kq[t * hkn + hk];
            const float v = vc[(t * hv + h) * S + j];
            float ka = 0.0f, kb = 0.0f, qa = 0.0f, qb = 0.0f;
#pragma clang loop unroll(full)
            for (int w = 0; w < 32; w += 2) {
                const float4 k0 = ks[tt][w], k1 = ks[tt][w + 1], q0 = qs[tt][w], q1 = qs[tt][w + 1];
                ka = fma(st[4 * w], k0.x, ka); ka = fma(st[4 * w + 1], k0.y, ka);
                ka = fma(st[4 * w + 2], k0.z, ka); ka = fma(st[4 * w + 3], k0.w, ka);
                kb = fma(st[4 * w + 4], k1.x, kb); kb = fma(st[4 * w + 5], k1.y, kb);
                kb = fma(st[4 * w + 6], k1.z, kb); kb = fma(st[4 * w + 7], k1.w, kb);
                qa = fma(st[4 * w], q0.x, qa); qa = fma(st[4 * w + 1], q0.y, qa);
                qa = fma(st[4 * w + 2], q0.z, qa); qa = fma(st[4 * w + 3], q0.w, qa);
                qb = fma(st[4 * w + 4], q1.x, qb); qb = fma(st[4 * w + 5], q1.y, qb);
                qb = fma(st[4 * w + 6], q1.z, qb); qb = fma(st[4 * w + 7], q1.w, qb);
            }
            const float g = g_b.x, kv = ka + kb, qv = qa + qb;
            const float delta = (v - g * kv) * g_b.y;
            o[(t * hv + h) * S + j] = fma(g, qv, delta * kqv) * scale;
#pragma clang loop unroll(full)
            for (int w = 0; w < 32; ++w) {
                const float4 k4 = ks[tt][w];
                st[4 * w] = fma(g, st[4 * w], k4.x * delta);
                st[4 * w + 1] = fma(g, st[4 * w + 1], k4.y * delta);
                st[4 * w + 2] = fma(g, st[4 * w + 2], k4.z * delta);
                st[4 * w + 3] = fma(g, st[4 * w + 3], k4.w * delta);
            }
        }
    }
#pragma clang loop unroll(full)
    for (int i = 0; i < S; ++i) state[(i * hv + h) * S + j] = st[i];
}

// a probe for the bench (function constant 6): 1 = no per-token exchange (wrong results), 2 = no state update
constant int GDN_PROBE [[function_constant(6)]];
constant int gdn_probe = is_function_constant_defined(GDN_PROBE) ? GDN_PROBE : 0;

// the rows of a column split over IB simdgroups (128 / IB rows a lane, fewer registers): grid (h_v * 4), IB * 32
// threads; lane = column (32 columns a threadgroup); when TB == 0 (single-pass), each simdgroup loads its R rows of q
// and k across its lanes and broadcasts them with simd_broadcast (one initial dot pass at t = 0, one barrier a token,
// only red[2][IB][32] in threadgroup memory); TB > 0 keeps the TB-token threadgroup-staged reference.
template <int IB, int TB>
kernel void gdn_scan_r(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]],
                       device const float* kn [[buffer(2)]], device const float* vc [[buffer(3)]],
                       device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]],
                       device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]],
                       uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
                       ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
    constexpr int S = 128, R = S / IB;
    threadgroup float2 red[2][IB][32];
    const int g_idx = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k);
    const int hk = (hv == 3 * hkn) ? (g_idx / 3) : (g_idx % hkn);
    const int h = (hv == 3 * hkn) ? (hk + (g_idx % 3) * hkn) : g_idx;
    const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * R;
    const int T = int(a.n_tok);
    const float scale = 1.0f / sqrt(float(S));
    const int rs = hv * S, s0 = (i0 * hv + h) * S + j;
    float st[R];
#pragma clang loop unroll(full)
    for (int i = 0; i < R; ++i) st[i] = (state + s0)[i * rs];
    if (TB == 0) {
        constexpr int C = (R + 31) / 32, R1 = R / C;
        const ushort lr = lane & (R1 - 1);
        const int ks = hkn * S, k0 = hk * S + i0 + int(lr), v0 = h * S + j;
        float kl[C], ql[C];
#pragma clang loop unroll(full)
        for (int c = 0; c < C; ++c) {
            kl[c] = kn[k0 + c * 32];
            ql[c] = qn[k0 + c * 32];
        }
        float2 g_b = gb[h];
        float v = vc[v0];
        float ka = 0.0f, qa = 0.0f;
#pragma clang loop unroll(full)
        for (int c = 0; c < C; ++c)
#pragma clang loop unroll(full)
            for (int r = 0; r < R1; ++r) {
                ka = fma(st[c * 32 + r], simd_broadcast(kl[c], ushort(r)), ka);
                qa = fma(st[c * 32 + r], simd_broadcast(ql[c], ushort(r)), qa);
            }
        for (int t = 0; t < T - 1; ++t) {
            const int slot = t & 1, tn = t + 1;
            const int kn_off = tn * ks + k0, vn_off = tn * rs + v0;
            float kn_l[C], qn_l[C];
#pragma clang loop unroll(full)
            for (int c = 0; c < C; ++c) {
                kn_l[c] = kn[kn_off + c * 32];
                qn_l[c] = qn[kn_off + c * 32];
            }
            const float2 g_b_n = gb[tn * hv + h];
            const float v_n = vc[vn_off];
            const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f;
            float2 s = float2(ka, qa);
            if (gdn_probe != 1) {
                red[slot][sg][lane] = s;
                threadgroup_barrier(mem_flags::mem_threadgroup);
                s = red[slot][0][lane];
#pragma clang loop unroll(full)
                for (int b = 1; b < IB; ++b) s += red[slot][b][lane];
            }
            const float g = g_b.x;
            const float delta = (v - g * s.x) * g_b.y;
            if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale;
            ka = 0.0f;
            qa = 0.0f;
            if (gdn_probe != 2)
#pragma clang loop unroll(full)
            for (int c = 0; c < C; ++c)
#pragma clang loop unroll(full)
                for (int r = 0; r < R1; ++r) {
                    const float kr = simd_broadcast(kl[c], ushort(r));
                    const float knr = simd_broadcast(kn_l[c], ushort(r)), qnr = simd_broadcast(qn_l[c], ushort(r));
                    st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta);
                    ka = fma(st[c * 32 + r], knr, ka);
                    qa = fma(st[c * 32 + r], qnr, qa);
                }
#pragma clang loop unroll(full)
            for (int c = 0; c < C; ++c) kl[c] = kn_l[c];
            g_b = g_b_n;
            v = v_n;
        }
        if (T > 0) {
            const int t = T - 1, slot = t & 1;
            const float kqv = sg == 0 ? kq[t * hkn + hk] : 0.0f;
            float2 s = float2(ka, qa);
            if (gdn_probe != 1) {
                red[slot][sg][lane] = s;
                threadgroup_barrier(mem_flags::mem_threadgroup);
                s = red[slot][0][lane];
#pragma clang loop unroll(full)
                for (int b = 1; b < IB; ++b) s += red[slot][b][lane];
            }
            const float g = g_b.x;
            const float delta = (v - g * s.x) * g_b.y;
            if (sg == 0) o[t * rs + v0] = fma(g, s.y, delta * kqv) * scale;
            if (gdn_probe != 2)
#pragma clang loop unroll(full)
            for (int c = 0; c < C; ++c)
#pragma clang loop unroll(full)
                for (int r = 0; r < R1; ++r) {
                    const float kr = simd_broadcast(kl[c], ushort(r));
                    st[c * 32 + r] = fma(g, st[c * 32 + r], kr * delta);
                }
        }
    } else {
        constexpr int R4 = R / 4, TBS = TB > 0 ? TB : 1;
        constexpr bool KC = R4 <= 8;
        threadgroup float4 qs[TBS][32], ks[TBS][32];
        threadgroup float vs[TBS][32];
        threadgroup float4 gk[TBS];
        device const float4* qn4_p = (device const float4*)qn;
        device const float4* kn4_p = (device const float4*)kn;
        for (int t0 = 0; t0 < T; t0 += TBS) {
            const int nb = min(TBS, T - t0);
            threadgroup_barrier(mem_flags::mem_threadgroup);
            for (int idx = tid; idx < nb * 32; idx += IB * 32) {
                const int tt = idx >> 5, w = idx & 31;
                qs[tt][w] = qn4_p[((t0 + tt) * hkn + hk) * 32 + w];
                ks[tt][w] = kn4_p[((t0 + tt) * hkn + hk) * 32 + w];
                vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w];
                if (w == 0) {
                    const float2 g_b = gb[(t0 + tt) * hv + h];
                    gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f);
                }
            }
            threadgroup_barrier(mem_flags::mem_threadgroup);
            float ka = 0.0f, qa = 0.0f;
            float4 kc[R4];
#pragma clang loop unroll(full)
            for (int w = 0; w < R4; ++w) {
                const float4 k4 = ks[0][i0 / 4 + w], q4 = qs[0][i0 / 4 + w];
                if (KC) kc[w] = k4;
                ka = fma(st[4 * w], k4.x, ka); ka = fma(st[4 * w + 1], k4.y, ka);
                ka = fma(st[4 * w + 2], k4.z, ka); ka = fma(st[4 * w + 3], k4.w, ka);
                qa = fma(st[4 * w], q4.x, qa); qa = fma(st[4 * w + 1], q4.y, qa);
                qa = fma(st[4 * w + 2], q4.z, qa); qa = fma(st[4 * w + 3], q4.w, qa);
            }
            for (int tt = 0; tt < nb; ++tt) {
                const int t = t0 + tt, slot = t & 1;
                const float4 g_b = gk[tt];
                const float kqv = g_b.z;
                const float v = vs[tt][lane];
                float2 s = float2(ka, qa);
                if (gdn_probe != 1) {
                    red[slot][sg][lane] = s;
                    threadgroup_barrier(mem_flags::mem_threadgroup);
                    s = red[slot][0][lane];
#pragma clang loop unroll(full)
                    for (int b = 1; b < IB; ++b) s += red[slot][b][lane];
                }
                const float g = g_b.x;
                const float delta = (v - g * s.x) * g_b.y;
                if (sg == 0) o[(t * hv + h) * S + j] = fma(g, s.y, delta * kqv) * scale;
                const int tn = min(tt + 1, nb - 1);
                ka = 0.0f;
                qa = 0.0f;
                if (gdn_probe != 2)
#pragma clang loop unroll(full)
                for (int w = 0; w < R4; ++w) {
                    const float4 k4 = KC ? kc[w] : ks[tt][i0 / 4 + w], kn4 = ks[tn][i0 / 4 + w], qn4 = qs[tn][i0 / 4 + w];
                    if (KC) kc[w] = kn4;
                    st[4 * w] = fma(g, st[4 * w], k4.x * delta);
                    st[4 * w + 1] = fma(g, st[4 * w + 1], k4.y * delta);
                    st[4 * w + 2] = fma(g, st[4 * w + 2], k4.z * delta);
                    st[4 * w + 3] = fma(g, st[4 * w + 3], k4.w * delta);
                    ka = fma(st[4 * w], kn4.x, ka); ka = fma(st[4 * w + 1], kn4.y, ka);
                    ka = fma(st[4 * w + 2], kn4.z, ka); ka = fma(st[4 * w + 3], kn4.w, ka);
                    qa = fma(st[4 * w], qn4.x, qa); qa = fma(st[4 * w + 1], qn4.y, qa);
                    qa = fma(st[4 * w + 2], qn4.z, qa); qa = fma(st[4 * w + 3], qn4.w, qa);
                }
            }
        }
    }
    device float* sp = state + s0;
#pragma clang loop unroll(full)
    for (int i = 0; i < R; ++i) sp[i * rs] = st[i];
}

// gdn_scan_r<8, 16> with the 16 state rows of a lane in four float4 variables (no array): a register-allocation check
kernel void gdn_scan_v(device float* state [[buffer(0)]], device const float4* qn [[buffer(1)]],
                       device const float4* kn [[buffer(2)]], device const float* vc [[buffer(3)]],
                       device const float2* gb [[buffer(4)]], device const float* kq [[buffer(5)]],
                       device float* o [[buffer(6)]], constant GdnPArgs& a [[buffer(7)]],
                       uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
                       ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
    constexpr int S = 128, IB = 8, TB = 16;
    threadgroup float4 qs[TB][32], ks[TB][32];
    threadgroup float2 red[2][IB][32];
    threadgroup float vs[TB][32];
    threadgroup float4 gk[TB];
    const int h = int(tg.x) / 4, hk = h % int(a.h_k), hv = int(a.h_v), hkn = int(a.h_k);
    const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * 16, w0 = sg * 4;
    const int T = int(a.n_tok);
    const float scale = 1.0f / sqrt(float(S));
    device float* sp = state + (i0 * hv + h) * S + j;
    const int rs = hv * S;
    float4 s0 = float4(sp[0], sp[rs], sp[2 * rs], sp[3 * rs]);
    float4 s1 = float4(sp[4 * rs], sp[5 * rs], sp[6 * rs], sp[7 * rs]);
    float4 s2 = float4(sp[8 * rs], sp[9 * rs], sp[10 * rs], sp[11 * rs]);
    float4 s3 = float4(sp[12 * rs], sp[13 * rs], sp[14 * rs], sp[15 * rs]);
    for (int t0 = 0; t0 < T; t0 += TB) {
        const int nb = min(TB, T - t0);
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (int idx = tid; idx < nb * 32; idx += IB * 32) {
            const int tt = idx >> 5, w = idx & 31;
            qs[tt][w] = qn[((t0 + tt) * hkn + hk) * 32 + w];
            ks[tt][w] = kn[((t0 + tt) * hkn + hk) * 32 + w];
            vs[tt][w] = vc[((t0 + tt) * hv + h) * S + (j - lane) + w];
            if (w == 0) {
                const float2 g_b = gb[(t0 + tt) * hv + h];
                gk[tt] = float4(g_b.x, g_b.y, kq[(t0 + tt) * hkn + hk], 0.0f);
            }
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
        for (int tt = 0; tt < nb; ++tt) {
            const int t = t0 + tt, slot = t & 1;
            const float4 k0 = ks[tt][w0], k1 = ks[tt][w0 + 1], k2 = ks[tt][w0 + 2], k3 = ks[tt][w0 + 3];
            const float4 q0 = qs[tt][w0], q1 = qs[tt][w0 + 1], q2 = qs[tt][w0 + 2], q3 = qs[tt][w0 + 3];
            const float ka = (dot(s0, k0) + dot(s1, k1)) + (dot(s2, k2) + dot(s3, k3));
            const float qa = (dot(s0, q0) + dot(s1, q1)) + (dot(s2, q2) + dot(s3, q3));
            red[slot][sg][lane] = float2(ka, qa);
            const float4 g_b = gk[tt];
            const float v = vs[tt][lane];
            threadgroup_barrier(mem_flags::mem_threadgroup);
            float2 r = red[slot][0][lane];
            for (int b = 1; b < IB; ++b) r += red[slot][b][lane];
            const float g = g_b.x;
            const float delta = (v - g * r.x) * g_b.y;
            if (sg == 0) o[(t * hv + h) * S + j] = fma(g, r.y, delta * g_b.z) * scale;
            s0 = fma(float4(g), s0, k0 * delta);
            s1 = fma(float4(g), s1, k1 * delta);
            s2 = fma(float4(g), s2, k2 * delta);
            s3 = fma(float4(g), s3, k3 * delta);
        }
    }
    sp[0] = s0.x; sp[rs] = s0.y; sp[2 * rs] = s0.z; sp[3 * rs] = s0.w;
    sp[4 * rs] = s1.x; sp[5 * rs] = s1.y; sp[6 * rs] = s1.z; sp[7 * rs] = s1.w;
    sp[8 * rs] = s2.x; sp[9 * rs] = s2.y; sp[10 * rs] = s2.z; sp[11 * rs] = s2.w;
    sp[12 * rs] = s3.x; sp[13 * rs] = s3.y; sp[14 * rs] = s3.z; sp[15 * rs] = s3.w;
}

// grid (h_v / 4, T) or (h_v, T), 128 threads: one simdgroup (32 lanes) per head h computes the RMSNorm and gated SiLU
// over the 128 channels (4 channels a lane) with zero threadgroup memory and zero barriers.
kernel void gdn_post(device const float* o [[buffer(0)]], device const float* z [[buffer(1)]],
                     device const float* gamma [[buffer(2)]], device half* yh [[buffer(3)]], device float* y [[buffer(4)]],
                     constant GdnPArgs& a [[buffer(5)]],
                     uint3 tg [[threadgroup_position_in_grid]], uint3 tgs [[threadgroups_per_grid]],
                     ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
    const int hv = int(a.h_v), t = int(tg.y);
    const int h = int(tgs.x) == hv ? (sg == 0 ? int(tg.x) : hv) : int(tg.x) * 4 + int(sg);
    if (h >= hv) return;
    const uint base = uint(t * hv + h) * 128 + lane;
    device const float* o_row = o + base;
    device const float* z_row = z + base;
    device const float* g_row = gamma + lane;
    device half* yh_row = yh + base;
    device float* y_row = y + base;
    const float ov0 = o_row[0], ov1 = o_row[32], ov2 = o_row[64], ov3 = o_row[96];
    // the gate and gamma do not depend on the norm: computed while the simd_sums run
    float sig[4];
#pragma clang loop unroll(full)
    for (uint k = 0; k < 4; ++k) sig[k] = 1.0f / (1.0f + exp(-z_row[k * 32]));
    const float gv[4] = {g_row[0], g_row[32], g_row[64], g_row[96]};
    const float4 pv = simd_sum(float4(ov0 * ov0, ov1 * ov1, ov2 * ov2, ov3 * ov3));
    const float tot = (pv.x + pv.y) + (pv.z + pv.w);
    const float rs = rsqrt(tot / 128.0f + a.eps);
    const float ov[4] = {ov0, ov1, ov2, ov3};
    float yv[4];
#pragma clang loop unroll(full)
    for (uint k = 0; k < 4; ++k) {
        yv[k] = ((rs * ov[k]) * gv[k]) * sig[k];
        yh_row[k * 32] = half(yv[k]);
    }
    if (a.f32) {
#pragma clang loop unroll(full)
        for (uint k = 0; k < 4; ++k) y_row[k * 32] = yv[k];
    }
}

template [[host_name("gdn_scan_1")]] kernel decltype(gdn_scan_t<1, 16>) gdn_scan_t<1, 16>;
template [[host_name("gdn_scan_2")]] kernel decltype(gdn_scan_t<2, 16>) gdn_scan_t<2, 16>;
template [[host_name("gdn_scan_4")]] kernel decltype(gdn_scan_t<4, 16>) gdn_scan_t<4, 16>;
template [[host_name("gdn_scanr_2")]] kernel decltype(gdn_scan_r<2, 0>) gdn_scan_r<2, 0>;
template [[host_name("gdn_scanr_4")]] kernel decltype(gdn_scan_r<4, 0>) gdn_scan_r<4, 0>;
template [[host_name("gdn_scanr_8")]] kernel decltype(gdn_scan_r<8, 0>) gdn_scan_r<8, 0>;
template [[host_name("gdn_scanr_4_8")]] kernel decltype(gdn_scan_r<4, 8>) gdn_scan_r<4, 8>;
template [[host_name("gdn_scanr_4_16")]] kernel decltype(gdn_scan_r<4, 16>) gdn_scan_r<4, 16>;

// ---- the chunked form (chunks of GC = 32 tokens), the WY / UT form of the gated delta rule.  With G_t the product
// of the chunk's gates up to t (log space: lg_t) and S0 the state entering the chunk:
//   (I + A) D = B,  A_ts = beta_t (k_t . k_s) G_t / G_s (s < t),  B = [beta v | beta G k]   ->   D = [U | W]
//   delta = U - W S0,   o_t = G_t S0^T q_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s,
//   S_L = G_L S0 + sum_s G_L / G_s k_s delta_s^T
// (the sequential recurrence's algebra; the rounding differs).  Three kernels:
//   gdn_cprep   per (head, chunk): lg, A, the forward substitution for U and W (a thread per column of B)
//   gdn_cstate  per (head, 32 state columns): the chunks in order - delta, the inter-chunk output, the state update
//   gdn_cout    per (head, chunk): o += the intra-chunk part, times 1/sqrt(128)
constant constexpr int GC = 32;

// grid (h_v, chunks), 256 threads.  U, W: [h][chunk][GC][128]; lgo: [h][chunk][GC]
kernel void gdn_cprep(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]],
                      device const float* vc [[buffer(2)]], device const float2* gb [[buffer(3)]],
                      device float* U [[buffer(4)]], device float* W [[buffer(5)]], device float* lgo [[buffer(6)]],
                      constant GdnPArgs& a [[buffer(7)]],
                      uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
    constexpr int S = 128;
    threadgroup float A[GC][GC];
    threadgroup float lg[GC], bt[GC];
    const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok);
    const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0);
    if (tid == 0) {
        float s = 0.0f;
        for (int t = 0; t < GC; ++t) {
            if (t < nb) s += max(log(gb[(t0 + t) * hv + h].x), -80.0f);
            lg[t] = s;
            bt[t] = t < nb ? gb[(t0 + t) * hv + h].y : 0.0f;
        }
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    for (int idx = tid; idx < GC * GC; idx += 256) {
        const int t = idx / GC, s = idx % GC;
        float v = 0.0f;
        if (s < t && t < nb) {
            device const float4* kt = (device const float4*)(kn + ((t0 + t) * hkn + hk) * S);
            device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S);
            float d = 0.0f;
            for (int i = 0; i < S / 4; ++i) d += dot(kt[i], ks[i]);
            v = bt[t] * d * exp(lg[t] - lg[s]);
        }
        A[t][s] = v;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    // forward substitution, column j of B a thread (threads 0..127: beta v; 128..255: beta G k)
    const int j = tid & 127;
    const bool isw = tid >= 128;
    float x[GC];
    for (int t = 0; t < GC; ++t) {
        float b = 0.0f;
        if (t < nb) b = isw ? bt[t] * exp(lg[t]) * kn[((t0 + t) * hkn + hk) * S + j] : bt[t] * vc[((t0 + t) * hv + h) * S + j];
        for (int s = 0; s < t; ++s) b = fma(-A[t][s], x[s], b);
        x[t] = b;
    }
    device float* o = (isw ? W : U) + (ulong(h) * NC + c) * GC * S + j;
    for (int t = 0; t < GC; ++t) o[t * S] = x[t];
    if (tid < GC) lgo[(h * NC + c) * GC + tid] = lg[tid];
}

// grid (h_v * 4), 256 threads = 8 row blocks of 16 x 32 columns; the state's 128 x 32 block in registers (16 a thread)
kernel void gdn_cstate(device float* state [[buffer(0)]], device const float* qn [[buffer(1)]],
                       device const float* kn [[buffer(2)]], device const float* U [[buffer(3)]],
                       device const float* W [[buffer(4)]], device const float* lgo [[buffer(5)]],
                       device float* delta [[buffer(6)]], device float* o [[buffer(7)]], constant GdnPArgs& a [[buffer(8)]],
                       uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
                       ushort lane [[thread_index_in_simdgroup]]) {
    constexpr int S = 128, RB = 16;
    threadgroup float red[8][16][32];
    threadgroup float dl[GC][32];
    const int h = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok);
    const int NC = (T + GC - 1) / GC;
    const int j = (int(tg.x) % 4) * 32 + lane, i0 = sg * RB;
    float st[RB];
    for (int r = 0; r < RB; ++r) st[r] = state[((i0 + r) * hv + h) * S + j];
    for (int c = 0; c < NC; ++c) {
        const int t0 = c * GC, nb = min(GC, T - t0);
        device const float* Uc = U + (ulong(h) * NC + c) * GC * S;
        device const float* Wc = W + (ulong(h) * NC + c) * GC * S;
        device const float* lg = lgo + (h * NC + c) * GC;
        // pass 0: delta = U - W S0;  pass 1: o_inter = G_t S0^T q_t;  16 tokens at a time
        for (int pass = 0; pass < 2; ++pass) {
            for (int hf = 0; hf < 2; ++hf) {
                float p[16];
                for (int u = 0; u < 16; ++u) {
                    const int t = hf * 16 + u;
                    float s = 0.0f;
                    if (t < nb) {
                        device const float* v = pass == 0 ? Wc + t * S + i0 : qn + ((t0 + t) * hkn + hk) * S + i0;
                        for (int r = 0; r < RB; ++r) s = fma(v[r], st[r], s);
                    }
                    p[u] = s;
                }
                threadgroup_barrier(mem_flags::mem_threadgroup);
                for (int u = 0; u < 16; ++u) red[sg][u][lane] = p[u];
                threadgroup_barrier(mem_flags::mem_threadgroup);
                for (int u = sg * 2; u < sg * 2 + 2; ++u) {
                    const int t = hf * 16 + u;
                    float s = 0.0f;
                    for (int b = 0; b < 8; ++b) s += red[b][u][lane];
                    if (t < nb) {
                        if (pass == 0) {
                            const float d = Uc[t * S + j] - s;
                            dl[t][lane] = d;
                            delta[((ulong(h) * NC + c) * GC + t) * S + j] = d;
                        } else {
                            o[((t0 + t) * hv + h) * S + j] = exp(lg[t]) * s;
                        }
                    } else if (pass == 0) {
                        dl[t][lane] = 0.0f;
                    }
                }
            }
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
        // S = G_L S0 + sum_s G_L / G_s k_s delta_s^T
        const float lgl = lg[nb - 1], gl = exp(lgl);
        for (int r = 0; r < RB; ++r) st[r] *= gl;
        for (int s = 0; s < nb; ++s) {
            const float f = exp(lgl - lg[s]) * dl[s][lane];
            device const float* ks = kn + ((t0 + s) * hkn + hk) * S + i0;
            for (int r = 0; r < RB; ++r) st[r] = fma(ks[r], f, st[r]);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }
    for (int r = 0; r < RB; ++r) state[((i0 + r) * hv + h) * S + j] = st[r];
}

// grid (h_v, chunks), 128 threads (column j): o_t = (o_inter_t + sum_{s <= t} (q_t . k_s) G_t / G_s delta_s) / sqrt(128)
kernel void gdn_cout(device const float* qn [[buffer(0)]], device const float* kn [[buffer(1)]],
                     device const float* delta [[buffer(2)]], device const float* lgo [[buffer(3)]],
                     device float* o [[buffer(4)]], constant GdnPArgs& a [[buffer(5)]],
                     uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
    constexpr int S = 128;
    threadgroup float M[GC][GC];
    const int h = int(tg.x), c = int(tg.y), hv = int(a.h_v), hkn = int(a.h_k), hk = h % hkn, T = int(a.n_tok);
    const int NC = (T + GC - 1) / GC, t0 = c * GC, nb = min(GC, T - t0);
    device const float* lg = lgo + (h * NC + c) * GC;
    for (int idx = tid; idx < GC * GC; idx += 128) {
        const int t = idx / GC, s = idx % GC;
        float v = 0.0f;
        if (s <= t && t < nb) {
            device const float4* qt = (device const float4*)(qn + ((t0 + t) * hkn + hk) * S);
            device const float4* ks = (device const float4*)(kn + ((t0 + s) * hkn + hk) * S);
            float d = 0.0f;
            for (int i = 0; i < S / 4; ++i) d += dot(qt[i], ks[i]);
            v = d * exp(lg[t] - lg[s]);
        }
        M[t][s] = v;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    const int j = tid;
    device const float* dc = delta + (ulong(h) * NC + c) * GC * S + j;
    float dv[GC];
    for (int s = 0; s < GC; ++s) dv[s] = s < nb ? dc[s * S] : 0.0f;
    const float scale = 1.0f / sqrt(float(S));
    for (int t = 0; t < nb; ++t) {
        float s = o[((t0 + t) * hv + h) * S + j];
        for (int u = 0; u <= t; ++u) s = fma(M[t][u], dv[u], s);
        o[((t0 + t) * hv + h) * S + j] = s * scale;
    }
}