tritium demo mingpt
transcript ·
python ·
c ·
rust
# TRITIUM-PROVENANCE sha256=6d7ce1815a5ea4eacf5df829ac23026deec566513c6bfb4d5c0fbfd4974218ee lang=python
# TRITIUM-HEIGHT high=0.19 mid=0.29 low=0.52 fns=17
# TRITIUM-PSEUDO fn=zip_add
#| fn zip_add(a, b, n, out):
#| return n
# TRITIUM-HEIGHT fn=zip_add high=0.00 mid=0.23 low=0.77 role=mid_helper
def zip_add(a, b, n, out):
i = 0
while i < n:
out[i] = a[i] + b[i]
i = i + 1
return n
# TRITIUM-PSEUDO fn=sum_f
#| fn sum_f(vals, n):
#| return s
# TRITIUM-HEIGHT fn=sum_f high=0.00 mid=0.41 low=0.59 role=mid_helper
def sum_f(vals, n):
s = 0.0
i = 0
while i < n:
s = s + vals[i]
i = i + 1
return s
# TRITIUM-PSEUDO fn=fmean
#| fn fmean(xs, n):
#| return sum_f(xs, n) / float(n)
# TRITIUM-HEIGHT fn=fmean high=0.75 mid=0.00 low=0.25
def fmean(xs: list[float], n: int) -> float:
return sum_f(xs, n) / float(n)
# TRITIUM-PSEUDO fn=fsqrt
#| fn fsqrt(x):
#| return g
# TRITIUM-HEIGHT fn=fsqrt high=0.00 mid=0.63 low=0.37 role=algo_mid
def fsqrt(x: float) -> float:
if x <= 0.0:
return 0.0
g = x
i = 0
while i < 8:
g = 0.5 * (g + x / g)
i = i + 1
return g
# TRITIUM-PSEUDO fn=layer_norm
#| fn layer_norm(xs, n, out, eps):
#| fmean(...)
#| float(...)
#| fsqrt(...)
#| return mu
# TRITIUM-HEIGHT fn=layer_norm high=0.24 mid=0.29 low=0.47 role=algo_mid
def layer_norm(xs: list[float], n: int, out: list[float], eps: float) -> float:
mu = fmean(xs, n)
v = 0.0
i = 0
while i < n:
d = xs[i] - mu
v = v + d * d
i = i + 1
v = v / float(n)
inv = 1.0 / fsqrt(v + eps)
i = 0
while i < n:
out[i] = (xs[i] - mu) * inv
i = i + 1
return mu
# TRITIUM-PSEUDO fn=ln_rows
#| fn ln_rows(x, n_tokens, d, out, tmp, nrm, eps):
#| layer_norm(...)
#| float(...)
#| return float(n_tokens) * 1.0
# TRITIUM-HEIGHT fn=ln_rows high=0.14 mid=0.19 low=0.66 role=algo_mid
def ln_rows(x: list[float], n_tokens: int, d: int, out: list[float], tmp: list[float], nrm: list[float], eps: float) -> float:
t = 0
while t < n_tokens:
i = 0
while i < d:
tmp[i] = x[t * d + i]
i = i + 1
layer_norm(tmp, d, nrm, eps)
i = 0
while i < d:
out[t * d + i] = nrm[i]
i = i + 1
t = t + 1
return float(n_tokens) * 1.0
# TRITIUM-PSEUDO fn=gemv_f
#| fn gemv_f(a, x, n, m, out):
#| return n
# TRITIUM-HEIGHT fn=gemv_f high=0.00 mid=0.26 low=0.74 role=algo_mid
def gemv_f(a: list[float], x: list[float], n: int, m: int, out: list[float]) -> int:
i = 0
while i < n:
s = 0.0
j = 0
while j < m:
s = s + a[i * m + j] * x[j]
j = j + 1
out[i] = s
i = i + 1
return n
# TRITIUM-PSEUDO fn=qkv_seq
#| fn qkv_seq(x, n_tokens, d, w, out, tmp, row):
#| gemv_f(...)
#| return n_tokens
# TRITIUM-HEIGHT fn=qkv_seq high=0.07 mid=0.20 low=0.73 role=algo_mid
def qkv_seq(x: list[float], n_tokens: int, d: int, w: list[float], out: list[float], tmp: list[float], row: list[float]) -> int:
t = 0
while t < n_tokens:
i = 0
while i < d:
tmp[i] = x[t * d + i]
i = i + 1
gemv_f(w, tmp, d, d, row)
i = 0
while i < d:
out[t * d + i] = row[i]
i = i + 1
t = t + 1
return n_tokens
# TRITIUM-PSEUDO fn=pack_heads
#| fn pack_heads(src, n_tokens, heads, dh, dest):
#| return heads * n_tokens * dh
# TRITIUM-HEIGHT fn=pack_heads high=0.00 mid=0.27 low=0.73 role=algo_mid
def pack_heads(src: list[float], n_tokens: int, heads: int, dh: int, dest: list[float]) -> int:
h = 0
while h < heads:
t = 0
while t < n_tokens:
u = 0
while u < dh:
dest[h * n_tokens * dh + t * dh + u] = src[t * heads * dh + h * dh + u]
u = u + 1
t = t + 1
h = h + 1
return heads * n_tokens * dh
# TRITIUM-PSEUDO fn=fexp
#| fn fexp(x):
#| float(...)
# TRITIUM-HEIGHT fn=fexp high=0.12 mid=0.66 low=0.22 role=algo_mid
def fexp(x: float) -> float:
ln2 = 0.6931471805599453
inv_ln2 = 1.4426950408889634
if x >= 0.0:
choice = 0.5
else:
choice = -0.5
n = int(x * inv_ln2 + choice)
if n > 60:
n = 60
if n < -60:
n = -60
r = x - float(n) * ln2
s = 1.0
t = 1.0
i = 1
while i <= 14:
t = t * r / float(i)
s = s + t
i = i + 1
p = 1.0
nn = n
if nn < 0:
nn = 0 - nn
j = 0
while j < nn:
p = p * 2.0
j = j + 1
return s / p
else:
j = 0
while j < nn:
p = p * 2.0
j = j + 1
return s * p
# TRITIUM-PSEUDO fn=mha_fwd
#| fn mha_fwd(q, k, v, scores, out, h, l, d, scale):
#| fexp(...)
#| return h * l * d
# TRITIUM-HEIGHT fn=mha_fwd high=0.03 mid=0.32 low=0.66 role=algo_mid
def mha_fwd(q: list[float], k: list[float], v: list[float], scores: list[float], out: list[float], h: int, l: int, d: int, scale: float) -> int:
heap_dist = 0
while heap_dist < h:
i = 0
while i < l:
j = 0
while j < l:
s = 0.0
t = 0
while t < d:
s = s + q[heap_dist * l * d + i * d + t] * k[heap_dist * l * d + j * d + t]
t = t + 1
scores[i * l + j] = s * scale
j = j + 1
i = i + 1
i = 0
while i < l:
m = scores[i * l]
j = 1
while j < l:
if scores[i * l + j] > m:
m = scores[i * l + j]
j = j + 1
ssum = 0.0
j = 0
while j < l:
e = fexp(scores[i * l + j] - m)
scores[i * l + j] = e
ssum = ssum + e
j = j + 1
j = 0
while j < l:
scores[i * l + j] = scores[i * l + j] / ssum
j = j + 1
i = i + 1
i = 0
while i < l:
t = 0
while t < d:
s = 0.0
j = 0
while j < l:
s = s + scores[i * l + j] * v[heap_dist * l * d + j * d + t]
j = j + 1
out[heap_dist * l * d + i * d + t] = s
t = t + 1
i = i + 1
heap_dist = heap_dist + 1
return h * l * d
# TRITIUM-PSEUDO fn=unpack_heads
#| fn unpack_heads(src, n_tokens, heads, dh, dest):
#| return n_tokens * heads * dh
# TRITIUM-HEIGHT fn=unpack_heads high=0.00 mid=0.27 low=0.73 role=algo_mid
def unpack_heads(src: list[float], n_tokens: int, heads: int, dh: int, dest: list[float]) -> int:
h = 0
while h < heads:
t = 0
while t < n_tokens:
u = 0
while u < dh:
dest[t * heads * dh + h * dh + u] = src[h * n_tokens * dh + t * dh + u]
u = u + 1
t = t + 1
h = h + 1
return n_tokens * heads * dh
# TRITIUM-PSEUDO fn=residual_add
#| fn residual_add(a, b, n, out):
#| return zip_add(a, b, n, out)
# TRITIUM-HEIGHT fn=residual_add high=0.33 mid=0.00 low=0.67
def residual_add(a: list[float], b: list[float], n: int, out: list[float]) -> int:
return zip_add(a, b, n, out)
# TRITIUM-PSEUDO fn=ftanh
#| fn ftanh(x):
#| fexp(...)
#| return (e - 1.0) / (e + 1.0)
# TRITIUM-HEIGHT fn=ftanh high=0.25 mid=0.75 low=0.00
def ftanh(x: float) -> float:
a = x
if a > 20.0:
a = 20.0
if a < -20.0:
a = -20.0
e = fexp(a + a)
return (e - 1.0) / (e + 1.0)
# TRITIUM-PSEUDO fn=gelu_tanh
#| fn gelu_tanh(x):
#| x3 = x * x * x
#| return 0.5 * x * (1.0 + ftanh(0.79788456 * (x + 0.044715 * x3)))
# TRITIUM-HEIGHT fn=gelu_tanh high=0.75 mid=0.25 low=0.00
def gelu_tanh(x: float) -> float:
x3 = x * x * x
return 0.5 * x * (1.0 + ftanh(0.79788456 * (x + 0.044715 * x3)))
# TRITIUM-PSEUDO fn=mlp_gelu_vec
#| fn mlp_gelu_vec(x, w1, w2, d, hidden, hidden_out, out):
#| gemv_f(...)
#| gelu_tanh(...)
#| return d
# TRITIUM-HEIGHT fn=mlp_gelu_vec high=0.30 mid=0.10 low=0.60
def mlp_gelu_vec(x: list[float], w1: list[float], w2: list[float], d: int, hidden: int, hidden_out: list[float], out: list[float]) -> int:
gemv_f(w1, x, hidden, d, hidden_out)
i = 0
while i < hidden:
hidden_out[i] = gelu_tanh(hidden_out[i])
i = i + 1
gemv_f(w2, hidden_out, d, hidden, out)
return d
# TRITIUM-PSEUDO fn=transformer_block
#| fn transformer_block(x, n_tokens, d, heads, dh, hidden, scale, wq, wk, wv, wo, w1, w2, q, k, v, qh, kh, vh, scores, attention, flat, normed, normed2, tmp, nrm, row, hidden_out, mlp_out, projected, residual, out):
#| ln_rows(...)
#| qkv_seq(...)
#| pack_heads(...)
#| mha_fwd(...)
#| unpack_heads(...)
#| residual_add(...)
#| gemv_f(...)
#| mlp_gelu_vec(...)
#| return n_tokens
# TRITIUM-HEIGHT fn=transformer_block high=0.28 mid=0.11 low=0.61
def transformer_block(x: list[float], n_tokens: int, d: int, heads: int, dh: int, hidden: int, scale: float, wq: list[float], wk: list[float], wv: list[float], wo: list[float], w1: list[float], w2: list[float], q: list[float], k: list[float], v: list[float], qh: list[float], kh: list[float], vh: list[float], scores: list[float], attention: list[float], flat: list[float], normed: list[float], normed2: list[float], tmp: list[float], nrm: list[float], row: list[float], hidden_out: list[float], mlp_out: list[float], projected: list[float], residual: list[float], out: list[float]) -> int:
ln_rows(x, n_tokens, d, normed, tmp, nrm, 0.0)
qkv_seq(normed, n_tokens, d, wq, q, tmp, row)
qkv_seq(normed, n_tokens, d, wk, k, tmp, row)
qkv_seq(normed, n_tokens, d, wv, v, tmp, row)
pack_heads(q, n_tokens, heads, dh, qh)
pack_heads(k, n_tokens, heads, dh, kh)
pack_heads(v, n_tokens, heads, dh, vh)
mha_fwd(qh, kh, vh, scores, attention, heads, n_tokens, dh, scale)
unpack_heads(attention, n_tokens, heads, dh, flat)
t = 0
while t < n_tokens:
i = 0
while i < d:
tmp[i] = flat[t * d + i]
i = i + 1
gemv_f(wo, tmp, d, d, row)
i = 0
while i < d:
projected[t * d + i] = row[i]
i = i + 1
t = t + 1
residual_add(x, projected, n_tokens * d, residual)
ln_rows(residual, n_tokens, d, normed2, tmp, nrm, 0.0)
t = 0
while t < n_tokens:
i = 0
while i < d:
tmp[i] = normed2[t * d + i]
i = i + 1
mlp_gelu_vec(tmp, w1, w2, d, hidden, hidden_out, row)
i = 0
while i < d:
mlp_out[t * d + i] = row[i]
i = i + 1
t = t + 1
residual_add(residual, mlp_out, n_tokens * d, out)
return n_tokens