Phase 5c + 5d: fused Pallas M2L/near-field + multi-GPU scaling in the distributed FMM

docs/north_star_phase3_scale_plan.md

@@ -0,0 +1,147 @@
+# North-star Phase 3 — scale the treecode-walk multi-GPU FMM to 100k–1M/GPU
+
+_Resume doc for a fresh session. Phases 1–2 are DONE and pushed; this is the remaining
+scale work. Written 2026-07-20._
+
+## Progress log
+
+- **2026-07-20 (later): branch reconciled onto main + Phase 3a DONE (local, not committed).**
+  Pulled `feat/phase5-multigpu-foldin` (was 83 behind); the fast-forward included merge
+  `53f0bc4` "Merge branch 'main'…" — this IS the §"Branch / merge note" reconciliation
+  (driver now imports M2L/L2L/L2P from `runtime.kernels.core`, treecode builder still from
+  `runtime._interaction_cache`; `local_walk` treecode work + the `far_pair_count` 0-pad seam
+  survived intact). Verified: CPU driver suite 6/6 green; GPU treecode parity ndev=2 per=1000
+  `self_ovf=0`, `overflow=False`, aggL2 7.77e-4 (better than the memo's 3.2e-3 — main's gains).
+  **Phase 3a (below) implemented + GPU/CPU validated.** NOT yet committed to the branch.
+
+## Where we are (Phases 1–2, DONE + pushed)
+
+**Goal of the North-star:** lift the distributed FMM's per-GPU-N ceiling by replacing the
+overflowing self dual-tree walk with the single-GPU fast-lane **treecode walk**. Diagnosis:
+the ceiling is `self_queue_overflow` — the dual-tree walk's transient *pair-queue*, not the
+final lists (at leaf=128/per=20000 the final lists are tiny yet the queue blows up).
+
+**Landed on `feat/phase5-multigpu-foldin` (origin, commit `4489d16` and its history):**
+- `DistributedFMMConfig.local_walk` ∈ {`dual_tree` (default), `treecode`}. The `treecode`
+  branch in `_make_fn`/`fn` (`jaccpot/distributed/fmm.py`) calls
+  `_build_treecode_artifacts_strict_streamed(tree, geom, theta, mac_type="dehnen", ...)` and
+  feeds its `compact_far_pairs` → the same real M2L (leaf-only targets → L2L no-op) and its
+  **self-excluded** `neighbor_list` → the same combined P2P. Cross-domain LET unchanged.
+- **Seam gotcha (already fixed):** treecode far pairs are **0-padded (not −1)** with the true
+  count in `far_pair_count`, so `s_active` uses `far_pair_count` on the treecode path (a `>=0`
+  test counts padding → degenerate delta=0 M2L → NaN).
+- **Validated:** CPU pytest `test_driver_local_walk_treecode`; GPU (A100) multi-GPU parity
+  ndev 2/3/4/5 (per=1000, `treecode`+`auto_scale_caps`): all overflow-free, `self_ovf=0`,
+  aggL2 vs direct 3.2e-3/7.2e-3/1.8e-2/2.6e-2 — vs the dual-tree walk's **0.53 garbage** (self
+  overflow) at ndev=2. The rising aggL2 with ndev is the **cross-domain far** (θ_cross=0.1 LET,
+  grows with domain count), NOT the treecode (self force is per-domain + exact).
+
+## How to resume (environment)
+
+- **Worktree:** `/export/home/tbuck/jaccpot-phase5-wt` on `feat/phase5-multigpu-foldin`
+  (isolated; shared checkout `/export/home/tbuck/jaccpot` stays on its own branch). jaccpot is
+  editable-installed from the shared checkout, so import it from the worktree via the
+  **sitecustomize finder-repoint**: `PYTHONPATH=<scratchpad>` where scratchpad has a
+  `sitecustomize.py` doing `__editable___jaccpot_0_0_1_finder.MAPPING['jaccpot'] =
+  '/export/home/tbuck/jaccpot-phase5-wt/jaccpot'` (see [[jaccpot-worktree-isolation]]).
+- **Python:** `/export/home/tbuck/micromamba/envs/odisseo/bin/python`. jaccpot forces x64.
+- **GPUs:** `autocvd -n N -l -o -q` → `CUDA_VISIBLE_DEVICES`. 8×A100-40GB box; free set varies.
+- **Pallas M2L:** export `JACCPOT_STATIC_STRICT_FUSED_M2L_PALLAS=1`. On Ampere the treecode
+  walk auto-uses its Pallas kernel; on CPU it falls back to JAX (so CPU tests work).
+- **Reusable scratchpad scripts (this session):** `treecode_swap_gpu.py` (GPU parity, argv
+  ndev/per/jit/config), `treecode_swap_parity.py` (CPU parity), `inspect_walks.py` (dump
+  walk outputs), `steady_time.py`/`scale_bench.py`/`leaf_sweep.py` (build-once timing +
+  cap-calibration). Copies of the earlier ones live in
+  `Odisseo/benchmark_a100/phase5_5c_multigpu/`.
+- **Test:** `XLA_FLAGS=--xla_force_host_platform_device_count=4 JAX_PLATFORMS=cpu pytest
+  tests/test_distributed_fmm_driver.py -o addopts="" -q` (CPU); drop the env + `autocvd` for GPU.
+
+## Phase 3 tasks (in order)
+
+### 3a — Right-size the treecode `near_cap` (DONE 2026-07-20, local; correctness + perf)
+**DONE.** `DistributedFMMConfig` gained `treecode_near_cap` / `treecode_far_cap`
+(`Optional[int]`, default None). In the driver's treecode branch (`_make_fn`/`fn`,
+`jaccpot/distributed/fmm.py`) the flat buffers are now right-sized from the STATIC local
+tree instead of the builder's fixed defaults: `near_cap = max(1<<14,
+max_neighbors_per_leaf * num_leaves)` and `far_cap = max(1<<14, max_interactions_per_node
+* num_leaves)` (both keyed off `with_scaled_caps`-scaled fields, so the auto-scale retry
+grows them). `_build_treecode_artifacts_strict_streamed` gained an optional `near_cap`
+kwarg (explicit override wins; None keeps the old env/1<<21 default → single-GPU fast lane
+byte-identical). The `_TreecodeWalkDiag` now surfaces REAL `far_overflow` (`far_pair_count
+>= far_cap`, clamped) / `near_overflow` (`near_pair_count > near_cap`, unclamped counts) —
+previously forced to 0, i.e. SILENT truncation with no diagnostic — so `_reduce_overflow`
+→ `auto_scale_caps` self-heals both, exactly like the dual-tree caps.
+VALIDATED (GPU A100 ndev=2 per=1000, `treecode_swap_gpu.py`): default caps give bit-identical
+parity to the 2M-buffer baseline (aggL2 7.7657e-4, `cap_retries=0`); forcing a tiny
+`treecode_near_cap=4000` → `near_overflow` fires → auto_scale retries 4000→8000→16000
+(`cap_retries=2`) → clears with identical forces; same for `treecode_far_cap=100`
+(`cap_retries=1`). CPU driver suite 6/6 green before AND after the edits.
+NOTE: the near-source-leaf table `S_near = max_neighbors_per_leaf + cross_max_neighbors_per_leaf`
+(Pallas near backend densification) is a SEPARATE truncation not touched here — watch it at scale.
+
+_(original 3a spec, for reference:)_
+### ~~3a — Right-size the treecode `near_cap`~~ (spec)
+The treecode near buffer defaults to `1<<21` (env `JACCPOT_STATIC_STRICT_FUSED_TREECODE_NEAR_CAP`,
+read in `_build_treecode_artifacts_strict_streamed`, `jaccpot/runtime/_interaction_cache.py`).
+Two problems: (1) `_combined_neighbors` then chews a 2M-edge array per device (slow); (2) **at
+~1M/GPU the actual near-pair count can EXCEED `1<<21` and the overflow guard is SKIPPED under
+trace** → silent truncation → wrong forces. Fix: size it to the actual bound and make it
+explicit. Options:
+- Add `DistributedFMMConfig.treecode_near_cap: Optional[int]` and, in the treecode branch, set
+  the env before the call (host-side, in `_make_fn`) OR — cleaner — call the lower-level
+  `build_treecode_far_pairs_and_neighbors` directly (as `tests/experimental/test_treecode_
+  graft_solidfmm.py` does) with explicit `near_capacity`/`far_pair_capacity`, replicating the
+  `_treecode_mac_extents` + topo-padding that `_build_treecode_artifacts_strict_streamed` does.
+- Safe bound: `near_cap ≈ num_leaves × max_neighbors_per_leaf` (per-device), with a factor of
+  safety; `far_cap ≈ num_leaves × (far-per-leaf)`. Validate empirically: assert the returned
+  `near_counts.sum()` / `far_pair_count` are **< cap** at the target N (no silent overflow).
+
+### 3b — Cross-walk: `auto_scale_caps` or treecode-swap it too
+The cross-domain LET walk (`dual_tree_walk_cross_impl`, via `yggdrax.distributed.cross_walk`)
+is STILL a dual-tree walk and overflows at higher ndev / connected ICs (that is why the ndev
+2–5 sweep used `auto_scale_caps`). For separated clusters the cross term is negligible so
+forces stay correct; for a **connected** IC the cross-domain far matters, so decide:
+- Cheapest: rely on `auto_scale_caps` (already works) — but the padded cross pair-queue is the
+  same overhead-scaling problem, so it may get slow/huge.
+- Proper: apply the SAME treecode-walk swap to the cross walk (a `local_walk`-style option for
+  the cross traversal). Bigger, but removes the cross ceiling too. Recommended before claiming
+  full scale on connected ICs.
+
+### 3c — A realistic connected IC
+Replace the separated-cluster benchmark IC with a **domain-decomposed galaxy disk** (the
+single-GPU fast lane's `notebooks/scalability/galaxy_disk_fmm_large_n.py` generates one;
+partition it across devices with the driver's `partition_for_devices`). This is what exercises
+the cross-domain far honestly and matches the single-GPU comparison. Adversarial single dense
+blobs are NOT representative (they inflate near-pairs for ANY near-field).
+
+### 3d — Scale runs + weak/strong scaling
+With 3a/3b/3c in place, push per-GPU N to 100k → 1M with `local_walk="treecode"`. Confirm
+`self_ovf=0` AND `not overflow` AND (assert final list sizes < caps) AND subsampled aggL2 vs
+direct < a few %. Then **build-once steady-state timing** (NOT `distributed_fmm_accelerations`,
+which recompiles ~250 s/call — use `make_force_evaluator(..., jit=True)` and call the compiled
+fn repeatedly; see `steady_time.py`). Weak scaling (fixed per-GPU N, ndev 2→5) + strong (fixed
+total, vary ndev); compare per-GPU throughput to the single-GPU O(N) curve (fast lane does
+4M/GPU). This is the payoff: demonstrate the ceiling is gone.
+
+## Gotchas already learned (don't rediscover)
+- **jit=True illegal-address crash is INTERMITTENT** (nondeterministic OOB); run each perf/probe
+  eval in its OWN process. Phase-2 correctness used jit=False. Probe jit=True + treecode
+  separately before relying on it for timing.
+- `distributed_fmm_accelerations` **rebuilds+recompiles every call** (~50–250 s). Build once for
+  timing.
+- The treecode uses **dehnen** MAC for accuracy-profile parity + multi-step stability; do NOT use
+  bh (dynamically unstable). θ_cross ≤ 0.1 (under-separated far-field goes garbage).
+- aggL2 is worst-weighted; the median (p50 ~2.8e-4 historically) is the fair accuracy metric.
+- Kill grid/sweep processes by PROCESS GROUP and verify zero before relaunch.
+
+## Verification
+- CPU pytest green (incl. `test_driver_local_walk_treecode`).
+- GPU: treecode ndev 2–5 parity vs direct, overflow-free, final list sizes < caps (assert).
+- Scale: 100k–1M/GPU overflow-free + accurate + weak/strong scaling curve; per-GPU throughput
+  vs the single-GPU 4M/GPU reference.
+
+## Branch / merge note
+`feat/phase5-multigpu-foldin` is on the OLD `jaccpot.runtime._fmm_impl` layout; main is on
+`jaccpot.runtime.kernels.core`. PR #47 (base main) has a merge conflict in `distributed/fmm.py`
+to resolve before landing (re-apply the 5c + treecode changes onto main's driver). The
+North-star wants main's newest fast lane, so reconcile onto main at some point.

docs/phase5_multigpu_pallas_foldin_plan.md

@@ -17,11 +17,48 @@ _Bookkeeping + plan, 2026-07-14._
     direct (beats bh's 0.24%), no overflow; solidfmm/bh kept behind explicit config + test.
   - Tests: `tests/test_distributed_fmm_driver.py` (default real+dehnen, real+bh, legacy
     solidfmm, jit==eager) + `tests/test_real_upward_sweep.py` — all green on CPU.
-- **NEXT — 5c (needs GPU):** set `JACCPOT_STATIC_STRICT_FUSED_M2L_PALLAS=1` (M2L already
-  routes through `_apply_real_m2l` → fused real M2L Pallas engages automatically) + swap the
-  near-field from `nearfield_mode="baseline"` to the fused Pallas P2P. Validate on Ampere.
-- **NEXT — 5d (needs GPU):** weak/strong multi-GPU scaling via `benchmark_multigpu/`; compare
-  per-device throughput to the single-GPU O(N) curve; watch LET/comm overhead.
+- **DONE — 5c (A100, 2026-07-19):**
+  - **Fused real M2L Pallas** engaged via `JACCPOT_STATIC_STRICT_FUSED_M2L_PALLAS=1` — no code
+    change; both self-M2L (`_accumulate_real_m2l_fullbatch`) and cross-M2L route through
+    `_apply_real_m2l`. Parity-neutral: flag-on vs flag-off forces match to **4.4e-9**.
+  - **Fused leafpair Pallas near-field**: new `DistributedFMMConfig.nearfield_backend`
+    {`auto`,`pallas`,`baseline`} (auto→pallas on sm_80+). The pallas branch computes the
+    intra-leaf self block via `_compute_leaf_p2p_prepared_large_n_self_only_impl` and the
+    cross-leaf pairs via `_radix_fast_lane_prepacked_pallas` over the `_combined_neighbors`
+    CSR densified to a padded `[u_leaves, S_near]` source-leaf table. Parity vs baseline P2P
+    = **2.1e-7** (CPU `interpret=True` de-risk 2e-16), aggL2 vs direct **1.78e-4** @ per=8000.
+  - **The near-field was the entire bottleneck.** Build-once steady-state whole-eval
+    (ndev=2, per=8000/leaf=128): baseline near **10.7 s → 43.5 ms with Pallas (~245x)**.
+    traversal + far-field are only ~40 ms — the ~10 s previously attributed to the
+    overhead-bound traversal was actually the pure-JAX baseline near-field leaf-pair P2P.
+    M2L-Pallas on/off is negligible at this near-field-dominated config.
+  - `jit=True` is STABLE with both backends (the earlier illegal-address crash did not
+    reproduce). NB: `distributed_fmm_accelerations` rebuilds+recompiles per call (~50-80 s) —
+    build once via `make_force_evaluator(...,jit=True)` for steady-state timing.
+- **DONE — 5d (A100, cap-calibrated build-once steady-state; 2026-07-20):**
+  - `distributed_fmm_accelerations` gained `auto_scale_caps` (retry with `with_scaled_caps`
+    on a traversal-buffer overflow) — the cross-domain LET grows with device count so the
+    fixed default caps overflow at ndev≥4. Test `test_driver_auto_scale_caps`.
+  - **Weak scaling** (per=8000/GPU, pallas near + fused M2L, overflow-free after calibration):
+
+    | ndev | N | cap× | min ms | throughput (part/s) |
+    |---|---|---|---|---|
+    | 2 | 16 000 | 1 | 41.6 | 3.8e5 |
+    | 3 | 24 000 | 1 | 46.6 | 5.2e5 |
+    | 4 | 32 000 | 2 | 59.6 | 5.4e5 |
+    | 5 | 40 000 | 2 | 64.3 | 6.2e5 |
+
+    Throughput RISES with GPU count (positive weak scaling); per-eval grows 42→64 ms from
+    LET comm + the ×2 padded-cap overhead at ndev≥4.
+  - **Strong scaling** (total N=40 000) is **density-limited**: at per-GPU N > ~8000 (ndev 2-4:
+    per=20000/13333/10000) the fixed-topology traversal caps explode and STILL overflow at
+    cap×64 (max retries) → forces truncated, 400-720 ms (padded pair-queue overhead dominates).
+    Only ndev=5 (per=8000) is valid (64.7 ms). **Healthy regime = per-GPU N ≈ 8000**; scale by
+    adding GPUs to keep per-GPU N there (= weak scaling).
+  - **CAVEAT: the jit=True illegal-address crash is INTERMITTENT** (nondeterministic OOB) — most
+    runs succeed but one recurred at weak ndev=2; run each eval in its own process. Root-cause +
+    a padded-pair-queue right-sizing (to lift the per-GPU-N ceiling for strong scaling) are the
+    two remaining follow-ups.
 
 ---
 

jaccpot/distributed/cap_presets.py

@@ -0,0 +1,117 @@
+"""Persistent capacity presets for the distributed FMM driver.
+
+The ``auto_scale_caps`` retry loop DISCOVERS the right traversal-buffer capacities for
+a given (per-GPU N, ndev, IC) but pays a ``shard_map`` recompile per retry. Persisting
+the converged caps keyed by problem size lets a repeat run at a known size START from
+them -- zero retries, a single compile. This is the cheap realisation of "size the caps
+from a pre-count": ``auto_scale`` *is* the pre-count, and we cache what it found.
+
+The overflow flags stay the safety net: if a preset undersizes (e.g. a denser IC at the
+same N), ``auto_scale`` still grows the caps and the caller can refresh the preset with
+the newly converged values. Caps are per-device, so the natural key is (per-GPU N, ndev);
+total N is recorded too. A preset is a validated STARTING POINT that transfers within an
+IC family (same morphology), not a guarantee across wildly different distributions.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import json
+import os
+from typing import Any, Optional
+
+from .fmm import DistributedFMMConfig
+
+# The traversal-buffer capacities the retry loop grows -- everything that affects a
+# static buffer shape. Order-independent; None means "driver right-sizes it".
+CAP_FIELDS = (
+    "max_interactions_per_node",
+    "max_neighbors_per_leaf",
+    "max_pair_queue",
+    "cross_max_interactions_per_node",
+    "cross_max_neighbors_per_leaf",
+    "cross_max_pair_queue",
+    "treecode_near_cap",
+    "treecode_far_cap",
+    "cross_far_cap",
+)
+
+
+def caps_of(config: DistributedFMMConfig) -> dict[str, Any]:
+    """Extract the cap values from a config (e.g. the converged ``result.config``)."""
+    return {f: getattr(config, f) for f in CAP_FIELDS}
+
+
+def apply_caps(
+    config: DistributedFMMConfig, caps: dict[str, Any]
+) -> DistributedFMMConfig:
+    """Return a copy of ``config`` with the cap fields present in ``caps`` applied."""
+    return dataclasses.replace(config, **{f: caps[f] for f in CAP_FIELDS if f in caps})
+
+
+def _key(per_gpu_n: int, ndev: int) -> str:
+    return f"{int(per_gpu_n)}:{int(ndev)}"
+
+
+def load_presets(path: Optional[str]) -> dict:
+    """Load the presets table from a JSON file (empty dict if missing/unset)."""
+    if path and os.path.exists(path):
+        with open(path) as fh:
+            return json.load(fh)
+    return {}
+
+
+def save_presets(path: str, presets: dict) -> None:
+    """Atomically write the presets table to ``path``."""
+    tmp = f"{path}.tmp"
+    with open(tmp, "w") as fh:
+        json.dump(presets, fh, indent=2, sort_keys=True)
+    os.replace(tmp, path)
+
+
+def _scale_caps(caps: dict[str, Any], num: int, den: int) -> dict[str, Any]:
+    """Scale integer caps by num/den (ceil); None stays None."""
+    return {
+        f: (None if caps.get(f) is None else int((int(caps[f]) * num + den - 1) // den))
+        for f in CAP_FIELDS
+    }
+
+
+def lookup(presets: dict, per_gpu_n: int, ndev: int) -> Optional[dict]:
+    """Caps for (per_gpu_n, ndev): exact match, else the nearest LARGER per-GPU N at the
+    same ndev (a safe over-estimate), else the nearest SMALLER preset SCALED UP by the
+    per-GPU-N ratio. The scaled seed is only a starting point -- auto_scale refines it (and
+    the caller can refresh the preset), so a slight under/over-estimate just costs a retry
+    or a little memory, not correctness. Extrapolating from a nearby preset is what makes
+    calibration cheap at an unmeasured N (few retries instead of the full doubling ladder
+    from the small defaults). None if no preset at this ndev at all."""
+    k = _key(per_gpu_n, ndev)
+    if k in presets:
+        return presets[k]["caps"]
+    same = [
+        (int(kk.split(":")[0]), v["caps"])
+        for kk, v in presets.items()
+        if kk.endswith(f":{int(ndev)}")
+    ]
+    if not same:
+        return None
+    larger = [(n, c) for n, c in same if n >= int(per_gpu_n)]
+    if larger:
+        return min(larger, key=lambda t: t[0])[1]
+    n0, c0 = max(same, key=lambda t: t[0])  # largest smaller
+    return _scale_caps(c0, int(per_gpu_n), n0)
+
+
+def record(
+    presets: dict, per_gpu_n: int, ndev: int, total_n: int, caps: dict[str, Any]
+) -> dict:
+    """Insert/update the (per_gpu_n, ndev) entry with the given caps (in place)."""
+    presets[_key(per_gpu_n, ndev)] = {
+        "per_gpu_n": int(per_gpu_n),
+        "ndev": int(ndev),
+        "total_n": int(total_n),
+        "caps": {
+            f: (None if caps.get(f) is None else int(caps[f])) for f in CAP_FIELDS
+        },
+    }
+    return presets