def forward_k

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

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

1030# token tok at position t (0-based), with the caches of positions 0..t-1 (one per layer) and room for cap positions;
1031# room is the proof that t < cap; keep: whether the result carries every layer's activations
1032def forward_k(+keep: Bool, +c: Cfg, -cap: Nat, +t: Nat, m: Model<c>, cs: List<&1, KV<cap, t>>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +tok: Nat) -> Out<c, cap, t>:
1033  match m:
1034    case Model{+d, emb, ls, +lnf}:
1035      forward.go(keep, c, cap, t, room, d, ls, lnf, cs, TA.Bands.read_row(Cfg.vocab(c), Cfg.dim(c), emb, tok))