H2O KV Cache - An Easy Explanation
Heavy Hitter Oracle, the basic idea behind it is that, it retains the tokens which have better attention and evicts the tokens which don't, it also keeps the recent tokens in cache and that is how it improves kv caching performance.
H2O Paper link: https://arxiv.org/abs/2306.14048
The blog explains in the easiest possible way how H2O works.
1. What is a token?
A language model breaks text into small pieces called tokens.
For example:
The cat sat
may contain three tokens:
[The, cat, sat]
The model reads these tokens and tries to predict the next one.
2. What are Q, K, and V?
For every token, attention creates three vectors:
Q (Query): What am I looking for?
K (Key): What information do I have?
V (Value): What information can I give?
The query is compared with the keys. This tells the model which tokens deserve more attention. The model then collects information from their values.
A simple way to remember it is:
Q searches K and gets information from V.
3. What does @ mean?
In Python and PyTorch, @ means matrix multiplication.
In attention, we use:
query @ keys.T
This compares one query with every key. The result contains one match score for each key.
4. What is attention?
Suppose the cache contains:
[The, cat, sat]
The model may produce these attention values:
The: 0.20
cat: 0.50
sat: 0.30
The values add up to 1. A larger value means the current token is paying more attention to that cached token.
5. What is a KV cache?
The model calculates a key and value for every token. Old keys and values do not change, so calculating them again would waste time.
The model saves them:
Key cache Value cache
[K1] [V1]
[K2] [V2]
[K3] [V3]
When a new token arrives, its key and value are appended as new rows:
key_cache = torch.cat([key_cache, new_key], dim=0)
value_cache = torch.cat([value_cache, new_value], dim=0)
These two saved tensors are the KV cache.
The problem is that the cache gets bigger after every token. A long text can use a lot of memory.
6. What is H2O?
H2O means Heavy-Hitter Oracle. It keeps the cache small by keeping two kinds of tokens:
Important old tokens, called heavy hitters.
The newest tokens, called recent tokens.
For example:
heavy_size = 2
recent_size = 2
The cache can hold four tokens in total.
7. How does H2O find important tokens?
Every token has an importance score. It starts at zero.
When a token receives attention, H2O adds that attention to its score:
old importance + new attention = new importance
Example:
Token: The cat sat
Old score: 1.0 0.7 0.2
New attention: 0.2 0.5 0.3
New score: 1.2 1.2 0.5
Our PyTorch code does this here:
h2o_scores = h2o_scores + attention.squeeze(0)
This is the main idea behind H2O:
A token becomes important when the model pays attention to it many times.
H2O does not understand the meaning of the word itself. It watches the model's attention numbers.
8. How does eviction work?
Suppose we have:
Tokens: [The, cat, sat, because, it]
Scores: [1.78, 1.58, 0.91, 0.51, 0.22]
We want two heavy tokens and two recent tokens.
The recent tokens are always protected:
[because, it]
The older tokens are:
The: 1.78
cat: 1.58
sat: 0.91
H2O keeps the two largest scores:
[The, cat]
The final cache becomes:
[The, cat, because, it]
sat is removed.
The same positions must be removed from keys, values, and scores:
key_cache = key_cache[keep_indices]
value_cache = value_cache[keep_indices]
h2o_scores = h2o_scores[keep_indices]
9. Why keep recent tokens?
Old tokens have had many chances to receive attention. New tokens have not.
Without a recent section, H2O might remove a new token before it gets a chance to become important.
Being recent is temporary. As new tokens arrive, an older recent token must either become a heavy hitter or be removed.
10. Full cache compared with H2O
Our last lesson compared both methods:
Full cache: 7 tokens and 28 KV numbers
H2O cache: 4 tokens and 16 KV numbers
The full cache keeps everything, so its attention is exact. H2O removes some tokens, so later outputs can be slightly different.
The tradeoff is:
Full cache = more memory, exact result
H2O cache = less memory, approximate result
11. What each lesson contains
learn_01_h2o.py: Choose heavy and recent tokens.learn_02_scores.py: Add attention to importance scores.learn_03_evict.py: Remove low-scoring old tokens.learn_04_online.py: Process tokens one at a time.lesson_01_attention.py: Calculate attention with PyTorch.learn_05_kv_cache.py: Build a real key and value cache.learn_06_h2o_tensor_cache.py: Remove real K and V rows using H2O.learn_07_compare_full_vs_h2o.py: Compare a full cache with H2O.
12. How to run the lessons
The first four lessons use normal Python. The later lessons need PyTorch.
Create a new environment and install the packages:
python3 -m venv .venv
source .venv/bin/activate
python -m pip install torch numpy
Then run a lesson:
python learn_06_h2o_tensor_cache.py
CPU is enough for these small examples.
Final summary
For every new token:
1. Create Q, K, and V.
2. Add K and V to the cache.
3. Use Q and K to calculate attention.
4. Use attention and V to calculate the output.
5. Add attention to the H2O importance scores.
6. Keep important old tokens and recent tokens.
7. Remove the other K and V rows.
That is the basic H2O KV-cache algorithm.
The lessons are here: https://github.com/vivek378521/h2o-kv-cache-lessons
