[ROCm][VMM] Combined L1-L4 ROCm VMM allocator (internal integration + validation)#930
Closed
phambinhfin wants to merge 1 commit into
Closed
[ROCm][VMM] Combined L1-L4 ROCm VMM allocator (internal integration + validation)#930phambinhfin wants to merge 1 commit into
phambinhfin wants to merge 1 commit into
Conversation
Combined branch bringing together the four layered ROCm VMM allocator PRs for end-to-end internal validation on AMD hardware: - Layer 1: RocmRawMemoryAllocation (openxla#43942) - Layer 2: RocmMemoryReservation + base MemoryReservation::Remap (openxla#43943) - Layer 3: RocmVmmAllocator (openxla#43946) - Layer 4: RocmDeviceAddressVmmAllocator + generic base + factory + tests (openxla#43947) This is the integration view of the series; the per-layer PRs above are how the work is proposed upstream. ROCm-only allocator targets; the generic base + factory are vendor-neutral and used by both CUDA and ROCm.
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.
Purpose
Internal integration branch that combines all four layered ROCm VMM allocator PRs into a single tree so fork CI can validate the whole stack together. The work is proposed upstream as four separate PRs (links below); this PR is the integration/validation view, not the upstream submission.
RocmRawMemoryAllocation(hipMemCreate/Release RAII)RocmMemoryReservation+ baseMemoryReservation::RemapRocmVmmAllocator(all-in-one create+reserve+map)RocmDeviceAddressVmmAllocator+ generic base + factory + testsCombined diff: +2432 / −17 across 21 files (allocator layers only — no
gpu_executable/PJRT enablement, no unrelated churn).Local validation (AMD MI300X, gfx942)
Built from
/workspace/rocm-jax/jaxwith--config=rocmagainst this working tree.End-to-end performance (why this stack exists)
These numbers come from the optimization that rides on top of this allocator (the selective copy-vs-remap pass in
gpu_executable.cc, #926). They reflect the full VMM stack (this allocator +NEVER_UPDATE+ the copy-vs-remap optimization) and are included here to motivate the allocator layers — the allocator alone is the enabler, not the source of the speedup.MI300X x8, 8 layers / 30 steps, median of steps ≥3; 0 unmap errors.
auto + collective leg (
coll_vmm): disabled (=0) vs auto (=1)=0)=1)vmm leg (no collectives): disabled (
=0) vs auto (=1)=0)=1)The per-step saving (~6–13 ms) is roughly fixed (work removed = remap of the small churning slices), a net win that never regresses;
automatches a hand-tuned fixed 4 KB cutoff.Notes