perf(rm): replace per-iteration sort in distributedAlloc with a min-heap#1826
Open
jonathan-meiri wants to merge 1 commit into
Open
perf(rm): replace per-iteration sort in distributedAlloc with a min-heap#1826jonathan-meiri wants to merge 1 commit into
jonathan-meiri wants to merge 1 commit into
Conversation
jonathan-meiri
force-pushed
the
optimize-distributedalloc-heap
branch
from
June 2, 2026 08:16
82b1b23 to
b9da143
Compare
jonathan-meiri
force-pushed
the
optimize-distributedalloc-heap
branch
from
June 30, 2026 15:35
b9da143 to
21cd040
Compare
jonathan-meiri
marked this pull request as ready for review
June 30, 2026 19:10
jonathan-meiri
force-pushed
the
optimize-distributedalloc-heap
branch
3 times, most recently
from
July 1, 2026 11:02
8333b12 to
5efcd18
Compare
Author
|
Hi @rajatchopra — following up on your comment in #1788 ("Lets attack the loop optimization in #1826 next"). This PR is out of draft and ready for review: a single commit with the heap-based |
Follow-up to the tie-break fix in PR NVIDIA#1788. The previous implementation sorted the full candidate list inside the allocation loop, paying O(n log n) per iteration for n iterations and giving O(n² log n) overall. Since all annotated replicas from the same underlying physical device share the same sort key, sorting at the replica granularity is wasted work — only m (the number of distinct physical devices contributing candidates) needs to be reordered. Refactor to: - Bucket candidates by their underlying physical device into a small gpuAllocState per device, holding `used`, `pickedFrom`, and the remaining annotated-ID candidates from that device. - Initialize a min-heap of these states ordered primarily by `used` (so devices with the fewest already-allocated replicas come first) and tie-broken by `pickedFrom` (so devices we have not touched in the current allocation are preferred when used counts match). - On each iteration pop the best device, take one of its remaining replicas, increment its counters, and push it back if more remain. Total cost drops to O(n log m). The tie-break semantics from PR NVIDIA#1788 are preserved unchanged; existing tests still pass without modification. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-Authored-By: runatom-ai <258621014+runatom-ai@users.noreply.github.com> Signed-off-by: Jonathan Meiri <33288957+Meiri28@users.noreply.github.com>
jonathan-meiri
force-pushed
the
optimize-distributedalloc-heap
branch
from
July 14, 2026 12:50
5efcd18 to
2319fe5
Compare
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
2 participants
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.
Summary
Follow-up optimization stacked on #1788. Until #1788 lands, this PR's cumulative diff includes that commit too — the new work here is in the second commit (the heap refactor). Not for review until #1788 lands.
Replaces the per-iteration
sort.SliceindistributedAllocwith a min-heap keyed by(used, pickedFrom). Brings the loop fromO(n² log n)toO(n log m)wherenis replicas requested andmis the number of physical GPUs touched in this allocation. Same correctness as #1788; same tests pass.Practically,
nandmare small in real configurations and the wall-clock impact is invisible — this is structural cleanliness, not a hot-path speedup.Opened as draft so it doesn't enter the review queue alongside #1788. Happy to mark it ready as a separate follow-up after #1788 lands, or fold the change into #1788 if that's preferred.
Contributed by @Meiri28 on behalf of @runatom-ai.