kernels: support conditions on use_kernel_forward_from_hub

docs/source/layers.md

@@ -27,6 +27,23 @@ The decorator does not change the behavior of the class -- it annotates
 the class with the given name (here `SiluAndMul`). The [`~kernels.kernelize`] function
 described below uses this name to look up kernels for the layer.
 
+Sometimes you only want to kernelize a layer depending on some state inside
+that layer. For instance, an MLP layer could support multiple activations, but
+a kernel that you want to register only supports one particular activation. In
+such cases, you can add a condition to a `use_kernel_forward_from_hub` decorator.
+The layer will then only be kernelized when the condition holds. The condition
+must be a callable that takes the instantiated layer and returns a `bool`. For
+example:
+
+```python
+@use_kernel_forward_from_hub(
+    "SwiGLUMLP",
+    condition=lambda module: module.config.hidden_act == "silu",
+)
+class MyMLP(nn.Module):
+    ...
+```
+
 ### External layers
 
 An existing layer that does not (yet) have the [`~kernels.use_kernel_forward_from_hub`]

kernels/src/kernels/layer/kernelize.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from copy import deepcopy
 from typing import TYPE_CHECKING
 
@@ -13,6 +14,8 @@
     import torch
     from torch import nn
 
+logger = logging.getLogger(__name__)
+
 
 def use_kernel_mapping(
     mapping: dict[
@@ -269,6 +272,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             )
 
         if hasattr(module_class, "kernel_layer_name"):
+            cond = getattr(module_class, "kernel_condition", None)
+            if cond and not cond(module):
+                logger.info(
+                    "Skipping kernelization for `%s` using `%s` due to kernel_condition.",
+                    module_class.__name__,
+                    module_class.kernel_layer_name,
+                )
+                continue
+
             kernelize_layer(module, mode=mode, device_type=device_type, use_fallback=use_fallback)
 
     return model

kernels/src/kernels/layer/layer.py

@@ -272,10 +272,7 @@ def __str__(self) -> str:
 _CACHED_LAYER: dict[RepositoryProtocol, Type["nn.Module"]] = {}
 
 
-def replace_kernel_forward_from_hub(
-    cls,
-    layer_name: str,
-):
+def replace_kernel_forward_from_hub(cls, layer_name: str, condition: Callable[["nn.Module"], bool] | None = None):
     """
     Function that prepares a layer class to use kernels from the Hugging Face Hub.
 
@@ -284,6 +281,15 @@ def replace_kernel_forward_from_hub(
     it is inherently fragile since the member variables and `forward` signature
     of such a layer can change.
 
+    Args:
+        layer_name (`str`):
+            The name of the layer to use for kernel lookup in registered mappings.
+        condition (`Callable[["nn.Module"], bool]`, *optional*):
+            Additional condition that is checked during kernelization. The callable
+            is passed the module instance and is evaluated for every instance of
+            the layer when [`~kernels.kernelize`] is called. If it returns `False`,
+            kernelization is skipped for that instance.
+
     Example:
         ```python
         from kernels import replace_kernel_forward_from_hub
@@ -293,9 +299,12 @@ def replace_kernel_forward_from_hub(
         ```
     """
     cls.kernel_layer_name = layer_name
+    # Wrap in `staticmethod` so that access through an instance does not bind
+    # it as a method (the condition takes the module as its only argument).
+    cls.kernel_condition = staticmethod(condition if condition is not None else lambda module: True)
 
 
-def use_kernel_forward_from_hub(layer_name: str):
+def use_kernel_forward_from_hub(layer_name: str, condition: Callable[["nn.Module"], bool] | None = None):
     """
     Decorator factory that makes a layer extensible using the specified layer name.
 
@@ -311,6 +320,11 @@ def use_kernel_forward_from_hub(layer_name: str):
     Args:
         layer_name (`str`):
             The name of the layer to use for kernel lookup in registered mappings.
+        condition (`Callable[["nn.Module"], bool]`, *optional*):
+            Additional condition that is checked during kernelization. The callable
+            is passed the module instance and is evaluated for every instance of
+            the layer when [`~kernels.kernelize`] is called. If it returns `False`,
+            kernelization is skipped for that instance.
 
     Returns:
         `Callable`: A decorator function that can be applied to layer classes.
@@ -356,10 +370,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     def decorator(ty):
         if inspect.isfunction(ty):
             Func = _create_func_module(ty)
-            replace_kernel_forward_from_hub(Func, layer_name)
+            replace_kernel_forward_from_hub(Func, layer_name, condition)
             return Func()
         elif inspect.isclass(ty):
-            replace_kernel_forward_from_hub(ty, layer_name)
+            replace_kernel_forward_from_hub(ty, layer_name, condition)
             return ty
         else:
             raise TypeError("@use_kernel_forward_from_hub can only be applied to classes or functions")

kernels/tests/test_layer.py

@@ -1,3 +1,4 @@
+import logging
 import sys
 from contextlib import nullcontext
 
@@ -423,6 +424,84 @@ class SiluAndMulWithKernelFallback(SiluAndMul):
     kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE)
 
 
+def test_kernel_condition_skips_kernelization(caplog):
+    @use_kernel_forward_from_hub("SiluAndMulNonExisting", condition=lambda module: False)
+    class SiluAndMulConditionSkipped(SiluAndMul):
+        pass
+
+    silu_and_mul = SiluAndMulConditionSkipped()
+
+    with caplog.at_level(logging.INFO, logger="kernels.layer.kernelize"):
+        kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE, use_fallback=False)
+
+    assert (
+        "Skipping kernelization for `SiluAndMulConditionSkipped` using `SiluAndMulNonExisting` due to kernel_condition."
+        in caplog.text
+    )
+
+    # The forward was not replaced...
+    assert "forward" not in silu_and_mul.__dict__
+
+    # ... and the original implementation is still used.
+    X = torch.randn(32, 64)
+    silu_and_mul(X)
+    assert silu_and_mul.n_calls == 1
+
+
+def test_kernel_condition_holds():
+    @use_kernel_forward_from_hub("SiluAndMulNonExisting", condition=lambda module: True)
+    class SiluAndMulConditionHolds(SiluAndMul):
+        pass
+
+    silu_and_mul = SiluAndMulConditionHolds()
+
+    # The condition holds, so kernelization is attempted and fails because
+    # there is no kernel mapping for the layer.
+    with pytest.raises(ValueError, match="No layer mapping for `SiluAndMulNonExisting`"):
+        kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE, use_fallback=False)
+
+
+def test_kernel_condition_receives_module_instance():
+    received = []
+
+    @use_kernel_forward_from_hub(
+        "SiluAndMulNonExisting",
+        condition=lambda module: received.append(module) or False,
+    )
+    class SiluAndMulConditionArg(SiluAndMul):
+        pass
+
+    silu_and_mul = SiluAndMulConditionArg()
+    # No mapping, but uses fallback.
+    kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE)
+
+    assert received == [silu_and_mul]
+
+
+def test_kernel_condition_per_instance():
+    @use_kernel_forward_from_hub("SiluAndMulNonExisting", condition=lambda module: module.allow_kernel)
+    class SiluAndMulConditional(SiluAndMul):
+        def __init__(self, allow_kernel: bool):
+            super().__init__()
+            self.allow_kernel = allow_kernel
+
+    skipped = SiluAndMulConditional(allow_kernel=False)
+    attempted = SiluAndMulConditional(allow_kernel=True)
+
+    # Only the instance for which the condition holds is kernelized.
+    kernelize(skipped, device="cuda", mode=Mode.INFERENCE, use_fallback=False)
+    with pytest.raises(ValueError, match="No layer mapping for `SiluAndMulNonExisting`"):
+        kernelize(attempted, device="cuda", mode=Mode.INFERENCE, use_fallback=False)
+
+
+def test_kernel_condition_defaults_to_true():
+    # Without an explicit condition, kernelization is never skipped. Accessing
+    # the condition through an instance must not bind it as a method.
+    silu_and_mul = SiluAndMulWithKernel()
+    assert SiluAndMulWithKernel.kernel_condition(silu_and_mul)
+    assert silu_and_mul.kernel_condition(silu_and_mul)
+
+
 def test_local_layer_repo(device):
     # Fetch a kernel to the local cache.
     path = install_kernel("kernels-test/backward-marker-test", revision="main")
@@ -601,7 +680,11 @@ def __init__(self, *args, **kwargs):
             self.foo = 42
 
     def stub_repo(layer):
-        return LayerRepository(repo_id="kernels-test/nonexisting", layer_name=layer.__name__, revision="main")
+        return LayerRepository(
+            repo_id="kernels-test/nonexisting",
+            layer_name=layer.__name__,
+            revision="main",
+        )
 
     with pytest.raises(
         TypeError,