[ROCm][VMM] Selective copy-vs-remap for small command-buffer slices + auto threshold#926
[ROCm][VMM] Selective copy-vs-remap for small command-buffer slices + auto threshold#926phambinhfin wants to merge 4 commits into
Conversation
…dow instead of VA-remapping Adds an opt-in size threshold (env XLA_VMM_TMP_COPY_THRESHOLD, bytes; 0 = disabled = current behavior) for the command-buffer VA remapping path. Slices whose size is <= threshold (the tiny scale/scalar/metric buffers that churn addresses every step) are excluded from the VA reservation and instead given a stable shadow device buffer whose address is baked into the command buffer once. Each step the slice is refreshed with a bidirectional device-to-device copy on the execution stream: real->shadow before execution (fresh inputs) and shadow->real after execution (propagate output scalars such as loss/token-count back to the buffers the host reads). This avoids the expensive hipMemUnmap/hipMemMap/hipMemSetAccess round-trip (which serializes across peer devices on ROCm) for those small buffers, while large slices still remap. This is the "balanced copy vs remap" idea discussed with the NVIDIA reviewer. llama2_7b, FSDP=8, 8xMI300X, coll_vmm leg (30 steps), loss matches baseline: threshold=0 : 132 ms (remap all) threshold=256B : 120 ms threshold=4096B : 121 ms threshold=65536B: 122 ms
…sized to free HBM
Turn XLA_VMM_TMP_COPY_THRESHOLD into a 3-way control:
0 -> disabled: remap every command-buffer slice (original behavior)
1 -> auto: derive the byte cutoff from this device's HBM bandwidth and
capacity AND the memory currently free (via DeviceMemoryUsage), so the
stable shadow pool is always guaranteed to fit (no searching, no OOM)
>1 -> force that exact byte cutoff (explicit override / sweeps)
The auto path admits the smallest slices first while bounding (a) per-slice copy
time, (b) aggregate per-step copy time, and (c) the shadow footprint by a
fraction of currently-free HBM. Large tensors keep stable addresses across steps
(their remap is already skipped/free), so they stay on the remap path; only the
tiny churning scalar/scale/metric slices are copied into a stable shadow.
ROCm-only: the entire path is gated on DeviceAddressVmmAllocator, so on NVIDIA
the cast is null, ExecuteThunksWithVaRemapping is never entered, and the env var
has no effect. NVIDIA behavior is unchanged.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| uint64_t copy_threshold = 0; | ||
| { | ||
| const char* e = std::getenv("XLA_VMM_TMP_COPY_THRESHOLD"); | ||
| // Unset defaults to auto (mode 1) so the smart selector is the default. | ||
| uint64_t mode = 1; | ||
| if (e != nullptr) { | ||
| mode = std::strtoull(e, nullptr, 10); | ||
| } | ||
| if (mode == 0) { | ||
| copy_threshold = 0; // disabled | ||
| } else if (mode > 1) { | ||
| copy_threshold = mode; // explicit byte cutoff | ||
| } else { | ||
| // mode == 1: auto-derive from device properties + free memory. | ||
| // | ||
| // Greedy admit-smallest-first: pull command-buffer slices into the copy | ||
| // set in increasing size order while (a) no single slice's bidirectional | ||
| // copy exceeds kPerBufferCopyUs, and (b) the aggregate per-step copy stays | ||
| // under a budget that is the MIN of a copy-time cap and a memory cap. The | ||
| // memory cap is taken from the currently-free HBM (kFreeHbmFraction of | ||
| // free), which guarantees the shadow pool always fits -- "always enough | ||
| // space" -- rather than from a blind fraction of total capacity. Large | ||
| // tensors blow these budgets and stay on the remap path; since they keep | ||
| // stable addresses across steps their remap is already skipped (free). | ||
| constexpr double kPerBufferCopyUs = 8.0; | ||
| constexpr double kTotalCopyUs = 15.0; | ||
| constexpr double kShadowHbmFraction = 0.005; // of total (fallback) | ||
| constexpr double kFreeHbmFraction = 0.01; // of free (primary cap) | ||
| constexpr uint64_t kFallbackThreshold = 65536; | ||
| const se::DeviceDescription& dd = executor->GetDeviceDescription(); | ||
| const int64_t bw = dd.memory_bandwidth(); // bytes/sec | ||
| const int64_t hbm = dd.device_memory_size(); // bytes | ||
| if (bw <= 0) { | ||
| // Bandwidth unknown: fall back to a conservative cutoff that still | ||
| // captures scalar/scale buffers. | ||
| copy_threshold = kFallbackThreshold; | ||
| } else { | ||
| // acc accumulates 2*size per admitted slice: bidirectional copy moves | ||
| // 2 bytes per shadow byte, and the shadow is allocated per VA set, so | ||
| // 2*size also bounds the double-buffered footprint. | ||
| const double per_buffer_cap = | ||
| kPerBufferCopyUs * 1e-6 * static_cast<double>(bw) / 2.0; | ||
| double total_budget = kTotalCopyUs * 1e-6 * static_cast<double>(bw); | ||
| // Memory cap: prefer the live free-memory reading so the shadow pool is | ||
| // bounded by what the device can actually spare right now. | ||
| int64_t free_mem = 0; | ||
| int64_t total_mem = 0; | ||
| double mem_cap = 0.0; | ||
| if (executor->DeviceMemoryUsage(&free_mem, &total_mem) && | ||
| free_mem > 0) { | ||
| mem_cap = kFreeHbmFraction * static_cast<double>(free_mem); | ||
| } else if (hbm > 0) { | ||
| mem_cap = kShadowHbmFraction * static_cast<double>(hbm); | ||
| } | ||
| if (mem_cap > 0.0) { | ||
| total_budget = std::min(total_budget, mem_cap); | ||
| } | ||
| std::vector<uint64_t> sizes; | ||
| sizes.reserve(command_buffer_allocation_indexes_.size()); | ||
| for (BufferAllocation::Index i : command_buffer_allocation_indexes_) { | ||
| se::DeviceAddressBase b = buffer_allocations.GetDeviceAddress(i); | ||
| if (vmm_allocator->GetRawAllocation(executor->device_ordinal(), b) == | ||
| nullptr) { | ||
| continue; | ||
| } | ||
| sizes.push_back(b.size()); | ||
| } | ||
| std::sort(sizes.begin(), sizes.end()); | ||
| double acc = 0.0; | ||
| for (uint64_t s : sizes) { | ||
| if (static_cast<double>(s) > per_buffer_cap) break; | ||
| if (acc + 2.0 * static_cast<double>(s) > total_budget) break; | ||
| acc += 2.0 * static_cast<double>(s); | ||
| copy_threshold = s; | ||
| } | ||
| XLA_VLOG_DEVICE(2, executor->device_ordinal()) << absl::StreamFormat( | ||
| "VA remapping: auto copy threshold=%d B (bw=%d GB/s hbm=%d GB " | ||
| "free=%d GB budget=%.0f B)", | ||
| copy_threshold, bw / 1000000000, hbm / 1000000000, | ||
| free_mem / 1000000000, total_budget); | ||
| } | ||
| } |
There was a problem hiding this comment.
bug (high): Threshold recomputed every step with non-deterministic inputs can cause VA reservation overflow.
copy_threshold is recalculated on every call because it depends on DeviceMemoryUsage(&free_mem, &total_mem), which changes as memory is allocated/freed. The VA reservation is sized once (line 1641 guard: va_ranges->va_reservation == nullptr) using the threshold from the first step. If free memory decreases on a later step, the auto threshold shrinks, fewer slices are classified as "small", and the offset-assignment loop (lines 1693–1711) computes a current_offset that exceeds the reserved VA size. The subsequent MapTo/Remap call would then reference offsets beyond the reservation — out-of-bounds VA mapping, potential memory corruption, or driver errors.
Suggested fix: Compute the threshold once and cache it in VaRanges, or re-validate that current_offset <= total_va_size and return an error / fall back if it overflows.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The threshold is now computed once in ComputeCopyThreshold() and frozen in VaRanges::copy_threshold via the initialized guard, so the VA reservation is sized exactly once and cannot be overflowed by later free-memory fluctuations.
There was a problem hiding this comment.
Resolved in this revision: threshold is now computed once in ComputeCopyThreshold() and frozen via the initialized guard, preventing the cross-step drift that could cause VA reservation overflow.
| absl::flat_hash_map<BufferAllocation::Index, se::DeviceAddressBase> | ||
| small_shadow; |
There was a problem hiding this comment.
bug (high): Device memory leak of shadow buffers on GpuExecutable destruction.
The small_shadow map stores DeviceAddressBase values (raw pointer + size, no RAII). Each shadow buffer is allocated with executor->Allocate(sz) (line 1798 of gpu_executable.cc) but is never deallocated when the GpuExecutable is destroyed — the destructor does not iterate module_va_ranges_ to free shadow buffers. Over the lifetime of an application that creates and destroys executables, this leaks device memory.
Suggested fix: Add cleanup in the destructor that iterates module_va_ranges_ and deallocates every DeviceAddressBase in small_shadow. Alternatively, store DeviceAddressHandle (which has RAII semantics) instead of raw DeviceAddressBase.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. VaRanges::~VaRanges() now deallocates all shadow buffers via the stored executor pointer, preventing device memory leaks on destruction.
There was a problem hiding this comment.
Resolved in this revision: ~VaRanges() destructor now iterates small_shadow and calls executor->Deallocate() on each shadow buffer.
| const char* e = std::getenv("XLA_VMM_TMP_COPY_THRESHOLD"); | ||
| // Unset defaults to auto (mode 1) so the smart selector is the default. | ||
| uint64_t mode = 1; | ||
| if (e != nullptr) { | ||
| mode = std::strtoull(e, nullptr, 10); |
There was a problem hiding this comment.
nit (low): std::getenv called on every step, and std::strtoull silently returns 0 on invalid input.
-
std::getenvis not thread-safe if any thread callssetenv/putenvconcurrently (UB per C standard). More practically, re-parsing the env var each step is unnecessary work — consider caching in astaticlocal. -
If
XLA_VMM_TMP_COPY_THRESHOLDis set to a non-numeric string (e.g.,"auto","true"),std::strtoullreturns 0, which silently disables the optimization. Consider logging a warning or usingabsl::SimpleAtoiwhich provides explicit error reporting.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. std::getenv is now called once per VA range (inside the !initialized guard) rather than every step, and absl::SimpleAtoi replaces std::strtoull with a proper parse-failure warning.
There was a problem hiding this comment.
Resolved in this revision: std::getenv is now called once per VA range init (not every step), and absl::SimpleAtoi replaces std::strtoull for safe parsing.
| if (is_small_copy_slice(original_buffer.size())) { | ||
| const uint64_t sz = original_buffer.size(); | ||
| auto it = va_ranges->small_shadow.find(i); | ||
| 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; | ||
| } | ||
| small_copies.push_back({/*shadow=*/it->second, /*real=*/original_buffer, | ||
| /*size=*/sz}); | ||
| small_mapped[i] = se::DeviceAddressBase(it->second.opaque(), sz); | ||
| continue; | ||
| } |
There was a problem hiding this comment.
nit (low): Shadow buffer lookup uses size() from DeviceAddressBase, which may not reflect the originally requested allocation size.
When checking it->second.size() < sz (line 1794), if executor->Allocate(sz) returns a buffer with a different (rounded-up) size than sz, the stored size reflects the allocator's rounded size, not the original sz. On the next call, this could prevent reuse of a perfectly adequate shadow or cause unnecessary reallocations. Consider storing the requested size separately or using >= with the original requested size.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The code now documents why size() is safe here: Allocate() returns a buffer whose size() is the allocated capacity (>= requested), so rounding only makes reuse more likely, and reallocating only triggers when capacity is genuinely too small.
There was a problem hiding this comment.
Resolved in this revision: the code now documents why size() comparison is safe (allocator rounding only makes reuse more likely, never less).
| // EXPERIMENT: automatic copy-vs-remap selection for small command-buffer | ||
| // slices. The tiny scale/scalar/metric buffers churn their physical address | ||
| // every step; remapping them costs an hipMemUnmap/Map/SetAccess round-trip | ||
| // that serializes across peer devices on ROCm. Instead we keep such a slice | ||
| // in a stable shadow buffer (fixed address baked into the command buffer once) | ||
| // and refresh it with a cheap device-to-device copy each step. | ||
| // | ||
| // The size cutoff is chosen AUTOMATICALLY from the device's HBM bandwidth and | ||
| // capacity (no manual tuning): admit command-buffer slices smallest-first | ||
| // into the copy set while (a) no single slice's copy exceeds kPerBufferCopyUs, | ||
| // (b) the total per-step copy stays under kTotalCopyUs, and (c) the shadow | ||
| // footprint stays under kShadowHbmFraction of HBM. Large tensors blow these | ||
| // budgets and stay on the remap path -- and since they keep stable addresses | ||
| // across steps, remapping them is already free (skipped). The per-step copy | ||
| // budget bounds the worst-case overhead, so "auto" is safe on any model. | ||
| // | ||
| // Control via env XLA_VMM_TMP_COPY_THRESHOLD (default = auto): | ||
| // "0" -> disabled: every slice stays on the remap path (original path). | ||
| // "1" -> auto: derive the byte cutoff from this device's HBM bandwidth and | ||
| // capacity AND the memory currently free, so the shadow pool is | ||
| // always guaranteed to fit (no searching, never risks OOM). | ||
| // ">1" -> force exactly that byte cutoff (explicit override / sweeps). | ||
| // We reserve "1" for auto because a literal 1-byte cutoff would admit nothing. | ||
| uint64_t copy_threshold = 0; | ||
| { | ||
| const char* e = std::getenv("XLA_VMM_TMP_COPY_THRESHOLD"); | ||
| // Unset defaults to auto (mode 1) so the smart selector is the default. | ||
| uint64_t mode = 1; | ||
| if (e != nullptr) { | ||
| mode = std::strtoull(e, nullptr, 10); | ||
| } | ||
| if (mode == 0) { | ||
| copy_threshold = 0; // disabled | ||
| } else if (mode > 1) { | ||
| copy_threshold = mode; // explicit byte cutoff | ||
| } else { | ||
| // mode == 1: auto-derive from device properties + free memory. | ||
| // | ||
| // Greedy admit-smallest-first: pull command-buffer slices into the copy | ||
| // set in increasing size order while (a) no single slice's bidirectional | ||
| // copy exceeds kPerBufferCopyUs, and (b) the aggregate per-step copy stays | ||
| // under a budget that is the MIN of a copy-time cap and a memory cap. The | ||
| // memory cap is taken from the currently-free HBM (kFreeHbmFraction of | ||
| // free), which guarantees the shadow pool always fits -- "always enough | ||
| // space" -- rather than from a blind fraction of total capacity. Large | ||
| // tensors blow these budgets and stay on the remap path; since they keep | ||
| // stable addresses across steps their remap is already skipped (free). | ||
| constexpr double kPerBufferCopyUs = 8.0; | ||
| constexpr double kTotalCopyUs = 15.0; | ||
| constexpr double kShadowHbmFraction = 0.005; // of total (fallback) | ||
| constexpr double kFreeHbmFraction = 0.01; // of free (primary cap) | ||
| constexpr uint64_t kFallbackThreshold = 65536; | ||
| const se::DeviceDescription& dd = executor->GetDeviceDescription(); | ||
| const int64_t bw = dd.memory_bandwidth(); // bytes/sec | ||
| const int64_t hbm = dd.device_memory_size(); // bytes | ||
| if (bw <= 0) { | ||
| // Bandwidth unknown: fall back to a conservative cutoff that still | ||
| // captures scalar/scale buffers. | ||
| copy_threshold = kFallbackThreshold; | ||
| } else { | ||
| // acc accumulates 2*size per admitted slice: bidirectional copy moves | ||
| // 2 bytes per shadow byte, and the shadow is allocated per VA set, so | ||
| // 2*size also bounds the double-buffered footprint. | ||
| const double per_buffer_cap = | ||
| kPerBufferCopyUs * 1e-6 * static_cast<double>(bw) / 2.0; | ||
| double total_budget = kTotalCopyUs * 1e-6 * static_cast<double>(bw); | ||
| // Memory cap: prefer the live free-memory reading so the shadow pool is | ||
| // bounded by what the device can actually spare right now. | ||
| int64_t free_mem = 0; | ||
| int64_t total_mem = 0; | ||
| double mem_cap = 0.0; | ||
| if (executor->DeviceMemoryUsage(&free_mem, &total_mem) && | ||
| free_mem > 0) { | ||
| mem_cap = kFreeHbmFraction * static_cast<double>(free_mem); | ||
| } else if (hbm > 0) { | ||
| mem_cap = kShadowHbmFraction * static_cast<double>(hbm); | ||
| } | ||
| if (mem_cap > 0.0) { | ||
| total_budget = std::min(total_budget, mem_cap); | ||
| } | ||
| std::vector<uint64_t> sizes; | ||
| sizes.reserve(command_buffer_allocation_indexes_.size()); | ||
| for (BufferAllocation::Index i : command_buffer_allocation_indexes_) { | ||
| se::DeviceAddressBase b = buffer_allocations.GetDeviceAddress(i); | ||
| if (vmm_allocator->GetRawAllocation(executor->device_ordinal(), b) == | ||
| nullptr) { | ||
| continue; | ||
| } | ||
| sizes.push_back(b.size()); | ||
| } | ||
| std::sort(sizes.begin(), sizes.end()); | ||
| double acc = 0.0; | ||
| for (uint64_t s : sizes) { | ||
| if (static_cast<double>(s) > per_buffer_cap) break; | ||
| if (acc + 2.0 * static_cast<double>(s) > total_budget) break; | ||
| acc += 2.0 * static_cast<double>(s); | ||
| copy_threshold = s; | ||
| } | ||
| XLA_VLOG_DEVICE(2, executor->device_ordinal()) << absl::StreamFormat( | ||
| "VA remapping: auto copy threshold=%d B (bw=%d GB/s hbm=%d GB " | ||
| "free=%d GB budget=%.0f B)", | ||
| copy_threshold, bw / 1000000000, hbm / 1000000000, | ||
| free_mem / 1000000000, total_budget); | ||
| } | ||
| } | ||
| } | ||
| auto is_small_copy_slice = [copy_threshold](uint64_t size) { | ||
| return copy_threshold > 0 && size <= copy_threshold; | ||
| }; |
There was a problem hiding this comment.
suggestion (low): ~110 lines of inline threshold logic in an already-long function.
Consider extracting the threshold computation into a separate helper method (e.g., ComputeCopyThreshold) to improve readability and testability. The helper could also be unit-tested independently with mock device descriptions.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The threshold logic is now extracted into ComputeCopyThreshold() (member function) and MeasureCopyVsRemapBreakeven() (free function in anonymous namespace), significantly reducing inline complexity.
There was a problem hiding this comment.
Resolved in this revision: threshold logic extracted into ComputeCopyThreshold() and calibration into MeasureCopyVsRemapBreakeven(), improving readability and testability.
| // a slot in the contiguous VA reservation. | ||
| if (is_small_copy_slice(buf.size())) { | ||
| continue; | ||
| } | ||
| total_va_size += round_up_to_granularity(buf.size()); | ||
| } | ||
|
|
||
| // Reserve a single large VA range for all command buffer allocations. | ||
| ASSIGN_OR_RETURN(va_ranges->va_reservation, | ||
| vmm_allocator->CreateReservation(executor, total_va_size)); | ||
| // With the copy-to-shadow experiment a large threshold can make every | ||
| // slice "small", leaving nothing to reserve; skip reservation in that case | ||
| // (all slices then flow through the shadow copy path). | ||
| ASSIGN_OR_RETURN(va_ranges->unmap_event, executor->CreateEvent()); | ||
|
|
||
| XLA_VLOG_DEVICE(3, executor->device_ordinal()) << absl::StreamFormat( | ||
| "VA remapping: Reserved single VA range for module %s " | ||
| "VA: %p total_size: %d granularity: %d", | ||
| module_name_, va_ranges->va_reservation->address().opaque(), | ||
| total_va_size, granularity); | ||
| if (total_va_size > 0) { | ||
| ASSIGN_OR_RETURN( | ||
| va_ranges->va_reservation, | ||
| vmm_allocator->CreateReservation(executor, total_va_size)); | ||
| XLA_VLOG_DEVICE(3, executor->device_ordinal()) << absl::StreamFormat( | ||
| "VA remapping: Reserved single VA range for module %s " | ||
| "VA: %p total_size: %d granularity: %d", | ||
| module_name_, va_ranges->va_reservation->address().opaque(), | ||
| total_va_size, granularity); | ||
| } | ||
| } | ||
| // NOTE: the actual unmap of a previously established mapping is deferred to | ||
| // the decision point below, so it can be skipped entirely when the physical |
There was a problem hiding this comment.
bug (medium): Re-initialization on every step when all slices are "small."
The initialization guard is if (va_ranges->va_reservation == nullptr). When the copy threshold makes all slices "small" (e.g., a very large explicit threshold), va_reservation is never set (total_va_size == 0, skips the reservation). On subsequent steps, the guard fires again, re-scanning all sizes, re-creating unmap_event via ASSIGN_OR_RETURN. The old unique_ptr<Event> is destroyed by the assignment, but there is a window where the previously-recorded event (line 1984) may still be in-flight on the GPU when it is destroyed — potentially undefined behavior on the HIP runtime.
Additionally, the redundant per-step work undermines the "compute once" design intent.
Suggested fix: Add a separate bool initialized flag to VaRanges, or use unmap_event != nullptr as the guard instead of va_reservation.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The init guard is now !va_ranges->initialized (an explicit bool) instead of va_reservation == nullptr, so the one-time setup block runs exactly once even when every slice is small and the reservation stays null.
There was a problem hiding this comment.
Resolved in this revision: init guard is now !va_ranges->initialized (a bool set once after setup) instead of va_reservation == nullptr, preventing re-initialization when all slices are small.
Review SummaryThis PR adds a copy-vs-remap optimization for the VMM command-buffer VA-remapping path. Small buffer slices are kept in stable "shadow" device buffers refreshed via D2D copies, avoiding expensive Key findings (see inline comments for details):
The most critical issue is that the auto-computed |
github-actions
Bot
removed
the
claude-review
…s, robustify Resolves the Claude review on PR #926: - Compute the copy-vs-remap threshold ONCE and cache it in VaRanges (copy_threshold), frozen for the VA range's lifetime. The VA reservation is sized exactly once from it, so a later step recomputing a smaller threshold (free HBM fluctuates) can no longer push more slices onto the remap path than the reservation can hold. [high: VA reservation overflow] - Add ~VaRanges() destructor that deallocates the small_shadow buffers via the captured executor. DeviceAddressBase has no RAII, so these leaked on GpuExecutable destruction. [high: shadow buffer leak] - Gate one-time setup on an explicit `initialized` flag instead of `va_reservation == nullptr`. When every slice is small the reservation legitimately stays null, which previously re-ran setup each step and re-created unmap_event while a recorded copy could still be in flight. [medium: re-init / in-flight event destruction] - Parse XLA_VMM_TMP_COPY_THRESHOLD once with absl::SimpleAtoi; a non-numeric value now logs a warning and falls back to auto instead of strtoull silently returning 0 (disable). Removes the per-step getenv. [low] - Extract the threshold logic into GpuExecutable::ComputeCopyThreshold for readability/testability. [low] - Clarify shadow-buffer reuse comment: small_shadow stores Allocate() capacity (>= requested), so rounding only aids reuse. [low] NVIDIA unaffected (path still gated on DeviceAddressVmmAllocator).
Fixes 1, 3 and 4 share the same root cause (recomputing per step), so caching the threshold in |
| static absl::Mutex* const kCalibMu = new absl::Mutex(); | ||
| static auto* const kCalibCache = new absl::flat_hash_map<int, uint64_t>(); | ||
| std::optional<uint64_t> measured; | ||
| { | ||
| absl::MutexLock l(kCalibMu); | ||
| auto it = kCalibCache->find(executor->device_ordinal()); | ||
| if (it != kCalibCache->end()) measured = it->second; | ||
| } | ||
| if (!measured.has_value()) { | ||
| absl::StatusOr<uint64_t> m = MeasureCopyVsRemapBreakeven( | ||
| executor, vmm_allocator, stream, granularity); | ||
| if (m.ok()) { | ||
| measured = *m; | ||
| absl::MutexLock l(kCalibMu); | ||
| (*kCalibCache)[executor->device_ordinal()] = *m; |
There was a problem hiding this comment.
nit (low): TOCTOU in calibration cache — benign but avoidable with try_emplace.
Two VA range indices on the same executor hold different va_ranges->mutex locks, so both can enter ComputeCopyThreshold concurrently for the same device ordinal. Both will see !measured.has_value(), both will run the full MeasureCopyVsRemapBreakeven probe (~53 remap iterations + 20 copy iterations), and both will write the result.
This is functionally correct (the cache write is idempotent), but wastes ~1-2 ms of startup time per extra VA range index and temporarily doubles probe memory. A simple fix is to hold kCalibMu across the measurement (the probe runs once per device, not per step, so contention is negligible):
| static absl::Mutex* const kCalibMu = new absl::Mutex(); | |
| static auto* const kCalibCache = new absl::flat_hash_map<int, uint64_t>(); | |
| std::optional<uint64_t> measured; | |
| { | |
| absl::MutexLock l(kCalibMu); | |
| auto it = kCalibCache->find(executor->device_ordinal()); | |
| if (it != kCalibCache->end()) measured = it->second; | |
| } | |
| if (!measured.has_value()) { | |
| absl::StatusOr<uint64_t> m = MeasureCopyVsRemapBreakeven( | |
| executor, vmm_allocator, stream, granularity); | |
| if (m.ok()) { | |
| measured = *m; | |
| absl::MutexLock l(kCalibMu); | |
| (*kCalibCache)[executor->device_ordinal()] = *m; | |
| static absl::Mutex* const kCalibMu = new absl::Mutex(); | |
| static auto* const kCalibCache = new absl::flat_hash_map<int, uint64_t>(); | |
| std::optional<uint64_t> measured; | |
| { | |
| absl::MutexLock l(kCalibMu); | |
| auto [it, inserted] = kCalibCache->try_emplace(executor->device_ordinal(), 0); | |
| if (!inserted) { | |
| measured = it->second; | |
| } else { | |
| absl::StatusOr<uint64_t> m = MeasureCopyVsRemapBreakeven( | |
| executor, vmm_allocator, stream, granularity); | |
| if (m.ok()) { | |
| it->second = *m; | |
| measured = *m; | |
| } else { | |
| kCalibCache->erase(it); | |
| } | |
| } | |
| } |
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The calibration cache now uses try_emplace under kCalibMu held across the probe, preventing duplicate measurements for the same device ordinal.
| constexpr double kMinRemapUs = 0.05; | ||
| constexpr double kMaxRemapUs = 5000.0; | ||
| if (!(c_us > kMinRemapUs && c_us < kMaxRemapUs)) { | ||
| return Internal("calibration: implausible remap cost %.3f us", c_us); |
There was a problem hiding this comment.
bug (medium): Internal() uses absl::StrCat, not printf-style formatting — %.3f will appear as a literal string.
Internal("calibration: implausible remap cost %.3f us", c_us) expands to absl::InternalError(absl::StrCat("calibration: implausible remap cost %.3f us", c_us)), producing something like "calibration: implausible remap cost %.3f us42.7" instead of the intended "calibration: implausible remap cost 42.700 us".
| return Internal("calibration: implausible remap cost %.3f us", c_us); | |
| return absl::InternalError(absl::StrFormat( | |
| "calibration: implausible remap cost %.3f us", c_us)); |
| return Internal( | ||
| "Failed to allocate %d-byte shadow for small command-buffer " | ||
| "slice %d", | ||
| sz, i); |
There was a problem hiding this comment.
bug (medium): Same Internal() format-string issue — %d specifiers are not expanded.
Internal() delegates to absl::StrCat, not printf. The error message will read something like "Failed to allocate %d-byte shadow for small command-buffer slice %d64016" instead of substituting the values. Use absl::StrFormat or absl::StrCat with separate arguments:
| return Internal( | |
| "Failed to allocate %d-byte shadow for small command-buffer " | |
| "slice %d", | |
| sz, i); | |
| return absl::InternalError(absl::StrFormat( | |
| "Failed to allocate %d-byte shadow for small command-buffer " | |
| "slice %d", | |
| sz, i)); |
| double total_budget = | ||
| measured.has_value() ? (mem_cap > 0.0 ? mem_cap : 1e18) |
There was a problem hiding this comment.
suggestion (medium): Unbounded shadow pool when device memory reporting fails.
When measured.has_value() is true but both DeviceMemoryUsage() and device_memory_size() return 0 (broken memory reporting), mem_cap stays 0 and total_budget is set to 1e18 — effectively infinite. The shadow pool then has no memory cap: every slice below the measured break-even is admitted regardless of aggregate footprint, which could cause OOM.
Consider falling back to a conservative fixed budget (e.g., 1 MB) rather than infinity, or at least logging a warning when neither memory source is available:
measured.has_value() ? (mem_cap > 0.0 ? mem_cap : 1048576.0) // 1 MB safety cap
: kTotalCopyUs * 1e-6 * static_cast<double>(bw);There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. When mem_cap == 0 (broken memory reporting), the code now falls back to copy_time_budget instead of leaving the budget unbounded.
| std::sort(sizes.begin(), sizes.end()); | ||
| uint64_t copy_threshold = 0; | ||
| double acc = 0.0; | ||
| for (uint64_t s : sizes) { | ||
| if (static_cast<double>(s) > per_buffer_cap) break; | ||
| if (acc + 2.0 * static_cast<double>(s) > total_budget) break; | ||
| acc += 2.0 * static_cast<double>(s); | ||
| copy_threshold = s; | ||
| } |
There was a problem hiding this comment.
nit (low): Greedy threshold can overrun the budget for duplicate-sized slices.
The loop sets copy_threshold = s after budgeting for one slice of size s. At runtime, is_small_copy_slice uses size <= copy_threshold, admitting all slices of that size — but the budget may have only afforded K out of N slices at the boundary size. The overrun is bounded by (N-K) * 2 * s bytes, which is likely negligible for the small scalars targeted here.
If you wanted strict budget adherence, you could count all slices at each distinct size before admitting that size bucket. Not urgent given the target use case.
There was a problem hiding this comment.
Resolved ✓ — addressed in this revision. The potential overrun is now explicitly documented (lines 1728-1732) as bounded by (count_at_size - 1) * 2 * size, which is negligible for the tiny scalar/scale slices this targets.
Re-review Summary (commit 30c2f2f)All 6 findings from the previous review have been addressed in this revision — resolved replies posted on each thread. Key improvements: threshold frozen at init via 4 new findings posted inline:
|
github-actions
Bot
removed
the
claude-review
…ation TOCTOU Resolves the actionable findings from the Claude re-review of PR #926 (commit 30c2f2f): - Unbounded shadow pool when device memory reporting fails [medium]: when a measured break-even exists but DeviceMemoryUsage() and device_memory_size() both report 0, total_budget was set to 1e18 (effectively infinite), admitting every below-break-even slice regardless of aggregate footprint (OOM risk). Now falls back to the copy-time budget (~kTotalCopyUs of bandwidth), the same principled bound the no-measurement branch already uses. - Calibration-cache TOCTOU [low]: two VA-range indices on the same executor could both run the expensive MeasureCopyVsRemapBreakeven probe concurrently. Now uses try_emplace under a single kCalibMu lock to reserve the slot before probing, erasing it on failure so a later call can retry. - Greedy duplicate-size budget overrun [low]: documented the bounded, negligible overrun when multiple slices share the boundary size (reviewer: not urgent). Not changed: the two "Internal() format specifiers won't expand" findings are false positives. xla::Internal() is the absl::StrFormat overload (util.h); the StrCat variant is the separate InternalStrCat type. %.3f/%d expand correctly, as the pre-existing Internal("...%d...", i) calls in this file show. Correctness/robustness only; steady-state timing and benchmark numbers unchanged. NVIDIA unaffected (path still gated on DeviceAddressVmmAllocator).
Re-review Summary (commit 8a464d0)All 11 findings from the prior review have been addressed in this revision:
No new issues found in this revision. The code quality has improved significantly. 🤖 Generated with Claude Code |
github-actions
Bot
removed
the
claude-review
| // median of real unmap/map/setaccess round-trips on a probe allocation, and bw | ||
| // as a real D2D copy -- so the cutoff is self-calibrating instead of a guess. | ||
| // Returns the break-even size in bytes, or an error (caller falls back). | ||
| absl::StatusOr<uint64_t> MeasureCopyVsRemapBreakeven( |
There was a problem hiding this comment.
This can be soo nondeterminitic. This should be some good enough treashold value with an escape hatch for user to sets its own/disable it.
There was a problem hiding this comment.
Thanks Dragan, let me try to find good threadhold.
2 participants
Dependency (read first)
This PR is on top of the ROCm VMM allocator (branch
phambinh/rocm-vmm-allocator, proposed upstream as openxla/xla openxla#40236). The base of this PR is intentionally that branch, notmain, so the diff here is only this optimization (gpu_executable.{cc,h}, +220/-8). Targetingmainwould drag in the whole VMM series, since it isn't inmainyet. It should merge after the VMM allocator lands.On NVIDIA the cast is
null,use_command_buffer_va_remappingisfalse, the function is never entered, andXLA_VMM_TMP_COPY_THRESHOLDhas no effect.Optimization
In the VMM command-buffer VA-remapping path, every slice whose physical allocation churns each step is re-mapped via
hipMemUnmap/hipMemMap/SetAccess. For the tiny scalar/scale/metric buffers that churn every step, that driver round-trip dominates and serializes across peer devices. Large tensors keep stable addresses across steps, so their remap is already skipped (free).This keeps the small churning slices in a stable shadow buffer (fixed address baked into the command buffer once) and refreshes them with a cheap device-to-device copy each step (copy-in before exec, copy-back after, so output scalars like loss/metrics propagate). Large slices stay on the remap path.
The cutoff is chosen automatically via
XLA_VMM_TMP_COPY_THRESHOLD:01DeviceMemoryUsage)>1Default (unset) is auto.
Auto algorithm (
=1)The threshold is derived once, per VA-range setup, with no searching. The goal is to copy only the small churning slices, with the extra per-step copy bounded to a negligible amount and the shadow pool guaranteed to fit in free memory.
bw(bytes/s) and capacity fromDeviceDescription, plus currently-free HBM viaexecutor->DeviceMemoryUsage(&free, &total). Ifbwis unknown, fall back to a fixed 64 KB cutoff.2 * sizebytes per admitted slice):size <= 8us * bw / 2. This excludes large tensors outright.min(~15us * bw, 1% of free HBM). The time term bounds the worst-case per-step overhead; the memory term guarantees the shadow pool always fits (never OOM), and it scales down automatically when the device is nearly full.2 * size; stop at the first slice that either exceeds the per-slice cap or would push the running total past the aggregate budget. The largest admitted size becomes the threshold.<= thresholdtakes the shadow-copy path; everything larger stays on the remap path.Why this is both safe and near-optimal: the smallest slices are exactly the churning scalars/scales whose per-step remap is the bottleneck, so admitting them captures essentially all the benefit. Large tensors are excluded twice over -- they blow the caps, and their remap was already free (stable address, skipped). The per-step copy overhead is bounded (~15 us) on any model, and the shadow footprint can't exceed the free-memory cap, so there is no manual tuning and no OOM risk. Empirically
autolands on the same plateau as a hand-tuned fixed cutoff (see below).Benchmarks (MI300X x8, 8 layers / 30 steps, median of steps >=3; 0 unmap errors everywhere)
Two speedup columns with different meanings:
speedup (vs disabled)—disabled (=0)→auto (=1), both on VMM. This isolates this PR's change (copy-vs-remap only) and is always a positive, ~6–13 ms win.auto vs base—auto (=1)vsbase (no VMM). This compares the whole VMM stack (VMM allocator +NEVER_UPDATE+ this optimization) against plain command buffers with no VMM, i.e. it folds in the cost/benefit of the underlying VMM series, not just this PR. A negative value (e.g.llama3_8b) means the VMM stack itself carries a small overhead on that model that this optimization narrows but doesn't fully erase.All
base (no VMM)numbers are medians of 3 runs (steps ≥3), same as the other columns. On the small models (llama2_7b,llama3.1-8b,gemma2-9b) the VMM legs are dramatically faster than no-VMM because the per-step command-buffer remap is far cheaper than re-recording the buffer every step; onllama3_8b/mixtral_8x1bbase and VMM are within noise.auto + collective leg (
coll_vmm): before (disabled) vs after (auto)=0)=1)vmm leg (no collectives): before (disabled) vs after (auto)
=0)=1)The saving is a roughly fixed per-step cost (~6-13 ms; the work removed is the remap of the small churning slices, whose count is set by command-buffer structure, not model size). It shows up as a larger percentage on short-step models and a smaller one on heavy models, but is always a net win and never regresses.
automatches a hand-tuned fixed 4 KB cutoff (llama2_7b: auto 121 ms vs fixed-4K 120 ms).Notes
[EXPERIMENT]- behind an env knob, default auto, ROCm-gated.