Open source

djinnax

Batch-native JAX simulation engines, with receipts.

My public benchmark repo. Simulation engines written in a strict batch-native, branchless discipline race the established JAX environment suites on identical tasks. Every engine passes step-for-step parity gates against its reference before a single timing runs, so the speedups measure the discipline and nothing else.

See the code on GitHub

Headline numbers

What the discipline buys.

240×

fastest engine vs. its reference

Median of five frozen-code fresh-process sweeps at batch 8,192 on a consumer RTX 4080 Super, measured against the reference implementation.

2.5B

env-steps per second, peak

Full environment rollouts, RNG included, in one persistent Triton kernel launch. Peak observed; the median under this host's permanent VRAM contention is 1.7B, and both figures are published.

numbers published before parity

Every engine is verified step-for-step against its reference before it is allowed near a benchmark. Known caveats are stated next to the numbers they qualify.

The spectrum

The win scales with control flow.

The speedup depends on how much branching the reference keeps in its hot path. An engine that is already a handful of array ops gains little. An engine full of conditionals, rewritten branchless, gains an order of magnitude, and moving the whole rollout into one persistent GPU kernel gains two. Within-run ratios vs. the reference, medians of an n=5 frozen-code sweep:

batch	light control flow	moderate	heavy, branchless	heavy, LUT	heavy, megakernel
64	1.1×	1.4×	33.9×	33.2×	239×
1,024	1.2×	1.5×	25.2×	31.7×	213×
8,192	1.1×	2.8×	22.7×	21.9×	240×
65,536	1.2×	1.2×	8.0×	13.8×	57׆

Three engines of increasing hot-path complexity, each parity-gated against pgx or jumanji. The heaviest appears three ways: branchless JAX, a lookup-table variant, and a persistent Triton megakernel that keeps the entire rollout in one kernel launch. †The batch 65,536 megakernel figure is a conservative lower bound under this host's VRAM contention. Full intervals, raw rows, and caveats live in the repo.

Newest receipt

One RNG site, up to 75× faster.

Many reference implementations hide a retry loop inside their randomness: pick a random slot, and if it is taken, roll again. Ported literally into a batched engine, that loop breaks kernel fusion and makes every simulation in the batch wait for the unluckiest one. I replaced it with a single draw from the exact conditional distribution, proved the two versions statistically identical across 40,000 draws each way, then raced full simulations on the standard harness.

75×

full simulations, batch 1,024

Median speedup for the collapsed draw over the faithful retry-loop port, spread 45× to 88× across fresh-process runs.

47×

still holds at batch 65,536

The retry loop pays the batch's worst case on every step, so the gap stays enormous at every batch size measured.

RNG site changed

Outcome distributions provably unchanged. The generalized samplers (retry, reroll, weighted tables) ship as djinnax/distributions.py.

Read the samplers and the measurement script, both on main.

Methodology

How the numbers are made.

Parity before timing

Each engine is checked step-for-step against the reference suite it races, and the megakernel is bit-verified against an XLA reference running the same step function. A benchmark only runs after the gate passes.

Medians with spread

Every figure is the median of multiple fresh-process runs with min and max reported alongside. Anything under 2× is measured interleaved and pairwise so machine drift cannot manufacture a win.

Null results kept

Experiments that failed to speed anything up stay in the repo, labeled as such. The 1.1× rows are published for the same reason the 240× is.

Caveats attached

RNG-stream differences, hardware specificity, and single-GPU-model evidence are stated next to the numbers they qualify, in the README and in the docs.

In the repo

Built to be picked up.

Beyond the engines and the measurement harness, the repo documents the discipline itself: a porting playbook, a guide to writing fast batch-native environments, a megakernel authoring recipe, and the evidence log explaining why each rule exists. Point a coding agent at it and the house style comes along. This is the same handover standard my client work ships with.

github.com/DonaldGallianoIII/djinnax

Contact

Have something slow?

Send me the problem: what it is, where it runs, and what it costs you in time. I come back with whether it can be made fast, how fast, and what the engagement looks like. You work with me directly, start to finish.

hello@crawdata.ai

I reply within two business days.

From the mud to models