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//3// Fused (flash-style) causal attention with GQA — forward + backward, f32.4// Nothing of size T² is ever written: the forward keeps the online-softmax5// state in registers and stores only L = m + log(l), one float per query6// row, which the backward uses to recompute P tile-by-tile. That is what7// lets context length scale — the unfused path in attention.metal needs8// B·H·T² floats of probabilities per layer (2.1 GB for batch 64 × T 1024 ×9// 8 heads), this needs B·H·T.10//11// Parallelization: one THREAD per output row —12// forward / dQ : one thread per (b, h, i) grid B·H·T13// dK,dV : one thread per (b, hkv, j) grid B·HKV·T14// so dK/dV accumulate privately per KV row (no atomics, deterministic) and15// each simdgroup's 32 threads are consecutive query rows of the same head,16// which makes their K/V reads a broadcast that the cache serves once.17//18// head_dim is a TEMPLATE parameter, not a runtime value: the per-thread19// q/o/accumulator arrays must be compile-time sized and fully unrolled or20// they land in thread-local (device-backed) memory instead of registers.21// Instantiated per supported head_dim below; the host falls back to the22// unfused kernel for other sizes.23//24// Online softmax (FlashAttention-2, RESEARCH.md §6):25// m_new = max(m, s) ; corr = exp(m − m_new) ; e = exp(s − m_new)26// l = l·corr + e ; o = o·corr + e·v ; divide by l once at the end27// m starts at −FLT_MAX, never −INFINITY: exp(−inf − (−inf)) is NaN.28#include <metal_stdlib>29using namespace metal;3031struct FlashParams {32 uint B, T, H, HKV;33 float scale;34 uint causal;35 uint window; // sliding window (keys kept, self incl.); 0 = full attention36};3738// ---------------------------------------------------------------- forward3940// Threadgroup = TGQ consecutive query rows of one (b, h). K/V arrive through41// threadgroup memory in blocks of BKV rows, staged cooperatively, so each42// K/V element is fetched from device once per threadgroup instead of once43// per thread. Causal blocks entirely above the diagonal are skipped outright44// (kb_lim), and only the diagonal block pays the per-element mask test.45enum : uint { TGQ = 64, BKV = 16 };4647template <uint HD>48kernel void flash_attn_fwd(device const float* Q [[buffer(0)]],49 device const float* K [[buffer(1)]],50 device const float* V [[buffer(2)]],51 device float* O [[buffer(3)]],52 device float* L [[buffer(4)]],53 constant FlashParams& p [[buffer(5)]],54 uint tgid [[threadgroup_position_in_grid]],55 uint tid [[thread_index_in_threadgroup]]) {56 threadgroup float Ks[BKV * HD];57 threadgroup float Vs[BKV * HD];5859 const uint q_blocks = (p.T + TGQ - 1) / TGQ;60 const uint qb = tgid % q_blocks;61 const uint h = (tgid / q_blocks) % p.H;62 const uint b = tgid / (q_blocks * p.H);6364 const uint i = qb * TGQ + tid;65 const bool active = (i < p.T) && (b < p.B);6667 const uint hkv = h / (p.H / p.HKV);68 const uint Cq = p.H * HD;69 const uint Ckv = p.HKV * HD;7071 float q[HD], o[HD];72 if (active) {73 device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;74#pragma clang loop unroll(full)75 for (uint d = 0; d < HD; ++d) {76 q[d] = qi[d];77 o[d] = 0.0f;78 }79 } else {80#pragma clang loop unroll(full)81 for (uint d = 0; d < HD; ++d) {82 q[d] = 0.0f;83 o[d] = 0.0f;84 }85 }8687 float m = -FLT_MAX;88 float l = 0.0f;8990 // Causal: this block's largest query index bounds the KV blocks we touch.91 const uint row_max = min(qb * TGQ + TGQ - 1, p.T - 1);92 const uint j_end = p.causal ? (row_max + 1) : p.T;93 const uint diag_start = p.causal ? (qb * TGQ) : p.T; // blocks below need no mask94 // Sliding window: this row's first visible key, and (uniform across the95 // threadgroup) the first KV block any row here can see — earlier blocks96 // are skipped outright, so compute scales with the window, not with T.97 const uint jmin = (p.window > 0 && i + 1 > p.window) ? i + 1 - p.window : 0;98 const uint tg_jmin =99 (p.window > 0 && qb * TGQ + 1 > p.window) ? qb * TGQ + 1 - p.window : 0;100 const uint jb_begin = (tg_jmin / BKV) * BKV;101102 for (uint jb = jb_begin; jb < j_end; jb += BKV) {103 const uint block_len = min(BKV, j_end - jb);104105 threadgroup_barrier(mem_flags::mem_threadgroup);106 for (uint e = tid; e < block_len * HD; e += TGQ) {107 const uint jr = e / HD;108 const uint d = e % HD;109 const ulong src = (ulong(b) * p.T + jb + jr) * Ckv + hkv * HD + d;110 Ks[jr * HD + d] = K[src];111 Vs[jr * HD + d] = V[src];112 }113 threadgroup_barrier(mem_flags::mem_threadgroup);114115 if (!active) continue;116 const bool needs_mask = p.causal && (jb + BKV > diag_start);117 for (uint jr = 0; jr < block_len; ++jr) {118 const uint j = jb + jr;119 if (needs_mask && j > i) break; // rest of the block is masked too120 if (j < jmin) continue; // below this row's window121122 threadgroup const float* kj = Ks + jr * HD;123 float s = 0.0f;124#pragma clang loop unroll(full)125 for (uint d = 0; d < HD; ++d) s = fma(q[d], kj[d], s);126 s *= p.scale;127128 const float m_new = max(m, s);129 const float corr = exp(m - m_new); // 0 on the first iteration130 const float e = exp(s - m_new);131 l = l * corr + e;132133 threadgroup const float* vj = Vs + jr * HD;134#pragma clang loop unroll(full)135 for (uint d = 0; d < HD; ++d) o[d] = fma(o[d], corr, e * vj[d]);136 m = m_new;137 }138 }139140 if (!active) return;141 const float inv = (l > 0.0f) ? (1.0f / l) : 0.0f;142 device float* oi = O + (ulong(b) * p.T + i) * Cq + h * HD;143#pragma clang loop unroll(full)144 for (uint d = 0; d < HD; ++d) oi[d] = o[d] * inv;145 // logsumexp; a fully-masked row would give -inf, guarded like the divide146 L[(ulong(b) * p.H + h) * p.T + i] = (l > 0.0f) ? (m + log(l)) : 0.0f;147}148149// --------------------------------------------------------------- backward150// D[b,h,i] = dO_i · O_i == rowsum(dP ∘ P) (FA2 identity), computed by151// attention_bwd_d_f32 in attention.metal and passed in here.152153template <uint HD>154kernel void flash_attn_bwd_dq(device const float* Q [[buffer(0)]],155 device const float* K [[buffer(1)]],156 device const float* V [[buffer(2)]],157 device const float* dO [[buffer(3)]],158 device const float* L [[buffer(4)]],159 device const float* D [[buffer(5)]],160 device float* dQ [[buffer(6)]],161 constant FlashParams& p [[buffer(7)]],162 uint gid [[thread_position_in_grid]]) {163 const uint i = gid % p.T;164 const uint h = (gid / p.T) % p.H;165 const uint b = gid / (p.T * p.H);166 if (b >= p.B) return;167168 const uint hkv = h / (p.H / p.HKV);169 const uint Cq = p.H * HD;170 const uint Ckv = p.HKV * HD;171172 device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;173 device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;174175 float q[HD], dq[HD], go[HD];176#pragma clang loop unroll(full)177 for (uint d = 0; d < HD; ++d) {178 q[d] = qi[d];179 go[d] = doi[d];180 dq[d] = 0.0f;181 }182183 const float li = L[(ulong(b) * p.H + h) * p.T + i];184 const float di = D[(ulong(b) * p.H + h) * p.T + i];185 const uint jmax = p.causal ? i : (p.T - 1);186 const uint jmin = (p.window > 0 && i + 1 > p.window) ? i + 1 - p.window : 0;187188 for (uint j = jmin; j <= jmax; ++j) {189 device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * HD;190 device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * HD;191192 float s = 0.0f, dp = 0.0f;193#pragma clang loop unroll(full)194 for (uint d = 0; d < HD; ++d) {195 s = fma(q[d], kj[d], s);196 dp = fma(go[d], vj[d], dp);197 }198 const float prob = exp(s * p.scale - li); // recomputed, never stored199 const float ds = prob * (dp - di) * p.scale;200#pragma clang loop unroll(full)201 for (uint d = 0; d < HD; ++d) dq[d] = fma(ds, kj[d], dq[d]);202 }203204 device float* dqi = dQ + (ulong(b) * p.T + i) * Cq + h * HD;205#pragma clang loop unroll(full)206 for (uint d = 0; d < HD; ++d) dqi[d] += dq[d];207}208209// dK and dV are separate kernels on purpose. Combined, one thread holds210// dk[HD] + dv[HD] and spills: measured 150 ms -> 118 ms for gpt-10m just by211// moving the read-only k/v out of registers, so the accumulators matter too.212// Split, each thread carries a single HD-sized accumulator; the price is213// recomputing the q·k dot in both kernels, which is cheaper than the spill.214215template <uint HD>216kernel void flash_attn_bwd_dv(device const float* Q [[buffer(0)]],217 device const float* K [[buffer(1)]],218 device const float* dO [[buffer(2)]],219 device const float* L [[buffer(3)]],220 device float* dV [[buffer(4)]],221 constant FlashParams& p [[buffer(5)]],222 uint gid [[thread_position_in_grid]]) {223 const uint j = gid % p.T;224 const uint hkv = (gid / p.T) % p.HKV;225 const uint b = gid / (p.T * p.HKV);226 if (b >= p.B) return;227228 const uint rep = p.H / p.HKV;229 const uint Cq = p.H * HD;230 const uint Ckv = p.HKV * HD;231 device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * HD;232233 float dv[HD];234#pragma clang loop unroll(full)235 for (uint d = 0; d < HD; ++d) dv[d] = 0.0f;236237 const uint imin = p.causal ? j : 0;238 // Sliding window: only queries within `window` of j ever saw it.239 const uint imax = (p.window > 0) ? min(p.T, j + p.window) : p.T;240 for (uint r = 0; r < rep; ++r) {241 const uint h = hkv * rep + r;242 device const float* Lh = L + (ulong(b) * p.H + h) * p.T;243 for (uint i = imin; i < imax; ++i) {244 device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;245 device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;246 float s = 0.0f;247#pragma clang loop unroll(full)248 for (uint d = 0; d < HD; ++d) s = fma(qi[d], kj[d], s);249 const float prob = exp(s * p.scale - Lh[i]);250#pragma clang loop unroll(full)251 for (uint d = 0; d < HD; ++d) dv[d] = fma(prob, doi[d], dv[d]);252 }253 }254255 device float* dvj = dV + (ulong(b) * p.T + j) * Ckv + hkv * HD;256#pragma clang loop unroll(full)257 for (uint d = 0; d < HD; ++d) dvj[d] += dv[d];258}259260template <uint HD>261kernel void flash_attn_bwd_dk(device const float* Q [[buffer(0)]],262 device const float* K [[buffer(1)]],263 device const float* V [[buffer(2)]],264 device const float* dO [[buffer(3)]],265 device const float* L [[buffer(4)]],266 device const float* D [[buffer(5)]],267 device float* dK [[buffer(6)]],268 constant FlashParams& p [[buffer(7)]],269 uint gid [[thread_position_in_grid]]) {270 const uint j = gid % p.T;271 const uint hkv = (gid / p.T) % p.HKV;272 const uint b = gid / (p.T * p.HKV);273 if (b >= p.B) return;274275 const uint rep = p.H / p.HKV;276 const uint Cq = p.H * HD;277 const uint Ckv = p.HKV * HD;278 device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * HD;279 device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * HD;280281 float dk[HD];282#pragma clang loop unroll(full)283 for (uint d = 0; d < HD; ++d) dk[d] = 0.0f;284285 const uint imin = p.causal ? j : 0;286 const uint imax = (p.window > 0) ? min(p.T, j + p.window) : p.T;287 for (uint r = 0; r < rep; ++r) {288 const uint h = hkv * rep + r;289 device const float* Lh = L + (ulong(b) * p.H + h) * p.T;290 device const float* Dh = D + (ulong(b) * p.H + h) * p.T;291 for (uint i = imin; i < imax; ++i) {292 device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;293 device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;294 float s = 0.0f, dp = 0.0f;295#pragma clang loop unroll(full)296 for (uint d = 0; d < HD; ++d) {297 s = fma(qi[d], kj[d], s);298 dp = fma(doi[d], vj[d], dp);299 }300 const float prob = exp(s * p.scale - Lh[i]);301 const float ds = prob * (dp - Dh[i]) * p.scale;302#pragma clang loop unroll(full)303 for (uint d = 0; d < HD; ++d) dk[d] = fma(ds, qi[d], dk[d]);304 }305 }306307 device float* dkj = dK + (ulong(b) * p.T + j) * Ckv + hkv * HD;308#pragma clang loop unroll(full)309 for (uint d = 0; d < HD; ++d) dkj[d] += dk[d];310}311312// ---------------------------------------------------------- instantiations313// Every config in configs/ uses head_dim 64 (d_model / n_heads); the others314// are here so the fused path covers common variants. Unsupported sizes fall315// back to the unfused kernels on the host side.316#define INSTANTIATE_FLASH(HD) \317 template [[host_name("flash_attn_fwd_f32_hd" #HD)]] kernel void \318 flash_attn_fwd<HD>(device const float*, device const float*, device const float*, \319 device float*, device float*, constant FlashParams&, uint, \320 uint); \321 template [[host_name("flash_attn_bwd_dq_f32_hd" #HD)]] kernel void \322 flash_attn_bwd_dq<HD>(device const float*, device const float*, \323 device const float*, device const float*, \324 device const float*, device const float*, device float*, \325 constant FlashParams&, uint); \326 template [[host_name("flash_attn_bwd_dv_f32_hd" #HD)]] kernel void \327 flash_attn_bwd_dv<HD>(device const float*, device const float*, \328 device const float*, device const float*, device float*, \329 constant FlashParams&, uint); \330 template [[host_name("flash_attn_bwd_dk_f32_hd" #HD)]] kernel void \331 flash_attn_bwd_dk<HD>(device const float*, device const float*, \332 device const float*, device const float*, \333 device const float*, device const float*, device float*, \334 constant FlashParams&, uint);335336INSTANTIATE_FLASH(16)337INSTANTIATE_FLASH(32)338INSTANTIATE_FLASH(48)339INSTANTIATE_FLASH(64)340INSTANTIATE_FLASH(80)341INSTANTIATE_FLASH(96)342INSTANTIATE_FLASH(128)343