# RESEARCH.md — Forge: Metal/C++ LLM Training Framework

Research phase completed 2026-07-31 (6 parallel web-research passes + local toolchain probes). This file records the API signatures, tile sizes, design decisions, and pitfalls that gate implementation. Update this file whenever new research is done during the project.

# 0. Local environment (verified empirically, 2026-07-31)

Decision: compile kernels with -std=metal3.2 (macOS 15+ baseline: full simdgroup_matrix, bfloat, atomic_thread_fence) and use the classic Metal 3 compute API via metal-cpp. Metal 4 (MTL4*, MTLTensor, cooperative tensors) is a parallel opt-in API — a future seam, not a dependency.

# 1. metal-cpp — setup and memory management

Source: 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.

In exactly one .cpp file:

#define NS_PRIVATE_IMPLEMENTATION
#define CA_PRIVATE_IMPLEMENTATION
#define MTL_PRIVATE_IMPLEMENTATION
#include <Foundation/Foundation.hpp>
#include <Metal/Metal.hpp>
#include <QuartzCore/QuartzCore.hpp>

Link frameworks: Foundation, Metal, QuartzCore (MetalKit not needed for compute).

Ownership (Cocoa rules, no ARC):

• new* / alloc / copy / Create* return retained objects → caller releases. That covers MTL::CreateSystemDefaultDevice(), newCommandQueue(), newBuffer(), newLibrary(), newFunction(), newComputePipelineState().
• Everything else is autoreleased — notably queue->commandBuffer() and cmdBuf->computeCommandEncoder(). They die when the enclosing NS::AutoreleasePool drains.
• Smart pointers: NS::TransferPtr(p) adopts retained results (no extra retain); NS::RetainPtr(p) retains autoreleased/borrowed objects.
• 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.

Key compute API (verified against metal-cpp headers):

MTL::Device* MTL::CreateSystemDefaultDevice();
CommandQueue* Device::newCommandQueue();
Buffer* Device::newBuffer(NS::UInteger length, MTL::ResourceOptions);          // MTL::ResourceStorageModeShared
Buffer* Device::newBuffer(const void* ptr, NS::UInteger len, MTL::ResourceOptions, void(^dealloc)(void*, NS::UInteger)); // NoCopy: page-aligned
Library* Device::newLibrary(const NS::String* filepath, NS::Error** err);      // load .metallib by path (CLI tool: do NOT rely on newDefaultLibrary)
Function* Library::newFunction(const NS::String* name, const MTL::FunctionConstantValues*, NS::Error**);
ComputePipelineState* Device::newComputePipelineState(const MTL::Function*, NS::Error**);
ComputeCommandEncoder* CommandBuffer::computeCommandEncoder(MTL::DispatchType);  // Serial | Concurrent
void Encoder::setComputePipelineState(const MTL::ComputePipelineState*);
void Encoder::setBuffer(const MTL::Buffer*, NS::UInteger offset, NS::UInteger index);
void Encoder::setBytes(const void* bytes, NS::UInteger length, NS::UInteger index);   // ≤ ~4 KB, for param structs
void Encoder::setThreadgroupMemoryLength(NS::UInteger length, NS::UInteger index);
void Encoder::dispatchThreads(MTL::Size grid, MTL::Size tg);          // non-uniform TGs, supported on all M-series — prefer for elementwise
void Encoder::dispatchThreadgroups(MTL::Size tgs, MTL::Size tg);      // for tiled kernels
void CommandBuffer::commit();  addCompletedHandler(const MTL::HandlerFunction&);   // std::function<void(CommandBuffer*)>
CFTimeInterval CommandBuffer::GPUStartTime() / GPUEndTime();          // per-command-buffer timing, read in completed handler
NS::UInteger ComputePipelineState::maxTotalThreadsPerThreadgroup();   // can be < 1024 (register pressure) — always query
NS::UInteger ComputePipelineState::threadExecutionWidth();            // 32 on Apple GPUs

# 2. Build system: .metal → .metallib in CMake

Two-step pipeline as CMake custom commands (Ninja/Makefiles generator):

