spb/forge Public MIT

1<!-- Author: Simon-Pierre Boucher — contact@spboucher.ai -->23# RESEARCH.md — Forge: Metal/C++ LLM Training Framework45Research phase completed 2026-07-31 (6 parallel web-research passes + local toolchain probes).6This file records the API signatures, tile sizes, design decisions, and pitfalls that gate7implementation. **Update this file whenever new research is done during the project.**89---1011## 0. Local environment (verified empirically, 2026-07-31)1213| Item | Value |14|---|---|15| Dev machine | Apple **M5 Max**, 40-core GPU, macOS **27.0** (build 26A5388g) |16| Primary deploy target | **M3 Ultra** (Apple9 family), 80-core GPU, ~28.3 TFLOPS FP32/FP16, 819 GB/s, 96 GB unified |17| Xcode | 26.6 (build 17F113), Apple clang 21.0.0 |18| Metal compiler | `xcrun -sdk macosx metal` — accepts `-std=metal3.0/3.1/3.2/4.0` (no 3.3); default compiles fine |19| CMake | 4.3.2 |20| Local probes (compiled OK from `-std=metal3.1` up) | `bfloat`, `simdgroup_bfloat8x8`, `simdgroup_half8x8` → `simdgroup_float8x8` mixed MMA, device `atomic_float` `fetch_add` |21| Local reference source | MLX checkout at `~/Desktop/other/mlx` (incl. vendored metal-cpp under `build/DerivedData/.../mlx-swift/Source/Cmlx/metal-cpp`) |2223**Decision: compile kernels with `-std=metal3.2`** (macOS 15+ baseline: full simdgroup_matrix,24bfloat, atomic_thread_fence) and use the **classic Metal 3 compute API** via metal-cpp. Metal 425(`MTL4*`, MTLTensor, cooperative tensors) is a parallel opt-in API — a future seam, not a dependency.2627---2829## 1. metal-cpp — setup and memory management3031Source: https://github.com/apple/metal-cpp (header-only, C++17+; releases track OS — use the macOS 26.x/27 release). Vendor into `third_party/metal-cpp`.3233In exactly **one** .cpp file:34```cpp35#define NS_PRIVATE_IMPLEMENTATION36#define CA_PRIVATE_IMPLEMENTATION37#define MTL_PRIVATE_IMPLEMENTATION38#include <Foundation/Foundation.hpp>39#include <Metal/Metal.hpp>40#include <QuartzCore/QuartzCore.hpp>41```42Link frameworks: **Foundation, Metal, QuartzCore** (MetalKit not needed for compute).4344**Ownership (Cocoa rules, no ARC):**45- `new*` / `alloc` / `copy` / `Create*` return **retained** objects → caller releases. That covers `MTL::CreateSystemDefaultDevice()`, `newCommandQueue()`, `newBuffer()`, `newLibrary()`, `newFunction()`, `newComputePipelineState()`.46- Everything else is **autoreleased** — notably `queue->commandBuffer()` and `cmdBuf->computeCommandEncoder()`. They die when the enclosing `NS::AutoreleasePool` drains.47- Smart pointers: `NS::TransferPtr(p)` adopts retained results (no extra retain); `NS::RetainPtr(p)` retains autoreleased/borrowed objects.48- **Wrap each training step (or N steps) in an `NS::AutoreleasePool`** (`alloc()->init()` … `drain()`), and one in `main()` — otherwise autoreleased command buffers accumulate unboundedly. Debug with `OBJC_DEBUG_MISSING_POOLS=YES`.4950**Key compute API (verified against metal-cpp headers):**51```cpp52MTL::Device* MTL::CreateSystemDefaultDevice();53CommandQueue* Device::newCommandQueue();54Buffer* Device::newBuffer(NS::UInteger length, MTL::ResourceOptions);          // MTL::ResourceStorageModeShared55Buffer* Device::newBuffer(const void* ptr, NS::UInteger len, MTL::ResourceOptions, void(^dealloc)(void*, NS::UInteger)); // NoCopy: page-aligned56Library* Device::newLibrary(const NS::String* filepath, NS::Error** err);      // load .metallib by path (CLI tool: do NOT rely on newDefaultLibrary)57Function* Library::newFunction(const NS::String* name, const MTL::FunctionConstantValues*, NS::Error**);58ComputePipelineState* Device::newComputePipelineState(const MTL::Function*, NS::Error**);59ComputeCommandEncoder* CommandBuffer::computeCommandEncoder(MTL::DispatchType);  // Serial | Concurrent60void Encoder::setComputePipelineState(const MTL::ComputePipelineState*);61void Encoder::setBuffer(const MTL::Buffer*, NS::UInteger offset, NS::UInteger index);62void Encoder::setBytes(const void* bytes, NS::UInteger length, NS::UInteger index);   // ≤ ~4 KB, for param structs63void Encoder::setThreadgroupMemoryLength(NS::UInteger length, NS::UInteger index);64void Encoder::dispatchThreads(MTL::Size grid, MTL::Size tg);          // non-uniform TGs, supported on all M-series — prefer for elementwise65void Encoder::dispatchThreadgroups(MTL::Size tgs, MTL::Size tg);      // for tiled kernels66void CommandBuffer::commit();  addCompletedHandler(const MTL::HandlerFunction&);   // std::function<void(CommandBuffer*)>67CFTimeInterval CommandBuffer::GPUStartTime() / GPUEndTime();          // per-command-buffer timing, read in completed handler68NS::UInteger ComputePipelineState::maxTotalThreadsPerThreadgroup();   // can be < 1024 (register pressure) — always query69NS::UInteger ComputePipelineState::threadExecutionWidth();            // 32 on Apple GPUs70```7172---7374## 2. Build system: .metal → .metallib in CMake7576Two-step pipeline as CMake custom commands (Ninja/Makefiles generator):77```78xcrun -sdk macosx metal -std=metal3.2 -O2 [-gline-tables-only -frecord-sources in Debug] -c foo.metal -o foo.air79xcrun -sdk macosx metallib foo.air bar.air ... -o forge.metallib80```81Load at runtime with `device->newLibrary(path, &err)`. Fast math is **ON by default**82(`-fmetal-math-mode=fast`: assumes no NaN/INF, reassociation, auto-FMA contraction). Keep it on83for throughput but use `metal::precise::` selectively in numerically sensitive spots and account84for FMA contraction in the 1e-4 CPU-parity tests.8586---8788## 3. MSL essentials (from MSL Specification v4.1, 2026-06-04)8990Prelude for all kernels: `#include <metal_stdlib>` + `using namespace metal;` (umbrella header91covers simdgroup, simdgroup_matrix, atomics, compute barriers).9293### simdgroup_matrix (the perf-critical path)94- Types: `simdgroup_half8x8`, `simdgroup_bfloat8x8` (MSL 3.1+), `simdgroup_float8x8`. **8×8 is the only shape.** Available on Apple7+ (all M-series). Each thread holds 2 elements (`vec<T,2>` via `.thread_elements()`).95- Signatures (verbatim; threadgroup and device overloads):96```metal97void simdgroup_load (thread simdgroup_matrix<T,8,8>& d, const {threadgroup|device} T* src,98                     ulong elements_per_row = 8, ulong2 matrix_origin = 0, bool transpose_matrix = false);99void simdgroup_store(thread simdgroup_matrix<T,8,8>  a, {threadgroup|device} T* dst,100                     ulong elements_per_row = 8, ulong2 matrix_origin = 0, bool transpose_matrix = false);101void simdgroup_multiply_accumulate(d, a, b, c);   // d = a*b + c102simdgroup_matrix<T,8,8> make_filled_simdgroup_matrix<T,8,8>(T value);103```104- **Mixed precision**: the spec documents only same-T operands, but **half×half→float accumulate compiles and is production-stable** (llama.cpp `kernel_mul_mm`: `simdgroup_half8x8` inputs, `simdgroup_float8x8` accumulators; confirmed compiling locally). MLX instead converts inputs to float fragments at load and runs all-float MMA (spec-clean). **Decision: MLX-style all-float fragments first (correctness), llama.cpp-style mixed MMA behind a compile-time switch (perf).**105- All simdgroup_matrix ops require **uniform SIMD-group control flow** — UB otherwise.106107### SIMD intrinsics (Apple7+ hardware reductions — all M-series)108`simd_sum, simd_product, simd_max, simd_min, simd_prefix_exclusive_sum` — one instruction, result broadcast to active lanes. `simd_shuffle(_up/_down/_rotate)`, `simd_broadcast(_first)`, `simd_ballot`, `simd_all/any`. SIMD width = **32**. Caveats: `simd_shuffle_xor` is **integer-only**; **bfloat is excluded** from all `simd_*` intrinsics; reductions cover *active* threads only.109110### Threadgroup memory & barriers111- Max **32 KB** per threadgroup (API limit; all Apple Silicon). Max 1024 threads/TG (but query the pipeline's `maxTotalThreadsPerThreadgroup`).112- Declare tiles inside the kernel (`threadgroup half As[...];`) or as `[[threadgroup(n)]]` arg (runtime-sized).113- `threadgroup_barrier(mem_flags::mem_threadgroup)` around shared-tile writes; `simdgroup_barrier(mem_flags::mem_none)` as a cheap scheduling fence between `simdgroup_load` batches and MMA (llama.cpp/MLX idiom).114115### Atomics116`atomic_float` (device memory only, Metal 3+): load/store/exchange/CAS/**fetch_add/sub** only, `memory_order_relaxed` only (below Metal 4.1). No float min/max atomics. Gradient accumulation via atomics is legal but nondeterministic → **use deterministic reduction trees / two-kernel splits instead** (required for deterministic mode anyway).117118### Numerics119- `bfloat`: MSL 3.1+, truncated f32; NOT implicitly convertible to/from half; no simd intrinsic support. Fine as storage + convert-to-float compute.120- `fast::exp` error grows as `3 + floor(|2x|)` ulp → **always subtract the row max before exp**. MLX folds `scale * M_LOG2E_F` and uses `fast::exp2` (exp2 is the HW primitive) — adopt this.121- `fma()` always correctly rounded; compiler auto-contracts `a*b+c` — a source of CPU/GPU last-ulp differences (tolerances absorb it).122123---124125## 4. Execution model & performance practices (Apple Silicon / Apple9)126127Architecture facts: 128 ALUs/core (4×32-wide), FP16 and FP32 FMA at the **same** 128 FMA/core/cycle rate — f16 wins via registers/bandwidth and (family 9) FP16∥FP32∥INT dual-issue "up to 2× ALU". M3+ **dynamic caching**: registers/threadgroup/stack share unified on-chip caches; register count no longer sets occupancy; HW auto-throttles occupancy. Consequences: (a) register-heavy small-K GEMM configs win on M3+ (see §5); (b) **threadgroup memory used as a mere read-cache of device data can be a pessimization on family 9** — use it for reuse + layout transformation only; (c) async copies (`simdgroup_event`) help M1/M2 but **hurt M3+** — don't use them.128129**Command submission (the rules Forge's device layer enforces):**1301. One `MTLCommandQueue`. Batch the whole training step (fwd+bwd+optimizer, 100s of dispatches) into **one command buffer**, few **serial** compute encoders. Within a serial encoder, dispatch N+1 sees dispatch N's writes — no barriers needed (`memoryBarrier` is ignored on serial encoders). MLX batches all ops of an eval this way; Candle measured ~13× per-dispatch overhead win.1312. **Never `waitUntilCompleted` per-op.** Sync at loss readback only: `addCompletedHandler` + semaphore/`MTLSharedEvent`, keep 1–2 steps in flight, double-buffer the batch-input/loss buffers.1323. `MTLStorageModeShared` everywhere. Coherency is defined **at command-buffer boundaries**: CPU writes before `commit` are visible; GPU writes are readable after completion. Never CPU-write a buffer an in-flight command buffer may touch.1334. **Buffer pool**: `newBuffer` is a kernel VM allocation — pool by size class, recycle **only after** the referencing command buffer completes (MLX issue #3689 use-after-free class). Suballocate at **256-byte alignment** (always valid; setBuffer offset needs 4 B on Apple GPUs but 256 is free insurance and vector-load friendly).1345. Dataloader: `mmap` the token file and wrap with `newBuffer(bytesNoCopy:)` — requires **16384-byte page** alignment/multiple (Apple Silicon page size). True zero-copy.1356. Threadgroup sizing: multiples of 32; 128–256 threads for GEMM tiles; 256–1024 for row reductions (simd_sum first, tiny cross-simd threadgroup step second — shuffle bandwidth is 2× threadgroup-memory bandwidth); vectorize memory-bound kernels with `float4`/`half4` (elementwise/AdamW/norm are pure-bandwidth, target ~800 GB/s).1367. Specialize kernels with **function constants** (`MTLFunctionConstantValues`) — alignment flags, has_bias, activation type, tile params — and cache pipelines by (kernel name + constants). No runtime branches for these.1378. Profiling: `GPUStartTime/GPUEndTime` per command buffer always-on for tokens/sec; `MTLCounterSampleBuffer` (timestamp counter set, encoder-boundary sampling — check `supportsCounterSampling`) for per-kernel tuning; `MTL_DEBUG_LAYER=1` in tests.138139---140141## 5. GEMM design (chosen tiles + justification)142143Reference implementations studied: MLX STEEL (`mlx/backend/metal/kernels/steel/gemm/`), llama.cpp `kernel_mul_mm`, Philip Turner's metal-flash-attention GEMM.144145**Structure (STEEL pattern):** `BlockLoader` stages device→threadgroup with one 16-byte vector read per thread per step, **+16 bytes row padding** against bank conflicts (`tgp_padding = 16/sizeof(T)`); tiles consumed as 8×8 `simdgroup` fragments, K-inner loop in steps of 8 with `simdgroup_barrier(mem_none)` fences; **f32 accumulator fragments always**; single-buffered tiles (no double buffering — multiple resident TGs hide latency); epilogue → predicated `store_result_safe` on edges.146147**Tile configs (starting points, to be autotuned on target):**148149| Case | BM×BN×BK | WM×WN (simdgroups) | Source/justification |150|---|---|---|---|151| f16/bf16 default (large M·N) | 64×64×16 | 1×2 (64 thr) … 2×2 (128 thr) | MLX large-device table |152| f16 NT (fwd `X·Wᵀ`) | 64×32×32 | 2×2 | MLX nt case |153| f16 NN (bwd `dX = dY·W`) | 64×64×16 or 32×64×16 | 1×2 | MLX nn cases |154| f32 | 32×64×16 (nt) / 64×32×32 (nn) | 1×2 / 2×2 | MLX f32 table |155| M3+ alternative to benchmark | 32×32×8, 1 simdgroup/TG | 1×1 | MFA Apple9 finding: dynamic caching favors register-heavy small-K |156157- **Transpose variants**: one templated kernel with `transpose_a/transpose_b` bools → nn/nt/tn/tt instantiations; detect stride-1 dim (MLX `check_transpose`), never materialize transposes. Training needs **nt (forward), nn (dX), tn (dW = dYᵀ·X)**.158- **Edges**: function constants `align_M/align_N/align_K` at pipeline creation; K-remainder tile handled FIRST with predicated `load_safe` (zero-fill OOB), then aligned main loop; per-element predicated fragment stores. No separate edge kernels.159- **Split-K** for dW-shaped GEMMs (small M,N, huge K = B·T): MLX routing `tiles(M,16)*tiles(N,16) ≤ 2048 (Ultra) && K/16 ≥ 8 && K ≥ max(M,N)`; ≤32 partitions accumulate into f32 scratch `[partitions, M, N]` + tiny reduce kernel. No atomics.160- **GEMV** for generation (min(M,N)=1): rows-per-simdgroup, float4 dots, `simd_shuffle_down` tree / `simd_sum` (llama.cpp `kernel_mul_mv` pattern).161- Grid swizzle for large matrices (`swizzle_log` 1–2), serpentine fragment iteration for register reuse.162- **Expected perf**: MFA/STEEL-class kernels hit 80–93% ALU on M-series → ~23–26 TFLOPS f16 on M3 Ultra; naive tiled non-simdgroup ≈ 35–40% — hence the roadmap naive → tiled → simdgroup.163- **Operand precision on the simdgroup path is a small win** (`tests/bench_precision.cpp`, M5 Max, f32 accumulator throughout): f16 and bf16 operands give **+18–22%** over f32 (e.g. 9.9 → 12.4 TFLOPs), not the 2× one might expect from Apple's "2× ALU on family 9" framing. This matches metal-benchmarks' claim that f16 and f32 FMA run at the same rate and the gain comes from register/bandwidth pressure. Mixed precision on *this* path is therefore a memory feature. Through MPP (§5b) the same operands are worth 4.9×.164- **Measured (this repo, `tests/bench_matmul.cpp`, M5 Max 40-core, f32)**: naive 1.3–1.5 · 16×16-tiled 2.5 · simdgroup 9.5–10.8 TFLOPs. The `nt` forward shapes and `nn` dX shape all land ≈10.6; `tn` (dW = dYᵀ·X, K = B·T = 65536) lags at 4.7 — few threadgroups × very long K, which is exactly the split-K case (MLX routes it that way; not yet implemented here).165- llama.cpp codegen gotcha: `sa[64*ib + ...]` (subscript) was "massively slower" than `*(sa + 64*ib + ...)` (pointer deref) — watch for it.166167---168169## 5b. Metal Performance Primitives (`matmul2d`) — measured 4.9× over hand-written MMA170171Probed and benchmarked locally 2026-07-31 (`src/kernels/matmul_mpp.metal`,172`tests/bench_precision.cpp`, `tests/mppcheck.cpp`). **This is the largest single173efficiency finding for this project and it changes the roadmap.**174175`MetalPerformancePrimitives.framework` ships in the macOS 26.5 SDK176(`Headers/MPPTensorOpsMatMul2d.h`, 642 lines) and compiles under `-std=metal4.0`.177`mpp::tensor_ops::matmul2d` is the cooperative-tensor matmul that targets the178per-core **neural accelerators** on M5-class hardware.179180Measured on M5 Max, 64×32 tile, 4 simdgroups, f32 accumulate, versus this repo's181own simdgroup_matrix GEMM at identical shapes:182183| shape | simdgroup f32 | simdgroup f16 | **MPP f32** | **MPP f16** |184|---|---|---|---|---|185| 2048³ | 9.3 T | 10.6 T | **14.9 T** | **51.6 T** |186| 4096³ | 9.9 T | 11.8 T | **14.1 T** | **44.1 T** |187| 65536×512×1408 | 9.9 T | 12.4 T | **15.3 T** | **22.8 T** |188189So **1.5× for f32 and 4.3–4.9× for f16**. Both verified numerically against the CPU190reference (`tests/mppcheck.cpp`): f32 is bit-exact, f16 is 3.8e-06 against a191reference fed the same rounded inputs, and every output element is nonzero (i.e. the192full K reduction really happens).193194Note this contradicts the "Rigel" paper's M4 Max finding that `matmul2d` still runs on195the shader cores and loses to a hand-fused GEMM — M5 evidently has the accelerator196hardware that M4 lacked. Re-measure per generation; do not assume.197198**API notes that cost time to work out (the header's own example is stale):**199- Bind ordinary buffers and build tensors *in-kernel* with the `tensor_inline`200  descriptor: `tensor<device T, dextents<int32_t,2>, tensor_inline> A(ptr, dextents<int32_t,2>(cols, rows))`. The default `tensor_handle` descriptor wraps an opaque handle that only a host-side `MTLTensor` can supply, so `tensor_inline` avoids all MTLTensor plumbing.201- Extents are **(columns, rows)**, and `slice(x, y)` takes column-then-row.202- Element type must be **non-const** and one of `uint8_t/int8_t/uint4b_format/int4b_format/float/half/bfloat` — a `const float` tensor fails a `static_assert`.203- The destination is a `cooperative_tensor`; zero it via `get_capacity()` +204  **`is_valid_element(i)`**. The header comment says `get_mask(i)`, which does not205  exist in this SDK.206- `#pragma unroll full` is not valid MSL here; use `#pragma clang loop unroll(full)`.207- Requires `-std=metal4.0`, so it is a per-file build flag (CMake keys off the208  `_mpp` filename suffix) and macOS 26+; the simdgroup_matrix path stays the209  portable fallback for macOS 15.210211**Transposes work, so this covers training, not just inference.** `transpose_left` /212`transpose_right` are descriptor fields (hence template parameters), and nn / nt / tn213are all bit-exact against the CPU reference — i.e. the forward `X·Wᵀ`, the `dX = dY·W`214and the `dW = dYᵀ·X` of the backward are all served. Two gotchas: the operand extents215must be swapped to match the transpose flag, and `slice()` is (column, row) so a216transposed operand is sliced the other way round. Also `op.template217get_destination_cooperative_tensor<...>()` once the kernel is itself a template.218219**Implication for the roadmap.** Mixed precision was previously judged a220memory-only win (f16 operands buy just +18–22% on the simdgroup path — measured,221see below). Through MPP the same f16 operands are worth **4.9×**. That reverses222the priority: f16/bf16 plumbing is now the highest-value remaining work, because223it is the gate to the accelerator path, not merely a way to halve activations.224225## 6. Fused attention design (flash-style)226227Algorithm: FlashAttention-2 (arXiv 2307.08691). Online softmax per Q-block over KV-blocks j:228```229m_new = max(m_old, rowmax(S_j))            S_j = (Q Kᵀ) * scale   (f32)230P~ = exp2(S_j - m_new)                     (scale folded: scale * M_LOG2E_F, fast::exp2)231factor = exp2(m_old - m_new)               (≤ 1; the paper's diag(...)^-1 is a known typo)232l = l * factor + rowsum(P~)233O = O * factor + P~ · V_j                  (O in f32 registers; divide by l ONCE after the loop)234L = m + log(l)                             (one f32/row, saved for backward)235```236**Layout (MLX steel_attention + llama.cpp flash_attn_ext, adapted):** one threadgroup per (Q-block, head, batch); 4 simdgroups, **split-Q** — each simdgroup owns its Q rows so row statistics need only `simd_max`/`simd_sum` (zero cross-warp softmax traffic). Tiles: **BQ=32, BK=32 (BK=16 at head_dim 128), head dims {64, 80, 128}** (MLX); K loaded via `simdgroup_load(..., transpose=true)`; PV matmul non-transposed. m/l/factors/S/O all **f32** even with f16 Q/K/V. K/V threadgroup tiles share one padded allocation.237238- **Causal masking**: `kb_lim = ceil((q_max)/BK)` skips fully-masked KV blocks entirely; only diagonal-band blocks (`kb >= kb_min_causal`) pay per-element masking. Mask value = **`-FLT_MAX` / finite_min, never -INFINITY** (`exp(-inf - -inf) = NaN`). Guard final divide: `l == 0 → output 0`.239- **GQA**: `kv_head = q_head / (n_heads / n_kv_heads)` index mapping — no repeat_kv copies.240- **The `D` term is free and worth a lot.** `D_i = rowsum(dP ∘ P)` — needed by both backward kernels — equals `dO_i · O_i` exactly, since `Σ_j P_ij (dO_i · V_j) = dO_i · (Σ_j P_ij V_j) = dO_i · O_i`. Computing it per (i,j) pair instead of once per query row makes the dk/dv kernel O(T³·hd). Measured in this repo: hoisting it into a one-thread-per-row preprocess kernel took a training step from **18.4k to 57.8k tokens/sec (3.1×)** with bit-identical losses, even in the *unfused* attention path.241- **Measured (this repo, `tests/bench_attention.cpp`, M5 Max, f32, causal), scalar vs simdgroup_matrix-tiled**, at gpt-10m shapes (B64 T512 H6 hd64) and gpt-25m shapes (B64 T1024 H8):242243| | fwd gpt-10m | fwd gpt-25m | bwd gpt-10m | bwd gpt-25m |244|---|---|---|---|---|245| scalar (one thread/row) | 8.0 ms (1.6 TF) | 40.9 ms (1.7 TF) | 107 ms | 561 ms |246| MMA-tiled | **2.3 ms (5.7 TF)** | **10.6 ms (6.5 TF)** | 9.8 ms | 45.5 ms |247| MMA-tiled, dK/dV split | — | — | **7.05 ms** | **32.9 ms** |248249  So 3.6× on the forward and **10.9–12.3× on the backward**. End-to-end that took gpt-10m training from 20.1k to 35.2k tokens/sec. The backward gains more than the forward because the scalar backward was also spilling (§8 6b) — MMA fixes the arithmetic *and* the register pressure at once, since an 8×8 fragment is 2 floats per lane where the scalar version held 4 arrays of head_dim.250- **Ecosystem survey (2026-07-31): almost nobody has a fused attention backward on Metal.** MLX's `ScaledDotProductAttentionVJP::use_fallback` returns `true` unconditionally and `eval_gpu` throws `"NYI"`; its Metal *forward* also refuses the fused path under grad tracing (`if (is_training) return true;`), and its fused VJP is CUDA/cuDNN-only. llama.cpp's Metal backend does not support `GGML_OP_FLASH_ATTN_BACK` (falls to `default: return false`). PyTorch MPS and Candle ship hand-written MSL forwards only. tinygrad's FA backward is AMD-only. Burn/CubeCL has a 3-kernel backward that is still a scalar scaffold (no tiling, no MMA). The only production prior art is philipturner's MFA and the third-party `mlx-mfa` package, plus an unmerged MLX PR (#3241).251- **Both production references use exactly the structure implemented here**: (i) `delta = rowsum(dO ∘ O)` precomputed in **fp32** outside the main kernels, (ii) a dQ kernel with grid `(NQ, H, B)`, one Q tile per threadgroup, looping KV and accumulating dQ in registers, (iii) a dKV kernel with grid `(NK, H_kv, B)`, one KV tile per threadgroup, looping Q. **No atomics anywhere.** They additionally work in the log2 domain (`P = exp2(S·scale·log2e − L_log2)`); the elementwise softmax step is ~1/head_dim of the MMA work here, so that is a micro-optimisation at these shapes.252- **Their tuned tile configs, and why BK=32 lost here.** MLX PR #3241 and mlx-mfa both tune D=64 dQ to `BQ=32, BK=32, WM=4` (mlx-mfa reports retuning dQ from BK=64 to 32 was worth 4–14%). Tried it: **fwd 2.27 → 3.63 ms, bwd 9.66 → 11.16 ms — worse.** BK=32 takes K+V tiles from 8.7 to 17.4 KB, so with the S/P tile the kernel crosses 22 KB and residency drops from 2 threadgroups per core to 1. Their configs fit because they alias temporally-disjoint threadgroup buffers (PR #3241 notes aliasing a reduction buffer over `Q_smem+dO_smem`: 23,040 → 14,848 bytes, doubling occupancy). **Conclusion: BK=32 is only reachable after implementing that aliasing — the tile size and the memory plan have to be tuned together, not independently.**253- **Register pressure is the binding constraint at head_dim 128**, not bandwidth: PR #3241 measures D=128 dKV at BQ=32/WM=2 needing ~338 registers/thread (over the 256 limit, spills) versus ~202 at BQ=16/WM=2. Worth knowing before extending past hd=64.254- **A fused backward is not automatically faster.** PR #3241's own numbers on M3 Max: fused beats MLX's unfused VJP only for causal, D≤96 (1.17–1.37×); dense is 0.36–0.70×, i.e. slower, because the backward needs ~2.5× the forward's FLOPs and 32×32 tiles underuse the MMA units next to large-tile GEMMs. The payoff there is memory (95% peak reduction at L=4096). This does not contradict the 10.9× measured here, which compares fused-scalar against fused-MMA — both already avoid the T² tensor.255- **Backward (training)**: forward saves `L`. **Key finding: MLX does not use fused attention for training** (fused VJP is `NYI`; comment says unfused is faster on Metal at their configs). Plan: **M4/M5 milestone = fused forward + unfused-composed backward validated vs CPU; stretch = MFA-style fused backward** — two kernels, no atomics: dQ-kernel parallel over Q rows `(3D+5)N²` FMAs, dK/dV-kernel parallel over KV rows `(4D+5)N²` FMAs, both recompute `P = exp(S - L)` per tile, plus a cheap preprocess `D = rowsum(dO ∘ O)`. MFA measured fwd 91–94% ALU on M3, fwd+bwd ~69–71%.256- Decode path (`forge generate`): sdpa_vector-style kernel — 1 Q row, 32-thread strided KV walk, scalar online update, split-K two-pass over long KV.257258---259260## 7. Training loop (nanoGPT / llm.c distilled)261262**Hyperparameters (defaults for configs):** AdamW β=(0.9, 0.95), eps 1e-8, weight_decay 0.1 **only on dim≥2 params** (weights/embeddings; never biases/norm gains); grad clip 1.0 (global norm); linear warmup `lr*(it+1)/(warmup+1)` → cosine to `min_lr = lr/10`; init: normal(0, 0.02), residual projections `0.02/sqrt(2*n_layers)`; tied wte/lm_head (one grad buffer — embedding and head grads accumulate into the same tensor); vocab padded to a multiple of 64 (matmul tile).263264**AdamW exact (llm.c, PyTorch-compatible):** t is 1-based;265```266m = β1·m + (1-β1)·g ;  v = β2·v + (1-β2)·g²267m̂ = m/(1-β1ᵗ) ;  v̂ = v/(1-β2ᵗ)268w -= lr · ( m̂/(sqrt(v̂) + eps) + wd·w )        # eps OUTSIDE sqrt; wd decoupled, never through m/v269```270Fused kernel: host precomputes bias corrections and `grad_scale = min(1, clip/global_norm) / loss_scale`; kernel reads grad → scales → m/v update → step → writes m, v, f32 master, and f16/bf16 working copy (llm.c uses **stochastic rounding** for bf16 working weights — low 16 mantissa bits vs random threshold — else sub-ulp updates stagnate).271272**Memory layout (llm.c pattern):** all params in **one flat f32 buffer** with a table of (offset, shape) per tensor; grads/m/v/master as identically-shaped flat buffers → optimizer is one kernel over `num_parameters`, and global-norm is a two-pass reduction (partials per threadgroup → single-TG reduce; fold the scale into the AdamW grad read, no third pass).273274**Mixed precision:** f32 master weights + f16/bf16 compute copies; grads accumulated in f32 across micro-batches. Loss scaling **for f16 only** (bf16 needs none): init 2^16, halve + **skip step** on inf/nan, double after 2000 clean steps; unscale before clipping. Keep softmax, norms, CE, and all reductions in f32.275276**Fused cross-entropy (llm.c fused_classifier):** one pass computes online (max, sumexp); `loss = max + log(sumexp) - logit[target]`; gradient `(softmax - onehot)·dloss` written **in place over the logits buffer** (dloss = 1/(B·T·grad_accum)); full softmax never materialized. Support `ignore_index=-1`.277278**Model formulas:**279- RoPE (θ=10000, **interleaved-pairs / GPT-J convention** — simplest kernel, stride-2 pairs; document it; backward = rotation by −mθ): `x'_{2k} = x_{2k}cos(mθ_k) − x_{2k+1}sin(mθ_k)`, `x'_{2k+1} = x_{2k}sin(mθ_k) + x_{2k+1}cos(mθ_k)`, `θ_k = 10000^{-2k/d_head}`; q and k only.280- RMSNorm (eps 1e-6): `y = w ⊙ x / sqrt(mean(x²) + eps)`; backward `dx_i = (1/rms)[g_i w_i − x_i·Σ_j(g_j w_j x_j)/(N·rms²)]`, `dw_i = Σ_rows g_i x_i / rms`.281- SwiGLU: `W2(silu(xW1) ⊙ xW3)`, `d_ff = round_up(2/3·4·d_model, multiple)`; `silu'(z) = σ(z)(1 + z(1−σ(z)))`.282- Data loader: flat uint16 `.bin`, random contiguous windows of `context+1`, y = x shifted by 1. Adopt llm.c's 256-int32 header `{magic 20240520, version 1, num_tokens}`. uint16 caps vocab at 65535 — fine.283- Checkpoint: step + RNG state(s) + params + m + v + master weights (+ dataloader position) → exact resume.284285**Reference points:** TinyStories = `roneneldan/TinyStories`, ~926M train / ~19M val tokens (GPT-2 tok); 4k–8k custom BPE vocab is the sweet spot for 10–50M models (llama2.c: 4096-vocab ≈ 32k-vocab compression on this domain). llama2.c 110M config: lr 4e-4, 131k tokens/step, dropout 0.1. BPE: minbpe algorithm (byte-level, lowest-merge-rank greedy encode); C++ needs a hand-rolled UTF-8 codepoint classifier if we adopt the GPT-2/GPT-4 regex pre-split (std::regex can't do `\p{L}`) — or skip the regex (BasicTokenizer) for small domain vocabs.286287---288289## 8. Consolidated pitfalls checklist290291**Metal/metal-cpp**2921. `NS_PRIVATE_IMPLEMENTATION` in exactly one TU; missing pool around per-step command buffers = leak.2932. Coherency only at command-buffer boundaries; never CPU-write in-flight buffers; never recycle a pooled buffer before its command buffer completes.2943. `setBytes` ≤ 4 KB; buffer-offset alignment → use 256 B in the allocator; bytesNoCopy needs 16 KB pages.2954. Query `maxTotalThreadsPerThreadgroup` per pipeline (register pressure can shrink it below 1024).2965. `memoryBarrier` is silently ignored on serial encoders — if we ever flip an encoder to Concurrent, the barriers must already be written.2976. Fast math default: no-NaN/no-INF assumptions, x/0 UB, auto-FMA. Subtract row max before every exp; use `precise::` where parity demands.298299**Kernels**3006f. **Register spill is directly measurable, headlessly, and worth checking on every kernel.** Xcode 26 ships `gpudebug`; `gpucapture start --pid <p> --boundary 1 --count N -o t.gputrace`, then `gpudebug -s <id> -c "profile run --gpu-state high"` / `"profile embed"` / `"profile load"`, then `go performance/shaders` + `info comp0` prints **Temp registers**, **Uniform registers** and **Spilled bytes** per kernel. This requires shader sources in the metallib: pass `-frecord-sources` at compile *and* link, and link with `metal`, not `metallib` (which rejects the flag). Measured here — scalar flash forward 126 registers / **368 spilled**; MMA forward 85 / **0** (that is the mechanism behind its 3.6×); MMA dQ 95 / 0; and the fused MMA dKV **111 / 4352 spilled**, a defect found only by this tooling. Splitting dKV into separate dV and dK kernels cut the backward a further 27% (9.66 → 7.05 ms).3016a. Hoist `D = rowsum(dO ∘ O)` out of the attention backward inner loops (see §6) — a 3.1× step-time factor, and it applies to unfused attention too.3026b. **Per-thread register pressure dominates the attention backward.** A one-thread-per-row backward that keeps `k`, `v`, `dk`, `dv` in registers is 4·head_dim floats — 1 KB/thread at hd=64 — and spills. Measured on gpt-10m (B64 T512 H6), backward per layer: **150 ms** (all four in registers) → **118 ms** (read-only `k`/`v` re-read from device; they are the same address every iteration, so L1 serves them) → **105 ms** (dK and dV split into separate kernels so each thread carries one accumulator). The split recomputes the `q·k` dot in both kernels and still wins. MFA's parameter tables encode the same tradeoff: at hd ≥ 56 it stops caching `K`, at hd ≥ 80 it drops to `dV`/`dK` only.3036d. **Applying a per-row scale to an accumulator held in fragments does not require knowing the lane mapping.** Online softmax needs `O <- diag(corr) · O` each KV block, but MSL leaves the element→lane mapping of `simdgroup_matrix` unspecified (MLX reverse-engineers it in `mma.h::get_coord`). Build the 8×8 diagonal in threadgroup memory and apply it with an ordinary MMA instead: spec-clean, costs head_dim/8 extra MMAs per block (~25% more MMA work), and keeps O in registers — staging O in threadgroup memory would have cost 8 KB and halved residency. The backward needs none of this: `L` from the forward already fixes the normalisation.3046e. Get S transposed by swapping the operand roles, not by transposing data: the dK/dV kernel wants `S^T`/`dS^T` (rows = KV, cols = query), which is exactly `S^T = K @ Q^T` and `dP^T = V @ dO^T`. `simdgroup_load(..., transpose=true)` supplies the `Q^T`/`dO^T` operands straight from the staged tiles at no cost.3056c. Staging K/V through threadgroup memory in the *forward* changed nothing (8.3 → 8.0 ms): the 32 threads of a simdgroup are consecutive query rows of one head, so they hit the same K/V addresses in the same cycle and the cache already broadcasts them. Tile staging pays off for reuse and layout transformation, not for reads that are already uniform across the simdgroup.3067. simdgroup_matrix requires uniform control flow; 8×8 only; mixed half→float MMA is undocumented (keep all-float fallback).3077a. **`constant constexpr uint TILE = 4` at program scope is NOT a compile-time constant.** In MSL `constant` is an *address-space qualifier* and program-scope variables are required to live there, so that declaration creates a constant-address-space **variable**. Loop bounds derived from it aren't compile-time known, the fragment loops don't unroll, `acc[i][j]` becomes dynamic indexing into an opaque `simdgroup_matrix` array, and the driver spills all 16 accumulators (256 B each) to the stack — every MMA then pays ~1 KB of stack traffic. Measured on the f32 GEMM: **0.82 → 10.2 TFLOPs (12×)** purely from declaring tile geometry as `enum : uint { ... }` plus `#pragma clang loop unroll(full)` on the fragment loops. Use enumerators (or `#define`) for tile constants, never `constant constexpr`.3087b. Do not diagnose this from `metal -S -emit-llvm`: at every `-O` level the AIR keeps the loops rolled and shows the same 3 allocas / 2 MMA intrinsics whether or not the fast path is hit. Unrolling and fragment promotion happen in the driver's AIR→ISA backend at pipeline creation, so **benchmark, don't read the IR**.3097c. Stage tiles with a transpose-aware thread→address map (walk the operand's contiguous axis). Mapping `idx` by the fixed logical axis leaves the `nt` case — the one the forward pass issues most — strided and uncoalesced.3107d. Budget threadgroup memory for occupancy: an early epilogue that staged the whole 64×64 output block cost 27 KB, capping residency at one threadgroup per core. Store one 8×8 fragment at a time through per-simdgroup scratch (1 KB total, simdgroup-local so no threadgroup barrier) → 10.5 KB, ~3 threadgroups resident.3118. Never `-INFINITY` masks (NaN via `exp(-inf − -inf)`); guard `l==0 → 0` on fully-masked rows.3129. Accumulate f32 always (matmul, softmax stats, norms, CE); f16 accumulators overflow (max 65504).31310. `simd_shuffle_xor` int-only; no `simd_*` on bfloat; reductions cover active lanes only.31411. No async copies on M3+; no threadgroup-memory-as-pure-read-cache on M3+.31512. +16 B row padding on threadgroup tiles; vectorized loads need alignment checks; pointer-deref vs subscript codegen trap.316317**Training**31811a. **Concurrent dispatch is the cheap encoder win, and it needs no Metal 4.** Metal's default compute encoder orders every dispatch against the previous one. The optimizer is the pathological case: one dispatch per parameter tensor (~100 for a 100M model) and a single-threadgroup `sumsq` per tensor, all mutually independent. Wrapping both optimizer passes in a concurrent-dispatch region took gpt-100m from 7.9k to **9.6k tokens/sec (+22%)** with identical losses. Switching encoder kind ends the encoder, and Metal orders tracked resources across encoders in a command buffer, so the switch is itself the barrier — no explicit `memoryBarrier` needed as long as everything inside a region is independent. An independent investigation measured ~15× on a synthetic batch of small independent dispatches and found Metal 4's encoding path offers no CPU-side advantage once invariant binds are hoisted: the concurrency (available since macOS 10.14) is the whole effect.31911b. Apple-silicon default math mode is `.relaxed`, which **honours Inf/NaN** — important, because f16 loss-scaling overflow detection depends on it. Setting the deprecated `fastMathEnabled = true` silently selects `.fast`, which does not. Passing no math flags (as here) keeps the safe default.32012a. **Buffer recycling is gated on sync, so the sync cadence sets peak memory.** Pooled buffers released while a command buffer is open park on the allocator's retire list and are only reusable after `sync()`. Syncing once per *optimizer* step means none of the gradient-accumulation micro-batches recycle anything, and peak activation memory scales with `grad_accum_steps` — `gpt-100m` at micro-batch 8 OOM'd on a 37 GB working set until the trainer began syncing once per micro-batch (a ~7 GB peak). Cost is one extra GPU sync per micro-batch.32112b. Activation memory, not parameter count, bounds model size without checkpointing: ~490 MB per layer at batch 8 × ctx 1024 × d_model 768, times 12 layers.32213. AdamW: eps outside sqrt, 1-based t, wd decoupled and dim≥2 only, clip before m/v (fold into grad read).32314. Warmup `(it+1)/(warmup+1)` (no lr=0 step); loss ÷ grad_accum inside micro-loop.32415. GradScaler: unscale before clip; skip step on inf/nan; f16 only.32516. bf16 working weights need f32 master or stochastic rounding.32617. RoPE convention (interleaved vs rotate_half) must be consistent everywhere — Forge uses interleaved (GPT-J).32718. Tied embeddings: single grad buffer for wte/lm_head.32819. CE in-place logit-gradient needs a sync between loss read and grad overwrite.32920. Deterministic mode: no atomic-add gradient accumulation — reduction trees / two-kernel splits.330331---332333## 9. Sources334335**Apple official**: [MSL Spec v4.1 PDF](https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf) · [Metal Feature Set Tables](https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf) · [metal-cpp](https://github.com/apple/metal-cpp) · [metal-cpp setup](https://developer.apple.com/metal/cpp/) · Tech Talks [111375 (M3/A17 GPU)](https://developer.apple.com/videos/play/tech-talks/111375/), [111373 (shader best practices)](https://developer.apple.com/videos/play/tech-talks/111373/) · WWDC [22-10159 (scale compute)](https://developer.apple.com/videos/play/wwdc2022/10159/), [20-10603 (GPU counters)](https://developer.apple.com/videos/play/wwdc2020/10603/), [25-205 (Metal 4)](https://developer.apple.com/videos/play/wwdc2025/205/) · [Metal Best Practices Guide (triple buffering)](https://developer.apple.com/library/archive/documentation/3DDrawing/Conceptual/MTLBestPracticesGuide/TripleBuffering.html) · [precompiling shader libraries](https://developer.apple.com/documentation/metal/building-a-shader-library-by-precompiling-source-files)336337**Reference kernels**: [MLX STEEL GEMM](https://github.com/ml-explore/mlx/tree/main/mlx/backend/metal/kernels/steel/gemm) · [MLX steel attention](https://github.com/ml-explore/mlx/tree/main/mlx/backend/metal/kernels/steel/attn) · [MLX matmul dispatch](https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/matmul.cpp) · [MLX sdpa dispatch](https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/scaled_dot_product_attention.cpp) · [llama.cpp ggml-metal.metal](https://github.com/ggml-org/llama.cpp/blob/master/ggml/src/ggml-metal/ggml-metal.metal) · [metal-flash-attention](https://github.com/philipturner/metal-flash-attention) · [metal-benchmarks](https://github.com/philipturner/metal-benchmarks) · [percisely.xyz/gemm](https://percisely.xyz/gemm)338339**Papers**: [FlashAttention-2](https://arxiv.org/abs/2307.08691) · [FlashAttention](https://arxiv.org/abs/2205.14135) · [Rigel/Metal-4 matmul](https://arxiv.org/abs/2606.12765)340341**Training references**: [nanoGPT](https://github.com/karpathy/nanoGPT) (train.py, model.py) · [llm.c](https://github.com/karpathy/llm.c) (train_gpt2.c/.cu, llmc/adamw.cuh, global_norm.cuh, fused_classifier.cuh, dataloader.h) · [minbpe](https://github.com/karpathy/minbpe) · [llama model.py (RoPE/RMSNorm/SwiGLU)](https://github.com/meta-llama/llama/blob/main/llama/model.py) · [llama2.c](https://github.com/karpathy/llama2.c) · [PyTorch AMP docs](https://docs.pytorch.org/docs/stable/amp.html) · [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) · [smollm-corpus](https://huggingface.co/datasets/HuggingFaceTB/smollm-corpus)342343**Community/pitfalls**: [dpogue.ca CMake+Metal](https://dpogue.ca/articles/cmake-metal.html) · [MLX #3689 (buffer use-after-free)](https://github.com/ml-explore/mlx/issues/3689) · [MLX #1864 (command-buffer batching)](https://github.com/ml-explore/mlx/pull/1864) · [MLX discussion #3209 (M3 throughput)](https://github.com/ml-explore/mlx/discussions/3209) · [HF #25199 (RoPE conventions)](https://github.com/huggingface/transformers/issues/25199) · [MoltenVK #1989 (offset alignment)](https://github.com/KhronosGroup/MoltenVK/issues/1989) · [llama.cpp #9094 (bf16 Metal status)](https://github.com/ggml-org/llama.cpp/issues/9094) · [tinygrad #7408 (bf16 simdgroup)](https://github.com/tinygrad/tinygrad/pull/7408)344