Integrate mamba SSM kernels from the hub

[romitjain]

Feature request

Currently, mamba kernels are imported via the main source package ex, for GraniteMoeHybrid

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.

[vasqu]

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

[romitjain]

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 kernel package on the hub. Is that expected?

[vasqu]

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

[MekkCyber]

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, 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.

[romitjain]

Thanks @MekkCyber
Let me implement it and raise a PR

[romitjain]

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

[romitjain]

I have also raised a draft PR (#41540), which is contingent on the hub PR being merged

[MekkCyber]

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

[romitjain]

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

[MekkCyber]

Thanks @romitjain will upstream the changes to kernels-community

[tanreinama]

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.

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?

[romitjain]

@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.

[romitjain]

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