spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Fused AdamW update (llm.c formulation): one thread per element; the host4// precomputes the bias corrections and folds the gradient-clip scale into5// grad_scale, so the kernel reads the raw gradient once and writes w/m/v.6// eps sits OUTSIDE the sqrt; weight decay is decoupled (never through m/v)7// and the host sets wd = 0 for dim<2 tensors.8//9// sumsq_f32: grid-stride sum of squares of one tensor into partials[tg] —10// used for the global grad-norm (host sums the partials at the sync point).11#include <metal_stdlib>12using namespace metal;1314struct AdamWParams {15    float lr, beta1, beta2, bc1, bc2, eps, wd, grad_scale;16};1718kernel void adamw_f32(device float*       w [[buffer(0)]],19                      device const float* g [[buffer(1)]],20                      device float*       m [[buffer(2)]],21                      device float*       v [[buffer(3)]],22                      constant AdamWParams& p [[buffer(4)]],23                      uint gid [[thread_position_in_grid]]) {24    const float grad = g[gid] * p.grad_scale;25    const float mi = p.beta1 * m[gid] + (1.0f - p.beta1) * grad;26    const float vi = p.beta2 * v[gid] + (1.0f - p.beta2) * grad * grad;27    m[gid] = mi;28    v[gid] = vi;29    const float mhat = mi / p.bc1;30    const float vhat = vi / p.bc2;31    w[gid] -= p.lr * (mhat / (sqrt(vhat) + p.eps) + p.wd * w[gid]);32}3334kernel void sumsq_f32(device const float* x        [[buffer(0)]],35                      device float*       partials [[buffer(1)]],36                      constant uint&      n        [[buffer(2)]],37                      uint gid      [[thread_position_in_grid]],38                      uint grid_sz  [[threads_per_grid]],39                      uint tg_id    [[threadgroup_position_in_grid]],40                      uint lane     [[thread_index_in_simdgroup]],41                      uint simd_idx [[simdgroup_index_in_threadgroup]],42                      uint n_simds  [[simdgroups_per_threadgroup]]) {43    float acc = 0.0f;44    for (uint i = gid; i < n; i += grid_sz) acc = fma(x[i], x[i], acc);4546    threadgroup float scratch[32];47    const float s = simd_sum(acc);48    if (simd_is_first()) scratch[simd_idx] = s;49    threadgroup_barrier(mem_flags::mem_threadgroup);50    const float mine = (lane < n_simds) ? scratch[lane] : 0.0f;51    const float total = simd_sum(mine);52    if (simd_idx == 0 && lane == 0) partials[tg_id] = total;53}5455// Single-threadgroup final sum (also reduces CE per-row losses).56kernel void sum_f32(device const float* x   [[buffer(0)]],57                    device float*       out [[buffer(1)]],58                    constant uint&      n   [[buffer(2)]],59                    constant float&     mul [[buffer(3)]],60                    uint lid      [[thread_index_in_threadgroup]],61                    uint tg_size  [[threads_per_threadgroup]],62                    uint lane     [[thread_index_in_simdgroup]],63                    uint simd_idx [[simdgroup_index_in_threadgroup]],64                    uint n_simds  [[simdgroups_per_threadgroup]]) {65    float acc = 0.0f;66    for (uint i = lid; i < n; i += tg_size) acc += x[i];6768    threadgroup float scratch[32];69    const float s = simd_sum(acc);70    if (simd_is_first()) scratch[simd_idx] = s;71    threadgroup_barrier(mem_flags::mem_threadgroup);72    const float mine = (lane < n_simds) ? scratch[lane] : 0.0f;73    const float total = simd_sum(mine);74    if (simd_idx == 0 && lane == 0) out[0] = total * mul;75}76