[Kernels] Refactor function handling

examples/modular-transformers/modeling_add_function.py

@@ -9,7 +9,7 @@
 import torch
 from torch import nn
 
-from ...integrations import use_kernel_func_from_hub
+from ...integrations import use_kernel_forward_from_hub
 
 
 def rotate_half(x):
@@ -19,7 +19,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
+@use_kernel_forward_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 

examples/modular-transformers/modeling_global_indexing.py

@@ -12,7 +12,7 @@
 from transformers.modeling_utils import AttentionInterface
 
 from ...cache_utils import Cache
-from ...integrations import use_kernel_func_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs
 from .configuration_global_indexing import GlobalIndexingConfig
@@ -25,7 +25,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
+@use_kernel_forward_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 

examples/modular-transformers/modeling_my_new_model2.py

@@ -12,7 +12,7 @@
 from ... import initialization as init
 from ...activations import ACT2FN
 from ...cache_utils import Cache
-from ...integrations import use_kernel_func_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
 from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
@@ -77,7 +77,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
+@use_kernel_forward_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 

examples/modular-transformers/modeling_super.py

@@ -14,7 +14,7 @@
 
 from ...activations import ACT2FN
 from ...cache_utils import Cache
-from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
@@ -134,7 +134,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
+@use_kernel_forward_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 

examples/modular-transformers/modeling_switch_function.py

@@ -11,7 +11,7 @@
 from torch import nn
 
 from ...cache_utils import Cache
-from ...integrations import use_kernel_func_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
 from ...processing_utils import Unpack
 from ...utils import TransformersKwargs
@@ -26,7 +26,7 @@ def rotate_half(x):
     return rot_x
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
+@use_kernel_forward_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 

examples/modular-transformers/modeling_test_suffix.py

@@ -11,7 +11,7 @@
 
 from ...activations import ACT2FN
 from ...cache_utils import Cache
-from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
 from ...processing_utils import Unpack
@@ -67,7 +67,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
+@use_kernel_forward_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 

src/transformers/integrations/hub_kernels.py

@@ -75,7 +75,7 @@
     from kernels import (
         use_kernel_forward_from_hub as _kernels_use_kernel_forward_from_hub,
     )
-    from kernels import use_kernel_func_from_hub as _kernels_use_kernel_func_from_hub
+    from kernels import use_kernelized_func as _kernels_use_kernelized_func
 
     _TRANSFORMERS_USE_HUB_KERNELS = os.environ.get("USE_HUB_KERNELS", "YES").upper()
     _kernels_enabled = _TRANSFORMERS_USE_HUB_KERNELS in ENV_VARS_TRUE_VALUES
@@ -89,14 +89,26 @@ def use_kernel_forward_from_hub(layer_name: str):
             )
             return lambda cls: cls
 
-    def use_kernel_func_from_hub(func_name: str):
+    def use_kernelized_func(module_names: Callable):
         if _kernels_enabled:
-            return _kernels_use_kernel_func_from_hub(func_name)
+            return _kernels_use_kernelized_func(module_names)
         else:
             logger.warning_once(
                 f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}"
             )
-            return lambda func: func
+            return lambda cls: cls
+
+    def use_kernel_func_from_hub(layer_name: str):
+        if _kernels_enabled:
+            logger.warning_once(
+                "`use_kernel_func_from_hub` and will be deprecated in transformers v5.16. Please use `use_kernel_forward_from_hub` instead."
+            )
+            return _kernels_use_kernel_forward_from_hub(layer_name)
+        else:
+            logger.warning_once(
+                f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}"
+            )
+            return lambda cls: cls
 
     # The default kernel mapping is built lazily (see `get_kernel_mapping_transformers`) so that simply
     # importing transformers (or `transformers.pipeline`) does not instantiate any `LayerRepository` /
@@ -312,34 +324,29 @@ def _build_kernel_mapping() -> dict:
                     )
                 }
             },
