hrx: charge the transient arena against the pool budget, so a launch whose scratch does not fit is refused (refs engine#1711)

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

@@ -1583,6 +1583,10 @@ static std::unique_ptr<ggml_backend_hrx_reg_context> create_registry_context() {
         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);
+        // engine#1711: the per-launch scratch is charged against the same budget as the tensor buffers. Without this the
+        // guardrail could only refuse a launch whose *tensors* did not fit, while the arena a large case actually grows
+        // into allocated uncharged.
+        device_ctx->transient_arena.set_budget(&device_ctx->memory_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.

ggml/src/ggml-hrx/runtime/transient-arena-budget.h

@@ -0,0 +1,80 @@
+// Copyright 2026 bong-water-water-bong
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// The transient arena's share of the pool budget (engine#1711).
+//
+// The pool guardrail (#680) charges the ggml tensor buffers and nothing else: the per-launch scratch an op actually grows
+// into is allocated straight from the runtime allocator, so "a launch that needs more is refused here rather than faulting
+// the device" was only true of one of the two ways a launch can run out of memory. This is the decision that makes it true
+// of the scratch too, kept pure and separate from the arena so it can be tested without a device.
+//
+// Two rules matter and both are here rather than in the caller:
+//   - the arena's charge is its CURRENT capacity, released before a new one is taken. A launch that inherits a buffer at
+//     least as large as it needs is charged nothing (the arena returns early), which is what keeps the accounting from
+//     double-counting a buffer that is already paid for;
+//   - a refusal is a decision about the BUDGET, not about the device: `refuse` means the pool cannot take the request,
+//     which the caller turns into a Status and the caller of that into a failed build (never a raw allocation that may
+//     fault).
+#pragma once
+
+#include <cstddef>
+#include <vector>
+
+namespace ggml::hrx {
+
+struct TransientArenaCharge {
+    bool   allocate      = false;  // the arena has to (re)allocate at required_size
+    bool   release_first = false;  // its current capacity must go back to the pool first (it is being replaced)
+    size_t release_bytes = 0;
+    size_t charge_bytes  = 0;  // what the pool must accept before the allocation may happen
+    bool   refuse        = false;  // a request larger than the WHOLE budget: a fast path, not the gate (see below)
+};
+
+// The order the arena must perform the steps in, so a test can pin it without an allocator (engine#1711's review: the
+// first version reserved AFTER allocating, which left the faulting allocation exactly where it was). Reserve comes
+// before allocate, and the replaced capacity is released before either.
+enum class TransientArenaStep { ReleaseOld, Reserve, Allocate };
+
+// capacity: the arena's current allocation (0 when it has none). required: the size the caller needs. budget: the pool
+// budget the device was admitted with (0 = unset, which the guardrail treats as unbounded).
+inline TransientArenaCharge plan_transient_arena_charge(size_t capacity, size_t required, size_t budget) {
+    TransientArenaCharge plan;
+    if (required == 0 || capacity >= required) {
+        return plan;  // nothing to do, and nothing to charge: this is the no-double-count case
+    }
+    plan.allocate      = true;
+    plan.release_first = capacity > 0;
+    plan.release_bytes = capacity;
+    plan.charge_bytes  = required;
+    plan.refuse        = budget != 0 && required > budget;
+    return plan;
+}
+
+// The steps for a planned grow, in order. `release_old` is dropped when there is nothing to replace, so a test can assert
+// the relative order of the two that matter rather than a fixed length.
+inline std::vector<TransientArenaStep> transient_arena_steps(const TransientArenaCharge & plan) {
+    std::vector<TransientArenaStep> steps;
+    if (!plan.allocate) {
+        return steps;
+    }
+    if (plan.release_first) {
+        steps.push_back(TransientArenaStep::ReleaseOld);
+    }
+    steps.push_back(TransientArenaStep::Reserve);
+    steps.push_back(TransientArenaStep::Allocate);
+    return steps;
+}
+
+}  // namespace ggml::hrx

ggml/src/ggml-hrx/runtime/transient-arena.cpp

@@ -1,5 +1,9 @@
 #include "transient-arena.h"
 
+#include "transient-arena-budget.h"
+
+#include "hrx-1bit-safety.h"
+
 #include "hrx-interop-utils.h"
 #include "hrx_runtime.h"
 
@@ -36,6 +40,7 @@ TransientArena::~TransientArena() {
 
 void TransientArena::clear() {
     std::lock_guard<std::mutex> lock(mutex_);
+    release_budget_locked();
     if (buffer_ != nullptr) {
         hrx_buffer_release(buffer_);
         buffer_ = nullptr;
@@ -48,6 +53,13 @@ void TransientArena::clear() {
     allocation_id_       = kInvalidTransientArenaAllocationId;
 }
 
+void TransientArena::release_budget_locked() {
+    if (budget_charged_ != 0 && budget_ != nullptr && *budget_ != 0) {
+        onebit::pool_release(budget_charged_);
+    }
+    budget_charged_ = 0;
+}
+
 uint64_t TransientArena::next_allocation_id() {
     const uint64_t id = next_allocation_id_++;
     if (next_allocation_id_ == kInvalidTransientArenaAllocationId) {
@@ -79,6 +91,20 @@ Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t
         status.log("missing HRX stream for transient arena allocation");
         return status;
     }
+    // engine#1711: the scratch is charged against the pool budget like a tensor buffer. The COARSE check here is a fast
+    // path for the obvious case only -- a request larger than the whole budget cannot fit however empty the pool is. It
+    // is not the gate: the gate is `pool_reserve` below, which is the only thing that can see what the tensor buffers
+    // have already taken, and it runs BEFORE the allocator (review of this PR: reserving after the allocation lets the
+    // allocation that can fault happen first, which is the whole failure this change exists to close).
+    if (budget_ != nullptr && *budget_ != 0) {
+        const TransientArenaCharge plan = plan_transient_arena_charge(allocation_capacity_, required_size, *budget_);
+        if (plan.refuse) {
+            status.log("transient arena: refusing a %zu byte scratch: the whole pool budget is %zu bytes (GGML_HRX_POOL_LIMIT_MB "
+                       "or the GTT ceiling, engine #680/#1711)",
+                       plan.charge_bytes, *budget_);
+            return status;
+        }
+    }
     if (buffer_ != nullptr) {
         if (diagnostic_fresh) {
             diagnostic_retired_buffers_.push_back(buffer_);
@@ -94,6 +120,20 @@ Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t
         allocation_id_       = kInvalidTransientArenaAllocationId;
     }
 
+    // The replaced capacity goes back to the pool first, then the request is admitted, and only an admitted request
+    // reaches the allocator: reserve-then-allocate is the order that makes "refused here rather than faulting the
+    // device" true. The reverse order was in the first version of this change and it left the faulting allocation
+    // exactly where it was.
+    release_budget_locked();
+    const bool charged = budget_ != nullptr && *budget_ != 0;
+    if (charged && !onebit::pool_reserve(*budget_, required_size, "transient arena scratch")) {
+        status.log("transient arena: refusing a %zu byte scratch: the pool budget is %zu bytes and the tensor buffers "
+                   "are already charged against it (GGML_HRX_POOL_LIMIT_MB or the GTT ceiling, engine #680/#1711)",
+                   required_size, *budget_);
+        return status;
+    }
+    budget_charged_ = charged ? required_size : 0;
+
     hrx_buffer_params_t params = {
         HRX_MEMORY_TYPE_DEVICE_LOCAL,
         HRX_MEMORY_ACCESS_ALL,
@@ -103,6 +143,7 @@ Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t
     hrx_buffer_t allocation = nullptr;
     if (ErrorResult error = take_status(
             hrx_allocator_allocate_buffer(hrx_device_allocator(device), params, required_size, &allocation))) {
+        release_budget_locked();  // the reservation is not kept for an allocation that did not happen
         status.log("allocate transient arena: %s", error->c_str());
         return status;
     }

ggml/src/ggml-hrx/runtime/transient-arena.h

@@ -39,6 +39,11 @@ class TransientArena {
     TransientArena() = default;
     ~TransientArena();
 
+    // engine#1711: the pool budget this arena's scratch is charged against. A pointer rather than a copy because the
+    // device context computes it at init and would otherwise have to push it in again; nullptr (the default) means the
+    // arena is uncharged, which is what every test that builds one directly keeps getting.
+    void set_budget(size_t * budget) { budget_ = budget; }
+
     TransientArena(const TransientArena &)             = delete;
     TransientArena & operator=(const TransientArena &) = delete;
 
@@ -53,13 +58,16 @@ class TransientArena {
     uint64_t allocation_id() const;
 
   private:
+    void                        release_budget_locked();
     uint64_t                    next_allocation_id();
     Status                      ensure_capacity_locked(hrx_device_t device, hrx_stream_t stream, size_t required_size);
     TransientArenaAllocationRef current_allocation_locked() const;
 
     mutable std::mutex mutex_;
     hrx_buffer_t       buffer_              = nullptr;
     std::vector<hrx_buffer_t> diagnostic_retired_buffers_;
+    size_t *           budget_              = nullptr;
+    size_t             budget_charged_      = 0;
     size_t             allocation_capacity_ = 0;
     uint64_t           allocation_id_       = kInvalidTransientArenaAllocationId;
     uint64_t           next_allocation_id_  = 1;

tests/CMakeLists.txt

@@ -365,6 +365,12 @@ if (TARGET ggml-hrx)
     target_include_directories(test-hrx-arena-refusal PRIVATE ../ggml/src/ggml-hrx)
     add_test(NAME test-hrx-arena-refusal COMMAND test-hrx-arena-refusal)
 
+    # The transient arena's share of the pool budget (engine#1711): the pure decision, no device. The property is that
+    # the arena is charged what it HOLDS, so a launch that fits inside an existing buffer is charged nothing.
+    add_executable(test-hrx-arena-budget test-hrx-arena-budget.cpp)
+    target_include_directories(test-hrx-arena-budget PRIVATE ../ggml/src/ggml-hrx)
+    add_test(NAME test-hrx-arena-budget COMMAND test-hrx-arena-budget)
+
     # MemAvailable parsing for the unbounded-residency warning: pure, no device and no /proc (engine #669)
     add_executable(test-hrx-host-memory-available test-hrx-host-memory-available.cpp
         ../ggml/src/ggml-hrx/runtime/host-memory-available.cpp)

tests/test-hrx-arena-budget.cpp

@@ -0,0 +1,118 @@
+// Copyright 2026 bong-water-water-bong
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Host test for the transient arena's share of the pool budget (engine#1711): the pure decision, no device, no HRX
+// header, no model. The property that matters is that the arena is charged what it HOLDS, not what it asks for each
+// launch: a launch that inherits a buffer at least as large as it needs must be charged nothing, or the accounting
+// double-counts a buffer that is already paid for and refuses launches that fit.
+
+#include "runtime/transient-arena-budget.h"
+
+#include <cstdio>
+
+static int failures = 0;
+
+#define REQUIRE(condition)                                                       \
+    do {                                                                         \
+        if (!(condition)) {                                                      \
+            std::printf("FAIL %s:%d: %s\n", __FILE__, __LINE__, #condition);     \
+            ++failures;                                                          \
+        }                                                                        \
+    } while (0)
+
+int main() {
+    using ggml::hrx::plan_transient_arena_charge;
+
+    {
+        // first launch: no arena yet, so the whole request is charged
+        const auto p = plan_transient_arena_charge(0, 1024, 4096);
+        REQUIRE(p.allocate);
+        REQUIRE(!p.release_first);
+        REQUIRE(p.release_bytes == 0);
+        REQUIRE(p.charge_bytes == 1024);
+        REQUIRE(!p.refuse);
+    }
+    {
+        // the no-double-count rule: a launch that fits inside what the arena already holds is charged NOTHING
+        const auto p = plan_transient_arena_charge(4096, 1024, 4096);
+        REQUIRE(!p.allocate);
+        REQUIRE(p.charge_bytes == 0);
+        REQUIRE(!p.refuse);
+        const auto equal = plan_transient_arena_charge(1024, 1024, 4096);
+        REQUIRE(!equal.allocate);
+        REQUIRE(equal.charge_bytes == 0);
+    }
+    {
+        // growing: the old capacity goes back to the pool and the new size is charged -- net, the delta
+        const auto p = plan_transient_arena_charge(1024, 4096, 8192);
+        REQUIRE(p.allocate);
+        REQUIRE(p.release_first);
+        REQUIRE(p.release_bytes == 1024);
+        REQUIRE(p.charge_bytes == 4096);
+        REQUIRE(!p.refuse);
+    }
+    {
+        // a request the budget cannot take is a refusal, and it is a refusal BEFORE the allocator is reached
+        const auto p = plan_transient_arena_charge(1024, 8192, 4096);
+        REQUIRE(p.allocate);
+        REQUIRE(p.refuse);
+        REQUIRE(p.charge_bytes == 8192);
+    }
+    {
+        // an unset budget (0) is unbounded, as pool_reserve() treats it -- the guardrail is opt-in
+        const auto p = plan_transient_arena_charge(0, 1u << 30, 0);
+        REQUIRE(p.allocate);
+        REQUIRE(!p.refuse);
+    }
+    {
+        // zero-size requests never allocate and never charge, whatever the budget
+        const auto p = plan_transient_arena_charge(0, 0, 64);
+        REQUIRE(!p.allocate);
+        REQUIRE(p.charge_bytes == 0);
+        REQUIRE(!p.refuse);
+    }
+    {
+        // exactly at the budget is allowed: the guardrail refuses what does not fit, not what exactly fits
+        const auto p = plan_transient_arena_charge(0, 4096, 4096);
+        REQUIRE(p.allocate);
+        REQUIRE(!p.refuse);
+    }
+
+    {
+        // THE ORDER (engine#1711 review): reserve before allocate, and the replaced capacity released before either.
+        // A first version of this change allocated first and reserved after, so the allocation that can fault ran
+        // before admission -- which is the failure the change exists to close, reintroduced by ordering.
+        using ggml::hrx::TransientArenaStep;
+        using ggml::hrx::transient_arena_steps;
+        const auto grow = transient_arena_steps(plan_transient_arena_charge(1024, 4096, 8192));
+        REQUIRE(grow.size() == 3);
+        REQUIRE(grow[0] == TransientArenaStep::ReleaseOld);
+        REQUIRE(grow[1] == TransientArenaStep::Reserve);
+        REQUIRE(grow[2] == TransientArenaStep::Allocate);
+        const auto first = transient_arena_steps(plan_transient_arena_charge(0, 4096, 8192));
+        REQUIRE(first.size() == 2);
+        REQUIRE(first[0] == TransientArenaStep::Reserve);
+        REQUIRE(first[1] == TransientArenaStep::Allocate);
+        const auto none = transient_arena_steps(plan_transient_arena_charge(4096, 1024, 8192));
+        REQUIRE(none.empty());
+    }
+
+    if (failures != 0) {
+        std::printf("test-hrx-arena-budget: %d failure(s)\n", failures);
+        return 1;
+    }
+    std::printf("test-hrx-arena-budget: all checks passed\n");
+    return 0;
+}