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.
- {kernels-0.17.0 → kernels-0.17.2}/Cargo.lock +4 -4
- {kernels-0.17.0 → kernels-0.17.2}/PKG-INFO +1 -1
- {kernels-0.17.0 → kernels-0.17.2}/kernels/Cargo.toml +2 -2
- kernels-0.17.2/kernels/tests/test_entry_points.py +19 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_layer.py +29 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/uv.lock +1 -1
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/Cargo.toml +1 -1
- {kernels-0.17.0 → kernels-0.17.2}/pyproject.toml +2 -2
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/__init__.py +2 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/__init__.py +2 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/kernelize.py +9 -5
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/layer.py +44 -8
- {kernels-0.17.0 → kernels-0.17.2}/Cargo.toml +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/README.md +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/README.md +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/pytest.ini +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/rust/config.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/rust/lib.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/rust/lock.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/rust/version.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/conftest.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/kernel_locking/kernels.lock +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/kernel_locking/pyproject.toml +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/layer_locking/kernels.lock +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/layer_locking/pyproject.toml +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_archs.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_basic.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_benchmarks.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_cli_info.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_data_config.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_data_digest.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_data_kernel_lock.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_data_metadata.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_data_version.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_deps.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_doctest.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_func.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_importer.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_interval_tree.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_kernel_locking.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_loaded_kernels.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_python_deps.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_resolver.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_status.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_tvm_ffi.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_user_agent.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_validate.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_variants.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels/tests/test_verify.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/compat.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/deps.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/git_url.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/kernel_deps.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/mod.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/name.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/parse.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/v1.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/v2.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/v3.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/v4.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/config/v5.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/digest.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/git.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/lib.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/lock.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/metadata.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/python_dependencies.json +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/src/version.rs +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/kernels-common/tests/data/legacy-sha-provenance.json +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/_rust.pyi +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/_system.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/_versions.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/_windows.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/archs.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/backends.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/benchmark.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/benchmarks/__init__.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/benchmarks/activation.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/benchmarks/attention.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/benchmarks/layer_norm.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/__init__.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/benchmark.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/benchmark_graphics.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/download.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/info.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/lock.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/verify_signature.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/cli/versions.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/compat.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/deps.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/hf_hub.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/importer.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/install.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/_interval_tree.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/device.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/func.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/globals.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/mode.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/layer/repos.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/load.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/locking.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/python_depends.json +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/python_deps.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/resolver.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/status.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/validate.py +0 -0
- {kernels-0.17.0 → kernels-0.17.2}/src/kernels/variants.py +0 -0
- {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.
|
|
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.
|
|
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.
|
|
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.
|
|
1575
|
+
version = "0.17.2"
|
|
1576
1576
|
dependencies = [
|
|
1577
1577
|
"kernels-common",
|
|
1578
1578
|
"pyo3",
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[package]
|
|
2
2
|
name = "kernels-python"
|
|
3
|
-
version = "0.17.
|
|
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.
|
|
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):
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "kernels"
|
|
3
|
-
version = "0.17.
|
|
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.
|
|
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
|
-
|
|
202
|
-
|
|
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
|
-
|
|
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:
|
|
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
|
|
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
|
|
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
|
|
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
|
|
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
|
|
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
|
|
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
|