spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Row-wise RMSNorm / LayerNorm. One threadgroup per row; simd reductions4// first (register shuffles beat threadgroup memory 2:1 on Apple GPUs),5// tiny threadgroup scratch only to cross simdgroups. LayerNorm is two-pass6// (mean, then variance) — the E[x²]−E[x]² shortcut loses precision and7// these kernels are held to 1e-4 parity vs the two-pass CPU reference.8#include <metal_stdlib>9using namespace metal;1011namespace {1213// Threadgroup-wide sum of one per-thread value; returns the same total to14// every thread. Scratch must hold 32 floats.15inline float tg_sum(float v, uint lane, uint simd_idx, uint n_simds,16                    threadgroup float* scratch) {17    const float s = simd_sum(v);18    if (simd_is_first()) scratch[simd_idx] = s;19    threadgroup_barrier(mem_flags::mem_threadgroup);20    const float mine = (lane < n_simds) ? scratch[lane] : 0.0f;21    return simd_sum(mine);22}2324} // namespace2526struct NormParams {27    uint C;28    float eps;29};3031kernel void rmsnorm_f32(device const float* X   [[buffer(0)]],32                        device const float* W   [[buffer(1)]],33                        device float*       OUT [[buffer(2)]],34                        constant NormParams& p  [[buffer(3)]],35                        uint row      [[threadgroup_position_in_grid]],36                        uint lid      [[thread_index_in_threadgroup]],37                        uint tg_size  [[threads_per_threadgroup]],38                        uint lane     [[thread_index_in_simdgroup]],39                        uint simd_idx [[simdgroup_index_in_threadgroup]],40                        uint n_simds  [[simdgroups_per_threadgroup]]) {41    device const float* x = X + ulong(row) * p.C;42    device float* out = OUT + ulong(row) * p.C;43    threadgroup float scratch[32];4445    float ss = 0.0f;46    for (uint j = lid; j < p.C; j += tg_size) ss = fma(x[j], x[j], ss);47    ss = tg_sum(ss, lane, simd_idx, n_simds, scratch);4849    const float inv_rms = rsqrt(ss / float(p.C) + p.eps);50    for (uint j = lid; j < p.C; j += tg_size) out[j] = W[j] * x[j] * inv_rms;51}5253kernel void layernorm_f32(device const float* X   [[buffer(0)]],54                          device const float* W   [[buffer(1)]],55                          device const float* B   [[buffer(2)]],56                          device float*       OUT [[buffer(3)]],57                          constant NormParams& p  [[buffer(4)]],58                          uint row      [[threadgroup_position_in_grid]],59                          uint lid      [[thread_index_in_threadgroup]],60                          uint tg_size  [[threads_per_threadgroup]],61                          uint lane     [[thread_index_in_simdgroup]],62                          uint simd_idx [[simdgroup_index_in_threadgroup]],63                          uint n_simds  [[simdgroups_per_threadgroup]]) {64    device const float* x = X + ulong(row) * p.C;65    device float* out = OUT + ulong(row) * p.C;66    threadgroup float scratch[32];6768    float s = 0.0f;69    for (uint j = lid; j < p.C; j += tg_size) s += x[j];70    const float mean = tg_sum(s, lane, simd_idx, n_simds, scratch) / float(p.C);71    threadgroup_barrier(mem_flags::mem_threadgroup); // scratch reuse7273    float var = 0.0f;74    for (uint j = lid; j < p.C; j += tg_size) {75        const float d = x[j] - mean;76        var = fma(d, d, var);77    }78    var = tg_sum(var, lane, simd_idx, n_simds, scratch) / float(p.C);7980    const float inv_std = rsqrt(var + p.eps);81    for (uint j = lid; j < p.C; j += tg_size)82        out[j] = fma(W[j], (x[j] - mean) * inv_std, B[j]);83}8485// ---- backwards -------------------------------------------------------------86// dx is row-parallel (one threadgroup per row) and saves the row statistics87// (inv_rms, or mean+inv_std) it computes anyway; dw/db are column-parallel88// (one thread per column) reading those saved stats — deterministic, no89// atomics, no per-column recomputation.9091kernel void rmsnorm_bwd_dx_f32(device const float* X       [[buffer(0)]],92                               device const float* W       [[buffer(1)]],93                               device const float* dOUT    [[buffer(2)]],94                               device float*       dX      [[buffer(3)]],95                               device float*       inv_rms [[buffer(4)]], // [rows]96                               constant NormParams& p      [[buffer(5)]],97                               uint row      [[threadgroup_position_in_grid]],98                               uint lid      [[thread_index_in_threadgroup]],99                               uint tg_size  [[threads_per_threadgroup]],100                               uint lane     [[thread_index_in_simdgroup]],101                               uint simd_idx [[simdgroup_index_in_threadgroup]],102                               uint n_simds  [[simdgroups_per_threadgroup]]) {103    device const float* x = X + ulong(row) * p.C;104    device const float* d = dOUT + ulong(row) * p.C;105    device float* dx = dX + ulong(row) * p.C;106    threadgroup float scratch[32];107108    float ss = 0.0f;109    for (uint j = lid; j < p.C; j += tg_size) ss = fma(x[j], x[j], ss);110    ss = tg_sum(ss, lane, simd_idx, n_simds, scratch);111    const float ir = rsqrt(ss / float(p.C) + p.eps);112    if (lid == 0) inv_rms[row] = ir;113    threadgroup_barrier(mem_flags::mem_threadgroup);114115    float dot = 0.0f; // sum_j g_j w_j x_j116    for (uint j = lid; j < p.C; j += tg_size) dot = fma(d[j] * W[j], x[j], dot);117    dot = tg_sum(dot, lane, simd_idx, n_simds, scratch);118119    const float coef = dot * ir * ir * ir / float(p.C);120    for (uint j = lid; j < p.C; j += tg_size)121        dx[j] += d[j] * W[j] * ir - x[j] * coef;122}123124// dw[j] += sum_i g[i,j] * x[i,j] * inv_rms[i]125kernel void rmsnorm_bwd_dw_f32(device const float* X       [[buffer(0)]],126                               device const float* dOUT    [[buffer(1)]],127                               device const float* inv_rms [[buffer(2)]],128                               device float*       dW      [[buffer(3)]],129                               constant NormParams& p      [[buffer(4)]],130                               constant uint&      rows    [[buffer(5)]],131                               uint j [[thread_position_in_grid]]) {132    float acc = 0.0f;133    for (uint i = 0; i < rows; ++i)134        acc = fma(dOUT[ulong(i) * p.C + j] * X[ulong(i) * p.C + j], inv_rms[i], acc);135    dW[j] += acc;136}137138kernel void layernorm_bwd_dx_f32(device const float* X    [[buffer(0)]],139                                 device const float* W    [[buffer(1)]],140                                 device const float* dOUT [[buffer(2)]],141                                 device float*       dX   [[buffer(3)]],142                                 device float*       mean_out [[buffer(4)]], // [rows]143                                 device float*       istd_out [[buffer(5)]], // [rows]144                                 constant NormParams& p   [[buffer(6)]],145                                 uint row      [[threadgroup_position_in_grid]],146                                 uint lid      [[thread_index_in_threadgroup]],147                                 uint tg_size  [[threads_per_threadgroup]],148                                 uint lane     [[thread_index_in_simdgroup]],149                                 uint simd_idx [[simdgroup_index_in_threadgroup]],150                                 uint n_simds  [[simdgroups_per_threadgroup]]) {151    device const float* x = X + ulong(row) * p.C;152    device const float* d = dOUT + ulong(row) * p.C;153    device float* dx = dX + ulong(row) * p.C;154    threadgroup float scratch[32];155156    float s = 0.0f;157    for (uint j = lid; j < p.C; j += tg_size) s += x[j];158    const float mean = tg_sum(s, lane, simd_idx, n_simds, scratch) / float(p.C);159    threadgroup_barrier(mem_flags::mem_threadgroup);160161    float var = 0.0f;162    for (uint j = lid; j < p.C; j += tg_size) {163        const float dv = x[j] - mean;164        var = fma(dv, dv, var);165    }166    var = tg_sum(var, lane, simd_idx, n_simds, scratch) / float(p.C);167    const float inv_std = rsqrt(var + p.eps);168    if (lid == 0) {169        mean_out[row] = mean;170        istd_out[row] = inv_std;171    }172    threadgroup_barrier(mem_flags::mem_threadgroup);173174    float m_dxhat = 0.0f, m_dxhat_xhat = 0.0f;175    for (uint j = lid; j < p.C; j += tg_size) {176        const float xhat = (x[j] - mean) * inv_std;177        const float dxhat = d[j] * W[j];178        m_dxhat += dxhat;179        m_dxhat_xhat = fma(dxhat, xhat, m_dxhat_xhat);180    }181    m_dxhat = tg_sum(m_dxhat, lane, simd_idx, n_simds, scratch) / float(p.C);182    threadgroup_barrier(mem_flags::mem_threadgroup);183    m_dxhat_xhat = tg_sum(m_dxhat_xhat, lane, simd_idx, n_simds, scratch) / float(p.C);184185    for (uint j = lid; j < p.C; j += tg_size) {186        const float xhat = (x[j] - mean) * inv_std;187        dx[j] += inv_std * (d[j] * W[j] - m_dxhat - xhat * m_dxhat_xhat);188    }189}190191// dw[j] += sum_i g[i,j] * xhat[i,j] ; db[j] += sum_i g[i,j]192kernel void layernorm_bwd_dwdb_f32(device const float* X    [[buffer(0)]],193                                   device const float* dOUT [[buffer(1)]],194                                   device const float* mean [[buffer(2)]],195                                   device const float* istd [[buffer(3)]],196                                   device float*       dW   [[buffer(4)]],197                                   device float*       dB   [[buffer(5)]],198                                   constant NormParams& p   [[buffer(6)]],199                                   constant uint&      rows [[buffer(7)]],200                                   uint j [[thread_position_in_grid]]) {201    float aw = 0.0f, ab = 0.0f;202    for (uint i = 0; i < rows; ++i) {203        const float g = dOUT[ulong(i) * p.C + j];204        const float xhat = (X[ulong(i) * p.C + j] - mean[i]) * istd[i];205        aw = fma(g, xhat, aw);206        ab += g;207    }208    dW[j] += aw;209    dB[j] += ab;210}211