Memory IO Efficiency of Multi-Query Attention

Multi-query attention can be much more efficient under large batch and context length.

Authors

Affiliations

Published

Feb. 1, 2023

Multi-query attention was first introduced in and was later used in PaLM for inference efficiency. In this blog, we will analyze why multi-query can be much more efficient than the traditional multi-head attention.

Multi-Query Attention at a Glance

The key difference of multi-query attention is to collapse all the heads of the projection matrices $P_{K}$ and $P_{V}$ to have only 1 output head instead of full $h$ heads. All other projection matrices ( $P_{Q}$ and $P_{O}$ ) still have sizes hdk. $P_{K}$ and $P_{V}$ have the size reduced from hdk to dk.

Note that given an input $x$ with hidden dimension $d$ , during incremental decoding, $x$ is still projected to many heads during to produce the query tensor (since the query has h heads). Since the query has many heads, the fact that key and value tensors have 1 head still leads to multiple head-interactions during logits and output computation. The single head in key and value tensors is broadcasted to perform attention with all the heads with $Q$ .

To see why such a simple change can lead to dramatically higher efficiency during incremental decoding, we provide background on counting the memory access and computation required for each tensor operation (einsum). Note: One can refer to The Illustrated Attention via Einstein Summation for the introduction to einsum.

Operation and Memory Access Counting (short version)

At a high level, the number operations and memory access for the tensor computation $⟨ A, B ⟩ \to C$ are:

Number of memory access: $O (| A | + | B | + | C |)$ where $| A |$ is the size of the tensor A (product of all dimensions). This is because to access each input or output, we need to either read from it or write to it at least once.*
Number of computations: $O (product (distinct dimensions in A and B)))$ .
For example, $⟨ b h n v, h d v ⟩ \to b h n d$ requires
- $O (b h n d v) = O (b n d^{2})$ number of operations
- and $O (b h n v + h d v + b h n d)$ memory access for both of the inputs as well as the output.

Operation and Memory Access Counting (longer version, can be skipped)

The number of operations for $A, B \to C$ is the number of duplicates * the number of base operations.
- Example 1: $b h n k, b h m k \to b h n m$ has $b h$ number of duplicates where the base operation is $n k, m k \to n m$ since $b h$ are the dimensions that are shared across all inputs and output. This matrix multiplication $n k, m k \to n m$ requires $n m k$ operations. Therefore, total number of operations is $O (b h * n m k)$ .
  - Note. for $n k, m k \to n m$ , $n$ and $m$ are the non-interacting dimensions and $k$ is the interacting dimension (getting summed over). The number of operations in general equals product(set(non-interacting dimensions)) * interacting dimension = nm * k.
- Example 2: $b h n v, h d v \to b n d$ . In this case, there’s no duplicate dimensions across inputs and output. Since this can be framed as $b n * h v, d * h v \to b n d$ , we see that bn and d are the non-interacting dimensions and hv are the interacting one. Therefore, the number of operations is $O (b n d * h v)$
- In general, this is equivalent to product(set(A, B)) where A and B here represent the dimensions.

Memory IO Cost

Now we can analyze the memory IO cost for multi-head and multi-query attention.

Incremental Decoding

Main Takeaway The calculations that incur the highest amount of memory access for normal multi-head attention are the logits and output calculations which involves the following tensor operation (for logits)

Multi Head $⟨ q, K ⟩ : b h k, b h m k \to b h m$
Here, there are bhmk number of operations but it requires bhmk memory access, which is the memory-bound regime (rather than the compute bound) and is inefficient. In contrast, for multi-query,
Multi Query $⟨ q, K ⟩ : b h k, b m k \to b h m$ which requires only bhk + bmk memory access.

Figure 2: Multi-Query Attention vs Multi-Head Attention. Multi-query is almost identical to multi-head except for 1 head for the key and value projection matrices.

Aditional Details

The following table provides analysis for number of operations and memory access cost (in terms of tight complexity bounds) for both the traditional multi-head attention versus multi-query attention.

