Merge branch 'main' into cleanup-huggingface-hub-integration

Author: Wauplin

flake.nix

@@ -158,8 +158,12 @@
               pytest
               pytest-benchmark
               pyyaml
+              tabulate
+              tomlkit
               torch
               types-pyyaml
+              types-requests
+              types-tabulate
               venvShellHook
             ]);
 

kernels/src/kernels/cli.py kernels/src/kernels/cli/__init__.pykernels/src/kernels/cli.py renamed to kernels/src/kernels/cli/__init__.py

@@ -6,16 +6,15 @@
 
 from kernels.compat import tomllib
 from kernels.lockfile import KernelLock, get_kernel_locks
-from kernels.upload import upload_kernels_dir
+from kernels.cli.upload import upload_kernels_dir
 from kernels.utils import (
     install_kernel,
     install_kernel_all_variants,
     KNOWN_BACKENDS,
 )
-from kernels.versions_cli import print_kernel_versions
-from kernels.init import run_init, parse_kernel_name
-
-from .doc import generate_readme_for_kernel
+from kernels.cli.init import run_init, parse_kernel_name
+from kernels.cli.versions import print_kernel_versions
+from kernels.cli.doc import generate_readme_for_kernel
 
 
 def main():
@@ -260,15 +259,15 @@ def check_kernel(
     *, macos: str, manylinux: str, python_abi: str, repo_id: str, revision: str
 ):
     try:
-        import kernels.check
+        from kernels.cli import check
     except ImportError:
         print(
             "`kernels check` requires the `kernel-abi-check` package: pip install kernel-abi-check",
             file=sys.stderr,
         )
         sys.exit(1)
 
