bend-ml-llama Releases (4): 0.2.1.0 0x171194e0c744bf1ffdcc47ea33747d37 indexed https://jend.leeeet.dev/files.txt?path=bend-ml-llama@0.2.1.0 0.2.0.0 0x7ac61588953ec1faa925002b4f879bcf unindexed https://jend.leeeet.dev/files.txt?path=bend-ml-llama@0.2.0.0 0.1.1.0 0x8ff6f11df9bf3e5a1f5eaa324b7fef8f unindexed https://jend.leeeet.dev/files.txt?path=bend-ml-llama@0.1.1.0 0.1.0.0 0x60692cfeaba1e21c25f3b68061e41fea unindexed https://jend.leeeet.dev/files.txt?path=bend-ml-llama@0.1.0.0 Definitions (213, 1 to 100 shown): BendMlLlama.add_zero (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=add_zero) law add_zero: for a: Nat {Nat.add(a, 0n) == a : Nat} BendMlLlama.add_succ (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=add_succ) law add_succ: for a: Nat for -b: Nat {1n+Nat.add(a, b) == Nat.add(a, 1n+b) : Nat} BendMlLlama.add_one (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=add_one) law add_one: for a: Nat {Nat.add(a, 1n) == 1n+a : Nat} BendMlLlama.add_assoc (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=add_assoc) law add_assoc: for a: Nat for -b: Nat for -c: Nat {Nat.add(Nat.add(a, b), c) == Nat.add(a, Nat.add(b, c)) : Nat} BendMlLlama.Cfg (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg) # kv: key/value heads; g: query heads per key/value head; half: half a head (a head has 2·half numbers, so its # size is even by construction); f: the MLP's inner width; layers; vocab: the vocabulary (the embedding has # vocab rows, and the LM head is the embedding); eps: RMSNorm's epsilon; theta: RoPE's base. type Cfg is Data: Cfg{kv: Nat, g: Nat, half: Nat, f: Nat, layers: Nat, vocab: Nat, eps: F32, theta: F32} BendMlLlama.head (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=head) # a head: 2·half numbers def head(+half: Nat) -> Nat: BendMlLlama.Cfg.half (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.half) def Cfg.half(+c: Cfg) -> Nat: BendMlLlama.Cfg.kv (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.kv) def Cfg.kv(+c: Cfg) -> Nat: BendMlLlama.Cfg.g (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.g) def Cfg.g(+c: Cfg) -> Nat: BendMlLlama.Cfg.ffn (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.ffn) def Cfg.ffn(+c: Cfg) -> Nat: BendMlLlama.Cfg.layers (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.layers) def Cfg.layers(+c: Cfg) -> Nat: BendMlLlama.Cfg.vocab (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.vocab) def Cfg.vocab(+c: Cfg) -> Nat: BendMlLlama.Cfg.eps (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.eps) def Cfg.eps(+c: Cfg) -> F32: BendMlLlama.Cfg.theta (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.theta) def Cfg.theta(+c: Cfg) -> F32: BendMlLlama.Cfg.head (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.head) # the size of a head def Cfg.head(+c: Cfg) -> Nat: BendMlLlama.Cfg.heads (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.heads) # the query heads: kv groups of g def Cfg.heads(+c: Cfg) -> Nat: BendMlLlama.Cfg.dim (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.dim) # the hidden size: one head per query head def Cfg.dim(+c: Cfg) -> Nat: BendMlLlama.Cfg.kvw (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.kvw) # the width of the keys (and of the values): one head per key/value head def Cfg.kvw(+c: Cfg) -> Nat: BendMlLlama.smollm2_135m (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=smollm2_135m) # SmolLM2-135M (HuggingFaceTB/SmolLM2-135M): 3 key/value heads of 3 query heads each, heads of 64 def smollm2_135m() -> Cfg: BendMlLlama.smollm2_360m (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=smollm2_360m) # SmolLM2-360M: 5 key/value heads of 3 query heads each, heads of 64 def smollm2_360m() -> Cfg: BendMlLlama.smollm2_1_7b (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=smollm2_1_7b) # SmolLM2-1.7B: 32 key/value heads of one query head each (plain multi-head attention), heads of 64 def smollm2_1_7b() -> Cfg: BendMlLlama.smollm2_135m_sizes (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=smollm2_135m_sizes) # LAW (the configs): SmolLM2-135M has 9 query heads of 64, a hidden size of 576 and keys of 192; SmolLM2-360M # 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 # size of 2048 and keys of 2048 (as many key/value heads as query heads), as their config.json files say. law smollm2_135m_sizes: {Cfg.dim(smollm2_135m()) == 576n : Nat} & {Cfg.heads(smollm2_135m()) == 9n : Nat} & {Cfg.head(smollm2_135m()) == 64n : Nat} & {Cfg.kvw(smollm2_135m()) == 192n : Nat} BendMlLlama.smollm2_360m_sizes (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=smollm2_360m_sizes) law smollm2_360m_sizes: {Cfg.dim(smollm2_360m()) == 960n : Nat} & {Cfg.heads(smollm2_360m()) == 15n : Nat} & {Cfg.head(smollm2_360m()) == 64n : Nat} & {Cfg.kvw(smollm2_360m()) == 320n : Nat} BendMlLlama.smollm2_1_7b_sizes (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=smollm2_1_7b_sizes) law smollm2_1_7b_sizes: {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} BendMlLlama.n32 (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=n32) def n32(+n: Nat) -> U32: BendMlLlama.Cfg.layer_size (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.layer_size) # the numbers of one layer, from the shapes of LW: two norms, q and o (dim × dim), k and v (kvw × dim), # gate and up (ffn × dim) and down (dim × ffn) def Cfg.layer_size(+c: Cfg) -> U32: BendMlLlama.Cfg.size (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.size) # the numbers of the model, from the shapes of Model: the embedding (vocab × dim), the layers, the final norm def Cfg.size(+c: Cfg) -> U32: BendMlLlama.layout_135m (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=layout_135m) # LAW (layout): the shapes the types give to SmolLM2-135M's weights add up to 134,515,008 numbers, # 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). # The loader checks every file's size against the shape in its type, so together the files cover the checkpoint # exactly. In U32, whose sums wrap at 2^32; every partial sum here is below it, so this is the exact count. law layout_135m: {Cfg.size(smollm2_135m()) == (134515008 : U32) : U32} BendMlLlama.layout_360m (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=layout_360m) law layout_360m: {Cfg.size(smollm2_360m()) == (361821120 : U32) : U32} BendMlLlama.layout_1_7b (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=layout_1_7b) law layout_1_7b: {Cfg.size(smollm2_1_7b()) == (1711376384 : U32) : U32} BendMlLlama.Head (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Head) # a head as its two halves: RoPE rotates number i of the first half with number i of the second type Head is Data: Head{lo: List<&2, F32>, hi: List<&2, F32>} BendMlLlama.split (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=split) def split(+half: Nat, +xs: List<&2, F32>) -> Head: BendMlLlama.join (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=join) def join(h: Head) -> List<&2, F32>: BendMlLlama.heads (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=heads) # n heads of 2·half numbers from xs, in order def heads(n: Nat, +half: Nat, +xs: List<&2, F32>) -> List<&2, Head>: BendMlLlama.chunks (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=chunks) # n plain heads (the values, which RoPE does not touch) def chunks(n: Nat, +w: Nat, +xs: List<&2, F32>) -> List<&2, List<&2, F32>>: BendMlLlama.joins (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=joins) def joins(hs: List<&2, Head>) -> List<&2, List<&2, F32>>: BendMlLlama.flat (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=flat) def flat(hs: List<&2, Head>) -> List<&2, F32>: BendMlLlama.head_split (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=head_split) # LAW (head_split): splitting a list at half and joining the two parts gives the list back, so cutting a head # into the halves that RoPE pairs up loses and reorders nothing. law head_split: for +xs: List<&2, F32> for +half: Nat {join(split(half, xs)) == xs : List<&2, F32>} BendMlLlama.list.add (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.add) # xs + ys, element by element (to the shorter length) def list.add(xs: List<&2, F32>, ys: List<&2, F32>) -> List<&2, F32>: BendMlLlama.list.mul (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.mul) # xs * ys, element by element (to the shorter length) def list.mul(xs: List<&2, F32>, ys: List<&2, F32>) -> List<&2, F32>: BendMlLlama.list.scale (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.scale) # k * xs def list.scale(+k: F32, xs: List<&2, F32>) -> List<&2, F32>: BendMlLlama.list.sum.go (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.sum.go) # acc + xs, a running sum from the left def list.sum.go(xs: List<&2, F32>, acc: F32) -> F32: BendMlLlama.list.max.go (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.max.go) def list.max.go(xs: List<&2, F32>, cur: F32) -> F32: BendMlLlama.list.max (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.max) # the maximum of xs (0 for an empty list) def list.max(xs: List<&2, F32>) -> F32: BendMlLlama.list.exp_shift (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.exp_shift) # exp(x - m) for each x def list.exp_shift(xs: List<&2, F32>, +m: F32) -> List<&2, F32>: BendMlLlama.list.div (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.div) # xs / s def list.div(xs: List<&2, F32>, +s: F32) -> List<&2, F32>: BendMlLlama.list.softmax.fin (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.softmax.fin) def list.softmax.fin(+es: List<&2, F32>) -> List<&2, F32>: BendMlLlama.list.softmax (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.softmax) # stable softmax: subtracts the maximum before exponentiating def list.softmax(+xs: List<&2, F32>) -> List<&2, F32>: BendMlLlama.squares (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=squares) # the squares of xs def squares(xs: List<&2, F32>) -> List<&2, F32>: BendMlLlama.pair_sums (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=pair_sums) # adjacent pairs added: [a, b, c, d, e] -> [a + b, c + d, e] def pair_sums(xs: List<&2, F32>) -> List<&2, F32>: BendMlLlama.tree_sum (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=tree_sum) # the sum as a balanced tree of additions (pairwise summation): its rounding error grows with log2(n), not n, as # in PyTorch's cascade sum. A running sum loses the small squares next to a large one (NOTES, 2026-10-10). def tree_sum(fuel: Nat, xs: List<&2, F32>) -> F32: BendMlLlama.rmsnorm (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=rmsnorm) # RMSNorm: g * (x * (1 / sqrt(mean(x²) + eps))), in this order (LlamaRMSNorm.forward) def rmsnorm(+c: Cfg, +x: List<&2, F32>, +g: List<&2, F32>) -> List<&2, F32>: BendMlLlama.inv_freq.go (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=inv_freq.go) # inv_freq[i] = 1 / (theta ^ (2i / head)), i = 0..half-1 (LlamaRotaryEmbedding, default rope) def inv_freq.go(n: Nat, +i: Nat, +theta: F32, +hd: F32) -> List<&2, F32>: BendMlLlama.inv_freq (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=inv_freq) def inv_freq(+c: Cfg) -> List<&2, F32>: BendMlLlama.angles (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=angles) def angles(+t: F32, fs: List<&2, F32>) -> List<&2, F32>: BendMlLlama.cos.l (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=cos.l) def cos.l(xs: List<&2, F32>) -> List<&2, F32>: BendMlLlama.sin.l (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=sin.l) def sin.l(xs: List<&2, F32>) -> List<&2, F32>: BendMlLlama.Rot (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Rot) # the cosines and sines of position t, one per frequency (shared by every head of every layer) type Rot is Data: Rot{cs: List<&2, F32>, sn: List<&2, F32>} BendMlLlama.rot_at (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=rot_at) def rot_at(+c: Cfg, +t: Nat) -> Rot: BendMlLlama.rot_lo (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=rot_lo) # first half: lo*cos + (-hi)*sin; second half: hi*cos + lo*sin (q*cos + rotate_half(q)*sin) def rot_lo(h1: List<&2, F32>, h2: List<&2, F32>, cs: List<&2, F32>, sn: List<&2, F32>) -> List<&2, F32>: BendMlLlama.rot_hi (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=rot_hi) def rot_hi(h1: List<&2, F32>, h2: List<&2, F32>, cs: List<&2, F32>, sn: List<&2, F32>) -> List<&2, F32>: BendMlLlama.rope (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=rope) def rope(rt: Rot, h: Head) -> Head: BendMlLlama.ropes (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=ropes) def ropes(+rt: Rot, hs: List<&2, Head>) -> List<&2, Head>: BendMlLlama.silu (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=silu) # silu(x) = x / (1 + exp(-x)), the form ATen's CPU kernel uses def silu(+x: F32) -> F32: BendMlLlama.swiglu (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=swiglu) # silu(gate) * up def swiglu(g: List<&2, F32>, u: List<&2, F32>) -> List<&2, F32>: BendMlLlama.KV (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=KV) # One pair of arrays per key/value head: the keys (after RoPE) and the values of positions 0..t-1, position p at # [p·head, p·head + head). The type indexes are the capacity cap (the positions the arrays have room for, set by # KV.empty) and t, the number of positions written. KV.push asks for a proof that t < cap, so a cache is never # written past its room (kv_room says why that is enough). type KV<-cap: Nat, -t: Nat> is Type: KV{ks: List<&1, Array>, vs: List<&1, Array>} BendMlLlama.arrays (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=arrays) # n arrays of 2^d slots def arrays(n: Nat, +d: Nat) -> List<&1, Array>: BendMlLlama.KV.empty (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=KV.empty) # a cache with no position yet, with room for cap positions per key/value head def KV.empty(+c: Cfg, +cap: Nat) -> KV: BendMlLlama.write_rows (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=write_rows) # writes row j of rs into array j, at off def write_rows(arrs: List<&1, Array>, rs: List<&2, List<&2, F32>>, +off: U32) -> List<&1, Array>: BendMlLlama.KV.push (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=KV.push) # position t: the keys and values of every key/value head written at t·head; room is the proof that t < cap def KV.push(+c: Cfg, -cap: Nat, +t: Nat, kv: KV, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +ks: List<&2, List<&2, F32>>, +vs: List<&2, List<&2, F32>>) -> KV: BendMlLlama.kv_count (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=kv_count) # the number of key/value heads a cache holds def kv_count(-cap: Nat, -t: Nat, kv: KV) -> Nat: BendMlLlama.arrays_len (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=arrays_len) # LAW (arrays_len): arrays(n, d) makes exactly n arrays. law arrays_len: for +n: Nat for +d: Nat {List.length(&1, Array, arrays(n, d)) == n : Nat} BendMlLlama.kv_heads (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=kv_heads) # LAW (kv_heads): an empty cache holds one array of keys (and one of values) per key/value head, so key/value head j # of the model reads column j and no other. law kv_heads: for +c: Cfg for +cap: Nat {kv_count(cap, 0n, KV.empty(c, cap)) == Cfg.kv(c) : Nat} BendMlLlama.leb_trans (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=leb_trans) # a <= b and b <= c -> a <= c def leb_trans(a: Nat, b: Nat, c: Nat, e1: {True{} == TA.leb(a, b) : Bool}, e2: {True{} == TA.leb(b, c) : Bool}) -> {True{} == TA.leb(a, c) : Bool}: BendMlLlama.leb_mul (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=leb_mul) # a <= b -> a·h <= b·h def leb_mul(a: Nat, b: Nat, +h: Nat, e: {True{} == TA.leb(a, b) : Bool}) -> {True{} == TA.leb(Nat.mul(a, h), Nat.mul(b, h)) : Bool}: BendMlLlama.kv_room (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=kv_room) # LAW (kv_room): KV.empty(c, cap) gives each key/value head arrays of 2^cap_depth(cap·head) slots, and KV.push at # position t writes head numbers from t·head. For a position below the capacity (t < cap) they all fall inside the # slots: (1 + t)·head <= 2^cap_depth(cap·head), so a push never wraps around onto position 0. (What stays trusted is # 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 # compile (tests/bad_cap.bend). law kv_room: for +cap: Nat for +h: Nat for +t: Nat for e: {True{} == TA.leb(1n+t, cap) : Bool} {True{} == TA.leb(Nat.mul(1n+t, h), TA.pow2(TA.cap_depth(Nat.mul(cap, h)))) : Bool} BendMlLlama.step_room (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=step_room) # For callers that walk the positions with a counter of steps left (fuel), keeping t + fuel = cap: with 1 + p steps # left, position t is below the capacity (step_room), and after the step t + 1 + p = cap still holds (step_inv). def step_room(+t: Nat, +p: Nat, -cap: Nat, e: {Nat.add(t, 1n+p) == cap : Nat}) -> {True{} == TA.leb(1n+t, cap) : Bool}: BendMlLlama.step_inv (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=step_inv) def step_inv(+t: Nat, -p: Nat, -cap: Nat, e: {Nat.add(t, 1n+p) == cap : Nat}) -> {Nat.add(1n+t, p) == cap : Nat}: BendMlLlama.groups (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=groups) # ===================================================================================================== # Grouped-query attention # ===================================================================================================== # the query heads in n groups of g, in order: group j holds query heads j·g .. j·g + g - 1 def groups(n: Nat, +g: Nat, +hs: List<&2, Head>) -> List<&2, List<&2, Head>>: BendMlLlama.groups_len (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=groups_len) # LAW (groups_len): the query heads make exactly as many groups as there are key/value heads, so attention # (which walks the groups and the columns together) pairs every group with a key/value head. law groups_len: for n: Nat for +g: Nat for +hs: List<&2, Head> {List.length(&2, List<&2, Head>, groups(n, g, hs)) == n : Nat} BendMlLlama.div_block (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=div_block) # one block of 1 + m numbers in Nat.divmod: the counter m runs down to 0, then the quotient d grows by one law div_block: for +m: Nat for -n: Nat for -d: Nat for -r: Nat {Nat.divmod.go(Nat.add(1n+m, n), m, d, r) == Nat.divmod.go(n, Nat.add(m, r), 1n+d, 0n) : Nat & Nat} BendMlLlama.div_blocks (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=div_blocks) # j whole blocks of 1 + s numbers add j to the quotient law div_blocks: for +j: Nat for +s: Nat for +n: Nat for +d: Nat {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} BendMlLlama.div_rest (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=div_rest) # fewer than 1 + s numbers after the last whole block leave the quotient alone law div_rest: for +i: Nat for -k: Nat for -d: Nat for -r: Nat {Pair.fst(Nat, Nat, Nat.divmod.go(i, Nat.add(i, k), d, r)) == d : Nat} BendMlLlama.gqa_head (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=gqa_head) # 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 # head i of group j, reads key/value head (j·g + i) / g = j. That division is how transformers maps query heads # to key/value heads (repeat_kv), so grouping the heads computes the same attention. law gqa_head: for +j: Nat for +i: Nat for +k: Nat {Nat.div(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat} BendMlLlama.SR (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=SR) # the scores of one query head against the keys of its key/value head, and the keys handed back type SR is Type: SR{k: Array, s: List<&2, F32>} BendMlLlama.sr_len (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=sr_len) def sr_len(sr: SR) -> Nat: BendMlLlama.sr_of (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=sr_of) def sr_of(k: Array, al: TA.AL) -> SR: BendMlLlama.read_len (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=read_len) # LAW (read_len): reading m + 1 numbers of an array onto acc gives m + 1 more scores. # PROOF: by induction on m. With 0 left it reads one number onto acc (1 + the length of acc, by the definitions); # 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 # leaf_len, for scores). law read_len: for k: Array for +m: Nat for +idx: U32 for +acc: List<&2, F32> for rc: Array & F32 {sr_len(sr_of(k, TA.read_l(m, idx, acc, rc))) == Nat.add(1n+m, List.length(&2, F32, acc)) : Nat} BendMlLlama.sr_read (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=sr_read) # read the 1 + t scores the gemm wrote (read_l reads from the last index down, so it starts at t) def sr_read(+t: Nat, g: TA.G) -> SR: BendMlLlama.scores (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=scores) # q·k for the keys of positions 0..t, oldest first: K (t + 1 rows of `head` numbers) times q, the product # Mat.matmul_nt computes for one row (bend-ml-tensor-array's gemm) def scores(+hd: Nat, +t: Nat, k: Array, +q: List<&2, F32>) -> SR: BendMlLlama.sr_len_g (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=sr_len_g) def sr_len_g(+t: Nat, g: TA.G) -> {sr_len(sr_read(t, g)) == 1n+t : Nat}: BendMlLlama.causal (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=causal) # LAW (causal): attention at position t scores exactly t + 1 keys, the positions 0..t: one score per position # written, none for a position not yet computed. # PROOF: the gemm's result has one constructor, so matching it opens it without computing the product # (sr_len_g); read_len counts the 1 + t numbers read back. law causal: for +hd: Nat for +t: Nat for k: Array for +q: List<&2, F32> {sr_len(scores(hd, t, k, q)) == 1n+t : Nat} BendMlLlama.WR (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=WR) # the weighted sum of the values of positions 0..t: w (t + 1 weights) times V (t + 1 rows of `head` numbers) type WR is Type: WR{v: Array, o: List<&2, F32>} BendMlLlama.wr_of (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=wr_of) def wr_of(v: Array, al: TA.AL) -> WR: BendMlLlama.wr_read (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=wr_read) def wr_read(+hd: Nat, g: TA.G) -> WR: BendMlLlama.weighted (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=weighted) def weighted(+hd: Nat, +t: Nat, v: Array, +w: List<&2, F32>) -> WR: BendMlLlama.HO (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=HO) # one query head: softmax(q·K · scale) · V, and the two arrays handed back type HO is Type: HO{k: Array, v: Array, o: List<&2, F32>} BendMlLlama.HO.then (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=HO.then) def HO.then(-R: Type, h: HO, then: Array -> Array -> List<&2, F32> -> R) -> R: BendMlLlama.head_wr (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=head_wr) def head_wr(k: Array, wr: WR) -> HO: BendMlLlama.head_sr (https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=head_sr) def head_sr(+hd: Nat, +t: Nat, +scale: F32, v: Array, sr: SR) -> HO: More: https://jend.leeeet.dev/packages/bend-ml-llama.txt?definitions_offset=100