NVIDIA Warp: low-level control, side by side

NVIDIA Warp: low-level control, side by side#

One task, written twice: a batch of projectiles integrated under gravity until each one hits the ground, each stopping at its own step.

Time: ~8 min · Runs on: GPU only (both tools need one) · You need: eagle, hawk, NumPy, CuPy (GPU only) and Warp; read first: Stop when you’re done

NVIDIA Warp expresses this with explicit thread-level control: a thread index, a done-flag array, an early return, and a guarded write-back. hawk expresses it by declaring what each argument is, with termination as part of the kernel’s signature. Both run on the GPU here, and the two results agree exactly. See hawk next to Warp for the narrative version of this comparison.

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 python -m raptor_docs.nb colab_setup (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

import numpy as np

G = 9.81
DT = 0.02
N = 2048
H_MAX = 20.0

rng = np.random.default_rng(3)
y0 = rng.uniform(2.0, H_MAX, size=N)   # initial height
x0 = np.zeros(N)                        # initial downrange distance
vx0 = rng.uniform(-3.0, 3.0, size=N)    # initial horizontal speed
vy0 = np.zeros(N)                       # initial vertical speed

# enough steps for the highest sample to reach the ground, plus a margin
MAX_STEPS = int(np.ceil(np.sqrt(2 * H_MAX / G) / DT)) + 20
MAX_STEPS

Warp: idiomatic#

Warp’s own vocabulary is thread-level: wp.tid() gets this thread’s index, and the kernel decides for itself what a “done” sample means. A done flag array, an early return for a finished sample, and a guarded write-back for one that crosses the ground this step — the arrays are updated in place.

try:
    import warp as wp
except ImportError as err:
    raise ImportError("this example compares hawk with NVIDIA Warp: install it with `pip install warp-lang`") from err

wp.config.log_level = wp.LOG_WARNING  # suppress the init banner (local cache path)
wp.init()
DEVICE = "cuda:0"

The kernel itself, thread-level and explicit:

@wp.kernel
def warp_step(x: wp.array(dtype=wp.float64), y: wp.array(dtype=wp.float64),
              vx: wp.array(dtype=wp.float64), vy: wp.array(dtype=wp.float64),
              done: wp.array(dtype=wp.int32), g: wp.float64, dt: wp.float64):
    i = wp.tid()
    if done[i] == 1:
        return  # early return: a finished sample costs nothing further
    new_vy = vy[i] - g * dt
    new_y = y[i] + new_vy * dt
    new_x = x[i] + vx[i] * dt
    if new_y <= wp.float64(0.0):
        new_y = wp.float64(0.0)
        done[i] = 1
    # guarded write-back: only this still-live-this-step sample's state moves
    x[i] = new_x
    y[i] = new_y
    vy[i] = new_vy

x_w = wp.array(x0, dtype=wp.float64, device=DEVICE)
y_w = wp.array(y0, dtype=wp.float64, device=DEVICE)
vx_w = wp.array(vx0, dtype=wp.float64, device=DEVICE)
vy_w = wp.array(vy0, dtype=wp.float64, device=DEVICE)
done_w = wp.zeros(N, dtype=wp.int32, device=DEVICE)

for _ in range(MAX_STEPS):
    wp.launch(warp_step, dim=N, inputs=[x_w, y_w, vx_w, vy_w, done_w, G, DT],
              device=DEVICE)
wp.synchronize_device(DEVICE)

print("all samples landed:", bool((done_w.numpy() == 1).all()))

all samples landed: True

hawk: termination in the signature#

terminated: Terminated is a declared argument role, not a value the body reads and branches on: a sample flagged terminated going into a launch keeps its output plane’s prior value, enforced by the runtime itself, regardless of what the body computed for it.

This kernel keeps SEPARATE in/out planes (x/x_out, and so on) and a host-side copy-back loop, on purpose: stop when you’re done’s Mutable-plus-eagle.simulate idiom is the one to reach for in your own code (it needs none of the loop below), but it hides exactly the in-place-update step this page is comparing against Warp’s own explicit loop — showing it here keeps the two tools at the same level of abstraction, one in-place buffer update per launch, same as Warp’s.

import cupy as cp
import hawk
from hawk import Mutable, Param, Scalar, Terminated
from hawk.math import maximum

@hawk.kernel
def hawk_step(x: Scalar, y: Scalar, vx: Scalar, vy: Scalar, g: Param, dt: Param,
              terminated: Terminated, x_out: Mutable[Scalar], y_out: Mutable[Scalar],
              vx_out: Mutable[Scalar], vy_out: Mutable[Scalar]):
    new_vy = vy - g * dt
    new_y = maximum(y + new_vy * dt, 0.0)   # clamp at the ground, Warp's guard too
    x_out = x + vx * dt
    y_out = new_y
    vx_out = vx
    vy_out = new_vy

eagle.deploy compiles it; the deployed plan binds persistent buffers:

import eagle

the_plan = eagle.deploy(hawk_step)

x_s, y_s, vx_s, vy_s = cp.asarray(x0), cp.asarray(y0), cp.asarray(vx0), cp.asarray(vy0)
x_o, y_o, vx_o, vy_o = cp.zeros(N), cp.zeros(N), cp.zeros(N), cp.zeros(N)
terminated_h = cp.zeros(N, dtype=cp.bool_)

bound = the_plan.bind(x=x_s, y=y_s, vx=vx_s, vy=vy_s, g=G, dt=DT,
                       terminated=terminated_h, x_out=x_o, y_out=y_o,
                       vx_out=vx_o, vy_out=vy_o)

for _ in range(MAX_STEPS):
    bound.launch()
    terminated_h[:] = terminated_h | (y_o <= 0.0)
    x_s[:], y_s[:], vx_s[:], vy_s[:] = x_o, y_o, vx_o, vy_o
cp.cuda.get_current_stream().synchronize()

print("all samples landed:", bool(cp.asnumpy(terminated_h).all()))

all samples landed: True

y_warp, x_warp = y_w.numpy(), x_w.numpy()
y_hawk, x_hawk = cp.asnumpy(y_o), cp.asnumpy(x_o)

max_dy = float(np.max(np.abs(y_warp - y_hawk)))
max_dx = float(np.max(np.abs(x_warp - x_hawk)))
np.testing.assert_array_equal(y_warp, y_hawk)
np.testing.assert_array_equal(x_warp, x_hawk)
msg = f"Warp and hawk agree exactly over {N} samples:"
print(msg, f"max|dy|={max_dy}, max|dx|={max_dx}")

Warp and hawk agree exactly over 2048 samples: max|dy|=0.0, max|dx|=0.0

They work together#

A wp.array handed to a hawk kernel through eagle’s zero-copy DLPack exchange (see Interoperability) — no copy, no intermediate framework. cp.from_dlpack views Warp’s own device memory; the pointer identity below proves it.

heights = np.array([10.0, 9.0, 8.0, 7.0, 6.0])
height_wp = wp.array(heights, dtype=wp.float64, device=DEVICE)
height_cp = cp.from_dlpack(height_wp)  # zero-copy view of Warp's own buffer
print("zero-copy (same pointer):", height_cp.data.ptr == height_wp.ptr)

n5 = height_wp.shape[0]
y_out5 = cp.zeros(n5)
handoff = the_plan.bind(
    x=cp.zeros(n5), y=height_cp, vx=cp.zeros(n5), vy=cp.zeros(n5), g=G, dt=DT,
    terminated=cp.zeros(n5, dtype=cp.bool_),
    x_out=cp.zeros(n5), y_out=y_out5, vx_out=cp.zeros(n5), vy_out=cp.zeros(n5))
handoff.launch()
cp.cuda.get_current_stream().synchronize()
print("one hawk step from a Warp-owned height plane:", y_out5.get())

zero-copy (same pointer): True
one hawk step from a Warp-owned height plane: [9.996076 8.996076 7.996076 6.996076 5.996076]

Where Warp shines#

Warp is the better choice when the job calls for its own strengths: a built-in geometry, mesh and simulation library, and reverse-mode gradients of a whole recorded simulation through its tape, rather than of one kernel at a time. Neither tool replaces the other — they pass arrays zero-copy and are reaching for different parts of the same problem.

What just happened#

The same physics, written against two different vocabularies, agrees exactly: Warp’s explicit thread-level control, and hawk’s declared terminated role that the runtime enforces for you.
eagle.deploy(...).bind(...) packs persistent device buffers once; .launch() replays the same compiled kernel every step with no new allocation — the same shape Warp’s own in-place arrays use.
A wp.array and a cupy array holding a hawk kernel’s output share memory with no copy, through the DLPack protocol both already speak.

Next#

Compiling and caching — what eagle.deploy’s cache actually keys on, and when a change forces a real rebuild.

NVIDIA Warp and hawk both let you write a GPU kernel in Python. They start from different places, and that shows in what each one lets you control. This page lines the two up. Warp is described as its own documentation describes it; for anything beyond that, read Warp’s docs. The two tools exchange arrays without a copy (the worked example above shows it), so the choice is per kernel.

Low-level control, side by side#

What you control#

	Warp, as documented by Warp	hawk and eagle
The unit you write	one thread, found with `wp.tid()`; you index arrays yourself	one lane (a sample, or a point of a declared `Grid`); planes declare what each argument is (`Scalar`, `Param`, `Mutable[...]`)
Reaching another sample	read any array element; update with atomics such as `wp.atomic_add`	by name: `u.at(i)` to read, `hawk.scatter(y)[key] += v` to accumulate, `hawk.sum() max() min() mean()` to reduce (Across lanes)
Sums over the threads	atomics, or Warp’s array and tile reductions	a reduction stored to a cell; `policy=` picks `"auto"`, `"compensated"` or `"ordered"` (the same bits on every run, CPU and GPU)
Launch and graph	`wp.launch`, streams, CUDA graph capture	`eagle.deploy`, `eagle.Graph`, graph capture and replay, `eagle.simulate` (Compiling)
Gradients	a recorded tape of launches, replayed backward	a derived `jvp`/`vjp` per kernel, `step_vjp` for a stepping kernel, `eagle.diff.reverse` over a recorded graph
Escape hatch	native code from Python	`hawk.raw_device`, for a block of device code

Where hawk and eagle now come close#

Cross-sample steps used to be separate plane kinds (Accum, Reduce). Now they are written in the body, as named forms, and each says what it costs. A read is u.at(i). A by-key += is the target’s atomic add, so its order is not fixed. A reduction is finished in a fixed order if you ask for policy="ordered". A worked example: One sample, one lane.

Where Warp still offers more#

Thread-level freedom. In Warp any thread can read or write any element, loop over neighbours or call a function on a mesh. hawk keeps one lane per body, and a read of another lane (.at) is of an input plane, so it cannot read a draw or a plane this launch writes.
Cooperation inside a kernel. Warp documents tile operations for block-level work. In hawk a reduced value read later in the body ends the launch: the kernel runs as launches in sequence (the Limits section of Across lanes).
More atomics. Warp has atomic min and max. hawk’s by-key store is += or =; take a max or min over all lanes with hawk.max(x) / hawk.min(x) into a cell, not by key.
Libraries. Warp ships geometry, meshes and simulation building blocks, and a tape over a whole simulation. hawk differentiates one kernel or one recorded graph at a time.

Neither list is a ranking. If a job needs thread-level freedom or Warp’s geometry, write it in Warp and hand the arrays to hawk; if it is many independent samples with a few named exchanges between them, hawk keeps the kernel small and the sums reproducible.