1/* @ajax/mesh-ref-gpt2: GPT-2 124M in the cluster float mode (mesh-ref gpt2.rs), the CPU reference of
2 * programs/gpt2x (the exact-mode GPU kernels) and the exact executor of programs/gpt2 on a CPU. Every
3 * value is a fixed function of the weights and the token prefix: the order of every sum is fixed and
4 * followed op for op by the GPU kernels, so the logits and the KV cache are bitwise the same on every
5 * executor, and a row computed alone (decode) equals the same row computed in a batch (prefill).
6 * Work is split over POSIX threads; the results never depend on the thread count. Host. */
7#ifndef MESH_REF_GPT2_H
8#define MESH_REF_GPT2_H
9
10#include "mesh/ref/gpt2_kv.h"
11#include "mesh/ref/gpt2_shape.h"
12
13#include <stdbool.h>
14#include <stddef.h>
15#include <stdint.h>
16
17/* The logits (VOCAB) of a normalized row h: wte times h in matmul's warp-strided GEMV order. */
18void mesh_gpt2_project_vocab(const float *wte, const float *h, size_t threads, float *out);
19
20/* Called once layer l's K and V are in the cache (a prefill streams them from there). */
21typedef void (*mesh_gpt2_after_layer)(size_t layer, const mesh_gpt2_kv *kv, void *ctx);
22
23/* The rows `tokens` at positions mesh_gpt2_kv_len(kv).. through the model; their K and V are appended to kv. The
24 * logits of the last row go to logits (VOCAB floats). False when t = 0, the context would overflow, or
25 * memory runs out. */
26bool mesh_gpt2_forward(const float *w, const mesh_gpt2_layout *lay, const uint32_t *tokens, size_t t,
27                       mesh_gpt2_kv *kv, size_t threads, float *logits);
28bool mesh_gpt2_forward_layers(const float *w, const mesh_gpt2_layout *lay, const uint32_t *tokens, size_t t,
29                              mesh_gpt2_kv *kv, size_t threads, mesh_gpt2_after_layer after, void *ctx,
30                              float *logits);
31
32/* A greedy generation: prefill `prompt`, then `steps` decode steps. toks gets steps + 1 tokens (the
33 * first from the prefill), *digest the digest of every step's logits (prefill's included), and kv
34 * (initialized here; free it) the cache. */
35bool mesh_gpt2_generate(const float *w, const mesh_gpt2_layout *lay, const uint32_t *prompt, size_t n,
36                        size_t steps, size_t threads, uint32_t *toks, uint64_t *digest, mesh_gpt2_kv *kv);
37
38#endif