chickadee » nanograd-components » multihead-attention-core

multihead-attention-core Q K V n-heads #!key (causal? #f)procedure
scaled-dot-product-attention Q K V head-size seq-len #!key (causal? #f)procedure

Attention without the input and output projections: multihead-attention-core works on plain matrices, and scaled-dot-product-attention on (len, n_heads, head_size) tensors.

The module also exports the helpers causal-mask, embedding-lookup, slice-columns, concat-columns, get-frequencies, rope-head-size, cache-kv-pair! and get-cached-kv, and the predicates transformer?, transformer-block?, multihead-attention?, ffn-block?, rope-embedding? and kv-cache?.