spb/forge Public MIT
Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.
C++ 61.2%
C 23%
Python 7.6%
TeX 7.2%
CMake 1.1%
1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#include "ops/cpu/cpu_ops.h"34#include <dispatch/dispatch.h>56#include <algorithm>7#include <cassert>8#include <cmath>9#include <cstdio>10#include <cstdlib>11#include <limits>1213namespace forge::cpu {1415namespace {1617void check(bool cond, const char* msg) {18 if (!cond) {19 std::fprintf(stderr, "forge/cpu: %s\n", msg);20 std::abort();21 }22}2324// Row-parallel helper (GCD). Serial below the threshold so tiny test25// tensors don't pay dispatch overhead.26template <typename F>27void parallel_rows(int64_t n, F&& body) {28 if (n < 32) {29 for (int64_t i = 0; i < n; ++i) body(i);30 } else {31 dispatch_apply(size_t(n), DISPATCH_APPLY_AUTO,32 ^(size_t i) { body(int64_t(i)); });33 }34}3536int32_t token_at(const Tensor& ids, int64_t i) {37 return ids.dtype() == DType::I32 ? ids.data<int32_t>()[i]38 : int32_t(ids.data<uint16_t>()[i]);39}4041} // namespace4243// ---- init -------------------------------------------------------------------4445void fill_normal(Tensor& t, float mean, float stddev, std::mt19937_64& rng) {46 std::normal_distribution<float> dist(mean, stddev);47 float* p = t.data<float>();48 for (int64_t i = 0; i < t.numel(); ++i) p[i] = dist(rng);49}5051void fill_uniform_int(Tensor& t, int64_t low, int64_t high, std::mt19937_64& rng) {52 std::uniform_int_distribution<int64_t> dist(low, high - 1);53 for (int64_t i = 0; i < t.numel(); ++i) t.set_item(i, float(dist(rng)));54}5556// ---- matmul -----------------------------------------------------------------5758void matmul(const Tensor& a, const Tensor& b, Tensor& c,59 bool transpose_a, bool transpose_b, bool accumulate) {60 check(a.dtype() == DType::F32 && b.dtype() == DType::F32 && c.dtype() == DType::F32,61 "matmul: f32 only");62 check(a.ndim() == 2 && b.ndim() == 2 && c.ndim() == 2, "matmul: 2-D only");63 check(a.is_contiguous() && b.is_contiguous() && c.is_contiguous(),64 "matmul: contiguous only");6566 const int64_t M = transpose_a ? a.size(1) : a.size(0);67 const int64_t K = transpose_a ? a.size(0) : a.size(1);68 const int64_t Kb = transpose_b ? b.size(1) : b.size(0);69 const int64_t N = transpose_b ? b.size(0) : b.size(1);70 check(K == Kb, "matmul: inner dims mismatch");71 check(c.size(0) == M && c.size(1) == N, "matmul: output shape mismatch");7273 const float* A = a.data<float>();74 const float* B = b.data<float>();75 float* C = c.data<float>();76 const int64_t lda = a.size(1);77 const int64_t ldb = b.size(1);7879 // i-k-j order streams contiguous rows of B and C in the common case.80 parallel_rows(M, [&](int64_t i) {81 float* crow = C + i * N;82 if (!accumulate)83 for (int64_t j = 0; j < N; ++j) crow[j] = 0.0f;84 for (int64_t k = 0; k < K; ++k) {85 const float aik = transpose_a ? A[k * lda + i] : A[i * lda + k];86 if (!transpose_b) {87 const float* brow = B + k * ldb;88 for (int64_t j = 0; j < N; ++j) crow[j] += aik * brow[j];89 } else {90 for (int64_t j = 0; j < N; ++j) crow[j] += aik * B[j * ldb + k];91 }92 }93 });94}9596// ---- elementwise ------------------------------------------------------------9798void add(const Tensor& a, const Tensor& b, Tensor& out) {99 check(a.numel() == b.numel() && a.numel() == out.numel(), "add: numel mismatch");100 const float* pa = a.data<float>();101 const float* pb = b.data<float>();102 float* po = out.data<float>();103 for (int64_t i = 0; i < a.numel(); ++i) po[i] = pa[i] + pb[i];104}105106void mul(const Tensor& a, const Tensor& b, Tensor& out) {107 check(a.numel() == b.numel() && a.numel() == out.numel(), "mul: numel mismatch");108 const float* pa = a.data<float>();109 const float* pb = b.data<float>();110 float* po = out.data<float>();111 for (int64_t i = 0; i < a.numel(); ++i) po[i] = pa[i] * pb[i];112}113114void scale(const Tensor& a, float s, Tensor& out) {115 check(a.numel() == out.numel(), "scale: numel mismatch");116 const float* pa = a.data<float>();117 float* po = out.data<float>();118 for (int64_t i = 0; i < a.numel(); ++i) po[i] = pa[i] * s;119}120121void add_bias(const Tensor& x, const Tensor& bias, Tensor& out) {122 const int64_t C = bias.numel();123 const int64_t N = x.numel() / C;124 const float* px = x.data<float>();125 const float* pb = bias.data<float>();126 float* po = out.data<float>();127 for (int64_t i = 0; i < N; ++i)128 for (int64_t j = 0; j < C; ++j) po[i * C + j] = px[i * C + j] + pb[j];129}130131void add_bias_backward(const Tensor& dout, Tensor& dbias) {132 const int64_t C = dbias.numel();133 const int64_t N = dout.numel() / C;134 const float* pd = dout.data<float>();135 float* pb = dbias.data<float>();136 for (int64_t i = 0; i < N; ++i)137 for (int64_t j = 0; j < C; ++j) pb[j] += pd[i * C + j];138}139140void silu(const Tensor& x, Tensor& out) {141 const float* px = x.data<float>();142 float* po = out.data<float>();143 for (int64_t i = 0; i < x.numel(); ++i) {144 const float v = px[i];145 po[i] = v / (1.0f + std::exp(-v));146 }147}148149void silu_backward(const Tensor& x, const Tensor& dout, Tensor& dx) {150 const float* px = x.data<float>();151 const float* pd = dout.data<float>();152 float* pdx = dx.data<float>();153 for (int64_t i = 0; i < x.numel(); ++i) {154 const float v = px[i];155 const float sig = 1.0f / (1.0f + std::exp(-v));156 pdx[i] += pd[i] * sig * (1.0f + v * (1.0f - sig));157 }158}159160void gelu(const Tensor& x, Tensor& out) {161 constexpr float k = 0.7978845608028654f; // sqrt(2/pi)162 const float* px = x.data<float>();163 float* po = out.data<float>();164 for (int64_t i = 0; i < x.numel(); ++i) {165 const float v = px[i];166 po[i] = 0.5f * v * (1.0f + std::tanh(k * (v + 0.044715f * v * v * v)));167 }168}169170void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx) {171 constexpr float k = 0.7978845608028654f;172 const float* px = x.data<float>();173 const float* pd = dout.data<float>();174 float* pdx = dx.data<float>();175 for (int64_t i = 0; i < x.numel(); ++i) {176 const float v = px[i];177 const float u = k * (v + 0.044715f * v * v * v);178 const float t = std::tanh(u);179 const float du = k * (1.0f + 3.0f * 0.044715f * v * v);180 pdx[i] += pd[i] * (0.5f * (1.0f + t) + 0.5f * v * (1.0f - t * t) * du);181 }182}183184void relu2(const Tensor& x, Tensor& out) {185 const float* px = x.data<float>();186 float* po = out.data<float>();187 for (int64_t i = 0; i < x.numel(); ++i) {188 const float v = std::max(px[i], 0.0f);189 po[i] = v * v;190 }191}192193void relu2_backward(const Tensor& x, const Tensor& dout, Tensor& dx) {194 const float* px = x.data<float>();195 const float* pd = dout.data<float>();196 float* pdx = dx.data<float>();197 for (int64_t i = 0; i < x.numel(); ++i)198 pdx[i] += pd[i] * 2.0f * std::max(px[i], 0.0f);199}200201void softcap(const Tensor& x, float cap, Tensor& out) {202 const float* px = x.data<float>();203 float* po = out.data<float>();204 for (int64_t i = 0; i < x.numel(); ++i) po[i] = cap * std::tanh(px[i] / cap);205}206207void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx) {208 const float* py = y.data<float>();209 const float* pd = dout.data<float>();210 float* pdx = dx.data<float>();211 for (int64_t i = 0; i < y.numel(); ++i) {212 const float t = py[i] / cap;213 pdx[i] += pd[i] * (1.0f - t * t);214 }215}216217// ---- norms --------------------------------------------------------------------218219void rmsnorm(const Tensor& x, const Tensor& w, float eps, Tensor& out) {220 const int64_t C = w.numel();221 const int64_t N = x.numel() / C;222 const float* pw = w.data<float>();223 const float* px = x.data<float>();224 float* po = out.data<float>();225 parallel_rows(N, [&](int64_t i) {226 const float* row = px + i * C;227 float ss = 0.0f;228 for (int64_t j = 0; j < C; ++j) ss += row[j] * row[j];229 const float inv_rms = 1.0f / std::sqrt(ss / float(C) + eps);230 for (int64_t j = 0; j < C; ++j) po[i * C + j] = pw[j] * row[j] * inv_rms;231 });232}233234void rmsnorm_backward(const Tensor& x, const Tensor& w, float eps,235 const Tensor& dout, Tensor& dx, Tensor& dw) {236 const int64_t C = w.numel();237 const int64_t N = x.numel() / C;238 const float* px = x.data<float>();239 const float* pw = w.data<float>();240 const float* pd = dout.data<float>();241 float* pdx = dx.data<float>();242 float* pdw = dw.data<float>();243 // dw is a cross-row reduction: keep it serial for determinism.244 for (int64_t i = 0; i < N; ++i) {245 const float* row = px + i * C;246 const float* drow = pd + i * C;247 float ss = 0.0f;248 for (int64_t j = 0; j < C; ++j) ss += row[j] * row[j];249 const float inv_rms = 1.0f / std::sqrt(ss / float(C) + eps);250 float dot = 0.0f; // sum_j g_j w_j x_j251 for (int64_t j = 0; j < C; ++j) dot += drow[j] * pw[j] * row[j];252 const float coef = dot * inv_rms * inv_rms * inv_rms / float(C);253 for (int64_t j = 0; j < C; ++j) {254 pdx[i * C + j] += drow[j] * pw[j] * inv_rms - row[j] * coef;255 pdw[j] += drow[j] * row[j] * inv_rms;256 }257 }258}259260void layernorm(const Tensor& x, const Tensor& w, const Tensor& b, float eps, Tensor& out) {261 const int64_t C = w.numel();262 const int64_t N = x.numel() / C;263 const float* px = x.data<float>();264 const float* pw = w.data<float>();265 const float* pb = b.data<float>();266 float* po = out.data<float>();267 parallel_rows(N, [&](int64_t i) {268 const float* row = px + i * C;269 float mean = 0.0f;270 for (int64_t j = 0; j < C; ++j) mean += row[j];271 mean /= float(C);272 float var = 0.0f;273 for (int64_t j = 0; j < C; ++j) var += (row[j] - mean) * (row[j] - mean);274 var /= float(C);275 const float inv_std = 1.0f / std::sqrt(var + eps);276 for (int64_t j = 0; j < C; ++j)277 po[i * C + j] = pw[j] * (row[j] - mean) * inv_std + pb[j];278 });279}280281void layernorm_backward(const Tensor& x, const Tensor& w, float eps,282 const Tensor& dout, Tensor& dx, Tensor& dw, Tensor& db) {283 const int64_t C = w.numel();284 const int64_t N = x.numel() / C;285 const float* px = x.data<float>();286 const float* pw = w.data<float>();287 const float* pd = dout.data<float>();288 float* pdx = dx.data<float>();289 float* pdw = dw.data<float>();290 float* pdb = db.data<float>();291 for (int64_t i = 0; i < N; ++i) {292 const float* row = px + i * C;293 const float* drow = pd + i * C;294 float mean = 0.0f;295 for (int64_t j = 0; j < C; ++j) mean += row[j];296 mean /= float(C);297 float var = 0.0f;298 for (int64_t j = 0; j < C; ++j) var += (row[j] - mean) * (row[j] - mean);299 var /= float(C);300 const float inv_std = 1.0f / std::sqrt(var + eps);301 // dxhat = g*w ; dx = inv_std * (dxhat − mean(dxhat) − xhat*mean(dxhat∘xhat))302 float mean_dxhat = 0.0f, mean_dxhat_xhat = 0.0f;303 for (int64_t j = 0; j < C; ++j) {304 const float xhat = (row[j] - mean) * inv_std;305 const float dxhat = drow[j] * pw[j];306 mean_dxhat += dxhat;307 mean_dxhat_xhat += dxhat * xhat;308 }309 mean_dxhat /= float(C);310 mean_dxhat_xhat /= float(C);311 for (int64_t j = 0; j < C; ++j) {312 const float xhat = (row[j] - mean) * inv_std;313 pdx[i * C + j] += inv_std * (drow[j] * pw[j] - mean_dxhat - xhat * mean_dxhat_xhat);314 pdw[j] += drow[j] * xhat;315 pdb[j] += drow[j];316 }317 }318}319320// ---- softmax ------------------------------------------------------------------321322void softmax(const Tensor& x, Tensor& out) {323 const int64_t C = x.shape().back();324 const int64_t N = x.numel() / C;325 const float* px = x.data<float>();326 float* po = out.data<float>();327 parallel_rows(N, [&](int64_t i) {328 const float* row = px + i * C;329 float m = row[0];330 for (int64_t j = 1; j < C; ++j) m = std::max(m, row[j]);331 float sum = 0.0f;332 for (int64_t j = 0; j < C; ++j) {333 const float e = std::exp(row[j] - m);334 po[i * C + j] = e;335 sum += e;336 }337 const float inv = 1.0f / sum;338 for (int64_t j = 0; j < C; ++j) po[i * C + j] *= inv;339 });340}341342void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx) {343 const int64_t C = p.shape().back();344 const int64_t N = p.numel() / C;345 const float* pp = p.data<float>();346 const float* pd = dout.data<float>();347 float* pdx = dx.data<float>();348 for (int64_t i = 0; i < N; ++i) {349 const float* prow = pp + i * C;350 const float* drow = pd + i * C;351 float dot = 0.0f;352 for (int64_t j = 0; j < C; ++j) dot += drow[j] * prow[j];353 for (int64_t j = 0; j < C; ++j) pdx[i * C + j] += prow[j] * (drow[j] - dot);354 }355}356357// ---- QAT / MoE ------------------------------------------------------------------358359void fake_quant(const Tensor& w, int mode, Tensor& out) {360 const int64_t R = w.size(0), C = w.size(1);361 const float* pw = w.data<float>();362 float* po = out.data<float>();363 for (int64_t i = 0; i < R; ++i) {364 const float* row = pw + i * C;365 float* orow = po + i * C;366 if (mode == 0) { // int8: symmetric absmax367 float amax = 0.0f;368 for (int64_t j = 0; j < C; ++j) amax = std::max(amax, std::fabs(row[j]));369 const float s = std::max(amax / 127.0f, 1e-12f);370 for (int64_t j = 0; j < C; ++j) orow[j] = std::rint(row[j] / s) * s;371 } else { // ternary: BitNet b1.58 absmean372 float asum = 0.0f;373 for (int64_t j = 0; j < C; ++j) asum += std::fabs(row[j]);374 const float s = std::max(asum / float(C), 1e-12f);375 for (int64_t j = 0; j < C; ++j)376 orow[j] = std::min(1.0f, std::max(-1.0f, std::rint(row[j] / s))) * s;377 }378 }379}380381void sigmoid(const Tensor& x, Tensor& out) {382 const float* px = x.data<float>();383 float* po = out.data<float>();384 for (int64_t i = 0; i < x.numel(); ++i) po[i] = 1.0f / (1.0f + std::exp(-px[i]));385}386387void sigmoid_backward(const Tensor& y, const Tensor& dout, Tensor& dx) {388 const float* py = y.data<float>();389 const float* pd = dout.data<float>();390 float* pdx = dx.data<float>();391 for (int64_t i = 0; i < y.numel(); ++i) pdx[i] += pd[i] * py[i] * (1.0f - py[i]);392}393394namespace {395// Kept set of the k largest (score + bias) entries, ties to the lower index396// (matches Metal). The returned sum is over the BIASLESS scores — the bias397// only steers selection (DeepSeek-V3 noaux routing).398void topk_select(const float* row, const float* bias, int64_t E, int64_t k,399 bool* kept, float* sum) {400 for (int64_t j = 0; j < E; ++j) kept[j] = false;401 float S = 0.0f;402 for (int64_t sel = 0; sel < k; ++sel) {403 float best = -std::numeric_limits<float>::max();404 int64_t arg = 0;405 for (int64_t j = 0; j < E; ++j) {406 const float v = row[j] + bias[j];407 if (!kept[j] && v > best) { best = v; arg = j; }408 }409 kept[arg] = true;410 S += row[arg];411 }412 *sum = S;413}414} // namespace415416void topk_renorm(const Tensor& p, const Tensor& bias, int64_t k, bool norm,417 Tensor& out) {418 const int64_t E = p.shape().back();419 const int64_t N = p.numel() / E;420 const float* pp = p.data<float>();421 const float* pb = bias.data<float>();422 float* po = out.data<float>();423 bool kbuf[64];424 for (int64_t i = 0; i < N; ++i) {425 const float* row = pp + i * E;426 float S = 0.0f;427 topk_select(row, pb, E, k, kbuf, &S);428 const float inv = norm ? 1.0f / std::max(S, 1e-12f) : 1.0f;429 for (int64_t j = 0; j < E; ++j) po[i * E + j] = kbuf[j] ? row[j] * inv : 0.0f;430 }431}432433void topk_renorm_backward(const Tensor& p, const Tensor& bias, const Tensor& dout,434 int64_t k, bool norm, Tensor& dp) {435 const int64_t E = p.shape().back();436 const int64_t N = p.numel() / E;437 const float* pp = p.data<float>();438 const float* pb = bias.data<float>();439 const float* pd = dout.data<float>();440 float* pdp = dp.data<float>();441 bool kbuf[64];442 for (int64_t i = 0; i < N; ++i) {443 const float* row = pp + i * E;444 const float* drow = pd + i * E;445 float S = 0.0f;446 topk_select(row, pb, E, k, kbuf, &S);447 if (!norm) {448 for (int64_t j = 0; j < E; ++j)449 if (kbuf[j]) pdp[i * E + j] += drow[j];450 continue;451 }452 const float inv = 1.0f / std::max(S, 1e-12f);453 float dot = 0.0f;454 for (int64_t j = 0; j < E; ++j)455 if (kbuf[j]) dot += drow[j] * row[j] * inv;456 for (int64_t j = 0; j < E; ++j)457 if (kbuf[j]) pdp[i * E + j] += (drow[j] - dot) * inv;458 }459}460461void expert_counts(const Tensor& gates, Tensor& counts) {462 const int64_t E = gates.shape().back();463 const int64_t N = gates.numel() / E;464 const float* pg = gates.data<float>();465 float* pc = counts.data<float>();466 for (int64_t i = 0; i < N; ++i)467 for (int64_t e = 0; e < E; ++e)468 if (pg[i * E + e] != 0.0f) pc[e] += 1.0f;469}470471void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out) {472 const int64_t C = x.shape().back();473 const int64_t N = x.numel() / C;474 const int64_t E = gates.shape().back();475 const float* px = x.data<float>();476 const float* pg = gates.data<float>();477 float* po = out.data<float>();478 for (int64_t i = 0; i < N; ++i) {479 const float s = pg[i * E + e];480 for (int64_t j = 0; j < C; ++j) po[i * C + j] = px[i * C + j] * s;481 }482}483484void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst) {485 const int64_t C = x.shape().back();486 const int64_t N = x.numel() / C;487 const int64_t E = gates.shape().back();488 const float* px = x.data<float>();489 const float* pg = gates.data<float>();490 float* pd = dst.data<float>();491 for (int64_t i = 0; i < N; ++i) {492 const float s = pg[i * E + e];493 for (int64_t j = 0; j < C; ++j) pd[i * C + j] += px[i * C + j] * s;494 }495}496497void row_scale_gate_backward(const Tensor& dout, const Tensor& x, int64_t e,498 Tensor& dgates) {499 const int64_t C = x.shape().back();500 const int64_t N = x.numel() / C;501 const int64_t E = dgates.shape().back();502 const float* pd = dout.data<float>();503 const float* px = x.data<float>();504 float* pg = dgates.data<float>();505 for (int64_t i = 0; i < N; ++i) {506 float acc = 0.0f;507 for (int64_t j = 0; j < C; ++j) acc += pd[i * C + j] * px[i * C + j];508 pg[i * E + e] += acc;509 }510}511512// ---- embedding ------------------------------------------------------------------513514void embedding(const Tensor& weight, const Tensor& ids, Tensor& out) {515 const int64_t C = weight.size(1);516 const int64_t N = ids.numel();517 const float* pw = weight.data<float>();518 float* po = out.data<float>();519 for (int64_t i = 0; i < N; ++i) {520 const int32_t tok = token_at(ids, i);521 const float* src = pw + int64_t(tok) * C;522 float* dst = po + i * C;523 for (int64_t j = 0; j < C; ++j) dst[j] = src[j];524 }525}526527void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight) {528 const int64_t C = dweight.size(1);529 const int64_t N = ids.numel();530 const float* pd = dout.data<float>();531 float* pw = dweight.data<float>();532 for (int64_t i = 0; i < N; ++i) {533 const int32_t tok = token_at(ids, i);534 float* dst = pw + int64_t(tok) * C;535 const float* src = pd + i * C;536 for (int64_t j = 0; j < C; ++j) dst[j] += src[j];537 }538}539540// ---- RoPE ----------------------------------------------------------------------541542namespace {543void rope_impl(const float* in, float* out, int64_t B, int64_t T, int64_t H, int64_t hd,544 const float* freqs, int64_t pos_offset, bool inverse, bool accumulate) {545 const int64_t C = H * hd;546 parallel_rows(B * T, [&](int64_t bt) {547 const int64_t t = bt % T;548 const float pos = float(t + pos_offset);549 const float* src = in + bt * C;550 float* dst = out + bt * C;551 for (int64_t h = 0; h < H; ++h) {552 for (int64_t k = 0; k < hd / 2; ++k) {553 const float angle = pos * freqs[k];554 const float c = std::cos(angle);555 const float s = inverse ? -std::sin(angle) : std::sin(angle);556 const int64_t i0 = h * hd + 2 * k;557 const float x0 = src[i0], x1 = src[i0 + 1];558 const float y0 = x0 * c - x1 * s;559 const float y1 = x0 * s + x1 * c;560 if (accumulate) {561 dst[i0] += y0;562 dst[i0 + 1] += y1;563 } else {564 dst[i0] = y0;565 dst[i0 + 1] = y1;566 }567 }568 }569 });570}571} // namespace572573void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset,574 Tensor& out) {575 check(x.ndim() == 3, "rope: expected [B,T,C]");576 const int64_t hd = x.size(2) / n_heads;577 check(hd % 2 == 0, "rope: head_dim must be even");578 check(freqs.numel() == hd / 2, "rope: freqs must have head_dim/2 entries");579 rope_impl(x.data<float>(), out.data<float>(), x.size(0), x.size(1), n_heads, hd,580 freqs.data<float>(), pos_offset, /*inverse=*/false, /*accumulate=*/false);581}582583void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,584 int64_t pos_offset, Tensor& dx) {585 const int64_t hd = dout.size(2) / n_heads;586 rope_impl(dout.data<float>(), dx.data<float>(), dout.size(0), dout.size(1), n_heads,587 hd, freqs.data<float>(), pos_offset, /*inverse=*/true, /*accumulate=*/true);588}589590// ---- attention -------------------------------------------------------------------591592void attention(const Tensor& q, const Tensor& k, const Tensor& v,593 int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,594 Tensor& out, Tensor* probs_out, int64_t window, float attn_softcap) {595 check(q.ndim() == 3 && k.ndim() == 3 && v.ndim() == 3, "attention: expected [B,T,C]");596 const int64_t B = q.size(0), T = q.size(1);597 const int64_t hd = q.size(2) / n_heads;598 const int64_t rep = n_heads / n_kv_heads;599 const int64_t Cq = n_heads * hd, Ckv = n_kv_heads * hd;600 check(probs_out != nullptr, "attention: probs_out required (reference impl)");601 check(probs_out->numel() == B * n_heads * T * T, "attention: probs shape");602603 const float* pq = q.data<float>();604 const float* pk = k.data<float>();605 const float* pv = v.data<float>();606 float* po = out.data<float>();607 float* pp = probs_out->data<float>();608609 parallel_rows(B * n_heads, [&](int64_t bh) {610 const int64_t b = bh / n_heads;611 const int64_t h = bh % n_heads;612 const int64_t hkv = h / rep;613 float* P = pp + bh * T * T;614615 for (int64_t i = 0; i < T; ++i) {616 const float* qi = pq + (b * T + i) * Cq + h * hd;617 const int64_t jmax = causal ? i : T - 1;618 // sliding window: attend only to the last `window` keys (incl. self)619 const int64_t jmin = window > 0 ? std::max<int64_t>(0, i - window + 1) : 0;620 // scores (masked positions never written; treated as prob 0)621 float m = -INFINITY;622 for (int64_t j = jmin; j <= jmax; ++j) {623 const float* kj = pk + (b * T + j) * Ckv + hkv * hd;624 float s = 0.0f;625 for (int64_t d = 0; d < hd; ++d) s += qi[d] * kj[d];626 s *= scale;627 if (attn_softcap > 0.0f)628 s = attn_softcap * std::tanh(s / attn_softcap);629 P[i * T + j] = s;630 m = std::max(m, s);631 }632 float sum = 0.0f;633 for (int64_t j = jmin; j <= jmax; ++j) {634 const float e = std::exp(P[i * T + j] - m);635 P[i * T + j] = e;636 sum += e;637 }638 const float inv = 1.0f / sum;639 for (int64_t j = 0; j < jmin; ++j) P[i * T + j] = 0.0f;640 for (int64_t j = jmin; j <= jmax; ++j) P[i * T + j] *= inv;641 for (int64_t j = jmax + 1; j < T; ++j) P[i * T + j] = 0.0f;642643 float* oi = po + (b * T + i) * Cq + h * hd;644 for (int64_t d = 0; d < hd; ++d) oi[d] = 0.0f;645 for (int64_t j = jmin; j <= jmax; ++j) {646 const float p = P[i * T + j];647 const float* vj = pv + (b * T + j) * Ckv + hkv * hd;648 for (int64_t d = 0; d < hd; ++d) oi[d] += p * vj[d];649 }650 }651 });652}653654void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,655 const Tensor& probs, const Tensor& out, const Tensor& dout,656 int64_t n_heads, int64_t n_kv_heads, float scale,657 Tensor& dq, Tensor& dk, Tensor& dv, float attn_softcap) {658 const int64_t B = q.size(0), T = q.size(1);659 const int64_t hd = q.size(2) / n_heads;660 const int64_t rep = n_heads / n_kv_heads;661 const int64_t Cq = n_heads * hd, Ckv = n_kv_heads * hd;662663 const float* pq = q.data<float>();664 const float* pk = k.data<float>();665 const float* pv = v.data<float>();666 const float* pp = probs.data<float>();667 const float* po = out.data<float>();668 const float* pd = dout.data<float>();669 float* pdq = dq.data<float>();670 float* pdk = dk.data<float>();671 float* pdv = dv.data<float>();672673 // Serial over heads: dk/dv rows are shared across q-heads under GQA.674 for (int64_t bh = 0; bh < B * n_heads; ++bh) {675 const int64_t b = bh / n_heads;676 const int64_t h = bh % n_heads;677 const int64_t hkv = h / rep;678 const float* P = pp + bh * T * T;679680 for (int64_t i = 0; i < T; ++i) {681 const float* qi = pq + (b * T + i) * Cq + h * hd;682 const float* doi = pd + (b * T + i) * Cq + h * hd;683 float* dqi = pdq + (b * T + i) * Cq + h * hd;684685 // dP_ij = dO_i · V_j ; dS = P ∘ (dP − D_i), and the FA2 identity686 // gives D_i = Σ_j dP_ij P_ij = dO_i · O_i in one pass.687 const float* oi = po + (b * T + i) * Cq + h * hd;688 float row_dot = 0.0f;689 for (int64_t d = 0; d < hd; ++d) row_dot += doi[d] * oi[d];690 for (int64_t j = 0; j < T; ++j) {691 const float p = P[i * T + j];692 if (p == 0.0f) continue;693 const float* vj = pv + (b * T + j) * Ckv + hkv * hd;694 const float* kj = pk + (b * T + j) * Ckv + hkv * hd;695 float* dvj = pdv + (b * T + j) * Ckv + hkv * hd;696 float* dkj = pdk + (b * T + j) * Ckv + hkv * hd;697698 float dp = 0.0f;699 for (int64_t d = 0; d < hd; ++d) dp += doi[d] * vj[d];700 float ds = p * (dp - row_dot) * scale;701 if (attn_softcap > 0.0f) {702 // chain through s' = cap·tanh(s/cap): recompute the raw703 // score, factor is 1 − tanh²704 float s = 0.0f;705 for (int64_t d = 0; d < hd; ++d) s += qi[d] * kj[d];706 const float t = std::tanh(s * scale / attn_softcap);707 ds *= 1.0f - t * t;708 }709710 for (int64_t d = 0; d < hd; ++d) {711 dvj[d] += p * doi[d];712 dqi[d] += ds * kj[d];713 dkj[d] += ds * qi[d];714 }715 }716 }717 }718}719720// ---- cross entropy ------------------------------------------------------------721722float cross_entropy(const Tensor& logits, const Tensor& targets, Tensor* dlogits) {723 const int64_t V = logits.size(1);724 const int64_t N = logits.size(0);725 const float* pl = logits.data<float>();726727 int64_t n_valid = 0;728 for (int64_t i = 0; i < N; ++i)729 if (token_at(targets, i) >= 0) ++n_valid;730 if (n_valid == 0) return 0.0f;731732 double total = 0.0;733 const float inv_n = 1.0f / float(n_valid);734 float* pd = dlogits ? dlogits->data<float>() : nullptr;735736 for (int64_t i = 0; i < N; ++i) {737 const int32_t tgt = token_at(targets, i);738 if (tgt < 0) continue;739 const float* row = pl + i * V;740 float m = row[0];741 for (int64_t j = 1; j < V; ++j) m = std::max(m, row[j]);742 double sum = 0.0;743 for (int64_t j = 0; j < V; ++j) sum += std::exp(double(row[j] - m));744 const double lse = double(m) + std::log(sum);745 total += lse - double(row[tgt]);746747 if (pd) {748 const float inv_sum = float(1.0 / sum);749 float* drow = pd + i * V;750 for (int64_t j = 0; j < V; ++j) {751 const float p = std::exp(row[j] - m) * inv_sum;752 drow[j] += (p - (j == tgt ? 1.0f : 0.0f)) * inv_n;753 }754 }755 }756 return float(total / double(n_valid));757}758759} // namespace forge::cpu760