--- date: '2025-08-21' description: self-attention from scores to cache layout id: attention modified: 2026-09-07 13:56:24 GMT-04:00 tags: - ml title: attention primer created: '2025-08-21' published: '2025-08-21' pageLayout: default slug: lectures/2/afp permalink: https://aarnphm.xyz/lectures/2/afp.md generator: quartz: v4.6.0 hostedProvider: Cloudflare baseUrl: aarnphm.xyz full: https://aarnphm.xyz/llms-full.txt --- Self-attention turns one residual sequence into a weighted average of learned value vectors. Queries choose which key rows receive mass; those weights then mix the values. The same equations explain the stable-softmax trick and the KV-cache accounting used at inference time. ## scores and normalization Let $X\in\mathbb{R}^{n\times d}$ be a sequence of $n$ residual vectors. For one head, $$ Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V, $$ with $Q,K,V\in\mathbb{R}^{n\times d_h}$. Scaled dot-product attention is $$ \operatorname{Attn}(Q,K,V) =\operatorname{softmax}_{\mathrm{row}}\!\left(\frac{QK^\top}{\sqrt{d_h}}+M\right)V, $$ where $M=0$ for unrestricted attention. A causal mask sets entries above the current position to $-\infty$ before the row-wise softmax. \[@vaswani2023attentionneed\] For temperature $T>0$, define $$ \operatorname{LSE}_T(z)=T\log\sum_j e^{z_j/T}. $$ Then $$ \nabla\operatorname{LSE}_T(z)=\operatorname{softmax}(z/T). $$ If $m=\max_jz_j$, the stable evaluation is $$ \operatorname{LSE}_T(z)=m+T\log\sum_j e^{(z_j-m)/T}. $$ Subtracting $m$ changes neither the softmax nor the normalized weights. The factor $1/\sqrt{d_h}$ serves a different purpose from temperature: under an isotropic initialization it keeps the variance of a query-key dot product near one. [[thoughts/RoPE|RoPE]] applies paired rotations to queries and keys before their dot product. The resulting inner product depends on relative position while preserving each rotated pair’s norm. \[@su2023roformerenhancedtransformerrotary\] ## two exact properties > \[!proposition\] Proposition 1. Permutation equivariance > > Remove positional information and any position-dependent mask. For a permutation matrix $P$, > > $$ > \operatorname{Attn}(PQ,PK,PV)=P\operatorname{Attn}(Q,K,V). > $$ The score matrix becomes $PQK^\top P^\top$. Row-wise softmax commutes with the matching row and column permutations, so $$ \operatorname{softmax}_{\mathrm{row}}(PZP^\top) =P\operatorname{softmax}_{\mathrm{row}}(Z)P^\top. $$ Multiplying by $PV$ gives the result. A causal mask is tied to sequence order and therefore removes this full permutation symmetry. > \[!proposition\] Proposition 2. Initialization variance > > Let $q,k\in\mathbb{R}^{d_h}$ be independent, centered random vectors with covariances $\Sigma_q$ and $\Sigma_k$. Then > > $$ > \operatorname{Var}(q^\top k)=\operatorname{tr}(\Sigma_q\Sigma_k). > $$ > > Under $\Sigma_q=\Sigma_k=I$, dividing by $\sqrt{d_h}$ changes the variance from $d_h$ to $1$. The independence assumption describes an initialization calculation. Queries and keys later come from the same residual stream, so it is not a distributional law for a trained model. ## Gaussian-kernel identity Fix a query $q$ and keys $k_1,\ldots,k_n$ of equal norm. Gaussian kernel weights satisfy $$ \exp\!\left(-\frac{\|q-k_j\|_2^2}{2\sigma^2}\right) =C(q)\exp\!\left(\frac{q^\top k_j}{\sigma^2}\right), $$ where $C(q)$ is independent of $j$. After normalization, the weights match dot-product attention when $$ \sigma^2=T\sqrt{d_h}. $$ The equal-key-norm condition does the work. Without it, $\|k_j\|_2^2$ contributes a key-dependent term, so ordinary dot-product attention is not exactly Gaussian Nadaraya-Watson regression. ## heads and KV groups Multi-head attention uses $h$ query heads and concatenates their outputs: $$ \operatorname{MHA}(X)=W_O[O_1;\ldots;O_h]. $$ MHA gives each query head its own key and value head. [[thoughts/GQA|Grouped-query attention]] uses $G #include __global__ void qk_scores_kernel( const float* __restrict__ Q, // [n, d] const float* __restrict__ K, // [n, d] float* __restrict__ S, // [n, n] int n, int d, float inv_sqrt_d) { int row = blockIdx.x * blockDim.x + threadIdx.x; // query index i int col = blockIdx.y * blockDim.y + threadIdx.y; // key index j if (row >= n || col >= n) return; const float* q = Q + row * d; const float* k = K + col * d; float acc = 0.f; for (int t = 0; t < d; ++t) acc += q[t] * k[t]; S[row * n + col] = acc * inv_sqrt_d; } extern "C" void qk_scores(const float* Q, const float* K, float* S, int n, int d) { dim3 block(16, 16); dim3 grid((n + block.x - 1) / block.x, (n + block.y - 1) / block.y); float inv_sqrt_d = 1.f / std::sqrtf((float)d); qk_scores_kernel<<>>(Q, K, S, n, d, inv_sqrt_d); } ``` - row-wise softmax: ```cuda title="row\_softmax.cu" path="lectures/2/row\_softmax.cu" #include #include #include __global__ void row_softmax_kernel(float* __restrict__ S, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= n) return; float m = -FLT_MAX; for (int j = 0; j < n; ++j) m = fmaxf(m, S[i*n + j]); float sum = 0.f; for (int j = 0; j < n; ++j) { float e = expf(S[i*n + j] - m); S[i*n + j] = e; sum += e; } float inv = 1.f / fmaxf(sum, 1e-12f); for (int j = 0; j < n; ++j) S[i*n + j] *= inv; } extern "C" void row_softmax(float* S, int n) { int block = 256; int grid = (n + block - 1) / block; row_softmax_kernel<<>>(S, n); } ``` - value aggregation: ```cuda title="apply\_values.cu" path="lectures/2/apply\_values.cu" #include __global__ void apply_values_kernel( const float* __restrict__ S, // [n, n] const float* __restrict__ V, // [n, d_v] float* __restrict__ O, // [n, d_v] int n, int d_v) { int i = blockIdx.x * blockDim.x + threadIdx.x; // row in S / O int t = blockIdx.y * blockDim.y + threadIdx.y; // value dim if (i >= n || t >= d_v) return; float acc = 0.f; for (int j = 0; j < n; ++j) acc += S[i*n + j] * V[j*d_v + t]; O[i*d_v + t] = acc; } extern "C" void apply_values(const float* S, const float* V, float* O, int n, int d_v) { dim3 block(16, 16); dim3 grid((n + block.x - 1) / block.x, (d_v + block.y - 1) / block.y); apply_values_kernel<<>>(S, V, O, n, d_v); } ``` [[thoughts/flash attention|FlashAttention]] computes the same attention result by tiling these operations and maintaining an online row maximum and normalizer. It avoids writing the full $n\times n$ score and probability matrices to high-bandwidth memory. \[@dao2022flashattentionfastmemoryefficientexact\] Triton reference: ```python title="attention\_triton.py" path="lectures/2/attention\_triton.py" # attention_triton.py import triton import triton.language as tl @triton.jit def attn_two_pass( Q, K, V, O, n, d, dv, stride_qm, stride_qd, stride_km, stride_kd, stride_vm, stride_vd, stride_om, stride_od, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_D: tl.constexpr, ): # Row block we compute row_offs = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) d_offs = tl.arange(0, BLOCK_D) n_offs = tl.arange(0, BLOCK_N) # ---------------------------- # Pass 1: compute per-row max & denom for stable softmax # ---------------------------- m_i = tl.full((BLOCK_M,), -1e9, dtype=tl.float32) l_i = tl.zeros((BLOCK_M,), dtype=tl.float32) for col_start in range(0, n, BLOCK_N): # Load Q-block [M, D] q = tl.load( Q + row_offs[:, None] * stride_qm + d_offs[None, :] * stride_qd, mask=(row_offs[:, None] < n) & (d_offs[None, :] < d), other=0.0, ) # Load K-block [N, D] k = tl.load( K + (col_start + n_offs)[:, None] * stride_km + d_offs[None, :] * stride_kd, mask=((col_start + n_offs)[:, None] < n) & (d_offs[None, :] < d), other=0.0, ) # scores [M, N] = q @ k^T / sqrt(d) scores = tl.dot(q, tl.trans(k)) * (1.0 / tl.sqrt(tl.float32(d))) # Online softmax: track max and rescale accumulated sum m_i_new = tl.maximum(m_i, tl.max(scores, axis=1)) # Rescale previous accumulator when max increases alpha = tl.exp(m_i - m_i_new) l_i = l_i * alpha # Add current tile's contribution p = tl.exp(scores - m_i_new[:, None]) l_i += tl.sum(p, axis=1) # Update running max m_i = m_i_new # ---------------------------- # Pass 2: accumulate O = softmax(S) @ V # ---------------------------- inv_l_i = 1.0 / l_i Oacc = tl.zeros((BLOCK_M, BLOCK_D), dtype=tl.float32) for col_start in range(0, n, BLOCK_N): q = tl.load( Q + row_offs[:, None] * stride_qm + d_offs[None, :] * stride_qd, mask=(row_offs[:, None] < n) & (d_offs[None, :] < d), other=0.0, ) k = tl.load( K + (col_start + n_offs)[:, None] * stride_km + d_offs[None, :] * stride_kd, mask=((col_start + n_offs)[:, None] < n) & (d_offs[None, :] < d), other=0.0, ) scores = tl.dot(q, tl.trans(k)) * (1.0 / tl.sqrt(tl.float32(d))) # Recenter using m_i computed in pass 1 scores = scores - m_i[:, None] p = tl.exp(scores) * inv_l_i[:, None] # [M, N] # Load V-block [N, Dv] v = tl.load( V + (col_start + n_offs)[:, None] * stride_vm + d_offs[None, :] * stride_vd, mask=((col_start + n_offs)[:, None] < n) & (d_offs[None, :] < dv), other=0.0, ) Oacc += tl.dot(p, v) # [M, Dv] # Store result tl.store( O + row_offs[:, None] * stride_om + d_offs[None, :] * stride_od, Oacc, mask=(row_offs[:, None] < n) & (d_offs[None, :] < dv), ) ```