feat(mutual): rungs on the distributed mesh, and a block-step entry point for it

[TobiBu]

Track C of docs/plan_2026-08_C_rungs.md, from a fresh origin/main (a07528a).

Needs TobiBu/yggdrax#53 first — the exact per-particle near-field predicate needs the
remote endpoint's rung, which has to ride the demand-driven halo import.

What landed

C1	rung / level_weights through the shard_map body and the driver, k_max on the config	mutual/distributed.py
C2	the weight applied correctly on the cross-domain far half	same
C3	DistributedBlockStepFMM, plus make_distributed_mutual_evaluator	nornax_adapter.py
C4	backend / pallas_interpret, use_pallas actually threaded	mutual/distributed.py

Before this, mutual/distributed.py had zero occurrences of rung,
level_weights or k_max, so block-step individual timesteps were single-device only.

C2 is the part that was not plumbing, and it needed a fourth test

The plan says a momentum test alone will not catch a weight applied to neither side. It
is worse than that. I injected each fault and measured three criteria — 2 devices,
N = 256, leaf 4, θ = cross_θ = 0.5, 33 cross far pairs, 991 cross near pairs:

fault	level partition	weight linearity	momentum
none	2.1e-16	2.9e-16	< 4e-17
far half left unweighted	3.9e-02	2.0e-02	< 4e-17
far weight applied twice	2.1e-16	9.4e-03	< 4e-17
far ignores the remote rung	2.1e-16	2.9e-16	< 4e-17

Momentum sees none of the three. A per-level residual at 1e-17 is not evidence the
levels are right — any single symmetric scalar per pair conserves momentum exactly,
whatever that scalar is. So:

• the partition test (one-hot rows must sum to the total) catches the first, and
  cannot catch squaring: w² == w for 0 and 1;
• a linearity test (a(u) + a(1−u) == a(1) at fractional rows) catches squaring —
  hence a separate test, with rows deliberately none of them 0 or 1;
• neither catches the third, which the plan did not anticipate: a level assignment
  that is wrong yet consistent is still one scalar per pair, so it passes partition,
  linearity and momentum. It needed
  test_the_cross_far_level_uses_the_REMOTE_endpoints_rung: two rung-uniform clumps,
  one per device, so every cross pair provably belongs to one level and clump A's
  level-0 force is an exact direct sum. Injected as f_level = node_rung[f_local] it
  moves 1.2e-3 against a 1e-13 bound while the other three stay at round-off.

Structurally, the weight is applied once, on the evaluating device, to both directions
of the batched M2L, before either leaves. Legal because everything downstream of an
expansion is linear in its coefficients (L2L re-centres, L2P evaluates), so weighting
the expansion is weighting the pair. The far payload's rung is read from the end of
the row, so the multipole width is stated once — indexing it at a computed
sh_size(order) offset would state it twice, and a mismatch would read a coefficient
as a rung, which survives the clamp as a perfectly valid pair level.

On the plan's second C2 concern — a cell straddling a boundary getting the same rung on
both devices — it holds by construction and is now asserted rather than assumed: a far
pair's remote endpoint is a coarse leaf holding exactly one frontier node
(accept_only_leaf_pairs + the degenerate contract check), so each endpoint's rung is
computed exactly once, by the device that owns the particles, and never recomputed
across the wire.

Which channel each rung travels on

• Near, exact per-particle → round B of the halo import, sized by the halo.
• Far, cell-level → the coarse frontier, one scalar per leaf appended to its
  multipole row.

Per-particle rungs could not go on the frontier: it is all_gather-ed, so
leaf_width columns there would be O(N_total) — invisible at cross_theta = 0 where
the import fetches everything anyway, and a hard cap the moment θ collapses the import
to a surface halo.

_node_rungs gives every node its cell rung as a range maximum expressed as one
prefix count per level — this module has only the inclusive node ranges, not the level
schedule force._cell_rungs walks. k_max is small, so it is O(n·k_max) with no
parent array and no propagation scan, and
tests/unit/mutual/test_distributed_node_rungs.py checks it against both a loop oracle
and _cell_rungs on a real tree.

C3 reaches the plan's literal criterion, and needed a traceable readout

block_kdk_rollout walks its base steps with a lax.scan, so it needs a force it can
trace. The driver assembled its result with a NumPy scatter and raised on a
host-read overflow flag — neither survives a trace. Both moved:

• the readout is a jnp scatter;
• the "every real particle appears exactly once" check moved to the partition,
  where it belongs — it depends only on the frozen gid layout, not on any evaluation;
• the overflow read is attempted rather than gated on isinstance(..., Tracer), the
  discipline BlockStepFMM._validate_rung documents: a concrete array closed over by a
  scan body is not a Tracer yet still cannot be read.

Result: nornax's own scanned block_kdk_rollout, 2 devices, multi-rung, 2 base steps
— momentum drift 1.8e-17, energy bounded, and jaccpot's own advance_base_step
matching nornax's to 1e-12.

And a compile that was being paid per force

shard_map wraps a fresh closure per call, so jax.jit sees a fresh cache key and
distributed_mutual_fmm recompiles every time — fine for one force, ruinous for the
n_sub + 1 evaluations a base step asks for. make_distributed_mutual_evaluator
partitions and compiles once. Measured, 2 forced CPU devices, N = 128:

build	< 0.1 s
first evaluation	20.5 s
each later one	~8 s

distributed_mutual_fmm is now that plus one call, behaviour unchanged.

What prepare() freezes is not the topology, and the docstring says so rather than
implying otherwise: the partition, padding layout, bounds, capacities and compiled
program are frozen; the per-device tree is rebuilt inside the program from the positions
handed in. That is a finer rebuild cadence than the single-device lane's, not a
coarser one, and legality is unchanged — every evaluation is internally self-consistent,
so its levels partition its own pairs and each level's momentum cancels exactly. The
consequence runs the other way: prepare need not be called once per base step, only
when the partition should change.

C4

build_mutual_state_device has accepted use_pallas all along and the driver never
passed it, so this lane ran pure JAX whatever the caller asked. Routing now mirrors the
single-device lane exactly, including its two measured decisions: the far field stays
pure JAX (both Pallas M2L shapes are slower), and every Pallas lane goes through its
custom_vjp wrapper, never the bare kernel. The cross-domain near field is pure JAX
either way — a different kernel (_tile_forces), and a Pallas lane for it is a new
kernel, not a wiring change.

Per the plan's trap list, the weighted path has its own Pallas test: level_weights[k]
is exactly the lowering that broke before, and the throughput benchmark runs unweighted.
interpret ∈ {True, False} is kept parametrized — the interpret=False case is the
only thing standing between a lowering regression and a broken backend="pallas", and
it skips off sm_80.

Verification

XLA_FLAGS="--xla_force_host_platform_device_count=2" JAX_PLATFORMS=cpu JAX_ENABLE_X64=1 \
  pytest -q -o addopts="" tests/integration/test_mutual_distributed.py \
                          tests/integration/test_mutual_distributed_nornax.py \
                          tests/integration/test_mutual_fmm.py \
                          tests/integration/test_mutual_fmm_nornax.py \
                          tests/unit/test_custom_vjp_parity.py \
                          tests/unit/mutual/

suite	where	result
test_mutual_distributed.py + tests/unit/mutual/	2 forced CPU devices	56 passed, 1 skipped (18m36s)
test_mutual_distributed_nornax.py	2 forced CPU devices	12 passed (77s)
the interpret=False case CPU must skip, + the rung subset	2x A100	6 passed (6m52s)
test_mutual_fmm.py, test_mutual_fmm_nornax.py, test_custom_vjp_parity.py	2x A100	104 passed (1h03m)

The single-device row is the regression check, and it is insurance rather than a gate:
the diff does not touch mutual/force.py, nearfield.py, farfield.py,
device_topology.py or topology.py, and the only deleted line in
nornax_adapter.py is its __all__, so BlockStepFMM is untouched by
construction. Run anyway, because "additive by inspection" is a claim.

The one skip on CPU is test_the_pallas_backend_matches_the_jax_backend[False] --
the real Triton lowering, which needs sm_80. It is green on 2x A100, which
matters more than the rest of the C4 row: interpret mode validates the kernel's
LOGIC, not its lowerability, and per this repo's own record two mutual kernels
passed every CPU interpret test and then failed on their first GPU run.

Wall-clock is long because the box was at load ~70 on 64 cores throughout, not
because of anything here.

CI will SKIP these tests until yggdrax#53 merges

