main.bend

bend-ml-llama@0.2.0.0 · 1,343 lines · 175 definitions

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
21import Base
22import bend-ml-tensor-array@0.1.14.0/main.bend as TA
23
24# =====================================================================================================
25# Arithmetic lemmas (used by the laws below)
26# =====================================================================================================
27
28law add_zero:
29  for a: Nat
30  {Nat.add(a, 0n) == a : Nat}
31
32def 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
40law add_succ:
41  for  a: Nat
42  for -b: Nat
43  {1n+Nat.add(a, b) == Nat.add(a, 1n+b) : Nat}
44
45def 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
53law add_one:
54  for a: Nat
55  {Nat.add(a, 1n) == 1n+a : Nat}
56
57def 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
65law 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
71def 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.
86type 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
90def head(+half: Nat) -> Nat:
91  Nat.add(half, half)
92
93def Cfg.half(+c: Cfg) -> Nat:
94  match c:
95    case Cfg{kv, g, half, f, l, v, e, t}:
96      half
97
98def Cfg.kv(+c: Cfg) -> Nat:
99  match c:
100    case Cfg{kv, g, half, f, l, v, e, t}:
101      kv
102
103def Cfg.g(+c: Cfg) -> Nat:
104  match c:
105    case Cfg{kv, g, half, f, l, v, e, t}:
106      g
107
108def Cfg.ffn(+c: Cfg) -> Nat:
109  match c:
110    case Cfg{kv, g, half, f, l, v, e, t}:
111      f
112
113def Cfg.layers(+c: Cfg) -> Nat:
114  match c:
115    case Cfg{kv, g, half, f, l, v, e, t}:
116      l
117
118def Cfg.vocab(+c: Cfg) -> Nat:
119  match c:
120    case Cfg{kv, g, half, f, l, v, e, t}:
121      v
122
123def Cfg.eps(+c: Cfg) -> F32:
124  match c:
125    case Cfg{kv, g, half, f, l, v, e, t}:
126      e
127
128def 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
134def Cfg.head(+c: Cfg) -> Nat:
135  head(Cfg.half(c))
136
137# the query heads: kv groups of g
138def Cfg.heads(+c: Cfg) -> Nat:
139  Nat.mul(Cfg.kv(c), Cfg.g(c))
140
141# the hidden size: one head per query head
142def 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
146def 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
150def 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
154def 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
158def 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.
164law 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
167def smollm2_135m_sizes():
168  ({==}, {==}, {==}, {==})
169
170law 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
173def smollm2_360m_sizes():
174  ({==}, {==}, {==}, {==})
175
176law 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
179def 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
184def 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)
189def 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
193def 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.
200law layout_135m:
201  {Cfg.size(smollm2_135m()) == (134515008 : U32) : U32}
202
203def layout_135m():
204  {==}
205
206law layout_360m:
207  {Cfg.size(smollm2_360m()) == (361821120 : U32) : U32}
208
209def layout_360m():
210  {==}
211
212law layout_1_7b:
213  {Cfg.size(smollm2_1_7b()) == (1711376384 : U32) : U32}
214
215def 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
223type Head is Data:
224  Head{lo: List<&2, F32>, hi: List<&2, F32>}
225
226def split(+half: Nat, +xs: List<&2, F32>) -> Head:
227  Head{List.take(&2, F32, xs, half), List.drop(&2, F32, xs, half)}
228
229def 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
235def 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)
243def 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
250def 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
257def 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.
262law 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.
270def 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)
287def 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)
295def 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
303def 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
311def 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
318def 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)
326def 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
334def 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
342def 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
349def 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
353def 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
357def 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]
365def 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).
374def 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)
386def 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)
390def 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
397def 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
400def 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
407def 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
414def 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)
422type Rot is Data:
423  Rot{cs: List<&2, F32>, sn: List<&2, F32>}
424
425def 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)
429def 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
436def 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
443def 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
448def 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
456def silu(+x: F32) -> F32:
457  (x / (1.0 + F32.exp(F32.neg(x)) : F32) : F32)
458
459# silu(gate) * up
460def 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).
475type 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
479def 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
487def 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
491def 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
499def 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
505def 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.
511law 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.
517def 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.
527law kv_heads:
528  for +c: Cfg
529  for +cap: Nat
530  {kv_count(cap, 0n, KV.empty(c, cap)) == Cfg.kv(c) : Nat}
531
532def 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
536def 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
550def 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).
565law 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).
574def 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
579def 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.
585def 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
592def 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.
601law 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
607def 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
616law 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
623def 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
632law 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
639def 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
651law 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
658def 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.
668law 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).
676def 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
686type SR is Type:
687  SR{k: Array<F32>, s: List<&2, F32>}
688
689def sr_len(sr: SR) -> Nat:
690  match sr:
691    case SR{k, s}:
692      List.length(&2, F32, s)
693
694def 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).
702law 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
710def 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
718def 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)
725def 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
728def 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.
737law 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
744def 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)
748type WR is Type:
749  WR{v: Array<F32>, o: List<&2, F32>}
750
751def wr_of(v: Array<F32>, al: TA.AL) -> WR:
752  match al:
753    case TA.AL{c, l}:
754      WR{v, l}
755
756def 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
761def 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
765type HO is Type:
766  HO{k: Array<F32>, v: Array<F32>, o: List<&2, F32>}
767
768def 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
773def head_wr(k: Array<F32>, wr: WR) -> HO:
774  match wr:
775    case WR{v, o}:
776      HO{k, v, o}
777
778def 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
783def 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
787def 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
797type AT is Type:
798  AT{ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>, o: List<&2, F32>}
799
800def 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
806type GW is Type:
807  GW{g: List<&2, Head>, k: Array<F32>, v: Array<F32>}
808
809def 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
817def 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)
829type GS is Type:
830  GS{l: List<&1, GW>, r: List<&1, GW>}
831
832def 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
840def 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
857def 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
868type KA<-cap: Nat, -t: Nat> is Type:
869  KA{kv: KV<cap, t>, o: List<&2, F32>}
870
871def 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
876def 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.
885def 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)
897type 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
902type 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
907type 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)
912def 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
917def 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)
923def 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)
928def 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.
951def 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
971def 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
983type 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
987type 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
991def 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
998def 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
1007type 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
1011def 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)
1017def 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
1023def 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
1030def 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)
1036def 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
1042def 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
1049def 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
1054def 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
1059def 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
1065def 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
1073type AM is Data:
1074  AM{i: Nat, v: F32}
1075
1076def AM.i(am: AM) -> Nat:
1077  match am:
1078    case AM{i, v}:
1079      i
1080
1081def AM.v(am: AM) -> F32:
1082  match am:
1083    case AM{i, v}:
1084      v
1085
1086def 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
1093def 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
1102def 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)
1115def bitcast(u: U32) -> F32:
1116  match u:
1117    case U32{w}:
1118      F32{w}
1119
1120def 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
1123def 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
1131def 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
1143def 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
1147def 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).
1157def depth_sel(few: Bool) -> Nat:
1158  match few:
1159    case True{}:
1160      0n
1161    case False{}:
1162      4n
1163
1164def depth_for(+threads: U32) -> Nat:
1165  depth_sel(U32.is_le(threads, 2))
1166
1167def 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
1172def 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
1179def 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)
1188def 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
1194def wpath(dir: String, name: String) -> String:
1195  dir ++ "/w/" ++ name ++ ".bin"
1196
1197def 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
1208def 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
1215def 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
1222def 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)
1236def 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
1243def 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
1255def 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)
1261def 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
1274def 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)
1284def 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
1295def 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
1302def 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
1308def 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
1316def 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
1329def u32s(+bs: List<&2, U32>) -> List<&2, U32>:
1330  u32s.go(bs, digit_head(bs), 0, False{}, Nil{})
1331
1332def 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)
1340def 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))