SPB Git

spb/forge Public MIT

Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.

llmtrainingcppmetalapple-silicon
C++ 61.2% C 23% Python 7.6% TeX 7.2% CMake 1.1%

Add architecture-variant waves 1+2: train Mistral/Qwen/Gemma/OLMo-class models by config

Wave 1 (module-level):
- rope refactor: kernels read a host-precomputed inv-freq table; unlocks
  HF-"llama3" rope scaling (rope_scale_*), per-layer theta, and NoPE
  layers (nope_every, SmolLM3)
- attention_bias (Qwen2.5 QKV bias), head_dim decoupled from
  d_model/n_heads (Qwen3), relu2 activation (nanoGPT-speedrun lineage),
  norm_placement pre|post|sandwich (OLMo2/Gemma)

Wave 2 (attention kernels):
- sliding_window + sliding_global_every (Mistral / Gemma3 local:global
  patterns) in the CPU reference, the unfused Metal kernels, and the
  fused scalar flash kernels — out-of-window KV blocks are skipped, so
  cost scales with the window; window > 0 auto-routes off the MMA kernel
- attn_softcap (Gemma2): cap*tanh on scores pre-softmax, unfused path,
  exact tanh' chain in all backwards
- rope_theta_global for dual-theta local/global layers (Gemma3)

Parity suites cover every knob (fused + unfused paths); gradcheck and
overfit stay green. New demo configs: gpt-50m-mistral, gpt-50m-gemma;
ARCHITECTURES.md documents the per-family config matrix.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed 5 days ago (Aug 5, 2026) parent 6d77a01

Showing 17 changed files with +572 and −139

