main.bend

bend-ml-llama@0.2.1.0 · 1,704 lines · 213 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# - 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
22import Base
23import bend-ml-tensor-array@0.1.19.0/main.bend as TA
24
25# =====================================================================================================
26# Arithmetic lemmas (used by the laws below)
27# =====================================================================================================
28
29law add_zero:
30  for a: Nat
31  {Nat.add(a, 0n) == a : Nat}
32
33def 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
41law add_succ:
42  for  a: Nat
43  for -b: Nat
44  {1n+Nat.add(a, b) == Nat.add(a, 1n+b) : Nat}
45
46def 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
54law add_one:
55  for a: Nat
56  {Nat.add(a, 1n) == 1n+a : Nat}
57
58def 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
66law 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
72def 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.
87type 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
91def head(+half: Nat) -> Nat:
92  Nat.add(half, half)
93
94def Cfg.half(+c: Cfg) -> Nat:
95  match c:
96    case Cfg{kv, g, half, f, l, v, e, t}:
97      half
98
99def Cfg.kv(+c: Cfg) -> Nat:
100  match c:
101    case Cfg{kv, g, half, f, l, v, e, t}:
102      kv
103
104def Cfg.g(+c: Cfg) -> Nat:
105  match c:
106    case Cfg{kv, g, half, f, l, v, e, t}:
107      g
108
109def Cfg.ffn(+c: Cfg) -> Nat:
110  match c:
111    case Cfg{kv, g, half, f, l, v, e, t}:
112      f
113
114def Cfg.layers(+c: Cfg) -> Nat:
115  match c:
116    case Cfg{kv, g, half, f, l, v, e, t}:
117      l
118
119def Cfg.vocab(+c: Cfg) -> Nat:
120  match c:
121    case Cfg{kv, g, half, f, l, v, e, t}:
122      v
123
124def Cfg.eps(+c: Cfg) -> F32:
125  match c:
126    case Cfg{kv, g, half, f, l, v, e, t}:
127      e
128
129def 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
135def Cfg.head(+c: Cfg) -> Nat:
136  head(Cfg.half(c))
137
138# the query heads: kv groups of g
139def Cfg.heads(+c: Cfg) -> Nat:
140  Nat.mul(Cfg.kv(c), Cfg.g(c))
141
142# the hidden size: one head per query head
143def 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
147def 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
151def 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
155def 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
159def 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.
165law 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
168def smollm2_135m_sizes():
169  ({==}, {==}, {==}, {==})
170
171law 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
174def smollm2_360m_sizes():
175  ({==}, {==}, {==}, {==})
176
177law 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
180def 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
185def 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)
190def 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
194def 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.
201law layout_135m:
202  {Cfg.size(smollm2_135m()) == (134515008 : U32) : U32}
203
204def layout_135m():
205  {==}
206
207law layout_360m:
208  {Cfg.size(smollm2_360m()) == (361821120 : U32) : U32}
209
210def layout_360m():
211  {==}
212
213law layout_1_7b:
214  {Cfg.size(smollm2_1_7b()) == (1711376384 : U32) : U32}
215
216def 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
224type Head is Data:
225  Head{lo: List<&2, F32>, hi: List<&2, F32>}
226
227def split(+half: Nat, +xs: List<&2, F32>) -> Head:
228  Head{List.take(&2, F32, xs, half), List.drop(&2, F32, xs, half)}
229
230def 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
236def 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)
244def 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
251def 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
258def 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.
263law 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.
271def 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)
288def 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)
296def 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
304def 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
312def 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
319def 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)
327def 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
335def 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
343def 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
350def 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
354def 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
358def 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]
366def 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).
375def 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)
387def 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)
391def 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
398def 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
401def 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
408def 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
415def 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)
423type Rot is Data:
424  Rot{cs: List<&2, F32>, sn: List<&2, F32>}
425
426def 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)
430def 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
437def 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
444def 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
449def 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
457def silu(+x: F32) -> F32:
458  (x / (1.0 + F32.exp(F32.neg(x)) : F32) : F32)
459
460# silu(gate) * up
461def 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).
476type 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
480def 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
488def 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
492def 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
500def 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
506def 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.
512law 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.
518def 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.
528law kv_heads:
529  for +c: Cfg
530  for +cap: Nat
531  {kv_count(cap, 0n, KV.empty(c, cap)) == Cfg.kv(c) : Nat}
532
533def 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
537def 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
551def 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).
566law 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).
575def 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
580def 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.
586def 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
593def 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.
602law 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
608def 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
617law 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
624def 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
633law 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
640def 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
652law 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
659def 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.
669law 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).
677def 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
687type SR is Type:
688  SR{k: Array<F32>, s: List<&2, F32>}
689
690def sr_len(sr: SR) -> Nat:
691  match sr:
692    case SR{k, s}:
693      List.length(&2, F32, s)
694
695def 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).
704law 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
712def 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)
720def 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)
727def 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
730def 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.
739law 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
746def 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)
750type WR is Type:
751  WR{v: Array<F32>, o: List<&2, F32>}
752
753def wr_of(v: Array<F32>, al: TA.AL) -> WR:
754  match al:
755    case TA.AL{c, l}:
756      WR{v, l}
757
758def 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
763def 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
767type HO is Type:
768  HO{k: Array<F32>, v: Array<F32>, o: List<&2, F32>}
769
770def 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
775def head_wr(k: Array<F32>, wr: WR) -> HO:
776  match wr:
777    case WR{v, o}:
778      HO{k, v, o}
779
780def 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
785def 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
789def 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
799type AT is Type:
800  AT{ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>, o: List<&2, F32>}
801
802def 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
808type GW is Type:
809  GW{g: List<&2, Head>, k: Array<F32>, v: Array<F32>}
810
811def 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
819def 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)
831type GS is Type:
832  GS{l: List<&1, GW>, r: List<&1, GW>}
833
834def 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
842def 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
859def 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
870type KA<-cap: Nat, -t: Nat> is Type:
871  KA{kv: KV<cap, t>, o: List<&2, F32>}
872
873def 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
878def 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.
887def 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)
899type 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
904type 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
909type 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)
914def 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
919def 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)
925def 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)
930def 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.
953def 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
973def 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
985type 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
989type 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
993def 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
1000def 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
1009type 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
1013def 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)
1019def 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
1025def 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
1032def 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)
1038def 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
1044def 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
1051def 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
1056def 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
1061def 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
1067def 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
1075type AM is Data:
1076  AM{i: Nat, v: F32}
1077
1078def AM.i(am: AM) -> Nat:
1079  match am:
1080    case AM{i, v}:
1081      i
1082
1083def AM.v(am: AM) -> F32:
1084  match am:
1085    case AM{i, v}:
1086      v
1087
1088def 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
1095def 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
1104def 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
1121def 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.
1130def 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
1134type Cand is Data:
1135  Cand{v: F32, i: Nat}
1136
1137def Cand.i(c: Cand) -> Nat:
1138  match c:
1139    case Cand{v, i}:
1140      i
1141
1142def 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)
1148def 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
1156def 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)
1166def 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)
1174def 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.
1183def 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
1195def 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)
1199def 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
1203def 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)
1211type 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)
1215def 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
1227def 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
1235def 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
1245def 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
1250def 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
1253def 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
1262def 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
1269def 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
1278def 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
1285def 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.
1295law 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
1305def 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.
1318law 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
1325def 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.
1337law 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
1350def 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.
1363law 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
1376def 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
1386def 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).
1395law 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
1404def 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)
1419def bitcast(u: U32) -> F32:
1420  match u:
1421    case U32{w}:
1422      F32{w}
1423
1424def 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
1427def 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
1435def 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
1447def 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)
1451def 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).
1461def depth_sel(few: Bool) -> Nat:
1462  match few:
1463    case True{}:
1464      0n
1465    case False{}:
1466      4n
1467
1468def 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)
1474def bf16_le(+b0: U32, +b1: U32) -> F32:
1475  bitcast(((b0 * 65536 : U32) + (b1 * 16777216 : U32) : U32))
1476
1477def 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
1486def 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
1493def 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
1500def 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)
1514def 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)
1524def 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
1538def 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
1545def 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
1554def 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
1562def 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
1569def 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
1575def 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
1583def 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
1596def u32s(+bs: List<&2, U32>) -> List<&2, U32>:
1597  u32s.go(bs, digit_head(bs), 0, False{}, Nil{})
1598
1599def 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)
1607def 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
1614def st_path(+dir: String) -> String:
1615  dir ++ "/hf/model.safetensors"
1616
1617def 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
1628def t_layer(+i: Nat, +j: Nat) -> Nat:
1629  Nat.add(Nat.add(1n, Nat.mul(9n, i)), j)
1630
1631def 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
1639def 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
1645def st_hi(+st: List<&2, U32>, +t: Nat) -> U32:
1646  st_nth(st, Nat.mul(2n, t))
1647
1648def 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
1652def 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)
1658def 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)
1667def 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)
1681def 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
1686def 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)
1696def 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}