What is an Attention mechanism?

How self-attention computes

Attention is a weighted average, and the interesting part is the weights. Given a sequence of vectors, three learned projections produce a query for each position, a key for each position and a value for each position. The score between position i and position j is the dot product of query i and key j, usually divided by the square root of the key dimension so the scores do not grow with the size of the projection. Those scores go through a softmax across the keys, giving a distribution over positions for every query. The output for position i is the weighted sum of the value vectors, using exactly that distribution.

Multi-head attention runs several of these in parallel with separate projections, so different heads can learn different relationships, then concatenates the results and projects once more. The point of the separate heads is not decoration: one softmax over all pairs is a single average, and the heads let the layer build several different averages at once.

Two properties follow from the formula and are worth stating explicitly, because they explain most of what follows. First, the operation is permutation-equivariant: shuffling the input shuffles the output, and nothing in it knows about order. That is why a model using it also needs positional information added to the representations. Second, every position is compared with every other position, so the work and the memory of the score matrix grow with the square of the sequence length. That quadratic term is the reason this one operation has shaped most of the engineering around it.

The score matrix is the whole problem

Almost every change to attention in the last few years is a change about that matrix.

Keep it out of memory. The score matrix is the largest tensor in the computation, and writing it to main memory and reading it back is what makes naive attention slow. FlashAttention is the best-known answer: our analysis of flash-attention describes it as an IO-aware exact attention kernel for PyTorch, shipped as four separate builds, v2, v3, v4 and the ROCm Triton backend. It tiles the computation so the scores never leave on-chip memory, and it returns exactly the same result as the reference implementation. The repository is described as the official implementation of FlashAttention and FlashAttention-2 from the papers behind them.

Make it smaller on purpose. Rather than computing every pair, latent attention compresses the keys and values into a smaller space and attends in that space. deepseek-ai/FlashMLA is the kernel-level expression of that idea: our analysis of FlashMLA describes it as a CUDA kernel library from DeepSeek providing optimised sparse and dense attention implementations for the H800 and B200 GPU architectures, powering the DeepSeek-V3 and later model families, and it is explicit that this is a low-level component rather than a model or an inference framework.

Look up instead of compute. Memory sparse attention retrieves the relevant part of the context during decoding rather than attending over everything. Our analysis of EverMind-AI/MSA describes it as an end-to-end trainable latent-memory framework that replaces full attention with sparse document retrieval, and records the README claiming near-linear complexity and under nine percent degradation from 16K to 100M tokens, while noting that the repository ships research code and a paper rather than a packaged inference server.

Package the kernels. mjun0812/flash-attention-prebuild-wheels publishes prebuilt flash-attn 2 and 3 wheels for Linux and Windows so a source build can be skipped, and our analysis of that repository is blunt that the wheel you need depends on an exact Python, CUDA and PyTorch combination.

When attention earns its cost

The operation is expensive in a specific way, and it is worth saying when that expense buys something.

Attention earns its cost when the relationship between inputs is not fixed in advance. A convolutional layer has a receptive field you choose; a recurrent layer has a fixed unrolling order. Attention lets every position decide what to read, which is what makes variable-length, reordered and recursively structured inputs trainable with one set of weights. Sequence transduction, document understanding and any task where a relevant token can be anywhere are the cases it was built for.

It does not earn its cost on fixed structure. If the answer is in a field, a table or a known slot, a lookup or a convolution is cheaper and easier to reason about. On short sequences the quadratic term never becomes the bottleneck and a simpler layer is the honest choice. And in a system that already retrieves the right passage before calling the model, a second mechanism for finding relevant tokens is largely duplicated work.

The bill arrives at serving time rather than training time. During autoregressive generation the layer keeps the keys and values it has already computed, so memory grows linearly with the context and generation per step reads the whole cache. Attention is a training-time idea with a long-tail inference consequence, and a team choosing it should choose it knowing that the cache is what they will be managing.

Cross-attention is the other cost. Encoder-decoder models pay a second attention computation per generated token against the source, on top of the causal one inside the decoder, which is one more reason the encoder-only and decoder-only designs spread.

Failure modes that survive into production

Five problems account for most of the attention bugs people hit, and none of them are exotic.

Missing positional information. Because the operation is permutation-equivariant, a model built without a positional signal can distinguish two tokens but not their order. This is a design error rather than a subtle one, and it is easy to make by hand-writing a layer.

Masking mistakes. A causal mask that is off by one position, or applied after the softmax instead of before it, produces a model that trains and then leaks. The symptom is usually good next-token prediction and bad generation.

Score scale. The division by the square root of the key dimension is not a detail. Without it, the softmax saturates, gradients vanish, and training stalls for reasons that look like a bad learning rate.

The quadratic term. At long context the score matrix stops fitting in memory, and the failure is an out-of-memory error rather than a slow run. The tiled kernels exist to move the pressure from memory to arithmetic, which is why the build requirements matter so much.

Attention weights read as explanations. The weights show how strongly one position was mixed into another at one layer. They are not a faithful account of what the model used, and treating them as one produces confident, wrong explanations of model behaviour.

How the implementations differ

The projects below are the ones you will actually touch, and they sit at different levels.

Kernels for the operation itself. flashinfer-ai/flashinfer is a library for LLM serving that provides unified APIs for attention, GEMM and MoE operations with multiple backend implementations including FlashAttention-2 and 3, cuDNN, CUTLASS and TensorRT-LLM. Our analysis of FlashInfer says it picks a backend per GPU, that the pip install is one line, and that the architecture coverage table is where the caveats start. Depending on a library that has already mapped cards to kernels is easier than maintaining your own.

Whole layers to read. lucidrains/x-transformers is described as a concise but complete full-attention transformer that packs encoder, decoder, cross-attention and vision transformer variants into a small set of composable classes. Our analysis of it puts the value in the feature flags rather than the defaults and calls the README the only real manual, which is the honest way to use it. xmu-xiaoma666/External-Attention-pytorch is the smaller version of the same resource: our analysis describes it as a collection of PyTorch implementations of attention, MLP, re-parameterization and convolution modules in standalone files, a reading aid and a module source rather than a training framework.

The older scoring functions, which are still the clearest explanation of the idea. philipperemy/keras-attention adds one Keras layer that returns a context vector from a sequence, with multiplicative scoring in the Luong style or additive scoring in the Bahdanau style. Our analysis describes it as installing from PyPI as attention, working with TensorFlow 2.x, and staying deliberately narrow. Comparing the two scoring functions side by side is the fastest way to understand what a learned weight is doing before you read a transformer's implementation.

In practice

Self-attention is a weighted average with learned weights, and everything difficult about it follows from two facts: it compares every pair of positions, and it has no notion of order. If you are implementing it, start with the score scale and the mask, then decide deliberately how the quadratic term will be handled, since the tiled kernels that do that job require a build matched to your exact Python, CUDA and PyTorch versions. If you are reading it, xmu-xiaoma666/External-Attention-pytorch and philipperemy/keras-attention show the parts in isolation, and lucidrains/x-transformers shows how many variants the design absorbs. If you are serving it, depend on a library that selects kernels per GPU rather than maintaining the build yourself.