added configs/gpt-50m-gemma.json +44 −0
@@ -0,0 +1,44 @@
1 +{
2 + "_comment": "Gemma-3-style variant of gpt-50m: 5 local (window 256, theta 10k) : 1 global (theta 1M) attention pattern, QK-norm, sandwich norm, embeddings scaled by sqrt(d_model), GELU. Same data/steps as gpt-50m.",
3 + "model": {
4 + "name": "gpt-50m-gemma",
5 + "n_layers": 12,
6 + "d_model": 576,
7 + "n_heads": 9,
8 + "n_kv_heads": 3,
9 + "d_ff": 1536,
10 + "vocab_size": 4096,
11 + "context_length": 1024,
12 + "tied_embeddings": true,
13 + "use_rope": true,
14 + "rope_theta": 10000.0,
15 + "rope_theta_global": 1000000.0,
16 + "norm": "rmsnorm",
17 + "norm_eps": 1e-06,
18 + "activation": "gelu",
19 + "dropout": 0.0,
20 + "qk_norm": true,
21 + "norm_placement": "sandwich",
22 + "scale_embeddings": true,
23 + "sliding_window": 256,
24 + "sliding_global_every": 6
25 + },
26 + "train": {
27 + "lr": 0.0005,
28 + "min_lr_ratio": 0.1,
29 + "warmup_steps": 117,
30 + "max_steps": 1170,
31 + "beta1": 0.9,
32 + "beta2": 0.95,
33 + "eps": 1e-08,
34 + "weight_decay": 0.1,
35 + "grad_clip": 1.0,
36 + "batch_size": 8,
37 + "grad_accum_steps": 8,
38 + "precision": "f32",
39 + "checkpoint_every": 200,
40 + "eval_every": 100,
41 + "eval_batches": 20,
42 + "seed": 1337
43 + }
44 +}
added configs/gpt-50m-mistral.json +39 −0
@@ -0,0 +1,39 @@
1 +{
2 + "_comment": "Mistral-style variant of gpt-50m: sliding-window attention (256 tokens, every layer), GQA 2:1, QKV bias off. Same data/steps as gpt-50m so the only variable is the attention pattern. The fused scalar flash kernel skips out-of-window KV blocks, so attention cost scales with the window, not the context.",
3 + "model": {
4 + "name": "gpt-50m-mistral",
5 + "n_layers": 10,
6 + "d_model": 640,
7 + "n_heads": 10,
8 + "n_kv_heads": 5,
9 + "d_ff": 1728,
10 + "vocab_size": 4096,
11 + "context_length": 1024,
12 + "tied_embeddings": true,
13 + "use_rope": true,
14 + "rope_theta": 10000.0,
15 + "norm": "rmsnorm",
16 + "norm_eps": 1e-06,
17 + "activation": "swiglu",
18 + "dropout": 0.0,
19 + "sliding_window": 256
20 + },
21 + "train": {
22 + "lr": 0.0005,
23 + "min_lr_ratio": 0.1,
24 + "warmup_steps": 117,
25 + "max_steps": 1170,
26 + "beta1": 0.9,
27 + "beta2": 0.95,
28 + "eps": 1e-08,
29 + "weight_decay": 0.1,
30 + "grad_clip": 1.0,
31 + "batch_size": 8,
32 + "grad_accum_steps": 8,
33 + "precision": "f32",
34 + "checkpoint_every": 200,
35 + "eval_every": 100,
36 + "eval_batches": 20,
37 + "seed": 1337
38 + }
39 +}
modified src/kernels/attention.metal +25 −6
@@ -18,6 +18,8 @@ struct AttnParams {
18 18 uint B, T, H, HKV, HD;
19 19 float scale;
20 20 uint causal;
21 + uint window; // sliding window (keys kept, self incl.); 0 = full
22 + float softcap; // cap*tanh(s/cap) pre-softmax; 0 = off
21 23 };
22 24
23 25 kernel void attention_fwd_f32(device const float* Q [[buffer(0)]],
@@ -40,29 +42,32 @@ kernel void attention_fwd_f32(device const float* Q [[buffer(0)]],
40 42 device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * p.HD;
41 43 device float* prow_base = P + ((ulong(b) * p.H + h) * p.T + i) * p.T;
42 44 const uint jmax = p.causal ? i : p.T - 1;
45 + const uint jmin = (p.window > 0 && i + 1 > p.window) ? i + 1 - p.window : 0;
43 46
44 47 float m = -FLT_MAX;
45 − for (uint j = 0; j <= jmax; ++j) {
48 + for (uint j = jmin; j <= jmax; ++j) {
46 49 device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
47 50 float s = 0.0f;
48 51 for (uint d = 0; d < p.HD; ++d) s = fma(qi[d], kj[d], s);
49 52 s *= p.scale;
53 + if (p.softcap > 0.0f) s = p.softcap * precise::tanh(s / p.softcap);
50 54 prow_base[j] = s;
51 55 m = max(m, s);
52 56 }
53 57 float sum = 0.0f;
54 − for (uint j = 0; j <= jmax; ++j) {
58 + for (uint j = jmin; j <= jmax; ++j) {
55 59 const float e = exp(prow_base[j] - m);
56 60 prow_base[j] = e;
57 61 sum += e;
58 62 }
59 63 const float inv = 1.0f / sum;
60 − for (uint j = 0; j <= jmax; ++j) prow_base[j] *= inv;
64 + for (uint j = 0; j < jmin; ++j) prow_base[j] = 0.0f;
65 + for (uint j = jmin; j <= jmax; ++j) prow_base[j] *= inv;
61 66 for (uint j = jmax + 1; j < p.T; ++j) prow_base[j] = 0.0f;
62 67
63 68 device float* oi = O + (ulong(b) * p.T + i) * Cq + h * p.HD;
64 69 for (uint d = 0; d < p.HD; ++d) oi[d] = 0.0f;
65 − for (uint j = 0; j <= jmax; ++j) {
70 + for (uint j = jmin; j <= jmax; ++j) {
66 71 const float prob = prow_base[j];
67 72 if (prob == 0.0f) continue;
68 73 device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
@@ -126,7 +131,14 @@ kernel void attention_bwd_dq_f32(device const float* Q [[buffer(0)]],
126 131 device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
127 132 float dp = 0.0f;
128 133 for (uint d = 0; d < p.HD; ++d) dp = fma(doi[d], vj[d], dp);
129 − const float ds = prob * (dp - row_dot) * p.scale;
134 + float ds = prob * (dp - row_dot) * p.scale;
135 + if (p.softcap > 0.0f) {
136 + device const float* qq = Q + (ulong(b) * p.T + i) * Cq + h * p.HD;
137 + float s = 0.0f;
138 + for (uint d = 0; d < p.HD; ++d) s = fma(qq[d], kj[d], s);
139 + const float t = precise::tanh(s * p.scale / p.softcap);
140 + ds *= 1.0f - t * t;
141 + }
130 142 for (uint d = 0; d < p.HD; ++d) dqi[d] = fma(ds, kj[d], dqi[d]);
131 143 }
132 144 }
@@ -167,7 +179,14 @@ kernel void attention_bwd_dkv_f32(device const float* Q [[buffer(0)]],
167 179 const float row_dot = D[(ulong(b) * p.H + h) * p.T + i];
168 180 float dp = 0.0f;
169 181 for (uint d = 0; d < p.HD; ++d) dp = fma(doi[d], vj[d], dp);
170 − const float ds = prob * (dp - row_dot) * p.scale;
182 + float ds = prob * (dp - row_dot) * p.scale;
183 + if (p.softcap > 0.0f) {
184 + device const float* kk = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
185 + float s = 0.0f;
186 + for (uint d = 0; d < p.HD; ++d) s = fma(qi[d], kk[d], s);
187 + const float t = precise::tanh(s * p.scale / p.softcap);
188 + ds *= 1.0f - t * t;
189 + }
171 190 for (uint d = 0; d < p.HD; ++d) {
172 191 dvj[d] = fma(prob, doi[d], dvj[d]);
173 192 dkj[d] = fma(ds, qi[d], dkj[d]);
modified src/kernels/elementwise.metal +22 −6
@@ -52,6 +52,21 @@ kernel void gelu_f32(device const float* x [[buffer(0)]],
52 52 out[gid] = 0.5f * v * (1.0f + precise::tanh(k * (v + 0.044715f * v * v * v)));
53 53 }
54 54
55 +// ReLU² (nanoGPT-speedrun lineage): out = max(x, 0)^2.
56 +kernel void relu2_f32(device const float* x [[buffer(0)]],
57 + device float* out [[buffer(1)]],
58 + uint gid [[thread_position_in_grid]]) {
59 + const float v = max(x[gid], 0.0f);
60 + out[gid] = v * v;
61 +}
62 +
63 +kernel void relu2_bwd_f32(device const float* x [[buffer(0)]],
64 + device const float* dout [[buffer(1)]],
65 + device float* dx [[buffer(2)]],
66 + uint gid [[thread_position_in_grid]]) {
67 + dx[gid] = fma(dout[gid], 2.0f * max(x[gid], 0.0f), dx[gid]);
68 +}
69 +
55 70 // Gemma-style logit softcap: out = cap * tanh(x / cap).
56 71 kernel void softcap_f32(device const float* x [[buffer(0)]],
57 72 device float* out [[buffer(1)]],
@@ -127,13 +142,15 @@ constant bool ROPE_INV [[function_constant(2)]];
127 142
128 143 struct RopeParams {
129 144 uint T, H, HD;
130 − float theta;
131 145 uint pos_offset;
132 146 };
133 147
134 −kernel void rope_f32(device const float* x [[buffer(0)]],
135 − device float* out [[buffer(1)]],
136 − constant RopeParams& p [[buffer(2)]],
148 +// freqs[k] is the per-pair inverse frequency (theta^(-2k/HD), possibly
149 +// rescaled — llama3 rope_scaling, per-layer theta). Host-precomputed.
150 +kernel void rope_f32(device const float* x [[buffer(0)]],
151 + device float* out [[buffer(1)]],
152 + constant RopeParams& p [[buffer(2)]],
153 + device const float* freqs [[buffer(3)]],
137 154 uint gid [[thread_position_in_grid]]) {
138 155 // gid indexes (bt, h, k) pairs: one thread per rotated pair
139 156 const uint pairs_per_row = p.H * (p.HD / 2);
@@ -143,8 +160,7 @@ kernel void rope_f32(device const float* x [[buffer(0)]],
143 160 const uint k = rem % (p.HD / 2);
144 161
145 162 const float pos = float(bt % p.T + p.pos_offset);
146 − const float freq = pow(p.theta, -2.0f * float(k) / float(p.HD));
147 − const float angle = pos * freq;
163 + const float angle = pos * freqs[k];
148 164 const float c = cos(angle);
149 165 const float s = ROPE_INV ? -sin(angle) : sin(angle);
150 166
modified src/kernels/flash_attention.metal +18 −5
@@ -32,6 +32,7 @@ struct FlashParams {
32 32 uint B, T, H, HKV;
33 33 float scale;
34 34 uint causal;
35 + uint window; // sliding window (keys kept, self incl.); 0 = full attention
35 36 };
36 37
37 38 // ---------------------------------------------------------------- forward
@@ -90,8 +91,15 @@ kernel void flash_attn_fwd(device const float* Q [[buffer(0)]],
90 91 const uint row_max = min(qb * TGQ + TGQ - 1, p.T - 1);
91 92 const uint j_end = p.causal ? (row_max + 1) : p.T;
92 93 const uint diag_start = p.causal ? (qb * TGQ) : p.T; // blocks below need no mask
93 −
94 − for (uint jb = 0; jb < j_end; jb += BKV) {
94 + // Sliding window: this row's first visible key, and (uniform across the
95 + // threadgroup) the first KV block any row here can see — earlier blocks
96 + // 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;
101 +
102 + for (uint jb = jb_begin; jb < j_end; jb += BKV) {
95 103 const uint block_len = min(BKV, j_end - jb);
96 104
97 105 threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -109,6 +117,7 @@ kernel void flash_attn_fwd(device const float* Q [[buffer(0)]],
109 117 for (uint jr = 0; jr < block_len; ++jr) {
110 118 const uint j = jb + jr;
111 119 if (needs_mask && j > i) break; // rest of the block is masked too
120 + if (j < jmin) continue; // below this row's window
112 121
113 122 threadgroup const float* kj = Ks + jr * HD;
114 123 float s = 0.0f;
@@ -174,8 +183,9 @@ kernel void flash_attn_bwd_dq(device const float* Q [[buffer(0)]],
174 183 const float li = L[(ulong(b) * p.H + h) * p.T + i];
175 184 const float di = D[(ulong(b) * p.H + h) * p.T + i];
176 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;
177 187
178 − for (uint j = 0; j <= jmax; ++j) {
188 + for (uint j = jmin; j <= jmax; ++j) {
179 189 device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * HD;
180 190 device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * HD;
181 191
@@ -225,10 +235,12 @@ kernel void flash_attn_bwd_dv(device const float* Q [[buffer(0)]],
225 235 for (uint d = 0; d < HD; ++d) dv[d] = 0.0f;
226 236
227 237 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;
228 240 for (uint r = 0; r < rep; ++r) {
229 241 const uint h = hkv * rep + r;
230 242 device const float* Lh = L + (ulong(b) * p.H + h) * p.T;
231 − for (uint i = imin; i < p.T; ++i) {
243 + for (uint i = imin; i < imax; ++i) {
232 244 device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;
233 245 device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;
234 246 float s = 0.0f;
@@ -271,11 +283,12 @@ kernel void flash_attn_bwd_dk(device const float* Q [[buffer(0)]],
271 283 for (uint d = 0; d < HD; ++d) dk[d] = 0.0f;
272 284
273 285 const uint imin = p.causal ? j : 0;
286 + const uint imax = (p.window > 0) ? min(p.T, j + p.window) : p.T;
274 287 for (uint r = 0; r < rep; ++r) {
275 288 const uint h = hkv * rep + r;
276 289 device const float* Lh = L + (ulong(b) * p.H + h) * p.T;
277 290 device const float* Dh = D + (ulong(b) * p.H + h) * p.T;
278 − for (uint i = imin; i < p.T; ++i) {
291 + for (uint i = imin; i < imax; ++i) {
279 292 device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;
280 293 device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;
281 294 float s = 0.0f, dp = 0.0f;
modified src/kernels/flash_attention_mma.metal +1 −0
@@ -36,6 +36,7 @@ struct FlashParams {
36 36 uint B, T, H, HKV;
37 37 float scale;
38 38 uint causal;
39 + uint window; // unused here: the host routes window > 0 to the scalar kernels
39 40 };
40 41
41 42 // Enumerators, never `constant constexpr`: see RESEARCH.md 7a — the latter is
modified src/nn/attention.h +39 −19
@@ -20,24 +20,31 @@ public:
20 20 };
21 21
22 22 // Causal multi-head self-attention with GQA and (optional) RoPE.
23 −// No biases (llama convention).
23 +// Per-layer flavour comes from `layer_idx` + config: NoPE layers skip RoPE
24 +// (SmolLM3), global layers may use a different rope theta (Gemma 3), and
25 +// head_dim may be decoupled from d_model/n_heads (Qwen3). Optional QKV bias
26 +// (Qwen2.5, on q/k/v only — llama convention keeps wo bias-free).
24 27 class CausalSelfAttention : public AttentionBase {
25 28 public:
26 − CausalSelfAttention(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)
29 + CausalSelfAttention(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng,
30 + int64_t layer_idx = 0)
27 31 : n_heads_(cfg.n_heads),
28 32 n_kv_heads_(cfg.n_kv_heads),
29 33 head_dim_(cfg.head_dim()),
30 − use_rope_(cfg.use_rope),
31 − rope_theta_(cfg.rope_theta),
34 + use_rope_(cfg.layer_uses_rope(layer_idx)),
32 35 qk_norm_(cfg.qk_norm),
33 − norm_eps_(cfg.norm_eps) {
36 + norm_eps_(cfg.norm_eps),
37 + window_(cfg.layer_is_global(layer_idx) ? 0 : cfg.sliding_window),
38 + attn_softcap_(cfg.attn_softcap) {
34 39 const float std = 0.02f;
35 40 const int64_t C = cfg.d_model;
36 − const int64_t Ckv = cfg.n_kv_heads * cfg.head_dim();
37 − wq_ = std::make_unique<Linear>(C, C, false, std, rng);
38 − wk_ = std::make_unique<Linear>(C, Ckv, false, std, rng);
39 − wv_ = std::make_unique<Linear>(C, Ckv, false, std, rng);
40 − wo_ = std::make_unique<Linear>(C, C, false, proj_std, rng); // residual projection
41 + const int64_t Cq = cfg.n_heads * head_dim_;
42 + const int64_t Ckv = cfg.n_kv_heads * head_dim_;
43 + const bool bias = cfg.attention_bias;
44 + wq_ = std::make_unique<Linear>(C, Cq, bias, std, rng);
45 + wk_ = std::make_unique<Linear>(C, Ckv, bias, std, rng);
46 + wv_ = std::make_unique<Linear>(C, Ckv, bias, std, rng);
47 + wo_ = std::make_unique<Linear>(Cq, C, false, proj_std, rng); // residual projection
41 48 const ops::QuantMode qm = quant_mode_from(cfg.quant);
42 49 wq_->set_quant(qm);
43 50 wk_->set_quant(qm);
@@ -49,18 +56,28 @@ public:
49 56 absorb("wo", *wo_);
50 57 if (qk_norm_) {
51 58 q_norm_w_ = register_param("q_norm.weight",
52 − Tensor::full({cfg.head_dim()}, 1.0f));
59 + Tensor::full({head_dim_}, 1.0f));
53 60 k_norm_w_ = register_param("k_norm.weight",
54 − Tensor::full({cfg.head_dim()}, 1.0f));
61 + Tensor::full({head_dim_}, 1.0f));
62 + }
63 + if (use_rope_) {
64 + const bool global = cfg.layer_is_global(layer_idx);
65 + const float theta = (global && cfg.rope_theta_global > 0.0f)
66 + ? cfg.rope_theta_global
67 + : cfg.rope_theta;
68 + freqs_ = ops::rope_freqs(head_dim_, theta, cfg.rope_scale_factor,
69 + cfg.rope_scale_low, cfg.rope_scale_high,
70 + cfg.rope_scale_orig_ctx);
55 71 }
56 72 }
57 73
58 74 Var forward(const Var& x) const override {
59 75 const int64_t B = x.value().size(0), T = x.value().size(1), C = x.value().size(2);
76 + const int64_t Cq = n_heads_ * head_dim_;
60 77 const int64_t Ckv = n_kv_heads_ * head_dim_;
61 78
62 79 Var x2d = x.reshaped({B * T, C});
63 − Var q = wq_->forward(x2d).reshaped({B, T, C});
80 + Var q = wq_->forward(x2d).reshaped({B, T, Cq});
64 81 Var k = wk_->forward(x2d).reshaped({B, T, Ckv});
65 82 Var v = wv_->forward(x2d).reshaped({B, T, Ckv});
66 83
@@ -68,29 +85,32 @@ public:
68 85 // Per-head RMSNorm before RoPE (Qwen3/Gemma3): normalize over
69 86 // head_dim by viewing heads as rows.
70 87 q = ops::rmsnorm(q.reshaped({B * T * n_heads_, head_dim_}), q_norm_w_,
71 − norm_eps_).reshaped({B, T, C});
88 + norm_eps_).reshaped({B, T, Cq});
72 89 k = ops::rmsnorm(k.reshaped({B * T * n_kv_heads_, head_dim_}), k_norm_w_,
73 90 norm_eps_).reshaped({B, T, Ckv});
74 91 }
75 92
76 93 if (use_rope_) {
77 − q = ops::rope(q, n_heads_, rope_theta_, 0);
78 − k = ops::rope(k, n_kv_heads_, rope_theta_, 0);
94 + q = ops::rope(q, n_heads_, freqs_, 0);
95 + k = ops::rope(k, n_kv_heads_, freqs_, 0);
79 96 }
80 97
81 98 const float scale = 1.0f / std::sqrt(float(head_dim_));
82 − Var o = ops::attention(q, k, v, n_heads_, n_kv_heads_, /*causal=*/true, scale);
83 − return wo_->forward(o.reshaped({B * T, C})).reshaped({B, T, C});
99 + Var o = ops::attention(q, k, v, n_heads_, n_kv_heads_, /*causal=*/true, scale,
100 + window_, attn_softcap_);
101 + return wo_->forward(o.reshaped({B * T, Cq})).reshaped({B, T, C});
84 102 }
85 103
86 104 private:
87 105 int64_t n_heads_, n_kv_heads_, head_dim_;
88 106 bool use_rope_;
89 − float rope_theta_;
90 107 bool qk_norm_;
91 108 float norm_eps_;
109 + int64_t window_; // 0 = full attention (global layer or SWA off)
110 + float attn_softcap_;
92 111 std::unique_ptr<Linear> wq_, wk_, wv_, wo_;
93 112 Var q_norm_w_, k_norm_w_; // defined only when qk_norm
113 + Tensor freqs_; // [head_dim/2], defined only when use_rope
94 114 };
95 115
96 116 } // namespace forge::nn
modified src/nn/config.h +61 −10
@@ -32,6 +32,20 @@ struct ModelConfig {
32 32 bool qk_norm = false; // RMSNorm on q/k per head before RoPE (Qwen3/Gemma3)
33 33 float final_softcap = 0.0f; // logits = cap*tanh(logits/cap) (Gemma2: 30)
34 34 bool scale_embeddings = false;// x *= sqrt(d_model) after embedding (Gemma)
35 + bool attention_bias = false; // bias on q/k/v projections only (Qwen2.5)
36 + int64_t head_dim_override = 0;// decouple head_dim from d_model/n_heads (Qwen3)
37 + int64_t nope_every = 0; // skip RoPE every Nth layer (SmolLM3: 4)
38 + std::string norm_placement = "pre"; // "pre" | "post" (OLMo2) | "sandwich" (Gemma)
39 + // RoPE scaling, HF "llama3" style (0 = off). Applied to the inv-freq table.
40 + float rope_scale_factor = 0.0f; // llama3: 32
41 + float rope_scale_low = 1.0f; // low_freq_factor
42 + float rope_scale_high = 4.0f; // high_freq_factor
43 + int64_t rope_scale_orig_ctx = 8192; // original_max_position_embeddings
44 + // Sliding-window attention (Mistral/Gemma3). 0 = full attention everywhere.
45 + int64_t sliding_window = 0;
46 + int64_t sliding_global_every = 0; // every Nth layer is global (Gemma3: 6); 0 = none
47 + float rope_theta_global = 0.0f; // theta for global layers (Gemma3: 1e6); 0 = theta
48 + float attn_softcap = 0.0f; // cap on attention scores (Gemma2: 50)
35 49 // Quantization-aware training: linear weights are fake-quantized each
36 50 // forward (per-row scales, straight-through estimator in backward).
37 51 // Master weights and the optimizer stay f32.
@@ -44,21 +58,36 @@ struct ModelConfig {
44 58 float moe_aux_weight = 0.01f; // load-balance loss: E * sum_e mean_gate_e^2
45 59 int64_t n_shared_experts = 0; // always-active experts (DeepSeek/Kimi style)
46 60
47 − int64_t head_dim() const { return d_model / n_heads; }
61 + int64_t head_dim() const {
62 + return head_dim_override > 0 ? head_dim_override : d_model / n_heads;
63 + }
64 +
65 + // True when layer i (0-based) attends globally; only meaningful with
66 + // sliding_window > 0.
67 + bool layer_is_global(int64_t i) const {
68 + if (sliding_window <= 0) return true;
69 + return sliding_global_every > 0 && (i + 1) % sliding_global_every == 0;
70 + }
71 + bool layer_uses_rope(int64_t i) const {
72 + return use_rope && !(nope_every > 0 && (i + 1) % nope_every == 0);
73 + }
48 74
49 − // Total parameter count (SwiGLU MLP has three mats, GELU has two).
75 + // Total parameter count (SwiGLU MLP has three mats, GELU/ReLU² have two).
50 76 int64_t num_params() const {
51 77 const int64_t hd = head_dim();
52 − const int64_t attn = d_model * d_model // wq
53 − + 2 * d_model * n_kv_heads * hd // wk, wv
54 − + d_model * d_model; // wo
78 + int64_t attn = d_model * n_heads * hd // wq
79 + + 2 * d_model * n_kv_heads * hd // wk, wv
80 + + n_heads * hd * d_model; // wo
81 + if (attention_bias) attn += (n_heads + 2 * n_kv_heads) * hd;
55 82 const int64_t mlp_one = (activation == "swiglu")
56 83 ? 3 * d_model * d_ff
57 84 : 2 * d_model * d_ff;
58 85 const int64_t mlp = n_experts > 0
59 86 ? (n_experts + n_shared_experts) * mlp_one + n_experts * d_model
60 87 : mlp_one; // experts + router
61 − int64_t norms = (norm == "layernorm" ? 2 : 1) * d_model * (2 * n_layers + 1);
88 + const int64_t norms_per_layer = norm_placement == "sandwich" ? 4 : 2;
89 + int64_t norms = (norm == "layernorm" ? 2 : 1) * d_model *
90 + (norms_per_layer * n_layers + 1);
62 91 if (qk_norm) norms += 2 * head_dim() * n_layers;
63 92 int64_t total = n_layers * (attn + mlp) + norms + vocab_size * d_model;
64 93 if (!tied_embeddings) total += vocab_size * d_model;
@@ -129,15 +158,37 @@ inline void from_json(const nlohmann::json& j, ModelConfig& c) {
129 158 c.qk_norm = j.value("qk_norm", c.qk_norm);
130 159 c.final_softcap = j.value("final_softcap", c.final_softcap);
131 160 c.scale_embeddings = j.value("scale_embeddings", c.scale_embeddings);
161 + c.attention_bias = j.value("attention_bias", c.attention_bias);
162 + c.head_dim_override = j.value("head_dim", c.head_dim_override);
163 + c.nope_every = j.value("nope_every", c.nope_every);
164 + c.norm_placement = j.value("norm_placement", c.norm_placement);
165 + c.rope_scale_factor = j.value("rope_scale_factor", c.rope_scale_factor);
166 + c.rope_scale_low = j.value("rope_scale_low", c.rope_scale_low);
167 + c.rope_scale_high = j.value("rope_scale_high", c.rope_scale_high);
168 + c.rope_scale_orig_ctx = j.value("rope_scale_orig_ctx", c.rope_scale_orig_ctx);
169 + c.sliding_window = j.value("sliding_window", c.sliding_window);
170 + c.sliding_global_every = j.value("sliding_global_every", c.sliding_global_every);
171 + c.rope_theta_global = j.value("rope_theta_global", c.rope_theta_global);
172 + c.attn_softcap = j.value("attn_softcap", c.attn_softcap);
132 173
133 − if (c.d_model % c.n_heads != 0)
134 − throw std::runtime_error("config: d_model must be divisible by n_heads");
174 + if (c.head_dim_override == 0 && c.d_model % c.n_heads != 0)
175 + throw std::runtime_error("config: d_model must be divisible by n_heads "
176 + "(or set head_dim explicitly)");
135 177 if (c.n_heads % c.n_kv_heads != 0)
136 178 throw std::runtime_error("config: n_heads must be divisible by n_kv_heads");
179 + if (c.head_dim() % 2 != 0)
180 + throw std::runtime_error("config: head_dim must be even (RoPE pairs)");
137 181 if (c.norm != "rmsnorm" && c.norm != "layernorm")
138 182 throw std::runtime_error("config: norm must be rmsnorm or layernorm");
139 − if (c.activation != "swiglu" && c.activation != "gelu")
140 − throw std::runtime_error("config: activation must be swiglu or gelu");
183 + if (c.activation != "swiglu" && c.activation != "gelu" && c.activation != "relu2")
184 + throw std::runtime_error("config: activation must be swiglu, gelu or relu2");
185 + if (c.norm_placement != "pre" && c.norm_placement != "post" &&
186 + c.norm_placement != "sandwich")
187 + throw std::runtime_error("config: norm_placement must be pre, post or sandwich");
188 + if (c.sliding_window < 0 || c.sliding_global_every < 0 || c.nope_every < 0)
189 + throw std::runtime_error("config: window/pattern values must be >= 0");
190 + if (c.attn_softcap < 0.0f)
191 + throw std::runtime_error("config: attn_softcap must be >= 0");
141 192 if (c.quant != "none" && c.quant != "int8" && c.quant != "ternary")
142 193 throw std::runtime_error("config: quant must be none, int8 or ternary");
143 194 if (c.n_experts < 0)
modified src/nn/mlp.h +5 −4
@@ -22,11 +22,11 @@ public:
22 22 virtual Var aux() const { return Var(); }
23 23 };
24 24
25 −// SwiGLU: w2( silu(x w1) ⊙ (x w3) ) | GELU: proj( gelu(fc(x)) )
25 +// SwiGLU: w2( silu(x w1) ⊙ (x w3) ) | GELU/ReLU²: w2( act(x w1) )
26 26 class MLP : public MLPBase {
27 27 public:
28 28 MLP(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)
29 − : swiglu_(cfg.activation == "swiglu") {
29 + : swiglu_(cfg.activation == "swiglu"), relu2_(cfg.activation == "relu2") {
30 30 const float std = 0.02f;
31 31 const int64_t C = cfg.d_model, F = cfg.d_ff;
32 32 w1_ = std::make_unique<Linear>(C, F, false, std, rng);
@@ -49,11 +49,12 @@ public:
49 49 Var up = w3_->forward(x);
50 50 return w2_->forward(ops::mul(gate, up));
51 51 }
52 − return w2_->forward(ops::gelu(w1_->forward(x)));
52 + Var h = w1_->forward(x);
53 + return w2_->forward(relu2_ ? ops::relu2(h) : ops::gelu(h));
53 54 }
54 55
55 56 private:
56 − bool swiglu_;
57 + bool swiglu_, relu2_;
57 58 std::unique_ptr<Linear> w1_, w2_, w3_;
58 59 };
59 60
modified src/nn/transformer.h +47 −8
@@ -30,12 +30,19 @@ private:
30 30 Var w_, b_;
31 31 };
32 32
33 −// Pre-norm block: x += attn(norm1(x)); x += mlp(norm2(x))
33 +// Residual block; norm placement is config-selected:
34 +// pre (default): x += attn(norm1(x)); x += mlp(norm2(x))
35 +// post (OLMo 2): x += norm1(attn(x)); x += norm2(mlp(x))
36 +// sandwich (Gemma): x += norm1p(attn(norm1(x))); x += norm2p(mlp(norm2(x)))
34 37 class TransformerBlock : public Module {
35 38 public:
36 − TransformerBlock(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)
37 − : norm1_(std::make_unique<Norm>(cfg)),
38 − attn_(std::make_unique<CausalSelfAttention>(cfg, proj_std, rng)),
39 + TransformerBlock(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng,
40 + int64_t layer_idx = 0)
41 + : placement_(cfg.norm_placement == "pre" ? Placement::Pre
42 + : cfg.norm_placement == "post" ? Placement::Post
43 + : Placement::Sandwich),
44 + norm1_(std::make_unique<Norm>(cfg)),
45 + attn_(std::make_unique<CausalSelfAttention>(cfg, proj_std, rng, layer_idx)),
39 46 norm2_(std::make_unique<Norm>(cfg)) {
40 47 if (cfg.n_experts > 0) mlp_ = std::make_unique<MoEMLP>(cfg, proj_std, rng);
41 48 else mlp_ = std::make_unique<MLP>(cfg, proj_std, rng);
@@ -43,23 +50,55 @@ public:
43 50 absorb("attn", *attn_);
44 51 absorb("norm2", *norm2_);
45 52 absorb("mlp", *mlp_);
53 + if (placement_ == Placement::Sandwich) {
54 + norm1_post_ = std::make_unique<Norm>(cfg);
55 + norm2_post_ = std::make_unique<Norm>(cfg);
56 + absorb("norm1_post", *norm1_post_);
57 + absorb("norm2_post", *norm2_post_);
58 + }
46 59 }
47 60
48 61 Var forward(const Var& x) const {
49 62 const int64_t B = x.value().size(0), T = x.value().size(1), C = x.value().size(2);
50 − Var h = ops::add(x, attn_->forward(norm1_->forward(x)));
51 − Var m = mlp_->forward(norm2_->forward(h).reshaped({B * T, C}));
52 − return ops::add(h, m.reshaped({B, T, C}));
63 + Var a;
64 + switch (placement_) {
65 + case Placement::Pre: a = attn_->forward(norm1_->forward(x)); break;
66 + case Placement::Post: a = norm1_->forward(attn_->forward(x)); break;
67 + case Placement::Sandwich:
68 + a = norm1_post_->forward(attn_->forward(norm1_->forward(x)));
69 + break;
70 + }
71 + Var h = ops::add(x, a);
72 + Var m;
73 + switch (placement_) {
74 + case Placement::Pre:
75 + m = mlp_->forward(norm2_->forward(h).reshaped({B * T, C}))
76 + .reshaped({B, T, C});
77 + break;
78 + case Placement::Post:
79 + m = norm2_->forward(
80 + mlp_->forward(h.reshaped({B * T, C})).reshaped({B, T, C}));
81 + break;
82 + case Placement::Sandwich:
83 + m = norm2_post_->forward(
84 + mlp_->forward(norm2_->forward(h).reshaped({B * T, C}))
85 + .reshaped({B, T, C}));
86 + break;
87 + }
88 + return ops::add(h, m);
53 89 }
54 90
55 91 // MoE load-balance loss of the last forward (undefined for dense MLP).
56 92 Var moe_aux() const { return mlp_->aux(); }
57 93
58 94 private:
95 + enum class Placement { Pre, Post, Sandwich };
96 + Placement placement_;
59 97 std::unique_ptr<Norm> norm1_;
60 98 std::unique_ptr<AttentionBase> attn_;
61 99 std::unique_ptr<Norm> norm2_;
62 100 std::unique_ptr<MLPBase> mlp_;
101 + std::unique_ptr<Norm> norm1_post_, norm2_post_; // sandwich only
63 102 };
64 103
65 104 // Decoder-only transformer, entirely shaped by ModelConfig.
@@ -79,7 +118,7 @@ public:
79 118 absorb("pos_emb", *pos_emb_);
80 119 }
81 120 for (int64_t i = 0; i < cfg.n_layers; ++i) {
82 − blocks_.push_back(std::make_unique<TransformerBlock>(cfg, proj_std, rng));
121 + blocks_.push_back(std::make_unique<TransformerBlock>(cfg, proj_std, rng, i));
83 122 absorb("blocks." + std::to_string(i), *blocks_.back());
84 123 }
85 124 final_norm_ = std::make_unique<Norm>(cfg);
modified src/ops/cpu/cpu_ops.cpp +46 −15
@@ -181,6 +181,23 @@ void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx) {
181 181 }
182 182 }
183 183
184 +void 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 +}
192 +
193 +void 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 +}
200 +
184 201 void softcap(const Tensor& x, float cap, Tensor& out) {
185 202 const float* px = x.data<float>();
186 203 float* po = out.data<float>();
@@ -487,7 +504,7 @@ void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight)
487 504
488 505 namespace {
489 506 void rope_impl(const float* in, float* out, int64_t B, int64_t T, int64_t H, int64_t hd,
490 − float theta, int64_t pos_offset, bool inverse, bool accumulate) {
507 + const float* freqs, int64_t pos_offset, bool inverse, bool accumulate) {
491 508 const int64_t C = H * hd;
492 509 parallel_rows(B * T, [&](int64_t bt) {
493 510 const int64_t t = bt % T;
@@ -496,8 +513,7 @@ void rope_impl(const float* in, float* out, int64_t B, int64_t T, int64_t H, int
496 513 float* dst = out + bt * C;
497 514 for (int64_t h = 0; h < H; ++h) {
498 515 for (int64_t k = 0; k < hd / 2; ++k) {
499 − const float freq = std::pow(theta, -2.0f * float(k) / float(hd));
500 − const float angle = pos * freq;
516 + const float angle = pos * freqs[k];
501 517 const float c = std::cos(angle);
502 518 const float s = inverse ? -std::sin(angle) : std::sin(angle);
503 519 const int64_t i0 = h * hd + 2 * k;
@@ -517,26 +533,28 @@ void rope_impl(const float* in, float* out, int64_t B, int64_t T, int64_t H, int
517 533 }
518 534 } // namespace
519 535
520 −void rope(const Tensor& x, int64_t n_heads, float theta, int64_t pos_offset, Tensor& out) {
536 +void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset,
537 + Tensor& out) {
521 538 check(x.ndim() == 3, "rope: expected [B,T,C]");
522 539 const int64_t hd = x.size(2) / n_heads;
523 540 check(hd % 2 == 0, "rope: head_dim must be even");
541 + check(freqs.numel() == hd / 2, "rope: freqs must have head_dim/2 entries");
524 542 rope_impl(x.data<float>(), out.data<float>(), x.size(0), x.size(1), n_heads, hd,
525 − theta, pos_offset, /*inverse=*/false, /*accumulate=*/false);
543 + freqs.data<float>(), pos_offset, /*inverse=*/false, /*accumulate=*/false);
526 544 }
527 545
528 −void rope_backward(const Tensor& dout, int64_t n_heads, float theta, int64_t pos_offset,
529 − Tensor& dx) {
546 +void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,
547 + int64_t pos_offset, Tensor& dx) {
530 548 const int64_t hd = dout.size(2) / n_heads;
531 549 rope_impl(dout.data<float>(), dx.data<float>(), dout.size(0), dout.size(1), n_heads,
532 − hd, theta, pos_offset, /*inverse=*/true, /*accumulate=*/true);
550 + hd, freqs.data<float>(), pos_offset, /*inverse=*/true, /*accumulate=*/true);
533 551 }
534 552
535 553 // ---- attention -------------------------------------------------------------------
536 554
537 555 void attention(const Tensor& q, const Tensor& k, const Tensor& v,
538 556 int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
539 − Tensor& out, Tensor* probs_out) {
557 + Tensor& out, Tensor* probs_out, int64_t window, float attn_softcap) {
540 558 check(q.ndim() == 3 && k.ndim() == 3 && v.ndim() == 3, "attention: expected [B,T,C]");
541 559 const int64_t B = q.size(0), T = q.size(1);
542 560 const int64_t hd = q.size(2) / n_heads;
@@ -560,29 +578,34 @@ void attention(const Tensor& q, const Tensor& k, const Tensor& v,
560 578 for (int64_t i = 0; i < T; ++i) {
561 579 const float* qi = pq + (b * T + i) * Cq + h * hd;
562 580 const int64_t jmax = causal ? i : T - 1;
581 + // sliding window: attend only to the last `window` keys (incl. self)
582 + const int64_t jmin = window > 0 ? std::max<int64_t>(0, i - window + 1) : 0;
563 583 // scores (masked positions never written; treated as prob 0)
564 584 float m = -INFINITY;
565 − for (int64_t j = 0; j <= jmax; ++j) {
585 + for (int64_t j = jmin; j <= jmax; ++j) {
566 586 const float* kj = pk + (b * T + j) * Ckv + hkv * hd;
567 587 float s = 0.0f;
568 588 for (int64_t d = 0; d < hd; ++d) s += qi[d] * kj[d];
569 589 s *= scale;
590 + if (attn_softcap > 0.0f)
591 + s = attn_softcap * std::tanh(s / attn_softcap);
570 592 P[i * T + j] = s;
571 593 m = std::max(m, s);
572 594 }
573 595 float sum = 0.0f;
574 − for (int64_t j = 0; j <= jmax; ++j) {
596 + for (int64_t j = jmin; j <= jmax; ++j) {
575 597 const float e = std::exp(P[i * T + j] - m);
576 598 P[i * T + j] = e;
577 599 sum += e;
578 600 }
579 601 const float inv = 1.0f / sum;
580 − for (int64_t j = 0; j <= jmax; ++j) P[i * T + j] *= inv;
602 + for (int64_t j = 0; j < jmin; ++j) P[i * T + j] = 0.0f;
603 + for (int64_t j = jmin; j <= jmax; ++j) P[i * T + j] *= inv;
581 604 for (int64_t j = jmax + 1; j < T; ++j) P[i * T + j] = 0.0f;
582 605
583 606 float* oi = po + (b * T + i) * Cq + h * hd;
584 607 for (int64_t d = 0; d < hd; ++d) oi[d] = 0.0f;
585 − for (int64_t j = 0; j <= jmax; ++j) {
608 + for (int64_t j = jmin; j <= jmax; ++j) {
586 609 const float p = P[i * T + j];
587 610 const float* vj = pv + (b * T + j) * Ckv + hkv * hd;
588 611 for (int64_t d = 0; d < hd; ++d) oi[d] += p * vj[d];
@@ -594,7 +617,7 @@ void attention(const Tensor& q, const Tensor& k, const Tensor& v,
594 617 void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
595 618 const Tensor& probs, const Tensor& out, const Tensor& dout,
596 619 int64_t n_heads, int64_t n_kv_heads, float scale,
597 − Tensor& dq, Tensor& dk, Tensor& dv) {
620 + Tensor& dq, Tensor& dk, Tensor& dv, float attn_softcap) {
598 621 const int64_t B = q.size(0), T = q.size(1);
599 622 const int64_t hd = q.size(2) / n_heads;
600 623 const int64_t rep = n_heads / n_kv_heads;
@@ -637,7 +660,15 @@ void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
637 660
638 661 float dp = 0.0f;
639 662 for (int64_t d = 0; d < hd; ++d) dp += doi[d] * vj[d];
640 − const float ds = p * (dp - row_dot) * scale;
663 + float ds = p * (dp - row_dot) * scale;
664 + if (attn_softcap > 0.0f) {
665 + // chain through s' = cap·tanh(s/cap): recompute the raw
666 + // score, factor is 1 − tanh²
667 + float s = 0.0f;
668 + for (int64_t d = 0; d < hd; ++d) s += qi[d] * kj[d];
669 + const float t = std::tanh(s * scale / attn_softcap);
670 + ds *= 1.0f - t * t;
671 + }
641 672
642 673 for (int64_t d = 0; d < hd; ++d) {
643 674 dvj[d] += p * doi[d];
modified src/ops/cpu/cpu_ops.h +14 −6
@@ -36,6 +36,8 @@ void silu(const Tensor& x, Tensor& out);
36 36 void silu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);
37 37 void gelu(const Tensor& x, Tensor& out); // tanh approximation (GPT-2)
38 38 void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);
39 +void relu2(const Tensor& x, Tensor& out); // max(x,0)^2
40 +void relu2_backward(const Tensor& x, const Tensor& dout, Tensor& dx);
39 41 void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap)
40 42 // dx += dout * (1 - (y/cap)^2) — takes the forward OUTPUT y
41 43 void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx);
@@ -62,24 +64,30 @@ void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight);
62 64 // ---- RoPE (interleaved-pairs / GPT-J convention; see RESEARCH.md §7) -------
63 65 // x: [B, T, H*hd]; rotates pairs (2k, 2k+1) inside each head. pos_offset
64 66 // shifts absolute positions (KV-cache generation).
65 −void rope(const Tensor& x, int64_t n_heads, float theta, int64_t pos_offset, Tensor& out);
67 +// freqs: [head_dim/2] per-pair inverse frequencies, precomputed by the caller.
68 +void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset,
69 + Tensor& out);
66 70 // Backward of a rotation is the inverse rotation applied to dout.
67 −void rope_backward(const Tensor& dout, int64_t n_heads, float theta, int64_t pos_offset,
68 − Tensor& dx);
71 +void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,
72 + int64_t pos_offset, Tensor& dx);
69 73
70 74 // ---- attention (composed reference: scores → mask → softmax → PV) ----------
71 75 // q: [B, T, H*hd], k/v: [B, T, Hkv*hd], out: [B, T, H*hd]. GQA via
72 76 // kv_head = h / (H / Hkv). If probs_out is non-null it receives the softmax
73 77 // probabilities [B, H, T, T] (needed for backward).
78 +// window > 0: sliding-window attention (attend to the last `window` keys,
79 +// self included). attn_softcap > 0: cap·tanh(score/cap) pre-softmax (Gemma 2).
74 80 void attention(const Tensor& q, const Tensor& k, const Tensor& v,
75 81 int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
76 − Tensor& out, Tensor* probs_out);
82 + Tensor& out, Tensor* probs_out, int64_t window = 0,
83 + float attn_softcap = 0.0f);
77 84 // Uses the FlashAttention-2 identity rowsum(dP ∘ P) == dO_i · O_i, so `out`
78 −// (the forward result) is required.
85 +// (the forward result) is required. Masked positions have prob 0, so the
86 +// window needs no explicit handling here.
79 87 void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
80 88 const Tensor& probs, const Tensor& out, const Tensor& dout,
81 89 int64_t n_heads, int64_t n_kv_heads, float scale,
82 − Tensor& dq, Tensor& dk, Tensor& dv);
90 + Tensor& dq, Tensor& dk, Tensor& dv, float attn_softcap = 0.0f);
83 91
84 92 // ---- QAT / MoE --------------------------------------------------------------
85 93 // Per-row fake quantization of a 2-D weight; mode 0 = int8 (absmax/127),
modified src/ops/metal/metal_ops.cpp +41 −20
@@ -218,6 +218,14 @@ void gelu(const Tensor& x, Tensor& out) {
218 218 encode_flat("gelu_f32", {&x, &out}, nullptr, 0, x.numel());
219 219 }
220 220
221 +void relu2(const Tensor& x, Tensor& out) {
222 + encode_flat("relu2_f32", {&x, &out}, nullptr, 0, x.numel());
223 +}
224 +
225 +void relu2_backward(const Tensor& x, const Tensor& dout, Tensor& dx) {
226 + encode_flat("relu2_bwd_f32", {&x, &dout, &dx}, nullptr, 0, x.numel());
227 +}
228 +
221 229 void softcap(const Tensor& x, float cap, Tensor& out) {
222 230 encode_flat("softcap_f32", {&x, &out}, &cap, sizeof(cap), x.numel());
223 231 }
@@ -437,14 +445,14 @@ void layernorm_backward(const Tensor& x, const Tensor& w, float eps,
437 445 namespace {
438 446 struct RopeParams {
439 447 uint32_t T, H, HD;
440 − float theta;
441 448 uint32_t pos_offset;
442 449 };
443 450
444 −void rope_encode(const Tensor& in, Tensor& out, int64_t n_heads, float theta,
451 +void rope_encode(const Tensor& in, Tensor& out, int64_t n_heads, const Tensor& freqs,
445 452 int64_t pos_offset, bool inverse) {
446 453 const int64_t B = in.size(0), T = in.size(1), C = in.size(2);
447 454 const int64_t hd = C / n_heads;
455 + check(freqs.numel() == hd / 2, "rope: freqs must have head_dim/2 entries");
448 456 Device& dev = Device::get();
449 457 MTL::FunctionConstantValues* constants = MTL::FunctionConstantValues::alloc()->init();
450 458 // slots 0/1 belong to matmul TA/TB; RoPE uses slot 2
@@ -457,21 +465,22 @@ void rope_encode(const Tensor& in, Tensor& out, int64_t n_heads, float theta,
457 465 enc->setComputePipelineState(pso);
458 466 enc->setBuffer(in.buffer(), in.buffer_offset(), 0);
459 467 enc->setBuffer(out.buffer(), out.buffer_offset(), 1);
460 − RopeParams p{uint32_t(T), uint32_t(n_heads), uint32_t(hd), theta,
461 − uint32_t(pos_offset)};
468 + RopeParams p{uint32_t(T), uint32_t(n_heads), uint32_t(hd), uint32_t(pos_offset)};
462 469 enc->setBytes(&p, sizeof(p), 2);
470 + enc->setBuffer(freqs.buffer(), freqs.buffer_offset(), 3);
463 471 const int64_t pairs = B * T * n_heads * (hd / 2);
464 472 enc->dispatchThreads(MTL::Size(NS::UInteger(pairs), 1, 1), MTL::Size(256, 1, 1));
465 473 }
466 474 } // namespace
467 475
468 −void rope(const Tensor& x, int64_t n_heads, float theta, int64_t pos_offset, Tensor& out) {
469 − rope_encode(x, out, n_heads, theta, pos_offset, /*inverse=*/false);
476 +void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset,
477 + Tensor& out) {
478 + rope_encode(x, out, n_heads, freqs, pos_offset, /*inverse=*/false);
470 479 }
471 480
472 −void rope_backward(const Tensor& dout, int64_t n_heads, float theta, int64_t pos_offset,
473 − Tensor& dx) {
474 − rope_encode(dout, dx, n_heads, theta, pos_offset, /*inverse=*/true);
481 +void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,
482 + int64_t pos_offset, Tensor& dx) {
483 + rope_encode(dout, dx, n_heads, freqs, pos_offset, /*inverse=*/true);
475 484 }
476 485
477 486 // ---- embedding -------------------------------------------------------------------
@@ -517,12 +526,14 @@ struct AttnParams {
517 526 uint32_t B, T, H, HKV, HD;
518 527 float scale;
519 528 uint32_t causal;
529 + uint32_t window;
530 + float softcap;
520 531 };
521 532 } // namespace
522 533
523 534 void attention(const Tensor& q, const Tensor& k, const Tensor& v,
524 535 int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
525 − Tensor& out, Tensor* probs_out) {
536 + Tensor& out, Tensor* probs_out, int64_t window, float attn_softcap) {
526 537 check(probs_out != nullptr, "attention: probs_out required (unfused path)");
527 538 const int64_t B = q.size(0), T = q.size(1);
528 539 const int64_t hd = q.size(2) / n_heads;
@@ -536,7 +547,8 @@ void attention(const Tensor& q, const Tensor& k, const Tensor& v,
536 547 enc->setBuffer(out.buffer(), out.buffer_offset(), 3);
537 548 enc->setBuffer(probs_out->buffer(), probs_out->buffer_offset(), 4);
538 549 AttnParams p{uint32_t(B), uint32_t(T), uint32_t(n_heads), uint32_t(n_kv_heads),
539 − uint32_t(hd), scale, causal ? 1u : 0u};
550 + uint32_t(hd), scale, causal ? 1u : 0u, uint32_t(window),
551 + attn_softcap};
540 552 enc->setBytes(&p, sizeof(p), 5);
541 553 enc->dispatchThreads(MTL::Size(NS::UInteger(B * n_heads * T), 1, 1),
542 554 MTL::Size(64, 1, 1));
@@ -545,12 +557,13 @@ void attention(const Tensor& q, const Tensor& k, const Tensor& v,
545 557 void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
546 558 const Tensor& probs, const Tensor& out, const Tensor& dout,
547 559 int64_t n_heads, int64_t n_kv_heads, float scale,
548 − Tensor& dq, Tensor& dk, Tensor& dv) {
560 + Tensor& dq, Tensor& dk, Tensor& dv, float attn_softcap) {
549 561 const int64_t B = q.size(0), T = q.size(1);
550 562 const int64_t hd = q.size(2) / n_heads;
551 563 Device& dev = Device::get();
564 + // window is irrelevant here: masked positions carry prob 0 in `probs`.
552 565 AttnParams p{uint32_t(B), uint32_t(T), uint32_t(n_heads), uint32_t(n_kv_heads),
553 − uint32_t(hd), scale, 1u};
566 + uint32_t(hd), scale, 1u, 0u, attn_softcap};
554 567
555 568 // D[b,h,i] = dO_i · O_i == rowsum(dP ∘ P): computed once per query row so
556 569 // neither backward kernel needs the O(T) inner recompute.
@@ -678,6 +691,8 @@ struct FlashParams {
678 691 uint32_t B, T, H, HKV;
679 692 float scale;
680 693 uint32_t causal;
694 + uint32_t window; // sliding window; 0 = full. Scalar kernels only — the
695 + // caller routes window > 0 away from MMA.
681 696 };
682 697
683 698 // Must match the INSTANTIATE_FLASH list in flash_attention.metal.
@@ -697,13 +712,15 @@ bool flash_supported(int64_t head_dim) {
697 712
698 713 void flash_attention(const Tensor& q, const Tensor& k, const Tensor& v,
699 714 int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
700 − Tensor& out, Tensor& lse, FlashKernel kernel) {
715 + Tensor& out, Tensor& lse, FlashKernel kernel, int64_t window) {
701 716 const int64_t B = q.size(0), T = q.size(1);
702 717 const int64_t hd = q.size(2) / n_heads;
703 718 check(flash_supported(hd), "flash_attention: unsupported head_dim");
704 719
705 − if (kernel == FlashKernel::Auto) kernel = FlashKernel::MMA;
720 + if (kernel == FlashKernel::Auto)
721 + kernel = window > 0 ? FlashKernel::Scalar : FlashKernel::MMA;
706 722 const bool mma = kernel == FlashKernel::MMA;
723 + check(!(mma && window > 0), "flash_attention: MMA kernel has no window support");
707 724
708 725 Device& dev = Device::get();
709 726 MTL::ComputePipelineState* pso = dev.pipeline(
@@ -716,7 +733,7 @@ void flash_attention(const Tensor& q, const Tensor& k, const Tensor& v,
716 733 enc->setBuffer(out.buffer(), out.buffer_offset(), 3);
717 734 enc->setBuffer(lse.buffer(), lse.buffer_offset(), 4);
718 735 FlashParams p{uint32_t(B), uint32_t(T), uint32_t(n_heads), uint32_t(n_kv_heads),
719 − scale, causal ? 1u : 0u};
736 + scale, causal ? 1u : 0u, uint32_t(window)};
720 737 enc->setBytes(&p, sizeof(p), 5);
721 738 // One threadgroup per (query block, head, batch). Block size and thread
722 739 // count must match the kernel's enums: the scalar kernel is one thread
@@ -736,22 +753,26 @@ void flash_attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
736 753 const Tensor& out, const Tensor& lse, const Tensor& dout,
737 754 int64_t n_heads, int64_t n_kv_heads, bool causal,
738 755 float scale, Tensor& dq, Tensor& dk, Tensor& dv,
739 − FlashKernel kernel) {
756 + FlashKernel kernel, int64_t window) {
740 757 const int64_t B = q.size(0), T = q.size(1);
741 758 const int64_t hd = q.size(2) / n_heads;
742 759 check(flash_supported(hd), "flash_attention_backward: unsupported head_dim");
743 − if (kernel == FlashKernel::Auto) kernel = FlashKernel::MMA;
760 + if (kernel == FlashKernel::Auto)
761 + kernel = window > 0 ? FlashKernel::Scalar : FlashKernel::MMA;
744 762 const bool mma = kernel == FlashKernel::MMA;
763 + check(!(mma && window > 0),
764 + "flash_attention_backward: MMA kernel has no window support");
745 765
746 766 Device& dev = Device::get();
747 767 FlashParams p{uint32_t(B), uint32_t(T), uint32_t(n_heads), uint32_t(n_kv_heads),
748 − scale, causal ? 1u : 0u};
768 + scale, causal ? 1u : 0u, uint32_t(window)};
749 769
750 770 // D[b,h,i] = dO_i . O_i (shared with the unfused path).
751 771 Tensor d_term = Tensor::empty({B, n_heads, T});
752 772 {
753 773 AttnParams ap{uint32_t(B), uint32_t(T), uint32_t(n_heads),
754 − uint32_t(n_kv_heads), uint32_t(hd), scale, causal ? 1u : 0u};
774 + uint32_t(n_kv_heads), uint32_t(hd), scale, causal ? 1u : 0u,
775 + 0u, 0.0f};
755 776 MTL::ComputePipelineState* pso = dev.pipeline("attention_bwd_d_f32");
756 777 MTL::ComputeCommandEncoder* enc = Stream::get().encoder();
757 778 enc->setComputePipelineState(pso);
modified src/ops/metal/metal_ops.h +19 −7
@@ -88,6 +88,8 @@ void scale(const Tensor& a, float s, Tensor& out);
88 88 void add_bias(const Tensor& x, const Tensor& bias, Tensor& out);
89 89 void silu(const Tensor& x, Tensor& out);
90 90 void gelu(const Tensor& x, Tensor& out);
91 +void relu2(const Tensor& x, Tensor& out); // max(x,0)^2
92 +void relu2_backward(const Tensor& x, const Tensor& dout, Tensor& dx);
91 93 void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap)
92 94 void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx);
93 95
@@ -106,20 +108,26 @@ void rmsnorm_backward(const Tensor& x, const Tensor& w, float eps,
106 108 void layernorm_backward(const Tensor& x, const Tensor& w, float eps,
107 109 const Tensor& dout, Tensor& dx, Tensor& dw, Tensor& db);
108 110
109 −void rope(const Tensor& x, int64_t n_heads, float theta, int64_t pos_offset, Tensor& out);
110 −void rope_backward(const Tensor& dout, int64_t n_heads, float theta, int64_t pos_offset,
111 − Tensor& dx);
111 +// freqs: [head_dim/2] per-pair inverse frequencies, host-precomputed
112 +// (fixed theta, llama3 rope-scaling, per-layer theta — all just tables).
113 +void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset,
114 + Tensor& out);
115 +void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,
116 + int64_t pos_offset, Tensor& dx);
112 117
113 118 void embedding(const Tensor& weight, const Tensor& ids, Tensor& out); // ids i32
114 119 void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight);
115 120
121 +// window > 0: sliding-window attention; attn_softcap > 0: cap·tanh pre-softmax
122 +// (unfused path only — the fused kernels don't support softcap).
116 123 void attention(const Tensor& q, const Tensor& k, const Tensor& v,
117 124 int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
118 − Tensor& out, Tensor* probs_out);
125 + Tensor& out, Tensor* probs_out, int64_t window = 0,
126 + float attn_softcap = 0.0f);
119 127 void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
120 128 const Tensor& probs, const Tensor& out, const Tensor& dout,
121 129 int64_t n_heads, int64_t n_kv_heads, float scale,
122 − Tensor& dq, Tensor& dk, Tensor& dv);
130 + Tensor& dq, Tensor& dk, Tensor& dv, float attn_softcap = 0.0f);
123 131
124 132 // Fused (flash) attention: stores no T x T probabilities, only the per-row
125 133 // logsumexp `lse` [B, n_heads, T] that the backward re-expands. Supported for
@@ -128,14 +136,18 @@ bool flash_supported(int64_t head_dim);
128 136 // Scalar: one thread per query row. MMA: simdgroup_matrix 8x8 tiles. Same
129 137 // outputs; Auto picks MMA where supported.
130 138 enum class FlashKernel { Auto, Scalar, MMA };
139 +// window > 0 requires the Scalar kernel (Auto routes there automatically);
140 +// out-of-window KV blocks are skipped, so cost scales with the window.
131 141 void flash_attention(const Tensor& q, const Tensor& k, const Tensor& v,
132 142 int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
133 − Tensor& out, Tensor& lse, FlashKernel kernel = FlashKernel::Auto);
143 + Tensor& out, Tensor& lse, FlashKernel kernel = FlashKernel::Auto,
144 + int64_t window = 0);
134 145 void flash_attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
135 146 const Tensor& out, const Tensor& lse, const Tensor& dout,
136 147 int64_t n_heads, int64_t n_kv_heads, bool causal,
137 148 float scale, Tensor& dq, Tensor& dk, Tensor& dv,
138 − FlashKernel kernel = FlashKernel::Auto);
149 + FlashKernel kernel = FlashKernel::Auto,
150 + int64_t window = 0);
139 151
140 152 // losses: [N] per-row buffer; loss_out: [1] mean over n_valid. dlogits
141 153 // optional (accumulated). ids must be i32.
modified src/ops/ops.cpp +87 −30
@@ -4,6 +4,8 @@
4 4 #include "ops/cpu/cpu_ops.h"
5 5 #include "ops/metal/metal_ops.h"
6 6
7 +#include <cmath>
8 +
7 9 namespace forge::ops {
8 10
9 11 namespace {
@@ -208,6 +210,23 @@ Var gelu(const Var& x) {
208 210 return result;
209 211 }
210 212
213 +Var relu2(const Var& x) {
214 + Tensor out = Tensor::empty(x.value().shape());
215 + if (gpu()) metal::relu2(x.value(), out);
216 + else cpu::relu2(x.value(), out);
217 +
218 + const bool needs = grad_needed({&x});
219 + Var result(std::move(out), needs);
220 + if (needs) {
221 + Tape::get().record([x, result]() {
222 + if (!x.requires_grad()) return;
223 + if (gpu()) metal::relu2_backward(x.value(), result.grad(), x.grad());
224 + else cpu::relu2_backward(x.value(), result.grad(), x.grad());
225 + });
226 + }
227 + return result;
228 +}
229 +
211 230 Var softcap(const Var& x, float cap) {
212 231 Tensor out = Tensor::empty(x.value().shape());
213 232 if (gpu()) metal::softcap(x.value(), cap, out);
@@ -370,25 +389,56 @@ Var embedding(const Var& weight, const Tensor& ids) {
370 389 return result;
371 390 }
372 391
373 −Var rope(const Var& x, int64_t n_heads, float theta, int64_t pos_offset) {
392 +Tensor rope_freqs(int64_t head_dim, float theta, float scale_factor,
393 + float low_freq_factor, float high_freq_factor, int64_t original_ctx) {
394 + Tensor freqs = Tensor::empty({head_dim / 2});
395 + float* f = freqs.data<float>();
396 + constexpr float kTwoPi = 6.28318530717958647692f;
397 + for (int64_t k = 0; k < head_dim / 2; ++k) {
398 + float inv = std::pow(theta, -2.0f * float(k) / float(head_dim));
399 + if (scale_factor > 0.0f) {
400 + // HF "llama3" rope scaling (modeling_rope_utils.py).
401 + const float wavelen = kTwoPi / inv;
402 + const float lo = float(original_ctx) / low_freq_factor; // long waves
403 + const float hi = float(original_ctx) / high_freq_factor; // short waves
404 + if (wavelen > lo) {
405 + inv /= scale_factor;
406 + } else if (wavelen > hi) {
407 + const float s = (float(original_ctx) / wavelen - low_freq_factor) /
408 + (high_freq_factor - low_freq_factor);
409 + inv = (1.0f - s) * inv / scale_factor + s * inv;
410 + }
411 + }
412 + f[k] = inv;
413 + }
414 + return freqs;
415 +}
416 +
417 +Var rope(const Var& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset) {
374 418 Tensor out = Tensor::empty(x.value().shape());
375 − if (gpu()) metal::rope(x.value(), n_heads, theta, pos_offset, out);
376 − else cpu::rope(x.value(), n_heads, theta, pos_offset, out);
419 + if (gpu()) metal::rope(x.value(), n_heads, freqs, pos_offset, out);
420 + else cpu::rope(x.value(), n_heads, freqs, pos_offset, out);
377 421
378 422 const bool needs = grad_needed({&x});
379 423 Var result(std::move(out), needs);
380 424 if (needs) {
381 − Tape::get().record([x, n_heads, theta, pos_offset, result]() {
425 + Tape::get().record([x, n_heads, freqs, pos_offset, result]() {
382 426 if (!x.requires_grad()) return;
383 − if (gpu()) metal::rope_backward(result.grad(), n_heads, theta, pos_offset, x.grad());
384 − else cpu::rope_backward(result.grad(), n_heads, theta, pos_offset, x.grad());
427 + if (gpu()) metal::rope_backward(result.grad(), n_heads, freqs, pos_offset, x.grad());
428 + else cpu::rope_backward(result.grad(), n_heads, freqs, pos_offset, x.grad());
385 429 });
386 430 }
387 431 return result;
388 432 }
389 433
434 +Var rope(const Var& x, int64_t n_heads, float theta, int64_t pos_offset) {
435 + const int64_t hd = x.value().size(2) / n_heads;
436 + return rope(x, n_heads, rope_freqs(hd, theta), pos_offset);
437 +}
438 +
390 439 Var attention(const Var& q, const Var& k, const Var& v,
391 − int64_t n_heads, int64_t n_kv_heads, bool causal, float scale) {
440 + int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
441 + int64_t window, float attn_softcap) {
392 442 const int64_t B = q.value().size(0), T = q.value().size(1);
393 443 const int64_t head_dim = q.value().size(2) / n_heads;
394 444 Tensor out = Tensor::empty(q.value().shape());
@@ -396,14 +446,17 @@ Var attention(const Var& q, const Var& k, const Var& v,
396 446
397 447 // Fused path: keeps only the per-row logsumexp instead of a [B,H,T,T]
398 448 // probability tensor, so memory is linear in T rather than quadratic.
399 − if (gpu() && metal::flash_supported(head_dim)) {
449 + // Softcap needs the materialized-probs path (the fused kernels' online
450 + // softmax has no cap hook yet).
451 + if (gpu() && metal::flash_supported(head_dim) && attn_softcap <= 0.0f) {
400 452 Tensor lse = Tensor::empty({B, n_heads, T});
401 453 metal::flash_attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads,
402 − causal, scale, out, lse);
454 + causal, scale, out, lse, metal::FlashKernel::Auto,
455 + window);
403 456 Var result(std::move(out), needs);
404 457 if (needs) {
405 458 Tape::get().record(
406 − [q, k, v, lse, n_heads, n_kv_heads, causal, scale, result]() {
459 + [q, k, v, lse, n_heads, n_kv_heads, causal, scale, window, result]() {
407 460 Tensor dq_scratch, dk_scratch, dv_scratch;
408 461 Tensor& dq = q.requires_grad() ? q.grad()
409 462 : (dq_scratch = Tensor::zeros(q.value().shape()));
@@ -414,7 +467,8 @@ Var attention(const Var& q, const Var& k, const Var& v,
414 467 metal::flash_attention_backward(q.value(), k.value(), v.value(),
415 468 result.value(), lse, result.grad(),
416 469 n_heads, n_kv_heads, causal, scale,
417 − dq, dk, dv);
470 + dq, dk, dv,
471 + metal::FlashKernel::Auto, window);
418 472 });
419 473 }
420 474 return result;
@@ -423,30 +477,33 @@ Var attention(const Var& q, const Var& k, const Var& v,
423 477 Tensor probs = Tensor::empty({B, n_heads, T, T});
424 478 if (gpu())
425 479 metal::attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads, causal,
426 − scale, out, &probs);
480 + scale, out, &probs, window, attn_softcap);
427 481 else
428 482 cpu::attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads, causal,
429 − scale, out, &probs);
483 + scale, out, &probs, window, attn_softcap);
430 484
431 485 Var result(std::move(out), needs);
432 486 if (needs) {
433 − Tape::get().record([q, k, v, probs, n_heads, n_kv_heads, scale, result]() {
434 − Tensor dq_scratch, dk_scratch, dv_scratch;
435 − Tensor& dq = q.requires_grad() ? q.grad()
436 − : (dq_scratch = Tensor::zeros(q.value().shape()));
437 − Tensor& dk = k.requires_grad() ? k.grad()
438 − : (dk_scratch = Tensor::zeros(k.value().shape()));
439 − Tensor& dv = v.requires_grad() ? v.grad()
440 − : (dv_scratch = Tensor::zeros(v.value().shape()));
441 − if (gpu())
442 − metal::attention_backward(q.value(), k.value(), v.value(), probs,
443 − result.value(), result.grad(), n_heads,
444 − n_kv_heads, scale, dq, dk, dv);
445 − else
446 − cpu::attention_backward(q.value(), k.value(), v.value(), probs,
447 − result.value(), result.grad(), n_heads,
448 − n_kv_heads, scale, dq, dk, dv);
449 − });
487 + Tape::get().record(
488 + [q, k, v, probs, n_heads, n_kv_heads, scale, attn_softcap, result]() {
489 + Tensor dq_scratch, dk_scratch, dv_scratch;
490 + Tensor& dq = q.requires_grad() ? q.grad()
491 + : (dq_scratch = Tensor::zeros(q.value().shape()));
492 + Tensor& dk = k.requires_grad() ? k.grad()
493 + : (dk_scratch = Tensor::zeros(k.value().shape()));
494 + Tensor& dv = v.requires_grad() ? v.grad()
495 + : (dv_scratch = Tensor::zeros(v.value().shape()));
496 + if (gpu())
497 + metal::attention_backward(q.value(), k.value(), v.value(), probs,
498 + result.value(), result.grad(), n_heads,
499 + n_kv_heads, scale, dq, dk, dv,
500 + attn_softcap);
501 + else
502 + cpu::attention_backward(q.value(), k.value(), v.value(), probs,
503 + result.value(), result.grad(), n_heads,
504 + n_kv_heads, scale, dq, dk, dv,
505 + attn_softcap);
506 + });
450 507 }
451 508 return result;
452 509 }
modified src/ops/ops.h +18 −3
@@ -31,6 +31,7 @@ Var scale(const Var& a, float s);
31 31
32 32 Var silu(const Var& x);
33 33 Var gelu(const Var& x);
34 +Var relu2(const Var& x); // max(x,0)^2 — nanoGPT-speedrun activation
34 35
35 36 // Gemma-style soft capping: cap * tanh(x / cap). Bounds logits smoothly.
36 37 Var softcap(const Var& x, float cap);
@@ -57,12 +58,26 @@ Var layernorm(const Var& x, const Var& w, const Var& b, float eps);
57 58 // weight [V,C], ids [B,T] (u16/i32) → [B,T,C]
58 59 Var embedding(const Var& weight, const Tensor& ids);
59 60
60 −// x [B,T,H*hd], interleaved-pairs RoPE per head
61 +// Build the [head_dim/2] inverse-frequency table: theta^(-2k/hd), optionally
62 +// rescaled HF-"llama3" style (scale_factor > 0): high-frequency components
63 +// untouched, low-frequency divided by factor, smooth blend between.
64 +Tensor rope_freqs(int64_t head_dim, float theta, float scale_factor = 0.0f,
65 + float low_freq_factor = 1.0f, float high_freq_factor = 4.0f,
66 + int64_t original_ctx = 8192);
67 +
68 +// x [B,T,H*hd], interleaved-pairs RoPE per head. freqs from rope_freqs()
69 +// (callers cache it — one table per layer flavour).
70 +Var rope(const Var& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset = 0);
71 +// Convenience: plain-theta table built per call (tests / one-offs).
61 72 Var rope(const Var& x, int64_t n_heads, float theta, int64_t pos_offset = 0);
62 73
63 −// q [B,T,H*hd], k/v [B,T,Hkv*hd] → [B,T,H*hd]; causal + GQA
74 +// q [B,T,H*hd], k/v [B,T,Hkv*hd] → [B,T,H*hd]; causal + GQA.
75 +// window > 0: sliding-window attention (Mistral/Gemma3) — fused scalar path.
76 +// attn_softcap > 0: cap·tanh(score/cap) pre-softmax (Gemma2) — routes to the
77 +// unfused path (probs materialized), so reserve it for short contexts.
64 78 Var attention(const Var& q, const Var& k, const Var& v,
65 − int64_t n_heads, int64_t n_kv_heads, bool causal, float scale);
79 + int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
80 + int64_t window = 0, float attn_softcap = 0.0f);
66 81
67 82 // logits [N,V], targets [N] (i32/u16, ignore_index=-1) → scalar mean loss
68 83 Var cross_entropy(const Var& logits, const Tensor& targets);
modified tests/test_ops.cpp +46 −0
@@ -600,6 +600,52 @@ int main() {
600 600 vcfg.tied_embeddings = true;
601 601 vcfg.qk_norm = true; vcfg.final_softcap = 30.0f; vcfg.scale_embeddings = true;
602 602 test_backend_parity_model("qk-norm + softcap + embed-scale", vcfg);
603 +
604 + // Wave-1 knobs: ReLU² + QKV bias + decoupled head_dim + NoPE + rope scaling
605 + forge::ModelConfig w1;
606 + w1.n_layers = 2; w1.d_model = 16; w1.n_heads = 2; w1.n_kv_heads = 1;
607 + w1.d_ff = 24; w1.vocab_size = 11; w1.context_length = 8;
608 + w1.tied_embeddings = true;
609 + w1.activation = "relu2"; w1.attention_bias = true;
610 + w1.head_dim_override = 10; w1.nope_every = 2;
611 + w1.rope_scale_factor = 8.0f; w1.rope_scale_orig_ctx = 4;
612 + test_backend_parity_model("relu2 + qkv-bias + head_dim10 + nope + rope-scale", w1);
613 +
614 + // Norm placements: OLMo2 post and Gemma sandwich
615 + forge::ModelConfig np;
616 + np.n_layers = 2; np.d_model = 16; np.n_heads = 2; np.n_kv_heads = 2;
617 + np.d_ff = 24; np.vocab_size = 11; np.context_length = 8;
618 + np.tied_embeddings = true;
619 + np.norm_placement = "post";
620 + test_backend_parity_model("post-norm (OLMo2)", np);
621 + np.norm_placement = "sandwich";
622 + test_backend_parity_model("sandwich norm (Gemma)", np);
623 +
624 + // Wave-2: sliding window (Mistral) through the fused scalar path
625 + // (head_dim 64), + local/global pattern, + Gemma3 dual rope theta
626 + forge::ModelConfig sw;
627 + sw.n_layers = 3; sw.d_model = 128; sw.n_heads = 2; sw.n_kv_heads = 1;
628 + sw.d_ff = 96; sw.vocab_size = 11; sw.context_length = 8;
629 + sw.tied_embeddings = true;
630 + sw.sliding_window = 3; sw.sliding_global_every = 2;
631 + sw.rope_theta_global = 1e6f;
632 + test_backend_parity_model("sliding-window 3 + global-every-2 [hd64]", sw);
633 +
634 + // Wave-2: attention softcap (Gemma2) — unfused path on both backends
635 + forge::ModelConfig sc;
636 + sc.n_layers = 2; sc.d_model = 16; sc.n_heads = 2; sc.n_kv_heads = 1;
637 + sc.d_ff = 24; sc.vocab_size = 11; sc.context_length = 8;
638 + sc.tied_embeddings = true;
639 + sc.attn_softcap = 20.0f;
640 + test_backend_parity_model("attn softcap 20 (Gemma2)", sc);
641 +
642 + // Sliding window through the UNFUSED kernel too (head_dim 8)
643 + forge::ModelConfig swu;
644 + swu.n_layers = 2; swu.d_model = 16; swu.n_heads = 2; swu.n_kv_heads = 2;
645 + swu.d_ff = 24; swu.vocab_size = 11; swu.context_length = 8;
646 + swu.tied_embeddings = true;
647 + swu.sliding_window = 3; swu.attn_softcap = 20.0f;
648 + test_backend_parity_model("sliding-window 3 + softcap [unfused hd8]", swu);
603 649 }
604 650 test_metal_moe_quant_ops();
605 651 test_flash_attention();
606 652