char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

hmirinrpc: avoid serializing buffers from other servers (#26500)a7cc83bba

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
1.6 KiB47 linesraw
1#include "ggml-alloc.h"
2#include "ggml-backend.h"
3#include "ggml-impl.h"
4#include "ggml-rpc.h"
5#include "ggml.h"
6
7int main(int argc, char ** argv) {
8    GGML_ASSERT(argc == 3);
9    ggml_backend_load_all();
10
11    const char * endpoint_a = argv[1];
12    const char * endpoint_b = argv[2];
13
14    ggml_backend_t backend_a = ggml_backend_rpc_init(endpoint_a, 0);
15    ggml_backend_t backend_b = ggml_backend_rpc_init(endpoint_b, 0);
16    GGML_ASSERT(backend_a != nullptr);
17    GGML_ASSERT(backend_b != nullptr);
18
19    ggml_init_params params = {
20        /* .mem_size   = */ ggml_tensor_overhead() + ggml_graph_overhead_custom(1, false),
21        /* .mem_buffer = */ nullptr,
22        /* .no_alloc   = */ true,
23    };
24    ggml_context * ctx = ggml_init(params);
25    GGML_ASSERT(ctx != nullptr);
26
27    ggml_tensor * tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
28    ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend_a);
29    GGML_ASSERT(buffer != nullptr);
30
31    // A remote pointer allocated by server A is not meaningful to server B.
32    ggml_cgraph * graph = ggml_new_graph_custom(ctx, 1, false);
33    graph->nodes[0] = tensor;
34    graph->n_nodes = 1;
35
36    GGML_ASSERT(ggml_backend_graph_compute(backend_b, graph) == GGML_STATUS_SUCCESS);
37    // Wait for server B to finish the graph before the script checks its log.
38    size_t free_mem;
39    size_t total_mem;
40    ggml_backend_rpc_get_device_memory(endpoint_b, 0, &free_mem, &total_mem);
41    GGML_ASSERT(total_mem > 0);
42    ggml_backend_buffer_free(buffer);
43    ggml_free(ctx);
44    ggml_backend_free(backend_b);
45    ggml_backend_free(backend_a);
46    return 0;
47}