Backends and batching

Two implementations of the same evaluation, NumPy and compiled, and what a batch buys on each. Written 2026-09-02. Numbers marked "measured here" were taken that day on a shared login node, OMP_NUM_THREADS=1, OMP_PROC_BIND=false, best of five, through the consumer route on the released v3 artifacts, with the compiled module aa-fanout/kernel-mode-kokkos/cpp/build-mrsafe (md5 48a74a07, built from commit 4336162d). Everything else says where it was measured.

Two backends, and which one you get

model = NRHJSurAA.load_h5(RELEASE_AA, backend="numpy")    # AA default
model = NRHJSurAA.load_h5(RELEASE_AA, backend="kokkos")
model = NRHJSurAdA.load(RELEASE_ADA)                       # "auto"

"numpy" is the reference path and needs nothing beyond the declared dependencies. "kokkos" routes the time-domain hot path, the tau inversion, the piecewise-Chebyshev synthesis of every mode at every sample and the e^{i m s Phi} rotation, through the compiled nrhjkokkos extension; the coefficient-space work stays in NumPy on both (nrhjsurrogate/execution/batched_evaluator.py module docstring). The adiabatic-angle class defaults to "auto": kokkos if the module imports, NumPy otherwise (nrhjsurrogate/driver/api.py:288-294); measured, the same call resolved to numpy without a module and kokkos with one.

Where the module comes from. $NRHJ_KOKKOS_DIR names one directory and nothing else is consulted; otherwise every cpp/build* beside the package holding an nrhjkokkos*.so is a candidate and the NEWEST wins (nrhjsurrogate/execution/compiled_module_loader.py:34). The first rule is why a shell that has sourced one of the checked-in env_*.sh files hands its compiled build to every process it starts, including one that believed it was testing a clean install. The wheel is py3-none-any and carries no compiled object, so a pip install user has the NumPy backend and only that, until checklist B1 (fetching and compiling Kokkos at install time) lands; the deposit gate proves the six files plus the wheel suffice on the NumPy route and refuses to pass if a compiled module is importable.

Agreement between the backends

Measured here on one waveform at (4, 0.3, -0.2), same 15775-sample grid:

The AA number is the tau-inversion class (kernel Newton against CubicSpline) and is what tests/test_aa_api.py::test_vs_numpy_reference gates at 2e-5 of peak (ACCURACY.md section 3). On the derivative side, the compiled and NumPy paths agree to 1.4e-14 relative, and a batched against an unbatched contraction to about 1e-10, which is not a defect but the conditioning of the contraction arriving at evaluation time: the worst element cancels by 8.3e6, so neither side is accurate below that.

A parity baseline is therefore backend specific. Measured here at the branch point of this page's own branch: the NumPy-only capture against parity_baseline_v4.npz fails on 143 of 447 arrays, all of them the 80 adiabatic-angle gradient arrays the NumPy route cannot produce, the 62 action-angle gradient arrays that differ at up to 1e-14, and the sentinel that records the skip; the capture with the compiled module above is 447 of 447 bitwise. Compare like with like.

Cost, one waveform and a batch

Measured here, milliseconds per waveform:

One thread on a shared node is a floor, not a throughput figure. The throughput figures, at scale, recorded where they were taken:

• Batched value evaluator, v3 artifact, 60 threads, exclusive node, measured 2026-08-27: 5.65 ms per waveform at B = 1, 0.618 at B = 64, 0.451 at B = 4096, against the v2 artifact's 0.356 at B = 4096, so v3 sits within 1.27x of v2. That branch is merged into main (5594afb7) and the plan builders on this tree accept symmetry_mode="kernel" (compiled_module_loader.py:421, :514). Before it, the shipping route refused the v3 artifacts' kernel mode in both plan builders and fell back to per-waveform NumPy regression, which is the 18x anti-scaling the earlier campaign measured and which must not be quoted as v3 throughput. The paper review's re-read of the same campaign at 127 threads gives 0.389 ms per waveform at B = 4096.
• A 120-thread cliff remains: v3 goes from 0.939 to 5.183 ms per waveform between 60 and 120 threads at B = 64 where v2 is flat, and it hits the Python fallback too. Run the batched arm at about 60 threads, not 128.
• Batched gradients, host-parallel route, 8192 samples, full span: 446.9 ms serial, 11.70 ms per waveform at B = 1024 on 60 processes. With the compiled merger-ringdown gradient kernel and the compiled theta stage, 6.894 ms per waveform at B = 1024 on 60 threads, 1.697x over host-parallel, at which point the merger-ringdown arm is 2.2 percent of the route and the inspiral arm 97.8 percent. Process guidance for the host route: about 32 processes below B = 256, about 60 above B = 512, never 128; the pool is built inside every call, 1.43 s at 60 processes, and Pool.map pipes 3.7 MB per row (checklist B6, left as is).

