| #include <stdint.h> |
| #include <string.h> |
| #include <math.h> |
|
|
| void tmk_linear_vec(const uint8_t* codes, int step, int m, int n, const float* tlut, const uint8_t* su, |
| const uint8_t* sv, float wscale, const float* x, float* y, float* tmp); |
| void tmk_gelu(float* x, long n); |
| void tmk_silu(float* x, long n); |
|
|
| typedef struct { const uint8_t* codes; const uint8_t* su; const uint8_t* sv; float ws; int m, n; } tmk_chunk; |
| typedef struct { const tmk_chunk* ch; int nch, step; const float* tlut; const float* bias; const float* dense; int m, n; } tmk_lin; |
|
|
| #define MAXL 8 |
| typedef struct { |
| int D, NL, NH, HD, FF, Pmax; |
| float eps; |
| const float* inv_freq; |
| const float* ln[MAXL][4]; |
| tmk_lin lin[MAXL][6]; |
| const float *lnf_w, *lnf_b; |
| tmk_lin in_proj; |
| const float* in_add; |
| } tmk_lm; |
|
|
| typedef struct { |
| int D, NL, LD; |
| float eps; |
| tmk_lin inp, cond, ada[MAXL], fc1[MAXL], fc2[MAXL], fada, fproj; |
| const float* norm[MAXL][2]; |
| const float* time; |
| } tmk_flow; |
|
|
| static void lin(const tmk_lin* L, const float* x, float* y, float* tmp) { |
| if (L->nch == 0) { |
| for (int i = 0; i < L->m; i++) { |
| const float* w = L->dense + (size_t)i * L->n; |
| float a = 0.f; |
| for (int j = 0; j < L->n; j++) a += w[j] * x[j]; |
| y[i] = a; |
| } |
| } else { |
| int r0 = 0; |
| for (int c = 0; c < L->nch; c++) { |
| const tmk_chunk* C = L->ch + c; |
| tmk_linear_vec(C->codes, L->step, C->m, C->n, L->tlut, C->su, C->sv, C->ws, x, y + r0, tmp); |
| r0 += C->m; |
| } |
| } |
| if (L->bias) for (int i = 0; i < L->m; i++) y[i] += L->bias[i]; |
| } |
|
|
| static void layer_norm(const float* x, float* y, int n, const float* w, const float* b, float eps) { |
| float mu = 0.f, var = 0.f; |
| for (int i = 0; i < n; i++) mu += x[i]; |
| mu /= n; |
| for (int i = 0; i < n; i++) { float d = x[i] - mu; var += d * d; } |
| float r = 1.0f / sqrtf(var / n + eps); |
| for (int i = 0; i < n; i++) { float v = (x[i] - mu) * r; y[i] = w ? v * w[i] + b[i] : v; } |
| } |
|
|
| static void rope(float* x, int H, int hd, int pos, const float* inv) { |
| const int h2 = hd / 2; |
| for (int i = 0; i < h2; i++) { |
| float f = (float)pos * inv[i], c = cosf(f), s = sinf(f); |
| for (int h = 0; h < H; h++) { |
| float* v = x + h * hd; |
| float a = v[i], b = v[i + h2]; |
| v[i] = a * c - b * s; v[i + h2] = b * c + a * s; |
| } |
| } |
| } |
|
|
| size_t tmk_lm_work(const tmk_lm* M) { return 6 * (size_t)M->D + 2 * (size_t)M->FF + M->Pmax + 4096; } |
|
|
| void tmk_lm_step(const tmk_lm* M, const float* lat, const float* x0, int pos, float* const* Kc, float* const* Vc, |
| float* hout, float* work) { |
| const int D = M->D, H = M->NH, hd = M->HD; |
| float *x = work, *h = x + D, *q = h + D, *k = q + D, *v = k + D, *a = v + D, *u = a + D, *s = u + 2 * M->FF, |
| *tmp = s + M->Pmax; |
| if (lat) { |
| lin(&M->in_proj, lat, x, tmp); |
| if (M->in_add) for (int i = 0; i < D; i++) x[i] += M->in_add[i]; |
| } else memcpy(x, x0, sizeof(float) * D); |
| const float scale = 1.0f / sqrtf((float)hd); |
| for (int l = 0; l < M->NL; l++) { |
| layer_norm(x, h, D, M->ln[l][0], M->ln[l][1], M->eps); |
| lin(&M->lin[l][0], h, q, tmp); lin(&M->lin[l][1], h, k, tmp); lin(&M->lin[l][2], h, v, tmp); |
| rope(q, H, hd, pos, M->inv_freq); rope(k, H, hd, pos, M->inv_freq); |
| for (int hh = 0; hh < H; hh++) { |
| float* Kh = Kc[l] + (size_t)hh * M->Pmax * hd; |
| float* Vh = Vc[l] + (size_t)hh * M->Pmax * hd; |
| memcpy(Kh + (size_t)pos * hd, k + hh * hd, sizeof(float) * hd); |
| memcpy(Vh + (size_t)pos * hd, v + hh * hd, sizeof(float) * hd); |
| const float* qh = q + hh * hd; |
| float mx = -3.0e38f; |
| for (int j = 0; j <= pos; j++) { |
| const float* kj = Kh + (size_t)j * hd; |
| float d = 0.f; |
| for (int e = 0; e < hd; e++) d += qh[e] * kj[e]; |
| s[j] = d * scale; |
| mx = s[j] > mx ? s[j] : mx; |
| } |
| float sum = 0.f; |
| for (int j = 0; j <= pos; j++) { s[j] = expf(s[j] - mx); sum += s[j]; } |
| float* ah = a + hh * hd; |
| for (int e = 0; e < hd; e++) ah[e] = 0.f; |
| for (int j = 0; j <= pos; j++) { |
| const float* vj = Vh + (size_t)j * hd; |
| const float w = s[j] / sum; |
| for (int e = 0; e < hd; e++) ah[e] += w * vj[e]; |
| } |
| } |
| lin(&M->lin[l][3], a, h, tmp); |
| for (int i = 0; i < D; i++) x[i] += h[i]; |
| layer_norm(x, h, D, M->ln[l][2], M->ln[l][3], M->eps); |
| lin(&M->lin[l][4], h, u, tmp); |
| tmk_gelu(u, M->FF); |
| lin(&M->lin[l][5], u, h, tmp); |
| for (int i = 0; i < D; i++) x[i] += h[i]; |
| } |
| layer_norm(x, hout, D, M->lnf_w, M->lnf_b, M->eps); |
| } |
|
|
| void tmk_flow_step(const tmk_flow* F, const float* hcond, const float* noise, float* out, float* work) { |
| const int D = F->D; |
| float *x = work, *sc = x + D, *ada = sc + D, *m = ada + 3 * D, *u = m + D, *t = u + D, *tmp = t + D; |
| lin(&F->inp, noise, x, tmp); |
| lin(&F->cond, hcond, sc, tmp); |
| for (int i = 0; i < D; i++) sc[i] += F->time[i]; |
| tmk_silu(sc, D); |
| for (int l = 0; l < F->NL; l++) { |
| lin(&F->ada[l], sc, ada, tmp); |
| const float *shift = ada, *scale = ada + D, *gate = ada + 2 * D; |
| layer_norm(x, m, D, F->norm[l][0], F->norm[l][1], F->eps); |
| for (int i = 0; i < D; i++) m[i] = m[i] * (1.f + scale[i]) + shift[i]; |
| lin(&F->fc1[l], m, u, tmp); |
| tmk_silu(u, D); |
| lin(&F->fc2[l], u, t, tmp); |
| for (int i = 0; i < D; i++) x[i] += gate[i] * t[i]; |
| } |
| lin(&F->fada, sc, ada, tmp); |
| layer_norm(x, m, D, NULL, NULL, F->eps); |
| for (int i = 0; i < D; i++) m[i] = m[i] * (1.f + ada[D + i]) + ada[i]; |
| lin(&F->fproj, m, out, tmp); |
| for (int i = 0; i < F->LD; i++) out[i] += noise[i]; |
| } |
|
|
