perf: DelaunayNN Sibson — unroll candidate loop, single concatenated pass, chunk as memory guard

[Jammy2211]

Overview

Post-#531 on the A100 (HST / Hilbert-1500 / MGE-60 / ConstantSplit) DelaunayNN costs 197.0 ms per likelihood vs 62.0 ms for barycentric Delaunay; the whole excess is the two Sibson passes (params→H prefix 144.8 ms vs 7.2 ms unbatched; 24.4 vs 5.1 ms per call at vmap 16). The compiled HLO of sibson_mappings_weights_from_tables has three nested loops (a lax.map over SIBSON_QUERY_CHUNK = 256 → 95 sequential chunks, the 32-trip cavity fori_loop, and a 3-trip candidate fori_loop): 1,244 kernel launches per chunk, 95 % in the cavity walk, ≈118k per likelihood at ~1.1 µs each — the A100 launch floor, serialised by the chunk loop. This issue ships Phase A (bit-identical launch-count cuts) and, after measurement, Phase B (cavity early exit). Phase C (loop-free k-ring cavity) is a separate prompt.

Witness: params→H prefix (regularization_matrix_prefix_s) ≤ 60 ms unbatched and ≤ 11 ms per call at vmap 16 after Phase A (baseline 144.789 / 24.424 in results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_walk_early_exit.json), pin EXPECTED_LOG_EVIDENCE_HST = 29144.581944 unchanged, delaunay_nn.py / delaunay_nn_caps.py jax_assertions pass.

Standing constraint: no slowdown for the numba CPU or JAX CPU likelihoods (NumPy Sibson path untouched; JAX CPU gated by a local before/after breakdown run).

Plan

1. Unroll the 3-trip candidate loop inside the cavity walk (bit-identical; launches per chunk
   1,244 → 892).
2. Locate and interpolate data grid + split points in one concatenated Sibson pass in
   jax_delaunay_nn, computing the circumcircles once.
3. Make the query chunk the only memory guard it was meant to be: sweep it on the A100 at
   vmap 16 with peak-VRAM readings and ship the largest safe value as the new default.
4. Add a breakdown stage that separates split-side Sibson from the 33-wide ConstantSplit
   assembly, so the remaining ~27 ms residual is attributed.
5. Prove parity (bit-identical outputs vs main), jit/vmap round-trips and the gradient checks
   in the workspace_test jax_assertions; CPU no-regression for the JAX CPU path; numba path is
   untouched by construction (NumPy Sibson path scipy_delaunay_nn is not changed).
6. Ship the library PR; A100 A/B (control at merge base, feature) incl. the chunk sweep; ship
   the profiling/workspace_test PR; then Phase B (early-exit cavity while_loop with the
   stop_gradient boundary) as a second library PR on the same task, measured the same way.

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary, library)
• autolens_workspace_test (jax_assertions extension)
• autolens_profiling (breakdown stage, submit scripts, A100 results)

A. PyAutoArray — autoarray/inversion/mesh/interpolator/sibson.py

A.1 Unroll the candidate loop. In _cavity_triangle_indexes_jax (156–208),
process_triangle (173) calls jax.lax.fori_loop(0, 3, add_candidate, …) (201). Replace
with a Python for edge in range(3): carry = add_candidate(edge, carry). Insertion order is
unchanged → stencil column order and fp summation order unchanged (the analyst verified
bit-identical mappings/sizes/weights/cavity_sizes/overflow/degenerate on CPU). Keep the
outer cavity fori_loop (203) as is in Phase A.

