🚨 [`Kernels`] Refactor all linear attn models & native kernels fallback by vasqu · Pull Request #47630 · huggingface/transformers

vasqu · 2026-07-29T21:10:27Z

Tl;dr: Allow kernels, og, and torch only to natively coexist along each other. This is also makes kernels as opt in, not as mandatory which caused very weird behavior where it suddenly took precedence even over the og kernels

Model types

GDN
Mamba2
Conv only (inkling, lfm2 etc)
Mamba1

Dependent on

Note:

There is one subtle difference between the og conv1ds and the hf kernels conv1ds
- Og fuses the activation into the kernel while the torch fallback and hf kernel do this afterwards (unfused)
- This is negligible and can be even explained by the difference complation of the kernel binaries
It can be negated by passing activation=None and manually applying the activation ...; act(hidden_states)

HuggingFaceDocBuilderDev · 2026-07-29T21:26:18Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Cyrilvallez

Very very nice, cannot wait for this PR! Just checked bamba for now to give a few early thoughts!

Cyrilvallez · 2026-07-31T03:19:13Z

+            # Only kernels can use this shortcircuit, fallback to normal torch otherwise
+            if fused_output is not None:
+                return fused_output


This makes it a bit awkward IMO... Any way to check if the function was found in kernels instead?

Hard to do because kernels can be swapped at runtime at any time

vasqu · 2026-08-04T13:15:17Z

run-slow: mamba, mamba2, falcon_mamba, jamba

github-actions · 2026-08-04T13:16:56Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/falcon_mamba", "models/jamba", "models/mamba", "models/mamba2"]
quantizations: []

github-actions · 2026-08-04T13:56:11Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	03ef5dcc	workflow commit (merge commit)
PR	c85e6220	branch commit (from PR)
main	410c4b17	base commit (on `main`)

Model CI Report

❌ 1 new failed tests from this PR 😭

falcon_mamba:
tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaIntegrationTests::test_batched_generation (✅ ⟹ ❌)

vasqu · 2026-08-04T14:06:24Z

run-slow: falcon_mamba

github-actions · 2026-08-04T14:08:07Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/falcon_mamba"]
quantizations: []

github-actions · 2026-08-04T14:12:33Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	54970429	workflow commit (merge commit)
PR	2e485f3f	branch commit (from PR)
main	49140dd1	base commit (on `main`)

⚠️ No test being reported (jobs are skipped or cancelled)!

vasqu · 2026-08-04T14:13:06Z

run-slow: falcon_mamba

vasqu · 2026-08-04T14:14:12Z

