#include "common.h"
constant bool FC_topk_moe_with_norm [[function_constant(FC_TOPK_MOE + 0)]];
constant int FC_topk_moe_n_expert [[function_constant(FC_TOPK_MOE + 1)]];
constant int FC_topk_moe_top_k [[function_constant(FC_TOPK_MOE + 2)]];
constant int FC_moe_reduce_n_expert_used [[function_constant(FC_MOE_REDUCE + 0)]];
// bitonic sort implementation following the CUDA kernels as reference
typedef void (argsort_t)(
constant ggml_metal_kargs_argsort & args,
device const char * src0,
device int32_t * dst,
threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort3 tpitg[[thread_position_in_threadgroup]],
ushort3 ntg[[threads_per_threadgroup]]);
template<ggml_sort_order order>
kernel void kernel_argsort_f32_i32(
constant ggml_metal_kargs_argsort & args,
device const char * src0,
device int32_t * dst,
threadgroup int32_t * shmem_i32 [[threadgroup(0)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort3 tpitg[[thread_position_in_threadgroup]],
ushort3 ntg[[threads_per_threadgroup]]) {
// bitonic sort
const int col = tpitg[0];
const int ib = tgpig[0] / args.ne01;
const int i00 = ib*ntg.x;
const int i01 = tgpig[0] % args.ne01;
const int i02 = tgpig[1];
const int i03 = tgpig[2];
device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03);
// initialize indices
shmem_i32[col] = i00 + col;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (int k = 2; k <= ntg.x; k *= 2) {
for (int j = k / 2; j > 0; j /= 2) {
int ixj = col ^ j;
if (ixj > col) {
if ((col & k) == 0) {
if (shmem_i32[col] >= args.ne00 ||
(shmem_i32[ixj] < args.ne00 && (order == GGML_SORT_ORDER_ASC ?
src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]] :
src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]]))
) {
SWAP(shmem_i32[col], shmem_i32[ixj]);
}
} else {
if (shmem_i32[ixj] >= args.ne00 ||
(shmem_i32[col] < args.ne00 && (order == GGML_SORT_ORDER_ASC ?
src0_row[shmem_i32[col]] < src0_row[shmem_i32[ixj]] :
src0_row[shmem_i32[col]] > src0_row[shmem_i32[ixj]]))
) {
SWAP(shmem_i32[col], shmem_i32[ixj]);
}
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
}
const int64_t i0 = ib*args.top_k;
// copy the result to dst without the padding
if (i0 + col < args.ne0 && col < args.top_k) {
dst += i0 + args.ne0*i01 + args.ne0*args.ne1*i02 + args.ne0*args.ne1*args.ne2*i03;
dst[col] = shmem_i32[col];
}
}
template [[host_name("kernel_argsort_f32_i32_asc")]] kernel argsort_t kernel_argsort_f32_i32<GGML_SORT_ORDER_ASC>;
template [[host_name("kernel_argsort_f32_i32_desc")]] kernel argsort_t kernel_argsort_f32_i32<GGML_SORT_ORDER_DESC>;
typedef void (argsort_merge_t)(
constant ggml_metal_kargs_argsort_merge & args,
device const char * src0,
device const int32_t * tmp,
device int32_t * dst,
uint3 tgpig[[threadgroup_position_in_grid]],
ushort3 tpitg[[thread_position_in_threadgroup]],
ushort3 ntg[[threads_per_threadgroup]]);
template<ggml_sort_order order>
kernel void kernel_argsort_merge_f32_i32(
constant ggml_metal_kargs_argsort_merge & args,
device const char * src0,
device const int32_t * tmp,
device int32_t * dst,
uint3 tgpig[[threadgroup_position_in_grid]],
ushort3 tpitg[[thread_position_in_threadgroup]],
ushort3 ntg[[threads_per_threadgroup]]) {
const int im = tgpig[0] / args.ne01;
const int i01 = tgpig[0] % args.ne01;
const int i02 = tgpig[1];
const int i03 = tgpig[2];
const int start = im * (2 * args.len);
const int len0 = MIN(args.len, MAX(0, args.ne0 - (int)(start)));
const int len1 = MIN(args.len, MAX(0, args.ne0 - (int)(start + args.len)));
const int total = len0 + len1;
device const int32_t * tmp0 = tmp + start
+ i01*args.ne0
+ i02*args.ne0*args.ne01
+ i03*args.ne0*args.ne01*args.ne02;
device const int32_t * tmp1 = tmp0 + args.len;
dst += start
+ i01*args.top_k
+ i02*args.top_k*args.ne01
+ i03*args.top_k*args.ne01*args.ne02;
device const float * src0_row = (device const float *)(src0
+ args.nb01*i01
+ args.nb02*i02
+ args.nb03*i03);
if (total == 0) {
return;
}
const int chunk = (total + ntg.x - 1) / ntg.x;
const int k0 = tpitg.x * chunk;
const int k1 = MIN(MIN(k0 + chunk, total), args.top_k);
if (k0 >= args.top_k) {
return;
}
if (k0 >= total) {
return;
}
int low = k0 > len1 ? k0 - len1 : 0;
int high = MIN(k0, len0);
// binary-search partition (i, j) such that i + j = k
while (low < high) {
const int mid = (low + high) >> 1;
const int32_t idx0 = tmp0[mid];
const int32_t idx1 = tmp1[k0 - mid - 1];
const float val0 = src0_row[idx0];
const float val1 = src0_row[idx1];
bool take_left;
if (order == GGML_SORT_ORDER_ASC) {
take_left = (val0 <= val1);
} else {
take_left = (val0 >= val1);
}
if (take_left) {
low = mid + 1;
} else {
high = mid;
}
}
int i = low;
int j = k0 - i;
// keep the merge fronts into registers
int32_t idx0 = 0;
float val0 = 0.0f;
if (i < len0) {
idx0 = tmp0[i];
val0 = src0_row[idx0];
}
int32_t idx1 = 0;
float val1 = 0.0f;
if (j < len1) {
idx1 = tmp1[j];
val1 = src0_row[idx1];
}
for (int k = k0; k < k1; ++k) {
int32_t out_idx;
if (i >= len0) {
while (k < k1) {
dst[k++] = tmp1[j++];
}
break;
} else if (j >= len1) {
while (k < k1) {
dst[k++] = tmp0[i++];
}
break;
} else {
bool take_left;
if (order == GGML_SORT_ORDER_ASC) {
take_left = (val0 <= val1);
} else {
take_left = (val0 >= val1);
}
if (take_left) {
out_idx = idx0;
++i;
if (i < len0) {
idx0 = tmp0[i];
val0 = src0_row[idx0];
}
} else {
out_idx = idx1;
++j;
if (j < len1) {
idx1 = tmp1[j];
val1 = src0_row[idx1];
}
}
}
dst[k] = out_idx;
}
}
template [[host_name("kernel_argsort_merge_f32_i32_asc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_ASC>;
template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_DESC>;
static inline uint ggml_top_k_f2ui(float x) {
uint y = as_type<uint>(x);
if ((y & 0x80000000u) != 0u) {
y ^= 0xFFFFFFFFu; // negative floats: flip all bits
} else {
y |= 0x80000000u; // positive floats: set the sign bit
}
return y;
}
kernel void kernel_top_k_f32_i32(
constant ggml_metal_kargs_top_k & args,
device const char * src0,
device int32_t * dst,
threadgroup atomic_uint * histo [[threadgroup(0)]],
threadgroup uint * sh_bucket [[threadgroup(1)]],
threadgroup uint * sh_above [[threadgroup(2)]],
threadgroup atomic_uint * out_count [[threadgroup(3)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort3 tpitg[[thread_position_in_threadgroup]],
ushort3 ntg[[threads_per_threadgroup]]) {
const uint ncols = args.ne00;
const uint top_k = args.top_k;
const uint i01 = tgpig[0];
const uint i02 = tgpig[1];
const uint i03 = tgpig[2];
device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03);
device int32_t * dst_row = dst + top_k*(i01 + args.ne01*i02 + args.ne01*args.ne02*i03);
const uint tid = tpitg.x;
const uint ntg_x = ntg.x;
uint prefix = 0; // fixed high bits of the threshold key
uint desired = top_k; // count still needed from the candidate range
for (int shift = 24; shift >= 0; shift -= 8) {
for (uint i = tid; i < 256; i += ntg_x) {
atomic_store_explicit(&histo[i], 0u, memory_order_relaxed);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint hi_mask = (shift + 8 >= 32) ? 0u : (0xFFFFFFFFu << uint(shift + 8));
const uint prefix_hi = prefix & hi_mask;
for (uint i = tid; i < ncols; i += ntg_x) {
const uint key = ggml_top_k_f2ui(src0_row[i]);
if ((key & hi_mask) == prefix_hi) {
atomic_fetch_add_explicit(&histo[(key >> uint(shift)) & 0xFFu], 1u, memory_order_relaxed);
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// top-down scan for the bucket holding the k-th value
if (tid == 0) {
uint acc = 0;
uint b = 0;
for (int bb = 255; bb >= 0; --bb) {
const uint c = atomic_load_explicit(&histo[bb], memory_order_relaxed);
if (acc + c >= desired) {
b = uint(bb);
break;
}
acc += c;
}
*sh_bucket = b;
*sh_above = acc;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
prefix |= *sh_bucket << uint(shift);
desired -= *sh_above;
// ensure every thread has consumed sh_bucket/sh_above before the next pass
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tid == 0) {
atomic_store_explicit(out_count, 0u, memory_order_relaxed);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// emit everything above the threshold, then fill the rest from ties
const uint threshold = prefix;
for (uint i = tid; i < ncols; i += ntg_x) {
if (ggml_top_k_f2ui(src0_row[i]) > threshold) {
const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed);
dst_row[pos] = (int32_t) i;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint i = tid; i < ncols; i += ntg_x) {
if (ggml_top_k_f2ui(src0_row[i]) == threshold) {
const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed);
if (pos < top_k) {
dst_row[pos] = (int32_t) i;
}
}
}
}
// fused SOFT_MAX + top-k + GET_ROWS (+ optional norm/scale) for MoE routing.
// One SIMDgroup handles one token row; n_expert is limited to 1024 by the host.
kernel void kernel_topk_moe_f32(
constant ggml_metal_kargs_topk_moe & args,
device const char * src0,
device float * weights,
device int32_t * ids,
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiisg[[thread_index_in_simdgroup]]) {
const int row = (int) tgpig.x;
if (row >= args.ne01) {
return;
}
const int n_expert = FC_topk_moe_n_expert;
const int top_k = FC_topk_moe_top_k;
const int lane = (int) tiisg;
const int n_per_lane = (n_expert + 31) / 32;
device const float * logits_row = (device const float *) (src0 + row * args.nb01);
device float * weights_row = weights + row * top_k;
device int32_t * ids_row = ids + row * (args.nb1_ids / sizeof(int32_t));
float wt[32];
float output_weights[32];
FOR_UNROLL (int i = 0; i < 32; ++i) {
wt[i] = -INFINITY;
output_weights[i] = 0.0f;
}
for (int i = lane; i < n_expert; i += 32) {
const float v = logits_row[i];
wt[i / 32] = isnan(v) ? -FLT_MAX : v;
}
// softmax over the expert logits
float max_val = -INFINITY;
FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
max_val = max(max_val, wt[i]);
}
max_val = simd_max(max_val);
float sum_val = 0.0f;
FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
wt[i] = exp(wt[i] - max_val);
sum_val += wt[i];
}
sum_val = simd_sum(sum_val);
const float inv_sum = 1.0f / sum_val;
FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
wt[i] *= inv_sum;
}
float wt_sum = 0.0f;
for (int k = 0; k < top_k; ++k) {
float best_val = -INFINITY;
int best_expert = -1;
FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
const int expert = lane + i * 32;
if (expert < n_expert && (wt[i] > best_val || (wt[i] == best_val && expert < best_expert))) {
best_val = wt[i];
best_expert = expert;
}
}
FOR_UNROLL (int mask = 16; mask > 0; mask >>= 1) {
const float val = simd_shuffle_xor(best_val, mask);
const int expert = simd_shuffle_xor(best_expert, mask);
if (val > best_val || (val == best_val && expert < best_expert)) {
best_val = val;
best_expert = expert;
}
}
if ((best_expert & 31) == lane) {
wt[best_expert / 32] = -INFINITY;
}
if ((k & 31) == lane) {
output_weights[k / 32] = best_val;
}
if ((best_expert & 31) == lane) {
ids_row[k] = best_expert;
if (FC_topk_moe_with_norm) {
wt_sum += best_val;
}
}
}
if (FC_topk_moe_with_norm) {
wt_sum = simd_sum(wt_sum);
wt_sum = max(wt_sum, args.clamp);
const float inv = 1.0f / wt_sum;
FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
output_weights[i] *= inv;
}
}
FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
const int idx = i * 32 + lane;
if (idx < top_k) {
weights_row[idx] = output_weights[i] * args.scale;
}
}
}
// fused MoE expert weighting + reduction: weighted = sum(experts[e] * weights[e]).
// The host guarantees all tensors are contiguous F32.
kernel void kernel_moe_reduce_f32(
constant ggml_metal_kargs_moe_reduce & args,
device const float * experts,
device const float * weights,
device float * dst,
uint3 tgpig[[threadgroup_position_in_grid]],
ushort3 tpitg[[thread_position_in_threadgroup]],
ushort3 ntg[[threads_per_threadgroup]]) {
const int64_t token = tgpig.x;
const int64_t col = (int64_t) tgpig.y * ntg.x + tpitg.x;
if (token >= args.ne02 || col >= args.ne00) {
return;
}
const int n_expert_used = FC_moe_reduce_n_expert_used;
const int64_t base = token * (int64_t) n_expert_used * args.ne00 + col;
float sum = 0.0f;
FOR_UNROLL (int e = 0; e < n_expert_used; ++e) {
sum += experts[base + e * args.ne00] * weights[token * n_expert_used + e];
}
dst[token * args.ne00 + col] = sum;
}