// metal/kernels/gdn.metal - one GDN (gated delta-net) layer's recurrence for a window of tokens, fused end to end:
//
// conv (4 taps, history [channels][3] oldest first) -> SiLU, on q (16 heads), k (16), v (48 heads) of 128
// q, k: L2 norm as native_gdn_l2_norm (rsqrt(sum/128 + eps/128), then x 1/sqrt(128), two roundings)
// gate = exp(softplus(alpha + dt) * a) (softplus = v > 20 ? v : log1p(exp(v))), beta = sigmoid(beta)
// kv_j = sum_i S[i][j] k[i]; delta_j = (v_j - gate * kv_j) * beta; S = gate * S + k[i] delta_j
// o_j = (sum_i S[i][j] q[i]) / sqrt(128); y_j = o_j * rsqrt(mean(o^2) + eps) * gamma_j * sigmoid(z_j)
//
// State layout as Strata's: S[(i * h_v + h) * 128 + j]. Value head h reads q/k head h % h_k.
// One threadgroup per value head, 512 threads = 16 simdgroups: simdgroup s holds columns j = 32 (s % 4) + lane and
// rows i = 32 (s / 4) .. + 32 of the head's 128 x 128 state in registers. GDN_NTOK tokens run in order; with
// `commit` the state and the conv history are written back after the last one (a verify window runs without commit,
// the fold of m accepted tokens is the same kernel with GDN_NTOK = m and commit). The conv history is read from
// `hist` and the committed one written to `hist_out` (ping-pong: the q/k channels are shared by three threadgroups).
constant int GDN_NTOK [[function_constant(2)]];
struct GdnArgs {
uint h_k; // 16
uint h_v; // 48
uint channels; // 10240 = (2 h_k + h_v) * 128
uint commit; // 0: verify (no write); 1: commit after the last token; 0x100 | k (1 <= k <= GDN_NTOK): commit
// the state and the conv history after the first k tokens, in place, while every token's y is
// still computed (a window whose first k rows are kept) - bitwise the commit of a k-token step
float eps; // 1e-6
uint plane; // 0: y only; else y is the raw plane of prepared activations (xprep layout, n = h_v * 128,
// `plane` floats between planes) and the other planes and the per-16 sums are written too
};
constant float gdn_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f,
0.00006103515625f};
constant float gdn_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f};
// The kernel is launched with IB * 128 threads (512 for gdn_step, 1024 for gdn_step8). Without the attribute the compiler
// is free to use as many registers as the 128-row state needs and a pipeline can end up with
// maxTotalThreadsPerThreadgroup below the launch size: an M2 Max gave 448 for the GDN_NTOK >= 2 specializations, and
// Metal then skips the dispatch without an error (#2323). The attribute makes the compiler fit the launch size.
template <int IB>
[[max_total_threads_per_threadgroup(IB * 128)]]
kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[buffer(1)]],
device const float* qkv [[buffer(2)]], device const float* alpha [[buffer(3)]],
device const float* beta [[buffer(4)]], device const float* z [[buffer(5)]],
device const float* conv_w [[buffer(6)]], device const float* dt [[buffer(7)]],
device const float* ssm_a [[buffer(8)]], device const float* gamma [[buffer(9)]],
device float* y [[buffer(10)]], constant GdnArgs& a [[buffer(11)]],
device float* hist_out [[buffer(12)]],
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;
threadgroup float qs[S], ks[S], vs[S];
constexpr int RPT = S / IB; // state rows per thread
// kv and o partials in their own buffers, the q/k and the output norm partials in theirs: no barrier between a
// token's phases for a buffer to be free again
threadgroup float red[IB][S], red2[IB - 1][S];
threadgroup float part[8], part2[4];
threadgroup float gb[2];
const int h = int(tg.x), hk = h % int(a.h_k), hv = int(a.h_v);
const int cb = sg & 3, ib = sg >> 2;
const int j = 32 * cb + lane, i0 = RPT * ib;
const float scale = 1.0f / sqrt(float(S));
// the channel this thread convolves: 0..127 q of head hk, 128..255 k of head hk, 256..383 v of head h
int ch = -1;
if (tid < 128) ch = hk * S + tid;
else if (tid < 256) ch = int(a.h_k) * S + hk * S + (tid - 128);
else if (tid < 384) ch = 2 * int(a.h_k) * S + h * S + (tid - 256);
float hst[3] = {0.0f, 0.0f, 0.0f}, cw[4] = {0.0f, 0.0f, 0.0f, 0.0f};
// the first token's inputs are issued with the state's loads (their latency is not paid again after it)
float in = ch >= 0 ? qkv[ch] : 0.0f;
if (ch >= 0) {
#pragma clang loop unroll(full)
for (int k = 0; k < 3; ++k) hst[k] = hist[ch * 3 + k];
const float4 cwv = *(device const float4*)(conv_w + ch * 4);
cw[0] = cwv.x; cw[1] = cwv.y; cw[2] = cwv.z; cw[3] = cwv.w;
}
const float gam = ib == 0 ? gamma[j] : 0.0f;
const float dth = tid == 384 ? dt[h] : 0.0f;
const float ssm_ah = tid == 384 ? ssm_a[h] : 0.0f;
float st[RPT];
const int st_stride = hv * S;
const int st_off = (i0 * hv + h) * S + j;
#pragma clang loop unroll(full)
for (int r = 0; r < RPT; ++r) st[r] = (state + st_off)[r * st_stride];
const int commit_at = a.commit == 1u ? GDN_NTOK : (a.commit >> 8) ? int(a.commit & 0xFFu) : 0;
for (int t = 0; t < GDN_NTOK; ++t) {
// this token's z and the next token's conv input, ahead of the barriers that need them
const float zz = ib == 0 ? z[t * hv * S + h * S + j] : 0.0f;
const float in_next = ch >= 0 && t + 1 < GDN_NTOK ? qkv[(t + 1) * int(a.channels) + ch] : 0.0f;
// conv + SiLU
float cv = 0.0f;
if (ch >= 0) {
const float vals[4] = {hst[0], hst[1], hst[2], in};
float s = 0.0f;
#pragma clang loop unroll(full)
for (int k = 0; k < 4; ++k) s = fma(vals[k], cw[k], s);
s = s + 0.0f;
cv = s / (1.0f + exp(-s));
hst[0] = hst[1];
hst[1] = hst[2];
hst[2] = in;
if (t + 1 == commit_at && (tid >= 256 || h < int(a.h_k))) {
#pragma clang loop unroll(full)
for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k];
}
}
in = in_next;
// raw q/k/v to shared memory + L2 norm partials of q (sg 0..3) and k (sg 4..7); gate and beta on idle lanes
if (tid < 256) {
const float ssq = simd_sum(cv * cv);
if (lane == 0) part[sg] = ssq;
if (tid < 128) qs[tid] = cv;
else ks[tid - 128] = cv;
} else if (tid < 384) {
vs[tid - 256] = cv;
} else if (tid == 384) {
const float v = alpha[t * hv + h] + dth;
// log1p(e^v) (Kahan: log(u) * e / (u - 1), u = 1 + e; exact where 1 + e rounds to 1)
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[0] = exp(sp * ssm_ah);
} else if (tid == 416) {
gb[1] = 1.0f / (1.0f + exp(-beta[t * hv + h]));
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const float tot_k = (part[4] + part[5]) + (part[6] + part[7]);
const float rsk = rsqrt(tot_k / float(S) + a.eps / float(S));
// the simdgroup's RPT rows of k and q one a lane, normalized on the fly and broadcast from the lane that holds them
const int qk_idx = i0 + (lane & (RPT - 1));
const float kl = (rsk * ks[qk_idx]) * scale;
// kv_j
float kv = 0.0f;
#pragma clang loop unroll(full)
for (int r = 0; r < RPT; ++r) kv = fma(st[r], simd_broadcast(kl, ushort(r)), kv);
red[ib][j] = kv;
const float tot_q = (part[0] + part[1]) + (part[2] + part[3]);
const float rsq = rsqrt(tot_q / float(S) + a.eps / float(S));
const float ql = (rsq * qs[qk_idx]) * scale;
const float g = gb[0], be = gb[1], vj = vs[j];
threadgroup_barrier(mem_flags::mem_threadgroup);
kv = 0.0f;
#pragma clang loop unroll(full)
for (int k = 0; k < IB; ++k) kv += red[k][j];
const float delta = (vj - g * kv) * be;
float o = 0.0f;
#pragma clang loop unroll(full)
for (int r = 0; r < RPT; ++r) {
st[r] = fma(g, st[r], simd_broadcast(kl, ushort(r)) * delta);
o = fma(st[r], simd_broadcast(ql, ushort(r)), o);
}
if (ib > 0) red2[ib - 1][j] = o;
if (t + 1 == commit_at) {
// the state and the conv history after token t (registers: st after this token's update, hst after its
// shift); the state is read only before the loop, so in place is safe. q/k history: the threadgroups
// h < h_k own their q/k head (h == hk); v history: every head its own - to hist_out, NOT in place: a q/k
// channel is read by the three threadgroups h % h_k at their start, and a threadgroup may start after its
// owner has finished (seen on the 10-core M4); the caller swaps the two
device float* st_base = state + st_off;
#pragma clang loop unroll(full)
for (int r = 0; r < RPT; ++r) st_base[r * st_stride] = st[r];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// the readout and the output norm: column j's o (only ib == 0 computes readout and writes y)
if (ib == 0) {
#pragma clang loop unroll(full)
for (int k = 0; k < IB - 1; ++k) o += red2[k][j];
o = o * scale;
const float osq = simd_sum(o * o);
if (lane == 0) part2[cb] = osq;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (ib == 0) {
const float tot = (part2[0] + part2[1]) + (part2[2] + part2[3]);
const float rs = rsqrt(tot / float(S) + a.eps);
const float yv = ((rs * o) * gam) * (1.0f / (1.0f + exp(-zz)));
y[t * hv * S + h * S + j] = yv;
if (a.plane) {
device float* raw = y + t * hv * S + h * S + j;
const int jj = j & 15;
raw[a.plane] = yv * gdn_q2s[jj & 7];
raw[2 * a.plane] = yv * gdn_p4s[jj & 3];
raw[3 * a.plane] = (jj & 1) ? yv * 0.00390625f : yv;
// the per-16 sums (even and odd columns apart, in column order) from the simdgroup's lanes
float s0 = 0.0f, s1 = 0.0f;
#pragma clang loop unroll(full)
for (int k = 0; k < 16; k += 2) {
s0 += simd_shuffle_down(yv, ushort(k));
s1 += simd_shuffle_down(yv, ushort(k + 1));
}
if (jj == 0) y[4 * a.plane + t * (hv * S / 16) + (h * S + j) / 16] = s0 + s1;
}
}
}
}
template [[host_name("gdn_step")]] kernel decltype(gdn_step_t<4>) gdn_step_t<4>;
template [[host_name("gdn_step8")]] kernel decltype(gdn_step_t<8>) gdn_step_t<8>;