kernelport 0.1.0__py3-none-any.whl

kernelport/backends/base.py

@@ -0,0 +1,138 @@
+"""Backend interface and registry.
+
+A backend lives in ``kernelport/backends/<name>/`` with three modules:
+
+- ``config.py``     CONFIG: BackendConfig (no heavy imports)
+- ``capability.py`` is_available() -> (bool, reason) and hardware_name() -> str
+                    (must import cleanly even when the backend's packages are missing)
+- ``runner.py``     prepare(port_path, spec, data) -> Prepared  (imports heavy packages lazily)
+
+The registry imports config and capability eagerly-on-demand and the runner
+only when something is actually run, so ``kernelport`` works on machines that
+have none of the GPU packages.
+"""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable
+
+from kernelport.harness.tolerances import Tolerance
+
+BACKEND_NAMES = ("triton", "cpu", "metal", "webgpu", "hip")
+
+
+@dataclass(frozen=True)
+class BackendConfig:
+    """Static description of a backend."""
+
+    name: str
+    extension: str  # port file extension, e.g. ".py"
+    description: str
+    dtypes: tuple[str, ...]  # output/compute dtypes this backend can run
+    tolerances: dict[str, Tolerance] = field(default_factory=dict)  # overrides, each with a reason
+    timing_comparable: bool = True  # False for CPU / browser numbers
+    timing_note: str = ""
+    hardware_hint: str = ""  # where verification can happen, e.g. "Apple Silicon Mac"
+    port_template: str = ""  # skeleton shown to agents
+    tune_space: dict[str, list[Any]] = field(default_factory=dict)
+
+
+@dataclass
+class Prepared:
+    """A port loaded and bound to one case's inputs, ready to launch repeatedly."""
+
+    launch: Callable[[], None]  # run once (may be async on GPUs)
+    sync: Callable[[], None]  # wait for completion
+    fetch: Callable[[], dict[str, Any]]  # output tensors as CPU torch tensors
+    timer: str = "wall_clock"  # "cuda_event" when launch runs on the current torch CUDA stream
+
+
+class Backend:
+    """Facade over a backend's config, capability and runner modules."""
+
+    def __init__(self, name: str) -> None:
+        if name not in BACKEND_NAMES:
+            raise KeyError(f"unknown backend {name!r}; known: {', '.join(BACKEND_NAMES)}")
+        self.name = name
+        self._pkg = f"kernelport.backends.{name}"
+        self.config: BackendConfig = importlib.import_module(f"{self._pkg}.config").CONFIG
+        self._cap = importlib.import_module(f"{self._pkg}.capability")
+
+    def is_available(self) -> tuple[bool, str]:
+        """(available, reason). Never raises."""
+        try:
+            return self._cap.is_available()
+        except Exception as exc:  # a broken install must not crash callers
+            return False, f"capability check failed: {type(exc).__name__}: {exc}"
+
+    def hardware_name(self) -> str:
+        """Name of the device ports run on here."""
+        try:
+            return self._cap.hardware_name()
+        except Exception:
+            return "unknown"
+
+    def prepare(self, port: Path, spec, data) -> Prepared:
+        """Load ``port`` and bind it to one case."""
+        runner = importlib.import_module(f"{self._pkg}.runner")
+        return runner.prepare(Path(port), spec, data)
+
+    def run(self, port: Path, spec, data) -> dict[str, Any]:
+        """Run once and return outputs as CPU tensors."""
+        p = self.prepare(port, spec, data)
+        p.launch()
+        p.sync()
+        return p.fetch()
+
+    def dtypes(self) -> tuple[str, ...]:
+        """dtypes runnable here: the config's, narrowed by the capability module if it knows better
+        (e.g. float16 only on WebGPU adapters with shader-f16)."""
+        fn = getattr(self._cap, "supported_dtypes", None)
+        if fn is None:
+            return self.config.dtypes
+        try:
+            return tuple(d for d in self.config.dtypes if d in fn())
+        except Exception:
+            return self.config.dtypes
+
+    def supports(self, dtype: str) -> bool:
+        """Whether this backend runs cases of ``dtype`` here."""
+        return dtype in self.dtypes()
+
+
+def bind_kwargs(fn, spec, tensors: dict, scalars: dict) -> dict:
+    """Match launch.json params to run()'s signature, by name."""
+    available = {p.name: (tensors[p.name] if p.is_tensor else scalars[p.name]) for p in spec.params}
+    accepted = inspect.signature(fn).parameters
+    if any(v.kind == v.VAR_KEYWORD for v in accepted.values()):
+        return available
+    missing = [k for k, v in accepted.items() if k not in available and v.default is inspect.Parameter.empty]
+    if missing:
+        raise TypeError(f"run() asks for unknown params {missing}; launch.json params are {list(available)}")
+    return {k: v for k, v in available.items() if k in accepted}
+
+
+_cache: dict[str, Backend] = {}
+
+
+def get_backend(name: str) -> Backend:
+    """Backend by name (cached)."""
+    if name not in _cache:
+        _cache[name] = Backend(name)
+    return _cache[name]
+
+
+def registered_backends() -> list[str]:
+    """Names of backends whose modules exist in this install (implemented or not)."""
+    names = []
+    for n in BACKEND_NAMES:
+        try:
+            importlib.import_module(f"kernelport.backends.{n}.config")
+            names.append(n)
+        except ModuleNotFoundError:
+            continue
+    return names

