perf(interactions): flat walk reports its peak wavefront and takes a MAC type; dual walk drops its emission conditionals - #74
Merged
Conversation
… mutual walk reports peak wavefront and takes the dual walk's MAC The lax.cond wrappers around the far/near emission in _dual_tree_walk_impl cost a device-to-host predicate sync per round and, because an identity branch cannot alias operand to result, a copy of the dense per-node buffers in and out every round; a masked update with an empty mask is an exact no-op, so they are gone (isolated walk at N=200k / leaf 64 / Q=2^20: 503 -> 390 ms). The MAC-extent gathers are issued once. dual_tree_walk_mutual -- the flat-emission walk -- gains peak_wavefront and rounds on MutualWalkResult (so a traced caller can size its queue from data), a static mac_type that routes acceptance through _compute_mac_ok (fed the dual walk's mac_extents it then reproduces the dual walk's far and near lists as SETS -- pinned by the new test for bh/dehnen at three thetas and a scaled radius; the strict default differs only on exact MAC equality), an index dtype that follows left_child_full (int32 buffers halve the queue traffic: 39.9 -> 26.4 ms at leaf 64), hoisted gathers and uniqueness promises on its scatters. bench/traversal_walk_bench.py times both walks and a bare scatter on one tree (autocvd picks the card); at leaf 64 the dual walk runs 503 ms and the mutual walk 39.9 ms for identical pair counts (858,488 far / 1,842,106 near). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…e, follows the input dtype; dual walk drops its emission conditionals dual_tree_walk_mutual: MutualWalkResult gains peak_wavefront (the largest push any round needed, before truncation) and rounds; a static mac_type routes the acceptance through _compute_mac_ok so a caller feeding the dual walk's own mac_extents gets its far and near lists as sets (the strict default differs on exact equality only); index dtype follows left_child_full; per-node gathers hoisted; compaction scatters carry unique_indices. _flat_append likewise. _dual_tree_walk_impl: the lax.cond around the far and near emissions replaced by the unconditional (exact no-op when empty) update -- the identity branch could not alias its dense operands, so XLA copied far_buffer/neighbor_buffer in and out every round and synced the predicate to the host. Measured on an A100, N=200k Plummer, leaf 64, queue 2^20, int64: 503 -> 390 ms per walk. The mac_extents gathers are reused for the split heuristic. bench/traversal_walk_bench.py times both walks in isolation on one tree (dual 503 ms vs mutual 39.9 ms at leaf 64 / 2^20 / int64; mutual 13.4 ms at int32 with the queue sized to the measured peak of 191,890). tests/unit/test_dual_tree_walk_mutual.py: set parity with the dual walk for bh/dehnen at three thetas and a radius scale, equality-only difference of the strict rule, jit, peak/rounds bounds, overflow flags, dtype. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…rnative (bench table, 2026-09-10) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
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
1 participant
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.
Part of the jaccpot "tree walk" plan (2026-09-10): the fused single-GPU lane's per-step cost at small leaves is the traced dual-tree walk, and this PR is the yggdrax half of moving that lane onto the flat-emission walk.
dual_tree_walk_mutualMutualWalkResultgainspeak_wavefront(the largest push any round needed, before truncation, so an overflowing probe still reports the queue it wanted) androunds.mac_type: Optional[MACType] = None:Nonekeeps the stricttheta d > r_a + r_brule (the mutual block-step lane's host-parity contract); aMACTyperoutes the test through_compute_mac_ok, so a caller feeding the dual walk's ownmac_extentsgets the dual walk's far and near lists as sets (the strict rule differs on exact equality only, which symmetric inputs do hit).left_child_full; per-node gathers hoisted; compaction scatters carryunique_indices=True(_flat_appendlikewise)._dual_tree_walk_impl: thelax.condaround the far and near emissions is replaced by the unconditional update (an exact no-op when the mask is empty: every index sinks out of bounds and is dropped, every increment is zero). The identity branch could not alias its dense operands, so XLA copiedfar_buffer/neighbor_bufferin and out every round and synced the predicate to the host. Themac_extentsgathers are reused for the split heuristic.Measured (
bench/traversal_walk_bench.py, idle A100, N=200k Plummer, theta 0.6, one tree, identical far/near pair counts for both walks):Tests: new
tests/unit/test_dual_tree_walk_mutual.py(11): set parity with the dual walk for bh/dehnen at theta 0.3/0.5/0.9 and a radius scale, the strict rule's equality-only difference, jit, peak/rounds bounds (a queue of exactlypeakfits,peak-1overflows), each overflow flag, dtype following the inputs. Existing traversal suites (test_traced_walk_capacity,test_tree_interactions,test_traversal_policy,test_interactions_backend_parity,distributed/test_cross_walk*) unchanged and green. pre-commit (black, isort, pydoclint, basedpyright) clean.The jaccpot side (the opt-in
JACCPOT_STATIC_STRICT_FUSED_FLAT_WALKlane) depends on this and follows as a jaccpot PR; jaccpot CI installs yggdrax from main, so this merges first.🤖 Generated with Claude Code