eagle inside your NVIDIA Warp program

eagle inside your NVIDIA Warp program#

You write kernels with @wp.kernel, launch them with wp.launch, and record gradients with wp.Tape. This page replaces one of those launches with a compiled kernel from eagle. Your arrays, your loop, your other Warp kernels, your tape and your CUDA graphs stay as they are.

Needs warp-lang (pip install warp-lang). Runs on the CPU too; the capture and stream sections need a CUDA GPU.

Colab or Kaggle: run this first cell (it installs the packages and downloads the helper files this notebook uses). Elsewhere it only checks the imports.

Show code cell source

Hide code cell source

# Colab setup, generated by docs/_tools/nb_colab_setup.py (do not edit).
import os

try:
    import google.colab  # noqa: F401
    HOSTED = True  # a Colab runtime
except ImportError:
    HOSTED = "KAGGLE_KERNEL_RUN_TYPE" in os.environ  # a Kaggle kernel

if HOSTED:
    import importlib.metadata, pathlib, re, shutil, subprocess, sys, urllib.request

    gpu = bool(shutil.which("nvidia-smi")) and subprocess.run(["nvidia-smi", "-L"], capture_output=True).returncode == 0
    extra = ""
    if gpu:
        # the CUDA major the runtime already has: a CuPy that is installed, else torch's, else the driver's
        major = None
        for m in ("13", "12"):
            try:
                importlib.metadata.version(f"cupy-cuda{m}x")
                major = m
                break
            except importlib.metadata.PackageNotFoundError:
                pass
        if major is None:
            try:
                import torch
                major = "13" if str(torch.version.cuda).startswith("13") else None
            except ImportError:
                pass
        if major is None:
            smi = re.search(r"CUDA Version: (\d+)", subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout)
            major = "13" if smi and int(smi.group(1)) >= 13 else "12"
        extra = f"[cuda{major}]"
        print(f"GPU found: installing the {extra} extras to match this runtime.")
    pre = ["--pre"] if os.environ.get("RAPTOR_PRE") else []  # optional: test a release candidate
    ref = os.environ.get("RAPTOR_REF", "main")  # optional: fetch the helper files from this git ref
    if pre or "RAPTOR_REF" in os.environ:
        print(f"Pre-release test: installing with --pre, helper files from {ref}.")
    # our packages with --upgrade, so a runtime that kept an older release gets the current one
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade", *pre, f"raptor-eagle{extra}", f"raptor-hawk{extra}"])
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "warp-lang"])
    if not gpu:
        print("Running on the CPU (no GPU in this runtime). For a GPU: " + ("Settings → Accelerator → GPU T4 x2." if "KAGGLE_KERNEL_RUN_TYPE" in os.environ else "Runtime → Change runtime type → T4 GPU."))

try:
    import eagle, hawk  # noqa: F401
except ImportError as err:
    raise ImportError(
        "eagle, hawk not found. Install with `pip install raptor-eagle raptor-hawk warp-lang`. For a GPU, add [cuda12] (or [cuda13]) to raptor-eagle and raptor-hawk."
    ) from err

1. Your program as it is#

A ball-flight step with air drag as a Warp kernel, rolled out 10 steps for 100,000 balls, a loss kernel, and a wp.Tape that gives the gradient with respect to the drag coefficients. This is plain Warp.

import warnings

import numpy as np

warnings.filterwarnings("ignore", category=UserWarning, message=r"(?s).*CuPy may not function correctly.*")
import warp as wp

wp.config.log_level = wp.LOG_WARNING
wp.init()
device = "cuda:0" if wp.is_cuda_available() else "cpu"
N, STEPS, DT, G = 100_000, 10, 0.05, 9.81
rng = np.random.default_rng(3)
v0 = wp.array(5.0 * rng.normal(size=(3, N)), dtype=wp.float64, device=device)
cd = wp.array(rng.uniform(0.0, 0.1, N), dtype=wp.float64, device=device, requires_grad=True)


@wp.kernel
def drag_step_wp(v: wp.array2d(dtype=wp.float64), cd: wp.array(dtype=wp.float64), dt: wp.float64, g: wp.float64,
                 v_next: wp.array2d(dtype=wp.float64)):
    i = wp.tid()
    x = v[0, i]
    y = v[1, i]
    z = v[2, i]
    k = cd[i] * wp.sqrt(x * x + y * y + z * z)
    v_next[0, i] = x - dt * k * x
    v_next[1, i] = y - dt * k * y
    v_next[2, i] = z + dt * (-g - k * z)


