Want to master LLM Cache Management? Start with these resources.

KV Cache: https://huggingface.co/docs/transformers/main/en/cache_explanation… https://hamzaelshafie.bearblog.dev/paged-attention-from-first-principles-a-view-inside-vllm…

Prefix Caching https://docs.vllm.ai/en/latest/features/automatic_prefix_caching.html…

Continuous Batching https://anyscale.com/blog/continuous-batching-llm-inference…

Speculative Decoding https://research.google/blog/looking-back-at-speculative-decoding/…

KV Cache Quantization https://huggingface.co/blog/kv-cache-quantization…

Research Papers:

1. FlashInfer (Attention Engine) https://arxiv.org/abs/2501.01005
2. Zipage (Compressed PagedAttention) https://arxiv.org/abs/2603.08743
3. IceCache (Memory-Efficient KV Cache) https://arxiv.org/abs/2604.10539

Caching · Hugging Face

Source: https://huggingface.co/docs/transformers/main/en/cache_explanation Imagine you’re having a conversation with someone, and instead of remembering what they previously said, they have to start from scratch every time you respond. This would be slow and inefficient, right?

You can extend this analogy to transformer models. Autoregressive model generation can be slow because it makes a prediction one token at a time. Each new prediction is dependent on all the previous context.

To predict the 1000th token, the model requires information from the previous 999 tokens. The information is represented as matrix multiplications across the token representations.

To predict the 1001th token, you need the same information from the previous 999 tokens in addition to any information from the 1000th token. This is a lot of matrix multiplications a model has to compute over and over for each token!

A key-value (KV) cache eliminates this inefficiency by storing kv pairs derived from the attention layers of previously processed tokens. The stored kv pairs are retrieved from the cache and reused for subsequent tokens, avoiding the need to recompute.

Caching should only be used forinference. It may cause unexpected errors if it’s enabled during training.

To better understand how and why caching works, let’s take a closer look at the structure of the attention matrices.

https://huggingface.co/docs/transformers/main/en/cache_explanation#attention-matricesAttention matrices

Thescaled dot-product attentionis calculated as shown below for a batch of sizeb, number of attention headsh, sequence length so farT, and dimension per attention headd\_head.Attention(Q,K,V)=softmax(QK⊤dhead×mask)V\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^\top}{\sqrt{d_{\text{head}}}} \times \text{mask} \right) V

The query (Q), key (K), and value (V) matrices are projections from the input embeddings of shape\(b, h, T, d\_head\).

For causal attention, the mask prevents the model from attending to future tokens. Once a token is processed, its representation never changes with respect to future tokens, which meansKpastK_{\text{past}}andVpastV_{\text{past}}can be cached and reused to compute the last token’s representation.Attention(qt,[k1,k2,…,kt−1⏟cached,kt],[v1,v2,…,vt−1⏟cached,vt])\text{Attention}(q_t, [\underbrace{k_1, k_2, \dots, k_{t-1}}_{\text{cached}}, k_{t}], [\underbrace{v_1, v_2, \dots, v_{t-1}}_{\text{cached}}, v_{t}])

At inference time, you only need the last token’s query to compute the representationxtx_tthat predicts the next token $ t+1 $. At each step, the new key and value vectors arestoredin the cache andappendedto the past keys and values.Kcache←concat(Kpast,kt),Vcache←concat(Vpast,vt)K_{\text{cache}} \leftarrow \text{concat}(K_{\text{past}}, k_t), \quad V_{\text{cache}} \leftarrow \text{concat}(V_{\text{past}}, v_t)

Attention is calculated independently in each layer of the model, and caching is done on a per-layer basis.

Refer to the table below to compare how caching improves efficiency.

without cachingwith cachingfor each step, recompute all previousKandVfor each step, only compute currentKandVattention cost per step isquadraticwith sequence lengthattention cost per step islinearwith sequence length (memory grows linearly, but compute/token remains low)## https://huggingface.co/docs/transformers/main/en/cache_explanation#cache-classCache class

A basic KV cache interface takes a key and value tensor for the current token and returns the updatedKandVtensors. This is internally managed by a model’sforwardmethod.

new_K, new_V = cache.update(k_t, v_t, layer_idx)
attn_output = attn_layer_idx_fn(q_t, new_K, new_V)

When you use Transformers’Cacheclass, the self-attention module performs several critical steps to integrate past and present information.

1. The attention module concatenates current kv pairs with past kv pairs stored in the cache. This creates attentions weights with the shape\(new\_tokens\_length, past\_kv\_length \+ new\_tokens\_length\). The current and past kv pairs are essentially combined to compute the attention scores, ensuring a model is aware of previous context and the current input.
2. When theforwardmethod is called iteratively, it’s crucial that the attention mask shape matches the combined length of the past and current kv pairs. The attention mask should have the shape\(batch\_size, past\_kv\_length \+ new\_tokens\_length\). This is typically handled internally ingenerate(), but if you want to implement your own generation loop withCache, keep this in mind! The attention mask should hold the past and current token values.

https://huggingface.co/docs/transformers/main/en/cache_explanation#cache-storage-implementationCache storage implementation

Caches are structured as a list of layers, where each layer contains a key and value cache. The key and value caches are tensors with the shape\[batch\_size, num\_heads, seq\_len, head\_dim\].

Layers can be of different types (e.g.DynamicLayer,StaticLayer,StaticSlidingWindowLayer), which mostly changes how sequence length is handled and how the cache is updated.

The simplest is aDynamicLayerthat grows as more tokens are processed. The sequence length dimension (seq\_len) increases with each new token:

cache.layers[idx].keys = torch.cat([cache.layers[idx].keys, key_states], dim=-2)
cache.layers[idx].values = torch.cat([cache.layers[idx].values, value_states], dim=-2)

Other layer types likeStaticLayerandStaticSlidingWindowLayerhave a fixed sequence length that is set when the cache is created. This makes them compatible withtorch\.compile. In the case ofStaticSlidingWindowLayer, existing tokens are shifted out of the cache when a new token is added.

The example below demonstrates how to create a generation loop withDynamicCache. As discussed, the attention mask is a concatenation of past and current token values.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache
from accelerate import Accelerator

device = Accelerator().device

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map=device)
tokenizer = AutoTokenizer.from_pretrained(model_id)

past_key_values = DynamicCache(config=model.config)
messages = [{"role": "user", "content": "Hello, what's your name."}]
inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)

generated_ids = inputs.input_ids
max_new_tokens = 10

for _ in range(max_new_tokens):
    outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
    
    next_token_ids = outputs.logits[:, -1:].argmax(-1)
    generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1)
    
    
    attention_mask = inputs["attention_mask"]
    attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
    inputs = {"input_ids": next_token_ids, "attention_mask": attention_mask}

print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])
"[INST] Hello, what's your name. [/INST]  Hello! My name is LLaMA,"

Updateon GitHub