perf(nearfield,m2l): fold the self-leaf block into the leaf-pair kernel; contention-free M2L chunk reduction - #339
Merged
Merged
Conversation
…-self' into perf/small-leaves
…ernel One extra unconditional pass per program over the program's own leaf, diagonal masked before the rsqrt, gated to chunk 0 on the source-chunk grid, replaces the lax.scan over leaves (_self_contributions) on the forward prepacked lane. JACCPOT_NEARFIELD_LEAFPAIR_FOLD_SELF=0 restores the scan; the cvjp lane keeps self_acc + cvjp(...) so the gradient path is byte-identical. Kernel name gains _self so the two variants never alias in the compilation cache. Tests: fold == pairs + _self_contributions to fp32 summation order on both grids, subtile lane indexing, softening 0 (no NaN, potential lane), wide accumulator, default-off byte parity, and a Triton (sm_80) run registered in the GPU gate. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
_chunk_segment_scatter_add ran once per 4096-pair chunk and its two scatters were pathological for XLA's atomic-add lowering: segment_sum sent every pair of a target to the same group row, and the final .at[].add sent every non-head slot (~3800 of 4096) to node 0 with a zero value. At leaf 64 (992k far pairs, 243 chunks/step) that was 169 ms of a 410 ms step, 686 us per launch for a 4096 x 25 reduction. Now a segmented associative_scan reduces within the sorted segments with no scatter, and non-head slots carry an out-of-bounds sink index dropped by mode='drop', so the one remaining scatter has unique in-bounds indices and no atomics. Pinned to np.add.at on random chunks (padded tail, all-same-target, nothing valid, degenerate width), node-0 untouched unless targeted, deterministic. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…R by target) One program per target node loops over its far-pair segment of a target-sorted source list and builds rotate -> z-translate -> rotate-back on chip from the two alignment angles: Dz as diag(cos |m| t) + A sin(|m| t) with a constant antisymmetric pattern, the constant B stack per degree, and a separable z-core Z = Zsf * outer(rinv^(n+1), rinv^k). Nothing per pair touches HBM except one (Cp,) multipole row; the program owns its output row, so the chunked lane's per-chunk argsort + segment_sum scatter disappear. Forward only. Interpret-mode parity vs m2l_rot_scale_real_batch: < 1e-10 (fp64) for orders 2..6, < 3e-4 (fp32); random CSR with empty targets and a padded -1 tail, active_pair_count truncation, on-axis deltas, and the wrapper under jit. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…as no n-ary stack Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…icrobench JACCPOT_STATIC_STRICT_FUSED_M2L_CSR=1 (Ampere+, or JACCPOT_M2L_CSR_INTERPRET=1 for CPU parity) routes _solidfmm_downward_accumulate_from_multipoles' flat real lanes to m2l_real_csr_pallas: one program per target owns its local row, so the chunked scan and its per-chunk scatter are gone. Wiring test counts the kernel entry and pins the force to the chunked lane (rel < 2e-5 at N=3000, p=3). bench/m2l_csr_microbench.py: ns/pair vs the pure-JAX chunked lane (degree-batched off/on) at 1M and 6M pairs, orders 4 and 6 -- the plan's G2a gate. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…k, degree-only z-core The first cut held two 128x32 one-hot pack/unpack matrices and a Cp x Cp z-core in registers per program and ran at ~93 ns/pair on a shared A100 (non-record), no faster than the pure-JAX lane. The multipole table is now packed into the centred (Bp, Wp) layout once in the wrapper and the output unpacked once; the z-translation preserves m, so its operator is a Bp x Bp degree matrix per column. Half the per-pair MACs, a third of the constants; 4 warps per program. Microbench: iterate the well-separation rejection (a coincident pair made the reference NaN), assert finiteness. CSR sm_80 tests registered in the GPU gate. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…n the slow list (28 s on CPU) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…(the knob is trace-time) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ixes, per-step table, capacity rules Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
… sm_80 tests run in the ordinary suite Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ernel One extra unconditional pass per program over the program's own leaf, diagonal masked before the rsqrt, gated to chunk 0 on the source-chunk grid, replaces the lax.scan over leaves (_self_contributions) on the forward prepacked lane. JACCPOT_NEARFIELD_LEAFPAIR_FOLD_SELF=0 restores the scan; the cvjp lane keeps self_acc + cvjp(...) so the gradient path is byte-identical. Kernel name gains _self so the two variants never alias in the compilation cache. Tests: fold == pairs + _self_contributions to fp32 summation order on both grids, subtile lane indexing, softening 0 (no NaN, potential lane), wide accumulator, default-off byte parity, and a Triton (sm_80) run registered in the GPU gate. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
_chunk_segment_scatter_add ran once per 4096-pair chunk and its two scatters were pathological for XLA's atomic-add lowering: segment_sum sent every pair of a target to the same group row, and the final .at[].add sent every non-head slot (~3800 of 4096) to node 0 with a zero value. At leaf 64 (992k far pairs, 243 chunks/step) that was 169 ms of a 410 ms step, 686 us per launch for a 4096 x 25 reduction. Now a segmented associative_scan reduces within the sorted segments with no scatter, and non-head slots carry an out-of-bounds sink index dropped by mode='drop', so the one remaining scatter has unique in-bounds indices and no atomics. Pinned to np.add.at on random chunks (padded tail, all-same-target, nothing valid, degenerate width), node-0 untouched unless targeted, deterministic. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
for more information, see https://pre-commit.ci
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
for more information, see https://pre-commit.ci
CI on #340: test_type_annotation_guard flagged the inner `one`/`body` closures; pydoclint wanted Raises sections on m2l_real_csr_tables / m2l_real_csr_pallas and full docstrings on _bapply/_dz; black reformatted four files. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
feat(m2l): target-tiled CSR Pallas M2L kernel with on-chip rotations (opt-in lane)
TobiBu
added this pull request to stack #342
September 10, 2026 19:34
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.
Two independent single-GPU fused-lane changes from the small-leaves plan (
docs/small_leaves_2026-09.mdlands with the follow-up CSR M2L PR).1. Self-leaf block folded into the leaf-pair Pallas kernel (
nearfield_fused_leaf.py,include_self; flagJACCPOT_NEARFIELD_LEAFPAIR_FOLD_SELF, default on). One extra unconditional pass per program over its own leaf with the diagonal masked before thersqrt, gated to chunk 0 on the source-chunk grid; replaces thelax.scanover leaves (_self_contributions) on the forward prepacked lane. The cvjp lane keepsself_acc + cvjp(...), so the gradient path is byte-identical. Kernel name gains_self. Time-neutral (the scan was already batched by #334): leaf-pair kernel 48.5 vs 48.2 ms at leaf 128, forces identical (aggL2 1.2011e-03 at leaf 64 either way).2.
_chunk_segment_scatter_addwithout duplicate-address atomics. The per-chunk M2L reduction sent every pair of a target to the samesegment_sumrow and every non-head slot (~3800 of 4096) to node 0 as a zero add; XLA's atomic scatter serialised both. At N=200k Plummer, leaf 64, theta 0.6 that was ONE fusion at 686 us x 247 launches = 169 ms of a 410 ms step. Now a segmentedassociative_scanreduces within the sorted segments with no scatter, and non-head slots carry an out-of-bounds index dropped bymode="drop", so the one remaining scatter has unique in-bounds indices.Per step (
strict_run_v2, N=200k, p=4, theta 0.6, idle A100):Tests: fold == pairs +
_self_contributionsto fp32 summation order on both grids, subtile lane indexing, softening 0, wide accumulator, default-off byte parity, Triton (sm_80) run; scatter pinned tonp.add.aton random chunks (padded tail, all-same-target, nothing valid), node 0 untouched unless targeted, deterministic, run on CPU and A100. Near-field suite 54/54 and strict-lane integration tests 17/17 on A100. The new sm_80 tests are not in_MUST_RUN_SM80: its registry test is bound to one module.🤖 Generated with Claude Code