kernelport/backends/cpu/__init__.py

@@ -0,0 +1 @@
+"""Numba-compiled CPU kernels backend."""

kernelport/backends/cpu/capability.py

@@ -0,0 +1,25 @@
+"""Capability check for the cpu (Numba) backend."""
+
+from __future__ import annotations
+
+import importlib.util
+
+
+def is_available() -> tuple[bool, str]:
+    """(available, reason)."""
+    if importlib.util.find_spec("numba") is None:
+        return False, "numba not installed (pip install kernelport[cpu])"
+    if importlib.util.find_spec("torch") is None:
+        return False, "torch not installed (needed to build inputs and run references)"
+    try:
+        import numba
+    except Exception as exc:  # e.g. numba built against a different NumPy
+        return False, f"numba failed to import: {exc}"
+    return True, f"Numba {numba.__version__} on CPU, {numba.config.NUMBA_NUM_THREADS} threads"
+
+
+def hardware_name() -> str:
+    """CPU model name, e.g. 'Intel(R) Core(TM) i7-13620H (16 threads)'."""
+    from kernelport.hardware import cpu_name
+
+    return cpu_name()

kernelport/backends/cpu/config.py

@@ -0,0 +1,36 @@
+"""cpu (Numba) backend configuration."""
+
+from kernelport.backends.base import BackendConfig
+
+PORT_TEMPLATE = '''"""CPU (Numba) port of <kernel>: <one line of what it computes>."""
+
+import numpy as np
+from numba import njit, prange
+
+
+@njit(parallel=True)
+def _kernel(x, y, n):
+    # CUDA: one thread per i. prange splits the loop across CPU threads.
+    for i in prange(n):
+        y[i] = x[i]
+
+
+def run(x, y, n):
+    """Same parameter names as launch.json. Arrays are NumPy with launch.json shapes; write outputs in place."""
+    _kernel(x, y, n)
+'''
+
+# No tolerance overrides: Numba compiles float32 arithmetic to IEEE float32 ops
+# (fastmath is off unless a port turns it on), and math.exp/log/sqrt on float32
+# call the C library's correctly-rounded-or-near versions.
+CONFIG = BackendConfig(
+    name="cpu",
+    extension=".py",
+    description="Numba-compiled CPU kernels (any machine)",
+    dtypes=("float32", "float64"),
+    tolerances={},
+    timing_comparable=False,
+    timing_note="CPU wall-clock time; not comparable to GPU kernel timings",
+    hardware_hint="any machine with Numba",
+    port_template=PORT_TEMPLATE,
+)

