Module ml::attention
Attention (scaled dot-product + multi-head).
Built on std/ml/tensor.h's Tensor. Every op here records its backward edges like tensor_matmul does, so attention layers train (trainer.h), and runs where its operands live (device.h): on a GPU the multi-head ops run as one fused kernel launch for every head -- QK^T, the softmax and the product with V, never forming the score matrix -- with a fused backward, and otherwise compose the device ops (transpose, matmul, masked softmax, column slices). Q/K/V are plain 2D tensors, [seqLen, dim] -- no batch dimension (call these once per sequence in a batch) and no positional encoding baked in (add it to Q/K yourself before calling, e.g. via tensor_add with a positional tensor).
Masks: query row i sits at position pos = i + (kvLen - seqLen) (the last seqLen keys are the queries' own tokens, any before them a KV cache). causal hides keys j > pos; a sliding window w > 0 hides keys j <= pos - w (the query sees its last w positions). A query with no visible key gets zeros. These are gpu_rocm.h / gpu_cuda_attention.h / gpu_mps.h's masks.
Scope note: standard scaled-dot-product/MHA, gated attention, gated MHA, linear attention, and flash attention are implemented and verified in this file. Multi-head Latent Attention (MLA)/EG-MLA, Shifted-Window (Swin) Attention, and Gated DeltaNet live in attention_advanced.h/.sc. std/ml/ROADMAP.md tracks anything still outstanding.
Inference kernels: gpu_cuda_attention.h (NVIDIA) and gpu_mps.h's mps_flash_attention / mps_linear_attention (Apple) -- batched multi-head flash and linear attention, GQA, causal and sliding-window masks, fp32/fp16/bf16, with the same argument forms -- and gpu_rocm.h's rocm_flash_attention / rocm_linear_attention (AMD, fp32), on their own [heads][seqLen][headDim] buffers, forward only.
Structs
| MultiHeadAttention | The projections around mha_forward as one layer: |
Functions
| tensor_set_cols | Writes 'src' into 'dst's column range [startCol, startCol+numCols) (numCols == src's own column count), row by row, in place -- not part of the graph (tensor_concat_cols is the differentiable way to assemble columns). |
| attention_forward | Single-head scaled dot-product attention: softmax(Q @ K^T / sqrt(headDim)) @ V. Q: [seqLen, headDim], K/V: [kvLen, headDim] -> result: [seqLen, headDim]. |
| mha_forward | Multi-head attention: splits Q/K/V's dModel columns into 'numHeads' equal headDim-wide slices (dModel must be divisible by numHeads), runs attention_forward() per head, and concatenates the per-head outputs back into one [seqLen, dModel] result — the standard MHA shape from "Attention Is All You Need", without a separate output-projection matrix (apply your own W_O via tensor_matmul on the result if needed). |
| attention_forward_mask | The two with a mask (causal 0/1, window 0 for none; see above). |
| mha_forward_mask | |
| attention_set_fused | Whether the multi-head ops on a GPU run as the fused kernels (1, the default) or as composed device ops (0: transpose, matmul, masked softmax per head -- for comparison and testing). Head dims over 128 always compose. Returns the previous setting. |
| gated_attention_forward | Output gating: sigmoid(gateLogits) elementwise-multiplied into the attention result before it's returned (the pattern used e.g. to modulate an attention block's contribution to the residual stream). 'gateLogits' must match the output shape: [seqLen, headDim] for gated_attention_forward, [seqLen, dModel] for gated_mha_forward. |
| gated_mha_forward | |
| linear_attention_forward | Kernel-trick reformulation (Katharopoulos et al., "Transformers are RNNs"): with feature map phi(x) = elu(x)+1 (elementwise, always positive), softmax attention's O(seqLenkvLenheadDim) cost becomes O((seqLen+kvLen)*headDim^2) by computing phi(K)^T @ V once ([headDim,headDim]) instead of the full [seqLen,kvLen] score matrix: |
| flash_attention_forward | Tiled, online-softmax attention (Dao et al.): processes K/V in blocks of 'kvBlockSize' rows, maintaining a running max/sum/weighted-output per query row and rescaling as each block's contribution arrives, so the full [seqLen,kvLen] score matrix is never materialized. Produces the same result as attention_forward() (up to floating-point rounding) — the point is the memory-access pattern, not a different answer. Q: [seqLen,headDim], K/V: [kvLen,headDim] -> result: [seqLen,headDim]. When it records gradients, or on a GPU, it is the fused attention op (one row per thread with an online softmax: also never the full score matrix), whose backward recomputes the scores from each row's log-sum-exp. |