xcrun -sdk macosx metal -std=metal3.2 -O2 [-gline-tables-only -frecord-sources in Debug] -c foo.metal -o foo.air
xcrun -sdk macosx metallib foo.air bar.air ... -o forge.metallib

Load at runtime with device->newLibrary(path, &err). Fast math is ON by default (-fmetal-math-mode=fast: assumes no NaN/INF, reassociation, auto-FMA contraction). Keep it on for throughput but use metal::precise:: selectively in numerically sensitive spots and account for FMA contraction in the 1e-4 CPU-parity tests.

# 3. MSL essentials (from MSL Specification v4.1, 2026-06-04)

Prelude for all kernels: #include <metal_stdlib> + using namespace metal; (umbrella header covers simdgroup, simdgroup_matrix, atomics, compute barriers).

# simdgroup_matrix (the perf-critical path)

• 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()).
• Signatures (verbatim; threadgroup and device overloads):

void simdgroup_load (thread simdgroup_matrix<T,8,8>& d, const {threadgroup|device} T* src,
                     ulong elements_per_row = 8, ulong2 matrix_origin = 0, bool transpose_matrix = false);
void simdgroup_store(thread simdgroup_matrix<T,8,8>  a, {threadgroup|device} T* dst,
                     ulong elements_per_row = 8, ulong2 matrix_origin = 0, bool transpose_matrix = false);
void simdgroup_multiply_accumulate(d, a, b, c);   // d = a*b + c
simdgroup_matrix<T,8,8> make_filled_simdgroup_matrix<T,8,8>(T value);

• 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).
• All simdgroup_matrix ops require uniform SIMD-group control flow — UB otherwise.

# SIMD intrinsics (Apple7+ hardware reductions — all M-series)

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.

# Threadgroup memory & barriers

• Max 32 KB per threadgroup (API limit; all Apple Silicon). Max 1024 threads/TG (but query the pipeline's maxTotalThreadsPerThreadgroup).
• Declare tiles inside the kernel (threadgroup half As[...];) or as [[threadgroup(n)]] arg (runtime-sized).
• 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).

# Atomics

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).

# Numerics

• bfloat: MSL 3.1+, truncated f32; NOT implicitly convertible to/from half; no simd intrinsic support. Fine as storage + convert-to-float compute.
• 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.
• fma() always correctly rounded; compiler auto-contracts a*b+c — a source of CPU/GPU last-ulp differences (tolerances absorb it).

# 4. Execution model & performance practices (Apple Silicon / Apple9)

Architecture 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.

Command submission (the rules Forge's device layer enforces):

1. 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.
2. 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.
3. 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.
4. 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).
5. Dataloader: mmap the token file and wrap with newBuffer(bytesNoCopy:) — requires 16384-byte page alignment/multiple (Apple Silicon page size). True zero-copy.
6. 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).
7. 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.
8. 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.

# 5. GEMM design (chosen tiles + justification)

Reference implementations studied: MLX STEEL (mlx/backend/metal/kernels/steel/gemm/), llama.cpp kernel_mul_mm, Philip Turner's metal-flash-attention GEMM.

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.

Tile configs (starting points, to be autotuned on target):

• 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).
• 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.
• 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.
• GEMV for generation (min(M,N)=1): rows-per-simdgroup, float4 dots, simd_shuffle_down tree / simd_sum (llama.cpp kernel_mul_mv pattern).
• Grid swizzle for large matrices (swizzle_log 1–2), serpentine fragment iteration for register reuse.
• 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.
• 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×.
• 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).
• llama.cpp codegen gotcha: sa[64*ib + ...] (subscript) was "massively slower" than *(sa + 64*ib + ...) (pointer deref) — watch for it.

# 5b. Metal Performance Primitives (matmul2d) — measured 4.9× over hand-written MMA

Probed and benchmarked locally 2026-07-31 (src/kernels/matmul_mpp.metal, tests/bench_precision.cpp, tests/mppcheck.cpp). This is the largest single efficiency finding for this project and it changes the roadmap.

