🚨 [Kernels] Refactor all linear attn models & native kernels fallback

src/transformers/conversion_mapping.py

@@ -1009,6 +1009,15 @@ def _build_checkpoint_conversion_mapping():
         "olmo_hybrid": [
             WeightRenaming("attention_layer_norm", "input_layernorm"),
             WeightRenaming("feedforward_layer_norm", "post_attention_layernorm"),
+            WeightConverter(
+                source_patterns=[
+                    "linear_attn.q_conv1d.weight",
+                    "linear_attn.k_conv1d.weight",
+                    "linear_attn.v_conv1d.weight",
+                ],
+                target_patterns="linear_attn.conv1d.weight",
+                operations=[Concatenate(dim=0)],
+            ),
         ],
         "qwen3_5_text": [PrefixChange(prefix_to_remove="language_model", model_prefix="model")],
         "sam3_tracker": [

src/transformers/integrations/__init__.py

@@ -81,6 +81,7 @@
         "replace_kernel_forward_from_hub",
         "use_kernel_forward_from_hub",
         "use_kernel_func_from_hub",
+        "use_kernel_func_from_hub_with_fallback",
         "use_kernelized_func",
     ],
     "integration_utils": [
@@ -237,6 +238,7 @@
         replace_kernel_forward_from_hub,
         use_kernel_forward_from_hub,
         use_kernel_func_from_hub,
+        use_kernel_func_from_hub_with_fallback,
         use_kernelized_func,
     )
     from .integration_utils import (

src/transformers/integrations/accelerate.py

@@ -920,31 +920,38 @@ def check_tied_parameters_on_same_device(tied_params, device_map):
             )
 
 
-def force_accelerate_hooks(child_module_name: str) -> Callable:
+def force_accelerate_hooks(child_module_names: str | list[str]) -> Callable:
     """
-    Decorator to forcefully fire the accelerate hooks of `child_module_name`, before entering the forward of the parent itself.
-    Indeed, the hooks of the child are only fired through the `forward` child's method, so if the child weights are used directly,
+    Decorator to forcefully fire the accelerate hooks of `child_module_names`, before entering the forward of the parent itself.
+    Indeed, the hooks of a child are only fired through the `forward` child's method, so if the child weights are used directly,
     as is the case inside `causal_conv1d_fn` and `causal_conv1d_update` for example, they will not be fired. This may cause device
     issues, especially in the case of offloading, that this decorator will correct.
     """
 
+    if isinstance(child_module_names, str):
+        child_module_names = [child_module_names]
+
     def decorator(forward_func: Callable) -> Callable:
         def wrapped(self, *args, **kwargs):
-            hooked_module = getattr(self, child_module_name)
-            hook = getattr(hooked_module, "_hf_hook", None)
-            if hook is not None:
-                # Note that here we only call the hook with the module, not `*args` not `**kwargs`, as we assume the `forward`
-                # on which this decorator is applied is responsible to move the args and kwargs with its own hook if any. This makes
-                # sense as the module decorated with this should have all internal modules on the same device
-                hook.pre_forward(hooked_module)
+            hooked_modules = []
+            for child_module_name in child_module_names:
+                hooked_module = getattr(self, child_module_name)
+                hook = getattr(hooked_module, "_hf_hook", None)
+                hooked_modules.append((hooked_module, hook))
+                if hook is not None:
+                    # Note that here we only call the hook with the module, not `*args` not `**kwargs`, as we assume the `forward`
+                    # on which this decorator is applied is responsible to move the args and kwargs with its own hook if any. This makes
+                    # sense as the module decorated with this should have all internal modules on the same device
+                    hook.pre_forward(hooked_module)
 
             output = forward_func(self, *args, **kwargs)
 
-            if hook is not None:
-                # Note that here we only call the hook with the module, not `output`, as we assume the `forward` on which
-                # this decorator is applied is responsible to move the output with its own hook if any. This makes sense
-                # as the module decorated with this should have all internal modules on the same device
-                hook.post_forward(hooked_module, ())
+            for hooked_module, hook in reversed(hooked_modules):
+                if hook is not None:
+                    # Note that here we only call the hook with the module, not `output`, as we assume the `forward` on which
+                    # this decorator is applied is responsible to move the output with its own hook if any. This makes sense
+                    # as the module decorated with this should have all internal modules on the same device
+                    hook.post_forward(hooked_module, ())
 
             return output
 

src/transformers/integrations/hub_kernels.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import functools
 import importlib
+import inspect
 import os
 import re
 import sys
@@ -32,6 +33,7 @@
     is_kernels_available,
     is_rocm_platform,
     is_torch_available,
+    resolve_internal_import,
 )
 from .flash_attention import flash_attention_forward
 
@@ -59,6 +61,18 @@
 _kernels_enabled = _TRANSFORMERS_USE_HUB_KERNELS in ENV_VARS_TRUE_VALUES
 
 
