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 3 mo ago (Aug 5, 2026) parent 6d77a01 Browse files

17 changed files +572 −139

added configs/gpt-50m-gemma.json +44 −0

++{
++  "_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.",
++  "model": {
++    "name": "gpt-50m-gemma",
++    "n_layers": 12,
++    "d_model": 576,
++    "n_heads": 9,
++    "n_kv_heads": 3,
++    "d_ff": 1536,
++    "vocab_size": 4096,
++    "context_length": 1024,
++    "tied_embeddings": true,
++    "use_rope": true,
++    "rope_theta": 10000.0,
++    "rope_theta_global": 1000000.0,
++    "norm": "rmsnorm",
++    "norm_eps": 1e-06,
++    "activation": "gelu",
++    "dropout": 0.0,
++    "qk_norm": true,
++    "norm_placement": "sandwich",
++    "scale_embeddings": true,
++    "sliding_window": 256,
++    "sliding_global_every": 6
++  },
++  "train": {
++    "lr": 0.0005,
++    "min_lr_ratio": 0.1,
++    "warmup_steps": 117,
++    "max_steps": 1170,
++    "beta1": 0.9,
++    "beta2": 0.95,
++    "eps": 1e-08,
++    "weight_decay": 0.1,
++    "grad_clip": 1.0,
++    "batch_size": 8,
++    "grad_accum_steps": 8,
++    "precision": "f32",
++    "checkpoint_every": 200,
++    "eval_every": 100,
++    "eval_batches": 20,
++    "seed": 1337
++  }
++}

added configs/gpt-50m-mistral.json +39 −0

++{
++  "_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.",
++  "model": {
++    "name": "gpt-50m-mistral",
++    "n_layers": 10,
++    "d_model": 640,
++    "n_heads": 10,
++    "n_kv_heads": 5,
++    "d_ff": 1728,
++    "vocab_size": 4096,
++    "context_length": 1024,
++    "tied_embeddings": true,
++    "use_rope": true,
++    "rope_theta": 10000.0,
++    "norm": "rmsnorm",
++    "norm_eps": 1e-06,
++    "activation": "swiglu",
++    "dropout": 0.0,
++    "sliding_window": 256
++  },
++  "train": {
++    "lr": 0.0005,
++    "min_lr_ratio": 0.1,
++    "warmup_steps": 117,
++    "max_steps": 1170,
++    "beta1": 0.9,
++    "beta2": 0.95,
++    "eps": 1e-08,
++    "weight_decay": 0.1,
++    "grad_clip": 1.0,
++    "batch_size": 8,
++    "grad_accum_steps": 8,
++    "precision": "f32",
++    "checkpoint_every": 200,
++    "eval_every": 100,
++    "eval_batches": 20,
++    "seed": 1337
++  }
++}

modified src/kernels/attention.metal +25 −6

				`@@ -18,6 +18,8 @@ struct AttnParams {`
