Prefer ruff formatting (#462)

* feat: prefer ruff formatting

* fix: run formatting

* fix: bump local and ci ruff version and reformat

* fix: bump python_depends check

* fix: run griffe with uvx

* feat: add make quality check

* nix: use ruff 0.15.10

---------

Co-authored-by: Daniël de Kok <me@danieldk.eu>

Author: drbh

.github/workflows/lint.yml

@@ -6,46 +6,39 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
-      - name: Run ruff
+      - name: Run ruff check
         uses: astral-sh/ruff-action@4919ec5cf1f49eff0871dbcea0da843445b837e6  # v3
-        src: >-
-          kernels
+        with:
+          src: kernels
+          version: "0.15.10"
+      - name: Run ruff format check
+        uses: astral-sh/ruff-action@4919ec5cf1f49eff0871dbcea0da843445b837e6  # v3
+        with:
+          src: kernels
+          version: "0.15.10"
+          args: format --check
 
-  black:
-    name: Run black check
+  griffe:
+    name: Check API compatibility
     runs-on: ubuntu-latest
     env:
       UV_PYTHON_PREFERENCE: only-managed
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
-
+        with:
+          fetch-depth: 0
       - name: Install uv and set the python version
         uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78  # v7
         with:
-          python-version: 3.12
-
-      - name: Install black
-        run: uv pip install black
-
-      - name: Check formatting
-        run: |
-          uv run black --check kernels
+          python-version: "3.12"
+      - name: Check for breaking changes
+        run: uvx griffe check kernels --search kernels/src -a main
 
   validate-dependencies:
     name: Validate python_depends.json
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
 
-      - name: Set up Python
-        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405  # v6
-        with:
-          python-version: "3.12"
-
       - name: Validate python_depends.json is up-to-date
-        run: |
-          python ( cd kernels && update_python_depends.py --validate ) || {
-            echo "Error: python_depends.json is out of date."
-            echo "Please run: python update_python_depends.py"
-            exit 1
-          }
+        run: diff kernels-data/src/python_dependencies.json kernels/src/kernels/python_depends.json

Makefile

@@ -1,15 +1,15 @@
-.PHONY: style kernel-builder-cli-docs
+.PHONY: style kernel-builder-cli-docs quality
+
 
 export check_dirs := kernels/src kernels/tests
 
 all: src/kernels/python_depends.json
 
-kernels/src/kernels/python_depends.json: kernel-builder/src/python_dependencies.json
+kernels/src/kernels/python_depends.json: kernels-data/src/python_dependencies.json
 	cp $< $@
 
 style:
-	black ${check_dirs}
-	isort ${check_dirs}
+	ruff format ${check_dirs}
 	ruff check ${check_dirs} --fix
 
 kernel-builder-cli-docs:
@@ -20,3 +20,7 @@ kernel-builder-cli-docs:
 	  | sed '/`--backends/,/^\*/{/^  Default value:/d;}' \
 	  > docs/source/builder-cli.md
 	@echo "Generated docs/source/builder-cli.md"
+
+quality:
+	ruff format --check ${check_dirs}
+	ruff check ${check_dirs}

kernels/pyproject.toml

@@ -56,11 +56,10 @@ kernels = "kernels.cli:main"
 [tool.setuptools.package-data]
 kernels = ["python_depends.json"]
 
-[tool.isort]
-profile = "black"
-line_length = 119
-
 [tool.ruff]
+# If the version is changed, apply the change in the Nix overlay
+# as well.
+required-version = "==0.15.10"
 exclude = [
   ".eggs",
   ".git",
@@ -85,4 +84,6 @@ line-length = 119
 # Ignored rules:
 # "E501" -> line length violation
 lint.ignore = ["E501"]
-lint.select = ["E", "F", "W"]
+lint.select = ["E", "F", "I", "W"]
+
+[tool.ruff.format]

kernels/src/kernels/_versions.py

@@ -82,9 +82,7 @@ def resolve_version_spec_as_ref(repo_id: str, version_spec: int | str) -> GitRef
         accepted_versions = sorted(requirement.filter(versions_old.keys()))
 
         if len(accepted_versions) == 0:
-            raise ValueError(
-                f"No version of `{repo_id}` satisfies requirement: {version_spec}"
-            )
+            raise ValueError(f"No version of `{repo_id}` satisfies requirement: {version_spec}")
 
         return versions_old[accepted_versions[-1]]
 

kernels/src/kernels/backends.py

@@ -241,9 +241,7 @@ def _select_backend(backend: str | None) -> Backend:
     if backend in supported:
         return supported[backend]
 
-    raise ValueError(
-        f"Invalid backend '{backend}', system supported backends: {', '.join(sorted(supported.keys()))}"
-    )
+    raise ValueError(f"Invalid backend '{backend}', system supported backends: {', '.join(sorted(supported.keys()))}")
 
 
 def _supported_backends() -> dict[str, Backend]:
@@ -267,9 +265,7 @@ def _get_cuda() -> Optional[CUDA]:
     runtime_version = ctypes.c_int(0)
     result = libcudart.cudaRuntimeGetVersion(ctypes.byref(runtime_version))
     if result != 0:
-        warnings.warn(
-            "System has CUDA runtime library, but cannot get runtime version."
-        )
+        warnings.warn("System has CUDA runtime library, but cannot get runtime version.")
         return None
 
     # cudaRuntimeGetVersion encodes the version as (major * 1000 + minor * 10).

kernels/src/kernels/benchmarks/attention.py

@@ -14,19 +14,15 @@ def _reference_attention(query, key, value, causal=False):
     """Reference implementation using PyTorch SDPA."""
     query, key, value = (x.transpose(1, 2).contiguous() for x in (query, key, value))
     with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
-        out = torch.nn.functional.scaled_dot_product_attention(
-            query, key, value, is_causal=causal
-        )
+        out = torch.nn.functional.scaled_dot_product_attention(query, key, value, is_causal=causal)
     return out.transpose(1, 2).contiguous()
 
 
 def _varlen_reference_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, causal=False):
     """Reference implementation for variable length attention."""
     batch_size = cu_seqlens_q.shape[0] - 1
     total_tokens_q = q.shape[0]
-    out = torch.zeros(
-        (total_tokens_q, q.shape[1], q.shape[2]), device=q.device, dtype=q.dtype
-    )
+    out = torch.zeros((total_tokens_q, q.shape[1], q.shape[2]), device=q.device, dtype=q.dtype)
 
     for b in range(batch_size):
         start_q, end_q = cu_seqlens_q[b], cu_seqlens_q[b + 1]
@@ -54,9 +50,7 @@ def setup_small(self):
         self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
 
     def benchmark_small(self):
-        self.out = _extract_output(
-            self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False)
-        )
+        self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False))
 
     def verify_small(self) -> torch.Tensor:
         return _reference_attention(self.q, self.k, self.v, causal=False)
