Skip to content

🚨 Generation cache preparation #43679

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zucchini-nlp merged 11 commits into from
Feb 4, 2026
Merged

🚨 Generation cache preparation #43679

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
dfad28c
do smth
zucchini-nlp Feb 2, 2026
b241557
how did this even happen O_o
zucchini-nlp Feb 2, 2026
bd31bb2
fix models
zucchini-nlp Feb 2, 2026
c042373
dia needs only one prefill, fix
zucchini-nlp Feb 2, 2026
d58e5b7
fxi repo again
zucchini-nlp Feb 2, 2026
524734f
comments, comments
zucchini-nlp Feb 3, 2026
4493bb2
almost forgot
zucchini-nlp Feb 3, 2026
e10455e
explicit keyword arg
zucchini-nlp Feb 3, 2026
57f2463
Merge branch 'main' into cache-prefill-chunk
zucchini-nlp Feb 3, 2026
0614fd0
Merge branch 'main' into cache-prefill-chunk
zucchini-nlp Feb 4, 2026
db40d72
Merge branch 'main' into cache-prefill-chunk
zucchini-nlp Feb 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions src/transformers/cache_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from .configuration_utils import PreTrainedConfig
from .utils import (
is_hqq_available,
is_optimum_quanto_available,
is_quanto_greater,
is_torch_greater_or_equal,
is_torchdynamo_compiling,
Expand Down Expand Up @@ -584,7 +585,12 @@ def __init__(
)

# We need to import quanto here to avoid circular imports due to optimum/quanto/models/transformers_models.py
if is_quanto_greater("0.2.5", accept_dev=True):
if not is_optimum_quanto_available():
raise ImportError(
"You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto "
"backend. Please install it via with `pip install optimum-quanto`"
)
elif is_quanto_greater("0.2.5", accept_dev=True):
from optimum.quanto import MaxOptimizer, qint2, qint4
else:
raise ImportError(
Expand Down Expand Up @@ -634,7 +640,10 @@ def __init__(
)

if not is_hqq_available():
raise ImportError("You need to install `hqq` to use `HQQQuantizedLayer`")
raise ImportError(
"You need to install `HQQ` in order to use KV cache quantization with HQQ backend. "
"Please install it via with `pip install hqq`"
)

if self.nbits not in [1, 2, 3, 4, 8]:
raise ValueError(
Expand Down
131 changes: 52 additions & 79 deletions src/transformers/generation/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,8 +48,6 @@
ModelOutput,
TransformersKwargs,
is_accelerate_available,
is_hqq_available,
is_optimum_quanto_available,
is_torchdynamo_exporting,
logging,
)
Expand Down Expand Up @@ -1858,20 +1856,19 @@ def _get_initial_cache_position(self, seq_length, device, model_kwargs):
model_kwargs["cache_position"] = cache_position
return model_kwargs

def _get_cache(self, cache_implementation: str, batch_size: int, max_cache_len: int, model_kwargs) -> Cache:
def _prepare_static_cache(
self, cache_implementation: str, batch_size: int, max_cache_len: int, model_kwargs
) -> Cache:

@zucchini-nlp zucchini-nlp Feb 2, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just a naming change, imo this is more descriptive

"""
Sets a cache for `generate`, that will persist across calls. A new cache will only be initialized a
new `generate` call requires a larger cache or uses a different batch size.

Returns the resulting cache object.
"""
requires_cross_attention_cache = (
self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None
)
offload_cache = "offloaded" in cache_implementation

if hasattr(self, "_cache"):
cache_to_check = self._cache.self_attention_cache if requires_cross_attention_cache else self._cache
cache_to_check = self._cache.self_attention_cache if self.config.is_encoder_decoder else self._cache

need_new_cache = (
not hasattr(self, "_cache")
Expand All @@ -1880,7 +1877,7 @@ def _get_cache(self, cache_implementation: str, batch_size: int, max_cache_len:
or cache_to_check.max_cache_len < max_cache_len
)

if requires_cross_attention_cache and hasattr(self, "_cache"):
if self.config.is_encoder_decoder and hasattr(self, "_cache"):
need_new_cache = (
need_new_cache
or self._cache.cross_attention_cache.max_cache_len != model_kwargs["encoder_outputs"][0].shape[1]
Expand All @@ -1893,7 +1890,7 @@ def _get_cache(self, cache_implementation: str, batch_size: int, max_cache_len:
"offloading": offload_cache,
}
self._cache = StaticCache(**self_attention_cache_kwargs)
if requires_cross_attention_cache:
if self.config.is_encoder_decoder:
cross_attention_cache_kwargs = {
"config": self.config.get_text_config(decoder=True),
"max_cache_len": model_kwargs["encoder_outputs"][0].shape[1],
Expand Down Expand Up @@ -1936,12 +1933,9 @@ def _prepare_cache_for_generation(
instantiated, writes it to `model_kwargs`, under the name expected by the model.
"""

is_hybrid_cache = any(class_name in self.__class__.__name__.lower() for class_name in ["mamba", "falconh1"])
cache_name = "past_key_values" if not is_hybrid_cache else "cache_params"

requires_cross_attention_cache = (
self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None
)
# TODO @raushan, unify cache arg naming for all models
is_linear_attn_cache = "mamba" in self.__class__.__name__.lower()
cache_name = "past_key_values" if not is_linear_attn_cache else "cache_params"

# Quick escape route 1: if the user specifies a cache, we only need to check for conflicting `generate` arguments
user_defined_cache = model_kwargs.get(cache_name)
Expand Down Expand Up @@ -1973,76 +1967,55 @@ def _prepare_cache_for_generation(

# Otherwise we NEED to prepare a cache, based on `generation_config.cache_implementation`

# TODO(joao): support static caches in assisted generation. assisted generation needs to roll back caches,
# which is only supported in dynamic caches atm
if (
generation_mode == GenerationMode.ASSISTED_GENERATION
and generation_config.cache_implementation is not None
):
logger.warning_once(
"An assistant model is provided, using a dynamic cache instead of a cache of type="
f"'{generation_config.cache_implementation}'."
)
generation_config.cache_implementation = None

# Assisted decoding and contrastive search require cache rollback, which is incompatible with sliding layers.
# To handle this, we skip passing the model config to DynamicCache (forcing a full-layer cache).
# The "dynamic_full" option is a shortcut for generate() users to avoid sliding layers on their own.
if (
generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH)
or generation_config.cache_implementation == "dynamic_full"
):
dynamic_cache_kwargs = {}
else:
dynamic_cache_kwargs = {"config": self.config.get_text_config(decoder=True)}
if generation_config.cache_implementation is not None:
if generation_config.cache_implementation in ALL_STATIC_CACHE_IMPLEMENTATIONS:
if generation_config.cache_implementation in DEPRECATED_STATIC_CACHE_IMPLEMENTATIONS:
logger.warning_once(
f"Using `cache_implementation='{generation_config.cache_implementation}' is deprecated. "
f"Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, and the layer structure will be "
"inferred automatically."
)
model_kwargs[cache_name] = self._get_cache(
cache_implementation=generation_config.cache_implementation,
batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size,
max_cache_len=max_cache_length,
model_kwargs=model_kwargs,
if generation_mode in (GenerationMode.ASSISTED_GENERATION, GenerationMode.CONTRASTIVE_SEARCH):
if generation_config.cache_implementation is not None:
logger.warning_once(
"An assistant model is provided, using a dynamic cache instead of a cache of type="
f"'{generation_config.cache_implementation}'."
)
elif generation_config.cache_implementation == "quantized":
if self.config.is_encoder_decoder or not self._supports_default_dynamic_cache():
raise ValueError(
"This model does not support the quantized cache. If you want your model to support quantized "
"cache, please open an issue and tag @zucchini-nlp."
)
generation_config.cache_implementation = "dynamic_full"

dynamic_cache_kwargs = {}
if generation_config.cache_implementation != "dynamic_full":
dynamic_cache_kwargs["config"] = self.config.get_text_config(decoder=True)

if generation_config.cache_implementation == "offloaded":
dynamic_cache_kwargs["offloading"] = True

if generation_config.cache_implementation in ALL_STATIC_CACHE_IMPLEMENTATIONS:
if generation_config.cache_implementation in DEPRECATED_STATIC_CACHE_IMPLEMENTATIONS:
logger.warning_once(
f"Using `cache_implementation='{generation_config.cache_implementation}' is deprecated "
f"and will be removed in v5.13. Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, "
"and the layer structure will be inferred automatically."
)
Comment thread
zucchini-nlp marked this conversation as resolved.
Resolved
model_kwargs["past_key_values"] = self._prepare_static_cache(
cache_implementation=generation_config.cache_implementation,
batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size,
max_cache_len=max_cache_length,
model_kwargs=model_kwargs,
)
elif generation_config.cache_implementation == "quantized":
if self.config.is_encoder_decoder or not self._supports_default_dynamic_cache():
raise ValueError(
"This model does not support the quantized cache. If you want your model to support quantized "
"cache, please open an issue and tag @zucchini-nlp."
Comment thread
zucchini-nlp marked this conversation as resolved.
Resolved
)

cache_config = generation_config.cache_config if generation_config.cache_config is not None else {}
cache_config.setdefault("config", self.config.get_text_config(decoder=True))
backend = cache_config.pop("backend", "quanto")
model_kwargs["past_key_values"] = QuantizedCache(backend=backend, **cache_config)
# i.e. `cache_implementation` in [None, "dynamic", "offloaded", "dynamic_full"]
# TODO: prepare linear cache from a single API, instead of creating in modeling code
Comment thread
zucchini-nlp marked this conversation as resolved.
Resolved
else:
model_kwargs["past_key_values"] = DynamicCache(**dynamic_cache_kwargs)

cache_config = generation_config.cache_config if generation_config.cache_config is not None else {}
# Add the config if it was not provided, as it's a required argument
if "config" not in cache_config:
cache_config["config"] = self.config.get_text_config()
# Pop the backend from the config (defaults to quanto if not defined)
backend = cache_config.pop("backend", "quanto")

if backend == "quanto" and not is_optimum_quanto_available():
raise ImportError(
"You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto "
"backend. Please install it via with `pip install optimum-quanto`"
)
elif backend == "HQQ" and not is_hqq_available():
raise ImportError(
"You need to install `HQQ` in order to use KV cache quantization with HQQ backend. "
"Please install it via with `pip install hqq`"
)
model_kwargs[cache_name] = QuantizedCache(backend=backend, **cache_config)
elif generation_config.cache_implementation == "offloaded":
model_kwargs[cache_name] = DynamicCache(**dynamic_cache_kwargs, offloading=True)
elif "dynamic" in generation_config.cache_implementation:
model_kwargs[cache_name] = DynamicCache(**dynamic_cache_kwargs)

# TODO (joao): this logic is incomplete, e.g. `offloaded` should apply to both caches. Refactor this function
# to correctly pass parameterization to both caches.
if (
requires_cross_attention_cache
self.config.is_encoder_decoder
and "past_key_values" in model_kwargs
and not isinstance(model_kwargs["past_key_values"], EncoderDecoderCache)
):
Expand Down
2 changes: 1 addition & 1 deletion src/transformers/models/afmoe/modeling_afmoe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -576,7 +576,7 @@ def forward(
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

if use_cache and past_key_values is None:
past_key_values = DynamicCache()
past_key_values = DynamicCache(config=self.config)

if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
Expand Down
2 changes: 1 addition & 1 deletion src/transformers/models/afmoe/modular_afmoe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -397,7 +397,7 @@ def forward(
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

if use_cache and past_key_values is None:
past_key_values = DynamicCache()
past_key_values = DynamicCache(config=self.config)

if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
Expand Down
19 changes: 13 additions & 6 deletions src/transformers/models/blt/modeling_blt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -930,7 +930,7 @@ def forward(
inputs_embeds = self.embed_tokens(input_ids)

if use_cache and past_key_values is None:
past_key_values = DynamicCache()
past_key_values = DynamicCache(config=self.config)

if cache_position is None:
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
Expand Down Expand Up @@ -1229,8 +1229,15 @@ def forward(
if (input_ids is None) ^ (inputs_embeds is not None):
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

if use_cache and past_key_values is None:
past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
if use_cache:
if past_key_values is None:
past_key_values = EncoderDecoderCache(
DynamicCache(config=self.config), DynamicCache(config=self.config)
)
elif not isinstance(past_key_values, EncoderDecoderCache):
# BLT uses an encoder-decoder cache even though it is not en encoder-decoder model. Create a cross-cache
# if not yet created by the user
past_key_values = EncoderDecoderCache(past_key_values, DynamicCache(config=self.config))
Comment thread
zucchini-nlp marked this conversation as resolved.
Resolved

# Extract input embeddings as early as possible
if inputs_embeds is not None:
Expand Down Expand Up @@ -1377,7 +1384,7 @@ class BltForCausalLM(BltPreTrainedModel, GenerationMixin):
_tied_weights_keys = {"model.local_encoder.embed_tokens.weight": "lm_head.weight"}

def __init__(self, config: BltConfig):
super().__init__(config.get_text_config())
super().__init__(config)
self.text_config = config.get_text_config()
Comment on lines -1380 to 1388

@zucchini-nlp zucchini-nlp Feb 2, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not a good practice if self.config is only the decoder-subconfig

self.vocab_size = config.vocab_size
self.model = BltModel(config)
Expand Down Expand Up @@ -1433,8 +1440,8 @@ def forward(
```python
>>> from transformers import AutoTokenizer, BltForCausalLM

>>> model = BltForCausalLM.from_pretrained("Llama-3.2-11B-Vision")
>>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision")
>>> model = BltForCausalLM.from_pretrained("itazap/blt-1b-hf")
>>> tokenizer = AutoTokenizer.from_pretrained("itazap/blt-1b-hf")

>>> prompt = "If I had to write a haiku, it would be:"
>>> inputs = tokenizer(prompt, return_tensors="pt")
Expand Down
Loading