The KV Cache: Memory Usage in Transformers

Key Takeaways

Transformers are used almost everywhere in natural language processing models like GPT can write pages of coherent text something that is really impressive but one fundamental limitation of Transformer language models is that as you generate more and more text you will use up more and more GPU memory and eventually you reach a point where your GPU runs out of memory your program crashes and you cannot generate any more text my name is bai I'm a machine learning engineer and a PhD in natural language processing and today I will explain why is this why is it that Transformer language models require so much more memory when it deals with longer text actually if you look at open ai's API pricing for GPT you see something interesting they charge you twice as much per input token to use the longer context model this is one of The Economic Consequences of the high memory usage when you need to handle large context links most of the memory usage is taken up by the KV cache or the key value cache in this video I explain exactly what this is and why we need it before I explain the KV cache let's quickly go over what happens in the self-attention mechanism when we generate a sentence at the beginning of a transformer layer each token corresponds to an embedding Vector X the first thing that happens is X is multiplied by three different matrices to generate the query key and value vectors these three matrices denoted by WQ w k and WV are learned from data during decoding these three are not the same size in fact the query q is usually a vector but the K and V are matrices this is how I like to think about it the query Vector represents the new token in this decoder step and since there is only one token this is a vector instead of a matrix the key Matrix represents all the previous contexts that the model should attend to and finally the value Matrix also represents all the previous context but is applied after softmax as a weighted sum during their attention mechanism we first take a DOT product between the query vector and the key Matrix then we take a soft Max and apply that as a weighted sum over the value Matrix in Auto regressive decoding we are generating one word at a time given all of the previous context so the K and V matrices contain information about the entire sequence but the query Vector only contains information about the last token that we have seen you can think of the dot product between q and K as doing attention between the current token that we care about and all of the previous tokens at the same time as we generate a sequence one token at a time the K and V matrices actually don't change very much this token corresponds to a column of the K Matrix and a row of the V Matrix and The crucial thing is that once we've computed the embedding for this word it's not going to change again no matter how many more words we generate but the model still has to do the heavy work of computing the key and the value vectors for this word on all subsequent steps this results in a quadratic number of Matrix Vector multiplications which is going to be really slow as an analogy imagine if you are a model writing a sentence one word at a time but each word you write you have to read every word that you've written before and then use that information to generate the next word obviously this is extremely inefficient and it will be much better if you could somehow remember what you wrote as you're writing it now we're finally ready to explain how the KV cache works when the model reads a new word it generates the query Vector as before but we cached the previous values for the key and value matrices so we no longer have to compute these vectors for the previous context instead we only have to compute one new column for the key Matrix and one new row for the value Matrix and then we proceed with the dot product and soft Max as usual to compute the scaled dot product attention by the way if you like this video so far please give me a thumbs up to feed the YouTube algorithm And subscribe to my channel now let's talk about how the KV cache fits in with the rest of the Transformer here we have the self-attention layer and instead of passing in a whole sequence of embeddings now we only pass in the previous knv cache and the embedding for the current token the self-attention layer computes to the new key and value vectors for the current token and appends them to the KV cache we will then need to store these key and value matrices somewhere in the gpu's memory so that we can retrieve them later when we're working on the next token notice that the only part of the model where the current token interacts with the previous token is the self-attention layer in every other layer such as the positional embedding the layer norm and the feed 4 neural network there is no interaction between the current token and the previous context so when we're using the KV cache we only have to do a constant amount of work for each new token and this work does not get bigger when the sequence gets longer now let's look at how much memory it takes to store the KV cache here is the equation for the memory usage first we have two because there are two matrices K and V that we need to explore Precision is the number of bytes per parameter for example in fp32 there are four bytes per parameter and layers is the number of layers in the model the model is the dimension of the embeddings in each layer the sequence length is the length that we need to generate at the end including all of the prompt tokens and everything that we generate and finally batch is the batch size we multiply these all together to get the total memory usage of the KVA cache let's walk through an example involving a 30 billion parameter model which nowadays is considered like medium large we have to store Two for the matrices K and V typically the position is 2 because influence is done in 16 bits and not 32. the number of layers in this model is 48 and the dimension size of this model is around 7000 and let's say that we capped the max sequence length at 102 full and we use a batch size of 128. if we multiply everything together we get that the KV cache for this model is 180 gigabytes and the model itself is 2 times 30 billion which is 60 gigabytes so you can see that the KV cache takes up three times as much memory as the model itself and this sort of ratio is pretty typical for inference scenarios and KV cache tends to be the dominant factor in memory usage during inference one more thing to be aware of is the difference in latency in processing that prompt versus subsequent tokens when the model is given the prompt and is deciding the first token to generate this has higher latency because there is no KV cache yet so it has to compute the K and V matrices for every token in the prompt but after this has been done each subsequent token will have lower latency because it only has to compute K and V for one token that's it for the KVA cache if you have any questions please don't hesitate to leave a comment below and if you found this content helpful please like And subscribe to my channel so you can get notified when I make new machine learning videos it will help me out a lot goodbye