1/* GPT-2's KV cache (mesh-ref gpt2.rs `Kv`). */
2#include "mesh/ref/gpt2_kv.h"
3
4#include "mesh/ref/gpt2_logits.h"
5
6#include <stdlib.h>
7
8void mesh_gpt2_kv_init(mesh_gpt2_kv *kv) {
9    for (int l = 0; l < MESH_GPT2_LAYERS; l++) kv->k[l] = kv->v[l] = NULL, kv->rows[l] = 0;
10    kv->cap = 0;
11}
12
13void mesh_gpt2_kv_free(mesh_gpt2_kv *kv) {
14    for (int l = 0; l < MESH_GPT2_LAYERS; l++) free(kv->k[l]), free(kv->v[l]);
15    mesh_gpt2_kv_init(kv);
16}
17
18bool mesh_gpt2_kv_reserve(mesh_gpt2_kv *kv, size_t positions) {
19    if (positions <= kv->cap) return true;
20    size_t cap = kv->cap ? kv->cap : 16;
21    while (cap < positions) cap *= 2;
22    for (int l = 0; l < MESH_GPT2_LAYERS; l++) {
23        float *k = realloc(kv->k[l], cap * MESH_GPT2_WIDTH * sizeof *k);
24        if (!k) return false;
25        kv->k[l] = k;
26        float *v = realloc(kv->v[l], cap * MESH_GPT2_WIDTH * sizeof *v);
27        if (!v) return false;
28        kv->v[l] = v;
29    }
30    kv->cap = cap;
31    return true;
32}
33
34uint64_t mesh_gpt2_kv_digest(const mesh_gpt2_kv *kv) {
35    uint64_t h = MESH_GPT2_DIGEST_SEED;
36    for (int l = 0; l < MESH_GPT2_LAYERS; l++) {
37        h = mesh_gpt2_digest_extend(h, kv->k[l], kv->rows[l] * MESH_GPT2_WIDTH);
38        h = mesh_gpt2_digest_extend(h, kv->v[l], kv->rows[l] * MESH_GPT2_WIDTH);
39    }
40    return h;
41}