From a22e5bcae5f36b332fce4e28257d69660a86e4e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Fri, 28 Aug 2026 11:42:10 +0000 Subject: [PATCH 1/3] kernels: support conditions on `use_kernel_forward_from_hub` Add and optional `condition` argument to `use_kernel_forward_from_hub`. This condition accepts an instance of the layer and returns whether kernelization of the layer should proceed. For instance, ``` @use_kernel_forward_from_hub( "SwiGLUMLP", condition=lambda module: module.config.hidden_act == "silu", ) class MyMLP(nn.Module): ... ``` would only kernelize the layer using a mapping registered for `SwiGLUMLP` its `hidden_act` configuration is `silu`. --- kernels/src/kernels/layer/kernelize.py | 12 ++++ kernels/src/kernels/layer/layer.py | 28 ++++++--- kernels/tests/test_layer.py | 78 +++++++++++++++++++++++++- 3 files changed, 110 insertions(+), 8 deletions(-) diff --git a/kernels/src/kernels/layer/kernelize.py b/kernels/src/kernels/layer/kernelize.py index 60d221d6a..4c4630490 100644 --- a/kernels/src/kernels/layer/kernelize.py +++ b/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 diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index 2e7aae205..74e205b6d 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/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") diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py index 699c4a119..a886ad6fd 100644 --- a/kernels/tests/test_layer.py +++ b/kernels/tests/test_layer.py @@ -423,6 +423,78 @@ class SiluAndMulWithKernelFallback(SiluAndMul): kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE) +def test_kernel_condition_skips_kernelization(): + @use_kernel_forward_from_hub("SiluAndMulNonExisting", condition=lambda module: False) + class SiluAndMulConditionSkipped(SiluAndMul): + pass + + silu_and_mul = SiluAndMulConditionSkipped() + + kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE, use_fallback=False) + + # 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 +673,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, From 8b006b67e5c601e839e3528ffcc5e91208a739b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 1 Sep 2026 09:30:21 +0000 Subject: [PATCH 2/3] Document the `condition` argument of `use_kernel_layer_forward_from_hub` --- docs/source/layers.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/docs/source/layers.md b/docs/source/layers.md index e11120417..c19c78993 100644 --- a/docs/source/layers.md +++ b/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`] From cdee0896f17c0032c541a9711011396b3d6f2a8f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Thu, 3 Sep 2026 06:58:49 +0000 Subject: [PATCH 3/3] Check log as well when condition fails --- kernels/src/kernels/layer/kernelize.py | 2 +- kernels/tests/test_layer.py | 11 +++++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/kernels/src/kernels/layer/kernelize.py b/kernels/src/kernels/layer/kernelize.py index 4c4630490..11188236c 100644 --- a/kernels/src/kernels/layer/kernelize.py +++ b/kernels/src/kernels/layer/kernelize.py @@ -275,7 +275,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: 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.", + "Skipping kernelization for `%s` using `%s` due to kernel_condition.", module_class.__name__, module_class.kernel_layer_name, ) diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py index a886ad6fd..4f28dc567 100644 --- a/kernels/tests/test_layer.py +++ b/kernels/tests/test_layer.py @@ -1,3 +1,4 @@ +import logging import sys from contextlib import nullcontext @@ -423,14 +424,20 @@ class SiluAndMulWithKernelFallback(SiluAndMul): kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE) -def test_kernel_condition_skips_kernelization(): +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() - kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE, use_fallback=False) + 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__