# Cache lazy initialization leads to spurious recompilation with chunked prefill
Author: dacorvo
Pull request: no
Labels: name: bug

### System Info

When #39797 simplified the cache management, it introduced the cache lazy optimization, assuming that there would always
be an eager forward call to trigger the initialization (typically the prefill).

This is however not true with chunked prefill, and it triggers a recompile on a second call to generate.

Cost on Llama-3.2-1B, A10G, bf16, 1024-token prompt, `prefill_chunk_size=1024`: the second `generate()` call takes **14.9 s and adds 18 Inductor artifacts** vs ~1.4 s / +0 without the issue.

Fortunately, there is an escape hatch: early_initialization can be called explicitly when it is required to compile all forward calls (like in the chunked prefill configuration).

## Fix

Have `generate()` call `cache.early_initialization(...)` from `_prepare_static_cache` whenever it is about to compile the prefill — i.e. when `generation_config.prefill_chunk_size is not None`. That is the only path inside `generate()` where the prefill ends up in a compiled region today, and the cache's batch / num_heads / head_dim / dtype / device are all derivable from the model config at that point.

Verified locally on the repro above: the second `generate()` collapses from 14.9 s / +18 artifacts to 1.4 s / +0.

The docstring at `cache_utils.py:302` should be updated at the same time — it currently tells users `early_initialization` is unnecessary for `compile` because "we internally don't compile the prefill", which is false once `prefill_chunk_size` is set.

### Who can help?

@ArthurZucker @Cyrilvallez 

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

https://github.com/dacorvo/transformers-generation-compile/blob/main/README.md

### Expected behavior

When using chunked prefill and a static cache, a second call to generate on a warm cache should not tirgger a recompilation.

Comments:
Hi! I'd like to work on this if it's still open.

**My understanding of the issue:**

The static cache uses lazy initialization - it sets up its tensors on the first `update()` call and registers them with `mark_static_address()` so the compiler knows not to recompile on the next `generate()` call. This registration is guarded by `if not is_torchdynamo_compiling()`, which is fine normally since the prefill runs eagerly.

But with `prefill_chunk_size` set, the prefill runs inside a compiled region - so `lazy_initialization()` fires during compilation, `mark_static_address()` is skipped, and the second `generate()` triggers a full recompile.

**Two possible fixes:**

**Option A (proposed):** Call `cache.early_initialization(...)` from `_prepare_static_cache` when `prefill_chunk_size is not None`. Simpler, minimal change.

**Option B (deeper):** Deferred registration inside `lazy_initialization()` itself, if called during compilation, set a `_pending_registration` flag and call `mark_static_address()` on the next eager `update()`. No special-casing in `generate()`, TP-safe, and handles any future path that compiles the prefill, not just chunked prefill.

The docstring at `cache_utils.py:302` needs updating either way.



> Hi! I'd like to work on this if it's still open.
> 
> **My understanding of the issue:**
> 
> The static cache uses lazy initialization - it sets up its tensors on the first `update()` call and registers them with `mark_static_address()` so the compiler knows not to recompile on the next `generate()` call. This registration is guarded by `if not is_torchdynamo_compiling()`, which is fine normally since the prefill runs eagerly.
> 
> But with `prefill_chunk_size` set, the prefill runs inside a compiled region - so `lazy_initialization()` fires during compilation, `mark_static_address()` is skipped, and the second `generate()` triggers a full recompile.
> 
> **Two possible fixes:**
> 
> **Option A (proposed):** Call `cache.early_initialization(...)` from `_prepare_static_cache` when `prefill_chunk_size is not None`. Simpler, minimal change.
> 
> **Option B (deeper):** Deferred registration inside `lazy_initialization()` itself, if called during compilation, set a `_pending_registration` flag and call `mark_static_address()` on the next eager `update()`. No special-casing in `generate()`, TP-safe, and handles any future path that compiles the prefill, not just chunked prefill.
> 
> The docstring at `cache_utils.py:302` needs updating either way.

Thanks for your interest, but I already have a fix (see the linked repo). Note that what you suggest in Option A is precisely what is described in the Fix paragraph of this pull-request.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Resolved in #46446