Deliberately, and verified. CI installs yggdrax from its default branch, so until #53
lands payload_sorted does not exist there -- and that is a TypeError raised deep
inside a shard_map trace on every weighted test. Both modules would go red rather
than skip, which is exactly the failure test_mutual_distributed.py's own docstring
records having been caught not preventing once before. So its guard now checks the
halo import's parameter alongside the two on the cross walk, and the nornax module
gets its own. Checked by running both against the OLD yggdrax checkout: 2 skipped,
each naming payload_sorted and #53. A guard nobody has seen fire is decoration.

pre-commit run --all-files (black, isort, pydoclint, flake8) green over the tree.

Why one PR

C1–C4 are interdependent: C1/C2's tests drive C3's evaluator (a 4× test-time win, and
also the claim under test — only shapes are static, so one compiled program must serve
every weight row and every rung assignment), and C4's threading lives inside the
function C3's refactor created. Reviewable commit by commit; the four commits are split
by concern.

Not in scope, flagged for track E

mutual/distributed.py's module docstring is stale and predates this PR: "The one
simplification, stated plainly: cross-domain interactions are EXACT. The cross walk is
driven with theta = 0" — untrue since the cross_theta lift (#208). Left alone
because track E owns module docstrings; the parameter-level docs it contradicts are
already correct.

🤖 Generated with Claude Code

[TobiBu]

Update: CI on this branch, and a scope fix it forced

test-distributed-mutual went red, in two independent ways — the first was hiding
the second.

Exit 5. The module-level guard for payload_sorted fired (CI installs yggdrax from
its default branch, which lacks it until #53), the whole file skipped, and in a job that
runs one file that is pytest's exit 5, "no tests were collected". Under set -o pipefail that failed the pytest step, so the next step — the one that exists to say
"the suite was skipped, not run" with an actionable message — never ran. It now tolerates
exit 5 specifically and still fails on every other code.

The requirement was too wide, and that is the substantive fix. payload_sorted
carries a cross-domain near pair's remote rung, so only the rung-weighted path needs
it — but it was passed as payload_sorted=None on the unweighted path too, which still
requires the parameter to exist. It is now passed conditionally, so the unweighted
lane runs against the published yggdrax and the guard becomes per-test. Measured against
the old checkout: the file now runs its pre-existing criteria and skips only the
weighted ones (4 passed / 1 skipped on the fast subset) instead of skipping wholesale at
0 passed — which was the vacuous-green state this job was added to prevent, reached from
the other direction.

A skip is no longer read as vacuity. Every skip in this file used to be a
whole-module one (old yggdrax, or < 2 devices), so "any skip" and "nothing ran" were
the same thing and ^SKIPPED was exact. They are not any more: the file now carries
per-test skips that are correct on a CPU runner — the sm_80-gated Pallas lowering,
and the weighted path above. Left alone, this job would have gone red again the moment
#53 merged, over a test that cannot run anywhere on a CPU box. The check now matches
nothing-collected, or one of the two whole-module guards firing.

That is a deliberate loosening, and worth flagging rather than burying: a future change
that skipped most of the suite for some other reason would not be caught, where
^SKIPPED would have caught it. Table-tested the five cases that matter:

log state	verdict
old yggdrax, whole module skipped	fails ✔
single device	fails ✔
no tests ran	fails ✔
sm_80 skip alone (post-#53)	passes ✔
today's mixed state (unweighted runs, weighted skips)	passes ✔

Re-verified after the change

suite	where	result
test_mutual_distributed.py + _nornax.py + tests/unit/mutual/	2 forced CPU devices, new yggdrax	68 passed, 1 skipped (20m29s)
the fast subset	2 forced CPU devices, old yggdrax	4 passed, 1 skipped (was 0 passed)
pre-commit run --all-files	—	green

Two caveats on CI coverage, neither new

• .github/workflows/ci.yml is not owned by any track in the 2026-08 plan. This PR is
  what made those paths reachable, so I changed it — but it is the most
  reviewable-away part of this branch.
• test_mutual_distributed_nornax.py cannot run in CI at all: nornax is not installed
  there, and the shard it lands in has one device. That matches the existing standing of
  test_mutual_fmm_nornax.py and is not a regression, but it does mean the C3 criterion
  is verified locally only (12 passed, 2 forced CPU devices) — not by CI.