// metal/kernels/moe.metal - routed experts of one layer for a verify window of 1..8 tokens, grouped by expert:
//
// moe_plan one threadgroup: the window's top-k ids -> groups (distinct experts, by id) with their entries
// (token, weight), each (token, k)'s entry index, and the indirect grid sizes of the next two kernels.
// Deterministic (no atomics).
// moe_gu per group, a block of FFN rows: gate and up rows read once, applied to every entry of the group
// (the shared input x is read from cache), h = silu(gate) * up per entry.
// moe_down per group, a block of output rows: ye[entry] = weight * (down . h[entry]).
// moe_combine y[token] = sum over k of ye[entry(token, k)], in k order.
//
// The plan is one int buffer (see P_* below). Expert addressing: MoeArgs (a pack blob or the GGUF tensors):
// [gate rows | up rows | down rows] in GGUF blocks (tools/iq_pack.py's experts.bin layout).
#define P_NG 0
#define P_IGU 4
#define P_IDN 8
#define P_MAXE 128
#define P_GEXP 16
#define P_GSTART (P_GEXP + P_MAXE)
#define P_ETOK (P_GSTART + P_MAXE + 4)
#define P_EW (P_ETOK + P_MAXE)
#define P_TOKENT (P_EW + P_MAXE)
#define P_SPLIT (P_TOKENT + P_MAXE) // function constant 26 only: [0] resident groups, then the indirect grids
#define P_IGU_R (P_SPLIT + 4) // ... of the resident groups (the GPU computes them during the host's turn)
#define P_IDN_R (P_SPLIT + 8)
#define P_IGU_M (P_SPLIT + 12) // ... and of the missing ones (after it)
#define P_IDN_M (P_SPLIT + 16)
#define P_GSLOT (P_SPLIT + 20)
#define P_INTS (P_GSLOT + P_MAXE)
// Expert e's matrices: gate at W + gate_off + e * gu_stride, up at W + up_off + e * gu_stride, down at
// W + down_off + e * d_stride (a pack's blobs: gate_off 0, strides = the blob; a GGUF: the three tensors).
struct MoeArgs {
ulong gate_off;
ulong up_off;
ulong down_off;
ulong gu_stride;
ulong d_stride;
uint gu_row; // bytes per gate/up row
uint d_row; // bytes per down row
uint n_embd;
uint n_ff;
uint n_tok;
uint topk;
uint n_expert;
uint gx_gu; // threadgroups along the rows for moe_gu (n_ff / rows per threadgroup)
uint gx_dn; // ... for moe_down (n_embd / rows per threadgroup)
uint pad;
};
// The SSD-streaming arena (metal/src/engine/stream.cpp): with function constant 20 set, expert e of the layer lies in
// an arena slot, smap[e] = (chunk << 28) | its 16 KiB page in that chunk; the chunks are W (buffer 0) and buffers 7..9
// (gate_off / up_off / down_off are then offsets inside the slot). The same block is in moe_mm.metal.
#ifndef STRATA_ARENA_METAL
#define STRATA_ARENA_METAL
constant bool kArenaFc [[function_constant(20)]];
constant bool kArena = is_function_constant_defined(kArenaFc) && kArenaFc;
static inline device const uchar* arena_slot(uint v, device const uchar* W, device const uchar* W1, device const uchar* W2,
device const uchar* W3) {
const uint c = v >> 28;
device const uchar* b = c == 0 ? W : (c == 1 ? W1 : (c == 2 ? W2 : W3));
return b + (ulong)(v & 0x0FFFFFFFu) * 16384ul;
}
#endif
// Function constant 26 (the SSD-streaming decode window, with smap): the groups whose expert is in the arena when the
// routing is made come first, the missing ones after them, each part with its own indirect grids (P_SPLIT): the
// resident part runs while the host reads the missing experts. The entries keep their (token, k) mapping, so the
// combine (and every number) is the same; without it the plan is the old one, bit for bit.
constant bool kSplitResFc [[function_constant(26)]];
constant bool kSplitRes = is_function_constant_defined(kSplitResFc) && kSplitResFc;
// s_grp, s_ent: 64 each (the upper halves for the missing groups)
static inline void plan_body(device const int* ids, device const float* w, device int* plan, constant MoeArgs& a,
ushort tid, ushort nth, ushort sg, ushort lane, threadgroup int* s_grp,
threadgroup int* s_ent, device const uint* smap = nullptr) {
const int ne = int(a.n_tok * a.topk);
// one thread per expert id (nth == n_expert, at most 1024)
const int e = tid;
int cnt = 0;
if (e < int(a.n_expert)) {
#pragma clang loop unroll_count(10)
for (int i = 0; i < ne; ++i) cnt += ids[i] == e;
}
const int has = cnt > 0;
const bool split = kSplitRes && smap != nullptr;
const int miss = (split && has && smap[e] == 0xFFFFFFFFu) ? 1 : 0;
const int h0 = has - miss, c0 = miss ? 0 : cnt, c1 = miss ? cnt : 0;
const int pg = simd_prefix_exclusive_sum(h0), pe = simd_prefix_exclusive_sum(c0);
const int pg1 = split ? simd_prefix_exclusive_sum(miss) : 0;
const int pe1 = split ? simd_prefix_exclusive_sum(c1) : 0;
if (lane == 31) {
s_grp[sg] = pg + h0; s_ent[sg] = pe + c0;
if (split) { s_grp[32 + sg] = pg1 + miss; s_ent[32 + sg] = pe1 + c1; }
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sg == 0) {
const int nsg = (nth + 31) / 32;
const int g0 = lane < nsg ? s_grp[lane] : 0, n0 = lane < nsg ? s_ent[lane] : 0;
const int og0 = simd_prefix_exclusive_sum(g0), oe0 = simd_prefix_exclusive_sum(n0);
if (lane < nsg) { s_grp[lane] = og0; s_ent[lane] = oe0; }
const int ng0 = simd_broadcast(og0 + g0, 31), ne0 = simd_broadcast(oe0 + n0, 31);
int og = ng0, oe = ne0;
if (split) {
const int g1 = lane < nsg ? s_grp[32 + lane] : 0, n1 = lane < nsg ? s_ent[32 + lane] : 0;
const int og1 = ng0 + simd_prefix_exclusive_sum(g1), oe1 = ne0 + simd_prefix_exclusive_sum(n1);
if (lane < nsg) { s_grp[32 + lane] = og1; s_ent[32 + lane] = oe1; }
og = simd_broadcast(og1 + g1, 31);
oe = simd_broadcast(oe1 + n1, 31);
}
if (lane == 0) {
plan[P_NG] = og;
plan[P_IGU] = int(a.gx_gu); plan[P_IGU + 1] = og; plan[P_IGU + 2] = 1;
plan[P_IDN] = int(a.gx_dn); plan[P_IDN + 1] = og; plan[P_IDN + 2] = 1;
plan[P_GSTART + og] = oe;
if (split) {
plan[P_SPLIT] = ng0;
plan[P_IGU_R] = int(a.gx_gu); plan[P_IGU_R + 1] = ng0; plan[P_IGU_R + 2] = 1;
plan[P_IDN_R] = int(a.gx_dn); plan[P_IDN_R + 1] = ng0; plan[P_IDN_R + 2] = 1;
plan[P_IGU_M] = int(a.gx_gu); plan[P_IGU_M + 1] = og - ng0; plan[P_IGU_M + 2] = 1;
plan[P_IDN_M] = int(a.gx_dn); plan[P_IDN_M + 1] = og - ng0; plan[P_IDN_M + 2] = 1;
}
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (has) {
const int gi = (miss ? s_grp[32 + sg] + pg1 : s_grp[sg] + pg), off = miss ? s_ent[32 + sg] + pe1 : s_ent[sg] + pe;
plan[P_GEXP + gi] = e;
plan[P_GSTART + gi] = off;
if (split && !miss) plan[P_GSLOT + gi] = int(smap[e]);
const bool topk10 = (a.topk == 10u);
int j = 0;
#pragma clang loop unroll_count(10)
for (int i = 0; i < ne; ++i)
if (ids[i] == e) {
plan[P_ETOK + off + j] = topk10 ? ((i * 205) >> 11) : (i / int(a.topk));
plan[P_EW + off + j] = as_type<int>(w[i]);
plan[P_TOKENT + i] = off + j;
++j;
}
}
}
kernel void moe_plan(device const int* ids [[buffer(0)]], device const float* w [[buffer(1)]],
device int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]],
ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
threadgroup int s_grp[64], s_ent[64];
plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent);
}
// The router's tail for the window, then the plan, in one threadgroup of n_expert (512) threads: simdgroup t routes
// token t exactly as native_router.cu (pinned llama.cpp topk-moe): softmax over the 512 logits, top 10 by
// probability (equal values: the lower expert index), weights = p / max(sum of the selected, 2^-14), the sum taken
// by winning lane in selection order and reduced with the same xor tree.
#ifndef STRATA_MOCE_FC_METAL
#define STRATA_MOCE_FC_METAL
constant int kMoceProtectFc [[function_constant(24)]];
constant int kMoceProtect = is_function_constant_defined(kMoceProtectFc) ? kMoceProtectFc : 3;
constant uint kMoceScaleR1BitsFc [[function_constant(25)]];
constant float kMoceScaleR1 = is_function_constant_defined(kMoceScaleR1BitsFc) ? as_type<float>(kMoceScaleR1BitsFc) : 1.0f;
#endif
kernel void moe_route(device const float* logits [[buffer(0)]], device int* ids [[buffer(1)]],
device float* w [[buffer(2)]], device int* plan [[buffer(3)]], constant MoeArgs& a [[buffer(4)]],
device const uint* smap [[buffer(5)]],
ushort tid [[thread_position_in_threadgroup]], ushort nth [[threads_per_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
threadgroup int s_grp[64], s_ent[64];
if (sg < a.n_tok) {
device const float* l = logits + sg * 512;
const bool e512 = (a.n_expert == 512u);
float v[16];
float mx = -INFINITY;
float sum = 0.0f;
if (e512) {
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) {
v[i] = l[lane + 32 * i];
mx = max(mx, v[i]);
}
#pragma clang loop unroll(full)
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];
}
#pragma clang loop unroll(full)
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 {
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY;
mx = max(mx, v[i]);
}
#pragma clang loop unroll(full)
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] = (lane + 32 * i < int(a.n_expert)) ? exp(v[i] - mx) : 0.0f;
sum += v[i];
}
#pragma clang loop unroll(full)
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) {
if (lane + 32 * i >= int(a.n_expert)) {
v[i] = -INFINITY;
} else {
v[i] *= rcp;
if (isnan(v[i])) v[i] = -FLT_MAX;
}
}
}
const bool draft_moce = (a.pad & 0x80000000u) != 0u && sg > 0;
const float base_moce = (a.pad & 0x7FFFFFFFu) != 0u ? as_type<float>(a.pad & 0x7FFFFFFFu) : 1.0f;
const float moce_scale = draft_moce ? (base_moce * base_moce) : base_moce;
const float rcp_moce = moce_scale > 1.0f ? (1.0f / moce_scale) : 1.0f;
const int r_prot = draft_moce ? 1 : kMoceProtect;
const float r1_scale = kMoceScaleR1;
const float rcp_r1 = r1_scale > 1.0f ? (1.0f / r1_scale) : 1.0f;
const float r_step_scale = (r1_scale > 1.0f && moce_scale > 1.0f) ? (moce_scale * rcp_r1) : moce_scale;
const bool has_r1 = (r1_scale > 1.0f && r_prot > 1);
const bool has_moce = (moce_scale > 1.0f || r1_scale > 1.0f) && smap != nullptr;
const int min_scaled_rank = has_r1 ? 1 : r_prot;
ushort res_mask = 0;
if (has_moce) {
if (e512) {
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i)
if (smap[lane + 32 * i] != 0xFFFFFFFFu) res_mask |= ushort(1u << i);
} else {
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
if (idx < int(a.n_expert) && smap[idx] != 0xFFFFFFFFu) res_mask |= ushort(1u << i);
}
}
}
ushort scaled_mask = 0;
float selected = 0.0f, ssum = 0.0f;
if (!has_moce) {
for (int rank = 0; rank < 10; ++rank) {
float best = v[0];
int ex = lane;
#pragma clang loop unroll(full)
for (int i = 1; i < 16; ++i)
if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
#pragma clang loop unroll(full)
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; }
}
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i];
if ((ex & 31) == lane) {
ids[sg * 10 + rank] = ex;
ssum += best;
}
if (rank == lane) selected = best;
}
} else {
for (int rank = 0; rank < 10; ++rank) {
if (has_r1) {
if (rank == 1) {
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) {
if ((res_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) {
v[i] *= r1_scale;
scaled_mask |= ushort(1u << i);
}
}
} else if (rank == r_prot && moce_scale > 1.0f) {
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) {
if ((scaled_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) {
v[i] *= r_step_scale;
}
}
}
} else if (rank == r_prot && moce_scale > 1.0f) {
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) {
if ((res_mask & ushort(1u << i)) != 0u && v[i] > 0.0f) {
v[i] *= moce_scale;
scaled_mask |= ushort(1u << i);
}
}
}
float best = v[0];
int ex = lane;
#pragma clang loop unroll(full)
for (int i = 1; i < 16; ++i)
if (v[i] > best) { best = v[i]; ex = lane + 32 * i; }
#pragma clang loop unroll(full)
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 >= min_scaled_rank && 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) {
// without a tail bonus, the rank-1 scale remains on unselected residents.
const bool r1_scaled = has_r1 && (rank < r_prot || moce_scale <= 1.0f);
true_prob = best * (r1_scaled ? rcp_r1 : rcp_moce);
}
}
#pragma clang loop unroll(full)
for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i];
if ((ex & 31) == lane) {
ids[sg * 10 + rank] = ex;
ssum += true_prob;
}
if (rank == lane) selected = true_prob;
}
}
#pragma clang loop unroll(full)
for (ushort m = 16; m > 0; m >>= 1) ssum += simd_shuffle_xor(ssum, m);
ssum = max(ssum, 6.103515625e-5f);
const float inv = 1.0f / ssum;
if (lane < 10) w[sg * 10 + lane] = selected * inv;
}
threadgroup_barrier(mem_flags::mem_device);
plan_body(ids, w, plan, a, tid, nth, sg, lane, s_grp, s_ent, smap);
}
// ---- gate + up + SwiGLU -------------------------------------------------------------------------------------------
// One pass over a group's rows for all its NE entries (1..8: an expert appears once per token); NE is uniform per
// threadgroup, so the kernel switches once into a body specialised for it.
//
// gu_fused: per row, gate and up decoded together and each entry's x slice loaded once for both (the 32-weight units,
// whose x load is two runs and their sums: the M5 Air, 2560 x 640, NE 3-4: Q4_K / Q5_K / Q6_K / Q8_0 -8..-25%,
// IQ4_XS -31..-36% from NE 2). The 16-weight units keep two x loads a row: the two decoded units held at once cost
// them more (IQ2_XS / IQ3_S at NR 1, NE 3: +15..+21%). The same sums in the same order either way.
template <int F, int NE> struct gu_fused {
enum { value = fmt<F>::UNIT == 32 && (NE >= 3 || F == T_IQ4_XS) };
};
template <int F>
static inline float q2_combine(ushort2 pk, float d, thread const float* xs, float xsum0) {
float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
#pragma clang loop unroll(full)
for (int j = 0; j < 8; j += 2) {
if (F == T_Q2_0R) {
const uint2 q = uint2(pk) >> (2 * j);
const uint2 bits = ((q | (q << 14)) & 0x00030003u) | 0x64006400u;
const float4 v = float4(as_type<half4>(bits) - 1025.0h);
s0 = fma(v.x, xs[j], s0);
s1 = fma(v.y, xs[j + 1], s1);
s2 = fma(v.z, xs[8 + j], s2);
s3 = fma(v.w, xs[8 + j + 1], s3);
} else {
s0 = fma(float(pk.x & ushort(3u << (2 * j))), xs[j], s0);
s1 = fma(float(pk.x & ushort(3u << (2 * j + 2))), xs[j + 1], s1);
s2 = fma(float(pk.y & ushort(3u << (2 * j))), xs[8 + j], s2);
s3 = fma(float(pk.y & ushort(3u << (2 * j + 2))), xs[8 + j + 1], s3);
}
}
const float dot = d * (s0 + s1) + d * (s2 + s3);
return F == T_Q2_0R ? (dot - 0.0f) : (dot - (d * xsum0 + 0.0f));
}
static inline float iq_combine(uint4 g, uint s16, float sc, thread const float* xs) {
float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;
#pragma clang loop unroll(full)
for (int k = 0; k < 2; ++k) {
const uchar4 v0 = as_type<uchar4>(g[k]), v1 = as_type<uchar4>(g[2 + k]);
const uint m0 = s16 >> (4 * k), m1 = s16 >> (8 + 4 * k);
#pragma clang loop unroll(full)
for (int j = 0; j < 4; j += 2) {
s0 = fma((m0 & (1u << j)) ? -float(v0[j]) : float(v0[j]), xs[4 * k + j], s0);
s1 = fma((m0 & (2u << j)) ? -float(v0[j + 1]) : float(v0[j + 1]), xs[4 * k + j + 1], s1);
s2 = fma((m1 & (1u << j)) ? -float(v1[j]) : float(v1[j]), xs[8 + 4 * k + j], s2);
s3 = fma((m1 & (2u << j)) ? -float(v1[j + 1]) : float(v1[j + 1]), xs[8 + 4 * k + j + 1], s3);
}
}
return sc * (s0 + s1) + sc * (s2 + s3) - 0.0f;
}
template <int F>
static inline float deq_iq_packed(device const uchar* b, int ui, thread uint4& g, thread uint& s16, threadgroup const uint* tab) {
const int ib = ui >> 1, h = ui & 1;
if (F == T_IQ2_S) {
const uint q2 = *(device const ushort*)(b + 2 + 2 * ui);
s16 = *(device const ushort*)(b + 34 + 2 * ui);
const uint qh = uint(b[66 + ib]) >> (4 * h);
const uint i0 = (q2 & 255) | ((qh << 8) & 0x300);
const uint i1 = (q2 >> 8) | ((qh << 6) & 0x300);
const uint2 g0 = ctab ? as_type<uint2>(iq2s_grid[i0]) : ((threadgroup const uint2*)tab)[i0];
const uint2 g1 = ctab ? as_type<uint2>(iq2s_grid[i1]) : ((threadgroup const uint2*)tab)[i1];
g = uint4(g0.x, g0.y, g1.x, g1.y);
const uint sc = b[74 + ib];
return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f);
} else if (F == T_IQ2_XXS) {
device const uchar* qs = b + 2 + 4 * ib + 2 * ui;
const uint a1 = rd32(b + 2 + 8 * ib + 4), s0 = a1 >> (14 * h);
const uint sa = s0 & 127, sb = (s0 >> 7) & 127;
s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12);
const uint2 g0 = ctab ? as_type<uint2>(iq2xxs_grid[qs[0]]) : ((threadgroup const uint2*)tab)[qs[0]];
const uint2 g1 = ctab ? as_type<uint2>(iq2xxs_grid[qs[1]]) : ((threadgroup const uint2*)tab)[qs[1]];
g = uint4(g0.x, g0.y, g1.x, g1.y);
return rdh(b) * fma(float(a1 >> 28), 0.25f, 0.125f);
} else if (F == T_IQ2_XS) {
device const ushort* qs = (device const ushort*)(b + 2) + 2 * ui;
const uint v0 = qs[0], v1 = qs[1];
const uint sa = v0 >> 9, sb = v1 >> 9;
s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12);
const uint2 g0 = ctab ? as_type<uint2>(iq2xs_grid[v0 & 511]) : ((threadgroup const uint2*)tab)[v0 & 511];
const uint2 g1 = ctab ? as_type<uint2>(iq2xs_grid[v1 & 511]) : ((threadgroup const uint2*)tab)[v1 & 511];
g = uint4(g0.x, g0.y, g1.x, g1.y);
const uint sc = b[66 + ib];
return rdh(b) * fma(float((sc >> (4 * h)) & 15), 0.25f, 0.125f);
} else if (F == T_IQ3_S) {
const int l0 = 2 * h, l1 = 2 * h + 1;
const uint qh = b[66 + ib];
s16 = *(device const ushort*)(b + 74 + 2 * ui);
const uint i0 = uint(b[2 + 8 * ib + 2 * l0]) | ((qh << (8 - 2 * l0)) & 256);
const uint i1 = uint(b[2 + 8 * ib + 2 * l0 + 1]) | ((qh << (7 - 2 * l0)) & 256);
const uint i2 = uint(b[2 + 8 * ib + 2 * l1]) | ((qh << (8 - 2 * l1)) & 256);
const uint i3 = uint(b[2 + 8 * ib + 2 * l1 + 1]) | ((qh << (7 - 2 * l1)) & 256);
g = uint4(g_iq3s(tab, i0), g_iq3s(tab, i1), g_iq3s(tab, i2), g_iq3s(tab, i3));
const uint sc = b[106 + (ib >> 1)];
return rdh(b) * fma(float((sc >> (4 * (ib & 1))) & 15), 2.0f, 1.0f);
} else {
const uint aux = rd32(b + 2 + 64 + 4 * ib), s0 = aux >> (14 * h);
const uint sa = s0 & 127, sb = (s0 >> 7) & 127;
s16 = (sa & 15u) | (sign4_iq2xs(sa, 1) << 4) | ((sb & 15u) << 8) | (sign4_iq2xs(sb, 1) << 12);
device const uchar* qs = b + 2 + 4 * ui;
g = uint4(g_iq3xxs(tab, qs[0]), g_iq3xxs(tab, qs[1]), g_iq3xxs(tab, qs[2]), g_iq3xxs(tab, qs[3]));
return rdh(b) * fma(float(aux >> 28), 0.5f, 0.25f);
}
}
template <int F, int LPR, int NR, int NE>
static inline void gu_body(device const uchar* g0, device const uchar* u0, device const float* x,
device const float* xs16, device float* h, device const int* etok, int c, int row0, int li,
constant MoeArgs& a, threadgroup const uint* tab) {
constexpr int UNIT = fmt<F>::UNIT;
float ag[NR][NE], au[NR][NE];
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r)
#pragma clang loop unroll(full)
for (int t = 0; t < NE; ++t) ag[r][t] = au[r][t] = 0.0f;
const int U = int(a.n_embd) / UNIT;
if (NE == 1) {
const int tok = etok[0];
x += tok * int(a.n_embd);
xs16 += tok * int(a.n_embd / 16);
}
for (int u = li; u < U; u += LPR) {
if (NE == 1 && UNIT == 16 && F != T_Q2_0 && F != T_Q2_0R) {
float xs[UNIT], xsum[2];
X<F>::load(x, xs16, u, xs, xsum);
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) {
float q[UNIT];
float4 c;
Dec<F>::run(g0 + r * a.gu_row, u, q, c, tab);
ag[r][0] += unit_combine<UNIT>(q, c, xs, xsum);
Dec<F>::run(u0 + r * a.gu_row, u, q, c, tab);
au[r][0] += unit_combine<UNIT>(q, c, xs, xsum);
}
} else if (NE == 1) {
float xs[UNIT], xsum[2];
X<F>::load(x, xs16, u, xs, xsum);
for (int r = 0; r < NR; ++r) {
float q[UNIT];
float4 c;
Dec<F>::run(g0 + r * a.gu_row, u, q, c, tab);
ag[r][0] += unit_combine<UNIT>(q, c, xs, xsum);
Dec<F>::run(u0 + r * a.gu_row, u, q, c, tab);
au[r][0] += unit_combine<UNIT>(q, c, xs, xsum);
}
} else if (F == T_Q2_0 || F == T_Q2_0R) {
const int boff = (u >> 2) * 18, qoff = 2 + 4 * (u & 3);
ushort2 pkg[NR], pku[NR];
float dg[NR], du[NR];
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) {
device const uchar* bg = g0 + r * a.gu_row + boff;
device const uchar* bu = u0 + r * a.gu_row + boff;
device const ushort* pg = (device const ushort*)(bg + qoff);
device const ushort* pu = (device const ushort*)(bu + qoff);
pkg[r] = ushort2(pg[0], pg[1]);
pku[r] = ushort2(pu[0], pu[1]);
dg[r] = rdh(bg);
du[r] = rdh(bu);
}
#pragma clang loop unroll(full)
for (int t = 0; t < NE; ++t) {
const int tok = etok[t];
float xs[UNIT], xsum[2];
X<F>::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum);
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) {
ag[r][t] += q2_combine<F>(pkg[r], dg[r], xs, xsum[0]);
au[r][t] += q2_combine<F>(pku[r], du[r], xs, xsum[0]);
}
}
} else if (gu_fused<F, NE>::value) {
for (int r = 0; r < NR; ++r) {
float qg[UNIT], qu[UNIT];
float4 cg, cu;
Dec<F>::run(g0 + r * a.gu_row, u, qg, cg, tab);
Dec<F>::run(u0 + r * a.gu_row, u, qu, cu, tab);
for (int t = 0; t < NE; ++t) {
const int tok = etok[t];
float xs[UNIT], xsum[2];
X<F>::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum);
ag[r][t] += unit_combine<UNIT>(qg, cg, xs, xsum);
au[r][t] += unit_combine<UNIT>(qu, cu, xs, xsum);
}
}
} else if (NR == 1 && NE == 4 && (F == T_IQ2_S || F == T_IQ2_XXS || F == T_IQ2_XS || F == T_IQ3_S || F == T_IQ3_XXS)) {
const int boff = (u >> 4) * fmt<F>::BB, ui = (LPR % 16 == 0) ? (li & 15) : (u & 15);
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) {
uint4 gg, gu;
uint sg, su;
const float dg = deq_iq_packed<F>(g0 + r * a.gu_row + boff, ui, gg, sg, tab);
const float du = deq_iq_packed<F>(u0 + r * a.gu_row + boff, ui, gu, su, tab);
#pragma clang loop unroll(full)
for (int t = 0; t < NE; ++t) {
const int tok = etok[t];
float xs[UNIT], xsum[2];
X<F>::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum);
ag[r][t] += iq_combine(gg, sg, dg, xs);
au[r][t] += iq_combine(gu, su, du, xs);
}
}
} else {
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) {
float q[UNIT];
float4 c;
Dec<F>::run(g0 + r * a.gu_row, u, q, c, tab);
#pragma clang loop unroll(full)
for (int t = 0; t < NE; ++t) {
const int tok = etok[t];
float xs[UNIT], xsum[2];
X<F>::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum);
ag[r][t] += unit_combine<UNIT>(q, c, xs, xsum);
}
Dec<F>::run(u0 + r * a.gu_row, u, q, c, tab);
#pragma clang loop unroll(full)
for (int t = 0; t < NE; ++t) {
const int tok = etok[t];
float xs[UNIT], xsum[2];
X<F>::load(x + tok * int(a.n_embd), xs16 + tok * int(a.n_embd / 16), u, xs, xsum);
au[r][t] += unit_combine<UNIT>(q, c, xs, xsum);
}
}
}
}
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r)
#pragma clang loop unroll(full)
for (int t = 0; t < NE; ++t) {
const float gv = lane_sum<LPR>(ag[r][t]), uv = lane_sum<LPR>(au[r][t]);
if (li == 0) h[(c + t) * int(a.n_ff) + row0 + r] = gv / (1.0f + exp(-gv)) * uv;
}
}
// function constant 27, with 26's plan: 1 = only the groups resident at routing time (indirect grid P_IGU_R / P_IDN_R),
// 2 = only the missing ones (P_IGU_M / P_IDN_M); 0 or unset = all (P_IGU / P_IDN)
constant int kPhaseFc [[function_constant(27)]];
constant int kPhase = is_function_constant_defined(kPhaseFc) ? kPhaseFc : 0;
// bench probe (function constant 22): 1 = no table load (wrong values for the i-quants)
constant int MOE_PROBE [[function_constant(22)]];
constant int moe_probe = is_function_constant_defined(MOE_PROBE) ? MOE_PROBE : 0;
template <int F, int LPR, int NR, int SG>
kernel void moe_gu(device const uchar* W [[buffer(0)]], device const float* x [[buffer(1)]],
device float* h [[buffer(2)]], device const int* plan [[buffer(3)]],
constant MoeArgs& a [[buffer(4)]], device const float* xs16 [[buffer(5)]],
device const uint* smap [[buffer(6)]], device const uchar* W1 [[buffer(7)]],
device const uchar* W2 [[buffer(8)]], device const uchar* W3 [[buffer(9)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) {
constexpr int RG = 32 / LPR;
threadgroup uint tab[fmt<F>::TAB];
if (fmt<F>::TAB > 1 && moe_probe != 1) {
load_tab<F>(tab, tid, SG * 32);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0);
if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return;
const int li = lane % LPR;
const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR;
const ulong ex = (ulong)plan[P_GEXP + g];
const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c;
device const int* etok = plan + P_ETOK + c;
const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u);
device const uchar* E = kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.gu_stride;
device const uchar* g0 = E + a.gate_off + (ulong)row0 * a.gu_row;
device const uchar* u0 = E + a.up_off + (ulong)row0 * a.gu_row;
switch (ne) {
case 1: gu_body<F, LPR, NR, 1>(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break;
case 2: gu_body<F, LPR, NR, 2>(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break;
case 3: gu_body<F, LPR, NR, 3>(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break;
case 4: gu_body<F, LPR, NR, 4>(g0, u0, x, xs16, h, etok, c, row0, li, a, tab); break;
default:
for (int k = 0; k < ne; k += 4) {
const int m = min(4, ne - k);
if (m == 4) gu_body<F, LPR, NR, 4>(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab);
else if (m == 3) gu_body<F, LPR, NR, 3>(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab);
else if (m == 2) gu_body<F, LPR, NR, 2>(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab);
else gu_body<F, LPR, NR, 1>(g0, u0, x, xs16, h, etok + k, c + k, row0, li, a, tab);
}
}
}
// ---- down ---------------------------------------------------------------------------------------------------------
template <int F, int LPR, int NR, int NE>
static inline void dn_body(device const uchar* d0, device const float* h, device const float* hs16, device float* ye,
device const int* plan, int c, int row0, int li, constant MoeArgs& a,
threadgroup const uint* tab) {
constexpr int UNIT = fmt<F>::UNIT;
float ac[NR][NE];
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r)
#pragma clang loop unroll(full)
for (int t = 0; t < NE; ++t) ac[r][t] = 0.0f;
const int U = int(a.n_ff) / UNIT;
for (int u = li; u < U; u += LPR) {
if (F == T_Q2_0 || F == T_Q2_0R) {
device const uchar* bu = d0 + (u >> 2) * 18;
const int qoff = 2 + 4 * (u & 3);
ushort2 pk[NR];
float d[NR];
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) {
device const uchar* b = bu + r * a.d_row;
device const ushort* p = (device const ushort*)(b + qoff);
pk[r] = ushort2(p[0], p[1]);
d[r] = rdh(b);
}
#pragma clang loop unroll(full)
for (int t = 0; t < NE; ++t) {
float xs[UNIT], xsum[2];
X<F>::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum);
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r)
ac[r][t] += q2_combine<F>(pk[r], d[r], xs, xsum[0]);
}
} else if (NE == 1) {
float xs[UNIT], xsum[2];
X<F>::load(h + c * int(a.n_ff), hs16 + c * int(a.n_ff / 16), u, xs, xsum);
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) {
float q[UNIT];
float4 sc;
Dec<F>::run(d0 + r * a.d_row, u, q, sc, tab);
ac[r][0] += unit_combine<UNIT>(q, sc, xs, xsum);
}
} else {
float q[NR][UNIT];
float4 sc[NR];
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) Dec<F>::run(d0 + r * a.d_row, u, q[r], sc[r], tab);
#pragma clang loop unroll(full)
for (int t = 0; t < NE; ++t) {
float xs[UNIT], xsum[2];
X<F>::load(h + (c + t) * int(a.n_ff), hs16 + (c + t) * int(a.n_ff / 16), u, xs, xsum);
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r) ac[r][t] += unit_combine<UNIT>(q[r], sc[r], xs, xsum);
}
}
}
#pragma clang loop unroll(full)
for (int r = 0; r < NR; ++r)
#pragma clang loop unroll(full)
for (int t = 0; t < NE; ++t) {
const float v = lane_sum<LPR>(ac[r][t]);
if (li == 0) ye[(c + t) * int(a.n_embd) + row0 + r] = as_type<float>(plan[P_EW + c + t]) * v;
}
}
template <int F, int LPR, int NR, int SG>
kernel void moe_down(device const uchar* W [[buffer(0)]], device const float* h [[buffer(1)]],
device float* ye [[buffer(2)]], device const int* plan [[buffer(3)]],
constant MoeArgs& a [[buffer(4)]], device const float* hs16 [[buffer(5)]],
device const uint* smap [[buffer(6)]], device const uchar* W1 [[buffer(7)]],
device const uchar* W2 [[buffer(8)]], device const uchar* W3 [[buffer(9)]],
uint3 tg [[threadgroup_position_in_grid]], ushort sg [[simdgroup_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]], ushort tid [[thread_index_in_threadgroup]]) {
constexpr int RG = 32 / LPR;
threadgroup uint tab[fmt<F>::TAB];
if (fmt<F>::TAB > 1) {
load_tab<F>(tab, tid, SG * 32);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const int g = int(tg.y) + (kPhase == 2 ? plan[P_SPLIT] : 0);
if (g >= (kPhase == 1 ? plan[P_SPLIT] : plan[P_NG])) return;
const int li = lane % LPR;
const int row0 = ((int(tg.x) * SG + sg) * RG + lane / LPR) * NR;
const ulong ex = (ulong)plan[P_GEXP + g];
const int c = plan[P_GSTART + g], ne = plan[P_GSTART + g + 1] - c;
const uint scode = (kArena && kPhase == 1) ? uint(plan[P_GSLOT + g]) : (kArena ? smap[ex] : 0u);
device const uchar* d0 = (kArena ? arena_slot(scode, W, W1, W2, W3) : W + ex * a.d_stride) + a.down_off + (ulong)row0 * a.d_row;
switch (ne) {
case 1: dn_body<F, LPR, NR, 1>(d0, h, hs16, ye, plan, c, row0, li, a, tab); break;
case 2: dn_body<F, LPR, NR, 2>(d0, h, hs16, ye, plan, c, row0, li, a, tab); break;
case 3: dn_body<F, LPR, NR, 3>(d0, h, hs16, ye, plan, c, row0, li, a, tab); break;
case 4: dn_body<F, LPR, NR, 4>(d0, h, hs16, ye, plan, c, row0, li, a, tab); break;
default:
for (int k = 0; k < ne; k += 4) {
const int m = min(4, ne - k);
if (m == 4) dn_body<F, LPR, NR, 4>(d0, h, hs16, ye, plan, c + k, row0, li, a, tab);
else if (m == 3) dn_body<F, LPR, NR, 3>(d0, h, hs16, ye, plan, c + k, row0, li, a, tab);
else if (m == 2) dn_body<F, LPR, NR, 2>(d0, h, hs16, ye, plan, c + k, row0, li, a, tab);
else dn_body<F, LPR, NR, 1>(d0, h, hs16, ye, plan, c + k, row0, li, a, tab);
}
}
}
kernel void moe_combine(device const float* ye [[buffer(0)]], device float* y [[buffer(1)]],
device const int* plan [[buffer(2)]], constant MoeArgs& a [[buffer(3)]],
uint gid [[thread_position_in_grid]]) {
const uint t = gid / a.n_embd, k = gid % a.n_embd;
if (t >= a.n_tok) return;
device const int* te = plan + P_TOKENT + t * a.topk;
float s = 0.0f;
if (a.topk == 10) {
#pragma clang loop unroll(full)
for (uint j = 0; j < 10; ++j) s += ye[te[j] * int(a.n_embd) + k];
} else {
for (uint j = 0; j < a.topk; ++j) s += ye[te[j] * int(a.n_embd) + k];
}
y[gid] = s;
}
// ---- instantiations: moe_gu_<type>_<LPR>_<NR>_<SG>, moe_down_<type>_<LPR>_<NR>_<SG> -------------------------------
#define GU(F, L, R, S) template [[host_name("moe_gu_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_gu<F, L, R, S>) moe_gu<F, L, R, S>;
#define DN(F, L, R, S) template [[host_name("moe_down_" #F "_" #L "_" #R "_" #S)]] kernel decltype(moe_down<F, L, R, S>) moe_down<F, L, R, S>;
#define GU_SET(F) GU(F, 32, 1, 4) GU(F, 32, 2, 4) GU(F, 32, 4, 4) GU(F, 32, 2, 2) GU(F, 32, 2, 8) GU(F, 32, 1, 8) \
GU(F, 16, 1, 4) GU(F, 16, 2, 4) GU(F, 8, 1, 4) GU(F, 8, 2, 4) GU(F, 8, 1, 2)
#define DN_SET(F) DN(F, 8, 1, 4) DN(F, 8, 2, 4) DN(F, 8, 4, 4) DN(F, 8, 1, 8) DN(F, 8, 2, 8) DN(F, 4, 1, 4) DN(F, 4, 2, 4) \
DN(F, 4, 4, 4) DN(F, 16, 1, 4) DN(F, 16, 2, 4) DN(F, 32, 1, 4) DN(F, 2, 2, 4)
GU_SET(T_Q2_0)
GU_SET(T_IQ3_XXS)
GU_SET(T_IQ2_XXS)
GU_SET(T_IQ2_XS)
GU_SET(T_IQ2_S)
GU_SET(T_IQ3_S)
GU_SET(T_IQ4_XS) // the IQ3_S pack's gate/up in one layer
GU_SET(T_Q4_K) GU_SET(T_Q5_K) GU_SET(T_Q6_K) GU_SET(T_Q8_0) // the UD-Q4/5/6_K_XL packs
DN_SET(T_Q5_1) DN_SET(T_Q8_0)
DN_SET(T_Q2_0)
DN_SET(T_Q2_0R)
DN_SET(T_IQ4_NL)