/* SPDX-License-Identifier: AGPL-3.0-or-later */ /* Copyright (C) 2026 David (Corey) Herrera */ /** * @file fm_kernel_moe.m * @brief Host dispatch and MSL source for the mixture-of-experts kernels. * * Written from a behaviour specification (docs/cleanroom/f_metal_spec.md); no ggml or * llama.cpp source was consulted. * * @status None * @issues See fm_kernel_moe.h (router -1 rounds and -3.0e38f seed, sort duplicate-expert * slot, no >= n_expert guard, grouped down reassociated, no overflow checks). * @todo None */ #include "metal/kernels/fm_kernel_moe.h" #include #include "metal/fm_metal_internal.h" #include "metal/kernels/fm_kernel_registry.h" /* Function-constant indices owned by this family (block 240-259). */ #define FM_MOE_FC_DOWN_TOPK 240 #define FM_MOE_FC_DOWN_HALVES 241 #define FM_MOE_FC_GATEUP_HALVES 242 /* Super-block geometry of the two weight formats the grouped kernels read. */ #define FM_MOE_BLOCK_WEIGHTS 256 #define FM_MOE_Q3_K_BLOCK_BYTES 110 #define FM_MOE_Q4_K_BLOCK_BYTES 144 /* ================================================================ * Router and sort * ================================================================ */ void fm_moe_router_topk_at(fm_bufref logits, fm_bufref top_idx, fm_bufref top_wt, int n_tokens, int n_expert, int expert_used, float expert_weights_scale) { if (!fm_metal_is_initialized() || !logits.buffer || !top_idx.buffer || !top_wt.buffer) return; if (n_tokens <= 0 || n_expert <= 0 || expert_used <= 0 || expert_used > n_expert) return; /* The scale travels as raw bits so every inline argument is a uint. */ uint32_t scale_bits = 0; memcpy(&scale_bits, &expert_weights_scale, sizeof(scale_bits)); fm_bufref refs[] = { logits, top_idx, top_wt }; uint32_t consts[] = { (uint32_t)n_tokens, (uint32_t)n_expert, (uint32_t)expert_used, scale_bits }; fm_dispatch_threadgroups_refs(@"fmmoe_router_topk", NULL, refs, 3, consts, 4, n_tokens, FM_MOE_REDUCE_THREADS); } void fm_moe_router_topk(fm_buffer *logits, fm_buffer *top_idx, fm_buffer *top_wt, int n_tokens, int n_expert, int expert_used, float expert_weights_scale) { if (!logits || !top_idx || !top_wt) return; fm_moe_router_topk_at(fm_bufref_at(logits, 0), fm_bufref_at(top_idx, 0), fm_bufref_at(top_wt, 0), n_tokens, n_expert, expert_used, expert_weights_scale); } void fm_moe_expert_sort_at(fm_bufref top_idx, fm_bufref expert_counts, fm_bufref expert_offsets, fm_bufref sorted_token, fm_bufref sorted_slot, fm_bufref sorted_expert, int n_tokens, int n_expert, int expert_used) { if (!fm_metal_is_initialized() || !top_idx.buffer || !expert_counts.buffer || !expert_offsets.buffer || !sorted_token.buffer || !sorted_slot.buffer || !sorted_expert.buffer) return; if (n_tokens <= 0 || n_expert <= 0 || expert_used <= 0) return; /* Count first, then offsets + scatter: the second pass needs every expert's * count, and a separate dispatch is the barrier that avoids atomics. */ fm_bufref count_refs[] = { top_idx, expert_counts }; uint32_t count_consts[] = { (uint32_t)n_tokens, (uint32_t)expert_used }; fm_dispatch_threadgroups_refs(@"fmmoe_expert_count", NULL, count_refs, 2, count_consts, 2, n_expert, FM_MOE_REDUCE_THREADS); fm_bufref scatter_refs[] = { top_idx, expert_counts, expert_offsets, sorted_token, sorted_slot, sorted_expert }; uint32_t scatter_consts[] = { (uint32_t)n_tokens, (uint32_t)n_expert, (uint32_t)expert_used }; fm_dispatch_threadgroups_refs(@"fmmoe_expert_scatter", NULL, scatter_refs, 6, scatter_consts, 3, n_expert, FM_MOE_REDUCE_THREADS); } void fm_moe_expert_sort(fm_buffer *top_idx, fm_buffer *expert_counts, fm_buffer *expert_offsets, fm_buffer *sorted_token, fm_buffer *sorted_slot, fm_buffer *sorted_expert, int n_tokens, int n_expert, int expert_used) { if (!top_idx || !expert_counts || !expert_offsets || !sorted_token || !sorted_slot || !sorted_expert) return; fm_moe_expert_sort_at(fm_bufref_at(top_idx, 0), fm_bufref_at(expert_counts, 0), fm_bufref_at(expert_offsets, 0), fm_bufref_at(sorted_token, 0), fm_bufref_at(sorted_slot, 0), fm_bufref_at(sorted_expert, 0), n_tokens, n_expert, expert_used); } /* ================================================================ * Grouped expert projections * ================================================================ */ /* Shared launch of both gate+up forms; `refs` already holds the form's buffers. */ static void fm_moe_gateup_launch(NSString *kernel_name, const fm_bufref *refs, int n_refs, int n_requests, int expert_used, int expert_ff, int cols, int row_stride_bytes) { if (!fm_metal_is_initialized()) return; if (n_requests <= 0 || expert_used <= 0 || expert_ff <= 0 || cols <= 0 || (cols % FM_MOE_BLOCK_WEIGHTS) != 0) return; if (row_stride_bytes < (cols / FM_MOE_BLOCK_WEIGHTS) * FM_MOE_Q3_K_BLOCK_BYTES) return; fm_fc_values fc; fm_fc_reset(&fc); if (fm_fc_set_int(&fc, FM_MOE_FC_GATEUP_HALVES, cols / FM_MOE_LANE_WEIGHTS) != 0) return; uint32_t consts[] = { (uint32_t)n_requests, (uint32_t)expert_used, (uint32_t)expert_ff, (uint32_t)cols, (uint32_t)row_stride_bytes }; fm_dispatch_threadgroups_refs(kernel_name, &fc, refs, n_refs, consts, 5, n_requests * expert_ff, FM_MOE_LANES); } void fm_moe_gateup_q3k_at(fm_bufref gate_all, fm_bufref up_all, fm_bufref x, fm_bufref top_idx, fm_bufref hid, int n_tokens, int expert_used, int expert_ff, int cols, int row_stride_bytes) { if (!gate_all.buffer || !up_all.buffer || !x.buffer || !top_idx.buffer || !hid.buffer) return; if (n_tokens <= 0) return; fm_bufref refs[] = { gate_all, up_all, x, top_idx, hid }; fm_moe_gateup_launch(@"fmmoe_gateup_q3k", refs, 5, n_tokens * expert_used, expert_used, expert_ff, cols, row_stride_bytes); } void fm_moe_gateup_q3k(fm_buffer *gate_all, fm_buffer *up_all, fm_buffer *x, fm_buffer *top_idx, fm_buffer *hid, int n_tokens, int expert_used, int expert_ff, int cols, int row_stride_bytes) { if (!gate_all || !up_all || !x || !top_idx || !hid) return; fm_moe_gateup_q3k_at(fm_bufref_at(gate_all, 0), fm_bufref_at(up_all, 0), fm_bufref_at(x, 0), fm_bufref_at(top_idx, 0), fm_bufref_at(hid, 0), n_tokens, expert_used, expert_ff, cols, row_stride_bytes); } void fm_moe_gateup_q3k_sorted_at(fm_bufref gate_all, fm_bufref up_all, fm_bufref x, fm_bufref sorted_token, fm_bufref sorted_slot, fm_bufref sorted_expert, fm_bufref hid, int n_requests, int expert_used, int expert_ff, int cols, int row_stride_bytes) { if (!gate_all.buffer || !up_all.buffer || !x.buffer || !sorted_token.buffer || !sorted_slot.buffer || !sorted_expert.buffer || !hid.buffer) return; fm_bufref refs[] = { gate_all, up_all, x, sorted_token, sorted_slot, sorted_expert, hid }; fm_moe_gateup_launch(@"fmmoe_gateup_q3k_sorted", refs, 7, n_requests, expert_used, expert_ff, cols, row_stride_bytes); } void fm_moe_gateup_q3k_sorted(fm_buffer *gate_all, fm_buffer *up_all, fm_buffer *x, fm_buffer *sorted_token, fm_buffer *sorted_slot, fm_buffer *sorted_expert, fm_buffer *hid, int n_requests, int expert_used, int expert_ff, int cols, int row_stride_bytes) { if (!gate_all || !up_all || !x || !sorted_token || !sorted_slot || !sorted_expert || !hid) return; fm_moe_gateup_q3k_sorted_at(fm_bufref_at(gate_all, 0), fm_bufref_at(up_all, 0), fm_bufref_at(x, 0), fm_bufref_at(sorted_token, 0), fm_bufref_at(sorted_slot, 0), fm_bufref_at(sorted_expert, 0), fm_bufref_at(hid, 0), n_requests, expert_used, expert_ff, cols, row_stride_bytes); } /* Shared launch of both down forms; block_bytes selects the stride check. */ static void fm_moe_down_launch(NSString *kernel_name, int block_bytes, fm_bufref down_all, fm_bufref hid, fm_bufref top_idx, fm_bufref top_wt, fm_bufref out, int n_tokens, int expert_used, int rows, int ff, int row_stride_bytes) { if (!fm_metal_is_initialized() || !down_all.buffer || !hid.buffer || !top_idx.buffer || !top_wt.buffer || !out.buffer) return; if (n_tokens <= 0 || expert_used <= 0 || rows <= 0 || ff <= 0 || (ff % FM_MOE_BLOCK_WEIGHTS) != 0) return; if (row_stride_bytes < (ff / FM_MOE_BLOCK_WEIGHTS) * block_bytes) return; fm_fc_values fc; fm_fc_reset(&fc); if (fm_fc_set_int(&fc, FM_MOE_FC_DOWN_TOPK, expert_used) != 0 || fm_fc_set_int(&fc, FM_MOE_FC_DOWN_HALVES, ff / FM_MOE_LANE_WEIGHTS) != 0) return; fm_bufref refs[] = { down_all, hid, top_idx, top_wt, out }; uint32_t consts[] = { (uint32_t)n_tokens, (uint32_t)rows, (uint32_t)ff, (uint32_t)row_stride_bytes }; fm_dispatch_threadgroups_refs(kernel_name, &fc, refs, 5, consts, 4, n_tokens * rows, FM_MOE_LANES); } void fm_moe_down_q3k_at(fm_bufref down_all, fm_bufref hid, fm_bufref top_idx, fm_bufref top_wt, fm_bufref out, int n_tokens, int expert_used, int rows, int ff, int row_stride_bytes) { fm_moe_down_launch(@"fmmoe_down_q3k", FM_MOE_Q3_K_BLOCK_BYTES, down_all, hid, top_idx, top_wt, out, n_tokens, expert_used, rows, ff, row_stride_bytes); } void fm_moe_down_q3k(fm_buffer *down_all, fm_buffer *hid, fm_buffer *top_idx, fm_buffer *top_wt, fm_buffer *out, int n_tokens, int expert_used, int rows, int ff, int row_stride_bytes) { if (!down_all || !hid || !top_idx || !top_wt || !out) return; fm_moe_down_q3k_at(fm_bufref_at(down_all, 0), fm_bufref_at(hid, 0), fm_bufref_at(top_idx, 0), fm_bufref_at(top_wt, 0), fm_bufref_at(out, 0), n_tokens, expert_used, rows, ff, row_stride_bytes); } void fm_moe_down_q4k_at(fm_bufref down_all, fm_bufref hid, fm_bufref top_idx, fm_bufref top_wt, fm_bufref out, int n_tokens, int expert_used, int rows, int ff, int row_stride_bytes) { fm_moe_down_launch(@"fmmoe_down_q4k", FM_MOE_Q4_K_BLOCK_BYTES, down_all, hid, top_idx, top_wt, out, n_tokens, expert_used, rows, ff, row_stride_bytes); } void fm_moe_down_q4k(fm_buffer *down_all, fm_buffer *hid, fm_buffer *top_idx, fm_buffer *top_wt, fm_buffer *out, int n_tokens, int expert_used, int rows, int ff, int row_stride_bytes) { if (!down_all || !hid || !top_idx || !top_wt || !out) return; fm_moe_down_q4k_at(fm_bufref_at(down_all, 0), fm_bufref_at(hid, 0), fm_bufref_at(top_idx, 0), fm_bufref_at(top_wt, 0), fm_bufref_at(out, 0), n_tokens, expert_used, rows, ff, row_stride_bytes); } /* ================================================================ * Elementwise and shared-expert gate * ================================================================ */ void fm_moe_silu_mul_at(fm_bufref a, fm_bufref b, fm_bufref out, int n) { if (!fm_metal_is_initialized() || !a.buffer || !b.buffer || !out.buffer || n <= 0) return; fm_bufref refs[] = { a, b, out }; uint32_t consts[] = { (uint32_t)n }; fm_dispatch_1d_refs(@"fmmoe_silu_mul", NULL, refs, 3, consts, 1, n); } void fm_moe_silu_mul(fm_buffer *a, fm_buffer *b, fm_buffer *out, int n) { if (!a || !b || !out) return; fm_moe_silu_mul_at(fm_bufref_at(a, 0), fm_bufref_at(b, 0), fm_bufref_at(out, 0), n); } void fm_moe_shexp_dot_gate_at(fm_bufref gate_w, fm_bufref x, fm_bufref shexp, int n_tokens, int n_embd) { if (!fm_metal_is_initialized() || !gate_w.buffer || !x.buffer || !shexp.buffer) return; if (n_tokens <= 0 || n_embd <= 0) return; fm_bufref refs[] = { gate_w, x, shexp }; uint32_t consts[] = { (uint32_t)n_tokens, (uint32_t)n_embd }; fm_dispatch_threadgroups_refs(@"fmmoe_shexp_dot_gate", NULL, refs, 3, consts, 2, n_tokens, FM_MOE_REDUCE_THREADS); } void fm_moe_shexp_dot_gate(fm_buffer *gate_w, fm_buffer *x, fm_buffer *shexp, int n_tokens, int n_embd) { if (!gate_w || !x || !shexp) return; fm_moe_shexp_dot_gate_at(fm_bufref_at(gate_w, 0), fm_bufref_at(x, 0), fm_bufref_at(shexp, 0), n_tokens, n_embd); } /* ================================================================ * MSL * * Split by kernel group so no one function grows past a screenful; the parts * are concatenated in order into the family's fragment. * ================================================================ */ /* Function constants and the threadgroup tree reductions. */ static NSString *fm_moe_msl_common(void) { return @ "/* ============================================================\n" " * Mixture of experts\n" " *\n" " * Fast math makes results depend on lane ownership, visit order and\n" " * reduction primitives, so every partition below is fixed. Trees are\n" " * halving trees over threadgroup memory with a barrier per level.\n" " * ============================================================ */\n" "\n" "constant int fmmoe_down_topk_fc [[function_constant(240)]];\n" "constant int fmmoe_down_topk = is_function_constant_defined(fmmoe_down_topk_fc)\n" " ? fmmoe_down_topk_fc : 8;\n" "constant int fmmoe_down_halves_fc [[function_constant(241)]];\n" "constant int fmmoe_down_halves = is_function_constant_defined(fmmoe_down_halves_fc)\n" " ? fmmoe_down_halves_fc : 4;\n" "constant int fmmoe_gateup_halves_fc [[function_constant(242)]];\n" "constant int fmmoe_gateup_halves = is_function_constant_defined(fmmoe_gateup_halves_fc)\n" " ? fmmoe_gateup_halves_fc : 16;\n" "\n" "/* Halving-tree sum; the caller has stored and barriered. Total in buf[0]. */\n" "static inline void fmmoe_tree_sum_float(threadgroup float* buf, uint tid, uint T) {\n" " for (uint s = T / 2u; s > 0u; s >>= 1) {\n" " if (tid < s) buf[tid] = buf[tid] + buf[tid + s];\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " }\n" "}\n" "\n" "static inline void fmmoe_tree_sum_int(threadgroup int* buf, uint tid, uint T) {\n" " for (uint s = T / 2u; s > 0u; s >>= 1) {\n" " if (tid < s) buf[tid] = buf[tid] + buf[tid + s];\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " }\n" "}\n" "\n" "static inline void fmmoe_tree_max_float(threadgroup float* buf, uint tid, uint T) {\n" " for (uint s = T / 2u; s > 0u; s >>= 1) {\n" " if (tid < s) buf[tid] = max(buf[tid], buf[tid + s]);\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " }\n" "}\n" "\n" "/* Best (value, index) into slot 0: larger value wins; equal values go to\n" " * the lower valid index. */\n" "static inline void fmmoe_tree_argmax(threadgroup float* sval, threadgroup int* sidx,\n" " uint tid, uint T) {\n" " for (uint s = T / 2u; s > 0u; s >>= 1) {\n" " if (tid < s) {\n" " float va = sval[tid];\n" " int ia = sidx[tid];\n" " float vb = sval[tid + s];\n" " int ib = sidx[tid + s];\n" " bool take = vb > va || (vb == va && ib >= 0 && (ia < 0 || ib < ia));\n" " sval[tid] = take ? vb : va;\n" " sidx[tid] = take ? ib : ia;\n" " }\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " }\n" "}\n" "\n" ; } /* Router: softmax, then topk rounds of argmax, then renormalize. */ static NSString *fm_moe_msl_router(void) { return @ "kernel void fmmoe_router_topk(\n" " device float* logits [[buffer(0)]],\n" " device int* top_idx [[buffer(1)]],\n" " device float* top_wt [[buffer(2)]],\n" " constant uint& n_tokens [[buffer(3)]],\n" " constant uint& n_expert [[buffer(4)]],\n" " constant uint& topk [[buffer(5)]],\n" " constant uint& scale_bits [[buffer(6)]],\n" " uint tok [[threadgroup_position_in_grid]],\n" " uint tid [[thread_position_in_threadgroup]],\n" " uint T [[threads_per_threadgroup]]\n" ") {\n" " threadgroup float sval[256];\n" " threadgroup int sidx[256];\n" " threadgroup float swt_sum;\n" " if (tok >= n_tokens) return;\n" " device float* lg = logits + (ulong)tok * n_expert;\n" " device int* ti = top_idx + (ulong)tok * topk;\n" " device float* tw = top_wt + (ulong)tok * topk;\n" "\n" " float m = -3.0e38f;\n" " for (uint e = tid; e < n_expert; e += T) m = max(m, lg[e]);\n" " sval[tid] = m;\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " fmmoe_tree_max_float(sval, tid, T);\n" " float row_max = sval[0];\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" "\n" " float sum = 0.0f;\n" " for (uint e = tid; e < n_expert; e += T) {\n" " float v = exp(lg[e] - row_max);\n" " lg[e] = v;\n" " sum += v;\n" " }\n" " sval[tid] = sum;\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " fmmoe_tree_sum_float(sval, tid, T);\n" " float inv_sum = 1.0f / sval[0];\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" "\n" " for (uint e = tid; e < n_expert; e += T) lg[e] = lg[e] * inv_sum;\n" " if (tid == 0u) swt_sum = 0.0f;\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" "\n" " for (uint sel = 0; sel < topk; sel++) {\n" " float best_val = -1.0f;\n" " int best = -1;\n" " for (uint e = tid; e < n_expert; e += T) {\n" " float p = lg[e];\n" " if (p > best_val) {\n" " best_val = p;\n" " best = int(e);\n" " }\n" " }\n" " sval[tid] = best_val;\n" " sidx[tid] = best;\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " fmmoe_tree_argmax(sval, sidx, tid, T);\n" " if (tid == 0u) {\n" " int winner = sidx[0];\n" " float winner_p = sval[0];\n" " ti[sel] = winner;\n" " tw[sel] = winner_p;\n" " swt_sum += winner_p;\n" " if (winner >= 0) lg[winner] = -1.0f;\n" " }\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " }\n" "\n" " if (tid == 0u) {\n" " float inv_wt = (swt_sum > 0.0f) ? 1.0f / swt_sum : 0.0f;\n" " float scale = as_type(scale_bits) * inv_wt;\n" " for (uint s = 0; s < topk; s++) tw[s] = tw[s] * scale;\n" " }\n" "}\n" "\n" ; } /* Expert sort: per-expert count, then offsets + ordered scatter. */ static NSString *fm_moe_msl_sort(void) { return @ "/* 1 + the slot where `expert` appears in one token's routing, 0 if absent\n" " * (a repeated expert sums its slots). */\n" "static inline int fmmoe_slot_plus_one(device const int* route, uint topk, int expert) {\n" " int sp1 = 0;\n" " for (uint s = 0; s < topk; s++) sp1 += int(route[s] == expert) * int(s + 1u);\n" " return sp1;\n" "}\n" "\n" "kernel void fmmoe_expert_count(\n" " device const int* top_idx [[buffer(0)]],\n" " device int* expert_counts [[buffer(1)]],\n" " constant uint& n_tokens [[buffer(2)]],\n" " constant uint& topk [[buffer(3)]],\n" " uint tg [[threadgroup_position_in_grid]],\n" " uint tid [[thread_position_in_threadgroup]],\n" " uint T [[threads_per_threadgroup]]\n" ") {\n" " threadgroup int scount[256];\n" " int expert = int(tg);\n" " int count = 0;\n" " for (uint t = tid; t < n_tokens; t += T)\n" " count += int(fmmoe_slot_plus_one(top_idx + (ulong)t * topk, topk, expert) != 0);\n" " scount[tid] = count;\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " fmmoe_tree_sum_int(scount, tid, T);\n" " if (tid == 0u) expert_counts[expert] = scount[0];\n" "}\n" "\n" "kernel void fmmoe_expert_scatter(\n" " device const int* top_idx [[buffer(0)]],\n" " device const int* expert_counts [[buffer(1)]],\n" " device int* expert_offsets [[buffer(2)]],\n" " device int* sorted_token [[buffer(3)]],\n" " device int* sorted_slot [[buffer(4)]],\n" " device int* sorted_expert [[buffer(5)]],\n" " constant uint& n_tokens [[buffer(6)]],\n" " constant uint& n_expert [[buffer(7)]],\n" " constant uint& topk [[buffer(8)]],\n" " uint tg [[threadgroup_position_in_grid]],\n" " uint tid [[thread_position_in_threadgroup]],\n" " uint T [[threads_per_threadgroup]]\n" ") {\n" " threadgroup int sred[256];\n" " threadgroup int sflag[256];\n" " int expert = int(tg);\n" "\n" " int partial = 0;\n" " for (uint j = tid; j < n_expert; j += T)\n" " partial += (int(j) < expert) ? expert_counts[j] : 0;\n" " sred[tid] = partial;\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " fmmoe_tree_sum_int(sred, tid, T);\n" " int base = sred[0];\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" "\n" " if (tid == 0u) {\n" " expert_offsets[expert] = base;\n" " if (tg == n_expert - 1u) expert_offsets[n_expert] = base + expert_counts[expert];\n" " }\n" "\n" " int running = 0;\n" " for (uint t0 = 0; t0 < n_tokens; t0 += T) {\n" " uint t = t0 + tid;\n" " int sp1 = (t < n_tokens) ? fmmoe_slot_plus_one(top_idx + (ulong)t * topk, topk,\n" " expert) : 0;\n" " sflag[tid] = int(sp1 != 0);\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " for (uint off = 1u; off < T; off <<= 1) {\n" " int add = (tid >= off) ? sflag[tid - off] : 0;\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " sflag[tid] += add;\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " }\n" " int inclusive = sflag[tid];\n" " int tile_total = sflag[T - 1u];\n" " if (sp1 != 0) {\n" " int pos = base + running + inclusive - 1;\n" " sorted_token[pos] = int(t);\n" " sorted_slot[pos] = sp1 - 1;\n" " sorted_expert[pos] = expert;\n" " }\n" " running += tile_total;\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " }\n" "}\n" "\n" ; } /* Half-super-block dots: the grouped kernels' lane unit. */ static NSString *fm_moe_msl_half_dots(void) { return @ "/* Six-bit Q3_K scale of sub-block s (i = s & 3, quarter q = s >> 2). */\n" "static inline int fmmoe_q3k_six_bit_scale(device const uchar* sc, uint s) {\n" " uint i = s & 3u;\n" " uint q = s >> 2;\n" " int lo_byte = int(sc[i + 4u * (q & 1u)]);\n" " int lo = (q < 2u) ? (lo_byte & 15) : (lo_byte >> 4);\n" " int hi = (int(sc[8u + i]) >> (2u * q)) & 3;\n" " return lo | (hi << 4);\n" "}\n" "\n" "/* Half u (128 weights) of a Q3_K block against xb, the block's 256\n" " * activations. One chain over level v, group g, weight l. */\n" "static inline float fmmoe_q3k_half_dot(device const uchar* blk, uint u,\n" " device const float* xb) {\n" " float d = float(*(device const half*)(blk + 108));\n" " float acc = 0.0f;\n" " for (uint v = 0; v < 4u; v++) {\n" " for (uint g = 0; g < 2u; g++) {\n" " float dl = d * float(fmmoe_q3k_six_bit_scale(blk + 96, 8u * u + 2u * v + g) - 32);\n" " device const uchar* qs = blk + 32u + 32u * u + 16u * g;\n" " device const uchar* hm = blk + 16u * g;\n" " device const float* xs = xb + 128u * u + 32u * v + 16u * g;\n" " uint shift = 2u * v;\n" " uint mask = 1u << (4u * u + v);\n" " for (uint l = 0; l < 16u; l++) {\n" " int q2 = int((uint(qs[l]) >> shift) & 3u);\n" " int quant = (uint(hm[l]) & mask) ? q2 : q2 - 4;\n" " acc += dl * float(quant) * xs[l];\n" " }\n" " }\n" " }\n" " return acc;\n" "}\n" "\n" "/* Six-bit Q4_K (scale, min) pair j from the 12 packed bytes sr. */\n" "static inline void fmmoe_q4k_scale_min(device const uchar* sr, uint j,\n" " thread int& sc, thread int& mn) {\n" " if (j < 4u) {\n" " sc = int(sr[j]) & 63;\n" " mn = int(sr[j + 4u]) & 63;\n" " } else {\n" " sc = (int(sr[j + 4u]) & 15) | ((int(sr[j - 4u]) >> 6) << 4);\n" " mn = (int(sr[j + 4u]) >> 4) | ((int(sr[j]) >> 6) << 4);\n" " }\n" "}\n" "\n" "/* Half u of a Q4_K block: groups 2u and 2u + 1, each 32 bytes whose low\n" " * nibbles are weights 64p + l and high nibbles 64p + 32 + l. */\n" "static inline float fmmoe_q4k_half_dot(device const uchar* blk, uint u,\n" " device const float* xb) {\n" " float d = float(*(device const half*)(blk));\n" " float dmin = float(*(device const half*)(blk + 2));\n" " device const uchar* sr = blk + 4;\n" " device const uchar* qs = blk + 16;\n" " float acc = 0.0f;\n" " for (uint p = 2u * u; p < 2u * u + 2u; p++) {\n" " int sc_lo, mn_lo, sc_hi, mn_hi;\n" " fmmoe_q4k_scale_min(sr, 2u * p, sc_lo, mn_lo);\n" " fmmoe_q4k_scale_min(sr, 2u * p + 1u, sc_hi, mn_hi);\n" " float d1 = d * float(sc_lo);\n" " float m1 = dmin * float(mn_lo);\n" " float d2 = d * float(sc_hi);\n" " float m2 = dmin * float(mn_hi);\n" " for (uint l = 0; l < 32u; l++) {\n" " uchar q = qs[32u * p + l];\n" " acc += (d1 * float(q & 15) - m1) * xb[64u * p + l];\n" " acc += (d2 * float(q >> 4) - m2) * xb[64u * p + 32u + l];\n" " }\n" " }\n" " return acc;\n" "}\n" "\n" ; } /* Grouped gate+up (slot-major and sorted) and grouped down (Q3_K, Q4_K). */ static NSString *fm_moe_msl_grouped(void) { return @ "/* One SIMD group per (request, row). Lane units are halves: the first HV\n" " * are gate halves, the rest up halves; both chains update every unit. */\n" "static inline void fmmoe_gateup_row(device const uchar* gate_all,\n" " device const uchar* up_all,\n" " device const float* x, device float* hid,\n" " uint request, uint token, int expert, uint row,\n" " uint lane, uint ff, uint cols, uint row_stride) {\n" " if (expert < 0) {\n" " if (lane == 0u) hid[(ulong)request * ff + row] = 0.0f;\n" " return;\n" " }\n" " ulong row_off = (ulong)expert * ((ulong)ff * row_stride) + (ulong)row * row_stride;\n" " device const uchar* gate_row = gate_all + row_off;\n" " device const uchar* up_row = up_all + row_off;\n" " device const float* xt = x + (ulong)token * cols;\n" " const uint HV = uint(fmmoe_gateup_halves);\n" "\n" " float ga = 0.0f;\n" " float ua = 0.0f;\n" " for (uint w = lane; w < 2u * HV; w += 32u) {\n" " bool is_gate = w < HV;\n" " uint h = is_gate ? w : w - HV;\n" " device const uchar* row_base = is_gate ? gate_row : up_row;\n" " float v = fmmoe_q3k_half_dot(row_base + (ulong)(h >> 1) * 110u, h & 1u,\n" " xt + (ulong)(h >> 1) * 256u);\n" " ga += is_gate ? v : 0.0f;\n" " ua += is_gate ? 0.0f : v;\n" " }\n" " float g = simd_sum(ga);\n" " float u = simd_sum(ua);\n" " if (lane == 0u) hid[(ulong)request * ff + row] = (g / (1.0f + exp(-g))) * u;\n" "}\n" "\n" "kernel void fmmoe_gateup_q3k(\n" " device const uchar* gate_all [[buffer(0)]],\n" " device const uchar* up_all [[buffer(1)]],\n" " device const float* x [[buffer(2)]],\n" " device const int* top_idx [[buffer(3)]],\n" " device float* hid [[buffer(4)]],\n" " constant uint& n_requests [[buffer(5)]],\n" " constant uint& topk [[buffer(6)]],\n" " constant uint& ff [[buffer(7)]],\n" " constant uint& cols [[buffer(8)]],\n" " constant uint& row_stride [[buffer(9)]],\n" " uint tg [[threadgroup_position_in_grid]],\n" " uint lane [[thread_position_in_threadgroup]]\n" ") {\n" " uint entry = tg / ff;\n" " if (entry >= n_requests) return;\n" " uint row = tg - entry * ff;\n" " uint request = entry;\n" " fmmoe_gateup_row(gate_all, up_all, x, hid, request, request / topk,\n" " top_idx[request], row, lane, ff, cols, row_stride);\n" "}\n" "\n" "kernel void fmmoe_gateup_q3k_sorted(\n" " device const uchar* gate_all [[buffer(0)]],\n" " device const uchar* up_all [[buffer(1)]],\n" " device const float* x [[buffer(2)]],\n" " device const int* sorted_token [[buffer(3)]],\n" " device const int* sorted_slot [[buffer(4)]],\n" " device const int* sorted_expert [[buffer(5)]],\n" " device float* hid [[buffer(6)]],\n" " constant uint& n_requests [[buffer(7)]],\n" " constant uint& topk [[buffer(8)]],\n" " constant uint& ff [[buffer(9)]],\n" " constant uint& cols [[buffer(10)]],\n" " constant uint& row_stride [[buffer(11)]],\n" " uint tg [[threadgroup_position_in_grid]],\n" " uint lane [[thread_position_in_threadgroup]]\n" ") {\n" " uint entry = tg / ff;\n" " if (entry >= n_requests) return;\n" " uint row = tg - entry * ff;\n" " int expert = sorted_expert[entry];\n" " uint token = uint(sorted_token[entry]);\n" " uint request = token * topk + uint(sorted_slot[entry]);\n" " fmmoe_gateup_row(gate_all, up_all, x, hid, request, token, expert, row, lane,\n" " ff, cols, row_stride);\n" "}\n" "\n" "/* One SIMD group per (token, row). Lane units are (slot, half) pairs across\n" " * all K slots; each unit's dot is weighted by its slot's router weight\n" " * before the single simd_sum. */\n" "template \n" "static inline void fmmoe_down_row(device const uchar* down_all, device const float* hid,\n" " device const int* top_idx, device const float* top_wt,\n" " device float* out, uint n_tokens, uint rows, uint ff,\n" " uint row_stride, uint tg, uint lane) {\n" " uint token = tg / rows;\n" " if (token >= n_tokens) return;\n" " uint row = tg - token * rows;\n" " const uint K = uint(fmmoe_down_topk);\n" " const uint HV = uint(fmmoe_down_halves);\n" " const uint B = IS_Q4K ? 144u : 110u;\n" "\n" " float acc = 0.0f;\n" " for (uint w = lane; w < K * HV; w += 32u) {\n" " uint slot = w / HV;\n" " uint h = w - slot * HV;\n" " uint request = token * K + slot;\n" " int expert = top_idx[request];\n" " if (expert < 0) continue;\n" " device const uchar* blk = down_all + (ulong)expert * ((ulong)rows * row_stride)\n" " + (ulong)row * row_stride + (ulong)(h >> 1) * B;\n" " device const float* xb = hid + (ulong)request * ff + (ulong)(h >> 1) * 256u;\n" " float dot = IS_Q4K ? fmmoe_q4k_half_dot(blk, h & 1u, xb)\n" " : fmmoe_q3k_half_dot(blk, h & 1u, xb);\n" " acc += top_wt[request] * dot;\n" " }\n" " float row_sum = simd_sum(acc);\n" " if (lane == 0u) out[(ulong)token * rows + row] = row_sum;\n" "}\n" "\n" "kernel void fmmoe_down_q3k(\n" " device const uchar* down_all [[buffer(0)]],\n" " device const float* hid [[buffer(1)]],\n" " device const int* top_idx [[buffer(2)]],\n" " device const float* top_wt [[buffer(3)]],\n" " device float* out [[buffer(4)]],\n" " constant uint& n_tokens [[buffer(5)]],\n" " constant uint& rows [[buffer(6)]],\n" " constant uint& ff [[buffer(7)]],\n" " constant uint& row_stride [[buffer(8)]],\n" " uint tg [[threadgroup_position_in_grid]],\n" " uint lane [[thread_position_in_threadgroup]]\n" ") {\n" " fmmoe_down_row(down_all, hid, top_idx, top_wt, out, n_tokens, rows, ff,\n" " row_stride, tg, lane);\n" "}\n" "\n" "kernel void fmmoe_down_q4k(\n" " device const uchar* down_all [[buffer(0)]],\n" " device const float* hid [[buffer(1)]],\n" " device const int* top_idx [[buffer(2)]],\n" " device const float* top_wt [[buffer(3)]],\n" " device float* out [[buffer(4)]],\n" " constant uint& n_tokens [[buffer(5)]],\n" " constant uint& rows [[buffer(6)]],\n" " constant uint& ff [[buffer(7)]],\n" " constant uint& row_stride [[buffer(8)]],\n" " uint tg [[threadgroup_position_in_grid]],\n" " uint lane [[thread_position_in_threadgroup]]\n" ") {\n" " fmmoe_down_row(down_all, hid, top_idx, top_wt, out, n_tokens, rows, ff,\n" " row_stride, tg, lane);\n" "}\n" "\n" ; } /* SiLU-multiply and the shared-expert sigmoid gate. */ static NSString *fm_moe_msl_shared(void) { return @ "kernel void fmmoe_silu_mul(\n" " device const float* a [[buffer(0)]],\n" " device const float* b [[buffer(1)]],\n" " device float* out [[buffer(2)]],\n" " constant uint& n [[buffer(3)]],\n" " uint i [[thread_position_in_grid]]\n" ") {\n" " if (i >= n) return;\n" " float v = a[i];\n" " out[i] = (v / (1.0f + exp(-v))) * b[i];\n" "}\n" "\n" "kernel void fmmoe_shexp_dot_gate(\n" " device const float* gate_w [[buffer(0)]],\n" " device const float* x [[buffer(1)]],\n" " device float* shexp [[buffer(2)]],\n" " constant uint& n_tokens [[buffer(3)]],\n" " constant uint& n_embd [[buffer(4)]],\n" " uint tok [[threadgroup_position_in_grid]],\n" " uint tid [[thread_position_in_threadgroup]],\n" " uint T [[threads_per_threadgroup]]\n" ") {\n" " threadgroup float sdot[256];\n" " if (tok >= n_tokens) return;\n" " device const float* xt = x + (ulong)tok * n_embd;\n" " float partial = 0.0f;\n" " for (uint i = tid; i < n_embd; i += T) partial += gate_w[i] * xt[i];\n" " sdot[tid] = partial;\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " fmmoe_tree_sum_float(sdot, tid, T);\n" " float sg = 1.0f / (1.0f + exp(-sdot[0]));\n" " device float* st = shexp + (ulong)tok * n_embd;\n" " for (uint i = tid; i < n_embd; i += T) st[i] = st[i] * sg;\n" "}\n" "\n" ; } NSString *fm_kernel_msl_moe(void) { NSMutableString *source = [NSMutableString string]; [source appendString:fm_moe_msl_common()]; [source appendString:fm_moe_msl_router()]; [source appendString:fm_moe_msl_sort()]; [source appendString:fm_moe_msl_half_dots()]; [source appendString:fm_moe_msl_grouped()]; [source appendString:fm_moe_msl_shared()]; return source; }