1/* @ajax/mesh-abi-view-gpt2x: programs/gpt2x's view: GPT-2 exact-mode kernels' parameter block, ops and epilogues.
2 *
3 * Part of alpha-c-mesh, a C port of alpha-rust-mesh (README.md says from where). */
4#ifndef MESH_ABI_VIEW_GPT2X_H
5#define MESH_ABI_VIEW_GPT2X_H
67#include <stdbool.h>
8#include <stddef.h>
9#include <stdint.h>
1011/* sha256(sha256(64 zero bytes)) = e2f61c3f71d1defd... (Python hashlib). */
12/* programs/gpt2x: GPT-2 124M's exact-mode kernels on the GPU (CUDA and
13 * Metal; mesh_ref::gpt2 is the reference, mesh-host `gpt2x` the host side).
14 * The weights, the KV cache and the activations live in the program's
15 * scratch as f32s; offsets below are in f32s of the scratch. A stage's
16 * operands are written to the parameter block before its tasks are
17 * published (the task's arg is only its unit of work); the tokens of an
18 * EMBED and the bytes of a PUT or COPY go through the view. */
19#define MESH_GPT2X_X_MAGIC 0x00u /* "GPX2" after OP_INIT */
20#define MESH_GPT2X_MAGIC 0x32585047u
21/* The parameter block: PARAM_WORDS u32s (word i at X_PARAMS + 4 i). */
22#define MESH_GPT2X_X_PARAMS 0x100u
23#define MESH_GPT2X_PARAM_WORDS 24u
24/* Tokens of an EMBED (u32 x CTX). */
25#define MESH_GPT2X_X_TOKENS 0x1000u
26/* The data window of PUT and COPY: WINDOW_FLOATS f32s. */
27#define MESH_GPT2X_X_WINDOW 0x10000u
28#define MESH_GPT2X_VIEW_BYTES 0x700000u
29#define MESH_GPT2X_WINDOW_FLOATS ((uint32_t)((MESH_GPT2X_VIEW_BYTES - MESH_GPT2X_X_WINDOW) / 4))
30/* f32s a PUT or COPY task moves (its arg is the chunk). */
31#define MESH_GPT2X_CHUNK 65536u
32/* Columns a GEMV task computes (mesh-rt linalg::GEMV_COLS). */
33#define MESH_GPT2X_GEMV_COLS 64u
34/* Vocabulary rows a LOGITS task computes (32 a warp, 4 at a time). */
35#define MESH_GPT2X_LOGIT_ROWS 256u
36/* Parameter words. */
37#define MESH_GPT2X_P_A 0u /* GEMM/GEMV A, NORM src, LOGITS h */
38#define MESH_GPT2X_P_B 1u /* GEMM/GEMV B, NORM gain, LOGITS wte */
39#define MESH_GPT2X_P_BIAS 2u /* GEMM/GEMV bias, NORM shift */
40#define MESH_GPT2X_P_C 3u /* GEMM/GEMV C, NORM dst, LOGITS out */
41#define MESH_GPT2X_P_M 4u
42#define MESH_GPT2X_P_K 5u
43#define MESH_GPT2X_P_N 6u
44#define MESH_GPT2X_P_EPI 7u /* EPI_* */
45#define MESH_GPT2X_P_T 8u /* rows of the stage (tokens) */
46#define MESH_GPT2X_P_P0 9u /* position of row 0 */
47#define MESH_GPT2X_P_QKV 10u /* KV, ATTN: the q|k|v rows (T x 3 WIDTH) */
48#define MESH_GPT2X_P_KC 11u /* KV, ATTN: this layer's K cache (CTX x WIDTH) */
49#define MESH_GPT2X_P_VC 12u /* KV, ATTN: this layer's V cache */
50#define MESH_GPT2X_P_Y 13u /* ATTN: out (T x WIDTH) */
51#define MESH_GPT2X_P_WTE 14u /* EMBED */
52#define MESH_GPT2X_P_WPE 15u /* EMBED */
53#define MESH_GPT2X_P_X 16u /* EMBED: out (T x WIDTH) */
54#define MESH_GPT2X_P_OFF 17u /* PUT, COPY: scratch offset */
55#define MESH_GPT2X_P_LEN 18u /* PUT, COPY: f32s */
56#define MESH_GPT2X_P_QB 19u /* ATTN: queries a task */
57/* Epilogues of GEMM and GEMV (mesh_ref::gpt2::Epi). */
58#define MESH_GPT2X_EPI_BIAS 0u
59#define MESH_GPT2X_EPI_BIAS_GELU 1u
60#define MESH_GPT2X_EPI_BIAS_RESIDUAL 2u
61/* window[chunk] -> scratch[P_OFF + CHUNK chunk ..] (P_LEN f32s in all). */
62#define MESH_GPT2X_OP_PUT 3u
63/* scratch[P_OFF + CHUNK chunk ..] -> window[chunk]. */
64#define MESH_GPT2X_OP_COPY 4u
65/* arg row r: x[r] = wte[tokens[r]] + wpe[P_P0 + r]. */
66#define MESH_GPT2X_OP_EMBED 5u
67/* arg row r: C[r] = layernorm(A[r]) * B + BIAS (rows of WIDTH). */
68#define MESH_GPT2X_OP_NORM 6u
69/* arg tile: C = A (M x K) B (K x N) through the epilogue (64 x 64 tiles). */
70#define MESH_GPT2X_OP_GEMM 7u
71/* arg r * ceil(N / GEMV_COLS) + column block: the same for one row. */
72#define MESH_GPT2X_OP_GEMV 8u
73/* arg row r: the K and V thirds of QKV row r -> the caches at P_P0 + r. */
74#define MESH_GPT2X_OP_KV 9u
75/* arg head * ceil(T / QB) + block: causal attention of QB query rows of one head. */
76#define MESH_GPT2X_OP_ATTN 10u
77/* arg block: out[256 block ..] = wte rows times h (the warp-strided GEMV). */
78#define MESH_GPT2X_OP_LOGITS 11u
79/* Counters: 0 tasks, 1 rows normalized, 2 tiles, 3 chunks moved. */
80#define MESH_GPT2X_C_TASKS 0u
8182#endif