tritium demo mingpt
transcript ·
python ·
c ·
rust
// TRITIUM-PROVENANCE sha256=6d7ce1815a5ea4eacf5df829ac23026deec566513c6bfb4d5c0fbfd4974218ee lang=rust
// TRITIUM-HEIGHT high=1.00 mid=0.00 low=0.00
fn fmean(vals: &mut [f64], n: i64) -> f64 {
let mut x = 0.0;
let mut i = 0;
while i < n {
x = x + vals[i as usize];
i += 1;
}
x / ((n) as f64)
}
fn fsqrt(x: f64) -> f64 {
let mut i: f64 = 0;
let mut i_2 = 0.0;
let mut x_2 = 0.0;
let mut fcmp4 = ((x) <= (0.0)) as i64;
if (fcmp4) != 0 {
return 0.0;
} else {
i_2 = x;
i = 0;
while i < 8 {
x_2 = 0.5 * (i_2 + (x / i_2));
i_2 = x_2;
i += 1;
}
return i_2;
}
}
fn layer_norm(vals: &mut [f64], n: i64, out: &mut [f64], eps: f64) -> f64 {
let mut x_2 = 0.0;
let mut mu = fmean(vals, n);
let mut x = 0.0;
let mut i = 0;
while i < n {
x_2 = vals[i as usize] - mu;
x = x + (x_2 * x_2);
i += 1;
}
x = (x / ((n) as f64));
let mut x_3 = 1.0 / (fsqrt(x + eps));
i = 0;
while i < n {
out[i as usize] = ((vals[i as usize] - mu) * x_3);
i += 1;
}
mu
}
fn ln_rows(x: &mut [f64], n_tokens: i64, d: i64, out: &mut [f64], tmp: &mut [f64], nrm: &mut [f64], eps: f64) -> f64 {
let mut i: f64 = 0;
let mut t = 0;
while t < n_tokens {
i = 0;
while i < d {
tmp[i as usize] = x[((((t * d) + i)) as usize)];
i += 1;
}
layer_norm(tmp, d, nrm, eps);
i = 0;
while i < d {
out[((((t * d) + i)) as usize)] = nrm[i as usize];
i += 1;
}
t += 1;
}
let mut x_2 = ((n_tokens) as f64) * 1.0;
x_2
}
fn gemv_f(a: &mut [f64], x: &mut [f64], n: i64, m: i64, out: &mut [f64]) -> i64 {
let mut j = 0;
let mut x_2 = 0.0;
let mut i = 0;
while i < n {
x_2 = 0.0;
j = 0;
while j < m {
x_2 = x_2 + (a[((((i * m) + j)) as usize)] * x[j as usize]);
j += 1;
}
out[i as usize] = x_2;
i += 1;
}
n
}
fn qkv_seq(x: &mut [f64], n_tokens: i64, d: i64, w: &mut [f64], out: &mut [f64], tmp: &mut [f64], row: &mut [f64]) -> i64 {
let mut i: f64 = 0;
let mut t = 0;
while t < n_tokens {
i = 0;
while i < d {
tmp[i as usize] = x[((((t * d) + i)) as usize)];
i += 1;
}
gemv_f(w, tmp, d, d, row);
i = 0;
while i < d {
out[((((t * d) + i)) as usize)] = row[i as usize];
i += 1;
}
t += 1;
}
n_tokens
}
fn pack_heads(src: &mut [f64], n_tokens: i64, heads: i64, dh: i64, dest: &mut [f64]) -> i64 {
let mut t = 0;
let mut u = 0;
let mut h = 0;
while h < heads {
t = 0;
while t < n_tokens {
u = 0;
while u < dh {
dest[((((((h * n_tokens) * dh) + (t * dh)) + u)) as usize)] = src[((((((t * heads) * dh) + (h * dh)) + u)) as usize)];
u += 1;
}
t += 1;
}
h += 1;
}
(heads * n_tokens) * dh
}
fn fexp(x: f64) -> f64 {
let mut j = 0;
let mut ifexp_2 = 0.0;
let mut x_2 = 0.0;
let mut ln2 = 0.6931471805599453;
let mut inv_ln2 = 1.4426950408889634;
let mut fcmp15 = ((x) >= (0.0)) as i64;
if (fcmp15) != 0 {
ifexp_2 = 0.5;
} else {
ifexp_2 = -0.5;
}
let mut ifexp = ifexp_2;
let mut xi = (((((x * inv_ln2) + ifexp) as i64)));
if xi > 60 {
xi = 60;
}
if xi < -60 {
xi = -60;
}
let mut x_5 = 1.0;
let mut x_4 = 1.0;
let mut i = 1;
while i <= 14 {
x_4 = (x_4 * (x - (((xi) as f64) * ln2))) / ((i) as f64);
x_5 = x_5 + x_4;
i += 1;
}
let mut x_3 = 1.0;
let mut i_2 = xi;
if i_2 < 0 {
i_2 = 0 - i_2;
j = 0;
while j < i_2 {
x_3 = x_3 * 2.0;
j += 1;
}
return x_5 / x_3;
} else {
j = 0;
while j < i_2 {
x_2 = x_3 * 2.0;
x_3 = x_2;
j += 1;
}
return x_5 * x_3;
}
}
fn mha_fwd(q: &mut [f64], k: &mut [f64], v: &mut [f64], scores: &mut [f64], out: &mut [f64], h: i64, l: i64, d: i64, scale: f64) -> i64 {
let mut i: f64 = 0;
let mut j = 0;
let mut t = 0;
let mut x_2 = 0.0;
let mut i_2 = 0.0;
let mut fcmp96 = 0;
let mut e = 0.0;
let mut x = 0.0;
let mut heap_dist = 0;
while heap_dist < h {
i = 0;
while i < l {
j = 0;
while j < l {
x_2 = 0.0;
t = 0;
while t < d {
x_2 = x_2 + (q[((((((heap_dist * l) * d) + (i * d)) + t)) as usize)] * k[((((((heap_dist * l) * d) + (j * d)) + t)) as usize)]);
t += 1;
}
scores[((((i * l) + j)) as usize)] = x_2 * scale;
j += 1;
}
i += 1;
}
i = 0;
while i < l {
i_2 = scores[((i * l) as usize)];
j = 1;
while j < l {
fcmp96 = ((scores[((((i * l) + j)) as usize)]) > (i_2)) as i64;
if (fcmp96) != 0 {
i_2 = scores[((((i * l) + j)) as usize)];
}
j += 1;
}
x = 0.0;
j = 0;
while j < l {
e = fexp((scores[((((i * l) + j)) as usize)] - i_2));
scores[((((i * l) + j)) as usize)] = e;
x = x + e;
j += 1;
}
j = 0;
while j < l {
scores[((((i * l) + j)) as usize)] = (scores[((((i * l) + j)) as usize)] / x);
j += 1;
}
i += 1;
}
i = 0;
while i < l {
t = 0;
while t < d {
x_2 = 0.0;
j = 0;
while j < l {
x_2 = (x_2 + (scores[((((i * l) + j)) as usize)] * v[((((((heap_dist * l) * d) + (j * d)) + t)) as usize)]));
j += 1;
}
out[((((((heap_dist * l) * d) + (i * d)) + t)) as usize)] = x_2;
t += 1;
}
i += 1;
}
heap_dist += 1;
}
(h * l) * d
}
fn unpack_heads(src: &mut [f64], n_tokens: i64, heads: i64, dh: i64, dest: &mut [f64]) -> i64 {
let mut t = 0;
let mut u = 0;
let mut h = 0;
while h < heads {
t = 0;
while t < n_tokens {
u = 0;
while u < dh {
dest[((((((t * heads) * dh) + (h * dh)) + u)) as usize)] = src[((((((h * n_tokens) * dh) + (t * dh)) + u)) as usize)];
u += 1;
}
t += 1;
}
h += 1;
}
(n_tokens * heads) * dh
}
fn residual_add(a: &mut [f64], b: &mut [f64], n: i64, out: &mut [f64]) -> i64 {
let mut i = 0;
while i < n {
out[i as usize] = a[i as usize] + b[i as usize];
i += 1;
}
n
}
fn ftanh(x: f64) -> f64 {
let mut i = x;
let mut fcmp5 = ((i) > (20.0)) as i64;
if (fcmp5) != 0 {
i = 20.0;
}
let mut fcmp7 = ((i) < (-20.0)) as i64;
if (fcmp7) != 0 {
i = -20.0;
}
let mut e = fexp(i + i);
let mut x_2 = e - 1.0;
let mut x_3 = e + 1.0;
x_2 / x_3
}
fn gelu_tanh(x: f64) -> f64 {
let mut x_5 = 0.5 * x;
let mut x_4 = 0.044715 * ((x * x) * x);
let mut x_2 = 0.79788456 * (x + x_4);
let mut x_3 = 1.0 + (ftanh(x_2));
x_5 * x_3
}
fn mlp_gelu_vec(x: &mut [f64], w1: &mut [f64], w2: &mut [f64], d: i64, hidden: i64, hidden_out: &mut [f64], out: &mut [f64]) -> i64 {
gemv_f(w1, x, hidden, d, hidden_out);
let mut i = 0;
while i < hidden {
hidden_out[i as usize] = (gelu_tanh(hidden_out[i as usize]));
i += 1;
}
gemv_f(w2, hidden_out, d, hidden, out);
d
}
fn transformer_block(x: &mut [f64], n_tokens: i64, d: i64, heads: i64, dh: i64, hidden: i64, scale: f64, wq: &mut [f64], wk: &mut [f64], wv: &mut [f64], wo: &mut [f64], w1: &mut [f64], w2: &mut [f64], q: &mut [f64], k: &mut [f64], v: &mut [f64], qh: &mut [f64], kh: &mut [f64], vh: &mut [f64], scores: &mut [f64], attention: &mut [f64], flat: &mut [f64], normed: &mut [f64], normed2: &mut [f64], tmp: &mut [f64], nrm: &mut [f64], row: &mut [f64], hidden_out: &mut [f64], mlp_out: &mut [f64], projected: &mut [f64], residual: &mut [f64], out: &mut [f64]) -> i64 {
let mut i: f64 = 0;
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);
let mut t = 0;
while t < n_tokens {
i = 0;
while i < d {
tmp[i as usize] = flat[((((t * d) + i)) as usize)];
i += 1;
}
gemv_f(wo, tmp, d, d, row);
i = 0;
while i < d {
projected[((((t * d) + i)) as usize)] = row[i as usize];
i += 1;
}
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 as usize] = normed2[((((t * d) + i)) as usize)];
i += 1;
}
mlp_gelu_vec(tmp, w1, w2, d, hidden, hidden_out, row);
i = 0;
while i < d {
mlp_out[((((t * d) + i)) as usize)] = row[i as usize];
i += 1;
}
t += 1;
}
residual_add(residual, mlp_out, n_tokens * d, out);
n_tokens
}