How KV Cache Makes LLMs Fast, and Why It Eats GPU Memory

Every LLM you've used, ChatGPT, Claude, Gemini, Llama, does the same thing under the hood. It predicts one token, appends it to the text, and runs again. It's your phone keyboard's next-word suggestion, in a loop, a few hundred times per answer.

Done naively, that loop gets slower with every token, because each run reprocesses everything that came before. Every serious inference server avoids this with one optimization, the KV cache.

This post follows one tiny example through the whole idea, shows the code change, and then measures it. On my laptop the cache made generation about 9x faster with identical output. After that we'll look at the bill, because the same cache is why long contexts eat GPU memory and why the order of your prompt changes your API cost.

Here's the difference, replayed from my real timings.

An LLM is autocomplete in a loop

Here's the whole generation loop, stripped down.

tokens = tokenize("The server returned") for _ in range(max_new_tokens): logits = model(tokens) # scores for every possible next token, at every position  next_token = logits[-1].argmax() # take the last position's highest score  tokens.append(next_token) 

Let's follow it for three steps. The prompt is "The server returned" and the model continues with "a 500 error". (Models work on tokens, which are words or pieces of words. In this example every word is one token.)

Look at the line with model(tokens). It gets the entire sequence every time. At step 3 it processes "The server returned a 500" just to produce one word.

What the model does with all those tokens

Inside, a model is a stack of identical layers. GPT-2 small has 12, Llama 3.1 8B has 32. In every layer, each token goes through attention, the step where a token pulls in information from the tokens before it.

If you've written a database query, you already know most of attention. Each token is turned into three vectors (lists of numbers) by multiplying it with weights the model learned in training.

• The query (Q) is what this token is looking for.
• The key (K) is what this token is about, like an indexed column.
• The value (V) is the information this token hands over when it matches.

When the model works on "500", its query is compared with the key of every token so far. The better the match, the more of that token's value gets mixed in. It's a lookup, but a fuzzy one. Instead of returning the single matching row, it returns a weighted blend of all the rows.

That's how "500" picks up that it's about a server response and not a price or a speed limit. Real models run several of these lookups side by side in every layer, each with its own weights. Each one is called a head.

The waste

Now put the loop and the lookup together. At every step, model(tokens) computes K and V for every token in the sequence, in every layer, uses them once, and throws them away.

Follow "server" down the left grid. Its K and V get computed at step 1, again at step 2, again at step 3, and at every step after. Same token, same position, same result, every time. Over five steps that's 25 computations where 7 would do.

It gets worse quickly. In my test below (a 41-token prompt, then 400 new tokens), the naive loop pushed 96,200 tokens through the model. With a cache, it was 440. The wasted work grows with the square of the length.

Why reusing is safe

You can't cache something that might change. So could the K and V for "server" change once "error" shows up?

A token's K and V depend on the token, its position, and (from the second layer on) whatever it picked up from the tokens it could see. So if "server" were allowed to look at "error", its numbers would shift every time the sentence grew.

It isn't allowed. GPT-style models use a causal mask. Every token can attend to itself and the tokens before it, never the ones after. That's the same rule that lets them write left to right in the first place.

So "server" sees exactly "The server" at every step, forever, and its K and V in every layer come out identical every time. It's a pure function being called with the same input over and over. You'd memoize that.

The fix is a few lines

Keep each layer's K and V, and on every later step feed in only the newest token.

tokens = tokenize("The server returned") logits, cache = model(tokens, cache=None) # prefill builds the cache for the prompt for _ in range(max_new_tokens): next_token = logits[-1].argmax() tokens.append(next_token) logits, cache = model([next_token], cache=cache) # only the new token goes in 

Inside each attention layer the change is just as small.

q, k, v = project(h_new) # Q, K and V for the new token's state in this layer only cache.k = concat(cache.k, k) # append, never recompute cache.v = concat(cache.v, v) out = attend(q, cache.k, cache.v) # the new query reads every saved key and value 

Queries don't get cached. A token's query is only needed while that token is being processed. Keys and values get read by every token that comes after it, so those are what you keep.

You rarely write this yourself. Hugging Face transformers does it by default (use_cache=True), and vLLM, TensorRT-LLM and llama.cpp all build on it.

I measured it

I wanted to see the effect, not just read about it. So I ran GPT-2 small (124 million parameters) on my own laptop, CPU only, in 32-bit floats. A 41-token prompt, then 400 new tokens, always taking the most likely next token. One warm-up pass, then the same generation with and without the cache, timing every single step.

First check, did the cache change the output? No. Both versions produced exactly the same 400 tokens in every run.

That's about 9x faster overall. Without the cache, each token costs more than the one before, and by token 400 it's 12 to 14 times slower than token 1. With the cache, the line stays almost flat. It still creeps up from about 45 ms to about 63 ms, because each new token has to read every saved key and value. Reading is just far cheaper than recomputing. (The spike near token 345 is my laptop doing something else for a moment.)

What this test doesn't show:

• It's a tiny model on a laptop CPU. On a GPU the absolute numbers are completely different. The shape, flat versus climbing, is what carries over, and the gap widens as the text gets longer.
• Same tokens doesn't mean bit-identical maths. The two paths compute in different shapes, so the internal numbers differ by tiny rounding amounts. In 32-bit that never flipped a pick. In 16-bit on a GPU it occasionally can.
• I also ran it without a warm-up once (cached run first). That run showed 7.5x, because the first cached step paid 2.5 s of setup. The two runs above had a warm-up.

The full script

