def kv_val

import bendtorch-gpu@1.3.0.0/main.bend as BendtorchGpu

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

672# the cell e = row·kvw + col of dK (a = G, b = Q) or of dV (a = P, b = dO): key j = row % t, over the group's g query
673# heads: the head's column j of a (rows j..t-1) · the head's column of b
674def kv_val(+e: U32, +c: EAC, a: Array<F32>, b: Array<F32>) -> T.D2:
675  match c:
676    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}:
677      +row = U32.div(e, kvw)
678      +col = U32.mod(e, kvw)
679      +bi = U32.div(row, t32)
680      +jj = U32.div(col, hd32)
681      +j = U32.mod(row, t32)
682      +ia = (toff(((bi * h : U32) + (jj * g : U32) : U32), j, tr) + j : U32)
683      +ib = ((((bi * t32 : U32) + j : U32) * w : U32) + (((jj * g : U32) * hd32 : U32) + U32.mod(col, hd32) : U32) : U32)
684      tsum_h(Nat.sub(gn, 1n), ia, ib, j, U32.to_nat((t32 - j : U32)), w, tr, hd32, 0.0, tdot(U32.to_nat((t32 - j : U32)), ia, ib, j, w, 0.0, Array.get(F32, a, ia), Array.get(F32, b, ib)))

Uses (16)

Used by (1)

bendtorch-gpu@1.3.0.0/main:b3_sel