update paper branch from main#39
Merged
Merged
Conversation
Adds a Docs workflow that builds the Sphinx HTML and deploys it to GitHub Pages via actions/upload-pages-artifact + actions/deploy-pages. It builds on every push and PR (catching doc breakage) but only deploys on pushes to main or a manual workflow_dispatch, so PRs never touch the live site. Links the hosted docs from the README. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…2 sorts) The static-radix fused fast-lane per step is launch/host-bound (GPU idle ~58-66% on A100 @200k): ~4200 host-launched micro-kernels, of which the dual-tree walk's `_per_key_prefix` contributes FOUR 262144-element argsorts inside the while-body (each lowering to ~150 bitonic device kernels). Within `_dual_tree_walk_impl`, the far (fwd/bwd) and near (fwd/bwd) per-key prefixes are INDEPENDENT of one another -- only their downstream slot-index use is sequential, not the argsort itself. Batch each pair into a single `jax.vmap(_per_key_prefix)` over a stacked (2, B) array, so XLA emits one batched bitonic sort per pair instead of two. `vmap` is semantically identical to the two sequential calls, so forces are bit-identical. Measured @200k / leaf256 / A100 (fused device-only lane, min-of-3 x 80 steps): - fused per-step HLO sort ops: 5 -> 3 - ms/step: 121.4 -> 112.5 (~7% faster) - energy conservation over 60 steps: max|dE/E0| = 2.48e-5 (unchanged) Not a null result (unlike the earlier off-critical-path M2L-count reduction): the walk argsorts are on the serial critical path, so cutting their launch count moves the wall. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
perf(walk): batch dual-tree walk per-key-prefix argsorts (4→2 sorts, ~7% faster/step)
ci(docs): publish the Sphinx site to GitHub Pages
…traversal Multi-GPU tree build and cross-walk now support all three backends. Distributed backend selector (radix + octree): - local_tree: DISTRIBUTED_TREE_TYPES + _validate_distributed_tree_type + shard_map-safe _build_local_tree dispatch; distributed_tree_moments / build_local_moments gain tree_type (root via argmin(parent)). - let: tree_type threaded to the local build in build_distributed_coarse_tree / classify_against_remote / distributed_let_import. Coarse tree stays radix (internal remote-COM representation; robust at coarse leaf_size=1). Leaf-only (bucket) median-split KD-tree (build_leaf_kdtree, LeafKDTree): - All points in leaves; internal nodes are pure split planes; compact contiguous node_ranges tiling [0,N) + nodes_by_level -> satisfies the same topology contract as radix/octree, giving complete pair coverage. The heap KD-tree (build_kdtree) used by the KNN kernels is untouched; the generic pipeline / tree_type="kdtree" / distributed now use the leaf-only build. Validated: 46 kd unit tests (heap KNN preserved + leaf-kd interactions parity vs radix); distributed CPU 9 passed incl. kd cross-walk complete disjoint source partition (the invariant the heap KD-tree violated). GPU validation in progress. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
for more information, see https://pre-commit.ci
feat(distributed): backend-selectable tree build + leaf-only KD-tree traversal
TobiBu
merged commit Jul 16, 2026
59a6626
into
paper/differentiable-applications
7 of 12 checks passed
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.
This pull request adds support for selecting the local tree build backend (
radix,octree, orkdtree) in distributed LET (Local Essential Tree) drivers and tests, and introduces optimizations to reduce kernel launch counts in the interaction implementation. The main changes include a new backend selector and validation, refactoring to use a unified local tree builder, and batching of sort operations for performance.Distributed tree backend selection and validation:
tree_typeparameter to distributed LET driver functions (build_distributed_coarse_tree,classify_against_remote,distributed_let_import) to allow selection of the local per-device tree backend ("radix","octree", or"kdtree"). The coarse tree is always built using the radix backend for robustness and consistency. ([[1]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-04eace56476fa086a5a1e0c082efb4aa00f09045f0449d5e94eaa7fbfdae8465R169-R182),[[2]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-04eace56476fa086a5a1e0c082efb4aa00f09045f0449d5e94eaa7fbfdae8465R338-R351),[[3]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-04eace56476fa086a5a1e0c082efb4aa00f09045f0449d5e94eaa7fbfdae8465R593-R602))DISTRIBUTED_TREE_TYPESand_validate_distributed_tree_typeinlocal_tree.pyto centralize supported backend validation and provide clear error messages for unsupported types. ([yggdrax/distributed/local_tree.pyR32-R104](https://github.com/TobiBu/yggdrax/pull/39/files#diff-cb6a9199e6e46d07a75e8f3e4a2b794fb5bc432f999899738e4a68a35d3b84ecR32-R104))_build_local_treeinstead of direct calls tobuild_tree, ensuring consistent backend dispatch and correct tracing insideshard_map. ([[1]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-04eace56476fa086a5a1e0c082efb4aa00f09045f0449d5e94eaa7fbfdae8465L203-R221),[[2]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-04eace56476fa086a5a1e0c082efb4aa00f09045f0449d5e94eaa7fbfdae8465L365-R389),[[3]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-04eace56476fa086a5a1e0c082efb4aa00f09045f0449d5e94eaa7fbfdae8465L612-R641))Testing improvements:
tests/distributed/test_local_tree_backends.pyto test the correctness and invariants of distributed tree builds across the supported backends, including a cross-walk test to verify the KD-tree's source partitioning. ([tests/distributed/test_local_tree_backends.pyR1-R214](https://github.com/TobiBu/yggdrax/pull/39/files#diff-f4c89e801e0e7e272e0a3e9fae2ecc137e3b82ac032b62fc9b4047b4b8ba04d5R1-R214))Performance optimizations:
_far_updateand_near_updateusingjax.vmap, halving the number of dominant sort-kernel launches during interaction computation. ([[1]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-09d47b4a1d31eae2e7ae4ddcba4640c00ff4b3e13bc7566707360fe68504a73cL1483-R1490),[[2]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-09d47b4a1d31eae2e7ae4ddcba4640c00ff4b3e13bc7566707360fe68504a73cL1508-R1515),[[3]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-09d47b4a1d31eae2e7ae4ddcba4640c00ff4b3e13bc7566707360fe68504a73cL1579-R1592),[[4]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-09d47b4a1d31eae2e7ae4ddcba4640c00ff4b3e13bc7566707360fe68504a73cL1609-R1622))Documentation and clarity:
[[1]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-04eace56476fa086a5a1e0c082efb4aa00f09045f0449d5e94eaa7fbfdae8465L109-R111),[[2]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-04eace56476fa086a5a1e0c082efb4aa00f09045f0449d5e94eaa7fbfdae8465L203-R221),[[3]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-04eace56476fa086a5a1e0c082efb4aa00f09045f0449d5e94eaa7fbfdae8465L365-R389),[[4]](https://github.com/TobiBu/yggdrax/pull/39/files#diff-04eace56476fa086a5a1e0c082efb4aa00f09045f0449d5e94eaa7fbfdae8465L612-R641))These changes make the distributed LET pipeline more flexible, robust, and performant, and improve test coverage for backend-specific invariants.