hawk

Write per-sample kernels in Python, with derivatives. Part of the RAPTOR family.

hawk: one kernel, two machines, and its derivatives

Write a numerical kernel for one sample as a plain Python function. hawk compiles it for the GPU or the CPU and derives its gradient (reverse mode) and tangent (forward mode).

The RAPTOR family

hawk is the kernel-authoring layer of the RAPTOR family: write a kernel once at the level of your application, and hawk derives its forward- and reverse-mode derivatives and compiles it for host or device through eagle. See where RAPTOR fits for the full four-part story.

Read the documentation — installation, a five-minute quickstart, tutorials, examples and the full Python API reference.

Write a kernel, get its gradient

import pathlib, tempfile

import hawk
from hawk import Kernel, Mutable, Scalar, Vector
from hawk.diff import vjp
from hawk.math import dot

@hawk.kernel
def energy(v: Vector[3], out: Mutable[Scalar]):
    out = 0.5 * dot(v, v)

energy_vjp = Kernel("energy_vjp", vjp(energy, wrt=("v",)))
work = pathlib.Path(tempfile.mkdtemp())
bundle = hawk.build([energy, energy_vjp], work, targets=("host",))

import numpy as np

v = np.array([[1.0, 2.0, 3.0, 4.0], [0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0]])
e = np.zeros(4)
hawk.run(hawk.load(work, "energy"), v=v, out=e)
print("energy:", e)
# energy: [0.5 2.5 5.  9. ]

bar_out = np.ones(4)             # seed: d(loss)/d(energy) = 1 for every sample
bar_v = np.zeros((3, 4))
hawk.run(hawk.load(work, "energy_vjp"), v=v, bar_out=bar_out, bar_v=bar_v)
print("gradient d(energy)/dv:")
print(bar_v)
# [[1. 2. 3. 4.]
#  [0. 1. 0. 1.]
#  [0. 0. 1. 1.]]

energy = 0.5 |v|^2, so d(energy)/dv = v exactly — the printed gradient above is the input v. Output above is the notebook's own: docs/content/tutorials/04_gradients_backward.ipynb. hawk.diff.jvp derives the forward-mode (tangent) twin the same way.

Five-minute quickstart · install: pip install raptor-hawk (CPU route; it brings aether-dsc along), or pip install "raptor-hawk[cuda12]" "raptor-eagle[cuda12]" for the GPU route

Writing kernels

A kernel is a plain Python function decorated @hawk.kernel, one sample at a time; its parameters are its declared planes (Scalar, Vector[3], Mutable, ...) and its body writes ordinary arithmetic:

import hawk
from hawk import Mutable, Scalar, Terminated, Vector
from hawk.math import maximum, norm

@hawk.kernel
def running_apogee(position: Vector[3], terminated: Terminated,
                    r_max: Mutable[Scalar]):
    r_max = maximum(r_max, norm(position))

from hawk import Param
from hawk.ext import Kind, Output

ACCEL = Kind("drag", output=Output.returned(Vector[3], slot="acc"))

@ACCEL
def drag(velocity: Vector[3], k: Param, speed: Mutable[Scalar]):
    speed = norm(velocity)          # an ordinary named sink
    return (-k * norm(velocity)) * velocity   # the returned slot, "acc"

from hawk.ext import KernelKind, Output

class Accel(KernelKind, slug="drag"):
    output = Output.returned(Vector[3], slot="acc")

@Accel
def drag(velocity: Vector[3], k: Param):
    return (-k * norm(velocity)) * velocity

Threading

Free-threaded CPython (3.13t/3.14t) is supported: the compiled modules declare GIL-free operation and the GIL stays disabled after import. Any number of threads may call module-level functions, build, compile, plan and cache concurrently. Distinct objects may be used from distinct threads without synchronisation. One stateful object (a stream, capture, launcher, graph, composer, plan, pipeline, active set, host kernel or arg block) shared by several threads is memory-safe — its calls serialise and a consumed object raises — but the ORDER of those calls is the caller's responsibility, exactly as for a NumPy array or a CuPy stream. CUDA adds two rules of its own: a stream capture is begun, filled and ended by one thread, and while any capture is open no thread may synchronise the whole device (stream-level synchronisation is fine). GPU routes are supported on free-threaded 3.14; on 3.13t the CPU route runs. hawk's compile cache is shared safely between threads, and a sealed-path kernel's cache key follows the user headers it includes.

