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)                       # "auto" (default)
model = NRHJSurAA.load_h5(RELEASE_AA, backend="kokkos")
model = NRHJSurAdA.load(RELEASE_ADA)                       # "auto"

Since 1.1 both classes default to "auto", which takes the compiled backend when it is usable and NumPy when it is not. NRHJSurAA defaulted to "numpy" before that.

"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