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	# - a token drawn with temperature, top-k and top-p is an id of the logits it was drawn from (sample_below).
18	#
19	# The order of every F32 operation follows transformers' modeling_llama.py (5.19), so the results match it to a
20	# few ulps and the argmax exactly (llama.bend's NOTES.md).
21	
22	import Base
23	import bend-ml-tensor-array@0.1.19.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	# ---- lists of F32: copied from bend-ml-tensor 0.1.2.0 (same author, MIT), each operation in the same order ----
286	
287	# xs + ys, element by element (to the shorter length)
288	def list.add(xs: List<&2, F32>, ys: List<&2, F32>) -> List<&2, F32>:
289	  match xs ys:
290	    case Con{+h, t} Con{+g, u}:
291	      (h + g : F32) <> list.add(t, u)
292	    case _ _:
293	      Nil{}
294	
295	# xs * ys, element by element (to the shorter length)
296	def list.mul(xs: List<&2, F32>, ys: List<&2, F32>) -> List<&2, F32>:
297	  match xs ys:
298	    case Con{+h, t} Con{+g, u}:
299	      (h * g : F32) <> list.mul(t, u)
300	    case _ _:
301	      Nil{}
302	
303	# k * xs
304	def list.scale(+k: F32, xs: List<&2, F32>) -> List<&2, F32>:
305	  match xs:
306	    case Nil{}:
307	      Nil{}
308	    case Con{+h, t}:
309	      (k * h : F32) <> list.scale(k, t)
310	
311	# acc + xs, a running sum from the left
312	def list.sum.go(xs: List<&2, F32>, acc: F32) -> F32:
313	  match xs:
314	    case Nil{}:
315	      acc
316	    case Con{+h, t}:
317	      list.sum.go(t, (acc + h : F32))
318	
319	def list.max.go(xs: List<&2, F32>, cur: F32) -> F32:
320	  match xs:
321	    case Nil{}:
322	      cur
323	    case Con{+h, t}:
324	      list.max.go(t, F32.max(cur, h))
325	
326	# the maximum of xs (0 for an empty list)
327	def list.max(xs: List<&2, F32>) -> F32:
328	  match xs:
329	    case Nil{}:
330	      0.0
331	    case Con{+h, t}:
332	      list.max.go(t, h)
333	
334	# exp(x - m) for each x
335	def list.exp_shift(xs: List<&2, F32>, +m: F32) -> List<&2, F32>:
336	  match xs:
337	    case Nil{}:
338	      Nil{}
339	    case Con{+h, t}:
340	      F32.exp((h - m : F32)) <> list.exp_shift(t, m)
341	
342	# xs / s
343	def list.div(xs: List<&2, F32>, +s: F32) -> List<&2, F32>:
344	  match xs:
345	    case Nil{}:
346	      Nil{}
347	    case Con{+h, t}:
348	      (h / s : F32) <> list.div(t, s)
349	
350	def list.softmax.fin(+es: List<&2, F32>) -> List<&2, F32>:
351	  list.div(es, list.sum.go(es, 0.0))
352	
353	# stable softmax: subtracts the maximum before exponentiating
354	def list.softmax(+xs: List<&2, F32>) -> List<&2, F32>:
355	  list.softmax.fin(list.exp_shift(xs, list.max(xs)))
356	
357	# the squares of xs
358	def squares(xs: List<&2, F32>) -> List<&2, F32>:
359	  match xs:
360	    case Nil{}:
361	      Nil{}
362	    case Con{+h, t}:
363	      (h * h : F32) <> squares(t)
364	
365	# adjacent pairs added: [a, b, c, d, e] -> [a + b, c + d, e]
366	def pair_sums(xs: List<&2, F32>) -> List<&2, F32>:
367	  match xs:
368	    case Con{+a, Con{+b, t}}:
369	      (a + b : F32) <> pair_sums(t)
370	    case _:
371	      xs
372	
373	# the sum as a balanced tree of additions (pairwise summation): its rounding error grows with log2(n), not n, as
374	# in PyTorch's cascade sum. A running sum loses the small squares next to a large one (NOTES, 2026-10-10).
375	def tree_sum(fuel: Nat, xs: List<&2, F32>) -> F32:
376	  match fuel xs:
377	    case _ Nil{}:
378	      0.0
379	    case _ Con{+a, Nil{}}:
380	      a
381	    case 0n Con{+a, t}:
382	      list.sum.go(t, a)
383	    case 1n+p Con{+a, t}:
384	      tree_sum(p, pair_sums(Con{a, t}))
385	
386	# RMSNorm: g * (x * (1 / sqrt(mean(x²) + eps))), in this order (LlamaRMSNorm.forward)
387	def rmsnorm(+c: Cfg, +x: List<&2, F32>, +g: List<&2, F32>) -> List<&2, F32>:
388	  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))
389	
390	# inv_freq[i] = 1 / (theta ^ (2i / head)), i = 0..half-1 (LlamaRotaryEmbedding, default rope)
391	def inv_freq.go(n: Nat, +i: Nat, +theta: F32, +hd: F32) -> List<&2, F32>:
392	  match n:
393	    case 0n:
394	      Nil{}
395	    case 1n+p:
396	      (1.0 / F32.pow(theta, (F32.from_nat(Nat.mul(2n, i)) / hd : F32)) : F32) <> inv_freq.go(p, 1n+i, theta, hd)
397	
398	def inv_freq(+c: Cfg) -> List<&2, F32>:
399	  inv_freq.go(Cfg.half(c), 0n, Cfg.theta(c), F32.from_nat(Cfg.head(c)))
400	
401	def angles(+t: F32, fs: List<&2, F32>) -> List<&2, F32>:
402	  match fs:
403	    case Nil{}:
404	      Nil{}
405	    case Con{+f, r}:
406	      (t * f : F32) <> angles(t, r)
407	
408	def cos.l(xs: List<&2, F32>) -> List<&2, F32>:
409	  match xs:
410	    case Nil{}:
411	      Nil{}
412	    case Con{+x, r}:
413	      F32.cos(x) <> cos.l(r)
414	
415	def sin.l(xs: List<&2, F32>) -> List<&2, F32>:
416	  match xs:
417	    case Nil{}:
418	      Nil{}
419	    case Con{+x, r}:
420	      F32.sin(x) <> sin.l(r)
421	
422	# the cosines and sines of position t, one per frequency (shared by every head of every layer)
423	type Rot is Data:
424	  Rot{cs: List<&2, F32>, sn: List<&2, F32>}
425	
426	def rot_at(+c: Cfg, +t: Nat) -> Rot:
427	  Rot{cos.l(angles(F32.from_nat(t), inv_freq(c))), sin.l(angles(F32.from_nat(t), inv_freq(c)))}
428	
429	# first half: lo*cos + (-hi)*sin; second half: hi*cos + lo*sin (q*cos + rotate_half(q)*sin)
430	def rot_lo(h1: List<&2, F32>, h2: List<&2, F32>, cs: List<&2, F32>, sn: List<&2, F32>) -> List<&2, F32>:
431	  match h1 h2 cs sn:
432	    case Con{+a, ta} Con{+b, tb} Con{+c, tc} Con{+s, ts}:
433	      ((a * c : F32) + (F32.neg(b) * s : F32) : F32) <> rot_lo(ta, tb, tc, ts)
434	    case _ _ _ _:
435	      Nil{}
436	
437	def rot_hi(h1: List<&2, F32>, h2: List<&2, F32>, cs: List<&2, F32>, sn: List<&2, F32>) -> List<&2, F32>:
438	  match h1 h2 cs sn:
439	    case Con{+a, ta} Con{+b, tb} Con{+c, tc} Con{+s, ts}:
440	      ((b * c : F32) + (a * s : F32) : F32) <> rot_hi(ta, tb, tc, ts)
441	    case _ _ _ _:
442	      Nil{}
443	
444	def rope(rt: Rot, h: Head) -> Head:
445	  match rt h:
446	    case Rot{+cs, +sn} Head{+lo, +hi}:
447	      Head{rot_lo(lo, hi, cs, sn), rot_hi(lo, hi, cs, sn)}
448	
449	def ropes(+rt: Rot, hs: List<&2, Head>) -> List<&2, Head>:
450	  match hs:
451	    case Nil{}:
452	      Nil{}
453	    case Con{h, t}:
454	      rope(rt, h) <> ropes(rt, t)
455	
456	# silu(x) = x / (1 + exp(-x)), the form ATen's CPU kernel uses
457	def silu(+x: F32) -> F32:
458	  (x / (1.0 + F32.exp(F32.neg(x)) : F32) : F32)
459	
460	# silu(gate) * up
461	def swiglu(g: List<&2, F32>, u: List<&2, F32>) -> List<&2, F32>:
462	  match g u:
463	    case Con{+a, ga} Con{+b, ub}:
464	      (silu(a) * b : F32) <> swiglu(ga, ub)
465	    case _ _:
466	      Nil{}
467	
468	# =====================================================================================================
469	# The key/value cache in flat arrays, indexed by the number of positions it holds
470	# =====================================================================================================
471	
472	# One pair of arrays per key/value head: the keys (after RoPE) and the values of positions 0..t-1, position p at
473	# [p·head, p·head + head). The type indexes are the capacity cap (the positions the arrays have room for, set by
474	# KV.empty) and t, the number of positions written. KV.push asks for a proof that t < cap, so a cache is never
475	# written past its room (kv_room says why that is enough).
476	type KV<-cap: Nat, -t: Nat> is Type:
477	  KV{ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>}
478	
479	# n arrays of 2^d slots
480	def arrays(n: Nat, +d: Nat) -> List<&1, Array<F32>>:
481	  match n:
482	    case 0n:
483	      Nil{}
484	    case 1n+p:
485	      Array.new(F32, d, 0.0) <> arrays(p, d)
486	
487	# a cache with no position yet, with room for cap positions per key/value head
488	def KV.empty(+c: Cfg, +cap: Nat) -> KV<cap, 0n>:
489	  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))))}
490	
491	# writes row j of rs into array j, at off
492	def write_rows(arrs: List<&1, Array<F32>>, rs: List<&2, List<&2, F32>>, +off: U32) -> List<&1, Array<F32>>:
493	  match arrs rs:
494	    case Con{a, at} Con{r, rt}:
495	      TA.fill_list(r, off, a) <> write_rows(at, rt, off)
496	    case _ _:
497	      arrs
498	
499	# position t: the keys and values of every key/value head written at t·head; room is the proof that t < cap
500	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>:
501	  match kv:
502	    case KV{ka, va}:
503	      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))))}
504	
505	# the number of key/value heads a cache holds
506	def kv_count(-cap: Nat, -t: Nat, kv: KV<cap, t>) -> Nat:
507	  match kv:
508	    case KV{ka, va}:
509	      List.length(&1, Array<F32>, ka)
510	
511	# LAW (arrays_len): arrays(n, d) makes exactly n arrays.
512	law arrays_len:
513	  for +n: Nat
514	  for +d: Nat
515	  {List.length(&1, Array<F32>, arrays(n, d)) == n : Nat}
516	
517	# PROOF: by induction on n: none for 0; one more than the rest for 1 + p.
518	def arrays_len(n, d):
519	  match n:
520	    case 0n:
521	      {==}
522	    case 1n+p:
523	      %arrays_len(p, d) : {1n+List.length(&1, Array<F32>, arrays(p, d)) == 1n+_ : Nat}
524	      {==}
525	
526	# LAW (kv_heads): an empty cache holds one array of keys (and one of values) per key/value head, so key/value head j
527	# of the model reads column j and no other.
528	law kv_heads:
529	  for +c: Cfg
530	  for +cap: Nat
531	  {kv_count(cap, 0n, KV.empty(c, cap)) == Cfg.kv(c) : Nat}
532	
533	def kv_heads(c, cap):
534	  arrays_len(Cfg.kv(c), TA.cap_depth(Nat.mul(cap, Cfg.head(c))))
535	
536	# a <= b  and  b <= c  ->  a <= c
537	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}:
538	  match a b c:
539	    case 0n _ _:
540	      {==}
541	    case 1n+p 0n _:
542	      %e1 : TA.BD(_, Unit, {True{} == TA.leb(1n+p, c) : Bool})
543	      Unit{}
544	    case 1n+p 1n+q 0n:
545	      %e2 : TA.BD(_, Unit, {True{} == TA.leb(1n+p, 0n) : Bool})
546	      Unit{}
547	    case 1n+p 1n+q 1n+r:
548	      leb_trans(p, q, r, e1, e2)
549	
550	# a <= b  ->  a·h <= b·h
551	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}:
552	  match a b:
553	    case 0n _:
554	      {==}
555	    case 1n+p 0n:
556	      %e : TA.BD(_, Unit, {True{} == TA.leb(Nat.mul(1n+p, h), 0n) : Bool})
557	      Unit{}
558	    case 1n++p 1n++q:
559	      TA.leb_add(h, h, Nat.mul(p, h), Nat.mul(q, h), TA.leb_refl(h), leb_mul(p, q, h, e))
560	
561	# LAW (kv_room): KV.empty(c, cap) gives each key/value head arrays of 2^cap_depth(cap·head) slots, and KV.push at
562	# position t writes head numbers from t·head. For a position below the capacity (t < cap) they all fall inside the
563	# slots: (1 + t)·head <= 2^cap_depth(cap·head), so a push never wraps around onto position 0. (What stays trusted is
564	# 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
565	# compile (tests/bad_cap.bend).
566	law kv_room:
567	  for +cap: Nat
568	  for +h: Nat
569	  for +t: Nat
570	  for e: {True{} == TA.leb(1n+t, cap) : Bool}
571	  {True{} == TA.leb(Nat.mul(1n+t, h), TA.pow2(TA.cap_depth(Nat.mul(cap, h)))) : Bool}
572	
573	# For callers that walk the positions with a counter of steps left (fuel), keeping t + fuel = cap: with 1 + p steps
574	# left, position t is below the capacity (step_room), and after the step t + 1 + p = cap still holds (step_inv).
575	def step_room(+t: Nat, +p: Nat, -cap: Nat, e: {Nat.add(t, 1n+p) == cap : Nat}) -> {True{} == TA.leb(1n+t, cap) : Bool}:
576	  %e : {True{} == TA.leb(1n+t, _) : Bool}
577	  %add_succ(t, p) : {True{} == TA.leb(1n+t, _) : Bool}
578	  TA.leb_up(t, t, p, TA.leb_refl(t))
579	
580	def step_inv(+t: Nat, -p: Nat, -cap: Nat, e: {Nat.add(t, 1n+p) == cap : Nat}) -> {Nat.add(1n+t, p) == cap : Nat}:
581	  %e : {1n+Nat.add(t, p) == _ : Nat}
582	  add_succ(t, p)
583	
584	# PROOF: (1 + t)·head <= cap·head by leb_mul, and cap·head <= 2^cap_depth(cap·head) by tensor-array's cap_ok; then
585	# leb_trans.
586	def kv_room(cap, h, t, e):
587	  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)))
588	
589	# =====================================================================================================
590	# Grouped-query attention
591	# =====================================================================================================
592	# the query heads in n groups of g, in order: group j holds query heads j·g .. j·g + g - 1
593	def groups(n: Nat, +g: Nat, +hs: List<&2, Head>) -> List<&2, List<&2, Head>>:
594	  match n:
595	    case 0n:
596	      Nil{}
597	    case 1n+p:
598	      List.take(&2, Head, hs, g) <> groups(p, g, List.drop(&2, Head, hs, g))
599	
600	# LAW (groups_len): the query heads make exactly as many groups as there are key/value heads, so attention
601	# (which walks the groups and the columns together) pairs every group with a key/value head.
602	law groups_len:
603	  for n: Nat
604	  for +g: Nat
605	  for +hs: List<&2, Head>
606	  {List.length(&2, List<&2, Head>, groups(n, g, hs)) == n : Nat}
607	
608	def groups_len(n, g, hs):
609	  match n:
610	    case 0n:
611	      {==}
612	    case 1n+p:
613	      %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}
614	      {==}
615	
616	# one block of 1 + m numbers in Nat.divmod: the counter m runs down to 0, then the quotient d grows by one
617	law div_block:
618	  for +m: Nat
619	  for -n: Nat
620	  for -d: Nat
621	  for -r: Nat
622	  {Nat.divmod.go(Nat.add(1n+m, n), m, d, r) == Nat.divmod.go(n, Nat.add(m, r), 1n+d, 0n) : Nat & Nat}
623	
624	def div_block(m, n, d, r):
625	  match m:
626	    case 0n:
627	      {==}
628	    case 1n+p:
629	      %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}
630	      div_block(p, n, d, 1n+r)
631	
632	# j whole blocks of 1 + s numbers add j to the quotient
633	law div_blocks:
634	  for +j: Nat
635	  for +s: Nat
636	  for +n: Nat
637	  for +d: Nat
638	  {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}
639	
640	def div_blocks(j, s, n, d):
641	  match j:
642	    case 0n:
643	      {==}
644	    case 1n+jp:
645	      %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}
646	      %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}
647	      %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}
648	      %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}
649	      div_blocks(jp, s, n, 1n+d)
650	
651	# fewer than 1 + s numbers after the last whole block leave the quotient alone
652	law div_rest:
653	  for +i: Nat
654	  for -k: Nat
655	  for -d: Nat
656	  for -r: Nat
657	  {Pair.fst(Nat, Nat, Nat.divmod.go(i, Nat.add(i, k), d, r)) == d : Nat}
658	
659	def div_rest(i, k, d, r):
660	  match i:
661	    case 0n:
662	      {==}
663	    case 1n+p:
664	      div_rest(p, k, d, 1n+r)
665	
666	# 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
667	# head i of group j, reads key/value head (j·g + i) / g = j. That division is how transformers maps query heads
668	# to key/value heads (repeat_kv), so grouping the heads computes the same attention.
669	law gqa_head:
670	  for +j: Nat
671	  for +i: Nat
672	  for +k: Nat
673	  {Nat.div(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}
674	
675	# PROOF: Nat.div counts how many whole blocks of g numbers fit. j·g is j whole blocks (div_blocks, which uses
676	# div_block for each block), and i < g numbers after them do not complete another one (div_rest).
677	def gqa_head(j, i, k):
678	  %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}
679	  %add_zero(j) : {Pair.fst(Nat, Nat, Nat.divmod.go(i, Nat.add(i, k), Nat.add(j, 0n), 0n)) == _ : Nat}
680	  div_rest(i, k, Nat.add(j, 0n), 0n)
681	
682	# =====================================================================================================
683	# Grouped-query attention over the arrays
684	# =====================================================================================================
685	
686	# the scores of one query head against the keys of its key/value head, and the keys handed back
687	type SR is Type:
688	  SR{k: Array<F32>, s: List<&2, F32>}
689	
690	def sr_len(sr: SR) -> Nat:
691	  match sr:
692	    case SR{k, s}:
693	      List.length(&2, F32, s)
694	
695	def sr_of(k: Array<F32>, al: TA.AL) -> SR:
696	  match al:
697	    case TA.AL{c, l}:
698	      SR{k, l}
699	
700	# LAW (read_len): reading m + 1 numbers of an array onto acc gives m + 1 more scores.
701	# PROOF: by induction on m. With 0 left it reads one number onto acc (1 + the length of acc, by the definitions);
702	# with 1 + p left it reads one and recurses one index down, and len_step moves the 1 out of the sum (tensor-array's
703	# leaf_len, for scores).
704	law read_len:
705	  for k: Array<F32>
706	  for +m: Nat
707	  for +idx: U32
708	  for +acc: List<&2, F32>
709	  for rc: Array<F32> & F32
710	  {sr_len(sr_of(k, TA.read_l(m, idx, acc, rc))) == Nat.add(1n+m, List.length(&2, F32, acc)) : Nat}
711	
712	def read_len(k, m, idx, acc, rc):
713	  match m rc:
714	    case 0n Tuple{d, +v}:
715	      {==}
716	    case 1n+p Tuple{d, +v}:
717	      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))))
718	
719	# read the 1 + t scores the gemm wrote (read_l reads from the last index down, so it starts at t)
720	def sr_read(+t: Nat, g: TA.G) -> SR:
721	  match g:
722	    case TA.G{qa, k, s}:
723	      sr_of(k, TA.read_l(t, U32.from_nat(t), Nil{}, Array.get(F32, s, U32.from_nat(t))))
724	
725	# q·k for the keys of positions 0..t, oldest first: K (t + 1 rows of `head` numbers) times q, the product
726	# Mat.matmul_nt computes for one row (bend-ml-tensor-array's gemm)
727	def scores(+hd: Nat, +t: Nat, k: Array<F32>, +q: List<&2, F32>) -> SR:
728	  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)))
729	
730	def sr_len_g(+t: Nat, g: TA.G) -> {sr_len(sr_read(t, g)) == 1n+t : Nat}:
731	  match g:
732	    case TA.G{qa, k, s}:
733	      TA.zero_out(sr_len(sr_of(k, TA.read_l(t, U32.from_nat(t), Nil{}, Array.get(F32, s, U32.from_nat(t))))), 1n+t, read_len(k, t, U32.from_nat(t), Nil{}, Array.get(F32, s, U32.from_nat(t))))
734	
735	# LAW (causal): attention at position t scores exactly t + 1 keys, the positions 0..t: one score per position
736	# written, none for a position not yet computed.
737	# PROOF: the gemm's result has one constructor, so matching it opens it without computing the product
738	# (sr_len_g); read_len counts the 1 + t numbers read back.
739	law causal:
740	  for +hd: Nat
741	  for +t: Nat
742	  for k: Array<F32>
743	  for +q: List<&2, F32>
744	  {sr_len(scores(hd, t, k, q)) == 1n+t : Nat}
745	
746	def causal(hd, t, k, q):
747	  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)))
748	
749	# the weighted sum of the values of positions 0..t: w (t + 1 weights) times V (t + 1 rows of `head` numbers)
750	type WR is Type:
751	  WR{v: Array<F32>, o: List<&2, F32>}
752	
753	def wr_of(v: Array<F32>, al: TA.AL) -> WR:
754	  match al:
755	    case TA.AL{c, l}:
756	      WR{v, l}
757	
758	def wr_read(+hd: Nat, g: TA.G) -> WR:
759	  match g:
760	    case TA.G{wa, v, o}:
761	      wr_of(v, TA.read_l(Nat.sub(hd, 1n), U32.from_nat(Nat.sub(hd, 1n)), Nil{}, Array.get(F32, o, U32.from_nat(Nat.sub(hd, 1n)))))
762	
763	def weighted(+hd: Nat, +t: Nat, v: Array<F32>, +w: List<&2, F32>) -> WR:
764	  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)))
765	
766	# one query head: softmax(q·K · scale) · V, and the two arrays handed back
767	type HO is Type:
768	  HO{k: Array<F32>, v: Array<F32>, o: List<&2, F32>}
769	
770	def HO.then(-R: Type, h: HO, then: Array<F32> -> Array<F32> -> List<&2, F32> -> R) -> R:
771	  match h:
772	    case HO{k, v, o}:
773	      then(k, v, o)
774	
775	def head_wr(k: Array<F32>, wr: WR) -> HO:
776	  match wr:
777	    case WR{v, o}:
778	      HO{k, v, o}
779	
780	def head_sr(+hd: Nat, +t: Nat, +scale: F32, v: Array<F32>, sr: SR) -> HO:
781	  match sr:
782	    case SR{k, s}:
783	      head_wr(k, weighted(hd, t, v, list.softmax(list.scale(scale, s))))
784	
785	def head_att(+hd: Nat, +t: Nat, +scale: F32, k: Array<F32>, v: Array<F32>, +q: List<&2, F32>) -> HO:
786	  head_sr(hd, t, scale, v, scores(hd, t, k, q))
787	
788	# the query heads of one group over its key/value head, in order; the arrays are threaded through
789	def group_att(hs: List<&2, Head>, +hd: Nat, +t: Nat, +scale: F32, k: Array<F32>, v: Array<F32>, +acc: List<&2, F32>) -> HO:
790	  match hs:
791	    case Nil{}:
792	      HO{k, v, acc}
793	    case Con{h, rest}:
794	      HO.then(HO, head_att(hd, t, scale, k, v, join(h)), k2 => v2 => +o =>
795	        # 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)
796	        group_att(rest, hd, t, scale, k2, v2, List.append(&2, F32, acc, o)))
797	
798	# every group, in order: the outputs of query heads 0, 1, ..., and the caches handed back
799	type AT is Type:
800	  AT{ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>, o: List<&2, F32>}
801	
802	def at_join(p: AT, q: AT) -> AT:
803	  match p q:
804	    case AT{k1, v1, o1} AT{k2, v2, o2}:
805	      AT{List.append(&1, Array<F32>, k1, k2), List.append(&1, Array<F32>, v1, v2), List.append(&2, F32, o1, o2)}
806	
807	# one key/value head's work: its group of query heads, its keys and its values
808	type GW is Type:
809	  GW{g: List<&2, Head>, k: Array<F32>, v: Array<F32>}
810	
811	def gws(gs: List<&2, List<&2, Head>>, ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>) -> List<&1, GW>:
812	  match gs ks vs:
813	    case Con{g, gt} Con{k, kt} Con{v, vt}:
814	      GW{g, k, v} <> gws(gt, kt, vt)
815	    case _ _ _:
816	      Nil{}
817	
818	# some groups one after another, in order
819	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:
820	  match ws:
821	    case Con{w, wt}:
822	      match w:
823	        case GW{g, k, v}:
824	          HO.then(AT, group_att(g, hd, t, scale, k, v, Nil{}), k2 => v2 => +o =>
825	            # 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)
826	            seq_att(wt, hd, t, scale, k2 <> aks, v2 <> avs, List.append(&2, F32, acc, o)))
827	    case Nil{}:
828	      AT{List.reverse(&1, Array<F32>, aks), List.reverse(&1, Array<F32>, avs), acc}
829	
830	# a list of groups cut after its first n (a tail loop: a helper in the fork tree must not hold the stack)
831	type GS is Type:
832	  GS{l: List<&1, GW>, r: List<&1, GW>}
833	
834	def cut(n: Nat, xs: List<&1, GW>, acc: List<&1, GW>) -> GS:
835	  match n xs:
836	    case 1n+m Con{x, xt}:
837	      cut(m, xt, x <> acc)
838	    case _ _:
839	      GS{List.reverse(&1, GW, acc), xs}
840	
841	# how many times n groups are halved (the larger half kept) until one or none is left
842	def halvings(fuel: Nat, +n: Nat) -> Nat:
843	  match fuel:
844	    case 0n:
845	      0n
846	    case 1n+f:
847	      match n:
848	        case 0n:
849	          0n
850	        case 1n+p:
851	          match p:
852	            case 0n:
853	              0n
854	            case 1n+q:
855	              1n+halvings(f, Nat.sub(2n+q, TA.half(2n+q)))
856	
857	# 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
858	# outputs and the caches come back in group order, so every head computes what it computed one after another
859	def par_att(d: Nat, gs: GS, +n: Nat, +hd: Nat, +t: Nat, +scale: F32) -> AT:
860	  match d gs:
861	    case 0n GS{l, r}:
862	      # 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)
863	      p q = seq_att(l, hd, t, scale, Nil{}, Nil{}, Nil{}) seq_att(r, hd, t, scale, Nil{}, Nil{}, Nil{})
864	      at_join(p, q)
865	    case 1n++e GS{l, r}:
866	      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)
867	      at_join(p, q)
868	
869	# the attention output and the cache handed back
870	type KA<-cap: Nat, -t: Nat> is Type:
871	  KA{kv: KV<cap, t>, o: List<&2, F32>}
872	
873	def ka_of(-cap: Nat, -t: Nat, at: AT) -> KA<cap, t>:
874	  match at:
875	    case AT{ks, vs, o}:
876	      KA{KV{ks, vs}, o}
877	
878	def KA.then(-R: Type, -cap: Nat, -t: Nat, ka: KA<cap, t>, then: KV<cap, t> -> List<&2, F32> -> R) -> R:
879	  match ka:
880	    case KA{kv, o}:
881	      then(kv, o)
882	
883	# attention of the (rotated) query heads at position t over the cache of positions 0..t (the newest included).
884	# d is the model's parallel depth (depth_for): with d = 0 (two threads or fewer) the key/value heads run one after
885	# another, since a fork tree on one thread is only overhead (+5% on the 1.7B, NOTES 2026-10-12); otherwise they run
886	# in parallel, one group per leaf. Both orders compute every head the same way, so the bits are the same.
887	def attention(+c: Cfg, +d: Nat, -cap: Nat, +t: Nat, qs: List<&2, Head>, kv: KV<cap, 1n+t>) -> KA<cap, 1n+t>:
888	  match d kv:
889	    case 0n KV{ks, vs}:
890	      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{}))
891	    case 1n+e KV{ks, vs}:
892	      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)))
893	
894	# =====================================================================================================
895	# A layer
896	# =====================================================================================================
897	
898	# the weights of one layer, with their shapes from the config (row = output, as PyTorch stores them)
899	type LW<-c: Cfg> is Type:
900	  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)>}
901	
902	# what a layer computes at one position, every intermediate: the tests compare them with transformers, and a
903	# backward pass would read them
904	type Act is Data:
905	  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>}
906	
907	# the layer's output at position t, its cache grown to t + 1 positions, its weights handed back, and the
908	# activations
909	type BR<-c: Cfg, -cap: Nat, -t: Nat> is Type:
910	  BR{y: List<&2, F32>, kv: KV<cap, 1n+t>, lw: LW<c>, act: Act}
911	
912	# continues with the fields of a layer's result (a computed value cannot be matched; a loop over positions
913	# recurses inside k)
914	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:
915	  match b:
916	    case BR{y, kv, lw, act}:
917	      k(y, kv, lw, act)
918	
919	def mv.go(-B: Type, -r: Nat, -k: Nat, bl: TA.BL<r, k>, then: TA.Bands<r, k> -> List<&2, F32> -> B) -> B:
920	  match bl:
921	    case TA.BL{w, y}:
922	      then(w, y)
923	
924	# y = W · x, then `then` with W handed back and y (an Array has one owner, so a product returns its matrix)
925	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:
926	  mv.go(B, r, k, TA.Bands.matvec_l(r, k, w, x), then)
927	
928	# 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
929	# parallel depth of attention (see attention)
930	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>:
931	  match lw:
932	    case LW{+ln1, wq, wk, wv, wo, +ln2, wg, wu, wd}:
933	      +n1 = rmsnorm(c, x, ln1)
934	      mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wq, n1, wq => +q =>
935	        mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wk, n1, wk => +k =>
936	          mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wv, n1, wv => +v =>
937	            +qs = ropes(rt, heads(Cfg.heads(c), Cfg.half(c), q))
938	            +ks = ropes(rt, heads(Cfg.kv(c), Cfg.half(c), k))
939	            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 =>
940	            mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wo, at, wo => +a =>
941	              +x1 = list.add(x, a)
942	              +n2 = rmsnorm(c, x1, ln2)
943	              mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wg, n2, wg => +g =>
944	                mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wu, n2, wu => +u =>
945	                  mv(BR<c, cap, t>, Cfg.dim(c), Cfg.ffn(c), wd, swiglu(g, u), wd => +m =>
946	                    +y = list.add(x1, m)
947	                    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}}))))))))
948	
949	# block for generation: the same computation, with the activations left empty. Generation never reads them, and
950	# building them (the copies of the rotated heads among them) cost about a fifth of each token on 16 threads (found by
951	# Mathews Fernando's bend-lint, unused_result). It is a copy of block's body, not a flag in it: a flag that chooses
952	# after the computation keeps every intermediate alive to the end, and that measured slower than building them.
953	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>:
954	  match lw:
955	    case LW{+ln1, wq, wk, wv, wo, +ln2, wg, wu, wd}:
956	      +n1 = rmsnorm(c, x, ln1)
957	      mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wq, n1, wq => +q =>
958	        mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wk, n1, wk => +k =>
959	          mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wv, n1, wv => +v =>
960	            +qs = ropes(rt, heads(Cfg.heads(c), Cfg.half(c), q))
961	            +ks = ropes(rt, heads(Cfg.kv(c), Cfg.half(c), k))
962	            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 =>
963	            mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wo, at, wo => +a =>
964	              +x1 = list.add(x, a)
965	              +n2 = rmsnorm(c, x1, ln2)
966	              mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wg, n2, wg => +g =>
967	                mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wu, n2, wu => +u =>
968	                  mv(BR<c, cap, t>, Cfg.dim(c), Cfg.ffn(c), wd, swiglu(g, u), wd => +m =>
969	                    +y = list.add(x1, m)
970	                    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{}}}))))))))
971	
972	# keep: block with the activations (the tests read them), or block_gen without them; chosen before anything is computed
973	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>:
974	  match keep:
975	    case True{}:
976	      block(c, d, lw, cap, t, kv, room, rt, x)
977	    case False{}:
978	      block_gen(c, d, lw, cap, t, kv, room, rt, x)
979	
980	# =====================================================================================================
981	# The model
982	# =====================================================================================================
983	
984	# emb: the embedding (vocab rows; also the LM head, tied); lnf: the final norm's gains
985	type Model<-c: Cfg> is Type:
986	  Model{d: Nat, emb: TA.Bands<Cfg.vocab(c), Cfg.dim(c)>, layers: List<&1, LW<c>>, lnf: List<&2, F32>}
987	
988	# the last layer's output, the layers and caches handed back, and every layer's activations
989	type Run<-c: Cfg, -cap: Nat, -t: Nat> is Type:
990	  Run{x: List<&2, F32>, layers: List<&1, LW<c>>, cs: List<&1, KV<cap, 1n+t>>, acts: List<&2, Act>}
991	
992	# pend is the previous layer's result: its output is the next layer's input
993	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>:
994	  match ls cs pend:
995	    case Con{l, lt} Con{k, kt} BR{+y, kv, l0, act}:
996	      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))
997	    case _ _ BR{+y, kv, l0, act}:
998	      Run{y, List.reverse(&1, LW<c>, l0 <> al), List.reverse(&1, KV<cap, 1n+t>, kv <> ac), List.reverse(&2, Act, act <> aa)}
999	
1000	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>:
1001	  match ls cs:
1002	    case Con{l, lt} Con{k, kt}:
1003	      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))
1004	    case _ _:
1005	      Run{x, Nil{}, Nil{}, Nil{}}
1006	
1007	# what one position produces: the logits, the final norm, every layer's activations, the caches grown to t + 1
1008	# positions and the model handed back
1009	type Out<-c: Cfg, -cap: Nat, -t: Nat> is Type:
1010	  Out{lg: List<&2, F32>, nf: List<&2, F32>, acts: List<&2, Act>, cs: List<&1, KV<cap, 1n+t>>, m: Model<c>}
1011	
1012	# continues with the fields of a position's result: logits, final norm, activations, caches, model
1013	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:
1014	  match o:
1015	    case Out{lg, nf, acts, cs, m}:
1016	      k(lg, nf, acts, cs, m)
1017	
1018	# the final norm, then logits = emb · nf (the LM head is the embedding matrix)
1019	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>:
1020	  match r:
1021	    case Run{+x, ls, cs, acts}:
1022	      +nf = rmsnorm(c, x, lnf)
1023	      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}})
1024	
1025	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>:
1026	  match r:
1027	    case TA.BRow{emb, row}:
1028	      head_of(c, cap, t, d, emb, lnf, layers(keep, c, d, cap, t, room, rot_at(c, t), ls, cs, row))
1029	
1030	# token tok at position t (0-based), with the caches of positions 0..t-1 (one per layer) and room for cap positions;
1031	# room is the proof that t < cap; keep: whether the result carries every layer's activations
1032	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>:
1033	  match m:
1034	    case Model{+d, emb, ls, +lnf}:
1035	      forward.go(keep, c, cap, t, room, d, ls, lnf, cs, TA.Bands.read_row(Cfg.vocab(c), Cfg.dim(c), emb, tok))
1036	
1037	# a position with every layer's activations (the tests read them)
1038	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>:
1039	  forward_k(True{}, c, cap, t, m, cs, room, tok)
1040	
1041	# a position for generation: the same logits, caches and model. With bands (d > 0, more than two threads) the
1042	# activations are left empty, about a fifth faster per token on 16 threads; with one band (d = 0) they are kept, since
1043	# leaving them out measured 4-11% slower per token on one thread (llama.bend NOTES, 2026-10-16), for no reason found yet
1044	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>:
1045	  match m:
1046	    case Model{+d, emb, ls, +lnf}:
1047	      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))
1048	
1049	# a prompt position, whose logits nobody reads: the layers and the caches as forward_gen gives them, without the final
1050	# norm and the product by the embedding (vocab x dim, the largest of a position), so lg and nf are empty
1051	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>:
1052	  match r:
1053	    case Run{x, ls, cs, acts}:
1054	      Out{Nil{}, Nil{}, acts, cs, Model{d, emb, ls, lnf}}
1055	
1056	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>:
1057	  match r:
1058	    case TA.BRow{emb, row}:
1059	      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))
1060	
1061	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>:
1062	  match m:
1063	    case Model{+d, emb, ls, +lnf}:
1064	      forward_pre.go(c, cap, t, room, d, ls, lnf, cs, TA.Bands.read_row(Cfg.vocab(c), Cfg.dim(c), emb, tok))
1065	
1066	# n empty caches, one per layer, each with room for cap positions
1067	def caches(+c: Cfg, n: Nat, +cap: Nat) -> List<&1, KV<cap, 0n>>:
1068	  match n:
1069	    case 0n:
1070	      Nil{}
1071	    case 1n+p:
1072	      KV.empty(c, cap) <> caches(c, p, cap)
1073	
1074	# the first index of the largest value (torch.argmax does the same), and the value
1075	type AM is Data:
1076	  AM{i: Nat, v: F32}
1077	
1078	def AM.i(am: AM) -> Nat:
1079	  match am:
1080	    case AM{i, v}:
1081	      i
1082	
1083	def AM.v(am: AM) -> F32:
1084	  match am:
1085	    case AM{i, v}:
1086	      v
1087	
1088	def gt_head(xs: List<&2, F32>, +v: F32) -> Bool:
1089	  match xs:
1090	    case Nil{}:
1091	      False{}
1092	    case Con{+h, t}:
1093	      F32.is_gt(h, v)
1094	
1095	def argmax.go(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> AM:
1096	  match xs g:
1097	    case Nil{} _:
1098	      AM{bi, bv}
1099	    case Con{+h, +t} True{}:
1100	      argmax.go(t, gt_head(t, h), 1n+i, i, h)
1101	    case Con{h, +t} False{}:
1102	      argmax.go(t, gt_head(t, bv), 1n+i, bi, bv)
1103	
1104	def argmax(xs: List<&2, F32>) -> AM:
1105	  match xs:
1106	    case Nil{}:
1107	      AM{0n, 0.0}
1108	    case Con{+h, +t}:
1109	      argmax.go(t, gt_head(t, h), 1n, 0n, h)
1110	
1111	# =====================================================================================================
1112	# Sampling: temperature, top-k and top-p, with a seeded generator written in Bend
1113	# =====================================================================================================
1114	# Every step is an F32 operation in a fixed order, so reference/sample_ref.py repeats it bit for bit from the same
1115	# logits: the k largest logits (ties: the lower id first, as a stable sort), each divided by the temperature; their
1116	# weights exp(v/temp - top/temp) (top: the largest), summed in order into Z; top-p keeps each candidate whose
1117	# cumulative weight before it is below p·Z (the first is always kept), Zp their sum; a uniform u in [0, 1) picks the
1118	# first kept candidate whose cumulative weight is above u·Zp (the last kept one if rounding leaves none).
1119	
1120	# Chris Wellons' lowbias32 hash (public domain): xor-shifts and two odd multipliers, a bijection on U32
1121	def hash32(+x: U32) -> U32:
1122	  +a = U32.xor(x, U32.shrn(x, 16n))
1123	  +b = U32.mul(a, 2146121005)
1124	  +c = U32.xor(b, U32.shrn(b, 15n))
1125	  +d = U32.mul(c, 2221713035)
1126	  U32.xor(d, U32.shrn(d, 16n))
1127	
1128	# draw n of a seed, a uniform F32 in [0, 1): the top 24 bits of hash32(hash32(seed) + n), times 2^-24 (exact). Nothing
1129	# is kept between draws, so a seed gives the same draws whatever ran before.
1130	def uniform(+seed: U32, +n: U32) -> F32:
1131	  F32.mul(U32.to_f32(U32.shrn(hash32(U32.add(hash32(seed), n)), 8n)), 5.9604644775390625e-8)
1132	
1133	# a candidate: a logit and its id
1134	type Cand is Data:
1135	  Cand{v: F32, i: Nat}
1136	
1137	def Cand.i(c: Cand) -> Nat:
1138	  match c:
1139	    case Cand{v, i}:
1140	      i
1141	
1142	def Cand.v(c: Cand) -> F32:
1143	  match c:
1144	    case Cand{v, i}:
1145	      v
1146	
1147	# whether a value v goes before the head of cs: the head is smaller (an equal one keeps its place, ahead of v)
1148	def goes_before(cs: List<&2, Cand>, +v: F32) -> Bool:
1149	  match cs:
1150	    case Nil{}:
1151	      True{}
1152	    case Con{Cand{hv, hi}, t}:
1153	      F32.is_lt(hv, v)
1154	
1155	# v (id i) into the candidates cs, largest first; g: whether it goes before the head
1156	def ins(cs: List<&2, Cand>, g: Bool, +v: F32, +i: Nat) -> List<&2, Cand>:
1157	  match cs g:
1158	    case Nil{} _:
1159	      Cand{v, i} <> Nil{}
1160	    case Con{h, t} True{}:
1161	      Cand{v, i} <> (h <> t)
1162	    case Con{h, +t} False{}:
1163	      h <> ins(t, goes_before(t, v), v, i)
1164	
1165	# the value of the last candidate (lo if there is none)
1166	def last_v(cs: List<&2, Cand>, +lo: F32) -> F32:
1167	  match cs:
1168	    case Nil{}:
1169	      lo
1170	    case Con{Cand{v, i}, t}:
1171	      last_v(t, v)
1172	
1173	# whether the head of xs enters the candidates: there is room, or it is above the smallest one (lo)
1174	def enters(xs: List<&2, F32>, room: Bool, +lo: F32) -> Bool:
1175	  match xs:
1176	    case Nil{}:
1177	      False{}
1178	    case Con{+h, t}:
1179	      Bool.or(room, F32.is_gt(h, lo))
1180	
1181	# the k largest of xs (the values from id i on) and of the candidates cs, largest first; cs holds m <= k of them,
1182	# lo the smallest; g: whether the head of xs enters. One pass, and a value at or below lo costs one comparison.
1183	def topk.go(xs: List<&2, F32>, g: Bool, +k: Nat, +i: Nat, +m: Nat, +lo: F32, +cs: List<&2, Cand>) -> List<&2, Cand>:
1184	  match xs g:
1185	    case Nil{} _:
1186	      cs
1187	    case Con{h, +t} False{}:
1188	      topk.go(t, enters(t, Nat.is_lt(m, k), lo), k, 1n+i, m, lo, cs)
1189	    case Con{+h, +t} True{}:
1190	      +cs2 = List.take(&2, Cand, ins(cs, goes_before(cs, h), h, i), k)
1191	      +m2 = Nat.min(1n+m, k)
1192	      +lo2 = last_v(cs2, h)
1193	      topk.go(t, enters(t, Nat.is_lt(m2, k), lo2), k, 1n+i, m2, lo2, cs2)
1194	
1195	def topk(+k: Nat, +xs: List<&2, F32>) -> List<&2, Cand>:
1196	  topk.go(xs, enters(xs, Nat.is_lt(0n, k), 0.0), k, 0n, 0n, 0.0, Nil{})
1197	
1198	# the weight of a logit v: exp(v/temp - top/temp), top the largest logit (its weight is 1)
1199	def weight(+v: F32, +temp: F32, +top: F32) -> F32:
1200	  F32.exp(F32.sub(F32.div(v, temp), F32.div(top, temp)))
1201	
1202	# acc plus the weights of cs, in order
1203	def total(cs: List<&2, Cand>, +temp: F32, +top: F32, +acc: F32) -> F32:
1204	  match cs:
1205	    case Nil{}:
1206	      acc
1207	    case Con{Cand{+v, i}, t}:
1208	      total(t, temp, top, F32.add(acc, weight(v, temp, top)))
1209	
1210	# how many candidates top-p keeps (n) and their weight (z)
1211	type Kept is Data:
1212	  Kept{n: Nat, z: F32}
1213	
1214	# g: whether the head of cs is kept (the weight acc before it is below lim = p·Z; the first is always kept)
1215	def kept.go(cs: List<&2, Cand>, g: Bool, +temp: F32, +top: F32, +lim: F32, +n: Nat, +acc: F32) -> Kept:
1216	  match cs g:
1217	    case Nil{} _:
1218	      Kept{n, acc}
1219	    case Con{h, t} False{}:
1220	      Kept{n, acc}
1221	    case Con{Cand{+v, i}, t} True{}:
1222	      +acc2 = F32.add(acc, weight(v, temp, top))
1223	      kept.go(t, F32.is_lt(acc2, lim), temp, top, lim, 1n+n, acc2)
1224	
1225	# whether the draw stops at the head of cs: it is the last kept one (q = 0 kept after it), or the cumulative
1226	# weight acc plus its own is above the target u·Zp
1227	def stops(cs: List<&2, Cand>, +q: Nat, +temp: F32, +top: F32, +target: F32, +acc: F32) -> Bool:
1228	  match cs:
1229	    case Nil{}:
1230	      True{}
1231	    case Con{Cand{+v, i}, t}:
1232	      Bool.or(Nat.is_eq(q, 0n), F32.is_lt(target, F32.add(acc, weight(v, temp, top))))
1233	
1234	# the candidate drawn; s: whether it is the head of cs
1235	def pick.go(cs: List<&2, Cand>, s: Bool, +q: Nat, +temp: F32, +top: F32, +target: F32, +acc: F32) -> Cand:
1236	  match cs s:
1237	    case Nil{} _:
1238	      Cand{0.0, 0n}
1239	    case Con{h, t} True{}:
1240	      h
1241	    case Con{Cand{+v, i}, +t} False{}:
1242	      +acc2 = F32.add(acc, weight(v, temp, top))
1243	      pick.go(t, stops(t, Nat.sub(q, 1n), temp, top, target, acc2), Nat.sub(q, 1n), temp, top, target, acc2)
1244	
1245	def sample.pick(+cs: List<&2, Cand>, kp: Kept, +temp: F32, +top: F32, +u: F32) -> Cand:
1246	  match kp:
1247	    case Kept{+n, +z}:
1248	      pick.go(cs, stops(cs, Nat.sub(n, 1n), temp, top, F32.mul(u, z), 0.0), Nat.sub(n, 1n), temp, top, F32.mul(u, z), 0.0)
1249	
1250	def sample.cs(+cs: List<&2, Cand>, +temp: F32, +top: F32, +p: F32, +u: F32) -> Cand:
1251	  sample.pick(cs, kept.go(cs, True{}, temp, top, F32.mul(p, total(cs, temp, top, 0.0)), 0n, 0.0), temp, top, u)
1252	
1253	def top_of(cs: List<&2, Cand>) -> F32:
1254	  match cs:
1255	    case Nil{}:
1256	      0.0
1257	    case Con{Cand{v, i}, t}:
1258	      v
1259	
1260	# a token drawn from the logits xs (its id and logit) with temperature temp > 0, the k > 0 largest logits and top-p
1261	# p, by the uniform draw u
1262	def sample(+xs: List<&2, F32>, +temp: F32, +k: Nat, +p: F32, +u: F32) -> Cand:
1263	  +cs = topk(k, xs)
1264	  sample.cs(cs, temp, top_of(cs), p, u)
1265	
1266	# ---- LAW (sample_below): the drawn id is an id of the logits ----
1267	
1268	# every candidate's id is below n
1269	def below(cs: List<&2, Cand>, +n: Nat) -> Bool:
1270	  match cs:
1271	    case Nil{}:
1272	      True{}
1273	    case Con{Cand{v, i}, t}:
1274	      # bend-lint: disable-next bend-lint-rules/performance/eager_branch -- below only appears in the laws' types; no program runs it not-expected-because: proofs only, nothing to measure
1275	      Bool.and(TA.leb(1n+i, n), below(t, n))
1276	
1277	# the two halves of a true Bool.and
1278	def and_l(+a: Bool, -b: Bool, e: {True{} == Bool.and(a, b) : Bool}) -> {True{} == a : Bool}:
1279	  match a:
1280	    case True{}:
1281	      {==}
1282	    case False{}:
1283	      e
1284	
1285	def and_r(+a: Bool, +b: Bool, e: {True{} == Bool.and(a, b) : Bool}) -> {True{} == b : Bool}:
1286	  match a:
1287	    case True{}:
1288	      e
1289	    case False{}:
1290	      %e : TA.BD(_, Unit, {True{} == b : Bool})
1291	      Unit{}
1292	
1293	# PROOF: by induction on cs. Inserting puts the new candidate (id i < n) at the front, or keeps the head and
1294	# inserts into the tail, whose ids are below n by the induction hypothesis.
1295	law ins_below:
1296	  for +cs: List<&2, Cand>
1297	  for +g: Bool
1298	  for +v: F32
1299	  for +i: Nat
1300	  for +n: Nat
1301	  for +e: {True{} == below(cs, n) : Bool}
1302	  for ei: {True{} == TA.leb(1n+i, n) : Bool}
1303	  {True{} == below(ins(cs, g, v, i), n) : Bool}
1304	
1305	def ins_below(cs, g, v, i, n, e, ei):
1306	  match cs g:
1307	    case Nil{} _:
1308	      %ei : {True{} == Bool.and(_, True{}) : Bool}
1309	      {==}
1310	    case Con{h, t} True{}:
1311	      %ei : {True{} == Bool.and(_, below(h <> t, n)) : Bool}
1312	      e
1313	    case Con{Cand{hv, hi}, +t} False{}:
1314	      %and_l(TA.leb(1n+hi, n), below(t, n), e) : {True{} == Bool.and(_, below(ins(t, goes_before(t, v), v, i), n)) : Bool}
1315	      ins_below(t, goes_before(t, v), v, i, n, and_r(TA.leb(1n+hi, n), below(t, n), e), ei)
1316	
1317	# PROOF: taking the first k candidates keeps a part of them.
1318	law take_below:
1319	  for +cs: List<&2, Cand>
1320	  for +k: Nat
1321	  for +n: Nat
1322	  for +e: {True{} == below(cs, n) : Bool}
1323	  {True{} == below(List.take(&2, Cand, cs, k), n) : Bool}
1324	
1325	def take_below(cs, k, n, e):
1326	  match cs k:
1327	    case Nil{} _:
1328	      {==}
1329	    case Con{h, t} 0n:
1330	      {==}
1331	    case Con{Cand{v, i}, +t} 1n+p:
1332	      %and_l(TA.leb(1n+i, n), below(t, n), e) : {True{} == Bool.and(_, below(List.take(&2, Cand, t, p), n)) : Bool}
1333	      take_below(t, p, n, and_r(TA.leb(1n+i, n), below(t, n), e))
1334	
1335	# PROOF: by induction on xs, with i + length(xs) = n: the head of xs has id i < n (step_room), and after it
1336	# 1 + i + length(tail) = n still holds (step_inv); ins_below and take_below keep the candidates' ids below n.
1337	law topk_below:
1338	  for +xs: List<&2, F32>
1339	  for +g: Bool
1340	  for +k: Nat
1341	  for +i: Nat
1342	  for +m: Nat
1343	  for +lo: F32
1344	  for +cs: List<&2, Cand>
1345	  for +n: Nat
1346	  for +inv: {Nat.add(i, List.length(&2, F32, xs)) == n : Nat}
1347	  for +e: {True{} == below(cs, n) : Bool}
1348	  {True{} == below(topk.go(xs, g, k, i, m, lo, cs), n) : Bool}
1349	
1350	def topk_below(xs, g, k, i, m, lo, cs, n, inv, e):
1351	  match xs g:
1352	    case Nil{} _:
1353	      e
1354	    case Con{h, +t} False{}:
1355	      topk_below(t, enters(t, Nat.is_lt(m, k), lo), k, 1n+i, m, lo, cs, n, step_inv(i, List.length(&2, F32, t), n, inv), e)
1356	    case Con{+h, +t} True{}:
1357	      +cs2 = List.take(&2, Cand, ins(cs, goes_before(cs, h), h, i), k)
1358	      +m2 = Nat.min(1n+m, k)
1359	      +lo2 = last_v(cs2, h)
1360	      topk_below(t, enters(t, Nat.is_lt(m2, k), lo2), k, 1n+i, m2, lo2, cs2, n, step_inv(i, List.length(&2, F32, t), n, inv), take_below(ins(cs, goes_before(cs, h), h, i), k, n, ins_below(cs, goes_before(cs, h), h, i, n, e, step_room(i, List.length(&2, F32, t), n, inv))))
1361	
1362	# PROOF: by induction on cs: the draw returns a candidate of cs, whose id is below n, or id 0 < n when cs is empty.
1363	law pick_below:
1364	  for +cs: List<&2, Cand>
1365	  for +s: Bool
1366	  for +q: Nat
1367	  for +temp: F32
1368	  for +top: F32
1369	  for +target: F32
1370	  for +acc: F32
1371	  for +n: Nat
1372	  for +e: {True{} == below(cs, n) : Bool}
1373	  for pos: {True{} == TA.leb(1n, n) : Bool}
1374	  {True{} == TA.leb(1n+Cand.i(pick.go(cs, s, q, temp, top, target, acc)), n) : Bool}
1375	
1376	def pick_below(cs, s, q, temp, top, target, acc, n, e, pos):
1377	  match cs s:
1378	    case Nil{} _:
1379	      pos
1380	    case Con{Cand{v, i}, t} True{}:
1381	      and_l(TA.leb(1n+i, n), below(t, n), e)
1382	    case Con{Cand{+v, +i}, +t} False{}:
1383	      +acc2 = F32.add(acc, weight(v, temp, top))
1384	      pick_below(t, stops(t, Nat.sub(q, 1n), temp, top, target, acc2), Nat.sub(q, 1n), temp, top, target, acc2, n, and_r(TA.leb(1n+i, n), below(t, n), e), pos)
1385	
1386	def sample.pick_below(+cs: List<&2, Cand>, kp: Kept, +temp: F32, +top: F32, +u: F32, +n: Nat, e: {True{} == below(cs, n) : Bool}, pos: {True{} == TA.leb(1n, n) : Bool}) -> {True{} == TA.leb(1n+Cand.i(sample.pick(cs, kp, temp, top, u)), n) : Bool}:
1387	  match kp:
1388	    case Kept{+c, +z}:
1389	      pick_below(cs, stops(cs, Nat.sub(c, 1n), temp, top, F32.mul(u, z), 0.0), Nat.sub(c, 1n), temp, top, F32.mul(u, z), 0.0, n, e, pos)
1390	
1391	# LAW (sample_below): from logits xs with at least one value, the id sample draws is below the length of xs: an id
1392	# of the vocabulary, whatever the temperature, k, p and draw.
1393	# PROOF: the top-k pass only stores ids of xs, each below its length (topk_below, from i = 0 and no candidates),
1394	# and the draw returns one of them (pick_below, through the top-p count, whatever it is).
1395	law sample_below:
1396	  for +xs: List<&2, F32>
1397	  for +temp: F32
1398	  for +k: Nat
1399	  for +p: F32
1400	  for +u: F32
1401	  for pos: {True{} == TA.leb(1n, List.length(&2, F32, xs)) : Bool}
1402	  {True{} == TA.leb(1n+Cand.i(sample(xs, temp, k, p, u)), List.length(&2, F32, xs)) : Bool}
1403	
1404	def sample_below(xs, temp, k, p, u, pos):
1405	  +cs = topk(k, xs)
1406	  +n = List.length(&2, F32, xs)
1407	  +top = top_of(cs)
1408	  sample.pick_below(cs, kept.go(cs, True{}, temp, top, F32.mul(p, total(cs, temp, top, 0.0)), 0n, 0.0), temp, top, u, n, topk_below(xs, enters(xs, Nat.is_lt(0n, k), 0.0), k, 0n, 0n, 0.0, Nil{}, n, {==}, {==}), pos)
1409	
1410	# =====================================================================================================
1411	# Loading, in plain Bend: every tensor read from the downloaded <dir>/hf/model.safetensors itself, BF16 at the
1412	# byte offsets that llama.bend's reference/prep.py writes into <dir>/w/st.txt, each number widened exactly to F32
1413	# (BF16 is the upper half of an F32); the config in <dir>/cfg.txt. llama/load.bend reads the same numbers with
1414	# foreign code, faster; this one uses File.read_at only, so --verdict checks it. The block
1415	# reader is bend-ml's (MIT).
1416	# =====================================================================================================
1417	
1418	# 4 little-endian bytes -> F32 (U32 and F32 hold the same Word(32n): it is a bitcast)
1419	def bitcast(u: U32) -> F32:
1420	  match u:
1421	    case U32{w}:
1422	      F32{w}
1423	
1424	def f32_le(+b0: U32, +b1: U32, +b2: U32, +b3: U32) -> F32:
1425	  bitcast((b0 + (b1 * 256 : U32) + (b2 * 65536 : U32) + (b3 * 16777216 : U32) : U32))
1426	
1427	def f32s(bs: List<&2, U32>) -> List<&2, F32>:
1428	  match bs:
1429	    case Con{b0, Con{b1, Con{b2, Con{b3, t}}}}:
1430	      f32_le(b0, b1, b2, b3) <> f32s(t)
1431	    case _:
1432	      Nil{}
1433	
1434	# reads the file in blocks of `blk` bytes; k blocks remaining; m is the block already read
1435	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>):
1436	  match k m:
1437	    case 0n Tuple{f, r}:
1438	      do IO<List<&2, F32>>:
1439	        File.close(f)
1440	        return List.concat(&2, F32, List.reverse(&2, List<&2, F32>, acc))
1441	    case 1n+p Tuple{f, r}:
1442	      do IO<List<&2, F32>>:
1443	        bs : List<&2, U32> <- IO.pass(List<&2, U32>, r)
1444	        m2 : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, (off + blk : U32), blk)
1445	        rf_loop(p, (off + blk : U32), blk, m2, f32s(bs) <> acc)
1446	
1447	def blocks_of(n: Nat, +per: Nat) -> Nat:
1448	  Nat.div(Nat.add(n, Nat.sub(per, 1n)), per)
1449	
1450	# n F32 numbers (little-endian) from a file (the tests' inputs; the weights are BF16, below)
1451	def read_floats(path: String, +n: Nat) -> IO(List<&2, F32>):
1452	  do IO<List<&2, F32>>:
1453	    f : File <- IO.try(File, File.open(path, "r"))
1454	    m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, 0, 1048576)
1455	    fl : List<&2, F32> <- rf_loop(blocks_of(Nat.mul(n, 4n), 1048576n), 0, 1048576, m, Nil{})
1456	    return List.take(&2, F32, fl, n)
1457	
1458	# the weight matrices are kept in 2^d row bands (TA.Bands); each matrix · vector product runs its bands in
1459	# parallel without copying the weights (bend-ml NOTES exp. 11). With 2 threads or fewer the bands cannot help,
1460	# so d = 0: a single band, which computes what the sequential Mat.matmul_nt computes (bend-ml exp. 14).
1461	def depth_sel(few: Bool) -> Nat:
1462	  match few:
1463	    case True{}:
1464	      0n
1465	    case False{}:
1466	      4n
1467	
1468	def depth_for(+threads: U32) -> Nat:
1469	  depth_sel(U32.is_le(threads, 2))
1470	
1471	# ---- BF16 tensors: matrices kept in bands, each band read from its own range of the file ----
1472	
1473	# 2 little-endian bytes of a BF16 number -> the F32 whose upper 16 bits they are (exact)
1474	def bf16_le(+b0: U32, +b1: U32) -> F32:
1475	  bitcast(((b0 * 65536 : U32) + (b1 * 16777216 : U32) : U32))
1476	
1477	def bf16s(bs: List<&2, U32>) -> List<&2, F32>:
1478	  match bs:
1479	    case Con{b0, Con{b1, t}}:
1480	      bf16_le(b0, b1) <> bf16s(t)
1481	    case _:
1482	      Nil{}
1483	
1484	# writes the BF16 numbers held in bs (2 bytes each), widened to F32, at index i, without building a list of F32
1485	# first
1486	def fill_bf16(bs: List<&2, U32>, +i: U32, a: Array<F32>) -> Array<F32>:
1487	  match bs:
1488	    case Con{+b0, Con{+b1, t}}:
1489	      fill_bf16(t, (i + 1 : U32), Array.set(F32, a, i, bf16_le(b0, b1)))
1490	    case _:
1491	      a
1492	
1493	def fill_bf16_mat(-r: Nat, -c: Nat, +i: U32, bs: List<&2, U32>, m: TA.Mat<r, c>) -> TA.Mat<r, c>:
1494	  match m:
1495	    case TA.Mat{d}:
1496	      TA.Mat{fill_bf16(bs, i, d)}
1497	
1498	# reads blocks of at most 1 MB (2^19 numbers) until rem bytes have been read, so a band never takes numbers that
1499	# belong to the next one
1500	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>):
1501	  match k m:
1502	    case 0n Tuple{f, rs}:
1503	      do IO<TA.Mat<r, c>>:
1504	        File.close(f)
1505	        return acc
1506	    case 1n+p Tuple{f, rs}:
1507	      do IO<TA.Mat<r, c>>:
1508	        bs : List<&2, U32> <- IO.pass(List<&2, U32>, rs)
1509	        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))))
1510	        stream_exact(p, r, c, Nat.sub(rem, 1048576n), (off + 1048576 : U32), (idx + 524288 : U32), m2, fill_bf16_mat(r, c, idx, bs, acc))
1511	
1512	# rr x c numbers from number `start` of the tensor at byte lo, into a band of r rows (rr == r, see load_tree; the
1513	# Array is sized for rr * c numbers as Mat.zeros would size it)
1514	def get_band(path: String, -r: Nat, +rr: Nat, +c: Nat, +lo: U32, +start: Nat) -> IO(TA.Mat<r, c>):
1515	  do IO<TA.Mat<r, c>>:
1516	    +off : U32 = (lo + U32.from_nat(Nat.mul(start, 2n)) : U32)
1517	    f : File <- IO.try(File, File.open(path, "r"))
1518	    m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, off, U32.from_nat(Nat.min(1048576n, Nat.mul(Nat.mul(rr, c), 2n))))
1519	    stream_exact(blocks_of(Nat.mul(Nat.mul(rr, c), 2n), 1048576n), r, c, Nat.mul(Nat.mul(rr, c), 2n), off, 0, m, TA.Mat{Array.new(F32, TA.cap_depth(Nat.mul(rr, c)), 0.0)})
1520	
1521	# fills a tree of 2^d bands (TA.Bands.zeros's shape: each node halves its rows with TA.half) in row order; rr is r at
1522	# run time. The tree is walked by its depth, not built first with TA.Bands.zeros and matched: that tree's zero-filled
1523	# arrays were dropped unread (NOTES.md, 2026-11-01)
1524	def load_tree(+c: Nat, -r: Nat, +d: Nat, +rr: Nat, -e: {rr == r : Nat}, +path: String, +lo: U32, +start: Nat) -> IO(TA.Bands<r, c>):
1525	  match d:
1526	    case 0n:
1527	      do IO<TA.Bands<r, c>>:
1528	        band : TA.Mat<r, c> <- get_band(path, r, rr, c, lo, start)
1529	        return TA.BLeaf{band}
1530	    case 1n+q:
1531	      do IO<TA.Bands<r, c>>:
1532	        x2 : TA.Bands<TA.half(r), c> <- load_tree(c, TA.half(r), q, TA.half(rr), TA.eq_half(rr, r, e), path, lo, start)
1533	        y2 : TA.Bands<Nat.sub(r, TA.half(r)), c> <- load_tree(c, Nat.sub(r, TA.half(r)), q, Nat.sub(rr, TA.half(rr)), TA.eq_rest(rr, r, e), path, lo, Nat.add(start, Nat.mul(TA.half(rr), c)))
1534	        return TA.BNode{x2, y2}
1535	
1536	# ---- the boundary: every tensor must lie inside the file, where File.read_at's U32 offset reaches ----
1537	
1538	def span_cmp(ok: Bool, +path: String, +t: Nat) -> IO(Unit):
1539	  match ok:
1540	    case True{}:
1541	      IO.pure(Unit, Unit{})
1542	    case False{}:
1543	      IO.die(Unit, 1, path ++ ": tensor " ++ Nat.show(t) ++ " of w/st.txt ends past the end of the file or starts beyond 4 GiB (File.read_at takes a U32 offset); run reference/prep.py --index")
1544	
1545	def span_ok(m: File & Result<&1, &1, U32 & String, U32>, +hi: U32, +lo: U32, +end: U32, +path: String, +t: Nat) -> IO(Unit):
1546	  match m:
1547	    case Tuple{f, r}:
1548	      do IO<Unit>:
1549	        +n : U32 <- IO.pass(U32, r)
1550	        File.close(f)
1551	        span_cmp(Bool.and(U32.is_eq(hi, 0), Bool.and(U32.is_le(lo, end), U32.is_le(end, n))), path, t)
1552	
1553	# stops the program unless tensor t, n BF16 numbers from byte hi * 2^32 + lo, ends inside the file with hi = 0
1554	def check_span(+path: String, +hi: U32, +lo: U32, +n: Nat, +t: Nat) -> IO(Unit):
1555	  do IO<Unit>:
1556	    f : File <- IO.try(File, File.open(path, "r"))
1557	    m : File & Result<&1, &1, U32 & String, U32> <- File.size(f)
1558	    span_ok(m, hi, lo, (lo + U32.from_nat(Nat.mul(n, 2n)) : U32), path, t)
1559	
1560	# ---- the config file: kv g half f layers vocab, then the bits of eps and theta as F32, in decimal ----
1561	
1562	def rb_done(m: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(List<&2, U32>):
1563	  (f, r) = m
1564	  do IO<List<&2, U32>>:
1565	    bs : List<&2, U32> <- IO.pass(List<&2, U32>, r)
1566	    File.close(f)
1567	    return bs
1568	
1569	def read_bytes(path: String) -> IO(List<&2, U32>):
1570	  do IO<List<&2, U32>>:
1571	    f : File <- IO.try(File, File.open(path, "r"))
1572	    m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_bytes(f, 1048576)
1573	    rb_done(m)
1574	
1575	def digit_head(xs: List<&2, U32>) -> Bool:
1576	  match xs:
1577	    case Nil{}:
1578	      False{}
1579	    case Con{+b, t}:
1580	      Bool.and(U32.is_le(48, b), U32.is_le(b, 57))
1581	
1582	# the decimal numbers in bs, separated by anything that is not a digit; dg says whether bs starts with a digit
1583	def u32s.go(bs: List<&2, U32>, dg: Bool, +cur: U32, have: Bool, acc: List<&2, U32>) -> List<&2, U32>:
1584	  match bs dg have:
1585	    case Nil{} _ True{}:
1586	      List.reverse(&2, U32, cur <> acc)
1587	    case Nil{} _ False{}:
1588	      List.reverse(&2, U32, acc)
1589	    case Con{+b, +t} True{} _:
1590	      u32s.go(t, digit_head(t), ((cur * 10 : U32) + U32.sub(b, 48) : U32), True{}, acc)
1591	    case Con{b, +t} False{} True{}:
1592	      u32s.go(t, digit_head(t), 0, False{}, cur <> acc)
1593	    case Con{b, +t} False{} False{}:
1594	      u32s.go(t, digit_head(t), 0, False{}, acc)
1595	
1596	def u32s(+bs: List<&2, U32>) -> List<&2, U32>:
1597	  u32s.go(bs, digit_head(bs), 0, False{}, Nil{})
1598	
1599	def cfg_of(+path: String, xs: List<&2, U32>) -> IO(Cfg):
1600	  match xs:
1601	    case Con{+kv, Con{+g, Con{+half, Con{+f, Con{+l, Con{+v, Con{+e, Con{+t, Nil{}}}}}}}}}:
1602	      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)})
1603	    case _:
1604	      IO.die(Cfg, 1, path ++ ": expected 8 numbers (kv g half f layers vocab eps_bits theta_bits); run make prep again")
1605	
1606	# the config of the model in dir (written by reference/prep.py)
1607	def read_cfg(+dir: String) -> IO(Cfg):
1608	  do IO<Cfg>:
1609	    bs : List<&2, U32> <- read_bytes(dir ++ "/cfg.txt")
1610	    cfg_of(dir ++ "/cfg.txt", u32s(bs))
1611	
1612	# ---- w/st.txt: tensor t's BF16 data starts at byte hi * 2^32 + lo of hf/model.safetensors (numbers 2t and 2t + 1) ----
1613	
1614	def st_path(+dir: String) -> String:
1615	  dir ++ "/hf/model.safetensors"
1616	
1617	def st_nth(xs: List<&2, U32>, k: Nat) -> U32:
1618	  match xs k:
1619	    case Con{+x, t} 0n:
1620	      x
1621	    case Con{x, t} 1n+p:
1622	      st_nth(t, p)
1623	    case Nil{} _:
1624	      0
1625	
1626	# the tensors in w/st.txt's order (reference/prep.py's weights()): emb, then layer i's ln1 q k v o ln2 gate up down
1627	# (j = 0..8), then lnf
1628	def t_layer(+i: Nat, +j: Nat) -> Nat:
1629	  Nat.add(Nat.add(1n, Nat.mul(9n, i)), j)
1630	
1631	def st_check(ok: Bool, +dir: String, xs: List<&2, U32>) -> IO(List<&2, U32>):
1632	  match ok:
1633	    case True{}:
1634	      IO.pure(List<&2, U32>, xs)
1635	    case False{}:
1636	      IO.die(List<&2, U32>, 1, dir ++ "/w/st.txt: not two numbers for each of the model's tensors; run reference/prep.py --index")
1637	
1638	# w/st.txt's numbers, two for each of the 2 + 9 * layers tensors
1639	def read_st(+dir: String, +layers: Nat) -> IO(List<&2, U32>):
1640	  do IO<List<&2, U32>>:
1641	    bs : List<&2, U32> <- read_bytes(dir ++ "/w/st.txt")
1642	    +xs : List<&2, U32> = u32s(bs)
1643	    st_check(Nat.is_eq(List.length(&2, U32, xs), Nat.mul(2n, Nat.add(2n, Nat.mul(9n, layers)))), dir, xs)
1644	
1645	def st_hi(+st: List<&2, U32>, +t: Nat) -> U32:
1646	  st_nth(st, Nat.mul(2n, t))
1647	
1648	def st_lo(+st: List<&2, U32>, +t: Nat) -> U32:
1649	  st_nth(st, 1n+Nat.mul(2n, t))
1650	
1651	# tensor t (w/st.txt's order) as bands of depth d
1652	def get_bands(+dir: String, +st: List<&2, U32>, +t: Nat, +d: Nat, +r: Nat, +c: Nat) -> IO(TA.Bands<r, c>):
1653	  do IO<TA.Bands<r, c>>:
1654	    check_span(st_path(dir), st_hi(st, t), st_lo(st, t), Nat.mul(r, c), t)
1655	    load_tree(c, r, d, r, {==}, st_path(dir), st_lo(st, t), 0n)
1656	
1657	# a norm's n weights as a list (n * 2 bytes, read at once)
1658	def get_norm(+dir: String, +st: List<&2, U32>, +t: Nat, +n: Nat) -> IO(List<&2, F32>):
1659	  do IO<List<&2, F32>>:
1660	    check_span(st_path(dir), st_hi(st, t), st_lo(st, t), n, t)
1661	    f : File <- IO.try(File, File.open(st_path(dir), "r"))
1662	    m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, st_lo(st, t), U32.from_nat(Nat.mul(n, 2n)))
1663	    bs : List<&2, U32> <- rb_done(m)
1664	    return List.take(&2, F32, bf16s(bs), n)
1665	
1666	# one layer's weights (tensors t_layer(i, 0..8)) from w/st.txt's numbers st; d: the band depth (depth_for)
1667	def load_lw_st(+c: Cfg, +dir: String, +st: List<&2, U32>, +i: Nat, +d: Nat) -> IO(LW<c>):
1668	  do IO<LW<c>>:
1669	    ln1 : List<&2, F32> <- get_norm(dir, st, t_layer(i, 0n), Cfg.dim(c))
1670	    q : TA.Bands<Cfg.dim(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 1n), d, Cfg.dim(c), Cfg.dim(c))
1671	    k : TA.Bands<Cfg.kvw(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 2n), d, Cfg.kvw(c), Cfg.dim(c))
1672	    v : TA.Bands<Cfg.kvw(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 3n), d, Cfg.kvw(c), Cfg.dim(c))
1673	    o : TA.Bands<Cfg.dim(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 4n), d, Cfg.dim(c), Cfg.dim(c))
1674	    ln2 : List<&2, F32> <- get_norm(dir, st, t_layer(i, 5n), Cfg.dim(c))
1675	    gate : TA.Bands<Cfg.ffn(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 6n), d, Cfg.ffn(c), Cfg.dim(c))
1676	    up : TA.Bands<Cfg.ffn(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 7n), d, Cfg.ffn(c), Cfg.dim(c))
1677	    down : TA.Bands<Cfg.dim(c), Cfg.ffn(c)> <- get_bands(dir, st, t_layer(i, 8n), d, Cfg.dim(c), Cfg.ffn(c))
1678	    return LW{ln1, q, k, v, o, ln2, gate, up, down}
1679	
1680	# layer i's weights alone (layer_test); d: the band depth (depth_for)
1681	def load_lw(+c: Cfg, +dir: String, +i: Nat, +d: Nat) -> IO(LW<c>):
1682	  do IO<LW<c>>:
1683	    +st : List<&2, U32> <- read_st(dir, Cfg.layers(c))
1684	    load_lw_st(c, dir, st, i, d)
1685	
1686	def load_layers(k: Nat, +c: Cfg, +dir: String, +st: List<&2, U32>, +d: Nat, +i: Nat, acc: List<&1, LW<c>>) -> IO(List<&1, LW<c>>):
1687	  match k:
1688	    case 0n:
1689	      IO.pure(List<&1, LW<c>>, List.reverse(&1, LW<c>, acc))
1690	    case 1n+p:
1691	      do IO<List<&1, LW<c>>>:
1692	        l : LW<c> <- load_lw_st(c, dir, st, i, d)
1693	        load_layers(p, c, dir, st, d, 1n+i, l <> acc)
1694	
1695	# the whole model; the embedding gets one more level of bands than the layers (2^5 with 16 threads)
1696	def load_model(+c: Cfg, +dir: String) -> IO(Model<c>):
1697	  do IO<Model<c>>:
1698	    threads : U32 <- IO.thread_count()
1699	    +d : Nat = depth_for(threads)
1700	    +st : List<&2, U32> <- read_st(dir, Cfg.layers(c))
1701	    emb : TA.Bands<Cfg.vocab(c), Cfg.dim(c)> <- get_bands(dir, st, 0n, Nat.add(d, Nat.min(d, 1n)), Cfg.vocab(c), Cfg.dim(c))
1702	    ls : List<&1, LW<c>> <- load_layers(Cfg.layers(c), c, dir, st, d, 0n, Nil{})
1703	    lnf : List<&2, F32> <- get_norm(dir, st, t_layer(Cfg.layers(c), 0n), Cfg.dim(c))
1704	    return Model{d, emb, ls, lnf}