From 203aac2cfb02b91a4993c6b1644486daf83a56ff Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Mon, 9 Nov 2026 18:36:39 -0400 Subject: [PATCH] ggml-hrx: the host-weight arena is charged to the pool, like the transient arena (engine #1711) Second slice of engine#1711. GGML_HRX_HOST_WEIGHT_LIMIT_MB limits how much residency the cache keeps, but nothing charged it to the pool ledger, so a large host-weight arena was invisible to pool_reserve and the guardrail over-committed by however much the cache held. The growth uses the same pure decision as the transient arena (plan_arena_charge, retires_old = false: the superseded arena is released, so its charge goes with it). The budget reaches it through CommandProgramExecutionContext, which already carries it since #268; the two acquire() call sites in the executor pass it down. Three arms added to tests/test-hrx-transient-arena-charge.cpp for the shape this path actually has: first growth charges the whole arena, a regrow to the same budget charges nothing (the common case - one arena serves many weights), and a growth past the budget swaps charge for charge. --- .../runtime/command-program-executor.cpp | 5 +-- ggml/src/ggml-hrx/runtime/host-memory.cpp | 35 +++++++++++++++++-- ggml/src/ggml-hrx/runtime/host-memory.h | 7 +++- tests/test-hrx-transient-arena-charge.cpp | 14 ++++++++ 4 files changed, 56 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp index 010a1684754d..e44ad94f4c97 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -711,7 +711,7 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec continue; } HostWeightAcquireResult resident = - context.host_weights->acquire(context.device, context.stream, *context.host_transfers, source); + context.host_weights->acquire(context.device, context.stream, *context.host_transfers, source, context.pool_budget); if (!resident.valid()) { // engine #1618: this line and the HSA error behind it carried no pool figure, // so an 89 GB model materialized against a 66.6 GiB budget read as a @@ -1019,7 +1019,8 @@ static Status rebind_prepared_host_weights(const CommandProgramExecutionContext // Evicted since this program was prepared (the cache is bounded, see HostWeightCache::limit): put it back and // point the prepared bindings at the new buffer. The old one is gone, so nothing else may still reference it. HostWeightAcquireResult materialized = - context.host_weights->acquire(context.device, context.stream, *context.host_transfers, resident.source); + context.host_weights->acquire(context.device, context.stream, *context.host_transfers, resident.source, + context.pool_budget); if (!materialized.valid()) { status.append(materialized.status); status.log("re-materialize resident host weight value %d failed", resident.value.value); diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index 58024b3db3bb..55422397fbbe 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -4,6 +4,8 @@ #include "ggml-quants.h" #include "hrx-interop-utils.h" #include "hrx_runtime.h" +#include "hrx-1bit-safety.h" // engine#1711: pool_reserve/pool_release +#include "transient-arena-charge.h" // engine#1711: the same charge decision as the transient arena #include "runtime/host-memory-available.h" #include "runtime/host-memory-ternary.h" @@ -1665,6 +1667,7 @@ bool HostWeightCache::arena_take_locked(size_t aligned, hrx_buffer_t & buffer, s bool HostWeightCache::arena_reserve(hrx_device_t device, hrx_stream_t stream, size_t size, + size_t pool_budget, hrx_buffer_t & buffer, size_t & offset, size_t & slot) { @@ -1738,9 +1741,28 @@ bool HostWeightCache::arena_reserve(hrx_device_t device, return false; } } + // engine#1711 (second slice): the host-weight arena is device-visible memory like any other, so its growth is + // charged to the pool before it happens -- the same decision the transient arena uses. Without this, a large + // residency is invisible to `pool_reserve` and the guardrail over-commits by however much the arena holds, which + // is what GGML_HRX_HOST_WEIGHT_LIMIT_MB *limits* but never accounts. + const ArenaChargePlan charge = plan_arena_charge(target, previous_capacity, false); + bool charge_taken = false; + if (charge.reserve_bytes > 0 && pool_budget > 0) { + if (!ggml::hrx::onebit::pool_reserve(pool_budget, charge.reserve_bytes, "the host-weight arena")) { + GGML_LOG_ERROR("%s: growing the host weight arena to %.1f MiB needs %zu more bytes and the pool has no room\n", + __func__, target / (1024.0 * 1024.0), charge.reserve_bytes); + return false; + } + charge_taken = true; + charged_bytes_ += charge.reserve_bytes; + } hrx_buffer_t grown = nullptr; const Status status = allocate_device_buffer(device, target, grown); if (!status.success()) { + if (charge_taken) { + ggml::hrx::onebit::pool_release(charge.reserve_bytes); + charged_bytes_ -= charged_bytes_ >= charge.reserve_bytes ? charge.reserve_bytes : charged_bytes_; + } GGML_LOG_ERROR("%s: growing the host weight arena to %.1f MiB failed (limit %.1f MiB, budget %.1f MiB, " "weight %.1f MiB)\n", __func__, target / (1024.0 * 1024.0), limit_ / (1024.0 * 1024.0), budget / (1024.0 * 1024.0), @@ -1751,6 +1773,10 @@ bool HostWeightCache::arena_reserve(hrx_device_t device, std::lock_guard lock(mutex_); if (arena_ != nullptr) { hrx_buffer_release(arena_); + if (charge.release_bytes > 0) { + ggml::hrx::onebit::pool_release(charge.release_bytes); + charged_bytes_ -= charged_bytes_ >= charge.release_bytes ? charge.release_bytes : charged_bytes_; + } } // Every slot in the old arena is gone, so no entry may keep pointing at it. entries_.clear(); @@ -1848,7 +1874,8 @@ void HostWeightCache::evict_for(hrx_stream_t stream, size_t needed) { HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device, hrx_stream_t stream, HostTransferManager & transfers, - const HostWeightSource & source) { + const HostWeightSource & source, + size_t pool_budget) { HostWeightAcquireResult result; if (stream == nullptr) { result.status.log("host weight residency requires an HRX stream"); @@ -1949,7 +1976,7 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device } } const bool reserved = - arena_reserve(device, stream, materialized_length, entry->buffer, entry->offset, entry->slot); + arena_reserve(device, stream, materialized_length, pool_budget, entry->buffer, entry->offset, entry->slot); if (!reserved) { size_t parked = 0; size_t resident = 0; @@ -2051,6 +2078,10 @@ void HostWeightCache::clear() { hrx_buffer_release(arena_); arena_ = nullptr; } + if (charged_bytes_ > 0) { + ggml::hrx::onebit::pool_release(charged_bytes_); // engine#1711 + charged_bytes_ = 0; + } arena_capacity_ = 0; arena_free_bytes_ = 0; arena_free_.clear(); diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h index 916045bdd1c9..a82d8fcd1ff8 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -122,7 +122,8 @@ class HostWeightCache { HostWeightAcquireResult acquire(hrx_device_t device, hrx_stream_t stream, HostTransferManager & transfers, - const HostWeightSource & source); + const HostWeightSource & source, + size_t pool_budget = 0); HostWeightCacheStats stats() const; void clear(); @@ -187,6 +188,9 @@ class HostWeightCache { // buffer. The arena grows by doubling up to the residency limit, and a growth drops every entry (the stream // is synchronized first), so the peak stays at the working set rather than at the knob. hrx_buffer_t arena_ = nullptr; + // Bytes this cache holds against the pool ledger (engine#1711). Tracked separately from arena_capacity_ so clear() + // returns exactly what was charged, including the case where a grow released a superseded arena. + size_t charged_bytes_ = 0; size_t arena_capacity_ = 0; size_t arena_free_bytes_ = 0; std::vector> arena_free_; @@ -194,6 +198,7 @@ class HostWeightCache { bool arena_reserve(hrx_device_t device, hrx_stream_t stream, size_t size, + size_t pool_budget, hrx_buffer_t & buffer, size_t & offset, size_t & slot); diff --git a/tests/test-hrx-transient-arena-charge.cpp b/tests/test-hrx-transient-arena-charge.cpp index 6e5d508ba5da..b0c7d6d5e8de 100644 --- a/tests/test-hrx-transient-arena-charge.cpp +++ b/tests/test-hrx-transient-arena-charge.cpp @@ -63,6 +63,20 @@ int main() { got = plan_arena_charge(1u << 20, 8u << 20, true); check(got.reserve_bytes == 0 && got.release_bytes == 0, "a smaller request charges nothing", got, {0, 0}); + // 6. The host-weight arena reuses this same decision (`retires_old = false`): it grows straight to its budget, so + // the first growth charges the whole arena and a later regrow to the same budget charges nothing. Getting this + // wrong double-charges a cache that keeps the same arena across many weights, which is the common case. + got = plan_arena_charge(4096u << 20, 0, false); + check(got.reserve_bytes == (4096u << 20) && got.release_bytes == 0, "host-weight arena: first growth charges the whole arena", got, + {4096u << 20, 0}); + got = plan_arena_charge(4096u << 20, 4096u << 20, false); + check(got.reserve_bytes == 0 && got.release_bytes == 0, "host-weight arena: a regrow to the same budget charges nothing", got, + {0, 0}); + // and a grow past the budget releases the superseded arena's charge in the same step (the old buffer is freed there) + got = plan_arena_charge(8192u << 20, 4096u << 20, false); + check(got.reserve_bytes == (4096u << 20) && got.release_bytes == (4096u << 20), "host-weight arena: growing past the budget swaps charge for charge", got, + {4096u << 20, 4096u << 20}); + if (failures != 0) { std::printf("%d arm(s) failed\n", failures); return 1;