1/* One head's attention for one query (mesh-ref gpt2.rs `attend_row`). */
2#include "mesh/ref/gpt2_attend.h"
3
4#include "mesh/ftz.h"
5#include "mesh/ref/block.h"
6#include "mesh/ref/gpt2_scalar.h"
7#include "mesh/ref/gpt2_shape.h"
8
9bool mesh_gpt2_attend_row(const float *q, size_t n, const float *k, size_t k_stride, const float *v,
10                          size_t v_stride, float *out) {
11    if (n > MESH_GPT2_CTX) return false;
12    float s[MESH_GPT2_CTX], e[MESH_GPT2_CTX], parts[MESH_GPT2_NT];
13    for (size_t j = 0; j < n; j++) {
14        const float *kj = k + j * k_stride;
15        float a = 0.0f;
16        for (int d = 0; d < MESH_GPT2_HEAD_DIM; d++) a = mesh_ftz_add(a, mesh_ftz_mul(q[d], kj[d]));
17        s[j] = mesh_ftz_mul(a, mesh_ftz_from_bits(MESH_GPT2_ATTN_SCALE_BITS));
18    }
19    for (size_t t = 0; t < MESH_GPT2_NT; t++) {
20        float m = mesh_block_max_f32_id();
21        for (size_t j = t; j < n; j += MESH_GPT2_NT) m = mesh_block_max_f32(m, s[j]);
22        parts[t] = m;
23    }
24    float m = mesh_block_reduce_f32(parts, MESH_GPT2_NT, mesh_block_max_f32_id(), mesh_block_max_f32);
25    for (size_t j = 0; j < n; j++) e[j] = mesh_gpt2_exp_poly(mesh_ftz_sub(s[j], m));
26    for (size_t t = 0; t < MESH_GPT2_NT; t++) {
27        float a = 0.0f;
28        for (size_t j = t; j < n; j += MESH_GPT2_NT) a = mesh_ftz_add(a, e[j]);
29        parts[t] = a;
30    }
31    float sum = mesh_block_reduce_f32(parts, MESH_GPT2_NT, mesh_block_sum_f32_id(), mesh_block_add_f32);
32    for (size_t j = 0; j < n; j++) e[j] = mesh_ftz_div(e[j], sum); /* e now holds p */
33    for (int d = 0; d < MESH_GPT2_HEAD_DIM; d++) {
34        float tot = 0.0f;
35        for (size_t g = 0; g < MESH_GPT2_PV_GROUPS; g++) {
36            float a = 0.0f;
37            for (size_t j = g; j < n; j += MESH_GPT2_PV_GROUPS) a = mesh_ftz_add(a, mesh_ftz_mul(e[j], v[j * v_stride + d]));
38            tot = g == 0 ? a : mesh_ftz_add(tot, a);
39        }
40        out[d] = tot;
41    }
42    return true;
43}