def pv_val

bendtorch-gpu@1.3.0.0 · main.bend:646

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

646# the cell e = row·w + col of o (a = P, b = V) or of dQ (a = G, b = K): row i of P (or G) for the query's head, its
647# i + 1 positions, · V's (or K's) column of the head's key/value head
648def pv_val(+e: U32, +c: EAC, a: Array<F32>, b: Array<F32>) -> T.D2:
649  match c:
650    case EAC{+n, +t, +t32, +h, +g, +gn, +hd, +hd32, +w, +kvw, +sc, +nw, +nk, +tr}:
651      +row = U32.div(e, w)
652      +col = U32.mod(e, w)
653      +bi = U32.div(row, t32)
654      +hi = U32.div(col, hd32)
655      +i = U32.mod(row, t32)
656      T.dot0(U32.to_nat((i + 1 : U32)), toff(((bi * h : U32) + hi : U32), i, tr), ((((bi * t32 : U32) * kvw : U32) + (U32.div(hi, g) * hd32 : U32) : U32) + U32.mod(col, hd32) : U32), 1, kvw, a, b)