feat: add crop() to StaticCache layers for assisted generation by ArkaD171717 · Pull Request #45745 · huggingface/transformers

ArkaD171717 · 2026-05-02T09:31:22Z

Assisted generation needs to roll back the KV cache when candidate tokens get
rejected
DynamicLayer.crop() does this by slicing, but StaticLayer can't slice
because its tensors have static memory addresses pinned by
mark_static_address() for torch.compile/cudagraphs

StaticLayer.crop() zeros out the evicted slots with fill_()/indexing and
updates cumulative_length in-place
Tensor identity is preserved so torch.compile doesn't break

StaticSlidingWindowLayer.crop() delegates to the above but raises ValueError
if the window has already been exceeded (same semantics as the dynamic version,
evicted tokens are gone)

Removes the hard ValueError in _assisted_decoding that was blocking
StaticCache entirely
Also propagates the main model's cache_implementation to the assistant model
in AssistedCandidateGenerator instead of always falling back to
"dynamic_full", per the reviewer request on #34797

11 unit tests for crop() (basic, negative, overflow clamp, tensor identity,
crop-then-update, sliding window)
8 integration tests for assisted generation with static cache (including
output-matching against dynamic cache)

Closes #32946

ArkaD171717 · 2026-05-02T18:57:21Z

I think @LagunaModelTest test_tp_generation failure is unrelated to this PR its a gloo transport layer mismatch in the tensor parallel test harness. All the other tests passed. Could a maintainer re-run the failed tests_tensor_parallel_ci job?

Qodo-Free-For-OSS · 2026-05-03T08:56:06Z

Hi, Assisted decoding always calls outputs.past_key_values.crop(...), but for encoder-decoder models past_key_values is an EncoderDecoderCache whose crop() rejects non-dynamic caches, so static cache assisted generation will raise TypeError at runtime.

Severity: action required | Category: correctness

How to fix: Support static in EncoderDecoderCache.crop

Agent prompt to fix - you can give this to your LLM of choice:

Issue description

Assisted decoding unconditionally calls outputs.past_key_values.crop(...). For encoder-decoder models, past_key_values is an EncoderDecoderCache, whose crop() currently enforces dynamic-only caches via check_dynamic_cache(). After this PR, assisted generation may select static caches, causing a runtime error.

Issue Context

_prepare_cache_for_generation now permits static cache implementations for assisted generation.

_assisted_decoding still calls .crop() on outputs.past_key_values.

EncoderDecoderCache.crop() must either support static self_attention_cache or assisted generation must continue to force dynamic caches for encoder-decoder models.

Fix Focus Areas

src/transformers/cache_utils.py[1620-1627]

src/transformers/generation/utils.py[1848-1859]

src/transformers/generation/utils.py[3614-3616]

Implementation guidance

Choose one:

Update EncoderDecoderCache.crop() to allow static self_attention_cache (e.g., call self.self_attention_cache.crop(...) regardless of cache class, and only enforce dynamic for methods that truly require it).

Gate static-cache allowance in _prepare_cache_for_generation for encoder-decoder assisted generation (keep forcing dynamic_full there).
Also add an encoder-decoder assisted-generation test with cache_implementation="static".

We noticed a couple of other issues in this PR as well - happy to share if helpful.

Found by Qodo code review

ArkaD171717 · 2026-05-03T11:02:51Z

Thank you for catching that, @Qodo-Free-For-OSS. It's now fixed and pushed. Please do share those other issues! I don't need agent prompts.

StaticLayer.crop() zeros out evicted KV slots in-place and updates cumulative_length via fill_(), preserving static tensor addresses for torch.compile/cudagraphs. StaticSlidingWindowLayer.crop() delegates to the above with a ValueError guard matching dynamic semantics. Remove the hard ValueError in _assisted_decoding that blocked StaticCache. Propagate the main model's cache_implementation to the assistant model in AssistedCandidateGenerator instead of always falling back to dynamic_full. 11 unit tests for crop(), 8 integration tests for assisted generation with static cache. Closes huggingface#32946

EncoderDecoderCache.crop() does not support static sub-caches, so encoder-decoder models must keep using dynamic_full for assisted generation. Only decoder-only models benefit from static cache here.

ArkaD171717 · 2026-05-07T15:15:26Z

Fixed Merge conflict

Rocketknight1 · 2026-05-08T11:18:47Z

cc @Cyrilvallez for generation

Cyrilvallez