kernelport/backends/cpu/runner.py

@@ -0,0 +1,67 @@
+"""Run a cpu (Numba) port.
+
+A cpu port is a Python file defining::
+
+    def run(<same parameter names as launch.json>):
+        ...call one or more @numba.njit functions...
+
+Tensor parameters arrive as C-contiguous NumPy arrays with the launch.json
+shapes (outputs pre-filled with NaN, or a sentinel for ints) and must be
+written in place. Scalars arrive as Python numbers. run() must return None.
+
+Numba compiles on the first call for each argument type, so the loaded run()
+is cached per file content: verification compiles once per dtype, not once
+per case, and benchmark warm-up excludes compilation.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from kernelport.backends.base import Prepared, bind_kwargs
+from kernelport.harness.inputs import CaseData
+from kernelport.harness.reference import file_hash, load_module
+from kernelport.harness.spec import LaunchSpec
+
+_loaded: dict[tuple[str, str], object] = {}
+
+
+def load_port(path: Path):
+    """Import the port (once per content hash) and return its run() function."""
+    path = Path(path).resolve()
+    key = (str(path), file_hash(path))
+    if key not in _loaded:
+        fn = getattr(load_module(path, "kp_cpu"), "run", None)
+        if fn is None:
+            raise AttributeError(f"{path.name} does not define run(...)")
+        _loaded[key] = fn
+    return _loaded[key]
+
+
+def _to_numpy(name: str, t: torch.Tensor):
+    if t.dtype == torch.bfloat16:
+        raise TypeError(f"{name}: bfloat16 has no NumPy dtype; the cpu backend does not run bfloat16 cases")
+    return t.clone().contiguous().numpy()
+
+
+def prepare(path: Path, spec: LaunchSpec, data: CaseData) -> Prepared:
+    """Load the port and bind it to one case's arrays."""
+    fn = load_port(path)
+    arrays = {k: _to_numpy(k, v) for k, v in data.tensors.items()}
+    kwargs = bind_kwargs(fn, spec, arrays, data.scalars)
+    outputs = [p.name for p in spec.params if p.is_output]
+
+    def launch() -> None:
+        ret = fn(**kwargs)
+        if ret is not None:
+            raise TypeError("run() returned a value; outputs must be written in place "
+                            "(e.g. y[:] = ... or y[i] = ...), and run() must return None")
+
+    return Prepared(
+        launch=launch,
+        sync=lambda: None,
+        fetch=lambda: {n: torch.from_numpy(arrays[n].copy()) for n in outputs},
+        timer="wall_clock",
+    )

kernelport/backends/hip/__init__.py

@@ -0,0 +1 @@
+"""HIP C++ kernels on AMD GPUs backend."""

kernelport/backends/hip/capability.py

@@ -0,0 +1,34 @@
+"""Capability check for the hip backend (ROCm torch + ROCm CuPy with hipRTC)."""
+
+from __future__ import annotations
+
+import importlib.util
+
+
+def is_available() -> tuple[bool, str]:
+    """(available, reason)."""
+    if importlib.util.find_spec("torch") is None:
+        return False, "torch not installed"
+    import torch
+
+    if not getattr(torch.version, "hip", None):
+        return False, "needs a ROCm build of torch on an AMD GPU"
+    if not torch.cuda.is_available():
+        return False, "ROCm torch sees no AMD GPU"
+    if importlib.util.find_spec("cupy") is None:
+        return False, "cupy (ROCm build, e.g. cupy-rocm-6-0) not installed; it provides hipRTC compilation"
+    try:
+        import cupy
+
+        if not getattr(cupy.cuda.runtime, "is_hip", False):
+            return False, "cupy is the CUDA build; install the ROCm build (e.g. cupy-rocm-6-0)"
+    except Exception as exc:
+        return False, f"cupy failed to import: {exc}"
+    return True, f"hipRTC via CuPy on ROCm {torch.version.hip}"
+
+
+def hardware_name() -> str:
+    """AMD GPU name with gfx arch."""
+    from kernelport.hardware import gpu_name
+
+    return gpu_name() or "no GPU"

