Skip to content

feat(strict): flat-emission wavefront walk as an opt-in traversal for the fused lane - #341

Merged
TobiBu merged 5 commits into
perf/small-leavesfrom
perf/flat-walk-lane
Sep 10, 2026
Merged

feat(strict): flat-emission wavefront walk as an opt-in traversal for the fused lane#341
TobiBu merged 5 commits into
perf/small-leavesfrom
perf/flat-walk-lane

Conversation

@TobiBu

@TobiBu TobiBu commented Sep 10, 2026

Copy link
Copy Markdown
Owner

Stacked on #340 (merged into perf/small-leaves, the base). Depends on yggdrax TobiBu/yggdrax#74 (mac_type, peak_wavefront on dual_tree_walk_mutual) -- merged. Plan: the "tree walk" plan of 2026-09-10; record in docs/tree_walk_2026-09.md.

Why. At N=200k, leaf 64, theta 0.6 a strict_run_v2 step costs 177 ms with the CSR M2L lane on, ~100 ms of it in the traced dual-tree walk and its list emission: dense per-node output rows (820 MB at leaf 64) loop-carried and copied through lax.cond, a sort per round, and a post-loop flatten over total_nodes x K slots. Measured in isolation the traced dual walk costs 503 ms per call and yggdrax's flat-emission dual_tree_walk_mutual 39.9 ms (13.4 ms with int32 indices and a queue sized to the measured peak) for identical pair counts.

What. The strict streamed split build routes BY DEFAULT (JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK=0 restores the dual walk) to dual_tree_walk_mutual fed the dual walk's own mac_extents and mac_type, so the far and near lists are the dual walk's as sets:

  • far pairs un-mutualised and interleaved so the live pairs stay a prefix (every M2L consumer masks idx < far_pair_count); capacity = the compact far-pair cap;
  • the leaf neighbour CSR from one stable argsort of the directed near pairs plus searchsorted offsets, width = the neighbour-edge cap, no per-leaf row buffer;
  • eager: queue ladder on queue_overflow; a far/near overflow doubles a cap the caller did NOT name in the env (floors 131072 / 2^21 directed pairs, ceilings 2^26 / 2^28) and raises naming a cap the caller did name; the settled widths travel through _strict_fused_capacity_handoff as floors to the traced refresh and later eager prepares, so scan carry shapes match and never shrink; traced: any overflow saturates far_pair_count to the capacity, which trips strict_run_v2's existing saturation guard (documented at the guard);
  • the capacity report is always emitted and carries peak_wavefront; _strict_fused_capacity_handoff sizes the traced queue as pow2(1.5 x peak) on this lane.
  • Defaults (second commit). Both JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK and JACCPOT_STATIC_STRICT_FUSED_M2L_CSR (feat(m2l): target-tiled CSR Pallas M2L kernel with on-chip rotations (opt-in lane) #340's CSR kernel, sm_80 gate unchanged) default to 1. A merely-defaulted flat-walk flag falls back to the dual walk QUIETLY when the configuration cannot take it (treecode walk requested, a MAC other than bh/dehnen, a solver-owned pair policy, the non-flat far-pair layout); an explicit 1 there still raises.
  • tests/conftest.py honours YGGDRAX_WORKTREE (it put the sibling checkout, a paper branch, at the front of sys.path unconditionally).

Checks. test_flat_walk_default_dispatch.py (CPU, 9): default routing, =0, the fallback/raise pairs, named-vs-unnamed caps, the floor handoff. Seam test (12): set parity vs the dual walk for bh/dehnen at three thetas and a radius scale, prefix-live far pairs with both directions, valid CSR, no self/duplicates, traced call keeps capacity width, overflows raise naming the knob, odd caps rejected, queue ladder grows, unnamed caps grow. GPU wiring test (3): lane entered by default, =0 takes the dual walk, forces within 2e-5 of the dual lane, both flags refused, treecode flag alone wins. End to end at N=200k leaf 64: aggL2 vs fp64 direct 1.2011e-03 (identical to the dual lane), trajectory-recovered force vs eager at steps 1-2 < 7e-4. test_strict_run_v2_refresh_capacity parametrised over the flag. Per step, clean idle A100, N=200k p4 th0.6 (flat walk + CSR + int32): leaf 256 96.2 ms, 128 71.3, 64 63.3 (from 409.6 at the start of the small-leaves work, 176.7 with CSR alone), 32 82.5; th0.8: leaf 64 41.6, 32 53.6. Per-order sweep at leaf 64 monotone and identical to 4 digits to the dual lane; 300-step dynamics flag on/off: no bias (docs/tree_walk_2026-09.md).

🤖 Generated with Claude Code

TobiBu and others added 4 commits September 10, 2026 15:57
… the fused lane

JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK=1 routes the strict streamed split build
to yggdrax's dual_tree_walk_mutual fed the dual walk's own mac_extents and
mac_type, so the far and near lists are the dual walk's as SETS (pinned in
tests/unit/runtime/test_flat_walk_production_seam.py for bh/dehnen at three
thetas and a radius scale). Far pairs are un-mutualised and INTERLEAVED so the
live pairs stay a prefix (every M2L consumer masks idx < far_pair_count); the
leaf neighbour CSR comes from one stable argsort of the directed near pairs
plus searchsorted offsets, width = the neighbour-edge cap, no per-leaf row
buffer. Eager: queue ladder on queue_overflow, far/near overflow raise naming
the cap. Traced: any overflow saturates far_pair_count to the capacity, which
trips strict_run_v2's existing saturation guard. The capacity report is always
emitted and carries peak_wavefront; _strict_fused_capacity_handoff sizes the
traced queue as pow2(1.5 x peak) on this lane instead of 2 x the eager rung.

