def attention

bend-ml-llama@0.2.0.0 · main.bend:881

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

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