🚨 [`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-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! 🤗

…e, adopt standard names - Drop the private RWKV7_WKV_FUNCTIONS registry and its config knob: it duplicated the kernel-substitution mechanism huggingface#47630 just unified across linear-attention models. The torch reference is called directly; an optimised kernel arrives via hub-kernels when one exists, not a per-model registry. The doc section demonstrating the registry goes with it. - Gate generate's input truncation on is_first_iteration instead of the state's existence: a pre-allocated state (the compile contract) or a warm state (a resumed chat turn) silently dropped every prompt token but the last one. Same gate Mamba uses; bare callers keep the old behaviour. - Rename cu_seq_lens to the ecosystem's cu_seq_lens_q, so padding-free collators reach the packed path instead of being swallowed by kwargs. - Refuse Rwkv7Cache.crop() loudly (the state is not invertible), fix keys_to_ignore_at_inference to name the output this model actually has, add the missing explicit models/__init__ re-export, and drop the attentions=None output field and ignored output_attentions parameter. - Tests: generate-with-state equalities, end-to-end left-padded batch generate, compiled decode step over allocate_state, the chunked-prefill mask-slicing contract, a multi-chunk numpy-oracle case; pin the numpy reference to its external sources and the rwkv package version used for the integration expectations; let the shared gradient-checkpointing trainings run via test_all_params_have_gradient. - Docs: compile guidance now reflects measurement (reduce-overhead with an eager prefill; max-autotune measured slower on this kernel chain).

#1102) * Vendored FLA: re-resolve kernel closures on modules imported before us Since huggingface/transformers#47630 the gated delta kernels are bound by use_kernel_func_from_hub_with_fallback, which resolves the implementation at decoration time and closes over it. A modeling module imported before fla is live freezes the pure-PyTorch fallback into that closure, and nothing done to fla afterwards can reach it. _repair_already_imported_modeling only rebinds module globals, which these modules no longer have. Re-apply the decorator to the undecorated function once fla is live, rather than writing to cell_contents: the wrapper also closes over the implementation's parameter names, and those differ between the fallback and the kernel, so patching the implementation alone would filter out arguments the kernel needs. * Test the repair against the live fla, not against the fallback A modeling module imported while a real fla install was live froze that install's kernel into the wrapper. UNSLOTH_FORCE_VENDORED_FLA and the Hopper #640 switch then purge the install, but the old check treated any non-fallback implementation as valid and skipped it, so the wrapper kept calling the kernel we replaced the install to avoid, which is the miscompiled backward. Compare against the kernel fla resolves to now instead. When there is no live replacement the rebuild lands on the pure-torch fallback, which is taken deliberately: torch beats calling into a purged install. The kernel-hub tests now skip on a transformers predating #47630, where the repair correctly returns () and there is nothing to patch.

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 in generate #47604

Merged

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

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

Hakureirm mentioned this pull request Aug 6, 2026

Add RWKV-7 (Goose) #47780

Open

5 tasks

sharmax-vikas mentioned this pull request Aug 10, 2026

Add Kimi Linear (Kimi Delta Attention) native support #47875

Open

2 tasks

DimensionSTP mentioned this pull request Aug 20, 2026

Installed nested optional kernels silently fall back to the Torch implementation #48148

Open

4 tasks

vasqu mentioned this pull request Aug 24, 2026

Fix inter-chunk recurrence in the Zamba2 / Nemotron-H Mamba2 slow path #47250

Closed

Rocketknight1 mentioned this pull request Aug 24, 2026

Zamba-7B-v1 runs slow during torch.compile due to use_associative_scan=False #48080

Open

danielhanchen mentioned this pull request Aug 25, 2026

Vendored FLA: re-resolve kernel closures on modules imported before us unslothai/unsloth-zoo#1102

Merged

		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!

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

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vasqu commented Jul 29, 2026 •

edited by github-actions Bot

Loading