From 5e7d3fab3783948ed08aa38c7ee1db548dda96cc Mon Sep 17 00:00:00 2001 From: medmekk Date: Mon, 8 Dec 2025 12:33:52 +0000 Subject: [PATCH 1/7] fix version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 4c99dc2..042dd38 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "kernels" -version = "0.11.1.dev0" +version = "0.11.4.dev0" description = "Download compute kernels" authors = [ { name = "OlivierDehaene", email = "olivier@huggingface.co" }, From 2b5059f8a4fb506680f3bc4dd6c49889b5694aa8 Mon Sep 17 00:00:00 2001 From: medmekk Date: Tue, 16 Dec 2025 16:10:24 +0000 Subject: [PATCH 2/7] initial commit --- src/kernels/cli.py | 250 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 250 insertions(+) diff --git a/src/kernels/cli.py b/src/kernels/cli.py index 0d95533..6d04d8a 100644 --- a/src/kernels/cli.py +++ b/src/kernels/cli.py @@ -15,6 +15,138 @@ BUILD_VARIANT_REGEX = re.compile(r"^(torch\d+\d+|torch-)") +# Default templates for kernel project initialization +DEFAULT_FLAKE_NIX = """\ +{ + description = "Flake for %(kernel_name)s kernel"; + inputs = { + kernel-builder.url = "github:huggingface/kernel-builder"; + }; + outputs = { self, kernel-builder }: + kernel-builder.lib.genFlakeOutputs { + path = ./.; + rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate; + }; +} +""" + +DEFAULT_BUILD_TOML = """\ +[general] +name = "%(kernel_name)s" +backends = ["cuda", "rocm", "metal", "xpu"] + +[torch] +src = [ + "torch-ext/torch_binding.cpp", + "torch-ext/torch_binding.h", +] + +[kernel.activation_cuda] +backend = "cuda" +# cuda-capabilities = ["9.0", "10.0", "12.0"] # if not specified, all capabilities will be used +depends = ["torch"] +src = [ + "%(kernel_name)s_cuda/kernel.cu", + "%(kernel_name)s_cuda/kernel.h" +] + +[kernel.activation_rocm] +backend = "rocm" +# rocm-archs = ["gfx906", "gfx908", "gfx90a", "gfx940", "gfx941", "gfx942", "gfx1030", "gfx1100", "gfx1101"] # if not specified, all architectures will be used +depends = ["torch"] +src = [ + "%(kernel_name)s_cuda/kernel.cu", + "%(kernel_name)s_cuda/kernel.h", +] + +[kernel.activation_xpu] +backend = "xpu" +depends = ["torch"] +src = [ + "%(kernel_name)s_xpu/kernel.cpp", + "%(kernel_name)s_xpu/kernel.hpp", +] + +[kernel.activation_metal] +backend = "metal" +depends = ["torch"] +src = [ + "%(kernel_name)s_metal/kernel.mm", + "%(kernel_name)s_metal/kernel.metal", +] +""" + +DEFAULT_GITATTRIBUTES = """\ +*.so filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +""" + +DEFAULT_README = """\ +# %(kernel_name)s + +A custom kernel for PyTorch. + +## Installation + +```bash +pip install kernels +``` + +## Usage + +```python +from kernels import get_kernel + +kernel = get_kernel("%(repo_id)s") + + +## License + +Apache-2.0 +""" + +DEFAULT_INIT_PY = """\ +# %(kernel_name)s kernel +# This file exports the kernel's public API + +import torch +from ._ops import ops + +def exported_kernel_function(x: torch.Tensor) -> torch.Tensor: + return ops.kernel_function(x) + +__all__ = ["exported_kernel_function"] +""" + +DEFAULT_TORCH_BINDING_CPP = """\ +#include + +#include "registration.h" +#include "torch_binding.h" + +TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + ops.def("kernel_function(Tensor input) -> (Tensor)"); +#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL) + ops.impl("kernel_function", torch::kCUDA, &kernel_function); +#elif defined(METAL_KERNEL) + ops.impl("kernel_function", torch::kMPS, kernel_function); +#elif defined(XPU_KERNEL) + ops.impl("kernel_function", torch::kXPU, &kernel_function); +#endif +} + +REGISTER_EXTENSION(TORCH_EXTENSION_NAME) + +""" + +DEFAULT_TORCH_BINDING_H = """\ +#pragma once + +#include + +torch::Tensor kernel_function(torch::Tensor input); +""" def main(): parser = argparse.ArgumentParser( @@ -113,6 +245,31 @@ def main(): ) ) + # Add init subcommand parser + init_parser = subparsers.add_parser( + "init", + help="Initialize a new kernel project structure", + ) + init_parser.add_argument( + "kernel_name", + type=str, + help="Name of the kernel (e.g., 'my-kernel')", + ) + init_parser.add_argument( + "--output-dir", + type=Path, + default=None, + help="Output directory for the kernel project (defaults to current directory)", + ) + init_parser.add_argument( + "--repo-id", + type=str, + default=None, + help="Repository ID for the kernel (e.g., 'kernels-community/my-kernel')", + ) + + init_parser.set_defaults(func=init_kernel_project) + args = parser.parse_args() args.func(args) @@ -237,3 +394,96 @@ def check_kernel( repo_id=repo_id, revision=revision, ) + + +def init_kernel_project(args): + """Initialize a new kernel project with the standard structure.""" + kernel_name = args.kernel_name + if "/" in kernel_name or "\\" in kernel_name: + raise ValueError("Kernel name cannot contain path separators, to specify an output directory use the --output-dir argument") + # Normalize kernel name (replace hyphens with underscores for Python compatibility) + kernel_name_normalized = kernel_name.replace("-", "_") + + # Determine output directory + if args.output_dir is not None: + output_dir = (Path(args.output_dir) / kernel_name).resolve() + else: + output_dir = Path.cwd() / kernel_name + + # Determine repo_id + repo_id = args.repo_id if args.repo_id else f"your-username/{kernel_name}" + + # Check if directory already exists + if output_dir.exists() and any(output_dir.iterdir()): + print( + f"Error: Directory '{output_dir}' already exists and is not empty.", + file=sys.stderr, + ) + sys.exit(1) + + # Create directory structure + dirs_to_create = [ + output_dir, + output_dir / f"{kernel_name_normalized}_cuda", + output_dir / f"{kernel_name_normalized}_rocm", + output_dir / f"{kernel_name_normalized}_metal", + output_dir / f"{kernel_name_normalized}_xpu", + output_dir / "torch-ext", + output_dir / "torch-ext" / kernel_name_normalized, + ] + + for dir_path in dirs_to_create: + dir_path.mkdir(parents=True, exist_ok=True) + + # Template substitution values + template_values = { + "kernel_name": kernel_name_normalized, + "repo_id": repo_id, + } + + # Create files + files_to_create = { + "flake.nix": DEFAULT_FLAKE_NIX % template_values, + "build.toml": DEFAULT_BUILD_TOML % template_values, + ".gitattributes": DEFAULT_GITATTRIBUTES, + "README.md": DEFAULT_README % template_values, + f"torch-ext/{kernel_name_normalized}/__init__.py": DEFAULT_INIT_PY % template_values, + f"torch-ext/torch_binding.cpp": DEFAULT_TORCH_BINDING_CPP % template_values, + f"torch-ext/torch_binding.h": DEFAULT_TORCH_BINDING_H % template_values, + } + + for file_path, content in files_to_create.items(): + full_path = output_dir / file_path + full_path.parent.mkdir(parents=True, exist_ok=True) + with open(full_path, "w") as f: + f.write(content) + + # Print success message + print(f"✅ Kernel project '{kernel_name}' initialized successfully at: {output_dir}") + print() + print("Project structure:") + _print_tree(output_dir, prefix="") + print() + print("Next steps:") + print(f" 1. cd {output_dir}") + print(f" 2. Add your kernel implementation in {kernel_name_normalized}/") + print(" 3. Update torch-ext/{kernel_name}/__init__.py to export your functions".format( + kernel_name=kernel_name_normalized + )) + print(" 4. Build with: nix run .#build-and-copy ") + print(f" 5. Upload with: kernels upload . --repo-id {repo_id}") + + +def _print_tree(directory: Path, prefix: str = ""): + """Print a directory tree structure.""" + entries = sorted(directory.iterdir(), key=lambda x: (x.is_file(), x.name)) + entries = [e for e in entries if not e.name.startswith(".git") or e.name == ".gitattributes"] + + for i, entry in enumerate(entries): + is_last = i == len(entries) - 1 + connector = "└── " if is_last else "├── " + print(f"{prefix}{connector}{entry.name}") + + if entry.is_dir(): + extension = " " if is_last else "│ " + _print_tree(entry, prefix + extension) From b9093ef44924592af5c9d3e73024a97efb377631 Mon Sep 17 00:00:00 2001 From: medmekk Date: Tue, 16 Dec 2025 16:15:11 +0000 Subject: [PATCH 3/7] fix style --- src/kernels/cli.py | 52 +++++++++++++++++++++++++++------------------- 1 file changed, 31 insertions(+), 21 deletions(-) diff --git a/src/kernels/cli.py b/src/kernels/cli.py index 6d04d8a..969a920 100644 --- a/src/kernels/cli.py +++ b/src/kernels/cli.py @@ -5,7 +5,7 @@ import sys from pathlib import Path -from huggingface_hub import create_repo, upload_folder, create_branch +from huggingface_hub import create_branch, create_repo, upload_folder from kernels.compat import tomllib from kernels.lockfile import KernelLock, get_kernel_locks @@ -55,7 +55,7 @@ # rocm-archs = ["gfx906", "gfx908", "gfx90a", "gfx940", "gfx941", "gfx942", "gfx1030", "gfx1100", "gfx1101"] # if not specified, all architectures will be used depends = ["torch"] src = [ - "%(kernel_name)s_cuda/kernel.cu", + "%(kernel_name)s_cuda/kernel.cu", "%(kernel_name)s_cuda/kernel.h", ] @@ -148,6 +148,7 @@ def exported_kernel_function(x: torch.Tensor) -> torch.Tensor: torch::Tensor kernel_function(torch::Tensor input); """ + def main(): parser = argparse.ArgumentParser( prog="kernel", description="Manage compute kernels" @@ -400,19 +401,21 @@ def init_kernel_project(args): """Initialize a new kernel project with the standard structure.""" kernel_name = args.kernel_name if "/" in kernel_name or "\\" in kernel_name: - raise ValueError("Kernel name cannot contain path separators, to specify an output directory use the --output-dir argument") + raise ValueError( + "Kernel name cannot contain path separators, to specify an output directory use the --output-dir argument" + ) # Normalize kernel name (replace hyphens with underscores for Python compatibility) kernel_name_normalized = kernel_name.replace("-", "_") - + # Determine output directory if args.output_dir is not None: output_dir = (Path(args.output_dir) / kernel_name).resolve() else: output_dir = Path.cwd() / kernel_name - + # Determine repo_id repo_id = args.repo_id if args.repo_id else f"your-username/{kernel_name}" - + # Check if directory already exists if output_dir.exists() and any(output_dir.iterdir()): print( @@ -420,7 +423,7 @@ def init_kernel_project(args): file=sys.stderr, ) sys.exit(1) - + # Create directory structure dirs_to_create = [ output_dir, @@ -431,33 +434,34 @@ def init_kernel_project(args): output_dir / "torch-ext", output_dir / "torch-ext" / kernel_name_normalized, ] - + for dir_path in dirs_to_create: dir_path.mkdir(parents=True, exist_ok=True) - + # Template substitution values template_values = { "kernel_name": kernel_name_normalized, "repo_id": repo_id, } - + # Create files files_to_create = { "flake.nix": DEFAULT_FLAKE_NIX % template_values, "build.toml": DEFAULT_BUILD_TOML % template_values, ".gitattributes": DEFAULT_GITATTRIBUTES, "README.md": DEFAULT_README % template_values, - f"torch-ext/{kernel_name_normalized}/__init__.py": DEFAULT_INIT_PY % template_values, - f"torch-ext/torch_binding.cpp": DEFAULT_TORCH_BINDING_CPP % template_values, - f"torch-ext/torch_binding.h": DEFAULT_TORCH_BINDING_H % template_values, + f"torch-ext/{kernel_name_normalized}/__init__.py": DEFAULT_INIT_PY + % template_values, + "torch-ext/torch_binding.cpp": DEFAULT_TORCH_BINDING_CPP % template_values, + "torch-ext/torch_binding.h": DEFAULT_TORCH_BINDING_H % template_values, } - + for file_path, content in files_to_create.items(): full_path = output_dir / file_path full_path.parent.mkdir(parents=True, exist_ok=True) with open(full_path, "w") as f: f.write(content) - + # Print success message print(f"✅ Kernel project '{kernel_name}' initialized successfully at: {output_dir}") print() @@ -467,9 +471,11 @@ def init_kernel_project(args): print("Next steps:") print(f" 1. cd {output_dir}") print(f" 2. Add your kernel implementation in {kernel_name_normalized}/") - print(" 3. Update torch-ext/{kernel_name}/__init__.py to export your functions".format( - kernel_name=kernel_name_normalized - )) + print( + " 3. Update torch-ext/{kernel_name}/__init__.py to export your functions".format( + kernel_name=kernel_name_normalized + ) + ) print(" 4. Build with: nix run .#build-and-copy ") print(f" 5. Upload with: kernels upload . --repo-id {repo_id}") @@ -477,13 +483,17 @@ def init_kernel_project(args): def _print_tree(directory: Path, prefix: str = ""): """Print a directory tree structure.""" entries = sorted(directory.iterdir(), key=lambda x: (x.is_file(), x.name)) - entries = [e for e in entries if not e.name.startswith(".git") or e.name == ".gitattributes"] - + entries = [ + e + for e in entries + if not e.name.startswith(".git") or e.name == ".gitattributes" + ] + for i, entry in enumerate(entries): is_last = i == len(entries) - 1 connector = "└── " if is_last else "├── " print(f"{prefix}{connector}{entry.name}") - + if entry.is_dir(): extension = " " if is_last else "│ " _print_tree(entry, prefix + extension) From 5f9de2c5edc3a6e2292dce111df8b9e0718f522d Mon Sep 17 00:00:00 2001 From: medmekk Date: Tue, 16 Dec 2025 16:31:22 +0000 Subject: [PATCH 4/7] black --- src/kernels/layer/func.py | 3 ++- src/kernels/layer/layer.py | 3 ++- src/kernels/layer/repos.py | 6 ++++-- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/kernels/layer/func.py b/src/kernels/layer/func.py index e2f1127..bd4724a 100644 --- a/src/kernels/layer/func.py +++ b/src/kernels/layer/func.py @@ -21,7 +21,8 @@ class FuncRepositoryProtocol(RepositoryProtocol, Protocol): @property - def func_name(self) -> str: ... + def func_name(self) -> str: + ... class FuncRepository: diff --git a/src/kernels/layer/layer.py b/src/kernels/layer/layer.py index 9965c64..579c464 100644 --- a/src/kernels/layer/layer.py +++ b/src/kernels/layer/layer.py @@ -32,7 +32,8 @@ class LayerRepositoryProtocol(RepositoryProtocol, Protocol): @property - def layer_name(self) -> str: ... + def layer_name(self) -> str: + ... class LayerRepository: diff --git a/src/kernels/layer/repos.py b/src/kernels/layer/repos.py index a947a15..67aabce 100644 --- a/src/kernels/layer/repos.py +++ b/src/kernels/layer/repos.py @@ -13,7 +13,8 @@ class RepositoryProtocol(Protocol): - def load(self) -> Type["nn.Module"]: ... + def load(self) -> Type["nn.Module"]: + ... class DeviceRepos(ABC): @@ -43,7 +44,8 @@ def create_repo(device: Device) -> "DeviceRepos": @abstractmethod def repos( self, - ) -> Optional[Dict[Mode, RepositoryProtocol]]: ... + ) -> Optional[Dict[Mode, RepositoryProtocol]]: + ... @abstractmethod def insert(self, device: Device, repos: Dict[Mode, RepositoryProtocol]): From ba65db1c4966dc89839ce1ee7565bf8cf8fa0517 Mon Sep 17 00:00:00 2001 From: medmekk Date: Tue, 16 Dec 2025 16:37:00 +0000 Subject: [PATCH 5/7] revert --- src/kernels/layer/func.py | 3 +-- src/kernels/layer/layer.py | 4 ++-- src/kernels/layer/repos.py | 4 ++-- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/src/kernels/layer/func.py b/src/kernels/layer/func.py index bd4724a..e2f1127 100644 --- a/src/kernels/layer/func.py +++ b/src/kernels/layer/func.py @@ -21,8 +21,7 @@ class FuncRepositoryProtocol(RepositoryProtocol, Protocol): @property - def func_name(self) -> str: - ... + def func_name(self) -> str: ... class FuncRepository: diff --git a/src/kernels/layer/layer.py b/src/kernels/layer/layer.py index 579c464..0b03b42 100644 --- a/src/kernels/layer/layer.py +++ b/src/kernels/layer/layer.py @@ -32,8 +32,8 @@ class LayerRepositoryProtocol(RepositoryProtocol, Protocol): @property - def layer_name(self) -> str: - ... + def layer_name(self) -> str: ... + class LayerRepository: diff --git a/src/kernels/layer/repos.py b/src/kernels/layer/repos.py index 67aabce..e4064b7 100644 --- a/src/kernels/layer/repos.py +++ b/src/kernels/layer/repos.py @@ -13,8 +13,8 @@ class RepositoryProtocol(Protocol): - def load(self) -> Type["nn.Module"]: - ... + def load(self) -> Type["nn.Module"]: ... + class DeviceRepos(ABC): From 882fca171c16cda0749b455a6f10bf3540f677e6 Mon Sep 17 00:00:00 2001 From: medmekk Date: Tue, 16 Dec 2025 16:40:09 +0000 Subject: [PATCH 6/7] fix --- src/kernels/layer/layer.py | 1 - src/kernels/layer/repos.py | 4 +--- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/src/kernels/layer/layer.py b/src/kernels/layer/layer.py index 0b03b42..9965c64 100644 --- a/src/kernels/layer/layer.py +++ b/src/kernels/layer/layer.py @@ -35,7 +35,6 @@ class LayerRepositoryProtocol(RepositoryProtocol, Protocol): def layer_name(self) -> str: ... - class LayerRepository: """ Repository and name of a layer for kernel mapping. diff --git a/src/kernels/layer/repos.py b/src/kernels/layer/repos.py index e4064b7..a947a15 100644 --- a/src/kernels/layer/repos.py +++ b/src/kernels/layer/repos.py @@ -16,7 +16,6 @@ class RepositoryProtocol(Protocol): def load(self) -> Type["nn.Module"]: ... - class DeviceRepos(ABC): """ Device-specific kernel layer repositories. @@ -44,8 +43,7 @@ def create_repo(device: Device) -> "DeviceRepos": @abstractmethod def repos( self, - ) -> Optional[Dict[Mode, RepositoryProtocol]]: - ... + ) -> Optional[Dict[Mode, RepositoryProtocol]]: ... @abstractmethod def insert(self, device: Device, repos: Dict[Mode, RepositoryProtocol]): From 665a8e25b63469a06ff546b5ba7a33e5a8772dc1 Mon Sep 17 00:00:00 2001 From: medmekk Date: Tue, 16 Dec 2025 16:43:46 +0000 Subject: [PATCH 7/7] fix black version --- src/kernels/cli.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/kernels/cli.py b/src/kernels/cli.py index 969a920..04ab1eb 100644 --- a/src/kernels/cli.py +++ b/src/kernels/cli.py @@ -463,7 +463,9 @@ def init_kernel_project(args): f.write(content) # Print success message - print(f"✅ Kernel project '{kernel_name}' initialized successfully at: {output_dir}") + print( + f"✅ Kernel project '{kernel_name}' initialized successfully at: {output_dir}" + ) print() print("Project structure:") _print_tree(output_dir, prefix="")