// metal/kernels/router.metal - the router GEMV and moe_route in ONE dispatch (after quant/units/moe/gemv.metal): the
// logits rows as gemv's (LPR = 32 lanes a row, KS = 1: bitwise gemv_<F>_32_1_*_1's), 16 rows a threadgroup of 512
// threads; every threadgroup bumps a device counter after its rows (device-scope fences around it), and the LAST one
// runs moe_route's body (the softmax / top-10 / weights, then the plan) with its 512 threads (= n_expert) and resets
// the counter. No threadgroup waits for another. ids / w / plan are bitwise moe_route's.
// Buffers: W, xp, s16, logits, args (GemvArgs: n_out = n_expert), ids, w, plan, MoeArgs, counter (one uint, zero).
#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
static inline void route_body(device const float* logits, device int* ids, device float* w, device int* plan,
constant MoeArgs& a, device const uint* smap, ushort tid, ushort nth, ushort sg, ushort lane,
threadgroup int* s_grp, threadgroup int* s_ent) {
if (sg < a.n_tok) {
device const float* l = logits + sg * 512;
float v[16];
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
v[i] = idx < int(a.n_expert) ? l[idx] : -INFINITY;
}
float mx = -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));
float sum = 0.0f;
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];
}
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 >= 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;
ushort scaled_mask = 0;
float selected = 0.0f, ssum = 0.0f;
for (int rank = 0; rank < 10; ++rank) {
if (r1_scale > 1.0f && r_prot > 1) {
if (rank == 1 && smap != nullptr) {
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
v[i] *= r1_scale;
scaled_mask |= ushort(1u << i);
}
}
} else if (rank == r_prot && moce_scale > 1.0f && smap != nullptr) {
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 && smap != nullptr) {
for (int i = 0; i < 16; ++i) {
const int idx = lane + 32 * i;
if (idx < int(a.n_expert) && v[i] > 0.0f && smap[idx] != 0xFFFFFFFFu) {
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;
const int min_scaled_rank = (r1_scale > 1.0f && r_prot > 1) ? 1 : r_prot;
if (rank >= min_scaled_rank && (moce_scale > 1.0f || r1_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) {
// without a tail bonus, the rank-1 scale remains on unselected residents.
const bool r1_scaled = r1_scale > 1.0f && r_prot > 1 &&
(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;
}
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);
}
template <int F>
kernel void router_route(device const uchar* W [[buffer(0)]], device const float* xp [[buffer(1)]],
device const float* s16 [[buffer(2)]], device float* logits [[buffer(3)]],
constant GemvArgs& ga [[buffer(4)]], device int* ids [[buffer(5)]], device float* w [[buffer(6)]],
device int* plan [[buffer(7)]], constant MoeArgs& a [[buffer(8)]],
device atomic_uint* counter [[buffer(9)]], device const uint* smap [[buffer(10)]],
uint tgx [[threadgroup_position_in_grid]], uint ngx [[threadgroups_per_grid]],
ushort tid [[thread_index_in_threadgroup]], ushort nth [[threads_per_threadgroup]],
ushort sg [[simdgroup_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) {
constexpr int UNIT = fmt<F>::UNIT;
threadgroup uint tab[fmt<F>::TAB];
threadgroup int s_grp[64], s_ent[64];
threadgroup uint last;
load_tab<F>(tab, tid, nth);
threadgroup_barrier(mem_flags::mem_threadgroup);
// the GEMV: simdgroup sg -> row tg.x * 16 + sg, 32 lanes over its units
const int row = int(tgx) * 16 + sg;
const int rr = min(row, int(ga.n_out) - 1);
device const uchar* rp = W + (ulong)rr * ga.row_bytes;
const int U = int(ga.n_in) / UNIT, s16n = int(ga.n_in) / 16;
float acc[8];
#pragma clang loop unroll(full)
for (int t = 0; t < NTOK; ++t) acc[t] = 0.0f;
for (int u = lane; u < U; u += 32) {
float q[UNIT];
float4 sc;
Dec<F>::run(rp, u, q, sc, tab);
#pragma clang loop unroll(full)
for (int t = 0; t < NTOK; ++t) {
float xs[UNIT], xsum[2];
X<F>::load(xp + t * int(ga.x_stride), s16 + t * s16n, u, xs, xsum);
acc[t] += unit_combine<UNIT>(q, sc, xs, xsum);
}
}
#pragma clang loop unroll(full)
for (int t = 0; t < NTOK; ++t) {
const float v = lane_sum<32>(acc[t]);
if (lane == 0 && row < int(ga.n_out)) logits[t * int(ga.y_stride) + row] = v;
}
// the last threadgroup to finish routes
atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);
threadgroup_barrier(mem_flags::mem_device);
if (tid == 0) last = atomic_fetch_add_explicit(counter, 1u, memory_order_relaxed) == ngx - 1 ? 1u : 0u;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (!last) return;
atomic_thread_fence(mem_flags::mem_device, memory_order_seq_cst, thread_scope_device);
route_body(logits, ids, w, plan, a, smap, tid, nth, sg, lane, s_grp, s_ent);
if (tid == 0) atomic_store_explicit(counter, 0u, memory_order_relaxed);
}
template [[host_name("router_route_30")]] kernel decltype(router_route<T_BF16>) router_route<T_BF16>;
template [[host_name("router_route_0")]] kernel decltype(router_route<T_F32>) router_route<T_F32>;
template [[host_name("router_route_8")]] kernel decltype(router_route<T_Q8_0>) router_route<T_Q8_0>;