The color red denote the change due to multi-query attention. Other operations are the same across multi-attention and multi-head if the difference is not stated explicitly.
Note: The number of operations are the same for multi-query and multi-attention

Table 1: Memory Access and Computation Complexities for Incremental Decoding with Multi-Head and Multi-Query Attention.

\begin{array}{llcc} Operation & Einsum & Memory Access & Computation \\ Input (x) : bd \\ q = ⟨ x, P_{q} ⟩ & b d, h d k \to b h k & b d + h d k = b d + d^{2} & b d h k = b d^{2} \\ K = ⟨ x, P_{k} ⟩ (+ K_{p r e v}) & [M H] b d, h d k \to b h k (+ b m h k) & b d + d^{2} & b d h k = b d^{2} \\ [M Q] b d, d k \to b k (+ b m k) & b d + d k \\ V = ⟨ x, P_{v} ⟩ (+ V_{p r e v}) & [M H] b d, h d v \to b h v (+ b m h v) & b d + d^{2} & b d h v = b d^{2} \\ [M Q] b d, d v \to b v (+ b m v) & b d + d v \\ logits = ⟨ q, K ⟩ & [M H] b h k, b h m k \to b h m & b h k + b h m k = b d + b m d & b h m k = b m d \\ [M Q] b h k, b m k \to b h m & b d + b m k + b h m \\ weights: softmax & b h m & b h m \\ out(O) = ⟨ weights, V ⟩ & [M H] b h m, b h m v \to b h v & b h m + b h m v = b h m + b m d & b h m v = d \\ [M Q] b h m, b m v \to b h v & b h m + b m v + b h v \\ y = ⟨ O, P_{O} ⟩ & b h v, h d v \to b d & b d + d^{2} & b d h v = b d^{2} \\ Total : Multi Head & b d + b m d + d^{2} & b h m + b m d + b d^{2} \approx b d^{2} \\ Total : Multi Query & b d + b m k + d^{2} \\ r : Multi Head & 1 / d + m / d + 1 / b \\ r : Multi Query & 1 / d + m / (d h) + 1 / b \end{array}

Note: $r$ is the ratio of memory access complexity versus computation complexity. A ratio close to 1 would indicate that there are 1-to-1 memory access per computation, which would be very inefficient. An unfused softmax or dropout is such examples of IO inefficienct operations.

Observations

for $b \sim 1$ or $m \sim d$ , the number of memory access is high compared to the number of operations
For multi-query, the offending term $m / d$ is reduced by $h$ to $m / (d h)$ .

Batch Computation Cost for Multi-Head Attention (can be skipped)

Batch computation in this case refers to when we compute attentions corresponding to n tokens. The analysis below shows that the number of memory access per operation is much less than 1-to-1 in which makes it quite efficient.

The table below shows the analysis per each operation. The memory access complexity are the same for both multi-head and multi-query. In practice, the multi-query setting is slightly faster due to lower constants. (In MQ, some $d^{2}$ terms are reduced to $d k$ , for example, but the total complexity is still bounded by $d^{2}$ )

Table 2: Memory Access and Computation Complexities for Batch Computation with Multi-Head and Multi-Query Attention. Note that we use n and m for final calculation of memory access and number of computations quite interchangeably since they are the same.

