spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "core/autograd.h"5#include "core/tensor.h"67#include <cstdint>89// Autograd-aware ops. Forward runs on the active backend (CPU reference or10// Metal stream) and, when the tape is enabled and an input requires grad,11// records a backward lambda that accumulates into input .grad tensors on12// the same backend.13//14// Metal-mode contract: values live in unified memory but are produced15// asynchronously — CPU code must not READ op outputs (incl. the loss16// scalar) before metal::sync(), and must not WRITE tensors that pending17// dispatches read. The trainer's step structure guarantees both.18namespace forge::ops {1920enum class Backend { CPU, Metal };21void set_backend(Backend b);22Backend backend();2324// c = a ⋅ b (2-D, optional transposes on the stored operands)25Var matmul(const Var& a, const Var& b, bool transpose_a = false, bool transpose_b = false);2627Var add(const Var& a, const Var& b);              // same shape28Var add_bias(const Var& x, const Var& bias);      // x [N,C] + bias [C]29Var mul(const Var& a, const Var& b);              // elementwise30Var scale(const Var& a, float s);3132Var silu(const Var& x);33Var gelu(const Var& x);34Var relu2(const Var& x); // max(x,0)^2 — nanoGPT-speedrun activation35Var sigmoid(const Var& x); // σ(x) — sigmoid MoE routing (DeepSeek V3)3637// Gemma-style soft capping: cap * tanh(x / cap). Bounds logits smoothly.38Var softcap(const Var& x, float cap);3940// Row softmax over the last dim (autograd; the attention op has its own41// fused softmax — this one is for the MoE router's small rows).42Var softmax(const Var& x);4344// QAT: per-row fake quantization of a 2-D weight. Forward quantizes45// (int8 absmax / ternary absmean); backward is the straight-through46// estimator — gradients flow unchanged into the f32 master weight.47enum class QuantMode { None, Int8, Ternary };48Var fake_quant(const Var& w, QuantMode mode);4950// MoE gating: keep top-k per row of probs, zero the rest. Selection ranks51// probs + bias (bias [E], requires_grad=false — the DeepSeek-V3 noaux52// balance bias; pass zeros for classic routing) but gate values are the53// biasless probs; norm renormalizes kept gates to sum 1. Gradients flow54// through kept entries only.55Var topk_renorm(const Var& probs, const Tensor& bias, int64_t k, bool norm = true);56// y[i,:] = x[i,:] * gates[i,e] — weight expert e's output by its gate column.57Var row_scale(const Var& x, const Var& gates, int64_t e);5859Var rmsnorm(const Var& x, const Var& w, float eps);60Var layernorm(const Var& x, const Var& w, const Var& b, float eps);6162// weight [V,C], ids [B,T] (u16/i32) → [B,T,C]63Var embedding(const Var& weight, const Tensor& ids);6465// Build the [head_dim/2] inverse-frequency table: theta^(-2k/hd), optionally66// rescaled HF-"llama3" style (scale_factor > 0): high-frequency components67// untouched, low-frequency divided by factor, smooth blend between.68Tensor rope_freqs(int64_t head_dim, float theta, float scale_factor = 0.0f,69                  float low_freq_factor = 1.0f, float high_freq_factor = 4.0f,70                  int64_t original_ctx = 8192);7172// x [B,T,H*hd], interleaved-pairs RoPE per head. freqs from rope_freqs()73// (callers cache it — one table per layer flavour).74Var rope(const Var& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset = 0);75// Convenience: plain-theta table built per call (tests / one-offs).76Var rope(const Var& x, int64_t n_heads, float theta, int64_t pos_offset = 0);7778// q [B,T,H*hd], k/v [B,T,Hkv*hd] → [B,T,H*hd]; causal + GQA.79// window > 0: sliding-window attention (Mistral/Gemma3) — fused scalar path.80// attn_softcap > 0: cap·tanh(score/cap) pre-softmax (Gemma2) — routes to the81// unfused path (probs materialized), so reserve it for short contexts.82Var attention(const Var& q, const Var& k, const Var& v,83              int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,84              int64_t window = 0, float attn_softcap = 0.0f);8586// logits [N,V], targets [N] (i32/u16, ignore_index=-1) → scalar mean loss87Var cross_entropy(const Var& logits, const Tensor& targets);8889// Shape helper: shares storage, no tape node needed (grad shapes follow value shapes).90Var reshape(const Var& x, std::vector<int64_t> shape);9192} // namespace forge::ops93