[xla:gpu][ROCm] VMM command-buffer: skip-remap + copy-into-shadow for launch-bound regimes#1001
[xla:gpu][ROCm] VMM command-buffer: skip-remap + copy-into-shadow for launch-bound regimes#1001phambinhfin wants to merge 10 commits into
Conversation
The VA-remapping command-buffer path (ExecuteWithVaRemapping) unmapped the entire reserved VA range and re-mapped every command-buffer allocation on every step, even though most slots (model weights) keep the same backing physical buffer step-to-step. On ROCm the VA unmap is the dominant per-step cost, so this full remap is wasteful. Wire the already-in-tree, unit-tested se::MemoryReservation::ScopedMapping:: Remap() per-slice skip primitive into the per-step path: cache each slot's source (BFC) address from the previous step and mark remap_required only for slots whose backing buffer moved. Unchanged slots keep their mapping (no unmap/map/SetAccess), and when nothing changed the unmap-event sync is skipped entirely (steady state issues zero driver calls). ROCm-gated (VmmRemapSkipEnabled); the CUDA path keeps its existing reset()+MapTo behavior. XLA_VMM_SKIP_REMAP=0 forces the legacy path for A/B. Measured (8xMI300X, FSDP=8, 8 layers, NEVER_UPDATE, 30 steps), median ms/step: llama2-7b: 1096.5 -> 552.0 (1.99x) llama3-8b: 1284.0 -> 616.0 (2.08x) Identical converged loss; no crash.
Small command-buffer slices (tiny scale/scalar/metric buffers that churn their address every step) are the worst case for VMM per-step remapping: a full hipMemUnmap/hipMemMap/hipMemSetAccess round-trip costs ~2-4ms even for a few bytes, and there can be many such slices per step. This change keeps those slices OUT of the contiguous VA reservation and instead gives each a stable shadow device buffer (allocated once, fixed address baked into the replayed command buffer). Each step the shadow is refreshed with a stream-ordered D2D copy (real->shadow before execute, shadow->real after) rather than a remap. Gated by XLA_VMM_TMP_COPY_THRESHOLD (bytes; 0 = disabled, default), ROCm-only, so the CUDA path and default ROCm behavior are unchanged. Combined with the per-slot skip-remap, this makes capturing collectives into the command buffer (NEVER_UPDATE) a net win in launch-bound (tensor-parallel) regimes. Loss/correctness unchanged.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| absl::flat_hash_map<BufferAllocation::Index, se::DeviceAddressBase> | ||
| small_shadow; |
There was a problem hiding this comment.
Bug: Shadow buffers are leaked on VaRanges destruction.
small_shadow holds DeviceAddressBase values (non-owning raw pointer+size) allocated via executor->Allocate() (.cc line 494), but VaRanges has no destructor and DeviceAddressBase does not free on destruction. When GpuExecutableBufferAllocator is destroyed, its module_va_ranges_ map destroys all VaRanges structs, but the device memory pointed to by small_shadow entries is never deallocated.
A destructor (or RAII wrapper like ScopedDeviceAddress) is needed to call executor->Deallocate() for each shadow buffer. This requires either storing the executor pointer in VaRanges or switching to a self-freeing wrapper type.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. VaRanges now has a destructor (lines 228-234 in the header) that iterates small_shadow and calls executor->Deallocate() for each entry, with the owning executor stored at line 225.
| } | ||
| se::DeviceAddressBase shadow = executor->Allocate(sz); |
There was a problem hiding this comment.
Concern: executor->Allocate(sz) uses the default memory_space=0.
The shadow buffer is always allocated in memory space 0, but the original buffer's allocation may have a non-zero color/memory_space (e.g., collective memory at color 1). If a small command-buffer slice happens to be a collective buffer, the shadow would be allocated in the wrong memory space.
This is unlikely for the "small scale/scalar/metric" use case described in the PR, but the code does not enforce the assumption. Consider either asserting allocation.color() == 0 for small-copy slices or passing the correct memory_space to Allocate.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The shadow allocation now reads the allocation's color and passes it as memory_space (lines 526-530 in .cc), so non-zero-color buffers are no longer silently placed in space 0.
| } | ||
| static const uint64_t threshold = []() -> uint64_t { | ||
| const char* v = std::getenv("XLA_VMM_TMP_COPY_THRESHOLD"); | ||
| return (v != nullptr) ? std::strtoull(v, nullptr, 10) : 0; |
There was a problem hiding this comment.
Nit: std::strtoull silently returns 0 for non-numeric strings.
std::strtoull("auto", nullptr, 10) returns 0 and std::strtoull("999999999999999999999", nullptr, 10) returns ULLONG_MAX -- both silently. The comment on line 113 mentions a future "auto" mode.
If a user sets XLA_VMM_TMP_COPY_THRESHOLD=auto today, they get threshold=0 (disabled) with no warning. Consider logging a warning when parsing fails, or at least documenting that non-numeric values silently disable the feature.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. ParseVmmCopyThresholdBytes now uses end-pointer and errno checking (lines 91-101 in .cc) and logs a WARNING for non-numeric values like "auto", making the failure visible instead of silent.
| bool VmmRemapSkipEnabled(absl::string_view platform_name) { | ||
| if (platform_name != "ROCM") { | ||
| return false; | ||
| } | ||
| static const bool enabled = [] { | ||
| const char* v = std::getenv("XLA_VMM_SKIP_REMAP"); | ||
| if (v == nullptr) { | ||
| return true; // default ON for ROCm |
There was a problem hiding this comment.
Nit: PR description says "opt-in via env vars" but skip-remap is default-ON for ROCm.
VmmRemapSkipEnabled returns true by default on ROCm (line 95), meaning XLA_VMM_SKIP_REMAP=0 must be set to disable it. This is the opposite of the "opt-in" framing in the PR description.
The "Notes for reviewers" section also says "gated and default-off-equivalent for CUDA / existing ROCm callers" which is only accurate for the copy-into-shadow path (threshold=0 by default). The skip-remap optimization is default-ON for all ROCm callers. Consider updating the PR description to match the code, or making this opt-in as well if the intent is cautious rollout.
There was a problem hiding this comment.
Resolved ✓ — the PR description now clearly states skip-remap is default ON for ROCm and only copy-into-shadow is opt-in, matching the code behavior.
| if (it == va_ranges_->small_shadow.end() || | ||
| it->second.size() < sz) { | ||
| if (it != va_ranges_->small_shadow.end()) { | ||
| executor->Deallocate(&it->second); | ||
| } | ||
| se::DeviceAddressBase shadow = executor->Allocate(sz); | ||
| if (shadow.is_null()) { | ||
| return Internal( | ||
| "Failed to allocate %d-byte shadow for small command-buffer " | ||
| "slice %d", | ||
| sz, i); | ||
| } | ||
| it = va_ranges_->small_shadow.insert_or_assign(i, shadow).first; |
There was a problem hiding this comment.
Potential issue: Shadow reallocation races with GPU access.
If it->second.size() < sz (shadow is too small), the old shadow is deallocated (line 491) and a new one is allocated immediately. However, the old shadow address may still be baked into a previously-captured command buffer that the GPU is executing. The Deallocate call frees the device memory while the GPU may still be accessing it.
In practice, under NEVER_UPDATE modes with static shapes, sizes should not change, so this path should not be hit. But if it ever is, it would need an unmap_event->Synchronize() (or equivalent) before deallocating the old shadow. Consider adding a DCHECK that this path is not taken after the first step, or adding the synchronization.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The shadow reallocation path (lines 534-544 in .cc) now synchronizes on unmap_event before deallocating the old shadow, preventing the GPU race. The code also adds a comment noting this branch is not expected under NEVER_UPDATE with static shapes.
Claude Review SummaryPR: Copy-into-shadow for small command-buffer VA slots (ROCm) This PR adds two ROCm-only optimizations to the VMM command-buffer VA remapping path: (1) skip-remap, which caches each slot's source address and skips unmap/map for unchanged slots, and (2) copy-into-shadow, which replaces VA remapping of small churning buffers with stable shadow buffers refreshed via D2D copies. The CUDA path is correctly left unchanged. Overall well-structured with clear comments. Key findings (5 inline comments):
No tests are included for the new helper functions ( Non-test delta: +224/−31 = 255 lines — well within review scope. |
github-actions
Bot
removed
the
claude-review
Fixes from the PR #1001 review of the copy-into-shadow path: - Fix shadow-buffer leak: VaRanges now stores its owning StreamExecutor and a destructor frees every small_shadow DeviceAddressBase (non-owning pointers were previously never deallocated). VaRanges is marked non-copyable. - Avoid a use-after-free race when a shadow slot grows: synchronize the unmap_event before deallocating the old (still possibly GPU-referenced) shadow buffer. - Allocate the shadow in the original buffer's memory space (color) instead of always memory_space=0, so a non-default-color slice is not silently misplaced. - XLA_VMM_TMP_COPY_THRESHOLD now warns on empty/non-numeric values (e.g. "auto") instead of silently parsing to 0 via strtoull. - Extract the env-var decision logic into pure, uncached vmm_internal helpers and add gpu_executable_buffer_allocator_test.cc covering them. ROCm-gated; CUDA path and default ROCm behavior unchanged.
| absl::string_view s(env_value); | ||
| return !(s == "0" || s == "false" || s == "off" || s == "FALSE" || | ||
| s == "OFF"); |
There was a problem hiding this comment.
Nit: Case-insensitive env var parsing is incomplete.
ParseVmmRemapSkipEnabled handles "false"/"FALSE" and "off"/"OFF" but not mixed-case variants like "False", "Off", "No", "no", or "NO". Consider using absl::AsciiStrToLower() or absl::EqualsIgnoreCase() for a more robust comparison:
| absl::string_view s(env_value); | |
| return !(s == "0" || s == "false" || s == "off" || s == "FALSE" || | |
| s == "OFF"); | |
| absl::string_view s(env_value); | |
| return !(s == "0" || absl::EqualsIgnoreCase(s, "false") || | |
| absl::EqualsIgnoreCase(s, "off")); |
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. ParseVmmRemapSkipEnabled now uses absl::EqualsIgnoreCase for all string comparisons (including "no"/"No"/"NO"), and unit tests confirm mixed-case variants are handled correctly.
| return Internal( | ||
| "Failed to allocate %d-byte shadow for small command-buffer " | ||
| "slice %d", | ||
| sz, i); |
There was a problem hiding this comment.
Bug: Internal() uses absl::StrCat, not printf-style formatting — %d placeholders are emitted literally.
Internal(fmt, sz, i) is an overload that calls absl::StrCat(fmt, sz, i), so the error message will read something like "Failed to allocate %d-byte shadow for small command-buffer slice %d655" — the %d tokens are not substituted.
Same issue on line 519 ("Command buffer allocation %d has null address").
Fix: use absl::StrCat/absl::StrFormat explicitly, or just concatenate:
| return Internal( | |
| "Failed to allocate %d-byte shadow for small command-buffer " | |
| "slice %d", | |
| sz, i); | |
| return Internal( | |
| absl::StrCat("Failed to allocate ", sz, | |
| "-byte shadow for small command-buffer slice ", | |
| i)); |
There was a problem hiding this comment.
Resolved — this was a false positive. The xla::Internal() used here is defined via XLA_ERROR_WITH_STRFORMAT_AND_BACKTRACE in xla/util.h (line 302), which uses absl::StrFormat, not absl::StrCat. The %d placeholders are correctly substituted. The tsl::errors::Internal in tsl/platform/errors.h (StrCat-based) is a different, deprecated overload that is not the one resolved here. Apologies for the noise.
| const int64_t memory_space = | ||
| (i < static_cast<BufferAllocation::Index>( | ||
| owner_->allocations_.size())) | ||
| ? static_cast<int64_t>(owner_->allocations_[i]->color()) | ||
| : 0; |
There was a problem hiding this comment.
Concern: memory_space silently falls back to 0 when index is out of range.
If i >= owner_->allocations_.size(), the code defaults memory_space to 0 instead of reporting an error. This would silently place a shadow buffer in the wrong memory space, masking a bug in the allocation index tracking. Consider returning an error instead:
| const int64_t memory_space = | |
| (i < static_cast<BufferAllocation::Index>( | |
| owner_->allocations_.size())) | |
| ? static_cast<int64_t>(owner_->allocations_[i]->color()) | |
| : 0; | |
| const int64_t memory_space = | |
| (i < static_cast<BufferAllocation::Index>( | |
| owner_->allocations_.size())) | |
| ? static_cast<int64_t>(owner_->allocations_[i]->color()) | |
| : [&]() -> int64_t { LOG(WARNING) << "Allocation index " << i << " out of range, defaulting memory_space to 0"; return 0; }(); |
Or more cleanly, return Internal(...) since this condition would indicate a logic bug.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. Out-of-range allocation index now returns an Internal error (lines 537-541 in .cc) instead of silently defaulting memory_space to 0.
| if (it == va_ranges_->small_shadow.end() || | ||
| it->second.size() < sz) { | ||
| if (it != va_ranges_->small_shadow.end()) { | ||
| // The old shadow address may still be baked into a | ||
| // previously-captured command buffer the GPU is replaying. This | ||
| // branch is only reachable on a later step (the slice grew), so | ||
| // unmap_event has been recorded; wait for that GPU work before | ||
| // freeing the old shadow. Under NEVER_UPDATE with static shapes | ||
| // this is not expected to be hit. | ||
| if (va_ranges_->unmap_event != nullptr) { | ||
| RETURN_IF_ERROR(va_ranges_->unmap_event->Synchronize()); | ||
| } | ||
| executor->Deallocate(&it->second); | ||
| } | ||
| se::DeviceAddressBase shadow = executor->Allocate(sz, memory_space); | ||
| if (shadow.is_null()) { | ||
| return Internal( | ||
| "Failed to allocate %d-byte shadow for small command-buffer " | ||
| "slice %d", | ||
| sz, i); | ||
| } | ||
| it = va_ranges_->small_shadow.insert_or_assign(i, shadow).first; |
There was a problem hiding this comment.
Concern: Shadow reallocation changes the baked-in command buffer address under NEVER_UPDATE.
The previous review flagged the GPU race on shadow deallocation (now fixed with the unmap_event sync on line 541-542). However, a second issue remains: when a shadow is reallocated (because the slice grew), executor->Allocate(sz, memory_space) returns a new device address. Under NEVER_UPDATE, the command buffer was captured once and has the old shadow address baked in — it will not be re-captured, so subsequent replays read/write the freed address.
The comment on line 539-540 acknowledges this ("Under NEVER_UPDATE with static shapes this is not expected to be hit"), which is correct in practice. But if it does get hit, the result is silent memory corruption rather than a clear failure. Consider returning an error in this branch when update_mode_ == NEVER_UPDATE, so a shape change is caught immediately rather than manifesting as a mysterious data corruption.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. Instead of reallocating to a new address (causing silent corruption under NEVER_UPDATE), the code now returns an Internal error when a small-copy slice grows (lines 553-565 in .cc), failing clearly rather than silently.
| for (SmallCopy& c : small_copies) { | ||
| RETURN_IF_ERROR(exec_stream->Memcpy(&c.shadow, c.real, c.size)); | ||
| } | ||
| XLA_VLOG_DEVICE(2, device_ordinal) << absl::StreamFormat( | ||
| "VA remapping: copied %d small slices to shadow (threshold=%d B)", | ||
| static_cast<int>(small_copies.size()), copy_threshold); | ||
| } | ||
|
|
||
| RETURN_IF_ERROR(execute(remapped_buffer_allocations)); | ||
|
|
||
| // Copy-BACK small-slice outputs (shadow->real), stream-ordered after the | ||
| // command buffer wrote results (e.g. loss / token-count scalars) into the | ||
| // stable shadow, propagating them to the real buffers the host/later thunks | ||
| // read. Runs before the unmap event is recorded. | ||
| if (!small_copies.empty()) { | ||
| se::Stream* exec_stream = run_options_->stream(); | ||
| for (SmallCopy& c : small_copies) { |
There was a problem hiding this comment.
Nit: Copy-in and copy-back run unconditionally for all small-copy slices.
Every small-copy slice gets both real->shadow (before execute) and shadow->real (after execute). For input-only buffers (e.g., scale factors), the copy-back is wasted work; for output-only buffers, the copy-in is wasted. With many small slices this could accumulate.
Consider checking BufferAllocation::is_readonly() or access mode to skip unnecessary copies. This may not matter with the expected small number of tiny slices today, but worth noting for future optimization or at least a comment documenting the decision.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The bidirectional copy is now documented as an intentional design choice (lines 480-486 in .cc), with a note that a future optimization could use BufferUse read/write to elide the unused direction.
Re-review summaryThe four issues from the initial review have been addressed well in this revision — shadow buffer lifecycle, memory-space color propagation, Five new findings posted inline:
The |
github-actions
Bot
removed
the
claude-review
- Fail loudly when a copy-into-shadow slice grows after first capture: under trace-once replay (NEVER_UPDATE) the shadow address is baked into the command buffer and cannot move, so reallocating would silently corrupt the replayed buffer. Return an error instead of reallocating (also removes the prior dealloc/realloc + GPU-race path). Static shapes never hit this. - Return an error if a command-buffer slice index is out of range instead of silently defaulting the shadow's memory space to 0. - Make XLA_VMM_SKIP_REMAP parsing case-insensitive for false/off/no via absl::EqualsIgnoreCase; extend the unit test accordingly. - Document that the small-slice copy is intentionally bidirectional (direction- agnostic, tiny buffers); note a future BufferUse-based optimization. Note: the reviewer's "Internal() uses StrCat" comment is a false positive -- xla::Internal is absl::StrFormat-based (xla/util.h), so %d is substituted. ROCm-gated; CUDA path and default ROCm behavior unchanged.
Re-review Summary (revision 2)Good progress on this revision — 5 of my previous findings are now addressed:
1 prior finding still open: Skip-remap is default-ON for ROCm while the PR description says "opt-in" — consider aligning the description. No new issues found in this revision. |
github-actions
Bot
removed
the
claude-review
| ASSIGN_OR_RETURN(se::MemoryReservation::ScopedMapping remapped, | ||
| std::move(*va_ranges_->scoped_mapping) | ||
| .Remap(absl::MakeSpan(remaps))); |
There was a problem hiding this comment.
Bug: ASSIGN_OR_RETURN after std::move leaves scoped_mapping in a moved-from state on failure.
If Remap() fails, ASSIGN_OR_RETURN returns early. At that point, *va_ranges_->scoped_mapping has already been consumed by std::move() — the optional still has_value() but the contained ScopedMapping is moved-from (empty/null). The scoped_mapping.reset() on the legacy path (line 660) is never reached because the error return exits the function.
On the next call to ExecuteWithVaRemapping, can_skip_remap (line 617) checks scoped_mapping.has_value() which is still true, so the skip-remap path tries std::move(*va_ranges_->scoped_mapping).Remap(...) on a null/moved-from ScopedMapping — likely UB or a crash.
Fix: reset the optional before returning the error:
auto remap_result = std::move(*va_ranges_->scoped_mapping)
.Remap(absl::MakeSpan(remaps));
if (!remap_result.ok()) {
va_ranges_->scoped_mapping.reset();
return remap_result.status();
}
va_ranges_->scoped_mapping = std::move(*remap_result);There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The Remap() failure path now explicitly resets scoped_mapping before returning the error (lines 664-670 in .cc), preventing the moved-from optional from being misinterpreted as valid on the next step.
| ~VaRanges() { | ||
| if (executor != nullptr) { | ||
| for (auto& [index, shadow] : small_shadow) { | ||
| executor->Deallocate(&shadow); | ||
| } | ||
| } |
There was a problem hiding this comment.
Bug: Destructor deallocates shadow buffers without GPU synchronization.
The destructor calls executor->Deallocate(&shadow) for each shadow buffer, but does not first synchronize with the GPU. The last step's command buffer may still be reading from or writing to these shadow buffers on the device (the unmap_event records after execute() returns, and nothing here waits for it).
On ROCm, hipFree on memory still in use by a running kernel is undefined behavior. This is distinct from the previously-flagged leak issue (#1) — that was about missing deallocation entirely; this is about deallocating while the GPU is still using the memory.
Fix: synchronize before the deallocation loop:
~VaRanges() {
if (executor != nullptr) {
if (unmap_event) {
unmap_event->Synchronize().IgnoreError();
}
for (auto& [index, shadow] : small_shadow) {
executor->Deallocate(&shadow);
}
}
}There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The VaRanges destructor now synchronizes on unmap_event before deallocating shadow buffers (lines 240-252 in the header), preventing use-after-free on device memory still in use by a running kernel.
| if (env_value == nullptr) { | ||
| return true; // default ON for ROCm | ||
| } |
There was a problem hiding this comment.
Nit: Empty-string handling is inconsistent with ParseVmmCopyThresholdBytes.
ParseVmmCopyThresholdBytes (line 89) treats an empty env var ("") as equivalent to unset (returns 0 / disabled). But here, ParseVmmRemapSkipEnabled does not check for empty string — if a user does export XLA_VMM_SKIP_REMAP= (setting to empty), the value "" does not match any of the falsey values ("0", "false", "off", "no"), so it returns true (enabled).
This is probably fine in practice (empty → enabled is a reasonable default), but it is inconsistent between the two sibling functions and the test file has no coverage for the empty-string case.
Consider adding an explicit empty-string check for symmetry:
if (env_value == nullptr || *env_value == '\0') {
return true; // default ON for ROCm
}There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. ParseVmmRemapSkipEnabled now explicitly checks for empty string on line 78 (*env_value == '\0') and treats it the same as unset (returns true / default ON), consistent with ParseVmmCopyThresholdBytes' empty-string handling. The test file also covers the empty-string case (RocmEmptyStringDefaultsOn).
| const bool can_skip_remap = | ||
| skip_remap_enabled && va_ranges_->scoped_mapping.has_value() && | ||
| !mapping_descriptors.empty() && | ||
| va_ranges_->last_mapped_src_addrs.size() == new_src_addrs.size(); |
There was a problem hiding this comment.
Minor: Skip-remap address comparison checks vector size but not allocation identity.
The can_skip_remap guard checks last_mapped_src_addrs.size() == new_src_addrs.size() to ensure the vectors are comparable element-by-element. However, it does not verify that the same allocation indices produced those vectors. If the set of VA-mapped allocation indices were to change between steps while keeping the same count (e.g., one allocation shrinks below the copy threshold while another grows above it), the per-slot comparison would compare addresses of different allocations, potentially concluding "no change" when the mapping is actually stale.
Under the NEVER_UPDATE / static-shapes contract this should not happen, but a cheap DCHECK that the allocation_va_offsets keys match the previous step would make the invariant explicit and catch violations early:
DCHECK(allocation_va_offsets_keys_match_previous_step);There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. A DCHECK now verifies per-slot reservation offsets match the previous step (lines 639-644 in .cc), catching any slot->allocation identity change that would invalidate the per-slot address comparison. The debug-only last_mapped_offsets vector tracks these offsets.
| // ROCm skip-remap booking. Source (BFC) device address mapped into each | ||
| // command-buffer slot during the previous step, in ascending | ||
| // reservation-offset order (same order the mapping descriptors are built). | ||
| // A slot whose source address is unchanged this step keeps its existing | ||
| // mapping instead of being unmapped+remapped, and when no slot changes the | ||
| // whole unmap/map/SetAccess (and the unmap-event sync) is skipped. Empty | ||
| // until the first mapping is established. ROCm-only; see | ||
| // VmmRemapSkipEnabled(). | ||
| std::vector<const void*> last_mapped_src_addrs; | ||
|
|
||
| // ROCm copy-into-shadow (env XLA_VMM_TMP_COPY_THRESHOLD): small | ||
| // command-buffer slices (tiny scale/scalar/metric buffers that churn their | ||
| // address every step) are kept OUT of the VA reservation and instead given |
There was a problem hiding this comment.
Nit: VaRanges mutable fields lack ABSL_GUARDED_BY(mutex) annotations.
The new fields added here (last_mapped_src_addrs, small_shadow) and the existing ones above (va_reservation, unmap_event, scoped_mapping) are all accessed under mutex in the .cc file, but none carry ABSL_GUARDED_BY(mutex) annotations. Adding them would let the thread-safety analysis catch any future unguarded access at compile time.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. All mutable VaRanges fields now carry ABSL_GUARDED_BY(mutex) annotations: va_reservation (line 189), unmap_event (line 194), scoped_mapping (line 199), last_mapped_src_addrs (line 209), last_mapped_offsets (line 217), and small_shadow (line 228).
Re-review Summary (Round 2)Most issues from the initial review have been addressed — resolved 1 remaining open thread. 5 new findings posted inline:
🤖 This review was generated by Claude. |
github-actions
Bot
removed
the
claude-review
…shadow Fixes from the second claude-review pass on PR #1001: - Skip-remap: if ScopedMapping::Remap() fails, reset scoped_mapping before returning the error. Remap() consumes the mapping via std::move, so leaving the (now moved-from) optional with has_value()==true would make the next step's skip path call Remap() on an empty mapping (UB). Resetting forces the legacy reset()+MapTo() recovery path instead. - VaRanges destructor: synchronize unmap_event before deallocating the shadow buffers. The last step's command buffer may still be using them on the device, and hipFree on in-use memory is UB on ROCm. Marked ABSL_NO_THREAD_SAFETY_ANALYSIS (runs at teardown, no other thread present). - Add ABSL_GUARDED_BY(mutex) to the mutex-protected VaRanges fields (va_reservation, unmap_event, scoped_mapping, last_mapped_src_addrs, small_shadow) so the thread-safety analyzer enforces the locking. - Debug-only DCHECK that the per-slot reservation offsets are stable across steps, making the skip-remap "same allocation per slot" invariant explicit (new last_mapped_offsets companion vector, populated only in non-NDEBUG). - ParseVmmRemapSkipEnabled: treat empty XLA_VMM_SKIP_REMAP the same as unset (default ON), matching ParseVmmCopyThresholdBytes; add a unit test. ROCm-gated; CUDA path and default ROCm behavior unchanged.
Re-review SummaryAll 15 previously flagged issues have been addressed in this revision. Posted "Resolved" on the 5 remaining open threads:
No new issues found. The skip-remap and copy-into-shadow implementations look solid — error recovery paths are correct, shadow buffer lifecycle is properly managed, and the env-var parsing has good test coverage. |
github-actions
Bot
removed
the
claude-review
Re-review SummaryAll 15 previously flagged issues have been addressed in this revision. No new findings. The skip-remap and copy-into-shadow optimizations look correct: shadow buffers are properly freed with GPU synchronization in the destructor, |
github-actions
Bot
removed
the
claude-review
Comment-only change: remove specific device timing figures and cross-vendor comparisons from the source comments, keeping the descriptions qualitative (what the path does and why) without quoting hardware-specific numbers. No functional change.
… flags Address upstream review: don't read environment variables directly. Replace the XLA_VMM_SKIP_REMAP / XLA_VMM_TMP_COPY_THRESHOLD getenv reads with standard XLA flags in xla.proto / DebugOptions. - xla.proto: add xla_gpu_command_buffer_vmm_skip_remap (bool, default true) and xla_gpu_command_buffer_vmm_copy_threshold_bytes (int64, default 0). - debug_options_flags.cc: set defaults and register both flags. - gpu_executable_buffer_allocator: drop the getenv-based env parsers and the process-static caches; read the flag values from the module's DebugOptions (owner_->debug_options_, with a null-guard fallback to the defaults). The vmm_internal helpers now apply only the ROCm-only gating (and clamp a negative threshold to 0) so they stay unit-testable. - Update the unit test to cover the flag-based helpers. ROCm-gated; CUDA path and default ROCm behavior unchanged.
…opy-shadow # Conflicts: # xla/xla.proto
…hVaRemapping Rename the two VMM command-buffer knobs with the xla_gpu_experimental_ prefix (xla_gpu_experimental_command_buffer_vmm_skip_remap and xla_gpu_experimental_command_buffer_vmm_copy_threshold_bytes); they are meant to fold into the default behavior and be removed, so they are not a stable contract. Field numbers are unchanged. Break the ~400-line ExecuteWithVaRemapping into a thin orchestrator plus five helpers on ExecutionScope: EnsureVaReservation, ComputeReservationOffsets, GetOrCreateSmallShadow, BuildStepMappingPlan and EstablishMapping. Behavior is unchanged; the per-executor mutex contract is preserved via ABSL_EXCLUSIVE_LOCKS_REQUIRED and the allocator unit test still passes.
…BUG) The per-slot reservation-offset bookkeeping used by the skip-remap DCHECK was guarded by #ifndef NDEBUG to keep the release hot path free of the work. Measured in isolation (-O2 -DNDEBUG) the maintenance is ~10-150 ns per step at typical command-buffer slot counts (tens to hundreds), <=0.3 us at 512 slots, and N*8 bytes of RAM -- negligible next to the microsecond-to-millisecond hipMemUnmap/hipMemMap it sits beside. Drop the preprocessor guard and always populate it for readability; the DCHECK is still compiled out in release.
| // ROCm-only by design: other platforms keep their existing full | ||
| // reset()+MapTo() behavior unchanged. Set XLA_VMM_SKIP_REMAP=0 to force the | ||
| // legacy path on ROCm for A/B measurement. | ||
| bool VmmRemapSkipEnabled(absl::string_view platform_name) { |
There was a problem hiding this comment.
Should not all of this go to debug options flags not bespoke env parsing?
2 participants
Summary (DRAFT — for review, not for merge yet)
Two ROCm-gated optimizations to the VMM command-buffer path (
ExecuteWithVaRemapping) that make capturing the graph once and replaying it (NEVER_UPDATE) a net win in launch-bound (tensor-parallel) regimes, where per-step VA remapping was previously eating the benefit.Both changes are ROCm-only; the CUDA path and default ROCm behavior are unchanged.
1. Skip remapping unchanged VA slots
The old path unmapped the whole reserved VA range and re-mapped every command-buffer allocation every step, even though most slots (weights) keep the same backing physical buffer step-to-step. On ROCm the unmap is the dominant per-step cost.
This caches each slot's source (BFC) address from the previous step and only marks
remap_requiredfor slots whose backing buffer actually moved, reusing the already-in-tree, unit-testedse::MemoryReservation::ScopedMapping::Remap()primitive. In steady state (nothing moved) it issues zero driver calls and skips the unmap-event sync.XLA_VMM_SKIP_REMAP— default ON for ROCm; set0to force the legacy path for A/B.2. Copy-into-shadow for tiny churning slots
Small slices (scale/scalar/metric buffers) change address every step and are the worst case for remapping (a full unmap/map/SetAccess round-trip costs milliseconds even for a few bytes, and there can be many per step). This keeps them out of the contiguous VA reservation and gives each a stable shadow buffer (fixed address baked into the replayed command buffer), refreshed each step with a stream-ordered D2D copy instead of a remap.
XLA_VMM_TMP_COPY_THRESHOLD(bytes;0= disabled = default), opt-in.Performance
Measured with MaxText (synthetic data, identical converged loss across configurations). Numbers are relative per-step speedups only (steady-state median, after warm-up), normalized to the default baseline configuration; absolute throughput is intentionally omitted.
Each model is compared as: the default baseline config (what people run today) vs. the VMM-optimized config (trace-once replay + collectives captured + per-step remap minimized). Baseline is normalized to
1.00×; higher is faster.8×MI300X
8×MI355
A genuine end-to-end per-step improvement across popular dense LLMs on both platforms, largest on the most launch-bound models. Loss and correctness are unchanged in every case; in the bandwidth-bound FSDP regime the same configuration is performance-neutral (no regression).
Notes for reviewers
XLA_VMM_TMP_COPY_THRESHOLD=0by default). Both are no-ops for CUDA and for existing default ROCm callers.gpu_executable_buffer_allocator_test).Test plan
XLA_VMM_SKIP_REMAP=0/1andXLA_VMM_TMP_COPY_THRESHOLD=0/65536on llama2 / llama3