18	uint B, T, H, HKV, HD;	18	uint B, T, H, HKV, HD;
19	float scale;	19	float scale;
20	uint causal;	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	kernel void attention_fwd_f32(device const float* Q [[buffer(0)]],	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	device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * p.HD;	42	device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * p.HD;
41	device float* prow_base = P + ((ulong(b) * p.H + h) * p.T + i) * p.T;	43	device float* prow_base = P + ((ulong(b) * p.H + h) * p.T + i) * p.T;
42	const uint jmax = p.causal ? i : p.T - 1;	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	float m = -FLT_MAX;	47	float m = -FLT_MAX;
45	− for (uint j = 0; j <= jmax; ++j) {	48	+ for (uint j = jmin; j <= jmax; ++j) {
46	device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;	49	device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
47	float s = 0.0f;	50	float s = 0.0f;
48	for (uint d = 0; d < p.HD; ++d) s = fma(qi[d], kj[d], s);	51	for (uint d = 0; d < p.HD; ++d) s = fma(qi[d], kj[d], s);
49	s *= p.scale;	52	s *= p.scale;
		53	+ if (p.softcap > 0.0f) s = p.softcap * precise::tanh(s / p.softcap);
50	prow_base[j] = s;	54	prow_base[j] = s;
51	m = max(m, s);	55	m = max(m, s);
52	}	56	}
53	float sum = 0.0f;	57	float sum = 0.0f;
54	− for (uint j = 0; j <= jmax; ++j) {	58	+ for (uint j = jmin; j <= jmax; ++j) {
55	const float e = exp(prow_base[j] - m);	59	const float e = exp(prow_base[j] - m);
56	prow_base[j] = e;	60	prow_base[j] = e;
57	sum += e;	61	sum += e;
58	}	62	}
59	const float inv = 1.0f / sum;	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	for (uint j = jmax + 1; j < p.T; ++j) prow_base[j] = 0.0f;	66	for (uint j = jmax + 1; j < p.T; ++j) prow_base[j] = 0.0f;
62		67
63	device float* oi = O + (ulong(b) * p.T + i) * Cq + h * p.HD;	68	device float* oi = O + (ulong(b) * p.T + i) * Cq + h * p.HD;
64	for (uint d = 0; d < p.HD; ++d) oi[d] = 0.0f;	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	const float prob = prow_base[j];	71	const float prob = prow_base[j];
67	if (prob == 0.0f) continue;	72	if (prob == 0.0f) continue;
68	device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;	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	device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;	131	device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * p.HD;
127	float dp = 0.0f;	132	float dp = 0.0f;
128	for (uint d = 0; d < p.HD; ++d) dp = fma(doi[d], vj[d], dp);	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	for (uint d = 0; d < p.HD; ++d) dqi[d] = fma(ds, kj[d], dqi[d]);	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	const float row_dot = D[(ulong(b) * p.H + h) * p.T + i];	179	const float row_dot = D[(ulong(b) * p.H + h) * p.T + i];
168	float dp = 0.0f;	180	float dp = 0.0f;
169	for (uint d = 0; d < p.HD; ++d) dp = fma(doi[d], vj[d], dp);	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	for (uint d = 0; d < p.HD; ++d) {	190	for (uint d = 0; d < p.HD; ++d) {
172	dvj[d] = fma(prob, doi[d], dvj[d]);	191	dvj[d] = fma(prob, doi[d], dvj[d]);
173	dkj[d] = fma(ds, qi[d], dkj[d]);	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	out[gid] = 0.5f * v * (1.0f + precise::tanh(k * (v + 0.044715f * v * v * v)));	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	// Gemma-style logit softcap: out = cap * tanh(x / cap).	70	// Gemma-style logit softcap: out = cap * tanh(x / cap).
56	kernel void softcap_f32(device const float* x [[buffer(0)]],	71	kernel void softcap_f32(device const float* x [[buffer(0)]],
57	device float* out [[buffer(1)]],	72	device float* out [[buffer(1)]],
				`@@ -127,13 +142,15 @@ constant bool ROPE_INV [[function_constant(2)]];`
127		142
128	struct RopeParams {	143	struct RopeParams {
129	uint T, H, HD;	144	uint T, H, HD;
130	− float theta;
131	uint pos_offset;	145	uint pos_offset;
132	};	146	};
133		147
134	−kernel void rope_f32(device const float* x [[buffer(0)]],	148	+// freqs[k] is the per-pair inverse frequency (theta^(-2k/HD), possibly
135	− device float* out [[buffer(1)]],	149	+// rescaled — llama3 rope_scaling, per-layer theta). Host-precomputed.
136	− constant RopeParams& p [[buffer(2)]],	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	uint gid [[thread_position_in_grid]]) {	154	uint gid [[thread_position_in_grid]]) {
138	// gid indexes (bt, h, k) pairs: one thread per rotated pair	155	// gid indexes (bt, h, k) pairs: one thread per rotated pair
139	const uint pairs_per_row = p.H * (p.HD / 2);	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	const uint k = rem % (p.HD / 2);	160	const uint k = rem % (p.HD / 2);
144		161
145	const float pos = float(bt % p.T + p.pos_offset);	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));	163	+ const float angle = pos * freqs[k];
147	− const float angle = pos * freq;
148	const float c = cos(angle);	164	const float c = cos(angle);
149	const float s = ROPE_INV ? -sin(angle) : sin(angle);	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	uint B, T, H, HKV;	32	uint B, T, H, HKV;
33	float scale;	33	float scale;
34	uint causal;	34	uint causal;
		35	+ uint window; // sliding window (keys kept, self incl.); 0 = full attention
35	};	36	};
36		37
37	// ---------------------------------------------------------------- forward	38	// ---------------------------------------------------------------- forward
				`@@ -90,8 +91,15 @@ kernel void flash_attn_fwd(device const float* Q [[buffer(0)]],`
90	const uint row_max = min(qb * TGQ + TGQ - 1, p.T - 1);	91	const uint row_max = min(qb * TGQ + TGQ - 1, p.T - 1);
91	const uint j_end = p.causal ? (row_max + 1) : p.T;	92	const uint j_end = p.causal ? (row_max + 1) : p.T;
92	const uint diag_start = p.causal ? (qb * TGQ) : p.T; // blocks below need no mask	93	const uint diag_start = p.causal ? (qb * TGQ) : p.T; // blocks below need no mask
93	−	94	+ // Sliding window: this row's first visible key, and (uniform across the
94	− for (uint jb = 0; jb < j_end; jb += BKV) {	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	const uint block_len = min(BKV, j_end - jb);	103	const uint block_len = min(BKV, j_end - jb);
96		104
97	threadgroup_barrier(mem_flags::mem_threadgroup);	105	threadgroup_barrier(mem_flags::mem_threadgroup);
				`@@ -109,6 +117,7 @@ kernel void flash_attn_fwd(device const float* Q [[buffer(0)]],`
109	for (uint jr = 0; jr < block_len; ++jr) {	117	for (uint jr = 0; jr < block_len; ++jr) {
110	const uint j = jb + jr;	118	const uint j = jb + jr;
111	if (needs_mask && j > i) break; // rest of the block is masked too	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	threadgroup const float* kj = Ks + jr * HD;	122	threadgroup const float* kj = Ks + jr * HD;
114	float s = 0.0f;	123	float s = 0.0f;
				`@@ -174,8 +183,9 @@ kernel void flash_attn_bwd_dq(device const float* Q [[buffer(0)]],`
174	const float li = L[(ulong(b) * p.H + h) * p.T + i];	183	const float li = L[(ulong(b) * p.H + h) * p.T + i];
175	const float di = D[(ulong(b) * p.H + h) * p.T + i];	184	const float di = D[(ulong(b) * p.H + h) * p.T + i];
176	const uint jmax = p.causal ? i : (p.T - 1);	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	device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * HD;	189	device const float* kj = K + (ulong(b) * p.T + j) * Ckv + hkv * HD;
180	device const float* vj = V + (ulong(b) * p.T + j) * Ckv + hkv * HD;	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	for (uint d = 0; d < HD; ++d) dv[d] = 0.0f;	235	for (uint d = 0; d < HD; ++d) dv[d] = 0.0f;
226		236
227	const uint imin = p.causal ? j : 0;	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	for (uint r = 0; r < rep; ++r) {	240	for (uint r = 0; r < rep; ++r) {
229	const uint h = hkv * rep + r;	241	const uint h = hkv * rep + r;
230	device const float* Lh = L + (ulong(b) * p.H + h) * p.T;	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	device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;	244	device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;
233	device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;	245	device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;
234	float s = 0.0f;	246	float s = 0.0f;
				`@@ -271,11 +283,12 @@ kernel void flash_attn_bwd_dk(device const float* Q [[buffer(0)]],`
271	for (uint d = 0; d < HD; ++d) dk[d] = 0.0f;	283	for (uint d = 0; d < HD; ++d) dk[d] = 0.0f;
272		284
273	const uint imin = p.causal ? j : 0;	285	const uint imin = p.causal ? j : 0;
		286	+ const uint imax = (p.window > 0) ? min(p.T, j + p.window) : p.T;
274	for (uint r = 0; r < rep; ++r) {	287	for (uint r = 0; r < rep; ++r) {
275	const uint h = hkv * rep + r;	288	const uint h = hkv * rep + r;
276	device const float* Lh = L + (ulong(b) * p.H + h) * p.T;	289	device const float* Lh = L + (ulong(b) * p.H + h) * p.T;
277	device const float* Dh = D + (ulong(b) * p.H + h) * p.T;	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	device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;	292	device const float* qi = Q + (ulong(b) * p.T + i) * Cq + h * HD;
280	device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;	293	device const float* doi = dO + (ulong(b) * p.T + i) * Cq + h * HD;
281	float s = 0.0f, dp = 0.0f;	294	float s = 0.0f, dp = 0.0f;

modified src/kernels/flash_attention_mma.metal +1 −0

				`@@ -36,6 +36,7 @@ struct FlashParams {`
36	uint B, T, H, HKV;	36	uint B, T, H, HKV;
37	float scale;	37	float scale;
38	uint causal;	38	uint causal;
		39	+ uint window; // unused here: the host routes window > 0 to the scalar kernels
39	};	40	};
40		41
41	// Enumerators, never `constant constexpr`: see RESEARCH.md 7a — the latter is	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	// Causal multi-head self-attention with GQA and (optional) RoPE.	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	class CausalSelfAttention : public AttentionBase {	27	class CausalSelfAttention : public AttentionBase {
25	public:	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	: n_heads_(cfg.n_heads),	31	: n_heads_(cfg.n_heads),
28	n_kv_heads_(cfg.n_kv_heads),	32	n_kv_heads_(cfg.n_kv_heads),
29	head_dim_(cfg.head_dim()),	33	head_dim_(cfg.head_dim()),
30	− use_rope_(cfg.use_rope),	34	+ use_rope_(cfg.layer_uses_rope(layer_idx)),
31	− rope_theta_(cfg.rope_theta),
32	qk_norm_(cfg.qk_norm),	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	const float std = 0.02f;	39	const float std = 0.02f;
35	const int64_t C = cfg.d_model;	40	const int64_t C = cfg.d_model;
36	− const int64_t Ckv = cfg.n_kv_heads * cfg.head_dim();	41	+ const int64_t Cq = cfg.n_heads * head_dim_;
37	− wq_ = std::make_unique<Linear>(C, C, false, std, rng);	42	+ const int64_t Ckv = cfg.n_kv_heads * head_dim_;
38	− wk_ = std::make_unique<Linear>(C, Ckv, false, std, rng);	43	+ const bool bias = cfg.attention_bias;
39	− wv_ = std::make_unique<Linear>(C, Ckv, false, std, rng);	44	+ wq_ = std::make_unique<Linear>(C, Cq, bias, std, rng);
40	− wo_ = std::make_unique<Linear>(C, C, false, proj_std, rng); // residual projection	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	const ops::QuantMode qm = quant_mode_from(cfg.quant);	48	const ops::QuantMode qm = quant_mode_from(cfg.quant);
42	wq_->set_quant(qm);	49	wq_->set_quant(qm);
43	wk_->set_quant(qm);	50	wk_->set_quant(qm);
				`@@ -49,18 +56,28 @@ public:`
49	absorb("wo", *wo_);	56	absorb("wo", *wo_);
50	if (qk_norm_) {	57	if (qk_norm_) {
51	q_norm_w_ = register_param("q_norm.weight",	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	k_norm_w_ = register_param("k_norm.weight",	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	Var forward(const Var& x) const override {	74	Var forward(const Var& x) const override {
59	const int64_t B = x.value().size(0), T = x.value().size(1), C = x.value().size(2);	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	const int64_t Ckv = n_kv_heads_ * head_dim_;	77	const int64_t Ckv = n_kv_heads_ * head_dim_;
61		78
62	Var x2d = x.reshaped({B * T, C});	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	Var k = wk_->forward(x2d).reshaped({B, T, Ckv});	81	Var k = wk_->forward(x2d).reshaped({B, T, Ckv});
65	Var v = wv_->forward(x2d).reshaped({B, T, Ckv});	82	Var v = wv_->forward(x2d).reshaped({B, T, Ckv});
66		83
				`@@ -68,29 +85,32 @@ public:`
68	// Per-head RMSNorm before RoPE (Qwen3/Gemma3): normalize over	85	// Per-head RMSNorm before RoPE (Qwen3/Gemma3): normalize over
69	// head_dim by viewing heads as rows.	86	// head_dim by viewing heads as rows.
70	q = ops::rmsnorm(q.reshaped({B * T * n_heads_, head_dim_}), q_norm_w_,	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	k = ops::rmsnorm(k.reshaped({B * T * n_kv_heads_, head_dim_}), k_norm_w_,	89	k = ops::rmsnorm(k.reshaped({B * T * n_kv_heads_, head_dim_}), k_norm_w_,
73	norm_eps_).reshaped({B, T, Ckv});	90	norm_eps_).reshaped({B, T, Ckv});
74	}	91	}
75		92
76	if (use_rope_) {	93	if (use_rope_) {
77	− q = ops::rope(q, n_heads_, rope_theta_, 0);	94	+ q = ops::rope(q, n_heads_, freqs_, 0);
78	− k = ops::rope(k, n_kv_heads_, rope_theta_, 0);	95	+ k = ops::rope(k, n_kv_heads_, freqs_, 0);
79	}	96	}
80		97
81	const float scale = 1.0f / std::sqrt(float(head_dim_));	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);	99	+ 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});	100	+ window_, attn_softcap_);
		101	+ return wo_->forward(o.reshaped({B * T, Cq})).reshaped({B, T, C});
84	}	102	}
85		103
86	private:	104	private:
87	int64_t n_heads_, n_kv_heads_, head_dim_;	105	int64_t n_heads_, n_kv_heads_, head_dim_;
88	bool use_rope_;	106	bool use_rope_;
89	− float rope_theta_;
90	bool qk_norm_;	107	bool qk_norm_;
91	float norm_eps_;	108	float norm_eps_;
		109	+ int64_t window_; // 0 = full attention (global layer or SWA off)
		110	+ float attn_softcap_;
92	std::unique_ptr<Linear> wq_, wk_, wv_, wo_;	111	std::unique_ptr<Linear> wq_, wk_, wv_, wo_;
93	Var q_norm_w_, k_norm_w_; // defined only when qk_norm	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	} // namespace forge::nn	116	} // namespace forge::nn

modified src/nn/config.h +61 −10

				`@@ -32,6 +32,20 @@ struct ModelConfig {`
32	bool qk_norm = false; // RMSNorm on q/k per head before RoPE (Qwen3/Gemma3)	32	bool qk_norm = false; // RMSNorm on q/k per head before RoPE (Qwen3/Gemma3)
33	float final_softcap = 0.0f; // logits = cap*tanh(logits/cap) (Gemma2: 30)	33	float final_softcap = 0.0f; // logits = cap*tanh(logits/cap) (Gemma2: 30)
34	bool scale_embeddings = false;// x *= sqrt(d_model) after embedding (Gemma)	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	// Quantization-aware training: linear weights are fake-quantized each	49	// Quantization-aware training: linear weights are fake-quantized each
36	// forward (per-row scales, straight-through estimator in backward).	50	// forward (per-row scales, straight-through estimator in backward).
37	// Master weights and the optimizer stay f32.	51	// Master weights and the optimizer stay f32.
				`@@ -44,21 +58,36 @@ struct ModelConfig {`
44	float moe_aux_weight = 0.01f; // load-balance loss: E * sum_e mean_gate_e^2	58	float moe_aux_weight = 0.01f; // load-balance loss: E * sum_e mean_gate_e^2
45	int64_t n_shared_experts = 0; // always-active experts (DeepSeek/Kimi style)	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	int64_t num_params() const {	76	int64_t num_params() const {
51	const int64_t hd = head_dim();	77	const int64_t hd = head_dim();
52	− const int64_t attn = d_model * d_model // wq	78	+ int64_t attn = d_model * n_heads * hd // wq
53	− + 2 * d_model * n_kv_heads * hd // wk, wv	79	+ + 2 * d_model * n_kv_heads * hd // wk, wv
54	− + d_model * d_model; // wo	80	+ + n_heads * hd * d_model; // wo
		81	+ if (attention_bias) attn += (n_heads + 2 * n_kv_heads) * hd;
55	const int64_t mlp_one = (activation == "swiglu")	82	const int64_t mlp_one = (activation == "swiglu")
56	? 3 * d_model * d_ff	83	? 3 * d_model * d_ff
57	: 2 * d_model * d_ff;	84	: 2 * d_model * d_ff;
58	const int64_t mlp = n_experts > 0	85	const int64_t mlp = n_experts > 0
59	? (n_experts + n_shared_experts) * mlp_one + n_experts * d_model	86	? (n_experts + n_shared_experts) * mlp_one + n_experts * d_model
60	: mlp_one; // experts + router	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	if (qk_norm) norms += 2 * head_dim() * n_layers;	91	if (qk_norm) norms += 2 * head_dim() * n_layers;
63	int64_t total = n_layers * (attn + mlp) + norms + vocab_size * d_model;	92	int64_t total = n_layers * (attn + mlp) + norms + vocab_size * d_model;
64	if (!tied_embeddings) total += vocab_size * d_model;	93	if (!tied_embeddings) total += vocab_size * d_model;
				`@@ -129,15 +158,37 @@ inline void from_json(const nlohmann::json& j, ModelConfig& c) {`
129	c.qk_norm = j.value("qk_norm", c.qk_norm);	158	c.qk_norm = j.value("qk_norm", c.qk_norm);
130	c.final_softcap = j.value("final_softcap", c.final_softcap);	159	c.final_softcap = j.value("final_softcap", c.final_softcap);
131	c.scale_embeddings = j.value("scale_embeddings", c.scale_embeddings);	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)	174	+ if (c.head_dim_override == 0 && c.d_model % c.n_heads != 0)
134	− throw std::runtime_error("config: d_model must be divisible by n_heads");	175	+ throw std::runtime_error("config: d_model must be divisible by n_heads "
		176	+ "(or set head_dim explicitly)");
135	if (c.n_heads % c.n_kv_heads != 0)	177	if (c.n_heads % c.n_kv_heads != 0)
136	throw std::runtime_error("config: n_heads must be divisible by n_kv_heads");	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	if (c.norm != "rmsnorm" && c.norm != "layernorm")	181	if (c.norm != "rmsnorm" && c.norm != "layernorm")
138	throw std::runtime_error("config: norm must be rmsnorm or layernorm");	182	throw std::runtime_error("config: norm must be rmsnorm or layernorm");
139	− if (c.activation != "swiglu" && c.activation != "gelu")	183	+ if (c.activation != "swiglu" && c.activation != "gelu" && c.activation != "relu2")
140	− throw std::runtime_error("config: activation must be swiglu or gelu");	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	if (c.quant != "none" && c.quant != "int8" && c.quant != "ternary")	192	if (c.quant != "none" && c.quant != "int8" && c.quant != "ternary")
142	throw std::runtime_error("config: quant must be none, int8 or ternary");	193	throw std::runtime_error("config: quant must be none, int8 or ternary");
143	if (c.n_experts < 0)	194	if (c.n_experts < 0)

modified src/nn/mlp.h +5 −4

				`@@ -22,11 +22,11 @@ public:`
22	virtual Var aux() const { return Var(); }	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	class MLP : public MLPBase {	26	class MLP : public MLPBase {
27	public:	27	public:
28	MLP(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)	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	const float std = 0.02f;	30	const float std = 0.02f;
31	const int64_t C = cfg.d_model, F = cfg.d_ff;	31	const int64_t C = cfg.d_model, F = cfg.d_ff;
32	w1_ = std::make_unique<Linear>(C, F, false, std, rng);	32	w1_ = std::make_unique<Linear>(C, F, false, std, rng);
				`@@ -49,11 +49,12 @@ public:`
49	Var up = w3_->forward(x);	49	Var up = w3_->forward(x);
50	return w2_->forward(ops::mul(gate, up));	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	private:	56	private:
56	− bool swiglu_;	57	+ bool swiglu_, relu2_;
57	std::unique_ptr<Linear> w1_, w2_, w3_;	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	Var w_, b_;	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	class TransformerBlock : public Module {	37	class TransformerBlock : public Module {
35	public:	38	public:
36	− TransformerBlock(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)	39	+ TransformerBlock(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng,
37	− : norm1_(std::make_unique<Norm>(cfg)),	40	+ int64_t layer_idx = 0)
38	− attn_(std::make_unique<CausalSelfAttention>(cfg, proj_std, rng)),	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	norm2_(std::make_unique<Norm>(cfg)) {	46	norm2_(std::make_unique<Norm>(cfg)) {
40	if (cfg.n_experts > 0) mlp_ = std::make_unique<MoEMLP>(cfg, proj_std, rng);	47	if (cfg.n_experts > 0) mlp_ = std::make_unique<MoEMLP>(cfg, proj_std, rng);
41	else mlp_ = std::make_unique<MLP>(cfg, proj_std, rng);	48	else mlp_ = std::make_unique<MLP>(cfg, proj_std, rng);
				`@@ -43,23 +50,55 @@ public:`
43	absorb("attn", *attn_);	50	absorb("attn", *attn_);
44	absorb("norm2", *norm2_);	51	absorb("norm2", *norm2_);
45	absorb("mlp", *mlp_);	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	Var forward(const Var& x) const {	61	Var forward(const Var& x) const {
49	const int64_t B = x.value().size(0), T = x.value().size(1), C = x.value().size(2);	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)));	63	+ Var a;
51	− Var m = mlp_->forward(norm2_->forward(h).reshaped({B * T, C}));	64	+ switch (placement_) {
52	− return ops::add(h, m.reshaped({B, T, C}));	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	// MoE load-balance loss of the last forward (undefined for dense MLP).	91	// MoE load-balance loss of the last forward (undefined for dense MLP).
56	Var moe_aux() const { return mlp_->aux(); }	92	Var moe_aux() const { return mlp_->aux(); }
57		93
58	private:	94	private:
		95	+ enum class Placement { Pre, Post, Sandwich };
		96	+ Placement placement_;
59	std::unique_ptr<Norm> norm1_;	97	std::unique_ptr<Norm> norm1_;
60	std::unique_ptr<AttentionBase> attn_;	98	std::unique_ptr<AttentionBase> attn_;
61	std::unique_ptr<Norm> norm2_;	99	std::unique_ptr<Norm> norm2_;
62	std::unique_ptr<MLPBase> mlp_;	100	std::unique_ptr<MLPBase> mlp_;
		101	+ std::unique_ptr<Norm> norm1_post_, norm2_post_; // sandwich only
63	};	102	};
64		103
65	// Decoder-only transformer, entirely shaped by ModelConfig.	104	// Decoder-only transformer, entirely shaped by ModelConfig.
				`@@ -79,7 +118,7 @@ public:`
79	absorb("pos_emb", *pos_emb_);	118	absorb("pos_emb", *pos_emb_);
80	}	119	}
81	for (int64_t i = 0; i < cfg.n_layers; ++i) {	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	absorb("blocks." + std::to_string(i), *blocks_.back());	122	absorb("blocks." + std::to_string(i), *blocks_.back());
84	}	123	}
85	final_norm_ = std::make_unique<Norm>(cfg);	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	void softcap(const Tensor& x, float cap, Tensor& out) {	201	void softcap(const Tensor& x, float cap, Tensor& out) {
185	const float* px = x.data<float>();	202	const float* px = x.data<float>();
186	float* po = out.data<float>();	203	float* po = out.data<float>();
				`@@ -487,7 +504,7 @@ void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight)`
487		504
488	namespace {	505	namespace {
489	void rope_impl(const float* in, float* out, int64_t B, int64_t T, int64_t H, int64_t hd,	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	const int64_t C = H * hd;	508	const int64_t C = H * hd;
492	parallel_rows(B * T, [&](int64_t bt) {	509	parallel_rows(B * T, [&](int64_t bt) {
493	const int64_t t = bt % T;	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	float* dst = out + bt * C;	513	float* dst = out + bt * C;
497	for (int64_t h = 0; h < H; ++h) {	514	for (int64_t h = 0; h < H; ++h) {
498	for (int64_t k = 0; k < hd / 2; ++k) {	515	for (int64_t k = 0; k < hd / 2; ++k) {
499	− const float freq = std::pow(theta, -2.0f * float(k) / float(hd));	516	+ const float angle = pos * freqs[k];
500	− const float angle = pos * freq;
501	const float c = std::cos(angle);	517	const float c = std::cos(angle);
502	const float s = inverse ? -std::sin(angle) : std::sin(angle);	518	const float s = inverse ? -std::sin(angle) : std::sin(angle);
503	const int64_t i0 = h * hd + 2 * k;	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	} // namespace	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	check(x.ndim() == 3, "rope: expected [B,T,C]");	538	check(x.ndim() == 3, "rope: expected [B,T,C]");
522	const int64_t hd = x.size(2) / n_heads;	539	const int64_t hd = x.size(2) / n_heads;
523	check(hd % 2 == 0, "rope: head_dim must be even");	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	rope_impl(x.data<float>(), out.data<float>(), x.size(0), x.size(1), n_heads, hd,	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,	546	+void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,
529	− Tensor& dx) {	547	+ int64_t pos_offset, Tensor& dx) {
530	const int64_t hd = dout.size(2) / n_heads;	548	const int64_t hd = dout.size(2) / n_heads;
531	rope_impl(dout.data<float>(), dx.data<float>(), dout.size(0), dout.size(1), n_heads,	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	// ---- attention -------------------------------------------------------------------	553	// ---- attention -------------------------------------------------------------------
536		554
537	void attention(const Tensor& q, const Tensor& k, const Tensor& v,	555	void attention(const Tensor& q, const Tensor& k, const Tensor& v,
538	int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,	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	check(q.ndim() == 3 && k.ndim() == 3 && v.ndim() == 3, "attention: expected [B,T,C]");	558	check(q.ndim() == 3 && k.ndim() == 3 && v.ndim() == 3, "attention: expected [B,T,C]");
541	const int64_t B = q.size(0), T = q.size(1);	559	const int64_t B = q.size(0), T = q.size(1);
542	const int64_t hd = q.size(2) / n_heads;	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	for (int64_t i = 0; i < T; ++i) {	578	for (int64_t i = 0; i < T; ++i) {
561	const float* qi = pq + (b * T + i) * Cq + h * hd;	579	const float* qi = pq + (b * T + i) * Cq + h * hd;
562	const int64_t jmax = causal ? i : T - 1;	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	// scores (masked positions never written; treated as prob 0)	583	// scores (masked positions never written; treated as prob 0)
564	float m = -INFINITY;	584	float m = -INFINITY;
565	− for (int64_t j = 0; j <= jmax; ++j) {	585	+ for (int64_t j = jmin; j <= jmax; ++j) {
566	const float* kj = pk + (b * T + j) * Ckv + hkv * hd;	586	const float* kj = pk + (b * T + j) * Ckv + hkv * hd;
567	float s = 0.0f;	587	float s = 0.0f;
568	for (int64_t d = 0; d < hd; ++d) s += qi[d] * kj[d];	588	for (int64_t d = 0; d < hd; ++d) s += qi[d] * kj[d];
569	s *= scale;	589	s *= scale;
		590	+ if (attn_softcap > 0.0f)
		591	+ s = attn_softcap * std::tanh(s / attn_softcap);
570	P[i * T + j] = s;	592	P[i * T + j] = s;
571	m = std::max(m, s);	593	m = std::max(m, s);
572	}	594	}
573	float sum = 0.0f;	595	float sum = 0.0f;
574	− for (int64_t j = 0; j <= jmax; ++j) {	596	+ for (int64_t j = jmin; j <= jmax; ++j) {
575	const float e = std::exp(P[i * T + j] - m);	597	const float e = std::exp(P[i * T + j] - m);
576	P[i * T + j] = e;	598	P[i * T + j] = e;
577	sum += e;	599	sum += e;
578	}	600	}
579	const float inv = 1.0f / sum;	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	for (int64_t j = jmax + 1; j < T; ++j) P[i * T + j] = 0.0f;	604	for (int64_t j = jmax + 1; j < T; ++j) P[i * T + j] = 0.0f;
582		605
583	float* oi = po + (b * T + i) * Cq + h * hd;	606	float* oi = po + (b * T + i) * Cq + h * hd;
584	for (int64_t d = 0; d < hd; ++d) oi[d] = 0.0f;	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	const float p = P[i * T + j];	609	const float p = P[i * T + j];
587	const float* vj = pv + (b * T + j) * Ckv + hkv * hd;	610	const float* vj = pv + (b * T + j) * Ckv + hkv * hd;
588	for (int64_t d = 0; d < hd; ++d) oi[d] += p * vj[d];	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	void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,	617	void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
595	const Tensor& probs, const Tensor& out, const Tensor& dout,	618	const Tensor& probs, const Tensor& out, const Tensor& dout,
596	int64_t n_heads, int64_t n_kv_heads, float scale,	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	const int64_t B = q.size(0), T = q.size(1);	621	const int64_t B = q.size(0), T = q.size(1);
599	const int64_t hd = q.size(2) / n_heads;	622	const int64_t hd = q.size(2) / n_heads;
600	const int64_t rep = n_heads / n_kv_heads;	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	float dp = 0.0f;	661	float dp = 0.0f;
639	for (int64_t d = 0; d < hd; ++d) dp += doi[d] * vj[d];	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	for (int64_t d = 0; d < hd; ++d) {	673	for (int64_t d = 0; d < hd; ++d) {
643	dvj[d] += p * doi[d];	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	void silu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);	36	void silu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);
37	void gelu(const Tensor& x, Tensor& out); // tanh approximation (GPT-2)	37	void gelu(const Tensor& x, Tensor& out); // tanh approximation (GPT-2)
38	void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);	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	void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap)	41	void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap)
40	// dx += dout * (1 - (y/cap)^2) — takes the forward OUTPUT y	42	// dx += dout * (1 - (y/cap)^2) — takes the forward OUTPUT y
41	void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx);	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	// ---- RoPE (interleaved-pairs / GPT-J convention; see RESEARCH.md §7) -------	64	// ---- RoPE (interleaved-pairs / GPT-J convention; see RESEARCH.md §7) -------
63	// x: [B, T, H*hd]; rotates pairs (2k, 2k+1) inside each head. pos_offset	65	// x: [B, T, H*hd]; rotates pairs (2k, 2k+1) inside each head. pos_offset
64	// shifts absolute positions (KV-cache generation).	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	// Backward of a rotation is the inverse rotation applied to dout.	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,	71	+void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,
68	− Tensor& dx);	72	+ int64_t pos_offset, Tensor& dx);
69		73
70	// ---- attention (composed reference: scores → mask → softmax → PV) ----------	74	// ---- attention (composed reference: scores → mask → softmax → PV) ----------
71	// q: [B, T, Hhd], k/v: [B, T, Hkvhd], out: [B, T, H*hd]. GQA via	75	// q: [B, T, Hhd], k/v: [B, T, Hkvhd], out: [B, T, H*hd]. GQA via
72	// kv_head = h / (H / Hkv). If probs_out is non-null it receives the softmax	76	// kv_head = h / (H / Hkv). If probs_out is non-null it receives the softmax
73	// probabilities [B, H, T, T] (needed for backward).	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	void attention(const Tensor& q, const Tensor& k, const Tensor& v,	80	void attention(const Tensor& q, const Tensor& k, const Tensor& v,
75	int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,	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	// Uses the FlashAttention-2 identity rowsum(dP ∘ P) == dO_i · O_i, so `out`	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	void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,	87	void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
80	const Tensor& probs, const Tensor& out, const Tensor& dout,	88	const Tensor& probs, const Tensor& out, const Tensor& dout,
81	int64_t n_heads, int64_t n_kv_heads, float scale,	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	// ---- QAT / MoE --------------------------------------------------------------	92	// ---- QAT / MoE --------------------------------------------------------------
85	// Per-row fake quantization of a 2-D weight; mode 0 = int8 (absmax/127),	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	encode_flat("gelu_f32", {&x, &out}, nullptr, 0, x.numel());	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	void softcap(const Tensor& x, float cap, Tensor& out) {	229	void softcap(const Tensor& x, float cap, Tensor& out) {
222	encode_flat("softcap_f32", {&x, &out}, &cap, sizeof(cap), x.numel());	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	namespace {	445	namespace {
438	struct RopeParams {	446	struct RopeParams {
439	uint32_t T, H, HD;	447	uint32_t T, H, HD;
440	− float theta;
441	uint32_t pos_offset;	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	int64_t pos_offset, bool inverse) {	452	int64_t pos_offset, bool inverse) {
446	const int64_t B = in.size(0), T = in.size(1), C = in.size(2);	453	const int64_t B = in.size(0), T = in.size(1), C = in.size(2);
447	const int64_t hd = C / n_heads;	454	const int64_t hd = C / n_heads;
		455	+ check(freqs.numel() == hd / 2, "rope: freqs must have head_dim/2 entries");
448	Device& dev = Device::get();	456	Device& dev = Device::get();
449	MTL::FunctionConstantValues* constants = MTL::FunctionConstantValues::alloc()->init();	457	MTL::FunctionConstantValues* constants = MTL::FunctionConstantValues::alloc()->init();
450	// slots 0/1 belong to matmul TA/TB; RoPE uses slot 2	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	enc->setComputePipelineState(pso);	465	enc->setComputePipelineState(pso);
458	enc->setBuffer(in.buffer(), in.buffer_offset(), 0);	466	enc->setBuffer(in.buffer(), in.buffer_offset(), 0);
459	enc->setBuffer(out.buffer(), out.buffer_offset(), 1);	467	enc->setBuffer(out.buffer(), out.buffer_offset(), 1);
460	− RopeParams p{uint32_t(T), uint32_t(n_heads), uint32_t(hd), theta,	468	+ RopeParams p{uint32_t(T), uint32_t(n_heads), uint32_t(hd), uint32_t(pos_offset)};
461	− uint32_t(pos_offset)};
462	enc->setBytes(&p, sizeof(p), 2);	469	enc->setBytes(&p, sizeof(p), 2);
		470	+ enc->setBuffer(freqs.buffer(), freqs.buffer_offset(), 3);
463	const int64_t pairs = B * T * n_heads * (hd / 2);	471	const int64_t pairs = B * T * n_heads * (hd / 2);
464	enc->dispatchThreads(MTL::Size(NS::UInteger(pairs), 1, 1), MTL::Size(256, 1, 1));	472	enc->dispatchThreads(MTL::Size(NS::UInteger(pairs), 1, 1), MTL::Size(256, 1, 1));
465	}	473	}
466	} // namespace	474	} // namespace
467		475
468	−void rope(const Tensor& x, int64_t n_heads, float theta, int64_t pos_offset, Tensor& out) {	476	+void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset,
469	− rope_encode(x, out, n_heads, theta, pos_offset, /inverse=/false);	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,	481	+void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,
473	− Tensor& dx) {	482	+ int64_t pos_offset, Tensor& dx) {
474	− rope_encode(dout, dx, n_heads, theta, pos_offset, /inverse=/true);	483	+ rope_encode(dout, dx, n_heads, freqs, pos_offset, /inverse=/true);
475	}	484	}
476		485
477	// ---- embedding -------------------------------------------------------------------	486	// ---- embedding -------------------------------------------------------------------
				`@@ -517,12 +526,14 @@ struct AttnParams {`
517	uint32_t B, T, H, HKV, HD;	526	uint32_t B, T, H, HKV, HD;
518	float scale;	527	float scale;
519	uint32_t causal;	528	uint32_t causal;
		529	+ uint32_t window;
		530	+ float softcap;
520	};	531	};
521	} // namespace	532	} // namespace
522		533
523	void attention(const Tensor& q, const Tensor& k, const Tensor& v,	534	void attention(const Tensor& q, const Tensor& k, const Tensor& v,
524	int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,	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	check(probs_out != nullptr, "attention: probs_out required (unfused path)");	537	check(probs_out != nullptr, "attention: probs_out required (unfused path)");
527	const int64_t B = q.size(0), T = q.size(1);	538	const int64_t B = q.size(0), T = q.size(1);
528	const int64_t hd = q.size(2) / n_heads;	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	enc->setBuffer(out.buffer(), out.buffer_offset(), 3);	547	enc->setBuffer(out.buffer(), out.buffer_offset(), 3);
537	enc->setBuffer(probs_out->buffer(), probs_out->buffer_offset(), 4);	548	enc->setBuffer(probs_out->buffer(), probs_out->buffer_offset(), 4);
538	AttnParams p{uint32_t(B), uint32_t(T), uint32_t(n_heads), uint32_t(n_kv_heads),	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	enc->setBytes(&p, sizeof(p), 5);	552	enc->setBytes(&p, sizeof(p), 5);
541	enc->dispatchThreads(MTL::Size(NS::UInteger(B * n_heads * T), 1, 1),	553	enc->dispatchThreads(MTL::Size(NS::UInteger(B * n_heads * T), 1, 1),
542	MTL::Size(64, 1, 1));	554	MTL::Size(64, 1, 1));
				`@@ -545,12 +557,13 @@ void attention(const Tensor& q, const Tensor& k, const Tensor& v,`
545	void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,	557	void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
546	const Tensor& probs, const Tensor& out, const Tensor& dout,	558	const Tensor& probs, const Tensor& out, const Tensor& dout,
547	int64_t n_heads, int64_t n_kv_heads, float scale,	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	const int64_t B = q.size(0), T = q.size(1);	561	const int64_t B = q.size(0), T = q.size(1);
550	const int64_t hd = q.size(2) / n_heads;	562	const int64_t hd = q.size(2) / n_heads;
551	Device& dev = Device::get();	563	Device& dev = Device::get();
		564	+ // window is irrelevant here: masked positions carry prob 0 in `probs`.
552	AttnParams p{uint32_t(B), uint32_t(T), uint32_t(n_heads), uint32_t(n_kv_heads),	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	// D[b,h,i] = dO_i · O_i == rowsum(dP ∘ P): computed once per query row so	568	// D[b,h,i] = dO_i · O_i == rowsum(dP ∘ P): computed once per query row so
556	// neither backward kernel needs the O(T) inner recompute.	569	// neither backward kernel needs the O(T) inner recompute.
				`@@ -678,6 +691,8 @@ struct FlashParams {`
678	uint32_t B, T, H, HKV;	691	uint32_t B, T, H, HKV;
679	float scale;	692	float scale;
680	uint32_t causal;	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	// Must match the INSTANTIATE_FLASH list in flash_attention.metal.	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	void flash_attention(const Tensor& q, const Tensor& k, const Tensor& v,	713	void flash_attention(const Tensor& q, const Tensor& k, const Tensor& v,
699	int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,	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	const int64_t B = q.size(0), T = q.size(1);	716	const int64_t B = q.size(0), T = q.size(1);
702	const int64_t hd = q.size(2) / n_heads;	717	const int64_t hd = q.size(2) / n_heads;
703	check(flash_supported(hd), "flash_attention: unsupported head_dim");	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	const bool mma = kernel == FlashKernel::MMA;	722	const bool mma = kernel == FlashKernel::MMA;
		723	+ check(!(mma && window > 0), "flash_attention: MMA kernel has no window support");
707		724
708	Device& dev = Device::get();	725	Device& dev = Device::get();
709	MTL::ComputePipelineState* pso = dev.pipeline(	726	MTL::ComputePipelineState* pso = dev.pipeline(
				`@@ -716,7 +733,7 @@ void flash_attention(const Tensor& q, const Tensor& k, const Tensor& v,`
716	enc->setBuffer(out.buffer(), out.buffer_offset(), 3);	733	enc->setBuffer(out.buffer(), out.buffer_offset(), 3);
717	enc->setBuffer(lse.buffer(), lse.buffer_offset(), 4);	734	enc->setBuffer(lse.buffer(), lse.buffer_offset(), 4);
718	FlashParams p{uint32_t(B), uint32_t(T), uint32_t(n_heads), uint32_t(n_kv_heads),	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	enc->setBytes(&p, sizeof(p), 5);	737	enc->setBytes(&p, sizeof(p), 5);
721	// One threadgroup per (query block, head, batch). Block size and thread	738	// One threadgroup per (query block, head, batch). Block size and thread
722	// count must match the kernel's enums: the scalar kernel is one thread	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	const Tensor& out, const Tensor& lse, const Tensor& dout,	753	const Tensor& out, const Tensor& lse, const Tensor& dout,
737	int64_t n_heads, int64_t n_kv_heads, bool causal,	754	int64_t n_heads, int64_t n_kv_heads, bool causal,
738	float scale, Tensor& dq, Tensor& dk, Tensor& dv,	755	float scale, Tensor& dq, Tensor& dk, Tensor& dv,
739	− FlashKernel kernel) {	756	+ FlashKernel kernel, int64_t window) {
740	const int64_t B = q.size(0), T = q.size(1);	757	const int64_t B = q.size(0), T = q.size(1);
741	const int64_t hd = q.size(2) / n_heads;	758	const int64_t hd = q.size(2) / n_heads;
742	check(flash_supported(hd), "flash_attention_backward: unsupported head_dim");	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	const bool mma = kernel == FlashKernel::MMA;	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	Device& dev = Device::get();	766	Device& dev = Device::get();
747	FlashParams p{uint32_t(B), uint32_t(T), uint32_t(n_heads), uint32_t(n_kv_heads),	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	// D[b,h,i] = dO_i . O_i (shared with the unfused path).	770	// D[b,h,i] = dO_i . O_i (shared with the unfused path).
751	Tensor d_term = Tensor::empty({B, n_heads, T});	771	Tensor d_term = Tensor::empty({B, n_heads, T});
752	{	772	{
753	AttnParams ap{uint32_t(B), uint32_t(T), uint32_t(n_heads),	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	MTL::ComputePipelineState* pso = dev.pipeline("attention_bwd_d_f32");	776	MTL::ComputePipelineState* pso = dev.pipeline("attention_bwd_d_f32");
756	MTL::ComputeCommandEncoder* enc = Stream::get().encoder();	777	MTL::ComputeCommandEncoder* enc = Stream::get().encoder();
757	enc->setComputePipelineState(pso);	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	void add_bias(const Tensor& x, const Tensor& bias, Tensor& out);	88	void add_bias(const Tensor& x, const Tensor& bias, Tensor& out);
89	void silu(const Tensor& x, Tensor& out);	89	void silu(const Tensor& x, Tensor& out);
90	void gelu(const Tensor& x, Tensor& out);	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	void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap)	93	void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap)
92	void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx);	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	void layernorm_backward(const Tensor& x, const Tensor& w, float eps,	108	void layernorm_backward(const Tensor& x, const Tensor& w, float eps,
107	const Tensor& dout, Tensor& dx, Tensor& dw, Tensor& db);	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);	111	+// freqs: [head_dim/2] per-pair inverse frequencies, host-precomputed
110	−void rope_backward(const Tensor& dout, int64_t n_heads, float theta, int64_t pos_offset,	112	+// (fixed theta, llama3 rope-scaling, per-layer theta — all just tables).
111	− Tensor& dx);	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	void embedding(const Tensor& weight, const Tensor& ids, Tensor& out); // ids i32	118	void embedding(const Tensor& weight, const Tensor& ids, Tensor& out); // ids i32
114	void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight);	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	void attention(const Tensor& q, const Tensor& k, const Tensor& v,	123	void attention(const Tensor& q, const Tensor& k, const Tensor& v,
117	int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,	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	void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,	127	void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
120	const Tensor& probs, const Tensor& out, const Tensor& dout,	128	const Tensor& probs, const Tensor& out, const Tensor& dout,
121	int64_t n_heads, int64_t n_kv_heads, float scale,	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	// Fused (flash) attention: stores no T x T probabilities, only the per-row	132	// Fused (flash) attention: stores no T x T probabilities, only the per-row
125	// logsumexp `lse` [B, n_heads, T] that the backward re-expands. Supported for	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	// Scalar: one thread per query row. MMA: simdgroup_matrix 8x8 tiles. Same	136	// Scalar: one thread per query row. MMA: simdgroup_matrix 8x8 tiles. Same
129	// outputs; Auto picks MMA where supported.	137	// outputs; Auto picks MMA where supported.
130	enum class FlashKernel { Auto, Scalar, MMA };	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	void flash_attention(const Tensor& q, const Tensor& k, const Tensor& v,	141	void flash_attention(const Tensor& q, const Tensor& k, const Tensor& v,
132	int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,	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	void flash_attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,	145	void flash_attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
135	const Tensor& out, const Tensor& lse, const Tensor& dout,	146	const Tensor& out, const Tensor& lse, const Tensor& dout,
136	int64_t n_heads, int64_t n_kv_heads, bool causal,	147	int64_t n_heads, int64_t n_kv_heads, bool causal,
137	float scale, Tensor& dq, Tensor& dk, Tensor& dv,	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	// losses: [N] per-row buffer; loss_out: [1] mean over n_valid. dlogits	152	// losses: [N] per-row buffer; loss_out: [1] mean over n_valid. dlogits
141	// optional (accumulated). ids must be i32.	153	// optional (accumulated). ids must be i32.

modified src/ops/ops.cpp +87 −30

				`@@ -4,6 +4,8 @@`
4	#include "ops/cpu/cpu_ops.h"	4	#include "ops/cpu/cpu_ops.h"
5	#include "ops/metal/metal_ops.h"	5	#include "ops/metal/metal_ops.h"
6		6
		7	+#include <cmath>
		8	+
7	namespace forge::ops {	9	namespace forge::ops {
8		10
9	namespace {	11	namespace {
				`@@ -208,6 +210,23 @@ Var gelu(const Var& x) {`
208	return result;	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	Var softcap(const Var& x, float cap) {	230	Var softcap(const Var& x, float cap) {
212	Tensor out = Tensor::empty(x.value().shape());	231	Tensor out = Tensor::empty(x.value().shape());
213	if (gpu()) metal::softcap(x.value(), cap, out);	232	if (gpu()) metal::softcap(x.value(), cap, out);
				`@@ -370,25 +389,56 @@ Var embedding(const Var& weight, const Tensor& ids) {`
370	return result;	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	Tensor out = Tensor::empty(x.value().shape());	418	Tensor out = Tensor::empty(x.value().shape());
375	− if (gpu()) metal::rope(x.value(), n_heads, theta, pos_offset, out);	419	+ if (gpu()) metal::rope(x.value(), n_heads, freqs, pos_offset, out);
376	− else cpu::rope(x.value(), n_heads, theta, pos_offset, out);	420	+ else cpu::rope(x.value(), n_heads, freqs, pos_offset, out);
377		421
378	const bool needs = grad_needed({&x});	422	const bool needs = grad_needed({&x});
379	Var result(std::move(out), needs);	423	Var result(std::move(out), needs);
380	if (needs) {	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	if (!x.requires_grad()) return;	426	if (!x.requires_grad()) return;
383	− if (gpu()) metal::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());
384	− else cpu::rope_backward(result.grad(), n_heads, theta, pos_offset, x.grad());	428	+ else cpu::rope_backward(result.grad(), n_heads, freqs, pos_offset, x.grad());
385	});	429	});
386	}	430	}
387	return result;	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	Var attention(const Var& q, const Var& k, const Var& v,	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	const int64_t B = q.value().size(0), T = q.value().size(1);	442	const int64_t B = q.value().size(0), T = q.value().size(1);
393	const int64_t head_dim = q.value().size(2) / n_heads;	443	const int64_t head_dim = q.value().size(2) / n_heads;
394	Tensor out = Tensor::empty(q.value().shape());	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	// Fused path: keeps only the per-row logsumexp instead of a [B,H,T,T]	447	// Fused path: keeps only the per-row logsumexp instead of a [B,H,T,T]
398	// probability tensor, so memory is linear in T rather than quadratic.	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	Tensor lse = Tensor::empty({B, n_heads, T});	452	Tensor lse = Tensor::empty({B, n_heads, T});
401	metal::flash_attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads,	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	Var result(std::move(out), needs);	456	Var result(std::move(out), needs);
404	if (needs) {	457	if (needs) {
405	Tape::get().record(	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	Tensor dq_scratch, dk_scratch, dv_scratch;	460	Tensor dq_scratch, dk_scratch, dv_scratch;
408	Tensor& dq = q.requires_grad() ? q.grad()	461	Tensor& dq = q.requires_grad() ? q.grad()
409	: (dq_scratch = Tensor::zeros(q.value().shape()));	462	: (dq_scratch = Tensor::zeros(q.value().shape()));
				`@@ -414,7 +467,8 @@ Var attention(const Var& q, const Var& k, const Var& v,`
414	metal::flash_attention_backward(q.value(), k.value(), v.value(),	467	metal::flash_attention_backward(q.value(), k.value(), v.value(),
415	result.value(), lse, result.grad(),	468	result.value(), lse, result.grad(),
416	n_heads, n_kv_heads, causal, scale,	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	return result;	474	return result;
				`@@ -423,30 +477,33 @@ Var attention(const Var& q, const Var& k, const Var& v,`
423	Tensor probs = Tensor::empty({B, n_heads, T, T});	477	Tensor probs = Tensor::empty({B, n_heads, T, T});
424	if (gpu())	478	if (gpu())
425	metal::attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads, causal,	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	else	481	else
428	cpu::attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads, causal,	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	Var result(std::move(out), needs);	485	Var result(std::move(out), needs);
432	if (needs) {	486	if (needs) {
433	− Tape::get().record([q, k, v, probs, n_heads, n_kv_heads, scale, result]() {	487	+ Tape::get().record(
434	− Tensor dq_scratch, dk_scratch, dv_scratch;	488	+ [q, k, v, probs, n_heads, n_kv_heads, scale, attn_softcap, result]() {
435	− Tensor& dq = q.requires_grad() ? q.grad()	489	+ Tensor dq_scratch, dk_scratch, dv_scratch;
436	− : (dq_scratch = Tensor::zeros(q.value().shape()));	490	+ Tensor& dq = q.requires_grad() ? q.grad()
437	− Tensor& dk = k.requires_grad() ? k.grad()	491	+ : (dq_scratch = Tensor::zeros(q.value().shape()));
438	− : (dk_scratch = Tensor::zeros(k.value().shape()));	492	+ Tensor& dk = k.requires_grad() ? k.grad()
439	− Tensor& dv = v.requires_grad() ? v.grad()	493	+ : (dk_scratch = Tensor::zeros(k.value().shape()));
440	− : (dv_scratch = Tensor::zeros(v.value().shape()));	494	+ Tensor& dv = v.requires_grad() ? v.grad()
441	− if (gpu())	495	+ : (dv_scratch = Tensor::zeros(v.value().shape()));
442	− metal::attention_backward(q.value(), k.value(), v.value(), probs,	496	+ if (gpu())
443	− result.value(), result.grad(), n_heads,	497	+ metal::attention_backward(q.value(), k.value(), v.value(), probs,
444	− n_kv_heads, scale, dq, dk, dv);	498	+ result.value(), result.grad(), n_heads,
445	− else	499	+ n_kv_heads, scale, dq, dk, dv,
446	− cpu::attention_backward(q.value(), k.value(), v.value(), probs,	500	+ attn_softcap);
447	− result.value(), result.grad(), n_heads,	501	+ else
448	− n_kv_heads, scale, dq, dk, dv);	502	+ cpu::attention_backward(q.value(), k.value(), v.value(), probs,
449	− });	503	+ result.value(), result.grad(), n_heads,
		504	+ n_kv_heads, scale, dq, dk, dv,
		505	+ attn_softcap);
		506	+ });
450	}	507	}
451	return result;	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	Var silu(const Var& x);	32	Var silu(const Var& x);
33	Var gelu(const Var& x);	33	Var gelu(const Var& x);
		34	+Var relu2(const Var& x); // max(x,0)^2 — nanoGPT-speedrun activation
34		35
35	// Gemma-style soft capping: cap * tanh(x / cap). Bounds logits smoothly.	36	// Gemma-style soft capping: cap * tanh(x / cap). Bounds logits smoothly.
36	Var softcap(const Var& x, float cap);	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	// weight [V,C], ids [B,T] (u16/i32) → [B,T,C]	58	// weight [V,C], ids [B,T] (u16/i32) → [B,T,C]
58	Var embedding(const Var& weight, const Tensor& ids);	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	Var rope(const Var& x, int64_t n_heads, float theta, int64_t pos_offset = 0);	72	Var rope(const Var& x, int64_t n_heads, float theta, int64_t pos_offset = 0);
62		73
63	−// q [B,T,Hhd], k/v [B,T,Hkvhd] → [B,T,H*hd]; causal + GQA	74	+// q [B,T,Hhd], k/v [B,T,Hkvhd] → [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	Var attention(const Var& q, const Var& k, const Var& v,	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	// logits [N,V], targets [N] (i32/u16, ignore_index=-1) → scalar mean loss	82	// logits [N,V], targets [N] (i32/u16, ignore_index=-1) → scalar mean loss
68	Var cross_entropy(const Var& logits, const Tensor& targets);	83	Var cross_entropy(const Var& logits, const Tensor& targets);

modified tests/test_ops.cpp +46 −0

+         vcfg.tied_embeddings = true;
+         vcfg.qk_norm = true; vcfg.final_softcap = 30.0f; vcfg.scale_embeddings = true;
+         test_backend_parity_model("qk-norm + softcap + embed-scale", vcfg);
++
++        // Wave-1 knobs: ReLU² + QKV bias + decoupled head_dim + NoPE + rope scaling
++        forge::ModelConfig w1;
++        w1.n_layers = 2; w1.d_model = 16; w1.n_heads = 2; w1.n_kv_heads = 1;
++        w1.d_ff = 24; w1.vocab_size = 11; w1.context_length = 8;
++        w1.tied_embeddings = true;
++        w1.activation = "relu2"; w1.attention_bias = true;
++        w1.head_dim_override = 10; w1.nope_every = 2;
++        w1.rope_scale_factor = 8.0f; w1.rope_scale_orig_ctx = 4;
++        test_backend_parity_model("relu2 + qkv-bias + head_dim10 + nope + rope-scale", w1);
++
++        // Norm placements: OLMo2 post and Gemma sandwich
++        forge::ModelConfig np;
++        np.n_layers = 2; np.d_model = 16; np.n_heads = 2; np.n_kv_heads = 2;
++        np.d_ff = 24; np.vocab_size = 11; np.context_length = 8;
++        np.tied_embeddings = true;
++        np.norm_placement = "post";
++        test_backend_parity_model("post-norm (OLMo2)", np);
++        np.norm_placement = "sandwich";
++        test_backend_parity_model("sandwich norm (Gemma)", np);
++
++        // Wave-2: sliding window (Mistral) through the fused scalar path
++        // (head_dim 64), + local/global pattern, + Gemma3 dual rope theta
++        forge::ModelConfig sw;
++        sw.n_layers = 3; sw.d_model = 128; sw.n_heads = 2; sw.n_kv_heads = 1;
++        sw.d_ff = 96; sw.vocab_size = 11; sw.context_length = 8;
++        sw.tied_embeddings = true;
++        sw.sliding_window = 3; sw.sliding_global_every = 2;
++        sw.rope_theta_global = 1e6f;
++        test_backend_parity_model("sliding-window 3 + global-every-2 [hd64]", sw);
++
++        // Wave-2: attention softcap (Gemma2) — unfused path on both backends
++        forge::ModelConfig sc;
++        sc.n_layers = 2; sc.d_model = 16; sc.n_heads = 2; sc.n_kv_heads = 1;
++        sc.d_ff = 24; sc.vocab_size = 11; sc.context_length = 8;
++        sc.tied_embeddings = true;
++        sc.attn_softcap = 20.0f;
++        test_backend_parity_model("attn softcap 20 (Gemma2)", sc);
++
++        // Sliding window through the UNFUSED kernel too (head_dim 8)
++        forge::ModelConfig swu;
++        swu.n_layers = 2; swu.d_model = 16; swu.n_heads = 2; swu.n_kv_heads = 2;
++        swu.d_ff = 24; swu.vocab_size = 11; swu.context_length = 8;
++        swu.tied_embeddings = true;
++        swu.sliding_window = 3; swu.attn_softcap = 20.0f;
++        test_backend_parity_model("sliding-window 3 + softcap [unfused hd8]", swu);
+     }
+     test_metal_moe_quant_ops();
+     test_flash_attention();

spb/forge

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

Theme