@wp.kernel
def loss_wp(v: wp.array2d(dtype=wp.float64), loss: wp.array(dtype=wp.float64)):
    i = wp.tid()
    wp.atomic_add(loss, 0, v[0, i] * v[0, i] + v[1, i] * v[1, i] + v[2, i] * v[2, i])


def warp_step(v, cd):
    v_next = wp.zeros((3, N), dtype=wp.float64, device=device, requires_grad=True)
    wp.launch(drag_step_wp, dim=N, inputs=[v, cd, DT, G], outputs=[v_next], device=device)
    return v_next


def run(step):
    loss = wp.zeros(1, dtype=wp.float64, device=device, requires_grad=True)
    tape = wp.Tape()
    with tape:
        v = v0
        for _ in range(STEPS):
            v = step(v, cd)
        wp.launch(loss_wp, dim=N, inputs=[v], outputs=[loss], device=device)
    return tape, loss, v


tape, loss, v_end = run(warp_step)
tape.backward(loss)
loss_warp, grad_warp = loss.numpy()[0], cd.grad.numpy().copy()
print(f"Warp step: loss {loss_warp:.6f}, sum of d loss / d cd {grad_warp.sum():.6f}")

Warp step: loss 6773513.693337, sum of d loss / d cd -51373813.198953

2. Drop the kernel in#

The same step as a hawk kernel, wrapped by eagle.frameworks.warp.function. It takes wp.arrays and returns wp.arrays, and runs on Warp’s current stream (on the host for CPU arrays). A Param is a one-cell array. run is not touched.

import eagle
import hawk
from eagle.frameworks import warp as eagle_warp
from hawk import Mutable, Param, Scalar, Vector
from hawk.math import norm, vec


@hawk.kernel
def drag_step(v: Vector[3], cd: Scalar, dt: Param, g: Param, v_next: Mutable[Vector[3]]):
    a = vec(0.0, 0.0, -g) - cd * norm(v) * v
    v_next = v + dt * a


step_fn = eagle_warp.function(drag_step)
dt_a = wp.array([DT], dtype=wp.float64, device=device)
g_a = wp.array([G], dtype=wp.float64, device=device)


def eagle_step(v, cd):  # the one changed line: the launch is eagle's
    return step_fn(v, cd, dt_a, g_a)


cd.grad.zero_()
tape, loss_e, v_end_e = run(eagle_step)
loss_eagle = loss_e.numpy()[0]
print(f"eagle step: loss {loss_eagle:.6f}")
np.testing.assert_allclose(loss_eagle, loss_warp, rtol=1e-10)
print("same loss as the Warp step")

eagle step: loss 6773513.693337
same loss as the Warp step

3. What stays the same#

wp.Tape records the eagle call like any Warp launch. tape.backward runs the reverse-mode kernel hawk derives from your kernel, and the gradient lands in cd.grad, next to the gradients of your own kernels.

tape.backward(loss_e)
np.testing.assert_allclose(cd.grad.numpy(), grad_warp, rtol=1e-9)
print(f"sum of d loss / d cd {cd.grad.numpy().sum():.6f}: the gradient of the Warp step")

sum of d loss / d cd -51373813.198953: the gradient of the Warp step

4. What crosses, and who owns it#

Your arrays are used as they are, with no copy. The result is a new wp.array; pass out= to write into one you own. When the arrays cannot be used as given, the call says so and names the argument.

A kernel with carried state (a Mutable plane it reads, then writes) takes that plane as an input and returns the next value as a new array; a finished sample keeps its value, and tape.backward carries its gradient through unchanged, so a loop needs no where.

mine = wp.zeros((3, N), dtype=wp.float64, device=device)
step_fn(v0, cd, dt_a, g_a, out=(mine,))
np.testing.assert_allclose(mine.numpy(), warp_step(v0, cd).numpy(), rtol=1e-10)
print("written into your wp.array in place")


def refused(label, call):
    try:
        call()
    except Exception as err:
        msg = str(err).splitlines()[0]
        print(f"{label}: {type(err).__name__}: {msg if len(msg) <= 150 else msg[:147] + '...'}")