MetalPerformancePrimitives.framework ships in the macOS 26.5 SDK (Headers/MPPTensorOpsMatMul2d.h, 642 lines) and compiles under -std=metal4.0. mpp::tensor_ops::matmul2d is the cooperative-tensor matmul that targets the per-core neural accelerators on M5-class hardware.

Measured on M5 Max, 64×32 tile, 4 simdgroups, f32 accumulate, versus this repo's own simdgroup_matrix GEMM at identical shapes:

So 1.5× for f32 and 4.3–4.9× for f16. Both verified numerically against the CPU reference (tests/mppcheck.cpp): f32 is bit-exact, f16 is 3.8e-06 against a reference fed the same rounded inputs, and every output element is nonzero (i.e. the full K reduction really happens).

Note this contradicts the "Rigel" paper's M4 Max finding that matmul2d still runs on the shader cores and loses to a hand-fused GEMM — M5 evidently has the accelerator hardware that M4 lacked. Re-measure per generation; do not assume.

API notes that cost time to work out (the header's own example is stale):

• Bind ordinary buffers and build tensors in-kernel with the tensor_inline 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.
• Extents are (columns, rows), and slice(x, y) takes column-then-row.
• 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.
• The destination is a cooperative_tensor; zero it via get_capacity() + is_valid_element(i). The header comment says get_mask(i), which does not exist in this SDK.
• #pragma unroll full is not valid MSL here; use #pragma clang loop unroll(full).
• Requires -std=metal4.0, so it is a per-file build flag (CMake keys off the _mpp filename suffix) and macOS 26+; the simdgroup_matrix path stays the portable fallback for macOS 15.

Transposes work, so this covers training, not just inference. transpose_left / transpose_right are descriptor fields (hence template parameters), and nn / nt / tn are all bit-exact against the CPU reference — i.e. the forward X·Wᵀ, the dX = dY·W and the dW = dYᵀ·X of the backward are all served. Two gotchas: the operand extents must be swapped to match the transpose flag, and slice() is (column, row) so a transposed operand is sliced the other way round. Also op.template get_destination_cooperative_tensor<...>() once the kernel is itself a template.

Implication for the roadmap. Mixed precision was previously judged a memory-only win (f16 operands buy just +18–22% on the simdgroup path — measured, see below). Through MPP the same f16 operands are worth 4.9×. That reverses the priority: f16/bf16 plumbing is now the highest-value remaining work, because it is the gate to the accelerator path, not merely a way to halve activations.

# 6. Fused attention design (flash-style)

Algorithm: FlashAttention-2 (arXiv 2307.08691). Online softmax per Q-block over KV-blocks j:

m_new = max(m_old, rowmax(S_j))            S_j = (Q Kᵀ) * scale   (f32)
P~ = exp2(S_j - m_new)                     (scale folded: scale * M_LOG2E_F, fast::exp2)
factor = exp2(m_old - m_new)               (≤ 1; the paper's diag(...)^-1 is a known typo)
l = l * factor + rowsum(P~)
O = O * factor + P~ · V_j                  (O in f32 registers; divide by l ONCE after the loop)
L = m + log(l)                             (one f32/row, saved for backward)

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.

• 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.
• GQA: kv_head = q_head / (n_heads / n_kv_heads) index mapping — no repeat_kv copies.
• 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.
• 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):

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.

• 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).
• 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.
• 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.
• 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.
• 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.
• 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%.
• 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.

# 7. Training loop (nanoGPT / llm.c distilled)

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).

AdamW exact (llm.c, PyTorch-compatible): t is 1-based;

m = β1·m + (1-β1)·g ;  v = β2·v + (1-β2)·g²
m̂ = m/(1-β1ᵗ) ;  v̂ = v/(1-β2ᵗ)
w -= lr · ( m̂/(sqrt(v̂) + eps) + wd·w )        # eps OUTSIDE sqrt; wd decoupled, never through m/v

Fused 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).

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).

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.

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.

Model formulas:

• 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.
• 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.
• SwiGLU: W2(silu(xW1) ⊙ xW3), d_ff = round_up(2/3·4·d_model, multiple); silu'(z) = σ(z)(1 + z(1−σ(z))).
• 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.
• Checkpoint: step + RNG state(s) + params + m + v + master weights (+ dataloader position) → exact resume.

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.