+                    ' I will be talking about the importance of the internet in our lives.\nThe internet is a global'
+                ],
+                ("cuda", 9): [
+                    ' I am going to talk about the “Theory of Relativity” by Albert Einstein.\n',


Not sure what caused the change on A10 but the output is the same across all devices now so I wouldn't see it as a regression as the outputs are aligned. I assumed some small matmul diff causing slightly different gens previously

github-actions · 2026-08-04T14:15:12Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/falcon_mamba"]
quantizations: []

github-actions · 2026-08-04T14:20:49Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	142930d3	workflow commit (merge commit)
PR	d81829f6	branch commit (from PR)
main	49140dd1	base commit (on `main`)

⚠️ No test being reported (jobs are skipped or cancelled)!

vasqu · 2026-08-04T14:27:24Z

run-slow: bamba, falcon_h1, falcon_mamba, granitemoehybrid, inkling, jamba, lfm2, lfm2_moe, mamba, mamba2, minimax, nemotron_h, olmo_hybrid, qwen3_5, qwen3_5_moe, qwen3_next

github-actions · 2026-08-04T14:29:14Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/bamba", "models/falcon_h1", "models/falcon_mamba", "models/granitemoehybrid", "models/inkling", "models/jamba", "models/lfm2", "models/lfm2_moe", "models/mamba", "models/mamba2", "models/minimax", "models/nemotron_h", "models/olmo_hybrid", "models/qwen3_5", "models/qwen3_5_moe", "models/qwen3_next"]
quantizations: []

github-actions · 2026-08-04T15:18:03Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	62f56867	workflow commit (merge commit)
PR	9dabdf29	branch commit (from PR)
main	49140dd1	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

Cyrilvallez

All right, very nice work! Still a few important details to fix, but should not be much efforts. Most notably:

Make sure we align the name of the torch functions and kernels function everywhere, as most are not the same. This makes it much harder to understand what we do. Having the same name makes it clear for anyone that it's supposed to do the same thing in pure torch
Correctly guard all conv parts with and not record_past for speculative decoding/mtp etc, and be coherent between all our implems

Cyrilvallez · 2026-08-05T04:46:09Z

+@use_kernel_func_from_hub_with_fallback(
+    "mamba_split_conv1d_scan_combined",
+    "mamba_ssm",
+)
+def mamba2_split_conv1d_scan_combined(
+    zxbcdt: torch.Tensor,
+    conv1d_weight: torch.Tensor,
+    conv1d_bias: torch.Tensor | None,
+    dt_bias: torch.Tensor,
+    A: torch.Tensor,
+    D: torch.Tensor,
+    chunk_size: int,
+    initial_states: torch.Tensor | None = None,
+    dt_limit: tuple[float, float] = (0.0, float("inf")),
+    return_final_states: bool = False,
+    activation: str = "silu",
+    rmsnorm_weight: torch.Tensor | None = None,
+    rmsnorm_eps: float = 1e-6,
+    outproj_weight: torch.Tensor | None = None,
+    outproj_bias: torch.Tensor | None = None,
+    headdim: int | None = None,
+    ngroups: int = 1,
+    norm_before_gate: bool = True,
+    **kwargs,
+):
+    return None


This is still a bit awkward IMO 🥲 We do need to have all the args and kwargs so that kernels can match the signature right? We cannot simply use mamba2_split_conv1d_scan_combined(*args, **kwargs)?

Yea not possible because we need to match the kernels signature

Cyrilvallez · 2026-08-05T05:14:20Z

-class JambaMambaMixer(MambaMixer):
+class JambaMambaMixer(FalconMambaMixer):


Is it really better to switch inheritance? Why do we want to add use_mambapy etc?

The main goal here is to show that it's nearly 1:1 in the forward minus mamba inner fn not being usable

Cyrilvallez · 2026-08-05T05:21:37Z

+        q, k, v = torch.split(
+            mixed_qkv.transpose(1, 2),
+            [self.key_dim, self.key_dim, self.value_dim],
+            dim=-1,
+        )


Humm actually, is it really better to merge them given that we need to cat and split?

Pretty sure it's worth it, the conv is a bigger limiting factor than the cat and split

(also just easier for our cache in general)

github-actions · 2026-08-05T08:30:16Z

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

run-slow: bamba, falcon_h1, falcon_mamba, granitemoehybrid, inkling, jamba, lfm2, lfm2_moe, mamba, mamba2, minimax, nemotron_h, olmo_hybrid, qwen3_5, qwen3_5_moe, qwen3_next

github-actions · 2026-08-05T08:59:42Z

CI recap

Dashboard: View test results in Grafana
Latest run: 30919052742:1
Result: success | Jobs: 1 | Tests: 1,786 | Failures: 24 | Duration: 2h 31m

Cyrilvallez

Alright, let's go now if you're done with all the changes after offline discussion!! Thanks again, super happy to land it and finally have fully clean mambas! 🤗

Cyrilvallez and others added 12 commits July 13, 2026 18:29

kernel native

f438a2e

fix warning

7c1335b

have to fix other modular to be coherent

f8ff0db

remove useless

9a8127e

use native lib as well

72b45ca

modular

54d80df

combine them

d2beaf6

doc

c256c47

ignore kwargs, allow og

30e0ff7

Merge remote-tracking branch 'upstream/main' into kernels-fallback

c57b520

fix

38c2ece

gdn like paths (missing conv of olmo hybrid)

7788cac

olmo hybrid fused conv style

cf49de6

vasqu mentioned this pull request Jul 29, 2026

Align OlmoHybrid to use a native cache #47604

Open

vasqu added 13 commits July 29, 2026 21:53

fix

73db89b

fix

8896696

poc mamba2, kernel must compile but torch seems to match

d3b27a6

kernels match --> hf kernels will be needed

f0aa581

style

263f307

quick fixes

a189da9

Merge remote-tracking branch 'upstream/main' into kernels-fallback

8134d38

mamba2 works

4ebeb11

mamba2 suite of models

1786c75

oops

0eca0d0

conv1ds across other models

82bfade

let's try this

520e340

mamba base implementation

f69acfe

Cyrilvallez reviewed Jul 31, 2026

View reviewed changes

Comment thread src/transformers/integrations/hub_kernels.py

fix bamba test

2b4fd75

update falcon mamba - aligned with all other devices

2e485f3

oops

d81829f

vasqu commented Aug 4, 2026

View reviewed changes

Merge branch 'main' into kernels-fallback

9dabdf2

vasqu mentioned this pull request Aug 4, 2026

🚨 Kernels=True ?! #47773

Draft

Cyrilvallez reviewed Aug 5, 2026

View reviewed changes

Cyrilvallez changed the title ~~[Kernels] Native fallback~~ 🚨 [Kernels] Refactor all mamba's recurrent part & Native kernels fallback Aug 5, 2026

Cyrilvallez changed the title ~~🚨 [Kernels] Refactor all mamba's recurrent part & Native kernels fallback~~ 🚨 [Kernels] Refactor all mamba models recurrent part & Native kernels fallback Aug 5, 2026

adress review

4c1864d

vasqu changed the title ~~🚨 [Kernels] Refactor all mamba models recurrent part & Native kernels fallback~~ 🚨 [Kernels] Refactor all linear attn models & native kernels fallback Aug 5, 2026

Cyrilvallez approved these changes Aug 5, 2026

View reviewed changes

vasqu added this pull request to the merge queue Aug 5, 2026

Merged via the queue into huggingface:main with commit e91f7ef Aug 5, 2026
111 checks passed

vasqu deleted the kernels-fallback branch August 5, 2026 09:44

		class JambaMambaMixer(MambaMixer):
		class JambaMambaMixer(FalconMambaMixer):

Uh oh!

Conversation

vasqu commented Jul 29, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Jul 29, 2026

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vasqu commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

CI Results

Commit Info

Model CI Report

Uh oh!

vasqu commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

CI Results

Commit Info

Uh oh!

vasqu commented Aug 4, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

CI Results

Commit Info

Uh oh!

vasqu commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 4, 2026

CI Results

Commit Info

Uh oh!

Cyrilvallez 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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot commented Aug 5, 2026

Uh oh!

github-actions Bot commented Aug 5, 2026

CI recap

vasqu commented Jul 29, 2026 •

edited by github-actions Bot

Loading