def attention

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

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

883# attention of the (rotated) query heads at position t over the cache of positions 0..t (the newest included).
884# d is the model's parallel depth (depth_for): with d = 0 (two threads or fewer) the key/value heads run one after
885# another, since a fork tree on one thread is only overhead (+5% on the 1.7B, NOTES 2026-10-12); otherwise they run
886# in parallel, one group per leaf. Both orders compute every head the same way, so the bits are the same.
887def attention(+c: Cfg, +d: Nat, -cap: Nat, +t: Nat, qs: List<&2, Head>, kv: KV<cap, 1n+t>) -> KA<cap, 1n+t>:
888  match d kv:
889    case 0n KV{ks, vs}:
890      ka_of(cap, 1n+t, seq_att(gws(groups(Cfg.kv(c), Cfg.g(c), qs), ks, vs), Cfg.head(c), t, (1.0 / F32.sqrt(F32.from_nat(Cfg.head(c))) : F32), Nil{}, Nil{}, Nil{}))
891    case 1n+e KV{ks, vs}:
892      ka_of(cap, 1n+t, par_att(Nat.sub(halvings(Cfg.kv(c), Cfg.kv(c)), 1n), cut(TA.half(Cfg.kv(c)), gws(groups(Cfg.kv(c), Cfg.g(c), qs), ks, vs), Nil{}), Cfg.kv(c), Cfg.head(c), t, (1.0 / F32.sqrt(F32.from_nat(Cfg.head(c))) : F32)))