ggml-hrx: the host-weight arena is charged to the pool, like the transient arena (engine #1711)

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);

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<std::mutex> 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();

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,13 +188,17 @@ 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<std::pair<size_t, size_t>> arena_free_;
 
     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);

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;