feat(m2l): target-tiled CSR Pallas M2L kernel with on-chip rotations (opt-in lane)

bench/gpu_gate.py

@@ -93,10 +93,12 @@
     "test_the_production_complex_fused_m2l_kernel_carries_the_axis_derivative",
 )
 # NOT registered here, by the registry's own contract: `tests/unit/test_gpu_gate_checks.py`
-# derives the gated set from ONE module (test_transverse_degeneracy_jvp.py) and fails on
-# any fragment outside it. The sm_80 test of the near-field self-fold
+# derives the gated set from ONE module (test_transverse_degeneracy_jvp.py) and
+# fails on any fragment outside it. The sm_80 tests of the near-field self-fold
 # (test_pallas_nearfield_fused.py::test_leafpair_include_self_gpu_matches_reference)
-# runs under the ordinary GPU suite; widening the registry is a separate change.
+# and of the CSR M2L kernel (test_m2l_real_csr_pallas.py::test_csr_pallas_gpu_matches_rot_scale)
+# therefore run under the ordinary GPU suite; widening the registry to several
+# modules is a separate change.
 
 # Measured on an A100 sm_80 / jax 0.10.2 and documented in ARCHITECTURE.md §9,
 # which carries the per-entry reasoning and the deterministic-ops column. A hit

bench/m2l_csr_microbench.py

