From 3685fbade32cd250a3a18ec4ae5c709bbd5f520d Mon Sep 17 00:00:00 2001 From: TobiBu Date: Sun, 19 Jul 2026 23:19:36 +0200 Subject: [PATCH 01/26] perf(distributed): fused Pallas leafpair near-field backend (5c) Add DistributedFMMConfig.nearfield_backend {auto,pallas,baseline}. On Ampere+ (auto default) route the combined [local;halo] near-field P2P through the fused leafpair Pallas kernel (the single-GPU fast-lane near-field) instead of the pure-JAX baseline: intra-leaf self block via _compute_leaf_p2p_prepared_large_n_self_only_impl + cross-leaf neighbour pairs via _radix_fast_lane_prepacked_pallas over the combined CSR densified to a padded [u_leaves, S_near] source-leaf-id table. Baseline path unchanged behind the knob. Validated on A100 (ndev=2, real+dehnen default): near-field parity vs baseline P2P = 2.1e-7 (CPU interpret de-risk 2e-16), aggL2 vs direct = 1.78e-4 @ per=8000. Build-once steady-state whole-eval 10.7 s -> 43.5 ms (~245x) at per=8000/leaf=128 -- the baseline near-field P2P, NOT the dual-tree traversal, was the dominant cost. jit=True stable with both backends. The fused real M2L Pallas (JACCPOT_STATIC_STRICT_FUSED_M2L_PALLAS=1) is already engaged via _apply_real_m2l for both self- and cross-M2L; parity-neutral (4.4e-9), negligible at this near-field-dominated config. Co-Authored-By: Claude Opus 4.8 (1M context) --- jaccpot/distributed/fmm.py | 118 +++++++++++++++++++++++++++++++------ 1 file changed, 101 insertions(+), 17 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 396069e..38bac47 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -66,7 +66,12 @@ from yggdrax.tree_moments import compute_tree_mass_moments from jaccpot.downward.local_expansions import LocalExpansionData -from jaccpot.nearfield.near_field import compute_leaf_p2p_accelerations +from jaccpot.nearfield.near_field import ( + _compute_leaf_p2p_prepared_large_n_self_only_impl, + _prepare_leaf_data_from_groups, + _radix_fast_lane_prepacked_pallas, + compute_leaf_p2p_accelerations, +) from jaccpot.operators.complex_ops import ( enforce_conjugate_symmetry_batch, m2l_complex_reference_batch, @@ -134,6 +139,11 @@ class DistributedFMMConfig: # JACCPOT_STATIC_STRICT_FUSED_M2L_PALLAS is set). Upward sweep + coarse M2M stay # complex; multipoles are converted to Dehnen real coeffs at the M2L boundary. basis: str = "real" + # Near-field backend: "baseline" (pure-JAX combined [local;halo] P2P) or + # "pallas" (fused leafpair Pallas kernel, sm_80+, the single-GPU fast-lane + # near-field). "auto" (DEFAULT) picks pallas on Ampere+, baseline elsewhere. + # Numerically equivalent to baseline (validated), only the cross-leaf P2P is fused. + nearfield_backend: str = "auto" # self dual-tree walk capacities max_interactions_per_node: int = 512 max_neighbors_per_leaf: int = 128 @@ -266,6 +276,25 @@ def _make_fn(config: DistributedFMMConfig, ndev: int, cap: int) -> Callable: theta_cross = config.theta_cross is_real = str(config.basis).strip().lower() == "real" + # Near-field backend resolution (trace-time; sm_80+ gate mirrors the single-GPU + # lane). "auto" -> fused leafpair Pallas on Ampere+, pure-JAX baseline elsewhere. + nf_backend = str(config.nearfield_backend).strip().lower() + if nf_backend not in {"auto", "pallas", "baseline"}: + raise ValueError( + "nearfield_backend must be 'auto', 'pallas' or 'baseline'; " + f"got {config.nearfield_backend!r}" + ) + if nf_backend == "auto": + from jaccpot.pallas.nearfield_fused_leaf import pallas_nearfield_fused_supported + + use_pallas_near = bool(pallas_nearfield_fused_supported()) + else: + use_pallas_near = nf_backend == "pallas" + # Padded source-leaf-id width for the fused leafpair kernel: per target leaf the + # combined CSR holds at most (local self-near + cross-near) neighbour leaves. + S_near = int(config.max_neighbors_per_leaf) + int(config.cross_max_neighbors_per_leaf) + soft2 = float(soft) ** 2 + C = sh_size(p) if cap % leaf != 0: raise ValueError(f"cap={cap} must be a multiple of leaf_size={leaf}") @@ -588,22 +617,77 @@ def _l2p(loc_coeffs): ) lp_idx = jnp.concatenate([loc_idx, halo_idx], axis=0) lp_mask = jnp.concatenate([loc_mask, halo_mask], axis=0) - near_full = compute_leaf_p2p_accelerations( - tree, - nbr, - concat_pos, - concat_mass, - G=G, - softening=soft, - nearfield_mode="baseline", - node_ranges_override=jnp.zeros((u_leaves + 1, 2), INDEX_DTYPE), - leaf_nodes_override=jnp.arange(u_leaves, dtype=INDEX_DTYPE), - neighbor_offsets_override=offsets, - neighbor_indices_override=src_s, - neighbor_counts_override=counts, - leaf_particle_indices_override=lp_idx, - leaf_particle_mask_override=lp_mask, - ) + if use_pallas_near: + # Fused leafpair Pallas near-field (single-GPU fast-lane kernel): the + # intra-leaf self block is computed separately and the cross-leaf + # neighbour pairs run through the leafpair kernel (sources gathered by + # id in-kernel). Mirrors compute_leaf_p2p_accelerations_radix_fast_lane; + # numerically equal to the baseline combined P2P (CPU-validated 2e-16). + soft2_a = jnp.asarray(soft2, concat_pos.dtype) + ( + leaf_positions, + leaf_masses, + leaf_mask_g, + safe_idx, + ) = _prepare_leaf_data_from_groups(lp_idx, lp_mask, concat_pos, concat_mass) + self_acc = _compute_leaf_p2p_prepared_large_n_self_only_impl( + concat_pos, + leaf_positions, + leaf_masses, + leaf_mask_g, + safe_idx, + G=G, + softening_sq=soft2_a, + ) + # Densify the combined CSR (offsets/src_s/counts, sorted by target) into a + # padded [u_leaves, S_near] source-leaf-id table + validity mask. + e = jnp.arange(src_s.shape[0], dtype=INDEX_DTYPE) + t_of_e = jnp.searchsorted(offsets, e, side="right") - 1 + rank = e - offsets[jnp.clip(t_of_e, 0, u_leaves)] + in_range = (t_of_e >= 0) & (t_of_e < u_leaves) & (rank < S_near) + flat = jnp.where(in_range, t_of_e * S_near + rank, -1) + sids = ( + jnp.zeros((u_leaves * S_near,), INDEX_DTYPE) + .at[flat] + .set(src_s, mode="drop") + .reshape(u_leaves, S_near) + ) + svalid = ( + jnp.zeros((u_leaves * S_near,), bool) + .at[flat] + .set(jnp.ones_like(flat, bool), mode="drop") + .reshape(u_leaves, S_near) + ) + pair_acc = _radix_fast_lane_prepacked_pallas( + sids.reshape(u_leaves, S_near, 1), + svalid.reshape(u_leaves, S_near, 1), + leaf_positions, + leaf_masses, + leaf_mask_g, + safe_idx, + concat_pos, + G=G, + softening_sq=soft2_a, + compute_potential=False, + ) + near_full = self_acc + pair_acc + else: + near_full = compute_leaf_p2p_accelerations( + tree, + nbr, + concat_pos, + concat_mass, + G=G, + softening=soft, + nearfield_mode="baseline", + node_ranges_override=jnp.zeros((u_leaves + 1, 2), INDEX_DTYPE), + leaf_nodes_override=jnp.arange(u_leaves, dtype=INDEX_DTYPE), + neighbor_offsets_override=offsets, + neighbor_indices_override=src_s, + neighbor_counts_override=counts, + leaf_particle_indices_override=lp_idx, + leaf_particle_mask_override=lp_mask, + ) # far_full is (cap,3) already; near_full is (cap + halo, 3) -> keep the # local rows. (The validated test slices BOTH to [:cap] on return.) From bd176d7b6306ece582060a5f3882866198cd93bb Mon Sep 17 00:00:00 2001 From: TobiBu Date: Sun, 19 Jul 2026 23:28:43 +0200 Subject: [PATCH 02/26] docs(phase5): mark 5c DONE (Pallas M2L flag + Pallas near-field backend) Near-field Pallas leafpair backend is the ~245x win (10.7s->43.5ms @8k, ndev=2); the baseline near-field P2P -- not the traversal -- was the ~10s cost. jit=True stable; weak-scales flat 2->5 GPUs but default caps overflow at ndev>=4 (5d needs cap calibration first). Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/phase5_multigpu_pallas_foldin_plan.md | 27 ++++++++++++++++++---- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/docs/phase5_multigpu_pallas_foldin_plan.md b/docs/phase5_multigpu_pallas_foldin_plan.md index f26f9c5..e59c3da 100644 --- a/docs/phase5_multigpu_pallas_foldin_plan.md +++ b/docs/phase5_multigpu_pallas_foldin_plan.md @@ -17,11 +17,28 @@ _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. +- **NEXT — 5d (now feasible):** weak-scaling of the pallas config (per=8000/GPU) is near-flat + ndev 2→5 (~42-48 ms), so the driver actually scales — BUT the default traversal caps + **overflow at ndev≥4** (cross-domain LET grows with device count → `cross_max_pair_queue`); + grow via `DistributedFMMConfig.with_scaled_caps` before a real strong/weak scaling study. --- From c42143ed52cde9d42db7dc93df6e75825b31836e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 19 Jul 2026 21:38:03 +0000 Subject: [PATCH 03/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaccpot/distributed/fmm.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 38bac47..578b115 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -292,7 +292,9 @@ def _make_fn(config: DistributedFMMConfig, ndev: int, cap: int) -> Callable: use_pallas_near = nf_backend == "pallas" # Padded source-leaf-id width for the fused leafpair kernel: per target leaf the # combined CSR holds at most (local self-near + cross-near) neighbour leaves. - S_near = int(config.max_neighbors_per_leaf) + int(config.cross_max_neighbors_per_leaf) + S_near = int(config.max_neighbors_per_leaf) + int( + config.cross_max_neighbors_per_leaf + ) soft2 = float(soft) ** 2 C = sh_size(p) From a0baec7ccfd3663592675c02f0065b47f05a27d4 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Sun, 19 Jul 2026 23:58:00 +0200 Subject: [PATCH 04/26] feat(distributed): auto_scale_caps retry for ndev>=4 LET overflow (5d prep) The cross-domain LET grows with device count, so the fixed default traversal caps overflow at ndev>=4 (measured: cross_max_pair_queue overflow at per=8000). Add auto_scale_caps/cap_scale_factor/max_cap_retries to distributed_fmm_accelerations: on a traversal-buffer overflow, scale the config caps via with_scaled_caps and rebuild/re-run until overflow-free; the returned result carries the config that produced the forces + diagnostics["cap_retries"]. Factor the overflow reduction into _reduce_overflow. Verified (weak scaling, per=8000/GPU): ndev 4/5 need cap_factor=2 to run overflow-free (throughput still rises 5.4e5->6.2e5 part/s). Test test_driver_auto_scale_caps: tiny caps overflow without retry, one x2 retry clears it, forces match direct (aggL2 1.9e-3, CPU 4-dev). Co-Authored-By: Claude Opus 4.8 (1M context) --- jaccpot/distributed/fmm.py | 68 ++++++++++++++++------------ tests/test_distributed_fmm_driver.py | 51 +++++++++++++++++++++ 2 files changed, 91 insertions(+), 28 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 578b115..3ca516a 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -745,6 +745,22 @@ def evaluate(pos_flat, mass_flat, gid_flat, counts): return jax.jit(evaluate) if jit else evaluate +_OVERFLOW_FIELDS = ( + "cross_queue_overflow", + "cross_far_overflow", + "cross_near_overflow", + "self_queue_overflow", + "self_far_overflow", + "self_near_overflow", +) + + +def _reduce_overflow(diag_o: np.ndarray) -> bool: + """True if any device flagged a traversal-buffer overflow.""" + idx = [DIAG_FIELDS.index(k) for k in _OVERFLOW_FIELDS] + return bool(np.any(diag_o[:, idx] > 0)) + + def distributed_fmm_accelerations( positions: np.ndarray, masses: np.ndarray, @@ -753,6 +769,9 @@ def distributed_fmm_accelerations( mesh=None, ndev: int | None = None, jit: bool = False, + auto_scale_caps: bool = False, + cap_scale_factor: float = 2.0, + max_cap_retries: int = 4, ) -> DistributedFMMResult: """Evaluate distributed FMM accelerations for all particles. @@ -760,6 +779,12 @@ def distributed_fmm_accelerations( runs the ``shard_map`` force pipeline, and scatters the per-device forces back into the original input order. + When ``auto_scale_caps`` is set, a traversal-buffer overflow (which grows + with device count as the cross-domain LET expands) triggers a retry with the + capacities scaled by ``cap_scale_factor`` (up to ``max_cap_retries`` times), + rebuilding the evaluator each time. The ``config`` on the returned result is + the one that actually produced the (non-overflowing) forces. + Returns a :class:`DistributedFMMResult`; ``.accelerations`` has shape ``(N, 3)`` in input order, ``.diagnostics`` includes the per-device pair counts / overflow flags and a reduced ``overflow`` bool. @@ -777,17 +802,22 @@ def distributed_fmm_accelerations( part = partition_for_devices(positions, masses, ndev, leaf_size=config.leaf_size) cap = part["cap"] counts_dev = jnp.asarray(part["counts"], INDEX_DTYPE) - - evaluate = make_force_evaluator(config, ndev, cap, mesh, jit=jit) - accel_o, gid_o, diag_o = evaluate( - jnp.asarray(part["pos_flat"]), - jnp.asarray(part["mass_flat"]), - jnp.asarray(part["gid_flat"]), - counts_dev, - ) + pos_f = jnp.asarray(part["pos_flat"]) + mass_f = jnp.asarray(part["mass_flat"]) + gid_f = jnp.asarray(part["gid_flat"]) + + attempt = 0 + while True: + evaluate = make_force_evaluator(config, ndev, cap, mesh, jit=jit) + accel_o, gid_o, diag_o = evaluate(pos_f, mass_f, gid_f, counts_dev) + diag_o = np.asarray(diag_o) + overflow = _reduce_overflow(diag_o) + if not overflow or not auto_scale_caps or attempt >= max_cap_retries: + break + config = config.with_scaled_caps(cap_scale_factor) + attempt += 1 accel_o = np.asarray(accel_o) gid_o = np.asarray(gid_o).reshape(-1).astype(np.int64) - diag_o = np.asarray(diag_o) n = part["n"] accel = np.zeros((n, 3), np.float64) @@ -804,26 +834,8 @@ def distributed_fmm_accelerations( ) diag = {name: diag_o[:, i] for i, name in enumerate(DIAG_FIELDS)} - overflow = bool( - np.any( - diag_o[ - :, - [ - DIAG_FIELDS.index(k) - for k in ( - "cross_queue_overflow", - "cross_far_overflow", - "cross_near_overflow", - "self_queue_overflow", - "self_far_overflow", - "self_near_overflow", - ) - ], - ] - > 0 - ) - ) diag["overflow"] = overflow + diag["cap_retries"] = attempt return DistributedFMMResult( accelerations=accel, diff --git a/tests/test_distributed_fmm_driver.py b/tests/test_distributed_fmm_driver.py index 6cdec85..4ab034a 100644 --- a/tests/test_distributed_fmm_driver.py +++ b/tests/test_distributed_fmm_driver.py @@ -11,6 +11,8 @@ pytest tests/test_distributed_fmm_driver.py -o addopts="" -q """ +import dataclasses + import jax.numpy as jnp import numpy as np import pytest @@ -150,3 +152,52 @@ def test_driver_jit_matches_eager(): ) print(f"jit-vs-eager aggL2 = {diff:.3e}") assert diff < 1e-6, f"jit path diverged from eager: {diff:.3e}" + + +def test_driver_auto_scale_caps(): + """``auto_scale_caps`` grows the traversal buffers on overflow until the + walk fits, without changing the forces. + + The cross-domain LET grows with device count, so the fixed default caps can + overflow at higher ``ndev``. Deliberately tiny caps force the overflow here; + the retry must clear it and still match direct N-body. + """ + ndev = min(4, device_count()) + mesh = make_mesh(ndev) + per = 64 + pts, mass = _separated_clusters(ndev, per) + tiny = dataclasses.replace( + DistributedFMMConfig(), + nearfield_backend="baseline", # pallas is GPU-only; CI runs on CPU + max_pair_queue=64, + cross_max_pair_queue=64, + max_interactions_per_node=16, + cross_max_interactions_per_node=16, + max_neighbors_per_leaf=16, + cross_max_neighbors_per_leaf=16, + ) + + r0 = distributed_fmm_accelerations(pts, mass, config=tiny, mesh=mesh, jit=False) + assert r0.overflow, "tiny caps should overflow without auto_scale_caps" + + r1 = distributed_fmm_accelerations( + pts, + mass, + config=tiny, + mesh=mesh, + jit=False, + auto_scale_caps=True, + cap_scale_factor=2.0, + max_cap_retries=6, + ) + assert not r1.overflow, "auto_scale_caps should clear the overflow" + assert r1.diagnostics["cap_retries"] >= 1 + + direct = np.asarray( + _direct(jnp.asarray(pts), jnp.asarray(mass), tiny.G, tiny.softening) + ) + err = float( + np.linalg.norm(r1.accelerations - direct) / (np.linalg.norm(direct) + 1e-30) + ) + print(f"auto_scale_caps aggL2 vs direct = {err:.6f} retries={r1.diagnostics['cap_retries']}") + assert err < 5e-2, f"forces wrong after cap retry: {err:.6f}" From 1e2e448e7f603cd5665d310aa85c16b46aecb410 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 20 Jul 2026 00:22:44 +0200 Subject: [PATCH 05/26] docs(phase5): 5d scaling results (weak scales, strong is density-limited) Weak (per=8000/GPU, cap-calibrated overflow-free): throughput rises 3.8e5->6.2e5 part/s for ndev 2->5 (per-eval 42->64ms). Strong (total=40000) is density-limited: per-GPU N>~8000 explodes the fixed-topology traversal caps (still overflow at x64), 400-720ms truncated -- only ndev=5/per=8000 valid. Healthy regime per-GPU N~8000. CAVEAT: jit=True illegal-address crash is INTERMITTENT (recurred once at ndev=2). Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/phase5_multigpu_pallas_foldin_plan.md | 28 ++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/docs/phase5_multigpu_pallas_foldin_plan.md b/docs/phase5_multigpu_pallas_foldin_plan.md index e59c3da..f31191c 100644 --- a/docs/phase5_multigpu_pallas_foldin_plan.md +++ b/docs/phase5_multigpu_pallas_foldin_plan.md @@ -35,10 +35,30 @@ _Bookkeeping + plan, 2026-07-14._ - `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. -- **NEXT — 5d (now feasible):** weak-scaling of the pallas config (per=8000/GPU) is near-flat - ndev 2→5 (~42-48 ms), so the driver actually scales — BUT the default traversal caps - **overflow at ndev≥4** (cross-domain LET grows with device count → `cross_max_pair_queue`); - grow via `DistributedFMMConfig.with_scaled_caps` before a real strong/weak scaling study. +- **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. --- From 272223aca04b1710732644729103d5a66b717960 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 19 Jul 2026 22:24:01 +0000 Subject: [PATCH 06/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_distributed_fmm_driver.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_distributed_fmm_driver.py b/tests/test_distributed_fmm_driver.py index 4ab034a..2367b86 100644 --- a/tests/test_distributed_fmm_driver.py +++ b/tests/test_distributed_fmm_driver.py @@ -199,5 +199,7 @@ def test_driver_auto_scale_caps(): err = float( np.linalg.norm(r1.accelerations - direct) / (np.linalg.norm(direct) + 1e-30) ) - print(f"auto_scale_caps aggL2 vs direct = {err:.6f} retries={r1.diagnostics['cap_retries']}") + print( + f"auto_scale_caps aggL2 vs direct = {err:.6f} retries={r1.diagnostics['cap_retries']}" + ) assert err < 5e-2, f"forces wrong after cap retry: {err:.6f}" From 4489d16fc4f44650567770de5a7217fa3e278d6f Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 20 Jul 2026 12:06:25 +0200 Subject: [PATCH 07/26] feat(distributed): fast-lane treecode self-walk option (North-star Phase 1) Add DistributedFMMConfig.local_walk {dual_tree, treecode}. The dual-tree self-walk's transient pair-queue caps per-GPU N (diagnosed self_queue_overflow: at per=1000/leaf=8 it overflows and gives garbage, aggL2 0.90 vs direct). The device-resident fast-lane treecode walk (_build_treecode_artifacts_strict_streamed, mac_type=dehnen for accuracy- profile parity + multi-step stability) streams far/near with no such queue and is a drop-in for the self-walk: its CompactTaggedFarPairs feed the same real M2L (leaf-only targets -> L2L no-op) and its self-excluded near CSR feeds the same combined P2P. Seam adapted: the 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 would count padding -> degenerate delta=0 M2L -> NaN). CPU-validated (ndev=4): treecode clears the self-overflow (self_ovf 4->0) and matches direct (per=1000 aggL2 4.0e-3 where dual-tree self-walk overflows to 0.90; per=256 1.3e-3). The cross-domain LET walk is still dual-tree, so auto_scale_caps keeps the whole eval overflow-free. Test test_driver_local_walk_treecode. Cross-walk treecode swap + shard_map/GPU scale validation (Phase 2/3) are follow-ups. Co-Authored-By: Claude Opus 4.8 (1M context) --- jaccpot/distributed/fmm.py | 84 ++++++++++++++++++++++++---- tests/test_distributed_fmm_driver.py | 40 +++++++++++++ 2 files changed, 114 insertions(+), 10 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 3ca516a..39adbff 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -32,7 +32,7 @@ import dataclasses from dataclasses import dataclass -from typing import Any, Callable +from typing import Any, Callable, NamedTuple import jax import jax.numpy as jnp @@ -84,6 +84,9 @@ _evaluate_local_expansions_for_particles, _propagate_solidfmm_locals_by_level, ) +from jaccpot.runtime._interaction_cache import ( + _build_treecode_artifacts_strict_streamed, +) from jaccpot.upward.real_tree_expansions import ( aggregate_m2m_real_by_level, prepare_real_upward_sweep, @@ -108,6 +111,21 @@ ) +class _TreecodeWalkDiag(NamedTuple): + """Minimal self-walk diagnostic shim for the treecode local walk. + + The treecode walk emits no ``DualTreeWalkResult``; it auto-sizes per-leaf caps + (overflow raises in the eager prepare pass, so the flags are 0 here) and exposes + the far/near counts we surface in the per-device diagnostic vector. + """ + + far_pair_count: Any + near_pair_count: Any + queue_overflow: Any + far_overflow: Any + near_overflow: Any + + @dataclass(frozen=True) class DistributedFMMConfig: """Static knobs for the distributed FMM force evaluation. @@ -144,6 +162,15 @@ class DistributedFMMConfig: # near-field). "auto" (DEFAULT) picks pallas on Ampere+, baseline elsewhere. # Numerically equivalent to baseline (validated), only the cross-leaf P2P is fused. nearfield_backend: str = "auto" + # Local self-interaction walk: "dual_tree" (yggdrax dual-tree walk, DEFAULT) or + # "treecode" (the single-GPU fast-lane device-resident treecode walk). The dual-tree + # walk's transient pair-queue caps per-GPU N (self_queue_overflow); the treecode walk + # streams far/near with no such queue, so per-GPU N scales like the single-GPU lane. + # Parity with dual_tree at mac_type="dehnen" (accuracy-profile parity, leaf-only far + # targets -> L2L no-op, self-excluded near CSR). Cross-domain LET is unchanged. + local_walk: str = "dual_tree" + # Sphere-radius scale for the treecode dehnen MAC (matches the dual-tree extents). + dehnen_radius_scale: float = 1.0 # self dual-tree walk capacities max_interactions_per_node: int = 512 max_neighbors_per_leaf: int = 128 @@ -276,6 +303,15 @@ def _make_fn(config: DistributedFMMConfig, ndev: int, cap: int) -> Callable: theta_cross = config.theta_cross is_real = str(config.basis).strip().lower() == "real" + # Local self-walk selection: dual-tree (default) or the fast-lane treecode walk. + lw = str(config.local_walk).strip().lower() + if lw not in {"dual_tree", "treecode"}: + raise ValueError( + f"local_walk must be 'dual_tree' or 'treecode'; got {config.local_walk!r}" + ) + use_treecode_local = lw == "treecode" + dehnen_radius_scale = float(config.dehnen_radius_scale) + # Near-field backend resolution (trace-time; sm_80+ gate mirrors the single-GPU # lane). "auto" -> fused leafpair Pallas on Ampere+, pure-JAX baseline elsewhere. nf_backend = str(config.nearfield_backend).strip().lower() @@ -407,14 +443,37 @@ def fn(pos, mass, gid, count): coeff_dtype = lp.dtype if is_real else cdtype packed_use = packed - inter, nbr, self_res = build_interactions_and_neighbors( - tree, - geom, - theta=theta, - traversal_config=cfg, - mac_type=mac, - return_result=True, - ) + if use_treecode_local: + # Fast-lane device-resident treecode walk in place of the dual-tree walk. + # Drop-in: compact_far_pairs (leaf-only far targets -> L2L no-op) feed the + # same real M2L; the self-excluded near CSR feeds the same combined P2P. + _art = _build_treecode_artifacts_strict_streamed( + tree=tree, + geometry=geom, + theta=theta, + mac_type=mac, + dehnen_radius_scale=dehnen_radius_scale, + compact_far_pair_capacity=None, + ) + inter = _art.compact_far_pairs + nbr = _art.neighbor_list + _z = jnp.zeros((), jnp.int64) + self_res = _TreecodeWalkDiag( + far_pair_count=jnp.asarray(inter.far_pair_count), + near_pair_count=jnp.sum(jnp.asarray(nbr.counts)), + queue_overflow=_z, + far_overflow=_z, + near_overflow=_z, + ) + else: + inter, nbr, self_res = build_interactions_and_neighbors( + tree, + geom, + theta=theta, + traversal_config=cfg, + mac_type=mac, + return_result=True, + ) # remote coarse tree over frontier (leaf COM + mass) mm = compute_tree_mass_moments(tree, lp, lm) @@ -557,7 +616,12 @@ def _l2p(loc_coeffs): zeros = jnp.zeros((int(total_nodes), C), dtype=coeff_dtype) s_src = jnp.asarray(inter.sources, INDEX_DTYPE) s_tgt = jnp.asarray(inter.targets, INDEX_DTYPE) - s_active = jnp.sum((s_tgt >= 0).astype(INDEX_DTYPE)) + # The treecode compact far pairs are 0-padded (not -1) with the true count in + # far_pair_count; the dual-tree list is trimmed, so a >=0 test recovers its count. + if use_treecode_local: + s_active = jnp.asarray(inter.far_pair_count, INDEX_DTYPE) + else: + s_active = jnp.sum((s_tgt >= 0).astype(INDEX_DTYPE)) if is_real: loc_self = _accumulate_real_m2l_fullbatch( zeros, diff --git a/tests/test_distributed_fmm_driver.py b/tests/test_distributed_fmm_driver.py index 2367b86..c17b8ee 100644 --- a/tests/test_distributed_fmm_driver.py +++ b/tests/test_distributed_fmm_driver.py @@ -203,3 +203,43 @@ def test_driver_auto_scale_caps(): f"auto_scale_caps aggL2 vs direct = {err:.6f} retries={r1.diagnostics['cap_retries']}" ) assert err < 5e-2, f"forces wrong after cap retry: {err:.6f}" + + +def test_driver_local_walk_treecode(): + """The fast-lane treecode self-walk replaces the dual-tree self-walk. + + The dual-tree self-walk's transient pair-queue caps per-GPU N (``self_queue_overflow``); + the device-resident treecode walk streams far/near with no such queue. This checks the + swap is a faithful drop-in: correct forces vs direct with ZERO self-overflow. (The + cross-domain LET walk is still dual-tree, so ``auto_scale_caps`` keeps the whole eval + overflow-free; the ceiling win itself is exercised at larger ``per`` off-CI.) + """ + ndev = min(4, device_count()) + mesh = make_mesh(ndev) + per = 256 + pts, mass = _separated_clusters(ndev, per) + base = DistributedFMMConfig() + cfg = dataclasses.replace(base, local_walk="treecode") + + r = distributed_fmm_accelerations( + pts, + mass, + config=cfg, + mesh=mesh, + jit=False, + auto_scale_caps=True, + cap_scale_factor=2.0, + max_cap_retries=6, + ) + d = r.diagnostics + self_ovf = float( + np.asarray(d["self_queue_overflow"]).sum() + + np.asarray(d["self_far_overflow"]).sum() + + np.asarray(d["self_near_overflow"]).sum() + ) + direct = np.asarray(_direct(jnp.asarray(pts), jnp.asarray(mass), base.G, base.softening)) + err = float(np.linalg.norm(r.accelerations - direct) / (np.linalg.norm(direct) + 1e-30)) + print(f"treecode local_walk: self_ovf={self_ovf} overflow={r.overflow} aggL2 vs direct={err:.4e}") + assert self_ovf == 0, "treecode self-walk must not overflow (that is the point of the swap)" + assert not r.overflow, "auto_scale_caps should clear the (dual-tree) cross-walk overflow" + assert err < 5e-2, f"treecode local force wrong vs direct: {err:.4e}" From 832772a98408caafedeb90ff4dae4e4bf52de2ff Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 20 Jul 2026 14:19:39 +0200 Subject: [PATCH 08/26] docs(north-star): Phase 3 scale plan (resume doc for a fresh session) Phases 1-2 done + pushed (treecode self-walk under shard_map, multi-GPU parity ndev 2-5 overflow-free). Phase 3 remaining: right-size the treecode near_cap (silent-overflow risk at ~1M/GPU), cross-walk treecode swap or auto_scale for connected ICs, a realistic domain-decomposed galaxy-disk IC, and 100k-1M/GPU scale + weak/strong timing (build-once) vs the single-GPU 4M/GPU curve. Includes resume environment (worktree + finder-repoint), reusable scripts, and gotchas. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/north_star_phase3_scale_plan.md | 114 +++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) create mode 100644 docs/north_star_phase3_scale_plan.md diff --git a/docs/north_star_phase3_scale_plan.md b/docs/north_star_phase3_scale_plan.md new file mode 100644 index 0000000..f7e6147 --- /dev/null +++ b/docs/north_star_phase3_scale_plan.md @@ -0,0 +1,114 @@ +# 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._ + +## 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=` 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` (do FIRST; correctness + perf) +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. From b98851fcdbfaec3fb1d6c6f16eb6a832749a2f73 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 20 Jul 2026 14:57:34 +0200 Subject: [PATCH 09/26] style(distributed): black-format test_driver_local_walk_treecode (line-length 88) Wrap the >88-char lines pre-commit.ci flagged (black 26.5.1). isort clean; fmm.py already conformant. No behaviour change. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_distributed_fmm_driver.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/tests/test_distributed_fmm_driver.py b/tests/test_distributed_fmm_driver.py index c17b8ee..fa2365b 100644 --- a/tests/test_distributed_fmm_driver.py +++ b/tests/test_distributed_fmm_driver.py @@ -237,9 +237,19 @@ def test_driver_local_walk_treecode(): + np.asarray(d["self_far_overflow"]).sum() + np.asarray(d["self_near_overflow"]).sum() ) - direct = np.asarray(_direct(jnp.asarray(pts), jnp.asarray(mass), base.G, base.softening)) - err = float(np.linalg.norm(r.accelerations - direct) / (np.linalg.norm(direct) + 1e-30)) - print(f"treecode local_walk: self_ovf={self_ovf} overflow={r.overflow} aggL2 vs direct={err:.4e}") - assert self_ovf == 0, "treecode self-walk must not overflow (that is the point of the swap)" - assert not r.overflow, "auto_scale_caps should clear the (dual-tree) cross-walk overflow" + direct = np.asarray( + _direct(jnp.asarray(pts), jnp.asarray(mass), base.G, base.softening) + ) + err = float( + np.linalg.norm(r.accelerations - direct) / (np.linalg.norm(direct) + 1e-30) + ) + print( + f"treecode local_walk: self_ovf={self_ovf} overflow={r.overflow} aggL2 vs direct={err:.4e}" + ) + assert ( + self_ovf == 0 + ), "treecode self-walk must not overflow (that is the point of the swap)" + assert ( + not r.overflow + ), "auto_scale_caps should clear the (dual-tree) cross-walk overflow" assert err < 5e-2, f"treecode local force wrong vs direct: {err:.4e}" From b6e07fc0cc809f6c4ff312a1041fa0c157eb7f81 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 20 Jul 2026 15:11:41 +0200 Subject: [PATCH 10/26] build: pin black==26.5.1 + isort==9.0.0b1 (align CI lint with pre-commit hooks) The CI `lint` job installs the [dev] extras and runs `black --check .` / `isort --check-only .`, but [dev] had floating ">=" pins while .pre-commit-config pinned the hooks. CI resolved a newer isort than the pinned 8.0.1 hook; the two disagree on single-import collapsing (no common fixed point), so `isort --check-only .` failed on files the hook considered clean (tests/unit/test_solver_api.py, jaccpot/runtime/_fmm_impl.py). Mirror main's fix: pin [dev] black/isort to the hook revs and bump the isort hook to 9.0.0b1, then reformat to 9.x style (collapse single parenthesized imports). isort + black now agree local/hook/CI. Co-Authored-By: Claude Opus 4.8 (1M context) --- .pre-commit-config.yaml | 2 +- jaccpot/runtime/_fmm_impl.py | 4 +--- pyproject.toml | 8 ++++++-- tests/unit/test_solver_api.py | 10 ++-------- 4 files changed, 10 insertions(+), 14 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 429d1a9..18badea 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -5,7 +5,7 @@ repos: - id: black - repo: https://github.com/PyCQA/isort - rev: 8.0.1 + rev: 9.0.0b1 hooks: - id: isort diff --git a/jaccpot/runtime/_fmm_impl.py b/jaccpot/runtime/_fmm_impl.py index 6dfb025..9a085cf 100644 --- a/jaccpot/runtime/_fmm_impl.py +++ b/jaccpot/runtime/_fmm_impl.py @@ -67,9 +67,7 @@ from jaccpot.downward.local_expansions import ( run_downward_sweep as run_tree_downward_sweep, ) -from jaccpot.downward.local_expansions import ( - translate_local_expansion, -) +from jaccpot.downward.local_expansions import translate_local_expansion from jaccpot.nearfield.near_field import ( compute_leaf_p2p_accelerations, compute_leaf_p2p_accelerations_large_n_accel_only, diff --git a/pyproject.toml b/pyproject.toml index e2b5bf2..c3b408d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,8 +40,12 @@ Documentation = "https://github.com/AstroAI-Lab/jaccpot#readme" [project.optional-dependencies] dev = [ - "black>=24.8.0", - "isort>=5.13.2", + # Pin black + isort to the .pre-commit-config.yaml hook revs so local, the + # pre-commit hooks, and CI's `black --check .` / `isort --check-only .` all + # format identically. A floating ">=" let CI resolve to a newer isort whose + # style disagreed with the pinned hook (single-import collapsing regression). + "black==26.5.1", + "isort==9.0.0b1", "pytest>=8.3.0", "pytest-cov>=5.0.0", "pytest-xdist>=3.6.0", diff --git a/tests/unit/test_solver_api.py b/tests/unit/test_solver_api.py index 9e7d4bd..84bceee 100644 --- a/tests/unit/test_solver_api.py +++ b/tests/unit/test_solver_api.py @@ -20,16 +20,10 @@ ) from jaccpot import FastMultipoleMethod from jaccpot import FastMultipoleMethod as ExpanseFMM -from jaccpot import ( - FMMAdvancedConfig, -) +from jaccpot import FMMAdvancedConfig from jaccpot import FMMPreset from jaccpot import FMMPreset as ExpansePreset -from jaccpot import ( - NearFieldConfig, - RuntimePolicyConfig, - TreeConfig, -) +from jaccpot import NearFieldConfig, RuntimePolicyConfig, TreeConfig def _sample_problem(n: int = 64): From 4b70aeedef597521707b8697b34d31d537e8b697 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 13:14:29 +0000 Subject: [PATCH 11/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaccpot/distributed/fmm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 8f40567..3488e37 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -77,15 +77,15 @@ m2l_complex_reference_batch, ) from jaccpot.operators.real_harmonics import sh_size +from jaccpot.runtime._interaction_cache import ( + _build_treecode_artifacts_strict_streamed, +) from jaccpot.runtime.kernels.core import ( _accumulate_m2l_fullbatch, _apply_real_m2l, _evaluate_local_expansions_for_particles, _propagate_solidfmm_locals_by_level, ) -from jaccpot.runtime._interaction_cache import ( - _build_treecode_artifacts_strict_streamed, -) from jaccpot.upward.real_tree_expansions import ( aggregate_m2m_real_by_level, prepare_real_upward_sweep, From e903664bdd50d9c34645f8388a3402cdcbd44be5 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 20 Jul 2026 21:15:13 +0200 Subject: [PATCH 12/26] feat(distributed): self-healing treecode caps + selective auto-scale (Phase 3a/3b) Lift the connected-IC scaling ceiling of the distributed FMM driver from ~50k to ~200k particles/GPU by making capacity self-heal without OOM. Phase 3a -- right-size the treecode local-walk flat buffers + surface overflow: - Add treecode_near_cap / treecode_far_cap to DistributedFMMConfig. When None, the driver sizes the near buffer to max_neighbors_per_leaf * num_leaves and the far buffer to max_interactions_per_node * num_leaves, instead of the builder's fixed 1<<21 (2M) near default -- which both wasted the combined-P2P neighbour build and could SILENTLY truncate at ~1M/GPU (the treecode overflow guard is eager-only, skipped under the shard_map trace). - Plumb an explicit near_cap through _build_treecode_artifacts_strict_streamed; when absent it keeps the env/1<<21 default so the single-GPU fast lane stays byte-identical. - Surface the treecode far/near overflow flags from the true counts vs the caps so _reduce_overflow / auto_scale_caps grow them like the dual-tree caps. Phase 3b -- selective auto-scale (with_selective_scaled_caps): - The retry loop grew EVERY cap uniformly (with_scaled_caps) on any overflow. On a connected IC only the cross-near buffer overflows, but uniform scaling also doubled cross_max_interactions_per_node (the cross-FAR cap) each retry -- and the cross-far list feeds the fused real-M2L, whose Pallas buffer pads to a power of two and OOMs (~6 GiB at N>=200k/2GPU) even though the true cross-far volume is ~10^3 pairs. This coupling was the real ceiling, not cross-far volume. - with_selective_scaled_caps grows only the caps whose per-device overflow flag fired (read from the diagnostic vector); the retry loop now uses it. Validated on 2x A100: per=100k (N=200k, was a hard OOM) and per=200k (N=400k) now run overflow-free; per=400k/GPU hits a separate fused-M2L memory wall (genuine ~1M cross-far pairs), addressed later by M2L chunking. CPU unit test covers the selective-scaling cases (cross-near-only grows only the near cap; self far+near; explicit treecode caps; per-device flags). Co-Authored-By: Claude Opus 4.8 --- docs/north_star_phase3_scale_plan.md | 35 +++++- jaccpot/distributed/fmm.py | 148 ++++++++++++++++++++++++-- jaccpot/runtime/_interaction_cache.py | 16 ++- 3 files changed, 185 insertions(+), 14 deletions(-) diff --git a/docs/north_star_phase3_scale_plan.md b/docs/north_star_phase3_scale_plan.md index f7e6147..d14b4c5 100644 --- a/docs/north_star_phase3_scale_plan.md +++ b/docs/north_star_phase3_scale_plan.md @@ -3,6 +3,17 @@ _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 @@ -47,7 +58,29 @@ final lists (at leaf=128/per=20000 the final lists are tiny yet the queue blows ## Phase 3 tasks (in order) -### 3a — Right-size the treecode `near_cap` (do FIRST; correctness + perf) +### 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 diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 3488e37..0387c21 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -32,7 +32,7 @@ import dataclasses from dataclasses import dataclass -from typing import Any, Callable, NamedTuple +from typing import Any, Callable, NamedTuple, Optional import jax import jax.numpy as jnp @@ -113,9 +113,12 @@ class _TreecodeWalkDiag(NamedTuple): """Minimal self-walk diagnostic shim for the treecode local walk. - The treecode walk emits no ``DualTreeWalkResult``; it auto-sizes per-leaf caps - (overflow raises in the eager prepare pass, so the flags are 0 here) and exposes - the far/near counts we surface in the per-device diagnostic vector. + The treecode walk emits no ``DualTreeWalkResult``; it auto-sizes per-leaf caps and + exposes the far/near counts we surface in the per-device diagnostic vector. There is + no transient pair-queue (``queue_overflow`` is always 0 -- that is the point of the + swap), but the flat far/near buffers can still overflow, so ``far_overflow`` / + ``near_overflow`` are set from the true counts vs the (right-sized) caps and drive + ``auto_scale_caps`` -- the eager-only builder guard is skipped under the trace. """ far_pair_count: Any @@ -170,6 +173,19 @@ class DistributedFMMConfig: local_walk: str = "dual_tree" # Sphere-radius scale for the treecode dehnen MAC (matches the dual-tree extents). dehnen_radius_scale: float = 1.0 + # Treecode local-walk flat buffer sizes (only used when local_walk="treecode"). + # The builder's own default is a fixed 1<<21 (2M) near-edge buffer, which makes the + # combined-P2P neighbour build chew a 2M-edge array per device AND can SILENTLY + # truncate the near list at ~1M/GPU (the treecode overflow guard is eager-only, so + # it is skipped under the shard_map trace -> wrong forces with no diagnostic). When + # these are None the driver right-sizes them from the local tree: the near buffer to + # ``max_neighbors_per_leaf * num_leaves`` (the same per-leaf near budget the dual-tree + # walk uses, and already grown by ``with_scaled_caps`` on the auto-scale retry) and + # the far buffer to ``treecode_far_cap`` or 131072. The true per-device far/near + # counts and an accurate overflow flag are surfaced in the ``self_*`` diagnostics so + # ``auto_scale_caps`` grows them on overflow exactly like the dual-tree caps. + treecode_near_cap: Optional[int] = None + treecode_far_cap: Optional[int] = None # self dual-tree walk capacities max_interactions_per_node: int = 512 max_neighbors_per_leaf: int = 128 @@ -199,6 +215,82 @@ def g(v: int) -> int: cross_max_interactions_per_node=g(self.cross_max_interactions_per_node), cross_max_neighbors_per_leaf=g(self.cross_max_neighbors_per_leaf), cross_max_pair_queue=g(self.cross_max_pair_queue), + # Grow the treecode buffers on retry too. When these are None the driver + # sizes the near buffer off ``max_neighbors_per_leaf`` (scaled above), so the + # retry already grows the effective near cap; scale any explicit values here. + treecode_near_cap=( + None if self.treecode_near_cap is None else g(self.treecode_near_cap) + ), + treecode_far_cap=( + None if self.treecode_far_cap is None else g(self.treecode_far_cap) + ), + ) + + def with_selective_scaled_caps( + self: "DistributedFMMConfig", diag_o: Any, factor: float + ) -> "DistributedFMMConfig": + """Return a copy scaling ONLY the caps whose overflow flag actually fired. + + ``diag_o`` is the per-device diagnostic matrix (shape ``(ndev, len(DIAG_FIELDS))``) + from the last evaluation; a cap is grown by ``factor`` iff its overflow flag is + nonzero on any device. + + Unlike :meth:`with_scaled_caps` (which grows *every* cap uniformly), this grows + only the buffer a device actually overflowed. Uniform scaling couples the + self/cross x far/near/queue caps, so a single overflowing buffer drags up an + unrelated, already-oversized cap. Concretely: on a connected IC only the + cross-near buffer overflows, but uniform scaling also doubles + ``cross_max_interactions_per_node`` (the cross-*far* cap) each retry -- and the + cross-far list feeds the fused real-M2L, whose Pallas buffer pads to a power of + two and OOMs (~6 GiB at N>=200k/2GPU) even though the true cross-far volume is + ~10^3 pairs. Selective scaling grows just the buffer that overflowed, lifting the + connected-IC ceiling without inflating the M2L. + """ + diag = np.asarray(diag_o) + + def flagged(name: str) -> bool: + return bool(np.any(diag[:, DIAG_FIELDS.index(name)] > 0)) + + def g(v: int) -> int: + return int(np.ceil(v * factor)) + + def maybe(name: str, value: int) -> int: + return g(value) if flagged(name) else value + + # Under the treecode local walk the self near/far caps derive from + # max_neighbors_per_leaf / max_interactions_per_node (times num_leaves) when the + # explicit treecode_*_cap overrides are None, so scaling those knobs grows the + # effective self buffers for both the dual-tree and treecode walks. + grow_self_far = flagged("self_far_overflow") + grow_self_near = flagged("self_near_overflow") + return dataclasses.replace( + self, + max_interactions_per_node=maybe( + "self_far_overflow", self.max_interactions_per_node + ), + max_neighbors_per_leaf=maybe( + "self_near_overflow", self.max_neighbors_per_leaf + ), + max_pair_queue=maybe("self_queue_overflow", self.max_pair_queue), + cross_max_interactions_per_node=maybe( + "cross_far_overflow", self.cross_max_interactions_per_node + ), + cross_max_neighbors_per_leaf=maybe( + "cross_near_overflow", self.cross_max_neighbors_per_leaf + ), + cross_max_pair_queue=maybe( + "cross_queue_overflow", self.cross_max_pair_queue + ), + treecode_near_cap=( + g(self.treecode_near_cap) + if (grow_self_near and self.treecode_near_cap is not None) + else self.treecode_near_cap + ), + treecode_far_cap=( + g(self.treecode_far_cap) + if (grow_self_far and self.treecode_far_cap is not None) + else self.treecode_far_cap + ), ) @@ -450,23 +542,56 @@ def fn( # Fast-lane device-resident treecode walk in place of the dual-tree walk. # Drop-in: compact_far_pairs (leaf-only far targets -> L2L no-op) feed the # same real M2L; the self-excluded near CSR feeds the same combined P2P. + # + # Right-size the flat far/near buffers from the (static) local tree instead + # of the builder's fixed 1<<21 near default: the near buffer is what the + # combined-P2P neighbour build iterates over, so an oversized buffer is pure + # overhead, and at ~1M/GPU the fixed 2M could be EXCEEDED (silent truncation: + # the treecode overflow guard is eager-only, skipped under this trace). The + # near budget mirrors the dual-tree walk's per-leaf bound (already grown by + # with_scaled_caps on the auto-scale retry). + num_internal = int(tree.topology.left_child.shape[0]) + num_leaves = int(total_nodes) - num_internal + if config.treecode_near_cap is not None: + tc_near_cap = int(config.treecode_near_cap) + else: + tc_near_cap = max( + 1 << 14, int(config.max_neighbors_per_leaf) * num_leaves + ) + # Far mirrors near: the compact far list holds <= (interaction-list size) far + # pairs per leaf, so budget max_interactions_per_node * num_leaves. Keyed off + # a with_scaled_caps-scaled field so the auto-scale retry grows it too, and + # num_leaves-proportional so it does not under-size at ~1M/GPU (the fixed + # 131072 could be exceeded by a spread IC -> silent far truncation). + if config.treecode_far_cap is not None: + tc_far_cap = int(config.treecode_far_cap) + else: + tc_far_cap = max( + 1 << 14, int(config.max_interactions_per_node) * num_leaves + ) _art = _build_treecode_artifacts_strict_streamed( tree=tree, geometry=geom, theta=theta, mac_type=mac, dehnen_radius_scale=dehnen_radius_scale, - compact_far_pair_capacity=None, + compact_far_pair_capacity=tc_far_cap, + near_cap=tc_near_cap, ) inter = _art.compact_far_pairs nbr = _art.neighbor_list _z = jnp.zeros((), jnp.int64) + # near_counts are UNCLAMPED true counts (see _compact_near), so > cap is exact + # truncation; far_pair_count is clamped to far_cap, so == cap flags a possible + # overflow. Surfaced so _reduce_overflow/auto_scale_caps grow the caps. + far_cnt = jnp.asarray(inter.far_pair_count) + near_cnt = jnp.sum(jnp.asarray(nbr.counts)) self_res = _TreecodeWalkDiag( - far_pair_count=jnp.asarray(inter.far_pair_count), - near_pair_count=jnp.sum(jnp.asarray(nbr.counts)), + far_pair_count=far_cnt, + near_pair_count=near_cnt, queue_overflow=_z, - far_overflow=_z, - near_overflow=_z, + far_overflow=(far_cnt >= tc_far_cap).astype(jnp.int64), + near_overflow=(near_cnt > tc_near_cap).astype(jnp.int64), ) else: inter, nbr, self_res = build_interactions_and_neighbors( @@ -888,7 +1013,10 @@ def distributed_fmm_accelerations( overflow = _reduce_overflow(diag_o) if not overflow or not auto_scale_caps or attempt >= max_cap_retries: break - config = config.with_scaled_caps(cap_scale_factor) + # Grow only the buffers that actually overflowed. Uniform scaling couples the + # cross-far cap (which feeds the fused-M2L Pallas buffer) to unrelated overflows + # and OOMs on connected ICs; see with_selective_scaled_caps. + config = config.with_selective_scaled_caps(diag_o, cap_scale_factor) attempt += 1 accel_o = np.asarray(accel_o) gid_o = np.asarray(gid_o).reshape(-1).astype(np.int64) diff --git a/jaccpot/runtime/_interaction_cache.py b/jaccpot/runtime/_interaction_cache.py index 64ef64a..edcefca 100644 --- a/jaccpot/runtime/_interaction_cache.py +++ b/jaccpot/runtime/_interaction_cache.py @@ -737,6 +737,7 @@ def _build_treecode_artifacts_strict_streamed( mac_type: MACType, dehnen_radius_scale: float, compact_far_pair_capacity: Optional[int], + near_cap: Optional[int] = None, ) -> _DualTreeArtifacts: """Strict fast-lane far/near build from the per-leaf treecode walk. @@ -863,9 +864,18 @@ def _env_int(name, default): "JACCPOT_STATIC_STRICT_FUSED_TREECODE_NEAR_PER_LEAF", auto_per_leaf ) max_stack = _env_int("JACCPOT_STATIC_STRICT_FUSED_TREECODE_STACK", 512) - near_cap = _env_int( - "JACCPOT_STATIC_STRICT_FUSED_TREECODE_NEAR_CAP", neighbor_edge_cap - ) + # An explicit ``near_cap`` (right-sized by the caller, e.g. the distributed driver + # to ~max_neighbors_per_leaf * num_leaves) overrides the env/1<<21 default. The + # 1<<21 default keeps the single-GPU fast lane byte-identical when no cap is passed; + # the fixed 2M buffer both wastes the downstream neighbour build and can SILENTLY + # truncate at very large N (the overflow guard below is eager-only). Callers that + # trace this (shard_map) should pass an explicit, validated ``near_cap``. + if near_cap is not None: + near_cap = int(near_cap) + else: + near_cap = _env_int( + "JACCPOT_STATIC_STRICT_FUSED_TREECODE_NEAR_CAP", neighbor_edge_cap + ) far_cap = int(compact_far_pair_capacity) if compact_far_pair_capacity else 131072 prod = build_treecode_far_pairs_and_neighbors( From 29b0be5235d3542f1f11d8fc238b49b746c3cdcf Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 20 Jul 2026 23:19:36 +0200 Subject: [PATCH 13/26] perf(distributed): right-size cross-far M2L input + cache converged caps as presets Two follow-ups to the Phase 3a/3b self-healing caps, both aimed at the connected-IC high-N path. 1. cross_far_cap -- slice the cross-far M2L input to the actual far volume. The cross walk sizes interaction_sources at t_total * cross_max_interactions_per_node but PACKS the valid far interactions into the front [0, far_pair_count); the tail is -1. The driver was feeding the whole (mostly-invalid) buffer to the fused real-M2L, which pads that length to a power of two and OOMs (~6 GiB at N>=400k/GPU on 2 A100) even though the true far volume is far smaller. Slice to cross_far_cap (default cross_max_interactions_per_node * num_target_leaves, ~2x tighter than the buffer and >= the far volume measured; override via the new config field). A far volume above the cap is OR'd into cross_far_overflow so auto_scale widens it -- the slice never silently drops interactions. Measured: per=100k/200k byte-identical (no regression); per=400k no longer OOMs (now bounded by the retry budget, not memory). 2. cap presets -- cache the converged caps keyed by (per-GPU N, ndev) and reuse them. auto_scale DISCOVERS the right capacities but pays a shard_map recompile per retry (up to 8 at per=200k). New jaccpot/distributed/cap_presets.py persists the converged caps to JSON; distributed_fmm_accelerations gains cap_presets_path= to seed a repeat run at a known size (skip the retries -> single compile) and save the converged caps after an overflow-free run. auto_scale stays the safety net; a preset is a validated starting point that transfers within an IC family. This is the cheap realisation of "size caps from a pre-count": auto_scale is the pre-count, cached by problem size. Co-Authored-By: Claude Opus 4.8 --- jaccpot/distributed/cap_presets.py | 101 +++++++++++++++++++++++++++++ jaccpot/distributed/fmm.py | 73 +++++++++++++++++++-- 2 files changed, 170 insertions(+), 4 deletions(-) create mode 100644 jaccpot/distributed/cap_presets.py diff --git a/jaccpot/distributed/cap_presets.py b/jaccpot/distributed/cap_presets.py new file mode 100644 index 0000000..5a1a95a --- /dev/null +++ b/jaccpot/distributed/cap_presets.py @@ -0,0 +1,101 @@ +"""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 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). None if nothing usable.""" + k = _key(per_gpu_n, ndev) + if k in presets: + return presets[k]["caps"] + larger = [ + (int(kk.split(":")[0]), v["caps"]) + for kk, v in presets.items() + if kk.endswith(f":{int(ndev)}") and int(kk.split(":")[0]) >= int(per_gpu_n) + ] + if larger: + return min(larger, key=lambda t: t[0])[1] + return None + + +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 diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 0387c21..60cc15f 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -195,6 +195,14 @@ class DistributedFMMConfig: cross_max_interactions_per_node: int = 512 cross_max_neighbors_per_leaf: int = 128 cross_max_pair_queue: int = 1 << 15 + # Static length of the cross-far M2L input. The cross walk sizes its far buffer at + # t_total * cross_max_interactions_per_node but PACKS the valid interactions into the + # front; feeding the whole (mostly-invalid) buffer to the fused real-M2L pads it to a + # power of two and OOMs (~6 GiB at N>=400k/GPU). When None the driver right-sizes the + # M2L input to cross_max_interactions_per_node * num_target_leaves (~2x tighter than + # the buffer, and >= the far volume on the ICs measured), overflowing into + # cross_far_overflow so auto_scale grows it. Set an explicit value to override. + cross_far_cap: Optional[int] = None def with_scaled_caps( self: "DistributedFMMConfig", factor: float @@ -215,6 +223,9 @@ def g(v: int) -> int: cross_max_interactions_per_node=g(self.cross_max_interactions_per_node), cross_max_neighbors_per_leaf=g(self.cross_max_neighbors_per_leaf), cross_max_pair_queue=g(self.cross_max_pair_queue), + cross_far_cap=( + None if self.cross_far_cap is None else g(self.cross_far_cap) + ), # Grow the treecode buffers on retry too. When these are None the driver # sizes the near buffer off ``max_neighbors_per_leaf`` (scaled above), so the # retry already grows the effective near cap; scale any explicit values here. @@ -263,6 +274,7 @@ def maybe(name: str, value: int) -> int: # effective self buffers for both the dual-tree and treecode walks. grow_self_far = flagged("self_far_overflow") grow_self_near = flagged("self_near_overflow") + grow_cross_far = flagged("cross_far_overflow") return dataclasses.replace( self, max_interactions_per_node=maybe( @@ -281,6 +293,14 @@ def maybe(name: str, value: int) -> int: cross_max_pair_queue=maybe( "cross_queue_overflow", self.cross_max_pair_queue ), + # cross_far_overflow also covers the M2L-input right-size (far volume > + # cross_far_cap); grow an explicit cap so the retry widens the M2L slice. When + # None the slice derives from cross_max_interactions_per_node (grown above). + cross_far_cap=( + g(self.cross_far_cap) + if (grow_cross_far and self.cross_far_cap is not None) + else self.cross_far_cap + ), treecode_near_cap=( g(self.treecode_near_cap) if (grow_self_near and self.treecode_near_cap is not None) @@ -777,9 +797,24 @@ def _l2p(loc_coeffs): total_nodes=int(total_nodes), ) - # cross M2L: separate source (coarse) / target (local) centres - x_src = jnp.asarray(cross.interaction_sources, INDEX_DTYPE) - x_tgt = jnp.asarray(cross.interaction_targets, INDEX_DTYPE) + # cross M2L: separate source (coarse) / target (local) centres. + # The cross walk sizes interaction_sources at t_total * KC but PACKS the valid far + # interactions into the front [0, far_pair_count); the tail is -1. Feeding the + # whole buffer to the fused real-M2L pads that length to a power of two and OOMs + # (~6 GiB at N>=400k/GPU). Slice to a right-sized cross-far cap so the M2L only + # spans the actual far volume. KC * num_target_leaves mirrors the self treecode + # far cap and is ~2x tighter than the t_total*KC buffer; cross_far_cap overrides. + # A far volume above the cap is surfaced (cross_far_overflow, below) so auto_scale + # widens it -- so the slice never silently drops interactions. + max_far = int(jnp.asarray(cross.interaction_sources).shape[0]) + num_tgt_leaves = int(jnp.asarray(nbr.leaf_indices).shape[0]) + if config.cross_far_cap is not None: + xfar_cap = int(config.cross_far_cap) + else: + xfar_cap = max(1 << 14, KC * num_tgt_leaves) + xfar_cap = min(xfar_cap, max_far) + x_src = jnp.asarray(cross.interaction_sources, INDEX_DTYPE)[:xfar_cap] + x_tgt = jnp.asarray(cross.interaction_targets, INDEX_DTYPE)[:xfar_cap] x_valid = x_tgt >= 0 xs = jnp.where(x_valid, x_src, 0) xt = jnp.where(x_valid, x_tgt, 0) @@ -889,12 +924,18 @@ def _l2p(loc_coeffs): # local rows. (The validated test slices BOTH to [:cap] on return.) accel = far_full[:cap] + near_full[:cap] + # cross-far overflow = per-node far-buffer overflow (KC) OR the M2L-input slice + # dropping interactions (far volume > xfar_cap); either way auto_scale grows it. + cross_far_ovf = jnp.maximum( + cross.far_overflow.astype(jnp.float64), + (cross.far_pair_count > xfar_cap).astype(jnp.float64), + ) diag = jnp.array( [ cross.far_pair_count.astype(jnp.float64), cross.near_pair_count.astype(jnp.float64), cross.queue_overflow.astype(jnp.float64), - cross.far_overflow.astype(jnp.float64), + cross_far_ovf, cross.near_overflow.astype(jnp.float64), self_res.far_pair_count.astype(jnp.float64), self_res.near_pair_count.astype(jnp.float64), @@ -971,6 +1012,7 @@ def distributed_fmm_accelerations( auto_scale_caps: bool = False, cap_scale_factor: float = 2.0, max_cap_retries: int = 4, + cap_presets_path: str | None = None, ) -> DistributedFMMResult: """Evaluate distributed FMM accelerations for all particles. @@ -1005,6 +1047,20 @@ def distributed_fmm_accelerations( mass_f = jnp.asarray(part["mass_flat"]) gid_f = jnp.asarray(part["gid_flat"]) + # Seed the caps from a saved preset for this (per-GPU N, ndev) so a repeat run at a + # known size starts already-sized and skips the auto_scale recompiles. auto_scale + # stays on as the safety net; the converged caps refresh the preset below. + total_n = int(part["n"]) + per_gpu_n = -(-total_n // ndev) # ceil + presets = {} + if cap_presets_path is not None: + from . import cap_presets as _cp + + presets = _cp.load_presets(cap_presets_path) + seed = _cp.lookup(presets, per_gpu_n, ndev) + if seed is not None: + config = _cp.apply_caps(config, seed) + attempt = 0 while True: evaluate = make_force_evaluator(config, ndev, cap, mesh, jit=jit) @@ -1039,6 +1095,15 @@ def distributed_fmm_accelerations( diag["overflow"] = overflow diag["cap_retries"] = attempt + # Persist the converged caps as the preset for this size (only when overflow-free, so + # we never cache a truncated/wrong config). Next run at this (per-GPU N, ndev) reuses + # them and skips the retries. + if cap_presets_path is not None and not overflow: + from . import cap_presets as _cp + + _cp.record(presets, per_gpu_n, ndev, total_n, _cp.caps_of(config)) + _cp.save_presets(cap_presets_path, presets) + return DistributedFMMResult( accelerations=accel, diagnostics=diag, From b871e38695169dcae88c1a4a22bf6f68cfca6971 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 05:51:20 +0000 Subject: [PATCH 14/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaccpot/distributed/cap_presets.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/jaccpot/distributed/cap_presets.py b/jaccpot/distributed/cap_presets.py index 5a1a95a..aa5d013 100644 --- a/jaccpot/distributed/cap_presets.py +++ b/jaccpot/distributed/cap_presets.py @@ -12,6 +12,7 @@ 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 @@ -45,9 +46,7 @@ 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} - ) + 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: From 95dcf75b120c35a93a2e9ab065dd14c7d1013645 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Tue, 21 Jul 2026 10:55:23 +0200 Subject: [PATCH 15/26] =?UTF-8?q?perf(distributed):=20optional=20fp32=20fa?= =?UTF-8?q?r-field=20M2L=20(=E2=89=882x=20per-GPU=20capacity)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-pair rotation blocks [N, Bp, mdp, mdp] built for the fused real M2L dominate its peak memory and are the cross-far OOM at high per-GPU N / domain count. Add DistributedFMMConfig.far_m2l_fp32 (default False): when set, the self + cross far-M2L inputs (multipoles, centres/deltas) are cast to fp32 so the kernel and its rotation blocks are fp32, halving M2L peak memory; the local-expansion accumulation stays coeff_dtype (fp64). Opt-in and distributed-only, so the single-GPU fast lane is untouched. Accuracy-neutral: the far field is bounded by the order-p truncation error (~1e-3), far above fp32 rounding (~1e-7). Measured per=50k/ndev=2 aggL2 1.6904e-3 == the fp64 baseline to 4 sig figs. Capacity: with fp32 + leaf_size=256, per=200k/GPU x4 (800k total) -- which OOM'd in fp64/leaf128 -- now runs overflow-free at aggL2 7.5e-3 (~2x the prior ~100k/GPU ceiling at ndev>=4). per=400k/GPU x4 still OOMs at theta_cross=0.1 (cross-far volume too large even halved); a looser theta_cross=0.3 fits but degrades accuracy to ~1e-1, so the accuracy-preserving path beyond ~200k/GPU is M2L chunking, not looser theta_cross. Co-Authored-By: Claude Opus 4.8 --- jaccpot/distributed/fmm.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 60cc15f..ed00e41 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -203,6 +203,12 @@ class DistributedFMMConfig: # the buffer, and >= the far volume on the ICs measured), overflowing into # cross_far_overflow so auto_scale grows it. Set an explicit value to override. cross_far_cap: Optional[int] = None + # Run the far-field M2L (self + cross) in fp32. The per-pair rotation blocks + # [N, Bp, mdp, mdp] dominate M2L peak memory; fp32 halves it (≈2x more particles/GPU) + # at negligible accuracy cost (fp32 ~1e-7 << the order-p FMM truncation error). The + # accumulation into the local expansions stays coeff_dtype (fp64). Opt-in; the + # single-GPU fast lane is unaffected (this field is distributed-only). + far_m2l_fp32: bool = False def with_scaled_caps( self: "DistributedFMMConfig", factor: float @@ -762,6 +768,9 @@ def _l2p(loc_coeffs): # tgt-minus-src delta convention; the real path routes through the same # Pallas-aware _apply_real_m2l kernel used by the single-GPU fast lane. zeros = jnp.zeros((int(total_nodes), C), dtype=coeff_dtype) + # Far-M2L compute dtype: fp32 halves the per-pair rotation-block memory (the M2L + # peak) at negligible accuracy cost; accumulation stays coeff_dtype (fp64). + m2l_dtype = jnp.float32 if config.far_m2l_fp32 else coeff_dtype s_src = jnp.asarray(inter.sources, INDEX_DTYPE) s_tgt = jnp.asarray(inter.targets, INDEX_DTYPE) # The treecode compact far pairs are 0-padded (not -1) with the true count in @@ -773,8 +782,8 @@ def _l2p(loc_coeffs): if is_real: loc_self = _accumulate_m2l_fullbatch( zeros, - packed_use, - centers, + packed_use.astype(m2l_dtype), + centers.astype(m2l_dtype), s_src, s_tgt, s_active, @@ -821,7 +830,10 @@ def _l2p(loc_coeffs): x_deltas = centers[xt] - c_centers[xs] if is_real: x_contribs = _apply_real_m2l( - coarse_packed_use[xs], x_deltas, order=p, m2l_impl="rot_scale" + coarse_packed_use[xs].astype(m2l_dtype), + x_deltas.astype(m2l_dtype), + order=p, + m2l_impl="rot_scale", ).astype(coeff_dtype) else: x_contribs = m2l_complex_reference_batch( From 71b6b798459f0f2154bcec5ba0a642acc634294d Mon Sep 17 00:00:00 2001 From: TobiBu Date: Tue, 21 Jul 2026 11:14:16 +0200 Subject: [PATCH 16/26] perf(distributed): blockwise cross-far M2L chunking (bounds peak memory) The fused real M2L builds per-pair rotation blocks [Npairs, Bp, mdp, mdp]; running the whole cross-far list in one batch is the >200k/GPU OOM (the block tensors, not the pair count, dominate M2L peak memory). Add DistributedFMMConfig.m2l_chunk (default None = one full batch, unchanged). When set, the cross-far M2L is accumulated over ceil(Npairs/chunk) fixed-size blocks via lax.scan: each block builds its own [chunk, Bp, mdp, mdp] tiles, runs the fused kernel, masks, and segment-sums into the local expansions. Peak M2L memory becomes ~chunk pairs, independent of the total -- removing the per-GPU ceiling at some throughput cost (one kernel launch per block). Numerics-identical to the full batch: the per-block contributions are masked and summed associatively into the same accumulator. Validated on 2x A100 (per=50k, m2l_chunk=65536): aggL2 1.6904e-3 == the full-batch fp32 baseline to all printed digits. Opt-in and distributed-only (real basis), so the single-GPU fast lane and the complex path are untouched. Combine with far_m2l_fp32 for the accuracy-preserving path past the ~200k/GPU ceiling. The high-N capacity confirmation (per=400k/GPU x4 @ theta_cross=0.1, which OOMs full-batch) is queued behind GPU availability; the memory bound is by construction. Co-Authored-By: Claude Opus 4.8 --- jaccpot/distributed/fmm.py | 89 +++++++++++++++++++++++++++++++++----- 1 file changed, 78 insertions(+), 11 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index ed00e41..79282fa 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -209,6 +209,12 @@ class DistributedFMMConfig: # accumulation into the local expansions stays coeff_dtype (fp64). Opt-in; the # single-GPU fast lane is unaffected (this field is distributed-only). far_m2l_fp32: bool = False + # Chunk the cross-far M2L over pairs in blocks of this many pairs (None = one full + # batch). The fused M2L builds per-pair rotation blocks [Npairs, Bp, mdp, mdp]; doing + # the whole cross-far list at once is the >200k/GPU OOM. Scanning fixed-size blocks + # bounds peak M2L memory to ~block pairs regardless of the total, removing the per-GPU + # ceiling at some throughput cost. Numerics-identical (associative accumulation). + m2l_chunk: Optional[int] = None def with_scaled_caps( self: "DistributedFMMConfig", factor: float @@ -404,6 +410,52 @@ def partition_for_devices( } +def _chunked_real_m2l_accumulate( + loc_init: jax.Array, + src_mult: jax.Array, + deltas: jax.Array, + tgt: jax.Array, + valid: jax.Array, + *, + order: int, + block: int, + total_nodes: int, + m2l_dtype: Any, +) -> jax.Array: + """Accumulate the real-basis M2L over pairs in fixed-size blocks (bounded peak mem). + + The fused M2L builds per-pair rotation blocks ``[Npairs, Bp, mdp, mdp]``; running the + whole cross-far list in one batch is the >200k/GPU OOM. Scan over ``ceil(Npairs/block)`` + blocks so peak M2L memory is ~``block`` pairs, independent of the total. The block M2Ls + are masked and segment-summed into ``loc_init`` -- associative, so numerics-identical to + the full batch. ``src_mult``/``deltas`` are gathered [Npairs, *] (cheap, ~C or 3 per pair); + only the per-pair rotation blocks (built inside the kernel) are memory-heavy, hence blocked. + """ + C = src_mult.shape[1] + npairs = src_mult.shape[0] + nblk = -(-npairs // block) # ceil + pad_n = nblk * block - npairs + out_dtype = loc_init.dtype + sm = jnp.pad(src_mult, ((0, pad_n), (0, 0))).reshape(nblk, block, C) + dl = jnp.pad(deltas, ((0, pad_n), (0, 0))).reshape(nblk, block, 3) + tg = jnp.pad(tgt, (0, pad_n)).reshape(nblk, block) + vd = jnp.pad(valid, (0, pad_n), constant_values=False).reshape(nblk, block) + + def body(loc: jax.Array, blk): + smb, dlb, tgb, vdb = blk + contribs = _apply_real_m2l( + smb.astype(m2l_dtype), dlb.astype(m2l_dtype), order=order, m2l_impl="rot_scale" + ).astype(out_dtype) + contribs = jnp.where(vdb[:, None], contribs, 0) + loc = loc + jax.ops.segment_sum( + contribs, jnp.where(vdb, tgb, 0), total_nodes + ) + return loc, None + + loc, _ = jax.lax.scan(body, loc_init, (sm, dl, tg, vd)) + return loc + + def _make_fn(config: DistributedFMMConfig, ndev: int, cap: int) -> Callable: """Build the per-device ``shard_map`` body closing over static shapes. @@ -828,19 +880,34 @@ def _l2p(loc_coeffs): xs = jnp.where(x_valid, x_src, 0) xt = jnp.where(x_valid, x_tgt, 0) x_deltas = centers[xt] - c_centers[xs] - if is_real: - x_contribs = _apply_real_m2l( - coarse_packed_use[xs].astype(m2l_dtype), - x_deltas.astype(m2l_dtype), + if is_real and config.m2l_chunk: + # Blockwise cross-far M2L: bounds peak memory to ~m2l_chunk pairs (removes the + # >200k/GPU OOM). Numerics-identical to the full batch (associative accumulate). + loc_full = _chunked_real_m2l_accumulate( + loc_self, + coarse_packed_use[xs], + x_deltas, + xt, + x_valid, order=p, - m2l_impl="rot_scale", - ).astype(coeff_dtype) + block=int(config.m2l_chunk), + total_nodes=int(total_nodes), + m2l_dtype=m2l_dtype, + ) else: - x_contribs = m2l_complex_reference_batch( - coarse_packed_use[xs], x_deltas, order=p, rotation=rot - ).astype(cdtype) - x_contribs = jnp.where(x_valid[:, None], x_contribs, 0) - loc_full = loc_self + jax.ops.segment_sum(x_contribs, xt, int(total_nodes)) + if is_real: + x_contribs = _apply_real_m2l( + coarse_packed_use[xs].astype(m2l_dtype), + x_deltas.astype(m2l_dtype), + order=p, + m2l_impl="rot_scale", + ).astype(coeff_dtype) + else: + x_contribs = m2l_complex_reference_batch( + coarse_packed_use[xs], x_deltas, order=p, rotation=rot + ).astype(cdtype) + x_contribs = jnp.where(x_valid[:, None], x_contribs, 0) + loc_full = loc_self + jax.ops.segment_sum(x_contribs, xt, int(total_nodes)) far_full = _l2p(loc_full) # near: leaf particle-index groups into [local ; halo] buffer From 586e29a6665a35ffbc3cbcd43334ff5ff8373e04 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 09:14:36 +0000 Subject: [PATCH 17/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaccpot/distributed/fmm.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 79282fa..5c9b381 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -444,12 +444,13 @@ def _chunked_real_m2l_accumulate( def body(loc: jax.Array, blk): smb, dlb, tgb, vdb = blk contribs = _apply_real_m2l( - smb.astype(m2l_dtype), dlb.astype(m2l_dtype), order=order, m2l_impl="rot_scale" + smb.astype(m2l_dtype), + dlb.astype(m2l_dtype), + order=order, + m2l_impl="rot_scale", ).astype(out_dtype) contribs = jnp.where(vdb[:, None], contribs, 0) - loc = loc + jax.ops.segment_sum( - contribs, jnp.where(vdb, tgb, 0), total_nodes - ) + loc = loc + jax.ops.segment_sum(contribs, jnp.where(vdb, tgb, 0), total_nodes) return loc, None loc, _ = jax.lax.scan(body, loc_init, (sm, dl, tg, vd)) From 730ad17431eeee28b902167ee5b96d21964b4b2f Mon Sep 17 00:00:00 2001 From: TobiBu Date: Tue, 21 Jul 2026 11:43:30 +0200 Subject: [PATCH 18/26] fix(distributed): annotate the chunked-M2L scan body for the type-annotation guard _chunked_real_m2l_accumulate's lax.scan body was missing param/return annotations, tripping tests/unit/test_type_annotation_guard.py (the only CI failure on the chunking commit; 428/428 other tests pass). Annotate it (carry + 4-tuple block -> (carry, None)). Co-Authored-By: Claude Opus 4.8 --- jaccpot/distributed/fmm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 5c9b381..8b331a6 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -441,7 +441,10 @@ def _chunked_real_m2l_accumulate( tg = jnp.pad(tgt, (0, pad_n)).reshape(nblk, block) vd = jnp.pad(valid, (0, pad_n), constant_values=False).reshape(nblk, block) - def body(loc: jax.Array, blk): + def body( + loc: jax.Array, + blk: tuple[jax.Array, jax.Array, jax.Array, jax.Array], + ) -> tuple[jax.Array, None]: smb, dlb, tgb, vdb = blk contribs = _apply_real_m2l( smb.astype(m2l_dtype), From f14e40e01f32456627c637698f81c5c2805e0560 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Tue, 21 Jul 2026 22:32:07 +0200 Subject: [PATCH 19/26] perf(distributed): chunk the self-far M2L too (removes the >1.2M/GPU M2L ceiling) The cross-far M2L was already chunked; the self-far M2L (_accumulate_m2l_fullbatch) was still full-batch and became the next OOM -- 14.24 GiB at per=1.6M/GPU x4 on the 4-GPU max-N probe (1.2M/GPU = 4.8M total fit; 1.6M/GPU OOMs at the self-far rotation blocks). Route the real self-far M2L through the same _chunked_real_m2l_accumulate (lax.scan over fixed-size blocks) when m2l_chunk is set. Pure refactor of the full-batch path: reproduce its gather exactly -- valid mask (idx < active & src>=0 & tgt>=0), safe src/tgt, and centers cast to m2l_dtype BEFORE the tgt-minus-src delta subtraction -- then accumulate in blocks (associative masked segment-sum). Same helper the cross-far path uses, which was GPU-parity-validated numerics-identical (per=50k aggL2 1.6904e-3 == full batch). Annotation guard passes. Now both far M2Ls are bounded by the block size, so per-GPU N is gated by the near-field / tree buffers rather than the M2L. High-N GPU capacity confirmation (per>=1.6M/GPU x4) is pending shared-node GPU availability. Co-Authored-By: Claude Opus 4.8 --- jaccpot/distributed/fmm.py | 35 +++++++++++++++++++++++++++++------ 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 8b331a6..89d12e4 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -209,11 +209,12 @@ class DistributedFMMConfig: # accumulation into the local expansions stays coeff_dtype (fp64). Opt-in; the # single-GPU fast lane is unaffected (this field is distributed-only). far_m2l_fp32: bool = False - # Chunk the cross-far M2L over pairs in blocks of this many pairs (None = one full - # batch). The fused M2L builds per-pair rotation blocks [Npairs, Bp, mdp, mdp]; doing - # the whole cross-far list at once is the >200k/GPU OOM. Scanning fixed-size blocks - # bounds peak M2L memory to ~block pairs regardless of the total, removing the per-GPU - # ceiling at some throughput cost. Numerics-identical (associative accumulation). + # Chunk the far-field M2L (both self and cross) over pairs in blocks of this many + # pairs (None = one full batch). The fused M2L builds per-pair rotation blocks + # [Npairs, Bp, mdp, mdp]; doing a whole far list at once is the M2L OOM (cross-far + # >200k/GPU; self-far >1.2M/GPU). Scanning fixed-size blocks bounds peak M2L memory to + # ~block pairs regardless of the total, removing the M2L per-GPU ceiling at some + # throughput cost. Numerics-identical (associative masked segment-sum accumulation). m2l_chunk: Optional[int] = None def with_scaled_caps( @@ -835,7 +836,29 @@ def _l2p(loc_coeffs): s_active = jnp.asarray(inter.far_pair_count, INDEX_DTYPE) else: s_active = jnp.sum((s_tgt >= 0).astype(INDEX_DTYPE)) - if is_real: + if is_real and config.m2l_chunk: + # Blockwise self-far M2L (same bound as the cross-far path): the per-pair + # rotation blocks are the >1.2M/GPU OOM here too. Reproduce the fullbatch + # gather (valid mask, safe src/tgt, tgt-minus-src deltas) then scan in blocks. + s_idx = jnp.arange(s_src.shape[0], dtype=INDEX_DTYPE) + s_valid = (s_idx < s_active) & (s_src >= 0) & (s_tgt >= 0) + s_safe_src = jnp.where(s_valid, s_src, 0) + s_safe_tgt = jnp.where(s_valid, s_tgt, 0) + # Match the full-batch path exactly (pure refactor): its centers are cast to + # m2l_dtype before the delta subtraction, so do the same here. + centers_m = centers.astype(m2l_dtype) + loc_self = _chunked_real_m2l_accumulate( + zeros, + packed_use[s_safe_src], + centers_m[s_safe_tgt] - centers_m[s_safe_src], + s_safe_tgt, + s_valid, + order=p, + block=int(config.m2l_chunk), + total_nodes=int(total_nodes), + m2l_dtype=m2l_dtype, + ) + elif is_real: loc_self = _accumulate_m2l_fullbatch( zeros, packed_use.astype(m2l_dtype), From 318d9613b14243b0314ea9cb0db98aee07cafbab Mon Sep 17 00:00:00 2001 From: TobiBu Date: Wed, 22 Jul 2026 16:05:11 +0200 Subject: [PATCH 20/26] perf(distributed): index-based chunked M2L (gather per block, drop the >1GB peak) The chunked M2L pre-gathered the full [Npairs, C] source multipoles and [Npairs, 3] deltas before the scan. At high per-GPU N that gather is itself the OOM: the 4-GPU probe cleared the self-far M2L blocks (14.24GiB) but then died on a 1.72GiB _gather (packed_use[s_safe_src], ~14M self-far pairs) on an already-full GPU. Refactor _chunked_real_m2l_accumulate to take the compact multipole array + index/centre arrays and gather [block, C] (and compute the tgt-minus-src delta) INSIDE the scan. Only the per-block rotation blocks are large now; the full [Npairs, *] arrays never materialise. Both far call sites updated (self: single centre array; cross: coarse src tree/centres). Deltas are fp64 tgt-minus-src then cast -- matches the prior cross convention exactly. Parity re-confirmed: per=50k aggL2 1.6904e-3 == baseline. Co-Authored-By: Claude Opus 4.8 --- jaccpot/distributed/fmm.py | 72 +++++++++++++++++++++----------------- 1 file changed, 40 insertions(+), 32 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 89d12e4..d4aacb5 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -413,9 +413,11 @@ def partition_for_devices( def _chunked_real_m2l_accumulate( loc_init: jax.Array, - src_mult: jax.Array, - deltas: jax.Array, - tgt: jax.Array, + multip: jax.Array, + src_idx: jax.Array, + tgt_idx: jax.Array, + src_centers: jax.Array, + tgt_centers: jax.Array, valid: jax.Array, *, order: int, @@ -425,39 +427,40 @@ def _chunked_real_m2l_accumulate( ) -> jax.Array: """Accumulate the real-basis M2L over pairs in fixed-size blocks (bounded peak mem). - The fused M2L builds per-pair rotation blocks ``[Npairs, Bp, mdp, mdp]``; running the - whole cross-far list in one batch is the >200k/GPU OOM. Scan over ``ceil(Npairs/block)`` - blocks so peak M2L memory is ~``block`` pairs, independent of the total. The block M2Ls - are masked and segment-summed into ``loc_init`` -- associative, so numerics-identical to - the full batch. ``src_mult``/``deltas`` are gathered [Npairs, *] (cheap, ~C or 3 per pair); - only the per-pair rotation blocks (built inside the kernel) are memory-heavy, hence blocked. + The fused M2L builds per-pair rotation blocks ``[Npairs, Bp, mdp, mdp]``; running a whole + far list in one batch is the M2L OOM. Scan over ``ceil(Npairs/block)`` blocks so peak M2L + memory is ~``block`` pairs, independent of the total. The block M2Ls are masked and + segment-summed into ``loc_init`` -- associative, so numerics-identical to the full batch. + + Index-based (not pre-gathered): only the compact ``multip`` (``[total_nodes, C]``), the + index arrays, and the centre arrays are resident; the per-pair source multipoles and deltas + are gathered ``[block, *]`` INSIDE the scan. This avoids materialising the full + ``[Npairs, C]`` gather (itself a >1GB peak at high N), so the only large intermediates are + the per-block rotation blocks. Deltas are ``tgt - src`` in the source dtype then cast. """ - C = src_mult.shape[1] - npairs = src_mult.shape[0] + npairs = src_idx.shape[0] nblk = -(-npairs // block) # ceil pad_n = nblk * block - npairs out_dtype = loc_init.dtype - sm = jnp.pad(src_mult, ((0, pad_n), (0, 0))).reshape(nblk, block, C) - dl = jnp.pad(deltas, ((0, pad_n), (0, 0))).reshape(nblk, block, 3) - tg = jnp.pad(tgt, (0, pad_n)).reshape(nblk, block) + si = jnp.pad(src_idx, (0, pad_n)).reshape(nblk, block) + ti = jnp.pad(tgt_idx, (0, pad_n)).reshape(nblk, block) vd = jnp.pad(valid, (0, pad_n), constant_values=False).reshape(nblk, block) def body( loc: jax.Array, - blk: tuple[jax.Array, jax.Array, jax.Array, jax.Array], + blk: tuple[jax.Array, jax.Array, jax.Array], ) -> tuple[jax.Array, None]: - smb, dlb, tgb, vdb = blk + sib, tib, vdb = blk + smb = multip[sib].astype(m2l_dtype) + dlb = (tgt_centers[tib] - src_centers[sib]).astype(m2l_dtype) contribs = _apply_real_m2l( - smb.astype(m2l_dtype), - dlb.astype(m2l_dtype), - order=order, - m2l_impl="rot_scale", + smb, dlb, order=order, m2l_impl="rot_scale" ).astype(out_dtype) contribs = jnp.where(vdb[:, None], contribs, 0) - loc = loc + jax.ops.segment_sum(contribs, jnp.where(vdb, tgb, 0), total_nodes) + loc = loc + jax.ops.segment_sum(contribs, jnp.where(vdb, tib, 0), total_nodes) return loc, None - loc, _ = jax.lax.scan(body, loc_init, (sm, dl, tg, vd)) + loc, _ = jax.lax.scan(body, loc_init, (si, ti, vd)) return loc @@ -844,14 +847,16 @@ def _l2p(loc_coeffs): s_valid = (s_idx < s_active) & (s_src >= 0) & (s_tgt >= 0) s_safe_src = jnp.where(s_valid, s_src, 0) s_safe_tgt = jnp.where(s_valid, s_tgt, 0) - # Match the full-batch path exactly (pure refactor): its centers are cast to - # m2l_dtype before the delta subtraction, so do the same here. - centers_m = centers.astype(m2l_dtype) + # Index-based: gather the source multipoles + deltas inside the scan (avoids the + # full [Npairs, C] gather that otherwise peaks >1GB at high N). Self tree = single + # centre array for both src and tgt. loc_self = _chunked_real_m2l_accumulate( zeros, - packed_use[s_safe_src], - centers_m[s_safe_tgt] - centers_m[s_safe_src], + packed_use, + s_safe_src, s_safe_tgt, + centers, + centers, s_valid, order=p, block=int(config.m2l_chunk), @@ -906,15 +911,17 @@ def _l2p(loc_coeffs): x_valid = x_tgt >= 0 xs = jnp.where(x_valid, x_src, 0) xt = jnp.where(x_valid, x_tgt, 0) - x_deltas = centers[xt] - c_centers[xs] if is_real and config.m2l_chunk: - # Blockwise cross-far M2L: bounds peak memory to ~m2l_chunk pairs (removes the - # >200k/GPU OOM). Numerics-identical to the full batch (associative accumulate). + # Blockwise cross-far M2L, index-based (gathers source multipoles + deltas per + # block, so no full [Npairs, C] materialisation). Bounds peak memory; numerics- + # identical to the full batch (associative accumulate). src=coarse tree/centres. loc_full = _chunked_real_m2l_accumulate( loc_self, - coarse_packed_use[xs], - x_deltas, + coarse_packed_use, + xs, xt, + c_centers, + centers, x_valid, order=p, block=int(config.m2l_chunk), @@ -922,6 +929,7 @@ def _l2p(loc_coeffs): m2l_dtype=m2l_dtype, ) else: + x_deltas = centers[xt] - c_centers[xs] if is_real: x_contribs = _apply_real_m2l( coarse_packed_use[xs].astype(m2l_dtype), From 292cac02f63dee6f940912146db21650c0dabf55 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Wed, 22 Jul 2026 16:06:44 +0200 Subject: [PATCH 21/26] perf(distributed): extrapolate cap presets to unmeasured N (scaled nearest seed) cap_presets.lookup returned None when no preset was >= the requested per-GPU N, so a new N calibrated from the small defaults -- the full auto_scale doubling ladder = ~10-12 shard_map RECOMPILES (~200s each, GPU idle), ~30-40 min per point on the capacity probe. Add a fallback: when only smaller presets exist, seed from the nearest smaller one SCALED UP by the per-GPU-N ratio (ceil, None-preserving). auto_scale still refines, so the seed only needs to be close -- a new N now converges in ~1 retry instead of ~12. Exact-match and nearest-larger behaviour unchanged. Co-Authored-By: Claude Opus 4.8 --- jaccpot/distributed/cap_presets.py | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/jaccpot/distributed/cap_presets.py b/jaccpot/distributed/cap_presets.py index aa5d013..ec3f55b 100644 --- a/jaccpot/distributed/cap_presets.py +++ b/jaccpot/distributed/cap_presets.py @@ -69,20 +69,37 @@ def save_presets(path: str, presets: dict) -> None: 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). None if nothing usable.""" + 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"] - larger = [ + same = [ (int(kk.split(":")[0]), v["caps"]) for kk, v in presets.items() - if kk.endswith(f":{int(ndev)}") and int(kk.split(":")[0]) >= int(per_gpu_n) + 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] - return None + n0, c0 = max(same, key=lambda t: t[0]) # largest smaller + return _scale_caps(c0, int(per_gpu_n), n0) def record( From 835db8e43b0e71f08309e74d62fccc93e99e8f51 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 14:07:40 +0000 Subject: [PATCH 22/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaccpot/distributed/fmm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index d4aacb5..51a5958 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -453,9 +453,9 @@ def body( sib, tib, vdb = blk smb = multip[sib].astype(m2l_dtype) dlb = (tgt_centers[tib] - src_centers[sib]).astype(m2l_dtype) - contribs = _apply_real_m2l( - smb, dlb, order=order, m2l_impl="rot_scale" - ).astype(out_dtype) + contribs = _apply_real_m2l(smb, dlb, order=order, m2l_impl="rot_scale").astype( + out_dtype + ) contribs = jnp.where(vdb[:, None], contribs, 0) loc = loc + jax.ops.segment_sum(contribs, jnp.where(vdb, tib, 0), total_nodes) return loc, None From 92c1add4d49bdd8bb2f59c3cb9916548fd73cbcd Mon Sep 17 00:00:00 2001 From: TobiBu Date: Wed, 22 Jul 2026 18:10:14 +0200 Subject: [PATCH 23/26] perf(distributed): int32 near-field densification indices (~2GB peak cut at high N) INDEX_DTYPE is int64, so the near-field CSR/densification materialised several [num_edges] int64 temporaries (~0.8GB each at >1M/GPU) plus the bincount scatter (the 1.27GiB alloc that OOM'd the 1.6M/GPU x4 probe). Narrow the pure index temporaries to int32 -- the argsort/bincount target tgt (leaf-row ids <= u_leaves) and the densification e/t_of_e/ rank/flat (< u_leaves*S_near) -- all << 2^31 at any feasible per-GPU N. Kernel-facing arrays (offsets/src_s/counts/sids) stay INDEX_DTYPE, so no kernel dtype risk. Numerically exact (int32 holds the same index values): per=50k parity aggL2 1.6904e-3 == baseline. Frees ~2GiB of the near-field peak; the 1.6M/GPU probe now clears the bincount and reaches a later 5GiB near-field-P2P alloc. Meaningfully exceeding ~1.2M/GPU still needs near-field leaf-block chunking (the whole near-field working set), not per-buffer narrowing. Co-Authored-By: Claude Opus 4.8 --- jaccpot/distributed/fmm.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 51a5958..96f3858 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -578,7 +578,9 @@ def _combined_neighbors(tree, nbr, cross, rct, halo, n_halo_rows, local_only=Fal src = jnp.concatenate([l_src_row, r_src_row]) valid = jnp.concatenate([l_valid, r_valid]) sentinel = jnp.asarray(u_leaves, INDEX_DTYPE) - tgt = jnp.where(valid, tgt, sentinel) + # tgt holds target leaf-row ids (<= u_leaves << 2^31), so int32 halves the argsort + # value array and the bincount scatter (the 1.27GiB near-field alloc at high N). + tgt = jnp.where(valid, tgt, sentinel).astype(jnp.int32) src = jnp.where(valid, src, sentinel) srt = jnp.argsort(tgt) @@ -986,11 +988,18 @@ def _l2p(loc_coeffs): ) # Densify the combined CSR (offsets/src_s/counts, sorted by target) into a # padded [u_leaves, S_near] source-leaf-id table + validity mask. - e = jnp.arange(src_s.shape[0], dtype=INDEX_DTYPE) - t_of_e = jnp.searchsorted(offsets, e, side="right") - 1 - rank = e - offsets[jnp.clip(t_of_e, 0, u_leaves)] + # int32 index math: e/t_of_e/rank/flat are pure [num_edges] index temporaries + # (values < u_leaves*S_near << 2^31 at any feasible per-GPU N), so int32 halves + # the near-field densification peak (the aggregate wall at >1.2M/GPU). Kernel- + # facing arrays (offsets/src_s/counts/sids) stay INDEX_DTYPE. + _i32 = jnp.int32 + e = jnp.arange(src_s.shape[0], dtype=_i32) + t_of_e = (jnp.searchsorted(offsets, e, side="right") - 1).astype(_i32) + rank = (e - offsets[jnp.clip(t_of_e, 0, u_leaves)].astype(_i32)).astype(_i32) in_range = (t_of_e >= 0) & (t_of_e < u_leaves) & (rank < S_near) - flat = jnp.where(in_range, t_of_e * S_near + rank, -1) + flat = jnp.where( + in_range, t_of_e * _i32(S_near) + rank, _i32(-1) + ) sids = ( jnp.zeros((u_leaves * S_near,), INDEX_DTYPE) .at[flat] From 01a9c81a41624958b37ca9b5be5c2dba4d6ca7a5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 16:11:01 +0000 Subject: [PATCH 24/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaccpot/distributed/fmm.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 96f3858..b394ed5 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -995,11 +995,11 @@ def _l2p(loc_coeffs): _i32 = jnp.int32 e = jnp.arange(src_s.shape[0], dtype=_i32) t_of_e = (jnp.searchsorted(offsets, e, side="right") - 1).astype(_i32) - rank = (e - offsets[jnp.clip(t_of_e, 0, u_leaves)].astype(_i32)).astype(_i32) - in_range = (t_of_e >= 0) & (t_of_e < u_leaves) & (rank < S_near) - flat = jnp.where( - in_range, t_of_e * _i32(S_near) + rank, _i32(-1) + rank = (e - offsets[jnp.clip(t_of_e, 0, u_leaves)].astype(_i32)).astype( + _i32 ) + in_range = (t_of_e >= 0) & (t_of_e < u_leaves) & (rank < S_near) + flat = jnp.where(in_range, t_of_e * _i32(S_near) + rank, _i32(-1)) sids = ( jnp.zeros((u_leaves * S_near,), INDEX_DTYPE) .at[flat] From 4149a51dddece1d1631aa62aa6aff68b8acad246 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Wed, 22 Jul 2026 20:35:01 +0200 Subject: [PATCH 25/26] perf(distributed): near-field leaf-block chunking (bound the near-field peak) The fused-Pallas near field densified the combined CSR into [u_leaves, S_near] tables + [num_edges] temporaries and ran the leafpair P2P over all u_leaves at once -- the ~40GB aggregate wall that capped the driver at ~1.2M particles/GPU (the M2L was already chunked; the near field was the remaining limiter). Chunk it over blocks of TARGET leaves, mirroring the M2L chunk (_chunked_real_m2l_accumulate). - DistributedFMMConfig.nearfield_chunk (Optional[int], default None): target-leaf block size for the pallas near field; None keeps the full-batch path. Applies only when use_pallas_near. - nearfield_leafpair_pallas_decoupled (pallas/nearfield_fused_leaf.py): new sibling wrapper that takes the target rows and the source gather pool as SEPARATE arrays (the kernel already reads them as separate refs; the old wrapper just bound the same array to both). Kernel body and all existing callers unchanged. - _radix_fast_lane_prepacked_pallas_decoupled (nearfield/near_field.py): sibling of the prepacked path calling the decoupled wrapper. - _chunked_pallas_nearfield_accumulate (distributed/fmm.py): lax.scan over ceil(n_lloc/block) target-leaf blocks. Per block: build [block, S_near] sids DIRECTLY from the CSR (sids[i,j]=src_s[offsets[b0+i]+j] for j1.2M/GPU) measured separately. Co-Authored-By: Claude Opus 4.8 --- jaccpot/distributed/fmm.py | 216 +++++++++++++++++++------ jaccpot/nearfield/near_field.py | 89 ++++++++++ jaccpot/pallas/nearfield_fused_leaf.py | 120 ++++++++++++++ 3 files changed, 377 insertions(+), 48 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index b394ed5..4e10a99 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -70,6 +70,7 @@ _compute_leaf_p2p_prepared_large_n_self_only_impl, _prepare_leaf_data_from_groups, _radix_fast_lane_prepacked_pallas, + _radix_fast_lane_prepacked_pallas_decoupled, compute_leaf_p2p_accelerations, ) from jaccpot.operators.complex_ops import ( @@ -216,6 +217,15 @@ class DistributedFMMConfig: # ~block pairs regardless of the total, removing the M2L per-GPU ceiling at some # throughput cost. Numerics-identical (associative masked segment-sum accumulation). m2l_chunk: Optional[int] = None + # Chunk the fused-Pallas NEAR field over blocks of this many TARGET leaves (None = one + # full batch). The near field densifies the combined CSR into [u_leaves, S_near] tables + # + [num_edges] temporaries and runs the leafpair P2P over all u_leaves at once -- the + # ~40GB aggregate wall at >1.2M/GPU. Scanning fixed-size target-leaf blocks bounds both + # the densification and the P2P peak to ~block*S_near (source pool stays resident), + # removing the near-field per-GPU ceiling at some throughput cost. Numerics-identical + # (disjoint target blocks + scatter-add). Applies ONLY when the pallas near-field is + # active (nearfield_backend "pallas", or "auto" on Ampere+); ignored for baseline. + nearfield_chunk: Optional[int] = None def with_scaled_caps( self: "DistributedFMMConfig", factor: float @@ -464,6 +474,95 @@ def body( return loc +def _chunked_pallas_nearfield_accumulate( + leaf_positions: jax.Array, + leaf_masses: jax.Array, + leaf_mask_g: jax.Array, + safe_idx: jax.Array, + concat_pos: jax.Array, + offsets: jax.Array, + counts: jax.Array, + src_s: jax.Array, + *, + n_lloc: int, + S_near: int, + block: int, + G: Any, + softening_sq: jax.Array, +) -> jax.Array: + """Fused-Pallas near field over blocks of TARGET leaves (bounded peak memory). + + Full-batch near field densifies the combined CSR into ``[u_leaves, S_near]`` tables (+ + ``[num_edges]`` temporaries) and runs the leafpair P2P over all ``u_leaves`` at once -- the + ~40GB aggregate wall at >1.2M/GPU. This scans ``ceil(n_lloc/block)`` target-leaf blocks, + building each block's ``[block, S_near]`` densification + running self + pairs on the block, + so peak drops from ``u_leaves*S_near`` to ``block*S_near``. The SOURCE gather pool + (``leaf_positions``/``leaf_masses``/``leaf_mask_g``) stays fully resident (sources span all + ``u_leaves`` rows). Only target rows ``[0, n_lloc)`` produce kept output; per-block partials + scatter-add into a ``[cap+halo, 3]`` accumulator, so it is numerics-identical to the full batch + (disjoint target blocks -> exactly identical, not merely associative). + + The per-block densification is built DIRECTLY from the CSR (``offsets``/``counts``/``src_s``): + ``sids_blk[i, j] = src_s[offsets[b0+i] + j]`` for ``j < min(counts[b0+i], S_near)`` -- the first + (up to) ``S_near`` sources of each target leaf, exactly what the full densification keeps (rank + ``< S_near``). This is ``[block, S_near]`` throughout (no ``[num_edges]`` temporaries) and is + bit-identical to the full batch: same ``src_s`` order, same per-leaf truncation, so raw per-leaf + counts exceeding ``S_near`` (offsets uncapped, counts from bincount) are truncated identically. + """ + nblk = -(-n_lloc // block) # ceil + num_edges = int(src_s.shape[0]) + near0 = jnp.zeros_like(concat_pos) + cols = jnp.arange(S_near, dtype=INDEX_DTYPE) # [S_near] + src_zero = jnp.asarray(0, src_s.dtype) + + def _body(near: jax.Array, k: jax.Array) -> tuple[jax.Array, None]: + b0 = k * block + rows = b0 + jnp.arange(block, dtype=INDEX_DTYPE) + valid_t = rows < n_lloc + safe_rows = jnp.where(valid_t, rows, 0) + tgt_pos = leaf_positions[safe_rows] # [block, W, 3] + tgt_mass = leaf_masses[safe_rows] # [block, W] + tgt_mask = leaf_mask_g[safe_rows] & valid_t[:, None] # [block, W] + tgt_idx = safe_idx[safe_rows] # [block, W] global scatter targets + + # self term (block): pure-JAX per-leaf; scatters to global-order [cap+halo, 3]. + self_blk = _compute_leaf_p2p_prepared_large_n_self_only_impl( + concat_pos, tgt_pos, tgt_mass, tgt_mask, tgt_idx, G=G, softening_sq=softening_sq + ) + + # per-block densification built DIRECTLY from the CSR (no [num_edges] window): + # sids_blk[i, j] = src_s[offsets[b0+i] + j] for j < min(counts[b0+i], S_near). + leaf_off = offsets[safe_rows] # [block] each target leaf's edge start + leaf_cnt = counts[safe_rows] # [block] raw near count per target leaf + keep_col = (cols[None, :] < jnp.minimum(leaf_cnt, S_near)[:, None]) & valid_t[ + :, None + ] # [block, S_near] + edge_idx = leaf_off[:, None] + cols[None, :] # [block, S_near] + safe_edge = jnp.clip(edge_idx, 0, num_edges - 1) + sids_blk = jnp.where(keep_col, src_s[safe_edge], src_zero) # [block, S_near] + svalid_blk = keep_col + + # pair term (block): full source pool resident, target restricted to the block. + pair_blk = _radix_fast_lane_prepacked_pallas_decoupled( + sids_blk[:, :, None], + svalid_blk[:, :, None], + tgt_pos, + tgt_mask, + tgt_idx, + leaf_positions, + leaf_masses, + leaf_mask_g, + concat_pos, + G=G, + softening_sq=softening_sq, + compute_potential=False, + ) + return near + self_blk + pair_blk, None + + near, _ = jax.lax.scan(_body, near0, jnp.arange(nblk, dtype=INDEX_DTYPE)) + return near + + def _make_fn(config: DistributedFMMConfig, ndev: int, cap: int) -> Callable: """Build the per-device ``shard_map`` body closing over static shapes. @@ -977,54 +1076,75 @@ def _l2p(loc_coeffs): leaf_mask_g, safe_idx, ) = _prepare_leaf_data_from_groups(lp_idx, lp_mask, concat_pos, concat_mass) - self_acc = _compute_leaf_p2p_prepared_large_n_self_only_impl( - concat_pos, - leaf_positions, - leaf_masses, - leaf_mask_g, - safe_idx, - G=G, - softening_sq=soft2_a, - ) - # Densify the combined CSR (offsets/src_s/counts, sorted by target) into a - # padded [u_leaves, S_near] source-leaf-id table + validity mask. - # int32 index math: e/t_of_e/rank/flat are pure [num_edges] index temporaries - # (values < u_leaves*S_near << 2^31 at any feasible per-GPU N), so int32 halves - # the near-field densification peak (the aggregate wall at >1.2M/GPU). Kernel- - # facing arrays (offsets/src_s/counts/sids) stay INDEX_DTYPE. - _i32 = jnp.int32 - e = jnp.arange(src_s.shape[0], dtype=_i32) - t_of_e = (jnp.searchsorted(offsets, e, side="right") - 1).astype(_i32) - rank = (e - offsets[jnp.clip(t_of_e, 0, u_leaves)].astype(_i32)).astype( - _i32 - ) - in_range = (t_of_e >= 0) & (t_of_e < u_leaves) & (rank < S_near) - flat = jnp.where(in_range, t_of_e * _i32(S_near) + rank, _i32(-1)) - sids = ( - jnp.zeros((u_leaves * S_near,), INDEX_DTYPE) - .at[flat] - .set(src_s, mode="drop") - .reshape(u_leaves, S_near) - ) - svalid = ( - jnp.zeros((u_leaves * S_near,), bool) - .at[flat] - .set(jnp.ones_like(flat, bool), mode="drop") - .reshape(u_leaves, S_near) - ) - pair_acc = _radix_fast_lane_prepacked_pallas( - sids.reshape(u_leaves, S_near, 1), - svalid.reshape(u_leaves, S_near, 1), - leaf_positions, - leaf_masses, - leaf_mask_g, - safe_idx, - concat_pos, - G=G, - softening_sq=soft2_a, - compute_potential=False, - ) - near_full = self_acc + pair_acc + if config.nearfield_chunk: + # Bounded-memory path: scan target-leaf blocks (densification + leafpair + # P2P per block) instead of materialising the full [u_leaves, S_near] tables + # and running all u_leaves at once. Numerics-identical (disjoint target + # blocks + scatter-add). See _chunked_pallas_nearfield_accumulate. + near_full = _chunked_pallas_nearfield_accumulate( + leaf_positions, + leaf_masses, + leaf_mask_g, + safe_idx, + concat_pos, + offsets, + counts, + src_s, + n_lloc=int(loc_idx.shape[0]), + S_near=S_near, + block=int(config.nearfield_chunk), + G=G, + softening_sq=soft2_a, + ) + else: + self_acc = _compute_leaf_p2p_prepared_large_n_self_only_impl( + concat_pos, + leaf_positions, + leaf_masses, + leaf_mask_g, + safe_idx, + G=G, + softening_sq=soft2_a, + ) + # Densify the combined CSR (offsets/src_s/counts, sorted by target) into a + # padded [u_leaves, S_near] source-leaf-id table + validity mask. + # int32 index math: e/t_of_e/rank/flat are pure [num_edges] index temporaries + # (values < u_leaves*S_near << 2^31 at any feasible per-GPU N), so int32 halves + # the near-field densification peak (the aggregate wall at >1.2M/GPU). Kernel- + # facing arrays (offsets/src_s/counts/sids) stay INDEX_DTYPE. + _i32 = jnp.int32 + e = jnp.arange(src_s.shape[0], dtype=_i32) + t_of_e = (jnp.searchsorted(offsets, e, side="right") - 1).astype(_i32) + rank = (e - offsets[jnp.clip(t_of_e, 0, u_leaves)].astype(_i32)).astype( + _i32 + ) + in_range = (t_of_e >= 0) & (t_of_e < u_leaves) & (rank < S_near) + flat = jnp.where(in_range, t_of_e * _i32(S_near) + rank, _i32(-1)) + sids = ( + jnp.zeros((u_leaves * S_near,), INDEX_DTYPE) + .at[flat] + .set(src_s, mode="drop") + .reshape(u_leaves, S_near) + ) + svalid = ( + jnp.zeros((u_leaves * S_near,), bool) + .at[flat] + .set(jnp.ones_like(flat, bool), mode="drop") + .reshape(u_leaves, S_near) + ) + pair_acc = _radix_fast_lane_prepacked_pallas( + sids.reshape(u_leaves, S_near, 1), + svalid.reshape(u_leaves, S_near, 1), + leaf_positions, + leaf_masses, + leaf_mask_g, + safe_idx, + concat_pos, + G=G, + softening_sq=soft2_a, + compute_potential=False, + ) + near_full = self_acc + pair_acc else: near_full = compute_leaf_p2p_accelerations( tree, diff --git a/jaccpot/nearfield/near_field.py b/jaccpot/nearfield/near_field.py index 624401e..0b1cfeb 100644 --- a/jaccpot/nearfield/near_field.py +++ b/jaccpot/nearfield/near_field.py @@ -3219,6 +3219,95 @@ def _radix_fast_lane_prepacked_pallas( return pair_acc +def _radix_fast_lane_prepacked_pallas_decoupled( + source_leaf_ids_padded: Array, + source_valid_mask_padded: Array, + target_positions: Array, + target_mask: Array, + target_particle_idx: Array, + source_positions: Array, + source_masses: Array, + source_mask: Array, + positions: Array, + *, + G: Union[float, Array], + softening_sq: Array, + compute_potential: bool = False, + num_warps: Optional[int] = None, + num_stages: int = 1, + target_subtile: Optional[int] = None, + interpret: bool = False, +) -> Union[Array, Tuple[Array, Array]]: + """Decoupled twin of :func:`_radix_fast_lane_prepacked_pallas`. + + The TARGET leaves (``target_positions``/``target_mask``/``target_particle_idx``, a block of + ``[num_targets, W, *]``) are separate from the full SOURCE gather pool + (``source_positions``/``source_masses``/``source_mask``, ``[num_sources, W, *]``). Source-leaf + ids in ``source_leaf_ids_padded`` reference global source rows in ``[0, num_sources)``. This + lets a caller run one block of target leaves against the full source pool (near-field leaf-block + chunking). Scatters the block's per-particle accel into a ``zeros_like(positions)`` buffer via + the global ``target_particle_idx`` (scatter-add) so per-block partials compose by ``+``. Target + masses are unused by the pair term. The intra-leaf self term is handled separately by the caller. + """ + from jaccpot.pallas.nearfield_fused_leaf import ( + nearfield_leafpair_pallas_decoupled, + ) + + dtype = positions.dtype + g_const = jnp.asarray(G, dtype=dtype) + + num_targets = int(source_leaf_ids_padded.shape[0]) + num_source_slots = int(source_leaf_ids_padded.shape[1]) * int( + source_leaf_ids_padded.shape[2] + ) + + accelerations = jnp.zeros_like(positions) + empty = ( + num_targets == 0 + or num_source_slots == 0 + or int(target_positions.shape[1]) == 0 + or int(source_positions.shape[0]) == 0 + ) + if empty: + if compute_potential: + return accelerations, jnp.zeros(positions.shape[:1], dtype=dtype) + return accelerations + + source_leaf_ids_flat = source_leaf_ids_padded.reshape( + (num_targets, num_source_slots) + ) + source_valid_flat = source_valid_mask_padded.reshape( + (num_targets, num_source_slots) + ) + + out = nearfield_leafpair_pallas_decoupled( + target_positions, + target_mask, + source_positions, + source_masses, + source_mask, + source_leaf_ids_flat, + source_valid_flat, + softening_sq=softening_sq, + G=g_const, + num_warps=num_warps, + num_stages=num_stages, + target_subtile=target_subtile, + interpret=interpret, + ) + + pair_acc = _scatter_contributions( + accelerations, target_particle_idx, out[..., :3], target_mask + ) + if compute_potential: + potentials = jnp.zeros(positions.shape[:1], dtype=dtype) + pair_pot = _scatter_scalar_contributions( + potentials, target_particle_idx, out[..., 3], target_mask + ) + return pair_acc, pair_pot + return pair_acc + + def compute_leaf_p2p_accelerations_radix_fast_lane( *, positions_sorted: Array, diff --git a/jaccpot/pallas/nearfield_fused_leaf.py b/jaccpot/pallas/nearfield_fused_leaf.py index 81b7cd2..f78d4e3 100644 --- a/jaccpot/pallas/nearfield_fused_leaf.py +++ b/jaccpot/pallas/nearfield_fused_leaf.py @@ -582,6 +582,125 @@ def _kernel(*refs): return out +def nearfield_leafpair_pallas_decoupled( + target_positions: Array, + target_mask: Array, + source_positions: Array, + source_masses: Array, + source_mask: Array, + source_leaf_ids: Array, + source_valid: Array, + *, + softening_sq: Array, + G: Array, + num_warps: int | None = None, + num_stages: int = 1, + target_subtile: int | None = None, + interpret: bool = False, +) -> Array: + """Leaf-pair near-field with the target set decoupled from the source gather pool. + + Identical kernel/math to :func:`nearfield_leafpair_pallas`, but the target rows + (``target_positions``/``target_mask``, shape ``[num_targets, W, 3]`` / ``[num_targets, W]``) + and the source gather tables (``source_positions``/``source_masses``/``source_mask``, shape + ``[num_sources, W, *]``) are SEPARATE arrays. ``source_leaf_ids``/``source_valid`` are + ``[num_targets, S]`` and reference source rows by global id in ``[0, num_sources)``. This + lets a caller compute a BLOCK of target leaves while keeping the full source pool resident + (the near-field leaf-block chunking used by the distributed driver). Passing the same array + as both target and source reproduces :func:`nearfield_leafpair_pallas` bit-for-bit. + + Returns ``[num_targets, W, _OUT_WIDTH]`` (accel lanes 0:3, potential lane 3). + """ + + if pl is None or plgpu is None: + raise RuntimeError("jax.experimental.pallas is not available") + + target_positions = jnp.asarray(target_positions) + dtype = target_positions.dtype + target_mask = jnp.asarray(target_mask, dtype=bool) + source_positions = jnp.asarray(source_positions, dtype=dtype) + source_masses = jnp.asarray(source_masses, dtype=dtype) + source_mask = jnp.asarray(source_mask, dtype=bool) + source_leaf_ids = jnp.asarray(source_leaf_ids) + source_valid = jnp.asarray(source_valid, dtype=bool) + softening_sq_arr = jnp.asarray([softening_sq], dtype=dtype) + g_arr = jnp.asarray([G], dtype=dtype) + + if target_positions.ndim != 3 or target_positions.shape[-1] != 3: + raise ValueError("target_positions must have shape (num_targets, W, 3)") + if source_positions.ndim != 3 or source_positions.shape[-1] != 3: + raise ValueError("source_positions must have shape (num_sources, W, 3)") + + num_targets = int(target_positions.shape[0]) + leaf_width = int(target_positions.shape[1]) + num_sources = int(source_positions.shape[0]) + num_source_slots = int(source_leaf_ids.shape[1]) + + if num_targets == 0 or leaf_width == 0 or num_source_slots == 0 or num_sources == 0: + return jnp.zeros((num_targets, leaf_width, _OUT_WIDTH), dtype=dtype) + + tgt_pos_padded = jnp.pad(target_positions, ((0, 0), (0, 0), (0, _POS_WIDTH - 3))) + src_pos_padded = jnp.pad(source_positions, ((0, 0), (0, 0), (0, _POS_WIDTH - 3))) + + bt = _resolve_subtile(target_subtile, leaf_width) + width_pad = ((leaf_width + bt - 1) // bt) * bt + n_sub = width_pad // bt + pad_t = width_pad - leaf_width + + tgt_pos_padded = ( + jnp.pad(tgt_pos_padded, ((0, 0), (0, pad_t), (0, 0))) if pad_t else tgt_pos_padded + ) + tgt_mask_padded = jnp.pad(target_mask, ((0, 0), (0, pad_t))) + + if num_warps is None: + num_warps = max(1, bt // 32) + + def _kernel(*refs: object) -> None: + return _nearfield_leafpair_kernel( + *refs, num_source_slots=num_source_slots, leaf_width=leaf_width + ) + + kernel = pl.pallas_call( + _kernel, + out_shape=jax.ShapeDtypeStruct((num_targets, width_pad, _OUT_WIDTH), dtype), + in_specs=[ + pl.BlockSpec((1, bt, _POS_WIDTH), lambda leaf, sub: (leaf, sub, 0)), + pl.BlockSpec((1, bt), lambda leaf, sub: (leaf, sub)), + # Full source gather tables (indexed by data-dependent global source leaf id). + pl.BlockSpec( + (num_sources, leaf_width, _POS_WIDTH), lambda leaf, sub: (0, 0, 0) + ), + pl.BlockSpec((num_sources, leaf_width), lambda leaf, sub: (0, 0)), + pl.BlockSpec((num_sources, leaf_width), lambda leaf, sub: (0, 0)), + pl.BlockSpec((1, num_source_slots), lambda leaf, sub: (leaf, 0)), + pl.BlockSpec((1, num_source_slots), lambda leaf, sub: (leaf, 0)), + pl.BlockSpec((1,), lambda leaf, sub: (0,)), + pl.BlockSpec((1,), lambda leaf, sub: (0,)), + ], + out_specs=pl.BlockSpec((1, bt, _OUT_WIDTH), lambda leaf, sub: (leaf, sub, 0)), + grid=(num_targets, n_sub), + compiler_params=plgpu.CompilerParams( + num_warps=int(num_warps), num_stages=int(num_stages) + ), + interpret=bool(interpret), + name=f"nearfield_leafpair_dec_t{bt}_s{num_source_slots}_w{leaf_width}", + ) + out = kernel( + tgt_pos_padded, + tgt_mask_padded, + src_pos_padded, + source_masses, + source_mask, + source_leaf_ids, + source_valid, + softening_sq_arr, + g_arr, + ) + if pad_t: + out = out[:, :leaf_width, :] + return out + + __all__ = [ "nearfield_fused_leaf", "nearfield_fused_leaf_backend", @@ -589,5 +708,6 @@ def _kernel(*refs): "nearfield_fused_leaf_pallas", "nearfield_leafpair_jax", "nearfield_leafpair_pallas", + "nearfield_leafpair_pallas_decoupled", "pallas_nearfield_fused_supported", ] From 861417e1c4d17ba79939012e011e3ed6bf76ddad Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 18:39:13 +0000 Subject: [PATCH 26/26] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- jaccpot/distributed/fmm.py | 8 +++++++- jaccpot/pallas/nearfield_fused_leaf.py | 4 +++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/jaccpot/distributed/fmm.py b/jaccpot/distributed/fmm.py index 4e10a99..38c4c42 100644 --- a/jaccpot/distributed/fmm.py +++ b/jaccpot/distributed/fmm.py @@ -527,7 +527,13 @@ def _body(near: jax.Array, k: jax.Array) -> tuple[jax.Array, None]: # self term (block): pure-JAX per-leaf; scatters to global-order [cap+halo, 3]. self_blk = _compute_leaf_p2p_prepared_large_n_self_only_impl( - concat_pos, tgt_pos, tgt_mass, tgt_mask, tgt_idx, G=G, softening_sq=softening_sq + concat_pos, + tgt_pos, + tgt_mass, + tgt_mask, + tgt_idx, + G=G, + softening_sq=softening_sq, ) # per-block densification built DIRECTLY from the CSR (no [num_edges] window): diff --git a/jaccpot/pallas/nearfield_fused_leaf.py b/jaccpot/pallas/nearfield_fused_leaf.py index f78d4e3..a88fd49 100644 --- a/jaccpot/pallas/nearfield_fused_leaf.py +++ b/jaccpot/pallas/nearfield_fused_leaf.py @@ -648,7 +648,9 @@ def nearfield_leafpair_pallas_decoupled( pad_t = width_pad - leaf_width tgt_pos_padded = ( - jnp.pad(tgt_pos_padded, ((0, 0), (0, pad_t), (0, 0))) if pad_t else tgt_pos_padded + jnp.pad(tgt_pos_padded, ((0, 0), (0, pad_t), (0, 0))) + if pad_t + else tgt_pos_padded ) tgt_mask_padded = jnp.pad(target_mask, ((0, 0), (0, pad_t)))