Multi-headed Self Attention

First lets understand the self attention layer. The self attention layer provides a mechanism for each token to attend to the previous tokens in the prompt/conversation. By doing so, it helps to frame the meaning of the current token given the context of previous tokens. The output of the self attention layer is a vector for each token, and each of those vectors frame the meaning of their respective tokens keeping the context of previous tokens in consideration. There can be multiple such output vectors, equal to the number of heads defined during training. The computation of these output heads is expensive if done in a trivial way.

Let's take a look at the code for the self attention block.

class MultiHeadSelfAttention(nn.Module):
  
  def __init__(self, n_emb, n_key, n_head):
    super().__init__()
    self.n_key = n_key
    self.n_head = n_head
    self.Qlin = nn.Linear(n_emb, n_head*n_key)
    self.Klin = nn.Linear(n_emb, n_head*n_key)
    self.Vlin = nn.Linear(n_emb, n_head*n_key)
    self.W0 = nn.Linear(n_head*n_key, n_emb)
  
  def forward(self, x):
    Q = self.Qlin(x)
    K = self.Klin(x)
    V = self.Vlin(x)

    Q = Q.view(Q.shape[0], Q.shape[1], self.n_head, self.n_key).transpose(1,2)
    K  = K.view(K.shape[0], K.shape[1], self.n_head, self.n_key).transpose(1,2)
    V = V.view(V.shape[0], V.shape[1], self.n_head, self.n_key).transpose(1,2)

    score = Q @ K.transpose(-2,-1)
    score = score / sqrt(self.n_key)
    n_new, n_total = score.shape[-2], score.shape[-1]
    mask = torch.tril(torch.ones(n_total, n_total, device=score.device))[-n_new:]
    score = score.masked_fill(mask == 0, float('-inf'))
    score = torch.softmax(score, -1)

    Z = score @ V
    Z = Z.transpose(1,2)
    Z = Z.reshape(Z.shape[0], Z.shape[1], -1)

    Z = self.W0(Z)

    return Z

  def reset_cache(self):
    self.k_cache = None
    self.v_cache = None

In the forward pass, we pass the input x. Since this is the uncached version of self attention block, all previous tokens need to be passed in x. It is required so that a token can attend to its previous tokens. Using the input x, a query, key and value vector is computed for each token with the help of linear layers. After that for each token attention scores are computed. It is done by taking the query vector of the token and obtaining its dot product with the key vector of every previous token. score = Q @ K.transpose(-2,-1) does exactly that.

Let us take a look at the dimensions of the vectors. Before being passed to the self attention layers in decoders, each token is passed through an embedding layer. Let us assume that each token is transformed to a vector of dimension e. So the input to the self attention block will have dimensions (b, t, e). Here b is the batch size, and t is the number of tokens. The query vector Q will end up having dimensions (b, h, t, k) where h is the number of heads configured, and k is the number of vector dimensions each of the linear layers for key, query and value transforms the input vector to. In score = Q @ K.transpose(-2,-1), since we are multiplying query by the transpose of key, the dimension of key transpose will be (b, h, k, t). Therefore the time complexity of this matrix multiplication will be \(O(b.h.k.t^2)\) . During inference b=1, and h and k are constants, so the overall time complexity becomes \(O(t^2)\).

So the number of computations done is proportional to the square of the number of tokens. While this is fine for small context of conversations, for a medium to large context window this value goes beyond what the cores of a traditional GPU can offer. And so it takes a lot of time to perform inference.

Here is how we could perform an optimisation. Notice that during inference we pass a set of tokens to get the next token, we append the new token to the set of tokens and pass it again through the network to get another token. This cycle repeats. In the self attention block the mask is causal. It means that when we compute how much to attend to each previous token, we only need the previous tokens keys and values vectors. When we would have computed the previous tokens before, we would already have computed its key and value vectors. It means we can reuse them in this forward pass instead of recomputing them all over again. That would save a lot of computations. We can cache the key and value vectors of the previous tokens and reuse them when computing a new token. This also means that we do not need to pass every token appearing before when computing a new token, since we have cached the necessary data. So the entire workflow would look like this - We cache the key and value vectors for previous tokens. When computing the new token through a forward pass, we only pass the one previous token t=1). In the self attention block, get the key and value vectors for previous tokens via the cache instead of recomputing them. Use the previous key and value vectors to obtain the attention values in the score matrix. Below is the updated code that uses key value cache.

