Commit 66963a8bc for llama.cpp
commit 66963a8bc711773c36c1bc342356bb9f59c10481
Author: Jess Sullivan <Jess@sulliwood.org>
Date: Fri Sep 25 03:35:09 2026 -0400
rpc: include nb in the get_alloc_size cache key and floor the result at ggml_nbytes (#29283)
* rpc : include nb in the get_alloc_size cache key and floor the result at ggml_nbytes
* cont : remove redundant comment
* cont : add TODO
---------
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp
index c24caad77..353b79b07 100644
--- a/ggml/src/ggml-rpc/ggml-rpc.cpp
+++ b/ggml/src/ggml-rpc/ggml-rpc.cpp
@@ -858,6 +858,11 @@ 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
+ // 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.
// 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 +876,7 @@ 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];
+ uint64_t nb[GGML_MAX_DIMS];
};
alloc_size_cache_key key = {};
@@ -880,6 +886,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 +900,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty
std::lock_guard<std::mutex> lock(cache_mutex);
auto it = cache.find(cache_hash);
if (it != cache.end()) {
- return it->second;
+ return std::max<size_t>(it->second, min_size);
}
}
@@ -915,7 +922,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<size_t>(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 4502e2ce7..f253b6e14 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);