BendMlLlama.block_gen: def of bend-ml-llama@0.2.1.0
Import: import bend-ml-llama@0.2.1.0/main.bend as BendMlLlama
Lines 949-970 of bend-ml-llama@0.2.1.0/main.bend: https://jend.leeeet.dev/read.txt?path=bend-ml-llama@0.2.1.0/main.bend&start_line=949&end_line=970

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

# block for generation: the same computation, with the activations left empty. Generation never reads them, and
# building them (the copies of the rotated heads among them) cost about a fifth of each token on 16 threads (found by
# Mathews Fernando's bend-lint, unused_result). It is a copy of block's body, not a flag in it: a flag that chooses
# after the computation keeps every intermediate alive to the end, and that measured slower than building them.
def 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>:
  match lw:
    case LW{+ln1, wq, wk, wv, wo, +ln2, wg, wu, wd}:
      +n1 = rmsnorm(c, x, ln1)
      mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wq, n1, wq => +q =>
        mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wk, n1, wk => +k =>
          mv(BR<c, cap, t>, Cfg.kvw(c), Cfg.dim(c), wv, n1, wv => +v =>
            +qs = ropes(rt, heads(Cfg.heads(c), Cfg.half(c), q))
            +ks = ropes(rt, heads(Cfg.kv(c), Cfg.half(c), k))
            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 =>
            mv(BR<c, cap, t>, Cfg.dim(c), Cfg.dim(c), wo, at, wo => +a =>
              +x1 = list.add(x, a)
              +n2 = rmsnorm(c, x1, ln2)
              mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wg, n2, wg => +g =>
                mv(BR<c, cap, t>, Cfg.ffn(c), Cfg.dim(c), wu, n2, wu => +u =>
                  mv(BR<c, cap, t>, Cfg.dim(c), Cfg.ffn(c), wd, swiglu(g, u), wd => +m =>
                    +y = list.add(x1, m)
                    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{}}}))))))))

Uses (29):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Bool
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=List
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Act
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=BR
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.dim
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.ffn
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.half
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.head
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.heads
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.kv
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.kvw
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=KA.then
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=KV
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=KV.push
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=LW
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Rot
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=attention
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=chunks
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=heads
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=joins
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=list.add
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=mv
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=rmsnorm
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=ropes
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=swiglu
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-array@0.1.19.0/main.bend&name=leb