kernelport/backends/hip/config.py

@@ -0,0 +1,28 @@
+"""hip backend configuration."""
+
+from kernelport.backends.base import BackendConfig
+
+PORT_TEMPLATE = '''// HIP port of <kernel>: <one line of what it computes>.
+// Same entry name, parameter order and launch (grid/block/shared mem) as launch.json.
+#include <hip/hip_fp16.h>
+#include <hip/hip_bf16.h>
+
+template <typename T>
+__global__ void vector_add(const T* __restrict__ a, const T* __restrict__ b, T* __restrict__ out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = a[i] + b[i];
+}
+'''
+
+# No tolerance overrides: HIP float arithmetic is IEEE single precision like CUDA.
+CONFIG = BackendConfig(
+    name="hip",
+    extension=".hip",
+    description="HIP C++ kernels on AMD GPUs, compiled at runtime with hipRTC",
+    dtypes=("float32", "float16", "bfloat16"),
+    tolerances={},
+    timing_comparable=True,
+    timing_note="",
+    hardware_hint="AMD GPU with ROCm (ROCm torch + cupy-rocm)",
+    port_template=PORT_TEMPLATE,
+)

kernelport/backends/hip/runner.py

@@ -0,0 +1,48 @@
+"""Run a hip port: HIP C++ compiled at runtime with hipRTC (through CuPy's ROCm build).
+
+A hip port is a ``.hip`` source file containing the kernel(s). It is launched
+exactly like the original CUDA kernel: same entry (``launch.json`` "entry",
+with ``{T}`` filled from HIP type names), same parameter order, same grid,
+block and shared memory from ``launch.json``. Tensors are device buffers on the
+AMD GPU (outputs pre-filled with NaN); scalars are passed by value.
+
+Requirements: a ROCm build of torch and a ROCm build of CuPy (e.g. cupy-rocm-6-0),
+which compiles RawModule sources with hipRTC.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from kernelport.backends.base import Prepared
+from kernelport.harness import run_cuda
+from kernelport.harness.inputs import CaseData
+from kernelport.harness.spec import LaunchSpec
+
+# C++ names substituted for {T} in the entry. HIP spells bf16 differently from CUDA.
+HIP_TYPE_NAMES = {
+    "float32": "float", "float64": "double", "float16": "__half", "bfloat16": "__hip_bfloat16",
+    "int8": "int8_t", "uint8": "uint8_t", "int16": "int16_t", "int32": "int", "int64": "long long",
+}
+
+# hipRTC, like NVRTC, has no host math headers. hip_runtime is implicit in hipRTC.
+HIPRTC_PRELUDE = run_cuda.NVRTC_PRELUDE
+HIPRTC_OPTIONS = ("-std=c++17",)
+
+
+def prepare(path: Path, spec: LaunchSpec, data: CaseData) -> Prepared:
+    """Compile (cached) and bind the port to one case on the GPU."""
+    kern = run_cuda.get_kernel(path, spec, data.dtype, type_names=HIP_TYPE_NAMES,
+                               prelude=HIPRTC_PRELUDE, options=HIPRTC_OPTIONS, device_only=False)
+    dev = {k: v.to("cuda") for k, v in data.tensors.items()}  # torch's ROCm build keeps the "cuda" name
+    grid, block, shmem = spec.grid_block(data.dims)
+    args = tuple(run_cuda.kernel_args(spec, data, dev))
+    outputs = [p.name for p in spec.params if p.is_output]
+    return Prepared(
+        launch=lambda: kern(grid, block, args, shared_mem=shmem),
+        sync=torch.cuda.synchronize,
+        fetch=lambda: {n: dev[n].cpu() for n in outputs},
+        timer="cuda_event",  # torch.cuda.Event maps to hipEvent on ROCm
+    )

kernelport/backends/metal/__init__.py

@@ -0,0 +1 @@
+"""MLX custom Metal kernels (mx.fast.metal_kernel) backend."""

kernelport/backends/metal/capability.py

@@ -0,0 +1,32 @@
+"""Capability check for the metal (MLX) backend."""
+
+from __future__ import annotations
+
+import importlib.util
+import platform
+import sys
+
+
+def is_available() -> tuple[bool, str]:
+    """(available, reason)."""
+    if sys.platform != "darwin" or platform.machine() != "arm64":
+        return False, "Metal needs macOS on Apple Silicon"
+    if importlib.util.find_spec("mlx") is None:
+        return False, "mlx not installed (pip install kernelport[metal])"
+    if importlib.util.find_spec("torch") is None:
+        return False, "torch not installed (needed to build inputs and run references)"
+    try:
+        import mlx.core as mx
+
+        if not mx.metal.is_available():
+            return False, "MLX reports no Metal device"
+    except Exception as exc:
+        return False, f"mlx failed to import: {exc}"
+    return True, "MLX custom Metal kernels"
+
+
+def hardware_name() -> str:
+    """Apple chip name, e.g. 'Apple M2 Pro (Metal)'."""
+    from kernelport.hardware import _cpu_model
+
+    return f"{_cpu_model() or 'Apple Silicon'} (Metal)"

kernelport/backends/metal/config.py

@@ -0,0 +1,35 @@
+"""metal (MLX custom Metal kernels) backend configuration."""
+
+from kernelport.backends.base import BackendConfig
+
+PORT_TEMPLATE = '''"""Metal (MLX) port of <kernel>: <one line of what it computes>."""
+
+# Kernel BODY only: MLX writes the signature. Inputs = input/inout tensors then scalars
+# (scalars are 1-element arrays: n[0]); outputs = output params, inouts as <name>_out.
+SOURCE = """
+    uint i = thread_position_in_grid.x;      // CUDA: blockIdx.x * blockDim.x + threadIdx.x
+    if (i < uint(n[0])) {
+        y[i] = x[i];
+    }
+"""
+THREADGROUP = (256, 1, 1)                     # CUDA: block
+
+
+def grid(n):
+    """TOTAL threads per dimension (Metal dispatchThreads), not threadgroup counts."""
+    return (n, 1, 1)
+'''
+
+# No tolerance overrides. Metal's compiler may use fast-math, so ports call
+# metal::precise:: versions of exp/log/sqrt/rsqrt (see skills/kernelport/backends/metal.md).
+CONFIG = BackendConfig(
+    name="metal",
+    extension=".py",
+    description="Metal kernels through MLX custom kernels (mx.fast.metal_kernel)",
+    dtypes=("float32", "float16", "bfloat16"),
+    tolerances={},
+    timing_comparable=False,
+    timing_note="wall-clock time including MLX graph evaluation; not comparable to device-event timings",
+    hardware_hint="Apple Silicon Mac",
+    port_template=PORT_TEMPLATE,
+)

kernelport/backends/metal/runner.py

@@ -0,0 +1,142 @@
+"""Run a metal port: a Metal kernel body compiled by MLX (``mx.fast.metal_kernel``).
+
+A metal port is a Python file defining::
+
+    SOURCE = '''
+        uint i = thread_position_in_grid.x;
+        if (i < n[0]) { out[i] = a[i] + b[i]; }
+    '''
+    THREADGROUP = (256, 1, 1)
+    def grid(n):                 # TOTAL threads per dimension (MLX/Metal dispatchThreads),
+        return (n, 1, 1)         # not threadgroup counts as in CUDA
+
+MLX writes the kernel signature from these names (launch.json order):
+
+- inputs: every ``input`` and ``inout`` tensor param, then every scalar param as a
+  one-element array (read it as ``n[0]``);
+- outputs: every ``output`` param by name, and every ``inout`` param as ``<name>_out``
+  (MLX outputs are fresh buffers: copy or update from the input ``<name>``).
+
+``T`` is a template parameter when launch.json uses dtype "T". Outputs are
+pre-filled with ``INIT_VALUE`` (default NaN, so unwritten elements are caught).
+``ATOMIC_OUTPUTS = True`` declares every output as ``device atomic<...>*``.
+``HEADER`` may hold helper functions placed before the kernel.
+"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from kernelport.backends.base import Prepared, bind_kwargs
+from kernelport.harness.inputs import CaseData
+from kernelport.harness.reference import file_hash, load_module
+from kernelport.harness.spec import LaunchSpec
+
+_loaded: dict[tuple[str, str], object] = {}
+_kernels: dict[tuple, object] = {}
+SCALAR_NP = {"int32": np.int32, "int64": np.int64, "uint32": np.uint32, "float32": np.float32,
+             "float64": np.float32, "float16": np.float16}
+
+
+def _mx():
+    import mlx.core as mx
+
+    return mx
+
+
+def mx_dtype(dtype: str):
+    mx = _mx()
+    return {"float32": mx.float32, "float16": mx.float16, "bfloat16": mx.bfloat16, "int32": mx.int32,
+            "int64": mx.int64, "uint32": mx.uint32, "int8": mx.int8, "uint8": mx.uint8,
+            "int16": mx.int16}[dtype]
+
+
+def load_port(path: Path):
+    path = Path(path).resolve()
+    key = (str(path), file_hash(path))
+    if key not in _loaded:
+        mod = load_module(path, "kp_metal")
+        if not isinstance(getattr(mod, "SOURCE", None), str):
+            raise AttributeError(f"{path.name} must define SOURCE = '''...''' (the Metal kernel body)")
+        if not callable(getattr(mod, "grid", None)):
+            raise AttributeError(f"{path.name} must define grid(...) returning total threads per dimension")
+        _loaded[key] = mod
+    return _loaded[key]
+
+
+def names(spec: LaunchSpec) -> tuple[list[str], list[str]]:
+    """(input_names, output_names) as the kernel sees them."""
+    ins = [p.name for p in spec.params if p.kind in ("input", "inout")]
+    ins += [p.name for p in spec.params if p.kind == "scalar"]
+    outs = [p.name if p.kind == "output" else f"{p.name}_out" for p in spec.params if p.is_output]
+    return ins, outs
+
+
+def _to_mx(t: torch.Tensor):
+    mx = _mx()
+    if t.dtype == torch.bfloat16:  # NumPy has no bf16: go through float32 (exact both ways)
+        return mx.array(t.float().numpy()).astype(mx.bfloat16)
+    return mx.array(t.contiguous().numpy())
+
+
+def _to_torch(a, dtype: str) -> torch.Tensor:
+    mx = _mx()
+    if dtype == "bfloat16":
+        return torch.from_numpy(np.array(a.astype(mx.float32))).to(torch.bfloat16)
+    return torch.from_numpy(np.array(a).copy())
+
+
+def _triple(v, what: str) -> tuple[int, int, int]:
+    v = (v,) if isinstance(v, int) else tuple(int(x) for x in v)
+    if not 1 <= len(v) <= 3 or any(x < 0 for x in v):
+        raise ValueError(f"{what} must be 1 to 3 non-negative ints, got {v}")
+    return v + (1,) * (3 - len(v))
+
+
+def prepare(path: Path, spec: LaunchSpec, data: CaseData) -> Prepared:
+    """Build the MLX kernel (cached) and bind it to one case."""
+    mx = _mx()
+    mod = load_port(path)
+    in_names, out_names = names(spec)
+    atomic = bool(getattr(mod, "ATOMIC_OUTPUTS", False))
+    header = getattr(mod, "HEADER", "") or ""
+    key = (id(mx), mod.SOURCE, header, tuple(in_names), tuple(out_names), atomic)
+    if key not in _kernels:
+        _kernels[key] = mx.fast.metal_kernel(
+            name=f"kp_{spec.name}", input_names=in_names, output_names=out_names,
+            source=mod.SOURCE, header=header, ensure_row_contiguous=True, atomic_outputs=atomic)
+    kernel = _kernels[key]
+
+    inputs = [_to_mx(data.tensors[p.name]) for p in spec.params if p.kind in ("input", "inout")]
+    for p in spec.params:
+        if p.kind == "scalar":
+            inputs.append(mx.array(np.array([data.scalars[p.name]], dtype=SCALAR_NP[p.dtype])))
+    outs = [p for p in spec.params if p.is_output]
+    out_dtypes = [p.resolved_dtype(data.dtype) for p in outs]
+    out_shapes = [tuple(data.tensors[p.name].shape) for p in outs]
+    template = [("T", mx_dtype(data.dtype))] if any(p.dtype == "T" for p in spec.params) else []
+
+    arrays = {p.name: data.tensors[p.name].numpy() if data.tensors[p.name].dtype != torch.bfloat16
+              else data.tensors[p.name] for p in spec.params if p.is_tensor}
+    g = _triple(mod.grid(**bind_kwargs(mod.grid, spec, arrays, data.scalars)), "grid()")
+    tg = _triple(getattr(mod, "THREADGROUP", (256, 1, 1)), "THREADGROUP")
+    init = getattr(mod, "INIT_VALUE", math.nan)
+    state: dict = {}
+
+    def launch() -> None:
+        state["outs"] = kernel(inputs=inputs, template=template, grid=g, threadgroup=tg,
+                               output_shapes=out_shapes, output_dtypes=[mx_dtype(d) for d in out_dtypes],
+                               init_value=init, verbose=False)
+
+    def sync() -> None:
+        if "outs" in state:
+            mx.eval(*state["outs"])
+
+    def fetch() -> dict:
+        return {p.name: _to_torch(a, d) for p, a, d in zip(outs, state["outs"], out_dtypes)}
+
+    return Prepared(launch=launch, sync=sync, fetch=fetch, timer="wall_clock")

