def fu3

bendtorch@1.3.0.0 · main.bend:2742

import Base

2742# a unit of the parallel work is one sequence's key/value group j: its g query heads' columns of Q (t x g·hd), head j's
2743# columns of K and V (t x hd), and an output of its own (t x g·hd). Inside a unit the heads keep their order and their
2744# arithmetic (head_cols with j = 0: the law kv_cols), so the units give the bits of the whole sequence computed at once.
2745def fu3(+t: Nat, +gw: Nat, qq: Array<F32>, kk: Array<F32>, xs: List<&1, FS>, q: Array<F32>, k: Array<F32>, r: A2) -> FL:
2746  match r:
2747    case A2{vv, v}:
2748      FL{qq, kk, vv, FS{q, k, v, zeros_a(Nat.mul(t, gw))} <> xs}