kernels 0.17.0__tar.gz → 0.17.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (108) hide show
  1. {kernels-0.17.0 → kernels-0.17.2}/Cargo.lock +4 -4
  2. {kernels-0.17.0 → kernels-0.17.2}/PKG-INFO +1 -1
  3. {kernels-0.17.0 → kernels-0.17.2}/kernels/Cargo.toml +2 -2
  4. kernels-0.17.2/kernels/tests/test_entry_points.py +19 -0
  5. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_layer.py +29 -0
  6. {kernels-0.17.0 → kernels-0.17.2}/kernels/uv.lock +1 -1
  7. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/Cargo.toml +1 -1
  8. {kernels-0.17.0 → kernels-0.17.2}/pyproject.toml +2 -2
  9. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/__init__.py +2 -0
  10. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/__init__.py +2 -0
  11. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/kernelize.py +9 -5
  12. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/layer.py +44 -8
  13. {kernels-0.17.0 → kernels-0.17.2}/Cargo.toml +0 -0
  14. {kernels-0.17.0 → kernels-0.17.2}/README.md +0 -0
  15. {kernels-0.17.0 → kernels-0.17.2}/kernels/README.md +0 -0
  16. {kernels-0.17.0 → kernels-0.17.2}/kernels/pytest.ini +0 -0
  17. {kernels-0.17.0 → kernels-0.17.2}/kernels/rust/config.rs +0 -0
  18. {kernels-0.17.0 → kernels-0.17.2}/kernels/rust/lib.rs +0 -0
  19. {kernels-0.17.0 → kernels-0.17.2}/kernels/rust/lock.rs +0 -0
  20. {kernels-0.17.0 → kernels-0.17.2}/kernels/rust/version.rs +0 -0
  21. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/conftest.py +0 -0
  22. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/kernel_locking/kernels.lock +0 -0
  23. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/kernel_locking/pyproject.toml +0 -0
  24. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/layer_locking/kernels.lock +0 -0
  25. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/layer_locking/pyproject.toml +0 -0
  26. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_archs.py +0 -0
  27. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_basic.py +0 -0
  28. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_benchmarks.py +0 -0
  29. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_cli_info.py +0 -0
  30. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_data_config.py +0 -0
  31. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_data_digest.py +0 -0
  32. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_data_kernel_lock.py +0 -0
  33. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_data_metadata.py +0 -0
  34. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_data_version.py +0 -0
  35. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_deps.py +0 -0
  36. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_doctest.py +0 -0
  37. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_func.py +0 -0
  38. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_importer.py +0 -0
  39. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_interval_tree.py +0 -0
  40. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_kernel_locking.py +0 -0
  41. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_loaded_kernels.py +0 -0
  42. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_python_deps.py +0 -0
  43. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_resolver.py +0 -0
  44. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_status.py +0 -0
  45. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_tvm_ffi.py +0 -0
  46. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_user_agent.py +0 -0
  47. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_validate.py +0 -0
  48. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_variants.py +0 -0
  49. {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_verify.py +0 -0
  50. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/compat.rs +0 -0
  51. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/deps.rs +0 -0
  52. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/git_url.rs +0 -0
  53. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/kernel_deps.rs +0 -0
  54. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/mod.rs +0 -0
  55. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/name.rs +0 -0
  56. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/parse.rs +0 -0
  57. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/v1.rs +0 -0
  58. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/v2.rs +0 -0
  59. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/v3.rs +0 -0
  60. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/v4.rs +0 -0
  61. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/v5.rs +0 -0
  62. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/digest.rs +0 -0
  63. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/git.rs +0 -0
  64. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/lib.rs +0 -0
  65. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/lock.rs +0 -0
  66. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/metadata.rs +0 -0
  67. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/python_dependencies.json +0 -0
  68. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/version.rs +0 -0
  69. {kernels-0.17.0 → kernels-0.17.2}/kernels-common/tests/data/legacy-sha-provenance.json +0 -0
  70. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/_rust.pyi +0 -0
  71. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/_system.py +0 -0
  72. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/_versions.py +0 -0
  73. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/_windows.py +0 -0
  74. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/archs.py +0 -0
  75. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/backends.py +0 -0
  76. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/benchmark.py +0 -0
  77. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/benchmarks/__init__.py +0 -0
  78. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/benchmarks/activation.py +0 -0
  79. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/benchmarks/attention.py +0 -0
  80. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/benchmarks/layer_norm.py +0 -0
  81. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/__init__.py +0 -0
  82. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/benchmark.py +0 -0
  83. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/benchmark_graphics.py +0 -0
  84. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/download.py +0 -0
  85. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/info.py +0 -0
  86. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/lock.py +0 -0
  87. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/verify_signature.py +0 -0
  88. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/versions.py +0 -0
  89. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/compat.py +0 -0
  90. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/deps.py +0 -0
  91. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/hf_hub.py +0 -0
  92. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/importer.py +0 -0
  93. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/install.py +0 -0
  94. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/_interval_tree.py +0 -0
  95. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/device.py +0 -0
  96. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/func.py +0 -0
  97. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/globals.py +0 -0
  98. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/mode.py +0 -0
  99. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/repos.py +0 -0
  100. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/load.py +0 -0
  101. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/locking.py +0 -0
  102. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/python_depends.json +0 -0
  103. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/python_deps.py +0 -0
  104. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/resolver.py +0 -0
  105. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/status.py +0 -0
  106. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/validate.py +0 -0
  107. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/variants.py +0 -0
  108. {kernels-0.17.0 → kernels-0.17.2}/src/kernels/verify.py +0 -0
@@ -1056,7 +1056,7 @@ dependencies = [
1056
1056
 
1057
1057
  [[package]]
1058
1058
  name = "hf-kernel-builder"
1059
- version = "0.17.0"
1059
+ version = "0.17.2"
1060
1060
  dependencies = [
1061
1061
  "base32",
1062
1062
  "built",
@@ -1537,7 +1537,7 @@ dependencies = [
1537
1537
 
1538
1538
  [[package]]
1539
1539
  name = "kernel-abi-check"
1540
- version = "0.17.0"
1540
+ version = "0.17.2"
1541
1541
  dependencies = [
1542
1542
  "cpp_demangle",
1543
1543
  "eyre",
@@ -1551,7 +1551,7 @@ dependencies = [
1551
1551
 
1552
1552
  [[package]]
1553
1553
  name = "kernels-common"
1554
- version = "0.17.0"
1554
+ version = "0.17.2"
1555
1555
  dependencies = [
1556
1556
  "base64",
1557
1557
  "digest",
@@ -1572,7 +1572,7 @@ dependencies = [
1572
1572
 
1573
1573
  [[package]]
1574
1574
  name = "kernels-python"
1575
- version = "0.17.0"
1575
+ version = "0.17.2"
1576
1576
  dependencies = [
1577
1577
  "kernels-common",
1578
1578
  "pyo3",
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: kernels
3
- Version: 0.17.0
3
+ Version: 0.17.2
4
4
  Requires-Dist: huggingface-hub>=1.10.0
5
5
  Requires-Dist: packaging>=20.0
6
6
  Requires-Dist: pyyaml>=6
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "kernels-python"
3
- version = "0.17.0"
3
+ version = "0.17.2"
4
4
  edition = "2024"
5
5
  description = "Kernels data structures (Python bindings)"
6
6
  homepage = "https://github.com/huggingface/kernels"
@@ -19,4 +19,4 @@ serde_json = "1"
19
19
 
20
20
  [dependencies.kernels-common]
21
21
  path = "../kernels-common"
22
- version = "0.17.0"
22
+ version = "0.17.2"
@@ -0,0 +1,19 @@
1
+ from importlib.metadata import distribution
2
+
3
+ import pytest
4
+
5
+
6
+ def test_egg_info_writer_entry_point_is_importable():
7
+ """The `egg_info.writers` entry point is loaded by setuptools during *any*
8
+ `egg_info` run in an environment where `kernels` is installed, including
9
+ builds of unrelated packages. If it points at a module that is not shipped,
10
+ every such build fails with `ModuleNotFoundError`."""
11
+ entry_points = [ep for ep in distribution("kernels").entry_points if ep.group == "egg_info.writers"]
12
+ assert entry_points, "no egg_info.writers entry point registered"
13
+
14
+ for ep in entry_points:
15
+ try:
16
+ writer = ep.load()
17
+ except ModuleNotFoundError as e:
18
+ pytest.fail(f"entry point {ep.name} = {ep.value!r} is not importable: {e}")
19
+ assert callable(writer)
@@ -13,6 +13,7 @@ from kernels import (
13
13
  CUDAProperties,
14
14
  Device,
15
15
  FuncRepository,
16
+ KernelizeFallback,
16
17
  LayerRepository,
17
18
  LocalLayerRepository,
18
19
  Mode,
@@ -425,6 +426,34 @@ def test_layer_fallback_works():
425
426
  kernelize(silu_and_mul, device="cuda", mode=Mode.INFERENCE)
426
427
 
427
428
 
429
+ def test_missing_kernel_build_falls_back(monkeypatch, caplog):
430
+ repo = LayerRepository(
431
+ repo_id="kernels-test/no-compatible-build",
432
+ layer_name="SiluAndMul",
433
+ version=1,
434
+ )
435
+
436
+ def unavailable_build():
437
+ raise FileNotFoundError("Cannot find a build variant for this system")
438
+
439
+ monkeypatch.setattr(repo, "load", unavailable_build)
440
+ layer = SiluAndMulWithKernel()
441
+ mapping = {"SiluAndMul": {"cuda": repo}}
442
+
443
+ with (
444
+ use_kernel_mapping(mapping, inherit_mapping=False),
445
+ caplog.at_level(logging.INFO, logger="kernels.layer.layer"),
446
+ ):
447
+ kernelize(layer, device="cuda", mode=Mode.INFERENCE, use_fallback=KernelizeFallback.ALL)
448
+ assert "Kernel for layer `SiluAndMul` on cuda could not be loaded" in caplog.text
449
+
450
+ with use_kernel_mapping(mapping, inherit_mapping=False):
451
+ with pytest.raises(FileNotFoundError, match="Cannot find a build variant"):
452
+ kernelize(layer, device="cuda", mode=Mode.INFERENCE)
453
+ with pytest.raises(FileNotFoundError, match="Cannot find a build variant"):
454
+ kernelize(layer, device="cuda", mode=Mode.INFERENCE, use_fallback=False)
455
+
456
+
428
457
  def test_kernel_condition_skips_kernelization(caplog):
429
458
  @use_kernel_forward_from_hub("SiluAndMulNonExisting", condition=lambda module: False)
430
459
  class SiluAndMulConditionSkipped(SiluAndMul):
@@ -1158,7 +1158,7 @@ wheels = [
1158
1158
 
1159
1159
  [[package]]
1160
1160
  name = "kernels"
1161
- version = "0.17.0"
1161
+ version = "0.17.2"
1162
1162
  source = { editable = "." }
1163
1163
  dependencies = [
1164
1164
  { name = "huggingface-hub" },
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "kernels-common"
3
- version = "0.17.0"
3
+ version = "0.17.2"
4
4
  edition = "2024"
5
5
  description = "Code shared by kernel-builder and kernels"
6
6
  homepage = "https://github.com/huggingface/kernels"
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "kernels"
3
- version = "0.17.0"
3
+ version = "0.17.2"
4
4
  description = "Download compute kernels"
5
5
  authors = [
6
6
  { name = "Daniel de Kok", email = "daniel@huggingface.co" },
@@ -63,7 +63,7 @@ docs = [
63
63
  kernels = "kernels.cli:main"
64
64
 
65
65
  [project.entry-points."egg_info.writers"]
66
- "kernels.lock" = "kernels.lockfile:write_egg_lockfile"
66
+ "kernels.lock" = "kernels.locking:write_egg_lockfile"
67
67
 
68
68
  [tool.maturin]
69
69
  module-name = "kernels._rust"
@@ -13,6 +13,7 @@ from kernels.layer import (
13
13
  CUDAProperties,
14
14
  Device,
15
15
  FuncRepository,
16
+ KernelizeFallback,
16
17
  LayerRepository,
17
18
  LocalFuncRepository,
18
19
  LocalLayerRepository,
@@ -48,6 +49,7 @@ __all__ = [
48
49
  "Benchmark",
49
50
  "CUDAProperties",
50
51
  "Device",
52
+ "KernelizeFallback",
51
53
  "ROCMProperties",
52
54
  "FuncRepository",
53
55
  "LayerRepository",
@@ -11,6 +11,7 @@ from .kernelize import (
11
11
  use_kernel_mapping,
12
12
  )
13
13
  from .layer import (
14
+ KernelizeFallback,
14
15
  LayerRepository,
15
16
  LocalLayerRepository,
16
17
  LockedLayerRepository,
@@ -23,6 +24,7 @@ from .mode import Mode
23
24
  __all__ = [
24
25
  "CUDAProperties",
25
26
  "Device",
27
+ "KernelizeFallback",
26
28
  "ROCMProperties",
27
29
  "FuncRepository",
28
30
  "LayerRepository",
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING
6
6
 
7
7
  from .device import Device
8
8
  from .globals import _KERNEL_MAPPING
9
- from .layer import kernelize_layer
9
+ from .layer import KernelizeFallback, kernelize_layer
10
10
  from .mode import Mode
11
11
  from .repos import DeviceRepos, RepositoryProtocol
12
12
 
@@ -180,7 +180,7 @@ def kernelize(
180
180
  *,
181
181
  mode: Mode,
182
182
  device: str | "torch.device" | None = None,
183
- use_fallback: bool = True,
183
+ use_fallback: bool | KernelizeFallback = True,
184
184
  ):
185
185
  """
186
186
  Replace layer forward methods with optimized kernel implementations.
@@ -197,9 +197,10 @@ def kernelize(
197
197
  device (`Union[str, torch.device]`, *optional*):
198
198
  The device type to load kernels for. Supported device types are: "cuda", "mps", "npu", "rocm", "tpu", "xpu".
199
199
  The device type will be inferred from the model parameters when not provided.
200
- use_fallback (`bool`, *optional*, defaults to `True`):
201
- Whether to use the original forward method of modules when no compatible kernel could be found.
202
- If set to `False`, an exception will be raised in such cases.
200
+ use_fallback (`bool | KernelizeFallback`, *optional*, defaults to `True`):
201
+ Cases in which to use the original forward method. `True` is equivalent to `KernelizeFallback.DEFAULT`: it allows
202
+ fallback when no compatible mapping or mode exists. `False` is equivalent to `KernelizeFallback.NONE` and raises
203
+ instead. `KernelizeFallback.ALL` also falls back when kernel loading fails.
203
204
 
204
205
  Returns:
205
206
  `nn.Module`: The kernelized model with optimized kernel implementations.
@@ -259,6 +260,9 @@ def kernelize(
259
260
 
260
261
  assert isinstance(device_type, Device)
261
262
 
263
+ if isinstance(use_fallback, bool):
264
+ use_fallback = KernelizeFallback.DEFAULT if use_fallback else KernelizeFallback.NONE
265
+
262
266
  for _, module in model.named_modules():
263
267
  module_class = type(module)
264
268
 
@@ -3,6 +3,7 @@ from __future__ import annotations
3
3
  import functools
4
4
  import inspect
5
5
  import logging
6
+ from enum import Flag, auto
6
7
  from inspect import Parameter, Signature
7
8
  from pathlib import Path
8
9
  from types import MethodType, ModuleType
@@ -36,6 +37,30 @@ if TYPE_CHECKING:
36
37
  logger = logging.getLogger(__name__)
37
38
 
38
39
 
40
+ class KernelizeFallback(Flag):
41
+ """Cases in which kernelization may keep the original layer forward.
42
+
43
+ - `NONE`: Raise instead of falling back.
44
+ - `NO_LAYER`: No kernel mapping exists for the layer.
45
+ - `NO_DEVICE`: The layer has no mapping for the requested device type.
46
+ - `NO_COMPATIBLE_PROPERTIES`: No mapping matches the device properties.
47
+ - `NO_COMPATIBLE_MODE`: No repository or loaded kernel supports the requested mode.
48
+ - `CANNOT_LOAD`: Kernel loading fails.
49
+ - `DEFAULT`: Fall back for missing mappings, properties, or mode; raise on a loading error.
50
+ - `ALL`: Fall back in every case above, including missing files during loading.
51
+ """
52
+
53
+ NONE = 0
54
+ NO_LAYER = auto()
55
+ NO_DEVICE = auto()
56
+ NO_COMPATIBLE_PROPERTIES = auto()
57
+ NO_COMPATIBLE_MODE = auto()
58
+ CANNOT_LOAD = auto()
59
+
60
+ DEFAULT = NO_LAYER | NO_DEVICE | NO_COMPATIBLE_PROPERTIES | NO_COMPATIBLE_MODE
61
+ ALL = DEFAULT | CANNOT_LOAD
62
+
63
+
39
64
  class LayerRepositoryProtocol(RepositoryProtocol, Protocol):
40
65
  @property
41
66
  def layer_name(self) -> str: ...
@@ -468,7 +493,7 @@ def use_kernelized_func(*args: Callable):
468
493
  return decorator
469
494
 
470
495
 
471
- def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use_fallback):
496
+ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use_fallback: KernelizeFallback):
472
497
  module_class = type(module)
473
498
  layer_name = module_class.kernel_layer_name # type: ignore[attr-defined]
474
499
 
@@ -485,7 +510,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
485
510
  f"Check if the layer name matches one of the kernels in the mapping or add the kernel "
486
511
  f"you want to use to the mapping. Defaulting to original forward implementation."
487
512
  )
488
- if not use_fallback:
513
+ if KernelizeFallback.NO_LAYER not in use_fallback:
489
514
  raise ValueError(f"No layer mapping for `{layer_name}`")
490
515
  _replace_forward(module, module_class)
491
516
  return
@@ -494,7 +519,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
494
519
  property_repos = kernel.get(device_type.type)
495
520
 
496
521
  if property_repos is None:
497
- if not use_fallback:
522
+ if KernelizeFallback.NO_DEVICE not in use_fallback:
498
523
  raise ValueError(f"No layer mapping for `{layer_name}` with device type `{device_type}`")
499
524
  _replace_forward(module, module_class)
500
525
  return
@@ -502,7 +527,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
502
527
  repos = property_repos.repos
503
528
 
504
529
  if repos is None:
505
- if not use_fallback:
530
+ if KernelizeFallback.NO_COMPATIBLE_PROPERTIES not in use_fallback:
506
531
  raise ValueError(f"No layer mapping for `{layer_name}` device `{device_type}` with the right properties")
507
532
  _replace_forward(module, module_class)
508
533
  return
@@ -513,7 +538,7 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
513
538
  )
514
539
 
515
540
  if repo_with_mode is None:
516
- if not use_fallback:
541
+ if KernelizeFallback.NO_COMPATIBLE_MODE not in use_fallback:
517
542
  raise ValueError(f"No repository for `{layer_name}` for configuration mode={mode}")
518
543
  _replace_forward(module, module_class)
519
544
  return
@@ -523,7 +548,18 @@ def kernelize_layer(module: "nn.Module", *, mode: Mode, device_type: Device, use
523
548
  logging.info(f"Using function/layer from repo {repo}")
524
549
  logging.debug(f"kernelize mode: {mode}, repo mode: {repo_mode}")
525
550
 
526
- layer = _get_layer_memoize(repo, module_class)
551
+ try:
552
+ layer = _get_layer_memoize(repo, module_class)
553
+ except FileNotFoundError:
554
+ if KernelizeFallback.CANNOT_LOAD not in use_fallback:
555
+ raise
556
+ logger.info(
557
+ "Kernel for layer `%s` on %s could not be loaded; using the original forward.",
558
+ layer_name,
559
+ device_type.type,
560
+ )
561
+ _replace_forward(module, module_class)
562
+ return
527
563
 
528
564
  # Ideally we would do validation on the mapping where we check that
529
565
  # e.g. if a repo class is registered for TRAINING | TORCH_COMPILE,
@@ -596,7 +632,7 @@ def _conditionally_replace_forward(
596
632
  module: "nn.Module",
597
633
  layer: Type["nn.Module"],
598
634
  mode: Mode,
599
- use_fallback: bool,
635
+ use_fallback: KernelizeFallback,
600
636
  ):
601
637
  module_class = type(module)
602
638
 
@@ -609,7 +645,7 @@ def _conditionally_replace_forward(
609
645
  needs_fallback_for_backward = Mode.TRAINING in mode and not getattr(layer, "has_backward", True)
610
646
 
611
647
  if needs_fallback_for_compile or needs_fallback_for_backward:
612
- if use_fallback:
648
+ if KernelizeFallback.NO_COMPATIBLE_MODE in use_fallback:
613
649
  if needs_fallback_for_compile:
614
650
  logging.info("Layer does not support torch.compile, using fallback")
615
651
  if needs_fallback_for_backward:
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes