From 4843ea80cae673ea287c3eba31c3d92e1c195b99 Mon Sep 17 00:00:00 2001 From: Jess Sullivan Date: Sat, 19 Sep 2026 17:06:23 -0400 Subject: [PATCH 1/3] rpc : include nb in the get_alloc_size cache key and floor the result at ggml_nbytes --- ggml/src/ggml-rpc/ggml-rpc.cpp | 12 ++++++++++-- tests/test-rpc-multi-server.cpp | 24 +++++++++++++++++++++++- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index c24caad77ee9..d2948ea1111f 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -858,6 +858,10 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty if (rpc_get) { ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context; + // the reported size must never be below ggml_nbytes: rpc_tensor stores nb[] as uint32_t, + // so a stride over 4 GiB is truncated on the wire and the remote size comes back too small + const size_t min_size = ggml_nbytes(tensor); + // Cache key for calls to read the alloc_size. // We deliberately exclude src tensor dimensions from the key because: // 1. For CPU backends, alloc_size = ggml_nbytes(output) regardless of src shapes @@ -871,6 +875,9 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty uint32_t op; int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; uint32_t ne[GGML_MAX_DIMS]; + // nb[] must be in the key: the map is indexed by the key hash and never compares the + // key, so two tensors that differ only in strides would share one entry + uint64_t nb[GGML_MAX_DIMS]; }; alloc_size_cache_key key = {}; @@ -880,6 +887,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty memcpy(key.op_params, tensor->op_params, sizeof(key.op_params)); for (int i = 0; i < GGML_MAX_DIMS; i++) { key.ne[i] = (uint32_t)tensor->ne[i]; + key.nb[i] = (uint64_t)tensor->nb[i]; } uint64_t cache_hash = fnv_hash((const uint8_t *)&key, sizeof(key)); @@ -893,7 +901,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty std::lock_guard lock(cache_mutex); auto it = cache.find(cache_hash); if (it != cache.end()) { - return it->second; + return std::max(it->second, min_size); } } @@ -915,7 +923,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty cache[cache_hash] = response.alloc_size; } - return response.alloc_size; + return std::max(response.alloc_size, min_size); } return ggml_nbytes(tensor); diff --git a/tests/test-rpc-multi-server.cpp b/tests/test-rpc-multi-server.cpp index 4502e2ce71fa..f253b6e144d6 100644 --- a/tests/test-rpc-multi-server.cpp +++ b/tests/test-rpc-multi-server.cpp @@ -17,7 +17,7 @@ int main(int argc, char ** argv) { GGML_ASSERT(backend_b != nullptr); ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead() + ggml_graph_overhead_custom(1, false), + /* .mem_size = */ 3*ggml_tensor_overhead() + ggml_graph_overhead_custom(1, false), /* .mem_buffer = */ nullptr, /* .no_alloc = */ true, }; @@ -40,6 +40,28 @@ int main(int argc, char ** argv) { ggml_backend_rpc_get_device_memory(endpoint_b, 0, &free_mem, &total_mem); GGML_ASSERT(total_mem > 0); ggml_backend_buffer_free(buffer); + + // Two tensors with the same ne[] but different nb[] must not share a cached alloc size. + // ref: https://github.com/ggml-org/llama.cpp/issues/28360 + ggml_backend_buffer_type_t buft = ggml_backend_rpc_buffer_type(endpoint_a, 0); + GGML_ASSERT(buft != nullptr); + + // MUL_MAT may need extra memory, so the size is read from the server [TAG_ALLOC_SIZE_EXPAND] + ggml_tensor * packed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 64, 64); + packed->op = GGML_OP_MUL_MAT; + + // same ne[], twice the row stride + ggml_tensor * strided = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 64, 64); + strided->op = GGML_OP_MUL_MAT; + strided->nb[1] = 2*strided->nb[1]; + strided->nb[2] = strided->ne[1]*strided->nb[1]; + strided->nb[3] = strided->nb[2]; + GGML_ASSERT(ggml_nbytes(strided) > ggml_nbytes(packed)); + + // ask for the packed tensor first, so a cache keyed without nb[] holds the smaller size + GGML_ASSERT(ggml_backend_buft_get_alloc_size(buft, packed) >= ggml_nbytes(packed)); + GGML_ASSERT(ggml_backend_buft_get_alloc_size(buft, strided) >= ggml_nbytes(strided)); + ggml_free(ctx); ggml_backend_free(backend_b); ggml_backend_free(backend_a); From 172011478c20058620867c3f7f15711058d0d51f Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 25 Sep 2026 10:33:32 +0300 Subject: [PATCH 2/3] cont : remove redundant comment --- ggml/src/ggml-rpc/ggml-rpc.cpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index d2948ea1111f..a58e251b10cf 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -875,8 +875,6 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty uint32_t op; int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; uint32_t ne[GGML_MAX_DIMS]; - // nb[] must be in the key: the map is indexed by the key hash and never compares the - // key, so two tensors that differ only in strides would share one entry uint64_t nb[GGML_MAX_DIMS]; }; From bb0eed026c4e10139572ff213a7aefddb36c8680 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 25 Sep 2026 10:34:51 +0300 Subject: [PATCH 3/3] cont : add TODO --- ggml/src/ggml-rpc/ggml-rpc.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index a58e251b10cf..353b79b07212 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -860,6 +860,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty // the reported size must never be below ggml_nbytes: rpc_tensor stores nb[] as uint32_t, // so a stride over 4 GiB is truncated on the wire and the remote size comes back too small + // TODO: change rpc_tensor nb to 64-bit int const size_t min_size = ggml_nbytes(tensor); // Cache key for calls to read the alloc_size.