GPT-2's attention for one head and query: scaled dots, block-tree max and sum, softmax, the 4-group value sum, in Z
$ sexc add @ajax/mesh-ref-gpt2-attend- total 3
- drwxr-xr-x-../
- drwxr-xr-x-Oct 27 01:59v/# 1 version
- drwxr-xr-x-Oct 27 01:59files/# the files of 0.1.0
- latest
- 0.1.0, published Oct 27 01:59 by ajax
- owners
- ajax
- license
- MIT
- dependencies
- @ajax/mesh-ftz ^0.1.0, @ajax/mesh-ref-block ^0.1.0, @ajax/mesh-ref-gpt2-shape ^0.1.0, @ajax/mesh-ref-gpt2-scalar ^0.1.0
- dependents
- @ajax/mesh-prog-gpt2x, @ajax/mesh-ref-gpt2
- keywords
- alphaos mesh reference gpt2 determinism
- homepage
- https://github.com/thomasdavis/alpha-c-mesh/tree/main/packages/mesh-ref-gpt2-attend
- repository
- https://github.com/thomasdavis/alpha-c-mesh
1# @ajax/mesh-ref-gpt2-attend 2 3One GPT-2 head's attention for one query, in the GPU node's order: thread t scores keys t, t + 256, … 4(each a sequential 64-term dot from +0, times 1/8), the max and the sum of exponentials by block trees, 5p_j = e_j / sum, and output lane d sums p_j·v_j[d] in four groups (j = g, g + 4, …) added in order. Keys 6and values are addressed by base pointer and stride, so the KV cache is read in place. 7 8**Ported from** `crates/mesh-ref/src/gpt2.rs` (`attend_row`). Freestanding (the scores live on the 9stack, at most 1024 keys). Compiled with `-O2` in every profile. 10 11**Tested** by gpt2.rs's `attention_of_one_key_is_its_value`, and golden data over 1 to 600 keys in 12`@ajax/mesh-ref-gpt2`'s tests. 13 14## Differences 15 16- Rust takes closures returning each key and value slice; here a pointer and a stride. More than 1024 17 keys returns false instead of growing a `Vec`.