GPU Acceleration Evaluation: JAX vs. Numba

This page documents a 2026-07-26 evaluation of whether PyFracVAL’s numba-jitted hot path would benefit from a JAX/GPU port. The conclusion is negative: the existing numba CPU kernels outperform JAX (both CPU and GPU) by one to four orders of magnitude at every problem size PyFracVAL operates at, and the gap does not close as problems grow. Measurements were taken on an NVIDIA TITAN X (Pascal, compute capability 6.1).

Scope

Every @jit-decorated function in pyfracval/ was catalogued (kernels in overlap.py, geometry.py, pca_kernels.py, cca_kernels.py, densify.py, and experimental/). The numerical hot path concentrates in two places:

• overlap.py: pairwise overlap/distance checks between two point sets, called up to ~18,000 times per aggregate (once per rotation attempt during PCA/CCA sticking);

• geometry._rodrigues_rotation_2d: rotation of an (N,3) cluster around an axis, called once per rotation attempt alongside the overlap check.

Both are, on their face, reasonable candidates for GPU acceleration: numerical array operations, called frequently, on a problem that nominally scales with particle count — the profile a JAX/GPU port typically targets.

Method

benchmarks/jax_kernels.py ports the two hot-path kernels to JAX (rodrigues_rotation_jax, max_overlap_pairwise_jax, max_overlap_single_vs_agg_jax, batch_check_overlaps_pca_jax — the last mirrors pca_kernels.batch_check_overlaps_pca, batching K candidate rotations into a single call, the shape best suited to JAX/GPU). jax_enable_x64 is enabled, since PyFracVAL uses float64 throughout and JAX’s float32 default would silently change simulation precision.

benchmarks/jax_vs_numba_benchmark.py times numba against JAX-CPU and JAX-GPU for each kernel across a size sweep. Each kernel/framework/size combination is warmed up once (JIT-compile cost, excluded from the timed loop) and then timed over 50 calls with fresh random data per call, matching how the kernels are invoked in the simulation: coordinates and radii differ on every rotation attempt, so nothing can remain resident on the GPU across calls without a substantially larger rewrite of the sticking loop. JAX timings include host–device transfer and block_until_ready() per call, since the caller needs the concrete result before the next step of the algorithm; this reflects actual usage rather than a benchmark artifact. Raw results: benchmark_results/jax_vs_numba_summary.json.

Results

Rodrigues rotation

No early exit, purely vectorized — the case most favorable to JAX:

N particles	numba	JAX (CPU)	JAX (GPU)
16	6.6 µs	11.5 µs	62.6 µs
64	7.0 µs	11.5 µs	60.4 µs
256	8.0 µs	12.2 µs	45.4 µs
1,024	10.7 µs	15.9 µs	46.5 µs
4,096	20.9 µs	29.6 µs	59.9 µs
8,192	34.1 µs	54.2 µs	58.6 µs

Even for this kernel, with no branching and nothing for JAX to lose algorithmically, numba is faster at every size tested. JAX-GPU’s ~45–60 µs floor reflects dispatch, kernel-launch, and PCIe transfer latency; it changes little with N because at these sizes the compute is negligible relative to that fixed cost. numba’s cached compiled-to-native dispatch has no equivalent floor.

Pairwise overlap

numba’s early exit and bounding-sphere pre-check against JAX’s dense O(N²) matrix:

N per cluster	numba	JAX (CPU)	JAX (GPU)
16	1.0 µs	13.2 µs	50.6 µs
64	0.8 µs	68.1 µs	55.8 µs
256	1.4 µs	438.2 µs	120.1 µs
1,024	3.5 µs	2,452 µs	412 µs
4,096	32.3 µs	45,647 µs	3,609 µs

The gap widens here, reaching roughly three orders of magnitude at N=4,096. numba’s calculate_max_overlap_cca_fast performs a cheap squared-distance vs. squared-radius-sum comparison before computing any sqrt, skipping the expensive arithmetic for pairs that cannot overlap, and returns as soon as it finds one pair that does. JAX/XLA has no equivalent: on SIMD/GPU hardware a per-element where does not skip work — both branches are computed for every lane — so JAX must materialize the full N×N distance matrix and reduce it unconditionally.

Two follow-up checks indicate the gap is architectural rather than a tuning artifact:

• No crossover exists even far past realistic sizes. Sweeping the sparse/scattered scenario up to N=32,768 (≈1.07 billion pairs), numba stays at 67–200 µs throughout (the pre-check continues to skip nearly everything) while JAX-GPU grows to 165 ms; the gap widens with N.

