bend-ml-llama

Search uses release 0.2.0.0 · 175 definitions

Releases (3)

0.2.0.0 indexed
0.1.1.0 unindexed
0.1.0.0 unindexed

Definitions (175)

main.bend

law add_zero
law add_succ
law add_one
law add_assoc
type Cfgkv: 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.
def heada head: 2·half numbers
def Cfg.half
def Cfg.kv
def Cfg.g
def Cfg.ffn
def Cfg.layers
def Cfg.vocab
def Cfg.eps
def Cfg.theta
def Cfg.headthe size of a head
def Cfg.headsthe query heads: kv groups of g
def Cfg.dimthe hidden size: one head per query head
def Cfg.kvwthe width of the keys (and of the values): one head per key/value head
def smollm2_135mSmolLM2-135M (HuggingFaceTB/SmolLM2-135M): 3 key/value heads of 3 query heads each, heads of 64
def smollm2_360mSmolLM2-360M: 5 key/value heads of 3 query heads each, heads of 64
def smollm2_1_7bSmolLM2-1.7B: 32 key/value heads of one query head each (plain multi-head attention), heads of 64
law smollm2_135m_sizesLAW (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_360m_sizes
law smollm2_1_7b_sizes
def n32
def Cfg.layer_sizethe numbers of one layer, from the shapes of LW<c>: two norms, q and o (dim × dim), k and v (kvw × dim), gate and up (ffn × dim) and down (dim × ffn)
def Cfg.sizethe numbers of the model, from the shapes of Model<c>: the embedding (vocab × dim), the layers, the final norm
law layout_135mLAW (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).
law layout_360m
law layout_1_7b
type Heada head as its two halves: RoPE rotates number i of the first half with number i of the second
def split
def join
def headsn heads of 2·half numbers from xs, in order
def chunksn plain heads (the values, which RoPE does not touch)
def joins
def flat
law head_splitLAW (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.
def list.addxs + ys, element by element (to the shorter length)
def list.mulxs * ys, element by element (to the shorter length)
def list.scalek * xs
def list.sum.goacc + xs, a running sum from the left
def list.max.go
def list.maxthe maximum of xs (0 for an empty list)
def list.exp_shiftexp(x - m) for each x
def list.divxs / s
def list.softmax.fin
def list.softmaxstable softmax: subtracts the maximum before exponentiating
def squaresthe squares of xs
def pair_sumsadjacent pairs added: [a, b, c, d, e] -> [a + b, c + d, e]
def tree_sumthe sum as a balanced tree of additions (pairwise summation): its rounding error grows with log2(n), not n, as in PyTorch's cascade sum.
def rmsnormRMSNorm: g * (x * (1 / sqrt(mean(x²) + eps))), in this order (LlamaRMSNorm.forward)
def inv_freq.goinv_freq[i] = 1 / (theta ^ (2i / head)), i = 0..half-1 (LlamaRotaryEmbedding, default rope)
def inv_freq
def angles
def cos.l
def sin.l
type Rotthe cosines and sines of position t, one per frequency (shared by every head of every layer)
def rot_at
def rot_lofirst half: lo*cos + (-hi)*sin; second half: hi*cos + lo*sin (q*cos + rotate_half(q)*sin)
def rot_hi
def rope
def ropes
def silusilu(x) = x / (1 + exp(-x)), the form ATen's CPU kernel uses
def swiglusilu(gate) * up
type KVOne 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).
def arraysn arrays of 2^d slots
def KV.emptya cache with no position yet, with room for cap positions per key/value head
def write_rowswrites row j of rs into array j, at off
def KV.pushposition t: the keys and values of every key/value head written at t·head; room is the proof that t < cap
def kv_countthe number of key/value heads a cache holds
law arrays_lenLAW (arrays_len): arrays(n, d) makes exactly n arrays.
law kv_headsLAW (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.
def leb_transa <= b and b <= c -> a <= c
def leb_mula <= b -> a·h <= b·h
law kv_roomLAW (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.
def step_roomFor 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_inv
def groups===================================================================================================== Grouped-query attention ===================================================================================================== the query heads in n groups of g, in order: group j holds query heads j·g ..
law groups_lenLAW (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 div_blockone block of 1 + m numbers in Nat.divmod: the counter m runs down to 0, then the quotient d grows by one
law div_blocksj whole blocks of 1 + s numbers add j to the quotient
law div_restfewer than 1 + s numbers after the last whole block leave the quotient alone
law gqa_headLAW (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.
type SRthe scores of one query head against the keys of its key/value head, and the keys handed back
def sr_len
def sr_of
law read_lenLAW (read_len): reading m + 1 numbers of an array onto acc gives m + 1 more scores.
def sr_readread the 1 + t scores the gemm wrote
def scoresq·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 sr_len_g
law causalLAW (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.
type WRthe weighted sum of the values of positions 0..t: w (t + 1 weights) times V (t + 1 rows of `head` numbers)
def wr_of
def wr_read
def weighted
type HOone query head: softmax(q·K · scale) · V, and the two arrays handed back
def HO.then
def head_wr
def head_sr