[generate] Stop setting the static cache as an attribute to save memory

[Cyrilvallez]

What does this PR do?

As per the title. Keeping the cache as an attribute self._cache inside generate is a very bad idea in general, because it will keep a big memory footprint without the user noticing, which can be very surprising and unwanted if we want to make 1 call to generate with static cache, and then use the model for something else.
It was previously done to avoid recompilation if we were making several calls to generate with several lengths, but simply keeping the old max_cache_len as attribute and re-using that later is enough to avoid recompilation.

The following snippet illustrates the issue:

import torch
from transformers import AutoModelForCausalLM

model_id = "meta-llama/Llama-3.2-1B-instruct"
device = 0
model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device)
print(f"Memory reserved before call to generate: {torch.cuda.memory_allocated(device) / 1024**3:.2f} GiB")

foo = model.generate(torch.randint(100, 200, (10, 8192), device=device), max_new_tokens=100, cache_implementation="static")
print(f"Memory reserved after call to generate: {torch.cuda.memory_allocated(device) / 1024**3:.2f} GiB")

Before this PR, it prints:

Memory reserved before call to generate: 2.30 GiB
Memory reserved after call to generate: 4.83 GiB

After this PR:

Memory reserved before call to generate: 2.30 GiB
Memory reserved after call to generate: 2.30 GiB

which makes a HUGE difference as can be easily seen.

Note that this does not incur any additional recompilation, since we still check the max size of the cache. This can be checked with the following snippet for example:

import os
os.environ["TORCH_LOGS"] = "recompiles"
import torch
from transformers import AutoModelForCausalLM

model_id = "meta-llama/Llama-3.2-1B-instruct"
device = 0
model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device)

foo = model.generate(torch.randint(100, 200, (10, 8192), device=device), max_new_tokens=100, cache_implementation="static")
# smaller size should not trigger any recompilation
foo = model.generate(torch.randint(100, 200, (10, 8000), device=device), max_new_tokens=100, cache_implementation="static")

[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.

[vasqu]

One question im a bit hesitant about because I dont follow the implications here 100% but other than that it's a nice optim 👀

[vasqu]

Do we have any test that checks recompilations in regards to that?

[Cyrilvallez]

We don't, but I confirmed it locally at least - will try to add one, but can be tricky to catch recompiles automatically

[vasqu]

It does mean that we recreate a cache in any case, no? Would there be any issue when we use multiple iterations with a chunked prefill? Just a thought, unsure

[Cyrilvallez]

Not sure I understand what you mean? Basically, now we create a new StaticCache every generate call vs keeping one attached to the model. To be honest, I have no idea why it was ever attached to the model... compile caches the graphs based on shapes and ops, not objects addresses, so we can just recreate the object

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 30888111653:1
Result: success | Jobs: 16 | Tests: 172,237 | Failures: 0 | Duration: 18h 20m

[vasqu]

Thanks, I meant that in some cases before we did not recreate the cache itself (nothing to do with the attached cache) but tbh it's weird to have the attached version so rather get that out