Attention Mechanism

Attention Mechanism

Self Attention

KV-Cache

Motivation

One of the major downsides of transformer architecture is that memory footprint and calculation grow quadratically in terms of sequence length.

During inference time, the model is essentially performing next token prediction task, i.e., it is trying to generate a next token conditioned on all previously generated tokens. Take a closer look at attention, one could notice that it only needs $Q_{T}$, $K$ and $V$ in order to perform necessary calculation for the next token.

Without KV-cache, during each forward pass, the self attention block goes through calculating $Q, K, V$ matrices for all previous tokens. However, there exist a huge amount of repeated calculations during the process. With KV-cache, we only use query for the last token, and append the corresponding key and value to KV-cache.

Memory usage for KV-cache is $2 \times \text{num\_layers}\times\text{sequence\_length}\times d_{k}$.

Tip

For the cross-attention, only the output K and V vectors from the last/top encoder layer are used by decoder layers, hence KV-Cache doesn’t have to be updated once computed. Whereas the decoder self-attention needs to update the cache in each forward pass by appending the new kv data.

read more

• More detailed calculation analysis for KV-Cache
• LLM Inference Series: 3. KV caching explained

Multi Head Attention

Multi Query Attention

In vanilla multi-head attention setting, there is a separate set of query, key and value vectors for each token. In this case, the memory usage goes up quickly for long sequence and incremental inference is often slow. Multi-query attention(MQA) instead shares keys and values across all different attention heads.

import torch

B = 4  # batch
T = 128  # sequence length
D = 512  # embeddings dimension
H = 8  # number of headscc

D_single = D // H  # single head dimension

torch.manual_seed(47)
X = torch.randn(B, T, D)

Wq = torch.nn.Linear(D, D)
Wk = torch.nn.Linear(D, D_single)
Wv = torch.nn.Linear(D, D_single)

Q, K, V = Wq(X), Wk(X), Wv(X)
print("Q: ", Q.shape)
print("K: ", K.shape)
print("V: ", V.shape)

>>> torch.Size([4, 128, 512])
>>> torch.Size([4, 128, 64])
>>> torch.Size([4, 128, 64])

During the calculation for block $i$, the $K$ and $V$ vectors are broadcast to multiple heads by Tensor.expand before doing multi-head attention. Once the output is computed, memory for the KV-cache is free before proceeding to the next block.

Q_ = Q.view(B, T, H, D_single).transpose(1, 2)
K_ = K.unsqueeze(1).expand(B, H, T, D_single).transpose(2, 3)
V_ = V.unsqueeze(1).expand(B, H, T, D_single)

print("Q: ", Q_.shape)
print("K: ", K_.shape)
print("V: ", V_.shape)

The output is

>>> Q: torch.Size([4, 8, 128, 64])
>>> K: torch.Size([4, 8, 64, 128])
>>> V: torch.Size([4, 8, 128, 64])

Note

In terms of memory usage comparison between MQA and MHA, there may not be a noticeable decrease in memory usage for generating relatively short sequences, one can refer to this discussion or conduct more experiments themselves.

Grouped Query Attention

Grouped query attention is somewhat an interpolation of MHA and MQA with the goal of balancing the pros and cons. Tokens are first grouped together according to positions and share key and value vectors within groups. GQA is a balance between MHA from vanilla transformers and MQA, leveraging the advantages from both sides.

The following table is a direct comparison between dimensions of query, key, value and output for MHA and MQA. $B$ is batch size, $T$ is sequence length, $d$ is embedding dimension, $\text{n\_heads}$ is number of heads, $d_{k}=d/\text{n\_heads}$ is query/key/value vector dimension, $\text{n\_kv\_heads}$ is the number of key/value heads and $G=\text{n\_heads/n\_kv\_heads}$ is the number of groups.

Comparison between MHA, MQA and GQA

$$\begin{array}{cccc} & \text{MHA} & \text{MQA} & \text{GQA} \\ \hline X & (B, T, d) & (B, T, d) & (B, T, d) \\ Q & (B, \text{n\_heads}, T, d_{k}) & (B, \text{n\_heads}, T, d_{k}) & (B, G, \text{n\_kv\_heads}, T, d_{k}) \\ K & (B, \text{n\_heads}, T, d_{k}) & (B, \mathbf{1}, T, d_{k}) & (B, \text{n\_kv\_heads}, T, d_{k}) \\ V & (B, \text{n\_heads}, T, d_{k}) & (B, \mathbf{1}, T, d_{k}) & (B, \text{n\_kv\_heads}, T, d_{k}) \\ O & (B, T, d) & (B, T, d) & (B, \text{n\_kv\_heads}, T, d_{k}) \\ W_{O} & (d, d) & (d, d) & (\text{n\_kv\_heads}, d) \\ Y & (B, T, d) & (B, T, d) & (B, T, d) \\ \end{array}$$

Other variants of Attention

Linear Attention

read more

• Linear Attention