-    kernels.check.check_kernel(
+    check.check_kernel(
         macos=macos,
         manylinux=manylinux,
         python_abi=python_abi,
@@ -278,7 +277,7 @@ def check_kernel(
 
 
 def run_benchmark(args):
-    from kernels import benchmark
+    from kernels.cli import benchmark
 
     benchmark.run_benchmark(
         repo_id=args.repo_id,

…/kernels/_vendored/convert_rst_to_mdx.py …nels/cli/_vendored/convert_rst_to_mdx.pykernels/src/kernels/_vendored/convert_rst_to_mdx.py renamed to kernels/src/kernels/cli/_vendored/convert_rst_to_mdx.py kernels/src/kernels/_vendored/convert_rst_to_mdx.py renamed to kernels/src/kernels/cli/_vendored/convert_rst_to_mdx.py

@@ -14,6 +14,7 @@
 
 from huggingface_hub.utils import build_hf_headers, disable_progress_bars, get_session, hf_raise_for_status
 
+from kernels.benchmark import Benchmark
 from kernels.utils import _get_hf_api, backend
 
 MISSING_DEPS: list[str] = []
@@ -63,43 +64,6 @@ def _calculate_iqr_and_outliers(
     return q1, q3, iqr, outliers
 
 
-class Benchmark:
-    """Base class for kernel benchmarks.
-
-    Subclass this to create a benchmark script with automatic timing,
-    verification, and reproducibility support. The kernel is loaded
-    automatically from the repo_id specified in the CLI command.
-
-    Example:
-        class MyBenchmark(Benchmark):
-            seed = 42
-
-            def setup(self):
-                self.x = torch.randn(128, 1024, device=self.device, dtype=torch.float16)
-                self.out = torch.empty(128, 512, device=self.device, dtype=torch.float16)
-
-            def benchmark_silu(self):
-                self.kernel.silu_and_mul(self.out, self.x)
-
-            def verify_silu(self) -> torch.Tensor:
-                # Return reference tensor; runner compares with self.out
-                return torch.nn.functional.silu(self.x[..., :512]) * self.x[..., 512:]
-
-    Run with: kernels benchmark <repo_id>
-    """
-
-    seed: int | None = None  # Optional: seed for reproducibility
-    device: str = "cpu"  # Set automatically by runner
-
-    def __init__(self) -> None:
-        self.kernel: Any = None
-        self.out: Any = None  # Output tensor, set by setup methods
-
-    def setup(self) -> None:
-        """Override to set up tensors as instance attributes."""
-        pass
-
-
 @dataclass
 class TimingResults:
     mean_ms: float
@@ -319,9 +283,7 @@ def _get_macos_gpu() -> tuple[str | None, int | None]:
         from ctypes import POINTER, byref, c_char_p, c_int, c_int64, c_uint32, c_void_p
 
         iokit = ctypes.CDLL("/System/Library/Frameworks/IOKit.framework/IOKit")
-        cf = ctypes.CDLL(
-            "/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation"
-        )
+        cf = ctypes.CDLL("/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation")
 
         iokit.IOServiceMatching.restype = c_void_p
         iokit.IOServiceMatching.argtypes = [c_char_p]
@@ -382,9 +344,7 @@ def _get_macos_gpu() -> tuple[str | None, int | None]:
             cf.CFRelease(key)
 
         # Get GPU core count
-        key = cf.CFStringCreateWithCString(
-            None, b"gpu-core-count", kCFStringEncodingUTF8
-        )
+        key = cf.CFStringCreateWithCString(None, b"gpu-core-count", kCFStringEncodingUTF8)
         if key:
             prop = iokit.IORegistryEntryCreateCFProperty(service, key, None, 0)
             if prop:
@@ -425,9 +385,7 @@ def collect_machine_info() -> MachineInfo:
             if hasattr(torch.version, "hip") and torch.version.hip:
                 backend_type = f"ROCm {torch.version.hip}"
             else:
-                backend_type = (
-                    f"CUDA {torch.version.cuda}" if torch.version.cuda else "CUDA"
-                )
+                backend_type = f"CUDA {torch.version.cuda}" if torch.version.cuda else "CUDA"
         elif backend_name == "xpu":
             gpu = torch.xpu.get_device_name(0)
             backend_type = "XPU"
@@ -479,16 +437,14 @@ def run_benchmark_class(
 
     # Find all benchmark_* methods
     benchmark_methods = [
-        name
-        for name in dir(benchmark_cls)
-        if name.startswith("benchmark_") and callable(getattr(benchmark_cls, name))
+        name for name in dir(benchmark_cls) if name.startswith("benchmark_") and callable(getattr(benchmark_cls, name))
     ]
 
     if not benchmark_methods:
         raise RuntimeError(f"No benchmark_* methods found in {benchmark_cls.__name__}")
 
     # Load kernel once for all workloads
-    from kernels import get_local_kernel, get_kernel
+    from kernels import get_kernel, get_local_kernel
 
     if is_local:
         kernel = get_local_kernel(Path(repo_id), "activation")
@@ -663,9 +619,7 @@ def run_benchmark_script(
         raise RuntimeError(f"No Benchmark subclasses found in {script_path}")
 
     machine_info = collect_machine_info()
-    gpu_cores_str = (
-        f" ({machine_info.gpu_cores} cores)" if machine_info.gpu_cores else ""
-    )
+    gpu_cores_str = f" ({machine_info.gpu_cores} cores)" if machine_info.gpu_cores else ""
     print(file=sys.stderr)
     print(f"  GPU      {machine_info.gpu}{gpu_cores_str}", file=sys.stderr)
     print(f"  CPU      {machine_info.cpu}", file=sys.stderr)
@@ -736,8 +690,7 @@ def run_benchmark(
     if is_local:
         if repo_id.count("/") == 1 and not repo_id.startswith(("./", "../")):
             warnings.warn(
-                f"'{repo_id}' exists locally but looks like a repo_id. "
-                f"Use './{repo_id}' to be explicit.",
+                f"'{repo_id}' exists locally but looks like a repo_id. Use './{repo_id}' to be explicit.",
                 stacklevel=2,
             )
         branch = "local"
@@ -765,7 +718,6 @@ def run_benchmark(
     assert revision is not None  # Guaranteed by parsing logic above
 
     print(f"Downloading {repo_id}@{revision}...", file=sys.stderr)
-
     if is_local:
         repo_path = repo_id_path.resolve()
     else:

kernels/src/kernels/benchmark.py kernels/src/kernels/cli/benchmark.pykernels/src/kernels/benchmark.py renamed to kernels/src/kernels/cli/benchmark.py

@@ -6,7 +6,7 @@
 import yaml
 
 from ._vendored.convert_rst_to_mdx import convert_rst_docstring_to_mdx
-from .utils import get_kernel
+from kernels.utils import get_kernel
 
 _RE_PARAMETERS = re.compile(
     r"<parameters>(((?!<parameters>).)*)</parameters>", re.DOTALL

kernels/src/kernels/check.py kernels/src/kernels/cli/check.pykernels/src/kernels/check.py renamed to kernels/src/kernels/cli/check.py

@@ -3,9 +3,10 @@
 import argparse
 import os
 
-from kernels.init import run_init, parse_kernel_name
+from kernels.cli.init import run_init, parse_kernel_name
 from kernels.utils import KNOWN_BACKENDS
 
+
 def e2e_init(backends: list[str]) -> None:
     kernel_name = "testuser/test-kernel"
     template_repo = "drbh/template"
@@ -16,7 +17,9 @@ def e2e_init(backends: list[str]) -> None:
         overwrite=False,
     )
     expected_normalized_name = "test_kernel"
-    expected_backend_dirs = {Path(f"{expected_normalized_name}_{backend}") for backend in args.backends}
+    expected_backend_dirs = {
+        Path(f"{expected_normalized_name}_{backend}") for backend in args.backends
+    }
 
     # Replacement logic
     # special case for "rocm" backend since it uses "cuda" source