Added kernels from kernel hub for Bamba model by romitjain · Pull Request #41540 · huggingface/transformers

romitjain · 2025-10-13T08:59:35Z

What does this PR do?

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

Fixes # (issue)

#41208

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
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 · 2025-10-13T09:03:19Z

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 · 2025-10-16T08:38:21Z

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 · 2025-10-16T08:47:44Z

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

MekkCyber · 2025-10-16T08:52:55Z

sure here is the pr : #41577

romitjain · 2025-10-17T06:05:46Z

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

MekkCyber · 2025-10-17T08:27:51Z

Yes, i did that for falcon models here : #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.

…eature-bamba-kernels-from-hub

Signed-off-by: romit <romit@ibm.com>

romitjain · 2025-11-10T17:10:09Z

@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 (state-spaces/mamba#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?

src/transformers/integrations/hub_kernels.py

romitjain · 2025-11-10T17:11:16Z

src/transformers/integrations/hub_kernels.py


 _HUB_KERNEL_MAPPING: dict[str, dict[str, str]] = {
    "causal-conv1d": {"repo_id": "kernels-community/causal-conv1d"},
+    "mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "revision": "clean-mamba-ssm"},


Copied over from: #41664

vasqu

I don't have many comments except we should align with #41664 on how to lazy load.

These kernels are essentially the same so we ought to standardize it. The mamba-ssm side should be fixed on their repo, thx for the PR there!

cc @MekkCyber if you can take look here since it's close to what you did for the other mamba related model (falcon)

vasqu · 2025-11-11T18:32:15Z

src/transformers/utils/import_utils.py

+@lru_cache
+def is_einops_available() -> bool:
+    return _is_package_available("einops")


Shouldnt be needed?

Now that I recall it, we would need to this check before trying to load mamba-ssm kernels from kernels hub according to README here: https://huggingface.co/kernels-community/mamba-ssm

Should I inject these requirements in _HUB_KERNEL_MAPPING in src/transformers/integrations/hub_kernels.py

But we don't use this check anywhere, no? It doesn't hurt to have it either way, just a bit confused about the usage

for the clean-mamba-ssm implementation I don't think we need this, let's just remove it for now

Sure, will remove this

Also, re: @vasqu's comment, I had added this earlier but forgot to remove this in my latest commit. Will remove it since clean-mamba-ssm won't need it

@MekkCyber clean-mamba-ssm has a issue. It does not expose all the kernel functions: mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined

see: #41540 (comment)

Hemmm, sure will export them in the branch

romitjain · 2025-11-12T05:38:27Z

@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 · 2025-11-12T12:38:17Z

@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

Thanks a lot @romitjain !

src/transformers/models/bamba/modular_bamba.py

src/transformers/models/jamba/modular_jamba.py

MekkCyber · 2025-11-12T13:05:15Z

src/transformers/utils/import_utils.py

+@lru_cache
+def is_einops_available() -> bool:
+    return _is_package_available("einops")


for the clean-mamba-ssm implementation I don't think we need this, let's just remove it for now

Signed-off-by: romit <romit@ibm.com>

romitjain · 2025-11-12T17:19:44Z

@MekkCyber I have addressed your comments, PTAL

PS: It would require resolution of this issue: #41540 (comment)

Signed-off-by: romit <romit@ibm.com>

romitjain · 2025-11-13T04:31:39Z

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 · 2025-11-13T13:20:16Z

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

…/feature-bamba-kernels-from-hub

MekkCyber

Thanks for fixing this @romitjain, I will find some time this week to rework how imports are done inside the kernel because it's really not optimal to have such nested imports

MekkCyber · 2025-12-09T14:29:40Z

src/transformers/models/bamba/modular_bamba.py

+mamba_ssm_triton = getattr(getattr(mamba_ssm, "ops", None), "triton", None)
+selective_state_update = getattr(
+    getattr(mamba_ssm_triton, "selective_state_update", None), "selective_state_update", None


I think we need to rework how are these imports nested. It doesn't make sense

I didn't get you @MekkCyber?
I am also not in favor of these nested imports, but that is what you had previously requested

yes i mean it shouldn't be done this way in the kernel, like we should have an init file with all necessary functions, and then we only use getattr once, and not in a nested way

Sure, let me know once you update the kernels init file, I can make the changes here.

Thanks a lot ! And sorry this is taking a lot of time 🙏, I’ll work on fixing things this week, at the latest.

No worries at all. Thank you for your continued review on this!

MekkCyber · 2025-12-10T16:40:00Z

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

…/feature-bamba-kernels-from-hub

Signed-off-by: romit <romit@ibm.com>

romitjain · 2025-12-15T08:33:03Z

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

MekkCyber

Thanks @romitjain ! lgtm, only one nit and good to go

MekkCyber · 2025-12-15T09:18:10Z

src/transformers/models/bamba/modeling_bamba.py

+causal_conv1d = lazy_load_kernel("causal-conv1d")
+causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
+causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)
+
+mamba_ssm = lazy_load_kernel("mamba-ssm")
+selective_state_update = getattr(mamba_ssm, "selective_state_update", None)
+mamba_chunk_scan_combined = getattr(mamba_ssm, "mamba_chunk_scan_combined", None)
+mamba_split_conv1d_scan_combined = getattr(mamba_ssm, "mamba_split_conv1d_scan_combined", None)


let's move this inside the Mixer to avoid loading the kernels at import time, the same we did here : https://github.com/romitjain/transformers/blob/03590f18cf59be7a9e215b946bd7ade3d8b12a7a/src/transformers/models/mamba/modeling_mamba.py#L201:L214

@MekkCyber Done for mamba2, jamba and bamba

Signed-off-by: romit <romit@ibm.com>

MekkCyber · 2025-12-15T12:43:02Z

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

Signed-off-by: romit <romit@ibm.com>

romitjain · 2025-12-15T13:03:12Z

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

MekkCyber

lgtm! thanks for iterating on this PR

romitjain · 2025-12-15T14:44:49Z

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

MekkCyber · 2025-12-15T21:21:29Z

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

romitjain · 2025-12-16T04:55:22Z

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

github-actions · 2025-12-16T11:02:29Z

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

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

github-actions · 2025-12-16T11:11:46Z

View the CircleCI Test Summary for this PR:

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

ArthurZucker

ty

* Added kernels from kernel hub for Bamba model * Updated kernel loading Signed-off-by: romit <romit@ibm.com> * Remove einops Signed-off-by: romit <romit@ibm.com> * Removed global vars Signed-off-by: romit <romit@ibm.com> * Fixed make style Signed-off-by: romit <romit@ibm.com> * Nit Signed-off-by: romit <romit@ibm.com> * Added modeling files Signed-off-by: romit <romit@ibm.com> * Fixed merge conflict Signed-off-by: romit <romit@ibm.com> * fixed lint Signed-off-by: romitjain <romit@ibm.com> * Removed global import * Small updates * Updated * Resolved merge conflicts * Fixed the nested import Signed-off-by: romit <romit@ibm.com> * Moved imports inside mixer Signed-off-by: romit <romit@ibm.com> * CI CD fix Signed-off-by: romit <romit@ibm.com> --------- Signed-off-by: romit <romit@ibm.com> Signed-off-by: romitjain <romit@ibm.com> Co-authored-by: Mohamed Mekkouri <93391238+MekkCyber@users.noreply.github.com>

Added kernels from kernel hub for Bamba model

4f3c38d

romitjain marked this pull request as draft October 13, 2025 08:59

romitjain mentioned this pull request Oct 13, 2025

Integrate mamba SSM kernels from the hub #41208

Closed

romitjain added 2 commits November 10, 2025 08:23

Merge branch 'main' of github.com:romitjain/transformers into romit/f…

d68a2e7

…eature-bamba-kernels-from-hub

Updated kernel loading

a892065

Signed-off-by: romit <romit@ibm.com>

romitjain marked this pull request as ready for review November 10, 2025 17:06

github-actions bot requested review from ArthurZucker and SunMarc November 10, 2025 17:07

romitjain commented Nov 10, 2025

View reviewed changes

src/transformers/integrations/hub_kernels.py Show resolved Hide resolved

romitjain commented Nov 10, 2025

View reviewed changes

vasqu reviewed Nov 11, 2025

View reviewed changes

MekkCyber reviewed Nov 12, 2025

View reviewed changes

romitjain added 3 commits November 12, 2025 15:41

Remove einops

5ed0fc1

Signed-off-by: romit <romit@ibm.com>

Removed global vars

0a4f79b

Signed-off-by: romit <romit@ibm.com>

Fixed make style

db594e9

Signed-off-by: romit <romit@ibm.com>

romitjain requested review from MekkCyber and vasqu November 12, 2025 17:19

Nit

2ef69e6

Signed-off-by: romit <romit@ibm.com>

romitjain added 3 commits December 9, 2025 19:42

Small updates

a3800e8

Updated

8e48976

Merge branch 'main' of github.com:huggingface/transformers into romit…

df3be53

…/feature-bamba-kernels-from-hub

MekkCyber reviewed Dec 9, 2025

View reviewed changes

Resolved merge conflicts

779486c

romitjain added 2 commits December 15, 2025 07:09

Merge branch 'main' of github.com:huggingface/transformers into romit…

bca8f06

…/feature-bamba-kernels-from-hub

Fixed the nested import

03590f1

Signed-off-by: romit <romit@ibm.com>

romitjain requested a review from MekkCyber December 15, 2025 08:33

MekkCyber reviewed Dec 15, 2025

View reviewed changes

Moved imports inside mixer

4867ef7

Signed-off-by: romit <romit@ibm.com>

romitjain requested a review from MekkCyber December 15, 2025 09:44

CI CD fix

8ab9724

Signed-off-by: romit <romit@ibm.com>

Merge branch 'main' into romit/feature-bamba-kernels-from-hub

db7d824

MekkCyber approved these changes Dec 15, 2025

View reviewed changes

Merge branch 'main' into romit/feature-bamba-kernels-from-hub

e555ef3

Merge branch 'main' into romit/feature-bamba-kernels-from-hub

0569a61

Merge branch 'main' into romit/feature-bamba-kernels-from-hub

65e283a

ArthurZucker approved these changes Dec 16, 2025

View reviewed changes

ArthurZucker merged commit 0f89661 into huggingface:main Dec 16, 2025
21 of 23 checks passed

romitjain deleted the romit/feature-bamba-kernels-from-hub branch December 16, 2025 11:42

Conversation

romitjain commented Oct 13, 2025

What does this PR do?

Before submitting

Who can review?

Uh oh!

romitjain commented Oct 13, 2025

Uh oh!

MekkCyber commented Oct 16, 2025

Uh oh!

romitjain commented Oct 16, 2025

Uh oh!

MekkCyber commented Oct 16, 2025

Uh oh!

romitjain commented Oct 17, 2025

Uh oh!

MekkCyber commented Oct 17, 2025

Uh oh!

romitjain commented Nov 10, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vasqu left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

romitjain commented Nov 12, 2025

Uh oh!

vasqu commented Nov 12, 2025

Uh oh!

MekkCyber left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

romitjain commented Nov 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

romitjain commented Nov 13, 2025

Uh oh!

SunMarc commented Nov 13, 2025

Uh oh!

MekkCyber left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

romitjain Dec 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MekkCyber commented Dec 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vasqu left a comment •

edited

Loading

romitjain commented Nov 12, 2025 •

edited

Loading

romitjain Dec 9, 2025 •

edited

Loading

MekkCyber commented Dec 10, 2025 •

edited

Loading

MekkCyber Dec 15, 2025 •

edited

Loading

MekkCyber commented Dec 15, 2025 •

edited

Loading