// metal/kernels/prompt.metal - the elementwise steps of a prompt chunk around the GEMMs (after quant/units.metal):
//
// route_topk per token: softmax over the 512 logits, top 10 (ties: the lower expert), weights renormalised
// over the selected with a 2^-14 floor - moe_route's rule; one simdgroup a token
// gr_epi_down after the HC down GEMM (V = xn . [w_down; w_inject], T x (hc_lr + hc)): lo_h = half(silu(v / hc)),
// inject = v
// gr_epi_up after the HC up GEMM (G = lo . w_up, T x hc*n): mixed = mean_c xn[c] * sigmoid(G[c]) -> float and
// half (the next GEMMs' operand)
// swiglu_h h = half(silu(g) * u)
// gemm_hc_down, gemm_swiglu_up (after gemm_hc_up) the HC down and shared-expert up GEMMs with gr_epi_down /
// swiglu_h as register epilogues
// moe_out bo = y_routed + sigmoid(sg) * y_shared (prompt_chain.inc; the prompt chunk folds it into the next
// read: gr_norm_h_moe, gr_apply_moe, ple_moe_gate_norm; with mm_combine folded in too: gr_norm_h_moec,
// gr_apply_moec)
kernel void route_topk(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
device float* w [[buffer(2)]], constant uint& n_tok_packed [[buffer(3)]],
device const uint* smap [[buffer(4)]], constant float& prompt_moce [[buffer(5)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]]) {
const uint n_tok = n_tok_packed & 0xFFFFu;
const int n_exp = (n_tok_packed >> 16) ? int(n_tok_packed >> 16) : 512;
const uint t = tg.x * 8 + sg;
if (t >= n_tok) return;
device const float* l = logits + ulong(t) * ulong(n_exp);
float v[16];
float mx = -INFINITY;
float sum = 0.0f;
if (n_exp == 512) {
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) {
v[i] = l[lane + 32 * i];
mx = max(mx, v[i]);
}
for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) {
v[i] = exp(v[i] - mx);
sum += v[i];
}
for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
const float rcp = 1.0f / sum;
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) {
v[i] *= rcp;
if (isnan(v[i])) v[i] = -FLT_MAX;
}
} else {
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
v[i] = idx < n_exp ? l[idx] : -INFINITY;
}
for (int i = 0; i < 16; ++i) mx = max(mx, v[i]);
for (ushort m = 16; m > 0; m >>= 1) mx = max(mx, simd_shuffle_xor(mx, m));
for (int i = 0; i < 16; ++i) {
v[i] = (lane + 32 * i < n_exp) ? exp(v[i] - mx) : 0.0f;
sum += v[i];
}
for (ushort m = 16; m > 0; m >>= 1) sum += simd_shuffle_xor(sum, m);
const float rcp = 1.0f / sum;
for (int i = 0; i < 16; ++i) {
if (lane + 32 * i >= n_exp) {
v[i] = -INFINITY;
} else {
v[i] *= rcp;
if (isnan(v[i])) v[i] = -FLT_MAX;
}
}
}
const int tail_exact = int(prompt_moce) / 10000;
const float rem_moce = prompt_moce - float(tail_exact * 10000);
const int body_prot = int(rem_moce) / 100;
const float body_bonus = rem_moce - float(body_prot * 100);
const bool in_tail = tail_exact > 0 && (t + uint(tail_exact) >= n_tok);
const int p_prot = in_tail ? 10 : body_prot;
const float p_bonus = in_tail ? 0.0f : body_bonus;
const float moce_scale = p_bonus > 0.0f ? exp(min(p_bonus, 40.0f)) : 1.0f;
const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f;
ushort scaled_mask = 0;
float selected = 0.0f, ssum = 0.0f;
for (int rank = 0; rank < 10; ++rank) {
if (rank == p_prot && moce_scale > 1.0f) {
if (n_exp == 512) {
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
if (v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
v[i] *= moce_scale;
scaled_mask |= ushort(1u << i);
}
}
} else {
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
if (idx < n_exp && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
v[i] *= moce_scale;
scaled_mask |= ushort(1u << i);
}
}
}
}
float best = v[0];
int ex = lane;
for (int i = 1; i < 16; ++i)
if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
for (ushort m = 16; m > 0; m >>= 1) {
const float ob = simd_shuffle_xor(best, m);
const int oe = simd_shuffle_xor(ex, m);
if (ob > best || (ob == best && oe < ex)) { best = ob; ex = oe; }
}
float true_prob = best;
if (rank >= p_prot && moce_scale > 1.0f && best > 0.0f) {
const ushort win_lane = ushort(ex & 31);
const ushort win_i = ushort(ex >> 5);
const ushort win_mask = simd_broadcast(scaled_mask, win_lane);
if ((win_mask & ushort(1u << win_i)) != 0u) true_prob = best * rcp_moce;
}
if ((ex & 31) == lane) {
// v stays in registers: no runtime index into it
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) v[i] = i == (ex >> 5) ? -INFINITY : v[i];
ids[t * 10 + rank] = ex;
ssum += true_prob;
}
if (rank == lane) selected = true_prob;
}
for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m);
if (prompt_moce < 0.0f) {
const int skip_tail = int(-prompt_moce);
const float tau = (-prompt_moce) - float(skip_tail);
const bool skip_active = tau > 0.0f && !(skip_tail > 0 && (t + uint(skip_tail) >= n_tok));
if (skip_active) {
float kept_prob = (lane < 10) ? selected : 0.0f;
if (lane >= 3 && lane < 10) {
const int ex_r = ids[t * 10 + lane];
if (smap[ex_r] == 0xFFFFFFFFu && selected < tau * ssum) {
ids[t * 10 + lane] = -1;
kept_prob = 0.0f;
}
}
selected = kept_prob;
ssum = kept_prob;
for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m);
}
}
ssum = max(ssum, 6.103515625e-5f);
if (lane < 10) w[t * 10 + lane] = selected * (1.0f / ssum);
}
struct EpiArgs {
uint n_tok;
uint n_embd; // 2560
uint hc; // 4
uint hc_lr; // 320
};
// grid (T * (hc_lr + hc) / 256)
kernel void gr_epi_down(device const float* V [[buffer(0)]], device half* lo_h [[buffer(1)]],
device float* inject [[buffer(2)]], constant EpiArgs& a [[buffer(3)]],
uint gid [[thread_position_in_grid]]) {
const uint nr = a.hc_lr + a.hc, t = gid / nr, r = gid % nr;
if (t >= a.n_tok) return;
const float v = V[gid];
if (r < a.hc_lr) {
const float z = v / float(a.hc);
lo_h[t * a.hc_lr + r] = half(z / (1.0f + exp(-z)));
} else {
inject[t * a.hc + r - a.hc_lr] = v;
}
}
// grid (T * n_embd / 256)
kernel void gr_epi_up(device const float* G [[buffer(0)]], device const float* xn [[buffer(1)]],
device float* mixed [[buffer(2)]], device half* mixed_h [[buffer(3)]],
constant EpiArgs& a [[buffer(4)]], uint gid [[thread_position_in_grid]]) {
const uint n = a.n_embd, t = gid / n, d = gid % n;
if (t >= a.n_tok) return;
float m = 0.0f;
for (uint c = 0; c < a.hc; ++c) {
const float g = G[ulong(t) * a.hc * n + c * n + d];
m = fma(xn[ulong(t) * a.hc * n + c * n + d], 1.0f / (1.0f + exp(-g)), m);
}
m = m / float(a.hc);
mixed[gid] = m;
mixed_h[gid] = half(m);
}
kernel void swiglu_h(device const float* g [[buffer(0)]], device const float* u [[buffer(1)]], device half* h [[buffer(2)]],
constant uint& n [[buffer(3)]], uint gid [[thread_position_in_grid]]) {
if (gid >= n) return;
const float gv = g[gid];
h[gid] = half(gv / (1.0f + exp(-gv)) * u[gid]);
}
// grid (T * n_embd / 256); sg: T values
kernel void moe_out(device const float* yr [[buffer(0)]], device const float* ys [[buffer(1)]],
device const float* sg [[buffer(2)]], device float* bo [[buffer(3)]], constant uint2& a [[buffer(4)]],
uint gid [[thread_position_in_grid]]) {
const uint t = gid / a.y;
if (t >= a.x) return;
bo[gid] = fma(1.0f / (1.0f + exp(-sg[t])), ys[gid], yr[gid]);
}
// ---- the HC read with fewer passes (after dense_mm.metal: GemmArgs, MetalPerformancePrimitives) ---------------------
// gr_norm_h: gr_norm with xn written as fp16 only (the GEMM operand; gemm_hc_up's epilogue reads it too).
// one threadgroup per (stream c, token t), 256 threads; args as gr_norm's (GrArgs, gr.metal)
kernel void gr_norm_h(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 half* xnh [[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], ss = 0.0f;
for (uint d = tid, i = 0; d < n; d += 256, ++i) {
float v = r[d];
if (a.apply) {
v = fma(bo[t * n + d], g, v);
r[d] = v;
}
rv[i] = 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 half* o = xnh + (t * a.hc + c) * n;
device const float* w = w_norm + c * n;
for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs);
}
// gr_norm_h's tail for one token's hc streams: the sums of squares ss[c] (per thread) to the norm, xnh = R' * w * rs
inline void gr_norm_h_out(thread const float (*rv)[10], device const float* w_norm, device half* xnh, uint t, uint hc,
uint n, float eps, thread const float* ss, threadgroup float (*part)[8], ushort tid,
ushort sg, ushort lane) {
for (uint c = 0; c < hc; ++c) {
const float s = simd_sum(ss[c]);
if (lane == 0) part[c][sg] = s;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint c = 0; c < hc; ++c) {
float tot = 0.0f;
#pragma clang loop unroll(full)
for (int i = 0; i < 8; ++i) tot += part[c][i];
const float rs = rsqrt(tot / float(n) + eps);
device half* o = xnh + (t * hc + c) * n;
device const float* w = w_norm + c * n;
for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[c][i] * w[d] * rs);
}
}
// stream-wise reduction keeps one residual stream plus bo live, rather than all four residual streams.
inline void gr_norm_h_stream_out(thread const float* rv, device const float* w, device half* o, uint n, float eps,
float ss, threadgroup float* part, ushort tid, ushort sg, ushort lane) {
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) + eps);
if (n == 2560) {
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) {
const uint d = tid + i * 256;
o[d] = half(rv[i] * w[d] * rs);
}
} else {
for (uint d = tid, i = 0; d < n; d += 256, ++i) o[d] = half(rv[i] * w[d] * rs);
}
}
// gr_norm_h4: gr_norm_h with one threadgroup per token over all hc streams, so bo is read once a token rather than
// hc times; R', xnh bitwise gr_norm_h's. grid (1, T), 256 threads; hc <= 4
kernel void gr_norm_h4(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 half* xnh [[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[4][8];
const uint t = tg.y, n = a.n_embd, hc = a.hc;
device float* r = R + t * hc * n;
if (!a.apply) {
for (uint c = 0; c < hc; ++c) {
float rv[10], ss = 0.0f;
if (n == 2560) {
device const float* rc = r + c * 2560 + tid;
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) {
const float v = rc[i * 256];
rv[i] = v;
ss = fma(v, v, ss);
}
} else {
for (uint d = tid, i = 0; d < n; d += 256, ++i) {
const float v = r[c * n + d];
rv[i] = v;
ss = fma(v, v, ss);
}
}
gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane);
}
return;
}
float bv[10];
if (n == 2560) {
device const float* brow = bo + t * 2560 + tid;
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) bv[i] = brow[i * 256];
} else {
for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = bo[t * n + d];
}
for (uint c = 0; c < hc; ++c) {
const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc)));
float rv[10], ss = 0.0f;
if (n == 2560) {
device float* rc = r + c * 2560 + tid;
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) {
const float v = fma(bv[i], g, rc[i * 256]);
rc[i * 256] = v;
rv[i] = v;
ss = fma(v, v, ss);
}
} else {
for (uint d = tid, i = 0; d < n; d += 256, ++i) {
const float v = fma(bv[i], g, r[c * n + d]);
r[c * n + d] = v;
rv[i] = v;
ss = fma(v, v, ss);
}
}
gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane);
}
}
// gr_norm_h_moe: gr_norm_h with the sublayer output bo formed on the fly as moe_out - bo = fma(sigmoid(sg[t]), ys, yr)
// - so the MoE's write never round-trips through bo; R', xnh bitwise moe_out + gr_norm_h. Always applies.
// One threadgroup per token t (all hc streams: bo is formed once a token), 256 threads; hc <= 4.
kernel void gr_norm_h_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]],
device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]],
device const float* inj [[buffer(4)]], device const float* w_norm [[buffer(5)]],
device half* xnh [[buffer(6)]], constant GrArgs& 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]]) {
threadgroup float part[4][8];
const uint t = tg.y, n = a.n_embd, hc = a.hc;
device float* r = R + t * hc * n;
float rv[4][10], g[4], ss[4] = {0.0f, 0.0f, 0.0f, 0.0f};
for (uint c = 0; c < hc; ++c) g[c] = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc)));
const float gs = 1.0f / (1.0f + exp(-sgv[t]));
for (uint d = tid, i = 0; d < n; d += 256, ++i) {
const float bo = fma(gs, ys[t * n + d], yr[t * n + d]);
for (uint c = 0; c < hc; ++c) {
const float v = fma(bo, g[c], r[c * n + d]);
r[c * n + d] = v;
rv[c][i] = v;
ss[c] = fma(v, v, ss[c]);
}
}
gr_norm_h_out(rv, w_norm, xnh, t, hc, n, a.eps, ss, part, tid, sg, lane);
}
// token t's routed slots for moe_comb: compacted valid YE row offsets (uint(r) * n) and their weights in threadgroup memory
inline uint moe_comb_slots(device const int* pos, device const float* srt_w, uint t, uint topk, uint n,
threadgroup uint* roff, threadgroup float* wk, ushort tid) {
int r = -1;
float w = 0.0f;
if (tid < topk) {
r = pos[t * topk + tid];
if (r >= 0) w = srt_w[r];
}
const bool valid = r >= 0;
const ushort prefix = simd_prefix_exclusive_sum(ushort(valid));
const ushort cnt = simd_sum(ushort(valid));
if (valid) {
roff[prefix] = uint(r) * n;
wk[prefix] = w;
}
if (tid == 0) roff[16] = cnt;
threadgroup_barrier(mem_flags::mem_threadgroup);
return roff[16];
}
// the routed experts' write at (t, d) as mm_combine (moe_mm.metal) forms it: sum_k w(t, k) * YE[pos(t, k)][d], k in order
inline float moe_comb(device const float* YE, threadgroup const uint* roff, threadgroup const float* wk, uint cnt,
uint d) {
float s = 0.0f;
for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s);
return s;
}
// gr_norm_h_moec: gr_norm_h_moe with mm_combine folded in - yr formed from the expert rows YE (pos, srt_w: mm_place's)
// as it is read, so it is never written; R', xnh bitwise mm_combine + gr_norm_h_moe. grid (1, T), 256 threads;
// hc <= 4, topk <= 16
kernel void gr_norm_h_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]],
device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]],
device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]],
device const float* inj [[buffer(6)]], device const float* w_norm [[buffer(7)]],
device half* xnh [[buffer(8)]], constant GrArgs& a [[buffer(9)]], constant uint& topk [[buffer(10)]],
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[4][8];
threadgroup uint roff[17];
threadgroup float wk[16];
const uint t = tg.y, n = a.n_embd, hc = a.hc;
const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid);
device float* r = R + t * hc * n;
float bv[10];
const float gs = 1.0f / (1.0f + exp(-sgv[t]));
if (n == 2560) {
// k outer: each slot's weight and row read once, its ten loads in flight together; s[i]'s fma order is moe_comb's
float s[10] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
for (uint k = 0; k < cnt; ++k) {
const float w = wk[k];
device const float* yk = YE + roff[k] + tid;
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) s[i] = fma(w, yk[i * 256], s[i]);
}
device const float* yrow = ys + t * 2560 + tid;
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) bv[i] = fma(gs, yrow[i * 256], s[i]);
} else {
for (uint d = tid, i = 0; d < n; d += 256, ++i) bv[i] = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d));
}
for (uint c = 0; c < hc; ++c) {
const float g = 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc)));
float rv[10], ss = 0.0f;
if (n == 2560) {
device float* rc = r + c * 2560 + tid;
#pragma clang loop unroll(full)
for (uint i = 0; i < 10; ++i) {
const float v = fma(bv[i], g, rc[i * 256]);
rc[i * 256] = v;
rv[i] = v;
ss = fma(v, v, ss);
}
} else {
for (uint d = tid, i = 0; d < n; d += 256, ++i) {
const float v = fma(bv[i], g, r[c * n + d]);
r[c * n + d] = v;
rv[i] = v;
ss = fma(v, v, ss);
}
}
gr_norm_h_stream_out(rv, w_norm + c * n, xnh + (t * hc + c) * n, n, a.eps, ss, part[c], tid, sg, lane);
}
}
// gr_apply_moe: the last layer's MoE write applied to R alone (R' of moe_out + gr_norm; its xn has no reader).
// grid (hc * n / 256, T), 256 threads
kernel void gr_apply_moe(device float* R [[buffer(0)]], device const float* yr [[buffer(1)]],
device const float* ys [[buffer(2)]], device const float* sgv [[buffer(3)]],
device const float* inj [[buffer(4)]], constant GrArgs& a [[buffer(5)]],
uint3 tg [[threadgroup_position_in_grid]], ushort tid [[thread_index_in_threadgroup]]) {
const uint ch = tg.x * 256 + tid, t = tg.y, n = a.n_embd;
if (ch >= a.hc * n) return;
const uint c = ch / n, d = ch - c * n;
const float g = 2.0f / (1.0f + exp(-inj[t * a.hc + c] / float(a.hc)));
const float gs = 1.0f / (1.0f + exp(-sgv[t]));
device float* r = R + t * a.hc * n + ch;
*r = fma(fma(gs, ys[t * n + d], yr[t * n + d]), g, *r);
}
// gr_apply_moec: gr_apply_moe with mm_combine folded in (as gr_norm_h_moec); R' bitwise mm_combine + gr_apply_moe.
// grid (ceil(n / 256), T), 256 threads: a thread's d over all hc streams (yr formed once); topk <= 16
kernel void gr_apply_moec(device float* R [[buffer(0)]], device const float* YE [[buffer(1)]],
device const int* pos [[buffer(2)]], device const float* srt_w [[buffer(3)]],
device const float* ys [[buffer(4)]], device const float* sgv [[buffer(5)]],
device const float* inj [[buffer(6)]], constant GrArgs& a [[buffer(7)]],
constant uint& topk [[buffer(8)]], uint3 tg [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]]) {
threadgroup uint roff[17];
threadgroup float wk[16];
const uint d = tg.x * 256 + tid, t = tg.y, n = a.n_embd, hc = a.hc;
const uint cnt = moe_comb_slots(pos, srt_w, t, topk, n, roff, wk, tid);
if (d >= n) return;
const float gs = 1.0f / (1.0f + exp(-sgv[t]));
const float bo = fma(gs, ys[t * n + d], moe_comb(YE, roff, wk, cnt, d));
device float* r = R + t * hc * n + d;
for (uint c = 0; c < hc; ++c) r[c * n] = fma(bo, 2.0f / (1.0f + exp(-inj[t * hc + c] / float(hc))), r[c * n]);
}
// gemm_hc_up: G = lo . w_up^T with w_up's rows interleaved by dq_f16 (perm = hc: row d * hc + c), so a TN-column
// tile holds TN / hc outputs d with all their streams; the epilogue mixed[t][d] = mean_c xn[t][c][d] *
// sigmoid(G[t][d * hc + c]) -> fp16 and, if mixed is bound, fp32. G is never written.
// grid (hc * n / TN, ceil(T / TM)), 32 * SG threads.
template <int TM, int TN, int SG, bool REG = false>
inline void hc_up_tile(device half* A, device half* B, device const half* xnh, device float* mixed,
device half* mixh, constant GemmArgs& g, constant EpiArgs& e,
int t0, int c0, ushort tid, threadgroup float* ot) {
tensor<device half, dextents<int32_t, 2>, tensor_inline> tA(A, dextents<int32_t, 2>(g.K, g.M));
tensor<device half, dextents<int32_t, 2>, tensor_inline> tB(B, dextents<int32_t, 2>(g.K, g.N));
constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v<int>, false, true, false);
matmul2d<desc, execution_simdgroups<SG>> op;
auto mA = tA.slice(0, t0);
auto mB = tB.slice(0, c0);
if constexpr (REG) {
auto acc = op.template get_destination_cooperative_tensor<decltype(mA), decltype(mB), float>();
op.run(mA, mB, acc);
const int n = int(e.n_embd), d0 = c0 / 4;
// On M5, each lane owns four consecutive output columns: one output's HC streams.
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < acc.get_capacity(); i += 4) {
const auto ix = acc.get_multidimensional_index(i);
const int t = t0 + ix[1], d = d0 + ix[0] / 4;
if (t >= int(e.n_tok)) continue;
device const half* xnh_t = xnh + ulong(t) * 4 * n + d;
float s = 0.0f;
#pragma clang loop unroll(full)
for (int c = 0; c < 4; ++c)
s = fma(float(xnh_t[c * n]), 1.0f / (1.0f + exp(-acc[i + c])), s);
s *= 0.25f;
const ulong out_idx = ulong(t) * n + d;
if (mixed) mixed[out_idx] = s;
mixh[out_idx] = half(s);
}
} else {
tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline> tO(ot, dextents<int32_t, 2>(TN, TM));
op.run(mA, mB, tO);
threadgroup_barrier(mem_flags::mem_threadgroup);
const int hc = int(e.hc), n = int(e.n_embd), DPT = TN / hc;
const int d0 = c0 / hc;
for (int idx = tid; idx < TM * DPT; idx += SG * 32) {
const int m = idx / DPT, dd = idx - m * DPT, t = t0 + m;
if (t >= int(e.n_tok)) continue;
float s = 0.0f;
for (int c = 0; c < hc; ++c)
s = fma(float(xnh[ulong(t) * hc * n + c * n + d0 + dd]), 1.0f / (1.0f + exp(-ot[m * TN + dd * hc + c])), s);
s = s / float(hc);
if (mixed) mixed[ulong(t) * n + d0 + dd] = s;
mixh[ulong(t) * n + d0 + dd] = half(s);
}
}
}
template <int TM, int TN, int SG, bool REG = false>
kernel void gemm_hc_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const half* xnh [[buffer(2)]],
device float* mixed [[buffer(3)]], device half* mixh [[buffer(4)]], constant GemmArgs& g [[buffer(5)]],
constant EpiArgs& e [[buffer(6)]], uint3 tg [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]]) {
threadgroup float ot[REG ? 1 : TM * TN];
const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN;
if constexpr (REG && TM > 16) {
if (t0 + 16 >= int(e.n_tok)) {
hc_up_tile<16, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot);
return;
}
}
if constexpr (REG && TM > 32) {
if (t0 + 32 >= int(e.n_tok)) {
hc_up_tile<32, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot);
return;
}
}
hc_up_tile<TM, TN, SG, REG>(A, B, xnh, mixed, mixh, g, e, t0, c0, tid, ot);
}
template [[host_name("gemm_hc_up_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4>) gemm_hc_up<64, 64, 4>;
template [[host_name("gemm_hc_up_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4>) gemm_hc_up<64, 32, 4>;
template [[host_name("gemm_hc_up_32_64_4")]] kernel decltype(gemm_hc_up<32, 64, 4>) gemm_hc_up<32, 64, 4>;
template [[host_name("gemm_hc_up_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4>) gemm_hc_up<128, 64, 4>;
template [[host_name("gemm_hc_up_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4>) gemm_hc_up<64, 128, 4>;
template [[host_name("gemm_hc_up_reg_64_64_4")]] kernel decltype(gemm_hc_up<64, 64, 4, true>) gemm_hc_up<64, 64, 4, true>;
template [[host_name("gemm_hc_up_reg_128_64_4")]] kernel decltype(gemm_hc_up<128, 64, 4, true>) gemm_hc_up<128, 64, 4, true>;
template [[host_name("gemm_hc_up_reg_64_128_4")]] kernel decltype(gemm_hc_up<64, 128, 4, true>) gemm_hc_up<64, 128, 4, true>;
template [[host_name("gemm_hc_up_reg_64_32_4")]] kernel decltype(gemm_hc_up<64, 32, 4, true>) gemm_hc_up<64, 32, 4, true>;
template [[host_name("gemm_hc_up_reg_32_128_4")]] kernel decltype(gemm_hc_up<32, 128, 4, true>) gemm_hc_up<32, 128, 4, true>;
// gemm_hc_down: the HC down GEMM V = xn . [w_down; w_inject]^T (T x (hc_lr + hc), K = hc * n_embd) with gr_epi_down
// as its register epilogue: lo_h = half(silu(v / hc)), inject = v. V is never written; lo_h, inject bitwise gemm_f16 +
// gr_epi_down. M5 layout only (coop_k4: a lane's elements 4j .. 4j + 3 are four consecutive columns of one row);
// hc_lr % 4 == 0, hc == 4. grid (ceil((hc_lr + hc) / TN), ceil(T / TM)), 32 * SG threads
template <int TM, int TN, int SG>
inline void hc_down_tile(device half* A, device half* B, device half* lo_h, device float* inject,
constant EpiArgs& e, int t0, int c0) {
const int K = int(e.hc * e.n_embd), nr = int(e.hc_lr + e.hc);
tensor<device half, dextents<int32_t, 2>, tensor_inline> tA(A, dextents<int32_t, 2>(K, int(e.n_tok)));
tensor<device half, dextents<int32_t, 2>, tensor_inline> tB(B, dextents<int32_t, 2>(K, nr));
constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v<int>, false, true, false);
matmul2d<desc, execution_simdgroups<SG>> op;
auto mA = tA.slice(0, t0);
auto mB = tB.slice(0, c0);
auto acc = op.template get_destination_cooperative_tensor<decltype(mA), decltype(mB), float>();
op.run(mA, mB, acc);
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < acc.get_capacity(); i += 4) {
const auto ix = acc.get_multidimensional_index(i);
const int r = c0 + ix[0], t = t0 + ix[1];
if (t >= int(e.n_tok) || r >= nr) continue;
if (r < int(e.hc_lr)) {
half4 o;
#pragma clang loop unroll(full)
for (int c = 0; c < 4; ++c) {
const float z = acc[i + c] * 0.25f; // / hc, exact: hc == 4
o[c] = half(z / (1.0f + exp(-z)));
}
*(device half4*)(lo_h + ulong(t) * e.hc_lr + r) = o;
} else {
*(device float4*)(inject + ulong(t) * e.hc) = float4(acc[i], acc[i + 1], acc[i + 2], acc[i + 3]);
}
}
}
template <int TM, int TN, int SG>
inline void hc_down_rows(device half* A, device half* B, device half* lo_h, device float* inject,
constant EpiArgs& e, int t0, int c0) {
// Trimming 32-row parents regresses full tiles on M5.
if constexpr (TM > 32) {
if (t0 + 16 >= int(e.n_tok)) {
hc_down_tile<16, TN, SG>(A, B, lo_h, inject, e, t0, c0);
return;
}
if (t0 + 32 >= int(e.n_tok)) {
hc_down_tile<32, TN, SG>(A, B, lo_h, inject, e, t0, c0);
return;
}
}
hc_down_tile<TM, TN, SG>(A, B, lo_h, inject, e, t0, c0);
}
template <int TM, int TN, int SG>
kernel void gemm_hc_down(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device half* lo_h [[buffer(2)]],
device float* inject [[buffer(3)]], constant EpiArgs& e [[buffer(4)]],
uint3 tg [[threadgroup_position_in_grid]]) {
const int c0 = int(tg.x) * TN, t0 = int(tg.y) * TM;
if constexpr (TN == 128) {
// A single 96-column tail beats serial 64+16 on M5, including SG=2.
if (e.hc == 4 && c0 + 64 == int(e.hc_lr)) {
hc_down_rows<TM, 96, SG>(A, B, lo_h, inject, e, t0, c0);
return;
}
}
hc_down_rows<TM, TN, SG>(A, B, lo_h, inject, e, t0, c0);
}
template [[host_name("gemm_hc_down_16_64_2")]] kernel decltype(gemm_hc_down<16, 64, 2>) gemm_hc_down<16, 64, 2>;
template [[host_name("gemm_hc_down_32_64_2")]] kernel decltype(gemm_hc_down<32, 64, 2>) gemm_hc_down<32, 64, 2>;
template [[host_name("gemm_hc_down_32_128_4")]] kernel decltype(gemm_hc_down<32, 128, 4>) gemm_hc_down<32, 128, 4>;
template [[host_name("gemm_hc_down_64_128_4")]] kernel decltype(gemm_hc_down<64, 128, 4>) gemm_hc_down<64, 128, 4>;
template [[host_name("gemm_hc_down_32_128_2")]] kernel decltype(gemm_hc_down<32, 128, 2>) gemm_hc_down<32, 128, 2>;
// gemm_swiglu_up: the shared expert's up GEMM u = x . w_up^T (M x N, K) with swiglu_h as its register epilogue against
// the gate GEMM's g: h = half(silu(g) * u). u is never written; h bitwise gemm_f16 + swiglu_h. One accumulator: the
// gate stays a plain gemm_f16 before it. M5 layout only (as gemm_hc_down); N % 4 == 0.
// grid (ceil(N / TN), ceil(M / TM)), 32 * SG threads
template <int TM, int TN, int SG>
inline void swiglu_up_tile(device half* A, device half* B, device const float* g, device half* h,
constant GemmArgs& a, int t0, int c0) {
tensor<device half, dextents<int32_t, 2>, tensor_inline> tA(A, dextents<int32_t, 2>(a.K, a.M));
tensor<device half, dextents<int32_t, 2>, tensor_inline> tB(B, dextents<int32_t, 2>(a.K, a.N));
constexpr auto desc = matmul2d_descriptor(TM, TN, dynamic_length_v<int>, false, true, false);
matmul2d<desc, execution_simdgroups<SG>> op;
auto mA = tA.slice(0, t0);
auto mB = tB.slice(0, c0);
auto acc = op.template get_destination_cooperative_tensor<decltype(mA), decltype(mB), float>();
op.run(mA, mB, acc);
#pragma clang loop unroll(full)
for (uint16_t i = 0; i < acc.get_capacity(); i += 4) {
const auto ix = acc.get_multidimensional_index(i);
const int c = c0 + ix[0], t = t0 + ix[1];
if (t >= a.M || c >= a.N) continue;
const ulong idx = ulong(t) * uint(a.N) + uint(c);
const float4 gv = *(device const float4*)(g + idx);
half4 o;
#pragma clang loop unroll(full)
for (int k = 0; k < 4; ++k) o[k] = half(gv[k] / (1.0f + exp(-gv[k])) * acc[i + k]);
*(device half4*)(h + idx) = o;
}
}
template <int TM, int TN, int SG>
kernel void gemm_swiglu_up(device half* A [[buffer(0)]], device half* B [[buffer(1)]], device const float* g [[buffer(2)]],
device half* h [[buffer(3)]], constant GemmArgs& a [[buffer(4)]],
uint3 tg [[threadgroup_position_in_grid]]) {
const int t0 = int(tg.y) * TM, c0 = int(tg.x) * TN;
if constexpr (TM > 32) {
if (t0 + 16 >= a.M) {
swiglu_up_tile<16, TN, SG>(A, B, g, h, a, t0, c0);
return;
}
if (t0 + 32 >= a.M) {
swiglu_up_tile<32, TN, SG>(A, B, g, h, a, t0, c0);
return;
}
}
swiglu_up_tile<TM, TN, SG>(A, B, g, h, a, t0, c0);
}
template [[host_name("gemm_swiglu_up_32_64_2")]] kernel decltype(gemm_swiglu_up<32, 64, 2>) gemm_swiglu_up<32, 64, 2>;
template [[host_name("gemm_swiglu_up_64_128_4")]] kernel decltype(gemm_swiglu_up<64, 128, 4>) gemm_swiglu_up<64, 128, 4>;