Measured in isolation (A100, N=200k, leaf 64, theta 0.6): the traced dual walk
503 ms per call at queue 2^20 int64; the flat walk 39.9 ms (26 int32, 13.4 at a
queue sized to the measured 191,890 peak) with identical pair counts. End to
end the strict scan gives aggL2 1.2011e-03 (identical) with the truncation
check < 7e-4 at steps 1-2.

tests/conftest.py honours YGGDRAX_WORKTREE: it used to put the sibling checkout
(a paper branch) at the front of sys.path unconditionally, so no jaccpot test
run could exercise a yggdrax branch. CI installs yggdrax from main and has no
sibling directory; unaffected.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…raps

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK and JACCPOT_STATIC_STRICT_FUSED_M2L_CSR
default to 1. Rules that make a default safe where the opt-in could simply
raise:

* A merely-defaulted flat-walk flag falls back to the traced dual walk
  QUIETLY when the configuration cannot take the flat walk (treecode walk
  requested, a MAC other than bh/dehnen, a solver-owned pair policy, the
  non-flat far-pair layout). An explicit "1" against such a configuration
  still raises: then the caller asked for a walk it cannot have.
* A capacity the caller did not NAME in the environment (no
  ..._COMPACT_FAR_PAIR_CAP / ..._NEIGHBOR_EDGE_PROFILE_FIXED_CAP) grows
  eagerly from its floor (131072 / 2^21 directed pairs; ceilings 2^26 /
  2^28); a named cap is exact and raises on overflow (#333). Cap growth no
  longer consumes the queue retry budget. Without this the old 2^21 near
  default raised at leaf 64 / N=200k (needs 2^22) -- a regression the dual
  walk never had because it sized from eager data.
* _strict_fused_capacity_handoff hands the settled far/near widths to the
  traced refresh AND to later eager prepares as floors, so lax.scan carry
  shapes match and never shrink.
* The CSR M2L gate keeps its sm_80 / interpret predicate.

Tests: test_flat_walk_default_dispatch.py (CPU: default routing, =0, the
three fallback/raise pairs, named-vs-unnamed caps, floor handoff); seam test
for unnamed-cap growth; the wiring tests now pin default-on and the treecode
flag winning over a defaulted flat walk. Docs updated.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@TobiBu
TobiBu force-pushed the perf/flat-walk-lane branch from 7874eea to b91d8bf Compare September 10, 2026 13:57
@TobiBu
TobiBu changed the base branch from perf/small-leaves-csr-m2l to perf/small-leaves September 10, 2026 13:57
@TobiBu
TobiBu marked this pull request as ready for review September 10, 2026 13:57
The runtime-typecheck CI job (jaxtyping + beartype) checks the return type of
_build_dual_tree_artifacts_split_strict_streamed, so the bare object() sentinel
failed there; the stub now hands back a dual-walk-built _DualTreeArtifacts.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@TobiBu
TobiBu added this pull request to stack #342 September 10, 2026 19:34
@TobiBu
TobiBu merged commit f970c98 into main Sep 10, 2026
20 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants