Fix several bugs in `cache_implementation=static` by dacorvo · Pull Request #46446 · huggingface/transformers

dacorvo · 2026-06-05T12:38:38Z

What does this PR do?

Static-cache fixes for generate(), one commit each:

Fix chunked-prefill recompilation: the lazy-init branch bakes the cache init into the first compiled graph, forcing a second graph; init eagerly instead (# Cache lazy initialization leads to spurious recompilation with chunked prefill #46421).
Honor cache_config["max_cache_len"] on the static path so the cache is sized for the worst case across calls (Cache config max_cache_length only honored for Quantized Cache #46424).
Make Cache.early_initialization skip linear-attention layers (it corrupted their conv/recurrent shapes → RuntimeError) (Static cache early_initialization corrupts linear-attention layers (wrong shapes → RuntimeError) #46439).
Let create_masks_for_generate defer to the model for layer types with no mask function, instead of KeyError: 'linear_attention' (Qwen3.5 crash with cache_implementation="static" — KeyError: 'linear_attention' #46441).
Make Cache.max_cache_len skip linear-attention layers (the static-cache reuse check AttributeError'd on hybrids).

Together these make static-cache + chunked-prefill generation work without recompiles for plain (Llama) and hybrid linear-attention (Qwen3.5) models. Each commit has a regression test.

I confirm that this is not a pure code agent PR.

Before submitting

Did you read the contributor guideline,
Pull Request section?
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@ArthurZucker @Cyrilvallez @vasqu

Copilot

Pull request overview

This PR fixes several correctness and stability issues when using cache_implementation="static" during generation, especially for chunked prefill and hybrid models that include linear-attention layers.

Changes:

Make static-cache sizing respect generation_config.cache_config["max_cache_len"] to avoid reallocations/recompiles across calls.
Fix static-cache eager initialization and cache metadata (is_initialized, max_cache_len) to safely handle hybrid linear-attention layers.
Make mask pre-building (create_masks_for_generate) defer to the model when a layer type has no registered mask builder, avoiding KeyError for hybrids.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
`src/transformers/generation/utils.py`	Adjusts static-cache sizing and adds eager cache initialization for chunked prefill.
`src/transformers/cache_utils.py`	Makes `early_initialization`, `is_initialized`, and `max_cache_len` robust to linear-attention layers; updates docs.
`src/transformers/masking_utils.py`	Defers mask pre-building when encountering unmapped `layer_types`.
`tests/utils/test_cache_utils.py`	Adds regression tests for eager-init behavior, hybrid linear-attention handling, and `max_cache_len` honoring.
`tests/utils/test_masking_utils.py`	Adds regression test ensuring unmapped layer types defer mask creation instead of raising.

HuggingFaceDocBuilderDev · 2026-06-05T12:53:24Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

With `prefill_chunk_size` set, `generate()` compiles the prefill (via `get_compiled_call()` per chunk). It hits a condition in cache.update that leads to a graph specialization between the first call to generate and the second (one has the cache init, the second doesn't), hence a recompilation. The fix is to force a call to early_initialization outside of the compile region when chunked prefill. Also corrects the `lazy_initialization` docstring, which claimed `early_initialization` was unnecessary for `compile` because the prefill is never compiled internally — not true once `prefill_chunk_size` is set. Fixes #46421 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

`GenerationConfig.cache_config` is documented as carrying arguments for the key-value cache, but `_prepare_cache_for_generation` read it only on the quantized branch. The static branch hardcoded `max_cache_len = max_length - 1`, so `cache_config={"max_cache_len": N}` with `cache_implementation="static"` was silently ignored. Fixes #46424 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

`early_initialization` initialized every layer via the `(num_heads, head_dim)` key/value layout, but linear attention layers (mamba/conv/linear_attention) don't use that layout — they track their own statically-shaped conv/recurrent states, initialized lazily on update. Initializing them through the key/value path is wrong and crashes on the first call to prefill. Fixes #46439 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

`generate()` pre-builds attention masks for compilable caches via helpers registered in `LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING`. No helpers exist for hybrid linear-attention models (Qwen3.5, Qwen3-Next, ...), since they build their own masks in the forward, so this raised a cryptic `KeyError: 'linear_attention'` and made `cache_implementation="static"` unusable with these models. This skips the early creation of the mask for these models. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The `max_cache_len` property iterated every layer, but linear attention layers have no `max_cache_len`, so the auto-allocated static-cache reuse check in `_prepare_static_cache` raised `AttributeError` on a hybrid model. Skip them. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…ti-device device_map

ArthurZucker

Ty!

ArthurZucker · 2026-06-09T10:27:25Z

    def is_initialized(self) -> bool:
        """Return whether the cache data is initialized"""
-        return len(self.layers) > 0 and all(layer.is_initialized for layer in self.layers)
+        # Layers without an `is_initialized` flag (e.g. linear attention layers, which track their own conv/recurrent


i think it might be better to have them set that flag in general no?

That's one option, yes, but this means we need to set the flag to True in LinearAttentionCacheLayerMixin, which I find a bit disturbing, because it's not really initialized.
Another option would be to have a class constant like supports_early_init.
The two checks become more readable:

if not supports_early_init or is_initialized: continue

And

return len(self.layers) > 0 and all(layer.is_initialized for layer in self.layers if layers.supports_early_init)

Implemented what i had suggested.

…ache

…nstead of getattr

github-actions · 2026-06-09T14:30:11Z

CI Dashboard: View test results in Grafana

…ked-prefill-recompile

Cyrilvallez

Nice, thanks for the PR! It indeed touches a few parts that we let a bit on the side until now (chunked prefill, linear layer early init).
IMO, the easiest and most foolproof way to proceed though would be to avoid compiling the first chunk of the prefill altogether, so that the cache will properly be early initialized as usual. This would allow early init even on multi gpu setup, and avoid checking dims in case of TP etc, which can be a hassle. I don't think it would harm performances in general. WDYT?

Cyrilvallez · 2026-06-15T02:55:11Z

+            # A user-provided `cache_config["max_cache_len"]` lets the static cache be sized for the worst case
+            # across calls, so that later calls with a longer prompt or a larger `max_new_tokens` (up to that
+            # ceiling) reuse the same cache instead of triggering a reallocation (and a `torch.compile`
+            # recompilation). Without this, the cache is sized to the current call's `max_length` only. See #46424.
+            cache_config = generation_config.cache_config or {}
+            if cache_config.get("max_cache_len") is not None:
+                max_cache_length = max(max_cache_length, cache_config["max_cache_len"])
+            cache_batch_size = max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size


Indeed, nice to have a mechanism to preset the length manually. I think it was not done before because cache_config is not really documented/well known, and I believe it was previously introduced mostly for quantized cache.
Maybe we could get this argument out of cache_config, wdyt? I believe it would probably make a bit more sense to have a properly documented max_cache_len directly? That way we can properly explain that it's used only with static cache, and that it avoids recompilation if several forward will be run with different lengths?

OK, reverted the change and added a max_cache_len GenerationConfig parameter directly.

Cyrilvallez · 2026-06-15T03:06:21Z

+            elif prefill_chunk_size is not None:
+                # Chunked prefill compiles the prefill, so eagerly init the fresh cache to avoid a recompile next call
+                # (#46421). Skipped (-> lazy init) when it can't be initialized on a single device.
+                init_shape = self._get_static_cache_init_shape()
+                if init_shape is not None:
+                    num_heads, head_dim = init_shape
+                    self._cache.early_initialization(
+                        batch_size=batch_size,
+                        num_heads=num_heads,
+                        head_dim=head_dim,
+                        dtype=self.dtype,
+                        device=self.device,
+                    )


Here, looks like we actually miss most of the cases, as it will only work for 1 gpu setup. The better way to proceed I believe would be to simply force the 1st chunk to NOT use compile, so that the cache is properly initialized before compile. WDYT? I don't think it would harm performance at all in general to have a single forward not compiled

Let's take a step back. The whole point of using chunked prefill and static cache in combination is to have fully static graphs for prefill and decode, regardless of the prompt length and max new tokens. This might look like a mild optimization for CUDA devices (would need proper benchmarks to confirm), but it is essential on devices that do not support dynamic shapes (Neuron and TPU), for which eager mode is just "fragmented compilation".
Currently, without this fix, the prefill is compiled twice on any configuration.
The code I submitted fixes it for the configurations where it matters most. For the others (large model not using TP but sharded on multiple devices, incompatible head counts), it keeps the existing behaviour.
Your suggestion would fix those edge cases, but by never compiling the first prefill chunk on any configuration, which is very harmful for Neuron and TPU.

This reverts commit 90b0b61.

Expose a documented top-level `GenerationConfig.max_cache_len` to pre-size the static cache for the worst case across calls, so later calls with a longer prompt or larger `max_new_tokens` (up to that ceiling) reuse the same cache instead of reallocating and recompiling. Replaces the previously silently-ignored `cache_config["max_cache_len"]` path with an explicit, validated flag. Fixes #46424 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…ked-prefill-recompile

github-actions · 2026-06-16T07:37:10Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46446&sha=6f39ad

ArthurZucker

lgtm, not sure it will generalize well (_get_static_cache_init_shape)

ArthurZucker · 2026-06-18T12:25:52Z

+        max_cache_len (`int`, *optional*):
+            Only used with static caches (`cache_implementation` set to `"static"` or `"offloaded_static"`).
+            Pre-sizes the cache to this length instead of the current call's `max_length`. Set it once to the
+            largest call you expect so that repeated `generate()` calls with a longer prompt or a larger
+            `max_new_tokens` (up to this ceiling) reuse the same cache instead of triggering a reallocation and a
+            `torch.compile` recompilation.


i am surprise this is not in generation config

ArthurZucker · 2026-06-18T12:29:47Z

+        num_key_value_heads = getattr(text_config, "num_key_value_heads", None) or text_config.num_attention_heads
+        if num_key_value_heads % tp_size != 0:
+            return None
+        head_dim = getattr(text_config, "head_dim", None) or text_config.hidden_size // text_config.num_attention_heads
+        return num_key_value_heads // tp_size, head_dim


this does not necessarily scale well with DSv4 for example and qk_head_dim

Added yet another exclusion

* Fix static-cache recompilation with chunked prefill With `prefill_chunk_size` set, `generate()` compiles the prefill (via `get_compiled_call()` per chunk). It hits a condition in cache.update that leads to a graph specialization between the first call to generate and the second (one has the cache init, the second doesn't), hence a recompilation. The fix is to force a call to early_initialization outside of the compile region when chunked prefill. Also corrects the `lazy_initialization` docstring, which claimed `early_initialization` was unnecessary for `compile` because the prefill is never compiled internally — not true once `prefill_chunk_size` is set. Fixes huggingface#46421 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Honor cache_config max_cache_len on the static cache path `GenerationConfig.cache_config` is documented as carrying arguments for the key-value cache, but `_prepare_cache_for_generation` read it only on the quantized branch. The static branch hardcoded `max_cache_len = max_length - 1`, so `cache_config={"max_cache_len": N}` with `cache_implementation="static"` was silently ignored. Fixes huggingface#46424 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Make StaticCache.early_initialization robust to linear attention layers `early_initialization` initialized every layer via the `(num_heads, head_dim)` key/value layout, but linear attention layers (mamba/conv/linear_attention) don't use that layout — they track their own statically-shaped conv/recurrent states, initialized lazily on update. Initializing them through the key/value path is wrong and crashes on the first call to prefill. Fixes huggingface#46439 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Let create_masks_for_generate defer for models with custom layer masks `generate()` pre-builds attention masks for compilable caches via helpers registered in `LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING`. No helpers exist for hybrid linear-attention models (Qwen3.5, Qwen3-Next, ...), since they build their own masks in the forward, so this raised a cryptic `KeyError: 'linear_attention'` and made `cache_implementation="static"` unusable with these models. This skips the early creation of the mask for these models. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Make StaticCache.max_cache_len robust to linear attention layers The `max_cache_len` property iterated every layer, but linear attention layers have no `max_cache_len`, so the auto-allocated static-cache reuse check in `_prepare_static_cache` raised `AttributeError` on a hybrid model. Skip them. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * review(generate): skip/shard eager static-cache init under TP and multi-device device_map * fix: typo * review(generate): move eager static-cache init into _prepare_static_cache * review(generate): isolate eager static-cache init args in a helper * review(cache): gate early-init via a supports_early_init layer flag instead of getattr * Revert "Honor cache_config max_cache_len on the static cache path" This reverts commit 90b0b61. * feat(generate): add max_cache_len generation flag for static caches Expose a documented top-level `GenerationConfig.max_cache_len` to pre-size the static cache for the worst case across calls, so later calls with a longer prompt or larger `max_new_tokens` (up to that ceiling) reuse the same cache instead of reallocating and recompiling. Replaces the previously silently-ignored `cache_config["max_cache_len"]` path with an explicit, validated flag. Fixes huggingface#46424 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * review: MLA models cannot be lazy initialized --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Copilot AI review requested due to automatic review settings June 5, 2026 12:38

Copilot started reviewing on behalf of dacorvo June 5, 2026 12:38

View session

Copilot AI reviewed Jun 5, 2026

View reviewed changes

Comment thread src/transformers/generation/utils.py Resolved

Comment thread src/transformers/cache_utils.py Resolved

dacorvo and others added 7 commits June 8, 2026 14:49

review(generate): skip/shard eager static-cache init under TP and mul…

40ddb52

…ti-device device_map

fix: typo

aa28be8

dacorvo force-pushed the fix-static-cache-chunked-prefill-recompile branch from 915ade8 to aa28be8 Compare June 8, 2026 12:50

ArthurZucker reviewed Jun 9, 2026

View reviewed changes

dacorvo added 4 commits June 9, 2026 13:33

review(generate): move eager static-cache init into _prepare_static_c…

f1593a9

…ache

review(generate): isolate eager static-cache init args in a helper

835e90f

Merge branch 'main' into fix-static-cache-chunked-prefill-recompile

58842f7

review(cache): gate early-init via a supports_early_init layer flag i…

6ec3342

…nstead of getattr

Merge remote-tracking branch 'origin/main' into fix-static-cache-chun…

78035b8

…ked-prefill-recompile

Cyrilvallez reviewed Jun 15, 2026

View reviewed changes

dacorvo and others added 2 commits June 15, 2026 11:25

Revert "Honor cache_config max_cache_len on the static cache path"

bf237fd

This reverts commit 90b0b61.

dacorvo requested review from ArthurZucker and Cyrilvallez June 15, 2026 12:27

dacorvo added 3 commits June 15, 2026 15:30

Merge remote-tracking branch 'origin/main' into fix-static-cache-chun…

8bd037d

…ked-prefill-recompile

Merge remote-tracking branch 'origin/main' into fix-static-cache-chun…

4c41a82

…ked-prefill-recompile

Merge remote-tracking branch 'origin/main' into fix-static-cache-chun…

6f39ad1

…ked-prefill-recompile

Rocketknight1 mentioned this pull request Jun 17, 2026

Fix: honor cache_config["max_cache_len"] on the static-cache path #46700

Closed

ArthurZucker approved these changes Jun 18, 2026

View reviewed changes

review: MLA models cannot be lazy initialized

736803a

dacorvo requested a review from ArthurZucker June 19, 2026 13:21

dacorvo added this pull request to the merge queue Jun 19, 2026

Merged via the queue into main with commit e7835fb Jun 19, 2026
123 checks passed

dacorvo deleted the fix-static-cache-chunked-prefill-recompile branch June 19, 2026 15:12

davidnichols-ops mentioned this pull request Jul 12, 2026

fix: restore dict return type in create_masks_for_generate (#46962) #47267

Closed

3 of 4 tasks

ydshieh mentioned this pull request Aug 22, 2026

[generate] Fix CUDAGraphs tensor overwrite for encoder-decoder static cache #48199

Closed

2 tasks done

ydshieh added a commit that referenced this pull request Aug 22, 2026

Fix issue reference: #46421 is the issue, #46446 is the PR that fixed it

a611fca

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

dacorvo mentioned this pull request Sep 3, 2026

Exception raised for Gemma4 with static cache and chunked prefill #48501

Closed

This was referenced Oct 2, 2026

generate() with a static cache: host-side inputs miss tensors nested in dicts, and TP undersizes the early cache #49261

Closed

Fix static-cache generate() with host-side inputs and under TP #49317

Merged

Conversation

dacorvo commented Jun 5, 2026

What does this PR do?

Before submitting

Who can review?

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Jun 5, 2026

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dacorvo Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Jun 9, 2026

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jun 16, 2026

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

dacorvo Jun 9, 2026 •

edited

Loading