perf(nearfield): source-chunk grid axis for the leafpair kernel; batched self-P2P scan - #334
Merged
Merged
Conversation
…hed self-P2P scan Profiling one fused eval at N=200k, leaf 256, theta 1.0 on an A100 (93 % device busy, so not launch-bound) put 40 of 47 ms in nearfield_leafpair and 7.5 ms in the intra-leaf self term. The kernel's wall time was bounded below by ONE single-warp program: the grid is (leaf, 32-target subtile) and each program loops over its leaf's whole neighbour row, and in a centrally concentrated distribution the longest row is num_leaves-1 -- a halo leaf so extended that the mutual MAC makes it near to every leaf -- so one warp serially sums all N sources at ~100 ns per lane-step (19 ms at 200k, 38 ms at 400k, theta- and order-independent). Truncating rows 781->150 slots removed 18 ms while dropping 21 % of the entries; a proportional model predicts 8. nearfield_leafpair_pallas gains source_chunk: the grid becomes (leaf, subtile, chunk), each program sums <= chunk source slots into a partial (kept in the wide dtype when the wide accumulator is on) and the partials are reduced afterwards. JACCPOT_NEARFIELD_LEAFPAIR_SOURCE_CHUNK (default 64; 0 = the old single pass, bit-for-bit) selects it. Measured on a shared card: theta 1.0 46 -> 31 ms, theta 0.8 62 -> 46 ms, theta 0.4 124 -> 115 ms; aggL2 identical to 4 s.f. at theta 0.8/1.0 and 3.3065e-5 -> 3.2528e-5 at theta 0.4 (summation order, at the fp32 accumulation floor). _self_contributions scanned one leaf per step -- 782 x 5 launches of ~2 us at leaf 256, 1563 x 5 at leaf 128 where the GPU sat 41 % idle. It now vmaps JACCPOT_NEARFIELD_SELF_BATCH leaves per step (default 32; 1 = old scan), which keeps the remat argument per step and cuts the launch count 32x. Forces unchanged to 1e-7. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
TobiBu
force-pushed
the
perf/nearfield-row-split-batched-self
branch
from
September 6, 2026 08:42
3f4ad18 to
9b67434
Compare
tests/unit/test_shared_env_switches.py keeps the set of raw os.environ readers outside runtime/ closed, and the batched self-P2P scan added one. Route it through the sanctioned reader instead: env_int reads at call time and falls back to the default on a malformed value, where the raw int(os.environ.get(...) or "1") silently turned a typo into batch=1. Fixes test-full (unit), test-smoke (3.11/3.12, unit) and test-runtime-typecheck, which all run that suite. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
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.
Why
Profiling one fused eval at N=200k, leaf 256, theta 1.0 (A100, 93 % device busy): 40 of 47 ms in
nearfield_leafpair_t32_s1024_w256, 7.5 ms in the intra-leaf self term (782 x 5 launches of ~2 us; 15 ms at leaf 128 with the GPU 41 % idle). The kernel's wall time was bounded below by ONE single-warp program: the grid is (leaf, subtile) and each program loops its leaf's whole neighbour row; in a Plummer halo the longest row isnum_leaves - 1, so one warp serially sums all N sources at ~100 ns per lane-step (19 ms at 200k, 38 at 400k, theta- and order-independent). Truncating rows 781 -> 150 slots removed 18 ms while dropping 21 % of entries.What
nearfield_leafpair_pallas(source_chunk=): grid (leaf, subtile, chunk), each program sums <= chunk slots into a partial (kept in the wide dtype when the wide accumulator is on), partials reduced afterwards.JACCPOT_NEARFIELD_LEAFPAIR_SOURCE_CHUNK(default 64; 0 = the old single pass, bit-for-bit)._self_contributionsvmapsJACCPOT_NEARFIELD_SELF_BATCH(default 32; 1 = old) leaves per scan step.Measured (shared card, so indicative; quiet numbers to follow)
theta 1.0: 46 -> 31 ms; theta 0.8: 62 -> 46 ms; theta 0.4: 124 -> 115 ms. aggL2 identical to 4 s.f. at theta 0.8/1.0; at theta 0.4 3.3065e-5 -> 3.2528e-5 (summation order, at the fp32 accumulation floor).
tests/unit/operators/test_pallas_nearfield_fused.pyandtests/unit/test_custom_vjp_parity.pygreen.🤖 Generated with Claude Code
Update 2026-09-07. Merged
main(which now carries #335's near-field axis contracts) and fixed the one CI failure: the batched self-P2P scan readJACCPOT_NEARFIELD_SELF_BATCHwith a rawos.environ.get, whichtests/unit/test_shared_env_switches.pyrejects — the set of raw env readers outsideruntime/is closed by that test. It now goes throughjaccpot._env.env_int(reads at call time, falls back to the default on a malformed value, whereint(os.environ.get(...) or "1")silently turned a typo intobatch=1). That one test was the whole oftest-full (unit),test-smoke (3.11/3.12, unit)andtest-runtime-typecheck.Verified on this tree after the merge: the env-switch suite passes, and
tests/unit/nearfield,tests/unit/pallas/test_fused_leaf_axis_contracts.py,tests/unit/operators/test_pallas_nearfield_fused.pyandtests/unit/test_custom_vjp_parity.pyall pass on CPU — so the new axis contracts accept the chunked-grid and batched-scan shapes. black/isort/flake8/pydoclint clean.