8. Transformers › 8.6 Accelerating Attention

Section 8.6
Accelerating Attention

One of the most expensive operations performed in a transformer is constructing the attention matrix, and a significant amount of research and engineering has gone into accelerating these operations. Two important ideas are caching and online compute.

8.6.1 Key-Value Caching

Recall that auto-regressive generation samples one next-token at a time using the conditioned on the system prompt, user prompt and previously sampled tokens, i.e., the context window. We can save compute by caching previous key and value calculations, \(k_1, \ldots, k_{t-1}\) and \(v_1, \ldots, v_{t-1}\), respectively, since these are needed to generate the \(t\)-th output

\begin{align*} y_t^T &= \textbf{softmax}\left(\frac{q_t^T k_1}{\sqrt{d}}, \frac{q_t^T k_2}{\sqrt{d}}, \ldots, \frac{q_t^T k_t}{\sqrt{d}}\right)^T \begin{bmatrix} \text{--- } v_1^T \text{ ---} \\ \text{--- } v_2^T \text{ ---} \\ \vdots \\ \text{--- } v_t^T \text{ ---} \end{bmatrix} \end{align*}

Here keys and values from prompt tokens pre-fill the cache. Keys and values from auto-regressively sampled tokens are iteratively added to the cache. Note that query vectors, \(q_t\), do not need to be cached since they have no effect on future tokens.

Key-value caching is computationally efficient but memory intensive. It grows linearly with number of transformer layers and sequence length. A very active area of research is concerned about compressing the context window and managing what information gets stored in the cache.

8.6.2 Flash Attention

Another issue that arises with growing context windows is the available memory on a GPU, which is limited and expensive. Indeed, there may be insufficient fast memory to store the full set of queries, keys and values making memory transfer between the CPU and GPU a bottleneck.

Recall from Lecture 7 that a numerically stable way to compute softmax,

\begin{align} \textbf{softmax}(z) &= \left(\frac{\exp (z_1 - z_{\text{max}})}{Z}, \ldots, \frac{\exp (z_t - z_{\text{max}})}{Z}\right) \tag{271}\end{align}

involves first determining the maximum logit value, \(z_{\text{max}} = \max \{z_1, \ldots, z_t\}\), then computing the exponential of logits offset by this value, \(\exp (z_1 - z_{\text{max}})\), and finally normalizing by the partition function, \(Z = \sum_{i=1}^{t} \exp (z_i - z_{\text{max}})\). This requires three passes through the data and is the right thing to do when memory access is not the bottleneck. When it is the bottleneck we can reduce the number of passes to two. Online softmax normalization [72] proceeds by keeping track of the maximum value found so far and renormalizing whenever a new maximum is found:

1:function OnlineSoftmax(\(z\))
2:\(z_{\text{max}}, Z = -\infty, 0\)▷ initialize
3:for all \(i = 1, \ldots, n\) do
4:\(z_{\text{tmp}} \gets z_{\text{max}}\)▷ potential new maximum
5:\(z_{\text{max}} \gets \max \{z_i, z_{\text{max}}\}\)
6:\(Z \gets Z \exp(z_{\text{tmp}} - z_{\text{max}}) + \exp(z_i - z_{\text{max}})\)▷ re-scale \(Z\) and accumulate
7:end for
8:return \(z_{\text{max}}\), \(Z\)
9:end function

The repeated multiplication of \(Z\) by \(\exp(z_{\text{tmp}} - z_{\text{max}})\) should raise alarm bells here. However, when \(z_{\text{tmp}} = z_{\text{max}}\) this results in multiplication by one (i.e., does nothing), otherwise a new maximum has been found so \(\exp(z_i - z_{\text{max}})\) dominates and not too much precision is lost. This may be more easily seen in the alternative implementation of for loop (Lines 4–6):

1:if \(z_i > z_{\text{max}}\) then▷ new max element
2:\(Z \gets Z \exp(z_{\text{max}} - z_i)\)▷ re-scale \(Z\)
3:\(z_{\text{max}} \gets z_i\)▷ update new maximum
4:end if
5:\(Z \gets Z + \exp(z_i - z_{\text{max}})\)▷ accumulate \(Z\)

Flash attention [19] combines smart IO-aware memory access (blocking) with online softmax normalization to compute output \(y_t\) while avoiding repeated loading and storing of keys and values. Pseudocode is shown below.

1:function FlashAttention(\(q_t\), KV-cache)
2:\(z_{\text{max}},\, Z = -\infty, 0\)
3:for all \(i = 1, \ldots, t\) do
4:\(z_i = q_t^T k_i / \sqrt{d}\)
5:\(z_{\text{tmp}},\, z_{\text{max}} \gets z_{\text{max}},\, \max \{z_i, z_{\text{max}}\}\)
6:\(Z_{\text{tmp}} = Z \exp(z_{\text{tmp}} - z_{\text{max}})\)
7:\(Z \gets Z_{\text{tmp}} + \exp(z_i - z_{\text{max}})\)
8:\(y_t \gets \frac{Z_{\text{tmp}}}{Z} y_t + \frac{\exp(z_i - z_{\text{max}})}{Z} v_i\)
9:end for
10:return \(y_t\)
11:end function