mamba generation throughput lower than original due to DecodingCGCache
Author: y1xia0w
Pull request: no
Labels: name: Feature request

name: Good Difficult Issue

name: Compilation

### System Info

Python 3.10.13, CUDA 12.1
GPU = NVIDIA GeForce RTX 2080 Ti. Max memory = 10.747 GB.

torch==2.2.1
torchaudio==2.1.0
torchvision==0.16.0
tokenizers==0.15.2
transformers ==git+https://github.com/huggingface/transformers@dd1c9052159ae824c8acef7c2552f9fad5ca020a
triton==2.2.0
causal_conv1d==git+https://github.com/Dao-AILab/causal-conv1d.git@96456720c00393a5c32872d8352d7a7ec31fb3db#egg=causal_conv1d
mamba_ssm==git+https://github.com/state-spaces/mamba.git@9127d1f47f367f5c9cc49c73ad73557089d02cb8#egg=mamba_ssm

### Who can help?

text models: @ArthurZucker and @younesbelkada
generate: @gante

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

The key model initialization and generation parts are given as below.

## Original code repo
In the original [code repo](https://github.com/state-spaces/mamba/blob/main/benchmarks/benchmark_generation_mamba_simple.py)

```
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-130m")
model.eval()

model.generate(
        input_ids=input_ids,
        max_length=max_length,
        **cg=True**
    )
```
Then throughput for generating 1K length is
```
Number of parameters: 129135360
Prompt length: 100, generation length: 1000
Prompt processing + decoding time: 1011 ms
```

## Using the HF library
```
from transformers import MambaForCausalLM
model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
model.eval()

model.generate(
        input_ids=input_ids,
        max_length=max_length
    )
```
Then throughput for generating 1K length is
```
Number of parameters: 129135360
Prompt length: 100, generation length: 1000
state-spaces/mamba-130m-hf prompt processing + decoding time: 15970ms
```


### Expected behavior

The "cg=True" is [confirmed to be the part has a significant impact on the generation performance](https://github.com/state-spaces/mamba/issues/90) for mamba.

I have tried:

1. Passing the "use_cache=True" as follows won't affect the results
```
model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf", use_cache=True)
or
model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf", cache_params={use_cache: True})
or
model.config.use_cache=True
```

2. Modifying the mamba model to force the argument "use_cache=True" in the [MambaModel](https://github.com/huggingface/transformers/blob/main/src/transformers/models/mamba/modeling_mamba.py), but still not working.


I assume this is related to the #29605, but modifying the argument directly seems not solving the problem.

Comments:
Hi! Thanks for raising this issue. I agree. I've realized the same since some weeks ago while testing Mamba, as previously I was using the mamba-ssm repo and was able to to generation much faster using the DecodingCGCache. I've been tracking that other issue too, as although it seems just variables refactor it includes a change in the prepare_inputs_for_generation function passing the kwargs (https://github.com/huggingface/transformers/pull/29605/files#diff-e1d4758c08973fdac2c23a8a3710872d943ce8509035205da4a681bc4dcaf1c3R694). I didn't created any issue as I wondered that PR will be merged soon due to its simplicity, but seems it is not. Also, I'm not sure if that PR solves all the issue.

Hey! Thanks both I'll dive a bit on this, contributions are also welcome. 
The `use_cache` should be using the cache, and in my test it was, I'll check.

Hi folks 👋 

It doesn't look like a caching issue, but a compilation one -- under the hood, the `cg` flag in the original repo triggers [compilation](https://github.com/state-spaces/mamba/blob/9127d1f47f367f5c9cc49c73ad73557089d02cb8/mamba_ssm/utils/generation.py#L340) (`cg` probably stands for cuda graphs).

Our implementation is not compatible with `fullgraph` compilation (the equivalent), which may explain the difference.

@ArthurZucker in case you want to reproduce:
```py
from transformers import MambaForCausalLM
import torch
from time import time

model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf", device_map="auto")
model.eval()
model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

input_ids = torch.arange(100, device="cuda").unsqueeze(0)
max_length = 1000

start = time()
model = torch.compile(model)
out = model.generate(
    input_ids=input_ids,
    max_length=max_length,
)
print(f"Time: {time() - start:.2f}s")
print(out.shape)
```

Might be the `einsum` that were all replaces by normal operation, saw some issue with this recently as well

@gante Nice catch and thank you for your effort to address it!
@ArthurZucker Could you please elaborate a bit more and also share the other related issues? Thank you very much!

#29544 is what made me think of that!

Hi! Thanks for the findings :)
@ArthurZucker is there any update on this issue? I saw https://github.com/huggingface/transformers/pull/29544 is merged, but  I'm not sure if the compilation issue described affecting the CG is fixed or not.
Thanks!

I did not test the compilation with mamba, I right now can't investigate 😢 would love if you can. 
Now to be fair, we need to compare greedy generation only, not sampling, not beam search, making sure apple are compared with apple. As said, the compilation is not currently supported

I'll set this as a feature request: add support for `torch.compile` for `Mamba`

Hi @ArthurZucker :)
I tested using exactly the same generation config, just greedy search with no sampling. Original Mamba repo was faster than HF during generation.
So then, just to understand, the only possible fix would be supporting compilation for Mamba so that de CG works as in the original repo?
Thanks!

Added it to our generate + torch.compile tracker, which we are actively tackling :)

yes, same issue. 

is there any update for the solution?

#30139 and #31247 should help! 

Nice to see that updates! Are there any other pending issues to complete this 29699 one? Thanks!

I think it's completed

Alright! Was the code above re-tested by anyone to compare speed again? Using which transformers release? Thanks!

@ArthurZucker Thank you for the updates on the nice job!

I have tested the throughput again, there seems some improvement, however, the difference remains significant.

Package info:
```
torch==2.2.1
transformers==4.44.0
triton==2.2.0
mamba-ssm==git+https://github.com/state-spaces/mamba.git@9127d1f47f367f5c9cc49c73ad73557089d02cb8
causal_conv1d==git+https://github.com/Dao-AILab/causal-conv1d.git@96456720c00393a5c32872d8352d7a7ec31fb3db
```

I run the generation with the same 100 prompt length and 1000 generation length 5 times and get the average:
```
# HF transformers v4.44.0, no flash_attention_2 support
# Time: 14.75s
# torch.Size([1, 1000])

# mamba_ssm, without cg=True
# Time: 13.99s
# torch.Size([1, 1000])

# mamba_ssm, with cg=True
# Time: 1.33s
# torch.Size([1, 1000])
```

I assume the reason of the generation speed difference would be:
- The original mamba adopt part of the flash attention mechanism in their implemetation, which significantly accelerate the generation speed.
- The flash-attn capability on mamba architecutre hasn't been supported in the v4.44.0 `transformers` library.

Reference:
Quote from [mamba document in HF library](https://huggingface.co/docs/transformers/en/model_doc/mamba)
> The current implementation leverages the original cuda kernels: the equivalent of flash attention for Mamba are hosted in the [mamba-ssm](https://github.com/state-spaces/mamba) and the [causal_conv1d](https://github.com/Dao-AILab/causal-conv1d) repositories. Make sure to install them if your hardware supports them!

When you said you tested, did you use `torch.compile(..., mode="reduce_overhead")`? 

Did you read this https://github.com/huggingface/transformers/pull/31247#issuecomment-2155723654 in the linked PR for compile support ? 

`cg=True` uses cuda graphs, to compare apples with apples you have to use the provided perf snippet


想问下基于tensformer调用模型的吞吐量测试代码是在哪里找的？


可以在这里找到: 
```python
@torch.no_grad
def perf():
    tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-1.4b-hf")
    tokenizer.pad_token = tokenizer.eos_token
    inputs = tokenizer("Hey how are you doing today ? " * 100, return_tensors="pt", padding=True).to('cuda')

    model = MambaForCausalLM.from_pretrained("state-spaces/mamba-1.4b-hf", torch_dtype=torch.float16)
    model.config.use_cache = True
    model.to('cuda')    

    input_ids = inputs.input_ids
    cache = MambaCache(model.config, 1, device=input_ids.device)
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    logits = model(input_ids, cache_params = cache).logits
    next_token = torch.argmax(logits[:, -1], dim=-1)[:, None]

    model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)
    torch.cuda.synchronize()
    for i in range(10):
        start.record()
        logits = model(next_token.clone(), cache_params = cache).logits
        next_token = torch.argmax(logits[:, -1], dim=-1)[:, None]

        end.record()
        torch.cuda.synchronize()
        print(f'Step {i}, Total time: {start.elapsed_time(end)} ms, next_token = {next_token.int()}')

```
(https://github.com/huggingface/transformers/pull/31247#issuecomment-2155723654)