1/* @ajax/mesh-ref-gpt2-kv: GPT-2's KV cache (mesh-ref gpt2.rs `Kv`): per layer, K and V rows of WIDTH
2 * (heads concatenated) for positions 0..len. Rows are only ever appended. Host. */
3#ifndef MESH_REF_GPT2_KV_H
4#define MESH_REF_GPT2_KV_H
5
6#include "mesh/ref/gpt2_shape.h"
7
8#include <stdbool.h>
9#include <stddef.h>
10#include <stdint.h>
11
12typedef struct {
13    float *k[MESH_GPT2_LAYERS], *v[MESH_GPT2_LAYERS]; /* rows x WIDTH each */
14    size_t rows[MESH_GPT2_LAYERS];                   /* positions held per layer (K and V alike) */
15    size_t cap;                                      /* positions allocated */
16} mesh_gpt2_kv;
17
18/* Positions in the cache: layer 0's, as the Rust's `Kv::len` (during a forward pass the later layers
19 * catch up one by one). */
20static inline size_t mesh_gpt2_kv_len(const mesh_gpt2_kv *kv) { return kv->rows[0]; }
21
22void mesh_gpt2_kv_init(mesh_gpt2_kv *kv);
23void mesh_gpt2_kv_free(mesh_gpt2_kv *kv);
24/* Room for `positions` positions in every layer (never shrinks). False when memory runs out. */
25bool mesh_gpt2_kv_reserve(mesh_gpt2_kv *kv, size_t positions);
26/* FNV-1a 64 of every layer's K then V rows (from the digest seed). */
27uint64_t mesh_gpt2_kv_digest(const mesh_gpt2_kv *kv);
28
29#endif