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