A.2 One concatenated pass in jax_delaunay_nn (629–720): the local
mappings_weights_for is called at 670 (data) and 699 (split). Call it once on
jnp.concatenate([query_points, split_points]) and slice all six outputs at n_query. Hoist
delaunay_circumcircles_from (computed inside sibson_mappings_weights_from_tables at 451)
so it runs once — add an optional circumcircles= argument to
sibson_mappings_weights_from_tables (default None → compute as now) so jax_sibson (505,
used by autolens_workspace_test/.../jax_assertions/delaunay_nn.py) and the NumPy path are
unaffected. Split points keep their own nearest-vertex fallback (the located
simplex_indexes/fallback come from the single pix_indexes_delaunay_walk_from call, which
now also runs once on the concatenated queries — the #531 pattern). scipy_delaunay_nn
(546–626) untouched.

A.3 Chunk as a tunable memory guard. SIBSON_QUERY_CHUNK (36) stays the module default
and DelaunayNN.query_chunk (mesh/mesh/delaunay_nn.py:47) the mesh attribute. Add an
environment override read once at import, PYAUTO_SIBSON_QUERY_CHUNK (documented next to the
constant), so the A100 sweep can vary it without code edits. The new default is set from the
sweep result in the same PR (decision rule: largest of 512/1024/2048/4096 whose peak VRAM at
vmap 16 stays under ~50 % of the 80 GB, and which is fastest per call; if 4096 wins on both,
ship 2048 unless the margin is > 15 %, to keep headroom for larger cells). Update the
SIBSON_QUERY_CHUNK comment to say what the chunk guards ((C,3,2) intermediates ≈ 15–25 kB
per query per lane) and cite the sweep.

A.4 Docstrings. sibson_mappings_weights_from_tables docstring: the loop structure, why
the chunk exists, launch-count reasoning in one paragraph. jax_delaunay_nn: the single-pass
note mirroring jax_delaunay's.

B. PyAutoArray — unit tests (NumPy only)

test_autoarray/inversion/pixelization/interpolator/test_sibson.py (and the mesh/mapper
DelaunayNN tests): all existing tests must pass unchanged. Add a NumPy-path test that
sibson_mappings_weights_from_tables(..., circumcircles=precomputed) equals the default call
bit-for-bit, and that scipy_delaunay_nn output is unchanged (regression on a fixed seed:
mappings, sizes, weights, split outputs). The JAX unroll itself is exercised only on the JAX
path (workspace_test), per the no-JAX-in-unit-tests rule.

C. autolens_workspace_test — scripts/misc/jax_assertions/delaunay_nn.py

Extend (not a new script): (1) a bit-identity check against a frozen reference is not
possible in CI without main, so instead add a parity check jax_delaunay_nn (concatenated)
vs jax_sibson called separately on data and split queries — same tables, same chunk — asserting
exact equality of mappings/sizes/weights/cavity sizes for both halves; (2) a chunk-invariance
check: query_chunk 64 vs 256 vs 1024 give identical outputs (proves the chunk is a pure
memory guard and the env override works); (3) keep the existing gradient/linear-precision/
flip-continuity checks; keep the runtime under the script's current budget (reuse its
SIBSON_* overrides). delaunay_nn_caps.py unchanged (caps untouched).

D. autolens_profiling — scripts/imaging/likelihood_breakdown/delaunay_nn.py

D.1 New _setup_prefix_fn stage (819) that stops right after
InterpolatorDelaunayNN._mappings_sizes_weights_split (returns the split mappings/weights
alongside step-6 outputs, strict superset like stage 11), timed in --split-setup and
reported as "Split-point Sibson" so the H row becomes "ConstantSplit assembly" alone. Update
the module docstring's row table and the H-row attribution paragraph (prefix-difference
semantics; judge on regularization_matrix_prefix_s).
D.2 Honour PYAUTO_SIBSON_QUERY_CHUNK in the recorded JSON metadata (xla_flags-style
provenance field) so sweep rows are self-describing; --output-dir naming
delaunay_nn_hpc_a100_fp64_chunk<N>.json.
D.3 Submit scripts: hpc/batch_gpu/submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_{control,launch_latency}
(the #531 pattern: private checkouts at merge base and feature on PYTHONPATH, provenance
printed in the SLURM log) plus a chunk-sweep array script (one task per chunk, nvidia-smi
peak VRAM sampled to the log).

E. Gates before the library PR

• pytest test_autoarray -q green; black --check on touched files.
• Workspace_test delaunay_nn.py and delaunay_nn_caps.py pass locally against the worktree
  autoarray (CPU); paste PASS lines and wall time.
• CPU no-regression (standing user constraint): likelihood_breakdown/delaunay_nn.py --split-setup locally on main vs branch, 3–5 warm reps, params→H not slower, pin
  unchanged; delaunay_numba.py control unchanged (NumPy Sibson path untouched; grep confirms
  no NumPy-path call site changed).
• Local micro-benchmark of jax_delaunay_nn (N=1500, Q=17,980 + 6,000) on CPU before/after
  as a fine-grained readout (expect the unroll + single pass to show; the chunk effect is a GPU
  story).
• Gradient: the Phase A change keeps the fori_loop→scan lowering, so reverse-mode is
  unchanged; the jax_assertions gradient checks are the proof.

F. Ship + A100 + Phase B

1. ship_library (Opus): PR on PyAutoArray, pending-release, issue + Mind updated.
2. A100 (Opus, RAL): control (merge base) vs feature delaunay_nn.py --split-setup --vmap-batch 16 at the current chunk, then the chunk sweep 512/1024/2048/4096 at vmap 16
   with peak VRAM; the runtime cell; DelaunayNN pin unchanged on every leg. Pick the default
   per A.3 and push it as a follow-up commit on the same library branch before the PR is
   merged (re-run the unit tests + jax_assertions on that commit). Write
   results/notes/delaunay_nn_launch_latency.md and the JSON/PNGs.