-        }
-
-        # Add function kernel mappings
-        _FUNCTION_KERNEL_MAPPING = {
             "rotary_pos_emb": {
                 "xpu": {
-                    Mode.INFERENCE: FuncRepository(
-                        repo_id="kernels-community/rotary", func_name="apply_rotary_transformers", version=1
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/rotary", layer_name="apply_rotary_transformers", version=1
                     )
                 },
-                "cuda": FuncRepository(
-                    repo_id="kernels-community/rotary", func_name="apply_rotary_transformers", version=1
+                "cuda": LayerRepository(
+                    repo_id="kernels-community/rotary", layer_name="apply_rotary_transformers", version=1
                 ),
                 "rocm": {
-                    Mode.INFERENCE: FuncRepository(
-                        repo_id="kernels-community/aiter-rope", func_name="apply_rotary_transformers", version=2
+                    Mode.INFERENCE: LayerRepository(
+                        repo_id="kernels-community/aiter-rope", layer_name="apply_rotary_transformers", version=2
                     )
                 },
             },
             "ForCausalLMLoss": {
                 "cuda": {
-                    Mode.TRAINING | Mode.TORCH_COMPILE: FuncRepository(
-                        repo_id="kernels-community/liger-kernels", func_name="LigerForCausalLMLoss", version=2
+                    Mode.TRAINING | Mode.TORCH_COMPILE: LayerRepository(
+                        repo_id="kernels-community/liger-kernels", layer_name="LigerForCausalLMLossLayer", version=2
                     ),
                 },
             },
         }
-        _KERNEL_MAPPING = _KERNEL_MAPPING | _FUNCTION_KERNEL_MAPPING
 
         return _KERNEL_MAPPING
 
@@ -366,9 +373,15 @@ def decorator(cls):
 
         return decorator
 
+    def use_kernelized_func(*args, **kwargs):
+        def decorator(cls):
+            return cls
+
+        return decorator
+
     def use_kernel_func_from_hub(*args, **kwargs):
-        def decorator(func):
-            return func
+        def decorator(cls):
+            return cls
 
         return decorator
 
@@ -553,42 +566,22 @@ def kernelize(model: "PreTrainedModel", mode: "Mode | None" = None):
     if not is_kernels_available():
         raise ImportError(_MISSING_KERNELS_MESSAGE)
 
-    def attach_hidden_kernels(module):
-        for name, fn in getattr(module, "_hidden_kernels", {}).items():
-            if name not in dict(module.named_children()):
-                if not isinstance(fn, nn.Module):
-                    raise ValueError(
-                        f"Attempted to register a kernel for {name}, but it was not a `torch.nn.Module`. "
-                        "This means the underlying function needs to be decorated with `@use_kernel_func_from_hub`. "
-                        "Please submit and issue to the transformers repo: `https://github.com/huggingface/transformers/issues`."
-                    )
-                module.register_module(name, fn)
-
-    def detach_hidden_kernels(module):
-        for name in getattr(module, "_hidden_kernels", {}):
-            # Skip deregistering if it failed to properly register,
-            # i.e. `ValueError` will be raised afterwards
-            if hasattr(module, name):
-                delattr(module, name)
-
-    try:
-        model.apply(attach_hidden_kernels)
-
-        mode = Mode.INFERENCE if not model.training else Mode.TRAINING if mode is None else mode
-        device_type = model.device.type
+    def get_device(device_type):
         if device_type == "cuda" and is_rocm_platform():
             device_type = "rocm"
-        device = Device(type=device_type)
-        if model.kernel_config is not None:
-            inherit_mapping = not model.kernel_config.use_local_kernel
-            with use_kernel_mapping(model.kernel_config.kernel_mapping, inherit_mapping=inherit_mapping):
-                _kernels_kernelize(model, device=device, mode=mode)
-        else:
+        return Device(type=device_type)
+
+    mode = Mode.INFERENCE if not model.training else Mode.TRAINING if mode is None else mode
+    device = get_device(model.device.type)
+
+    if model.kernel_config is not None:
+        inherit_mapping = not model.kernel_config.use_local_kernel
+        with use_kernel_mapping(model.kernel_config.kernel_mapping, inherit_mapping=inherit_mapping):
             _kernels_kernelize(model, device=device, mode=mode)
+    else:
+        _kernels_kernelize(model, device=device, mode=mode)
 
-        model._use_kernels = True
-    finally:
-        model.apply(detach_hidden_kernels)
+    model._use_kernels = True
 
 
 def get_kernel(
@@ -608,45 +601,6 @@ def get_kernel(
     )
 
 
-def use_kernelized_func(module_names: list[Callable] | Callable):
-    """
-    This decorator attaches the target function within the module as a plain attribute (not as a submodule).
-    Keep in mind that this registration is only meant for `kernelize` to recognize its target modules (i.e.
-    function exchanged for a weightless `nn.Module` with the same forward) to then exchange to the kernel
-    variation (in-place) if the conditions are met.
-
-    We cache each of these function-based registrations: After proper registration and exchange it is removed
-    from the module's `_modules` dict as it does not really act as `nn.Module` but a base function.
-    """
-    if isinstance(module_names, Callable):
-        module_names = [module_names]
-
-    def decorator(cls):
-        orig_init = cls.__init__
-
-        def new_init(self, *args, **kwargs):
-            orig_init(self, *args, **kwargs)
-
-            # Register new function as non-submodule within the modules dict
-            hidden_kernels = self.__dict__.setdefault("_hidden_kernels", {})
-            for fn in module_names:
-                name = (
-                    getattr(fn, "__name__", None)
-                    or getattr(fn, "kernel_layer_name", None)
-                    or getattr(fn, "func_name", None)
-                )
-                if name is None:
-                    raise ValueError(f"Could not infer kernel function name for {fn!r}")
-
-                # Do not register as submodule! Hide it behind a dict to be removed later after registering it
-                hidden_kernels[name] = fn
-
-        cls.__init__ = new_init
-        return cls
-
-    return decorator
-
-
 # Whether to allow hub kernels coming from untrusted repos, i.e. repos outside `kernels-community`
 ALLOW_ALL_KERNELS = False
 

src/transformers/models/afmoe/modeling_afmoe.py

@@ -28,12 +28,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import (
-    use_experts_implementation,
-    use_kernel_forward_from_hub,
-    use_kernel_func_from_hub,
-    use_kernelized_func,
-)
+from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernelized_func
 from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
@@ -256,7 +251,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
+@use_kernel_forward_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 

src/transformers/models/apertus/modeling_apertus.py

@@ -27,7 +27,7 @@
 from ...activations import ACT2CLS, ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
 from ...masking_utils import create_causal_mask
 from ...modeling_layers import GenericForTokenClassification, GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
@@ -149,7 +149,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
+@use_kernel_forward_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 

src/transformers/models/arcee/modeling_arcee.py

@@ -29,7 +29,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
 from ...masking_utils import create_causal_mask
 from ...modeling_layers import (
     GenericForQuestionAnswering,
@@ -154,7 +154,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
+@use_kernel_forward_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.
 

src/transformers/models/aria/modeling_aria.py

@@ -29,7 +29,7 @@
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache
 from ...generation import GenerationMixin
-from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
 from ...masking_utils import create_causal_mask
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import GradientCheckpointingLayer
@@ -384,7 +384,7 @@ def rotate_half(x):
     return torch.cat((-x2, x1), dim=-1)
 
 
-@use_kernel_func_from_hub("rotary_pos_emb")
+@use_kernel_forward_from_hub("rotary_pos_emb")
 def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
     """Applies Rotary Position Embedding to the query and key tensors.