From d026f2d80045701f44e47adc60985eef95304bd2 Mon Sep 17 00:00:00 2001
From: bong-water-water-bong <bong-water-water-bong@1bit.gg>
Date: Tue, 10 Nov 2026 02:32:39 -0400
Subject: [PATCH] hrx: charge the per-weight host-weight buffer, which is the
 DEFAULT mode (engine#1823)

#1711 charged three allocators -- the transient arena, the host-weight arena and host staging -- and the host-weight
charge went on the **arena** slot path. The `else` branch beside it is the **unbounded** mode, `GGML_HRX_HOST_WEIGHT_LIMIT_MB=0`,
which is the default: each weight gets its own device buffer, and nothing charged it. So the ledger under-reported by
however much that mode holds, in the configuration nearly every run uses, and any budget arithmetic built on
`pool_allocated()` was optimistic by exactly that amount.

The charge is taken through `pool_reserve`, the same function the arena uses, and returned through `pool_release` with the
amount stored on the entry:

* **before** the allocation, so a refusal leaves nothing to refund;
* refunded if the allocation then fails;
* released in `~Entry` from the stored `charged` field, not a recomputed size -- a refund that recomputes is how a ledger
  drifts, and returning it twice is what the staging-buffer test exists to prevent.

`charged == 0` when `pool_budget` is 0, which is the same "no budget to charge against" case the arena path already
treats as uncharged.

Verified: `g++-16 -std=c++26 -fsyntax-only` on the translation unit, against the pinned hrx headers
(`third_party/hrx-system/libhrx/include`, the engine's pin `7f4236924a`) -- clean. **No device arm yet**: `Entry` is
internal to this file and the observable (`pool_allocated()`) only moves when a weight is actually materialised, so the
test needs either an exposed seam or a device run; I would rather say that plainly than add an arm that proves the
struct's destructor runs. A `gpu-hold` run with `GGML_HRX_HOST_WEIGHT_LIMIT_MB=0` reading the ledger before and after
materialising is the arm I would add next, and it is cheap on the repaired `build-hrx`.

Fixes engine#1823. Refs engine#1711, #268, #269, #270.
---
 ggml/src/ggml-hrx/runtime/host-memory.cpp | 25 ++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp
index 55422397fbbe..934aa30b794d 100644
--- a/ggml/src/ggml-hrx/runtime/host-memory.cpp
+++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp
@@ -1384,8 +1384,14 @@ struct HostWeightLease::Entry {
         if (owns_buffer && buffer != nullptr) {
             hrx_buffer_release(buffer);
         }
+        // engine#1823: the per-weight buffer above was charged on allocate; return exactly what was charged.
+        if (charged != 0) {
+            ggml::hrx::onebit::pool_release(charged);
+            charged = 0;
+        }
     }
-    bool owns_buffer = false;
+    bool   owns_buffer = false;
+    size_t charged     = 0;   // engine#1823: bytes this entry holds against the pool, 0 when unbounded charging is off
 
     hrx_buffer_t buffer = nullptr;
     size_t       offset = 0;
@@ -2007,8 +2013,25 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t             device
             return result;
         }
     } else {
+        // engine#1823: this branch is the DEFAULT (GGML_HRX_HOST_WEIGHT_LIMIT_MB=0 selects per-weight residency), and a
+        // per-weight device buffer is device memory exactly like the arena's slots, so it is charged the same way. Without
+        // this the ledger ignores the mode nearly every run uses, and every budget calculation built on it is optimistic
+        // by however much this mode holds. Charged BEFORE the allocation so a refusal leaves nothing to refund, and
+        // returned in ~Entry with the SAME stored amount: a refund that recomputes its size is how a ledger drifts, and
+        // returning it twice is what the staging-buffer test exists to prevent.
+        if (pool_budget > 0) {
+            if (!ggml::hrx::onebit::pool_reserve(pool_budget, materialized_length, "a per-weight host-weight buffer")) {
+                result.status.log("a per-weight host-weight buffer of %zu bytes does not fit the pool", materialized_length);
+                return result;
+            }
+            entry->charged = materialized_length;
+        }
         result.status = allocate_device_buffer(device, materialized_length, entry->buffer);
         if (!result.status.success()) {
+            if (entry->charged != 0) {
+                ggml::hrx::onebit::pool_release(entry->charged);
+                entry->charged = 0;
+            }
             return result;
         }
         entry->owns_buffer = true;