MambaMixer cache not working for chunks
Author: OleguerCanal
Pull request: no

### System Info

I'm trying to run inference on a finetuned MambaModel. The input are audio features and I need output features from Mamba. The standard way would be to simply run inference like so:

```python
      n_tokens = x.shape[1]
      cache_params = None
      last_hidden_states = []
      for i in range(n_tokens):
          mamba_output = mamba_backbone.forward(
              input_ids=None,
              inputs_embeds=x[:, i:i+1, :],
              cache_params=cache_params,
              use_cache=True,
          )
          cache_params = mamba_output.cache_params
          last_hidden_states.append(mamba_output.last_hidden_state)
```

However, since n_tokens can be very large. This is very slow. I thought of processing it in chunks (as large as memory allows), taking advantage of the scan operation. However, if I simply do this:

```python
        chunk=1000
        n_tokens = x.shape[1]
        cache_params = None
        last_hidden_states = []
        for i in range(0, n_tokens, chunk):
            mamba_output = mamba_backbone.forward(
                input_ids=None,
                inputs_embeds=x[:, i:i+chunk, :],
                cache_params=cache_params,
                use_cache=True,
            )
            cache_params = mamba_output.cache_params
            last_hidden_states.append(mamba_output.last_hidden_state)
```

I get the following error by the causal_conv1d function:

```BASH
    mamba_output = mamba_backbone.forward(
  File "/home/mary/.virtualenvs/benchmarking/lib/python3.8/site-packages/transformers/models/mamba/modeling_mamba.py", line 579, in forward
    hidden_states = mixer_block(hidden_states, cache_params=cache_params)
  File "/home/mary/.virtualenvs/benchmarking/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/mary/.virtualenvs/benchmarking/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/mary/.virtualenvs/benchmarking/lib/python3.8/site-packages/transformers/models/mamba/modeling_mamba.py", line 344, in forward
    hidden_states = self.mixer(hidden_states, cache_params=cache_params)
  File "/home/mary/.virtualenvs/benchmarking/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/mary/.virtualenvs/benchmarking/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/mary/.virtualenvs/benchmarking/lib/python3.8/site-packages/transformers/models/mamba/modeling_mamba.py", line 308, in forward
    return self.cuda_kernels_forward(hidden_states, cache_params)
  File "/home/mary/.virtualenvs/benchmarking/lib/python3.8/site-packages/transformers/models/mamba/modeling_mamba.py", line 177, in cuda_kernels_forward
    hidden_states = causal_conv1d_update(
  File "/home/mary/.virtualenvs/benchmarking/lib/python3.8/site-packages/causal_conv1d/causal_conv1d_interface.py", line 187, in causal_conv1d_update
    return causal_conv1d_cuda.causal_conv1d_update(
RuntimeError: x must have shape (batch_size, dim)
```

I don't fully understand it, since the cache has the same shape either way. Any ideas of what could be the issue? (it doesn't crash if I do the chunk first, and then iterate one by one).

cc: @sanchit-gandhi @SunMarc 

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

This code reproduces the error:

```python
import torch
from transformers import MambaConfig, MambaModel

x = torch.randn(2, 1_000, 20).to("cuda")

config = MambaConfig(
    vocab_size=50,
    hidden_size=20,
    state_size=16,
    num_hidden_layers=3,
    d_conv=4,
    expand=2,
)

mamba_backbone = MambaModel(config).to("cuda")

chunk=100
n_tokens = x.shape[1]
cache_params = None
last_hidden_states = []
for i in range(0, n_tokens, chunk):
    print(i)
    mamba_output = mamba_backbone.forward(
        input_ids=None,
        inputs_embeds=x[:, i:i+chunk, :],
        cache_params=cache_params,
        use_cache=True,
    )
    cache_params = mamba_output.cache_params
    last_hidden_states.append(mamba_output.last_hidden_state)
```

### Expected behavior

The output should be the same either processing the inference in chunks or element-wise.

Comments:
cc @ArthurZucker 

Hey @OleguerCanal, Mamba(1) doesn't support an initial hidden state like you think it does. What you describe in essence is sequence parallelism (we can cut the input on the sequence dimension and apply the inference subsequently chunk for chunk).

When you pass a cache, we expect the input (i.e. x) to have a shape of `(bsz, 1, hidden_size)` as Mamba decodes on a one-by-one basis. I assume that due to `seq_len` being greater than 1, 100 in this case, it leads to those shape mismatches (possibly due to something like `squeeze()` not reducing the dimension as expected). 

And even if that worked, you would still have numerical issues as (rms) normalisations and linear projections would also have to be handled properly. 

Edit: Could you show what you did with "Any ideas of what could be the issue? (it doesn't crash if I do the chunk first, and then iterate one by one).". ~~Does it mean you didn't use the cache and went one-by-one?~~ I think it refers to your first code snippet which aligns with the one-by-one input I described above.

You are right @vasqu
I initially thought it was a problem with the huggingface implementation and now I see it is simply not supported by the current Mamba. I'll close the issue, thanks for the response.