Update Mamba types and pass through use_cache attr to MambaModel
Author: koayon
Pull request: yes

# What does this PR do?

1. Previously the cache_params variable was inconsistently typed as `MambaCache`, `torch.Tensor` or `list[torch.Tensor]`. This PR updates this to `MambaCache` everywhere which is inline with the attributes that are being accessed in the logic. 
2. Previously if you ran the forward method with use_cache=True, the cache_params in the output would still be None as this argument wasn't being passed through to the `MambaModel`. This PR fixes this as below:
3. Also updates the docstrings in line with this change

Allowed the use_cache information to be passed through so that you can do: 
```python
import torch as t

from transformers import AutoTokenizer

if __name__ == "__main__":
    model = MambaForCausalLM(MambaConfig())
    tokeniser = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")
    input_ids: t.Tensor = tokeniser("Hey how are you doing?", return_tensors="pt")[  # type: ignore
        "input_ids"
    ]

    out: MambaCausalLMOutput = model(input_ids=input_ids, use_cache=True)
    assert out.cache_params is not None
    print(out.cache_params.ssm_states)
```

And get back the ssm_states, which was not previously possible

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).


## Who can review?

@ArthurZucker
@gante 

Comments:
Thanks for adding this @koayon 

Pinging @gante for first review of the cache logic, as @ArthurZucker is off this week

Hey @ArthurZucker! Hope you had a great holiday 🙌 
Is it possible for you to take a little look at the docstring here? 

Hey @ArthurZucker, thanks for your review! 🙌 
I've reverted the styling changes, thanks 🙏 

In terms of the image that you were sending, it's unfortunately not showing up for me. But without the change to pass in use_cache I don't see the cache_params being returned. If there's a difference for you, I've just thought that it might be how it's running on CUDA vs MPS/CPU?

I append the following to the file: 

```python
import torch as t

from transformers import AutoTokenizer

if __name__ == "__main__":
    model = MambaForCausalLM(MambaConfig())
    tokeniser = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")
    input_ids: t.Tensor = tokeniser("Hey how are you doing?", return_tensors="pt")["input_ids"]  # type: ignore

    out: MambaCausalLMOutput = model(input_ids=input_ids, use_cache=True)
    assert out.cache_params is not None
    print(out.cache_params.ssm_states)
 ```
 
 If the use_cache argument isn't passed through to the backbone (either with kwargs or separately as in the newer version), there is no cache_params returned and I get the error: 
 ```bash
python src/transformers/models/mamba/modeling_mamba.py      
...
Traceback (most recent call last):
  File "/[PATH_TO_TRANSFORMERS]/transformers/src/transformers/models/mamba/modeling_mamba.py", line 688, in <module>
    assert out.cache_params is not None
AssertionError
```

whereas with the use_cache argument being passed through I get a tensor returned: 
```bash
{0: tensor([[[-5.5237e-04,  9.6599e-04,  6.6771e-04,  ..., -5.3982e-04,
          -4.6061e-04, -7.0508e-04],
         [-3.7170e-05, -2.2089e-04, -1.0218e-04,  ...,  7.1232e-05, 
         ...
```

It does seem like this would be required for the expected behaviour.  Please let me know if you have any questions! 😄 

Alright, when using from_pretrained, the cache is used and passed subsequently, but not when using the initialization

@ArthurZucker great suggestion, didn't realise that was an attribute of the Config 👌 
Think it's ready to merge now 🚀 

The failing test seems new, but it's because when training the use_cache should be disabled by the model

I'll have a look