kernels 0.4.4__tar.gz → 0.5.0.dev0__tar.gz

{kernels-0.4.4/src/kernels.egg-info → kernels-0.5.0.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: kernels
-Version: 0.4.4
+Version: 0.5.0.dev0
 Summary: Download compute kernels
 Author-email: OlivierDehaene <olivier@huggingface.co>, Daniel de Kok <daniel@huggingface.co>, David Holtz <david@huggingface.co>, Nicolas Patry <nicolas@huggingface.co>
 License: Apache-2.0
@@ -12,6 +12,7 @@ Requires-Dist: packaging>=20.0
 Requires-Dist: tomli>=2.0; python_version < "3.11"
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
+Dynamic: license-file
 
 # kernels
 

{kernels-0.4.4 → kernels-0.5.0.dev0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernels"
-version = "0.4.4"
+version = "0.5.0.dev0"
 description = "Download compute kernels"
 authors = [
   { name = "OlivierDehaene", email = "olivier@huggingface.co" },

{kernels-0.4.4 → kernels-0.5.0.dev0}/src/kernels/layer.py

@@ -138,6 +138,8 @@ def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool
             return fallback_forward(self, x, *args, **kwargs)
 
         needs_backward = self.training
+        is_compiling = _is_torchdynamo_compiling()
+
         kernel = _KERNEL_MAPPING.get().get(layer_name)
         if kernel is None:
             warnings.warn(
@@ -165,7 +167,14 @@ def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool
         # Short-circuit if we already loaded the layer.
         layer = cached_layer.get(repo, None)
         if layer is not None:
-            if needs_backward and not getattr(layer, "has_backward", True):
+            # Switch to fallback when the layer does not support:
+            # compilation/compile when needed.
+            # backward when needed
+            needs_fallback = needs_backward and not getattr(layer, "has_backward", True)
+            needs_fallback |= is_compiling and not getattr(
+                layer, "can_torch_compile", False
+            )
+            if needs_fallback:
                 return fallback_forward(self, x, *args, **kwargs)
             return layer.forward(self, x, *args, **kwargs)
 
@@ -185,8 +194,15 @@ def replace_kernel_forward_from_hub(cls, layer_name: str, *, use_fallback: bool
 
         cached_layer[repo] = layer
 
-        if needs_backward and not getattr(layer, "has_backward", True):
+        # Switch to fallback when the layer does not support
+        # compilation/compile when needed.
+        needs_fallback = needs_backward and not getattr(layer, "has_backward", True)
+        needs_fallback |= is_compiling and not getattr(
+            layer, "can_torch_compile", False
+        )
+        if needs_fallback:
             return fallback_forward(self, x, *args, **kwargs)
+
         return layer.forward(self, x, *args, **kwargs)
 
     cls.forward = forward
@@ -245,7 +261,8 @@ def _validate_layer(*, check_cls, cls):
     torch_module_members = {name for name, _ in inspect.getmembers(nn.Module)}
     cls_members = {name for name, _ in inspect.getmembers(cls)}
     difference = cls_members - torch_module_members
-    if difference != set() and difference != {"has_backward"}:
+    # verify if : difference ⊄ {"can_torch_compile", "has_backward"}
+    if not difference <= {"can_torch_compile", "has_backward"}:
         raise TypeError("Layer must not contain additional members.")
 
     # Check whether the forward signatures are similar.
@@ -262,3 +279,19 @@ def _validate_layer(*, check_cls, cls):
             raise TypeError(
                 f"Forward signature does not match: different kind of arguments ({param} ({param.kind}) and {ref_param} ({ref_param.kind})"
             )
+
+
+def _is_torchdynamo_compiling():
+    # Importing torch._dynamo causes issues with PyTorch profiler (https://github.com/pytorch/pytorch/issues/130622)
+    # hence rather relying on `torch.compiler.is_compiling()` when possible (torch>=2.3)
+    try:
+        import torch
+
+        return torch.compiler.is_compiling()
+    except Exception:
+        try:
+            import torch._dynamo as dynamo  # noqa: F401
+
+            return dynamo.is_compiling()
+        except Exception:
+            return False

{kernels-0.4.4 → kernels-0.5.0.dev0/src/kernels.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: kernels
-Version: 0.4.4
+Version: 0.5.0.dev0
 Summary: Download compute kernels
 Author-email: OlivierDehaene <olivier@huggingface.co>, Daniel de Kok <daniel@huggingface.co>, David Holtz <david@huggingface.co>, Nicolas Patry <nicolas@huggingface.co>
 License: Apache-2.0
@@ -12,6 +12,7 @@ Requires-Dist: packaging>=20.0
 Requires-Dist: tomli>=2.0; python_version < "3.11"
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
+Dynamic: license-file
 
 # kernels
 

{kernels-0.4.4 → kernels-0.5.0.dev0}/tests/test_layer.py

@@ -19,6 +19,12 @@ kernel_layer_mapping = {
             revision="layers",
         )
     },
+    "SiluAndMulNoCompile": {
+        "cuda": LayerRepository(
+            repo_id="kernels-test/op-without-fake-test",
+            layer_name="SiluAndMul",
+        )
+    },
     "SiluAndMulStringDevice": {
         "cuda": LayerRepository(
             repo_id="kernels-community/activation",
@@ -43,6 +49,11 @@ class SiluAndMul(nn.Module):
         return F.silu(input[..., :d]) * input[..., d:]
 
 
+@use_kernel_forward_from_hub("SiluAndMulNoCompile")
+class SiluAndMulNoCompileKernel(SiluAndMul):
+    pass
+
+
 @use_kernel_forward_from_hub("SiluAndMul")
 class SiluAndMulWithKernel(SiluAndMul):
     pass
@@ -101,8 +112,29 @@ def test_layer_fallback_works():
     SiluAndMulWithKernelFallback()
 
 
+@pytest.mark.parametrize("cls", [SiluAndMulWithKernel, SiluAndMulNoCompileKernel])
+@pytest.mark.parametrize("device", ["cuda", "cpu"])
+def test_torch_compile_layer(cls, device):
+    silu_and_mul = SiluAndMul()
+
+    X = torch.randn((32, 64), dtype=torch.float32, device=device)
+    Y = silu_and_mul(X)
+
+    silu_and_mul_with_kernel = cls()
+    silu_and_mul_with_kernel.eval()
+    silu_and_mul_compiled = torch.compile(silu_and_mul_with_kernel)
+
+    Y_compiled = silu_and_mul_compiled(X)
+
+    torch.testing.assert_close(Y_compiled, Y)
+
+
 def test_mapping_contexts():
-    assert set(_KERNEL_MAPPING.get().keys()) == {"SiluAndMul", "SiluAndMulStringDevice"}
+    assert set(_KERNEL_MAPPING.get().keys()) == {
+        "SiluAndMul",
+        "SiluAndMulStringDevice",
+        "SiluAndMulNoCompile",
+    }
 
     extra_mapping1 = {
         "TestKernel": {
@@ -118,6 +150,7 @@ def test_mapping_contexts():
         assert set(_KERNEL_MAPPING.get().keys()) == {
             "SiluAndMul",
             "SiluAndMulStringDevice",
+            "SiluAndMulNoCompile",
             "TestKernel",
         }
 
@@ -135,6 +168,7 @@ def test_mapping_contexts():
             assert set(_KERNEL_MAPPING.get().keys()) == {
                 "SiluAndMul",
                 "SiluAndMulStringDevice",
+                "SiluAndMulNoCompile",
                 "TestKernel",
             }
             assert (
@@ -145,6 +179,7 @@ def test_mapping_contexts():
         assert set(_KERNEL_MAPPING.get().keys()) == {
             "SiluAndMul",
             "SiluAndMulStringDevice",
+            "SiluAndMulNoCompile",
             "TestKernel",
         }
         assert (
@@ -164,6 +199,7 @@ def test_mapping_contexts():
         assert set(_KERNEL_MAPPING.get().keys()) == {
             "SiluAndMul",
             "SiluAndMulStringDevice",
+            "SiluAndMulNoCompile",
             "TestKernel",
         }
         assert (
@@ -174,6 +210,7 @@ def test_mapping_contexts():
     assert set(_KERNEL_MAPPING.get().keys()) == {
         "SiluAndMul",
         "SiluAndMulStringDevice",
+        "SiluAndMulNoCompile",
     }