Added kernels from kernel hub for Bamba model
Author: romitjain
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [X] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? 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](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@vasqu @MekkCyber @drbh

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @zach-huggingface @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker @S1ro1
- CIs: @ydshieh

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc @zach-huggingface
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Comments:
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.

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`

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

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

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

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.

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


@vasqu In the PR that you referenced, the `lazy_load_kernel` function will be called for every forward step [ref](https://github.com/huggingface/transformers/blob/9939e585c0c6d6bdc1eec90c5dfeab2e2e54db11/src/transformers/models/falcon_mamba/modular_falcon_mamba.py#L302)

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 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



@MekkCyber I have addressed your comments, PTAL

PS: It would require resolution of this issue: https://github.com/huggingface/transformers/pull/41540#discussion_r2518310902

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?

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

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

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

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

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.

Thanks @MekkCyber 👍🏽 

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

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

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

@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

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

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

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

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

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

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

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

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

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

View the CircleCI Test Summary for this PR:

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