3. ship_workspace (Opus): workspace_test PR (C) and profiling PR (D + results), library-first
   gate.
4. Phase B on the same task after the numbers: cavity fori_loop (203) → lax.while_loop
   with cond = (position < max_cavity_triangles) & jnp.any(position < count); overflow
   detection (198, 418–423) unchanged; stop_gradient on query, circumcentres,
   circumradii_squared only where they enter _contains_query inside the loop; the
   circumcentres[safe_cavity] gather (~259) stays traced. Same gates, second library PR,
   same A100 A/B. Expected ~7 ms after Phase A. Decision point for you after Phase A's numbers.

Expected after Phase A: params→H ≈ 50 ms unbatched (2.9×), ≈ 9 ms per call at vmap 16 (2.7×);
with Phase B ≈ 45 / 8. Floor with Phase C ≈ 30 / 6 (Delaunay is 7.2 / 5.1).

Branch Survey

Repository	Branch	Dirty?	Claims
./PyAutoArray	main (d7c96762, #531 merged)	clean	none; orphan worktree delaunay-area-magnification-audit still present (unregistered, warn only)
./autolens_workspace_test	main (d818cfb)	clean	none
./autolens_profiling	main (99f4b53)	clean	retire-gpu1-mig-exclusion (hpc/ submits it deletes, README, activate.sh) — disjoint → parallel-claim note as before

Task name delaunay-nn-launch-latency; branch feature/delaunay-nn-launch-latency; worktree
~/Code/PyAutoLabs-wt/delaunay-nn-launch-latency/; active.md status library-dev.

Suggested branch: feature/delaunay-nn-launch-latency

Key Files

• PyAutoArray/autoarray/inversion/mesh/interpolator/sibson.py — _cavity_triangle_indexes_jax, sibson_mappings_weights_from_tables, jax_delaunay_nn, SIBSON_QUERY_CHUNK
• PyAutoArray/autoarray/inversion/mesh/mesh/delaunay_nn.py — mesh-level query_chunk binding
• PyAutoArray/test_autoarray/inversion/pixelization/interpolator/test_sibson.py — NumPy tests
• autolens_workspace_test/scripts/misc/jax_assertions/delaunay_nn.py — JAX parity / gradient gate
• autolens_profiling/scripts/imaging/likelihood_breakdown/delaunay_nn.py — breakdown stages, pin
• autolens_profiling/results/notes/delaunay_walk_early_exit.md — the post-feat: JAX Delaunay walk — early-exit while_loop, chunk only the seed argmin #531 baseline and attribution caveat

Original Prompt

Click to expand starting prompt

DelaunayNN (Sibson) on the A100: kill the kernel-launch latency in the cavity walk

Type: feature
Target: autoarray
Repos:

• PyAutoArray
• autolens_profiling
• autolens_workspace_test
  Themes:
• jax-gpu
• delaunay
• profiling
• performance
  Difficulty: medium
  Autonomy: supervised
  Priority: high
  Status: draft
  Consequence: judge
  Witness: the A100 DelaunayNN breakdown's params→H prefix (regularization_matrix_prefix_s, results/breakdown/imaging/delaunay_nn_hpc_a100_fp64_walk_early_exit.json is the 2026-09-07 baseline: 144.789 ms unbatched / 24.424 ms per call at vmap 16) drops to at most 60 ms unbatched and 11 ms per call at vmap 16 after Phase A, with EXPECTED_LOG_EVIDENCE_HST = 29144.581944 unchanged and the delaunay_nn.py / delaunay_nn_caps.py jax_assertions passing
  Review-minutes: 40
  Unattended: ready
  Filed: 2026-09-07
  Supersedes: draft/feature/autoarray/sibson_single_concatenated_walk.md

Original request (verbatim):

ok then prm, then HPCPullPyAuto, then go on to do the work on delaunayNN. We recently did a likelihood_breakdown of delaunay_nn, so check that out and then work out if on the A100 we can make it really fast overall

The measurement (A100, HST / Hilbert-1500 / MGE-60 / ConstantSplit, post PyAutoArray#531)

DelaunayNN costs 197.0 ms per likelihood unbatched against 62.0 ms for barycentric Delaunay,
and the whole excess sits in the two Sibson passes: the params→H prefix is 144.8 ms (NN) vs
7.2 ms (Delaunay); per call at vmap 16 it is 24.4 ms vs 5.1 ms. The four-way split charges
92 ms to "Triangulation + interpolation" (the data-side pass over 17,980 queries) and 45 ms to
the H row (the split-side pass over the 6,000 ConstantSplit points, which re-runs the whole
Sibson pipeline including a second circumcircle computation). The downstream rows are
identical to Delaunay's (blurred mapping matrix 8.36 vs 8.34 ms at vmap 16), so the 32-wide
mapper costs nothing after the setup.

The compiled HLO of sibson_mappings_weights_from_tables
(PyAutoArray/autoarray/inversion/mesh/interpolator/sibson.py) has three nested loops: the
lax.map over SIBSON_QUERY_CHUNK = 256 queries (95 sequential chunks: 71 data + 24 split),
the 32-trip cavity fori_loop (MAX_CAVITY_TRIANGLES, the loop always runs to the cap), and a
3-trip candidate fori_loop inside it. That is 32 × (4 + 3 × 11) + 60 = 1,244 kernel launches
per chunk, 95 % of them in the cavity walk, ≈ 118,000 launches per likelihood. Against the
measured 144.8 ms that is ~1.1 µs per launch: the A100 launch floor. Corroboration: 16× the
lanes (vmap 16) cost 2.7× the time. The cost is dispatch latency serialised by the chunk loop,
the same disease DELAUNAY_LOCATE_CHUNK had, one level down. Observed cavity sizes on this
cell are mean 3.5, max 9 (audit maxima across 101 traced meshes: 25 main / 19 split).

Phase A: one PR, all launch-count reductions that keep bit-identical output

1. Unroll the 3-trip candidate fori_loop in process_triangle (sibson.py ~179–201) into a
   Python for edge in range(3). Insertion order is preserved, so stencil column order and
   fp summation order are unchanged; verified bit-identical on CPU. Launches per chunk
   1,244 → 892.
2. Raise SIBSON_QUERY_CHUNK (sibson.py:36, bound at mesh/mesh/delaunay_nn.py:47). The
   chunk is only a memory guard; sweep 512 / 1024 / 2048 / 4096 on the A100 at vmap 16 and
   record peak VRAM (a vmap-64 OOM on this cell is already on record in
   delaunay_nn_hpc_a100_fp64_vmap64.json). Ship the largest value with comfortable headroom
   as the new default; keep the constant overridable.
3. Locate and interpolate the data grid and the split points in one concatenated
   mappings_weights_for call in jax_delaunay_nn (sibson.py ~670 and ~699), slicing at
   n_query, and hoist delaunay_circumcircles_from so the circumcircles are computed once.
   Split points keep their own nearest-vertex fallback. Small on its own (~1 chunk), worth it
   once the chunk is large.
4. Add a _setup_prefix_fn stage to
   autolens_profiling/scripts/imaging/likelihood_breakdown/delaunay_nn.py that stops after
   _mappings_sizes_weights_split, so the H row separates split-side Sibson from the
   33-wide ConstantSplit assembly (the ~27 ms non-chunk residual is currently unattributed).

Expected: params→H ≈ 50 ms unbatched (2.9×) and ≈ 9 ms per call at vmap 16 (2.7×).

Phase B: early-exit the cavity walk (after Phase A is measured)

Convert the cavity fori_loop (sibson.py ~203–208) to a lax.while_loop with
cond = (position < cap) & any(position < count), cap kept as the safety bound, overflow
detection (sibson.py ~198, ~418–423) unchanged. Under vmap the exit is at the global max
(~25), so the win is ~20 % of the cavity part: worth ~7 ms after Phase A. Gradient contract:
today the fori_loop lowers to scan and is reverse-mode differentiable; while_loop is
not, so the float inputs to _contains_query (query, circumcentres, circumradii_squared)
must be stop_gradient-wrapped as they enter the loop (integer/bool outputs only, zero a.e.
derivative), while the circumcentres[safe_cavity] gather at ~259 that feeds the Sibson
weights stays traced. Re-run the DelaunayNN gradient checks after.

Phase C (separate prompt, not this task): loop-free cavity via a fixed k-ring gather

Replace the cavity walk with a one-shot containment test over the seed simplex's k-ring
(3-ring 21 / 4-ring 45 candidates), compact, overflow → NaN as now. Removes the loop
entirely (launches per chunk → ~80) for an estimated floor of ~30 ms unbatched / ~6 ms per
call at vmap 16, but changes candidate ordering and therefore fp summation order in the
mapping matrix, so the pin needs a relative tolerance and the cap audit must be re-run.
File after Phase B's numbers.

Contracts (do not relax)

• Pin EXPECTED_LOG_EVIDENCE_HST = 29144.581944 in delaunay_nn.py passes unchanged through
  Phases A and B; a shift means a mapping or summation order changed and is a bug.
• Judge on the params→H prefix (regularization_matrix_prefix_s), not on the Tri+interp or H
  rows: both are prefix differences and step 3 moves work between them exactly as feat: JAX Delaunay walk — early-exit while_loop, chunk only the seed argmin #531 did
  for Delaunay (results/notes/delaunay_walk_early_exit.md).
• SIBSON_MAX_NEIGHBORS / MAX_CAVITY_TRIANGLES stay at 32 (cap audit
  results/notes/delaunay_nn_cap_audit.md); any cap change re-runs
  autolens_workspace_test/scripts/misc/jax_assertions/delaunay_nn_caps.py.
• jax_sibson (sibson.py ~505) is exercised by
  autolens_workspace_test/scripts/misc/jax_assertions/delaunay_nn.py (sibson_tables), so
  it stays; apply the unroll to the shared sibson_mappings_weights_from_tables so both
  entry points benefit.
• Unit tests stay NumPy-only; JAX-path parity, jit/vmap and gradient checks live in the
  workspace_test jax_assertions scripts.

Verification on the A100

Same-node, same-session A/B as PyAutoArray#531 (control = private checkout at the merge
base, feature = branch; shared /mnt/ral/jnightin/PyAuto untouched while science jobs run):
scripts/imaging/likelihood_breakdown/delaunay_nn.py --config-name hpc_a100_fp64 --split-setup --vmap-batch 16, plus the chunk sweep for step 2 with peak-VRAM readings, and the runtime
cell. Report every row unbatched and per call at vmap 16, the params→H prefix, the new
ConstantSplit stage, and the single-JIT total, against the 2026-09-07 post-#531 baseline.

Folds in draft/feature/autoarray/sibson_single_concatenated_walk.md (filed 2026-09-07 as the
follow-up of complete/2026/09/delaunay-walk-early-exit.md); the analysis above shows the
concatenation alone is worth ~1 chunk, so it ships as step A.3 rather than on its own.

[Jammy2211]

Library PR Created

PyAutoArray: #533 (pending-release, commit 9e9a7d50)

Phase A of this issue is shipped to a PR. The A100 evidence said the DelaunayNN params→H prefix (144.8 ms unbatched / 24.4 ms per call at vmap 16, against Delaunay's 7.2 / 5.1) is a launch floor — 1,244 kernel launches per 256-query chunk × 95 chunks ≈ 118k launches at ≈1.1 µs each — so the PR cuts launches, not flops: (A.1) the 3-trip candidate-edge fori_loop in _cavity_triangle_indexes_jax is unrolled at trace time behind _sibson_unroll_candidates() (env PYAUTO_SIBSON_UNROLL_CANDIDATES "1"/"0" wins, else jax.default_backend() != "cpu"), removing ~28% of the launches per chunk (1,244 → 892) on accelerators while CPU keeps the rolled loop, which measures ~8–9% faster there; (A.2) jax_delaunay_nn locates and interpolates the data grid and the 4N split points in ONE concatenated pass (walk once, Sibson once, slice at n_query), with delaunay_circumcircles_from hoisted through a new optional circumcircles= kwarg on sibson_mappings_weights_from_tables; (A.3) SIBSON_QUERY_CHUNK (default still 256) is overridable by PYAUTO_SIBSON_QUERY_CHUNK and documented as a memory guard on the (C, 3, 2) per-cavity intermediates rather than a speed knob; (A.4) docstrings for the loop structure and the launch-count reasoning. Everything is bit-identical: all 16 jax_delaunay_nn outputs (data and split halves) match a frozen main reference after each step and on the gated / env=1 / env=0 paths, and outputs are invariant at chunk 64 / 256 / 1024. Gates: pytest test_autoarray -q 1456 passed (4 new NumPy-only tests), black clean, jax_assertions/delaunay_nn.py and delaunay_nn_caps.py exit 0 locally, and an interleaved in-process paired A/B reads control-vs-feature 0.993 on CPU with the gate (0.927 with the unroll forced on) — no measurable CPU regression. Heart was YELLOW at ship time (score 70, no RED; "workspace validation not passing (5 failed, 2 timeout, cloud#34099198772)" plus the stale "release validation incomplete: no rehearsal for current source") — organism-scope, nothing in this branch is in the release chain.

API impact: none beyond one new optional kwarg — jax_delaunay_nn, jax_sibson, scipy_delaunay_nn, InterpolatorDelaunayNN and aa.mesh.DelaunayNN keep their signatures and outputs; a workspace grep finds no caller of sibson_mappings_weights_from_tables.

Still to land on this branch before merge: the A100 chunk sweep (512 / 1024 / 2048 / 4096 at vmap 16 with peak VRAM) that chooses the new SIBSON_QUERY_CHUNK default, plus the A100 witness leg — params→H prefix (regularization_matrix_prefix_s) ≤ 60 ms unbatched / ≤ 11 ms per call at vmap 16, judged on the prefix rather than the Tri+interp or H rows.

Next: /start_workspace — the autolens_workspace_test jax_assertions/delaunay_nn.py extension (single-pass parity, chunk invariance, env override) and the autolens_profiling breakdown stage / submits / results ship as separate workspace PRs and do not gate this merge. Phase B (cavity early exit) is a second library PR on this issue after the A100 numbers.

[Jammy2211]

Workspace PR opened — 2026-09-07

PRs

• delaunay_nn: assert single-pass parity, chunk invariance and the chunk env override autolens_workspace_test#307

Upstream PR

#533

Summary

The JAX assertion gate scripts/misc/jax_assertions/delaunay_nn.py timed the new jax_delaunay_nn path but asserted nothing about the two claims Phase A rests on. It now checks single-pass parity — the one concatenated data-grid + split-point Sibson pass reproduces two separate jax_sibson passes element-for-element on both halves, across all six outputs — chunk invariance at query_chunk 64/256/1024 via the kwarg, and the PYAUTO_SIBSON_QUERY_CHUNK env override, checked in a fresh subprocess because the library reads it once at import. All existing checks retained; +8.4 s runtime (direct run 29.0 s, harness PASS (31.0s)).

Session Notes

• Parity is compared NaN-aware (np.array_equal(..., equal_nan=True) on float tables): overflow and Watson degeneracy deliberately produce NaN weights, so a plain == comparison would pass on any pair of NaN tables and fail on a legitimately NaN one — i.e. vacuously green in exactly the regime the check exists for.
• The split points fed to the separate jax_sibson pass are taken from jax_delaunay_nn's own return (index 5), so both sides see identical inputs and the check cannot drift with the split-cross construction.
• The parent-side assert compares aa.mesh.DelaunayNN.query_chunk against the library's own SIBSON_QUERY_CHUNK rather than the literal 256, so the accelerator sweep can change the default without breaking this gate. The parent also asserts PYAUTO_SIBSON_QUERY_CHUNK is unset in its own environment — an inherited value would silence the override check.
• Library-first merge gate: perf: speed up JAX NNLS solver under vmap at production-fiducial size #307 must merge after sibson: cut DelaunayNN kernel launches — gated candidate unroll, single concatenated pass, chunk as memory guard #533. Against main autoarray the script fails (no env override, no concatenated pass).
• Heart is YELLOW on organism-scope workspace validation (cloud run 34099198772), unrelated; acknowledged before shipping.
• Still to come on this task: the autolens_profiling A100 results PR, shipped separately.

[Jammy2211]

A100 verification — 2026-09-08

Phase A re-measured on euclid-ral-gpu-2, all ten job units inside one 28-minute window (jobs 342321–342325, 342329, 342330): control = private checkout at the merge base d7c96762, feature = 9e9a7d50 at chunk 256 and the shipped e5d31d74 at the new 4096 default. HST / Hilbert-1500 / MGE-60 / ConstantSplit, fp64, dense inversion, --split-setup --vmap-batch 16. The shared /mnt/ral/jnightin/PyAuto install was not touched (private PyAutoArray checkouts prepended to PYTHONPATH; live CPU runs 342299/342311/342314 were unaffected).

Witness — half met

• Unbatched half (params→H ≤ 60 ms): MET. 143.90 → 28.21 ms at the shipped default (5.10×); 24.66 ms on the sweep's own 4096 row. Met at every chunk ≥ 512.
• Batched half (≤ 11 ms per call at vmap 16): NOT met. 24.32 → 16.44 ms (1.48×).

Shipped default (e5d31d74, chunk 4096) vs control

	control	shipped	Δ
params→H prefix, unbatched	143.903	28.212	−115.691 (5.10×)
params→H prefix, per call @vmap 16	24.317	16.435	−7.9 (1.48×)
Triangulation + interpolation (unbatched)	97.274	13.365	−83.909
Regularization matrix H (unbatched)	45.103	13.768	−31.335
Total step-by-step (unbatched)	196.357	80.716	−115.641

Whole-likelihood runtime cell

	control	feature @256	shipped @4096	Δ (control → shipped)
Full pipeline (single JIT)	201.320	136.082	76.322	−125.0 (2.64×)
vmap batch 16, per call	58.022	54.209	52.618	−5.4 (1.10×)

Chunk sweep (PYAUTO_SIBSON_QUERY_CHUNK)

chunk	params→H unbatched	per call @vmap 16	Sibson share @vmap 16	ConstantSplit assembly @vmap 16	peak sampled VRAM	pin
256 (9e9a7d50, 342322)	88.287	23.072	10.488	10.152	41,495 MiB	PASS
512 (342323_0)	52.205	18.243	8.309	10.014	41,503 MiB	PASS
1024 (342323_1)	35.324	19.567	7.062	12.557	41,503 MiB	PASS
2048 (342323_2)	27.218	19.127	6.602	10.058	41,503 MiB	PASS
4096 (342323_3)	24.664	16.451	6.398	10.028	41,503 MiB	PASS
4096 shipped (e5d31d74, 342329)	28.212	16.435	6.434	10.001	not sampled	PASS
control 256 (342321)	143.903	24.317	11.695	9.885	41,495 MiB	PASS

The 256 row is the same-chunk A/B — unroll + single concatenated pass alone give 143.903 → 88.287 ms unbatched (1.63×); the rest is the chunk.

VRAM note. The ~41.5 GiB plateau (50.7 % of 81,920 MiB) is identical at every chunk and on the control leg, and the GPU read 0 MiB with no other process at the start of every job — so it is this cell's vmap-16 dense inversion block, not the cavity intermediates. A 16× change of chunk moves it by 8 MiB. The pre-registered VRAM clause was therefore uninformative and was not used to break the tie; the decision rests on speed alone (4096 is 14.0 % faster per call at vmap 16 than 2048, 9.4 % faster unbatched). Guard arithmetic: (C, 3, 2) fp64 intermediates at ~25 kB per query per lane over 4096 × 16 lanes ≈ 1.6 GB, ~2 % of an 80 GB card. PYAUTO_SIBSON_QUERY_CHUNK stays as the escape hatch.

Pins

EXPECTED_LOG_EVIDENCE_HST = 29144.581944 passed unchanged on every leg and every chunk — both breakdown A/B legs, all four sweep chunks, the shipped re-measure, and pinned_drift: [] on all four runtime rows. Five different lax.map block sizes, one log evidence: the chunk is bit-neutral end-to-end.

The finding: the ConstantSplit assembly, not Sibson

The new --split-setup split-Sibson stage cuts the params→H interval in two. At vmap 16 and the shipped chunk:

piece	per call @vmap 16	comparison
Sibson (locate + both interpolation passes)	~6.4 ms	Delaunay's whole interpolator prefix: 5.07 ms
H, ConstantSplit assembly	~10.0 ms	Delaunay's H row: 0.07 ms

The assembly row reads 9.9–10.2 ms at every chunk and on the control — untouched by Phase A and untouchable by any change to the Sibson chunking. After Phase A roughly 60 % of the per-call params→H prefix is the 33-wide split-stencil regularization assembly (reg_split_from plus the block-diagonal build), 1.6× the entire Sibson share and 143× Delaunay's equivalent row. That is why the ≤ 11 ms batched half could not be met: ~10 ms of it was never Sibson's to give.

(The "Split-point Sibson" row's 35.4 → ~0 ms collapse is an attribution shift, not a saving — before the change a prefix stopping at step 6 never asked for the split points, so XLA eliminated the second Sibson pass out of it. Only the params→H prefix says whether work went away, and it fell 5.10×.)

Phase B — decision pending the human

Numbers only, from the note:

• Phase B (cavity fori_loop → lax.while_loop early exit) targets the Sibson share only; its headroom at vmap 16 is ~6.4 ms per call, and the plan's own ~20 % estimate puts it at ≈ 1.3 ms per call at vmap 16 (~2.5 % of a 52.6 ms production evaluation).
• A ConstantSplit-assembly rework is ~8× larger on the batched number (~10.0 ms, ~19 % of a production evaluation) but has no prompt and would need its own investigation.
• Phase C (loop-free k-ring cavity) is now the smallest of the three on the batched number and should be re-costed before it is planned.

Phase B is cheaper and lower-risk (one loop, stop_gradient boundary already scoped, no numerical reordering) and is the only one with a written plan. Which goes first is a call about risk appetite, not about which is larger — left to the human.

Links

• Profiling PR (results + the new breakdown stage): A100 re-measure: DelaunayNN Sibson launch latency, Phase A (PyAutoArray#533) autolens_profiling#227
• Full note: results/notes/delaunay_nn_launch_latency.md in that PR
• Library PR: sibson: cut DelaunayNN kernel launches — gated candidate unroll, single concatenated pass, chunk as memory guard #533
• Workspace PR: delaunay_nn: assert single-pass parity, chunk invariance and the chunk env override autolens_workspace_test#307

[Jammy2211]

Shipped

Phase A of this issue is merged.

PRs

• PyAutoArray — sibson: cut DelaunayNN kernel launches — gated candidate unroll, single concatenated pass, chunk as memory guard #533 (pending-release, merged 2026-09-08)
• autolens_workspace_test — delaunay_nn: assert single-pass parity, chunk invariance and the chunk env override autolens_workspace_test#307
• autolens_profiling — A100 re-measure: DelaunayNN Sibson launch latency, Phase A (PyAutoArray#533) autolens_profiling#227

What shipped

1. Gated candidate unroll — the 3-trip candidate-edge fori_loop in the Sibson cavity walk is unrolled at trace time when _sibson_unroll_candidates() is true (PYAUTO_SIBSON_UNROLL_CANDIDATES wins, else jax.default_backend() != "cpu"). Same add_candidate order for edges 0/1/2, so the same cavity insertion order, stencil column order and fp summation order. ~28 % fewer kernel launches per chunk (1,244 → 892). CPU stays rolled.
2. Single concatenated pass — jax_delaunay_nn locates and interpolates the data grid and the 6,000 ConstantSplit cross points in one Sibson pass, computing the circumcircles once (the feat: JAX Delaunay walk — early-exit while_loop, chunk only the seed argmin #531 pattern, one boundary further in).
3. SIBSON_QUERY_CHUNK is a memory guard again, default 4096, with PYAUTO_SIBSON_QUERY_CHUNK as an import-time override.
4. Split-Sibson breakdown stage (autolens_profiling) and an extended delaunay_nn.py jax_assertions gate (autolens_workspace_test) asserting single-pass parity, chunk invariance and the chunk env override.

A100 numbers (euclid-ral-gpu-2, ten job units in one 28-minute window; full note: autolens_profiling/results/notes/delaunay_nn_launch_latency.md)

	control	shipped
params→H prefix, unbatched	143.90 ms	28.21 ms (5.10×)
params→H prefix, per call @vmap 16	24.32 ms	16.44 ms (1.48×)
whole likelihood, single JIT	201.32 ms	76.32 ms (2.64×)
whole likelihood, per call @vmap 16	58.02 ms	52.62 ms (1.10×)

Witness — half met, and the note says why

• Unbatched half (params→H ≤ 60 ms): met at every chunk ≥ 512 (28.21 ms shipped).
• Batched half (≤ 11 ms per call at vmap 16): not met — 16.44 ms. The new split-Sibson stage attributes it: the ConstantSplit regularization assembly is ~10.0 ms per call, unchanged at every chunk and on the control (143× barycentric Delaunay's 0.07 ms H row, ~19 % of the 52.6 ms batched likelihood). ~60 % of the remaining per-call prefix was never Sibson's to give.

Defaults, gates and pins

• Default chunk 4096. Pre-registered rule A.3 had a VRAM clause; it turned out vacuous — the ~41.5 GiB plateau is identical at every chunk and on the control (it is the vmap-16 dense inversion block, not the cavity intermediates; a 16× chunk change moves it by 8 MiB). With that leg dead the decision rests on speed alone: 4096 is 14.0 % faster per call and 9.4 % faster unbatched than 2048. The constant's comment carries the table, the rule and the arithmetic.
• CPU is gated off the unroll: the paired CPU A/B ratio was 0.993 (rolled is marginally faster there).
• EXPECTED_LOG_EVIDENCE_HST = 29144.581944 held exactly on all ten legs and every chunk; pinned_drift: [] on all four runtime rows.
• 1,456 PyAutoArray unit tests pass.

Not re-filed

• Phase B (cavity early exit) is deliberately not filed. Its headroom is the ~6.4 ms Sibson share, worth ≈ 1.3 ms per call at vmap 16 (~2.5 % of a production evaluation) against the assembly's 10.0 ms. The follow-up points at the assembly instead: PyAutoMind/draft/feature/autoarray/delaunay_nn_constant_split_assembly.md.
• Phase C (loop-free k-ring cavity) stays unfiled — it is the smallest of the three levers on the batched number and should be re-costed against these figures before it is planned.

Record: PyAutoMind/complete/2026/09/delaunay-nn-launch-latency.md.