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>

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);

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);