Make mamba use cache
Author: zucchini-nlp
Pull request: yes

# What does this PR do?

Fixes #30849. As it was reported in the issue, mamba was not using "use_cache" from generation kwargs. This PR adds it into kwargs and also returns the "cache_params", so that it can actually be used if `use_cache=True`

I tried to run Mamba test for generation locally, but it's failing for me even on `main` prob due to hardware differences (it has assertion on logits level)

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

Hope you don't mind the nits :D Just a mamba enthusiast over here.

Also tested the changes locally with commit a3cdff417bb3db94d5402dc009a07df97954a863 and it worked fine for `pytest tests/models/mamba -k "test_simple_generate" -s -vvvvv` (the logit levels test only). Haven't tried the others.

(GPU: 3080ti-16GB, CPU: 12th Gen Intel(R) Core(TM) i7-12700H)

Forgot to ask about benchmarking the generation speed btw, does it improve it (because now it's actually usingthe cache) ? 

@ArthurZucker also forgot, will do a small evaluation to verify tomorrow :)

Out of curiosity, don't we have the same issue with Jamba?

@ylacombe hey! I just checked jamba and it seems to be passing the updated cache every step, so there should be no such problem. But let me know if you noticed something else :)

@ArthurZucker done! The evaluation shows that caching is faster and more memory efficient than without cache. Below are the results. no special setup and installations for mamba were done, I just tested in my usual transformer env.

Peak at: 27 MiB for cached generation
Time taken: 28.82830 s for cached generation

Peak at: 532 MiB for generation wo cache
Time taken: 829.49494 s for cached generation

```python
from transformers import AutoTokenizer, MambaForCausalLM
import torch
from time import perf_counter

model_id = "state-spaces/mamba-2.8b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token
model = MambaForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
input_ids = tokenizer("Hey how are you doing?", return_tensors="pt")["input_ids"].to(model.device)

# Get init mem used up by model params
init_mem = torch.cuda.memory_allocated()
torch.cuda.reset_max_memory_allocated()

# Eval for cached generation
start = perf_counter()
out = model.generate(input_ids, do_sample=False, use_cache=True, max_new_tokens=500)
torch.cuda.synchronize()

total_time = perf_counter() - start
peak = (torch.cuda.max_memory_allocated() - init_mem) // 1024 ** 2
print(f"Peak at: {peak} MiB for cached generation")
print(f"Time taken: {total_time:.05f} s for cached generation")

# Reset and start eval on use_cache=False
torch.cuda.reset_max_memory_allocated()

start = perf_counter()
out = model.generate(input_ids, do_sample=False, use_cache=False, max_new_tokens=500)
torch.cuda.synchronize()

total_time = perf_counter() - start
peak = (torch.cuda.max_memory_allocated() - init_mem) // 1024 ** 2
print(f"Peak at: {peak} MiB for generation wo cache")
print(f"Time taken: {total_time:.05f} s for cached generation")
```

💘 great work