// metal/kernels/gr.metal - the hyper-connection (gated residual) read for a window of tokens, with the previous
// sublayer's write folded in. The math is fused_gr.hpp's (native contract: FP32 activations, BF16 weights):
//
// gr_norm R'[c] = R[c] + bo * 2 sigmoid(inj_prev[c] / hc) (when apply; written back in place)
// xn[c] = R'[c] * w_norm[c] * rsqrt(mean(R'[c]^2) + eps)
// gr_down lo[k] = silu((w_down[k] . xn) / hc) k < hc_lr; inject[c] = w_inject[c] . xn
// gr_up mixed[d] = mean_c xn[c][d] * sigmoid(w_up[c n + d] . lo), written as prepared activations (xprep's
// planes and per-16 sums) for the GEMVs that read it
//
// Shapes: n_embd 2560, hc 4, hc_lr 320. Per token t: R[t] (hc * n), bo[t] (n), inj_prev[t] (hc), xn[t] (hc * n),
// lo[t] (hc_lr), inject[t] (hc), XP planes of n_tok rows.
constant int GR_NTOK [[function_constant(1)]];
// the prepared activations' per-position plane scales (units.metal's xprep): q2 planes 4^-(j%8), p4 planes 256^-(j%4)
constant float gr_q2s[8] = {1.0f, 0.25f, 0.0625f, 0.015625f, 0.00390625f, 0.0009765625f, 0.000244140625f,
0.00006103515625f};
constant float gr_p4s[4] = {1.0f, 0.00390625f, 1.52587890625e-05f, 5.9604644775390625e-08f};
struct GrArgs {
uint n_embd;
uint hc;
uint hc_lr;
uint n_tok;
uint apply;
float eps;
};
// one threadgroup per (stream c, token t), 256 threads
kernel void gr_norm(device float* R [[buffer(0)]], device const float* bo [[buffer(1)]],
device const float* inj [[buffer(2)]], device const float* w_norm [[buffer(3)]],
device float* xn [[buffer(4)]], constant GrArgs& a [[buffer(5)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
threadgroup float part[8];
const uint c = tg.x, t = tg.y, n = a.n_embd;
device float* r = R + (t * a.hc + c) * n;
const float g = a.apply ? 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc))) : 0.0f;
float rv[10];
float ss = 0.0f;
if (n == 2560) {
device float* rp = r + tid;
if (a.apply) {
device const float* bp = bo + t * 2560u + tid;
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) {
const float v = fma(bp[i * 256], g, rp[i * 256]);
rp[i * 256] = v;
rv[i] = v;
ss = fma(v, v, ss);
}
} else {
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) {
rv[i] = rp[i * 256];
ss = fma(rv[i], rv[i], ss);
}
}
} else {
for (uint d = tid; d < n; d += 256) {
float v = r[d];
if (a.apply) {
v = fma(bo[t * n + d], g, v);
r[d] = v;
}
if (n <= 2560) rv[d / 256] = v;
ss = fma(v, v, ss);
}
}
ss = simd_sum(ss);
if (lane == 0) part[sg] = ss;
threadgroup_barrier(mem_flags::mem_threadgroup);
float tot = 0.0f;
#pragma clang loop unroll(full)
for (int i = 0; i < 8; ++i) tot += part[i];
const float rs = rsqrt(tot / float(n) + a.eps);
device float* o = xn + (t * a.hc + c) * n;
device const float* w = w_norm + c * n;
if (n == 2560) {
o += tid;
w += tid;
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) o[i * 256] = rv[i] * w[i * 256] * rs;
} else {
for (uint d = tid; d < n; d += 256)
o[d] = (n <= 2560 ? rv[d / 256] : r[d]) * w[d] * rs;
}
}
// rows: 0 .. hc_lr - 1 = w_down, hc_lr .. hc_lr + hc - 1 = w_inject (both bf16 rows of hc * n_embd).
// SG simdgroups split the row's units KS ways; LPR = 32 (a simdgroup per row); RS = SG / KS simdgroups per
// threadgroup, NR rows each.
// one token's dot with a decoded unit (gr_up): x is the token's slice of lo
template <int UNIT>
static inline float gr_dot(device const float* x, thread const float* q, float4 sc) {
float xs[UNIT], xsum[2] = {0.0f, 0.0f};
#pragma clang loop unroll(full)
for (int k = 0; k < UNIT; k += 16) load16(x + k, xs + k);
return unit_combine<UNIT>(q, sc, xs, xsum);
}
// a unit's decoded rows for grd_acc, read four codes at a time
template <int F, int NR>
struct GrdUnit {
float q[NR][fmt<F>::UNIT];
float4 sc[NR];
void load(device const uchar* w, uint rb, int u, threadgroup const uint* notab) thread {
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) Dec<F>::run(w + ulong(r) * rb, u, q[r], sc[r], notab);
}
float4 v(int r, int i) const thread { return float4(q[r][4 * i], q[r][4 * i + 1], q[r][4 * i + 2], q[r][4 * i + 3]); }
};
// Q8_0 two rows a simdgroup: the codes stay packed (8 registers a row, not 32 floats) and are widened per token.
// Measured 10-16% faster past 6 tokens (where the token loop is rolled), no faster below, and slower for one row.
template <>
struct GrdUnit<T_Q8_0, 2> {
char4 q[2][8];
float4 sc[2];
void load(device const uchar* w, uint rb, int u, threadgroup const uint*) thread {
#pragma clang loop unroll(full)
for (int r = 0; r < 2; ++r) {
device const uchar* b = w + ulong(r) * rb + 34 * u;
device const ushort* p = (device const ushort*)(b + 2);
#pragma clang loop unroll(full)
for (int k = 0; k < 8; ++k) q[r][k] = as_type<char4>(uint(p[2 * k]) | (uint(p[2 * k + 1]) << 16));
const float d = rdh(b);
sc[r] = float4(d, d, 0.0f, 0.0f);
}
}
float4 v(int r, int i) const thread { return float4(q[r][i]); }
};
// NR consecutive rows' K-split partials over xn (units of format F; the rows rb bytes apart), acc[r * 8 + t].
// unit_combine's sums in its order, with each xn float4 read straight from memory once for all NR rows: xn is
// 15-30x the weights' bytes at 4-8 tokens, so NR = 2 halves the kernel's traffic. GR_NTOK is a function constant,
// which the compiler does not unroll over by itself: the pragma keeps acc in registers instead of the stack. Past 6
// tokens the unrolled loop measured slower than the rolled one, so it stays rolled there; ROLL keeps it rolled at
// every count (the packed Q8_0 pair at KS = 4: 1-3% faster at 4-6 tokens, where KS = 8 measured 2-6% slower).
template <int F, int NR, bool ROLL = false>
static inline void grd_acc(device const uchar* w, uint rb, int nin, int ks, int KS, ushort lane, device const float* xn,
thread float* acc, threadgroup const uint* notab) {
constexpr int UNIT = fmt<F>::UNIT, H = UNIT / 2;
const int U = nin / UNIT, u0 = ks * U / KS, u1 = (ks + 1) * U / KS;
for (int u = u0 + lane; u < u1; u += 32) {
GrdUnit<F, NR> d;
d.load(w, rb, u, notab);
auto tok = [&](int t) {
device const float4* x = (device const float4*)(xn + t * nin + UNIT * u);
float s[NR][4];
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) s[r][0] = s[r][1] = s[r][2] = s[r][3] = 0.0f;
#pragma clang loop unroll(full)
for (int j = 0; j < H; j += 4) {
const float4 a = x[j / 4], b = x[(H + j) / 4];
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) {
const float4 qa = d.v(r, j / 4), qb = d.v(r, (H + j) / 4);
s[r][0] = fma(qa.x, a.x, s[r][0]);
s[r][1] = fma(qa.y, a.y, s[r][1]);
s[r][2] = fma(qb.x, b.x, s[r][2]);
s[r][3] = fma(qb.y, b.y, s[r][3]);
s[r][0] = fma(qa.z, a.z, s[r][0]);
s[r][1] = fma(qa.w, a.w, s[r][1]);
s[r][2] = fma(qb.z, b.z, s[r][2]);
s[r][3] = fma(qb.w, b.w, s[r][3]);
}
}
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r)
acc[r * 8 + t] += d.sc[r].x * (s[r][0] + s[r][1]) + d.sc[r].y * (s[r][2] + s[r][3]) -
(d.sc[r].z * 0.0f + d.sc[r].w * 0.0f);
};
if (GR_NTOK <= 6 && !ROLL) {
#pragma clang loop unroll(full)
for (int t = 0; t < GR_NTOK; ++t) tok(t);
} else {
for (int t = 0; t < GR_NTOK; ++t) tok(t);
}
}
}
// FI: the inject rows' format (the Unsloth files: Q8_0 down, F32 inject); FI = F is the original kernel. NR must
// divide hc_lr and hc, so that a simdgroup's rows are all w_down or all w_inject. The sums do not depend on SG or NR:
// every SG / NR with the same KS is bitwise the same kernel.
template <int SG, int KS, int NR, int F, int FI>
kernel void gr_down_f(device const uchar* w_down [[buffer(0)]], device const uchar* w_inject [[buffer(1)]],
device const float* xn [[buffer(2)]], device float* lo [[buffer(3)]],
device float* inject [[buffer(4)]], constant GrArgs& a [[buffer(5)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]]) {
constexpr int RS = SG / KS;
threadgroup float red[KS * RS * NR * 8];
threadgroup uint notab[1];
const int ks = sg / RS, rs = sg % RS;
const int row0 = (int(tg.x) * RS + rs) * NR;
const int nrows = int(a.hc_lr + a.hc);
const int rr = min(row0, nrows - NR);
const int nin = int(a.hc * a.n_embd);
const uint rb = uint(nin / fmt<F>::QK * fmt<F>::BB), rbi = uint(nin / fmt<FI>::QK * fmt<FI>::BB);
float acc[NR * 8];
#pragma clang loop unroll(full)
for (int t = 0; t < NR * 8; ++t) acc[t] = 0.0f;
constexpr bool roll = F == T_Q8_0 && NR == 2 && KS == 4;
if (rr < int(a.hc_lr)) grd_acc<F, NR, roll>(w_down + ulong(rr) * rb, rb, nin, ks, KS, lane, xn, acc, notab);
else grd_acc<FI, NR>(w_inject + ulong(rr - int(a.hc_lr)) * rbi, rbi, nin, ks, KS, lane, xn, acc, notab);
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) {
#pragma clang loop unroll(full)
for (int t = 0; t < GR_NTOK; ++t) {
const float v = simd_sum(acc[r * 8 + t]);
if (lane == 0) red[((ks * RS + rs) * NR + r) * 8 + t] = v;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (ks == 0 && lane < NR * 8) {
const int r = lane / 8, t = lane % 8, row = row0 + r;
if (t < GR_NTOK && row < nrows) {
float v = 0.0f;
#pragma clang loop unroll(full)
for (int k = 0; k < KS; ++k) v += red[((k * RS + rs) * NR + r) * 8 + t];
if (row < int(a.hc_lr)) {
const float z = v / float(a.hc);
lo[t * int(a.hc_lr) + row] = z / (1.0f + exp(-z));
} else {
inject[t * int(a.hc) + row - int(a.hc_lr)] = v;
}
}
}
}
// A threadgroup owns 16 consecutive d for all hc = 4 streams: 8 simdgroups, simdgroup s -> stream s / 2 and
// d0 + 8 (s % 2) .. + 8; LPR = 4 lanes per row (320 = 20 units of 16, 5 per lane).
// XP: [4 kinds][n_tok][n] then S16 [n_tok][n / 16] (units.metal's xprep layout).
template <int F>
kernel void gr_up_f(device const uchar* w_up [[buffer(0)]], device const float* lo [[buffer(1)]],
device const float* xn [[buffer(2)]], device float* xp [[buffer(3)]],
constant GrArgs& a [[buffer(4)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) {
threadgroup float gate[8][4][16];
threadgroup uint notab[1];
const int n = int(a.n_embd), c = sg >> 1, dl = 8 * (sg & 1) + lane / 4, li = lane & 3;
const int d = int(tg.x) * 16 + dl;
constexpr int UNIT = fmt<F>::UNIT;
device const uchar* w = w_up + ulong(c * n + d) * (a.hc_lr / fmt<F>::QK * fmt<F>::BB);
const int U = int(a.hc_lr) / UNIT;
float acc[8];
#pragma clang loop unroll(full)
for (int t = 0; t < 8; ++t) acc[t] = 0.0f;
for (int u = li; u < U; u += 4) {
float q[UNIT];
float4 sc;
Dec<F>::run(w, u, q, sc, notab);
#pragma clang loop unroll(full)
for (int t = 0; t < GR_NTOK; ++t) acc[t] += gr_dot<UNIT>(lo + t * int(a.hc_lr) + UNIT * u, q, sc);
}
#pragma clang loop unroll(full)
for (int t = 0; t < GR_NTOK; ++t) {
const float v = lane_sum<4>(acc[t]);
if (li == 0) gate[t][c][dl] = 1.0f / (1.0f + exp(-v));
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// 16 threads per token: thread j -> d0 + j
const int t = tid / 16, j = tid & 15;
if (t < GR_NTOK) {
const int dd = int(tg.x) * 16 + j;
const int hc = int(a.hc);
float m = 0.0f;
if (hc == 4) {
#pragma clang loop unroll(full)
for (int cc = 0; cc < 4; ++cc) m = fma(xn[(t * 4 + cc) * n + dd], gate[t][cc][j], m);
} else {
for (int cc = 0; cc < hc; ++cc) m = fma(xn[(t * hc + cc) * n + dd], gate[t][cc][j], m);
}
m = m / float(hc);
const uint plane = a.n_tok * uint(n);
device float* raw = xp + t * n + dd;
raw[0] = m;
raw[plane] = m * gr_q2s[j & 7];
raw[2 * plane] = m * gr_p4s[j & 3];
raw[3 * plane] = (j & 1) ? m * 0.00390625f : m;
float s0 = 0.0f, s1 = 0.0f;
#pragma clang loop unroll(full)
for (int k = 0; k < 16; k += 2) {
s0 += simd_shuffle_down(m, ushort(k));
s1 += simd_shuffle_down(m, ushort(k + 1));
}
if (j == 0) xp[4 * a.n_tok * uint(n) + t * (n / 16) + tg.x] = s0 + s1;
}
}
#define GRD1(S, K, R, SFX) \
template [[host_name("gr_down_" #S "_" #K SFX)]] kernel decltype(gr_down_f<S, K, R, T_BF16, T_BF16>) gr_down_f<S, K, R, T_BF16, T_BF16>; \
template [[host_name("gr_down_q8_" #S "_" #K SFX)]] kernel decltype(gr_down_f<S, K, R, T_Q8_0, T_Q8_0>) gr_down_f<S, K, R, T_Q8_0, T_Q8_0>; \
template [[host_name("gr_down_q8f_" #S "_" #K SFX)]] kernel decltype(gr_down_f<S, K, R, T_Q8_0, T_F32>) gr_down_f<S, K, R, T_Q8_0, T_F32>;
#define GRD(S, K) GRD1(S, K, 1, "")
GRD(8, 8) GRD(8, 4) GRD(4, 4) GRD(4, 2) GRD(16, 16) GRD(16, 8)
// two rows a simdgroup (gr_read past 4 tokens)
GRD1(8, 8, 2, "_r2") GRD1(4, 4, 2, "_r2")
template [[host_name("gr_up")]] kernel decltype(gr_up_f<T_BF16>) gr_up_f<T_BF16>;
template [[host_name("gr_up_q8")]] kernel decltype(gr_up_f<T_Q8_0>) gr_up_f<T_Q8_0>;
template [[host_name("gr_down_q4_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q4_0, T_Q8_0>;
template [[host_name("gr_up_q4")]] kernel decltype(gr_up_f<T_Q4_0>) gr_up_f<T_Q4_0>;
template [[host_name("gr_down_q5_8_8")]] kernel decltype(gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>) gr_down_f<8, 8, 1, T_Q5_0, T_Q8_0>;
template [[host_name("gr_up_q5")]] kernel decltype(gr_up_f<T_Q5_0>) gr_up_f<T_Q5_0>;