@@ -70,9 +64,7 @@ def setup_medium(self):
         self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
 
     def benchmark_medium(self):
-        self.out = _extract_output(
-            self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False)
-        )
+        self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False))
 
     def verify_medium(self) -> torch.Tensor:
         return _reference_attention(self.q, self.k, self.v, causal=False)
@@ -86,9 +78,7 @@ def setup_large(self):
         self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
 
     def benchmark_large(self):
-        self.out = _extract_output(
-            self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False)
-        )
+        self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=False))
 
     def verify_large(self) -> torch.Tensor:
         return _reference_attention(self.q, self.k, self.v, causal=False)
@@ -106,9 +96,7 @@ def setup_small(self):
         self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
 
     def benchmark_small(self):
-        self.out = _extract_output(
-            self.kernel.flash_attn_func(self.q, self.k, self.v, causal=True)
-        )
+        self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=True))
 
     def verify_small(self) -> torch.Tensor:
         return _reference_attention(self.q, self.k, self.v, causal=True)
@@ -122,9 +110,7 @@ def setup_medium(self):
         self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
 
     def benchmark_medium(self):
-        self.out = _extract_output(
-            self.kernel.flash_attn_func(self.q, self.k, self.v, causal=True)
-        )
+        self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=True))
 
     def verify_medium(self) -> torch.Tensor:
         return _reference_attention(self.q, self.k, self.v, causal=True)
@@ -138,9 +124,7 @@ def setup_large(self):
         self.out = torch.empty(B, S, H, D, device="cuda", dtype=torch.float16)
 
     def benchmark_large(self):
-        self.out = _extract_output(
-            self.kernel.flash_attn_func(self.q, self.k, self.v, causal=True)
-        )
+        self.out = _extract_output(self.kernel.flash_attn_func(self.q, self.k, self.v, causal=True))
 
     def verify_large(self) -> torch.Tensor:
         return _reference_attention(self.q, self.k, self.v, causal=True)
@@ -180,9 +164,7 @@ def benchmark_small(self):
         )
 
     def verify_small(self) -> torch.Tensor:
-        return _varlen_reference_attention(
-            self.q, self.k, self.v, self.cu_seqlens, self.cu_seqlens, causal=False
-        )
+        return _varlen_reference_attention(self.q, self.k, self.v, self.cu_seqlens, self.cu_seqlens, causal=False)
 
     # Workload: medium (5 sequences, max_seqlen=256)
     def setup_medium(self):
@@ -214,9 +196,7 @@ def benchmark_medium(self):
         )
 
     def verify_medium(self) -> torch.Tensor:
-        return _varlen_reference_attention(
-            self.q, self.k, self.v, self.cu_seqlens, self.cu_seqlens, causal=False
-        )
+        return _varlen_reference_attention(self.q, self.k, self.v, self.cu_seqlens, self.cu_seqlens, causal=False)
 
     # Workload: large (8 sequences, max_seqlen=512)
     def setup_large(self):
@@ -248,6 +228,4 @@ def benchmark_large(self):
         )
 
     def verify_large(self) -> torch.Tensor:
-        return _varlen_reference_attention(
-            self.q, self.k, self.v, self.cu_seqlens, self.cu_seqlens, causal=False
-        )
+        return _varlen_reference_attention(self.q, self.k, self.v, self.cu_seqlens, self.cu_seqlens, causal=False)

kernels/src/kernels/benchmarks/layer_norm.py

@@ -129,9 +129,7 @@ def benchmark_small(self):
         )[0].view(self.B, self.S, self.D)
 
     def verify_small(self) -> torch.Tensor:
-        return torch.nn.functional.layer_norm(
-            self.x, [self.D], self.weight, eps=self.eps
-        )
+        return torch.nn.functional.layer_norm(self.x, [self.D], self.weight, eps=self.eps)
 
     # Workload: medium (B=4, S=512, D=2048)
     def setup_medium(self):
@@ -160,9 +158,7 @@ def benchmark_medium(self):
         )[0].view(self.B, self.S, self.D)
 
     def verify_medium(self) -> torch.Tensor:
-        return torch.nn.functional.layer_norm(
-            self.x, [self.D], self.weight, eps=self.eps
-        )
+        return torch.nn.functional.layer_norm(self.x, [self.D], self.weight, eps=self.eps)
 
     # Workload: large (B=8, S=1024, D=4096)
     def setup_large(self):
@@ -191,6 +187,4 @@ def benchmark_large(self):
         )[0].view(self.B, self.S, self.D)
 
     def verify_large(self) -> torch.Tensor:
-        return torch.nn.functional.layer_norm(
-            self.x, [self.D], self.weight, eps=self.eps
-        )
+        return torch.nn.functional.layer_norm(self.x, [self.D], self.weight, eps=self.eps)

kernels/src/kernels/cli/__init__.py

@@ -15,9 +15,7 @@
 
 
 def main():
-    parser = argparse.ArgumentParser(
-        prog="kernel", description="Manage compute kernels"
-    )
+    parser = argparse.ArgumentParser(prog="kernel", description="Manage compute kernels")
     subparsers = parser.add_subparsers(required=True)
 
     check_parser = subparsers.add_parser("check", help="Check a kernel for compliance")
@@ -29,12 +27,8 @@ def main():
         help="The kernel revision (branch, tag, or commit SHA, defaults to 'main')",
     )
     check_parser.add_argument("--macos", type=str, help="macOS version", default="15.0")
-    check_parser.add_argument(
-        "--manylinux", type=str, help="Manylinux version", default="manylinux_2_28"
-    )
-    check_parser.add_argument(
-        "--python-abi", type=str, help="Python ABI version", default="3.9"
-    )
+    check_parser.add_argument("--manylinux", type=str, help="Manylinux version", default="manylinux_2_28")
+    check_parser.add_argument("--python-abi", type=str, help="Python ABI version", default="3.9")
     check_parser.set_defaults(
         func=lambda args: check_kernel(
             macos=args.macos,
@@ -107,12 +101,8 @@ def main():
         type=str,
         help="Kernel repo ID (e.g., kernels-community/activation)",
     )
-    benchmark_parser.add_argument(
-        "--branch", type=str, help="Kernel branch to benchmark"
-    )
-    benchmark_parser.add_argument(
-        "--version", type=int, help="Kernel version to benchmark"
-    )
+    benchmark_parser.add_argument("--branch", type=str, help="Kernel branch to benchmark")
+    benchmark_parser.add_argument("--version", type=int, help="Kernel version to benchmark")
     benchmark_parser.add_argument(
         "--output",
         type=str,
@@ -230,9 +220,7 @@ def default(self, o):
         return super().default(o)
 
 
-def check_kernel(
-    *, macos: str, manylinux: str, python_abi: str, repo_id: str, revision: str
-):
+def check_kernel(*, macos: str, manylinux: str, python_abi: str, repo_id: str, revision: str):
     try:
         from kernels.cli import check
     except ImportError: