1/* Golden data from the Rust (mesh-ref gpt2.rs) on synthetic weights: the row operations on hashed data
2 * (layer norm, linear with each epilogue and a column range, attention over 1..600 keys, argmax, the
3 * digest), the synthetic weights, a greedy generation (prompt 11 7, one step), and gpt2.rs's
4 * decode_rows_equal_prefill_rows: three tokens one at a time give the KV and logits of the three at
5 * once, both the Rust's. */
6#include "mesh/ref/gpt2.h"
7
8#include "mesh/ftz.h"
9#include "mesh/ref/gpt2_attend.h"
10#include "mesh/ref/gpt2_linear.h"
11#include "mesh/ref/gpt2_logits.h"
12#include "mesh/ref/gpt2_norm.h"
13#include "mesh/ref/gpt2_weights.h"
14#include "mesh/ref/matmul.h"
15
16#include <stdio.h>
17#include <stdlib.h>
18
19static int failures;
20
21static void want(const char *what, uint64_t got, uint64_t want) {
22    if (got != want) {
23        fprintf(stderr, "%s: %#018llx, want %#018llx\n", what, (unsigned long long)got, (unsigned long long)want);
24        failures++;
25    }
26}
27
28static uint64_t fh(uint64_t h, uint64_t w) { return (h ^ w) * 0x100000001b3ull; }
29
30static float *mat(uint32_t seed, size_t n) {
31    float *m = malloc(n * sizeof *m);
32    mesh_matmul_matrix(seed, n, m);
33    return m;
34}
35
36static void rows(void) {
37    uint64_t h = 0xcbf29ce484222325ull;
38    float *x = mat(1, 768), *g = mat(2, 768), *b = mat(3, 768), out[768];
39    mesh_gpt2_normalize_row(x, g, b, out);
40    for (int i = 0; i < 768; i++) h = fh(h, mesh_ftz_bits(out[i]));
41    float *w = mat(4, 768 * 96), *bias = mat(5, 96);
42    const mesh_gpt2_epi epis[3] = {MESH_GPT2_EPI_BIAS, MESH_GPT2_EPI_BIAS_GELU, MESH_GPT2_EPI_BIAS_RESIDUAL};
43    for (int e = 0; e < 3; e++) {
44        float *o = mat(6, 96), *o2 = mat(6, 40);
45        mesh_gpt2_project_row(x, 768, w, bias, 96, epis[e], o);
46        for (int i = 0; i < 96; i++) h = fh(h, mesh_ftz_bits(o[i]));
47        mesh_gpt2_project_cols(x, 768, w, bias, 96, epis[e], 50, 40, o2);
48        for (int i = 0; i < 40; i++) h = fh(h, mesh_ftz_bits(o2[i]));
49        free(o), free(o2);
50    }
51    const size_t ns[6] = {1, 2, 7, 256, 257, 600};
52    for (int t = 0; t < 6; t++) {
53        float *q = mat(7, 64), *ks = mat(8, ns[t] * 64), *vs = mat(9, ns[t] * 64), o[64];
54        mesh_gpt2_attend_row(q, ns[t], ks, 64, vs, 64, o);
55        for (int i = 0; i < 64; i++) h = fh(h, mesh_ftz_bits(o[i]));
56        free(q), free(ks), free(vs);
57    }
58    float *lg = mat(10, 1000);
59    h = fh(h, mesh_gpt2_argmax(lg, 1000));
60    h = fh(h, mesh_gpt2_digest_extend(MESH_GPT2_DIGEST_SEED, lg, 1000));
61    free(x), free(g), free(b), free(w), free(bias), free(lg);
62    want("rows", h, 0x2212298fd900ec07ull);
63}
64
65int main(void) {
66    rows();
67    mesh_gpt2_layout lay;
68    mesh_gpt2_layout_init(&lay);
69    float *w = mesh_gpt2_synthetic(5);
70    if (!w) {
71        fprintf(stderr, "no memory for the weights\n");
72        return 1;
73    }
74    uint64_t h = 0xcbf29ce484222325ull;
75    for (size_t i = 0; i < lay.total; i += 997) h = fh(h, mesh_ftz_bits(w[i]));
76    want("synthetic", h, 0x86cd90d5c4c081b6ull);
77
78    const uint32_t prompt[2] = {11, 7};
79    uint32_t toks[2];
80    uint64_t digest;
81    mesh_gpt2_kv kv;
82    if (!mesh_gpt2_generate(w, &lay, prompt, 2, 1, 8, toks, &digest, &kv)) return 1;
83    if (toks[0] != 48384 || toks[1] != 48384) fprintf(stderr, "tokens %u %u\n", toks[0], toks[1]), failures++;
84    want("generate", digest, 0x9fc013029dba45d5ull);
85    want("generate kv", mesh_gpt2_kv_digest(&kv), 0xfc7a3e32394ce039ull);
86    mesh_gpt2_kv_free(&kv);
87
88    /* decode_rows_equal_prefill_rows, three tokens */
89    const uint32_t three[3] = {11, 7, 50256};
90    float *la = malloc(MESH_GPT2_VOCAB * sizeof *la), *lb = malloc(MESH_GPT2_VOCAB * sizeof *lb);
91    mesh_gpt2_kv a, b;
92    mesh_gpt2_kv_init(&a);
93    mesh_gpt2_kv_init(&b);
94    if (!mesh_gpt2_forward(w, &lay, three, 3, &a, 8, la)) return 1;
95    for (int t = 0; t < 3; t++)
96        if (!mesh_gpt2_forward(w, &lay, &three[t], 1, &b, 3, lb)) return 1;
97    want("prefill kv", mesh_gpt2_kv_digest(&a), 0xdddcc2e98df017bfull);
98    want("decode kv", mesh_gpt2_kv_digest(&b), 0xdddcc2e98df017bfull);
99    want("prefill logits", mesh_gpt2_digest_extend(MESH_GPT2_DIGEST_SEED, la, MESH_GPT2_VOCAB), 0xa7b4303cf98c5170ull);
100    want("decode logits", mesh_gpt2_digest_extend(MESH_GPT2_DIGEST_SEED, lb, MESH_GPT2_VOCAB), 0xa7b4303cf98c5170ull);
101    mesh_gpt2_kv_free(&a);
102    mesh_gpt2_kv_free(&b);
103    free(la), free(lb), free(w);
104    return failures != 0;
105}