spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Mixture-of-experts kernels. All operate on router-sized rows (E experts,4// E <= 64 in practice), so every kernel is one thread per row with a plain5// loop — no threadgroup reductions needed. row_scale kernels are flat6// elementwise over [N, C] activations.7#include <metal_stdlib>8using namespace metal;910// Generic row softmax backward: dx += p ∘ (dout − dot(dout, p)), row = last11// dim. Thread-per-row; meant for small C (the router's E columns).12kernel void softmax_bwd_f32(device const float* P    [[buffer(0)]],13                            device const float* DOUT [[buffer(1)]],14                            device float*       DX   [[buffer(2)]],15                            constant uint&      C    [[buffer(3)]],16                            uint row [[thread_position_in_grid]]) {17    device const float* prob = P + ulong(row) * C;18    device const float* dout = DOUT + ulong(row) * C;19    device float* dx = DX + ulong(row) * C;20    float dot = 0.0f;21    for (uint j = 0; j < C; ++j) dot += dout[j] * prob[j];22    for (uint j = 0; j < C; ++j) dx[j] += prob[j] * (dout[j] - dot);23}2425// Top-k gating with a selection bias (DeepSeek-V3 "noaux"): the k experts26// are chosen by score+BIAS, but gate VALUES come from the biasless score —27// the bias only steers routing. norm=1 renormalizes kept gates to sum 128// (classic top-k softmax); norm=0 keeps raw scores (sigmoid routing, V3).29// Ties broken by lower index, which keeps forward and backward selections30// identical. p = (E, k, norm)31struct TopkParams { uint E, K, norm; };3233static inline void topk_select_biased(device const float* in,34                                      device const float* bias, uint E, uint K,35                                      thread bool* kept, thread float* S) {36    for (uint j = 0; j < E; ++j) kept[j] = false;37    float acc = 0.0f;38    for (uint sel = 0; sel < K; ++sel) {39        float best = -FLT_MAX;40        uint arg = 0;41        for (uint j = 0; j < E; ++j) {42            const float v = in[j] + bias[j];43            if (!kept[j] && v > best) { best = v; arg = j; }44        }45        kept[arg] = true;46        acc += in[arg]; // gate mass is biasless47    }48    *S = acc;49}5051kernel void topk_renorm_f32(device const float* P    [[buffer(0)]],52                            device const float* BIAS [[buffer(1)]],53                            device float*       OUT  [[buffer(2)]],54                            constant TopkParams& p   [[buffer(3)]],55                            uint row [[thread_position_in_grid]]) {56    device const float* in = P + ulong(row) * p.E;57    device float* out = OUT + ulong(row) * p.E;5859    bool kept[64];60    float S;61    topk_select_biased(in, BIAS, p.E, p.K, kept, &S);62    const float inv = p.norm ? 1.0f / max(S, 1e-12f) : 1.0f;63    for (uint j = 0; j < p.E; ++j) out[j] = kept[j] ? in[j] * inv : 0.0f;64}6566// Backward: with norm, g_i = p_i / S over kept entries so67// dp_i += (dg_i − Σ_j dg_j g_j) / S; without norm, g_i = p_i so dp_i += dg_i.68// The kept set (incl. bias) is recomputed with the same tie-breaking.69kernel void topk_renorm_bwd_f32(device const float* P    [[buffer(0)]],70                                device const float* BIAS [[buffer(1)]],71                                device const float* DOUT [[buffer(2)]],72                                device float*       DP   [[buffer(3)]],73                                constant TopkParams& p   [[buffer(4)]],74                                uint row [[thread_position_in_grid]]) {75    device const float* in = P + ulong(row) * p.E;76    device const float* dout = DOUT + ulong(row) * p.E;77    device float* dp = DP + ulong(row) * p.E;7879    bool kept[64];80    float S;81    topk_select_biased(in, BIAS, p.E, p.K, kept, &S);82    if (!p.norm) {83        for (uint j = 0; j < p.E; ++j)84            if (kept[j]) dp[j] += dout[j];85        return;86    }87    const float inv = 1.0f / max(S, 1e-12f);88    float dot = 0.0f;89    for (uint j = 0; j < p.E; ++j)90        if (kept[j]) dot += dout[j] * in[j] * inv;91    for (uint j = 0; j < p.E; ++j)92        if (kept[j]) dp[j] += (dout[j] - dot) * inv;93}9495// counts[e] += number of rows whose gate for expert e is nonzero — the load96// statistic driving the noaux bias update. One thread per expert.97kernel void expert_counts_f32(device const float* G      [[buffer(0)]],98                              device float*       COUNTS [[buffer(1)]],99                              constant uint2&     p      [[buffer(2)]], // (E, N)100                              uint e [[thread_position_in_grid]]) {101    float acc = 0.0f;102    for (uint i = 0; i < p.y; ++i)103        if (G[ulong(i) * p.x + e] != 0.0f) acc += 1.0f;104    COUNTS[e] += acc;105}106107// 12-byte layout matching the host-side struct (uint3 would pad to 16).108struct RowScaleParams { uint C, E, e; };109110// out[i, c] = x[i, c] * G[i, e] — scale each row of X by one gate column.111// Flat over N*C.112kernel void row_scale_f32(device const float*       X   [[buffer(0)]],113                          device const float*       G   [[buffer(1)]],114                          device float*             OUT [[buffer(2)]],115                          constant RowScaleParams&  p   [[buffer(3)]],116                          uint gid [[thread_position_in_grid]]) {117    const uint i = gid / p.C;118    OUT[gid] = X[gid] * G[ulong(i) * p.E + p.e];119}120121// dst[i, c] += x[i, c] * G[i, e] — accumulating variant (dx of row_scale).122kernel void row_scale_acc_f32(device const float*       X   [[buffer(0)]],123                              device const float*       G   [[buffer(1)]],124                              device float*             DST [[buffer(2)]],125                              constant RowScaleParams&  p   [[buffer(3)]],126                              uint gid [[thread_position_in_grid]]) {127    const uint i = gid / p.C;128    DST[gid] = fma(X[gid], G[ulong(i) * p.E + p.e], DST[gid]);129}130131// dG[i, e] += dot(dout[i, :], x[i, :]) — gate gradient, one thread per row.132kernel void row_scale_gate_bwd_f32(device const float*       DOUT [[buffer(0)]],133                                   device const float*       X    [[buffer(1)]],134                                   device float*             DG   [[buffer(2)]],135                                   constant RowScaleParams&  p    [[buffer(3)]],136                                   uint row [[thread_position_in_grid]]) {137    device const float* dout = DOUT + ulong(row) * p.C;138    device const float* x = X + ulong(row) * p.C;139    float acc = 0.0f;140    for (uint j = 0; j < p.C; ++j) acc += dout[j] * x[j];141    DG[ulong(row) * p.E + p.e] += acc;142}143