def bg_fin

bendtorch@1.3.0.0 · main.bend:2701

import Base

2701# the group's summed dK and dV into key/value head j's columns of dk and dv
2702def bg_fin(+t: Nat, +hd: Nat, +kvw: Nat, +j: Nat, dk: Array<F32>, dv: Array<F32>, s: BS) -> BG:
2703  match s:
2704    case BS{q, k, v, d, dq, dka, dva}:
2705      BG{q, k, v, d, dq, cols_in(t, hd, U32.from_nat(kvw), U32.from_nat(Nat.mul(j, hd)), dka, dk), cols_in(t, hd, U32.from_nat(kvw), U32.from_nat(Nat.mul(j, hd)), dva, dv)}