AutoRegressive Generation
Unlike training, where the transformer processes an entire sequence in parallel, autoregressive generation produces text one token at a time. Suppose the prompt is “I love”. The input tokens are first converted into token embeddings and positional embeddings, which are added together to form the input to the first transformer layer. At each transformer layer, every token is independently projected into Query (Q), Key (K), and Value (V) vectors using learned linear projections. Attention is then computed using causal masking, meaning a token can attend only to itself and the tokens before it. For example, the token “I” can attend only to “I”, while “love” can attend to “I” and “love”, but never to future tokens. After attention, the output passes through the feed-forward network, producing a new hidden state for each token. This process repeats across all transformer layers. At the final layer, only the hidden state of the last token is passed through the language modeling (LM) head to produce logits over the vocabulary, from which the next token is selected (e.g., “transformers”). That predicted token is appended to the sequence, giving “I love transformers”, and the entire process repeats to generate the next token. We are generating tokens in a Autoregressive fashion of appending the predicted next tokens to the end of sequence and running the inference loop again.
Why inference becomes slow without KV cache
A straightforward implementation performs a complete forward pass over the entire sequence at every generation step. After generating “transformers”, the model now receives “I love transformers” as input. It again computes embeddings, projects every token into new queries, keys, and values, performs attention for every token at every layer, runs every feed-forward network, and produces new hidden states, even for tokens that were already processed in the previous iteration.
This repeated computation is unnecessary because of the causal mask. A token’s hidden state depends only on itself and the tokens before it. Future tokens can never influence it. For example, the hidden state of “I” depends only on “I”, while the hidden state of “love” depends only on “I” and “love”. Adding the future token “transformers” does not change either of these hidden states. Since the hidden states remain unchanged, their corresponding Q, K, and V vectors also remain unchanged. Nevertheless, without KV cache, the model recomputes these identical vectors during every forward pass. As the sequence grows longer, more and more previously processed tokens are recomputed, causing inference latency to increase with every generated token. The model is repeatedly performing the same calculations and obtaining exactly the same results, making inference increasingly inefficient.
For example, for a sequence length of 500, once the next token is predicted the sequence length increases to 501 and all the past 500 tokens Q, K, V projects are generated again across all layers even though its value haven’t changed. Ideally a optimal algorithm should generate Q, K, V projects for only the new token (501th token). This optimization will be implemented in our tiny transformer.
What exactly gets cached (K and V, not Q)
The redundancy described above motivates KV caching. Once the hidden state of a token has been computed at a particular transformer layer, its corresponding Key (K) and Value (V) vectors are fixed because neither the hidden state nor the model weights change during inference. These K and V vectors are therefore computed once and stored in a cache for each transformer layer. When a new token is generated, the model computes only the hidden state, query, key, and value for this new token. The new key and value are appended to the cache, while the cached keys and values of all previous tokens are reused directly for attention.
The Query (Q) is not cached because it serves a fundamentally different purpose. A query represents the current token asking the question, “Which previous tokens should I attend to?” Once that question has been answered and the token’s hidden state has been computed, that query is never needed again. Future tokens do not reuse old queries; instead, each newly generated token creates its own new query because it must decide how to attend to the accumulated context. In contrast, the keys and values of previous tokens continue to act as the memory that every future token consults during attention. Thus, queries are transient and used only once, whereas keys and values remain useful throughout the remainder of the generation process, which is why only K and V are cached.
Implementing KV cache in a tiny transformer
We are breaking down our process into 2 steps,
- Prefill
- Decode
Once we receive the user prompt, we will go through the prefill phase and generated K, V for all the tokens present in it and store in cache. Since the cache is now populated, we can perform next token prediction (decode) a single token and append the new K, V to the cache.
These changes takes place only in the MHA layer and transformer’s forward call. In MHA we create and populate the kv cache for that layer while the transformer combines the cache for all layers and return all model output. This output is later passed as parameter to the model to update the cache for newly generated tokens.
class MultiHeadAttention(nn.Module):
def __init__(self, heads, d_model):
super().__init__()
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
self.d_model = d_model
self.heads = heads
assert self.d_model % self.heads == 0, "d_model must be divisible by heads"
self.head_dim = self.d_model // self.heads
def prefill(self, x):
B, S, d_model = x.shape
Q = self.W_q(x)
K = self.W_k(x)
V = self.W_v(x)
Q = Q.view(B, S, self.heads, self.head_dim).transpose(1, 2)
K = K.view(B, S, self.heads, self.head_dim).transpose(1, 2)
V = V.view(B, S, self.heads, self.head_dim).transpose(1, 2)
kv_cache = (K, V)
attention_weights = Q @ K.transpose(-1, -2)
attention_weights = attention_weights / (self.head_dim ** 0.5)
mask = torch.tril(torch.ones(S, S, device=x.device)).view(1, 1, S, S)
attention_weights = attention_weights.masked_fill(mask == 0, float('-inf'))
attention_probs = torch.softmax(attention_weights, dim=-1)
attention = attention_probs @ V
attention = attention.transpose(1, 2).contiguous().view(B, S, d_model)
return self.W_o(attention), kv_cache
def decode(self, x, kv_cache):
B, S, d_model = x.shape
Q = self.W_q(x)
K = self.W_k(x)
V = self.W_v(x)
Q = Q.view(B, S, self.heads, self.head_dim).transpose(1, 2)
K = K.view(B, S, self.heads, self.head_dim).transpose(1, 2)
V = V.view(B, S, self.heads, self.head_dim).transpose(1, 2)
K = torch.cat([kv_cache[0], K], dim=2)
V = torch.cat([kv_cache[1], V], dim=2)
kv_cache = (K, V) # updated cache
attention_weights = Q @ K.transpose(-1, -2)
attention_weights = attention_weights / (self.head_dim ** 0.5)
# we don't need the mask since we obviouly can't know the future token
attention_probs = torch.softmax(attention_weights, dim=-1)
attention = attention_probs @ V
attention = attention.transpose(1, 2).contiguous().view(B, S, d_model)
return self.W_o(attention), kv_cache
def forward(self, x, kv_cache=None):
if kv_cache is None :
return self.prefill(x)
else:
return self.decode(x, kv_cache)
Based on the presence of cache param, we identify the flow. In the prefill stage we create the kv cache and return it while in decode phase, we update the cache. The input to the prefill phase will be more than one token ( Ex : training / inference on first step) but the decode phase receives only one token.
class Transformer(nn.Module):
def prefill(self, x):
B, S = x.shape
out = token_emb + position_emb # (B, S, d_model)
kv_cache = {}
for layer_no, layer in enumerate(self.decoder):
out, layer_kv_cache = layer(out) # (B, S, d_model)
kv_cache[layer_no] = layer_kv_cache
return out, kv_cache
def decode(self, x, kv_cache):
assert x.shape[1] == 1, "Decode expects exactly one token."
B, S = x.shape
out = token_emb + position_emb
for layer_no, layer in enumerate(self.decoder):
out, layer_kv_cache = layer(out, kv_cache[layer_no])
kv_cache[layer_no] = layer_kv_cache
return out, kv_cache
def forward(self, x, kv_cache=None):
if kv_cache is None:
return self.prefill(x)
else:
return self.decode(x, kv_cache)
In our transformer layer, we create the placeholder to store layer-wise cache and layer for decode phase we pass the cache for that layer to use instead of generating again.
Inference
Once our model is trained this way, we need to perform our inference utilizing this optimization. Lets first see infernce without kv cache and later induce the kv cache algorithm.
def inference(prompt, max_new_tokens=50):
ids = []
for token in prompt:
if token not in char_to_id:
print(f"Unknown character in prompt : {token}")
ids.append(char_to_id[token])
with torch.inference_mode():
for _ in range(max_new_tokens):
context = ids[-seq_len:]
inputs = torch.tensor([context], device=device).long() # (1, S)
logits, kv_cache = model(inputs)
logits = logits[-1]
next_token = int(logits.argmax().item())
ids.append(next_token)
return "".join([id_to_char[token_id] for token_id in ids])
Inference is just three simple steps,
- Run the model on token ids
- Fetch the logits of the last token
- Using decoding strategy(here we used greedy decoding), fetch the next token and append it to prompt Repeat the process till we reach max tokens
Inference with kv cache
def inference_from_cache(prompt, max_new_tokens=50):
ids = []
for token in prompt:
if token not in char_to_id:
print(f"Unknown character in prompt : {token}")
ids.append(char_to_id[token])
kv_cache = None
with torch.inference_mode():
# Prefill phase
context = ids[-seq_len:]
inputs = torch.tensor([context], device=device).long()
logits, kv_cache = model.prefill(inputs)
logits = logits[-1]
next_token = int(logits.argmax().item())
ids.append(next_token)
# Decode phase
for _ in range(max_new_tokens-1):
inputs = torch.tensor([[next_token]], device=device).long()
logits, kv_cache = model.decode(inputs, kv_cache)
logits = logits[-1]
next_token = int(logits.argmax().item())
ids.append(next_token)
return "".join([id_to_char[token_id] for token_id in ids]), last_logits_kv_cache
- Prefill phase - Run the entire prompt and generate the next token. Also prepare the kv cache
- Decode phase - For the recent token, run the model and use the populated cache to avoid reduntant calc and use stored memory
We all know that our optimization is a design choice and better engineering. We didn’t change the mathematics of the model or the way its trained. We simply reduced the number of operations which saves us memory, compute and time. To justify this, I did a sanity tesing passing the same prompt to both the inference and the model returned same results. Also verified the logits of the last token to ensure more strictly.
inference("The cat sat on") -> 'The cat sat on the start of the season . The second half of the '
inference_from_cache("The cat sat on") -> 'The cat sat on the start of the season . The second half of the '
With KV cache, we are just doing a highly efficient way of inference. Instead of recalculating Q, K, V for tokens that are already calculated we are optimizing. This is just like taking a code, refactoring and defrencing a few variables, to make it fast and optimized.
The idea is: “Did my optimization change the model’s behavior?” It shouldn’t. KV cache is purely an optimization.
Measuring Improvements and Speedup
Since this is an optimization, we expect results, fastness and less usage of memory. We have to verify our optimization, When we measure the time
There are three common metrics:
- Time to First Token (TTFT) — How long prefill() takes.
- Time Per Output Token (TPOT) — How long each decode() step takes.
- Total generation time — Time to generate N tokens.
KV cache primarily improves TPOT.
But here we are going to measure the total generation time, which combines the first two metrics
import time
for tokens_size in [20, 40, 60, 80, 100, 110]:
print(f"For token size : {tokens_size}")
torch.cuda.synchronize()
start = time.perf_counter()
out, last_logits_non_kv_cache = inference("The cat sat on", max_new_tokens=tokens_size)
torch.cuda.synchronize()
elapsed_baseline = time.perf_counter() - start
torch.cuda.synchronize()
start = time.perf_counter()
out, last_logits_kv_cache = inference_from_cache("The cat sat on", max_new_tokens=tokens_size)
torch.cuda.synchronize()
elapsed_kv = time.perf_counter() - start
print(f"baseline : {round(elapsed_baseline * 1000, 2)} ms || kv cache : {round(elapsed_kv * 1000, 2)} ms")
On GPU, you should synchronize because CUDA is asynchronous. Otherwise you’ll measure kernel launch time rather than execution time.
The results are
For token size : 20
Baseline : 22.63 ms || KV cache : 12.39 ms
For token size : 40
Baseline : 30.37 ms || KV cache : 24.57 ms
For token size : 60
Baseline : 45.64 ms || KV cache : 36.65 ms
For token size : 80
Baseline : 61.18 ms || KV cache : 49.2 ms
For token size : 100
Baseline : 78.02 ms || KV cache : 60.74 ms
For token size : 110
Baseline : 86.73 ms || KV cache : 66.47 ms
speedup % = ((Avg kv sec - Avg org sec) / Avg kv sec) * 100
This gives the latency reduction and we’re seeing about 18–23% latency reduction.
KV cache memory formula
Lets also measure the memory kv cache takes,
For one layer, memory = 2 * B * S * d_model * bytes per element
The factor of 2 is due to one cache for keys and one cache for values.
For our model, B = 32, S = 128, d_model = 512, dtype = FP32 which is 4 bytes per element. memory = 16, 777, 216 bytes = 16 MB
We are not including heads since (B, heads, S, head_dim) == (B, S, d_model) in terms of counts of elements
For all the layers = L * 2 * B * S * d_model * bytes per element memory = 4 * 16 MB = 64 MB
How cache grows with sequence length
KV cache grows linearly with sequence length, we can also observe it from formula that if we increase the sequence length, the memory increases all well.
Seq len - Relative memory 512 - 1 x 1024 - 2 x 2048 - 4 x 4096 - 8 x
Double the sequence length will also double the memory
Why long context models become expensive
Because KV cache grows linearly with the context length. Every new token stores one key and one value for every decoder layer. So,
KV Cache Memory is directly propertional to L * S * d_model
where (S) is the context length.
Example (12M model)
- 4K context → ~ 524 MB KV cache
- 8K context → ~1 GB
- 16K context → ~2 GB
- 32K context → ~4 GB
- 64K context → ~8 GB
Just doubling the context doubles the KV cache memory.
Long-context models require:
- More GPU memory (larger KV cache)
- More memory bandwidth (reading a larger cache every token)
- More attention computation per generated token
That’s why techniques like sliding-window attention, paged KV cache, and GQA are used to reduce memory usage or improve throughput for long-context inference.
As I researched,
Production LLMs handle very long generation using several strategies:
- Sliding-window attention: Keep only the most recent (N) tokens.
- Paged KV cache: Store the cache in memory-efficient blocks.
- Cache eviction/compression: Discard or compress older parts of the cache.
- Very large context models: Allocate enough memory for long contexts.
Difficulty faced during this experiment
In my Measuring Improvements and Speedup section the token sizes I used are less than 128, Apparently when I gave longer max_new_tokens, the model didn’t have learned positional encodings.
The key limitation. With learned positional embeddings:
- The cache keeps growing:
Prompt: 100 tokens
Decode 1 → cache = 101 Decode 2 → cache = 102 … Decode 50 → cache = 150
But your positional embedding table only contains positions = 0 … 127 So eventually you ask for position = 128, the embedding lookup fails. That’s exactly what happened.
RoPE removes the embedding lookup limit.
Learned positional embeddings
Limitation: Embedding table 0, 1, 2 … 127 Position 128 doesn’t exist.
RoPE
There is no embedding table. Position is computed mathematically: Position 128 -> rotation by angle θ(128) So position 128, 500, 2000, etc., are all computable. That’s why RoPE doesn’t crash due to an embedding lookup.
RoPE addresses the position representation problem
Tip : When this limitation hit, there was a assertion ( device-side assert )
Once CUDA hits an assertion, the CUDA context is poisoned. Every later CUDA call (even torch.cuda.synchronize()) will fail.
What to do ? Restart the Jupyter kernel (or Python process). Reload the model. Run the benchmark again.
The assertion we hit : indexSelectSmallIndex:
Assertion srcIndex < srcSelectDimSize failed