Adaptive RK7(8): results#
The RK4 oscillator above is a fixed-step problem: every accepted step costs
the same arithmetic, and the four friction points (launch overhead, host
round trips, finished samples, idle lanes) show up against a constant
per-step cost. The second card runs the same class of problem through an
adaptive solver instead: a two-body Kepler orbit (float64, eccentricity
uniform in [0, 0.9]) integrated with Fehlberg’s embedded RKF7(8) pair to
rtol = atol = 1e-10, so each sample also decides, step by step, whether to
accept or shrink – a second, data-dependent source of per-sample
divergence on top of each sample’s own stopping time. eagle.simulate
appears here too, running the same attempt kernel the hand-built hawk + eagle graph arm uses, under its automatic launch policy.
Pick a device with the tabs; the table of every arm and every N is collapsed under the chart. How to read the numbers: Reading the numbers.
GPU#
Run on: driver 580.178.04; CUDA runtime 12.9; host CPU Intel Xeon Platinum 8481C CPU @ 2.70GHz; compiler x86_64-conda-linux-gnu-g++ (conda-forge gcc 14.4.0-7) 14.4.0; Python 3.12.14; Warp 1.17.0, JAX 0.11.2, PyTorch 2.12.1+cu129, CuPy 14.2.0, NumPy 2.5.3.
N = 1,000,000
Log scale, from 27.3 ms to 14.5 s.
Table of every arm and every N
Device: NVIDIA H100 80GB HBM3 (compute capability 9.0, 132 SMs). Peak FP64: 33.5 T FLOP/s. Medians of 5 interleaved repetitions (3 for an arm whose warm-up says 5 would take over 60 s); IQR in parentheses.
Workload: two-body Kepler orbits (mu = 1, float64) from periapsis, a = 1, e uniform in [0, 0.9], random orientation; adaptive RKF7(8) (Fehlberg’s tableau, a mixed absolute/relative max-norm error control and a standard step-size controller). Tolerance rtol = atol = 1e-10; final time 64 (10.2 periods). 1106 FLOP per attempt.
arm |
loop structure |
lines of code |
|---|---|---|
eagle graph |
one CUDA graph: device WHILE loop of the attempt kernel, which marks and counts its finished samples and runs the attempts per launch eagle’s policy picks; finished samples skipped by the Terminated guard |
22 |
eagle graph + compaction |
as eagle graph, the attempt launched over the live samples only (index map rebuilt every 16 attempts) |
23 |
eagle graph + compaction + reorder |
as + compaction, live samples moved to the front of every plane when spread thin; restored at the end |
23 |
hawk + eagle persistent |
ONE launch, no graph, no map, no policy kernel: a grid sized from the SM count and the kernel’s occupancy, each lane fetching base + atomicAdd(counter, 1) whenever its sample is done |
22 |
eagle.simulate (eagle picks the launch mode) |
eagle.simulate takes the kernel, its state and its parameters by name; eagle’s automatic policy picks the launch mode from the batch (one launch for a small batch, the persistent launch above it); the row records the mode taken |
21 |
CuPy masked |
host loop: one masked whole-batch attempt (array expressions), host check per attempt |
25 |
PyTorch masked |
host loop: one masked whole-batch attempt (tensor expressions), host check per attempt |
27 |
PyTorch + CUDA graph |
host loop: a CUDA graph of 16 masked attempts replayed, host check per block |
53 |
JAX jit(vmap(while_loop)) |
one compiled executable: per-sample while loop, batched (runs until the slowest sample finishes) |
25 |
Warp per-thread kernel |
one launch: each thread loops over its own sample’s attempts to the final time |
74 |
eagle CPU (OpenMP) |
host loop: one attempt over the batch (OpenMP), the kernel marks and counts its finished samples |
22 |
Lines of code: non-blank, non-comment lines between the # >>> code:NAME and # <<< code:NAME markers of the script; every arm but Warp also uses the shared scheme block (the RKF7(8) stages, the right-hand side, the error norm), counted once in shared_scheme_lines (25 lines).
N |
arm |
wall (IQR) |
kernel-only |
end-to-end |
attempts·samples/s |
useful FLOP/s (% peak) |
device memory (× minimum) |
host memory |
|---|---|---|---|---|---|---|---|---|
1,000 |
eagle graph |
3.3 ms (32.7 µs) |
– |
3.61 ms |
153 M |
169 G (0.51%) |
2 MiB (29.1× min.) |
2.31 MiB |
1,000 |
eagle graph + compaction |
6.31 ms (20.2 µs) |
– |
6.57 ms |
80 M |
88.5 G (0.26%) |
2 MiB (29.1× min.) |
2.23 MiB |
1,000 |
eagle graph + compaction + reorder |
6.44 ms (6.45 µs) |
– |
6.7 ms |
78.4 M |
86.7 G (0.26%) |
2 MiB (29.1× min.) |
2.49 MiB |
1,000 |
hawk + eagle persistent |
3.17 ms (28.1 µs) |
– |
3.34 ms |
159 M |
176 G (0.53%) |
0 MiB (0× min.) |
2.22 MiB |
1,000 |
eagle.simulate (eagle picks the launch mode) |
2.67 ms (5.53 µs) |
– |
2.85 ms |
189 M |
209 G (0.62%) |
0 MiB (0× min.) |
2.19 MiB |
1,000 |
CuPy masked |
8.09 s (180 ms) |
– |
8.09 s |
62.4 k |
69 M (0.00021%) |
0 MiB (0× min.) |
7.25 MiB |
1,000 |
PyTorch masked |
6.33 s (126 ms) |
– |
6.33 s |
79.8 k |
88.2 M (0.00026%) |
92 MiB (1.34e+03× min.) |
0.207 MiB |
1,000 |
PyTorch + CUDA graph |
1.71 s (138 µs) |
– |
1.71 s |
296 k |
327 M (0.00098%) |
204 MiB (2.97e+03× min.) |
19.9 MiB |
1,000 |
JAX jit(vmap(while_loop)) |
35.1 ms (232 µs) |
– |
37 ms |
14.4 M |
15.9 G (0.048%) |
160 MiB (2.33e+03× min.) |
0.043 MiB |
1,000 |
Warp per-thread kernel |
3.59 ms (28.8 µs) |
– |
4.26 ms |
141 M |
155 G (0.46%) |
30 MiB (437× min.) |
0 MiB |
1,000 |
eagle CPU (OpenMP) |
7.3 ms (71.5 µs) |
– |
7.39 ms |
69.2 M |
76.5 G (–) |
– |
0.727 MiB |
10,000 |
eagle graph |
3.34 ms (17 µs) |
– |
3.99 ms |
1.48 G |
1.64 T (4.9%) |
2 MiB (2.91× min.) |
4.29 MiB |
10,000 |
eagle graph + compaction |
6.44 ms (7.5 µs) |
– |
7 ms |
770 M |
851 G (2.5%) |
2 MiB (2.91× min.) |
4.31 MiB |
10,000 |
eagle graph + compaction + reorder |
7.14 ms (1.99 µs) |
– |
7.73 ms |
694 M |
768 G (2.3%) |
2 MiB (2.91× min.) |
5.01 MiB |
10,000 |
hawk + eagle persistent |
3.23 ms (15.6 µs) |
– |
3.71 ms |
1.53 G |
1.7 T (5.1%) |
0 MiB (0× min.) |
4.12 MiB |
10,000 |
eagle.simulate (eagle picks the launch mode) |
2.72 ms (9.19 µs) |
– |
3.23 ms |
1.83 G |
2.02 T (6%) |
0 MiB (0× min.) |
4.16 MiB |
10,000 |
CuPy masked |
8.01 s (18 ms) |
– |
8.01 s |
619 k |
685 M (0.002%) |
0 MiB (0× min.) |
8.15 MiB |
10,000 |
PyTorch masked |
6.25 s (3.46 ms) |
– |
6.25 s |
793 k |
877 M (0.0026%) |
92 MiB (134× min.) |
1.91 MiB |
10,000 |
PyTorch + CUDA graph |
1.84 s (65.9 µs) |
– |
1.84 s |
2.7 M |
2.98 G (0.0089%) |
212 MiB (309× min.) |
21.7 MiB |
10,000 |
JAX jit(vmap(while_loop)) |
40.5 ms (419 µs) |
– |
43.2 ms |
123 M |
136 G (0.41%) |
160 MiB (233× min.) |
0.109 MiB |
10,000 |
Warp per-thread kernel |
3.63 ms (11 µs) |
– |
4.59 ms |
1.37 G |
1.51 T (4.5%) |
30 MiB (43.7× min.) |
1.39 MiB |
10,000 |
eagle CPU (OpenMP) |
67.5 ms (1.74 ms) |
– |
67.8 ms |
73.4 M |
81.2 G (–) |
– |
2.78 MiB |
100,000 |
eagle graph |
10.3 ms (15.4 µs) |
– |
14.2 ms |
4.84 G |
5.35 T (16%) |
10 MiB (1.46× min.) |
25.4 MiB |
100,000 |
eagle graph + compaction |
12.7 ms (23.7 µs) |
– |
16 ms |
3.92 G |
4.33 T (13%) |
12 MiB (1.75× min.) |
25.5 MiB |
100,000 |
eagle graph + compaction + reorder |
13 ms (29.4 µs) |
– |
16 ms |
3.84 G |
4.25 T (13%) |
14 MiB (2.04× min.) |
26.5 MiB |
100,000 |
hawk + eagle persistent |
7.3 ms (21.8 µs) |
– |
10.1 ms |
6.83 G |
7.55 T (23%) |
8 MiB (1.17× min.) |
25.2 MiB |
100,000 |
eagle.simulate (eagle picks the launch mode) |
7.2 ms (46.8 µs) |
– |
9.96 ms |
6.92 G |
7.65 T (23%) |
8 MiB (1.17× min.) |
25.3 MiB |
100,000 |
CuPy masked |
8.13 s (149 ms) |
– |
8.13 s |
6.13 M |
6.78 G (0.02%) |
8 MiB (1.17× min.) |
23.1 MiB |
100,000 |
PyTorch masked |
6.56 s (48.7 ms) |
– |
6.58 s |
7.59 M |
8.39 G (0.025%) |
100 MiB (14.6× min.) |
19.6 MiB |
100,000 |
PyTorch + CUDA graph |
2.26 s (428 µs) |
– |
2.27 s |
22 M |
24.3 G (0.073%) |
316 MiB (46× min.) |
39 MiB |
100,000 |
JAX jit(vmap(while_loop)) |
70.1 ms (1.11 ms) |
– |
78.6 ms |
710 M |
786 G (2.3%) |
160 MiB (23.3× min.) |
28.2 MiB |
100,000 |
Warp per-thread kernel |
10.5 ms (35.8 µs) |
– |
14 ms |
4.75 G |
5.25 T (16%) |
30 MiB (4.37× min.) |
18.5 MiB |
100,000 |
eagle CPU (OpenMP) |
660 ms (1.79 ms) |
– |
662 ms |
75.5 M |
83.5 G (–) |
– |
25.9 MiB |
1,000,000 |
eagle graph |
95.4 ms (729 µs) |
– |
135 ms |
5.22 G |
5.78 T (17%) |
72 MiB (1.05× min.) |
230 MiB |
1,000,000 |
eagle graph + compaction |
117 ms (236 µs) |
– |
155 ms |
4.28 G |
4.73 T (14%) |
88 MiB (1.28× min.) |
230 MiB |
1,000,000 |
eagle graph + compaction + reorder |
108 ms (267 µs) |
– |
144 ms |
4.61 G |
5.1 T (15%) |
116 MiB (1.69× min.) |
231 MiB |
1,000,000 |
hawk + eagle persistent |
43.7 ms (282 µs) |
– |
72.2 ms |
11.4 G |
12.6 T (38%) |
70 MiB (1.02× min.) |
230 MiB |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
43.6 ms (27.3 µs) |
– |
80.9 ms |
11.4 G |
12.6 T (38%) |
70 MiB (1.02× min.) |
230 MiB |
1,000,000 |
CuPy masked |
9.06 s (19.9 ms) |
– |
9.07 s |
55 M |
60.8 G (0.18%) |
78 MiB (1.14× min.) |
183 MiB |
1,000,000 |
PyTorch masked |
8.61 s (6.27 ms) |
– |
8.64 s |
57.9 M |
64 G (0.19%) |
190 MiB (2.77× min.) |
191 MiB |
1,000,000 |
PyTorch + CUDA graph |
7.17 s (3.41 ms) |
– |
7.19 s |
69.5 M |
76.9 G (0.23%) |
1.37e+03 MiB (19.9× min.) |
240 MiB |
1,000,000 |
JAX jit(vmap(while_loop)) |
1.11 s (5.96 ms) |
– |
1.19 s |
449 M |
496 G (1.5%) |
800 MiB (11.7× min.) |
320 MiB |
1,000,000 |
Warp per-thread kernel |
101 ms (550 µs) |
– |
125 ms |
4.91 G |
5.43 T (16%) |
94 MiB (1.37× min.) |
183 MiB |
1,000,000 |
eagle CPU (OpenMP) |
6.51 s (20.1 ms) |
– |
6.54 s |
76.6 M |
84.7 G (–) |
– |
253 MiB |
Agreement with the analytic orbit
Every arm, every sample: final time reached exactly; position error against the analytic Kepler state (Newton on Kepler’s equation to machine precision) below the stated bound n_acc tol (1 + v_p)(1 + 6 pi n_orbits v_p), derived from the controller’s per-step budget and the Kepler along-track amplification (see truth_bound; not fitted). Step counts are compared with the reference arm (eagle graph): any difference comes from a floating-point difference flipping one accept/reject decision (contraction into fused multiply-adds, pow and sqrt implementations differ between compilers and libraries); the table reports how many samples differ and by how much.
N |
arm |
max position error |
max error / bound |
samples with different step counts |
max Δ accepted / rejected |
|---|---|---|---|---|---|
1,000 |
eagle graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction + reorder |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
hawk + eagle persistent |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle.simulate (eagle picks the launch mode) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
CuPy masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch + CUDA graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
JAX jit(vmap(while_loop)) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
Warp per-thread kernel |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle CPU (OpenMP) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
10,000 |
eagle graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction + reorder |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
hawk + eagle persistent |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle.simulate (eagle picks the launch mode) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
CuPy masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch + CUDA graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
JAX jit(vmap(while_loop)) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
Warp per-thread kernel |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle CPU (OpenMP) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
100,000 |
eagle graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction + reorder |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
hawk + eagle persistent |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle.simulate (eagle picks the launch mode) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
CuPy masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch + CUDA graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
JAX jit(vmap(while_loop)) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
Warp per-thread kernel |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle CPU (OpenMP) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction + reorder |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
hawk + eagle persistent |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
CuPy masked |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
PyTorch masked |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
PyTorch + CUDA graph |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
Warp per-thread kernel |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
eagle CPU (OpenMP) |
2.26e-07 |
6.64e-03 |
2 |
1 / 1 |
Steps per sample at N = 1,000,000 (reference arm): accepted min 322 / median 371 / p90 571 / max 693; rejected min 1 / median 59 / max 264.
Where each tool fits
Each tool here solves the same 13-stage adaptive problem to the same tolerance and agrees with the analytic orbit; they differ in how the per-sample control flow is written and run. Numbers below are at N = 1,000,000; the fastest arm there is eagle.simulate (eagle picks the launch mode) (43.6 ms).
eagle graph. the attempt is per-sample scalar code (22 lines with its launch; the kernel finishes its own samples), one kernel per attempt looped on the device: 95.4 ms.
eagle graph + compaction. the same kernel launched over the live samples only: 117 ms; the plain graph takes 0.819× as long.
eagle graph + compaction + reorder. adds the occasional reorder of the planes: 108 ms; the compacted graph takes 1.08× as long.
hawk + eagle persistent. the plain (not active-set) attempt kernel, forced through hawk’s persist entry: one launch, no graph, no map, no policy kernel – a grid of SMs blocks of 256 lanes steals the next unfinished sample off a counter when their own finishes: 43.7 ms; the eagle graph takes 2.18× as long.
eagle.simulate (eagle picks the launch mode). the same kernel through eagle.simulate, its state and parameters by name (21 lines with the kernel): 43.6 ms; the eagle graph takes 2.19× as long.
CuPy masked. array code in the NumPy style (25 lines), the control flow as masks over the batch: 9.06 s.
PyTorch masked. the same array formulation in PyTorch (27 lines), running where a PyTorch model already runs: 8.61 s.
PyTorch + CUDA graph. records 16 attempts as a CUDA graph and replays them, removing per-kernel launch cost: 7.17 s; eager PyTorch takes 1.2× as long.
JAX jit(vmap(while_loop)). writes the per-sample loop directly (25 lines) and compiles the whole batch into one executable: 1.11 s.
Warp per-thread kernel. an explicit per-thread loop in a Python-embedded kernel language (74 lines, its own stages): 101 ms.
eagle CPU (OpenMP). the same hawk kernel on the host’s cores, no GPU needed: 6.51 s.
Compile time
Its own pass, never part of any wall. Each measurement is a fresh process: imports and the GPU context run untimed, then the timer spans the first call to a ready kernel (trace + code generation + compile + load, or the cache lookup). Cold: every cache empty (HAWK_CACHE_DIR, CUPY_CACHE_DIR, CUDA_CACHE_PATH, WARP_CACHE_PATH and JAX’s persistent compilation cache, each a fresh directory); warm: a new process on the caches the first cold run populated. eagle GPU: eagle.deploy of the hawk attempt kernel and of its active-set variant (host and device builds each), and one eagle graph of each built and run on 64 samples (eagle.simulate deploys the first again, a cache hit). eagle CPU: the same eagle.deploy of the attempt kernel (host and device builds) and one host run. CuPy: one masked attempt on 64 samples (every elementwise kernel compiled on first use). JAX: lower + compile from abstract shapes at N. Warp: the kernel module (code generation, compile, load) and one launch. Medians of 3 cold and 3 warm processes.
arm |
compiled |
cold (median of 3) |
warm (median of 3) |
|---|---|---|---|
eagle graph, eagle graph + compaction, eagle graph + compaction + reorder, hawk + eagle persistent, eagle.simulate (eagle picks the launch mode) |
eagle.deploy of the attempt kernel and its active-set variant |
2.48 s |
333 ms |
eagle CPU (OpenMP) |
eagle.deploy of the attempt kernel, host run |
2.04 s |
246 ms |
CuPy masked |
CuPy elementwise kernels |
2.58 s |
1.05 s |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000 |
2.23 s |
272 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 10,000 |
2.44 s |
274 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 100,000 |
2.92 s |
270 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000,000 |
4.22 s |
267 ms |
Warp per-thread kernel |
Warp kernel module |
9.28 s |
32.6 ms |
No compile step: PyTorch masked, PyTorch + CUDA graph (PyTorch’s eager kernels ship prebuilt; the CUDA graph is recorded, not compiled).
Memory method
A separate pass after the timing pass. Each (arm, N) runs in a fresh process: imports, the CUDA context created by one tiny device operation, then the baseline read, before the arm builds or allocates anything; then the arm’s kernels built and a warm-up run on a 64-sample batch (every compile happens here; JAX also compiles for shape N), so every library’s pool reserved during the warm-up counts in the reading; CuPy’s pools and PyTorch’s cache trimmed after the warm-up; then the arm built for N (eagle’s graph, PyTorch’s captured graph) and one full run with upload and download; then the reading again, the allocators still holding their high-water reservation (the same measure for every arm, no sampling needed). Device: the memory NVML accounts to the process on the timed GPU, after minus baseline. Caveats: each allocator’s rounding and growth policy count (JAX’s allocator grows in regions that can exceed the request; PyTorch rounds blocks up and a captured CUDA graph keeps a private pool; Warp’s stream-ordered pool reserves in chunks of tens of MiB, its release threshold set to 0 after the warm-up); the driver accounts in pages of about 2 MiB, so small-N rows read 0 or one page. Host: the peak resident set (VmHWM, reset to the current RSS at the host baseline, taken after the warm-up) minus the baseline RSS. Minimum: the planes the workload needs, N x (8 float64 planes (state, time, step) x 8 bytes + 2 counters x 4 bytes) = N x 72 bytes.
Run on: driver 580.178.04; CUDA runtime 12.9; host CPU Intel Xeon CPU @ 2.20GHz; compiler x86_64-conda-linux-gnu-g++ (conda-forge gcc 14.4.0-7) 14.4.0; Python 3.12.14; Warp 1.17.0, JAX 0.11.2, PyTorch 2.12.1+cu129, CuPy 14.2.0, NumPy 2.5.3.
N = 1,000,000
Log scale, from 1.45 s to 88.0 s.
Table of every arm and every N
Device: NVIDIA L4 (compute capability 8.9, 58 SMs). Peak FP64: 473 G FLOP/s. Medians of 5 interleaved repetitions (3 for an arm whose warm-up says 5 would take over 60 s); IQR in parentheses.
Workload: two-body Kepler orbits (mu = 1, float64) from periapsis, a = 1, e uniform in [0, 0.9], random orientation; adaptive RKF7(8) (Fehlberg’s tableau, a mixed absolute/relative max-norm error control and a standard step-size controller). Tolerance rtol = atol = 1e-10; final time 64 (10.2 periods). 1106 FLOP per attempt.
arm |
loop structure |
lines of code |
|---|---|---|
eagle graph |
one CUDA graph: device WHILE loop of the attempt kernel, which marks and counts its finished samples and runs the attempts per launch eagle’s policy picks; finished samples skipped by the Terminated guard |
22 |
eagle graph + compaction |
as eagle graph, the attempt launched over the live samples only (index map rebuilt every 16 attempts) |
23 |
eagle graph + compaction + reorder |
as + compaction, live samples moved to the front of every plane when spread thin; restored at the end |
23 |
hawk + eagle persistent |
ONE launch, no graph, no map, no policy kernel: a grid sized from the SM count and the kernel’s occupancy, each lane fetching base + atomicAdd(counter, 1) whenever its sample is done |
22 |
eagle.simulate (eagle picks the launch mode) |
eagle.simulate takes the kernel, its state and its parameters by name; eagle’s automatic policy picks the launch mode from the batch (one launch for a small batch, the persistent launch above it); the row records the mode taken |
21 |
CuPy masked |
host loop: one masked whole-batch attempt (array expressions), host check per attempt |
25 |
PyTorch masked |
host loop: one masked whole-batch attempt (tensor expressions), host check per attempt |
27 |
PyTorch + CUDA graph |
host loop: a CUDA graph of 16 masked attempts replayed, host check per block |
53 |
JAX jit(vmap(while_loop)) |
one compiled executable: per-sample while loop, batched (runs until the slowest sample finishes) |
25 |
Warp per-thread kernel |
one launch: each thread loops over its own sample’s attempts to the final time |
74 |
eagle CPU (OpenMP) |
host loop: one attempt over the batch (OpenMP), the kernel marks and counts its finished samples |
22 |
Lines of code: non-blank, non-comment lines between the # >>> code:NAME and # <<< code:NAME markers of the script; every arm but Warp also uses the shared scheme block (the RKF7(8) stages, the right-hand side, the error norm), counted once in shared_scheme_lines (25 lines).
N |
arm |
wall (IQR) |
kernel-only |
end-to-end |
attempts·samples/s |
useful FLOP/s (% peak) |
device memory (× minimum) |
host memory |
|---|---|---|---|---|---|---|---|---|
1,000 |
eagle graph |
13.6 ms (54 µs) |
– |
14.1 ms |
37.1 M |
41 G (8.7%) |
0 MiB (0× min.) |
2.65 MiB |
1,000 |
eagle graph + compaction |
16.2 ms (38.2 µs) |
– |
16.6 ms |
31.2 M |
34.6 G (7.3%) |
0 MiB (0× min.) |
2.52 MiB |
1,000 |
eagle graph + compaction + reorder |
16.4 ms (29.2 µs) |
– |
16.8 ms |
30.8 M |
34.1 G (7.2%) |
0 MiB (0× min.) |
2.52 MiB |
1,000 |
hawk + eagle persistent |
51.6 ms (6.93 ms) |
– |
52 ms |
9.78 M |
10.8 G (2.3%) |
0 MiB (0× min.) |
2.21 MiB |
1,000 |
eagle.simulate (eagle picks the launch mode) |
12.7 ms (15.2 µs) |
– |
13.1 ms |
39.6 M |
43.8 G (9.3%) |
0 MiB (0× min.) |
2.2 MiB |
1,000 |
CuPy masked |
15 s (13.1 ms) |
– |
15 s |
33.7 k |
37.3 M (0.0079%) |
0 MiB (0× min.) |
0.00781 MiB |
1,000 |
PyTorch masked |
9.73 s (35 ms) |
– |
9.73 s |
51.9 k |
57.4 M (0.012%) |
28 MiB (408× min.) |
0.195 MiB |
1,000 |
PyTorch + CUDA graph |
1.64 s (508 µs) |
– |
1.64 s |
308 k |
340 M (0.072%) |
176 MiB (2.56e+03× min.) |
14.5 MiB |
1,000 |
JAX jit(vmap(while_loop)) |
91.3 ms (159 µs) |
– |
93.8 ms |
5.53 M |
6.12 G (1.3%) |
138 MiB (2.01e+03× min.) |
0.0391 MiB |
1,000 |
Warp per-thread kernel |
64.3 ms (222 µs) |
– |
65.3 ms |
7.85 M |
8.68 G (1.8%) |
30 MiB (437× min.) |
0 MiB |
1,000 |
eagle CPU (OpenMP) |
23.2 ms (28.1 µs) |
– |
23.4 ms |
21.7 M |
24 G (–) |
– |
0.52 MiB |
10,000 |
eagle graph |
46.8 ms (122 µs) |
– |
47.9 ms |
106 M |
117 G (25%) |
0 MiB (0× min.) |
4.26 MiB |
10,000 |
eagle graph + compaction |
36.4 ms (46.6 µs) |
– |
37.4 ms |
136 M |
150 G (32%) |
0 MiB (0× min.) |
4.28 MiB |
10,000 |
eagle graph + compaction + reorder |
37.2 ms (41 µs) |
– |
38.2 ms |
133 M |
148 G (31%) |
0 MiB (0× min.) |
4.99 MiB |
10,000 |
hawk + eagle persistent |
59.6 ms (548 µs) |
– |
60.5 ms |
83.1 M |
92 G (19%) |
0 MiB (0× min.) |
4.15 MiB |
10,000 |
eagle.simulate (eagle picks the launch mode) |
45.2 ms (13 µs) |
– |
46.2 ms |
110 M |
121 G (26%) |
0 MiB (0× min.) |
4.16 MiB |
10,000 |
CuPy masked |
15 s (30.3 ms) |
– |
15 s |
330 k |
365 M (0.077%) |
0 MiB (0× min.) |
7.8 MiB |
10,000 |
PyTorch masked |
9.85 s (31.3 ms) |
– |
9.85 s |
503 k |
557 M (0.12%) |
28 MiB (40.8× min.) |
1.91 MiB |
10,000 |
PyTorch + CUDA graph |
1.58 s (511 µs) |
– |
1.58 s |
3.13 M |
3.47 G (0.73%) |
184 MiB (268× min.) |
16.3 MiB |
10,000 |
JAX jit(vmap(while_loop)) |
218 ms (132 µs) |
– |
221 ms |
22.7 M |
25.2 G (5.3%) |
138 MiB (201× min.) |
0.0469 MiB |
10,000 |
Warp per-thread kernel |
65.4 ms (122 µs) |
– |
66.8 ms |
75.9 M |
83.9 G (18%) |
30 MiB (43.7× min.) |
1.39 MiB |
10,000 |
eagle CPU (OpenMP) |
172 ms (291 µs) |
– |
172 ms |
28.9 M |
31.9 G (–) |
– |
2.63 MiB |
100,000 |
eagle graph |
407 ms (109 µs) |
– |
413 ms |
122 M |
135 G (29%) |
8 MiB (1.17× min.) |
25.4 MiB |
100,000 |
eagle graph + compaction |
256 ms (104 µs) |
– |
261 ms |
195 M |
215 G (46%) |
10 MiB (1.46× min.) |
25.4 MiB |
100,000 |
eagle graph + compaction + reorder |
257 ms (103 µs) |
– |
262 ms |
194 M |
215 G (45%) |
12 MiB (1.75× min.) |
26.3 MiB |
100,000 |
hawk + eagle persistent |
276 ms (1.36 ms) |
– |
282 ms |
180 M |
199 G (42%) |
8 MiB (1.17× min.) |
25.3 MiB |
100,000 |
eagle.simulate (eagle picks the launch mode) |
269 ms (1.59 ms) |
– |
275 ms |
185 M |
205 G (43%) |
8 MiB (1.17× min.) |
25.3 MiB |
100,000 |
CuPy masked |
15.1 s (28.8 ms) |
– |
15.1 s |
3.3 M |
3.65 G (0.77%) |
8 MiB (1.17× min.) |
23.1 MiB |
100,000 |
PyTorch masked |
10.3 s (41 ms) |
– |
10.3 s |
4.82 M |
5.33 G (1.1%) |
36 MiB (5.24× min.) |
19.7 MiB |
100,000 |
PyTorch + CUDA graph |
3.64 s (7.76 ms) |
– |
3.65 s |
13.7 M |
15.1 G (3.2%) |
288 MiB (41.9× min.) |
33.7 MiB |
100,000 |
JAX jit(vmap(while_loop)) |
830 ms (1.33 ms) |
– |
841 ms |
60 M |
66.3 G (14%) |
138 MiB (20.1× min.) |
27.4 MiB |
100,000 |
Warp per-thread kernel |
443 ms (437 µs) |
– |
449 ms |
113 M |
124 G (26%) |
30 MiB (4.37× min.) |
18.5 MiB |
100,000 |
eagle CPU (OpenMP) |
1.65 s (4.07 ms) |
– |
1.65 s |
30.2 M |
33.4 G (–) |
– |
26 MiB |
1,000,000 |
eagle graph |
3.92 s (165 µs) |
– |
3.98 s |
127 M |
141 G (30%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
eagle graph + compaction |
2.36 s (134 µs) |
– |
2.42 s |
211 M |
234 G (49%) |
86 MiB (1.25× min.) |
231 MiB |
1,000,000 |
eagle graph + compaction + reorder |
2.34 s (158 µs) |
– |
2.4 s |
213 M |
236 G (50%) |
114 MiB (1.66× min.) |
231 MiB |
1,000,000 |
hawk + eagle persistent |
2.32 s (1.26 ms) |
– |
2.38 s |
215 M |
237 G (50%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
2.32 s (849 µs) |
– |
2.38 s |
215 M |
237 G (50%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
CuPy masked |
52.8 s (11.9 ms) |
– |
52.8 s |
9.45 M |
10.4 G (2.2%) |
78 MiB (1.14× min.) |
183 MiB |
1,000,000 |
PyTorch masked |
55 s (20.2 ms) |
– |
55 s |
9.07 M |
10 G (2.1%) |
126 MiB (1.84× min.) |
191 MiB |
1,000,000 |
PyTorch + CUDA graph |
55.1 s (819 µs) |
– |
55.2 s |
9.04 M |
10 G (2.1%) |
1.34e+03 MiB (19.5× min.) |
241 MiB |
1,000,000 |
JAX jit(vmap(while_loop)) |
7.27 s (3.54 ms) |
– |
7.39 s |
68.6 M |
75.9 G (16%) |
266 MiB (3.87× min.) |
307 MiB |
1,000,000 |
Warp per-thread kernel |
4.27 s (364 µs) |
– |
4.32 s |
117 M |
129 G (27%) |
94 MiB (1.37× min.) |
183 MiB |
1,000,000 |
eagle CPU (OpenMP) |
16.4 s (48.2 ms) |
– |
16.5 s |
30.3 M |
33.5 G (–) |
– |
252 MiB |
Agreement with the analytic orbit
Every arm, every sample: final time reached exactly; position error against the analytic Kepler state (Newton on Kepler’s equation to machine precision) below the stated bound n_acc tol (1 + v_p)(1 + 6 pi n_orbits v_p), derived from the controller’s per-step budget and the Kepler along-track amplification (see truth_bound; not fitted). Step counts are compared with the reference arm (eagle graph): any difference comes from a floating-point difference flipping one accept/reject decision (contraction into fused multiply-adds, pow and sqrt implementations differ between compilers and libraries); the table reports how many samples differ and by how much.
N |
arm |
max position error |
max error / bound |
samples with different step counts |
max Δ accepted / rejected |
|---|---|---|---|---|---|
1,000,000 |
eagle graph |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction + reorder |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
hawk + eagle persistent |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
CuPy masked |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
PyTorch masked |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
PyTorch + CUDA graph |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
Warp per-thread kernel |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
eagle CPU (OpenMP) |
2.26e-07 |
6.64e-03 |
2 |
1 / 1 |
1,000 |
eagle graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction + reorder |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
hawk + eagle persistent |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle.simulate (eagle picks the launch mode) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
CuPy masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch + CUDA graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
JAX jit(vmap(while_loop)) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
Warp per-thread kernel |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle CPU (OpenMP) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
10,000 |
eagle graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction + reorder |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
hawk + eagle persistent |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle.simulate (eagle picks the launch mode) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
CuPy masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch + CUDA graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
JAX jit(vmap(while_loop)) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
Warp per-thread kernel |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle CPU (OpenMP) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
100,000 |
eagle graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction + reorder |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
hawk + eagle persistent |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle.simulate (eagle picks the launch mode) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
CuPy masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch + CUDA graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
JAX jit(vmap(while_loop)) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
Warp per-thread kernel |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle CPU (OpenMP) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
Steps per sample at N = 1,000,000 (reference arm): accepted min 322 / median 371 / p90 571 / max 693; rejected min 1 / median 59 / max 264.
Where each tool fits
Each tool here solves the same 13-stage adaptive problem to the same tolerance and agrees with the analytic orbit; they differ in how the per-sample control flow is written and run. Numbers below are at N = 1,000,000; the fastest arm there is hawk + eagle persistent (2.32 s).
eagle graph. the attempt is per-sample scalar code (22 lines with its launch; the kernel finishes its own samples), one kernel per attempt looped on the device: 3.92 s.
eagle graph + compaction. the same kernel launched over the live samples only: 2.36 s; the plain graph takes 1.66× as long.
eagle graph + compaction + reorder. adds the occasional reorder of the planes: 2.34 s; the compacted graph takes 1.01× as long.
hawk + eagle persistent. the plain (not active-set) attempt kernel, forced through hawk’s persist entry: one launch, no graph, no map, no policy kernel – a grid of SMs blocks of 256 lanes steals the next unfinished sample off a counter when their own finishes: 2.32 s; the eagle graph takes 1.69× as long.
eagle.simulate (eagle picks the launch mode). the same kernel through eagle.simulate, its state and parameters by name (21 lines with the kernel): 2.32 s; the eagle graph takes 1.69× as long.
CuPy masked. array code in the NumPy style (25 lines), the control flow as masks over the batch: 52.8 s.
PyTorch masked. the same array formulation in PyTorch (27 lines), running where a PyTorch model already runs: 55 s.
PyTorch + CUDA graph. records 16 attempts as a CUDA graph and replays them, removing per-kernel launch cost: 55.1 s; eager PyTorch takes 0.997× as long.
JAX jit(vmap(while_loop)). writes the per-sample loop directly (25 lines) and compiles the whole batch into one executable: 7.27 s.
Warp per-thread kernel. an explicit per-thread loop in a Python-embedded kernel language (74 lines, its own stages): 4.27 s.
eagle CPU (OpenMP). the same hawk kernel on the host’s cores, no GPU needed: 16.4 s.
Compile time
Its own pass, never part of any wall. Each measurement is a fresh process: imports and the GPU context run untimed, then the timer spans the first call to a ready kernel (trace + code generation + compile + load, or the cache lookup). Cold: every cache empty (HAWK_CACHE_DIR, CUPY_CACHE_DIR, CUDA_CACHE_PATH, WARP_CACHE_PATH and JAX’s persistent compilation cache, each a fresh directory); warm: a new process on the caches the first cold run populated. eagle GPU: eagle.deploy of the hawk attempt kernel and of its active-set variant (host and device builds each), and one eagle graph of each built and run on 64 samples (eagle.simulate deploys the first again, a cache hit). eagle CPU: the same eagle.deploy of the attempt kernel (host and device builds) and one host run. CuPy: one masked attempt on 64 samples (every elementwise kernel compiled on first use). JAX: lower + compile from abstract shapes at N. Warp: the kernel module (code generation, compile, load) and one launch. Medians of 3 cold and 3 warm processes.
arm |
compiled |
cold (median of 3) |
warm (median of 3) |
|---|---|---|---|
eagle graph, eagle graph + compaction, eagle graph + compaction + reorder, hawk + eagle persistent, eagle.simulate (eagle picks the launch mode) |
eagle.deploy of the attempt kernel and its active-set variant |
3.81 s |
593 ms |
eagle CPU (OpenMP) |
eagle.deploy of the attempt kernel, host run |
3.39 s |
427 ms |
CuPy masked |
CuPy elementwise kernels |
4.03 s |
1.65 s |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000 |
3.5 s |
417 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 10,000 |
4.46 s |
413 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 100,000 |
4.61 s |
400 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000,000 |
4.68 s |
397 ms |
Warp per-thread kernel |
Warp kernel module |
12.3 s |
62.3 ms |
No compile step: PyTorch masked, PyTorch + CUDA graph (PyTorch’s eager kernels ship prebuilt; the CUDA graph is recorded, not compiled).
Memory method
A separate pass after the timing pass. Each (arm, N) runs in a fresh process: imports, the CUDA context created by one tiny device operation, then the baseline read, before the arm builds or allocates anything; then the arm’s kernels built and a warm-up run on a 64-sample batch (every compile happens here; JAX also compiles for shape N), so every library’s pool reserved during the warm-up counts in the reading; CuPy’s pools and PyTorch’s cache trimmed after the warm-up; then the arm built for N (eagle’s graph, PyTorch’s captured graph) and one full run with upload and download; then the reading again, the allocators still holding their high-water reservation (the same measure for every arm, no sampling needed). Device: the memory NVML accounts to the process on the timed GPU, after minus baseline. Caveats: each allocator’s rounding and growth policy count (JAX’s allocator grows in regions that can exceed the request; PyTorch rounds blocks up and a captured CUDA graph keeps a private pool; Warp’s stream-ordered pool reserves in chunks of tens of MiB, its release threshold set to 0 after the warm-up); the driver accounts in pages of about 2 MiB, so small-N rows read 0 or one page. Host: the peak resident set (VmHWM, reset to the current RSS at the host baseline, taken after the warm-up) minus the baseline RSS. Minimum: the planes the workload needs, N x (8 float64 planes (state, time, step) x 8 bytes + 2 counters x 4 bytes) = N x 72 bytes.
Fastest arm by median wall time:
N = 1,000: eagle.simulate (eagle picks the launch mode) (12.7 ms), then eagle graph (13.6 ms, 1.07× the time)
N = 10,000: eagle graph + compaction (36.4 ms), then eagle graph + compaction + reorder (37.2 ms, 1.02× the time)
N = 100,000: eagle graph + compaction (256 ms), then eagle graph + compaction + reorder (257 ms, 1× the time)
N = 1,000,000: hawk + eagle persistent (2.32 s), then eagle.simulate (eagle picks the launch mode) (2.32 s, 1× the time)
eagle_simulate against its hand-built twin eagle_graph:
N = 1,000: wall 0.935× the twin’s time
N = 10,000: wall 0.967× the twin’s time
N = 100,000: wall 0.661× the twin’s time
N = 1,000,000: wall 0.593× the twin’s time
Run on: driver 580.178.04; CUDA runtime 12.9; host CPU Intel Xeon CPU @ 2.20GHz; compiler x86_64-conda-linux-gnu-g++ (conda-forge gcc 14.4.0-7) 14.4.0; Python 3.12.14; Warp 1.17.0, JAX 0.11.2, PyTorch 2.12.1+cu129, CuPy 14.2.0, NumPy 2.5.3.
N = 1,000,000
Log scale, from 66.6 ms to 24.7 s.
Table of every arm and every N
Device: NVIDIA A100-SXM4-40GB (compute capability 8.0, 108 SMs). Peak FP64: 9.75 T FLOP/s. Medians of 5 interleaved repetitions (3 for an arm whose warm-up says 5 would take over 60 s); IQR in parentheses.
Workload: two-body Kepler orbits (mu = 1, float64) from periapsis, a = 1, e uniform in [0, 0.9], random orientation; adaptive RKF7(8) (Fehlberg’s tableau, a mixed absolute/relative max-norm error control and a standard step-size controller). Tolerance rtol = atol = 1e-10; final time 64 (10.2 periods). 1106 FLOP per attempt.
arm |
loop structure |
lines of code |
|---|---|---|
eagle graph |
one CUDA graph: device WHILE loop of the attempt kernel, which marks and counts its finished samples and runs the attempts per launch eagle’s policy picks; finished samples skipped by the Terminated guard |
22 |
eagle graph + compaction |
as eagle graph, the attempt launched over the live samples only (index map rebuilt every 16 attempts) |
23 |
eagle graph + compaction + reorder |
as + compaction, live samples moved to the front of every plane when spread thin; restored at the end |
23 |
hawk + eagle persistent |
ONE launch, no graph, no map, no policy kernel: a grid sized from the SM count and the kernel’s occupancy, each lane fetching base + atomicAdd(counter, 1) whenever its sample is done |
22 |
eagle.simulate (eagle picks the launch mode) |
eagle.simulate takes the kernel, its state and its parameters by name; eagle’s automatic policy picks the launch mode from the batch (one launch for a small batch, the persistent launch above it); the row records the mode taken |
21 |
CuPy masked |
host loop: one masked whole-batch attempt (array expressions), host check per attempt |
25 |
PyTorch masked |
host loop: one masked whole-batch attempt (tensor expressions), host check per attempt |
27 |
PyTorch + CUDA graph |
host loop: a CUDA graph of 16 masked attempts replayed, host check per block |
53 |
JAX jit(vmap(while_loop)) |
one compiled executable: per-sample while loop, batched (runs until the slowest sample finishes) |
25 |
Warp per-thread kernel |
one launch: each thread loops over its own sample’s attempts to the final time |
74 |
eagle CPU (OpenMP) |
host loop: one attempt over the batch (OpenMP), the kernel marks and counts its finished samples |
22 |
Lines of code: non-blank, non-comment lines between the # >>> code:NAME and # <<< code:NAME markers of the script; every arm but Warp also uses the shared scheme block (the RKF7(8) stages, the right-hand side, the error norm), counted once in shared_scheme_lines (25 lines).
N |
arm |
wall (IQR) |
kernel-only |
end-to-end |
attempts·samples/s |
useful FLOP/s (% peak) |
device memory (× minimum) |
host memory |
|---|---|---|---|---|---|---|---|---|
1,000 |
eagle graph |
4.96 ms (91.1 µs) |
– |
5.49 ms |
102 M |
113 G (1.2%) |
0 MiB (0× min.) |
2.32 MiB |
1,000 |
eagle graph + compaction |
8.87 ms (245 µs) |
– |
9.35 ms |
56.9 M |
63 G (0.65%) |
0 MiB (0× min.) |
2.11 MiB |
1,000 |
eagle graph + compaction + reorder |
9.12 ms (21.3 µs) |
– |
9.59 ms |
55.4 M |
61.2 G (0.63%) |
0 MiB (0× min.) |
2.4 MiB |
1,000 |
hawk + eagle persistent |
5.66 ms (44.7 µs) |
– |
5.98 ms |
89.3 M |
98.7 G (1%) |
0 MiB (0× min.) |
2.19 MiB |
1,000 |
eagle.simulate (eagle picks the launch mode) |
3.98 ms (3.7 µs) |
– |
4.32 ms |
127 M |
140 G (1.4%) |
0 MiB (0× min.) |
2.2 MiB |
1,000 |
CuPy masked |
14.7 s (27.4 ms) |
– |
14.7 s |
34.4 k |
38 M (0.00039%) |
0 MiB (0× min.) |
6.96 MiB |
1,000 |
PyTorch masked |
10.2 s (21.5 ms) |
– |
10.2 s |
49.6 k |
54.8 M (0.00056%) |
74 MiB (1.08e+03× min.) |
0.199 MiB |
1,000 |
PyTorch + CUDA graph |
1.92 s (19.3 ms) |
– |
1.92 s |
263 k |
290 M (0.003%) |
222 MiB (3.23e+03× min.) |
14.5 MiB |
1,000 |
JAX jit(vmap(while_loop)) |
39.5 ms (219 µs) |
– |
42.1 ms |
12.8 M |
14.1 G (0.15%) |
138 MiB (2.01e+03× min.) |
0.0469 MiB |
1,000 |
Warp per-thread kernel |
6.29 ms (63.2 µs) |
– |
7.3 ms |
80.3 M |
88.8 G (0.91%) |
30 MiB (437× min.) |
0 MiB |
1,000 |
eagle CPU (OpenMP) |
17.3 ms (973 µs) |
– |
17.5 ms |
29.2 M |
32.3 G (–) |
– |
0.805 MiB |
10,000 |
eagle graph |
7.08 ms (6.11 µs) |
– |
8.17 ms |
700 M |
774 G (7.9%) |
0 MiB (0× min.) |
4.27 MiB |
10,000 |
eagle graph + compaction |
9.74 ms (40.3 µs) |
– |
10.7 ms |
509 M |
563 G (5.8%) |
0 MiB (0× min.) |
4.22 MiB |
10,000 |
eagle graph + compaction + reorder |
10.7 ms (33.2 µs) |
– |
11.6 ms |
464 M |
513 G (5.3%) |
0 MiB (0× min.) |
4.88 MiB |
10,000 |
hawk + eagle persistent |
5.7 ms (11.9 µs) |
– |
6.46 ms |
869 M |
961 G (9.9%) |
0 MiB (0× min.) |
4.14 MiB |
10,000 |
eagle.simulate (eagle picks the launch mode) |
5.7 ms (37.6 µs) |
– |
6.47 ms |
870 M |
962 G (9.9%) |
0 MiB (0× min.) |
4.13 MiB |
10,000 |
CuPy masked |
14.9 s (25.6 ms) |
– |
14.9 s |
334 k |
369 M (0.0038%) |
0 MiB (0× min.) |
7.8 MiB |
10,000 |
PyTorch masked |
10.3 s (26.4 ms) |
– |
10.3 s |
479 k |
530 M (0.0054%) |
74 MiB (108× min.) |
1.91 MiB |
10,000 |
PyTorch + CUDA graph |
1.93 s (8.12 ms) |
– |
1.93 s |
2.57 M |
2.84 G (0.029%) |
230 MiB (335× min.) |
16.3 MiB |
10,000 |
JAX jit(vmap(while_loop)) |
65.7 ms (97.3 µs) |
– |
69 ms |
75.5 M |
83.5 G (0.86%) |
138 MiB (201× min.) |
0.133 MiB |
10,000 |
Warp per-thread kernel |
6.24 ms (15.2 µs) |
– |
7.68 ms |
794 M |
879 G (9%) |
30 MiB (43.7× min.) |
1.39 MiB |
10,000 |
eagle CPU (OpenMP) |
137 ms (411 µs) |
– |
138 ms |
36.2 M |
40 G (–) |
– |
2.92 MiB |
100,000 |
eagle graph |
23.1 ms (273 µs) |
– |
28.9 ms |
2.16 G |
2.38 T (24%) |
8 MiB (1.17× min.) |
25.3 MiB |
100,000 |
eagle graph + compaction |
22.8 ms (50.4 µs) |
– |
28.4 ms |
2.18 G |
2.41 T (25%) |
10 MiB (1.46× min.) |
25.2 MiB |
100,000 |
eagle graph + compaction + reorder |
23.5 ms (15.7 µs) |
– |
28.9 ms |
2.12 G |
2.35 T (24%) |
12 MiB (1.75× min.) |
26.3 MiB |
100,000 |
hawk + eagle persistent |
14.8 ms (51.7 µs) |
– |
19.9 ms |
3.36 G |
3.72 T (38%) |
8 MiB (1.17× min.) |
25.3 MiB |
100,000 |
eagle.simulate (eagle picks the launch mode) |
14.6 ms (96.5 µs) |
– |
19.8 ms |
3.42 G |
3.78 T (39%) |
8 MiB (1.17× min.) |
25.3 MiB |
100,000 |
CuPy masked |
15 s (69.2 ms) |
– |
15.1 s |
3.31 M |
3.66 G (0.038%) |
8 MiB (1.17× min.) |
23 MiB |
100,000 |
PyTorch masked |
10.8 s (42.8 ms) |
– |
10.9 s |
4.59 M |
5.08 G (0.052%) |
82 MiB (11.9× min.) |
19.6 MiB |
100,000 |
PyTorch + CUDA graph |
2.61 s (6.48 ms) |
– |
2.61 s |
19.1 M |
21.1 G (0.22%) |
334 MiB (48.6× min.) |
33.6 MiB |
100,000 |
JAX jit(vmap(while_loop)) |
149 ms (166 µs) |
– |
160 ms |
335 M |
371 G (3.8%) |
138 MiB (20.1× min.) |
27.8 MiB |
100,000 |
Warp per-thread kernel |
24 ms (23.2 µs) |
– |
30.2 ms |
2.08 G |
2.3 T (24%) |
30 MiB (4.37× min.) |
18.5 MiB |
100,000 |
eagle CPU (OpenMP) |
1.31 s (4.56 ms) |
– |
1.32 s |
37.9 M |
42 G (–) |
– |
26.1 MiB |
1,000,000 |
eagle graph |
245 ms (44.3 ms) |
– |
313 ms |
2.03 G |
2.25 T (23%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
eagle graph + compaction |
209 ms (155 µs) |
– |
273 ms |
2.38 G |
2.63 T (27%) |
86 MiB (1.25× min.) |
231 MiB |
1,000,000 |
eagle graph + compaction + reorder |
199 ms (215 µs) |
– |
251 ms |
2.51 G |
2.77 T (28%) |
114 MiB (1.66× min.) |
231 MiB |
1,000,000 |
hawk + eagle persistent |
107 ms (41.7 µs) |
– |
158 ms |
4.67 G |
5.17 T (53%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
107 ms (39.9 µs) |
– |
158 ms |
4.68 G |
5.17 T (53%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
CuPy masked |
15.4 s (15.3 ms) |
– |
15.5 s |
32.3 M |
35.7 G (0.37%) |
78 MiB (1.14× min.) |
183 MiB |
1,000,000 |
PyTorch masked |
15 s (16.9 ms) |
– |
15.1 s |
33.2 M |
36.7 G (0.38%) |
172 MiB (2.5× min.) |
191 MiB |
1,000,000 |
PyTorch + CUDA graph |
13.7 s (1.28 ms) |
– |
13.7 s |
36.4 M |
40.3 G (0.41%) |
1.39e+03 MiB (20.2× min.) |
241 MiB |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.19 s (3.79 ms) |
– |
2.31 s |
227 M |
251 G (2.6%) |
788 MiB (11.5× min.) |
319 MiB |
1,000,000 |
Warp per-thread kernel |
218 ms (42.1 µs) |
– |
263 ms |
2.28 G |
2.52 T (26%) |
94 MiB (1.37× min.) |
183 MiB |
1,000,000 |
eagle CPU (OpenMP) |
13.1 s (4.16 ms) |
– |
13.1 s |
38 M |
42.1 G (–) |
– |
253 MiB |
Agreement with the analytic orbit
Every arm, every sample: final time reached exactly; position error against the analytic Kepler state (Newton on Kepler’s equation to machine precision) below the stated bound n_acc tol (1 + v_p)(1 + 6 pi n_orbits v_p), derived from the controller’s per-step budget and the Kepler along-track amplification (see truth_bound; not fitted). Step counts are compared with the reference arm (eagle graph): any difference comes from a floating-point difference flipping one accept/reject decision (contraction into fused multiply-adds, pow and sqrt implementations differ between compilers and libraries); the table reports how many samples differ and by how much.
N |
arm |
max position error |
max error / bound |
samples with different step counts |
max Δ accepted / rejected |
|---|---|---|---|---|---|
1,000 |
eagle graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction + reorder |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
hawk + eagle persistent |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle.simulate (eagle picks the launch mode) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
CuPy masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch + CUDA graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
JAX jit(vmap(while_loop)) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
Warp per-thread kernel |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle CPU (OpenMP) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
10,000 |
eagle graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction + reorder |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
hawk + eagle persistent |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle.simulate (eagle picks the launch mode) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
CuPy masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch + CUDA graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
JAX jit(vmap(while_loop)) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
Warp per-thread kernel |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle CPU (OpenMP) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
100,000 |
eagle graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction + reorder |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
hawk + eagle persistent |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle.simulate (eagle picks the launch mode) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
CuPy masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch + CUDA graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
JAX jit(vmap(while_loop)) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
Warp per-thread kernel |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle CPU (OpenMP) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction + reorder |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
hawk + eagle persistent |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
CuPy masked |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
PyTorch masked |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
PyTorch + CUDA graph |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
Warp per-thread kernel |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
eagle CPU (OpenMP) |
2.26e-07 |
6.64e-03 |
2 |
1 / 1 |
Steps per sample at N = 1,000,000 (reference arm): accepted min 322 / median 371 / p90 571 / max 693; rejected min 1 / median 59 / max 264.
Where each tool fits
Each tool here solves the same 13-stage adaptive problem to the same tolerance and agrees with the analytic orbit; they differ in how the per-sample control flow is written and run. Numbers below are at N = 1,000,000; the fastest arm there is eagle.simulate (eagle picks the launch mode) (107 ms).
eagle graph. the attempt is per-sample scalar code (22 lines with its launch; the kernel finishes its own samples), one kernel per attempt looped on the device: 245 ms.
eagle graph + compaction. the same kernel launched over the live samples only: 209 ms; the plain graph takes 1.17× as long.
eagle graph + compaction + reorder. adds the occasional reorder of the planes: 199 ms; the compacted graph takes 1.05× as long.
hawk + eagle persistent. the plain (not active-set) attempt kernel, forced through hawk’s persist entry: one launch, no graph, no map, no policy kernel – a grid of SMs blocks of 256 lanes steals the next unfinished sample off a counter when their own finishes: 107 ms; the eagle graph takes 2.3× as long.
eagle.simulate (eagle picks the launch mode). the same kernel through eagle.simulate, its state and parameters by name (21 lines with the kernel): 107 ms; the eagle graph takes 2.3× as long.
CuPy masked. array code in the NumPy style (25 lines), the control flow as masks over the batch: 15.4 s.
PyTorch masked. the same array formulation in PyTorch (27 lines), running where a PyTorch model already runs: 15 s.
PyTorch + CUDA graph. records 16 attempts as a CUDA graph and replays them, removing per-kernel launch cost: 13.7 s; eager PyTorch takes 1.1× as long.
JAX jit(vmap(while_loop)). writes the per-sample loop directly (25 lines) and compiles the whole batch into one executable: 2.19 s.
Warp per-thread kernel. an explicit per-thread loop in a Python-embedded kernel language (74 lines, its own stages): 218 ms.
eagle CPU (OpenMP). the same hawk kernel on the host’s cores, no GPU needed: 13.1 s.
Compile time
Its own pass, never part of any wall. Each measurement is a fresh process: imports and the GPU context run untimed, then the timer spans the first call to a ready kernel (trace + code generation + compile + load, or the cache lookup). Cold: every cache empty (HAWK_CACHE_DIR, CUPY_CACHE_DIR, CUDA_CACHE_PATH, WARP_CACHE_PATH and JAX’s persistent compilation cache, each a fresh directory); warm: a new process on the caches the first cold run populated. eagle GPU: eagle.deploy of the hawk attempt kernel and of its active-set variant (host and device builds each), and one eagle graph of each built and run on 64 samples (eagle.simulate deploys the first again, a cache hit). eagle CPU: the same eagle.deploy of the attempt kernel (host and device builds) and one host run. CuPy: one masked attempt on 64 samples (every elementwise kernel compiled on first use). JAX: lower + compile from abstract shapes at N. Warp: the kernel module (code generation, compile, load) and one launch. Medians of 3 cold and 3 warm processes.
arm |
compiled |
cold (median of 3) |
warm (median of 3) |
|---|---|---|---|
eagle graph, eagle graph + compaction, eagle graph + compaction + reorder, hawk + eagle persistent, eagle.simulate (eagle picks the launch mode) |
eagle.deploy of the attempt kernel and its active-set variant |
3.77 s |
572 ms |
eagle CPU (OpenMP) |
eagle.deploy of the attempt kernel, host run |
3.41 s |
422 ms |
CuPy masked |
CuPy elementwise kernels |
4.05 s |
1.64 s |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000 |
3.5 s |
422 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 10,000 |
3.51 s |
421 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 100,000 |
4.62 s |
412 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000,000 |
6.14 s |
412 ms |
Warp per-thread kernel |
Warp kernel module |
11.9 s |
44.6 ms |
No compile step: PyTorch masked, PyTorch + CUDA graph (PyTorch’s eager kernels ship prebuilt; the CUDA graph is recorded, not compiled).
Memory method
A separate pass after the timing pass. Each (arm, N) runs in a fresh process: imports, the CUDA context created by one tiny device operation, then the baseline read, before the arm builds or allocates anything; then the arm’s kernels built and a warm-up run on a 64-sample batch (every compile happens here; JAX also compiles for shape N), so every library’s pool reserved during the warm-up counts in the reading; CuPy’s pools and PyTorch’s cache trimmed after the warm-up; then the arm built for N (eagle’s graph, PyTorch’s captured graph) and one full run with upload and download; then the reading again, the allocators still holding their high-water reservation (the same measure for every arm, no sampling needed). Device: the memory NVML accounts to the process on the timed GPU, after minus baseline. Caveats: each allocator’s rounding and growth policy count (JAX’s allocator grows in regions that can exceed the request; PyTorch rounds blocks up and a captured CUDA graph keeps a private pool; Warp’s stream-ordered pool reserves in chunks of tens of MiB, its release threshold set to 0 after the warm-up); the driver accounts in pages of about 2 MiB, so small-N rows read 0 or one page. Host: the peak resident set (VmHWM, reset to the current RSS at the host baseline, taken after the warm-up) minus the baseline RSS. Minimum: the planes the workload needs, N x (8 float64 planes (state, time, step) x 8 bytes + 2 counters x 4 bytes) = N x 72 bytes.
Fastest arm by median wall time:
N = 1,000: eagle.simulate (eagle picks the launch mode) (3.98 ms), then eagle graph (4.96 ms, 1.25× the time)
N = 10,000: eagle.simulate (eagle picks the launch mode) (5.7 ms), then hawk + eagle persistent (5.7 ms, 1× the time)
N = 100,000: eagle.simulate (eagle picks the launch mode) (14.6 ms), then hawk + eagle persistent (14.8 ms, 1.02× the time)
N = 1,000,000: eagle.simulate (eagle picks the launch mode) (107 ms), then hawk + eagle persistent (107 ms, 1× the time)
eagle_simulate against its hand-built twin eagle_graph:
N = 1,000: wall 0.803× the twin’s time
N = 10,000: wall 0.805× the twin’s time
N = 100,000: wall 0.631× the twin’s time
N = 1,000,000: wall 0.435× the twin’s time
Run on: driver 580.178.04; CUDA runtime 12.9; host CPU Intel Xeon CPU @ 2.00GHz; compiler x86_64-conda-linux-gnu-g++ (conda-forge gcc 14.4.0-7) 14.4.0; Python 3.12.14; Warp 1.17.0, JAX 0.11.2, PyTorch 2.12.1+cu129, CuPy 14.2.0, NumPy 2.5.3.
N = 1,000,000
Log scale, from 2.66 s to 152 s.
Table of every arm and every N
Device: Tesla T4 (compute capability 7.5, 40 SMs). Peak FP64: 254 G FLOP/s. Medians of 5 interleaved repetitions (3 for an arm whose warm-up says 5 would take over 60 s); IQR in parentheses.
Workload: two-body Kepler orbits (mu = 1, float64) from periapsis, a = 1, e uniform in [0, 0.9], random orientation; adaptive RKF7(8) (Fehlberg’s tableau, a mixed absolute/relative max-norm error control and a standard step-size controller). Tolerance rtol = atol = 1e-10; final time 64 (10.2 periods). 1106 FLOP per attempt.
arm |
loop structure |
lines of code |
|---|---|---|
eagle graph |
one CUDA graph: device WHILE loop of the attempt kernel, which marks and counts its finished samples and runs the attempts per launch eagle’s policy picks; finished samples skipped by the Terminated guard |
22 |
eagle graph + compaction |
as eagle graph, the attempt launched over the live samples only (index map rebuilt every 16 attempts) |
23 |
eagle graph + compaction + reorder |
as + compaction, live samples moved to the front of every plane when spread thin; restored at the end |
23 |
hawk + eagle persistent |
ONE launch, no graph, no map, no policy kernel: a grid sized from the SM count and the kernel’s occupancy, each lane fetching base + atomicAdd(counter, 1) whenever its sample is done |
22 |
eagle.simulate (eagle picks the launch mode) |
eagle.simulate takes the kernel, its state and its parameters by name; eagle’s automatic policy picks the launch mode from the batch (one launch for a small batch, the persistent launch above it); the row records the mode taken |
21 |
CuPy masked |
host loop: one masked whole-batch attempt (array expressions), host check per attempt |
25 |
PyTorch masked |
host loop: one masked whole-batch attempt (tensor expressions), host check per attempt |
27 |
PyTorch + CUDA graph |
host loop: a CUDA graph of 16 masked attempts replayed, host check per block |
53 |
JAX jit(vmap(while_loop)) |
one compiled executable: per-sample while loop, batched (runs until the slowest sample finishes) |
25 |
Warp per-thread kernel |
one launch: each thread loops over its own sample’s attempts to the final time |
74 |
eagle CPU (OpenMP) |
host loop: one attempt over the batch (OpenMP), the kernel marks and counts its finished samples |
22 |
Lines of code: non-blank, non-comment lines between the # >>> code:NAME and # <<< code:NAME markers of the script; every arm but Warp also uses the shared scheme block (the RKF7(8) stages, the right-hand side, the error norm), counted once in shared_scheme_lines (25 lines).
N |
arm |
wall (IQR) |
kernel-only |
end-to-end |
attempts·samples/s |
useful FLOP/s (% peak) |
device memory (× minimum) |
host memory |
|---|---|---|---|---|---|---|---|---|
1,000 |
eagle graph |
17.2 ms (118 µs) |
– |
17.7 ms |
29.4 M |
32.5 G (13%) |
0 MiB (0× min.) |
2.71 MiB |
1,000 |
eagle graph + compaction |
20.6 ms (157 µs) |
– |
21 ms |
24.5 M |
27.1 G (11%) |
0 MiB (0× min.) |
2.64 MiB |
1,000 |
eagle graph + compaction + reorder |
20.9 ms (65.4 µs) |
– |
21.4 ms |
24.1 M |
26.7 G (10%) |
0 MiB (0× min.) |
2.38 MiB |
1,000 |
hawk + eagle persistent |
55.8 ms (26 ms) |
– |
56.1 ms |
9.05 M |
10 G (3.9%) |
0 MiB (0× min.) |
2.21 MiB |
1,000 |
eagle.simulate (eagle picks the launch mode) |
16.3 ms (25.5 µs) |
– |
16.6 ms |
31.1 M |
34.4 G (14%) |
0 MiB (0× min.) |
2.25 MiB |
1,000 |
CuPy masked |
14 s (298 ms) |
– |
14 s |
36.2 k |
40 M (0.016%) |
0 MiB (0× min.) |
0.00391 MiB |
1,000 |
PyTorch masked |
9.23 s (120 ms) |
– |
9.23 s |
54.7 k |
60.5 M (0.024%) |
14 MiB (204× min.) |
0.316 MiB |
1,000 |
PyTorch + CUDA graph |
2.71 s (20 ms) |
– |
2.71 s |
186 k |
206 M (0.081%) |
160 MiB (2.33e+03× min.) |
14.6 MiB |
1,000 |
JAX jit(vmap(while_loop)) |
100 ms (53 µs) |
– |
103 ms |
5.05 M |
5.59 G (2.2%) |
12 MiB (175× min.) |
0.0586 MiB |
1,000 |
Warp per-thread kernel |
81.8 ms (254 µs) |
– |
82.8 ms |
6.18 M |
6.83 G (2.7%) |
30 MiB (437× min.) |
0.125 MiB |
1,000 |
eagle CPU (OpenMP) |
44.4 ms (5.32 ms) |
– |
44.6 ms |
11.4 M |
12.6 G (–) |
– |
0.406 MiB |
10,000 |
eagle graph |
78.5 ms (96.2 µs) |
– |
79.6 ms |
63.2 M |
69.9 G (27%) |
0 MiB (0× min.) |
4.31 MiB |
10,000 |
eagle graph + compaction |
60.5 ms (19.8 µs) |
– |
61.4 ms |
81.9 M |
90.6 G (36%) |
0 MiB (0× min.) |
4.14 MiB |
10,000 |
eagle graph + compaction + reorder |
61.3 ms (72.3 µs) |
– |
62.2 ms |
80.8 M |
89.4 G (35%) |
0 MiB (0× min.) |
4.94 MiB |
10,000 |
hawk + eagle persistent |
76.1 ms (170 µs) |
– |
76.8 ms |
65.2 M |
72.1 G (28%) |
0 MiB (0× min.) |
4.2 MiB |
10,000 |
eagle.simulate (eagle picks the launch mode) |
76 ms (21.2 µs) |
– |
76.8 ms |
65.2 M |
72.1 G (28%) |
0 MiB (0× min.) |
3.29 MiB |
10,000 |
CuPy masked |
14.1 s (61.2 ms) |
– |
14.1 s |
352 k |
389 M (0.15%) |
0 MiB (0× min.) |
7.78 MiB |
10,000 |
PyTorch masked |
9.34 s (52.1 ms) |
– |
9.34 s |
531 k |
587 M (0.23%) |
14 MiB (20.4× min.) |
1.97 MiB |
10,000 |
PyTorch + CUDA graph |
2.78 s (9.26 ms) |
– |
2.78 s |
1.79 M |
1.98 G (0.78%) |
168 MiB (245× min.) |
16.3 MiB |
10,000 |
JAX jit(vmap(while_loop)) |
148 ms (560 µs) |
– |
151 ms |
33.6 M |
37.1 G (15%) |
12 MiB (17.5× min.) |
0.109 MiB |
10,000 |
Warp per-thread kernel |
82.8 ms (66.2 µs) |
– |
84.2 ms |
59.9 M |
66.2 G (26%) |
30 MiB (43.7× min.) |
1.39 MiB |
10,000 |
eagle CPU (OpenMP) |
342 ms (622 µs) |
– |
342 ms |
14.5 M |
16 G (–) |
– |
2.73 MiB |
100,000 |
eagle graph |
900 ms (6.3 ms) |
– |
905 ms |
55.3 M |
61.2 G (24%) |
8 MiB (1.17× min.) |
25.3 MiB |
100,000 |
eagle graph + compaction |
472 ms (2.19 ms) |
– |
478 ms |
106 M |
117 G (46%) |
10 MiB (1.46× min.) |
25.2 MiB |
100,000 |
eagle graph + compaction + reorder |
462 ms (2.2 ms) |
– |
469 ms |
108 M |
119 G (47%) |
12 MiB (1.75× min.) |
26.3 MiB |
100,000 |
hawk + eagle persistent |
484 ms (1.75 ms) |
– |
490 ms |
103 M |
114 G (45%) |
8 MiB (1.17× min.) |
25.3 MiB |
100,000 |
eagle.simulate (eagle picks the launch mode) |
471 ms (971 µs) |
– |
476 ms |
106 M |
117 G (46%) |
8 MiB (1.17× min.) |
24.5 MiB |
100,000 |
CuPy masked |
14 s (20.1 ms) |
– |
14 s |
3.55 M |
3.93 G (1.5%) |
8 MiB (1.17× min.) |
22.9 MiB |
100,000 |
PyTorch masked |
9.78 s (87.5 ms) |
– |
9.79 s |
5.09 M |
5.63 G (2.2%) |
22 MiB (3.2× min.) |
19.8 MiB |
100,000 |
PyTorch + CUDA graph |
7.78 s (9.92 ms) |
– |
7.79 s |
6.4 M |
7.08 G (2.8%) |
272 MiB (39.6× min.) |
33.7 MiB |
100,000 |
JAX jit(vmap(while_loop)) |
1.37 s (2.84 ms) |
– |
1.38 s |
36.4 M |
40.2 G (16%) |
44 MiB (6.41× min.) |
27.9 MiB |
100,000 |
Warp per-thread kernel |
809 ms (3.58 ms) |
– |
815 ms |
61.6 M |
68.1 G (27%) |
30 MiB (4.37× min.) |
18.6 MiB |
100,000 |
eagle CPU (OpenMP) |
3.33 s (70.9 ms) |
– |
3.33 s |
15 M |
16.6 G (–) |
– |
26 MiB |
1,000,000 |
eagle graph |
7.39 s (185 ms) |
– |
7.45 s |
67.4 M |
74.6 G (29%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
eagle graph + compaction |
4.41 s (79.3 ms) |
– |
4.47 s |
113 M |
125 G (49%) |
86 MiB (1.25× min.) |
231 MiB |
1,000,000 |
eagle graph + compaction + reorder |
4.3 s (27 ms) |
– |
4.36 s |
116 M |
128 G (50%) |
114 MiB (1.66× min.) |
231 MiB |
1,000,000 |
hawk + eagle persistent |
4.26 s (1.9 ms) |
– |
4.33 s |
117 M |
129 G (51%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
4.26 s (1.47 ms) |
– |
4.33 s |
117 M |
129 G (51%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
CuPy masked |
89.9 s (115 ms) |
– |
90 s |
5.54 M |
6.13 G (2.4%) |
78 MiB (1.14× min.) |
183 MiB |
1,000,000 |
PyTorch masked |
95.1 s (47.3 ms) |
– |
95.2 s |
5.24 M |
5.79 G (2.3%) |
112 MiB (1.63× min.) |
191 MiB |
1,000,000 |
PyTorch + CUDA graph |
95.4 s (14.5 ms) |
– |
95.4 s |
5.23 M |
5.78 G (2.3%) |
1.32e+03 MiB (19.3× min.) |
241 MiB |
1,000,000 |
JAX jit(vmap(while_loop)) |
13.1 s (26 ms) |
– |
13.2 s |
38.2 M |
42.2 G (17%) |
268 MiB (3.9× min.) |
309 MiB |
1,000,000 |
Warp per-thread kernel |
7.88 s (1.31 ms) |
– |
7.93 s |
63.2 M |
70 G (27%) |
94 MiB (1.37× min.) |
183 MiB |
1,000,000 |
eagle CPU (OpenMP) |
36.1 s (1.65 s) |
– |
36.1 s |
13.8 M |
15.3 G (–) |
– |
253 MiB |
Agreement with the analytic orbit
Every arm, every sample: final time reached exactly; position error against the analytic Kepler state (Newton on Kepler’s equation to machine precision) below the stated bound n_acc tol (1 + v_p)(1 + 6 pi n_orbits v_p), derived from the controller’s per-step budget and the Kepler along-track amplification (see truth_bound; not fitted). Step counts are compared with the reference arm (eagle graph): any difference comes from a floating-point difference flipping one accept/reject decision (contraction into fused multiply-adds, pow and sqrt implementations differ between compilers and libraries); the table reports how many samples differ and by how much.
N |
arm |
max position error |
max error / bound |
samples with different step counts |
max Δ accepted / rejected |
|---|---|---|---|---|---|
1,000,000 |
eagle graph |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction + reorder |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
hawk + eagle persistent |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
CuPy masked |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
PyTorch masked |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
PyTorch + CUDA graph |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
Warp per-thread kernel |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
eagle CPU (OpenMP) |
2.26e-07 |
6.64e-03 |
2 |
1 / 1 |
1,000 |
eagle graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction + reorder |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
hawk + eagle persistent |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle.simulate (eagle picks the launch mode) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
CuPy masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch + CUDA graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
JAX jit(vmap(while_loop)) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
Warp per-thread kernel |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle CPU (OpenMP) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
10,000 |
eagle graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction + reorder |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
hawk + eagle persistent |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle.simulate (eagle picks the launch mode) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
CuPy masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch + CUDA graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
JAX jit(vmap(while_loop)) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
Warp per-thread kernel |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle CPU (OpenMP) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
100,000 |
eagle graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction + reorder |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
hawk + eagle persistent |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle.simulate (eagle picks the launch mode) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
CuPy masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch + CUDA graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
JAX jit(vmap(while_loop)) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
Warp per-thread kernel |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle CPU (OpenMP) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
Steps per sample at N = 1,000,000 (reference arm): accepted min 322 / median 371 / p90 571 / max 693; rejected min 1 / median 59 / max 264.
Where each tool fits
Each tool here solves the same 13-stage adaptive problem to the same tolerance and agrees with the analytic orbit; they differ in how the per-sample control flow is written and run. Numbers below are at N = 1,000,000; the fastest arm there is eagle.simulate (eagle picks the launch mode) (4.26 s).
eagle graph. the attempt is per-sample scalar code (22 lines with its launch; the kernel finishes its own samples), one kernel per attempt looped on the device: 7.39 s.
eagle graph + compaction. the same kernel launched over the live samples only: 4.41 s; the plain graph takes 1.68× as long.
eagle graph + compaction + reorder. adds the occasional reorder of the planes: 4.3 s; the compacted graph takes 1.03× as long.
hawk + eagle persistent. the plain (not active-set) attempt kernel, forced through hawk’s persist entry: one launch, no graph, no map, no policy kernel – a grid of SMs blocks of 256 lanes steals the next unfinished sample off a counter when their own finishes: 4.26 s; the eagle graph takes 1.73× as long.
eagle.simulate (eagle picks the launch mode). the same kernel through eagle.simulate, its state and parameters by name (21 lines with the kernel): 4.26 s; the eagle graph takes 1.73× as long.
CuPy masked. array code in the NumPy style (25 lines), the control flow as masks over the batch: 89.9 s.
PyTorch masked. the same array formulation in PyTorch (27 lines), running where a PyTorch model already runs: 95.1 s.
PyTorch + CUDA graph. records 16 attempts as a CUDA graph and replays them, removing per-kernel launch cost: 95.4 s; eager PyTorch takes 0.998× as long.
JAX jit(vmap(while_loop)). writes the per-sample loop directly (25 lines) and compiles the whole batch into one executable: 13.1 s.
Warp per-thread kernel. an explicit per-thread loop in a Python-embedded kernel language (74 lines, its own stages): 7.88 s.
eagle CPU (OpenMP). the same hawk kernel on the host’s cores, no GPU needed: 36.1 s.
Compile time
Its own pass, never part of any wall. Each measurement is a fresh process: imports and the GPU context run untimed, then the timer spans the first call to a ready kernel (trace + code generation + compile + load, or the cache lookup). Cold: every cache empty (HAWK_CACHE_DIR, CUPY_CACHE_DIR, CUDA_CACHE_PATH, WARP_CACHE_PATH and JAX’s persistent compilation cache, each a fresh directory); warm: a new process on the caches the first cold run populated. eagle GPU: eagle.deploy of the hawk attempt kernel and of its active-set variant (host and device builds each), and one eagle graph of each built and run on 64 samples (eagle.simulate deploys the first again, a cache hit). eagle CPU: the same eagle.deploy of the attempt kernel (host and device builds) and one host run. CuPy: one masked attempt on 64 samples (every elementwise kernel compiled on first use). JAX: lower + compile from abstract shapes at N. Warp: the kernel module (code generation, compile, load) and one launch. Medians of 3 cold and 3 warm processes.
arm |
compiled |
cold (median of 3) |
warm (median of 3) |
|---|---|---|---|
eagle graph, eagle graph + compaction, eagle graph + compaction + reorder, hawk + eagle persistent, eagle.simulate (eagle picks the launch mode) |
eagle.deploy of the attempt kernel and its active-set variant |
4.85 s |
657 ms |
eagle CPU (OpenMP) |
eagle.deploy of the attempt kernel, host run |
3.85 s |
456 ms |
CuPy masked |
CuPy elementwise kernels |
4.43 s |
1.8 s |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000 |
4.6 s |
422 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 10,000 |
3.54 s |
408 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 100,000 |
4.84 s |
426 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000,000 |
4.82 s |
420 ms |
Warp per-thread kernel |
Warp kernel module |
12.7 s |
94 ms |
No compile step: PyTorch masked, PyTorch + CUDA graph (PyTorch’s eager kernels ship prebuilt; the CUDA graph is recorded, not compiled).
Memory method
A separate pass after the timing pass. Each (arm, N) runs in a fresh process: imports, the CUDA context created by one tiny device operation, then the baseline read, before the arm builds or allocates anything; then the arm’s kernels built and a warm-up run on a 64-sample batch (every compile happens here; JAX also compiles for shape N), so every library’s pool reserved during the warm-up counts in the reading; CuPy’s pools and PyTorch’s cache trimmed after the warm-up; then the arm built for N (eagle’s graph, PyTorch’s captured graph) and one full run with upload and download; then the reading again, the allocators still holding their high-water reservation (the same measure for every arm, no sampling needed). Device: the memory NVML accounts to the process on the timed GPU, after minus baseline. Caveats: each allocator’s rounding and growth policy count (JAX’s allocator grows in regions that can exceed the request; PyTorch rounds blocks up and a captured CUDA graph keeps a private pool; Warp’s stream-ordered pool reserves in chunks of tens of MiB, its release threshold set to 0 after the warm-up); the driver accounts in pages of about 2 MiB, so small-N rows read 0 or one page. Host: the peak resident set (VmHWM, reset to the current RSS at the host baseline, taken after the warm-up) minus the baseline RSS. Minimum: the planes the workload needs, N x (8 float64 planes (state, time, step) x 8 bytes + 2 counters x 4 bytes) = N x 72 bytes.
Fastest arm by median wall time:
N = 1,000: eagle.simulate (eagle picks the launch mode) (16.3 ms), then eagle graph (17.2 ms, 1.06× the time)
N = 10,000: eagle graph + compaction (60.5 ms), then eagle graph + compaction + reorder (61.3 ms, 1.01× the time)
N = 100,000: eagle graph + compaction + reorder (462 ms), then eagle.simulate (eagle picks the launch mode) (471 ms, 1.02× the time)
N = 1,000,000: eagle.simulate (eagle picks the launch mode) (4.26 s), then hawk + eagle persistent (4.26 s, 1× the time)
eagle_simulate against its hand-built twin eagle_graph:
N = 1,000: wall 0.946× the twin’s time
N = 10,000: wall 0.969× the twin’s time
N = 100,000: wall 0.523× the twin’s time
N = 1,000,000: wall 0.576× the twin’s time
Run on: driver 560.35.03; CUDA runtime 12.9; host CPU Intel Xeon W-2125 CPU @ 4.00GHz; compiler g++ (GCC) 11.5.0 20240719 (Red Hat 11.5.0-5.0.1); Python 3.12.14; Warp 1.17.0, JAX 0.11.2, PyTorch 2.12.1, CuPy 14.2.0, NumPy 2.5.3.
N = 1,000,000
Log scale, from 8.41 s to 294 s.
Table of every arm and every N
Device: Quadro P2000 (compute capability 6.1, 8 SMs). Peak FP64: 94.8 G FLOP/s. Medians of 5 interleaved repetitions (3 for an arm whose warm-up says 5 would take over 60 s); IQR in parentheses.
Workload: two-body Kepler orbits (mu = 1, float64) from periapsis, a = 1, e uniform in [0, 0.9], random orientation; adaptive RKF7(8) (Fehlberg’s tableau, a mixed absolute/relative max-norm error control and a standard step-size controller). Tolerance rtol = atol = 1e-10; final time 64 (10.2 periods). 1106 FLOP per attempt.
arm |
loop structure |
lines of code |
|---|---|---|
eagle graph |
one CUDA graph: device WHILE loop of the attempt kernel, which marks and counts its finished samples and runs the attempts per launch eagle’s policy picks; finished samples skipped by the Terminated guard |
22 |
eagle graph + compaction |
as eagle graph, the attempt launched over the live samples only (index map rebuilt every 16 attempts) |
23 |
eagle graph + compaction + reorder |
as + compaction, live samples moved to the front of every plane when spread thin; restored at the end |
23 |
hawk + eagle persistent |
ONE launch, no graph, no map, no policy kernel: a grid sized from the SM count and the kernel’s occupancy, each lane fetching base + atomicAdd(counter, 1) whenever its sample is done |
22 |
eagle.simulate (eagle picks the launch mode) |
eagle.simulate takes the kernel, its state and its parameters by name; eagle’s automatic policy picks the launch mode from the batch (one launch for a small batch, the persistent launch above it); the row records the mode taken |
21 |
CuPy masked |
host loop: one masked whole-batch attempt (array expressions), host check per attempt |
25 |
PyTorch masked |
host loop: one masked whole-batch attempt (tensor expressions), host check per attempt |
27 |
PyTorch + CUDA graph |
host loop: a CUDA graph of 16 masked attempts replayed, host check per block |
53 |
JAX jit(vmap(while_loop)) |
one compiled executable: per-sample while loop, batched (runs until the slowest sample finishes) |
25 |
Warp per-thread kernel |
one launch: each thread loops over its own sample’s attempts to the final time |
74 |
eagle CPU (OpenMP) |
host loop: one attempt over the batch (OpenMP), the kernel marks and counts its finished samples |
22 |
Lines of code: non-blank, non-comment lines between the # >>> code:NAME and # <<< code:NAME markers of the script; every arm but Warp also uses the shared scheme block (the RKF7(8) stages, the right-hand side, the error norm), counted once in shared_scheme_lines (25 lines).
N |
arm |
wall (IQR) |
kernel-only |
end-to-end |
attempts·samples/s |
useful FLOP/s (% peak) |
device memory (× minimum) |
host memory |
|---|---|---|---|---|---|---|---|---|
1,000 |
eagle graph |
30.9 ms (143 µs) |
30.1 ms |
31.3 ms |
16.3 M |
18 G (19%) |
0 MiB (0× min.) |
2.38 MiB |
1,000 |
eagle graph + compaction |
34.9 ms (127 µs) |
33.3 ms |
35.2 ms |
14.5 M |
16 G (17%) |
0 MiB (0× min.) |
2.25 MiB |
1,000 |
eagle graph + compaction + reorder |
35.3 ms (940 µs) |
34.7 ms |
35.6 ms |
14.3 M |
15.8 G (17%) |
0 MiB (0× min.) |
2.5 MiB |
1,000 |
hawk + eagle persistent |
49.3 ms (442 µs) |
49.4 ms |
49.5 ms |
10.3 M |
11.3 G (12%) |
0 MiB (0× min.) |
2.25 MiB |
1,000 |
eagle.simulate (eagle picks the launch mode) |
29.8 ms (93.7 µs) |
29.6 ms |
30.1 ms |
17 M |
18.8 G (20%) |
0 MiB (0× min.) |
2.5 MiB |
1,000 |
CuPy masked |
8.76 s (676 ms) |
1.9 s |
8.76 s |
57.6 k |
63.8 M (0.067%) |
0 MiB (0× min.) |
6.88 MiB |
1,000 |
PyTorch masked |
6.08 s (240 ms) |
2.75 s |
6.16 s |
83.1 k |
91.9 M (0.097%) |
4 MiB (58.3× min.) |
0.125 MiB |
1,000 |
PyTorch + CUDA graph |
2.96 s (1.86 ms) |
3.23 s |
2.96 s |
171 k |
189 M (0.2%) |
182 MiB (2.65e+03× min.) |
13.4 MiB |
1,000 |
JAX jit(vmap(while_loop)) |
107 ms (1.2 ms) |
91.6 ms |
109 ms |
4.72 M |
5.22 G (5.5%) |
12 MiB (175× min.) |
0 MiB |
1,000 |
Warp per-thread kernel |
55.5 ms (220 µs) |
55.4 ms |
56.4 ms |
9.11 M |
10.1 G (11%) |
30 MiB (437× min.) |
0.125 MiB |
1,000 |
eagle CPU (OpenMP) |
70 ms (13.1 ms) |
– |
70.1 ms |
7.21 M |
7.98 G (–) |
– |
0.375 MiB |
10,000 |
eagle graph |
248 ms (91.2 µs) |
247 ms |
249 ms |
20 M |
22.1 G (23%) |
0 MiB (0× min.) |
4.25 MiB |
10,000 |
eagle graph + compaction |
162 ms (39.9 µs) |
162 ms |
163 ms |
30.6 M |
33.8 G (36%) |
0 MiB (0× min.) |
4.25 MiB |
10,000 |
eagle graph + compaction + reorder |
163 ms (91.7 µs) |
162 ms |
164 ms |
30.4 M |
33.6 G (35%) |
0 MiB (0× min.) |
4.75 MiB |
10,000 |
hawk + eagle persistent |
179 ms (1.41 ms) |
178 ms |
179 ms |
27.7 M |
30.7 G (32%) |
0 MiB (0× min.) |
3.75 MiB |
10,000 |
eagle.simulate (eagle picks the launch mode) |
174 ms (3.25 ms) |
175 ms |
175 ms |
28.4 M |
31.4 G (33%) |
0 MiB (0× min.) |
3.88 MiB |
10,000 |
CuPy masked |
8.76 s (728 ms) |
2.85 s |
8.76 s |
566 k |
626 M (0.66%) |
0 MiB (0× min.) |
7.62 MiB |
10,000 |
PyTorch masked |
5.98 s (511 ms) |
3.32 s |
6.05 s |
829 k |
917 M (0.97%) |
4 MiB (5.83× min.) |
2.25 MiB |
10,000 |
PyTorch + CUDA graph |
3.59 s (2 ms) |
3.8 s |
3.59 s |
1.38 M |
1.53 G (1.6%) |
190 MiB (277× min.) |
14.4 MiB |
10,000 |
JAX jit(vmap(while_loop)) |
428 ms (345 µs) |
417 ms |
430 ms |
11.6 M |
12.8 G (14%) |
12 MiB (17.5× min.) |
0.25 MiB |
10,000 |
Warp per-thread kernel |
276 ms (1.4 ms) |
275 ms |
277 ms |
18 M |
19.9 G (21%) |
30 MiB (43.7× min.) |
2 MiB |
10,000 |
eagle CPU (OpenMP) |
831 ms (179 ms) |
– |
831 ms |
5.97 M |
6.6 G (–) |
– |
2.62 MiB |
100,000 |
eagle graph |
2.38 s (1.03 ms) |
2.4 s |
2.38 s |
20.9 M |
23.1 G (24%) |
8 MiB (1.17× min.) |
25 MiB |
100,000 |
eagle graph + compaction |
1.43 s (13.9 ms) |
1.45 s |
1.44 s |
34.8 M |
38.5 G (41%) |
10 MiB (1.46× min.) |
25 MiB |
100,000 |
eagle graph + compaction + reorder |
1.43 s (14.8 ms) |
1.45 s |
1.43 s |
34.9 M |
38.6 G (41%) |
12 MiB (1.75× min.) |
25.7 MiB |
100,000 |
hawk + eagle persistent |
1.38 s (15.1 ms) |
1.4 s |
1.39 s |
36 M |
39.8 G (42%) |
8 MiB (1.17× min.) |
24.7 MiB |
100,000 |
eagle.simulate (eagle picks the launch mode) |
1.38 s (16.4 ms) |
1.4 s |
1.39 s |
36 M |
39.8 G (42%) |
8 MiB (1.17× min.) |
24.9 MiB |
100,000 |
CuPy masked |
19.2 s (69 ms) |
19.3 s |
19.2 s |
2.6 M |
2.88 G (3%) |
8 MiB (1.17× min.) |
22.7 MiB |
100,000 |
PyTorch masked |
17.3 s (23.3 ms) |
17.1 s |
17.4 s |
2.88 M |
3.18 G (3.4%) |
12 MiB (1.75× min.) |
21.6 MiB |
100,000 |
PyTorch + CUDA graph |
17.2 s (20.6 ms) |
17.4 s |
17.2 s |
2.9 M |
3.2 G (3.4%) |
294 MiB (42.8× min.) |
30.6 MiB |
100,000 |
JAX jit(vmap(while_loop)) |
4.22 s (16 ms) |
4.26 s |
4.23 s |
11.8 M |
13.1 G (14%) |
44 MiB (6.41× min.) |
27 MiB |
100,000 |
Warp per-thread kernel |
2.64 s (2.89 ms) |
2.67 s |
2.65 s |
18.8 M |
20.8 G (22%) |
30 MiB (4.37× min.) |
19.3 MiB |
100,000 |
eagle CPU (OpenMP) |
3.18 s (2.37 s) |
– |
3.19 s |
15.6 M |
17.3 G (–) |
– |
25.5 MiB |
1,000,000 |
eagle graph |
23.7 s (7.59 ms) |
23.7 s |
23.8 s |
21.1 M |
23.3 G (25%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
eagle graph + compaction |
14.1 s (5.22 ms) |
14.2 s |
14.5 s |
35.2 M |
39 G (41%) |
86 MiB (1.25× min.) |
231 MiB |
1,000,000 |
eagle graph + compaction + reorder |
14.1 s (3.9 ms) |
14.4 s |
14.6 s |
35.3 M |
39.1 G (41%) |
114 MiB (1.66× min.) |
231 MiB |
1,000,000 |
hawk + eagle persistent |
13.5 s (4.62 ms) |
13.6 s |
13.7 s |
37 M |
41 G (43%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
13.4 s (2.95 ms) |
13.6 s |
13.7 s |
37.1 M |
41 G (43%) |
70 MiB (1.02× min.) |
231 MiB |
1,000,000 |
CuPy masked |
182 s (8.26 ms) |
183 s |
182 s |
2.73 M |
3.02 G (3.2%) |
78 MiB (1.14× min.) |
183 MiB |
1,000,000 |
PyTorch masked |
184 s (7.12 ms) |
183 s |
184 s |
2.71 M |
3 G (3.2%) |
102 MiB (1.49× min.) |
202 MiB |
1,000,000 |
PyTorch + CUDA graph |
184 s (2.5 ms) |
185 s |
184 s |
2.7 M |
2.99 G (3.2%) |
1.35e+03 MiB (19.6× min.) |
192 MiB |
1,000,000 |
JAX jit(vmap(while_loop)) |
40.1 s (5.46 ms) |
40.3 s |
40.3 s |
12.4 M |
13.7 G (14%) |
268 MiB (3.9× min.) |
310 MiB |
1,000,000 |
Warp per-thread kernel |
26.3 s (2.46 ms) |
26.4 s |
26.4 s |
18.9 M |
20.9 G (22%) |
94 MiB (1.37× min.) |
191 MiB |
1,000,000 |
eagle CPU (OpenMP) |
38.9 s (10.2 s) |
– |
39.4 s |
12.8 M |
14.2 G (–) |
– |
252 MiB |
Agreement with the analytic orbit
Every arm, every sample: final time reached exactly; position error against the analytic Kepler state (Newton on Kepler’s equation to machine precision) below the stated bound n_acc tol (1 + v_p)(1 + 6 pi n_orbits v_p), derived from the controller’s per-step budget and the Kepler along-track amplification (see truth_bound; not fitted). Step counts are compared with the reference arm (eagle graph): any difference comes from a floating-point difference flipping one accept/reject decision (contraction into fused multiply-adds, pow and sqrt implementations differ between compilers and libraries); the table reports how many samples differ and by how much.
N |
arm |
max position error |
max error / bound |
samples with different step counts |
max Δ accepted / rejected |
|---|---|---|---|---|---|
1,000,000 |
eagle graph |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction + reorder |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
hawk + eagle persistent |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
CuPy masked |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
PyTorch masked |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
PyTorch + CUDA graph |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
Warp per-thread kernel |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
eagle CPU (OpenMP) |
2.26e-07 |
6.64e-03 |
1 |
1 / 0 |
1,000 |
eagle graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction + reorder |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
hawk + eagle persistent |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle.simulate (eagle picks the launch mode) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
CuPy masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch + CUDA graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
JAX jit(vmap(while_loop)) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
Warp per-thread kernel |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle CPU (OpenMP) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
10,000 |
eagle graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction + reorder |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
hawk + eagle persistent |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle.simulate (eagle picks the launch mode) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
CuPy masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch + CUDA graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
JAX jit(vmap(while_loop)) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
Warp per-thread kernel |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle CPU (OpenMP) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
100,000 |
eagle graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction + reorder |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
hawk + eagle persistent |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle.simulate (eagle picks the launch mode) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
CuPy masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch + CUDA graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
JAX jit(vmap(while_loop)) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
Warp per-thread kernel |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle CPU (OpenMP) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
Steps per sample at N = 1,000,000 (reference arm): accepted min 322 / median 371 / p90 571 / max 693; rejected min 1 / median 59 / max 264.
Where each tool fits
Each tool here solves the same 13-stage adaptive problem to the same tolerance and agrees with the analytic orbit; they differ in how the per-sample control flow is written and run. Numbers below are at N = 1,000,000; the fastest arm there is eagle.simulate (eagle picks the launch mode) (13.4 s).
eagle graph. the attempt is per-sample scalar code (22 lines with its launch; the kernel finishes its own samples), one kernel per attempt looped on the device: 23.7 s.
eagle graph + compaction. the same kernel launched over the live samples only: 14.1 s; the plain graph takes 1.67× as long.
eagle graph + compaction + reorder. adds the occasional reorder of the planes: 14.1 s; the compacted graph takes 1× as long.
hawk + eagle persistent. the plain (not active-set) attempt kernel, forced through hawk’s persist entry: one launch, no graph, no map, no policy kernel – a grid of SMs blocks of 256 lanes steals the next unfinished sample off a counter when their own finishes: 13.5 s; the eagle graph takes 1.76× as long.
eagle.simulate (eagle picks the launch mode). the same kernel through eagle.simulate, its state and parameters by name (21 lines with the kernel): 13.4 s; the eagle graph takes 1.76× as long.
CuPy masked. array code in the NumPy style (25 lines), the control flow as masks over the batch: 182 s.
PyTorch masked. the same array formulation in PyTorch (27 lines), running where a PyTorch model already runs: 184 s.
PyTorch + CUDA graph. records 16 attempts as a CUDA graph and replays them, removing per-kernel launch cost: 184 s; eager PyTorch takes 0.996× as long.
JAX jit(vmap(while_loop)). writes the per-sample loop directly (25 lines) and compiles the whole batch into one executable: 40.1 s.
Warp per-thread kernel. an explicit per-thread loop in a Python-embedded kernel language (74 lines, its own stages): 26.3 s.
eagle CPU (OpenMP). the same hawk kernel on the host’s cores, no GPU needed: 38.9 s.
Compile time
Its own pass, never part of any wall. Each measurement is a fresh process: imports and the GPU context run untimed, then the timer spans the first call to a ready kernel (trace + code generation + compile + load, or the cache lookup). Cold: every cache empty (HAWK_CACHE_DIR, CUPY_CACHE_DIR, CUDA_CACHE_PATH, WARP_CACHE_PATH and JAX’s persistent compilation cache, each a fresh directory); warm: a new process on the caches the first cold run populated. eagle GPU: eagle.deploy of the hawk attempt kernel and of its active-set variant (host and device builds each), and one eagle graph of each built and run on 64 samples (eagle.simulate deploys the first again, a cache hit). eagle CPU: the same eagle.deploy of the attempt kernel (host and device builds) and one host run. CuPy: one masked attempt on 64 samples (every elementwise kernel compiled on first use). JAX: lower + compile from abstract shapes at N. Warp: the kernel module (code generation, compile, load) and one launch. Medians of 3 cold and 3 warm processes.
arm |
compiled |
cold (median of 3) |
warm (median of 3) |
|---|---|---|---|
eagle graph, eagle graph + compaction, eagle graph + compaction + reorder, hawk + eagle persistent, eagle.simulate (eagle picks the launch mode) |
eagle.deploy of the attempt kernel and its active-set variant |
4.48 s |
621 ms |
eagle CPU (OpenMP) |
eagle.deploy of the attempt kernel, host run |
3.07 s |
499 ms |
CuPy masked |
CuPy elementwise kernels |
4.04 s |
1.71 s |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000,000 |
4.26 s |
502 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000 |
4.12 s |
545 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 10,000 |
3.93 s |
482 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 100,000 |
4.75 s |
403 ms |
Warp per-thread kernel |
Warp kernel module |
9.35 s |
60.4 ms |
No compile step: PyTorch masked, PyTorch + CUDA graph (PyTorch’s eager kernels ship prebuilt; the CUDA graph is recorded, not compiled).
Memory method
A separate pass after the timing pass. Each (arm, N) runs in a fresh process: imports, the CUDA context created by one tiny device operation, then the baseline read, before the arm builds or allocates anything; then the arm’s kernels built and a warm-up run on a 64-sample batch (every compile happens here; JAX also compiles for shape N), so every library’s pool reserved during the warm-up counts in the reading; CuPy’s pools and PyTorch’s cache trimmed after the warm-up; then the arm built for N (eagle’s graph, PyTorch’s captured graph) and one full run with upload and download; then the reading again, the allocators still holding their high-water reservation (the same measure for every arm, no sampling needed). Device: the memory NVML accounts to the process on the timed GPU, after minus baseline. Caveats: each allocator’s rounding and growth policy count (JAX’s allocator grows in regions that can exceed the request; PyTorch rounds blocks up and a captured CUDA graph keeps a private pool; Warp’s stream-ordered pool reserves in chunks of tens of MiB, its release threshold set to 0 after the warm-up); the driver accounts in pages of about 2 MiB, so small-N rows read 0 or one page. Host: the peak resident set (VmHWM, reset to the current RSS at the host baseline, taken after the warm-up) minus the baseline RSS. Minimum: the planes the workload needs, N x (8 float64 planes (state, time, step) x 8 bytes + 2 counters x 4 bytes) = N x 72 bytes.
torch.compile is not run: its GPU code generator (Inductor, through Triton) needs compute capability 7.0 or newer, and this GPU is 6.1. The PyTorch + CUDA graph arm removes the same per-kernel launch cost without generating code.
Fastest arm by median wall time:
N = 1,000: eagle.simulate (eagle picks the launch mode) (29.8 ms), then eagle graph (30.9 ms, 1.04× the time)
N = 10,000: eagle graph + compaction (162 ms), then eagle graph + compaction + reorder (163 ms, 1.01× the time)
N = 100,000: eagle.simulate (eagle picks the launch mode) (1.38 s), then hawk + eagle persistent (1.38 s, 1× the time)
N = 1,000,000: eagle.simulate (eagle picks the launch mode) (13.4 s), then hawk + eagle persistent (13.5 s, 1× the time)
eagle_simulate against its hand-built twin eagle_graph:
N = 1,000: wall 0.963×, kernel-only 0.985× the twin’s time
N = 10,000: wall 0.703×, kernel-only 0.71× the twin’s time
N = 100,000: wall 0.581×, kernel-only 0.581× the twin’s time
N = 1,000,000: wall 0.568×, kernel-only 0.575× the twin’s time
Run on: driver 580.178.04; CUDA runtime 12.9; host CPU AMD EPYC 9B45; compiler x86_64-conda-linux-gnu-g++ (conda-forge gcc 14.4.0-7) 14.4.0; Python 3.12.14; Warp 1.17.0, JAX 0.11.2, PyTorch 2.12.1+cu129, CuPy 14.2.0, NumPy 2.5.3.
N = 1,000,000
Log scale, from 392 ms to 14.9 s.
Table of every arm and every N
Device: NVIDIA RTX PRO 6000 Blackwell Server Edition (compute capability 12.0, 188 SMs). Peak FP64: – FLOP/s. Medians of 5 interleaved repetitions (3 for an arm whose warm-up says 5 would take over 60 s); IQR in parentheses.
Workload: two-body Kepler orbits (mu = 1, float64) from periapsis, a = 1, e uniform in [0, 0.9], random orientation; adaptive RKF7(8) (Fehlberg’s tableau, a mixed absolute/relative max-norm error control and a standard step-size controller). Tolerance rtol = atol = 1e-10; final time 64 (10.2 periods). 1106 FLOP per attempt.
arm |
loop structure |
lines of code |
|---|---|---|
eagle graph |
one CUDA graph: device WHILE loop of the attempt kernel, which marks and counts its finished samples and runs the attempts per launch eagle’s policy picks; finished samples skipped by the Terminated guard |
22 |
eagle graph + compaction |
as eagle graph, the attempt launched over the live samples only (index map rebuilt every 16 attempts) |
23 |
eagle graph + compaction + reorder |
as + compaction, live samples moved to the front of every plane when spread thin; restored at the end |
23 |
hawk + eagle persistent |
ONE launch, no graph, no map, no policy kernel: a grid sized from the SM count and the kernel’s occupancy, each lane fetching base + atomicAdd(counter, 1) whenever its sample is done |
22 |
eagle.simulate (eagle picks the launch mode) |
eagle.simulate takes the kernel, its state and its parameters by name; eagle’s automatic policy picks the launch mode from the batch (one launch for a small batch, the persistent launch above it); the row records the mode taken |
21 |
CuPy masked |
host loop: one masked whole-batch attempt (array expressions), host check per attempt |
25 |
PyTorch masked |
host loop: one masked whole-batch attempt (tensor expressions), host check per attempt |
27 |
PyTorch + CUDA graph |
host loop: a CUDA graph of 16 masked attempts replayed, host check per block |
53 |
JAX jit(vmap(while_loop)) |
one compiled executable: per-sample while loop, batched (runs until the slowest sample finishes) |
25 |
Warp per-thread kernel |
one launch: each thread loops over its own sample’s attempts to the final time |
74 |
eagle CPU (OpenMP) |
host loop: one attempt over the batch (OpenMP), the kernel marks and counts its finished samples |
22 |
Lines of code: non-blank, non-comment lines between the # >>> code:NAME and # <<< code:NAME markers of the script; every arm but Warp also uses the shared scheme block (the RKF7(8) stages, the right-hand side, the error norm), counted once in shared_scheme_lines (25 lines).
N |
arm |
wall (IQR) |
kernel-only |
end-to-end |
attempts·samples/s |
useful FLOP/s (% peak) |
device memory (× minimum) |
host memory |
|---|---|---|---|---|---|---|---|---|
1,000 |
eagle graph |
24.7 ms (57.2 µs) |
– |
24.9 ms |
20.5 M |
22.6 G (–) |
2 MiB (29.1× min.) |
2.64 MiB |
1,000 |
eagle graph + compaction |
26.8 ms (82.4 µs) |
– |
27 ms |
18.8 M |
20.8 G (–) |
2 MiB (29.1× min.) |
2.52 MiB |
1,000 |
eagle graph + compaction + reorder |
26.9 ms (87.7 µs) |
– |
27.1 ms |
18.8 M |
20.8 G (–) |
2 MiB (29.1× min.) |
2.64 MiB |
1,000 |
hawk + eagle persistent |
49.3 ms (854 µs) |
– |
49.4 ms |
10.2 M |
11.3 G (–) |
0 MiB (0× min.) |
2.2 MiB |
1,000 |
eagle.simulate (eagle picks the launch mode) |
24.2 ms (0.801 µs) |
– |
24.3 ms |
20.9 M |
23.1 G (–) |
0 MiB (0× min.) |
2.19 MiB |
1,000 |
CuPy masked |
5.28 s (26 ms) |
– |
5.28 s |
95.6 k |
106 M (–) |
0 MiB (0× min.) |
0.00781 MiB |
1,000 |
PyTorch masked |
3.92 s (16.3 ms) |
– |
3.93 s |
129 k |
142 M (–) |
88 MiB (1.28e+03× min.) |
0.199 MiB |
1,000 |
PyTorch + CUDA graph |
1.6 s (18 ms) |
– |
1.6 s |
316 k |
350 M (–) |
204 MiB (2.97e+03× min.) |
22.9 MiB |
1,000 |
JAX jit(vmap(while_loop)) |
104 ms (94 µs) |
– |
105 ms |
4.88 M |
5.39 G (–) |
160 MiB (2.33e+03× min.) |
0.102 MiB |
1,000 |
Warp per-thread kernel |
54.1 ms (31.1 µs) |
– |
54.6 ms |
9.33 M |
10.3 G (–) |
30 MiB (437× min.) |
0 MiB |
1,000 |
eagle CPU (OpenMP) |
3.45 ms (7.53 µs) |
– |
3.5 ms |
146 M |
162 G (–) |
– |
0.762 MiB |
10,000 |
eagle graph |
25.8 ms (29.8 µs) |
– |
26.1 ms |
192 M |
212 G (–) |
2 MiB (2.91× min.) |
4.32 MiB |
10,000 |
eagle graph + compaction |
28.1 ms (36.5 µs) |
– |
28.3 ms |
177 M |
195 G (–) |
2 MiB (2.91× min.) |
4.51 MiB |
10,000 |
eagle graph + compaction + reorder |
28.6 ms (16.8 µs) |
– |
28.8 ms |
173 M |
192 G (–) |
2 MiB (2.91× min.) |
5.19 MiB |
10,000 |
hawk + eagle persistent |
52 ms (450 µs) |
– |
52.1 ms |
95.4 M |
106 G (–) |
0 MiB (0× min.) |
4.11 MiB |
10,000 |
eagle.simulate (eagle picks the launch mode) |
25.2 ms (2.77 µs) |
– |
25.4 ms |
196 M |
217 G (–) |
0 MiB (0× min.) |
4.02 MiB |
10,000 |
CuPy masked |
5.3 s (14.9 ms) |
– |
5.3 s |
935 k |
1.03 G (–) |
0 MiB (0× min.) |
8.25 MiB |
10,000 |
PyTorch masked |
3.97 s (15.7 ms) |
– |
3.97 s |
1.25 M |
1.38 G (–) |
88 MiB (128× min.) |
1.91 MiB |
10,000 |
PyTorch + CUDA graph |
1.7 s (139 µs) |
– |
1.7 s |
2.92 M |
3.23 G (–) |
212 MiB (309× min.) |
24.9 MiB |
10,000 |
JAX jit(vmap(while_loop)) |
149 ms (348 µs) |
– |
151 ms |
33.2 M |
36.7 G (–) |
160 MiB (233× min.) |
0.0742 MiB |
10,000 |
Warp per-thread kernel |
57 ms (31.8 µs) |
– |
57.6 ms |
86.9 M |
96.2 G (–) |
30 MiB (43.7× min.) |
1.3 MiB |
10,000 |
eagle CPU (OpenMP) |
12.7 ms (50.6 µs) |
– |
12.7 ms |
392 M |
433 G (–) |
– |
2.75 MiB |
100,000 |
eagle graph |
151 ms (28.5 µs) |
– |
152 ms |
330 M |
365 G (–) |
10 MiB (1.46× min.) |
25.4 MiB |
100,000 |
eagle graph + compaction |
102 ms (110 µs) |
– |
103 ms |
488 M |
539 G (–) |
12 MiB (1.75× min.) |
25.7 MiB |
100,000 |
eagle graph + compaction + reorder |
103 ms (29.8 µs) |
– |
104 ms |
485 M |
536 G (–) |
14 MiB (2.04× min.) |
26.6 MiB |
100,000 |
hawk + eagle persistent |
91.9 ms (109 µs) |
– |
92.8 ms |
542 M |
599 G (–) |
8 MiB (1.17× min.) |
24.9 MiB |
100,000 |
eagle.simulate (eagle picks the launch mode) |
90.4 ms (858 µs) |
– |
91.3 ms |
551 M |
609 G (–) |
8 MiB (1.17× min.) |
25.3 MiB |
100,000 |
CuPy masked |
5.32 s (19.8 ms) |
– |
5.32 s |
9.37 M |
10.4 G (–) |
8 MiB (1.17× min.) |
22.9 MiB |
100,000 |
PyTorch masked |
4.17 s (12.6 ms) |
– |
4.18 s |
11.9 M |
13.2 G (–) |
96 MiB (14× min.) |
20.4 MiB |
100,000 |
PyTorch + CUDA graph |
1.87 s (90.5 µs) |
– |
1.87 s |
26.6 M |
29.5 G (–) |
316 MiB (46× min.) |
41.7 MiB |
100,000 |
JAX jit(vmap(while_loop)) |
411 ms (146 µs) |
– |
414 ms |
121 M |
134 G (–) |
160 MiB (23.3× min.) |
29.9 MiB |
100,000 |
Warp per-thread kernel |
164 ms (56.1 µs) |
– |
166 ms |
303 M |
335 G (–) |
30 MiB (4.37× min.) |
18.8 MiB |
100,000 |
eagle CPU (OpenMP) |
106 ms (2.14 ms) |
– |
106 ms |
471 M |
521 G (–) |
– |
26.2 MiB |
1,000,000 |
eagle graph |
1.02 s (20.5 µs) |
– |
1.03 s |
489 M |
541 G (–) |
72 MiB (1.05× min.) |
228 MiB |
1,000,000 |
eagle graph + compaction |
607 ms (15.4 µs) |
– |
621 ms |
821 M |
908 G (–) |
88 MiB (1.28× min.) |
229 MiB |
1,000,000 |
eagle graph + compaction + reorder |
607 ms (47.8 µs) |
– |
621 ms |
821 M |
908 G (–) |
116 MiB (1.69× min.) |
229 MiB |
1,000,000 |
hawk + eagle persistent |
628 ms (507 µs) |
– |
640 ms |
794 M |
878 G (–) |
70 MiB (1.02× min.) |
229 MiB |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
627 ms (1.03 ms) |
– |
639 ms |
795 M |
879 G (–) |
70 MiB (1.02× min.) |
229 MiB |
1,000,000 |
CuPy masked |
9.31 s (7.65 ms) |
– |
9.32 s |
53.5 M |
59.2 G (–) |
78 MiB (1.14× min.) |
183 MiB |
1,000,000 |
PyTorch masked |
9.94 s (1.09 ms) |
– |
9.95 s |
50.1 M |
55.5 G (–) |
186 MiB (2.71× min.) |
191 MiB |
1,000,000 |
PyTorch + CUDA graph |
8.55 s (208 µs) |
– |
8.56 s |
58.3 M |
64.5 G (–) |
1.37e+03 MiB (19.9× min.) |
239 MiB |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.79 s (14.9 ms) |
– |
2.82 s |
179 M |
198 G (–) |
948 MiB (13.8× min.) |
319 MiB |
1,000,000 |
Warp per-thread kernel |
1.12 s (12.1 µs) |
– |
1.13 s |
446 M |
494 G (–) |
94 MiB (1.37× min.) |
184 MiB |
1,000,000 |
eagle CPU (OpenMP) |
969 ms (247 µs) |
– |
980 ms |
514 M |
569 G (–) |
– |
247 MiB |
Agreement with the analytic orbit
Every arm, every sample: final time reached exactly; position error against the analytic Kepler state (Newton on Kepler’s equation to machine precision) below the stated bound n_acc tol (1 + v_p)(1 + 6 pi n_orbits v_p), derived from the controller’s per-step budget and the Kepler along-track amplification (see truth_bound; not fitted). Step counts are compared with the reference arm (eagle graph): any difference comes from a floating-point difference flipping one accept/reject decision (contraction into fused multiply-adds, pow and sqrt implementations differ between compilers and libraries); the table reports how many samples differ and by how much.
N |
arm |
max position error |
max error / bound |
samples with different step counts |
max Δ accepted / rejected |
|---|---|---|---|---|---|
1,000,000 |
eagle graph |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle graph + compaction + reorder |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
hawk + eagle persistent |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
eagle.simulate (eagle picks the launch mode) |
2.26e-07 |
6.64e-03 |
0 |
0 / 0 |
1,000,000 |
CuPy masked |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
PyTorch masked |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
PyTorch + CUDA graph |
2.26e-07 |
6.64e-03 |
2 |
1 / 2 |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.26e-07 |
6.64e-03 |
4 |
1 / 2 |
1,000,000 |
Warp per-thread kernel |
2.26e-07 |
6.64e-03 |
2 |
0 / 2 |
1,000,000 |
eagle CPU (OpenMP) |
2.26e-07 |
6.64e-03 |
2 |
1 / 1 |
1,000 |
eagle graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle graph + compaction + reorder |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
hawk + eagle persistent |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle.simulate (eagle picks the launch mode) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
CuPy masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch masked |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
PyTorch + CUDA graph |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
JAX jit(vmap(while_loop)) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
Warp per-thread kernel |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
1,000 |
eagle CPU (OpenMP) |
1.85e-07 |
6.35e-03 |
0 |
0 / 0 |
10,000 |
eagle graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle graph + compaction + reorder |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
hawk + eagle persistent |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle.simulate (eagle picks the launch mode) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
CuPy masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch masked |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
PyTorch + CUDA graph |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
JAX jit(vmap(while_loop)) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
Warp per-thread kernel |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
10,000 |
eagle CPU (OpenMP) |
2.02e-07 |
6.57e-03 |
0 |
0 / 0 |
100,000 |
eagle graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle graph + compaction + reorder |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
hawk + eagle persistent |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle.simulate (eagle picks the launch mode) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
CuPy masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch masked |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
PyTorch + CUDA graph |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
JAX jit(vmap(while_loop)) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
Warp per-thread kernel |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
100,000 |
eagle CPU (OpenMP) |
2.25e-07 |
6.62e-03 |
0 |
0 / 0 |
Steps per sample at N = 1,000,000 (reference arm): accepted min 322 / median 371 / p90 571 / max 693; rejected min 1 / median 59 / max 264.
Where each tool fits
Each tool here solves the same 13-stage adaptive problem to the same tolerance and agrees with the analytic orbit; they differ in how the per-sample control flow is written and run. Numbers below are at N = 1,000,000; the fastest arm there is eagle graph + compaction (607 ms).
eagle graph. the attempt is per-sample scalar code (22 lines with its launch; the kernel finishes its own samples), one kernel per attempt looped on the device: 1.02 s.
eagle graph + compaction. the same kernel launched over the live samples only: 607 ms; the plain graph takes 1.68× as long.
eagle graph + compaction + reorder. adds the occasional reorder of the planes: 607 ms; the compacted graph takes 1× as long.
hawk + eagle persistent. the plain (not active-set) attempt kernel, forced through hawk’s persist entry: one launch, no graph, no map, no policy kernel – a grid of SMs blocks of 256 lanes steals the next unfinished sample off a counter when their own finishes: 628 ms; the eagle graph takes 1.62× as long.
eagle.simulate (eagle picks the launch mode). the same kernel through eagle.simulate, its state and parameters by name (21 lines with the kernel): 627 ms; the eagle graph takes 1.62× as long.
CuPy masked. array code in the NumPy style (25 lines), the control flow as masks over the batch: 9.31 s.
PyTorch masked. the same array formulation in PyTorch (27 lines), running where a PyTorch model already runs: 9.94 s.
PyTorch + CUDA graph. records 16 attempts as a CUDA graph and replays them, removing per-kernel launch cost: 8.55 s; eager PyTorch takes 1.16× as long.
JAX jit(vmap(while_loop)). writes the per-sample loop directly (25 lines) and compiles the whole batch into one executable: 2.79 s.
Warp per-thread kernel. an explicit per-thread loop in a Python-embedded kernel language (74 lines, its own stages): 1.12 s.
eagle CPU (OpenMP). the same hawk kernel on the host’s cores, no GPU needed: 969 ms.
Compile time
Its own pass, never part of any wall. Each measurement is a fresh process: imports and the GPU context run untimed, then the timer spans the first call to a ready kernel (trace + code generation + compile + load, or the cache lookup). Cold: every cache empty (HAWK_CACHE_DIR, CUPY_CACHE_DIR, CUDA_CACHE_PATH, WARP_CACHE_PATH and JAX’s persistent compilation cache, each a fresh directory); warm: a new process on the caches the first cold run populated. eagle GPU: eagle.deploy of the hawk attempt kernel and of its active-set variant (host and device builds each), and one eagle graph of each built and run on 64 samples (eagle.simulate deploys the first again, a cache hit). eagle CPU: the same eagle.deploy of the attempt kernel (host and device builds) and one host run. CuPy: one masked attempt on 64 samples (every elementwise kernel compiled on first use). JAX: lower + compile from abstract shapes at N. Warp: the kernel module (code generation, compile, load) and one launch. Medians of 3 cold and 3 warm processes.
arm |
compiled |
cold (median of 3) |
warm (median of 3) |
|---|---|---|---|
eagle graph, eagle graph + compaction, eagle graph + compaction + reorder, hawk + eagle persistent, eagle.simulate (eagle picks the launch mode) |
eagle.deploy of the attempt kernel and its active-set variant |
1.72 s |
290 ms |
eagle CPU (OpenMP) |
eagle.deploy of the attempt kernel, host run |
1.37 s |
175 ms |
CuPy masked |
CuPy elementwise kernels |
1.82 s |
724 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000 |
1.52 s |
170 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 10,000 |
1.62 s |
171 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 100,000 |
2.14 s |
170 ms |
JAX jit(vmap(while_loop)) |
XLA executable, N = 1,000,000 |
2.78 s |
165 ms |
Warp per-thread kernel |
Warp kernel module |
4.38 s |
50.7 ms |
No compile step: PyTorch masked, PyTorch + CUDA graph (PyTorch’s eager kernels ship prebuilt; the CUDA graph is recorded, not compiled).
Memory method
A separate pass after the timing pass. Each (arm, N) runs in a fresh process: imports, the CUDA context created by one tiny device operation, then the baseline read, before the arm builds or allocates anything; then the arm’s kernels built and a warm-up run on a 64-sample batch (every compile happens here; JAX also compiles for shape N), so every library’s pool reserved during the warm-up counts in the reading; CuPy’s pools and PyTorch’s cache trimmed after the warm-up; then the arm built for N (eagle’s graph, PyTorch’s captured graph) and one full run with upload and download; then the reading again, the allocators still holding their high-water reservation (the same measure for every arm, no sampling needed). Device: the memory NVML accounts to the process on the timed GPU, after minus baseline. Caveats: each allocator’s rounding and growth policy count (JAX’s allocator grows in regions that can exceed the request; PyTorch rounds blocks up and a captured CUDA graph keeps a private pool; Warp’s stream-ordered pool reserves in chunks of tens of MiB, its release threshold set to 0 after the warm-up); the driver accounts in pages of about 2 MiB, so small-N rows read 0 or one page. Host: the peak resident set (VmHWM, reset to the current RSS at the host baseline, taken after the warm-up) minus the baseline RSS. Minimum: the planes the workload needs, N x (8 float64 planes (state, time, step) x 8 bytes + 2 counters x 4 bytes) = N x 72 bytes.
Fastest arm by median wall time:
N = 1,000: eagle CPU (OpenMP) (3.45 ms), then eagle.simulate (eagle picks the launch mode) (24.2 ms, 7× the time)
N = 10,000: eagle CPU (OpenMP) (12.7 ms), then eagle.simulate (eagle picks the launch mode) (25.2 ms, 1.99× the time)
N = 100,000: eagle.simulate (eagle picks the launch mode) (90.4 ms), then hawk + eagle persistent (91.9 ms, 1.02× the time)
N = 1,000,000: eagle graph + compaction (607 ms), then eagle graph + compaction + reorder (607 ms, 1× the time)
eagle_simulate against its hand-built twin eagle_graph:
N = 1,000: wall 0.979× the twin’s time
N = 10,000: wall 0.978× the twin’s time
N = 100,000: wall 0.599× the twin’s time
N = 1,000,000: wall 0.615× the twin’s time
device |
compute capability / CPU model |
date |
|---|---|---|
NVIDIA H100 80GB HBM3 |
9.0 |
2026-10-13 |
NVIDIA L4 |
8.9 |
2026-10-14 |
NVIDIA A100-SXM4-40GB |
8.0 |
2026-10-13 |
Tesla T4 |
7.5 |
2026-10-13 |
Quadro P2000 |
6.1 |
2026-10-13 |
NVIDIA RTX PRO 6000 Blackwell Server Edition |
12.0 |
2026-10-14 |
CPU#
Run on: host CPU AMD EPYC 9B45; compiler x86_64-conda-linux-gnu-g++ (conda-forge gcc 14.4.0-7) 14.4.0; Python 3.12.14; JAX 0.11.2, PyTorch 2.12.1+cpu, NumPy 2.4.2.
N = 1,000,000
Log scale, from 2.35 s to 96.9 s.
Table of every arm and every N
CPU: AMD EPYC 9B45, 16 logical CPUs. Medians of 5 interleaved repetitions; IQR in parentheses. Workload as the GPU card (1106 FLOP per attempt); masked arms up to N = 100,000.
N |
arm |
threads |
wall (IQR) |
attempts·samples/s |
useful FLOP/s |
host memory |
compile (cold) |
lines of code |
|---|---|---|---|---|---|---|---|---|
1,000 |
eagle CPU, 8 threads |
8 |
7.01 ms (221 µs) |
72 M |
79.7 G |
0.281 MiB |
1.48 s |
22 |
1,000 |
eagle CPU, 1 thread |
1 |
30.7 ms (19.7 µs) |
16.4 M |
18.2 G |
0 MiB |
1.42 s |
22 |
1,000 |
Numba prange |
8 |
34.4 ms (12.9 ms) |
14.7 M |
16.2 G |
0 MiB |
2.27 s |
58 |
1,000 |
JAX jit(vmap(while_loop)) |
8 |
86.6 ms (122 µs) |
5.83 M |
6.45 G |
0.68 MiB |
1.2 s |
25 |
1,000 |
PyTorch masked |
8 |
1.16 s (10.7 ms) |
436 k |
482 M |
1.62 MiB |
– |
27 |
1,000 |
NumPy masked |
1 |
625 ms (3.29 ms) |
808 k |
894 M |
0.824 MiB |
– |
23 |
10,000 |
eagle CPU, 8 threads |
8 |
40.3 ms (623 µs) |
123 M |
136 G |
1.29 MiB |
1.43 s |
22 |
10,000 |
eagle CPU, 1 thread |
1 |
300 ms (470 µs) |
16.5 M |
18.3 G |
1.14 MiB |
1.44 s |
22 |
10,000 |
Numba prange |
8 |
198 ms (1.91 ms) |
25 M |
27.6 G |
0.824 MiB |
2.32 s |
58 |
10,000 |
JAX jit(vmap(while_loop)) |
8 |
399 ms (200 µs) |
12.4 M |
13.7 G |
1.84 MiB |
1.11 s |
25 |
10,000 |
PyTorch masked |
8 |
3.05 s (649 ms) |
1.63 M |
1.8 G |
12.4 MiB |
– |
27 |
10,000 |
NumPy masked |
1 |
3.24 s (467 ms) |
1.53 M |
1.69 G |
7.97 MiB |
– |
23 |
100,000 |
eagle CPU, 8 threads |
8 |
381 ms (2.41 ms) |
131 M |
145 G |
11.7 MiB |
1.43 s |
22 |
100,000 |
eagle CPU, 1 thread |
1 |
2.99 s (7.41 ms) |
16.6 M |
18.4 G |
14.2 MiB |
1.43 s |
22 |
100,000 |
Numba prange |
8 |
1.97 s (19.4 ms) |
25.2 M |
27.9 G |
11.6 MiB |
2.28 s |
58 |
100,000 |
JAX jit(vmap(while_loop)) |
8 |
3.44 s (7.52 ms) |
14.5 M |
16 G |
30.2 MiB |
1.14 s |
25 |
100,000 |
PyTorch masked |
8 |
44 s (368 ms) |
1.13 M |
1.25 G |
122 MiB |
– |
27 |
100,000 |
NumPy masked |
1 |
29.7 s (569 ms) |
1.68 M |
1.85 G |
85.4 MiB |
– |
23 |
1,000,000 |
eagle CPU, 8 threads |
8 |
3.76 s (7.93 ms) |
133 M |
147 G |
115 MiB |
1.44 s |
22 |
1,000,000 |
eagle CPU, 1 thread |
1 |
29.9 s (13 ms) |
16.6 M |
18.4 G |
146 MiB |
1.43 s |
22 |
1,000,000 |
Numba prange |
8 |
19.9 s (168 ms) |
25 M |
27.7 G |
121 MiB |
2.27 s |
58 |
1,000,000 |
JAX jit(vmap(while_loop)) |
8 |
60.5 s (189 ms) |
8.23 M |
9.11 G |
312 MiB |
1.12 s |
25 |
N |
arm |
max position error |
max error / bound |
samples with different step counts |
|---|---|---|---|---|
1,000 |
eagle CPU, 8 threads |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
eagle CPU, 1 thread |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
Numba prange |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
JAX jit(vmap(while_loop)) |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
PyTorch masked |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
NumPy masked |
1.85e-07 |
6.35e-03 |
0 |
10,000 |
eagle CPU, 8 threads |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
eagle CPU, 1 thread |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
Numba prange |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
JAX jit(vmap(while_loop)) |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
PyTorch masked |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
NumPy masked |
2.02e-07 |
6.57e-03 |
0 |
100,000 |
eagle CPU, 8 threads |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
eagle CPU, 1 thread |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
Numba prange |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
JAX jit(vmap(while_loop)) |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
PyTorch masked |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
NumPy masked |
2.25e-07 |
6.62e-03 |
0 |
1,000,000 |
eagle CPU, 8 threads |
2.26e-07 |
6.64e-03 |
0 |
1,000,000 |
eagle CPU, 1 thread |
2.26e-07 |
6.64e-03 |
0 |
1,000,000 |
Numba prange |
2.26e-07 |
6.64e-03 |
2 |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.26e-07 |
6.64e-03 |
4 |
Run on: host CPU AMD EPYC 7B13; compiler x86_64-conda-linux-gnu-g++ (conda-forge gcc 14.4.0-7) 14.4.0; Python 3.12.14; JAX 0.11.2, PyTorch 2.12.1+cpu, NumPy 2.4.2.
N = 1,000,000
Log scale, from 16.4 s to 99.5 s.
Table of every arm and every N
CPU: AMD EPYC 7B13, 16 logical CPUs. Medians of 5 interleaved repetitions; IQR in parentheses. Workload as the GPU card (1106 FLOP per attempt); masked arms up to N = 100,000.
N |
arm |
threads |
wall (IQR) |
attempts·samples/s |
useful FLOP/s |
host memory |
compile (cold) |
lines of code |
|---|---|---|---|---|---|---|---|---|
1,000 |
eagle CPU, 8 threads |
8 |
37.4 ms (9.3 ms) |
13.5 M |
14.9 G |
0.199 MiB |
2.15 s |
22 |
1,000 |
eagle CPU, 1 thread |
1 |
212 ms (84.3 µs) |
2.38 M |
2.63 G |
0 MiB |
2.13 s |
22 |
1,000 |
Numba prange |
8 |
61.7 ms (8.8 ms) |
8.18 M |
9.04 G |
0 MiB |
3.56 s |
58 |
1,000 |
JAX jit(vmap(while_loop)) |
8 |
245 ms (2.66 ms) |
2.06 M |
2.28 G |
0.73 MiB |
2.01 s |
25 |
1,000 |
PyTorch masked |
8 |
2.02 s (4.02 ms) |
250 k |
277 M |
1.56 MiB |
– |
27 |
1,000 |
NumPy masked |
1 |
1.08 s (4.01 ms) |
465 k |
515 M |
0.836 MiB |
– |
23 |
10,000 |
eagle CPU, 8 threads |
8 |
273 ms (6.8 ms) |
18.2 M |
20.1 G |
1.2 MiB |
2.15 s |
22 |
10,000 |
eagle CPU, 1 thread |
1 |
2.09 s (3.1 ms) |
2.37 M |
2.62 G |
1.07 MiB |
2.14 s |
22 |
10,000 |
Numba prange |
8 |
319 ms (543 µs) |
15.5 M |
17.2 G |
0.598 MiB |
3.53 s |
58 |
10,000 |
JAX jit(vmap(while_loop)) |
8 |
830 ms (3.3 ms) |
5.97 M |
6.61 G |
2.04 MiB |
1.86 s |
25 |
10,000 |
PyTorch masked |
8 |
8.25 s (116 ms) |
601 k |
665 M |
12.2 MiB |
– |
27 |
10,000 |
NumPy masked |
1 |
7.7 s (751 ms) |
644 k |
712 M |
8.25 MiB |
– |
23 |
100,000 |
eagle CPU, 8 threads |
8 |
2.63 s (2.97 ms) |
18.9 M |
20.9 G |
11.6 MiB |
2.15 s |
22 |
100,000 |
eagle CPU, 1 thread |
1 |
21 s (15.8 ms) |
2.37 M |
2.62 G |
14.2 MiB |
2.14 s |
22 |
100,000 |
Numba prange |
8 |
3.16 s (3.36 ms) |
15.8 M |
17.4 G |
11.3 MiB |
3.54 s |
58 |
100,000 |
JAX jit(vmap(while_loop)) |
8 |
6.15 s (11.9 ms) |
8.09 M |
8.95 G |
30.3 MiB |
2.32 s |
25 |
100,000 |
PyTorch masked |
8 |
42.6 s (2.89 s) |
1.17 M |
1.29 G |
121 MiB |
– |
27 |
100,000 |
NumPy masked |
1 |
43.5 s (1.34 s) |
1.15 M |
1.27 G |
85.1 MiB |
– |
23 |
1,000,000 |
eagle CPU, 8 threads |
8 |
26.2 s (21.7 ms) |
19 M |
21 G |
115 MiB |
2.16 s |
22 |
1,000,000 |
eagle CPU, 1 thread |
1 |
210 s (241 ms) |
2.38 M |
2.63 G |
146 MiB |
2.14 s |
22 |
1,000,000 |
Numba prange |
8 |
31.9 s (94.9 ms) |
15.6 M |
17.3 G |
121 MiB |
3.55 s |
58 |
1,000,000 |
JAX jit(vmap(while_loop)) |
8 |
62.2 s (17 ms) |
8.01 M |
8.86 G |
313 MiB |
2.31 s |
25 |
N |
arm |
max position error |
max error / bound |
samples with different step counts |
|---|---|---|---|---|
1,000 |
eagle CPU, 8 threads |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
eagle CPU, 1 thread |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
Numba prange |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
JAX jit(vmap(while_loop)) |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
PyTorch masked |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
NumPy masked |
1.85e-07 |
6.35e-03 |
0 |
10,000 |
eagle CPU, 8 threads |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
eagle CPU, 1 thread |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
Numba prange |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
JAX jit(vmap(while_loop)) |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
PyTorch masked |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
NumPy masked |
2.02e-07 |
6.57e-03 |
0 |
100,000 |
eagle CPU, 8 threads |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
eagle CPU, 1 thread |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
Numba prange |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
JAX jit(vmap(while_loop)) |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
PyTorch masked |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
NumPy masked |
2.25e-07 |
6.62e-03 |
0 |
1,000,000 |
eagle CPU, 8 threads |
2.26e-07 |
6.64e-03 |
0 |
1,000,000 |
eagle CPU, 1 thread |
2.26e-07 |
6.64e-03 |
0 |
1,000,000 |
Numba prange |
2.26e-07 |
6.64e-03 |
0 |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.26e-07 |
6.64e-03 |
2 |
Fastest arm by median wall time:
N = 1,000: eagle CPU, 8 threads (37.4 ms), then Numba prange (61.7 ms, 1.65× the time)
N = 10,000: eagle CPU, 8 threads (273 ms), then Numba prange (319 ms, 1.17× the time)
N = 100,000: eagle CPU, 8 threads (2.63 s), then Numba prange (3.16 s, 1.2× the time)
N = 1,000,000: eagle CPU, 8 threads (26.2 s), then Numba prange (31.9 s, 1.22× the time)
Run on: host CPU aarch64; compiler aarch64-conda-linux-gnu-g++ (conda-forge gcc 14.4.0-7) 14.4.0; Python 3.12.14; JAX 0.11.2, PyTorch 2.12.1+cpu, NumPy 2.4.2.
N = 1,000,000
Log scale, from 8.32 s to 66.6 s.
Table of every arm and every N
CPU: aarch64, 16 logical CPUs. Medians of 5 interleaved repetitions; IQR in parentheses. Workload as the GPU card (1106 FLOP per attempt); masked arms up to N = 100,000.
N |
arm |
threads |
wall (IQR) |
attempts·samples/s |
useful FLOP/s |
host memory |
compile (cold) |
lines of code |
|---|---|---|---|---|---|---|---|---|
1,000 |
eagle CPU, 8 threads |
8 |
14.7 ms (17.3 µs) |
34.4 M |
38.1 G |
0.309 MiB |
1.08 s |
22 |
1,000 |
eagle CPU, 1 thread |
1 |
108 ms (34.1 µs) |
4.7 M |
5.19 G |
0 MiB |
1.09 s |
22 |
1,000 |
Numba prange |
8 |
23.8 ms (256 µs) |
21.2 M |
23.4 G |
0 MiB |
2.67 s |
58 |
1,000 |
JAX jit(vmap(while_loop)) |
8 |
182 ms (665 µs) |
2.78 M |
3.07 G |
0.777 MiB |
1.46 s |
25 |
1,000 |
PyTorch masked |
8 |
1.48 s (1.54 ms) |
341 k |
377 M |
1.01 MiB |
– |
27 |
1,000 |
NumPy masked |
1 |
742 ms (1.46 ms) |
681 k |
753 M |
0.824 MiB |
– |
23 |
10,000 |
eagle CPU, 8 threads |
8 |
139 ms (62.1 µs) |
35.8 M |
39.6 G |
1.24 MiB |
1.09 s |
22 |
10,000 |
eagle CPU, 1 thread |
1 |
1.06 s (18.1 µs) |
4.68 M |
5.17 G |
1.07 MiB |
1.08 s |
22 |
10,000 |
Numba prange |
8 |
229 ms (4.17 ms) |
21.6 M |
23.9 G |
0.52 MiB |
2.69 s |
58 |
10,000 |
JAX jit(vmap(while_loop)) |
8 |
582 ms (531 µs) |
8.52 M |
9.42 G |
1.72 MiB |
1.43 s |
25 |
10,000 |
PyTorch masked |
8 |
5 s (6.6 ms) |
991 k |
1.1 G |
16 MiB |
– |
27 |
10,000 |
NumPy masked |
1 |
4.19 s (888 ms) |
1.18 M |
1.31 G |
8.38 MiB |
– |
23 |
100,000 |
eagle CPU, 8 threads |
8 |
1.33 s (202 µs) |
37.3 M |
41.3 G |
14.5 MiB |
1.08 s |
22 |
100,000 |
eagle CPU, 1 thread |
1 |
10.6 s (301 µs) |
4.69 M |
5.19 G |
14.3 MiB |
1.08 s |
22 |
100,000 |
Numba prange |
8 |
2.25 s (5.67 ms) |
22.1 M |
24.5 G |
11.7 MiB |
2.69 s |
58 |
100,000 |
JAX jit(vmap(while_loop)) |
8 |
4.34 s (2.41 ms) |
11.5 M |
12.7 G |
30.5 MiB |
1.44 s |
25 |
100,000 |
PyTorch masked |
8 |
12.9 s (20.7 ms) |
3.85 M |
4.26 G |
136 MiB |
– |
27 |
100,000 |
NumPy masked |
1 |
35.4 s (762 ms) |
1.41 M |
1.55 G |
85.3 MiB |
– |
23 |
1,000,000 |
eagle CPU, 8 threads |
8 |
13.3 s (599 µs) |
37.4 M |
41.4 G |
146 MiB |
1.08 s |
22 |
1,000,000 |
eagle CPU, 1 thread |
1 |
106 s (5.65 ms) |
4.69 M |
5.19 G |
146 MiB |
1.09 s |
22 |
1,000,000 |
Numba prange |
8 |
22.5 s (83.6 ms) |
22.2 M |
24.5 G |
122 MiB |
2.67 s |
58 |
1,000,000 |
JAX jit(vmap(while_loop)) |
8 |
41.6 s (32.1 ms) |
12 M |
13.2 G |
313 MiB |
1.45 s |
25 |
N |
arm |
max position error |
max error / bound |
samples with different step counts |
|---|---|---|---|---|
1,000 |
eagle CPU, 8 threads |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
eagle CPU, 1 thread |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
Numba prange |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
JAX jit(vmap(while_loop)) |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
PyTorch masked |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
NumPy masked |
1.85e-07 |
6.35e-03 |
0 |
10,000 |
eagle CPU, 8 threads |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
eagle CPU, 1 thread |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
Numba prange |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
JAX jit(vmap(while_loop)) |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
PyTorch masked |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
NumPy masked |
2.02e-07 |
6.57e-03 |
0 |
100,000 |
eagle CPU, 8 threads |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
eagle CPU, 1 thread |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
Numba prange |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
JAX jit(vmap(while_loop)) |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
PyTorch masked |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
NumPy masked |
2.25e-07 |
6.62e-03 |
0 |
1,000,000 |
eagle CPU, 8 threads |
2.26e-07 |
6.64e-03 |
0 |
1,000,000 |
eagle CPU, 1 thread |
2.26e-07 |
6.64e-03 |
0 |
1,000,000 |
Numba prange |
2.26e-07 |
6.64e-03 |
0 |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.26e-07 |
6.64e-03 |
1 |
Fastest arm by median wall time:
N = 1,000: eagle CPU, 8 threads (14.7 ms), then Numba prange (23.8 ms, 1.63× the time)
N = 10,000: eagle CPU, 8 threads (139 ms), then Numba prange (229 ms, 1.65× the time)
N = 100,000: eagle CPU, 8 threads (1.33 s), then Numba prange (2.25 s, 1.69× the time)
N = 1,000,000: eagle CPU, 8 threads (13.3 s), then Numba prange (22.5 s, 1.69× the time)
Run on: host CPU INTEL XEON PLATINUM 8581C CPU @ 2.30GHz; compiler x86_64-conda-linux-gnu-g++ (conda-forge gcc 14.4.0-7) 14.4.0; Python 3.12.14; JAX 0.11.2, PyTorch 2.12.1+cpu, NumPy 2.4.2.
N = 1,000,000
Log scale, from 3.04 s to 82.1 s.
Table of every arm and every N
CPU: INTEL(R) XEON(R) PLATINUM 8581C CPU @ 2.30GHz, 16 logical CPUs. Medians of 5 interleaved repetitions; IQR in parentheses. Workload as the GPU card (1106 FLOP per attempt); masked arms up to N = 100,000.
N |
arm |
threads |
wall (IQR) |
attempts·samples/s |
useful FLOP/s |
host memory |
compile (cold) |
lines of code |
|---|---|---|---|---|---|---|---|---|
1,000 |
eagle CPU, 8 threads |
8 |
5.59 ms (84.8 µs) |
90.4 M |
99.9 G |
0.281 MiB |
2.51 s |
22 |
1,000 |
eagle CPU, 1 thread |
1 |
38.7 ms (1.09 ms) |
13 M |
14.4 G |
0 MiB |
1.66 s |
22 |
1,000 |
Numba prange |
8 |
35.2 ms (13.2 ms) |
14.3 M |
15.9 G |
0 MiB |
3.9 s |
58 |
1,000 |
JAX jit(vmap(while_loop)) |
8 |
159 ms (1.18 ms) |
3.17 M |
3.51 G |
0.66 MiB |
2.84 s |
25 |
1,000 |
PyTorch masked |
8 |
1.27 s (9.08 ms) |
398 k |
441 M |
1.63 MiB |
– |
27 |
1,000 |
NumPy masked |
1 |
728 ms (8.46 ms) |
693 k |
767 M |
0.824 MiB |
– |
23 |
10,000 |
eagle CPU, 8 threads |
8 |
50.9 ms (114 µs) |
97.4 M |
108 G |
1.29 MiB |
1.66 s |
22 |
10,000 |
eagle CPU, 1 thread |
1 |
389 ms (2.18 ms) |
12.8 M |
14.1 G |
1.14 MiB |
1.66 s |
22 |
10,000 |
Numba prange |
8 |
341 ms (9.59 ms) |
14.6 M |
16.1 G |
0.609 MiB |
2.59 s |
58 |
10,000 |
JAX jit(vmap(while_loop)) |
8 |
719 ms (3.58 ms) |
6.89 M |
7.62 G |
1.93 MiB |
1.33 s |
25 |
10,000 |
PyTorch masked |
8 |
4.2 s (196 ms) |
1.18 M |
1.31 G |
12.4 MiB |
– |
27 |
10,000 |
NumPy masked |
1 |
5.48 s (699 ms) |
905 k |
1 G |
8.3 MiB |
– |
23 |
100,000 |
eagle CPU, 8 threads |
8 |
488 ms (2.35 ms) |
102 M |
113 G |
11.7 MiB |
1.66 s |
22 |
100,000 |
eagle CPU, 1 thread |
1 |
3.88 s (9.05 ms) |
12.8 M |
14.2 G |
14.2 MiB |
1.66 s |
22 |
100,000 |
Numba prange |
8 |
3.31 s (6.75 ms) |
15 M |
16.6 G |
11.8 MiB |
2.6 s |
58 |
100,000 |
JAX jit(vmap(while_loop)) |
8 |
5.54 s (3.12 ms) |
8.98 M |
9.94 G |
30.3 MiB |
1.35 s |
25 |
100,000 |
PyTorch masked |
8 |
24.3 s (1.69 s) |
2.05 M |
2.27 G |
122 MiB |
– |
27 |
100,000 |
NumPy masked |
1 |
56 s (223 ms) |
889 k |
983 M |
85.3 MiB |
– |
23 |
1,000,000 |
eagle CPU, 8 threads |
8 |
4.86 s (13.5 ms) |
103 M |
113 G |
115 MiB |
1.68 s |
22 |
1,000,000 |
eagle CPU, 1 thread |
1 |
39 s (12.7 ms) |
12.8 M |
14.2 G |
146 MiB |
1.66 s |
22 |
1,000,000 |
Numba prange |
8 |
33.2 s (53 ms) |
15 M |
16.6 G |
122 MiB |
2.59 s |
58 |
1,000,000 |
JAX jit(vmap(while_loop)) |
8 |
51.3 s (42.2 ms) |
9.71 M |
10.7 G |
312 MiB |
1.33 s |
25 |
N |
arm |
max position error |
max error / bound |
samples with different step counts |
|---|---|---|---|---|
1,000 |
eagle CPU, 8 threads |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
eagle CPU, 1 thread |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
Numba prange |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
JAX jit(vmap(while_loop)) |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
PyTorch masked |
1.85e-07 |
6.35e-03 |
0 |
1,000 |
NumPy masked |
1.85e-07 |
6.35e-03 |
0 |
10,000 |
eagle CPU, 8 threads |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
eagle CPU, 1 thread |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
Numba prange |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
JAX jit(vmap(while_loop)) |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
PyTorch masked |
2.02e-07 |
6.57e-03 |
0 |
10,000 |
NumPy masked |
2.02e-07 |
6.57e-03 |
0 |
100,000 |
eagle CPU, 8 threads |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
eagle CPU, 1 thread |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
Numba prange |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
JAX jit(vmap(while_loop)) |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
PyTorch masked |
2.25e-07 |
6.62e-03 |
0 |
100,000 |
NumPy masked |
2.25e-07 |
6.62e-03 |
0 |
1,000,000 |
eagle CPU, 8 threads |
2.26e-07 |
6.64e-03 |
0 |
1,000,000 |
eagle CPU, 1 thread |
2.26e-07 |
6.64e-03 |
0 |
1,000,000 |
Numba prange |
2.26e-07 |
6.64e-03 |
2 |
1,000,000 |
JAX jit(vmap(while_loop)) |
2.26e-07 |
6.64e-03 |
4 |
Fastest arm by median wall time:
N = 1,000: eagle CPU, 8 threads (5.59 ms), then Numba prange (35.2 ms, 6.3× the time)
N = 10,000: eagle CPU, 8 threads (50.9 ms), then Numba prange (341 ms, 6.7× the time)
N = 100,000: eagle CPU, 8 threads (488 ms), then Numba prange (3.31 s, 6.78× the time)
N = 1,000,000: eagle CPU, 8 threads (4.86 s), then Numba prange (33.2 s, 6.84× the time)
device |
compute capability / CPU model |
date |
|---|---|---|
AMD EPYC 9B45 |
AMD EPYC 9B45 |
2026-10-13 |
AMD EPYC 7B13 |
AMD EPYC 7B13 |
2026-10-13 |
Google Axion (Arm Neoverse-V2) |
aarch64 |
2026-10-13 |
Intel Xeon Platinum 8581C |
INTEL(R) XEON(R) PLATINUM 8581C CPU @ 2.30GHz |
2026-10-13 |