[Cache] Support compilable cache reuse with smaller batch sizes

[gante]

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:

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")

[github-actions]

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.

[gante]

Note: adding the feature to MambaCache would imply adding model-level changes. And it wouldn't work with the fast kernels.

[gante]

The rest of the diff in this file is to make mamba + compile work again (MambaIntegrationTests::test_compile_mamba_cache was red)

In general, models with unique caches are messy to use with generate, and need some work. A model should be able to tell generate "hey, I can only use this cache class"

[gante]

Not all tests here are @slow. Tests that take >1s kept the decorator

[gante]

generation is not meant to work well with right-padding, no need to spend resources testing it

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[zucchini-nlp]

Great, thanks for adding support on this and happy to see it can be done with minimal changes!

[zucchini-nlp]

we didn't change the signature, why is this needed?

[gante]

MambaCache + mamba was broken ☠️ This is needed to fix it

[gante]

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

[gante]

(caches have been refactored, better start from scratch)