[rfc] Prototype to make torch.compile work with DynamicCache
Author: anijain2305
Pull request: yes
Labels: name: Core: Modeling

name: Compilation

# What does this PR do?

## Background on StaticCache vs DynamicCache
A static KV cache must be preallocated for the model’s maximum context window (e.g., 32k tokens), even if you only generate a few tokens. That leads to large memory pressure and wastage. A dynamic cache grows with the sequence actually in use: with P​ prefill tokens and D decode steps, memory scales with P + D instead of max_seq_length. For example, if you prefill 2,000 tokens and generate 50, the dynamic cache holds ~2,050 tokens’ worth of K/V per layer, versus 32k for a static allocation, significantly reducing memory pressure.

Lets do a quick experiment for Qwen3

<img width="724" height="120" alt="image" src="https://github.com/user-attachments/assets/1d37a49e-98a7-4447-8d20-ed4e43fb11d0" />


In all cases, the output generated text was exactly the same.

As expected, StaticCache has a higher memory footprint. But the generation time is also higher. This is because the main tensors involved in [SDPA](https://github.com/huggingface/transformers/blob/242bb2cafccec9f90479f5f688bca9d240b1031f/src/transformers/integrations/sdpa_attention.py#L83) operation - (key, value, attention_mask) have max seq length shape. Even though attention_mask fires up only a portion of the actual compute, there is still more work to be done here compared to Dynamic KV cache. (Side chat - Driss and Boyuan told me that we might be able to use better eager code or flex decoding to make this better, but thats not the focus of this doc.)



## torch.compile and Dynamic Cache
transformers introduced StaticCache just for torch.compile. For Qwen3, torch.compile latency is 20 seconds, achieving 1.2x speedup and 25 seconds of compile time ([tlparse](https://manifold.edge.x2p.facebook.net/v0/read/tree/logs/.tmpifaWBm/index.html?bucketName=tlparse_reports&apiKey=tlparse_reports-key&withPayload=1&timeoutMsec=10000)). 

But as pointed out earlier, DynamicCache is a better default setting for eager PyTorch. So, we will look at the changes required to make torch.compile work with DynamicCache. WIth this PR,  I see torch.compile improve **generation latency to 3.6 (while eager was 8.88 seconds), leading to 2.4x speedup** and a cold start latency of 36 seconds ([tlparse](https://manifold.edge.x2p.facebook.net/v0/read/tree/logs/.tmpq4CUzG/index.html?bucketName=tlparse_reports&apiKey=tlparse_reports-key&withPayload=1&timeoutMsec=10000)) and warm start of 11 seconds.


For this RFC
* We will use fullgraph=True and do full model compilation (no regional compilation yet).
* We will compile only the decode step to keep the discussion simple.

There are two main changes
**Avoid recompilations (symbolic seq length)** : Tell torch.compile to treat sequence length as dynamic for all KV-cache–related tensors - key, value, and attention_mask, using torch._dynamo.mark_dynamic. To ensure a single shared symbolic size is used everywhere (since mark_dynamic alone can yield distinct symbols), wrap the model and insert explicit size-equality checks (e.g., via torch._check) that tie the symbols together. This nudges the compiler to unify on one seq-len symbol, improving inductor codegen and in future compile time.

**CUDA graphs**: Disable cudagraphs for DynamicCache. Because the sequence length changes step-to-step, cudagraphs would constantly re-record, which undercuts their benefit and can hurt latency. (They remain useful when shapes are fixed, e.g., with a static cache.)


### Update 

The PR now extends torch.compile to `DynamicSlidingWindowLayer`. There is one extra recompilation on the integer value of `cumulative_length` that I hope can be prevented by converting the integer into a scalar tensor.

Speedup - Pytorch eager has 6.66 seconds, compile takes 5.10 seconds for Gemma2. Although the compile time is high (currently due to 2 compilations).

Comments:
cc @ArthurZucker @gante 

@anijain2305 we are happy to help you / push this to merge it, tell us if you need anything! 

Yes sorry everyone, I told @anijain2305 I would take over and wanted to do it this week, but did not find the time 🥲 Will do next week!

@anijain2305 we are happy to push this further!