# 8. Consolidated pitfalls checklist

Metal/metal-cpp

1. NS_PRIVATE_IMPLEMENTATION in exactly one TU; missing pool around per-step command buffers = leak.
2. Coherency only at command-buffer boundaries; never CPU-write in-flight buffers; never recycle a pooled buffer before its command buffer completes.
3. setBytes ≤ 4 KB; buffer-offset alignment → use 256 B in the allocator; bytesNoCopy needs 16 KB pages.
4. Query maxTotalThreadsPerThreadgroup per pipeline (register pressure can shrink it below 1024).
5. memoryBarrier is silently ignored on serial encoders — if we ever flip an encoder to Concurrent, the barriers must already be written.
6. Fast math default: no-NaN/no-INF assumptions, x/0 UB, auto-FMA. Subtract row max before every exp; use precise:: where parity demands.

Kernels 6f. 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). 6a. 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. 6b. 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. 6d. 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. 6e. 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. 6c. 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. 7. simdgroup_matrix requires uniform control flow; 8×8 only; mixed half→float MMA is undocumented (keep all-float fallback). 7a. 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. 7b. 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. 7c. 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. 7d. 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. 8. Never -INFINITY masks (NaN via exp(-inf − -inf)); guard l==0 → 0 on fully-masked rows. 9. Accumulate f32 always (matmul, softmax stats, norms, CE); f16 accumulators overflow (max 65504). 10. simd_shuffle_xor int-only; no simd_* on bfloat; reductions cover active lanes only. 11. No async copies on M3+; no threadgroup-memory-as-pure-read-cache on M3+. 12. +16 B row padding on threadgroup tiles; vectorized loads need alignment checks; pointer-deref vs subscript codegen trap.

Training 11a. 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. 11b. 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. 12a. 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. 12b. 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. 13. AdamW: eps outside sqrt, 1-based t, wd decoupled and dim≥2 only, clip before m/v (fold into grad read). 14. Warmup (it+1)/(warmup+1) (no lr=0 step); loss ÷ grad_accum inside micro-loop. 15. GradScaler: unscale before clip; skip step on inf/nan; f16 only. 16. bf16 working weights need f32 master or stochastic rounding. 17. RoPE convention (interleaved vs rotate_half) must be consistent everywhere — Forge uses interleaved (GPT-J). 18. Tied embeddings: single grad buffer for wte/lm_head. 19. CE in-place logit-gradient needs a sync between loss read and grad overwrite. 20. Deterministic mode: no atomic-add gradient accumulation — reduction trees / two-kernel splits.

# 9. Sources

Apple official: MSL Spec v4.1 PDF · Metal Feature Set Tables · metal-cpp · metal-cpp setup · Tech Talks 111375 (M3/A17 GPU), 111373 (shader best practices) · WWDC 22-10159 (scale compute), 20-10603 (GPU counters), 25-205 (Metal 4) · Metal Best Practices Guide (triple buffering) · precompiling shader libraries

Reference kernels: MLX STEEL GEMM · MLX steel attention · MLX matmul dispatch · MLX sdpa dispatch · llama.cpp ggml-metal.metal · metal-flash-attention · metal-benchmarks · percisely.xyz/gemm

Papers: FlashAttention-2 · FlashAttention · Rigel/Metal-4 matmul

Training references: nanoGPT (train.py, model.py) · llm.c (train_gpt2.c/.cu, llmc/adamw.cuh, global_norm.cuh, fused_classifier.cuh, dataloader.h) · minbpe · llama model.py (RoPE/RMSNorm/SwiGLU) · llama2.c · PyTorch AMP docs · TinyStories · smollm-corpus

Community/pitfalls: dpogue.ca CMake+Metal · MLX #3689 (buffer use-after-free) · MLX #1864 (command-buffer batching) · MLX discussion #3209 (M3 throughput) · HF #25199 (RoPE conventions) · MoltenVK #1989 (offset alignment) · llama.cpp #9094 (bf16 Metal status) · tinygrad #7408 (bf16 simdgroup)