feat: add crop() to StaticCache layers for assisted generation
Author: ArkaD171717
Pull request: yes

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

Comments:
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?

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:
> 1) 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).
> 2) 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](https://qodo.ai) code review*

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. 

Fixed Merge conflict

cc @cyrilvallez for generation

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

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