feat(strict): flat-emission wavefront walk as an opt-in traversal for the fused lane - #341
Merged
Merged
Conversation
… 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>
for more information, see https://pre-commit.ci
…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
force-pushed
the
perf/flat-walk-lane
branch
from
September 10, 2026 13:57
7874eea to
b91d8bf
Compare
TobiBu
changed the base branch from
perf/small-leaves-csr-m2l
to
perf/small-leaves
September 10, 2026 13:57
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
added this pull request to stack #342
September 10, 2026 19:34
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #340 (merged into
perf/small-leaves, the base). Depends on yggdrax TobiBu/yggdrax#74 (mac_type,peak_wavefrontondual_tree_walk_mutual) -- merged. Plan: the "tree walk" plan of 2026-09-10; record indocs/tree_walk_2026-09.md.Why. At N=200k, leaf 64, theta 0.6 a
strict_run_v2step 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 throughlax.cond, a sort per round, and a post-loop flatten overtotal_nodes x Kslots. Measured in isolation the traced dual walk costs 503 ms per call and yggdrax's flat-emissiondual_tree_walk_mutual39.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=0restores the dual walk) todual_tree_walk_mutualfed the dual walk's ownmac_extentsandmac_type, so the far and near lists are the dual walk's as sets:idx < far_pair_count); capacity = the compact far-pair cap;searchsortedoffsets, width = the neighbour-edge cap, no per-leaf row buffer;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_handoffas floors to the traced refresh and later eager prepares, so scan carry shapes match and never shrink; traced: any overflow saturatesfar_pair_countto the capacity, which tripsstrict_run_v2's existing saturation guard (documented at the guard);peak_wavefront;_strict_fused_capacity_handoffsizes the traced queue as pow2(1.5 x peak) on this lane.JACCPOT_STATIC_STRICT_FUSED_FLAT_WALKandJACCPOT_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 explicit1there still raises.tests/conftest.pyhonoursYGGDRAX_WORKTREE(it put the sibling checkout, a paper branch, at the front ofsys.pathunconditionally).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,=0takes 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_capacityparametrised 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