refused("wrong length", lambda: step_fn(v0, cd[:5], dt_a, g_a))
refused("a plain number for a Param", lambda: step_fn(v0, cd, DT, g_a))
refused("a vec3 array where a (3, n) plane is declared",
        lambda: step_fn(wp.zeros(N, dtype=wp.vec3d, device=device), cd, dt_a, g_a))

written into your wp.array in place
wrong length: ValueError: drag_step: 'cd' has 5 samples, others 100000
a plain number for a Param: TypeError: drag_step: 'dt' must be a wp.array, got float
a vec3 array where a (3, n) plane is declared: ValueError: drag_step: 'v' has shape (100000,), expected (3, 100000)

5. Your Warp kernels inside an eagle batch#

The other direction: your wp.launch becomes a node of an eagle graph, next to eagle’s kernels. Here a Warp kernel feeds a hawk kernel whose weights are a Warp array:

import tempfile
import warnings

warnings.filterwarnings("ignore", category=UserWarning, message=r"(?s).*CuPy may not function correctly.*")

import numpy as np

import eagle, eagle.diff, eagle.plan
import hawk
from eagle.graph import opaque
from eagle.interop import import_buffer
from hawk import Mutable, Scalar, Table
import cupy as cp

gpu = device != "cpu"
xp = cp if gpu else np
@hawk.kernel
def lin(w: Table[Scalar], z: Scalar, y: Mutable[Scalar]):
    y = w.at(0) * z + w.at(1)       # w is ONE buffer, read by every sample

(plan,) = eagle.plan.auto([lin], targets=("cuda" if gpu else "host",), cache_dir=tempfile.mkdtemp())
x, z, y = xp.linspace(-1, 1, 8), xp.zeros(8), xp.zeros(8)
w = wp.array([2.0, 0.5], dtype=wp.float64, device=device)


@wp.kernel
def squash(x: wp.array(dtype=wp.float64), z: wp.array(dtype=wp.float64)):
    i = wp.tid()
    z[i] = wp.tanh(wp.float64(0.7) * x[i])


xw, zw = wp.from_dlpack(x), wp.from_dlpack(z)
S = wp.Stream(device) if gpu else None


def model():                                      # your Warp launch
    with wp.ScopedStream(S) if gpu else wp.ScopedDevice(device):
        wp.launch(squash, dim=8, inputs=[xw, zw], device=device)

to_np = (lambda a: a.get()) if gpu else np.asarray
g = (eagle.Graph()
     .add_node(opaque(model, reads={"x": x}, writes={"z": z}, **({"stream": S} if gpu else {})), name="model")
     .add(plan.bind(w=import_buffer(w), z=z, y=y), name="lin").build(record=True))
w.assign(wp.array([6.0, 1.5], dtype=wp.float64, device=device))   # in place: the next launch reads it
g.launch()
print("a launch after the build sees the new weights:", bool(np.allclose(to_np(y), 6 * np.tanh(0.7 * to_np(x)) + 1.5)))

a launch after the build sees the new weights: True

eagle fences the stream S against the stages before and after it, so Warp and eagle stay in order with nothing to synchronize (the three tiers has the picture).

Your Warp code	Write it as	Runs
a graph you captured (`wp.ScopedCapture`)	`child(wp_graph, reads=..., writes=...)`	captured as one node
launches on a stream of your own	`opaque(fn, stream=S)`	stream-ordered, the host never waits
launches plus a host readback	`opaque(fn, eager=True)`	host-synchronized

child takes the wp.Graph object, not its handle, so eagle can keep it alive; rebuild the eagle graph after you re-capture it. A stream or host node cannot sit inside a stepping loop (see Limits).

Weights are one buffer. import_buffer(w) of a wp.array binds it to the Table; update it in place, as above.

Gradients. Your vjp runs host-synchronized, so a wp.Tape inside it is fine:

def vjp(primals, cotangents):
    xa = wp.from_dlpack(xp.asarray(primals["x"]).copy())
    xa.requires_grad = True
    za = wp.zeros(8, dtype=wp.float64, device=device, requires_grad=True)
    with wp.Tape() as tape:
        wp.launch(squash, dim=8, inputs=[xa, za], device=device)
    tape.backward(grads={za: wp.from_dlpack(xp.asarray(cotangents["z"]).copy())})
    return {"x": xp.asarray(xa.grad.numpy())}


