WP2: a cheaper near field, a cheaper partition assembly, and a metric that was lying - #62
Merged
Merged
Conversation
This was referenced Sep 4, 2026
TobiBu
force-pushed
the
perf/svgd-cheaper-gradient
branch
from
September 5, 2026 19:29
f571f38 to
fd2eff9
Compare
Three changes to the near field, which Gate 0 says is where the Stein update's time actually is (the far field is 0.45 % of the pair terms). **One evaluation per unordered pair.** The walk emits every near leaf pair in both directions and the old accumulation evaluated both, although the kernel value is shared: the contribution to i from j is k*s_j + k*(x_i - x_j)/h^2 and to j from i is k*s_i - k*(x_i - x_j)/h^2, one exp for both. The partition now keeps one entry per unordered pair and the accumulation scatters both directions from one evaluation. num_near_leaf_pairs on the topology still reports the *directed* count so the bench numbers stay comparable across work packages. Symmetry of the walk's near list is a precondition of the old code too -- without it a particle in B never receives from A -- so it is now checked and raised on rather than assumed. **Rematerialised chunks.** The (pairs, ml, ml, d) tensor was built whole and stored by reverse mode, which is why the gradient costs 13x the forward at N = 1e4 (Gate 0, float64). It is now a lax.scan over chunks under jax.checkpoint: peak memory is one chunk, capped at 64 MB by default, and the backward pass recomputes rather than reads. chunk_pairs is a static parameter if the default is wrong for a problem. **The far-entry expansion is vectorised.** One np.arange per far pair in a Python list comprehension was 136 ms of a 236 ms build at N = 2e4 -- more than the entire device half -- and grew superlinearly. cumsum/repeat gives the same array; checked against the loop on 500 random cases including the empty one. Verified: at theta = 0 the update still equals the exact update to 3e-16, and the answer is invariant to chunk_pairs (1, 7, and one-chunk agree to 9e-17).
…umber The 64 MiB default made the forward pass 4.6x slower at N = 1e4 (10.12 ms against 2.19 ms at 256 MiB) and in exchange produced the best-looking forward-to-gradient ratio in the sweep, 1.72. That ratio is Gate 2's metric, and it was being met by making the denominator worse -- the absolute forward-plus-gradient cost at 64 MiB is 17.46 ms against 9.97 ms at 256 MiB. 256 MiB is the measured optimum: best forward, near-best total, and it still bounds memory, which one chunk does not (4996 MiB by N = 3e4). The table is in the constant's docstring so the next person does not have to re-derive it.
M is by far the largest array in the partition -- 89,555,008 entries at N = 1e5 -- and building it on the host cost 1074 ms of a 1438 ms build, nearly all of it allocating ~2 GB of numpy and copying it back. Gate 2 made that the largest single cost anywhere in this plan. The expansion now runs on device. Only two scalars cross to the host: the kept far-pair count and M, both needed to give the device arrays a shape. Everything else -- the tag filter, the cumsum, the two repeats -- stays where the data is. The leaf blocks and the near-pair rows stay on the host: they are O(L) and O(P) and were never the problem.
Assembling M on device is a 3.6x win on the whole build at N = 1e5 (1456 -> 409 ms, its host half 1091 -> 48 ms) and a 17 % *loss* at N = 1e4 (83 -> 96 ms), because the device path pays ~10 eager dispatches and two synchronisations whatever the size while the host path pays per entry. So neither is right everywhere. The path is chosen by the actual entry count -- which costs nothing to know, since node_ranges is already on the host for the leaf blocks -- with the threshold at 2^21 entries, where the two costs cross. Both branches produce the same arrays.
TobiBu
force-pushed
the
perf/svgd-cheaper-gradient
branch
from
September 5, 2026 19:36
fd2eff9 to
27f3c97
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #60.
Four changes to the Stein update's cost under differentiation:
both directions and the old accumulation evaluated both, though the kernel value is shared.
Symmetry of the near list was already a precondition — without it a particle in B never
receives from A — so it is now checked rather than assumed.
(pairs, ml, ml, d)tensor was built whole and stored byreverse mode. It is now a
lax.scanunderjax.checkpoint, peak memory one chunk.np.arange-per-pair loop: 22.8 → 8.3 mson the host half at N = 10⁴.
N = 10⁵ — with the crossover measured (
_DEVICE_FAR_EXPANSION_MIN_ENTRIES), because below~10⁶ entries the dispatch costs more than the copy it avoids.
The chunk default is 256 MiB, and the story of how it got there matters. 64 MiB — what the
plan specified — produced the best forward-to-gradient ratio in the sweep (1.72) entirely by
making the forward 4.6× slower, at nearly double the absolute cost. That is the same failure
mode as [D-016]: a ratio improved by degrading its denominator. The table is in the constant's
docstring.
Ratio at N = 10⁴ falls 12.79 → 2.80 and absolute forward-plus-gradient falls 33.0 → 14.6 ms.
The forward regressed 2.58 → 5.20 ms, which is the honest cost of the scan, and is why later
PRs in this stack revisit it.
Evidence:
reports/YGGDRAX_perf_report.md§5, §5.1, §5.4.🤖 Generated with Claude Code