diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 705dcc54155a..67263e46332b 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,14 @@ 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. 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) { if (device->copy_staging_size >= size) { return true; @@ -265,13 +274,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 +297,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..8628b4a9e1f9 100644 --- a/tests/test-hrx-buffer.cpp +++ b/tests/test-hrx-buffer.cpp @@ -353,6 +353,50 @@ 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. +// 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 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 + 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); + + 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_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 +432,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) {