File size: 2,076 Bytes
2166d7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | # Integer FlashAttention model: bit-exact twin of rtl/src/fa_engine.sv and
# software/fa_int_model.py (see the latter for the number formats).
#
# Q, K, V : n*d row-major arrays of words; only the low 8 bits (signed, Q4.4) are used.
# Result : n*d signed integers with 8 fractional bits.
import std/math
const
FaDMax* = 16
FaBC = 4
FaNegInf = -(1 shl 30)
FaMaxN* = 1_048_576
ExpLut = [65536, 62757, 60097, 57549, 55109, 52773, 50535, 48393, 46341,
44376, 42495, 40693, 38968, 37316, 35734, 34219, 32768]
ScaleQ = [0, 256, 181, 148, 128, 114, 105, 97, 91, 85, 81, 77, 74, 71, 68, 66, 64]
proc s8(x: int): int =
let v = x and 0xFF
if v >= 128: v - 256 else: v
proc asr(x, n: int): int = floorDiv(x, 1 shl n)
proc expNeg*(delta: int): int =
## exp(-delta/256) in Q0.16, delta >= 0
let dl = min(delta, 4095)
let y8 = (dl * 369) shr 8
let n = y8 shr 8
let f = y8 and 255
let i = f shr 4
let r = f and 15
let t0 = ExpLut[i]
let t1 = ExpLut[i + 1]
let base = t0 - (((t0 - t1) * r) shr 4)
if n >= 16: 0 else: base shr n
proc faAttention*(q, k, v: seq[int]; n, d: int): seq[int] =
doAssert n >= 1 and d >= 1 and d <= FaDMax
result = newSeq[int](n * d)
for i in 0 ..< n:
var m = FaNegInf
var l = 0
var o = newSeq[int](d)
var j0 = 0
while j0 < n:
let keys = min(FaBC, n - j0)
var logits = newSeq[int](keys)
for jj in 0 ..< keys:
var dot = 0
for c in 0 ..< d:
dot += s8(q[i * d + c]) * s8(k[(j0 + jj) * d + c])
logits[jj] = asr(dot * ScaleQ[d], 8)
var tileMax = logits[0]
for x in logits: tileMax = max(tileMax, x)
let mNew = max(m, tileMax)
let alpha = expNeg(mNew - m)
l = (l * alpha) shr 16
for c in 0 ..< d: o[c] = asr(o[c] * alpha, 16)
m = mNew
for jj in 0 ..< keys:
let p = expNeg(m - logits[jj])
l += p
for c in 0 ..< d: o[c] += p * s8(v[(j0 + jj) * d + c])
j0 += FaBC
for c in 0 ..< d:
result[i * d + c] = (o[c] * 16) div l # div truncates toward zero
|