Install

pip install raptor-hawk                                     # CPU only: no GPU, driver or CUDA packages needed
pip install "raptor-hawk[cuda12]" "raptor-eagle[cuda12]"    # GPU: hawk compiles, eagle runs; [cuda13] on both for CUDA 13

Linux x86_64, CPython 3.9-3.14 (including free-threaded 3.13t and 3.14t), and a host g++ 11 or newer. raptor-hawk pulls aether-dsc (the sealed C++ headers hawk compiles against) automatically. No nvcc or CUDA toolkit is needed. Tested with CUDA 12.6 and CUDA 13.0 (CUDA 12.6 or newer).

NVIDIA packages come only through the extras. raptor-hawk[cuda12] pulls cuda-bindings 12, nvidia-cuda-nvrtc-cu12 and nvidia-cuda-cccl-cu12; raptor-hawk[cuda13] pulls cuda-bindings 13, nvidia-cuda-nvrtc 13 and nvidia-cuda-cccl 13 (CUDA 13's wheels have no -cu13 suffix); raptor-eagle[cuda12] / [cuda13] pull CuPy (cupy-cuda12x / cupy-cuda13x, with the CUDA headers CuPy compiles against). Without an extra pip installs no NVIDIA package: you get the CPU route, or the GPU route through a CUDA setup you already have. Pick the extra matching the CUDA version your driver reports (nvidia-smi, top right).

Platforms: built and tested on Linux x86_64 only so far (CPython 3.9–3.14, including free-threaded 3.13t and 3.14t), on NVIDIA GPUs from Pascal (Quadro P2000) and Turing (Tesla T4). There are no wheels for macOS, Windows or ARM yet, and WSL2 is untested. raptor-core and aether-dsc are pure Python and install anywhere. On free-threaded 3.13t and 3.14t, hawk runs GIL-free.

See installation for the CPU-only route (no GPU, driver or NVRTC needed at all) and every other detail.

Performance

Every number in this section is read from eagle's committed cards (benchmarks/perf_card/card_quadro-p2000.md, benchmarks/rk78_card/card_quadro-p2000.md). Kernels authored with hawk are what eagle's performance card times: hawk + eagle (eagle.simulate) is 3.9×–47× faster than NVIDIA Warp (3.9×), JAX (8.7×), CuPy and PyTorch (47×) at 1,000,000 RK4 oscillators finishing at different times — 156 ms on a Quadro P2000, using 3.8×–5.2× less GPU memory than CuPy and PyTorch (a million samples in 50 MiB, 1.31× the bare minimum state; CuPy 5.1×, PyTorch 6.9×; NVIDIA Warp 64 MiB, 1.7×) — and the same hawk kernel, unmodified, integrates 1,000,000 adaptive RK7(8) orbits in 13.0 s on the GPU or 7.15 s on 8 CPU threads alone. This is the regime hawk + eagle's compaction targets (many samples stopping at different times); in other regimes the other tools hold their own — on a dense workload where nothing finishes early Warp ties it (693 ms vs 699 ms at N = 1,000,000) and is level at N = 10,000 (7.22 ms vs 7.22 ms), and on this FP64-weak development card the CPU arm is the faster one there. See eagle's performance page for every number, every arm, including the ones where it doesn't win.

Going deeper

from hawk.compile import cubin, cubin_available

if cubin_available():
    image = cubin(source)          # DeviceImage(image=<SASS bytes>, target="cubin", arch="sm_XX")

pip install ../aether/dsc
pip install -e .[test]
pytest -m "not gpu"

hawk write it (this repo) · eagle run it · aether the numerics core · raptor the shared contracts — the RAPTOR family

Apache-2.0 (see LICENSE and NOTICE) · cite via "Cite this repository" (CITATION.cff; every tagged release is archived on Zenodo: doi:10.5281/zenodo.23250242) · built to make GPU computing accessible on modest hardware, for research and education. Collaboration is the point, and a citation is the currency — get in touch.