law kv_cols

bendtorch@1.3.0.0 · LAWS.bend:248

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

248# LAW (kv_cols): the attention code takes the columns each query head reads from T.head_cols, in both passes (the head
249# loops pass no head index further, so they cannot compute other columns). With g = 1 + i + k query heads per key/value
250# head, query head j·g + i reads the key and value columns of head j: j·hd onwards.
251law kv_cols:
252  for +t: Nat
253  for +hd: Nat
254  for +w: Nat
255  for +kvw: Nat
256  for +sc: F32
257  for +j: Nat
258  for +i: Nat
259  for +k: Nat
260  {T.hc_kc(T.head_cols(T.Dm{t, 1n+Nat.add(i, k), hd, w, kvw, sc}, Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i))) == U32.from_nat(Nat.mul(j, hd)) : U32}