OSS4AI AI Stack · Season 1 · #2
by Yujian Tang
Cover from the original ten-slide carousel; this article adapts selected panels. Open full-size image
The core idea: a weighted mix
Attention lets a token combine information from other tokens. In scaled dot-product attention, each token produces a query, compares that query with keys, converts the resulting scores into weights, and takes a weighted mix of values.
The short version is “attention is a weighted average.” The important part is where the weights come from: learned projections and comparisons, not a person deciding which word matters.
A search-engine analogy can help: a query represents what is being sought, a key represents what can be matched, and a value is the information passed along. But this is only an analogy. Queries, keys, and values are numerical vectors; they are not literal search terms, labels, or explanations of the model's intent.
Project, score, mix
Start with a matrix of token representations, X. Learned weight matrices create three projections:
Q = XW_Q
K = XW_K
V = XW_V
Next, compare queries with keys using dot products. Scale those scores by the square root of the key dimension, d_k, and apply softmax to each row. Softmax turns a row of scores into nonnegative weights that sum to one.
Finally, multiply those weights by the values:
Attention(Q, K, V) = softmax(QKᵀ / √d_k) V
This is the scaled dot-product attention formula from Attention Is All You Need. It describes one attention operation, not an entire transformer. The calculation below follows one query so that each step is visible.
Learned projections, scaled scores, softmax weights and a weighted mix of values. Open full-size image
Work through “the cat sat”
Suppose we want to calculate the output for “sat” in a three-token toy example, with d_k = 2.
Choose its query as:
q(sat) = [2, 0]
Use these keys and values:
Token | Key | Value |
|---|---|---|
the | [0, 1] | [1, 0] |
cat | [1, 0] | [0, 1] |
sat | [0.5, 0.5] | [1, 1] |
The query's dot products with the three keys are 0, 2, and 1. Divide by √2 to get approximately 0, 1.41, and 0.71. Softmax produces these rounded weights:
the cat sat
0.14 0.58 0.28
Now mix the values:
0.14 × [1, 0] + 0.58 × [0, 1] + 0.28 × [1, 1]
≈ [0.42, 0.86]
“Cat” receives the largest weight in this calculation. That does not establish a linguistic rule that “sat” always attends mostly to “cat.” These vectors were deliberately chosen for teaching. The numbers illustrate the arithmetic; they were not extracted from a trained model. Real models learn the projection matrices.
![Three-token toy attention example with query [2,0], weights0.14,0.58,0.28 and output approximately[0.42,0.86]](https://beehiiv-images-production.s3.amazonaws.com/uploads/asset/file/3e61c471-e159-4d55-997d-da71b884e897/hatchery-s1-02-B-05.png?t=1794267202)
Invented teaching vectors, not measurements from a trained model. The weights and output are rounded. Open full-size image
Why context length changes the workload
In dense self-attention, the conceptual score matrix contains a comparison for every query-key pair: n × n, per head, per layer.
Using decimal token counts:
1,000 tokens: 1 million score positions.
8,000 tokens: 64 million.
128,000 tokens: about 16.4 billion.
Doubling the sequence length gives four times as many pairwise score positions. These are arithmetic counts, not measured latency figures. Causal attention also restricts access to future tokens; the full grid is a useful way to understand the scale, not a claim that every position is usable in every model.
Dense attention score-position counts, not API billing or measured runtime. “k” uses decimal token counts. Open full-size image
FlashAttention and the KV cache solve different problems
FlashAttention computes exact attention with fewer GPU memory reads and writes. It does not turn dense attention's pairwise arithmetic into a linear-time calculation, and it need not store the entire score matrix in GPU memory.
During generation, the KV cache retains earlier tokens' keys and values so they do not need to be recomputed for every new token. That connects attention inside the model to memory management in a serving engine.
The practical takeaway is: long context increases attention work and cache demands. It is not literally “billed in n².” Pair counts explain a computational scaling issue, not a provider's pricing formula or a guaranteed fourfold increase in end-to-end latency.
Sources and scope
Adapted from OSS4AI AI Stack Season 1 #2. The numerical example and sequence-length counts are teaching calculations, not experimental results.
