1/* @ajax/mesh-ref-gpt2-shape: GPT-2 124M's shape and the flat weight buffer's layout (mesh-ref gpt2.rs):
2 * every tensor in checkpoint order and shape (wte, wpe, the blocks, ln_f), the same on the CPU and in
3 * the GPU's scratch. Weights are [in][out] as the checkpoint stores them (HF Conv1D). Freestanding. */
4#ifndef MESH_REF_GPT2_SHAPE_H
5#define MESH_REF_GPT2_SHAPE_H
6
7#include <stdbool.h>
8#include <stddef.h>
9
10#define MESH_GPT2_LAYERS 12
11#define MESH_GPT2_HEADS 12
12#define MESH_GPT2_HEAD_DIM 64
13#define MESH_GPT2_WIDTH 768
14#define MESH_GPT2_FF 3072
15#define MESH_GPT2_VOCAB 50257
16#define MESH_GPT2_CTX 1024
17/* The threads a row's reductions are split over (the GPU node: 8 warps). */
18#define MESH_GPT2_NT 256
19/* The groups of the attention output's sum over keys (NT / HEAD_DIM). */
20#define MESH_GPT2_PV_GROUPS (MESH_GPT2_NT / MESH_GPT2_HEAD_DIM)
21/* KV values per token: layers x (K, V) x width. */
22#define MESH_GPT2_KV_PER_TOKEN (MESH_GPT2_LAYERS * 2 * MESH_GPT2_WIDTH)
23/* Tensors in the checkpoint: wte, wpe, 12 per block, ln_f weight and bias. */
24#define MESH_GPT2_TENSORS (2 + 12 * MESH_GPT2_LAYERS + 2)
25
26/* Offsets (in floats) of one block's tensors in the flat weight buffer. */
27typedef struct {
28    size_t ln1_g, ln1_b, attn_w, attn_b, proj_w, proj_b, ln2_g, ln2_b, fc_w, fc_b, fc2_w, fc2_b;
29} mesh_gpt2_block_off;
30
31typedef struct {
32    size_t wte, wpe;
33    mesh_gpt2_block_off blocks[MESH_GPT2_LAYERS];
34    size_t lnf_g, lnf_b, total;
35} mesh_gpt2_layout;
36
37void mesh_gpt2_layout_init(mesh_gpt2_layout *l);
38/* Tensor i (0..MESH_GPT2_TENSORS) in buffer order: its checkpoint name ("h.3.attn.c_attn.weight", at
39 * most 31 characters plus NUL), offset and element count. False if i is out of range. */
40bool mesh_gpt2_tensor(const mesh_gpt2_layout *l, size_t i, char name[32], size_t *off, size_t *count);
41
42#endif