DO NOT MERGE: generate compatible with torch.compile(fullgraph=True)

[gante]

What does this PR do?

This PR is a 🔪 mangled🔪 version of generate where torch.compile(model.generate, fullgraph=True) works and returns the same values. It should NOT be merged, but rather be used as a reference -- other PRs will be created that push the needed changes, once at a time, to ensure we don't break other features.

---

Script to test correctness

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
import torch
import copy

torch_device = "cuda"

EXPECTED_GENERATION = [
    "The best color is the one that complements the skin tone of the",
    "We should not undermind the issues at hand.\nWe should not undermind the issues",
]

tokenizer = AutoTokenizer.from_pretrained(
    "NousResearch/Llama-2-7b-chat-hf", padding_side="left", pad_token="<s>"
)
model = AutoModelForCausalLM.from_pretrained(
    "NousResearch/Llama-2-7b-chat-hf",
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to(torch_device)
inputs = tokenizer(
    ["The best color is", "We should not undermind the issues at hand"], padding=True, return_tensors="pt"
).to(model.device)

generation_kwargs = {
    "do_sample": False,
    "max_new_tokens": 10,
}

print("Dynamic cache")
gen_out = model.generate(**inputs, **generation_kwargs)
decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
print(decoded)
assert decoded == EXPECTED_GENERATION

print("Static cache")
model.generation_config.cache_implementation = "static"
gen_out = model.generate(**inputs, **generation_kwargs)
decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
assert decoded == EXPECTED_GENERATION
print(decoded)

print("Compiled static cache")
generation_config = copy.deepcopy(model.generation_config)
generation_config.update(**generation_kwargs)
compiled_generate = torch.compile(model.generate, fullgraph=True, mode="reduce-overhead")
gen_out = compiled_generate(**inputs, generation_config=generation_config)
decoded = tokenizer.batch_decode(gen_out, skip_special_tokens=True)
print(decoded)
assert decoded == EXPECTED_GENERATION

fixes #27837

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

[github-actions]

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 are likely to be ignored.

[amyeroberts]

@gante If this PR is going to be long lived - you can add the WIP label and it will stop the bot closing if stale

[gante]

Closed in favor of #30788