diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 705dcc54155a..e24b7ad23653 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -1583,6 +1583,10 @@ static std::unique_ptr 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. diff --git a/ggml/src/ggml-hrx/runtime/transient-arena-budget.h b/ggml/src/ggml-hrx/runtime/transient-arena-budget.h new file mode 100644 index 000000000000..3a7342db551d --- /dev/null +++ b/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 +#include + +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 transient_arena_steps(const TransientArenaCharge & plan) { + std::vector 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 diff --git a/ggml/src/ggml-hrx/runtime/transient-arena.cpp b/ggml/src/ggml-hrx/runtime/transient-arena.cpp index a12f89e0ab38..c6f9dccbe7aa 100644 --- a/ggml/src/ggml-hrx/runtime/transient-arena.cpp +++ b/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 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; } diff --git a/ggml/src/ggml-hrx/runtime/transient-arena.h b/ggml/src/ggml-hrx/runtime/transient-arena.h index a50639962c7c..79ace60452e8 100644 --- a/ggml/src/ggml-hrx/runtime/transient-arena.h +++ b/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,6 +58,7 @@ 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; @@ -60,6 +66,8 @@ class TransientArena { mutable std::mutex mutex_; hrx_buffer_t buffer_ = nullptr; std::vector 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; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c592e6b41034..46b3ebc81c76 100644 --- a/tests/CMakeLists.txt +++ b/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) diff --git a/tests/test-hrx-arena-budget.cpp b/tests/test-hrx-arena-budget.cpp new file mode 100644 index 000000000000..6822aca3d0f2 --- /dev/null +++ b/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 + +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; +}