Integrate mamba SSM kernels from the hub
Author: romitjain
Pull request: no
Labels: name: Feature request

### Feature request

Currently, mamba kernels are imported via the main source package ex, for [GraniteMoeHybrid](https://github.com/huggingface/transformers/blob/main/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py#L44-L46)

Can we migrate this to use the kernels-hub (`kernels-community/mamba-ssm`) variation instead?

### Motivation

Removes the external dependency. Kernel hub is also integrated at several other places throughout the library.

### Your contribution

I can submit a PR for migrating from the PyPi `mamba_ssm` package to the `kernels` package for mamba ops.

Comments:
Hey @romitjain, I think that would be very cool. We should keep the option to have the og package as well tho. We had a similar PR for causal-conv1d over here #40765 

cc @MekkCyber to confirm if that's ok 

Thanks for the reference @vasqu 
If you are good with the proposal, I will raise a PR.

PS: I can't seem to open the [causal conv1d](https://huggingface.co/kernels-community/causal-conv1d) kernel package on the hub. Is that expected?

Me neither xD might be set to private atm, let's wait for @MekkCyber to respond as he's more involved with kernels 

Yup sorry it's public now. Yes it's very cool to have `mamba-ssm` as well used in transformers from the hub. We already have [kernels-community/mamba-ssm](https://huggingface.co/kernels-community/mamba-ssm), You can definitely try to make it work and ping me when it's done 🤗 . If you have any issues with the kernel feel free to open an issue on the hub.

Thanks @MekkCyber 
Let me implement it and raise a PR

The `mamba_ssm` kernel is not compatible with the latest transformers repo. I have raised a PR to fix it, can someone here please take a look or tag an appropriate member?

https://huggingface.co/kernels-community/mamba-ssm/discussions/3

(Maybe @MekkCyber)

Thank you

I have also raised a draft PR (https://github.com/huggingface/transformers/pull/41540), which is contingent on the hub PR being merged

Thanks @romitjain ! will look into it and rebuild the kernel

@MekkCyber Can you please rebuild the kernel? The issue has been solved in the source repo: https://github.com/state-spaces/mamba/pull/807

Thanks @romitjain will upstream the changes to `kernels-community`

I'm honestly fed up with the current mamba implementation.
I was hoping that mamba would consume less GPU memory when the token window was large, but in fact it does the opposite. It consumes a significantly larger amount of memory.
I tried it with 4.57.3.
```python
import json
import torch
import transformers
from transformers import GraniteMoeHybridForCausalLM, AutoTokenizer, GraniteMoeHybridConfig
test = """{ 
"architectures": [ 
"GraniteMoeHybridForCausalLM" 
], 
"attention_bias": false, 
"attention_dropout": 0.0, 
"attention_multiplier": 0.015625, 
"bos_token_id": 100257, 
"embedding_multiplier": 12, 
"eos_token_id": 100257, 
"hidden_act": "silu", 
"hidden_size": 768, 
"init_method": "mup", 
"initializer_range": 0.1, "intermediate_size": 2048, 
"layer_types": [ 
"mamba", 
"mamba", 
"mamba", 
"mamba", 
"mamba", 
"mamba" 
], 
"logits_scaling": 3, 
"mamba_chunk_size": 256, 
"mamba_conv_bias": true, 
"mamba_d_conv": 4, 
"mamba_d_head": 32, 
"mamba_d_state": 128, 
"mamba_expand": 2, 
"mamba_n_groups": 1, 
"mamba_n_heads": 48, 
"mamba_proj_bias": false, 
"max_position_embeddings": 2000, 
"model_type": "granitemoehybrid", 
"normalization_function": "rmsnorm", 
"num_attention_heads": 12, 
"num_experts_per_tok": 0, 
"num_hidden_layers": 6, 
"num_key_value_heads": 4, 
"num_local_experts": 0, 
"output_router_logits": false, 
"pad_token_id": 100256, 
"position_embedding_type": "nope", 
"residual_multiplier": 0.246, 
"rms_norm_eps": 1e-05, 
"rope_scaling": null, 
"rope_theta": 10000, 
"router_aux_loss_coef": 0.01, 
"shared_intermediate_size": 2048, 
"tie_word_embeddings": true, 
"torch_dtype": "bfloat16", 
"transformers_version": "4.56.0", 
"use_cache": true, 
"vocab_size": 100352
}"""
conf = GraniteMoeHybridConfig.from_dict(json.loads(test))
model = GraniteMoeHybridForCausalLM(conf).to("cuda:1")
model = model.bfloat16()
tokens = [100264, 9125, 100265, 2675, 527, 264, 11190, 18328, 13, 
5321, 6106, 14847, 527, 6721, 11, 13687, 11, 323, 6220, 13, 100257, 198, 100264, 882, 100265, 5618, 1160,
832, 29022, 8483, 27692, 7559, 304, 279, 3723, 4273,
13, 1472, 1288, 1193, 2612, 1202, 836, 323, 3813,
13, 100257, 198, 100264, 78191]+[2675, 527, 264, 11190, 18328, 13,
5321, 6106, 14847, 527, 6721, 11, 13687, 11, 323,
6220, 13, 100257, 198, 100264, 882, 100265, 5618, 1160,
832, 29022, 8483, 27692, 7559, 304, 279, 3723, 4273,
13, 1472, 1288, 1193, 2612, 1202, 836, 323, 3813,
13, 100257, 198, 100264, 78191]*10000+[100265]
tokens = [tokens[:2000]]
mask = [[1]*2000]
out = model(input_ids=torch.tensor(tokens).long().to("cuda:1"), attention_mask=torch.tensor(mask).long().to("cuda:1"))
```
This actually uses 23GB of GPU memory! That seems like an insane amount of memory.
Changing mamba to attention in the layer_types reduces it to 1.9GB, which is a reasonable size.
Does it get fixed with SSM kernels?

@tanreinama 
You can try two options:
1. Install mamba-ssm pypi package `pip install mamba-ssm[causal-conv1d]` and try running the above code again. In STDOUT, you should see something like - fastpath available.. using mamba implementation on GPU
2. Build transformers from source and try again. In the attached PR, we have integrated the mamba kernels natively in transformers and you will not need to install anything extra.



@MekkCyber I am closing this issue, since it got resolved in #41540 

@tanreinama You likely use the slow path which is expected to explode in memory. If you want to run it efficiently, there is no way around the kernels. @romitjain has already pointed both options for this

Thank you  @romitjain , @vasqu  . Even when using the mamba-ssm kernel, it works with device="cuda", but not with device="cuda:1". This seems to be a problem with Triton, but currently there seems to be no way to efficiently run mamba on a GPU other than device#0.