+# Maps from func name to the internal module path
+_KERNELS_INTERNAL_PATH_MAPPINGS = {
+    "chunk_gated_delta_rule": "ops.gated_delta_rule",
+    "recurrent_gated_delta_rule": "ops.gated_delta_rule",
+    "mamba_split_conv1d_scan_combined": "ops.triton.ssd_combined",
+    "selective_state_update": "ops.triton.selective_state_update",
+    "mamba_chunk_scan_combined": "ops.triton.ssd_combined",
+    "mamba_inner_fn": "ops.selective_scan_interface",
+    "selective_scan_fn": "ops.selective_scan_interface",
+}
+
+
 if is_kernels_available():
     from kernels import (
         CUDAProperties,
@@ -154,12 +168,12 @@ def _build_kernel_mapping() -> dict:
                     Mode.TRAINING: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="causal_conv1d_fn",
-                        version=1,
+                        version=2,
                     ),
                     Mode.INFERENCE: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="causal_conv1d_fn",
-                        version=1,
+                        version=2,
                     ),
                 },
             },
@@ -168,15 +182,113 @@ def _build_kernel_mapping() -> dict:
                     Mode.TRAINING: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="causal_conv1d_update",
-                        version=1,
+                        version=2,
                     ),
                     Mode.INFERENCE: LayerRepository(
                         repo_id="kernels-community/mamba-ssm",
                         layer_name="causal_conv1d_update",
+                        version=2,
+                    ),
+                },
+            },
+            "chunk_gated_delta_rule": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="chunk_gated_delta_rule",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="chunk_gated_delta_rule",
+                        version=1,
+                    ),
+                },
+            },
+            "recurrent_gated_delta_rule": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="recurrent_gated_delta_rule",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="recurrent_gated_delta_rule",
                         version=1,
                     ),
                 },
             },
+            "mamba_chunk_scan_combined": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_chunk_scan_combined",
+                        version=2,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_chunk_scan_combined",
+                        version=2,
+                    ),
+                },
+            },
+            "mamba_split_conv1d_scan_combined": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_split_conv1d_scan_combined",
+                        version=2,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_split_conv1d_scan_combined",
+                        version=2,
+                    ),
+                },
+            },
+            "mamba_inner_fn": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_inner_fn",
+                        version=2,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="mamba_inner_fn",
+                        version=2,
+                    ),
+                },
+            },
+            "selective_scan_fn": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="selective_scan_fn",
+                        version=2,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="selective_scan_fn",
+                        version=2,
+                    ),
+                },
+            },
+            "selective_state_update": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="selective_state_update",
+                        version=2,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/mamba-ssm",
+                        layer_name="selective_state_update",
+                        version=2,
+                    ),
+                },
+            },
             "SwiGLUMLP": {
                 "cuda": {
                     Mode.INFERENCE | Mode.TORCH_COMPILE: LayerRepository(
@@ -267,6 +379,20 @@ def _build_kernel_mapping() -> dict:
                     ),
                 },
             },
+            "RMSNormGated": {
+                "cuda": {
+                    Mode.TRAINING: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="FusedRMSNormGated",
+                        version=1,
+                    ),
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/fla",
+                        layer_name="FusedRMSNormGated",
+                        version=1,
+                    ),
+                },
+            },
             "MegaBlocksMoeMLP": {
                 "cuda": {
                     Mode.TRAINING: LayerRepository(
@@ -451,9 +577,6 @@ def register_kernel_mapping_transformers(*args, **kwargs):
 
 
 _HUB_KERNEL_MAPPING: dict[str, dict[str, str]] = {
-    "causal-conv1d": {"repo_id": "kernels-community/causal-conv1d", "version": 1},
-    "mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
-    "falcon_mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
     "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4},
     "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 2},
     "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"},
@@ -638,6 +761,46 @@ def get_kernel(
     )
 
 
+def use_kernel_func_from_hub_with_fallback(func_name: str, package: str, internal_path: str | None = None):
+    """
+    The same as `use_kernel_forward_from_hub` but with the optional fallback to an original package if it exists, e.g.,
+    FLA for Gated Delta Rule, mamba-ssm for mamba2, etc.
+
+    This combines all options with kernels, enabling kernels on top of the original package if requested as well.
+    The order of priority is
+        1. Hf kernels (if requested)
+        2. Original package
+        3. Torch only path
+    """
+    kernel_wrapper_decorator = use_kernel_forward_from_hub(func_name)
+
+    # Allow internal path prefix if given to resolve non __init__ imports
+    internal_path = _KERNELS_INTERNAL_PATH_MAPPINGS.get(func_name, internal_path)  # defaults
+    full_path = func_name if internal_path is None else f"{internal_path}.{func_name}"
+
+    def decorator(torch_function: Callable) -> Callable:
+        implementation = None
+        try:
+            module = importlib.import_module(package)
+            implementation = resolve_internal_import(module, full_path)
+        except Exception:
+            implementation = torch_function
+        finally:
+            implementation = torch_function if implementation is None else implementation
+
+        # Make it "frozen" like to let dynamo not try to look into any ordering
+        applicable_params = tuple(inspect.signature(implementation).parameters)
+
+        @functools.wraps(torch_function)
+        def wrapped(*args, **kwargs):
+            kwargs = {k: v for k, v in kwargs.items() if k in applicable_params}
+            return implementation(*args, **kwargs)
+
+        return kernel_wrapper_decorator(wrapped)
+
+    return decorator
+
+
 # Whether to allow hub kernels coming from untrusted repos, i.e. repos outside `kernels-community`
 ALLOW_ALL_KERNELS = False