BendMlLlama.attention: def of bend-ml-llama@0.2.1.0
Import: import bend-ml-llama@0.2.1.0/main.bend as BendMlLlama
Lines 883-892 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=883&end_line=892

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

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

Uses (22):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32.div
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32.from_nat
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32.sqrt
  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=Base/base.bend&name=Nat.sub
  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.g
  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.kv
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Head
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=KA
  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=cut
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=groups
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=gws
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=halvings
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=ka_of
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=par_att
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=seq_att
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-array@0.1.19.0/main.bend&name=half