Timing on the ms-per-waveform convention throughout; a batched figure is amortised over the batch it names.

Threads

Importing the package sets OMP_PROC_BIND=close and OMP_PLACES=cores if neither is already set, never OMP_NUM_THREADS (nrhjsurrogate/__init__.py:30). The runtime reads the binding once, at its first thread pool, so a value set after the import is silently ignored; nrhjsurrogate.thread_binding() reads back what the runtime actually did rather than echoing what was asked. Measured here with OMP_PROC_BIND=false already in the environment, the import applied nothing and the readback said false, which is the contract. NRHJ_NO_THREAD_DEFAULTS opts out.

The gain is modest and was restated once: 1.08 to 1.11x on the whole call at 32 threads on an exclusive AMD EPYC 7742, with a demonstrated 1.3 percent noise floor; an earlier 1.25 to 1.28x "one stage" figure is WITHDRAWN because that run was not exclusive (INSTALL.md section 5). Two traps: a forked worker pool under the default binding pins EVERY process to core 0, so the checked-in preamble exports OMP_PROC_BIND=false for multi-process work, and srun without -n 1 fans a command out silently.

What each backend refuses or lacks

• kokkos, dynamics=True on the result dictionary raises dynamics output is not available on backend='kokkos', measured, rather than dropping the key. get_orbital_dynamics, the remnant and the error estimate all work on a kokkos-backed model, measured, because they are host-side.
• NumPy, adiabatic-angle derivatives: NRHJSurAdA.modes_grad raises ImportError without the module, measured, and strain_grad and time_derivative are always kernel-fused (api.py:699-716). The paper's compiled-only performance claims were flagged for exactly this. The install-verification suite records that check as SKIP on a NumPy-only install (install-verification.md).
• The compiled merger-ringdown gradient stage is opt-in and the action-angle scalar modes_grad takes it when a module is present, so the gradient arrays are the ones a parity baseline sees move between backends.
• The normaliser guard. A compiled module is stamped with the post-Newtonian normaliser it was built for and a stamped artifact is refused through a module that does not match or does not say (compiled_module_loader.py:273, assert_normaliser_matches); the hole it closes was a stale .so serving a silent 1.5e-2 error on (3, 3).
• The source guard is broken. assert_module_matches_source (compiled_module_loader.py:188), the measurement preflight that ties a binary to a source commit, raises NameError: name 'os' is not defined when called (measured directly). Until it is fixed a published timing cannot use it to prove which build it ran on; this page names the build's md5 and provenance commit by hand instead.
• safe is a compound of two levers, safe_elementfits and safe_anchor (aa_api.py:546-565). True, the default, is the canonical route: a correctly rounded anchor contraction and the bitwise Chebyshev loop. False selects the BLAS anchor and the fast stacked evaluation through libaa_fastpath.so when it is present, which the compiler may contract into fused multiply-adds; it is not -ffast-math, and it moves results at the roundoff class, 0 on the waveform and 1e-11 to 6.8e-10 on derivatives and error estimates (install-verification.md). The rename planned for this keyword has not happened; it is still safe.
• physical_batch(..., single_precision=True) opts into float32 device arithmetic at up to 1.9e-4 of peak strain, about 1e-8 in mismatch, for likelihood exploration and never for a gate (driver/aa_batch.py).

Batching, the shapes

• Array input on model(...), both backends: t stays (N,) and each mode becomes (B, N), measured with B = 2. On the NumPy backend that is a loop.
• physical_batch / AABatch.physical: rows of (m1, m2, chi1z, chi2z [, distance, inclination, phase, t_c]) on ONE shared physical grid, each row zero-padded outside its own span, the shape a batched likelihood consumes; a requested window shorter than a row's span is refused rather than truncated (driver/aa_batch.py module docstring, grid policy).
• modes_grad_batch (driver/aa_batch_gradients.py): a host-parallel fork pool over the scalar gradient chain, bitwise against it (parameter-derivatives.md).
• Below-reach f_low in a batch on kokkos is a row loop over the scalar arm (hybridisation.md).
• driver/bilby_batch.py is the batched-likelihood adapter, 1.5e-14 relative in log-likelihood against the stock serial path (ACCURACY.md).

Where the code is