perf(tree): node depths by pointer doubling — the KD-tree traversal was dispatch-bound

[TobiBu]

The finding

The KD-tree traversal was 80× the radix backend's at N = 10⁶ (the worst number in Yggdrax I), and it was not in the walk.

The KD-tree's dual-tree walk runs the same number of wavefront generations as radix's — 16 against 16 at N = 10⁴, 22 against 24 at 10⁵ — at the same mean occupancy and the same wavefront capacity (138 150 against 140 634). On device it is faster: 61.6 ms against radix's 64.6 ms at N = 10⁵.

All of the excess was get_node_levels, called on every traversal by _result_to_interactions. When a topology carries a node_level field it hands it back; when it does not, it fell back to a Python loop of num_nodes - 1 max-relaxation rounds. Radix and octree carry the field; the KD-tree's topology does not, so it was the only backend taking the fallback.

That is O(num_nodes) eagerly dispatched primitives:

N	KD-tree nodes	loop rounds	get_node_levels	per round
10³	31	30	46.9 ms	1.56 ms
10⁴	511	510	799.0 ms	1.57 ms
10⁵	4 095	4 094	6 385.3 ms	1.56 ms

cProfile of one warmed traversal at N = 10⁵: 32 896 dispatched primitives for the KD-tree against 134 for radix. Same disease as the octree build (d775f87), in a different place — not more work, just more launches.

The change

Depths now come from a device-side pointer-doubling lax.while_loop: each node keeps a shortcut into its ancestor chain and the distance covered, a round doubles every shortcut's reach, so the walk to the root converges in O(log depth) rounds inside one dispatched computation whatever the tree's size.

_interactions_impl._compute_node_depths already did exactly this inside the jitted walk; it becomes a thin alias for the shared helper, so the walk's depths and the interaction list's levels have one implementation and cannot drift apart.

Result

bench/differentiability/scaling.py --sizes 1000 10000 100000 1000000, leaf 64, θ = 0.6, one A100, jax 0.9.0. Traversal, ms:

N	radix	kdtree before	after	kdtree/radix	speed-up
10³	26.11	74.42	26.22	1.00×	2.8×
10⁴	41.86	851.62	46.27	1.11×	18.4×
10⁵	82.77	6 703.50	80.74	0.98×	83.0×
10⁶	710.27	107 471.15	491.05	0.69×	219×

At N = 10⁶ the KD-tree traversal is now the fastest of the three backends. That is not an artefact: it emits 3 320 390 far and 2 996 204 near entries against radix's 3 842 956 and 6 154 042 — barely half the near-field work, because median splits give tighter nodes.

Nothing about the tree moved. Node counts, far-list lengths and near-list lengths are identical to the pre-change run at every N and for every backend.

Correctness

Equal to the relaxation it replaces element for element: heap-shaped trees, single chains (the worst case for depth) and random forests with unattached nodes, at 18 sizes from 1 to 257; and on all three backends' real topologies, where the radix and octree results also agree with those topologies' own node_level field. Pinned in tests/unit/test_node_levels.py (43 cases), which tests against the definition — distance to the root along the parent chain — not against either implementation.

get_node_levels is also used by the level-major geometry path and by grouped_interactions, which get the same fix for free.

Checks

• pytest tests/unit tests/applications — 689 passed on a single CPU device (what CI runs).
• On a pinned GPU, 688/689: test_svgd_pallas_nearfield.py::test_the_twin_is_chunk_invariant fails identically on main (same digit, 1.130039386627306e-16). It asserts exact float64 equality across a chunked reduction, so it cannot hold on a GPU as written — a tolerance is the fix, and it is not this PR's file.
• tests/unit/test_kd_tree_parity.py, test_backend_conformance.py, test_interactions_backend_parity.py — 21 passed.
• black, isort clean; basedpyright 0 errors on both touched files.

Report: reports/YGGDRAX_kdtree_traversal_report.md in the programme repo.

🤖 Generated with Claude Code