kernelport/backends/triton/__init__.py

@@ -0,0 +1 @@
+"""Triton backend (NVIDIA CUDA and AMD ROCm)."""

kernelport/backends/triton/capability.py

@@ -0,0 +1,29 @@
+"""Can Triton ports run here? Imports nothing heavy at module level."""
+
+from __future__ import annotations
+
+import importlib.util
+
+
+def is_available() -> tuple[bool, str]:
+    """(available, reason)."""
+    if importlib.util.find_spec("torch") is None:
+        return False, "torch not installed (pip install kernelport[triton])"
+    if importlib.util.find_spec("triton") is None:
+        return False, "triton not installed (pip install kernelport[triton]; triton-windows on Windows)"
+    import torch
+
+    if not torch.cuda.is_available():
+        return False, "no CUDA or ROCm GPU visible to torch"
+    try:
+        import triton  # noqa: F401
+    except Exception as exc:
+        return False, f"triton failed to import: {exc}"
+    return True, "Triton on " + ("ROCm" if getattr(torch.version, "hip", None) else "CUDA")
+
+
+def hardware_name() -> str:
+    """GPU name with arch, e.g. 'NVIDIA GeForce RTX 5060 Laptop GPU (sm_120)'."""
+    from kernelport.hardware import gpu_name
+
+    return gpu_name() or "no GPU"