From 5c218b314b3b2732971801fe9bfc2cfe31f53c85 Mon Sep 17 00:00:00 2001
From: Marc Sun <marc@huggingface.co>
Date: Thu, 24 Sep 2026 15:29:14 +0000
Subject: [PATCH 1/4] fix fetch build fallback

---
 kernels/src/kernels/layer/layer.py |  8 +++++++-
 kernels/tests/test_layer.py        | 22 ++++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py
index c02d5946..35582bc4 100644
--- a/kernels/src/kernels/layer/layer.py
+++ b/kernels/src/kernels/layer/layer.py
@@ -517,7 +517,13 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
     logging.info(f"Using function/layer from repo {repo}")
     logging.debug(f"kernelize mode: {mode}, repo mode: {repo_mode}")
 
-    layer = _get_layer_memoize(repo, module_class)
+    try:
+        layer = _get_layer_memoize(repo, module_class)
+    except FileNotFoundError:
+        if not use_fallback:
+            raise
+        _replace_forward(module, module_class)
+        return
 
     # Ideally we would do validation on the mapping where we check that
     # e.g. if a repo class is registered for TRAINING | TORCH_COMPILE,
diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py
index 41186af6..1d6f8805 100644
--- a/kernels/tests/test_layer.py
+++ b/kernels/tests/test_layer.py
@@ -425,6 +425,28 @@ class SiluAndMulWithKernelFallback(SiluAndMul):
     kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE)
 
 
+def test_missing_kernel_build_falls_back(monkeypatch):
+    repo = LayerRepository(
+        repo_id="kernels-test/no-compatible-build",
+        layer_name="SiluAndMul",
+        version=1,
+    )
+
+    def unavailable_build():
+        raise FileNotFoundError("Cannot find a build variant for this system")
+
+    monkeypatch.setattr(repo, "load", unavailable_build)
+    layer = SiluAndMulWithKernel()
+    mapping = {"SiluAndMul": {"cuda": repo}}
+
+    with use_kernel_mapping(mapping, inherit_mapping=False):
+        kernelize(layer, device="cuda", mode=Mode.INFERENCE)
+
+    with use_kernel_mapping(mapping, inherit_mapping=False):
+        with pytest.raises(FileNotFoundError, match="Cannot find a build variant"):
+            kernelize(layer, device="cuda", mode=Mode.INFERENCE, use_fallback=False)
+
+
 def test_kernel_condition_skips_kernelization(caplog):
     @use_kernel_forward_from_hub("SiluAndMulNonExisting", condition=lambda module: False)
     class SiluAndMulConditionSkipped(SiluAndMul):

From ec2f04480546aee951af077742a1ec6a15379f7d Mon Sep 17 00:00:00 2001
From: Marc Sun <marc@huggingface.co>
Date: Thu, 24 Sep 2026 16:08:25 +0000
Subject: [PATCH 2/4] fallback but butter

---
 kernels/src/kernels/__init__.py        |  2 ++
 kernels/src/kernels/layer/__init__.py  |  2 ++
 kernels/src/kernels/layer/kernelize.py | 14 +++++---
 kernels/src/kernels/layer/layer.py     | 46 +++++++++++++++++++++-----
 kernels/tests/test_layer.py            | 13 ++++++--
 5 files changed, 61 insertions(+), 16 deletions(-)

diff --git a/kernels/src/kernels/__init__.py b/kernels/src/kernels/__init__.py
index 7d9bc7d0..5c301ba2 100644
--- a/kernels/src/kernels/__init__.py
+++ b/kernels/src/kernels/__init__.py
@@ -12,6 +12,7 @@
 from kernels.layer import (
     CUDAProperties,
     Device,
+    Fallback,
     FuncRepository,
     LayerRepository,
     LocalFuncRepository,
@@ -48,6 +49,7 @@
     "Benchmark",
     "CUDAProperties",
     "Device",
+    "Fallback",
     "ROCMProperties",
     "FuncRepository",
     "LayerRepository",
diff --git a/kernels/src/kernels/layer/__init__.py b/kernels/src/kernels/layer/__init__.py
index 217e3d0f..8ca217c7 100644
--- a/kernels/src/kernels/layer/__init__.py
+++ b/kernels/src/kernels/layer/__init__.py
@@ -11,6 +11,7 @@
     use_kernel_mapping,
 )
 from .layer import (
+    Fallback,
     LayerRepository,
     LocalLayerRepository,
     LockedLayerRepository,
@@ -23,6 +24,7 @@
 __all__ = [
     "CUDAProperties",
     "Device",
+    "Fallback",
     "ROCMProperties",
     "FuncRepository",
     "LayerRepository",
diff --git a/kernels/src/kernels/layer/kernelize.py b/kernels/src/kernels/layer/kernelize.py
index 11188236..c6a23b1c 100644
--- a/kernels/src/kernels/layer/kernelize.py
+++ b/kernels/src/kernels/layer/kernelize.py
@@ -6,7 +6,7 @@
 
 from .device import Device
 from .globals import _KERNEL_MAPPING
-from .layer import kernelize_layer
+from .layer import Fallback, kernelize_layer
 from .mode import Mode
 from .repos import DeviceRepos, RepositoryProtocol
 
@@ -180,7 +180,7 @@ def kernelize(
     *,
     mode: Mode,
     device: str | "torch.device" | None = None,
-    use_fallback: bool = True,
+    use_fallback: bool | Fallback = True,
 ):
     """
     Replace layer forward methods with optimized kernel implementations.
@@ -197,9 +197,10 @@ def kernelize(
         device (`Union[str, torch.device]`, *optional*):
             The device type to load kernels for. Supported device types are: "cuda", "mps", "npu", "rocm", "tpu", "xpu".
             The device type will be inferred from the model parameters when not provided.
-        use_fallback (`bool`, *optional*, defaults to `True`):
-            Whether to use the original forward method of modules when no compatible kernel could be found.
-            If set to `False`, an exception will be raised in such cases.
+        use_fallback (`bool | Fallback`, *optional*, defaults to `True`):
+            Cases in which to use the original forward method. `True` is equivalent to `Fallback.DEFAULT`: it allows
+            fallback when no compatible mapping or mode exists. `False` is equivalent to `Fallback.NONE` and raises
+            instead. `Fallback.ALL` also falls back when a selected kernel repository has no compatible build.
 
     Returns:
         `nn.Module`: The kernelized model with optimized kernel implementations.
@@ -259,6 +260,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
     assert isinstance(device_type, Device)
 
+    if isinstance(use_fallback, bool):
+        use_fallback = Fallback.DEFAULT if use_fallback else Fallback.NONE
+
     for _, module in model.named_modules():
         module_class = type(module)
 
diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py
index 35582bc4..05bc7873 100644
--- a/kernels/src/kernels/layer/layer.py
+++ b/kernels/src/kernels/layer/layer.py
@@ -2,6 +2,7 @@
 
 import inspect
 import logging
+from enum import Flag, auto
 from inspect import Parameter, Signature
 from pathlib import Path
 from types import MethodType, ModuleType
@@ -34,6 +35,30 @@
 logger = logging.getLogger(__name__)
 
 
+class Fallback(Flag):
+    """Cases in which kernelization may keep the original layer forward.
+
+    - `NONE`: Raise instead of falling back.
+    - `NO_LAYER`: No kernel mapping exists for the layer.
+    - `NO_DEVICE`: The layer has no mapping for the requested device type.
+    - `NO_COMPATIBLE_PROPERTIES`: No mapping matches the device properties.
+    - `NO_COMPATIBLE_MODE`: No repository or loaded kernel supports the requested mode.
+    - `COMPATIBLE_BUILD_NOT_FOUND`: The selected repository has no build for this environment.
+    - `DEFAULT`: Fall back for missing mappings, properties, or mode; raise for a missing build.
+    - `ALL`: Fall back in every case above, including a missing build.
+    """
+
+    NONE = 0
+    NO_LAYER = auto()
+    NO_DEVICE = auto()
+    NO_COMPATIBLE_PROPERTIES = auto()
+    NO_COMPATIBLE_MODE = auto()
+    COMPATIBLE_BUILD_NOT_FOUND = auto()
+
+    DEFAULT = NO_LAYER | NO_DEVICE | NO_COMPATIBLE_PROPERTIES | NO_COMPATIBLE_MODE
+    ALL = DEFAULT | COMPATIBLE_BUILD_NOT_FOUND
+
+
 class LayerRepositoryProtocol(RepositoryProtocol, Protocol):
     @property
     def layer_name(self) -> str: ...
@@ -462,7 +487,7 @@ def new_init(self, *args, **kwargs):
     return decorator
 
 
-def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use_fallback):
+def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use_fallback: Fallback):
     module_class = type(module)
     layer_name = module_class.kernel_layer_name  # type: ignore[attr-defined]
 
@@ -479,7 +504,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
             f"Check if the layer name matches one of the kernels in the mapping or add the kernel "
             f"you want to use to the mapping. Defaulting to original forward implementation."
         )
-        if not use_fallback:
+        if Fallback.NO_LAYER not in use_fallback:
             raise ValueError(f"No layer mapping for `{layer_name}`")
         _replace_forward(module, module_class)
         return
@@ -488,7 +513,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
     property_repos = kernel.get(device_type.type)
 
     if property_repos is None:
-        if not use_fallback:
+        if Fallback.NO_DEVICE not in use_fallback:
             raise ValueError(f"No layer mapping for `{layer_name}` with device type `{device_type}`")
         _replace_forward(module, module_class)
         return
@@ -496,7 +521,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
     repos = property_repos.repos
 
     if repos is None:
-        if not use_fallback:
+        if Fallback.NO_COMPATIBLE_PROPERTIES not in use_fallback:
             raise ValueError(f"No layer mapping for `{layer_name}` device `{device_type}` with the right properties")
         _replace_forward(module, module_class)
         return
@@ -507,7 +532,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
     )
 
     if repo_with_mode is None:
-        if not use_fallback:
+        if Fallback.NO_COMPATIBLE_MODE not in use_fallback:
             raise ValueError(f"No repository for `{layer_name}` for configuration mode={mode}")
         _replace_forward(module, module_class)
         return
@@ -520,8 +545,13 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
     try:
         layer = _get_layer_memoize(repo, module_class)
     except FileNotFoundError:
-        if not use_fallback:
+        if Fallback.COMPATIBLE_BUILD_NOT_FOUND not in use_fallback:
             raise
+        logger.info(
+            "No compatible kernel build for layer `%s` on %s; using the original forward.",
+            layer_name,
+            device_type.type,
+        )
         _replace_forward(module, module_class)
         return
 
@@ -596,7 +626,7 @@ def _conditionally_replace_forward(
     module: "nn.Module",
     layer: Type["nn.Module"],
     mode: Mode,
-    use_fallback: bool,
+    use_fallback: Fallback,
 ):
     module_class = type(module)
 
@@ -609,7 +639,7 @@ def _conditionally_replace_forward(
     needs_fallback_for_backward = Mode.TRAINING in mode and not getattr(layer, "has_backward", True)
 
     if needs_fallback_for_compile or needs_fallback_for_backward:
-        if use_fallback:
+        if Fallback.NO_COMPATIBLE_MODE in use_fallback:
             if needs_fallback_for_compile:
                 logging.info("Layer does not support torch.compile, using fallback")
             if needs_fallback_for_backward:
diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py
index 1d6f8805..99afc4b5 100644
--- a/kernels/tests/test_layer.py
+++ b/kernels/tests/test_layer.py
@@ -12,6 +12,7 @@
 from kernels import (
     CUDAProperties,
     Device,
+    Fallback,
     FuncRepository,
     LayerRepository,
     LocalLayerRepository,
@@ -425,7 +426,7 @@ class SiluAndMulWithKernelFallback(SiluAndMul):
     kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE)
 
 
-def test_missing_kernel_build_falls_back(monkeypatch):
+def test_missing_kernel_build_falls_back(monkeypatch, caplog):
     repo = LayerRepository(
         repo_id="kernels-test/no-compatible-build",
         layer_name="SiluAndMul",
@@ -439,10 +440,16 @@ def unavailable_build():
     layer = SiluAndMulWithKernel()
     mapping = {"SiluAndMul": {"cuda": repo}}
 
-    with use_kernel_mapping(mapping, inherit_mapping=False):
-        kernelize(layer, device="cuda", mode=Mode.INFERENCE)
+    with (
+        use_kernel_mapping(mapping, inherit_mapping=False),
+        caplog.at_level(logging.INFO, logger="kernels.layer.layer"),
+    ):
+        kernelize(layer, device="cuda", mode=Mode.INFERENCE, use_fallback=Fallback.ALL)
+    assert "No compatible kernel build for layer `SiluAndMul` on cuda" in caplog.text
 
     with use_kernel_mapping(mapping, inherit_mapping=False):
+        with pytest.raises(FileNotFoundError, match="Cannot find a build variant"):
+            kernelize(layer, device="cuda", mode=Mode.INFERENCE)
         with pytest.raises(FileNotFoundError, match="Cannot find a build variant"):
             kernelize(layer, device="cuda", mode=Mode.INFERENCE, use_fallback=False)
 

From df7d1574924f2103f0c60a2fe7cf4328680d8fec Mon Sep 17 00:00:00 2001
From: Marc Sun <marc@huggingface.co>
Date: Thu, 24 Sep 2026 16:15:32 +0000
Subject: [PATCH 3/4] better naming

---
 kernels/src/kernels/layer/kernelize.py |  2 +-
 kernels/src/kernels/layer/layer.py     | 14 +++++++-------
 kernels/tests/test_layer.py            |  2 +-
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/kernels/src/kernels/layer/kernelize.py b/kernels/src/kernels/layer/kernelize.py
index c6a23b1c..53aefb04 100644
--- a/kernels/src/kernels/layer/kernelize.py
+++ b/kernels/src/kernels/layer/kernelize.py
@@ -200,7 +200,7 @@ def kernelize(
         use_fallback (`bool | Fallback`, *optional*, defaults to `True`):
             Cases in which to use the original forward method. `True` is equivalent to `Fallback.DEFAULT`: it allows
             fallback when no compatible mapping or mode exists. `False` is equivalent to `Fallback.NONE` and raises
-            instead. `Fallback.ALL` also falls back when a selected kernel repository has no compatible build.
+            instead. `Fallback.ALL` also falls back when kernel loading fails.
 
     Returns:
         `nn.Module`: The kernelized model with optimized kernel implementations.
diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py
index 05bc7873..67d8d12b 100644
--- a/kernels/src/kernels/layer/layer.py
+++ b/kernels/src/kernels/layer/layer.py
@@ -43,9 +43,9 @@ class Fallback(Flag):
     - `NO_DEVICE`: The layer has no mapping for the requested device type.
     - `NO_COMPATIBLE_PROPERTIES`: No mapping matches the device properties.
     - `NO_COMPATIBLE_MODE`: No repository or loaded kernel supports the requested mode.
-    - `COMPATIBLE_BUILD_NOT_FOUND`: The selected repository has no build for this environment.
-    - `DEFAULT`: Fall back for missing mappings, properties, or mode; raise for a missing build.
-    - `ALL`: Fall back in every case above, including a missing build.
+    - `CANNOT_LOAD`: Kernel loading fails.
+    - `DEFAULT`: Fall back for missing mappings, properties, or mode; raise on a loading error.
+    - `ALL`: Fall back in every case above, including missing files during loading.
     """
 
     NONE = 0
@@ -53,10 +53,10 @@ class Fallback(Flag):
     NO_DEVICE = auto()
     NO_COMPATIBLE_PROPERTIES = auto()
     NO_COMPATIBLE_MODE = auto()
-    COMPATIBLE_BUILD_NOT_FOUND = auto()
+    CANNOT_LOAD = auto()
 
     DEFAULT = NO_LAYER | NO_DEVICE | NO_COMPATIBLE_PROPERTIES | NO_COMPATIBLE_MODE
-    ALL = DEFAULT | COMPATIBLE_BUILD_NOT_FOUND
+    ALL = DEFAULT | CANNOT_LOAD
 
 
 class LayerRepositoryProtocol(RepositoryProtocol, Protocol):
@@ -545,10 +545,10 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
     try:
         layer = _get_layer_memoize(repo, module_class)
     except FileNotFoundError:
-        if Fallback.COMPATIBLE_BUILD_NOT_FOUND not in use_fallback:
+        if Fallback.CANNOT_LOAD not in use_fallback:
             raise
         logger.info(
-            "No compatible kernel build for layer `%s` on %s; using the original forward.",
+            "Kernel for layer `%s` on %s could not be loaded; using the original forward.",
             layer_name,
             device_type.type,
         )
diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py
index 99afc4b5..79666501 100644
--- a/kernels/tests/test_layer.py
+++ b/kernels/tests/test_layer.py
@@ -445,7 +445,7 @@ def unavailable_build():
         caplog.at_level(logging.INFO, logger="kernels.layer.layer"),
     ):
         kernelize(layer, device="cuda", mode=Mode.INFERENCE, use_fallback=Fallback.ALL)
-    assert "No compatible kernel build for layer `SiluAndMul` on cuda" in caplog.text
+    assert "Kernel for layer `SiluAndMul` on cuda could not be loaded" in caplog.text
 
     with use_kernel_mapping(mapping, inherit_mapping=False):
         with pytest.raises(FileNotFoundError, match="Cannot find a build variant"):

From c46af312072b5056db8ad5ee73ba5d0be0f1eab1 Mon Sep 17 00:00:00 2001
From: Marc Sun <marc@huggingface.co>
Date: Tue, 29 Sep 2026 13:00:35 +0000
Subject: [PATCH 4/4] update

---
 kernels/src/kernels/__init__.py        |  4 ++--
 kernels/src/kernels/layer/__init__.py  |  4 ++--
 kernels/src/kernels/layer/kernelize.py | 14 +++++++-------
 kernels/src/kernels/layer/layer.py     | 18 +++++++++---------
 kernels/tests/test_layer.py            |  4 ++--
 5 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/kernels/src/kernels/__init__.py b/kernels/src/kernels/__init__.py
index 5c301ba2..b0fb477d 100644
--- a/kernels/src/kernels/__init__.py
+++ b/kernels/src/kernels/__init__.py
@@ -12,8 +12,8 @@
 from kernels.layer import (
     CUDAProperties,
     Device,
-    Fallback,
     FuncRepository,
+    KernelizeFallback,
     LayerRepository,
     LocalFuncRepository,
     LocalLayerRepository,
@@ -49,7 +49,7 @@
     "Benchmark",
     "CUDAProperties",
     "Device",
-    "Fallback",
+    "KernelizeFallback",
     "ROCMProperties",
     "FuncRepository",
     "LayerRepository",
diff --git a/kernels/src/kernels/layer/__init__.py b/kernels/src/kernels/layer/__init__.py
index 8ca217c7..a46f23be 100644
--- a/kernels/src/kernels/layer/__init__.py
+++ b/kernels/src/kernels/layer/__init__.py
@@ -11,7 +11,7 @@
     use_kernel_mapping,
 )
 from .layer import (
-    Fallback,
+    KernelizeFallback,
     LayerRepository,
     LocalLayerRepository,
     LockedLayerRepository,
@@ -24,7 +24,7 @@
 __all__ = [
     "CUDAProperties",
     "Device",
-    "Fallback",
+    "KernelizeFallback",
     "ROCMProperties",
     "FuncRepository",
     "LayerRepository",
diff --git a/kernels/src/kernels/layer/kernelize.py b/kernels/src/kernels/layer/kernelize.py
index 53aefb04..a06d4df2 100644
--- a/kernels/src/kernels/layer/kernelize.py
+++ b/kernels/src/kernels/layer/kernelize.py
@@ -6,7 +6,7 @@
 
 from .device import Device
 from .globals import _KERNEL_MAPPING
-from .layer import Fallback, kernelize_layer
+from .layer import KernelizeFallback, kernelize_layer
 from .mode import Mode
 from .repos import DeviceRepos, RepositoryProtocol
 
@@ -180,7 +180,7 @@ def kernelize(
     *,
     mode: Mode,
     device: str | "torch.device" | None = None,
-    use_fallback: bool | Fallback = True,
+    use_fallback: bool | KernelizeFallback = True,
 ):
     """
     Replace layer forward methods with optimized kernel implementations.
@@ -197,10 +197,10 @@ def kernelize(
         device (`Union[str, torch.device]`, *optional*):
             The device type to load kernels for. Supported device types are: "cuda", "mps", "npu", "rocm", "tpu", "xpu".
             The device type will be inferred from the model parameters when not provided.
-        use_fallback (`bool | Fallback`, *optional*, defaults to `True`):
-            Cases in which to use the original forward method. `True` is equivalent to `Fallback.DEFAULT`: it allows
-            fallback when no compatible mapping or mode exists. `False` is equivalent to `Fallback.NONE` and raises
-            instead. `Fallback.ALL` also falls back when kernel loading fails.
+        use_fallback (`bool | KernelizeFallback`, *optional*, defaults to `True`):
+            Cases in which to use the original forward method. `True` is equivalent to `KernelizeFallback.DEFAULT`: it allows
+            fallback when no compatible mapping or mode exists. `False` is equivalent to `KernelizeFallback.NONE` and raises
+            instead. `KernelizeFallback.ALL` also falls back when kernel loading fails.
 
     Returns:
         `nn.Module`: The kernelized model with optimized kernel implementations.
@@ -261,7 +261,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     assert isinstance(device_type, Device)
 
     if isinstance(use_fallback, bool):
-        use_fallback = Fallback.DEFAULT if use_fallback else Fallback.NONE
+        use_fallback = KernelizeFallback.DEFAULT if use_fallback else KernelizeFallback.NONE
 
     for _, module in model.named_modules():
         module_class = type(module)
diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py
index 67d8d12b..23b5a138 100644
--- a/kernels/src/kernels/layer/layer.py
+++ b/kernels/src/kernels/layer/layer.py
@@ -35,7 +35,7 @@
 logger = logging.getLogger(__name__)
 
 
-class Fallback(Flag):
+class KernelizeFallback(Flag):
     """Cases in which kernelization may keep the original layer forward.
 
     - `NONE`: Raise instead of falling back.
@@ -487,7 +487,7 @@ def new_init(self, *args, **kwargs):
     return decorator
 
 
-def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use_fallback: Fallback):
+def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use_fallback: KernelizeFallback):
     module_class = type(module)
     layer_name = module_class.kernel_layer_name  # type: ignore[attr-defined]
 
@@ -504,7 +504,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
             f"Check if the layer name matches one of the kernels in the mapping or add the kernel "
             f"you want to use to the mapping. Defaulting to original forward implementation."
         )
-        if Fallback.NO_LAYER not in use_fallback:
+        if KernelizeFallback.NO_LAYER not in use_fallback:
             raise ValueError(f"No layer mapping for `{layer_name}`")
         _replace_forward(module, module_class)
         return
@@ -513,7 +513,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
     property_repos = kernel.get(device_type.type)
 
     if property_repos is None:
-        if Fallback.NO_DEVICE not in use_fallback:
+        if KernelizeFallback.NO_DEVICE not in use_fallback:
             raise ValueError(f"No layer mapping for `{layer_name}` with device type `{device_type}`")
         _replace_forward(module, module_class)
         return
@@ -521,7 +521,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
     repos = property_repos.repos
 
     if repos is None:
-        if Fallback.NO_COMPATIBLE_PROPERTIES not in use_fallback:
+        if KernelizeFallback.NO_COMPATIBLE_PROPERTIES not in use_fallback:
             raise ValueError(f"No layer mapping for `{layer_name}` device `{device_type}` with the right properties")
         _replace_forward(module, module_class)
         return
@@ -532,7 +532,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
     )
 
     if repo_with_mode is None:
-        if Fallback.NO_COMPATIBLE_MODE not in use_fallback:
+        if KernelizeFallback.NO_COMPATIBLE_MODE not in use_fallback:
             raise ValueError(f"No repository for `{layer_name}` for configuration mode={mode}")
         _replace_forward(module, module_class)
         return
@@ -545,7 +545,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
     try:
         layer = _get_layer_memoize(repo, module_class)
     except FileNotFoundError:
-        if Fallback.CANNOT_LOAD not in use_fallback:
+        if KernelizeFallback.CANNOT_LOAD not in use_fallback:
             raise
         logger.info(
             "Kernel for layer `%s` on %s could not be loaded; using the original forward.",
@@ -626,7 +626,7 @@ def _conditionally_replace_forward(
     module: "nn.Module",
     layer: Type["nn.Module"],
     mode: Mode,
-    use_fallback: Fallback,
+    use_fallback: KernelizeFallback,
 ):
     module_class = type(module)
 
@@ -639,7 +639,7 @@ def _conditionally_replace_forward(
     needs_fallback_for_backward = Mode.TRAINING in mode and not getattr(layer, "has_backward", True)
 
     if needs_fallback_for_compile or needs_fallback_for_backward:
-        if Fallback.NO_COMPATIBLE_MODE in use_fallback:
+        if KernelizeFallback.NO_COMPATIBLE_MODE in use_fallback:
             if needs_fallback_for_compile:
                 logging.info("Layer does not support torch.compile, using fallback")
             if needs_fallback_for_backward:
diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py
index 79666501..3ad253b4 100644
--- a/kernels/tests/test_layer.py
+++ b/kernels/tests/test_layer.py
@@ -12,8 +12,8 @@
 from kernels import (
     CUDAProperties,
     Device,
-    Fallback,
     FuncRepository,
+    KernelizeFallback,
     LayerRepository,
     LocalLayerRepository,
     Mode,
@@ -444,7 +444,7 @@ def unavailable_build():
         use_kernel_mapping(mapping, inherit_mapping=False),
         caplog.at_level(logging.INFO, logger="kernels.layer.layer"),
     ):
-        kernelize(layer, device="cuda", mode=Mode.INFERENCE, use_fallback=Fallback.ALL)
+        kernelize(layer, device="cuda", mode=Mode.INFERENCE, use_fallback=KernelizeFallback.ALL)
     assert "Kernel for layer `SiluAndMul` on cuda could not be loaded" in caplog.text
 
     with use_kernel_mapping(mapping, inherit_mapping=False):