1	# bend-ml-llama: Llama-architecture language models (SmolLM2 and the rest of the family) in Bend 2, with the
2	# structure in the types and laws checked by Bend's kernel.
3	#
4	# The sizes of a model live in a config (Cfg), and every weight matrix has its shape in its type, computed from
5	# the config: a model of other sizes is the same code with another config. The numbers (F32) are tested against
6	# PyTorch / transformers by the scripts of the llama.bend repository; the structure is proved here:
7	#
8	# - a head has 2·half numbers and is split into two halves of half numbers (RoPE rotates each number with the
9	#   one half a head later), and splitting a head and joining the halves gives it back (head_split);
10	# - the query heads are cut into kv groups of g heads, and group j reads key/value head j: there are exactly kv
11	#   groups (groups_len) and an empty cache holds one array of keys and one of values per key/value head
12	#   (kv_heads, arrays_len), and query head j·g + i reads key/value head (j·g + i) / g = j, transformers'
13	#   repeat_kv mapping (gqa_head);
14	# - the key/value cache at position t has the type KV<t>: t positions written in flat arrays, so attention at
15	#   position t scores exactly the positions 0..t and nothing after (causal, read_len);
16	# - the shapes of the weights in the types add up to the number of values in the weight files (layout laws).
17	#
18	# The order of every F32 operation follows transformers' modeling_llama.py (5.19), so the results match it to a
19	# few ulps and the argmax exactly (llama.bend's NOTES.md).
20	
21	import Base
22	import bend-ml-tensor@0.1.2.0/main.bend as T
23	import bend-ml-tensor-array@0.1.12.0/main.bend as TA
24	
25	# =====================================================================================================
26	# Arithmetic lemmas (used by the laws below)
27	# =====================================================================================================
28	
29	law add_zero:
30	  for a: Nat
31	  {Nat.add(a, 0n) == a : Nat}
32	
33	def add_zero(a):
34	  match a:
35	    case 0n:
36	      {==}
37	    case 1n+p:
38	      %add_zero(p) : {1n+Nat.add(p, 0n) == 1n+_ : Nat}
39	      {==}
40	
41	law add_succ:
42	  for  a: Nat
43	  for -b: Nat
44	  {1n+Nat.add(a, b) == Nat.add(a, 1n+b) : Nat}
45	
46	def add_succ(a, b):
47	  match a:
48	    case 0n:
49	      {==}
50	    case 1n+p:
51	      %add_succ(p, b) : {2n+Nat.add(p, b) == 1n+_ : Nat}
52	      {==}
53	
54	law add_one:
55	  for a: Nat
56	  {Nat.add(a, 1n) == 1n+a : Nat}
57	
58	def add_one(a):
59	  match a:
60	    case 0n:
61	      {==}
62	    case 1n+p:
63	      %add_one(p) : {1n+Nat.add(p, 1n) == 1n+_ : Nat}
64	      {==}
65	
66	law add_assoc:
67	  for  a: Nat
68	  for -b: Nat
69	  for -c: Nat
70	  {Nat.add(Nat.add(a, b), c) == Nat.add(a, Nat.add(b, c)) : Nat}
71	
72	def add_assoc(a, b, c):
73	  match a:
74	    case 0n:
75	      {==}
76	    case 1n+p:
77	      %add_assoc(p, b, c) : {1n+Nat.add(Nat.add(p, b), c) == 1n+_ : Nat}
78	      {==}
79	
80	# =====================================================================================================
81	# The configuration
82	# =====================================================================================================
83	
84	# kv: key/value heads; g: query heads per key/value head; half: half a head (a head has 2·half numbers, so its
85	# size is even by construction); f: the MLP's inner width; layers; vocab: the vocabulary (the embedding has
86	# vocab rows, and the LM head is the embedding); eps: RMSNorm's epsilon; theta: RoPE's base.
87	type Cfg is Data:
88	  Cfg{kv: Nat, g: Nat, half: Nat, f: Nat, layers: Nat, vocab: Nat, eps: F32, theta: F32}
89	
90	# a head: 2·half numbers
91	def head(+half: Nat) -> Nat:
92	  Nat.add(half, half)
93	
94	def Cfg.half(+c: Cfg) -> Nat:
95	  match c:
96	    case Cfg{kv, g, half, f, l, v, e, t}:
97	      half
98	
99	def Cfg.kv(+c: Cfg) -> Nat:
100	  match c:
101	    case Cfg{kv, g, half, f, l, v, e, t}:
102	      kv
103	
104	def Cfg.g(+c: Cfg) -> Nat:
105	  match c:
106	    case Cfg{kv, g, half, f, l, v, e, t}:
107	      g
108	
109	def Cfg.ffn(+c: Cfg) -> Nat:
110	  match c:
111	    case Cfg{kv, g, half, f, l, v, e, t}:
112	      f
113	
114	def Cfg.layers(+c: Cfg) -> Nat:
115	  match c:
116	    case Cfg{kv, g, half, f, l, v, e, t}:
117	      l
118	
119	def Cfg.vocab(+c: Cfg) -> Nat:
120	  match c:
121	    case Cfg{kv, g, half, f, l, v, e, t}:
122	      v
123	
124	def Cfg.eps(+c: Cfg) -> F32:
125	  match c:
126	    case Cfg{kv, g, half, f, l, v, e, t}:
127	      e
128	
129	def Cfg.theta(+c: Cfg) -> F32:
130	  match c:
131	    case Cfg{kv, g, half, f, l, v, e, t}:
132	      t
133	
134	# the size of a head
135	def Cfg.head(+c: Cfg) -> Nat:
136	  head(Cfg.half(c))
137	
138	# the query heads: kv groups of g
139	def Cfg.heads(+c: Cfg) -> Nat:
140	  Nat.mul(Cfg.kv(c), Cfg.g(c))
141	
142	# the hidden size: one head per query head
143	def Cfg.dim(+c: Cfg) -> Nat:
144	  Nat.mul(Cfg.heads(c), Cfg.head(c))
145	
146	# the width of the keys (and of the values): one head per key/value head
147	def Cfg.kvw(+c: Cfg) -> Nat:
148	  Nat.mul(Cfg.kv(c), Cfg.head(c))
149	
150	# SmolLM2-135M (HuggingFaceTB/SmolLM2-135M): 3 key/value heads of 3 query heads each, heads of 64
151	def smollm2_135m() -> Cfg:
152	  Cfg{3n, 3n, 32n, 1536n, 30n, 49152n, 0.00001, 100000.0}
153	
154	# SmolLM2-360M: 5 key/value heads of 3 query heads each, heads of 64
155	def smollm2_360m() -> Cfg:
156	  Cfg{5n, 3n, 32n, 2560n, 32n, 49152n, 0.00001, 100000.0}
157	
158	# SmolLM2-1.7B: 32 key/value heads of one query head each (plain multi-head attention), heads of 64
159	def smollm2_1_7b() -> Cfg:
160	  Cfg{32n, 1n, 32n, 8192n, 24n, 49152n, 0.00001, 130000.0}
161	
162	# LAW (the configs): SmolLM2-135M has 9 query heads of 64, a hidden size of 576 and keys of 192; SmolLM2-360M
163	# has 15 query heads of 64, a hidden size of 960 and keys of 320; SmolLM2-1.7B has 32 query heads of 64, a hidden
164	# size of 2048 and keys of 2048 (as many key/value heads as query heads), as their config.json files say.
165	law smollm2_135m_sizes:
166	  {Cfg.dim(smollm2_135m()) == 576n : Nat} & {Cfg.heads(smollm2_135m()) == 9n : Nat} & {Cfg.head(smollm2_135m()) == 64n : Nat} & {Cfg.kvw(smollm2_135m()) == 192n : Nat}
167	
168	def smollm2_135m_sizes():
169	  ({==}, {==}, {==}, {==})
170	
171	law smollm2_360m_sizes:
172	  {Cfg.dim(smollm2_360m()) == 960n : Nat} & {Cfg.heads(smollm2_360m()) == 15n : Nat} & {Cfg.head(smollm2_360m()) == 64n : Nat} & {Cfg.kvw(smollm2_360m()) == 320n : Nat}
173	
174	def smollm2_360m_sizes():
175	  ({==}, {==}, {==}, {==})
176	
177	law smollm2_1_7b_sizes:
178	  {Cfg.dim(smollm2_1_7b()) == 2048n : Nat} & {Cfg.heads(smollm2_1_7b()) == 32n : Nat} & {Cfg.head(smollm2_1_7b()) == 64n : Nat} & {Cfg.kvw(smollm2_1_7b()) == 2048n : Nat}
179	
180	def smollm2_1_7b_sizes():
181	  ({==}, {==}, {==}, {==})
182	
183	# ---- the layout: the shapes in the types add up to the number of values in the checkpoint ----
184	
185	def n32(+n: Nat) -> U32:
186	  U32.from_nat(n)
187	
188	# the numbers of one layer, from the shapes of LW<c>: two norms, q and o (dim × dim), k and v (kvw × dim),
189	# gate and up (ffn × dim) and down (dim × ffn)
190	def Cfg.layer_size(+c: Cfg) -> U32:
191	  ((2 : U32) * n32(Cfg.dim(c)) + (2 : U32) * n32(Cfg.dim(c)) * n32(Cfg.dim(c)) + (2 : U32) * n32(Cfg.kvw(c)) * n32(Cfg.dim(c)) + (3 : U32) * n32(Cfg.ffn(c)) * n32(Cfg.dim(c)) : U32)
192	
193	# the numbers of the model, from the shapes of Model<c>: the embedding (vocab × dim), the layers, the final norm
194	def Cfg.size(+c: Cfg) -> U32:
195	  (n32(Cfg.vocab(c)) * n32(Cfg.dim(c)) + n32(Cfg.layers(c)) * Cfg.layer_size(c) + n32(Cfg.dim(c)) : U32)
196	
197	# LAW (layout): the shapes the types give to SmolLM2-135M's weights add up to 134,515,008 numbers,
198	# SmolLM2-360M's to 361,821,120 and SmolLM2-1.7B's to 1,711,376,384: the number of values in each model.safetensors (reference/prep.py counts them).
199	# The loader checks every file's size against the shape in its type, so together the files cover the checkpoint
200	# exactly. In U32, whose sums wrap at 2^32; every partial sum here is below it, so this is the exact count.
201	law layout_135m:
202	  {Cfg.size(smollm2_135m()) == (134515008 : U32) : U32}
203	
204	def layout_135m():
205	  {==}
206	
207	law layout_360m:
208	  {Cfg.size(smollm2_360m()) == (361821120 : U32) : U32}
209	
210	def layout_360m():
211	  {==}
212	
213	law layout_1_7b:
214	  {Cfg.size(smollm2_1_7b()) == (1711376384 : U32) : U32}
215	
216	def layout_1_7b():
217	  {==}
218	
219	# =====================================================================================================
220	# Heads
221	# =====================================================================================================
222	
223	# a head as its two halves: RoPE rotates number i of the first half with number i of the second
224	type Head is Data:
225	  Head{lo: List<&2, F32>, hi: List<&2, F32>}
226	
227	def split(+half: Nat, +xs: List<&2, F32>) -> Head:
228	  Head{List.take(&2, F32, xs, half), List.drop(&2, F32, xs, half)}
229	
230	def join(h: Head) -> List<&2, F32>:
231	  match h:
232	    case Head{lo, hi}:
233	      List.append(&2, F32, lo, hi)
234	
235	# n heads of 2·half numbers from xs, in order
236	def heads(n: Nat, +half: Nat, +xs: List<&2, F32>) -> List<&2, Head>:
237	  match n:
238	    case 0n:
239	      Nil{}
240	    case 1n+p:
241	      split(half, List.take(&2, F32, xs, head(half))) <> heads(p, half, List.drop(&2, F32, xs, head(half)))
242	
243	# n plain heads (the values, which RoPE does not touch)
244	def chunks(n: Nat, +w: Nat, +xs: List<&2, F32>) -> List<&2, List<&2, F32>>:
245	  match n:
246	    case 0n:
247	      Nil{}
248	    case 1n+p:
249	      List.take(&2, F32, xs, w) <> chunks(p, w, List.drop(&2, F32, xs, w))
250	
251	def joins(hs: List<&2, Head>) -> List<&2, List<&2, F32>>:
252	  match hs:
253	    case Nil{}:
254	      Nil{}
255	    case Con{h, t}:
256	      join(h) <> joins(t)
257	
258	def flat(hs: List<&2, Head>) -> List<&2, F32>:
259	  List.concat(&2, F32, joins(hs))
260	
261	# LAW (head_split): splitting a list at half and joining the two parts gives the list back, so cutting a head
262	# into the halves that RoPE pairs up loses and reorders nothing.
263	law head_split:
264	  for +xs: List<&2, F32>
265	  for +half: Nat
266	  {join(split(half, xs)) == xs : List<&2, F32>}
267	
268	# PROOF: by induction on the list and half together: an empty list stays empty; at half 0 the first part is
269	# empty and the second is the whole list; otherwise the first number is kept in front and the rest is the
270	# induction hypothesis.
271	def head_split(xs, half):
272	  match xs half:
273	    case Nil{} _:
274	      {==}
275	    case Con{h, t} 0n:
276	      {==}
277	    case Con{h, t} 1n+p:
278	      %head_split(t, p) : {h <> List.append(&2, F32, List.take(&2, F32, t, p), List.drop(&2, F32, t, p)) == h <> _ : List<&2, F32>}
279	      {==}
280	
281	# =====================================================================================================
282	# The numeric pieces (tested against transformers, never proved: F32 is not a real number)
283	# =====================================================================================================
284	
285	# the squares of xs
286	def squares(xs: List<&2, F32>) -> List<&2, F32>:
287	  match xs:
288	    case Nil{}:
289	      Nil{}
290	    case Con{+h, t}:
291	      (h * h : F32) <> squares(t)
292	
293	# adjacent pairs added: [a, b, c, d, e] -> [a + b, c + d, e]
294	def pair_sums(xs: List<&2, F32>) -> List<&2, F32>:
295	  match xs:
296	    case Con{+a, Con{+b, t}}:
297	      (a + b : F32) <> pair_sums(t)
298	    case _:
299	      xs
300	
301	# the sum as a balanced tree of additions (pairwise summation): its rounding error grows with log2(n), not n, as
302	# in PyTorch's cascade sum. A running sum loses the small squares next to a large one (NOTES, 2026-10-10).
303	def tree_sum(fuel: Nat, xs: List<&2, F32>) -> F32:
304	  match fuel xs:
305	    case _ Nil{}:
306	      0.0
307	    case _ Con{+a, Nil{}}:
308	      a
309	    case 0n Con{+a, t}:
310	      T.sum.go(t, a)
311	    case 1n+p Con{+a, t}:
312	      tree_sum(p, pair_sums(Con{a, t}))
313	
314	# RMSNorm: g * (x * (1 / sqrt(mean(x²) + eps))), in this order (LlamaRMSNorm.forward)
315	def rmsnorm(+c: Cfg, +x: List<&2, F32>, +g: List<&2, F32>) -> List<&2, F32>:
316	  T.zip_with.mul(g, T.scale.l((1.0 / F32.sqrt(((tree_sum(Cfg.dim(c), squares(x)) / F32.from_nat(Cfg.dim(c)) : F32) + Cfg.eps(c) : F32)) : F32), x))
317	
318	# inv_freq[i] = 1 / (theta ^ (2i / head)), i = 0..half-1 (LlamaRotaryEmbedding, default rope)
319	def inv_freq.go(n: Nat, +i: Nat, +theta: F32, +hd: F32) -> List<&2, F32>:
320	  match n:
321	    case 0n:
322	      Nil{}
323	    case 1n+p:
324	      (1.0 / F32.pow(theta, (F32.from_nat(Nat.mul(2n, i)) / hd : F32)) : F32) <> inv_freq.go(p, 1n+i, theta, hd)
325	
326	def inv_freq(+c: Cfg) -> List<&2, F32>:
327	  inv_freq.go(Cfg.half(c), 0n, Cfg.theta(c), F32.from_nat(Cfg.head(c)))
328	
329	def angles(+t: F32, fs: List<&2, F32>) -> List<&2, F32>:
330	  match fs:
331	    case Nil{}:
332	      Nil{}
333	    case Con{+f, r}:
334	      (t * f : F32) <> angles(t, r)
335	
336	def cos.l(xs: List<&2, F32>) -> List<&2, F32>:
337	  match xs:
338	    case Nil{}:
339	      Nil{}
340	    case Con{+x, r}:
341	      F32.cos(x) <> cos.l(r)
342	
343	def sin.l(xs: List<&2, F32>) -> List<&2, F32>:
344	  match xs:
345	    case Nil{}:
346	      Nil{}
347	    case Con{+x, r}:
348	      F32.sin(x) <> sin.l(r)
349	
350	# the cosines and sines of position t, one per frequency (shared by every head of every layer)
351	type Rot is Data:
352	  Rot{cs: List<&2, F32>, sn: List<&2, F32>}
353	
354	def rot_at(+c: Cfg, +t: Nat) -> Rot:
355	  Rot{cos.l(angles(F32.from_nat(t), inv_freq(c))), sin.l(angles(F32.from_nat(t), inv_freq(c)))}
356	
357	# first half: lo*cos + (-hi)*sin; second half: hi*cos + lo*sin (q*cos + rotate_half(q)*sin)
358	def rot_lo(h1: List<&2, F32>, h2: List<&2, F32>, cs: List<&2, F32>, sn: List<&2, F32>) -> List<&2, F32>:
359	  match h1 h2 cs sn:
360	    case Con{+a, ta} Con{+b, tb} Con{+c, tc} Con{+s, ts}:
361	      ((a * c : F32) + (F32.neg(b) * s : F32) : F32) <> rot_lo(ta, tb, tc, ts)
362	    case _ _ _ _:
363	      Nil{}
364	
365	def rot_hi(h1: List<&2, F32>, h2: List<&2, F32>, cs: List<&2, F32>, sn: List<&2, F32>) -> List<&2, F32>:
366	  match h1 h2 cs sn:
367	    case Con{+a, ta} Con{+b, tb} Con{+c, tc} Con{+s, ts}:
368	      ((b * c : F32) + (a * s : F32) : F32) <> rot_hi(ta, tb, tc, ts)
369	    case _ _ _ _:
370	      Nil{}
371	
372	def rope(rt: Rot, h: Head) -> Head:
373	  match rt h:
374	    case Rot{+cs, +sn} Head{+lo, +hi}:
375	      Head{rot_lo(lo, hi, cs, sn), rot_hi(lo, hi, cs, sn)}
376	
377	def ropes(+rt: Rot, hs: List<&2, Head>) -> List<&2, Head>:
378	  match hs:
379	    case Nil{}:
380	      Nil{}
381	    case Con{h, t}:
382	      rope(rt, h) <> ropes(rt, t)
383	
384	# silu(x) = x / (1 + exp(-x)), the form ATen's CPU kernel uses
385	def silu(+x: F32) -> F32:
386	  (x / (1.0 + F32.exp(F32.neg(x)) : F32) : F32)
387	
388	# silu(gate) * up
389	def swiglu(g: List<&2, F32>, u: List<&2, F32>) -> List<&2, F32>:
390	  match g u:
391	    case Con{+a, ga} Con{+b, ub}:
392	      (silu(a) * b : F32) <> swiglu(ga, ub)
393	    case _ _:
394	      Nil{}
395	
396	# =====================================================================================================
397	# The key/value cache in flat arrays, indexed by the number of positions it holds
398	# =====================================================================================================
399	
400	# One pair of arrays per key/value head: the keys (after RoPE) and the values of positions 0..t-1, position p at
401	# [p·head, p·head + head). The type indexes are the capacity cap (the positions the arrays have room for, set by
402	# KV.empty) and t, the number of positions written. KV.push asks for a proof that t < cap, so a cache is never
403	# written past its room (kv_room says why that is enough).
404	type KV<-cap: Nat, -t: Nat> is Type:
405	  KV{ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>}
406	
407	# n arrays of 2^d slots
408	def arrays(n: Nat, +d: Nat) -> List<&1, Array<F32>>:
409	  match n:
410	    case 0n:
411	      Nil{}
412	    case 1n+p:
413	      Array.new(F32, d, 0.0) <> arrays(p, d)
414	
415	# a cache with no position yet, with room for cap positions per key/value head
416	def KV.empty(+c: Cfg, +cap: Nat) -> KV<cap, 0n>:
417	  KV{arrays(Cfg.kv(c), TA.cap_depth(Nat.mul(cap, Cfg.head(c)))), arrays(Cfg.kv(c), TA.cap_depth(Nat.mul(cap, Cfg.head(c))))}
418	
419	# writes row j of rs into array j, at off
420	def write_rows(arrs: List<&1, Array<F32>>, rs: List<&2, List<&2, F32>>, +off: U32) -> List<&1, Array<F32>>:
421	  match arrs rs:
422	    case Con{a, at} Con{r, rt}:
423	      TA.fill_list(r, off, a) <> write_rows(at, rt, off)
424	    case _ _:
425	      arrs
426	
427	# position t: the keys and values of every key/value head written at t·head; room is the proof that t < cap
428	def KV.push(+c: Cfg, -cap: Nat, +t: Nat, kv: KV<cap, t>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +ks: List<&2, List<&2, F32>>, +vs: List<&2, List<&2, F32>>) -> KV<cap, 1n+t>:
429	  match kv:
430	    case KV{ka, va}:
431	      KV{write_rows(ka, ks, U32.from_nat(Nat.mul(t, Cfg.head(c)))), write_rows(va, vs, U32.from_nat(Nat.mul(t, Cfg.head(c))))}
432	
433	# the number of key/value heads a cache holds
434	def kv_count(-cap: Nat, -t: Nat, kv: KV<cap, t>) -> Nat:
435	  match kv:
436	    case KV{ka, va}:
437	      List.length(&1, Array<F32>, ka)
438	
439	# LAW (arrays_len): arrays(n, d) makes exactly n arrays.
440	law arrays_len:
441	  for +n: Nat
442	  for +d: Nat
443	  {List.length(&1, Array<F32>, arrays(n, d)) == n : Nat}
444	
445	# PROOF: by induction on n: none for 0; one more than the rest for 1 + p.
446	def arrays_len(n, d):
447	  match n:
448	    case 0n:
449	      {==}
450	    case 1n+p:
451	      %arrays_len(p, d) : {1n+List.length(&1, Array<F32>, arrays(p, d)) == 1n+_ : Nat}
452	      {==}
453	
454	# LAW (kv_heads): an empty cache holds one array of keys (and one of values) per key/value head, so key/value head j
455	# of the model reads column j and no other.
456	law kv_heads:
457	  for +c: Cfg
458	  for +cap: Nat
459	  {kv_count(cap, 0n, KV.empty(c, cap)) == Cfg.kv(c) : Nat}
460	
461	def kv_heads(c, cap):
462	  arrays_len(Cfg.kv(c), TA.cap_depth(Nat.mul(cap, Cfg.head(c))))
463	
464	# a <= b  and  b <= c  ->  a <= c
465	def leb_trans(a: Nat, b: Nat, c: Nat, e1: {True{} == TA.leb(a, b) : Bool}, e2: {True{} == TA.leb(b, c) : Bool}) -> {True{} == TA.leb(a, c) : Bool}:
466	  match a b c:
467	    case 0n _ _:
468	      {==}
469	    case 1n+p 0n _:
470	      %e1 : TA.BD(_, Unit, {True{} == TA.leb(1n+p, c) : Bool})
471	      Unit{}
472	    case 1n+p 1n+q 0n:
473	      %e2 : TA.BD(_, Unit, {True{} == TA.leb(1n+p, 0n) : Bool})
474	      Unit{}
475	    case 1n+p 1n+q 1n+r:
476	      leb_trans(p, q, r, e1, e2)
477	
478	# a <= b  ->  a·h <= b·h
479	def leb_mul(a: Nat, b: Nat, +h: Nat, e: {True{} == TA.leb(a, b) : Bool}) -> {True{} == TA.leb(Nat.mul(a, h), Nat.mul(b, h)) : Bool}:
480	  match a b:
481	    case 0n _:
482	      {==}
483	    case 1n+p 0n:
484	      %e : TA.BD(_, Unit, {True{} == TA.leb(Nat.mul(1n+p, h), 0n) : Bool})
485	      Unit{}
486	    case 1n++p 1n++q:
487	      TA.leb_add(h, h, Nat.mul(p, h), Nat.mul(q, h), TA.leb_refl(h), leb_mul(p, q, h, e))
488	
489	# LAW (kv_room): KV.empty(c, cap) gives each key/value head arrays of 2^cap_depth(cap·head) slots, and KV.push at
490	# position t writes head numbers from t·head. For a position below the capacity (t < cap) they all fall inside the
491	# slots: (1 + t)·head <= 2^cap_depth(cap·head), so a push never wraps around onto position 0. (What stays trusted is
492	# tensor-array's: Array.new allocates 2^d slots.) The type does not stop a caller from pushing at t >= cap; the
493	# callers size cap to the positions they run.
494	law kv_room:
495	  for +cap: Nat
496	  for +h: Nat
497	  for +t: Nat
498	  for e: {True{} == TA.leb(1n+t, cap) : Bool}
499	  {True{} == TA.leb(Nat.mul(1n+t, h), TA.pow2(TA.cap_depth(Nat.mul(cap, h)))) : Bool}
500	
501	# For callers that walk the positions with a counter of steps left (fuel), keeping t + fuel = cap: with 1 + p steps
502	# left, position t is below the capacity (step_room), and after the step t + 1 + p = cap still holds (step_inv).
503	def step_room(+t: Nat, +p: Nat, -cap: Nat, e: {Nat.add(t, 1n+p) == cap : Nat}) -> {True{} == TA.leb(1n+t, cap) : Bool}:
504	  %e : {True{} == TA.leb(1n+t, _) : Bool}
505	  %add_succ(t, p) : {True{} == TA.leb(1n+t, _) : Bool}
506	  TA.leb_up(t, t, p, TA.leb_refl(t))
507	
508	def step_inv(+t: Nat, -p: Nat, -cap: Nat, e: {Nat.add(t, 1n+p) == cap : Nat}) -> {Nat.add(1n+t, p) == cap : Nat}:
509	  %e : {1n+Nat.add(t, p) == _ : Nat}
510	  add_succ(t, p)
511	
512	# PROOF: (1 + t)·head <= cap·head by leb_mul, and cap·head <= 2^cap_depth(cap·head) by tensor-array's cap_ok; then
513	# leb_trans.
514	def kv_room(cap, h, t, e):
515	  leb_trans(Nat.mul(1n+t, h), Nat.mul(cap, h), TA.pow2(TA.cap_depth(Nat.mul(cap, h))), leb_mul(1n+t, cap, h, e), TA.cap_ok(Nat.mul(cap, h)))
516	
517	# =====================================================================================================
518	# Grouped-query attention
519	# =====================================================================================================
520	# the query heads in n groups of g, in order: group j holds query heads j·g .. j·g + g - 1
521	def groups(n: Nat, +g: Nat, +hs: List<&2, Head>) -> List<&2, List<&2, Head>>:
522	  match n:
523	    case 0n:
524	      Nil{}
525	    case 1n+p:
526	      List.take(&2, Head, hs, g) <> groups(p, g, List.drop(&2, Head, hs, g))
527	
528	# LAW (groups_len): the query heads make exactly as many groups as there are key/value heads, so attention
529	# (which walks the groups and the columns together) pairs every group with a key/value head.
530	law groups_len:
531	  for n: Nat
532	  for +g: Nat
533	  for +hs: List<&2, Head>
534	  {List.length(&2, List<&2, Head>, groups(n, g, hs)) == n : Nat}
535	
536	def groups_len(n, g, hs):
537	  match n:
538	    case 0n:
539	      {==}
540	    case 1n+p:
541	      %groups_len(p, g, List.drop(&2, Head, hs, g)) : {1n+List.length(&2, List<&2, Head>, groups(p, g, List.drop(&2, Head, hs, g))) == 1n+_ : Nat}
542	      {==}
543	
544	# one block of 1 + m numbers in Nat.divmod: the counter m runs down to 0, then the quotient d grows by one
545	law div_block:
546	  for +m: Nat
547	  for -n: Nat
548	  for -d: Nat
549	  for -r: Nat
550	  {Nat.divmod.go(Nat.add(1n+m, n), m, d, r) == Nat.divmod.go(n, Nat.add(m, r), 1n+d, 0n) : Nat & Nat}
551	
552	def div_block(m, n, d, r):
553	  match m:
554	    case 0n:
555	      {==}
556	    case 1n+p:
557	      %Equal.sym(Nat, 1n+Nat.add(p, r), Nat.add(p, 1n+r), add_succ(p, r)) : {Nat.divmod.go(Nat.add(1n+p, n), p, d, 1n+r) == Nat.divmod.go(n, _, 1n+d, 0n) : Nat & Nat}
558	      div_block(p, n, d, 1n+r)
559	
560	# j whole blocks of 1 + s numbers add j to the quotient
561	law div_blocks:
562	  for +j: Nat
563	  for +s: Nat
564	  for +n: Nat
565	  for +d: Nat
566	  {Nat.divmod.go(Nat.add(Nat.mul(j, 1n+s), n), s, d, 0n) == Nat.divmod.go(n, s, Nat.add(j, d), 0n) : Nat & Nat}
567	
568	def div_blocks(j, s, n, d):
569	  match j:
570	    case 0n:
571	      {==}
572	    case 1n+jp:
573	      %Equal.sym(Nat, Nat.add(Nat.add(1n+s, Nat.mul(jp, 1n+s)), n), Nat.add(1n+s, Nat.add(Nat.mul(jp, 1n+s), n)), add_assoc(1n+s, Nat.mul(jp, 1n+s), n)) : {Nat.divmod.go(_, s, d, 0n) == Nat.divmod.go(n, s, Nat.add(1n+jp, d), 0n) : Nat & Nat}
574	      %Equal.sym(Nat & Nat, Nat.divmod.go(Nat.add(1n+s, Nat.add(Nat.mul(jp, 1n+s), n)), s, d, 0n), Nat.divmod.go(Nat.add(Nat.mul(jp, 1n+s), n), Nat.add(s, 0n), 1n+d, 0n), div_block(s, Nat.add(Nat.mul(jp, 1n+s), n), d, 0n)) : {_ == Nat.divmod.go(n, s, Nat.add(1n+jp, d), 0n) : Nat & Nat}
575	      %Equal.sym(Nat, Nat.add(s, 0n), s, add_zero(s)) : {Nat.divmod.go(Nat.add(Nat.mul(jp, 1n+s), n), _, 1n+d, 0n) == Nat.divmod.go(n, s, Nat.add(1n+jp, d), 0n) : Nat & Nat}
576	      %Equal.sym(Nat, 1n+Nat.add(jp, d), Nat.add(jp, 1n+d), add_succ(jp, d)) : {Nat.divmod.go(Nat.add(Nat.mul(jp, 1n+s), n), s, 1n+d, 0n) == Nat.divmod.go(n, s, _, 0n) : Nat & Nat}
577	      div_blocks(jp, s, n, 1n+d)
578	
579	# fewer than 1 + s numbers after the last whole block leave the quotient alone
580	law div_rest:
581	  for +i: Nat
582	  for -k: Nat
583	  for -d: Nat
584	  for -r: Nat
585	  {Pair.fst(Nat, Nat, Nat.divmod.go(i, Nat.add(i, k), d, r)) == d : Nat}
586	
587	def div_rest(i, k, d, r):
588	  match i:
589	    case 0n:
590	      {==}
591	    case 1n+p:
592	      div_rest(p, k, d, 1n+r)
593	
594	# LAW (gqa_head): with g = 1 + i + k query heads per key/value head (g > 0 and i < g), query head j·g + i, the
595	# head i of group j, reads key/value head (j·g + i) / g = j. That division is how transformers maps query heads
596	# to key/value heads (repeat_kv), so grouping the heads computes the same attention.
597	law gqa_head:
598	  for +j: Nat
599	  for +i: Nat
600	  for +k: Nat
601	  {Nat.div(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}
602	
603	# PROOF: Nat.div counts how many whole blocks of g numbers fit. j·g is j whole blocks (div_blocks, which uses
604	# div_block for each block), and i < g numbers after them do not complete another one (div_rest).
605	def gqa_head(j, i, k):
606	  %Equal.sym(Nat & Nat, Nat.divmod.go(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), Nat.add(i, k), 0n, 0n), Nat.divmod.go(i, Nat.add(i, k), Nat.add(j, 0n), 0n), div_blocks(j, Nat.add(i, k), i, 0n)) : {Pair.fst(Nat, Nat, _) == j : Nat}
607	  %add_zero(j) : {Pair.fst(Nat, Nat, Nat.divmod.go(i, Nat.add(i, k), Nat.add(j, 0n), 0n)) == _ : Nat}
608	  div_rest(i, k, Nat.add(j, 0n), 0n)
609	
610	# =====================================================================================================
611	# Grouped-query attention over the arrays
612	# =====================================================================================================
613	
614	# the scores of one query head against the keys of its key/value head, and the keys handed back
615	type SR is Type:
616	  SR{k: Array<F32>, s: List<&2, F32>}
617	
618	def sr_len(sr: SR) -> Nat:
619	  match sr:
620	    case SR{k, s}:
621	      List.length(&2, F32, s)
622	
623	def sr_of(k: Array<F32>, al: TA.AL) -> SR:
624	  match al:
625	    case TA.AL{c, l}:
626	      SR{k, l}
627	
628	# LAW (read_len): reading m + 1 numbers of an array onto acc gives m + 1 more scores.
629	# PROOF: by induction on m. With 0 left it reads one number and reverses the list (rev_len counts it); with 1 + p
630	# left it reads one and recurses, and len_step moves the 1 out of the sum (tensor-array's leaf_len, for scores).
631	law read_len:
632	  for k: Array<F32>
633	  for +m: Nat
634	  for +idx: U32
635	  for +acc: List<&2, F32>
636	  for rc: Array<F32> & F32
637	  {sr_len(sr_of(k, TA.read_l(m, idx, acc, rc))) == Nat.add(1n+m, List.length(&2, F32, acc)) : Nat}
638	
639	def read_len(k, m, idx, acc, rc):
640	  match m rc:
641	    case 0n Tuple{d, +v}:
642	      TA.zero_out(List.length(&2, F32, List.reverse.go(&2, F32, v <> acc, Nil{})), 1n+List.length(&2, F32, acc), TA.rev_len(v <> acc, Nil{}))
643	    case 1n+p Tuple{d, +v}:
644	      TA.len_step(sr_len(sr_of(k, TA.read_l(p, (idx + 1 : U32), v <> acc, Array.get(F32, d, (idx + 1 : U32))))), p, List.length(&2, F32, acc), read_len(k, p, (idx + 1 : U32), v <> acc, Array.get(F32, d, (idx + 1 : U32))))
645	
646	# read the 1 + t scores the gemm wrote
647	def sr_read(+t: Nat, g: TA.G) -> SR:
648	  match g:
649	    case TA.G{qa, k, s}:
650	      sr_of(k, TA.read_l(t, 0, Nil{}, Array.get(F32, s, 0)))
651	
652	# q·k for the keys of positions 0..t, oldest first: K (t + 1 rows of `head` numbers) times q, the product
653	# Mat.matmul_nt computes for one row (bend-ml-tensor-array's gemm)
654	def scores(+hd: Nat, +t: Nat, k: Array<F32>, +q: List<&2, F32>) -> SR:
655	  sr_read(t, TA.gemm(1n, 1n+t, hd, U32.from_nat(hd), 1, 1, U32.from_nat(hd), TA.fill_list(q, 0, Array.new(F32, TA.cap_depth(hd), 0.0)), k, Array.new(F32, TA.cap_depth(1n+t), 0.0)))
656	
657	def sr_len_g(+t: Nat, g: TA.G) -> {sr_len(sr_read(t, g)) == 1n+t : Nat}:
658	  match g:
659	    case TA.G{qa, k, s}:
660	      TA.zero_out(sr_len(sr_of(k, TA.read_l(t, 0, Nil{}, Array.get(F32, s, 0)))), 1n+t, read_len(k, t, 0, Nil{}, Array.get(F32, s, 0)))
661	
662	# LAW (causal): attention at position t scores exactly t + 1 keys, the positions 0..t: one score per position
663	# written, none for a position not yet computed.
664	# PROOF: the gemm's result has one constructor, so matching it opens it without computing the product
665	# (sr_len_g); read_len counts the 1 + t numbers read back.
666	law causal:
667	  for +hd: Nat
668	  for +t: Nat
669	  for k: Array<F32>
670	  for +q: List<&2, F32>
671	  {sr_len(scores(hd, t, k, q)) == 1n+t : Nat}
672	
673	def causal(hd, t, k, q):
674	  sr_len_g(t, TA.gemm(1n, 1n+t, hd, U32.from_nat(hd), 1, 1, U32.from_nat(hd), TA.fill_list(q, 0, Array.new(F32, TA.cap_depth(hd), 0.0)), k, Array.new(F32, TA.cap_depth(1n+t), 0.0)))
675	
676	# the weighted sum of the values of positions 0..t: w (t + 1 weights) times V (t + 1 rows of `head` numbers)
677	type WR is Type:
678	  WR{v: Array<F32>, o: List<&2, F32>}
679	
680	def wr_of(v: Array<F32>, al: TA.AL) -> WR:
681	  match al:
682	    case TA.AL{c, l}:
683	      WR{v, l}
684	
685	def wr_read(+hd: Nat, g: TA.G) -> WR:
686	  match g:
687	    case TA.G{wa, v, o}:
688	      wr_of(v, TA.read_l(Nat.sub(hd, 1n), 0, Nil{}, Array.get(F32, o, 0)))
689	
690	def weighted(+hd: Nat, +t: Nat, v: Array<F32>, +w: List<&2, F32>) -> WR:
691	  wr_read(hd, TA.gemm(1n, hd, 1n+t, U32.from_nat(1n+t), 1, U32.from_nat(hd), 1, TA.fill_list(w, 0, Array.new(F32, TA.cap_depth(1n+t), 0.0)), v, Array.new(F32, TA.cap_depth(hd), 0.0)))
692	
693	# one query head: softmax(q·K · scale) · V, and the two arrays handed back
694	type HO is Type:
695	  HO{k: Array<F32>, v: Array<F32>, o: List<&2, F32>}
696	
697	def HO.then(-R: Type, h: HO, then: Array<F32> -> Array<F32> -> List<&2, F32> -> R) -> R:
698	  match h:
699	    case HO{k, v, o}:
700	      then(k, v, o)
701	
702	def head_wr(k: Array<F32>, wr: WR) -> HO:
703	  match wr:
704	    case WR{v, o}:
705	      HO{k, v, o}
706	
707	def head_sr(+hd: Nat, +t: Nat, +scale: F32, v: Array<F32>, sr: SR) -> HO:
708	  match sr:
709	    case SR{k, s}:
710	      head_wr(k, weighted(hd, t, v, T.softmax.l(T.scale.l(scale, s))))
711	
712	def head_att(+hd: Nat, +t: Nat, +scale: F32, k: Array<F32>, v: Array<F32>, +q: List<&2, F32>) -> HO:
713	  head_sr(hd, t, scale, v, scores(hd, t, k, q))
714	
715	# the query heads of one group over its key/value head, in order; the arrays are threaded through
716	def group_att(hs: List<&2, Head>, +hd: Nat, +t: Nat, +scale: F32, k: Array<F32>, v: Array<F32>, +acc: List<&2, F32>) -> HO:
717	  match hs:
718	    case Nil{}:
719	      HO{k, v, acc}
720	    case Con{h, rest}:
721	      HO.then(HO, head_att(hd, t, scale, k, v, join(h)), k2 => v2 => +o =>
722	        group_att(rest, hd, t, scale, k2, v2, List.append(&2, F32, acc, o)))
723	
724	# every group, in order: the outputs of query heads 0, 1, ..., and the caches handed back
725	type AT is Type:
726	  AT{ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>, o: List<&2, F32>}
727	
728	def at_join(p: AT, q: AT) -> AT:
729	  match p q:
730	    case AT{k1, v1, o1} AT{k2, v2, o2}:
731	      AT{List.append(&1, Array<F32>, k1, k2), List.append(&1, Array<F32>, v1, v2), List.append(&2, F32, o1, o2)}
732	
733	# one key/value head's work: its group of query heads, its keys and its values
734	type GW is Type:
735	  GW{g: List<&2, Head>, k: Array<F32>, v: Array<F32>}
736	
737	def gws(gs: List<&2, List<&2, Head>>, ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>) -> List<&1, GW>:
738	  match gs ks vs:
739	    case Con{g, gt} Con{k, kt} Con{v, vt}:
740	      GW{g, k, v} <> gws(gt, kt, vt)
741	    case _ _ _:
742	      Nil{}
743	
744	# some groups one after another, in order
745	def seq_att(ws: List<&1, GW>, +hd: Nat, +t: Nat, +scale: F32, aks: List<&1, Array<F32>>, avs: List<&1, Array<F32>>, +acc: List<&2, F32>) -> AT:
746	  match ws:
747	    case Con{w, wt}:
748	      match w:
749	        case GW{g, k, v}:
750	          HO.then(AT, group_att(g, hd, t, scale, k, v, Nil{}), k2 => v2 => +o =>
751	            seq_att(wt, hd, t, scale, k2 <> aks, v2 <> avs, List.append(&2, F32, acc, o)))
752	    case Nil{}:
753	      AT{List.reverse(&1, Array<F32>, aks), List.reverse(&1, Array<F32>, avs), acc}
754	
755	# a list of groups cut after its first n (a tail loop: a helper in the fork tree must not hold the stack)
756	type GS is Type:
757	  GS{l: List<&1, GW>, r: List<&1, GW>}
758	
759	def cut(n: Nat, xs: List<&1, GW>, acc: List<&1, GW>) -> GS:
760	  match n xs:
761	    case 1n+m Con{x, xt}:
762	      cut(m, xt, x <> acc)
763	    case _ _:
764	      GS{List.reverse(&1, GW, acc), xs}
765	
766	# how many times n groups are halved (the larger half kept) until one or none is left
767	def halvings(fuel: Nat, +n: Nat) -> Nat:
768	  match fuel:
769	    case 0n:
770	      0n
771	    case 1n+f:
772	      match n:
773	        case 0n:
774	          0n
775	        case 1n+p:
776	          match p:
777	            case 0n:
778	              0n
779	            case 1n+q:
780	              1n+halvings(f, Nat.sub(2n+q, TA.half(2n+q)))
781	
782	# the n groups of gs (cut in two halves) in a tree of parallel lets of depth d, each leaf a half of a node; the
783	# outputs and the caches come back in group order, so every head computes what it computed one after another
784	def par_att(d: Nat, gs: GS, +n: Nat, +hd: Nat, +t: Nat, +scale: F32) -> AT:
785	  match d gs:
786	    case 0n GS{l, r}:
787	      p q = seq_att(l, hd, t, scale, Nil{}, Nil{}, Nil{}) seq_att(r, hd, t, scale, Nil{}, Nil{}, Nil{})
788	      at_join(p, q)
789	    case 1n++e GS{l, r}:
790	      p q = par_att(e, cut(TA.half(TA.half(n)), l, Nil{}), TA.half(n), hd, t, scale) par_att(e, cut(TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), hd, t, scale)
791	      at_join(p, q)
792	
793	# the attention output and the cache handed back
794	type KA<-cap: Nat, -t: Nat> is Type:
795	  KA{kv: KV<cap, t>, o: List<&2, F32>}
796	
797	def ka_of(-cap: Nat, -t: Nat, at: AT) -> KA<cap, t>:
798	  match at:
799	    case AT{ks, vs, o}:
800	      KA{KV{ks, vs}, o}
801	
802	def KA.then(-R: Type, -cap: Nat, -t: Nat, ka: KA<cap, t>, then: KV<cap, t> -> List<&2, F32> -> R) -> R:
803	  match ka:
804	    case KA{kv, o}:
805	      then(kv, o)
806	
807	# attention of the (rotated) query heads at position t over the cache of positions 0..t (the newest included).
808	# d is the model's parallel depth (depth_for): with d = 0 (two threads or fewer) the key/value heads run one after
809	# another, since a fork tree on one thread is only overhead (+5% on the 1.7B, NOTES 2026-10-12); otherwise they run
810	# in parallel, one group per leaf. Both orders compute every head the same way, so the bits are the same.
811	def attention(+c: Cfg, +d: Nat, -cap: Nat, +t: Nat, qs: List<&2, Head>, kv: KV<cap, 1n+t>) -> KA<cap, 1n+t>:
812	  match d kv:
813	    case 0n KV{ks, vs}:
814	      ka_of(cap, 1n+t, seq_att(gws(groups(Cfg.kv(c), Cfg.g(c), qs), ks, vs), Cfg.head(c), t, (1.0 / F32.sqrt(F32.from_nat(Cfg.head(c))) : F32), Nil{}, Nil{}, Nil{}))
815	    case 1n+e KV{ks, vs}:
816	      ka_of(cap, 1n+t, par_att(Nat.sub(halvings(Cfg.kv(c), Cfg.kv(c)), 1n), cut(TA.half(Cfg.kv(c)), gws(groups(Cfg.kv(c), Cfg.g(c), qs), ks, vs), Nil{}), Cfg.kv(c), Cfg.head(c), t, (1.0 / F32.sqrt(F32.from_nat(Cfg.head(c))) : F32)))
817	
818	# =====================================================================================================
819	# A layer
820	# =====================================================================================================
821	
822	# the weights of one layer, with their shapes from the config (row = output, as PyTorch stores them)
823	type LW<-c: Cfg> is Type:
824	  LW{ln1: List<&2, F32>, q: TA.Bands<Cfg.dim(c), Cfg.dim(c)>, k: TA.Bands<Cfg.kvw(c), Cfg.dim(c)>, v: TA.Bands<Cfg.kvw(c), Cfg.dim(c)>, o: TA.Bands<Cfg.dim(c), Cfg.dim(c)>, ln2: List<&2, F32>, gate: TA.Bands<Cfg.ffn(c), Cfg.dim(c)>, up: TA.Bands<Cfg.ffn(c), Cfg.dim(c)>, down: TA.Bands<Cfg.dim(c), Cfg.ffn(c)>}
825	
826	# what a layer computes at one position, every intermediate: the tests compare them with transformers, and a
827	# backward pass would read them
828	type Act is Data:
829	  Act{n1: List<&2, F32>, q: List<&2, F32>, k: List<&2, F32>, v: List<&2, F32>, qr: List<&2, F32>, kr: List<&2, F32>, attn: List<&2, F32>, a: List<&2, F32>, n2: List<&2, F32>, m: List<&2, F32>, y: List<&2, F32>}
830	
831	# the layer's output at position t, its cache grown to t + 1 positions, its weights handed back, and the
832	# activations
833	type BR<-c: Cfg, -cap: Nat, -t: Nat> is Type:
834	  BR{y: List<&2, F32>, kv: KV<cap, 1n+t>, lw: LW<c>, act: Act}
835	
836	# continues with the fields of a layer's result (a computed value cannot be matched; a loop over positions
837	# recurses inside k)
838	def BR.then(-R: Type, -c: Cfg, -cap: Nat, -t: Nat, b: BR<c, cap, t>, k: List<&2, F32> -> KV<cap, 1n+t> -> LW<c> -> Act -> R) -> R:
839	  match b:
840	    case BR{y, kv, lw, act}:
841	      k(y, kv, lw, act)
842	
843	def mv.go(-B: Type, -r: Nat, -k: Nat, bl: TA.BL<r, k>, then: TA.Bands<r, k> -> List<&2, F32> -> B) -> B:
844	  match bl:
845	    case TA.BL{w, y}:
846	      then(w, y)
847	
848	# y = W · x, then `then` with W handed back and y (an Array has one owner, so a product returns its matrix)
849	def mv(-B: Type, +r: Nat, +k: Nat, w: TA.Bands<r, k>, +x: List<&2, F32>, then: TA.Bands<r, k> -> List<&2, F32> -> B) -> B:
850	  mv.go(B, r, k, TA.Bands.matvec_l(r, k, w, x), then)
851	
852	# the layer at position t (0-based) on input x, with the cache of positions 0..t-1 and the rotation of t; d is the
853	# parallel depth of attention (see attention)
854	def block(+c: Cfg, +d: Nat, lw: LW<c>, -cap: Nat, +t: Nat, kv: KV<cap, t>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +rt: Rot, +x: List<&2, F32>) -> BR<c, cap, t>:
855	  match lw:
856	    case LW{+ln1, wq, wk, wv, wo, +ln2, wg, wu, wd}:
857	      +n1 = rmsnorm(c, x, ln1)
858	      mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wq, n1, wq => +q =>
859	        mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wk, n1, wk => +k =>
860	          mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wv, n1, wv => +v =>
861	            +qs = ropes(rt, heads(Cfg.heads(c), Cfg.half(c), q))
862	            +ks = ropes(rt, heads(Cfg.kv(c), Cfg.half(c), k))
863	            KA.then(BR<c, cap, t>, cap, 1n+t, attention(c, d, cap, t, qs, KV.push(c, cap, t, kv, room, joins(ks), chunks(Cfg.kv(c), Cfg.head(c), v))), kv2 => +at =>
864	            mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wo, at, wo => +a =>
865	              +x1 = T.zip_with.add(x, a)
866	              +n2 = rmsnorm(c, x1, ln2)
867	              mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wg, n2, wg => +g =>
868	                mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wu, n2, wu => +u =>
869	                  mv(BR<c, cap, t>, Cfg.dim(c), Cfg.ffn(c), wd, swiglu(g, u), wd => +m =>
870	                    +y = T.zip_with.add(x1, m)
871	                    BR{y, kv2, LW{ln1, wq, wk, wv, wo, ln2, wg, wu, wd}, Act{n1, q, k, v, flat(qs), flat(ks), at, a, n2, m, y}}))))))))
872	
873	# block for generation: the same computation, with the activations left empty. Generation never reads them, and
874	# building them (the copies of the rotated heads among them) cost about a fifth of each token on 16 threads (found by
875	# Mathews Fernando's bend-lint, unused_result). It is a copy of block's body, not a flag in it: a flag that chooses
876	# after the computation keeps every intermediate alive to the end, and that measured slower than building them.
877	def block_gen(+c: Cfg, +d: Nat, lw: LW<c>, -cap: Nat, +t: Nat, kv: KV<cap, t>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +rt: Rot, +x: List<&2, F32>) -> BR<c, cap, t>:
878	  match lw:
879	    case LW{+ln1, wq, wk, wv, wo, +ln2, wg, wu, wd}:
880	      +n1 = rmsnorm(c, x, ln1)
881	      mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wq, n1, wq => +q =>
882	        mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wk, n1, wk => +k =>
883	          mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wv, n1, wv => +v =>
884	            +qs = ropes(rt, heads(Cfg.heads(c), Cfg.half(c), q))
885	            +ks = ropes(rt, heads(Cfg.kv(c), Cfg.half(c), k))
886	            KA.then(BR<c, cap, t>, cap, 1n+t, attention(c, d, cap, t, qs, KV.push(c, cap, t, kv, room, joins(ks), chunks(Cfg.kv(c), Cfg.head(c), v))), kv2 => +at =>
887	            mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wo, at, wo => +a =>
888	              +x1 = T.zip_with.add(x, a)
889	              +n2 = rmsnorm(c, x1, ln2)
890	              mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wg, n2, wg => +g =>
891	                mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wu, n2, wu => +u =>
892	                  mv(BR<c, cap, t>, Cfg.dim(c), Cfg.ffn(c), wd, swiglu(g, u), wd => +m =>
893	                    +y = T.zip_with.add(x1, m)
894	                    BR{y, kv2, LW{ln1, wq, wk, wv, wo, ln2, wg, wu, wd}, Act{Nil{}, Nil{}, Nil{}, Nil{}, Nil{}, Nil{}, Nil{}, Nil{}, Nil{}, Nil{}, Nil{}}}))))))))
895	
896	# keep: block with the activations (the tests read them), or block_gen without them; chosen before anything is computed
897	def block_k(keep: Bool, +c: Cfg, +d: Nat, lw: LW<c>, -cap: Nat, +t: Nat, kv: KV<cap, t>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +rt: Rot, +x: List<&2, F32>) -> BR<c, cap, t>:
898	  match keep:
899	    case True{}:
900	      block(c, d, lw, cap, t, kv, room, rt, x)
901	    case False{}:
902	      block_gen(c, d, lw, cap, t, kv, room, rt, x)
903	
904	# =====================================================================================================
905	# The model
906	# =====================================================================================================
907	
908	# emb: the embedding (vocab rows; also the LM head, tied); lnf: the final norm's gains
909	type Model<-c: Cfg> is Type:
910	  Model{d: Nat, emb: TA.Bands<Cfg.vocab(c), Cfg.dim(c)>, layers: List<&1, LW<c>>, lnf: List<&2, F32>}
911	
912	# the last layer's output, the layers and caches handed back, and every layer's activations
913	type Run<-c: Cfg, -cap: Nat, -t: Nat> is Type:
914	  Run{x: List<&2, F32>, layers: List<&1, LW<c>>, cs: List<&1, KV<cap, 1n+t>>, acts: List<&2, Act>}
915	
916	# pend is the previous layer's result: its output is the next layer's input
917	def run(+keep: Bool, +c: Cfg, +d: Nat, -cap: Nat, +t: Nat, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +rt: Rot, ls: List<&1, LW<c>>, cs: List<&1, KV<cap, t>>, al: List<&1, LW<c>>, ac: List<&1, KV<cap, 1n+t>>, aa: List<&2, Act>, pend: BR<c, cap, t>) -> Run<c, cap, t>:
918	  match ls cs pend:
919	    case Con{l, lt} Con{k, kt} BR{+y, kv, l0, act}:
920	      run(keep, c, d, cap, t, room, rt, lt, kt, l0 <> al, kv <> ac, act <> aa, block_k(keep, c, d, l, cap, t, k, room, rt, y))
921	    case _ _ BR{+y, kv, l0, act}:
922	      Run{y, List.reverse(&1, LW<c>, l0 <> al), List.reverse(&1, KV<cap, 1n+t>, kv <> ac), List.reverse(&2, Act, act <> aa)}
923	
924	def layers(+keep: Bool, +c: Cfg, +d: Nat, -cap: Nat, +t: Nat, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +rt: Rot, ls: List<&1, LW<c>>, cs: List<&1, KV<cap, t>>, +x: List<&2, F32>) -> Run<c, cap, t>:
925	  match ls cs:
926	    case Con{l, lt} Con{k, kt}:
927	      run(keep, c, d, cap, t, room, rt, lt, kt, Nil{}, Nil{}, Nil{}, block_k(keep, c, d, l, cap, t, k, room, rt, x))
928	    case _ _:
929	      Run{x, Nil{}, Nil{}, Nil{}}
930	
931	# what one position produces: the logits, the final norm, every layer's activations, the caches grown to t + 1
932	# positions and the model handed back
933	type Out<-c: Cfg, -cap: Nat, -t: Nat> is Type:
934	  Out{lg: List<&2, F32>, nf: List<&2, F32>, acts: List<&2, Act>, cs: List<&1, KV<cap, 1n+t>>, m: Model<c>}
935	
936	# continues with the fields of a position's result: logits, final norm, activations, caches, model
937	def Out.then(-R: Type, -c: Cfg, -cap: Nat, -t: Nat, o: Out<c, cap, t>, k: List<&2, F32> -> List<&2, F32> -> List<&2, Act> -> List<&1, KV<cap, 1n+t>> -> Model<c> -> R) -> R:
938	  match o:
939	    case Out{lg, nf, acts, cs, m}:
940	      k(lg, nf, acts, cs, m)
941	
942	# the final norm, then logits = emb · nf (the LM head is the embedding matrix)
943	def head_of(+c: Cfg, -cap: Nat, -t: Nat, +d: Nat, emb: TA.Bands<Cfg.vocab(c), Cfg.dim(c)>, +lnf: List<&2, F32>, r: Run<c, cap, t>) -> Out<c, cap, t>:
944	  match r:
945	    case Run{+x, ls, cs, acts}:
946	      +nf = rmsnorm(c, x, lnf)
947	      mv(Out<c, cap, t>, Cfg.vocab(c), Cfg.dim(c), emb, nf, emb => +lg => Out{lg, nf, acts, cs, Model{d, emb, ls, lnf}})
948	
949	def forward.go(+keep: Bool, +c: Cfg, -cap: Nat, +t: Nat, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +d: Nat, ls: List<&1, LW<c>>, +lnf: List<&2, F32>, cs: List<&1, KV<cap, t>>, r: TA.BRow<Cfg.vocab(c), Cfg.dim(c)>) -> Out<c, cap, t>:
950	  match r:
951	    case TA.BRow{emb, row}:
952	      head_of(c, cap, t, d, emb, lnf, layers(keep, c, d, cap, t, room, rot_at(c, t), ls, cs, row))
953	
954	# token tok at position t (0-based), with the caches of positions 0..t-1 (one per layer) and room for cap positions;
955	# room is the proof that t < cap; keep: whether the result carries every layer's activations
956	def forward_k(+keep: Bool, +c: Cfg, -cap: Nat, +t: Nat, m: Model<c>, cs: List<&1, KV<cap, t>>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +tok: Nat) -> Out<c, cap, t>:
957	  match m:
958	    case Model{+d, emb, ls, +lnf}:
959	      forward.go(keep, c, cap, t, room, d, ls, lnf, cs, TA.Bands.read_row(Cfg.vocab(c), Cfg.dim(c), emb, tok))
960	
961	# a position with every layer's activations (the tests read them)
962	def forward(+c: Cfg, -cap: Nat, +t: Nat, m: Model<c>, cs: List<&1, KV<cap, t>>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +tok: Nat) -> Out<c, cap, t>:
963	  forward_k(True{}, c, cap, t, m, cs, room, tok)
964	
965	# a position for generation: the same logits, caches and model. With bands (d > 0, more than two threads) the
966	# activations are left empty, about a fifth faster per token on 16 threads; with one band (d = 0) they are kept, since
967	# leaving them out measured 4-11% slower per token on one thread (llama.bend NOTES, 2026-10-16), for no reason found yet
968	def forward_gen(+c: Cfg, -cap: Nat, +t: Nat, m: Model<c>, cs: List<&1, KV<cap, t>>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +tok: Nat) -> Out<c, cap, t>:
969	  match m:
970	    case Model{+d, emb, ls, +lnf}:
971	      forward.go(Nat.is_eq(d, 0n), c, cap, t, room, d, ls, lnf, cs, TA.Bands.read_row(Cfg.vocab(c), Cfg.dim(c), emb, tok))
972	
973	# a prompt position, whose logits nobody reads: the layers and the caches as forward_gen gives them, without the final
974	# norm and the product by the embedding (vocab x dim, the largest of a position), so lg and nf are empty
975	def pre_of(-c: Cfg, -cap: Nat, -t: Nat, +d: Nat, emb: TA.Bands<Cfg.vocab(c), Cfg.dim(c)>, +lnf: List<&2, F32>, r: Run<c, cap, t>) -> Out<c, cap, t>:
976	  match r:
977	    case Run{x, ls, cs, acts}:
978	      Out{Nil{}, Nil{}, acts, cs, Model{d, emb, ls, lnf}}
979	
980	def forward_pre.go(+c: Cfg, -cap: Nat, +t: Nat, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +d: Nat, ls: List<&1, LW<c>>, +lnf: List<&2, F32>, cs: List<&1, KV<cap, t>>, r: TA.BRow<Cfg.vocab(c), Cfg.dim(c)>) -> Out<c, cap, t>:
981	  match r:
982	    case TA.BRow{emb, row}:
983	      pre_of(c, cap, t, d, emb, lnf, layers(Nat.is_eq(d, 0n), c, d, cap, t, room, rot_at(c, t), ls, cs, row))
984	
985	def forward_pre(+c: Cfg, -cap: Nat, +t: Nat, m: Model<c>, cs: List<&1, KV<cap, t>>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +tok: Nat) -> Out<c, cap, t>:
986	  match m:
987	    case Model{+d, emb, ls, +lnf}:
988	      forward_pre.go(c, cap, t, room, d, ls, lnf, cs, TA.Bands.read_row(Cfg.vocab(c), Cfg.dim(c), emb, tok))
989	
990	# n empty caches, one per layer, each with room for cap positions
991	def caches(+c: Cfg, n: Nat, +cap: Nat) -> List<&1, KV<cap, 0n>>:
992	  match n:
993	    case 0n:
994	      Nil{}
995	    case 1n+p:
996	      KV.empty(c, cap) <> caches(c, p, cap)
997	
998	# the first index of the largest value (torch.argmax does the same), and the value
999	type AM is Data:
1000	  AM{i: Nat, v: F32}
1001	
1002	def AM.i(am: AM) -> Nat:
1003	  match am:
1004	    case AM{i, v}:
1005	      i
1006	
1007	def AM.v(am: AM) -> F32:
1008	  match am:
1009	    case AM{i, v}:
1010	      v
1011	
1012	def gt_head(xs: List<&2, F32>, +v: F32) -> Bool:
1013	  match xs:
1014	    case Nil{}:
1015	      False{}
1016	    case Con{+h, t}:
1017	      F32.is_gt(h, v)
1018	
1019	def argmax.go(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> AM:
1020	  match xs g:
1021	    case Nil{} _:
1022	      AM{bi, bv}
1023	    case Con{+h, +t} True{}:
1024	      argmax.go(t, gt_head(t, h), 1n+i, i, h)
1025	    case Con{h, +t} False{}:
1026	      argmax.go(t, gt_head(t, bv), 1n+i, bi, bv)
1027	
1028	def argmax(xs: List<&2, F32>) -> AM:
1029	  match xs:
1030	    case Nil{}:
1031	      AM{0n, 0.0}
1032	    case Con{+h, +t}:
1033	      argmax.go(t, gt_head(t, h), 1n, 0n, h)
1034	
1035	# =====================================================================================================
1036	# Loading (the weights as written by llama.bend's reference/prep.py: <dir>/w/<name>.bin, little-endian F32,
1037	# one file per tensor, row = output; the config in <dir>/cfg.txt). The block reader is bend-ml's (MIT).
1038	# =====================================================================================================
1039	
1040	# 4 little-endian bytes -> F32 (U32 and F32 hold the same Word(32n): it is a bitcast)
1041	def bitcast(u: U32) -> F32:
1042	  match u:
1043	    case U32{w}:
1044	      F32{w}
1045	
1046	def f32_le(+b0: U32, +b1: U32, +b2: U32, +b3: U32) -> F32:
1047	  bitcast((b0 + (b1 * 256 : U32) + (b2 * 65536 : U32) + (b3 * 16777216 : U32) : U32))
1048	
1049	def f32s(bs: List<&2, U32>) -> List<&2, F32>:
1050	  match bs:
1051	    case Con{b0, Con{b1, Con{b2, Con{b3, t}}}}:
1052	      f32_le(b0, b1, b2, b3) <> f32s(t)
1053	    case _:
1054	      Nil{}
1055	
1056	# reads the file in blocks of `blk` bytes; k blocks remaining; m is the block already read
1057	def rf_loop(k: Nat, +off: U32, +blk: U32, m: File & Result<&1, &1, U32 & String, List<&2, U32>>, acc: List<&2, List<&2, F32>>) -> IO(List<&2, F32>):
1058	  match k m:
1059	    case 0n Tuple{f, r}:
1060	      do IO<List<&2, F32>>:
1061	        File.close(f)
1062	        return List.concat(&2, F32, List.reverse(&2, List<&2, F32>, acc))
1063	    case 1n+p Tuple{f, r}:
1064	      do IO<List<&2, F32>>:
1065	        bs : List<&2, U32> <- IO.pass(List<&2, U32>, r)
1066	        m2 : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, (off + blk : U32), blk)
1067	        rf_loop(p, (off + blk : U32), blk, m2, f32s(bs) <> acc)
1068	
1069	def blocks_of(n: Nat, +per: Nat) -> Nat:
1070	  Nat.div(Nat.add(n, Nat.sub(per, 1n)), per)
1071	
1072	# n F32 numbers (little-endian) from a file
1073	def read_floats(path: String, +n: Nat) -> IO(List<&2, F32>):
1074	  do IO<List<&2, F32>>:
1075	    f : File <- IO.try(File, File.open(path, "r"))
1076	    m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, 0, 1048576)
1077	    fl : List<&2, F32> <- rf_loop(blocks_of(Nat.mul(n, 4n), 1048576n), 0, 1048576, m, Nil{})
1078	    return List.take(&2, F32, fl, n)
1079	
1080	# the weight matrices are kept in 2^d row bands (TA.Bands); each matrix · vector product runs its bands in
1081	# parallel without copying the weights (bend-ml NOTES exp. 11). With 2 threads or fewer the bands cannot help,
1082	# so d = 0: a single band, which computes what the sequential Mat.matmul_nt computes (bend-ml exp. 14).
1083	def depth_sel(few: Bool) -> Nat:
1084	  match few:
1085	    case True{}:
1086	      0n
1087	    case False{}:
1088	      4n
1089	
1090	def depth_for(+threads: U32) -> Nat:
1091	  depth_sel(U32.is_le(threads, 2))
1092	
1093	def nm(+i: Nat, f: String) -> String:
1094	  "h" ++ Nat.show(i) ++ "." ++ f
1095	
1096	# ---- the boundary: every weight file must hold exactly the numbers its shape needs ----
1097	
1098	def size_cmp(ok: Bool, +path: String, +got: U32, +want: Nat) -> IO(Unit):
1099	  match ok:
1100	    case True{}:
1101	      IO.pure(Unit, Unit{})
1102	    case False{}:
1103	      IO.die(Unit, 1, path ++ ": expected " ++ Nat.show(want) ++ " bytes, found " ++ U32.show(got) ++ "; run make prep again")
1104	
1105	def size_ok(m: File & Result<&1, &1, U32 & String, U32>, +want: Nat, +path: String) -> IO(Unit):
1106	  match m:
1107	    case Tuple{f, r}:
1108	      do IO<Unit>:
1109	        +n : U32 <- IO.pass(U32, r)
1110	        File.close(f)
1111	        size_cmp(U32.is_eq(n, U32.from_nat(want)), path, n, want)
1112	
1113	# stops the program unless the file holds exactly n F32 numbers (n * 4 bytes)
1114	def check_size(+path: String, +n: Nat) -> IO(Unit):
1115	  do IO<Unit>:
1116	    f : File <- IO.try(File, File.open(path, "r"))
1117	    m : File & Result<&1, &1, U32 & String, U32> <- File.size(f)
1118	    size_ok(m, Nat.mul(n, 4n), path)
1119	
1120	def wpath(dir: String, name: String) -> String:
1121	  dir ++ "/w/" ++ name ++ ".bin"
1122	
1123	def get_list(+dir: String, +name: String, +n: Nat) -> IO(List<&2, F32>):
1124	  do IO<List<&2, F32>>:
1125	    check_size(wpath(dir, name), n)
1126	    read_floats(wpath(dir, name), n)
1127	
1128	# reads the file in 1 MB blocks and writes each block straight into the matrix, so only one
1129	# block is a list at any time
1130	# ---- matrices kept in bands: each band is read from its own range of the file ----
1131	
1132	# writes the little-endian F32 numbers held in bs (4 bytes each) at index i, without building a
1133	# list of F32 first
1134	def fill_bytes(bs: List<&2, U32>, +i: U32, a: Array<F32>) -> Array<F32>:
1135	  match bs:
1136	    case Con{+b0, Con{+b1, Con{+b2, Con{+b3, t}}}}:
1137	      fill_bytes(t, (i + 1 : U32), Array.set(F32, a, i, f32_le(b0, b1, b2, b3)))
1138	    case _:
1139	      a
1140	
1141	def fill_bytes_mat(-r: Nat, -c: Nat, +i: U32, bs: List<&2, U32>, m: TA.Mat<r, c>) -> TA.Mat<r, c>:
1142	  match m:
1143	    case TA.Mat{d}:
1144	      TA.Mat{fill_bytes(bs, i, d)}
1145	
1146	# reads blocks of at most 1 MB until rem bytes have been read, so a band never takes numbers
1147	# that belong to the next one
1148	def stream_exact(k: Nat, -r: Nat, -c: Nat, +rem: Nat, +off: U32, +idx: U32, m: File & Result<&1, &1, U32 & String, List<&2, U32>>, acc: TA.Mat<r, c>) -> IO(TA.Mat<r, c>):
1149	  match k m:
1150	    case 0n Tuple{f, rs}:
1151	      do IO<TA.Mat<r, c>>:
1152	        File.close(f)
1153	        return acc
1154	    case 1n+p Tuple{f, rs}:
1155	      do IO<TA.Mat<r, c>>:
1156	        bs : List<&2, U32> <- IO.pass(List<&2, U32>, rs)
1157	        m2 : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, (off + 1048576 : U32), U32.from_nat(Nat.min(1048576n, Nat.sub(rem, 1048576n))))
1158	        stream_exact(p, r, c, Nat.sub(rem, 1048576n), (off + 1048576 : U32), (idx + 262144 : U32), m2, fill_bytes_mat(r, c, idx, bs, acc))
1159	
1160	# rr x c numbers starting at number `start` of the file, into a band of r rows (rr == r, see
1161	# load_tree; the Array is sized for rr * c numbers as Mat.zeros would size it)
1162	def get_band(path: String, -r: Nat, +rr: Nat, +c: Nat, +start: Nat) -> IO(TA.Mat<r, c>):
1163	  do IO<TA.Mat<r, c>>:
1164	    f : File <- IO.try(File, File.open(path, "r"))
1165	    m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, U32.from_nat(Nat.mul(start, 4n)), U32.from_nat(Nat.min(1048576n, Nat.mul(Nat.mul(rr, c), 4n))))
1166	    stream_exact(blocks_of(Nat.mul(Nat.mul(rr, c), 4n), 1048576n), r, c, Nat.mul(Nat.mul(rr, c), 4n), U32.from_nat(Nat.mul(start, 4n)), 0, m, TA.Mat{Array.new(F32, TA.cap_depth(Nat.mul(rr, c)), 0.0)})
1167	
1168	# fills the bands of b (whose shape comes from TA.Bands.zeros) in row order; rr is r at run time
1169	def load_tree(+c: Nat, -r: Nat, b: TA.Bands<r, c>, +rr: Nat, -e: {rr == r : Nat}, +path: String, +start: Nat) -> IO(TA.Bands<r, c>):
1170	  match b:
1171	    case TA.BLeaf{m}:
1172	      do IO<TA.Bands<r, c>>:
1173	        band : TA.Mat<r, c> <- get_band(path, r, rr, c, start)
1174	        return TA.BLeaf{band}
1175	    case TA.BNode{x, y}:
1176	      do IO<TA.Bands<r, c>>:
1177	        x2 : TA.Bands<TA.half(r), c> <- load_tree(c, TA.half(r), x, TA.half(rr), TA.eq_half(rr, r, e), path, start)
1178	        y2 : TA.Bands<Nat.sub(r, TA.half(r)), c> <- load_tree(c, Nat.sub(r, TA.half(r)), y, Nat.sub(rr, TA.half(rr)), TA.eq_rest(rr, r, e), path, Nat.add(start, Nat.mul(TA.half(rr), c)))
1179	        return TA.BNode{x2, y2}
1180	
1181	def get_bands(+dir: String, +name: String, +d: Nat, +r: Nat, +c: Nat) -> IO(TA.Bands<r, c>):
1182	  do IO<TA.Bands<r, c>>:
1183	    check_size(wpath(dir, name), Nat.mul(r, c))
1184	    load_tree(c, r, TA.Bands.zeros(d, r, c), r, {==}, wpath(dir, name), 0n)
1185	
1186	# one layer's weights; d: the band depth (depth_for)
1187	def load_lw(+c: Cfg, +dir: String, +i: Nat, +d: Nat) -> IO(LW<c>):
1188	  do IO<LW<c>>:
1189	    ln1 : List<&2, F32> <- get_list(dir, nm(i, "ln1"), Cfg.dim(c))
1190	    q : TA.Bands<Cfg.dim(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "q"), d, Cfg.dim(c), Cfg.dim(c))
1191	    k : TA.Bands<Cfg.kvw(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "k"), d, Cfg.kvw(c), Cfg.dim(c))
1192	    v : TA.Bands<Cfg.kvw(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "v"), d, Cfg.kvw(c), Cfg.dim(c))
1193	    o : TA.Bands<Cfg.dim(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "o"), d, Cfg.dim(c), Cfg.dim(c))
1194	    ln2 : List<&2, F32> <- get_list(dir, nm(i, "ln2"), Cfg.dim(c))
1195	    gate : TA.Bands<Cfg.ffn(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "gate"), d, Cfg.ffn(c), Cfg.dim(c))
1196	    up : TA.Bands<Cfg.ffn(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "up"), d, Cfg.ffn(c), Cfg.dim(c))
1197	    down : TA.Bands<Cfg.dim(c), Cfg.ffn(c)> <- get_bands(dir, nm(i, "down"), d, Cfg.dim(c), Cfg.ffn(c))
1198	    return LW{ln1, q, k, v, o, ln2, gate, up, down}
1199	
1200	def load_layers(k: Nat, +c: Cfg, +dir: String, +d: Nat, +i: Nat, acc: List<&1, LW<c>>) -> IO(List<&1, LW<c>>):
1201	  match k:
1202	    case 0n:
1203	      IO.pure(List<&1, LW<c>>, List.reverse(&1, LW<c>, acc))
1204	    case 1n+p:
1205	      do IO<List<&1, LW<c>>>:
1206	        l : LW<c> <- load_lw(c, dir, i, d)
1207	        load_layers(p, c, dir, d, 1n+i, l <> acc)
1208	
1209	# the whole model; the embedding gets one more level of bands than the layers (2^5 with 16 threads)
1210	def load_model(+c: Cfg, +dir: String) -> IO(Model<c>):
1211	  do IO<Model<c>>:
1212	    threads : U32 <- IO.thread_count()
1213	    +d : Nat = depth_for(threads)
1214	    emb : TA.Bands<Cfg.vocab(c), Cfg.dim(c)> <- get_bands(dir, "emb", Nat.add(d, Nat.min(d, 1n)), Cfg.vocab(c), Cfg.dim(c))
1215	    ls : List<&1, LW<c>> <- load_layers(Cfg.layers(c), c, dir, d, 0n, Nil{})
1216	    lnf : List<&2, F32> <- get_list(dir, "lnf", Cfg.dim(c))
1217	    return Model{d, emb, ls, lnf}
1218	
1219	# ---- the config file: kv g half f layers vocab, then the bits of eps and theta as F32, in decimal ----
1220	
1221	def rb_done(m: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(List<&2, U32>):
1222	  (f, r) = m
1223	  do IO<List<&2, U32>>:
1224	    bs : List<&2, U32> <- IO.pass(List<&2, U32>, r)
1225	    File.close(f)
1226	    return bs
1227	
1228	def read_bytes(path: String) -> IO(List<&2, U32>):
1229	  do IO<List<&2, U32>>:
1230	    f : File <- IO.try(File, File.open(path, "r"))
1231	    m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_bytes(f, 1048576)
1232	    rb_done(m)
1233	
1234	def digit_head(xs: List<&2, U32>) -> Bool:
1235	  match xs:
1236	    case Nil{}:
1237	      False{}
1238	    case Con{+b, t}:
1239	      Bool.and(U32.is_le(48, b), U32.is_le(b, 57))
1240	
1241	# the decimal numbers in bs, separated by anything that is not a digit; dg says whether bs starts with a digit
1242	def u32s.go(bs: List<&2, U32>, dg: Bool, +cur: U32, have: Bool, acc: List<&2, U32>) -> List<&2, U32>:
1243	  match bs dg have:
1244	    case Nil{} _ True{}:
1245	      List.reverse(&2, U32, cur <> acc)
1246	    case Nil{} _ False{}:
1247	      List.reverse(&2, U32, acc)
1248	    case Con{+b, +t} True{} _:
1249	      u32s.go(t, digit_head(t), ((cur * 10 : U32) + U32.sub(b, 48) : U32), True{}, acc)
1250	    case Con{b, +t} False{} True{}:
1251	      u32s.go(t, digit_head(t), 0, False{}, cur <> acc)
1252	    case Con{b, +t} False{} False{}:
1253	      u32s.go(t, digit_head(t), 0, False{}, acc)
1254	
1255	def u32s(+bs: List<&2, U32>) -> List<&2, U32>:
1256	  u32s.go(bs, digit_head(bs), 0, False{}, Nil{})
1257	
1258	def cfg_of(+path: String, xs: List<&2, U32>) -> IO(Cfg):
1259	  match xs:
1260	    case Con{+kv, Con{+g, Con{+half, Con{+f, Con{+l, Con{+v, Con{+e, Con{+t, Nil{}}}}}}}}}:
1261	      IO.pure(Cfg, Cfg{U32.to_nat(kv), U32.to_nat(g), U32.to_nat(half), U32.to_nat(f), U32.to_nat(l), U32.to_nat(v), bitcast(e), bitcast(t)})
1262	    case _:
1263	      IO.die(Cfg, 1, path ++ ": expected 8 numbers (kv g half f layers vocab eps_bits theta_bits); run make prep again")
1264	
1265	# the config of the model in dir (written by reference/prep.py)
1266	def read_cfg(+dir: String) -> IO(Cfg):
1267	  do IO<Cfg>:
1268	    bs : List<&2, U32> <- read_bytes(dir ++ "/cfg.txt")
1269	    cfg_of(dir ++ "/cfg.txt", u32s(bs))