Cache

Purpose

Autoregressive decoding caches the keys and values of every attention layer so that each new token only requires a single-token forward pass. The KV cache is the largest piece of state during generation and the right cache implementation can drop memory or unlock torch.compile. The cache hierarchy lives in src/transformers/cache_utils.py (1,574 LOC).

Key abstractions

Class	Where	When to use
Cache (base)	src/transformers/cache_utils.py	Abstract base; defines update, get_seq_length, reorder_cache
DynamicCache	same	Default; grows token-by-token. Used by greedy/sample without compile.
StaticCache	same	Pre-allocated tensors of fixed shape. Required for torch.compile.
SlidingWindowCache	same	Truncates to a sliding window (e.g., Mistral).
HybridCache	same	Mix of static + sliding (e.g., Gemma2, Llama4).
EncoderDecoderCache	same	Pairs an encoder cache with a decoder cache.
QuantizedCache, QuantizedCacheConfig	same	int4/int8 quantized KV state.
OffloadedCache, OffloadedStaticCache	same	Offload past layers to CPU.
MambaCache, RecurrentCache	src/transformers/models/<arch>/modeling_*.py	State-space models keep their own cache.

Why the abstraction exists

Before December 2023, every model maintained its own past_key_values tuple of tensors and reimplemented growth, sliding-window cropping, beam reordering, and serialization. PR #26681 consolidated that logic into cache_utils.py and a Cache base class. Models now read/write through cache.update(key, value, layer_idx) and let the cache class decide how to grow.

Choosing a cache

generate picks based on config.cache_implementation:

Setting	Resulting cache	Notes
None / "dynamic"	DynamicCache	Default
"static"	StaticCache	Compatible with torch.compile
"sliding_window"	SlidingWindowCache	For models like Mistral, Mixtral
"hybrid"	HybridCache	For Gemma2, Llama4, Phi3.5-mini
"quantized"	QuantizedCache	Memory savings
"offloaded"	OffloadedCache	Long-context, single-GPU
"offloaded_static"	OffloadedStaticCache	Long-context + compile

How a cache is used in forward

def forward(..., past_key_values: Cache | None = None, ...):
    ...
    for layer_idx, layer in enumerate(self.layers):
        ...
        key_states, value_states = past_key_values.update(key_states, value_states, layer_idx)
        ...

Beam search reorders the cache after each step:

past_key_values.reorder_cache(beam_idx)

Continuous batching uses a paged cache

src/transformers/generation/continuous_batching/ allocates a fixed pool of KV blocks (each cb_block_size tokens) and packs sequences from multiple concurrent requests into the same flattened tensor. The block table lives alongside the cache and the attention kernels (src/transformers/integrations/flash_paged.py, eager_paged.py, sdpa_paged.py) read from it.

graph LR
    R1[Request 1] -->|2 blocks| Pool[Block pool]
    R2[Request 2] -->|3 blocks| Pool
    R3[Request 3] -->|1 block| Pool
    Pool --> PagedAttn[Paged attention kernel]
    PagedAttn --> Logits[Logits per request]

torch.compile compatibility

Only fixed-shape caches are compatible with torch.compile because compile cannot recompile on shape changes. To compile a model:

1. Set model.config.cache_implementation = "static" (or hybrid/sliding_window if applicable).
2. Allocate via cache = StaticCache(config, max_batch_size=..., max_cache_len=..., device=..., dtype=...).
3. model.forward = torch.compile(model.forward, mode="reduce-overhead").
4. Pass the cache explicitly: model.generate(..., past_key_values=cache).

Quantized caches

QuantizedCache stores keys and values in int4/int8 and dequantizes per-step. Backed by quanto or hqq. Configuration:

from transformers import QuantizedCacheConfig, AutoModelForCausalLM
cfg = QuantizedCacheConfig(backend="quanto", nbits=4)
model.generate(..., cache_implementation="quantized", cache_config=cfg)

Offloaded caches

OffloadedCache keeps only the active layer's KV in GPU memory and rotates past layers to CPU. Useful for long context with a single GPU.

Integration points

• Generation — generate constructs caches.
• Attention — kernels read from caches via update.
• Modeling — forward accepts past_key_values: Cache.
• Continuous batching — paged caches.

Entry points for modification

• A new cache type → subclass Cache, implement update, get_seq_length, reorder_cache, optionally crop. Register in CACHE_MAP.
• For torch.compile support, ensure shapes are static and avoid Python-side conditionals.
• For paged variants, also update src/transformers/integrations/*_paged.py and the continuous-batching scheduler.