Hey! In general I'm not against, but could you please share your motivations for this? Since we need to rollback all the time, we basically lose the advantages of a StaticCache, so I'm not super sure it makes sense.
Also, please do not add new test files, add tests to existing test files!

ArkaD171717 · 2026-05-13T17:21:01Z

Hello @Cyrilvallez thanks for the reply,

The motivation is torch.compile / CUDA graphs; DynamicCache allocates new tensors on every grow/slice which forces graph recompilation and StaticCache pins tensor addresses via mark_static_address() so the compiled graph stays valid across calls. Assisted gen needs crop() to roll back rejected candidates but because of all the above slicing would create new tensors and break the compiled graph
Thats why StaticLayer.crop() zeroes out evicted slots in place (fill_/indexing) and updates cumulative_length with fill_(). B/c id(tensor) is preserved torch.compile doesnt see new obj and doesnt retrace
So advantage of StaticCache isnt lost but the compiled forward pass stays intact (the rollback itself is O(evicted_tokens) zeroing). The rollback cost is just a memset on evicted region and not a recompilation

Sorry about the test files issue, I moved all of them into existing files and will push that shortly

ArkaD171717 · 2026-05-15T19:57:27Z

accidentally closed wrong PR 😭😭Just to be clear this one is still open

Cyrilvallez

Hey! Overall, I don't think this is worth the effort as it is in the curent state... We could ad the method, but it would lead us nowhere... The big issues are the following:

assisted decoding does not trigger compilation under static cache anyway
crop does not work with sliding window, so all static caches cannot be used

Cyrilvallez · 2026-06-22T14:46:21Z

+        # We need to roll back the cache in assisted generation. Static caches now support
+        # rollback via StaticLayer.crop(), so propagate the main model's cache_implementation
+        # when it is a static type; otherwise fall back to "dynamic_full".
+        if generation_config.cache_implementation in ALL_STATIC_CACHE_IMPLEMENTATIONS:
+            self.generation_config.cache_implementation = generation_config.cache_implementation


This will not work with sliding window layers

ArkaD171717 added 3 commits May 7, 2026 10:15

style: fix ruff formatting in test_cache_utils.py

672951d

fix: gate static cache for encoder-decoder assisted generation

2b3e8b2

EncoderDecoderCache.crop() does not support static sub-caches, so encoder-decoder models must keep using dynamic_full for assisted generation. Only decoder-only models benefit from static cache here.

ArkaD171717 force-pushed the static-cache-assisted-gen branch from 96a1f3f to 2b3e8b2 Compare May 7, 2026 15:15

Cyrilvallez reviewed May 12, 2026

View reviewed changes

Merge branch 'main' into static-cache-assisted-gen

c41a4ba

ArkaD171717 closed this May 15, 2026

ArkaD171717 reopened this May 15, 2026

Merge branch 'main' into static-cache-assisted-gen

a7a55d0

ArkaD171717 requested a review from Cyrilvallez June 8, 2026 01:45

Cyrilvallez reviewed Jun 22, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: add crop() to StaticCache layers for assisted generation#45745

feat: add crop() to StaticCache layers for assisted generation#45745
ArkaD171717 wants to merge 5 commits into
huggingface:mainfrom
ArkaD171717:static-cache-assisted-gen

ArkaD171717 commented May 2, 2026

Uh oh!

ArkaD171717 commented May 2, 2026

Uh oh!

Qodo-Free-For-OSS commented May 3, 2026

Issue description

Issue Context

Fix Focus Areas

Implementation guidance

Uh oh!

ArkaD171717 commented May 3, 2026

Uh oh!

ArkaD171717 commented May 7, 2026

Uh oh!

Rocketknight1 commented May 8, 2026

Uh oh!

Cyrilvallez left a comment

Uh oh!

ArkaD171717 commented May 13, 2026

Uh oh!

ArkaD171717 commented May 15, 2026

Uh oh!

Cyrilvallez left a comment

Uh oh!

Cyrilvallez Jun 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

ArkaD171717 commented May 2, 2026

Uh oh!

ArkaD171717 commented May 2, 2026

Uh oh!

Qodo-Free-For-OSS commented May 3, 2026

Issue description

Issue Context

Fix Focus Areas

Implementation guidance

Uh oh!

ArkaD171717 commented May 3, 2026

Uh oh!

ArkaD171717 commented May 7, 2026

Uh oh!

Rocketknight1 commented May 8, 2026

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

ArkaD171717 commented May 13, 2026

Uh oh!

ArkaD171717 commented May 15, 2026

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Cyrilvallez Jun 22, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants