[Cache] Support compilable cache reuse with smaller batch sizes
Author: gante
Pull request: yes

# What does this PR do?

⚠️ this PR needs to be rebased, don't review/merge

Supercedes #37389 
Partially solves #35444

This PR makes our `max_cache_size` argument in compilable caches finally true: we can now use a cache object with a batch size smaller than the one defined in the cache. Compile once and run with multiple input shapes -- particularly useful for export, as mentioned in #35444.

Adds other minor related fixes (see PR comments).

__________________

We can see in the following test script that this does not degrade compiled performance:

```py
from transformers import AutoModelForCausalLM, AutoTokenizer, StaticCache
import torch
import time

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", device_map="auto", torch_dtype=torch.float16)

input_ids = tokenizer(["The quick brown"], return_tensors="pt").input_ids.to(model.device)
cache_position = torch.arange(input_ids.shape[1]).to(model.device)

with torch.no_grad():
    #------------------------------------------------------------------------------------------------
    # OLD, cache batch size = input batch size
    # Measured on an RTX 4090: `main` = 0.223ms; this PR = 0.223ms
    cache = StaticCache(
        config=model.config,
        max_batch_size=1,
        max_cache_len=100,
        device=model.device,
        dtype=model.dtype
    )
    model.forward = torch.compile(model.forward, fullgraph=True, mode="reduce-overhead")

    # warmup
    for _ in range(3):
        outputs = model(input_ids, cache_position=cache_position, past_key_values=cache)

    # measure
    start = time.time()
    for _ in range(100):
        outputs = model(input_ids, cache_position=cache_position, past_key_values=cache)
    end = time.time()
    print(f"[Old] Average time taken: {((end - start) / 100) * 1000} ms")

    #------------------------------------------------------------------------------------------------
    # clear torch compile cache
    torch._dynamo.reset()

    #------------------------------------------------------------------------------------------------
    # NEW, cache batch size > input batch size
    # Measured on an RTX 4090: `main` = Doesn't work; this PR = 0.224ms
    cache = StaticCache(
        config=model.config,
        max_batch_size=16,  # 16 >> 1
        max_cache_len=100,
        device=model.device,
        dtype=model.dtype
    )
    model.forward = torch.compile(model.forward, fullgraph=True, mode="reduce-overhead")

    # warmup
    for _ in range(3):
        outputs = model(input_ids, cache_position=cache_position, past_key_values=cache)

    # measure
    start = time.time()
    for _ in range(100):
        outputs = model(input_ids, cache_position=cache_position, past_key_values=cache)
    end = time.time()
    print(f"[New] Average time taken: {((end - start) / 100) * 1000} ms")

```



Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page). This will assign reviewers and trigger CI.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_37394). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

(PR on hold: some slow cache tests are failing due to reasons unrelated to this PR, fixing them first before re-requesting a review)

(caches have been refactored, better start from scratch)