kernels 0.18.2__tar.gz → 0.18.3__tar.gz

{kernels-0.18.2 → kernels-0.18.3}/Cargo.lock

@@ -1058,7 +1058,7 @@ dependencies = [
 
 [[package]]
 name = "hf-kernel-builder"
-version = "0.18.2"
+version = "0.18.3"
 dependencies = [
  "base32",
  "built",
@@ -1550,7 +1550,7 @@ dependencies = [
 
 [[package]]
 name = "kernel-abi-check"
-version = "0.18.2"
+version = "0.18.3"
 dependencies = [
  "cpp_demangle",
  "eyre",
@@ -1581,7 +1581,7 @@ dependencies = [
 
 [[package]]
 name = "kernels-common"
-version = "0.18.2"
+version = "0.18.3"
 dependencies = [
  "base64",
  "digest",
@@ -1603,7 +1603,7 @@ dependencies = [
 
 [[package]]
 name = "kernels-python"
-version = "0.18.2"
+version = "0.18.3"
 dependencies = [
  "eyre",
  "kernels-common",

{kernels-0.18.2 → kernels-0.18.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernels
-Version: 0.18.2
+Version: 0.18.3
 Requires-Dist: huggingface-hub>=1.32.0
 Requires-Dist: packaging>=20.0
 Requires-Dist: pyyaml>=6

{kernels-0.18.2 → kernels-0.18.3}/kernels/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "kernels-python"
-version = "0.18.2"
+version = "0.18.3"
 edition = "2024"
 description = "Kernels data structures (Python bindings)"
 homepage = "https://github.com/huggingface/kernels"
@@ -20,4 +20,4 @@ serde_json = "1"
 
 [dependencies.kernels-common]
 path = "../kernels-common"
-version = "0.18.2"
+version = "0.18.3"

{kernels-0.18.2 → kernels-0.18.3}/kernels/tests/test_func.py

@@ -300,6 +300,7 @@ def test_kernel_func_nested_dispatch(compile):
         _bind_forward(inner._kernel_funcs["double_me"], TimesThree.forward)
 
     if compile:
+        torch._dynamo.reset()
         # We only need to know whether it's safe around get/set so the backend is not relevant
         model = torch.compile(model, backend="eager", fullgraph=True)
 
@@ -318,9 +319,69 @@ def test_kernel_func_nested_compile_context(monkeypatch, num_inner, expected):
         _bind_forward(inner._kernel_funcs["double_me"], TimesThree.forward)
 
     # Mimic different compiled boundaries
-    compiling = iter([True, True, False])
+    # > Outer: 5 checks (get/set/reset + 2 dispatches); each Inner: 4 (get/set/reset + 1 dispatch)
+    # > Each final boundary is set to false (last inner, outer dispatch, outer reset)
+    compiling = iter([True] * (4 * num_inner + 2) + [False] * 3)
     monkeypatch.setattr(torch.compiler, "is_compiling", lambda: next(compiling))
 
     # Same as in `test_kernel_func_nested_dispatch` (just *3 once or twice)
     x = torch.tensor(1.0)
     torch.testing.assert_close(model(x), torch.tensor(expected))
+
+
+def test_kernel_func_compile_reuses_forward_across_instances():
+    models = [SurpriseMe() for _ in range(24)]  # "num_hidden_layers"
+
+    for i, model in enumerate(models):
+        forward = AddOne.forward if i % 2 == 0 else TimesThree.forward
+        _bind_forward(model._kernel_funcs["surprise_me"], forward)
+
+    # Catch the graphs dynamo would create (including new ones created during recompiles)
+    graphs = []
+
+    def backend(gm, inputs):
+        graphs.append(gm)
+        return gm.forward
+
+    torch._dynamo.reset()
+    compiled = torch.compile(SurpriseMe.forward, backend=backend, fullgraph=True)
+
+    x = torch.arange(4).float()
+    for _ in range(2):
+        for i, model in enumerate(models):
+            expected = x + 1 if i % 2 == 0 else x * 3
+            torch.testing.assert_close(compiled(model, x), expected)
+
+    # One graph per implementation, not per private wrapper
+    assert len(graphs) == 2
+
+
+def test_kernel_func_compile_forward_switch():
+    a, b = SurpriseMe(), SurpriseMe()
+    a_func = a._kernel_funcs["surprise_me"]
+    b_func = b._kernel_funcs["surprise_me"]
+
+    _bind_forward(a_func, AddOne.forward)
+    _bind_forward(b_func, TimesThree.forward)
+
+    torch._dynamo.reset()
+    compiled = torch.compile(SurpriseMe.forward, backend="eager", fullgraph=True)
+    x = torch.arange(4).float()
+
+    torch.testing.assert_close(compiled(a, x), x + 1)
+    torch.testing.assert_close(compiled(b, x), x * 3)
+
+    # Switch A to B's implementation without affecting B
+    _bind_forward(a_func, TimesThree.forward)
+    torch.testing.assert_close(compiled(a, x), x * 3)
+    torch.testing.assert_close(compiled(b, x), x * 3)
+
+    # Restore the original forward
+    _bind_forward(a_func, type(a_func).forward)
+    torch.testing.assert_close(compiled(a, x), x)
+    torch.testing.assert_close(compiled(b, x), x * 3)
+
+    # Switch A back to its original kernel
+    _bind_forward(a_func, AddOne.forward)
+    torch.testing.assert_close(compiled(a, x), x + 1)
+    torch.testing.assert_close(compiled(b, x), x * 3)

{kernels-0.18.2 → kernels-0.18.3}/kernels/tests/test_layer.py

@@ -1521,13 +1521,15 @@ def test_get_kernelized_layers_func(stub_kernel):
 
 @pytest.fixture
 def heterogeneous_gpus(monkeypatch):
-    """Pretend that device 0 has capability 7.5, device 1 has 9.0, and the current device is 0."""
+    """Pretend that we are on CUDA, device 0 has capability 7.5, device 1 has 9.0, and the current device is 0."""
     state = {"current": 0}
     capabilities = {0: (7, 5), 1: (9, 0)}
 
     def get_device_capability(device=None):
         return capabilities[state["current"] if device is None else device]
 
+    monkeypatch.setattr(torch.version, "hip", None)
+    monkeypatch.setattr(torch.version, "cuda", "12.8")
     monkeypatch.setattr(torch.cuda, "get_device_capability", get_device_capability)
     return state
 
@@ -1579,6 +1581,23 @@ def test_kernelize_device_index(stub_kernel, heterogeneous_gpus):
         assert model[0]._kernel_repo == expected
 
 
+@pytest.mark.parametrize("is_rocm", [False, True])
+def test_kernelize_torch_device_rocm(stub_kernel, monkeypatch, is_rocm):
+    monkeypatch.setattr(torch.version, "hip", "7.0" if is_rocm else None)
+    monkeypatch.setattr(torch.version, "cuda", None if is_rocm else "12.8")
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (9, 4))
+
+    repo_cuda = LayerRepository(repo_id="kernels-test/cuda", layer_name="StubLayer", version=1)
+    repo_rocm = LayerRepository(repo_id="kernels-test/rocm", layer_name="StubLayer", version=1)
+    mapping = {"StubLayer": {"cuda": repo_cuda, "rocm": repo_rocm}}
+
+    # ROCm devices use the `cuda` device type in Torch.
+    model = nn.Sequential(StubLayer())
+    with use_kernel_mapping(mapping, inherit_mapping=False):
+        kernelize(model, mode=Mode.INFERENCE, device=torch.device("cuda", 0))
+    assert model[0]._kernel_repo == (repo_rocm if is_rocm else repo_cuda)
+
+
 def test_locked_repositories_are_picklable():
     lockfile = Path(__file__).parent / "layer_locking" / "kernels.lock"
     layer_repo = LockedLayerRepository("kernels-test/versions", layer_name="Version", lockfile=lockfile)

{kernels-0.18.2 → kernels-0.18.3}/kernels/uv.lock

@@ -1457,7 +1457,7 @@ wheels = [
 
 [[package]]
 name = "kernels"
-version = "0.18.2"
+version = "0.18.3"
 source = { editable = "." }
 dependencies = [
     { name = "huggingface-hub" },

{kernels-0.18.2 → kernels-0.18.3}/kernels-common/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "kernels-common"
-version = "0.18.2"
+version = "0.18.3"
 edition = "2024"
 description = "Code shared by kernel-builder and kernels"
 homepage = "https://github.com/huggingface/kernels"

{kernels-0.18.2 → kernels-0.18.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernels"
-version = "0.18.2"
+version = "0.18.3"
 description = "Download compute kernels"
 authors = [
   { name = "Daniel de Kok", email = "daniel@huggingface.co" },

{kernels-0.18.2 → kernels-0.18.3}/src/kernels/layer/kernelize.py

@@ -256,7 +256,7 @@ def kernelize(
         _validate_device_type(device)
         device_type = Device(type=device)
     else:
-        device_type = Device(device.type)
+        device_type = _device_type_from_torch(device.type)
 
     assert isinstance(device_type, Device)
 
@@ -314,9 +314,12 @@ def _find_device(model: "nn.Module") -> Device:
     except StopIteration:
         raise ValueError("Cannot determine model device, provide as `device` argument to `kernelize`.")
 
-    dev_type = param.device.type
+    return _device_type_from_torch(param.device.type)
+
+
+def _device_type_from_torch(dev_type: str) -> Device:
+    """Map a Torch device type to a kernels device type. ROCm devices use the `cuda` device type in Torch."""
     if dev_type == "cuda":
-        # Refine based on actual platform
         if _is_rocm_platform():
             return Device(type="rocm")
         elif _is_cuda_platform():

{kernels-0.18.2 → kernels-0.18.3}/src/kernels/layer/layer.py

@@ -576,7 +576,9 @@ def use_kernelized_func(*args: Callable):
                 name = getattr(fn, "__name__", None) or getattr(fn, "kernel_layer_name", None)
                 assert name is not None
 
-                active_kernel_funcs[id(fn)] = self._kernel_funcs[name]
+                # Private wrapper for eager, selected forward for compile
+                kernel_func = self._kernel_funcs[name]
+                active_kernel_funcs[id(fn)] = (kernel_func, kernel_func.forward.__func__)
 
             token = _ACTIVE_KERNEL_FUNCS.set(active_kernel_funcs)
             try:
@@ -828,6 +830,7 @@ def _rebuild_kernel_func(module_name: str, func_name: str):
 
 
 def _create_func_module(func: Callable) -> Type["nn.Module"]:
+    import torch
     from torch import nn
 
     class Func(nn.Module):
@@ -838,8 +841,14 @@ def _create_func_module(func: Callable) -> Type["nn.Module"]:
         def forward(self, *args, **kwargs):
             # Dispatch global calls to the active model's private wrapper
             if (active_kernel_funcs := _ACTIVE_KERNEL_FUNCS.get()) is not None:
-                if (kernel_func := active_kernel_funcs.get(id(self))) is not None:
-                    return kernel_func(*args, **kwargs)
+                if (active_kernel_func := active_kernel_funcs.get(id(self))) is not None:
+                    # Call the selected forward directly to avoid dynamo guards (private wrapper ids)
+                    kernel_func, forward_func = active_kernel_func
+                    if torch.compiler.is_compiling():
+                        if forward_func is not type(self).forward:
+                            return forward_func(self, *args, **kwargs)
+                    else:
+                        return kernel_func(*args, **kwargs)
 
             return func(*args, **kwargs)