hrx: the pool budget takes the host memory as a second term, so admission stops over-promising (engine #1502)

ggml/src/ggml-hrx/ggml-hrx.cpp

@@ -2,6 +2,7 @@
 #include "dispatch_registration/common/moe-placement-guard.h"
 #include "ggml-hrx.h"
 #include "hrx-1bit-safety.h"
+#include "runtime/host-memory-available.h"
 #include "ggml-hrx-dmabuf.h"
 
 #include "backend-buffer-binding.h"
@@ -1608,7 +1609,13 @@ static std::unique_ptr<ggml_backend_hrx_reg_context> create_registry_context() {
         size_t device_free            = device_ctx->memory_total;
         size_t device_total           = device_ctx->memory_total;
         device_free_total(device_ctx.get(), device_free, device_total);
-        device_ctx->memory_budget = ggml::hrx::onebit::pool_budget(device_ctx->memory_total, device_free);
+        // A failed /proc read is "unknown", not "zero bytes left": kHostUnknown disables the term, while a measured zero must
+        // refuse - a host with nothing left is the worst moment to admit a launch (engine #1502).
+        const size_t host_raw       = ggml::hrx::host_memory_available_bytes();
+        const size_t host_available = host_raw == 0 ? ggml::hrx::onebit::kHostUnknown : host_raw;
+        const ggml::hrx::onebit::pool_budget_terms_t terms =
+            ggml::hrx::onebit::pool_budget_terms(device_ctx->memory_total, device_free, host_available);
+        device_ctx->memory_budget = terms.budget;
         if (device_ctx->memory_budget == 0 && device_ctx->memory_total != 0) {
             // Nothing fits under the ceiling: refuse the device here, loudly, instead of planning a launch that can
             // reach the fault line (#680). Report only, do not try to be clever about it.
@@ -1625,13 +1632,16 @@ static std::unique_ptr<ggml_backend_hrx_reg_context> create_registry_context() {
         if (device_ctx->memory_budget != 0) {
             const size_t reserve = ggml::hrx::onebit::pool_reserve_bytes();
             GGML_LOG_WARN("ggml_hrx: %s guardrail: GTT total %.1f GiB, in use %.1f GiB, ceiling %.1f GiB "
-                          "(reserve %.1f GiB), pool budget %.1f GiB - a launch that needs more is refused here "
-                          "rather than faulting the device (engine #680)\n",
+                          "(reserve %.1f GiB), pool budget %.1f GiB (GTT term %.1f GiB, host MemAvailable %.1f GiB) - a launch that needs more is refused here "
+                          "rather than faulting the device (engine #680; bound by %s)\n",
                           device_ctx->name.c_str(), device_total / (1024.0 * 1024.0 * 1024.0),
                           (device_total > device_free ? device_total - device_free : 0) / (1024.0 * 1024.0 * 1024.0),
                           (device_total > reserve ? device_total - reserve : 0) / (1024.0 * 1024.0 * 1024.0),
                           reserve / (1024.0 * 1024.0 * 1024.0),
-                          device_ctx->memory_budget / (1024.0 * 1024.0 * 1024.0));
+                          terms.budget / (1024.0 * 1024.0 * 1024.0),
+                          terms.gtt_term / (1024.0 * 1024.0 * 1024.0),
+                          host_raw == 0 ? -1.0 : terms.host_available / (1024.0 * 1024.0 * 1024.0),
+                          terms.host_bound ? "the host" : (host_raw == 0 ? "the GTT counter (host unreadable)" : "the GTT counter"));
         }
         device_ctx->description       = *name + " (" + *architecture + ")";
         device_ctx->architecture      = *architecture;

ggml/src/ggml-hrx/hrx-1bit-safety.cpp

@@ -88,21 +88,41 @@ bool pool_budget_is_fixed_limit() {
     return env != nullptr && std::atoll(env) > 0;
 }
 
-size_t pool_budget(size_t memory_total, size_t memory_free) {
+pool_budget_terms_t pool_budget_terms(size_t memory_total, size_t memory_free, size_t memory_available) {
+    pool_budget_terms_t t;
     const char * env = std::getenv("GGML_HRX_POOL_LIMIT_MB");
     if (env != nullptr) {
         const long long mb = std::atoll(env);
         if (mb > 0) {
-            return static_cast<size_t>(mb) * kMiB;
+            // An operator-set limit is the budget outright, not a min() of terms; keeping it distinguishable is what
+            // pool_budget_is_fixed_limit() reports so a caller can word the line correctly (llama.cpp #255).
+            t.host_available = memory_available;
+            t.gtt_term       = static_cast<size_t>(mb) * kMiB;
+            t.budget         = t.gtt_term;
+            return t;
         }
     }
     if (memory_total == 0) {
-        return 0;  // unknown device size: keep the unfettered behaviour
+        return t;  // unknown device size: keep the unfettered behaviour
     }
     const size_t reserve = pool_reserve_bytes();
     const size_t ceiling = memory_total > reserve ? memory_total - reserve : 0;
     const size_t in_use  = memory_total > memory_free ? memory_total - memory_free : 0;
-    return ceiling > in_use ? ceiling - in_use : 0;
+    t.gtt_term       = ceiling > in_use ? ceiling - in_use : 0;
+    t.host_available = memory_available;
+    // engine #1502: on an APU the GTT counter is not the allocator's ceiling, so the budget is the smaller of the two terms -
+    // the same model the engine's own admission check has always used (app/mem_admit.h).
+    if (memory_available != kHostUnknown && memory_available < t.gtt_term) {
+        t.budget     = memory_available;
+        t.host_bound = true;
+    } else {
+        t.budget = t.gtt_term;
+    }
+    return t;
+}
+
+size_t pool_budget(size_t memory_total, size_t memory_free, size_t memory_available) {
+    return pool_budget_terms(memory_total, memory_free, memory_available).budget;
 }
 
 // engine #1417: the budget above is a snapshot taken at device open, and the ledger is per-process, so another process
@@ -199,14 +219,15 @@ pool_budget_state_t pool_budget_note_refusal(pool_budget_state_t state, size_t s
     return state;
 }
 
-pool_budget_state_t pool_budget_observe(size_t memory_total, size_t free_bytes, size_t allocated, pool_budget_state_t state) {
+pool_budget_state_t pool_budget_observe(size_t memory_total, size_t free_bytes, size_t allocated, pool_budget_state_t state,
+                                        size_t memory_available) {
     if (state.cap == 0) {
         return state;  // no cap in force: the registration-time budget stands
     }
     if (free_bytes > state.cap_free) {
         // The pool has grown past the figure the refusal saw, so the cap's evidence no longer holds: a transient refusal must not
         // become a permanent ceiling (engine #1502 review). Recompute from the current figures and drop the cap.
-        const size_t uncapped = pool_budget(memory_total, free_bytes + allocated);
+        const size_t uncapped = pool_budget(memory_total, free_bytes + allocated, memory_available);
         GGML_LOG_INFO("%s: the pool grew from %.1f MiB to %.1f MiB free since the last refusal; lifting the %.1f MiB cap to "
                       "%.1f MiB\n",
                       __func__, state.cap_free / (double) kMiB, free_bytes / (double) kMiB, state.cap / (double) kMiB,
@@ -218,7 +239,7 @@ pool_budget_state_t pool_budget_observe(size_t memory_total, size_t free_bytes,
     }
     // The pool is no larger than it was at the refusal: keep the cap AND its evidence, and bound only this call's budget by it.
     // Clearing it here would freeze the low reading for the rest of the process (engine #1502 review).
-    const size_t capped = pool_budget(memory_total, free_bytes + allocated);
+    const size_t capped = pool_budget(memory_total, free_bytes + allocated, memory_available);
     state.budget = state.cap < capped ? state.cap : capped;
     return state;
 }

ggml/src/ggml-hrx/hrx-1bit-safety.h

@@ -28,7 +28,24 @@ namespace ggml::hrx::onebit {
 // budget is the GTT ceiling (memory_total minus GGML_HRX_GTT_RESERVE_MB, default 16384) less what other processes
 // already hold (memory_total - memory_free). A GTT plan past this budget is refused at allocation instead of
 // faulting the device: #680 is a gfxhub page fault plus a MODE2 reset at GTT 106.5 of 120 GiB.
-size_t pool_budget(size_t memory_total, size_t memory_free);
+// The two terms the budget is the minimum of, and the budget itself. Returned together so a caller can PRINT what the
+// decision used instead of recomputing it, and so the log can name the term that bound (engine #1502).
+struct pool_budget_terms_t {
+    size_t gtt_term       = 0;   // (total - reserve) - in_use, from the driver counters
+    size_t host_available = 0;   // MemAvailable as passed in, or kHostUnknown
+    size_t budget         = 0;   // min(gtt_term, host_available), or gtt_term when the host figure is unknown
+    bool   host_bound     = false;
+};
+
+// |memory_available| is the host's MemAvailable. On an APU the GTT ceiling and DRAM are the same pool, so what the allocator
+// can serve is bounded by both: this is the term that differs between two loads with identical GTT counters and opposite
+// outcomes (engine #1502). kHostUnknown means "could not be read" and only then is the term ignored - a measured zero means the
+// host has nothing left, the worst moment to admit a launch, so it must refuse. Defaulted so a diagnostic-only caller keeps
+// meaning "GTT arithmetic only" (llama.cpp #255) and every existing call site still compiles.
+inline constexpr size_t kHostUnknown = static_cast<size_t>(-1);
+
+pool_budget_terms_t pool_budget_terms(size_t memory_total, size_t memory_free, size_t memory_available = kHostUnknown);
+size_t pool_budget(size_t memory_total, size_t memory_free, size_t memory_available = kHostUnknown);
 // true when GGML_HRX_POOL_LIMIT_MB is in force: pool_budget() then returns that fixed limit rather than a
 // remainder, so a caller that prints the figure has to word it differently.
 bool   pool_budget_is_fixed_limit();
@@ -67,7 +84,8 @@ pool_budget_state_t pool_budget_note_refusal(pool_budget_state_t state, size_t s
 //    pool_remaining() subtracts the ledger again further down the admission path;
 //  - a cap is cleared only when it is LIFTED (the pool grew past the figure the refusal saw). When the pool is smaller the cap and
 //    its evidence are kept, and only this call's budget is bounded by it - otherwise one low reading becomes a permanent ceiling.
-pool_budget_state_t pool_budget_observe(size_t memory_total, size_t free_bytes, size_t allocated, pool_budget_state_t state);
+pool_budget_state_t pool_budget_observe(size_t memory_total, size_t free_bytes, size_t allocated, pool_budget_state_t state,
+                                        size_t memory_available = kHostUnknown);
 
 // The reserve the ceiling keeps for the driver and for kernel/transient spikes, in bytes.
 size_t pool_reserve_bytes();

tests/test-hrx-pool-refusal-budget.cpp

@@ -43,6 +43,9 @@ static void expect_eq(const char * what, size_t got, size_t want) {
 
 int main() {
     using ggml::hrx::onebit::pool_budget_after_refusal;
+    using ggml::hrx::onebit::pool_budget;
+    using ggml::hrx::onebit::pool_budget_terms;
+    using ggml::hrx::onebit::kHostUnknown;
 
     // The issue's own case: the guardrail admitted 15.35 GiB out of 29.2 GiB and the pool refused it.
     const size_t budget   = 29 * GiB + 205 * MiB;   // 29.2 GiB
@@ -145,6 +148,27 @@ int main() {
                 !pool_live_fits(1 * MiB, gtt, 0, SIZE_MAX));
     expect_true("a zero-size request with room passes", pool_live_fits(0, gtt, gtt - 1 * GiB, margin_256));
 
+    // ---- the host term (this lane's half of #1502) ---------------------------------------------------------------
+    // `total` already declared above by the refusal section
+    const size_t issue_free = total - 74 * GiB - 819 * MiB;                      // in use 74.8 GiB, as the issue measured
+    const size_t issue_budget = 29 * GiB + 205 * MiB;                            // 29.2 GiB, as the guardrail printed
+    expect_eq("the GTT arithmetic reproduces the issue's 29.2 GiB budget",
+              pool_budget(total, issue_free, kHostUnknown), issue_budget);
+    const size_t host_tight = 15 * GiB;
+    expect_eq("a tight host replaces the GTT term", pool_budget(total, issue_free, host_tight), host_tight);
+    expect_true("so the refused 15.35 GiB is no longer admissible",
+                pool_budget(total, issue_free, host_tight) < 15 * GiB + 359 * MiB);
+    expect_true("while the GTT-only arithmetic still called it admissible",
+                pool_budget(total, issue_free, kHostUnknown) > 15 * GiB + 359 * MiB);
+    expect_eq("a roomy host leaves the GTT term alone (no behaviour change)",
+              pool_budget(total, issue_free, 80 * GiB), issue_budget);
+    expect_eq("a host figure larger than the GTT term never inflates the budget",
+              pool_budget(total, issue_free, 900 * GiB), issue_budget);
+    expect_eq("host unknown keeps the GTT arithmetic", pool_budget(total, issue_free, kHostUnknown), issue_budget);
+    expect_eq("a measured zero host refuses", pool_budget(total, issue_free, 0), 0);
+    expect_true("and the terms say the host bound it", pool_budget_terms(total, issue_free, 0).host_bound);
+    expect_eq("unknown device stays unfettered", pool_budget(0, 0, host_tight), 0);
+
     printf(failures == 0 ? "test-hrx-pool-refusal-budget: all requirements passed\n"
                          : "test-hrx-pool-refusal-budget: %d FAILED\n", failures);
     return failures == 0 ? 0 : 1;