law gqa_head

bend-ml-llama@0.2.1.0 · main.bend:666

import Base

666# LAW (gqa_head): with g = 1 + i + k query heads per key/value head (g > 0 and i < g), query head j·g + i, the
667# head i of group j, reads key/value head (j·g + i) / g = j. That division is how transformers maps query heads
668# to key/value heads (repeat_kv), so grouping the heads computes the same attention.
669law gqa_head:
670  for +j: Nat
671  for +i: Nat
672  for +k: Nat
673  {Nat.div(Nat.add(Nat.mul(j, 1n+Nat.add(i, k)), i), 1n+Nat.add(i, k)) == j : Nat}