From 5b8c6c9e7f25c413b515e994aff780964541caf1 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Sat, 7 Nov 2026 22:41:23 -0400 Subject: [PATCH 1/2] hrx: cap the host copy staging buffer and copy through it in chunks (engine #856) The per-device staging buffer for ordered host copies (engine #696) grew to the largest set_tensor and was never released, so after a load it held a host-mapped copy of the largest weight for the life of the process (995 MiB on a 27B, 304 MiB on a 4B). Cap it at 64 MiB and move larger copies through it in chunks, each synchronized before the buffer is reused, so the copy ordering is unchanged. test-hrx-buffer gains chunked_copy: 2*64 MiB + 12345 bytes at an unaligned offset round-trips intact and the staging buffer stays at the cap. Co-Authored-By: Claude Opus 5.5 --- ggml/src/ggml-hrx/ggml-hrx.cpp | 44 ++++++++++++++++++++++++--------- tests/test-hrx-buffer.cpp | 45 ++++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 705dcc54155a..d7dc3733bf1b 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -24,6 +24,7 @@ #include "runtime/prepared-command-program-cache.h" #include "runtime/transient-arena.h" +#include #include #include #include @@ -241,6 +242,12 @@ static void buffer_memset(ggml_backend_buffer_t buffer, // one dispatch queue, and the synchronize that follows completes them before the caller continues. The staging // buffer is kept per device (few distinct sizes) instead of being allocated per copy: per-copy allocate/release is // what libhrx's per-size pools retain (engine #669). +// engine #856: the staging buffer is capped, and a larger copy goes through it in chunks. Sizing it to the largest +// copy kept one host-mapped buffer the size of the largest weight (995 MiB on a 27B, usually the token embedding) +// resident in GTT for the life of the process, only to stage the load. Each chunk is synchronized before the +// buffer is reused, so the ordering of engine #696 is unchanged. +static constexpr size_t kCopyStagingChunk = 64ull << 20; + static bool ensure_copy_staging(ggml_backend_hrx_device_context * device, size_t size) { if (device->copy_staging_size >= size) { return true; @@ -265,13 +272,21 @@ static bool ordered_upload(ggml_backend_hrx_device_context * device, hrx_buffer_t destination, size_t offset) { std::lock_guard lock(device->buffer_stream_mutex); - if (!ensure_copy_staging(device, size)) { + const size_t chunk = std::min(size, kCopyStagingChunk); + if (!ensure_copy_staging(device, chunk)) { return false; } - std::memcpy(device->copy_staging_mapped, data, size); - return HRX_CHECK(hrx_stream_copy_buffer(device->buffer_stream, device->copy_staging, 0, destination, offset, - size)) && - HRX_CHECK(hrx_stream_synchronize(device->buffer_stream)); + const auto * source = static_cast(data); + for (size_t done = 0; done < size; done += chunk) { + const size_t n = std::min(chunk, size - done); + std::memcpy(device->copy_staging_mapped, source + done, n); + if (!HRX_CHECK(hrx_stream_copy_buffer(device->buffer_stream, device->copy_staging, 0, destination, + offset + done, n)) || + !HRX_CHECK(hrx_stream_synchronize(device->buffer_stream))) { + return false; + } + } + return true; } static bool ordered_download(ggml_backend_hrx_device_context * device, @@ -280,16 +295,21 @@ static bool ordered_download(ggml_backend_hrx_device_context * device, void * data, size_t size) { std::lock_guard lock(device->buffer_stream_mutex); - if (!ensure_copy_staging(device, size)) { + const size_t chunk = std::min(size, kCopyStagingChunk); + if (!ensure_copy_staging(device, chunk)) { return false; } - const bool ok = HRX_CHECK(hrx_stream_copy_buffer(device->buffer_stream, source, offset, device->copy_staging, 0, - size)) && - HRX_CHECK(hrx_stream_synchronize(device->buffer_stream)); - if (ok) { - std::memcpy(data, device->copy_staging_mapped, size); + auto * destination = static_cast(data); + for (size_t done = 0; done < size; done += chunk) { + const size_t n = std::min(chunk, size - done); + if (!HRX_CHECK(hrx_stream_copy_buffer(device->buffer_stream, source, offset + done, device->copy_staging, 0, + n)) || + !HRX_CHECK(hrx_stream_synchronize(device->buffer_stream))) { + return false; + } + std::memcpy(destination + done, device->copy_staging_mapped, n); } - return ok; + return true; } static void buffer_set(ggml_backend_buffer_t buffer, diff --git a/tests/test-hrx-buffer.cpp b/tests/test-hrx-buffer.cpp index e281a3199f09..f489c46da34c 100644 --- a/tests/test-hrx-buffer.cpp +++ b/tests/test-hrx-buffer.cpp @@ -353,6 +353,46 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { REQUIRE(weight_stats.resident_bytes == 0); } +// engine #856: a copy larger than the staging cap goes through the capped buffer in chunks. The bytes must arrive +// intact across chunk edges (a non-multiple tail, a non-zero offset) and the staging buffer must stay at the cap +// instead of growing to the copy size. +static void run_chunked_copy_checks(ggml_backend_t backend) { + ggml_backend_hrx_context * hrx = backend_context(backend); + + constexpr size_t kCap = 64ull << 20; + constexpr size_t kOffset = 4096 + 12; + constexpr size_t kBytes = 2 * kCap + 12345; + ggml_init_params params = {}; + params.mem_size = 16 * 1024; + params.no_alloc = true; + ggml_context * context = ggml_init(params); + REQUIRE(context != nullptr); + ggml_tensor * tensor = ggml_new_tensor_1d(context, GGML_TYPE_I8, static_cast(kOffset + kBytes)); + ggml_backend_buffer_t buffer = ggml_backend_alloc_buffer(backend, kOffset + kBytes + 4096); + REQUIRE(buffer != nullptr); + tensor->buffer = buffer; + tensor->data = ggml_backend_buffer_get_base(buffer); + REQUIRE(ggml_backend_buffer_init_tensor(buffer, tensor) == GGML_STATUS_SUCCESS); + + std::vector input(kBytes); + uint32_t x = 0x9e3779b9u; + for (uint8_t & b : input) { + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + b = static_cast(x); + } + ggml_backend_tensor_set(tensor, input.data(), kOffset, kBytes); + REQUIRE(hrx->device->copy_staging_size <= kCap); + std::vector output(kBytes, 0); + ggml_backend_tensor_get(tensor, output.data(), kOffset, kBytes); + REQUIRE(hrx->device->copy_staging_size <= kCap); + REQUIRE(std::memcmp(output.data(), input.data(), kBytes) == 0); + + ggml_backend_buffer_free(buffer); + ggml_free(context); +} + template static void run_with_hrx_backend(F && f) { ggml_backend_t backend = ggml_backend_hrx_init(0); @@ -388,6 +428,11 @@ static void register_hrx_buffer_cases(test_runner::Suite & suite) { run_host_weight_cache_checks(context); }); }); + suite.device_case("chunked_copy", [] { + run_with_hrx_backend([](ggml_backend_t backend, ggml_backend_hrx_context *) { + run_chunked_copy_checks(backend); + }); + }); } int main(int argc, char ** argv) { From 287f7c23bfb60d25802cf5865a24c495fb0231af Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Mon, 9 Nov 2026 16:10:22 -0400 Subject: [PATCH 2/2] hrx: state the staging lock and residual, add cap boundary arms (engine #856) Review follow-up on the chunked copy staging change: - the comment states that buffer_stream_mutex covers all chunks of one copy, and that the buffer stays at most 64 MiB per device, never shrunk, on purpose; - chunked_copy also round-trips exactly the cap and the cap plus one byte. Co-Authored-By: Claude Opus 5.5 --- ggml/src/ggml-hrx/ggml-hrx.cpp | 4 +++- tests/test-hrx-buffer.cpp | 36 +++++++++++++++++++--------------- 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index d7dc3733bf1b..67263e46332b 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -245,7 +245,9 @@ static void buffer_memset(ggml_backend_buffer_t buffer, // engine #856: the staging buffer is capped, and a larger copy goes through it in chunks. Sizing it to the largest // copy kept one host-mapped buffer the size of the largest weight (995 MiB on a 27B, usually the token embedding) // resident in GTT for the life of the process, only to stage the load. Each chunk is synchronized before the -// buffer is reused, so the ordering of engine #696 is unchanged. +// buffer is reused, so the ordering of engine #696 is unchanged. The callers hold buffer_stream_mutex across all +// chunks of one copy, so no other staged copy can reallocate the buffer between them. The buffer still never +// shrinks: it stays at most kCopyStagingChunk per device for the life of the process, kept for reuse on purpose. static constexpr size_t kCopyStagingChunk = 64ull << 20; static bool ensure_copy_staging(ggml_backend_hrx_device_context * device, size_t size) { diff --git a/tests/test-hrx-buffer.cpp b/tests/test-hrx-buffer.cpp index f489c46da34c..8628b4a9e1f9 100644 --- a/tests/test-hrx-buffer.cpp +++ b/tests/test-hrx-buffer.cpp @@ -356,38 +356,42 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { // engine #856: a copy larger than the staging cap goes through the capped buffer in chunks. The bytes must arrive // intact across chunk edges (a non-multiple tail, a non-zero offset) and the staging buffer must stay at the cap // instead of growing to the copy size. +// Copies through the capped staging buffer: exactly the cap, one byte over it, and two chunks plus a partial tail, +// all at an unaligned offset. Each must arrive intact and leave the staging buffer at or below the cap. static void run_chunked_copy_checks(ggml_backend_t backend) { ggml_backend_hrx_context * hrx = backend_context(backend); constexpr size_t kCap = 64ull << 20; constexpr size_t kOffset = 4096 + 12; - constexpr size_t kBytes = 2 * kCap + 12345; + constexpr size_t kMax = 2 * kCap + 12345; ggml_init_params params = {}; params.mem_size = 16 * 1024; params.no_alloc = true; ggml_context * context = ggml_init(params); REQUIRE(context != nullptr); - ggml_tensor * tensor = ggml_new_tensor_1d(context, GGML_TYPE_I8, static_cast(kOffset + kBytes)); - ggml_backend_buffer_t buffer = ggml_backend_alloc_buffer(backend, kOffset + kBytes + 4096); + ggml_tensor * tensor = ggml_new_tensor_1d(context, GGML_TYPE_I8, static_cast(kOffset + kMax)); + ggml_backend_buffer_t buffer = ggml_backend_alloc_buffer(backend, kOffset + kMax + 4096); REQUIRE(buffer != nullptr); tensor->buffer = buffer; tensor->data = ggml_backend_buffer_get_base(buffer); REQUIRE(ggml_backend_buffer_init_tensor(buffer, tensor) == GGML_STATUS_SUCCESS); - std::vector input(kBytes); - uint32_t x = 0x9e3779b9u; - for (uint8_t & b : input) { - x ^= x << 13; - x ^= x >> 17; - x ^= x << 5; - b = static_cast(x); + uint32_t x = 0x9e3779b9u; + for (const size_t bytes : { kCap, kCap + 1, kMax }) { + std::vector input(bytes); + for (uint8_t & b : input) { + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + b = static_cast(x); + } + ggml_backend_tensor_set(tensor, input.data(), kOffset, bytes); + REQUIRE(hrx->device->copy_staging_size <= kCap); + std::vector output(bytes, 0); + ggml_backend_tensor_get(tensor, output.data(), kOffset, bytes); + REQUIRE(hrx->device->copy_staging_size <= kCap); + REQUIRE(std::memcmp(output.data(), input.data(), bytes) == 0); } - ggml_backend_tensor_set(tensor, input.data(), kOffset, kBytes); - REQUIRE(hrx->device->copy_staging_size <= kCap); - std::vector output(kBytes, 0); - ggml_backend_tensor_get(tensor, output.data(), kOffset, kBytes); - REQUIRE(hrx->device->copy_staging_size <= kCap); - REQUIRE(std::memcmp(output.data(), input.data(), kBytes) == 0); ggml_backend_buffer_free(buffer); ggml_free(context);