• Densely-packed data, where numba’s early return fires almost immediately, widens the gap further: numba drops to ~0.5–0.7 µs while JAX-GPU still pays 100 µs–2.9 ms for N=256–4,096, computing the full matrix regardless.

No data distribution in this problem favors JAX: sparse data lets the pre-check skip almost everything, and dense data lets the early return fire almost immediately, while JAX pays the full O(N²) cost in either case.

Batched overlap check

K candidate rotations per call, mirroring pca_kernels.batch_check_overlaps_pca — the comparison expected to favor JAX/GPU most, since per-call dispatch overhead is amortized over K:

N_agg	K	numba	JAX (CPU)	JAX (GPU)
64	1	5.8 µs	13.5 µs	43.2 µs
64	360	9.0 µs	219.9 µs	60.2 µs
512	360	29.9 µs	640.2 µs	131.4 µs
2,048	1	10.3 µs	40.8 µs	50.6 µs
2,048	360	100.5 µs	1,532.6 µs	332.4 µs

Batching improves JAX’s relative standing: the numba/JAX-GPU ratio narrows from ~7× (K=1) to ~3.3× (K=360, N_agg=2,048), JAX’s closest result in this evaluation, though still a clear loss in absolute terms. numba’s batch_check_overlaps_pca is itself parallelized (prange across candidates) and cached, leaving little margin to close.

Compilation caching

Both frameworks’ persistent caches were verified directly (empty vs. warm cache directory, same kernel, same input size):

Framework	Cold compile (empty cache)	Cache hit (warm .devenv/state/)
numba (cache=True)	501 ms	299 ms
JAX (jax_compilation_cache_dir)	155 ms	19 ms

JAX’s cache produces the larger relative improvement (~8×), since XLA compilation is heavier than numba’s LLVM path for these small kernels; JAX additionally pays a one-time ~300 ms CUDA context initialization per process, unrelated to caching. Both caches function as intended and matter for the startup latency of short-lived processes (a benchmark run, a single CLI invocation). Neither changes the steady-state figures above, which determine the outcome: the simulation loop runs each kernel thousands of times per aggregate within one long-lived process, so warm-cache startup cost is amortized away, and the per-call dispatch/launch overhead that caching does not address decides the comparison.

Conclusion

A JAX/GPU port of overlap.py or geometry.rodrigues_rotation is not being pursued. The existing numba implementation is close to optimal for this problem’s shape: many small, branch-heavy, early-exit-friendly pairwise geometry computations, called in a tight sequential loop, on modest per-call array sizes. Branching that skips work is precisely what GPU-style parallelism cannot exploit, and the fixed per-call dispatch/launch/transfer floor (tens of microseconds, independent of framework maturity or caching) dominates at every problem size PyFracVAL generates aggregates at.

The result parallels the CCA sticking retrospective (experiments.md): a second case in which additional compute does not improve the outcome, because the binding constraint is not throughput but the structure of the problem. The devenv CUDA setup and the JAX kernel ports are retained (benchmarks/jax_kernels.py, benchmarks/jax_vs_numba_benchmark.py) in case a different parameterization — generating a large batch of independent aggregates in parallel on the GPU, rather than accelerating one aggregate’s inner rotation loop — is revisited later.

Implementation notes

devenv.nix mirrors ../YASF-new/devenv.nix’s CUDA setup: a cudaPackages_12 toolkit via symlinkJoin, CUDA_HOME, and NUMBA_CUDA_DRIVER pointing at /run/opengl-driver/lib/libcuda.so, with /run/opengl-driver/lib appended to LD_LIBRARY_PATH (JAX’s CUDA plugin requires the driver directory on the standard dynamic-linker search path; numba does not, since it reads NUMBA_CUDA_DRIVER directly). devenv.yaml sets allow_unfree: true for the CUDA toolkit’s license. Both NUMBA_CACHE_DIR and JAX_COMPILATION_CACHE_DIR point at .devenv/state/, so compiled kernels persist across process restarts for both frameworks.

jax[cuda12] is not synced by default, since it pulls approximately 1.5 GB of CUDA wheels (cudnn, nccl, nvshmem); install it explicitly with uv sync --group test --group docs --group jax_bench (all three groups together, since uv sync --group X alone replaces rather than extends the synced set). With this in place:

>>> import jax; jax.devices()
[CudaDevice(id=0)]