Added kernels from kernel hub for Bamba model

[romitjain]

What does this PR do?

Adds support for mamba_ssm and causal_conv1d kernels from the kernel-hub in bamba models.

Fixes # (issue)

https://github.com/huggingface/transformers/issues/41208

Before submitting

• [ ] 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, Pull Request section?
• [X] Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case.
• [ ] Did you make sure to update the documentation with your changes? Here are the documentation guidelines, and here are tips on formatting docstrings.
• [ ] Did you write any new necessary tests?

Who can review?

@vasqu @MekkCyber @drbh

[romitjain]

Downstream changes have not been done (eg: modeling_granitemoehybrid.py) Since it was mentioned I can't update those files, not sure how to go about them.

[MekkCyber]

Hi @romitjain I'm working on a pr to make the kernel function mapping easier to do so we don't have to use new functions like lazy_load_mamba_ssm in the modeling files. Once merged we can refactor your PR to make it work with the new API, and then we can apply the changes to all other models using make fix-copies

[romitjain]

@MekkCyber Sure, no worries. Let me know (or you can share your PR here) once done, and I will udpate my PR

[MekkCyber]

sure here is the pr : https://github.com/huggingface/transformers/pull/41577

[romitjain]

@MekkCyber I believe I can refactor my PR by using your new mapping function now?

[MekkCyber]

Yes, i did that for falcon models here : https://github.com/huggingface/transformers/pull/41664, you can do that for bamba models then using the same API, however the mamba-ssm kernel needs to be fixed now before merging the PRs.

[romitjain]

@MekkCyber @vasqu

The mamba_ssm kernels were breaking for me because it was not able to import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined. This was working earlier but has broken in the latest revision. The default revision uses the incorrect output class (https://github.com/state-spaces/mamba/pull/807).

For my local testing, I made local changes to the mamba_ssm repo and tested the flow. However, we would need to fix it on kernels hub for this PR to work end to end.

Other than that, structurally, this is ready for review. Can you please have a look?

[romitjain]

@vasqu In the PR that you referenced, the lazy_load_kernel function will be called for every forward step ref

Since it is not cached, IMO either the approach in this PR or adding lru_cache decorator to lazy_load_kernel would be a better solve.

WDYT?

[vasqu]

@vasqu Since it is not cached, IMO either the approach in this PR or adding lru_cache decorator to lazy_load_kernel would be a better solve.

SGTM. I will still leave it to @MekkCyber tho as he's the main guy for anything kernels. In essence, it just should be the same way everywhere to avoid unnecessary tech debt

[romitjain]

@MekkCyber I have addressed your comments, PTAL

PS: It would require resolution of this issue: https://github.com/huggingface/transformers/pull/41540#discussion_r2518310902

[romitjain]

I did not run make fix-copies since I believe I should not edit the modeling files. But due to that the CI is failing, lmk what should be the steps?

[SunMarc]

In your case, you indeed need to run make fix-copies to propagate the modification you did in the modular files to the real modeling files

[romitjain]

@SunMarc Done, can you please review? cc: @MekkCyber

[romitjain]

@ArthurZucker That is working for me! If kernels is not installed, it falls back to look at default system packages

[romitjain]

@MekkCyber Gentle reminder. What would be the next steps?

[MekkCyber]

Hi @romitjain, nothing to do on your side for now. I need to refactor the kernel function calls to make them less verbose and less visually intrusive.

[romitjain]

Thanks @MekkCyber 👍🏽

[romitjain]

@MekkCyber Is there any ETA for this merge? I can help out with other changes if required in lazy_load_kernel.

[MekkCyber]

Thanks @romitjain ! We are currently working on adding a decorator to use function kernels from the hub. Once it's ready we can integrate it. Will let you know

[romitjain]

@MekkCyber Would really appreciate an update on this. I am happy to help

I see two functions use_kernel_forward_from_hub and use_kernel_func_from_hub, do you suggest that they should be used instead?

[romitjain]

@MekkCyber I have resolved your comments. You might need to fix the upstream mamba-ssm package in kernel hub too.

https://github.com/huggingface/transformers/pull/41540#discussion_r2518310902

[MekkCyber]

Hi @romitjain ! It should be good now, you can use the kernels from this version : https://huggingface.co/kernels-community/mamba-ssm/tree/v0.0.4 instead of the clean-mamba-ssm one, let me know if you have any problems when you test it

[romitjain]

@MekkCyber Thanks for making the upstream fix, all the imports have been improved now to avoid nested calls. PTAL

[MekkCyber]

Thanks @romitjain ! can you fix the styling issues with make style ?

[romitjain]

Oops, I forgot to run that @MekkCyber Done now.

[romitjain]

Thanks @MekkCyber What would be the next steps for the merge?

[MekkCyber]

Failling tests seem unrelated! we just need to wait for a green CI to merge

[romitjain]

Seems like it is failing for the latest commits: https://github.com/huggingface/transformers/commits/main/

[github-actions[bot]]

[For maintainers] Suggested jobs to run (before merge)

run-slow: bamba, granitemoehybrid, jamba, mamba2, qwen3_next

[github-actions[bot]]

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=41540&sha=65e283