spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Elementwise kernels. Flat 1-D dispatch via dispatchThreads (non-uniform4// threadgroups — no bounds-check tail), threadgroup size a multiple of 32.5// Purely bandwidth-bound; float4-vectorized variants come with the M56// fusion pass.7#include <metal_stdlib>8using namespace metal;910kernel void add_f32(device const float* a   [[buffer(0)]],11                    device const float* b   [[buffer(1)]],12                    device float*       out [[buffer(2)]],13                    uint gid [[thread_position_in_grid]]) {14    out[gid] = a[gid] + b[gid];15}1617kernel void mul_f32(device const float* a   [[buffer(0)]],18                    device const float* b   [[buffer(1)]],19                    device float*       out [[buffer(2)]],20                    uint gid [[thread_position_in_grid]]) {21    out[gid] = a[gid] * b[gid];22}2324kernel void scale_f32(device const float* a   [[buffer(0)]],25                      device float*       out [[buffer(1)]],26                      constant float&     s   [[buffer(2)]],27                      uint gid [[thread_position_in_grid]]) {28    out[gid] = a[gid] * s;29}3031// out[i] = x[i] + bias[i % C] — row-broadcast bias32kernel void add_bias_f32(device const float* x    [[buffer(0)]],33                         device const float* bias [[buffer(1)]],34                         device float*       out  [[buffer(2)]],35                         constant uint&      C    [[buffer(3)]],36                         uint gid [[thread_position_in_grid]]) {37    out[gid] = x[gid] + bias[gid % C];38}3940kernel void silu_f32(device const float* x   [[buffer(0)]],41                     device float*       out [[buffer(1)]],42                     uint gid [[thread_position_in_grid]]) {43    const float v = x[gid];44    out[gid] = v / (1.0f + exp(-v));45}4647kernel void gelu_f32(device const float* x   [[buffer(0)]],48                     device float*       out [[buffer(1)]],49                     uint gid [[thread_position_in_grid]]) {50    const float v = x[gid];51    const float k = 0.7978845608028654f; // sqrt(2/pi)52    out[gid] = 0.5f * v * (1.0f + precise::tanh(k * (v + 0.044715f * v * v * v)));53}5455kernel void sigmoid_f32(device const float* x   [[buffer(0)]],56                        device float*       out [[buffer(1)]],57                        uint gid [[thread_position_in_grid]]) {58    out[gid] = 1.0f / (1.0f + exp(-x[gid]));59}6061// dx += dout * y * (1 - y), using the forward output y.62kernel void sigmoid_bwd_f32(device const float* y    [[buffer(0)]],63                            device const float* dout [[buffer(1)]],64                            device float*       dx   [[buffer(2)]],65                            uint gid [[thread_position_in_grid]]) {66    const float v = y[gid];67    dx[gid] = fma(dout[gid], v * (1.0f - v), dx[gid]);68}6970// ReLU² (nanoGPT-speedrun lineage): out = max(x, 0)^2.71kernel void relu2_f32(device const float* x   [[buffer(0)]],72                      device float*       out [[buffer(1)]],73                      uint gid [[thread_position_in_grid]]) {74    const float v = max(x[gid], 0.0f);75    out[gid] = v * v;76}7778kernel void relu2_bwd_f32(device const float* x    [[buffer(0)]],79                          device const float* dout [[buffer(1)]],80                          device float*       dx   [[buffer(2)]],81                          uint gid [[thread_position_in_grid]]) {82    dx[gid] = fma(dout[gid], 2.0f * max(x[gid], 0.0f), dx[gid]);83}8485// Gemma-style logit softcap: out = cap * tanh(x / cap).86kernel void softcap_f32(device const float* x   [[buffer(0)]],87                        device float*       out [[buffer(1)]],88                        constant float&     cap [[buffer(2)]],89                        uint gid [[thread_position_in_grid]]) {90    out[gid] = cap * precise::tanh(x[gid] / cap);91}9293// dx += dout * (1 - (y/cap)^2), using the forward output y.94kernel void softcap_bwd_f32(device const float* y    [[buffer(0)]],95                            device const float* dout [[buffer(1)]],96                            device float*       dx   [[buffer(2)]],97                            constant float&     cap  [[buffer(3)]],98                            uint gid [[thread_position_in_grid]]) {99    const float t = y[gid] / cap;100    dx[gid] = fma(dout[gid], 1.0f - t * t, dx[gid]);101}102103// ---- backward / accumulation kernels (all ACCUMULATE into their outputs) ----104105kernel void accum_f32(device float*       dst [[buffer(0)]],106                      device const float* src [[buffer(1)]],107                      uint gid [[thread_position_in_grid]]) {108    dst[gid] += src[gid];109}110111// dst += src * s ; s in a 1-element buffer so it can be produced on-GPU112// (e.g. d(loss) scaling) without a sync.113kernel void axpy_f32(device float*       dst [[buffer(0)]],114                     device const float* src [[buffer(1)]],115                     device const float* s   [[buffer(2)]],116                     uint gid [[thread_position_in_grid]]) {117    dst[gid] = fma(src[gid], s[0], dst[gid]);118}119120kernel void silu_bwd_f32(device const float* x    [[buffer(0)]],121                         device const float* dout [[buffer(1)]],122                         device float*       dx   [[buffer(2)]],123                         uint gid [[thread_position_in_grid]]) {124    const float v = x[gid];125    const float sig = 1.0f / (1.0f + exp(-v));126    dx[gid] = fma(dout[gid], sig * (1.0f + v * (1.0f - sig)), dx[gid]);127}128129kernel void gelu_bwd_f32(device const float* x    [[buffer(0)]],130                         device const float* dout [[buffer(1)]],131                         device float*       dx   [[buffer(2)]],132                         uint gid [[thread_position_in_grid]]) {133    const float v = x[gid];134    const float k = 0.7978845608028654f;135    const float u = k * (v + 0.044715f * v * v * v);136    const float t = precise::tanh(u);137    const float du = k * (1.0f + 3.0f * 0.044715f * v * v);138    dx[gid] = fma(dout[gid], 0.5f * (1.0f + t) + 0.5f * v * (1.0f - t * t) * du, dx[gid]);139}140141// dbias[j] += sum_rows dout[i,j] — one thread per column, strided rows.142// Column-major walk is uncoalesced but this kernel is a tiny fraction of a143// step; revisit in the M5 fusion pass if it ever shows in a trace.144kernel void add_bias_bwd_f32(device const float* dout  [[buffer(0)]],145                             device float*       dbias [[buffer(1)]],146                             constant uint2&     nc    [[buffer(2)]], // (N, C)147                             uint j [[thread_position_in_grid]]) {148    float acc = 0.0f;149    for (uint i = 0; i < nc.x; ++i) acc += dout[ulong(i) * nc.y + j];150    dbias[j] += acc;151}152153// ---- RoPE (interleaved pairs; INV=true applies the inverse rotation and154// accumulates — the backward pass) ----155156constant bool ROPE_INV [[function_constant(2)]];157158struct RopeParams {159    uint T, H, HD;160    uint pos_offset;161};162163// freqs[k] is the per-pair inverse frequency (theta^(-2k/HD), possibly164// rescaled — llama3 rope_scaling, per-layer theta). Host-precomputed.165kernel void rope_f32(device const float* x     [[buffer(0)]],166                     device float*       out   [[buffer(1)]],167                     constant RopeParams& p    [[buffer(2)]],168                     device const float* freqs [[buffer(3)]],169                     uint gid [[thread_position_in_grid]]) {170    // gid indexes (bt, h, k) pairs: one thread per rotated pair171    const uint pairs_per_row = p.H * (p.HD / 2);172    const uint bt = gid / pairs_per_row;173    const uint rem = gid % pairs_per_row;174    const uint h = rem / (p.HD / 2);175    const uint k = rem % (p.HD / 2);176177    const float pos = float(bt % p.T + p.pos_offset);178    const float angle = pos * freqs[k];179    const float c = cos(angle);180    const float s = ROPE_INV ? -sin(angle) : sin(angle);181182    const ulong i0 = ulong(bt) * (p.H * p.HD) + h * p.HD + 2 * k;183    const float x0 = x[i0], x1 = x[i0 + 1];184    if (ROPE_INV) {185        out[i0]     += x0 * c - x1 * s;186        out[i0 + 1] += x0 * s + x1 * c;187    } else {188        out[i0]     = x0 * c - x1 * s;189        out[i0 + 1] = x0 * s + x1 * c;190    }191}192