==> src/gpt2x.c <==1/* gpt2x: GPT-2 124M's exact-mode kernels (program.toml). Every op follows
2 * mesh_ref::gpt2 operation for operation -- the same sums in the same order,
3 * every f32 op correctly rounded, FTZ and unfused, exp / tanh / sqrt from
4 * + - * / -- so the logits and the KV cache are bitwise the reference's on
5 * every GPU. A node is 8 warps (256 threads). The host (meshd's gpt2x) runs
6 * the model as stages of tasks: it writes a stage's operands to the parameter
7 * block, publishes its tasks (the arg is the unit of work: a row, a tile, a
8 * head and query block), and waits for all of them before the next stage.
9 * Layouts: @ajax/mesh-abi-view-gpt2x.
10 * Ported from alpha-rust-mesh programs/gpt2x/src/lib.rs. */
11#include "mesh/rt.h"
12#include "mesh/rt/linalg.h"
13
14#define GX_WIDTH 768u
15#define GX_HEAD_DIM 64u
16#define GX_HEADS 12u
17#define GX_VOCAB 50257u
18#define GX_CTX 1024u
19#define GX_PV_GROUPS 4u
20
21/* mesh_ref::gpt2's constants (the same bits). */
22#define GX_ATTN_SCALE 0.125f
23#define GX_LN_EPS_BITS 0x3727c5acu
24#define GX_GELU_C_BITS 0x3f4c422au
25#define GX_GELU_A_BITS 0x3d372713u
26#define GX_LOG2E_BITS 0x3fb8aa3bu
27#define GX_LN2_HI_BITS 0x3f317200u
28#define GX_LN2_LO_BITS 0x35bfbe8eu
29static const uint32_t GX_EXP_POLY[11] = {
30 0x3493f27eu, 0x3638ef1du, 0x37d00d01u, 0x39500d01u, 0x3ab60b61u, 0x3c088889u, 0x3d2aaaabu, 0x3e2aaaabu, 0x3f000000u, 0x3f800000u, 0x3f800000u,
31};
32
33/* Shared memory (bytes): the parameter block, then the op's own. */
34#define GX_SH_PARAMS 0u
35#define GX_SH_OP 128u
36/* ATTN: q (64), scores (CTX), the PV partials (256). */
37#define GX_SH_Q GX_SH_OP
38#define GX_SH_S (GX_SH_Q + 4u * GX_HEAD_DIM)
39#define GX_SH_PART (GX_SH_S + 4u * GX_CTX)
40
41#define GX_INLINE static inline __attribute__((always_inline))
42
43/* mesh_ref::gpt2::exp_poly. */
44GX_INLINE float gx_exp_poly(float x) {
45 if (x < -87.0f) return 0.0f;
46 if (x > 88.0f) return __builtin_inff();
47 float t = x * mesh_f32_from_bits(GX_LOG2E_BITS);
48 float s = t >= 0.0f ? t + 0.5f : t - 0.5f;
49 int32_t k = (int32_t)s;
50 float kf = (float)k;
51 float r = (x - kf * mesh_f32_from_bits(GX_LN2_HI_BITS)) - kf * mesh_f32_from_bits(GX_LN2_LO_BITS);
52 float p = mesh_f32_from_bits(GX_EXP_POLY[0]);
53 for (int i = 1; i < 11; i++) p = p * r + mesh_f32_from_bits(GX_EXP_POLY[i]);
54 int32_t k1 = k / 2, k2 = k - k1;
55 return p * mesh_f32_from_bits((uint32_t)(k1 + 127) << 23) * mesh_f32_from_bits((uint32_t)(k2 + 127) << 23);
56}
57/* mesh_ref::gpt2::tanh_poly. */
58GX_INLINE float gx_tanh_poly(float x) {
59 if (x > 10.0f) return 1.0f;
60 if (x < -10.0f) return -1.0f;
61 return 1.0f - 2.0f / (gx_exp_poly(x * 2.0f) + 1.0f);
62}
63/* mesh_ref::gpt2::gelu_tanh. */
64GX_INLINE float gx_gelu_tanh(float x) {
65 float inner = mesh_f32_from_bits(GX_GELU_C_BITS) * (x + mesh_f32_from_bits(GX_GELU_A_BITS) * x * x * x);
66 return 0.5f * x * (1.0f + gx_tanh_poly(inner));
67}
68/* mesh_ref::gpt2::sqrt_newton. */
69GX_INLINE float gx_sqrt_newton(float a) {
70 float y = mesh_f32_from_bits((mesh_f32_bits(a) >> 1) + 0x1fbd1df5u);
71 for (int i = 0; i < 4; i++) y = 0.5f * (y + a / y);
72 return y;
73}
74
75/* The scratch as f32s, and its length. */
76typedef struct {
77 float *p;
78 uint32_t n;
79} gx_s;
80/* [off, off + len) inside the scratch. */
81GX_INLINE bool gx_fits(const gx_s *s, uint32_t off, uint32_t len) { return (uint64_t)off + len <= s->n; }
82GX_INLINE float gx_rd(const gx_s *s, uint32_t i) { return s->p[i]; }
83GX_INLINE void gx_wr(const gx_s *s, uint32_t i, float v) { s->p[i] = v; }
84
85static void gx_init(const mesh_ctx *ctx) {
86 if (mesh_ctx_leader(ctx)) mesh_ctx_set_u32(ctx, MESH_GPT2X_X_MAGIC, MESH_GPT2X_MAGIC);
87}
88
89/* The parameter block: thread 0 reads the view once (on a PCIe GPU every read
90 * of the mailbox crosses the bus), the node reads shared memory. */
91static void gx_params(const mesh_ctx *ctx, mesh_shared sh, uint32_t p[MESH_GPT2X_PARAM_WORDS]) {
92 if (mesh_ctx_leader(ctx)) {
93 mesh_view v = mesh_ctx_view(ctx);
94 for (uint32_t i = 0; i < MESH_GPT2X_PARAM_WORDS; i++) mesh_shared_set_unchecked_u32(sh, i, mesh_view_u32(v, MESH_GPT2X_X_PARAMS + 4 * i));
95 }
96 mesh_ctx_block_barrier(ctx);
97 for (uint32_t i = 0; i < MESH_GPT2X_PARAM_WORDS; i++) p[i] = mesh_shared_get_unchecked_u32(sh, i);
98}
99
100/* The epilogue (mesh_ref::gpt2::Epi) of output (r, c) of C (ldc n). */
101typedef struct {
102 const gx_s *s;
103 uint32_t epi, bias, c_off, n, row;
104} gx_finish;
105GX_INLINE void gx_finish_at(const gx_finish *f, uint32_t r, uint32_t c, float acc) {
106 float v = acc + gx_rd(f->s, f->bias + c);
107 uint32_t at = f->c_off + r * f->n + c;
108 float out = f->epi == MESH_GPT2X_EPI_BIAS_GELU ? gx_gelu_tanh(v) : f->epi == MESH_GPT2X_EPI_BIAS_RESIDUAL ? gx_rd(f->s, at) + v : v;
109 gx_wr(f->s, at, out);
110}
111static void gx_finish2(uint32_t r, uint32_t c, float acc, void *user) { gx_finish_at(user, r, c, acc); }
112static void gx_finish1(uint32_t c, float acc, void *user) {
113 const gx_finish *f = user;
114 gx_finish_at(f, f->row, c, acc);
115}
116
117/* Row layer normalization (mesh_ref::gpt2::normalize_row). */
118static void gx_normalize(const mesh_ctx *ctx, const gx_s *s, uint32_t src, uint32_t dst, uint32_t g, uint32_t b) {
119 uint32_t t = mesh_ctx_thread(ctx);
120 float part = 0.0f;
121 for (uint32_t i = t; i < GX_WIDTH; i += MESH_LINALG_TILE_THREADS) part = part + gx_rd(s, src + i);
122 float mean = mesh_ctx_block_sum_f32(ctx, part) / (float)GX_WIDTH;
123 float q = 0.0f;
124 for (uint32_t i = t; i < GX_WIDTH; i += MESH_LINALG_TILE_THREADS) {
125 float d = gx_rd(s, src + i) - mean;
126 q = q + d * d;
127 }
128 float var = mesh_ctx_block_sum_f32(ctx, q) / (float)GX_WIDTH;
129 float inv = 1.0f / gx_sqrt_newton(var + mesh_f32_from_bits(GX_LN_EPS_BITS));
130 for (uint32_t i = t; i < GX_WIDTH; i += MESH_LINALG_TILE_THREADS)
131 gx_wr(s, dst + i, (gx_rd(s, src + i) - mean) * inv * gx_rd(s, g + i) + gx_rd(s, b + i));
132}
133
134/* One head's causal attention for query row r at position pos (mesh_ref::gpt2::attend_row). */
135static void gx_attend(const mesh_ctx *ctx, const gx_s *s, const uint32_t p[MESH_GPT2X_PARAM_WORDS], uint32_t h, uint32_t r, uint32_t pos,
136 mesh_shared sq, mesh_shared ss, mesh_shared sp) {
137 uint32_t qkv = p[MESH_GPT2X_P_QKV], kc = p[MESH_GPT2X_P_KC], vc = p[MESH_GPT2X_P_VC], y = p[MESH_GPT2X_P_Y];
138 uint32_t t = mesh_ctx_thread(ctx), n = pos + 1, hd = h * GX_HEAD_DIM;
139 if (t < GX_HEAD_DIM) mesh_shared_set_unchecked_f32(sq, t, gx_rd(s, qkv + r * 3 * GX_WIDTH + hd + t));
140 mesh_ctx_block_barrier(ctx);
141 float lmax = mesh_f32_from_bits(MESH_RT_MAX_F32_ID_BITS);
142 for (uint32_t j = t; j < n; j += MESH_LINALG_TILE_THREADS) {
143 uint32_t k = kc + j * GX_WIDTH + hd;
144 float a = 0.0f;
145 for (uint32_t d = 0; d < GX_HEAD_DIM; d++) a = a + mesh_shared_get_unchecked_f32(sq, d) * gx_rd(s, k + d);
146 float sc = a * GX_ATTN_SCALE;
147 mesh_shared_set_unchecked_f32(ss, j, sc);
148 lmax = mesh_rt_max_f32(lmax, sc);
149 }
150 float m = mesh_ctx_block_max_f32(ctx, lmax);
151 float lsum = 0.0f;
152 for (uint32_t j = t; j < n; j += MESH_LINALG_TILE_THREADS) {
153 float e = gx_exp_poly(mesh_shared_get_unchecked_f32(ss, j) - m);
154 mesh_shared_set_unchecked_f32(ss, j, e);
155 lsum = lsum + e;
156 }
157 float sum = mesh_ctx_block_sum_f32(ctx, lsum);
158 for (uint32_t j = t; j < n; j += MESH_LINALG_TILE_THREADS) mesh_shared_set_unchecked_f32(ss, j, mesh_shared_get_unchecked_f32(ss, j) / sum);
159 mesh_ctx_block_barrier(ctx);
160 uint32_t g = t / GX_HEAD_DIM, d = t % GX_HEAD_DIM;
161 float a = 0.0f;
162 for (uint32_t j = g; j < n; j += GX_PV_GROUPS) a = a + mesh_shared_get_unchecked_f32(ss, j) * gx_rd(s, vc + j * GX_WIDTH + hd + d);
163 mesh_shared_set_unchecked_f32(sp, t, a);
164 mesh_ctx_block_barrier(ctx);
165 if (t < GX_HEAD_DIM) {
166 float tot = mesh_shared_get_unchecked_f32(sp, t);
167 for (uint32_t gg = 1; gg < GX_PV_GROUPS; gg++) tot = tot + mesh_shared_get_unchecked_f32(sp, gg * GX_HEAD_DIM + t);
168 gx_wr(s, y + r * GX_WIDTH + hd + t, tot);
169 }
170 mesh_ctx_block_barrier(ctx);
171}
172
173static uint32_t gx_div_ceil(uint32_t a, uint32_t b) { return a / b + (a % b != 0); }
174
175static uint32_t gx_task(const mesh_ctx *ctx, uint32_t op, uint32_t arg) {
176 if (op == MESH_OP_INIT) {
177 gx_init(ctx);
178 return MESH_R_NONE;
179 }
180 uint32_t len;
181 uint8_t *sp = mesh_ctx_scratch_ptr(ctx, &len);
182 if (!sp || mesh_ctx_threads(ctx) != MESH_LINALG_TILE_THREADS) return MESH_R_BAD_OP;
183 gx_s s = {(float *)(void *)sp, len / 4};
184 mesh_shared shp;
185 if (!mesh_ctx_shared(ctx, GX_SH_PARAMS, MESH_GPT2X_PARAM_WORDS, 4, &shp)) return MESH_R_BAD_OP;
186 uint32_t p[MESH_GPT2X_PARAM_WORDS];
187 gx_params(ctx, shp, p);
188 bool ok;
189 switch (op) {
190 case MESH_GPT2X_OP_PUT:
191 case MESH_GPT2X_OP_COPY: {
192 uint32_t off = p[MESH_GPT2X_P_OFF], total = p[MESH_GPT2X_P_LEN];
193 uint32_t lo = arg * MESH_GPT2X_CHUNK;
194 if (lo >= total || total > MESH_GPT2X_WINDOW_FLOATS || !gx_fits(&s, off, total)) return MESH_R_BAD_OP;
195 uint32_t hi = lo + MESH_GPT2X_CHUNK < total ? lo + MESH_GPT2X_CHUNK : total;
196 mesh_view v = mesh_ctx_view(ctx);
197 if (op == MESH_GPT2X_OP_PUT) {
198 MESH_CTX_FOR_STRIDED(ctx, hi - lo, i) gx_wr(&s, off + lo + i, mesh_view_f32(v, MESH_GPT2X_X_WINDOW + 4 * (lo + i)));
199 } else {
200 MESH_CTX_FOR_STRIDED(ctx, hi - lo, i) mesh_view_set_f32(v, MESH_GPT2X_X_WINDOW + 4 * (lo + i), gx_rd(&s, off + lo + i));
201 }
202 if (mesh_ctx_leader(ctx)) mesh_ctx_count(ctx, 3, 1);
203 ok = true;
204 break;
205 }
206 case MESH_GPT2X_OP_EMBED: {
207 uint32_t tt = p[MESH_GPT2X_P_T], p0 = p[MESH_GPT2X_P_P0], wte = p[MESH_GPT2X_P_WTE], wpe = p[MESH_GPT2X_P_WPE], x = p[MESH_GPT2X_P_X];
208 if (arg >= tt || p0 + tt > GX_CTX || !gx_fits(&s, wte, GX_VOCAB * GX_WIDTH) || !gx_fits(&s, wpe, GX_CTX * GX_WIDTH) || !gx_fits(&s, x, tt * GX_WIDTH))
209 return MESH_R_BAD_OP;
210 uint32_t tok = mesh_view_u32(mesh_ctx_view(ctx), MESH_GPT2X_X_TOKENS + 4 * arg);
211 if (tok >= GX_VOCAB) return MESH_R_BAD_OP;
212 MESH_CTX_FOR_STRIDED(ctx, GX_WIDTH, i)
213 gx_wr(&s, x + arg * GX_WIDTH + i, gx_rd(&s, wte + tok * GX_WIDTH + i) + gx_rd(&s, wpe + (p0 + arg) * GX_WIDTH + i));
214 ok = true;
215 break;
216 }
217 case MESH_GPT2X_OP_NORM: {
218 uint32_t src = p[MESH_GPT2X_P_A], g = p[MESH_GPT2X_P_B], b = p[MESH_GPT2X_P_BIAS], dst = p[MESH_GPT2X_P_C], tt = p[MESH_GPT2X_P_T];
219 if (arg >= tt || !gx_fits(&s, src, tt * GX_WIDTH) || !gx_fits(&s, dst, tt * GX_WIDTH) || !gx_fits(&s, g, GX_WIDTH) || !gx_fits(&s, b, GX_WIDTH))
220 return MESH_R_BAD_OP;
221 gx_normalize(ctx, &s, src + arg * GX_WIDTH, dst + arg * GX_WIDTH, g, b);
222 if (mesh_ctx_leader(ctx)) mesh_ctx_count(ctx, 1, 1);
223 ok = true;
224 break;
225 }
226 case MESH_GPT2X_OP_GEMM:
227 case MESH_GPT2X_OP_GEMV: {
228 uint32_t a = p[MESH_GPT2X_P_A], b = p[MESH_GPT2X_P_B], bias = p[MESH_GPT2X_P_BIAS], c = p[MESH_GPT2X_P_C];
229 uint32_t m = p[MESH_GPT2X_P_M], k = p[MESH_GPT2X_P_K], n = p[MESH_GPT2X_P_N], epi = p[MESH_GPT2X_P_EPI];
230 if (m == 0 || k == 0 || n == 0 || !gx_fits(&s, a, m * k) || !gx_fits(&s, b, k * n) || !gx_fits(&s, bias, n) || !gx_fits(&s, c, m * n) ||
231 epi > MESH_GPT2X_EPI_BIAS_RESIDUAL)
232 return MESH_R_BAD_OP;
233 gx_finish f = {&s, epi, bias, c, n, 0};
234 if (op == MESH_GPT2X_OP_GEMM) {
235 mesh_shared sa, sb;
236 if (!mesh_ctx_shared(ctx, GX_SH_OP, MESH_LINALG_TILE_SHARED / 2, 4, &sa) ||
237 !mesh_ctx_shared(ctx, GX_SH_OP + 4 * (MESH_LINALG_TILE_SHARED / 2), MESH_LINALG_TILE_SHARED / 2, 4, &sb))
238 return MESH_R_BAD_OP;
239 if (!mesh_linalg_gemm_tile(ctx, s.p + a, s.p + b, m, k, n, arg, sa, sb, gx_finish2, &f)) return MESH_R_BAD_OP;
240 if (mesh_ctx_leader(ctx)) mesh_ctx_count(ctx, 2, 1);
241 } else {
242 uint32_t nb = gx_div_ceil(n, MESH_GPT2X_GEMV_COLS);
243 uint32_t r = arg / nb, cb = arg % nb;
244 if (r >= m) return MESH_R_BAD_OP;
245 mesh_shared sb, sx;
246 if (!mesh_ctx_shared(ctx, GX_SH_OP, MESH_LINALG_GEMV_SLAB * MESH_LINALG_GEMV_COLS, 4, &sb) ||
247 !mesh_ctx_shared(ctx, GX_SH_OP + 4 * MESH_LINALG_GEMV_SLAB * MESH_LINALG_GEMV_COLS, MESH_LINALG_GEMV_SLAB, 4, &sx))
248 return MESH_R_BAD_OP;
249 f.row = r;
250 mesh_linalg_gemv_cols(ctx, s.p + a, s.p + b, k, n, r, cb * MESH_GPT2X_GEMV_COLS, sb, sx, gx_finish1, &f);
251 }
252 ok = true;
253 break;
254 }
255 case MESH_GPT2X_OP_KV: {
256 uint32_t tt = p[MESH_GPT2X_P_T], p0 = p[MESH_GPT2X_P_P0], qkv = p[MESH_GPT2X_P_QKV], kc = p[MESH_GPT2X_P_KC], vc = p[MESH_GPT2X_P_VC];
257 if (arg >= tt || p0 + tt > GX_CTX || !gx_fits(&s, qkv, tt * 3 * GX_WIDTH) || !gx_fits(&s, kc, GX_CTX * GX_WIDTH) || !gx_fits(&s, vc, GX_CTX * GX_WIDTH))
258 return MESH_R_BAD_OP;
259 uint32_t src = qkv + arg * 3 * GX_WIDTH, dst = (p0 + arg) * GX_WIDTH;
260 MESH_CTX_FOR_STRIDED(ctx, GX_WIDTH, i) {
261 gx_wr(&s, kc + dst + i, gx_rd(&s, src + GX_WIDTH + i));
262 gx_wr(&s, vc + dst + i, gx_rd(&s, src + 2 * GX_WIDTH + i));
263 }
264 ok = true;
265 break;
266 }
267 case MESH_GPT2X_OP_ATTN: {
268 uint32_t tt = p[MESH_GPT2X_P_T], p0 = p[MESH_GPT2X_P_P0], qb = p[MESH_GPT2X_P_QB];
269 if (qb == 0 || p0 + tt > GX_CTX || !gx_fits(&s, p[MESH_GPT2X_P_QKV], tt * 3 * GX_WIDTH) || !gx_fits(&s, p[MESH_GPT2X_P_KC], GX_CTX * GX_WIDTH) ||
270 !gx_fits(&s, p[MESH_GPT2X_P_VC], GX_CTX * GX_WIDTH) || !gx_fits(&s, p[MESH_GPT2X_P_Y], tt * GX_WIDTH))
271 return MESH_R_BAD_OP;
272 uint32_t nq = gx_div_ceil(tt, qb);
273 uint32_t h = arg / nq, blk = arg % nq;
274 if (h >= GX_HEADS) return MESH_R_BAD_OP;
275 mesh_shared sq, ss, spart;
276 if (!mesh_ctx_shared(ctx, GX_SH_Q, GX_HEAD_DIM, 4, &sq) || !mesh_ctx_shared(ctx, GX_SH_S, GX_CTX, 4, &ss) ||
277 !mesh_ctx_shared(ctx, GX_SH_PART, MESH_LINALG_TILE_THREADS, 4, &spart))
278 return MESH_R_BAD_OP;
279 for (uint32_t r = blk * qb; r < tt && r < (blk + 1) * qb; r++) gx_attend(ctx, &s, p, h, r, p0 + r, sq, ss, spart);
280 ok = true;
281 break;
282 }
283 case MESH_GPT2X_OP_LOGITS: {
284 uint32_t hh = p[MESH_GPT2X_P_A], wte = p[MESH_GPT2X_P_B], out = p[MESH_GPT2X_P_C];
285 if (!gx_fits(&s, hh, GX_WIDTH) || !gx_fits(&s, wte, GX_VOCAB * GX_WIDTH) || !gx_fits(&s, out, GX_VOCAB) || arg * MESH_GPT2X_LOGIT_ROWS >= GX_VOCAB)
286 return MESH_R_BAD_OP;
287 /* Each warp: 32 rows, 4 at a time (more loads in flight; each row's order
288 * is unchanged: lane l sums k = l, l + 32, .., then the butterfly). */
289 uint32_t l = mesh_ctx_lane(ctx);
290 uint32_t base = arg * MESH_GPT2X_LOGIT_ROWS + 32 * mesh_ctx_warp(ctx);
291 for (uint32_t q = 0; q < 8; q++) {
292 uint32_t r0 = base + 4 * q;
293 if (r0 >= GX_VOCAB) break;
294 /* Lane l's k = l, l + 32, .. (24 of them) in rounds of 4: 16 loads in flight, the adds in k order. */
295 float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
296 uint32_t rows[4];
297 for (uint32_t i = 0; i < 4; i++) rows[i] = r0 + i < GX_VOCAB - 1 ? r0 + i : GX_VOCAB - 1;
298 for (uint32_t k = l; k < GX_WIDTH; k += 128) {
299 float hk[4], w[4][4];
300 for (uint32_t j = 0; j < 4; j++) hk[j] = gx_rd(&s, hh + k + 32 * j);
301 for (uint32_t i = 0; i < 4; i++)
302 for (uint32_t j = 0; j < 4; j++) w[i][j] = gx_rd(&s, wte + rows[i] * GX_WIDTH + k + 32 * j);
303 for (uint32_t j = 0; j < 4; j++)
304 for (uint32_t i = 0; i < 4; i++) acc[i] = acc[i] + w[i][j] * hk[j];
305 }
306 for (uint32_t i = 0; i < 4; i++) {
307 float v = mesh_rt_warp_reduce_f32(acc[i], mesh_rt_add_f32);
308 if (l == 0 && r0 + i < GX_VOCAB) gx_wr(&s, out + r0 + i, v);
309 }
310 }
311 ok = true;
312 break;
313 }
314 default: ok = false; break;
315 }
316 if (!ok) return MESH_R_BAD_OP;
317 if (mesh_ctx_leader(ctx)) mesh_ctx_count(ctx, MESH_GPT2X_C_TASKS, 1);
318 return MESH_R_NONE;
319}
320
321MESH_PROGRAM(gx_init, gx_task);