WP5: a segmented near field, because float32 cannot afford a scatter - #63
Merged
Merged
Conversation
This was referenced Sep 4, 2026
…catter
Splitting the near field into its parts at N=1e5 on an A100 says the arithmetic
was never the problem:
gather+arithmetic the scatter alone
float64 22.60 ms 5.99 ms
float32 20.47 ms 258.48 ms
float32 is *faster* than float64 at gather and arithmetic, as it should be -- a
dense kernel of the same shape is 0.19 vs 0.39 ms. The whole penalty is one
.at[].add(): the scatter indices repeat ~62x on average, so float32 lowers to a
contended atomicAdd where float64 takes another path.
Leaves tile [0, n) disjointly, so summing each target leaf's directed pairs into
(L, ml, d) and then *placing* it is a permutation -- no atomics anywhere. Five
strategies measured; the two that survive are the two that ship:
float64 float32
halved pairs, two scatters 28.58 ms 278.94 ms
directed pairs, segment_sum, chunked 56.15 ms 32.48 ms
Neither wins in both columns, so accumulate= names the choice: "scatter"
(default), "segment", or "auto" (segment for float32). Rejected on the way: the
sorted-index hint (helps float32 6.4x, costs float64 4.5x) and the per-leaf
gather (worse in both -- neighbour counts are skewed mean 621 / max 3124, so
padding wastes 5x).
**The default is deliberately not the fastest forward.** "segment" is 4.0x
faster forward in float32 at N=1e4 (11.64 -> 2.93 ms) and slower under reverse
mode (fwd+grad 2392 -> 3172 ms at 1e5), because the transpose of a gather is a
scatter: reverse mode reintroduces exactly what the forward removed, and the
directed list evaluates each kernel twice besides. So svgd_phi_from_topology
keeps the gradient-safe path and only the forward-only entry points ask for
"auto". Nothing regresses; the win is available where it is real.
float64 is untouched: 5.72 ms forward at N=1e4 and 112.67 at 1e5, as before.
This is the baseline the tree update is measured against, so a weak one
flatters the tree -- and the old one was weak in two ways at once. It built an
(n, n, d) tensor of per-pair terms, which is d times the memory and runs at
elementwise rather than GEMM throughput, and it simply could not be evaluated
beyond N ~ 2e4, which is why every crossover claim so far rests on N <= 2e4.
The sum contracts:
phi_i = [ (K S)_i + ( x_i (K 1)_i - (K X)_i ) / h^2 ] / N
with K_ij = exp(-|x_i - x_j|^2 / 2h^2). One kernel matrix and two matmuls. An
optional block_size caps the kernel matrix at (block_size, n) via lax.map, so
large N costs bounded memory instead of being impossible.
Verified against the materialised form at n = 7, 64, 300 and 1001 for block
sizes None, 1, 13, 64 and 4096: max relative difference 5.0e-15.
TB, 2026-09-04: efficient implementations, not only working ones -- and that
has to include the thing we compare against.
TobiBu
force-pushed
the
perf/svgd-gather-accumulation
branch
from
September 5, 2026 19:36
5631c21 to
f8afdeb
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 #62. Its diff also shows #61's three tree-core commits until #61 merges; they
are not new here.
Localising the 5× float32 penalty that survived #62. Timing the near field's parts separately
at N = 10⁵ on an A100:
A dense kernel of the same shape with no indexing is 0.19 / 0.39 ms — the expected 2:1. So the
arithmetic is fine in both and the whole float32 penalty is one
phi.at[slots].add(...):the indices repeat ~62× on average, so float32 lowers to a contended
atomicAddwhile float64takes another path. This is a scatter problem, not a float32 problem — float64 pays 21 % of
its near field for the same line.
So
accumulate=gets a second option: leaves tile[0, n)disjointly, so summing each targetleaf's directed pairs into
(L, ml, d)and placing it is a permutation, and no atomic isneeded anywhere. float32 forward 586 → 321 ms at N = 10⁵.
The default stays
"scatter", deliberately. The segmented path is slower under reversemode — 2392 → 3172 ms at N = 10⁵ — because the transpose of a gather is a scatter: autodiff
reintroduces exactly the operation the forward removed, and doubles the arithmetic besides.
"auto"picks by dtype for forward-only callers; a caller that differentiates should leave thedefault alone. The table is in the function's docstring so the trade cannot be lost.
Also here:
exact_phicontracts to two matmuls instead of an(n, n, d)tensor, so thereference baseline reaches the same N the tree does instead of stopping at 2·10⁴.
Evidence:
reports/YGGDRAX_perf_report.md§5.2, §5.3.🤖 Generated with Claude Code