def block_gen

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

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

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