Skip to content

inkling: shared experts to the GPU, 1.75 to 2.46 tok/s on Apple Silicon #757

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JustVugg merged 2 commits into from
Aug 4, 2026
Merged

inkling: shared experts to the GPU, 1.75 to 2.46 tok/s on Apple Silicon #757

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cda2e4e
inkling: overlap shared experts with the last GPU round
rgbkrk Aug 1, 2026
b9ee0ed
metal+inkling: shared experts on the GPU as a bf16 block
rgbkrk Aug 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions c/backend_metal.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,13 +200,23 @@ kernel void moe_gemv(device const ulong* waddr [[buffer(0)]], device const ulong
acc += (dot(m0*sA,xs[l*2]) + dot(m1*sB,xs[l*2+1])) * (0.5f*db);
}
}
} else if (fmt == 5) { // bf16 raw rows: no scales, upper half of f32
device const ushort* w=(device const ushort*)(waddr[e])+(long)o*K;
device const ushort4* w4=(device const ushort4*)w;
for(int c=slane;c<K8;c+=32){ ushort4 a=w4[2*c], b=w4[2*c+1];
float4 w0=float4(as_type<float>(uint(a.x)<<16), as_type<float>(uint(a.y)<<16),
as_type<float>(uint(a.z)<<16), as_type<float>(uint(a.w)<<16));
float4 w1=float4(as_type<float>(uint(b.x)<<16), as_type<float>(uint(b.y)<<16),
as_type<float>(uint(b.z)<<16), as_type<float>(uint(b.w)<<16));
acc+=dot(w0,x4[2*c])+dot(w1,x4[2*c+1]); }
for(int i=K8*8+slane;i<K;i+=32) acc += as_type<float>(uint(w[i])<<16)*xr[i];
} else { device const char* w=(device const char*)(waddr[e])+(long)o*K;
device const char4* w4=(device const char4*)w;
for(int c=slane;c<K8;c+=32) acc+=dot(float4(w4[2*c]),x4[2*c])+dot(float4(w4[2*c+1]),x4[2*c+1]);
for(int i=K8*8+slane;i<K;i+=32) acc+=float(w[i])*xr[i];
}
acc=simd_sum(acc);
if(slane==0) yout[row] = (fmt==6) ? acc : acc*sc[o]; // fmt=6: scales live in-block
if(slane==0) yout[row] = (fmt==6||fmt==5) ? acc : acc*sc[o]; // fmt 6: in-block scales; fmt 5: none
}

