Inference Attention 1: The Full Score Matrix

Lesson

An eager framework runs attention as three separate operations, and so does this lab. A scores kernel multiplies Q by K transposed into a T x T matrix per head, a softmax kernel normalizes every row of that matrix in place, and a weighted-sum kernel multiplies the probabilities by V. Between the kernels the whole matrix lives in global memory. It is the oracle made literal: for each query, score every key, softmax the scores, take the weighted sum of the values.

The tensors follow the attention-v1 contract the whole chapter shares: Llama-3.2-1B attention per layer, 32 query heads, 8 KV heads, head size 64. q is [T][32][64], k and v are [T][8][64], out is [T][32][64], and query head h reads KV head h / 4. The prompt is causal: row i sees keys j <= i, so the scores kernel writes -INFINITY above the diagonal and the softmax writes 0 there. The scale is 1/8 and exact.

The driver allocates the score workspace, 32 x T x T floats, and passes it to your runNaiveAttention. The launches are already written in attention.cu; the three kernel bodies are yours. The scores kernel is your tiled matmul from the Matrix multiplication path, one head per blockIdx.z. The softmax is Reduction 5's: subtract the row max before exponentiating, because planted positions score 128 against each other and expf overflows near 88. The weighted sum is a second matmul.

What the matrix costs is the point of the lab. At 512 tokens it is 32 MiB. At 2,048 tokens it is 512 MiB, 13 times the 40 MiB of Q, K, V and output together, and the pipeline moves it through memory at least four times: written by the scores kernel, read and written by the softmax, read by the weighted sum. Every doubling of the context multiplies it by four: 2 GiB at 4,096 tokens, 128 GiB at 32,768. The driver prints the matrix size and the attention time at 512, 1,024 and 2,048 tokens so you can watch it grow, and the Trace timeline shows the three kernels of every call.

The driver checks every value against the attention-v1 oracle, double sums with tolerance 1e-4 plus 1e-3 relative, after proving the oracle against the contract's golden case. Prompts of 1, 17, 33, 257 and 512 tokens are checked row by row: 1 is the single-token row whose probability must be exactly 1, and 17 and 33 break any tile multiple. The output buffer is poisoned, so a row your kernels never write fails by itself. There is no speed gate. This lab is about what the matrix costs, and the next lab removes it.

What you will learn

• State attention as a soft lookup: score every key, softmax the scores, take the weighted sum of the values.
• Compose a tiled matmul and a stable row softmax into a three-kernel attention pipeline with an explicit score workspace.
• Apply causal masking and the grouped-query head mapping by index, without copying K or V.
• Explain why the score matrix grows with the square of the context and what that costs in memory and traffic.