def kv_val_f

bendtorch-gpu@1.3.0.0 · main.bend:1093

import Base
import bendtorch@1.3.0.0/main.bend as T

1093# the cell e = row·kvw + col of dK (a = G, b = Q) or of dV (a = P, b = dO), dense: key j = row % t, over the group's
1094# g query heads, each the head's column j of a from row 0 (t apart a row) · the head's column of b (w apart a row)
1095def kv_val_f(+e: U32, +c: EAC, a: Array<F32>, b: Array<F32>) -> T.D2:
1096  match c:
1097    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}:
1098      +row = U32.div(e, kvw)
1099      +col = U32.mod(e, kvw)
1100      +bi = U32.div(row, t32)
1101      +jj = U32.div(col, hd32)
1102      +j = U32.mod(row, t32)
1103      +ia = (toff_f(((bi * h : U32) + (jj * g : U32) : U32), 0, t32, tr) + j : U32)
1104      +ib = (((bi * t32 : U32) * w : U32) + (((jj * g : U32) * hd32 : U32) + U32.mod(col, hd32) : U32) : U32)
1105      dsum_h(Nat.sub(gn, 1n), ia, ib, t, t32, w, tr, hd32, 0.0, ddot(t, ia, ib, t32, w, 0.0, Array.get(F32, a, ia), Array.get(F32, b, ib)))