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 |
jonathan-meiri
force-pushed
the
optimize-distributedalloc-heap
branch
from
July 14, 2026 12:50
5efcd18 to
2319fe5
Compare
Contributor
|
Heap logic looks good. Though this PR must wait for #1621 first. Will need a rebase after that. |
jonathan-meiri
force-pushed
the
optimize-distributedalloc-heap
branch
from
July 22, 2026 08:24
2319fe5 to
6848b06
Compare
Author
|
Rebased on top of #1621 (now merged). The heap now sits on top of the new Ready for another look when you have a moment, @rajatchopra. |
jonathan-meiri
force-pushed
the
optimize-distributedalloc-heap
branch
from
July 26, 2026 06:28
6848b06 to
06c33eb
Compare
Follow-up on top of NVIDIA#1621, which introduced the shared greedyAlloc loop with a pluggable replicaComparator (distributed vs packed). The loop still sorts the full candidate slice 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 greedyAlloc to bucket candidates by their underlying physical device into a small gpuAllocState per device, holding a shared *replicaCount, the pickedFrom counter, and the remaining candidate IDs. A gpuPriorityQueue defers to the caller-supplied replicaComparator on allocated() for primary ordering and to pickedFrom for the tie-break (unchanged semantics). Each iteration pops the best device, takes one of its remaining replicas, updates counters, and pushes it back if any remain. Total cost drops to O(n log m). Both allocation policies (distributed and packed) benefit; no behavior change — the existing test suite (TestDistributedAlloc, TestPackedAlloc, TestPackedVsDistributedContrast, TestDistributedAlloc_PartiallyAllocated_DistributesAcrossDistinctGPUs, etc.) passes unchanged. 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 28, 2026 09:45
06c33eb to
36e7ad1
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
3 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.