Mutual FMM: two deferred enhancements (on-chip-rotation Pallas M2L, per-rung multipoles)

[TobiBu]

Two enhancements to jaccpot/mutual that were deliberately not shipped in #64.
Both are recorded here so the reasoning does not have to be rediscovered — neither
is a defect, and neither should be picked up without first re-testing the premise
that made it not worth doing.

Background: docs/momentum_conserving_fmm.md.

---

1. An on-chip-rotation Pallas M2L kernel

Status today. backend="pallas" routes the near field to
pallas/nearfield_mutual.py (2.2–3.6× forward, 3.1–4.1× reverse on an A100) but
routes the far field to pure JAX, because both existing Pallas M2L shapes
measured slower:

lane	fwd (N=10⁴)	rev (N=10⁴)	fwd (N=10⁵)
pure JAX	17.11 ms	29.64 ms	75.16 ms
Pallas fused	20.03 ms (0.85×)	48.70 ms (0.61×)	116.33 ms (0.65×)
Pallas z-core sandwich	20.28 ms (0.84×)	34.81 ms (0.85×)	129.66 ms (0.58×)

Both lanes are wired, differentiable and CI-covered in interpret mode;
JACCPOT_MUTUAL_M2L=fused|zcore|jax forces one on hardware.

Why they lose, and what would not fix it. The fused kernel takes the world↔z
rotations as explicit (pairs, Bp, mdp, mdp) arrays — at order 4 in float64 that
is 32 KB per pair against the sandwich's 0.39 KB, an 82× amplification (5× of
it pure power-of-two padding, 5×9×9 → 8×16×16). The reverse is worse (0.61×)
because its VJP writes bto_bar/bfr_bar, the same 32 KB per pair back out.

Two facts rule out the obvious explanations:

• At N=10⁴ the two Pallas shapes are within 1% of each other on the forward.
  They could not be, if the rotation vmaps dominated — so the premise the fused
  lane was built on ("fusing the rotations away is the cheapest win") was simply
  wrong.
• The regression deepens at N=10⁵ (0.65×/0.58×), exactly where the M2L grows
  from 20.5% to 80.5% of the far field. So "it needs a bigger problem to
  amortise" is also ruled out — the two sizes together close that door.

The actual enhancement. A kernel that takes deltas and builds the rotations
on chip avoids the traffic entirely. That means Wigner-d recurrences in
Triton — a new kernel, not a wiring change, which is why it was out of scope.

Do not start by re-wiring the existing lanes. They are already wired; the
measurement above is what they are worth. Start by estimating whether an on-chip
build can beat pure JAX's 17.11 ms / 75.16 ms at all, then write the kernel.

Reproduce the numbers with:

python -m bench.bench_mutual_backends --sizes 10000 100000 --theta 0.7 --order 4
python -m bench.bench_mutual_farfield_stages --sizes 10000 100000

(1× A100-PCIE-40GB sm_80, jax 0.9.0.1, float64, θ=0.7, order 4, leaf 32, best of 3
after a discarded compile call.)

---

2. Per-rung multipoles (strategy B1) for an exact per-particle far-field split

Status today. The near field splits interactions by the exact per-particle
predicate max(rung_i, rung_j). The far field assigns each cell the rung of its
finest particle and splits at cell granularity — falcON activity gating, strategy
B2 in the design.

Both are genuine partitions, so they sum to the total force and every level
conserves momentum exactly. What B2 does not do is reproduce a direct-sum
oracle's per-level decomposition: it over-refines, treating a coarse particle that
shares a cell with a fine one at the fine level. Cross-checks against
MutualDirectSumGravity are therefore made on total force, momentum and energy
rather than level by level.
test_uniform_rung_reproduces_the_oracle_level_split pins the boundary case where
the two splits coincide.

Why B1 was not shipped. An exact per-particle far-field split needs one
multipole expansion per rung per cell — (k_max + 1)× the multipole memory, plus
up to (k_max + 1)² weighted M2L combinations for mixed-rung cell pairs. That
cost lands on the large-N path, which is the path this work exists to make viable.
B2 keeps one multipole set per cell and costs nothing.

When it would be worth building. As a small-N debugging aid, not a
shipping lane: if a per-level discrepancy against the oracle ever needs
localising, an exact per-rung-multipole variant (masked P2M, reusing
update_multipoles_only/refresh_prepared_state) would say whether the cause is
the cell-level approximation or a real bug. Gate it so it can never be selected on
the large-N path.

Open question worth settling first: nobody has measured how much B2's
over-refinement actually costs in a real run. If a rung-mixed cell is rare at
production leaf sizes, B1 has no value even as a diagnostic. That measurement is
cheap and should precede any implementation.

---

Neither item blocks anything. The momentum-conserving FMM is complete and merged;
these are the two places where a measured decision was made to stop, and both
decisions have an expiry condition attached.