def tk_src

bendtorch-gpu@1.3.0.0 · main.bend:686

import Base

686# K and V transposed for the score cells (kᵀ: b·kv·hd·t, by sequence, key/value head, column c, position j), so
687# that neighbouring lanes, which take neighbouring keys j, read neighbouring cells. dot0 reads kᵀ with stride t: the same
688# numbers in the same order as reading k's row, so the same bits.
689def tk_src(+e: U32, +t32: U32, +kv: U32, +hd32: U32, +kvw: U32) -> U32:
690  +r = U32.div(e, t32)
691  +bj = U32.div(r, hd32)
692  +row = ((U32.div(bj, kv) * t32 : U32) + U32.mod(e, t32) : U32)
693  ((row * kvw : U32) + ((U32.mod(bj, kv) * hd32 : U32) + U32.mod(r, hd32) : U32) : U32)