grads = eagle.diff.reverse(g.record, {"lin.y": xp.ones(8)}, vjps={"model": vjp})
want = 6 * 0.7 * (1 - np.tanh(0.7 * to_np(x)) ** 2)
print("d y / d x through your Warp kernel:", bool(np.allclose(to_np(grads["x"]), want)))

d y / d x through your Warp kernel: True

6. On the GPU: CUDA graphs and streams#

A CUDA graph captures the whole rollout, loss and backward in one replay, eagle’s launches included. Build before you capture (fn.prepare): a compile is host work, and a capture records device work only.

step_fn.prepare("float64", device, backward=True)
with wp.ScopedCapture(device=device) as capture:
    tape_c, loss_c, _ = run(eagle_step)
    tape_c.backward(loss_c)

cd.assign(wp.array(0.5 * cd.numpy(), dtype=wp.float64, device=device))  # new values: the replay follows them
cd.grad.zero_()
wp.capture_launch(capture.graph)
wp.synchronize_device(device)
replayed = cd.grad.numpy().copy()
cd.grad.zero_()
t, l, _ = run(warp_step)
t.backward(l)
np.testing.assert_allclose(replayed, cd.grad.numpy(), rtol=1e-9)
print("one graph replay: the same gradient as the Warp step, for the new coefficients")

one graph replay: the same gradient as the Warp step, for the new coefficients

The launch rides Warp’s current stream and never syncs, so Warp, eagle, Warp on one stream is in order with nothing between them:

@wp.kernel
def slow_copy(x: wp.array2d(dtype=wp.float64), y: wp.array2d(dtype=wp.float64)):
    i = wp.tid()
    for r in range(3):
        s = x[r, i]
        for _ in range(2000):
            s = s + wp.float64(1.0e-300)
        y[r, i] = s


stream = wp.Stream(device)
with wp.ScopedStream(stream):
    slow = wp.empty((3, N), dtype=wp.float64, device=device)
    wp.launch(slow_copy, dim=N, inputs=[v0, slow], device=device)  # Warp, on its stream
    stepped = eagle_step(slow, cd)  # eagle, on the same stream
    back = wp.empty_like(stepped)
    wp.copy(back, stepped)  # Warp again
wp.synchronize_stream(stream)
np.testing.assert_allclose(back.numpy(), warp_step(v0, cd).numpy(), rtol=1e-12)
print("Warp, eagle, Warp on one stream, no sync between: in order")

Warp, eagle, Warp on one stream, no sync between: in order

7. Limits#

What does not work today, one line each:

Gradients are reverse mode only: `Tape.record_func`, which records eagle's backward, has no forward mode and no second derivative (Warp's design): differentiate once with `tape.backward`, or use the [JAX](jax.ipynb) bridge for forward mode and higher orders.
Planes are component-major, `(3, n)` for a vector, so arrays of Warp vector types (`wp.vec3d`) are refused (our scope): pass a `(3, n)` array.
A call runs the float32 or the float64 build, chosen by its first float array, and mixing them is refused (the kernel is built once per dtype, our scope): create every array in one float dtype.
A kernel is built on its first call, and a graph capture records device work only, so a build cannot happen inside one (CUDA's design): call `fn.prepare(mode, device)` (or the function once) before `wp.capture_begin`.
Warp keeps one global tape, creates its CUDA context on first use and holds the current stream per device, not per thread (Warp's design), so threads pass `tape=` to each call, the device is touched once (`wp.zeros(1, device=device)`) before they start, and CUDA `.numpy()`, `ScopedStream` and `set_stream` are not called from several threads at once: read back on one thread, or copy with an explicit stream.
Warp kernels cannot call an eagle kernel, because a Warp kernel is device code and an eagle launch is a call from Python (their design): launch eagle between your Warp launches.
One GPU per process until 0.7 (our scope): run one process per GPU, or choose it with `CUDA_VISIBLE_DEVICES`.
A stream or host node cannot run inside a stepping loop's device graph (our scope): capture it (`opaque(fn)`, or `child`), or call `runner.run()` with `max_steps=k`, then your code, in your own loop.

8. Where next#

Interoperability: every framework, and the certification table that lists the Warp rows.
How the interop contract works: what crosses, and who owns it.
PyTorch and CuPy and NumPy: the same one-line change in those workflows.