feat(m2l): target-tiled CSR Pallas M2L kernel with on-chip rotations (opt-in lane) - #340
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>
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>
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 #339 (base =
perf/small-leaves); retarget tomainonce #339 merges. Record:docs/small_leaves_2026-09.md.The kernel (
jaccpot/pallas/m2l_real_csr.py): one Pallas program per target node loops over its far-pair segment of a target-sorted source list (CSR built in the wrapper with one stable argsort, padding andactive_pair_counthonoured) and builds rotate -> z-translate -> rotate-back on chip from the two alignment angles, in the centred (degree, m) layout:Dz(t) = diag(cos|m|t) + A sin(|m|t)with a constant antisymmetric pattern, the constantB_Ustack per degree, and a degree-only z-core ((n+k)!matrix,rinv^(n+1)/rinv^kasexp(deg*log rinv)) because the z-translation preserves m. Nothing per pair touches HBM except one multipole row; the program owns its output row, so the chunked lane's per-chunk argsort +segment_sumscatter are gone. Forward only; the pure-JAX lane stays the differentiable path.Wiring: an opt-in fifth flat lane in
_solidfmm_downward_accumulate_from_multipolesbehindJACCPOT_STATIC_STRICT_FUSED_M2L_CSR=1(Ampere+;JACCPOT_M2L_CSR_INTERPRET=1for CPU parity). Default off, so nothing changes without the flag.Measured (idle A100,
bench/m2l_csr_microbench.py):In
strict_run_v2at N=200k Plummer, theta 0.6, p=4 the kernel costs 12.9 ms per step at leaf 64 where the pure-JAX M2L cost ~360; per step 236.5 -> 176.7 (leaf 64), 145.2 -> 122.3 (128), 649 -> 509 (32), 126.4 -> 120.3 (256), on top of #339. Forces identical to the chunked lane (aggL2 to 4 digits at every leaf; per-order sweep p=2..6 monotone and equal to 4 digits), strict-lane truncation check < 7e-4, 300-step dynamics with the lane on vs off: identical dKE/dLz to 4 digits, round-off-level divergence only.Tests: interpret parity vs
m2l_rot_scale_real_batch< 1e-10 (fp64, orders 2-6) and < 3e-4 (fp32); random CSR with empty targets and a padded tail;active_pair_counttruncation; on-axis deltas (rho = 0); wrapper under jit; axis contracts; pack/unpack round trip; Triton run for orders 2/4/6 on A100; a wiring test that counts the kernel entry and pins the force (rel < 2e-5); strict integration tests 17/17 on A100. The wiring test (28 s on CPU) is on the slow list.The per-step leaf optimum still sits at leaf 128-256: with M2L at 13 ms the remaining ~100 ms of a leaf-64 step is the traced dual-tree walk and list compaction, which scale with the leaf count and the traced pair queue. That is the next lever, and it is outside this PR.
🤖 Generated with Claude Code