class MultiHeadSelfAttention(nn.Module):
  
  def __init__(self, n_emb, n_key, n_head):
    super().__init__()
    self.n_key = n_key
    self.n_head = n_head
    self.Qlin = nn.Linear(n_emb, n_head*n_key)
    self.Klin = nn.Linear(n_emb, n_head*n_key)
    self.Vlin = nn.Linear(n_emb, n_head*n_key)
    self.W0 = nn.Linear(n_head*n_key, n_emb)
    self.register_buffer("k_cache", None)
    self.register_buffer("v_cache", None)
  
  def forward(self, x, use_cache = False):
    Q = self.Qlin(x)
    K = self.Klin(x)
    V = self.Vlin(x)

    Q = Q.view(Q.shape[0], Q.shape[1], self.n_head, self.n_key).transpose(1,2)
    K  = K.view(K.shape[0], K.shape[1], self.n_head, self.n_key).transpose(1,2)
    V = V.view(V.shape[0], V.shape[1], self.n_head, self.n_key).transpose(1,2)

    if use_cache:
      if self.k_cache is not None:
          K = torch.cat([self.k_cache, K], dim=2)
          V = torch.cat([self.v_cache, V], dim=2)
      self.k_cache = K
      self.v_cache = V

    score = Q @ K.transpose(-2,-1)
    score = score / sqrt(self.n_key)
    n_new, n_total = score.shape[-2], score.shape[-1]
    mask = torch.tril(torch.ones(n_total, n_total, device=score.device))[-n_new:]
    score = score.masked_fill(mask == 0, float('-inf'))
    score = torch.softmax(score, -1)

    Z = score @ V
    Z = Z.transpose(1,2)
    Z = Z.reshape(Z.shape[0], Z.shape[1], -1)

    Z = self.W0(Z)

    return Z

  def reset_cache(self):
    self.k_cache = None
    self.v_cache = None

Let's calculate how many matrix multiplications-additions we need to do with the KV cached version of the self attention block. Because we are using cache to store the previous key and value vectors, it is sufficient to only pass the previously computed token into the network, so the number of tokens t=1. Since batch size if 1 for inference, the dimension of the input vector x is (1, 1, e). Computing the query, key and value vector of the single passed token will take \(O(ek)\) time. Since e and k are constants, it is reduced to \(O(1)\). This is the first performance improvement. Earlier the same operation would have taken \(O(t)\). The second improvement comes from the attention score computation score = Q @ K.transpose(-2,-1). Since we have passed only one token in the forward pass, the dimension of Q will be (1, h, 1, k). As can be seen from the code the dimension of K.transpose(-2,-1) will be (1, h, k, t). So the total number of computations will be \(h.k.t\) . Since h and k are constants, the overall time complexity becomes \(O(t)\). This is a much improvement from \(O(t^2)\) previously.

Due to these performance improvements it becomes possible to serve LLM responses much faster. KV caching is a fundamental component of most popular inference engines like vLLM and SGLang. I did a performance test comparing both non cached and cached inference. The complete code for the same can be found here - https://github.com/pratyushranjan2/KV-cache-optimisation . With the non cached version the network was able to produce 283.5 tok/s. With the KV cached version it jumped to 1278.3 tok/s. That's a 4.5x improvement !

KV cache is one of many optimisations that is done to serve LLM responses faster. In this blog we saw an in memory version of KV cache. Real world inference engines like vLLM use a more advanced form called paged attention which is a different approach to store the KV cache inspired from virtual pages of operating system.

To conclude, AI inference is a new and emerging field within AI that is really exciting as it combines software engineering principles to serve AI workloads faster. There are many interesting things to learn. More on that in later blogs.