main.bend
bend-ml-llama@0.1.0.0 · 1,209 lines · 0 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@0.1.2.0/main.bend as T
23import bend-ml-tensor-array@0.1.7.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# the squares of xs
286def squares(xs: List<&2, F32>) -> List<&2, F32>:
287 match xs:
288 case Nil{}:
289 Nil{}
290 case Con{+h, t}:
291 (h * h : F32) <> squares(t)
292
293# adjacent pairs added: [a, b, c, d, e] -> [a + b, c + d, e]
294def pair_sums(xs: List<&2, F32>) -> List<&2, F32>:
295 match xs:
296 case Con{+a, Con{+b, t}}:
297 (a + b : F32) <> pair_sums(t)
298 case _:
299 xs
300
301# the sum as a balanced tree of additions (pairwise summation): its rounding error grows with log2(n), not n, as
302# in PyTorch's cascade sum. A running sum loses the small squares next to a large one (NOTES, 2026-10-10).
303def tree_sum(fuel: Nat, xs: List<&2, F32>) -> F32:
304 match fuel xs:
305 case _ Nil{}:
306 0.0
307 case _ Con{+a, Nil{}}:
308 a
309 case 0n Con{+a, t}:
310 T.sum.go(t, a)
311 case 1n+p Con{+a, t}:
312 tree_sum(p, pair_sums(Con{a, t}))
313
314# RMSNorm: g * (x * (1 / sqrt(mean(x²) + eps))), in this order (LlamaRMSNorm.forward)
315def rmsnorm(+c: Cfg, +x: List<&2, F32>, +g: List<&2, F32>) -> List<&2, F32>:
316 T.zip_with.mul(g, T.scale.l((1.0 / F32.sqrt(((tree_sum(Cfg.dim(c), squares(x)) / F32.from_nat(Cfg.dim(c)) : F32) + Cfg.eps(c) : F32)) : F32), x))
317
318# inv_freq[i] = 1 / (theta ^ (2i / head)), i = 0..half-1 (LlamaRotaryEmbedding, default rope)
319def inv_freq.go(n: Nat, +i: Nat, +theta: F32, +hd: F32) -> List<&2, F32>:
320 match n:
321 case 0n:
322 Nil{}
323 case 1n+p:
324 (1.0 / F32.pow(theta, (F32.from_nat(Nat.mul(2n, i)) / hd : F32)) : F32) <> inv_freq.go(p, 1n+i, theta, hd)
325
326def inv_freq(+c: Cfg) -> List<&2, F32>:
327 inv_freq.go(Cfg.half(c), 0n, Cfg.theta(c), F32.from_nat(Cfg.head(c)))
328
329def angles(+t: F32, fs: List<&2, F32>) -> List<&2, F32>:
330 match fs:
331 case Nil{}:
332 Nil{}
333 case Con{+f, r}:
334 (t * f : F32) <> angles(t, r)
335
336def cos.l(xs: List<&2, F32>) -> List<&2, F32>:
337 match xs:
338 case Nil{}:
339 Nil{}
340 case Con{+x, r}:
341 F32.cos(x) <> cos.l(r)
342
343def sin.l(xs: List<&2, F32>) -> List<&2, F32>:
344 match xs:
345 case Nil{}:
346 Nil{}
347 case Con{+x, r}:
348 F32.sin(x) <> sin.l(r)
349
350# the cosines and sines of position t, one per frequency (shared by every head of every layer)
351type Rot is Data:
352 Rot{cs: List<&2, F32>, sn: List<&2, F32>}
353
354def rot_at(+c: Cfg, +t: Nat) -> Rot:
355 Rot{cos.l(angles(F32.from_nat(t), inv_freq(c))), sin.l(angles(F32.from_nat(t), inv_freq(c)))}
356
357# first half: lo*cos + (-hi)*sin; second half: hi*cos + lo*sin (q*cos + rotate_half(q)*sin)
358def rot_lo(h1: List<&2, F32>, h2: List<&2, F32>, cs: List<&2, F32>, sn: List<&2, F32>) -> List<&2, F32>:
359 match h1 h2 cs sn:
360 case Con{+a, ta} Con{+b, tb} Con{+c, tc} Con{+s, ts}:
361 ((a * c : F32) + (F32.neg(b) * s : F32) : F32) <> rot_lo(ta, tb, tc, ts)
362 case _ _ _ _:
363 Nil{}
364
365def rot_hi(h1: List<&2, F32>, h2: List<&2, F32>, cs: List<&2, F32>, sn: List<&2, F32>) -> List<&2, F32>:
366 match h1 h2 cs sn:
367 case Con{+a, ta} Con{+b, tb} Con{+c, tc} Con{+s, ts}:
368 ((b * c : F32) + (a * s : F32) : F32) <> rot_hi(ta, tb, tc, ts)
369 case _ _ _ _:
370 Nil{}
371
372def rope(rt: Rot, h: Head) -> Head:
373 match rt h:
374 case Rot{+cs, +sn} Head{+lo, +hi}:
375 Head{rot_lo(lo, hi, cs, sn), rot_hi(lo, hi, cs, sn)}
376
377def ropes(+rt: Rot, hs: List<&2, Head>) -> List<&2, Head>:
378 match hs:
379 case Nil{}:
380 Nil{}
381 case Con{h, t}:
382 rope(rt, h) <> ropes(rt, t)
383
384# silu(x) = x / (1 + exp(-x)), the form ATen's CPU kernel uses
385def silu(+x: F32) -> F32:
386 (x / (1.0 + F32.exp(F32.neg(x)) : F32) : F32)
387
388# silu(gate) * up
389def swiglu(g: List<&2, F32>, u: List<&2, F32>) -> List<&2, F32>:
390 match g u:
391 case Con{+a, ga} Con{+b, ub}:
392 (silu(a) * b : F32) <> swiglu(ga, ub)
393 case _ _:
394 Nil{}
395
396# =====================================================================================================
397# The key/value cache in flat arrays, indexed by the number of positions it holds
398# =====================================================================================================
399
400# One pair of arrays per key/value head: the keys (after RoPE) and the values of positions 0..t-1, position p at
401# [p·head, p·head + head). The type indexes are the capacity cap (the positions the arrays have room for, set by
402# KV.empty) and t, the number of positions written. KV.push asks for a proof that t < cap, so a cache is never
403# written past its room (kv_room says why that is enough).
404type KV<-cap: Nat, -t: Nat> is Type:
405 KV{ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>}
406
407# n arrays of 2^d slots
408def arrays(n: Nat, +d: Nat) -> List<&1, Array<F32>>:
409 match n:
410 case 0n:
411 Nil{}
412 case 1n+p:
413 Array.new(F32, d, 0.0) <> arrays(p, d)
414
415# a cache with no position yet, with room for cap positions per key/value head
416def KV.empty(+c: Cfg, +cap: Nat) -> KV<cap, 0n>:
417 KV{arrays(Cfg.kv(c), TA.cap_depth(Nat.mul(cap, Cfg.head(c)))), arrays(Cfg.kv(c), TA.cap_depth(Nat.mul(cap, Cfg.head(c))))}
418
419# writes row j of rs into array j, at off
420def write_rows(arrs: List<&1, Array<F32>>, rs: List<&2, List<&2, F32>>, +off: U32) -> List<&1, Array<F32>>:
421 match arrs rs:
422 case Con{a, at} Con{r, rt}:
423 TA.fill_list(r, off, a) <> write_rows(at, rt, off)
424 case _ _:
425 arrs
426
427# position t: the keys and values of every key/value head written at t·head; room is the proof that t < cap
428def KV.push(+c: Cfg, -cap: Nat, +t: Nat, kv: KV<cap, t>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +ks: List<&2, List<&2, F32>>, +vs: List<&2, List<&2, F32>>) -> KV<cap, 1n+t>:
429 match kv:
430 case KV{ka, va}:
431 KV{write_rows(ka, ks, U32.from_nat(Nat.mul(t, Cfg.head(c)))), write_rows(va, vs, U32.from_nat(Nat.mul(t, Cfg.head(c))))}
432
433# the number of key/value heads a cache holds
434def kv_count(-cap: Nat, -t: Nat, kv: KV<cap, t>) -> Nat:
435 match kv:
436 case KV{ka, va}:
437 List.length(&1, Array<F32>, ka)
438
439# LAW (arrays_len): arrays(n, d) makes exactly n arrays.
440law arrays_len:
441 for +n: Nat
442 for +d: Nat
443 {List.length(&1, Array<F32>, arrays(n, d)) == n : Nat}
444
445# PROOF: by induction on n: none for 0; one more than the rest for 1 + p.
446def arrays_len(n, d):
447 match n:
448 case 0n:
449 {==}
450 case 1n+p:
451 %arrays_len(p, d) : {1n+List.length(&1, Array<F32>, arrays(p, d)) == 1n+_ : Nat}
452 {==}
453
454# LAW (kv_heads): an empty cache holds one array of keys (and one of values) per key/value head, so key/value head j
455# of the model reads column j and no other.
456law kv_heads:
457 for +c: Cfg
458 for +cap: Nat
459 {kv_count(cap, 0n, KV.empty(c, cap)) == Cfg.kv(c) : Nat}
460
461def kv_heads(c, cap):
462 arrays_len(Cfg.kv(c), TA.cap_depth(Nat.mul(cap, Cfg.head(c))))
463
464# a <= b and b <= c -> a <= c
465def leb_trans(a: Nat, b: Nat, c: Nat, e1: {True{} == TA.leb(a, b) : Bool}, e2: {True{} == TA.leb(b, c) : Bool}) -> {True{} == TA.leb(a, c) : Bool}:
466 match a b c:
467 case 0n _ _:
468 {==}
469 case 1n+p 0n _:
470 %e1 : TA.BD(_, Unit, {True{} == TA.leb(1n+p, c) : Bool})
471 Unit{}
472 case 1n+p 1n+q 0n:
473 %e2 : TA.BD(_, Unit, {True{} == TA.leb(1n+p, 0n) : Bool})
474 Unit{}
475 case 1n+p 1n+q 1n+r:
476 leb_trans(p, q, r, e1, e2)
477
478# a <= b -> a·h <= b·h
479def leb_mul(a: Nat, b: Nat, +h: Nat, e: {True{} == TA.leb(a, b) : Bool}) -> {True{} == TA.leb(Nat.mul(a, h), Nat.mul(b, h)) : Bool}:
480 match a b:
481 case 0n _:
482 {==}
483 case 1n+p 0n:
484 %e : TA.BD(_, Unit, {True{} == TA.leb(Nat.mul(1n+p, h), 0n) : Bool})
485 Unit{}
486 case 1n++p 1n++q:
487 TA.leb_add(h, h, Nat.mul(p, h), Nat.mul(q, h), TA.leb_refl(h), leb_mul(p, q, h, e))
488
489# LAW (kv_room): KV.empty(c, cap) gives each key/value head arrays of 2^cap_depth(cap·head) slots, and KV.push at
490# position t writes head numbers from t·head. For a position below the capacity (t < cap) they all fall inside the
491# slots: (1 + t)·head <= 2^cap_depth(cap·head), so a push never wraps around onto position 0. (What stays trusted is
492# tensor-array's: Array.new allocates 2^d slots.) The type does not stop a caller from pushing at t >= cap; the
493# callers size cap to the positions they run.
494law kv_room:
495 for +cap: Nat
496 for +h: Nat
497 for +t: Nat
498 for e: {True{} == TA.leb(1n+t, cap) : Bool}
499 {True{} == TA.leb(Nat.mul(1n+t, h), TA.pow2(TA.cap_depth(Nat.mul(cap, h)))) : Bool}
500
501# For callers that walk the positions with a counter of steps left (fuel), keeping t + fuel = cap: with 1 + p steps
502# left, position t is below the capacity (step_room), and after the step t + 1 + p = cap still holds (step_inv).
503def step_room(+t: Nat, +p: Nat, -cap: Nat, e: {Nat.add(t, 1n+p) == cap : Nat}) -> {True{} == TA.leb(1n+t, cap) : Bool}:
504 %e : {True{} == TA.leb(1n+t, _) : Bool}
505 %add_succ(t, p) : {True{} == TA.leb(1n+t, _) : Bool}
506 TA.leb_up(t, t, p, TA.leb_refl(t))
507
508def step_inv(+t: Nat, -p: Nat, -cap: Nat, e: {Nat.add(t, 1n+p) == cap : Nat}) -> {Nat.add(1n+t, p) == cap : Nat}:
509 %e : {1n+Nat.add(t, p) == _ : Nat}
510 add_succ(t, p)
511
512# PROOF: (1 + t)·head <= cap·head by leb_mul, and cap·head <= 2^cap_depth(cap·head) by tensor-array's cap_ok; then
513# leb_trans.
514def kv_room(cap, h, t, e):
515 leb_trans(Nat.mul(1n+t, h), Nat.mul(cap, h), TA.pow2(TA.cap_depth(Nat.mul(cap, h))), leb_mul(1n+t, cap, h, e), TA.cap_ok(Nat.mul(cap, h)))
516
517# =====================================================================================================
518# Grouped-query attention
519# =====================================================================================================
520# the query heads in n groups of g, in order: group j holds query heads j·g .. j·g + g - 1
521def groups(n: Nat, +g: Nat, +hs: List<&2, Head>) -> List<&2, List<&2, Head>>:
522 match n:
523 case 0n:
524 Nil{}
525 case 1n+p:
526 List.take(&2, Head, hs, g) <> groups(p, g, List.drop(&2, Head, hs, g))
527
528# LAW (groups_len): the query heads make exactly as many groups as there are key/value heads, so attention
529# (which walks the groups and the columns together) pairs every group with a key/value head.
530law groups_len:
531 for n: Nat
532 for +g: Nat
533 for +hs: List<&2, Head>
534 {List.length(&2, List<&2, Head>, groups(n, g, hs)) == n : Nat}
535
536def groups_len(n, g, hs):
537 match n:
538 case 0n:
539 {==}
540 case 1n+p:
541 %groups_len(p, g, List.drop(&2, Head, hs, g)) : {1n+List.length(&2, List<&2, Head>, groups(p, g, List.drop(&2, Head, hs, g))) == 1n+_ : Nat}
542 {==}
543
544# one block of 1 + m numbers in Nat.divmod: the counter m runs down to 0, then the quotient d grows by one
545law div_block:
546 for +m: Nat
547 for -n: Nat
548 for -d: Nat
549 for -r: Nat
550 {Nat.divmod.go(Nat.add(1n+m, n), m, d, r) == Nat.divmod.go(n, Nat.add(m, r), 1n+d, 0n) : Nat & Nat}
551
552def div_block(m, n, d, r):
553 match m:
554 case 0n:
555 {==}
556 case 1n+p:
557 %Equal.sym(Nat, 1n+Nat.add(p, r), Nat.add(p, 1n+r), add_succ(p, r)) : {Nat.divmod.go(Nat.add(1n+p, n), p, d, 1n+r) == Nat.divmod.go(n, _, 1n+d, 0n) : Nat & Nat}
558 div_block(p, n, d, 1n+r)
559
560# j whole blocks of 1 + s numbers add j to the quotient
561law div_blocks:
562 for +j: Nat
563 for +s: Nat
564 for +n: Nat
565 for +d: Nat
566 {Nat.divmod.go(Nat.add(Nat.mul(j, 1n+s), n), s, d, 0n) == Nat.divmod.go(n, s, Nat.add(j, d), 0n) : Nat & Nat}
567
568def div_blocks(j, s, n, d):
569 match j:
570 case 0n:
571 {==}
572 case 1n+jp:
573 %Equal.sym(Nat, Nat.add(Nat.add(1n+s, Nat.mul(jp, 1n+s)), n), Nat.add(1n+s, Nat.add(Nat.mul(jp, 1n+s), n)), add_assoc(1n+s, Nat.mul(jp, 1n+s), n)) : {Nat.divmod.go(_, s, d, 0n) == Nat.divmod.go(n, s, Nat.add(1n+jp, d), 0n) : Nat & Nat}
574 %Equal.sym(Nat & Nat, Nat.divmod.go(Nat.add(1n+s, Nat.add(Nat.mul(jp, 1n+s), n)), s, d, 0n), Nat.divmod.go(Nat.add(Nat.mul(jp, 1n+s), n), Nat.add(s, 0n), 1n+d, 0n), div_block(s, Nat.add(Nat.mul(jp, 1n+s), n), d, 0n)) : {_ == Nat.divmod.go(n, s, Nat.add(1n+jp, d), 0n) : Nat & Nat}
575 %Equal.sym(Nat, Nat.add(s, 0n), s, add_zero(s)) : {Nat.divmod.go(Nat.add(Nat.mul(jp, 1n+s), n), _, 1n+d, 0n) == Nat.divmod.go(n, s, Nat.add(1n+jp, d), 0n) : Nat & Nat}
576 %Equal.sym(Nat, 1n+Nat.add(jp, d), Nat.add(jp, 1n+d), add_succ(jp, d)) : {Nat.divmod.go(Nat.add(Nat.mul(jp, 1n+s), n), s, 1n+d, 0n) == Nat.divmod.go(n, s, _, 0n) : Nat & Nat}
577 div_blocks(jp, s, n, 1n+d)
578
579# fewer than 1 + s numbers after the last whole block leave the quotient alone
580law div_rest:
581 for +i: Nat
582 for -k: Nat
583 for -d: Nat
584 for -r: Nat
585 {Pair.fst(Nat, Nat, Nat.divmod.go(i, Nat.add(i, k), d, r)) == d : Nat}
586
587def div_rest(i, k, d, r):
588 match i:
589 case 0n:
590 {==}
591 case 1n+p:
592 div_rest(p, k, d, 1n+r)
593
594# LAW (gqa_head): with g = 1 + i + k query heads per key/value head (g > 0 and i < g), query head j·g + i, the
595# head i of group j, reads key/value head (j·g + i) / g = j. That division is how transformers maps query heads
596# to key/value heads (repeat_kv), so grouping the heads computes the same attention.
597law gqa_head:
598 for +j: Nat
599 for +i: Nat
600 for +k: Nat
601 {Nat.div(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}
602
603# PROOF: Nat.div counts how many whole blocks of g numbers fit. j·g is j whole blocks (div_blocks, which uses
604# div_block for each block), and i < g numbers after them do not complete another one (div_rest).
605def gqa_head(j, i, k):
606 %Equal.sym(Nat & Nat, Nat.divmod.go(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), Nat.add(i, k), 0n, 0n), Nat.divmod.go(i, Nat.add(i, k), Nat.add(j, 0n), 0n), div_blocks(j, Nat.add(i, k), i, 0n)) : {Pair.fst(Nat, Nat, _) == j : Nat}
607 %add_zero(j) : {Pair.fst(Nat, Nat, Nat.divmod.go(i, Nat.add(i, k), Nat.add(j, 0n), 0n)) == _ : Nat}
608 div_rest(i, k, Nat.add(j, 0n), 0n)
609
610# =====================================================================================================
611# Grouped-query attention over the arrays
612# =====================================================================================================
613
614# the scores of one query head against the keys of its key/value head, and the keys handed back
615type SR is Type:
616 SR{k: Array<F32>, s: List<&2, F32>}
617
618def sr_len(sr: SR) -> Nat:
619 match sr:
620 case SR{k, s}:
621 List.length(&2, F32, s)
622
623def sr_of(k: Array<F32>, al: TA.AL) -> SR:
624 match al:
625 case TA.AL{c, l}:
626 SR{k, l}
627
628# LAW (read_len): reading m + 1 numbers of an array onto acc gives m + 1 more scores.
629# PROOF: by induction on m. With 0 left it reads one number and reverses the list (rev_len counts it); with 1 + p
630# left it reads one and recurses, and len_step moves the 1 out of the sum (tensor-array's leaf_len, for scores).
631law read_len:
632 for k: Array<F32>
633 for +m: Nat
634 for +idx: U32
635 for +acc: List<&2, F32>
636 for rc: Array<F32> & F32
637 {sr_len(sr_of(k, TA.read_l(m, idx, acc, rc))) == Nat.add(1n+m, List.length(&2, F32, acc)) : Nat}
638
639def read_len(k, m, idx, acc, rc):
640 match m rc:
641 case 0n Tuple{d, +v}:
642 TA.zero_out(List.length(&2, F32, List.reverse.go(&2, F32, v <> acc, Nil{})), 1n+List.length(&2, F32, acc), TA.rev_len(v <> acc, Nil{}))
643 case 1n+p Tuple{d, +v}:
644 TA.len_step(sr_len(sr_of(k, TA.read_l(p, (idx + 1 : U32), v <> acc, Array.get(F32, d, (idx + 1 : U32))))), p, List.length(&2, F32, acc), read_len(k, p, (idx + 1 : U32), v <> acc, Array.get(F32, d, (idx + 1 : U32))))
645
646# read the 1 + t scores the gemm wrote
647def sr_read(+t: Nat, g: TA.G) -> SR:
648 match g:
649 case TA.G{qa, k, s}:
650 sr_of(k, TA.read_l(t, 0, Nil{}, Array.get(F32, s, 0)))
651
652# q·k for the keys of positions 0..t, oldest first: K (t + 1 rows of `head` numbers) times q, the product
653# Mat.matmul_nt computes for one row (bend-ml-tensor-array's gemm)
654def scores(+hd: Nat, +t: Nat, k: Array<F32>, +q: List<&2, F32>) -> SR:
655 sr_read(t, TA.gemm(1n, 1n+t, hd, U32.from_nat(hd), 1, 1, U32.from_nat(hd), TA.fill_list(q, 0, Array.new(F32, TA.cap_depth(hd), 0.0)), k, Array.new(F32, TA.cap_depth(1n+t), 0.0)))
656
657def sr_len_g(+t: Nat, g: TA.G) -> {sr_len(sr_read(t, g)) == 1n+t : Nat}:
658 match g:
659 case TA.G{qa, k, s}:
660 TA.zero_out(sr_len(sr_of(k, TA.read_l(t, 0, Nil{}, Array.get(F32, s, 0)))), 1n+t, read_len(k, t, 0, Nil{}, Array.get(F32, s, 0)))
661
662# LAW (causal): attention at position t scores exactly t + 1 keys, the positions 0..t: one score per position
663# written, none for a position not yet computed.
664# PROOF: the gemm's result has one constructor, so matching it opens it without computing the product
665# (sr_len_g); read_len counts the 1 + t numbers read back.
666law causal:
667 for +hd: Nat
668 for +t: Nat
669 for k: Array<F32>
670 for +q: List<&2, F32>
671 {sr_len(scores(hd, t, k, q)) == 1n+t : Nat}
672
673def causal(hd, t, k, q):
674 sr_len_g(t, TA.gemm(1n, 1n+t, hd, U32.from_nat(hd), 1, 1, U32.from_nat(hd), TA.fill_list(q, 0, Array.new(F32, TA.cap_depth(hd), 0.0)), k, Array.new(F32, TA.cap_depth(1n+t), 0.0)))
675
676# the weighted sum of the values of positions 0..t: w (t + 1 weights) times V (t + 1 rows of `head` numbers)
677type WR is Type:
678 WR{v: Array<F32>, o: List<&2, F32>}
679
680def wr_of(v: Array<F32>, al: TA.AL) -> WR:
681 match al:
682 case TA.AL{c, l}:
683 WR{v, l}
684
685def wr_read(+hd: Nat, g: TA.G) -> WR:
686 match g:
687 case TA.G{wa, v, o}:
688 wr_of(v, TA.read_l(Nat.sub(hd, 1n), 0, Nil{}, Array.get(F32, o, 0)))
689
690def weighted(+hd: Nat, +t: Nat, v: Array<F32>, +w: List<&2, F32>) -> WR:
691 wr_read(hd, TA.gemm(1n, hd, 1n+t, U32.from_nat(1n+t), 1, U32.from_nat(hd), 1, TA.fill_list(w, 0, Array.new(F32, TA.cap_depth(1n+t), 0.0)), v, Array.new(F32, TA.cap_depth(hd), 0.0)))
692
693# one query head: softmax(q·K · scale) · V, and the two arrays handed back
694type HO is Type:
695 HO{k: Array<F32>, v: Array<F32>, o: List<&2, F32>}
696
697def HO.then(-R: Type, h: HO, then: Array<F32> -> Array<F32> -> List<&2, F32> -> R) -> R:
698 match h:
699 case HO{k, v, o}:
700 then(k, v, o)
701
702def head_wr(k: Array<F32>, wr: WR) -> HO:
703 match wr:
704 case WR{v, o}:
705 HO{k, v, o}
706
707def head_sr(+hd: Nat, +t: Nat, +scale: F32, v: Array<F32>, sr: SR) -> HO:
708 match sr:
709 case SR{k, s}:
710 head_wr(k, weighted(hd, t, v, T.softmax.l(T.scale.l(scale, s))))
711
712def head_att(+hd: Nat, +t: Nat, +scale: F32, k: Array<F32>, v: Array<F32>, +q: List<&2, F32>) -> HO:
713 head_sr(hd, t, scale, v, scores(hd, t, k, q))
714
715# the query heads of one group over its key/value head, in order; the arrays are threaded through
716def group_att(hs: List<&2, Head>, +hd: Nat, +t: Nat, +scale: F32, k: Array<F32>, v: Array<F32>, +acc: List<&2, F32>) -> HO:
717 match hs:
718 case Nil{}:
719 HO{k, v, acc}
720 case Con{h, rest}:
721 HO.then(HO, head_att(hd, t, scale, k, v, join(h)), k2 => v2 => +o =>
722 group_att(rest, hd, t, scale, k2, v2, List.append(&2, F32, acc, o)))
723
724# every group, in order: the outputs of query heads 0, 1, ..., and the caches handed back
725type AT is Type:
726 AT{ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>, o: List<&2, F32>}
727
728def at_join(p: AT, q: AT) -> AT:
729 match p q:
730 case AT{k1, v1, o1} AT{k2, v2, o2}:
731 AT{List.append(&1, Array<F32>, k1, k2), List.append(&1, Array<F32>, v1, v2), List.append(&2, F32, o1, o2)}
732
733# one key/value head's work: its group of query heads, its keys and its values
734type GW is Type:
735 GW{g: List<&2, Head>, k: Array<F32>, v: Array<F32>}
736
737def gws(gs: List<&2, List<&2, Head>>, ks: List<&1, Array<F32>>, vs: List<&1, Array<F32>>) -> List<&1, GW>:
738 match gs ks vs:
739 case Con{g, gt} Con{k, kt} Con{v, vt}:
740 GW{g, k, v} <> gws(gt, kt, vt)
741 case _ _ _:
742 Nil{}
743
744# some groups one after another, in order
745def seq_att(ws: List<&1, GW>, +hd: Nat, +t: Nat, +scale: F32, aks: List<&1, Array<F32>>, avs: List<&1, Array<F32>>, +acc: List<&2, F32>) -> AT:
746 match ws:
747 case Con{w, wt}:
748 match w:
749 case GW{g, k, v}:
750 HO.then(AT, group_att(g, hd, t, scale, k, v, Nil{}), k2 => v2 => +o =>
751 seq_att(wt, hd, t, scale, k2 <> aks, v2 <> avs, List.append(&2, F32, acc, o)))
752 case Nil{}:
753 AT{List.reverse(&1, Array<F32>, aks), List.reverse(&1, Array<F32>, avs), acc}
754
755# a list of groups cut after its first n (a tail loop: a helper in the fork tree must not hold the stack)
756type GS is Type:
757 GS{l: List<&1, GW>, r: List<&1, GW>}
758
759def cut(n: Nat, xs: List<&1, GW>, acc: List<&1, GW>) -> GS:
760 match n xs:
761 case 1n+m Con{x, xt}:
762 cut(m, xt, x <> acc)
763 case _ _:
764 GS{List.reverse(&1, GW, acc), xs}
765
766# how many times n groups are halved (the larger half kept) until one or none is left
767def halvings(fuel: Nat, +n: Nat) -> Nat:
768 match fuel:
769 case 0n:
770 0n
771 case 1n+f:
772 match n:
773 case 0n:
774 0n
775 case 1n+p:
776 match p:
777 case 0n:
778 0n
779 case 1n+q:
780 1n+halvings(f, Nat.sub(2n+q, TA.half(2n+q)))
781
782# the n groups of gs (cut in two halves) in a tree of parallel lets of depth d, each leaf a half of a node; the
783# outputs and the caches come back in group order, so every head computes what it computed one after another
784def par_att(d: Nat, gs: GS, +n: Nat, +hd: Nat, +t: Nat, +scale: F32) -> AT:
785 match d gs:
786 case 0n GS{l, r}:
787 p q = seq_att(l, hd, t, scale, Nil{}, Nil{}, Nil{}) seq_att(r, hd, t, scale, Nil{}, Nil{}, Nil{})
788 at_join(p, q)
789 case 1n++e GS{l, r}:
790 p q = par_att(e, cut(TA.half(TA.half(n)), l, Nil{}), TA.half(n), hd, t, scale) par_att(e, cut(TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), hd, t, scale)
791 at_join(p, q)
792
793# the attention output and the cache handed back
794type KA<-cap: Nat, -t: Nat> is Type:
795 KA{kv: KV<cap, t>, o: List<&2, F32>}
796
797def ka_of(-cap: Nat, -t: Nat, at: AT) -> KA<cap, t>:
798 match at:
799 case AT{ks, vs, o}:
800 KA{KV{ks, vs}, o}
801
802def KA.then(-R: Type, -cap: Nat, -t: Nat, ka: KA<cap, t>, then: KV<cap, t> -> List<&2, F32> -> R) -> R:
803 match ka:
804 case KA{kv, o}:
805 then(kv, o)
806
807# attention of the (rotated) query heads at position t over the cache of positions 0..t (the newest included).
808# d is the model's parallel depth (depth_for): with d = 0 (two threads or fewer) the key/value heads run one after
809# another, since a fork tree on one thread is only overhead (+5% on the 1.7B, NOTES 2026-10-12); otherwise they run
810# in parallel, one group per leaf. Both orders compute every head the same way, so the bits are the same.
811def attention(+c: Cfg, +d: Nat, -cap: Nat, +t: Nat, qs: List<&2, Head>, kv: KV<cap, 1n+t>) -> KA<cap, 1n+t>:
812 match d kv:
813 case 0n KV{ks, vs}:
814 ka_of(cap, 1n+t, seq_att(gws(groups(Cfg.kv(c), Cfg.g(c), qs), ks, vs), Cfg.head(c), t, (1.0 / F32.sqrt(F32.from_nat(Cfg.head(c))) : F32), Nil{}, Nil{}, Nil{}))
815 case 1n+e KV{ks, vs}:
816 ka_of(cap, 1n+t, par_att(Nat.sub(halvings(Cfg.kv(c), Cfg.kv(c)), 1n), cut(TA.half(Cfg.kv(c)), gws(groups(Cfg.kv(c), Cfg.g(c), qs), ks, vs), Nil{}), Cfg.kv(c), Cfg.head(c), t, (1.0 / F32.sqrt(F32.from_nat(Cfg.head(c))) : F32)))
817
818# =====================================================================================================
819# A layer
820# =====================================================================================================
821
822# the weights of one layer, with their shapes from the config (row = output, as PyTorch stores them)
823type LW<-c: Cfg> is Type:
824 LW{ln1: List<&2, F32>, q: TA.Bands<Cfg.dim(c), Cfg.dim(c)>, k: TA.Bands<Cfg.kvw(c), Cfg.dim(c)>, v: TA.Bands<Cfg.kvw(c), Cfg.dim(c)>, o: TA.Bands<Cfg.dim(c), Cfg.dim(c)>, ln2: List<&2, F32>, gate: TA.Bands<Cfg.ffn(c), Cfg.dim(c)>, up: TA.Bands<Cfg.ffn(c), Cfg.dim(c)>, down: TA.Bands<Cfg.dim(c), Cfg.ffn(c)>}
825
826# what a layer computes at one position, every intermediate: the tests compare them with transformers, and a
827# backward pass would read them
828type Act is Data:
829 Act{n1: List<&2, F32>, q: List<&2, F32>, k: List<&2, F32>, v: List<&2, F32>, qr: List<&2, F32>, kr: List<&2, F32>, attn: List<&2, F32>, a: List<&2, F32>, n2: List<&2, F32>, m: List<&2, F32>, y: List<&2, F32>}
830
831# the layer's output at position t, its cache grown to t + 1 positions, its weights handed back, and the
832# activations
833type BR<-c: Cfg, -cap: Nat, -t: Nat> is Type:
834 BR{y: List<&2, F32>, kv: KV<cap, 1n+t>, lw: LW<c>, act: Act}
835
836# continues with the fields of a layer's result (a computed value cannot be matched; a loop over positions
837# recurses inside k)
838def BR.then(-R: Type, -c: Cfg, -cap: Nat, -t: Nat, b: BR<c, cap, t>, k: List<&2, F32> -> KV<cap, 1n+t> -> LW<c> -> Act -> R) -> R:
839 match b:
840 case BR{y, kv, lw, act}:
841 k(y, kv, lw, act)
842
843def mv.go(-B: Type, -r: Nat, -k: Nat, bl: TA.BL<r, k>, then: TA.Bands<r, k> -> List<&2, F32> -> B) -> B:
844 match bl:
845 case TA.BL{w, y}:
846 then(w, y)
847
848# y = W · x, then `then` with W handed back and y (an Array has one owner, so a product returns its matrix)
849def mv(-B: Type, +r: Nat, +k: Nat, w: TA.Bands<r, k>, +x: List<&2, F32>, then: TA.Bands<r, k> -> List<&2, F32> -> B) -> B:
850 mv.go(B, r, k, TA.Bands.matvec_l(r, k, w, x), then)
851
852# the layer at position t (0-based) on input x, with the cache of positions 0..t-1 and the rotation of t; d is the
853# parallel depth of attention (see attention)
854def block(+c: Cfg, +d: Nat, lw: LW<c>, -cap: Nat, +t: Nat, kv: KV<cap, t>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +rt: Rot, +x: List<&2, F32>) -> BR<c, cap, t>:
855 match lw:
856 case LW{+ln1, wq, wk, wv, wo, +ln2, wg, wu, wd}:
857 +n1 = rmsnorm(c, x, ln1)
858 mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wq, n1, wq => +q =>
859 mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wk, n1, wk => +k =>
860 mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wv, n1, wv => +v =>
861 +qs = ropes(rt, heads(Cfg.heads(c), Cfg.half(c), q))
862 +ks = ropes(rt, heads(Cfg.kv(c), Cfg.half(c), k))
863 KA.then(BR<c, cap, t>, cap, 1n+t, attention(c, d, cap, t, qs, KV.push(c, cap, t, kv, room, joins(ks), chunks(Cfg.kv(c), Cfg.head(c), v))), kv2 => +at =>
864 mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wo, at, wo => +a =>
865 +x1 = T.zip_with.add(x, a)
866 +n2 = rmsnorm(c, x1, ln2)
867 mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wg, n2, wg => +g =>
868 mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wu, n2, wu => +u =>
869 mv(BR<c, cap, t>, Cfg.dim(c), Cfg.ffn(c), wd, swiglu(g, u), wd => +m =>
870 +y = T.zip_with.add(x1, m)
871 BR{y, kv2, LW{ln1, wq, wk, wv, wo, ln2, wg, wu, wd}, Act{n1, q, k, v, flat(qs), flat(ks), at, a, n2, m, y}}))))))))
872
873# =====================================================================================================
874# The model
875# =====================================================================================================
876
877# emb: the embedding (vocab rows; also the LM head, tied); lnf: the final norm's gains
878type Model<-c: Cfg> is Type:
879 Model{d: Nat, emb: TA.Bands<Cfg.vocab(c), Cfg.dim(c)>, layers: List<&1, LW<c>>, lnf: List<&2, F32>}
880
881# the last layer's output, the layers and caches handed back, and every layer's activations
882type Run<-c: Cfg, -cap: Nat, -t: Nat> is Type:
883 Run{x: List<&2, F32>, layers: List<&1, LW<c>>, cs: List<&1, KV<cap, 1n+t>>, acts: List<&2, Act>}
884
885# pend is the previous layer's result: its output is the next layer's input
886def run(+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>:
887 match ls cs pend:
888 case Con{l, lt} Con{k, kt} BR{+y, kv, l0, act}:
889 run(c, d, cap, t, room, rt, lt, kt, l0 <> al, kv <> ac, act <> aa, block(c, d, l, cap, t, k, room, rt, y))
890 case _ _ BR{+y, kv, l0, act}:
891 Run{y, List.reverse(&1, LW<c>, l0 <> al), List.reverse(&1, KV<cap, 1n+t>, kv <> ac), List.reverse(&2, Act, act <> aa)}
892
893def layers(+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>:
894 match ls cs:
895 case Con{l, lt} Con{k, kt}:
896 run(c, d, cap, t, room, rt, lt, kt, Nil{}, Nil{}, Nil{}, block(c, d, l, cap, t, k, room, rt, x))
897 case _ _:
898 Run{x, Nil{}, Nil{}, Nil{}}
899
900# what one position produces: the logits, the final norm, every layer's activations, the caches grown to t + 1
901# positions and the model handed back
902type Out<-c: Cfg, -cap: Nat, -t: Nat> is Type:
903 Out{lg: List<&2, F32>, nf: List<&2, F32>, acts: List<&2, Act>, cs: List<&1, KV<cap, 1n+t>>, m: Model<c>}
904
905# continues with the fields of a position's result: logits, final norm, activations, caches, model
906def 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:
907 match o:
908 case Out{lg, nf, acts, cs, m}:
909 k(lg, nf, acts, cs, m)
910
911# the final norm, then logits = emb · nf (the LM head is the embedding matrix)
912def 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>:
913 match r:
914 case Run{+x, ls, cs, acts}:
915 +nf = rmsnorm(c, x, lnf)
916 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}})
917
918def forward.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>:
919 match r:
920 case TA.BRow{emb, row}:
921 head_of(c, cap, t, d, emb, lnf, layers(c, d, cap, t, room, rot_at(c, t), ls, cs, row))
922
923# token tok at position t (0-based), with the caches of positions 0..t-1 (one per layer) and room for cap positions;
924# room is the proof that t < cap
925def 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>:
926 match m:
927 case Model{+d, emb, ls, +lnf}:
928 forward.go(c, cap, t, room, d, ls, lnf, cs, TA.Bands.read_row(Cfg.vocab(c), Cfg.dim(c), emb, tok))
929
930# n empty caches, one per layer, each with room for cap positions
931def caches(+c: Cfg, n: Nat, +cap: Nat) -> List<&1, KV<cap, 0n>>:
932 match n:
933 case 0n:
934 Nil{}
935 case 1n+p:
936 KV.empty(c, cap) <> caches(c, p, cap)
937
938# the first index of the largest value (torch.argmax does the same), and the value
939type AM is Data:
940 AM{i: Nat, v: F32}
941
942def AM.i(am: AM) -> Nat:
943 match am:
944 case AM{i, v}:
945 i
946
947def AM.v(am: AM) -> F32:
948 match am:
949 case AM{i, v}:
950 v
951
952def gt_head(xs: List<&2, F32>, +v: F32) -> Bool:
953 match xs:
954 case Nil{}:
955 False{}
956 case Con{+h, t}:
957 F32.is_gt(h, v)
958
959def argmax.go(xs: List<&2, F32>, g: Bool, +i: Nat, +bi: Nat, +bv: F32) -> AM:
960 match xs g:
961 case Nil{} _:
962 AM{bi, bv}
963 case Con{+h, +t} True{}:
964 argmax.go(t, gt_head(t, h), 1n+i, i, h)
965 case Con{h, +t} False{}:
966 argmax.go(t, gt_head(t, bv), 1n+i, bi, bv)
967
968def argmax(xs: List<&2, F32>) -> AM:
969 match xs:
970 case Nil{}:
971 AM{0n, 0.0}
972 case Con{+h, +t}:
973 argmax.go(t, gt_head(t, h), 1n, 0n, h)
974
975# =====================================================================================================
976# Loading (the weights as written by llama.bend's reference/prep.py: <dir>/w/<name>.bin, little-endian F32,
977# one file per tensor, row = output; the config in <dir>/cfg.txt). The block reader is bend-ml's (MIT).
978# =====================================================================================================
979
980# 4 little-endian bytes -> F32 (U32 and F32 hold the same Word(32n): it is a bitcast)
981def bitcast(u: U32) -> F32:
982 match u:
983 case U32{w}:
984 F32{w}
985
986def f32_le(+b0: U32, +b1: U32, +b2: U32, +b3: U32) -> F32:
987 bitcast((b0 + (b1 * 256 : U32) + (b2 * 65536 : U32) + (b3 * 16777216 : U32) : U32))
988
989def f32s(bs: List<&2, U32>) -> List<&2, F32>:
990 match bs:
991 case Con{b0, Con{b1, Con{b2, Con{b3, t}}}}:
992 f32_le(b0, b1, b2, b3) <> f32s(t)
993 case _:
994 Nil{}
995
996# reads the file in blocks of `blk` bytes; k blocks remaining; m is the block already read
997def 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>):
998 match k m:
999 case 0n Tuple{f, r}:
1000 do IO<List<&2, F32>>:
1001 File.close(f)
1002 return List.concat(&2, F32, List.reverse(&2, List<&2, F32>, acc))
1003 case 1n+p Tuple{f, r}:
1004 do IO<List<&2, F32>>:
1005 bs : List<&2, U32> <- IO.pass(List<&2, U32>, r)
1006 m2 : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, (off + blk : U32), blk)
1007 rf_loop(p, (off + blk : U32), blk, m2, f32s(bs) <> acc)
1008
1009def blocks_of(n: Nat, +per: Nat) -> Nat:
1010 Nat.div(Nat.add(n, Nat.sub(per, 1n)), per)
1011
1012# n F32 numbers (little-endian) from a file
1013def read_floats(path: String, +n: Nat) -> IO(List<&2, F32>):
1014 do IO<List<&2, F32>>:
1015 f : File <- IO.try(File, File.open(path, "r"))
1016 m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, 0, 1048576)
1017 fl : List<&2, F32> <- rf_loop(blocks_of(Nat.mul(n, 4n), 1048576n), 0, 1048576, m, Nil{})
1018 return List.take(&2, F32, fl, n)
1019
1020# the weight matrices are kept in 2^d row bands (TA.Bands); each matrix · vector product runs its bands in
1021# parallel without copying the weights (bend-ml NOTES exp. 11). With 2 threads or fewer the bands cannot help,
1022# so d = 0: a single band, which computes what the sequential Mat.matmul_nt computes (bend-ml exp. 14).
1023def depth_sel(few: Bool) -> Nat:
1024 match few:
1025 case True{}:
1026 0n
1027 case False{}:
1028 4n
1029
1030def depth_for(+threads: U32) -> Nat:
1031 depth_sel(U32.is_le(threads, 2))
1032
1033def nm(+i: Nat, f: String) -> String:
1034 "h" ++ Nat.show(i) ++ "." ++ f
1035
1036# ---- the boundary: every weight file must hold exactly the numbers its shape needs ----
1037
1038def size_cmp(ok: Bool, +path: String, +got: U32, +want: Nat) -> IO(Unit):
1039 match ok:
1040 case True{}:
1041 IO.pure(Unit, Unit{})
1042 case False{}:
1043 IO.die(Unit, 1, path ++ ": expected " ++ Nat.show(want) ++ " bytes, found " ++ U32.show(got) ++ "; run make prep again")
1044
1045def size_ok(m: File & Result<&1, &1, U32 & String, U32>, +want: Nat, +path: String) -> IO(Unit):
1046 match m:
1047 case Tuple{f, r}:
1048 do IO<Unit>:
1049 +n : U32 <- IO.pass(U32, r)
1050 File.close(f)
1051 size_cmp(U32.is_eq(n, U32.from_nat(want)), path, n, want)
1052
1053# stops the program unless the file holds exactly n F32 numbers (n * 4 bytes)
1054def check_size(+path: String, +n: Nat) -> IO(Unit):
1055 do IO<Unit>:
1056 f : File <- IO.try(File, File.open(path, "r"))
1057 m : File & Result<&1, &1, U32 & String, U32> <- File.size(f)
1058 size_ok(m, Nat.mul(n, 4n), path)
1059
1060def wpath(dir: String, name: String) -> String:
1061 dir ++ "/w/" ++ name ++ ".bin"
1062
1063def get_list(+dir: String, +name: String, +n: Nat) -> IO(List<&2, F32>):
1064 do IO<List<&2, F32>>:
1065 check_size(wpath(dir, name), n)
1066 read_floats(wpath(dir, name), n)
1067
1068# reads the file in 1 MB blocks and writes each block straight into the matrix, so only one
1069# block is a list at any time
1070# ---- matrices kept in bands: each band is read from its own range of the file ----
1071
1072# writes the little-endian F32 numbers held in bs (4 bytes each) at index i, without building a
1073# list of F32 first
1074def fill_bytes(bs: List<&2, U32>, +i: U32, a: Array<F32>) -> Array<F32>:
1075 match bs:
1076 case Con{+b0, Con{+b1, Con{+b2, Con{+b3, t}}}}:
1077 fill_bytes(t, (i + 1 : U32), Array.set(F32, a, i, f32_le(b0, b1, b2, b3)))
1078 case _:
1079 a
1080
1081def fill_bytes_mat(-r: Nat, -c: Nat, +i: U32, bs: List<&2, U32>, m: TA.Mat<r, c>) -> TA.Mat<r, c>:
1082 match m:
1083 case TA.Mat{d}:
1084 TA.Mat{fill_bytes(bs, i, d)}
1085
1086# reads blocks of at most 1 MB until rem bytes have been read, so a band never takes numbers
1087# that belong to the next one
1088def 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>):
1089 match k m:
1090 case 0n Tuple{f, rs}:
1091 do IO<TA.Mat<r, c>>:
1092 File.close(f)
1093 return acc
1094 case 1n+p Tuple{f, rs}:
1095 do IO<TA.Mat<r, c>>:
1096 bs : List<&2, U32> <- IO.pass(List<&2, U32>, rs)
1097 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))))
1098 stream_exact(p, r, c, Nat.sub(rem, 1048576n), (off + 1048576 : U32), (idx + 262144 : U32), m2, fill_bytes_mat(r, c, idx, bs, acc))
1099
1100# rr x c numbers starting at number `start` of the file, into a band of r rows (rr == r, see
1101# load_tree; the Array is sized for rr * c numbers as Mat.zeros would size it)
1102def get_band(path: String, -r: Nat, +rr: Nat, +c: Nat, +start: Nat) -> IO(TA.Mat<r, c>):
1103 do IO<TA.Mat<r, c>>:
1104 f : File <- IO.try(File, File.open(path, "r"))
1105 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))))
1106 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)})
1107
1108# fills the bands of b (whose shape comes from TA.Bands.zeros) in row order; rr is r at run time
1109def 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>):
1110 match b:
1111 case TA.BLeaf{m}:
1112 do IO<TA.Bands<r, c>>:
1113 band : TA.Mat<r, c> <- get_band(path, r, rr, c, start)
1114 return TA.BLeaf{band}
1115 case TA.BNode{x, y}:
1116 do IO<TA.Bands<r, c>>:
1117 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)
1118 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)))
1119 return TA.BNode{x2, y2}
1120
1121def get_bands(+dir: String, +name: String, +d: Nat, +r: Nat, +c: Nat) -> IO(TA.Bands<r, c>):
1122 do IO<TA.Bands<r, c>>:
1123 check_size(wpath(dir, name), Nat.mul(r, c))
1124 load_tree(c, r, TA.Bands.zeros(d, r, c), r, {==}, wpath(dir, name), 0n)
1125
1126# one layer's weights; d: the band depth (depth_for)
1127def load_lw(+c: Cfg, +dir: String, +i: Nat, +d: Nat) -> IO(LW<c>):
1128 do IO<LW<c>>:
1129 ln1 : List<&2, F32> <- get_list(dir, nm(i, "ln1"), Cfg.dim(c))
1130 q : TA.Bands<Cfg.dim(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "q"), d, Cfg.dim(c), Cfg.dim(c))
1131 k : TA.Bands<Cfg.kvw(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "k"), d, Cfg.kvw(c), Cfg.dim(c))
1132 v : TA.Bands<Cfg.kvw(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "v"), d, Cfg.kvw(c), Cfg.dim(c))
1133 o : TA.Bands<Cfg.dim(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "o"), d, Cfg.dim(c), Cfg.dim(c))
1134 ln2 : List<&2, F32> <- get_list(dir, nm(i, "ln2"), Cfg.dim(c))
1135 gate : TA.Bands<Cfg.ffn(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "gate"), d, Cfg.ffn(c), Cfg.dim(c))
1136 up : TA.Bands<Cfg.ffn(c), Cfg.dim(c)> <- get_bands(dir, nm(i, "up"), d, Cfg.ffn(c), Cfg.dim(c))
1137 down : TA.Bands<Cfg.dim(c), Cfg.ffn(c)> <- get_bands(dir, nm(i, "down"), d, Cfg.dim(c), Cfg.ffn(c))
1138 return LW{ln1, q, k, v, o, ln2, gate, up, down}
1139
1140def load_layers(k: Nat, +c: Cfg, +dir: String, +d: Nat, +i: Nat, acc: List<&1, LW<c>>) -> IO(List<&1, LW<c>>):
1141 match k:
1142 case 0n:
1143 IO.pure(List<&1, LW<c>>, List.reverse(&1, LW<c>, acc))
1144 case 1n+p:
1145 do IO<List<&1, LW<c>>>:
1146 l : LW<c> <- load_lw(c, dir, i, d)
1147 load_layers(p, c, dir, d, 1n+i, l <> acc)
1148
1149# the whole model; the embedding gets one more level of bands than the layers (2^5 with 16 threads)
1150def load_model(+c: Cfg, +dir: String) -> IO(Model<c>):
1151 do IO<Model<c>>:
1152 threads : U32 <- IO.thread_count()
1153 +d : Nat = depth_for(threads)
1154 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))
1155 ls : List<&1, LW<c>> <- load_layers(Cfg.layers(c), c, dir, d, 0n, Nil{})
1156 lnf : List<&2, F32> <- get_list(dir, "lnf", Cfg.dim(c))
1157 return Model{d, emb, ls, lnf}
1158
1159# ---- the config file: kv g half f layers vocab, then the bits of eps and theta as F32, in decimal ----
1160
1161def rb_done(m: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(List<&2, U32>):
1162 (f, r) = m
1163 do IO<List<&2, U32>>:
1164 bs : List<&2, U32> <- IO.pass(List<&2, U32>, r)
1165 File.close(f)
1166 return bs
1167
1168def read_bytes(path: String) -> IO(List<&2, U32>):
1169 do IO<List<&2, U32>>:
1170 f : File <- IO.try(File, File.open(path, "r"))
1171 m : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_bytes(f, 1048576)
1172 rb_done(m)
1173
1174def digit_head(xs: List<&2, U32>) -> Bool:
1175 match xs:
1176 case Nil{}:
1177 False{}
1178 case Con{+b, t}:
1179 Bool.and(U32.is_le(48, b), U32.is_le(b, 57))
1180
1181# the decimal numbers in bs, separated by anything that is not a digit; dg says whether bs starts with a digit
1182def u32s.go(bs: List<&2, U32>, dg: Bool, +cur: U32, have: Bool, acc: List<&2, U32>) -> List<&2, U32>:
1183 match bs dg have:
1184 case Nil{} _ True{}:
1185 List.reverse(&2, U32, cur <> acc)
1186 case Nil{} _ False{}:
1187 List.reverse(&2, U32, acc)
1188 case Con{+b, +t} True{} _:
1189 u32s.go(t, digit_head(t), ((cur * 10 : U32) + U32.sub(b, 48) : U32), True{}, acc)
1190 case Con{b, +t} False{} True{}:
1191 u32s.go(t, digit_head(t), 0, False{}, cur <> acc)
1192 case Con{b, +t} False{} False{}:
1193 u32s.go(t, digit_head(t), 0, False{}, acc)
1194
1195def u32s(+bs: List<&2, U32>) -> List<&2, U32>:
1196 u32s.go(bs, digit_head(bs), 0, False{}, Nil{})
1197
1198def cfg_of(+path: String, xs: List<&2, U32>) -> IO(Cfg):
1199 match xs:
1200 case Con{+kv, Con{+g, Con{+half, Con{+f, Con{+l, Con{+v, Con{+e, Con{+t, Nil{}}}}}}}}}:
1201 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)})
1202 case _:
1203 IO.die(Cfg, 1, path ++ ": expected 8 numbers (kv g half f layers vocab eps_bits theta_bits); run make prep again")
1204
1205# the config of the model in dir (written by reference/prep.py)
1206def read_cfg(+dir: String) -> IO(Cfg):
1207 do IO<Cfg>:
1208 bs : List<&2, U32> <- read_bytes(dir ++ "/cfg.txt")
1209 cfg_of(dir ++ "/cfg.txt", u32s(bs))