def par_att

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

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

857# 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
858# outputs and the caches come back in group order, so every head computes what it computed one after another
859def par_att(d: Nat, gs: GS, +n: Nat, +hd: Nat, +t: Nat, +scale: F32) -> AT:
860  match d gs:
861    case 0n GS{l, r}:
862      # 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)
863      p q = seq_att(l, hd, t, scale, Nil{}, Nil{}, Nil{}) seq_att(r, hd, t, scale, Nil{}, Nil{}, Nil{})
864      at_join(p, q)
865    case 1n++e GS{l, r}:
866      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)
867      at_join(p, q)

Uses (10)

Used by (1)

bend-ml-llama@0.2.1.0/main:attention