import time, torch from transformers import AutoTokenizer, AutoModelForCausalLM tok = AutoTokenizer.from_pretrained("openai-community/gpt2") model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2").eval() torch.set_num_threads(4) ids = tok("The quick brown fox jumps over the lazy dog. " * 4, return_tensors="pt").input_ids def generate(use_cache, new_tokens=400): x, past, times = ids, None, [] with torch.no_grad(): for _ in range(new_tokens): t0 = time.perf_counter() if use_cache: # only the newest token goes in, everything older comes from the cache  out = model(x if past is None else x[:, -1:], past_key_values=past, use_cache=True) past = out.past_key_values else: # the whole text goes in again, every single step  out = model(x, use_cache=False) next_id = out.logits[:, -1, :].argmax(-1, keepdim=True) x = torch.cat([x, next_id], dim=1) times.append(time.perf_counter() - t0) return x, times with torch.no_grad(): model(ids) # warm-up, so the first timed run isn't paying for setup  b, t_nocache = generate(False) a, t_cache = generate(True) print("same output:", torch.equal(a, b)) print(f"with cache {sum(t_cache):.1f}s, without {sum(t_nocache):.1f}s") 

The bill

The cache holds a key and a value for every token, in every layer, for every KV head. A KV head is one set of keys and values inside a layer, and each one is a short vector (64 numbers in GPT-2 small, 128 in Llama 3.1 8B). Multiply it out and you get a simple formula.

cache bytes per token = 2 (K and V) × layers × KV heads × numbers per head × bytes per number

For GPT-2 small in 32-bit, that's 2 × 12 × 12 × 64 × 4 = 73,728 bytes, which is exactly what I measured from the cache tensors in my test.

Now a model people actually deploy. Llama 3.1 8B in 16-bit, going by its published config, is 2 × 32 × 8 × 128 × 2 = 128 KB per token. An 8K-token conversation needs about 1 GB of cache. Its full 128K-token context needs about 16 GB. The model's own weights are also about 16 GB in 16-bit. So a single full-length conversation needs as much GPU memory as the model itself. (These are binary units, so treat them as round numbers.)

Put a few of those conversations on one big GPU and watch what happens.

If you've tuned Postgres, this will feel familiar. Every connection holds its own memory, so RAM caps how many connections you can open. Or think of browser tabs, each one holding its own memory until the laptop starts swapping. Here every conversation holds its own cache, so GPU memory caps how many people one server can answer at the same time. That's a big part of why long context costs more, and one of the reasons providers put a limit on it.

What this means when you build with LLMs

1. The pause and the stream are two different phases.

Prefill runs your whole prompt through the model in one go and builds its cache. That's the pause before the first token, usually measured as time to first token. Decode is everything after, one token per step, each step reading the whole cache. A huge prompt makes the pause longer, and a long conversation makes every decode step a little slower too.

2. Your API calls are stateless, so the cache usually doesn't survive them. Every chat request sends the whole conversation again. The provider builds the cache on its GPU, uses it while your response streams, and in the simplest setup throws it away when the response ends. The next request starts prefill from scratch.

That's what prompt caching fixes. Many providers can keep the cache for the start of a request and reuse it when your next request begins with exactly the same tokens. Some do it automatically, some need you to switch it on, and most only cache prefixes above a minimum length (often around 1,024 tokens). Cached input tokens are usually billed at a lower rate, though some providers charge a little extra the first time the cache is written. It works like a cache key on your prompt's prefix, so it has to match from the very first token. Put the stable parts first (system prompt, tool definitions, long documents) and the changing parts last (the user's new message). Put a timestamp at the top and every request misses.

3. If you self-host, memory usually runs out before compute does. One lever comes with the model you pick. The rest are serving settings.

• Grouped-query attention (a model design choice). Several query heads share one set of keys and values. Llama 3.1 8B has 32 query heads per layer but only 8 KV heads, so its cache is a quarter of what it would otherwise be. That's the 8 in the formula above. You get this by choosing a model that has it, not by flipping a setting.
• Paged KV cache. Older serving systems reserved one big block per request, sized for the longest answer it might give. The vLLM paper measured that only 20.4% to 38.2% of that cache memory held real token data. vLLM hands memory out in small pages instead, like an operating system does with RAM.
• Prefix caching. The server side of prompt caching. Requests that share a beginning share its cache.
• KV cache quantization. Store keys and values in 8 bits instead of 16, which roughly halves the cache for a small accuracy cost. vLLM and TensorRT-LLM both support it.

Recap

• An LLM generates in a loop, and a naive loop reprocesses the whole sequence for every new token.
• Attention needs the key and value of every earlier token, in every layer.
• The causal mask means those never change once computed, so they can be cached and reused.
• On my laptop that was about 9x faster with exactly the same output.
• The price is memory per token per layer, which is why long contexts and many users at once get expensive.

Next in this series is batching, and why answering 50 people at once can cost less per person than answering one.

Go deeper

• Understanding and Coding the KV Cache in LLMs from Scratch by Sebastian Raschka, a full implementation walkthrough
• The Illustrated GPT-2 by Jay Alammar, for a visual tour of masked self-attention
• Efficient Memory Management for LLM Serving with PagedAttention, the vLLM paper

About me. I'm Ashwani Kumar, a backend-heavy full-stack engineer from Bengaluru, and I run Beyond CRUD. It's a free site for developers who want to go past CRUD apps and understand what real systems deal with, like caching, queues, databases, retries, reliability, and AI in production. Many lessons come with a lab you can clone and run. This post is the first in a series on how LLMs work under the hood, and the whole series will live there too. If that's your kind of thing, come take a look.