def par_att

import bend-ml-llama@0.2.0.0/main.bend as BendMlLlama

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

855# the n groups of gs (cut in two halves) in a tree of parallel lets of depth d, each leaf a half of a node; the
856# outputs and the caches come back in group order, so every head computes what it computed one after another
857def par_att(d: Nat, gs: GS, +n: Nat, +hd: Nat, +t: Nat, +scale: F32) -> AT:
858  match d gs:
859    case 0n GS{l, r}:
860      # bend-lint: disable-next bend-lint-rules/performance/deep_call_in_fork -- seq_att runs a handful of groups per leaf not-expected-because: its tail-loop rewrite measured as no change (2026-10-16, NOTES)
861      p q = seq_att(l, hd, t, scale, Nil{}, Nil{}, Nil{}) seq_att(r, hd, t, scale, Nil{}, Nil{}, Nil{})
862      at_join(p, q)
863    case 1n++e GS{l, r}:
864      p q = par_att(e, cut(TA.half(TA.half(n)), l, Nil{}), TA.half(n), hd, t, scale) par_att(e, cut(TA.half(Nat.sub(n, TA.half(n))), r, Nil{}), Nat.sub(n, TA.half(n)), hd, t, scale)
865      at_join(p, q)

Uses (10)

Used by (1)

bend-ml-llama@0.2.0.0/main:attention