\begin{array}{llcc} Operation & Einsum & Memory Access & Computation \\ Input M, N : bmd, bnd \\ q = ⟨ N, P_{q} ⟩ & b n d, d h k \to b h n k & b n d + d h k = b n d + d^{2} & b n d h k = b n d^{2} \\ K = ⟨ M, P_{k} ⟩ & [M H] b m d, d h k \to b h m k & b m d + d^{2} & b m d h k = b m d^{2} \\ [M Q] b m d, d k \to b m k & b m d + d k \\ V = ⟨ M, P_{v} ⟩ & [M H] b m d, d h v \to b h m v & b m d + d^{2} & b m d h v = b d^{2} \\ [M Q] b m d, d v \to b m v & b m d + d v \\ logits = ⟨ Q, K ⟩ & [M H] b h n k, b h m k \to b h n m & b n d + b m d + b h n^{2} & b h m n k = b m n d = b n^{2} d \\ [M Q] b h n k, b m k \to b h n m & b n d + b m k + b h n^{2} \\ weights: softmax & b h n m & b h n m \\ out(O) = ⟨ weights, V ⟩ & [M H] b h n m, b h m v \to b h n v & b h n m + b h m v = b h n m + b m d & b h n m v = b m n d = b n^{2} d \\ [M Q] b h n m, b m v \to b h n v & b h n m + b m v + b n d \\ y = ⟨ O, P_{O} ⟩ & b h n v, h v d \to b n d & b n d + d^{2} & b n d h v = b n d^{2} \\ Total : Multi Head & \approx b n d + b h n^{2} + d^{2} & b n d^{2} + b n^{2} d \approx b n d^{2} \\ Total : Multi Query & \approx b n d + b h n^{2} + d^{2} \\ r : Multi Head & 1 / d + 1 / k + 1 / (b n) << 1 \\ r : Multi Query & 1 / d + 1 / k + 1 / (b n) << 1 \end{array}

Explanation

At the end of the calculations, we use $n = m$ for the usual context encoding case (where the query and key inputs are the same).
Note: We perform some approximations such as (1) $d k < d^{2}$ and (2) $b n k < b n d$ to arrive at the total memory access.
To approximate the total computation, we assume that $d >> n$ which means that $b n d^{2} >> b n^{2} d$ , so the latter can be ignored.
Both MQ and MH have the same memory access complexity in the batch case, leading to the same efficiency for context encoding.

Implications

The context encoding is the compute-bound regime where all query and key interact over all positions at once. Typically, for a ~10B model, this context encoding latency on a single GPU can be around 400 ms for 2000 input length. This equates to roughly 0.1 ms per token on average. In contrast, the per token latency of such a model would typically be around ~10 ms at best. We can see that the incremental decoding is roughly 100 times (10 ms / 0.1 ms) less efficient.
One can typically perform incremental decoding with similar latency while increasing batch size from 1 up to a certain batch size where GPU memory would hit the limit. Increasing batch size increases inference efficiency since the model parameters are used to compute over many samples rather than just 1.
Multi-query can help reduce the memory consumption during incremental decoding quite significantly, and also help flatten the inference latency to increase much slower than in the MH case when batch size b or context length m increase.
Note - The dimensionality reduction of $P_{K}$ and $P_{V}$ leads to lower number of parameters (for example, 13B multi-head attention model becomes 10.5B multi-query model, fixing all other configurations constant). In order to scale up the multi-query attention model to be of similar size, one can increase other configurations.
Plot on latency and memory consumption – coming soon!

References

Fast Transformer Decoding: One Write-Head is All You Need [PDF]
Shazeer, N., 2019. CoRR, Vol abs/1911.02150.
PaLM: Scaling Language Modeling with Pathways [link]
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H.W., Sutton, C., Gehrmann, S., Schuh, P., Shi, K., Tsvyashchenko, S., Maynez, J., Rao, A., Barnes, P., Tay, Y., Shazeer, N., Prabhakaran, V., Reif, E., Du, N., Hutchinson, B., Pope, R., Bradbury, J., Austin, J., Isard, M., Gur{-}Ari, G., Yin, P., Duke, T., Levskaya, A., Ghemawat, S., Dev, S., Michalewski, H., Garcia, X., Misra, V., Robinson, K., Fedus, L., Zhou, D., Ippolito, D., Luan, D., Lim, H., Zoph, B., Spiridonov, A., Sepassi, R., Dohan, D., Agrawal, S., Omernick, M., Dai, A.M., Pillai, T.S., Pellat, M., Lewkowycz, A., Moreira, E., Child, R., Polozov, O., Lee, K., Zhou, Z., Wang, X., Saeta, B., Diaz, M., Firat, O., Catasta, M., Wei, J., Meier{-}Hellstern, K., Eck, D., Dean, J., Petrov, S. and Fiedel, N., 2022. CoRR, Vol abs/2204.02311. DOI: 10.48550/arXiv.2204.02311