@@ -0,0 +1,170 @@
+"""Go/no-go microbench for the CSR M2L Pallas kernel (plan "small leaves", gate G2a).
+
+Synthetic far-pair lists of 1M and 6M directed pairs over ``n`` nodes (the leaf
+64 / leaf 32 counts at N=200k), orders 4 and 6, fp32:
+
+* ``m2l_real_csr_pallas`` (one program per target, rotations on chip), jitted;
+* the pure-JAX chunked lane as production runs it -- ``m2l_rot_scale_real_batch``
+  over 4096-pair chunks in a ``lax.scan`` with ``_chunk_segment_scatter_add`` --
+  with ``JACCPOT_M2L_DEGREE_BATCHED`` off and on (the plan's 2.0 candidate, and
+  the honest pure-JAX baseline).
+
+Reports ns per directed pair (min of the timed calls) and the fp32 rel-L2 of
+the kernel against the pure-JAX lane. Gate: <= 15 ns/pair at p=4 and parity
+<= 1e-5 rel is "go" for wiring; the pure-JAX lane sits at ~130-375 ns/pair.
+
+    CUDA_VISIBLE_DEVICES=<idle> python bench/m2l_csr_microbench.py [--pairs 1000000 6000000] [--orders 4 6]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+
+import numpy as np
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--pairs", type=int, nargs="+", default=[1_000_000, 6_000_000])
+    ap.add_argument("--orders", type=int, nargs="+", default=[4, 6])
+    ap.add_argument("--nodes", type=int, default=12_500)
+    ap.add_argument("--repeats", type=int, default=5)
+    ap.add_argument("--chunk", type=int, default=4096)
+    ap.add_argument("--out", default=None)
+    args = ap.parse_args()
+    os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
+
+    import jax
+    import jax.numpy as jnp
+    from jax import lax
+
+    from jaccpot.operators.m2l_real_rot_scale import m2l_rot_scale_real_batch
+    from jaccpot.operators.real_harmonics import sh_size
+    from jaccpot.pallas.m2l_real_csr import m2l_real_csr_pallas
+    from jaccpot.runtime.kernels._m2l import _chunk_segment_scatter_add
+
+    dev = jax.devices()[0]
+    print(f"device {dev} cc={getattr(dev, 'compute_capability', '?')}", flush=True)
+    rng = np.random.default_rng(0)
+    n = int(args.nodes)
+    centers = jnp.asarray(rng.uniform(-1, 1, (n, 3)).astype(np.float32))
+    results = []
+
+    def timed(fn, *a):
+        out = jax.block_until_ready(fn(*a))
+        ts = []
+        for _ in range(args.repeats):
+            t0 = time.perf_counter()
+            out = jax.block_until_ready(fn(*a))
+            ts.append(time.perf_counter() - t0)
+        return out, min(ts), float(np.median(ts))
+
+    for order in args.orders:
+        C = sh_size(order)
+        mult = jnp.asarray(rng.standard_normal((n, C)).astype(np.float32))
+        for P in args.pairs:
+            # random well-separated pairs: reject |delta| < 0.3 (the MAC would too)
+            src = rng.integers(0, n, P).astype(np.int32)
+            tgt = rng.integers(0, n, P).astype(np.int32)
+            cn = np.asarray(centers)
+            for _ in range(50):  # resample until every pair is well separated
+                bad = np.linalg.norm(cn[tgt] - cn[src], axis=1) < 0.3
+                if not bad.any():
+                    break
+                src = np.where(bad, rng.integers(0, n, P), src).astype(np.int32)
+            assert not bad.any()
+            src_j, tgt_j = jnp.asarray(src), jnp.asarray(tgt)
+            counts = np.bincount(tgt, minlength=n)
+
+            csr = jax.jit(
+                lambda m, c, s, t: m2l_real_csr_pallas(m, c, s, t, order=order)
+            )
+            out_k, t_k, med_k = timed(csr, mult, centers, src_j, tgt_j)
+
+            chunk = int(args.chunk)
+            n_chunks = -(-P // chunk)
+            pad = n_chunks * chunk - P
+            src_p = jnp.pad(src_j, (0, pad), constant_values=0)
+            tgt_p = jnp.pad(tgt_j, (0, pad), constant_values=0)
+
+            def pure(m, c, s, t, P=P, chunk=chunk, n_chunks=n_chunks):
+                def body(acc, i):
+                    idx = i * chunk + jnp.arange(chunk, dtype=jnp.int32)
+                    valid = idx < P
+                    sc = s[idx]
+                    tc = t[idx]
+                    contrib = m2l_rot_scale_real_batch(
+                        m[sc], c[tc] - c[sc], order=order
+                    )
+                    return (
+                        _chunk_segment_scatter_add(
+                            acc, contrib, tc, valid, chunk_size=chunk
+                        ),
+                        None,
+                    )
+
+                acc, _ = lax.scan(
+                    body,
+                    jnp.zeros((n, C), jnp.float32),
+                    jnp.arange(n_chunks, dtype=jnp.int32),
+                )
+                return acc
+
+            rows = {}
+            for db in ("0", "1"):
+                os.environ["JACCPOT_M2L_DEGREE_BATCHED"] = db
+                # the knob is read at trace time and the cascade's inner jits are
+                # cached across variants -- without this the second variant reuses
+                # the first's trace (measured: bit-identical output and time)
+                jax.clear_caches()
+                fn = jax.jit(pure)
+                out_p, t_p, med_p = timed(fn, mult, centers, src_p, tgt_p)
+                rows[db] = (out_p, t_p, med_p)
+                del fn
+            ref = np.asarray(rows["0"][0], np.float64)
+            assert np.all(np.isfinite(ref)), "pure-JAX reference has non-finite rows"
+            assert np.all(
+                np.isfinite(np.asarray(out_k))
+            ), "CSR kernel produced non-finite rows"
+            rel = float(
+                np.linalg.norm(np.asarray(out_k, np.float64) - ref)
+                / np.linalg.norm(ref)
+            )
+            rel_db = float(
+                np.linalg.norm(np.asarray(rows["1"][0], np.float64) - ref)
+                / np.linalg.norm(ref)
+            )
+            row = dict(
+                order=order,
+                pairs=P,
+                nodes=n,
+                longest_row=int(counts.max()),
+                csr_ns_per_pair=1e9 * t_k / P,
+                csr_ms=1e3 * t_k,
+                csr_median_ms=1e3 * med_k,
+                pure_ns_per_pair=1e9 * rows["0"][1] / P,
+                pure_ms=1e3 * rows["0"][1],
+                pure_db_ns_per_pair=1e9 * rows["1"][1] / P,
+                pure_db_ms=1e3 * rows["1"][1],
+                rel_l2_csr_vs_pure=rel,
+                rel_l2_db_vs_pure=rel_db,
+                chunk=chunk,
+            )
+            results.append(row)
+            print(
+                f"p={order} pairs={P/1e6:.0f}M longest_row={counts.max()}: CSR {row['csr_ns_per_pair']:.1f} ns/pair "
+                f"({row['csr_ms']:.1f} ms) | pure-JAX {row['pure_ns_per_pair']:.1f} ns/pair | degree-batched "
+                f"{row['pure_db_ns_per_pair']:.1f} ns/pair | rel-L2 csr {rel:.2e}, db {rel_db:.2e}",
+                flush=True,
+            )
+    if args.out:
+        with open(args.out, "w") as fh:
+            json.dump(results, fh, indent=2)
+        print("wrote", args.out)
+
+
+if __name__ == "__main__":
+    main()

docs/small_leaves_2026-09.md

@@ -0,0 +1,91 @@
+# Small leaves on the fused single-GPU lane (2026-09-07 .. 2026-09-10)
+
+Plan: `~/.claude/plans/ok-please-draft-a-effervescent-crown.md` (Odisseo box). Branch
+`perf/small-leaves-p1`. Harness: `Odisseo-bench-multigpu/benchmark_multigpu/codes/{fit_smallleaf_caps,
+smallleaf_baseline,smallleaf_dynamics}.py`, `FAST_LANE_ENV_BY_LEAF` in `compare_force.py`; results under
+`artifacts/smallleaf/`.
+
+## Why
+
+At N=200k Plummer, p=4, theta 0.6 the fused lane at leaf 256 sums 58 % of N directly per target
+(pkdgrav3: 0.24 %). The direct share falls as ~W^0.9 with the leaf size, so smaller leaves are the
+only route to pkdgrav3's arithmetic -- and before this work leaf 64 was 3.1x SLOWER per step than
+leaf 256 with 3x fewer pair evaluations.
+
+## What the attributed baseline said
+
+`strict_run_v2` per step, N=200k, theta 0.6, idle A100 (`baseline_base_*.json`):
+
+| leaf | step ms | eval-only ms | leaf-pair kernel | downward (M2L+L2L) | launches/step |
+|---|---|---|---|---|---|
+| 256 | 131.8 | 70.0 | 79 | 41 | 7.8k |
+| 128 | 187.5 | 43.2 | 48 | 130 | 15.8k |
+| 64 | 409.6 | 26.5 | 22 | 372 | 32k |
+| 32 | 1168 | 19.6 | 13 | ~1100 | -- |
+
+The Pallas leaf-pair kernel has **no small-W cliff** (its time is exactly the direct-volume ratio),
+the self-leaf scan was already negligible after #334, and the whole penalty was the far field:
+one XLA scatter fusion launched once per 4096-pair M2L chunk at 686 us (`_chunk_segment_scatter_add`:
+`segment_sum` with hundreds of duplicates per address plus ~3800 zero-adds onto node 0 per chunk,
+both serialised by the atomic-add lowering) = 169 ms of the 410 ms leaf-64 step, and behind it a
+21k-launch/step storm from the chunked rotation cascade.
+
+## What was built
+
+1. **Self-leaf block folded into the leaf-pair Pallas kernel** (`nearfield_fused_leaf.py`,
+   `include_self`; flag `JACCPOT_NEARFIELD_LEAFPAIR_FOLD_SELF`, default on, forward prepacked lane
+   only). Time-neutral (the scan was already batched), removes the per-leaf launch family.
+2. **Contention-free chunk reduction** (`_m2l.py::_chunk_segment_scatter_add`): segmented
+   `associative_scan` + out-of-bounds sink (`mode="drop"`, unique in-bounds indices). Leaf 64:
+   410 -> 237 ms/step; leaf 128: 188 -> 145; leaf 32: 1168 -> 649.
+3. **Target-tiled CSR M2L Pallas kernel** (`jaccpot/pallas/m2l_real_csr.py`, flag
+   `JACCPOT_STATIC_STRICT_FUSED_M2L_CSR=1`): one program per target owns its local row; rotations
+   assembled on chip from the two alignment angles in the centred (degree, m) layout, degree-only
+   z-core, no per-pair operand in HBM. Microbench on an idle A100 (`m2l_csr_microbench.json`):
+
+   | pairs | p | CSR ns/pair | pure-JAX | degree-batched |
+   |---|---|---|---|---|
+   | 1M | 4 | 19.8 | 69.4 | 55.6 |
+   | 6M | 4 | 12.1 | 68.9 | 55.3 |
+   | 1M | 6 | 12.4 | 126.1 | 86.2 |
+   | 6M | 6 | 12.0 | 124.9 | 85.6 |
+
+   Parity 1e-6 (fp32) against `m2l_rot_scale_real_batch`; interpret parity < 1e-10 (fp64) for orders
+   2-6. In the step at leaf 64 the kernel costs 12.9 ms where the pure-JAX lane cost ~360.
+
+## Where it landed (per step, theta 0.6, N=200k, idle A100, foreign-process-free rows only)
+
+| leaf | baseline | fold + scatter | + CSR M2L | eval-only | aggL2 vs fp64 direct |
+|---|---|---|---|---|---|
+| 256 | 131.8 | 126.4 | **120.3** | 69 | 8.27e-4 |
+| 128 | 187.5 | 145.2 | **122.3** | 42 | 1.15e-3 |
+| 64 | 409.6 | 236.5 | **176.7** | 24 | 1.20e-3 |
+| 32 | 1168 | 649.0 | **509.4** | 15 | 1.30e-3 |
+
+Theta 0.8 with CSR: leaf 64 113.5 (baseline 192.0), leaf 32 364.1 (baseline OOM'd on a shared card).
+
+**The far field is solved; the per-step leaf optimum did not move.** With M2L at 13 ms, the leaf-64
+step is ~100 ms of traced dual-tree walk and list compaction (pair-queue-sized scatter fusions at
+~800 us x ~55 launches, ~200 memcpys, sorts) that scale with the leaf count and the traced
+`max_pair_queue` (1M at leaf 64, 2M at leaf 32). Eval-only (force at fixed lists) is 3-5x faster at
+leaf 32-64 than at 256, so a workload that refreshes lists every k steps, or a walk that costs
+O(leaves) rather than O(queue), would collect the gain. Gate G3 (<= 35 ms per step at 200k) is
+missed; G1, G2 and G2a are met. Next lever: the traced walk's queue buffers (yggdrax side).
+
+## Capacity rules for small leaves (fitted, `FAST_LANE_ENV_BY_LEAF`)
+
+* compact far-pair cap: pow2 >= 1.5x the far-pair count (101k / 363k / 992k / 2.34M at 256/128/64/32).
+* neighbour-EDGE cap: the traced refresh pads to `num_leaves x traced_neighbour_cap`, and the #333
+  carry-over sets that cap to pow2(1.5 x longest eager row + 1) with the longest row =
+  `num_leaves - 1` -> 2^21 already fails at leaf 128 (2^23), 2^25 at 64, 2^27 at 32.
+* `max_neighbors_per_leaf` must be explicit above the 2048 clamp; `max_interactions_per_node`
+  needs 16384 at leaf 32 (and at leaf 64 for theta 0.4).
+
+## Traps recorded
+
+* The perfetto trace of a leaf-32 step hung for 44 h after the timing (use `--no-trace` there).
+* Back-to-back configs on one card get rejected by the guard on the previous process's stale
+  utilisation reading; the queue scripts wait for 0 % and retry.
+* `JACCPOT_M2L_DEGREE_BATCHED` is read at trace time; a microbench that flips it must
+  `jax.clear_caches()` or the second variant reuses the first trace (bit-identical output).
+* The `l2l_only` detail diag mode is inconsistent with the cumulative modes -- do not attribute from it.