Introduction to KV Caches
Transformers: Concept of self-Attention
Key Idea: Compare each input with others to determine relevance in context.
Step 1: Compute scores by comparing each input with all others using dot
product: score(xi , x j ) = xi · x j ( Higher score → More relevance)
Step 2: Normalize with softmax: αi j = softmax(score(xi , x j )) ∀ j ≤ i
Step 3: Compute output as weighted sum of inputs:
Concept of Self-Attention
Using multiple weight matrices
Query is current input
Key is previous input compared to query.;
Value used to compute output.
WQ, WK, WV : Weight matrices
Self Attention and K, V, Q Matrices
For Single token/word:
Vectors of dimension d
(embedding vector size)
For ALL the tokens in the
Sequence window
matrices of dimension n*d (n= no of
together: tokens in sequence)
• The n vector dot products are consolidated into n*d matrix dot products
Understanding KV Caching
Autoregressive generation in LLMs involves the model generates one word at a time,
using all the previously generated words.
When generating one token at a time:
• The model gets one new token at each step t.
• Naively, it recomputes K, V for all tokens 1...t every time.
That means lots of repeated work!
For example:
•Step 1: compute attention for token 1
•Step 2: recompute K and V for token 1 again + new ones for token 2
•Step 3: recompute K and V for token 1, 2 again + new for token 3 …
💡 Clearly, most computations are repeated.
Core Idea of KV Caching
Instead of recomputing everything, store the computed K and V vectors from previous
steps.
Caches -> KV Caches
Small, fast memory located close to the
CPU that stores frequently accessed data
and instructions.
• Reduces access time compared to
main memory.
• Improved CPU Performance
Principles of Locality
•Temporal Locality: Recently accessed data is
likely to be accessed again soon.
➤ Example: Loop variables.
•Spatial Locality: Nearby memory locations are
likely to be accessed soon. KV caches typically sit in GPU VRAMs
➤ Example: Accessing elements in an array.
KV computations without Cache
Assume we have word sequence“Time flies” and want to generate next word after this
To calculate the output of
second word we need the
key and values of both
words but need the query
of second word as we are
generating the output of
second word
Using the query of second
word we can calculate the
α (weights) for calculating
the final weighted sum of
attention layer
KV computations without Cache
Lets say we generated the new token “fast” after “Time Flies”, we want to generate next
word after “fast”
Now we will need the key and values of all the three words and query of the third word
“fast” to calculate the next word. See that we need to calculate KV values of first & second
word again. This is repetition.
KV computations with Cache
Lets say we generated the new token “fast” after “Time Flies”, we want to generate next
word after “fast”.
In this case we have cached
the KV values of first and
second word when we were
calculating the output of the
second word.
Now to generate the output
of the third word we use the
KV values of the first tow
words from cache and
compute the QKV values of
the third word only.
KV Cache Optimizations
The number of cached Key-Value (KV) pairs greatly affects inference speed and
memory usage. => Need to optimize the KV cache management
Token Level Optimization
Focus: Fine-grained KV management without changing model architecture
•Selection: Store only the most relevant tokens.
•Budget Allocation: Dynamically distribute cache memory across tokens.
(Each layer’s KV-cache contributing differently to model performance.
Allocates memory intelligently based on each layer’s contribution to
prediction accuracy, optimizing overall memory utilization.)
•Merging : Combine similar KV pairs to reduce redundancy.
•Quantization: Lower KV precision to save memory.
•Low-Rank Decomposition: Compress cache via matrix factorization.
Token Level Optimization Example
SnapKV (Token Selection Method) retrieves only the most important tokens based
on importance scores (attention scores of the tokens when used for generating next
tokens), showing that small subset of prompt tokens holds key information needed
for response generation.
Baseline models (with KV cache
& no optimization) latency
increases as sequence length
increases because as input gets
longer, the KV cache keeps
growing making decoding
slower
Blue lines indicating SnapKV optimization remains constant latency but increasing with the
batch size (number of user requests) as we are selecting important tokens to cache.
KV Cache Optimizations
Model Level Optimization
Focus: New architectures or mechanisms are designed for transformers to
allow more efficient reuse of KV cache
• Attention Grouping & Sharing: Reuse or share KV within/across layers.
• Architecture Alteration: New attention designs or auxiliary modules.
(Eg: Block Architecture with global and local attention blocks reducing
the cache memory requirement)
• Non-Transformer Architectures: Integrate RNN-style or hybrid models
for efficient memory use. (As transformers face KV cache challenges,
researchers are revisiting RNN principles that process sequences
efficiently without the need for KV caches.)
Model Level Optimization Example
•Multihead Attention (MHA): Every
head has its own key and value
•Multi-Query Attention (MQA): Uses
one shared key and value for all
attention heads, making decoding much
faster with only a small drop in quality.
•Grouped Query Attention (GQA):
Improves on MQA by grouping heads to
share keys and values within each
group, reducing training instability while
keeping good speed.
MQA & GQA reduces the time significantly compared to MHA
KV Cache Optimizations
System Level Optimization
Focus: Hardware and runtime-level cache efficiency
•Memory Management: Techniques like virtual memory adaptation, prefix
sharing, and layer-aware allocation.
•Scheduling: Prefix-aware reuse of cache content, and layer-specific cache
and scheduling for optimal memory utilization.
•Hardware Accelerators: GPU, I/O, SSD, and heterogeneous solutions for
faster KV access.
System Level Optimization Example
vLLM
•Existing LLM servers store each request’s
KV cache in large contiguous blocks, wasting
memory as token lengths vary.
•This causes high fragmentation, with much
of the pre-allocated space staying unused.
•vLLM fixes this using Paged Attention, which
splits memory into small blocks (pages) that can be
reused efficiently like a computer’s virtual
memory.
•This allows flexible memory sharing, so multiple
users or requests can run together without wasting
GPU memory.
vLLM reduces memory usage and increases throughput