def block

bend-ml-llama@0.2.1.0 · main.bend:928

import Base
import bend-ml-tensor-array@0.1.19.0/main.bend as TA

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}}))))))))