spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "core/tensor.h"56#include <cstdint>7#include <random>89// CPU reference implementations (forward + backward). These are the ground10// truth every Metal kernel is validated against — clarity beats speed here.11// All tensors f32 and contiguous unless stated otherwise. Backward12// functions ACCUMULATE into their d* outputs (autograd semantics).13namespace forge::cpu {1415// ---- init -----------------------------------------------------------------16void fill_normal(Tensor& t, float mean, float stddev, std::mt19937_64& rng);17void fill_uniform_int(Tensor& t, int64_t low, int64_t high, std::mt19937_64& rng);1819// ---- matmul ---------------------------------------------------------------20// C[M,N] = A ⋅ B with optional transposes (A stored [M,K] or [K,M], B stored21// [K,N] or [N,K]). accumulate=true adds into C instead of overwriting22// (backward passes accumulate into gradients).23void matmul(const Tensor& a, const Tensor& b, Tensor& c,24            bool transpose_a = false, bool transpose_b = false,25            bool accumulate = false);2627// ---- elementwise ----------------------------------------------------------28void add(const Tensor& a, const Tensor& b, Tensor& out);29void mul(const Tensor& a, const Tensor& b, Tensor& out);30void scale(const Tensor& a, float s, Tensor& out);31// x: [N, C], bias: [C], out[i,j] = x[i,j] + bias[j]32void add_bias(const Tensor& x, const Tensor& bias, Tensor& out);33void add_bias_backward(const Tensor& dout, Tensor& dbias);3435void silu(const Tensor& x, Tensor& out);36void silu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);37void gelu(const Tensor& x, Tensor& out); // tanh approximation (GPT-2)38void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);39void relu2(const Tensor& x, Tensor& out); // max(x,0)^240void relu2_backward(const Tensor& x, const Tensor& dout, Tensor& dx);41void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap)42// dx += dout * (1 - (y/cap)^2) — takes the forward OUTPUT y43void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx);4445// ---- norms (row-wise over the last dim) -----------------------------------46// x: [N, C], w: [C] (b: [C] for layernorm), out: [N, C]47void rmsnorm(const Tensor& x, const Tensor& w, float eps, Tensor& out);48void rmsnorm_backward(const Tensor& x, const Tensor& w, float eps,49                      const Tensor& dout, Tensor& dx, Tensor& dw);50void layernorm(const Tensor& x, const Tensor& w, const Tensor& b, float eps, Tensor& out);51void layernorm_backward(const Tensor& x, const Tensor& w, float eps,52                        const Tensor& dout, Tensor& dx, Tensor& dw, Tensor& db);5354// ---- softmax (row-wise over the last dim, max-subtracted) ------------------55void softmax(const Tensor& x, Tensor& out);56// dx = P ∘ (dout − rowsum(dout ∘ P))57void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx);5859// ---- embedding -------------------------------------------------------------60// weight: [V, C] f32; ids: [N] i32; out: [N, C]61void embedding(const Tensor& weight, const Tensor& ids, Tensor& out);62void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight);6364// ---- RoPE (interleaved-pairs / GPT-J convention; see RESEARCH.md §7) -------65// x: [B, T, H*hd]; rotates pairs (2k, 2k+1) inside each head. pos_offset66// shifts absolute positions (KV-cache generation).67// freqs: [head_dim/2] per-pair inverse frequencies, precomputed by the caller.68void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset,69          Tensor& out);70// Backward of a rotation is the inverse rotation applied to dout.71void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,72                   int64_t pos_offset, Tensor& dx);7374// ---- attention (composed reference: scores → mask → softmax → PV) ----------75// q: [B, T, H*hd], k/v: [B, T, Hkv*hd], out: [B, T, H*hd]. GQA via76// kv_head = h / (H / Hkv). If probs_out is non-null it receives the softmax77// probabilities [B, H, T, T] (needed for backward).78// window > 0: sliding-window attention (attend to the last `window` keys,79// self included). attn_softcap > 0: cap·tanh(score/cap) pre-softmax (Gemma 2).80void attention(const Tensor& q, const Tensor& k, const Tensor& v,81               int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,82               Tensor& out, Tensor* probs_out, int64_t window = 0,83               float attn_softcap = 0.0f);84// Uses the FlashAttention-2 identity rowsum(dP ∘ P) == dO_i · O_i, so `out`85// (the forward result) is required. Masked positions have prob 0, so the86// window needs no explicit handling here.87void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,88                        const Tensor& probs, const Tensor& out, const Tensor& dout,89                        int64_t n_heads, int64_t n_kv_heads, float scale,90                        Tensor& dq, Tensor& dk, Tensor& dv, float attn_softcap = 0.0f);9192// ---- QAT / MoE --------------------------------------------------------------93// Per-row fake quantization of a 2-D weight; mode 0 = int8 (absmax/127),94// mode 1 = ternary (absmean, BitNet-style). Backward is STE — identity.95void fake_quant(const Tensor& w, int mode, Tensor& out);96void sigmoid(const Tensor& x, Tensor& out);97// dx += dout * y * (1-y) — takes the forward OUTPUT y98void sigmoid_backward(const Tensor& y, const Tensor& dout, Tensor& dx);99// Keep top-k per row of p [N,E] — selected by p+bias, gated by p alone100// (bias [E], all-zero for classic routing); norm renormalizes kept gates to101// sum 1. Ties broken by lower index, same as the Metal kernel.102void topk_renorm(const Tensor& p, const Tensor& bias, int64_t k, bool norm,103                 Tensor& out);104void topk_renorm_backward(const Tensor& p, const Tensor& bias, const Tensor& dout,105                          int64_t k, bool norm, Tensor& dp);106// counts[e] += #rows with nonzero gate for expert e (noaux load statistic)107void expert_counts(const Tensor& gates, Tensor& counts);108// out[i,:] = x[i,:] * gates[i,e]; _accumulate does dst += (dx path);109// gate_backward does dgates[i,e] += dot(dout[i,:], x[i,:]).110void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out);111void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst);112void row_scale_gate_backward(const Tensor& dout, const Tensor& x, int64_t e,113                             Tensor& dgates);114115// ---- fused softmax cross-entropy -------------------------------------------116// logits: [N, V]; targets: [N] i32 (ignore_index = -1). Returns mean loss117// over valid rows. dlogits (if non-null) receives (softmax − onehot)/n_valid118// — the gradient for dloss = 1, ACCUMULATED.119float cross_entropy(const Tensor& logits, const Tensor& targets, Tensor* dlogits);120121} // namespace forge::cpu122