// fmt=6 activation rotation for the GPU-resident down-projection input: one FWHT
Expand Down Expand Up @@ -1143,7 +1153,7 @@ static bool resolve_attn(const AttnW *W, float *Lc, float *Rc,
const float *const *gs, const float *const *us, const float *const *ds,
const float *xg, const int *xoff, const int *nr, int R,
id<MTLBuffer> xg_buf, id<MTLBuffer> gg_buf, id<MTLBuffer> uu_buf, id<MTLBuffer> hh_buf) {
if (!g_dev || (fmt != 1 && fmt != 2 && fmt != 6)) return nil;
if (!g_dev || (fmt != 1 && fmt != 2 && fmt != 5 && fmt != 6)) return nil;
if (fmt == 6) { /* e8 kernel assumes clean block tiling, and every FWHT tile of the
* down input (CPU tiling rule, e8_rot_rows) must fit threadgroup mem */
if ((D & 255) || (Iinter & 31)) return nil;
Expand Down
128 changes: 115 additions & 13 deletions c/inkling.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -795,6 +795,28 @@ static void model_init(Model *m, const char *snap, int cap, int bits) {
LDW(sh_g, "mlp.shared_experts.gate_proj");
LDW(sh_u, "mlp.shared_experts.up_proj");
LDW(sh_d, "mlp.shared_experts.down_proj");
#ifdef COLI_METAL
/* Move the bf16 shared-expert weights into one page-aligned slab
* (gates, then ups, then downs) and repoint the Wt handles into
* it: the CPU path reads the same bytes it always did, and the
* slab registers once so the GPU fmt=5 path can resolve it. */
if (g_metal && l->sh_g.h && l->sh_u.h && l->sh_d.h) {
int64_t I = c->moe_inter, ns = c->n_shared;
size_t one = (size_t)ns*I*D*2, pg = 16384;
size_t len = (3*one + pg - 1) / pg * pg;
void *sl = NULL;
if (!posix_memalign(&sl, pg, len)) {
memcpy((char*)sl, l->sh_g.h, one);
memcpy((char*)sl + one, l->sh_u.h, one);
memcpy((char*)sl + 2*one, l->sh_d.h, one);
free(l->sh_g.h); free(l->sh_u.h); free(l->sh_d.h);
l->sh_g.h = (uint16_t*)sl;
l->sh_u.h = (uint16_t*)((char*)sl + one);
l->sh_d.h = (uint16_t*)((char*)sl + 2*one);
coli_metal_register(sl, len);
}
}
#endif
}
#undef LD
#undef LDW
Expand Down Expand Up @@ -1128,6 +1150,28 @@ static void dense_mlp(Model *m, Layer *l, float *x, int S, float *out) {
* (2) fill ALL missing experts in one parallel burst (the NVMe wants queue
* depth — during prefill this batches the whole sequence's misses), then
* (3) compute. */
/* shared experts for all S positions: gamma inside (before down_proj is
* linear, so applied at the end). Factored out so the Metal path can run it
* on the CPU while the last routed-expert round is in flight on the GPU. */
static void shared_experts_cpu(Model *m, Layer *l, const float *x, int S,
float *out, const float *wgt,
float *g, float *u, float *hh) {
Cfg *c = &m->c;
int D = c->hidden, K = c->topk, I = c->moe_inter, ns = c->n_shared;
for (int s = 0; s < S; s++) {
const float *xs = x + (int64_t)s*D;
float *os = out + (int64_t)s*D;
const float *w = wgt + (int64_t)s*(K+ns);
for (int j = 0; j < ns; j++) {
matmul_w(g, xs, wt_off_i(l->sh_g, (int64_t)j*I*D, D), 1, D, I);
matmul_w(u, xs, wt_off_i(l->sh_u, (int64_t)j*I*D, D), 1, D, I);
for (int i = 0; i < I; i++) g[i] = siluf(g[i]) * u[i];
matmul_w(hh, g, wt_off_i(l->sh_d, (int64_t)j*D*I, I), 1, I, D);
for (int d = 0; d < D; d++) os[d] += w[K+j] * hh[d];
}
}
}

static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
Cfg *c = &m->c;
int D = c->hidden, E = c->n_experts, K = c->topk, I = c->moe_inter, ns = c->n_shared;
Expand Down Expand Up @@ -1201,6 +1245,7 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
int cap = m->cache[layer].cap; if (cap < 1) cap = 1;
float *g = falloc(2*I), *u = g + I, *hh = falloc(D);
int q4 = m->xq && m->rb13*2 == D; /* packed int4 vs int8 container */
int shared_done = 0; /* set when overlapped with the last GPU round */
int64_t npair = (int64_t)S*K;
#ifdef COLI_METAL
/* per-round scratch for the batched GPU submit: pairs grouped by expert,
Expand All @@ -1222,6 +1267,39 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
mxoff = malloc((cap+1)*sizeof(int)); mnr = malloc(cap*sizeof(int));
mrows = malloc(cap*sizeof(int)); mgi = malloc(cap*sizeof(int)); mfp = malloc(cap*sizeof(int));
}
/* Shared experts as ONE bf16 (fmt=5) block, submitted BEFORE the routed
* rounds so it rides the GPU while the CPU routes, fills, and packs.
* Every token uses every shared expert, so group j is simply all S rows
* with weight w[K+j]. Falls back to the CPU loop if begin refuses. */
ColiMetalMoeHandle *sh_h = NULL;
float *sxg = NULL, *srw = NULL;
if (mxg && ns > 0 && ns <= 8 && l->sh_g.h && l->sh_u.h && l->sh_d.h &&
!(getenv("INK_METAL_SHARED") && *getenv("INK_METAL_SHARED") == '0')) {
double ts = now_s();
sxg = falloc((int64_t)ns*S*D); srw = falloc((int64_t)ns*S);
const void *sgp[8], *sup[8], *sdp[8];
const float *sscale[8];
int sxoff[9], snr[8];
int *srows = malloc((size_t)ns*S*sizeof(int));
for (int j = 0; j < ns; j++) {
sgp[j] = l->sh_g.h + (int64_t)j*I*D;
sup[j] = l->sh_u.h + (int64_t)j*I*D;
sdp[j] = l->sh_d.h + (int64_t)j*D*I;
sscale[j] = (const float*)sgp[j]; /* fmt=5 never reads scales */
sxoff[j] = j*S; snr[j] = S;
for (int s = 0; s < S; s++) {
memcpy(sxg + ((int64_t)j*S + s)*D, x + (int64_t)s*D, (size_t)D*sizeof(float));
srows[j*S + s] = s;
srw[j*S + s] = wgt[(int64_t)s*(K+ns) + K + j];
}
}
sxoff[ns] = ns*S;
sh_h = coli_metal_moe_block_begin(ns, D, I, 5, sgp, sup, sdp,
sscale, sscale, sscale,
sxg, sxoff, snr, srows, srw);
free(srows);
m->t_shared += now_s() - ts;
}
#endif
for (int64_t base = 0; base < npair; base += cap) {
int64_t end = base + cap < npair ? base + cap : npair;
Expand Down Expand Up @@ -1279,6 +1357,28 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
mgp[j] = e->p13; mup[j] = e->p13 + (int64_t)I*m->rb13; mdp[j] = e->p2;
mgs[j] = e->s13; mus[j] = e->s13 + I; mds[j] = e->s2;
}
/* Last round: submit async and run the shared experts on the
* CPU while the GPU computes the routed ones — they are
* independent (both read x, both accumulate into out, and the
* GPU only touches out in _end's scatter-add, after the wait).
* On a GPU fault _end returns 0 and the CPU loop below redoes
* the round; shared_done stays set either way. */
if (base + cap >= npair && !sh_h) {
ColiMetalMoeHandle *h = coli_metal_moe_block_begin(
nb, D, I, q4 ? 2 : 1, mgp, mup, mdp, mgs, mus, mds,
mxg, mxoff, mnr, mrows, mrw);
if (h) {
double ts = now_s();
shared_experts_cpu(m, l, x, S, out, wgt, g, u, hh);
double sh = now_s() - ts;
m->t_shared += sh;
shared_done = 1;
int ok = coli_metal_moe_block_end(h, out);
m->t_expert += (now_s() - te) - sh;
if (ok) continue; /* round done on the GPU */
te = now_s(); /* fault: CPU redo below */
}
}
if (coli_metal_moe_block(nb, D, I, q4 ? 2 : 1, mgp, mup, mdp, mgs, mus, mds,
mxg, mxoff, mnr, mrows, mrw, out, S)) {
m->t_expert += now_s() - te;
Expand Down Expand Up @@ -1317,20 +1417,21 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
}
m->t_expert += now_s() - te;
}
/* shared experts: una volta per token, fuori dai giri (non usano la cache) */
for (int s = 0; s < S; s++) {
const float *xs = x + (int64_t)s*D;
float *os = out + (int64_t)s*D;
float *w = wgt + (int64_t)s*(K+ns);
#ifdef COLI_METAL
/* GPU shared block: wait + scatter-add. A fault falls through to CPU. */
if (sh_h) {
double ts = now_s();
/* gamma inside (before down_proj is linear, so applied at the end) */
for (int j = 0; j < ns; j++) {
matmul_w(g, xs, wt_off_i(l->sh_g, (int64_t)j*I*D, D), 1, D, I);
matmul_w(u, xs, wt_off_i(l->sh_u, (int64_t)j*I*D, D), 1, D, I);
for (int i = 0; i < I; i++) g[i] = siluf(g[i]) * u[i];
matmul_w(hh, g, wt_off_i(l->sh_d, (int64_t)j*D*I, I), 1, I, D);
for (int d = 0; d < D; d++) os[d] += w[K+j] * hh[d];
}
if (coli_metal_moe_block_end(sh_h, out)) shared_done = 1;
m->t_shared += now_s() - ts;
}
#endif
/* shared experts: una volta per token, fuori dai giri (non usano la cache).
* Under Metal these run on the GPU as an fmt=5 block overlapped with the
* routed rounds (or on the CPU during the last GPU round); either path
* sets shared_done. */
if (!shared_done) {
double ts = now_s();
shared_experts_cpu(m, l, x, S, out, wgt, g, u, hh);
m->t_shared += now_s() - ts;
}
free(logits); free(idx); free(keff); free(wgt); free(use); free(fill); free(fl);
Expand All @@ -1341,6 +1442,7 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
free(mgs); free(mus); free(mds); free(mslot);
free(mxoff); free(mnr); free(mrows); free(mgi); free(mfp);
}
free(sxg); free(srw);
#endif
}

Expand Down
Loading