Team Ai
Datasetpublic

Brunobkr/llama.cpp_AlgMor24_github

ΩFFFΣLLIa • llama.cpp • AlgMor24 ██████╗ ███████╗███████╗███████╗██╗ ██╗ ██╗ █████╗ ██╔═══██╗██╔════╝██╔════╝██╔════╝██║ ██║ ██║██╔══██╗ ██║ ██║█████╗ █████╗ █████╗ ██║ ██║ ██║███████║ ██║ ██║██╔══╝ ██╔══╝ ██╔══╝ ██║ ██║ ██║██╔══██║ ╚██████╔╝██║ ██║ ███████╗███████╗███████╗██║██║ ██║ ╚═════╝ ╚═╝ ╚═╝ ╚══════╝╚══════╝╚══════╝╚═╝╚═╝ ╚═╝ High-Performance LLM / VLM Inference & Autonomous Agentic Ecosystem… See the full description on the dataset page: https://huggingface.co/datasets/Brunobkr/llama.cpp_AlgMor24_github.

sourceHugging Faceupdated 2mo agoView on Hugging Face
0likes3.1kdownloads
ggml-rpc.cpp2056 linesDownload Raw Back to ggml-rpc
1#include "ggml-rpc.h"2#include "ggml-impl.h"3#include "ggml-backend-impl.h"4#include "ggml-cpp.h"5#include "transport.h"6 7#include <array>8#include <cinttypes>9#include <optional>10#include <string>11#include <vector>12#include <memory>13#include <mutex>14#include <unordered_map>15#include <unordered_set>16#include <cstring>17#include <fstream>18#include <filesystem>19#include <algorithm>20 21static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG");22 23#define LOG_DBG(...) \24    do { if (RPC_DEBUG) GGML_LOG_DEBUG(__VA_ARGS__); } while (0)25 26 27namespace fs = std::filesystem;28 29// macro for nicer error messages on server crash30#define RPC_STATUS_ASSERT(x) if (!(x)) GGML_ABORT("Remote RPC server crashed or returned malformed response")31 32// all RPC structures must be packed33#pragma pack(push, 1)34// ggml_tensor is serialized into rpc_tensor35struct rpc_tensor {36    uint64_t id;37    uint32_t type;38    uint64_t buffer;39    uint32_t ne[GGML_MAX_DIMS];40    uint32_t nb[GGML_MAX_DIMS];41    uint32_t op;42    int32_t  op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)];43    int32_t  flags;44    uint64_t src[GGML_MAX_SRC];45    uint64_t view_src;46    uint64_t view_offs;47    uint64_t data;48    char name[GGML_MAX_NAME];49 50    char padding[4];51};52 53static_assert(sizeof(rpc_tensor) % 8 == 0, "rpc_tensor size must be multiple of 8");54 55// RPC commands56enum rpc_cmd {57    RPC_CMD_ALLOC_BUFFER = 0,58    RPC_CMD_GET_ALIGNMENT,59    RPC_CMD_GET_MAX_SIZE,60    RPC_CMD_BUFFER_GET_BASE,61    RPC_CMD_FREE_BUFFER,62    RPC_CMD_BUFFER_CLEAR,63    RPC_CMD_SET_TENSOR,64    RPC_CMD_SET_TENSOR_HASH,65    RPC_CMD_GET_TENSOR,66    RPC_CMD_COPY_TENSOR,67    RPC_CMD_GRAPH_COMPUTE,68    RPC_CMD_GET_DEVICE_MEMORY,69    RPC_CMD_INIT_TENSOR,70    RPC_CMD_GET_ALLOC_SIZE,71    RPC_CMD_HELLO,72    RPC_CMD_DEVICE_COUNT,73    RPC_CMD_GRAPH_RECOMPUTE,74    RPC_CMD_MEMSET_TENSOR,75    RPC_CMD_COUNT,76};77 78static_assert(RPC_CMD_HELLO == 14, "RPC_CMD_HELLO must be always 14");79 80// Try RPC_CMD_SET_TENSOR_HASH first when data size is larger than this threshold81const size_t HASH_THRESHOLD = 10 * 1024 * 1024;82 83struct rpc_msg_hello_req {84    uint8_t conn_caps[RPC_CONN_CAPS_SIZE];85};86 87struct rpc_msg_hello_rsp {88    uint8_t major;89    uint8_t minor;90    uint8_t patch;91    uint8_t padding;92    uint8_t conn_caps[RPC_CONN_CAPS_SIZE];93};94 95struct rpc_msg_device_count_rsp {96    uint32_t device_count;97};98 99struct rpc_msg_get_alloc_size_req {100    uint32_t   device;101    rpc_tensor tensor;102    rpc_tensor srcs[GGML_MAX_SRC];103};104 105struct rpc_msg_get_alloc_size_rsp {106    uint64_t alloc_size;107};108 109struct rpc_msg_init_tensor_req {110    rpc_tensor tensor;111};112 113struct rpc_msg_alloc_buffer_req {114    uint32_t device;115    uint64_t size;116};117 118struct rpc_msg_alloc_buffer_rsp {119    uint64_t remote_ptr;120    uint64_t remote_size;121};122 123struct rpc_msg_get_alignment_req {124    uint32_t device;125};126 127struct rpc_msg_get_alignment_rsp {128    uint64_t alignment;129};130 131struct rpc_msg_get_max_size_req {132    uint32_t device;133};134 135struct rpc_msg_get_max_size_rsp {136    uint64_t max_size;137};138 139struct rpc_msg_buffer_get_base_req {140    uint64_t remote_ptr;141};142 143struct rpc_msg_buffer_get_base_rsp {144    uint64_t base_ptr;145};146 147struct rpc_msg_free_buffer_req {148    uint64_t remote_ptr;149};150 151struct rpc_msg_buffer_clear_req {152    uint64_t remote_ptr;153    uint8_t value;154};155 156struct rpc_msg_memset_tensor_req {157    rpc_tensor tensor;158    uint64_t offset;159    uint64_t size;160    uint8_t value;161};162 163struct rpc_msg_set_tensor_hash_req {164    rpc_tensor tensor;165    uint64_t offset;166    uint64_t hash;167};168 169struct rpc_msg_set_tensor_hash_rsp {170    uint8_t result;171};172 173struct rpc_msg_get_tensor_req {174    rpc_tensor tensor;175    uint64_t offset;176    uint64_t size;177};178 179struct rpc_msg_copy_tensor_req {180    rpc_tensor src;181    rpc_tensor dst;182};183 184struct rpc_msg_copy_tensor_rsp {185    uint8_t result;186};187 188struct rpc_msg_get_device_memory_req {189    uint32_t device;190};191 192struct rpc_msg_get_device_memory_rsp {193    uint64_t free_mem;194    uint64_t total_mem;195};196 197struct rpc_msg_graph_recompute_req {198    uint32_t device;199};200 201#pragma pack(pop)202 203// RPC data structures204 205static ggml_guid_t ggml_backend_rpc_guid() {206    static ggml_guid guid = {0x99, 0x68, 0x5b, 0x6c, 0xd2, 0x83, 0x3d, 0x24, 0x25, 0x36, 0x72, 0xe1, 0x5b, 0x0e, 0x14, 0x03};207    return &guid;208}209 210struct ggml_backend_rpc_device_context {211    std::string endpoint;212    uint32_t    device;213    std::string name;214    std::string description;215    uint64_t    last_graph_uid;216};217 218struct ggml_backend_rpc_buffer_type_context {219    std::string endpoint;220    uint32_t    device;221    std::string name;222    size_t      alignment;223    size_t      max_size;224};225 226struct ggml_backend_rpc_context {227    std::string endpoint;228    uint32_t    device;229    std::string name;230};231 232struct ggml_backend_rpc_buffer_context {233    std::shared_ptr<socket_t> sock;234    void * base_ptr;235    uint64_t remote_ptr;236};237 238// RPC helper functions239 240// Computes FNV-1a hash of the data241static uint64_t fnv_hash(const uint8_t * data, size_t len) {242    const uint64_t fnv_prime = 0x100000001b3ULL;243    uint64_t hash = 0xcbf29ce484222325ULL;244 245    for (size_t i = 0; i < len; ++i) {246        hash ^= data[i];247        hash *= fnv_prime;248    }249    return hash;250}251 252static bool send_msg(socket_ptr sock, const void * msg, size_t msg_size) {253    if (!sock->send_data(&msg_size, sizeof(msg_size))) {254        return false;255    }256    return sock->send_data(msg, msg_size);257}258 259static bool recv_msg(socket_ptr sock, void * msg, size_t msg_size) {260    uint64_t size;261    if (!sock->recv_data(&size, sizeof(size))) {262        return false;263    }264    if (size != msg_size) {265        return false;266    }267    return sock->recv_data(msg, msg_size);268}269 270static bool recv_msg(socket_ptr sock, std::vector<uint8_t> & input) {271    uint64_t size;272    if (!sock->recv_data(&size, sizeof(size))) {273        return false;274    }275    try {276        input.resize(size);277    } catch (const std::bad_alloc & e) {278        GGML_LOG_ERROR("Failed to allocate input buffer of size %" PRIu64 "\n", size);279        return false;280    }281    return sock->recv_data(input.data(), size);282}283 284static bool parse_endpoint(const std::string & endpoint, std::string & host, int & port) {285    size_t pos = endpoint.find(':');286    if (pos == std::string::npos) {287        return false;288    }289    host = endpoint.substr(0, pos);290    try {291        port = std::stoi(endpoint.substr(pos + 1));292    } catch (...) {293        return false;294    }295    return true;296}297 298// RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) |299// No response300static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) {301    uint8_t cmd_byte = cmd;302    if (!sock->send_data(&cmd_byte, sizeof(cmd_byte))) {303        return false;304    }305    if (!sock->send_data(&input_size, sizeof(input_size))) {306        return false;307    }308    if (!sock->send_data(input, input_size)) {309        return false;310    }311    return true;312}313 314// RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) |315// RPC response: | response_size (8 bytes) | response_data (response_size bytes) |316static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size, void * output, size_t output_size) {317    if (!send_rpc_cmd(sock, cmd, input, input_size)) {318        return false;319    }320    uint64_t out_size;321    if (!sock->recv_data(&out_size, sizeof(out_size))) {322        return false;323    }324    if (out_size != output_size) {325        return false;326    }327    if (!sock->recv_data(output, output_size)) {328        return false;329    }330    return true;331}332 333// RPC client-side implementation334 335// Performs HELLO handshake with transport auto-negotiation.336// Advertises local capabilities via conn_caps; if the server responds with337// matching capabilities, the socket is upgraded transparently.338static bool negotiate_hello(const std::shared_ptr<socket_t> & sock) {339    rpc_msg_hello_req request = {};340    rpc_msg_hello_rsp response = {};341 342    sock->get_caps(request.conn_caps);343 344    bool status = send_rpc_cmd(sock, RPC_CMD_HELLO, &request, sizeof(request), &response, sizeof(response));345    RPC_STATUS_ASSERT(status);346 347    if (response.major != RPC_PROTO_MAJOR_VERSION || response.minor > RPC_PROTO_MINOR_VERSION) {348        GGML_LOG_ERROR("RPC server version mismatch: %d.%d.%d\n",349                       response.major, response.minor, response.patch);350        return false;351    }352 353    sock->update_caps(response.conn_caps);354    return true;355}356 357static std::shared_ptr<socket_t> get_socket(const std::string & endpoint) {358    static std::mutex mutex;359    std::lock_guard<std::mutex> lock(mutex);360    static std::unordered_map<std::string, std::weak_ptr<socket_t>> sockets;361 362    auto it = sockets.find(endpoint);363    if (it != sockets.end()) {364        if (auto sock = it->second.lock()) {365            return sock;366        }367    }368    std::string host;369    int port;370    if (!parse_endpoint(endpoint, host, port)) {371        GGML_LOG_ERROR("Failed to parse endpoint: %s\n", endpoint.c_str());372        return nullptr;373    }374 375    if (!rpc_transport_init()) {376        return nullptr;377    }378    auto sock = socket_t::connect(host.c_str(), port);379    if (sock == nullptr) {380        return nullptr;381    }382    if (!negotiate_hello(sock)) {383        return nullptr;384    }385    LOG_DBG("[%s] connected to %s\n", __func__, endpoint.c_str());386    sockets[endpoint] = sock;387    return sock;388}389 390static void ggml_backend_rpc_buffer_free_buffer(ggml_backend_buffer_t buffer) {391    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;392    rpc_msg_free_buffer_req request = {ctx->remote_ptr};393    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_FREE_BUFFER, &request, sizeof(request), nullptr, 0);394    RPC_STATUS_ASSERT(status);395    delete ctx;396}397 398static void * ggml_backend_rpc_buffer_get_base(ggml_backend_buffer_t buffer) {399    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;400    if (ctx->base_ptr != nullptr) {401        return ctx->base_ptr;402    }403    rpc_msg_buffer_get_base_req request = {ctx->remote_ptr};404    rpc_msg_buffer_get_base_rsp response;405    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_BUFFER_GET_BASE, &request, sizeof(request), &response, sizeof(response));406    RPC_STATUS_ASSERT(status);407    ctx->base_ptr = reinterpret_cast<void *>(response.base_ptr);408    return ctx->base_ptr;409}410 411static bool ggml_backend_buffer_is_rpc(ggml_backend_buffer_t buffer) {412    return buffer->iface.free_buffer == ggml_backend_rpc_buffer_free_buffer;413}414 415static rpc_tensor serialize_tensor(const ggml_tensor * tensor) {416    rpc_tensor result;417    if (!tensor) {418        memset(&result, 0, sizeof(result));419        return result;420    }421 422    result.id = reinterpret_cast<uint64_t>(tensor);423    result.type = tensor->type;424    if (tensor->buffer && ggml_backend_buffer_is_rpc(tensor->buffer)) {425        ggml_backend_buffer_t buffer = tensor->buffer;426        ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;427        result.buffer = ctx != nullptr ? ctx->remote_ptr : 0;428        result.data = reinterpret_cast<uint64_t>(tensor->data);429    } else {430        result.buffer = 0;431        result.data   = 0;432    }433    for (uint32_t i = 0; i < GGML_MAX_DIMS; i++) {434        result.ne[i] = tensor->ne[i];435        result.nb[i] = tensor->nb[i];436    }437    result.op = tensor->op;438    for (uint32_t i = 0; i < GGML_MAX_OP_PARAMS / sizeof(int32_t); i++) {439        result.op_params[i] = tensor->op_params[i];440    }441    result.flags = tensor->flags;442    for (uint32_t i = 0; i < GGML_MAX_SRC; i++) {443        result.src[i] = reinterpret_cast<uint64_t>(tensor->src[i]);444    }445    result.view_src = reinterpret_cast<uint64_t>(tensor->view_src);446    result.view_offs = tensor->view_offs;447 448    // Avoid sending uninitialized data over the wire449    memset(result.name, 0, sizeof(result.name));450    memset(result.padding, 0, sizeof(result.padding));451 452    snprintf(result.name, GGML_MAX_NAME, "%s", tensor->name);453    return result;454}455 456static enum ggml_status ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) {457    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;458 459    // CUDA backend on the server pads everything to 512 due to CUDA limitations.460    // Due to bandwidth constraints, we only call the server init tensor functions if necessary.461    // In particular, only quantized tensors need padding462    if (ggml_is_quantized(tensor->type) && (tensor->ne[0] % 512 != 0) && (tensor->view_src == nullptr)) {463        rpc_msg_init_tensor_req request;464 465        request.tensor = serialize_tensor(tensor);466 467        bool status = send_rpc_cmd(ctx->sock, RPC_CMD_INIT_TENSOR, &request, sizeof(request), nullptr, 0);468        RPC_STATUS_ASSERT(status);469    }470    return GGML_STATUS_SUCCESS;471}472 473static void ggml_backend_rpc_buffer_memset_tensor(474        ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {475    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;476    rpc_msg_memset_tensor_req request = {477        /* .tensor = */ serialize_tensor(tensor),478        /* .offset = */ offset,479        /* .size   = */ size,480        /* .value  = */ value,481    };482    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_MEMSET_TENSOR, &request, sizeof(request), nullptr, 0);483    RPC_STATUS_ASSERT(status);484}485 486static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {487    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;488    rpc_tensor rpc_tensor = serialize_tensor(tensor);489    if (size > HASH_THRESHOLD) {490        rpc_msg_set_tensor_hash_req request;491        request.tensor = rpc_tensor;492        request.offset = offset;493        request.hash = fnv_hash((const uint8_t*)data, size);494        rpc_msg_set_tensor_hash_rsp response;495        bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR_HASH, &request, sizeof(request), &response, sizeof(response));496        RPC_STATUS_ASSERT(status);497        if (response.result) {498            // the server has the same data, no need to send it499            return;500        }501    }502    // input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes)503    size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + size;504    std::vector<uint8_t> input(input_size, 0);505    memcpy(input.data(), &rpc_tensor, sizeof(rpc_tensor));506    memcpy(input.data() + sizeof(rpc_tensor), &offset, sizeof(offset));507    memcpy(input.data() + sizeof(rpc_tensor) + sizeof(offset), data, size);508    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR, input.data(), input.size());509    RPC_STATUS_ASSERT(status);510}511 512static void ggml_backend_rpc_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {513    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;514    rpc_msg_get_tensor_req request;515    request.tensor = serialize_tensor(tensor);516    request.offset = offset;517    request.size = size;518    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_GET_TENSOR, &request, sizeof(request), data, size);519    RPC_STATUS_ASSERT(status);520}521 522static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) {523    if (ggml_backend_buffer_is_rpc(src->buffer)) {524        // check if src and dst are on the same server525        ggml_backend_buffer_t src_buffer = src->buffer;526        ggml_backend_rpc_buffer_context * src_ctx = (ggml_backend_rpc_buffer_context *)src_buffer->context;527        ggml_backend_buffer_t dst_buffer = dst->buffer;528        ggml_backend_rpc_buffer_context * dst_ctx = (ggml_backend_rpc_buffer_context *)dst_buffer->context;529        if (src_ctx->sock != dst_ctx->sock) {530            return false;531        }532        ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;533        rpc_msg_copy_tensor_req request;534        request.src = serialize_tensor(src);535        request.dst = serialize_tensor(dst);536        rpc_msg_copy_tensor_rsp response;537        bool status = send_rpc_cmd(ctx->sock, RPC_CMD_COPY_TENSOR, &request, sizeof(request), &response, sizeof(response));538        RPC_STATUS_ASSERT(status);539        return response.result;540    }541    return false;542}543 544static void ggml_backend_rpc_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {545    ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;546    rpc_msg_buffer_clear_req request = {ctx->remote_ptr, value};547    bool status = send_rpc_cmd(ctx->sock, RPC_CMD_BUFFER_CLEAR, &request, sizeof(request), nullptr, 0);548    RPC_STATUS_ASSERT(status);549}550 551static ggml_backend_buffer_i ggml_backend_rpc_buffer_interface = {552    /* .free_buffer     = */ ggml_backend_rpc_buffer_free_buffer,553    /* .get_base        = */ ggml_backend_rpc_buffer_get_base,554    /* .init_tensor     = */ ggml_backend_rpc_buffer_init_tensor,555    /* .memset_tensor   = */ ggml_backend_rpc_buffer_memset_tensor,556    /* .set_tensor      = */ ggml_backend_rpc_buffer_set_tensor,557    /* .get_tensor      = */ ggml_backend_rpc_buffer_get_tensor,558    /* .set_tensor_2d   = */ NULL,559    /* .get_tensor_2d   = */ NULL,560    /* .cpy_tensor      = */ ggml_backend_rpc_buffer_cpy_tensor,561    /* .clear           = */ ggml_backend_rpc_buffer_clear,562    /* .reset           = */ NULL,563};564 565static const char * ggml_backend_rpc_buffer_type_name(ggml_backend_buffer_type_t buft) {566    ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context;567    return buft_ctx->name.c_str();568}569 570static ggml_backend_buffer_t ggml_backend_rpc_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {571    ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context;572    rpc_msg_alloc_buffer_req request = {buft_ctx->device, size};573    rpc_msg_alloc_buffer_rsp response;574    auto sock = get_socket(buft_ctx->endpoint);575    bool status = send_rpc_cmd(sock, RPC_CMD_ALLOC_BUFFER, &request, sizeof(request), &response, sizeof(response));576    RPC_STATUS_ASSERT(status);577    if (response.remote_ptr != 0) {578        ggml_backend_buffer_t buffer = ggml_backend_buffer_init(buft,579            ggml_backend_rpc_buffer_interface,580            new ggml_backend_rpc_buffer_context{sock, nullptr, response.remote_ptr},581            response.remote_size);582        return buffer;583    } else {584        return nullptr;585    }586}587 588static size_t get_alignment(const std::shared_ptr<socket_t> & sock, uint32_t device) {589    rpc_msg_get_alignment_req request = {device};590    rpc_msg_get_alignment_rsp response;591    bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALIGNMENT, &request, sizeof(request), &response, sizeof(response));592    RPC_STATUS_ASSERT(status);593    return response.alignment;594}595 596static size_t ggml_backend_rpc_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {597    ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context;598    return buft_ctx->alignment;599}600 601static size_t get_max_size(const std::shared_ptr<socket_t> & sock, uint32_t device) {602    rpc_msg_get_max_size_req request = {device};603    rpc_msg_get_max_size_rsp response;604    bool status = send_rpc_cmd(sock, RPC_CMD_GET_MAX_SIZE, &request, sizeof(request), &response, sizeof(response));605    RPC_STATUS_ASSERT(status);606    return response.max_size;607}608 609static size_t ggml_backend_rpc_get_max_size(ggml_backend_buffer_type_t buft) {610    ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context;611    return buft_ctx->max_size;612}613 614static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {615    // should we query the remote server for the actual size616    bool rpc_get = false;617 618    // See comments in init_tensor.619    rpc_get |= ggml_is_quantized(tensor->type) && (tensor->ne[0] % 512 != 0) && (tensor->view_src == nullptr);620 621    // ops that require additional memory for fleeting data on certain backends622    // ref: https://github.com/ggml-org/llama.cpp/pull/15966623    rpc_get |= tensor->op == GGML_OP_FLASH_ATTN_EXT;624    rpc_get |= tensor->op == GGML_OP_MUL_MAT_ID;625 626    if (rpc_get) {627        ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context;628        auto sock = get_socket(buft_ctx->endpoint);629 630        rpc_msg_get_alloc_size_req request = {631            /*.device =*/ buft_ctx->device,632            /*.tensor =*/ serialize_tensor(tensor),633            /*.srcs   =*/ {},634        };635 636        // .get_alloc_size could be a function of the tensor's srcs, so we must serialize them as well637        for (int i = 0; i < GGML_MAX_SRC; i++) {638            request.srcs[i] = serialize_tensor(tensor->src[i]);639        }640 641        // TODO: cache the alloc responses to avoid extra RPC calls?642        rpc_msg_get_alloc_size_rsp response;643        bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALLOC_SIZE, &request, sizeof(request), &response, sizeof(response));644        RPC_STATUS_ASSERT(status);645 646        return response.alloc_size;647    }648 649    return ggml_nbytes(tensor);650}651 652static ggml_backend_buffer_type_i ggml_backend_rpc_buffer_type_interface = {653    /* .get_name         = */ ggml_backend_rpc_buffer_type_name,654    /* .alloc_buffer     = */ ggml_backend_rpc_buffer_type_alloc_buffer,655    /* .get_alignment    = */ ggml_backend_rpc_buffer_type_get_alignment,656    /* .get_max_size     = */ ggml_backend_rpc_get_max_size,657    /* .get_alloc_size   = */ ggml_backend_rpc_buffer_type_get_alloc_size,658    /* .is_host          = */ NULL,659};660 661static const char * ggml_backend_rpc_name(ggml_backend_t backend) {662    ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context;663 664    return rpc_ctx->name.c_str();665}666 667static void ggml_backend_rpc_free(ggml_backend_t backend) {668    ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context;669    delete rpc_ctx;670    delete backend;671}672 673static void ggml_backend_rpc_synchronize(ggml_backend_t backend) {674    GGML_UNUSED(backend);675    // this is no-op because we don't have any async operations676}677 678static void add_tensor(ggml_tensor * tensor, std::vector<rpc_tensor> & tensors, std::unordered_set<ggml_tensor*> & visited) {679    if (tensor == nullptr) {680        return;681    }682    if (visited.find(tensor) != visited.end()) {683        return;684    }685    visited.insert(tensor);686    for (int i = 0; i < GGML_MAX_SRC; i++) {687        add_tensor(tensor->src[i], tensors, visited);688    }689    add_tensor(tensor->view_src, tensors, visited);690    tensors.push_back(serialize_tensor(tensor));691}692 693static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::vector<uint8_t> & output) {694    uint32_t n_nodes = cgraph->n_nodes;695    std::vector<rpc_tensor> tensors;696    std::unordered_set<ggml_tensor*> visited;697    for (uint32_t i = 0; i < n_nodes; i++) {698        add_tensor(cgraph->nodes[i], tensors, visited);699    }700    // serialization format:701    // | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) |702    uint32_t n_tensors = tensors.size();703    int output_size = 2*sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor);704    output.resize(output_size, 0);705    uint8_t * dest = output.data();706    memcpy(dest, &device, sizeof(device));707    dest += sizeof(device);708    memcpy(dest, &n_nodes, sizeof(n_nodes));709    dest += sizeof(n_nodes);710    for (uint32_t i = 0; i < n_nodes; i++) {711        memcpy(dest + i * sizeof(uint64_t), &cgraph->nodes[i], sizeof(uint64_t));712    }713    dest += n_nodes * sizeof(uint64_t);714    memcpy(dest, &n_tensors, sizeof(n_tensors));715    dest += sizeof(n_tensors);716    rpc_tensor * out_tensors = (rpc_tensor *)dest;717    memcpy(out_tensors, tensors.data(), n_tensors * sizeof(rpc_tensor));718}719 720static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {721    ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context;722    ggml_backend_dev_t rpc_dev = ggml_backend_get_device(backend);723    ggml_backend_rpc_device_context * rpc_dev_ctx = (ggml_backend_rpc_device_context *)rpc_dev->context;724 725    GGML_ASSERT(cgraph->n_nodes > 0);726    bool reuse = cgraph->uid != 0 && rpc_dev_ctx->last_graph_uid == cgraph->uid;727    if (reuse) {728        rpc_msg_graph_recompute_req request;729        request.device = rpc_ctx->device;730        auto sock = get_socket(rpc_ctx->endpoint);731        bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request));732        RPC_STATUS_ASSERT(status);733    } else {734        rpc_dev_ctx->last_graph_uid = cgraph->uid;735        std::vector<uint8_t> input;736        serialize_graph(rpc_ctx->device, cgraph, input);737        auto sock = get_socket(rpc_ctx->endpoint);738        bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size());739        RPC_STATUS_ASSERT(status);740    }741    return GGML_STATUS_SUCCESS;742}743 744static ggml_backend_i ggml_backend_rpc_interface = {745    /* .get_name                = */ ggml_backend_rpc_name,746    /* .free                    = */ ggml_backend_rpc_free,747    /* .set_tensor_async        = */ NULL,748    /* .get_tensor_async        = */ NULL,749    /* .set_tensor_2d_async     = */ NULL,750    /* .get_tensor_2d_async     = */ NULL,751    /* .cpy_tensor_async        = */ NULL,752    /* .synchronize             = */ ggml_backend_rpc_synchronize,753    /* .graph_plan_create       = */ NULL,754    /* .graph_plan_free         = */ NULL,755    /* .graph_plan_update       = */ NULL,756    /* .graph_plan_compute      = */ NULL,757    /* .graph_compute           = */ ggml_backend_rpc_graph_compute,758    /* .event_record            = */ NULL,759    /* .event_wait              = */ NULL,760    /* .graph_optimize          = */ NULL,761};762 763ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, uint32_t device) {764    static std::mutex mutex;765    std::lock_guard<std::mutex> lock(mutex);766    std::string buft_name = "RPC" + std::to_string(device) + "[" + std::string(endpoint) + "]";767    // NOTE: buffer types are allocated and never freed; this is by design768    static std::unordered_map<std::string, ggml_backend_buffer_type_t> buft_map;769    auto it = buft_map.find(buft_name);770    if (it != buft_map.end()) {771        return it->second;772    }773    auto sock = get_socket(endpoint);774    if (sock == nullptr) {775        GGML_LOG_ERROR("Failed to connect to %s\n", endpoint);776        return nullptr;777    }778    size_t alignment = get_alignment(sock, device);779    size_t max_size = get_max_size(sock, device);780    ggml_backend_rpc_buffer_type_context * buft_ctx = new ggml_backend_rpc_buffer_type_context {781        /* .endpoint  = */ endpoint,782        /* .device    = */ device,783        /* .name      = */ buft_name,784        /* .alignment = */ alignment,785        /* .max_size  = */ max_size786    };787    auto reg = ggml_backend_rpc_add_server(endpoint);788    ggml_backend_buffer_type_t buft = new ggml_backend_buffer_type {789        /* .iface   = */ ggml_backend_rpc_buffer_type_interface,790        /* .device  = */ ggml_backend_reg_dev_get(reg, device),791        /* .context = */ buft_ctx792    };793    buft_map[buft_name] = buft;794    return buft;795}796 797ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device) {798    std::string dev_name = "RPC" + std::to_string(device) + "[" + std::string(endpoint) + "]";799    ggml_backend_rpc_context * ctx = new ggml_backend_rpc_context {800        /* .endpoint       = */ endpoint,801        /* .device         = */ device,802        /* .name           = */ dev_name,803    };804    auto reg = ggml_backend_rpc_add_server(endpoint);805    ggml_backend_t backend = new ggml_backend {806        /* .guid    = */ ggml_backend_rpc_guid(),807        /* .iface   = */ ggml_backend_rpc_interface,808        /* .device  = */ ggml_backend_reg_dev_get(reg, device),809        /* .context = */ ctx810    };811    return backend;812}813 814bool ggml_backend_is_rpc(ggml_backend_t backend) {815    return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_rpc_guid());816}817 818static void get_device_memory(const std::shared_ptr<socket_t> & sock, uint32_t device, size_t * free, size_t * total) {819    rpc_msg_get_device_memory_req request;820    request.device = device;821    rpc_msg_get_device_memory_rsp response;822    bool status = send_rpc_cmd(sock, RPC_CMD_GET_DEVICE_MEMORY, &request, sizeof(request), &response, sizeof(response));823    RPC_STATUS_ASSERT(status);824    *free = response.free_mem;825    *total = response.total_mem;826}827 828void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total) {829    auto sock = get_socket(endpoint);830    if (sock == nullptr) {831        *free = 0;832        *total = 0;833        return;834    }835    get_device_memory(sock, device, free, total);836}837 838// RPC server-side implementation839 840class rpc_server {841public:842    rpc_server(std::vector<ggml_backend_t> all_backends, const char * cache_dir)843        : backends(std::move(all_backends)), cache_dir(cache_dir) {844        stored_graphs.resize(backends.size());845    }846    ~rpc_server();847 848    void hello(rpc_msg_hello_rsp & response);849    bool alloc_buffer(const rpc_msg_alloc_buffer_req & request, rpc_msg_alloc_buffer_rsp & response);850    bool get_alignment(const rpc_msg_get_alignment_req & request, rpc_msg_get_alignment_rsp & response);851    bool get_max_size(const rpc_msg_get_max_size_req & request, rpc_msg_get_max_size_rsp & response);852    bool buffer_get_base(const rpc_msg_buffer_get_base_req & request, rpc_msg_buffer_get_base_rsp & response);853    bool free_buffer(const rpc_msg_free_buffer_req & request);854    bool buffer_clear(const rpc_msg_buffer_clear_req & request);855    bool memset_tensor(const rpc_msg_memset_tensor_req & request);856    bool set_tensor(const std::vector<uint8_t> & input);857    bool set_tensor_hash(const rpc_msg_set_tensor_hash_req & request, rpc_msg_set_tensor_hash_rsp & response);858    bool get_tensor(const rpc_msg_get_tensor_req & request, std::vector<uint8_t> & response);859    bool copy_tensor(const rpc_msg_copy_tensor_req & request, rpc_msg_copy_tensor_rsp & response);860    bool graph_compute(const std::vector<uint8_t> & input);861    bool graph_recompute(const rpc_msg_graph_recompute_req & request);862    bool init_tensor(const rpc_msg_init_tensor_req & request);863    bool get_alloc_size(const rpc_msg_get_alloc_size_req & request, rpc_msg_get_alloc_size_rsp & response);864    bool get_device_memory(const rpc_msg_get_device_memory_req & request, rpc_msg_get_device_memory_rsp & response);865 866    struct stored_graph {867        std::vector<uint8_t>   buffer;868        ggml_cgraph          * graph;869    };870 871private:872    bool get_cached_file(uint64_t hash, std::vector<uint8_t> & data);873    ggml_tensor * deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor);874    ggml_tensor * create_node(uint64_t id,875                              struct ggml_context * ctx,876                              const std::unordered_map<uint64_t, const rpc_tensor*> & tensor_ptrs,877                              std::unordered_map<uint64_t, struct ggml_tensor*> & tensor_map);878 879 880    std::vector<ggml_backend_t> backends;881    const char * cache_dir;882    std::unordered_set<ggml_backend_buffer_t> buffers;883    // store the last computed graph for each backend884    std::vector<stored_graph> stored_graphs;885};886 887void rpc_server::hello(rpc_msg_hello_rsp & response) {888    response.major = RPC_PROTO_MAJOR_VERSION;889    response.minor = RPC_PROTO_MINOR_VERSION;890    response.patch = RPC_PROTO_PATCH_VERSION;891    LOG_DBG("[%s] version: %d.%d.%d\n", __func__, response.major, response.minor, response.patch);892}893 894bool rpc_server::get_alloc_size(const rpc_msg_get_alloc_size_req & request, rpc_msg_get_alloc_size_rsp & response) {895    uint32_t dev_id = request.device;896    if (dev_id >= backends.size()) {897        return false;898    }899    ggml_backend_buffer_type_t buft;900    struct ggml_init_params params {901        /*.mem_size   =*/ ggml_tensor_overhead()*(1 + GGML_MAX_SRC),902        /*.mem_buffer =*/ NULL,903        /*.no_alloc   =*/ true,904    };905 906    ggml_context_ptr ctx_ptr { ggml_init(params) };907    GGML_ASSERT(ctx_ptr != nullptr);908    ggml_context * ctx = ctx_ptr.get();909 910    ggml_tensor * tensor = deserialize_tensor(ctx, &request.tensor);911    if (tensor == nullptr) {912        GGML_LOG_ERROR("Null tensor pointer passed to server get_alloc_size function.\n");913        return false;914    }915    for (int i = 0; i < GGML_MAX_SRC; i++) {916        if (request.srcs[i].id != 0) {917            tensor->src[i] = deserialize_tensor(ctx, &request.srcs[i]);918        }919    }920 921    LOG_DBG("[%s] device: %d, buffer: %p, data: %p\n", __func__, dev_id, (void*)tensor->buffer, tensor->data);922    if (tensor->buffer == nullptr) {923        //No buffer allocated.924        buft = ggml_backend_get_default_buffer_type(backends[dev_id]);925    } else {926        buft = tensor->buffer->buft;927    }928 929    response.alloc_size = ggml_backend_buft_get_alloc_size(buft, tensor);930 931    return true;932}933 934bool rpc_server::alloc_buffer(const rpc_msg_alloc_buffer_req & request, rpc_msg_alloc_buffer_rsp & response) {935    uint32_t dev_id = request.device;936    if (dev_id >= backends.size()) {937        return false;938    }939    ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backends[dev_id]);940    ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, request.size);941    response.remote_ptr = 0;942    response.remote_size = 0;943    if (buffer != nullptr) {944        response.remote_ptr = reinterpret_cast<uint64_t>(buffer);945        response.remote_size = buffer->size;946        LOG_DBG("[%s] device: %d, size: %" PRIu64 " -> remote_ptr: %" PRIx64 ", remote_size: %" PRIu64 "\n",947            __func__, dev_id, request.size, response.remote_ptr, response.remote_size);948        buffers.insert(buffer);949    } else {950        LOG_DBG("[%s] device: %d, size: %" PRIu64 " -> failed\n", __func__, dev_id, request.size);951    }952    return true;953}954 955bool rpc_server::get_alignment(const rpc_msg_get_alignment_req & request, rpc_msg_get_alignment_rsp & response) {956    uint32_t dev_id = request.device;957    if (dev_id >= backends.size()) {958        return false;959    }960    ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backends[dev_id]);961    size_t alignment = ggml_backend_buft_get_alignment(buft);962    LOG_DBG("[%s] device: %d, alignment: %lu\n", __func__, dev_id, alignment);963    response.alignment = alignment;964    return true;965}966 967bool rpc_server::get_max_size(const rpc_msg_get_max_size_req & request, rpc_msg_get_max_size_rsp & response) {968    uint32_t dev_id = request.device;969    if (dev_id >= backends.size()) {970        return false;971    }972    ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backends[dev_id]);973    size_t max_size = ggml_backend_buft_get_max_size(buft);974    LOG_DBG("[%s] device: %d, max_size: %lu\n", __func__, dev_id, max_size);975    response.max_size = max_size;976    return true;977}978 979bool rpc_server::buffer_get_base(const rpc_msg_buffer_get_base_req & request, rpc_msg_buffer_get_base_rsp & response) {980    LOG_DBG("[%s] remote_ptr: %" PRIx64 "\n", __func__, request.remote_ptr);981    ggml_backend_buffer_t buffer = reinterpret_cast<ggml_backend_buffer_t>(request.remote_ptr);982    if (buffers.find(buffer) == buffers.end()) {983        GGML_LOG_ERROR("[%s] buffer not found\n", __func__);984        return false;985    }986    void * base = ggml_backend_buffer_get_base(buffer);987    response.base_ptr = reinterpret_cast<uint64_t>(base);988    return true;989}990 991bool rpc_server::free_buffer(const rpc_msg_free_buffer_req & request) {992    LOG_DBG("[%s] remote_ptr: %" PRIx64 "\n", __func__, request.remote_ptr);993    ggml_backend_buffer_t buffer = reinterpret_cast<ggml_backend_buffer_t>(request.remote_ptr);994    if (buffers.find(buffer) == buffers.end()) {995        GGML_LOG_ERROR("[%s] buffer not found\n", __func__);996        return false;997    }998    ggml_backend_buffer_free(buffer);999    buffers.erase(buffer);1000    return true;1001}1002 1003bool rpc_server::buffer_clear(const rpc_msg_buffer_clear_req & request) {1004    LOG_DBG("[%s] remote_ptr: %" PRIx64 ", value: %u\n", __func__, request.remote_ptr, request.value);1005    ggml_backend_buffer_t buffer = reinterpret_cast<ggml_backend_buffer_t>(request.remote_ptr);1006    if (buffers.find(buffer) == buffers.end()) {1007        GGML_LOG_ERROR("[%s] buffer not found\n", __func__);1008        return false;1009    }1010    ggml_backend_buffer_clear(buffer, request.value);1011    return true;1012}1013 1014bool rpc_server::memset_tensor(const rpc_msg_memset_tensor_req & request) {1015    struct ggml_init_params params {1016        /*.mem_size   =*/ ggml_tensor_overhead(),1017        /*.mem_buffer =*/ NULL,1018        /*.no_alloc   =*/ true,1019    };1020    ggml_context_ptr ctx_ptr { ggml_init(params) };1021    GGML_ASSERT(ctx_ptr != nullptr);1022    ggml_context * ctx = ctx_ptr.get();1023    ggml_tensor * tensor = deserialize_tensor(ctx, &request.tensor);1024    if (tensor == nullptr || tensor->buffer == nullptr) {1025        GGML_LOG_ERROR("[%s] error deserializing tensor\n", __func__);1026        return false;1027    }1028 1029    const uint64_t tensor_size = ggml_nbytes(tensor);1030    if (request.offset > tensor_size || request.size > tensor_size - request.offset) {1031        GGML_LOG_ERROR("[%s] tensor region (offset=%" PRIu64 ", size=%" PRIu64 ") out of tensor bounds [0, %" PRIu64 ")\n",1032                       __func__, request.offset, request.size, tensor_size);1033        return false;1034    }1035 1036    const uint64_t buffer_start = (uint64_t) ggml_backend_buffer_get_base(tensor->buffer);1037    const uint64_t buffer_size = ggml_backend_buffer_get_size(tensor->buffer);1038    if (request.tensor.data < buffer_start) {1039        GGML_LOG_ERROR("[%s] tensor data before buffer start\n", __func__);1040        return false;1041    }1042    const uint64_t data_offset = request.tensor.data - buffer_start;1043    if (data_offset > buffer_size ||1044        request.offset > buffer_size - data_offset ||1045        request.size > buffer_size - data_offset - request.offset) {1046        GGML_LOG_ERROR("[%s] tensor region out of buffer bounds\n", __func__);1047        return false;1048    }1049    if (tensor->buffer->iface.memset_tensor == nullptr) {1050        GGML_LOG_ERROR("[%s] memset not implemented by backend buffer\n", __func__);1051        return false;1052    }1053 1054    LOG_DBG("[%s] buffer: %p, data: %p, offset: %" PRIu64 ", size: %" PRIu64 ", value: %u\n",1055            __func__, (void *) tensor->buffer, tensor->data, request.offset, request.size, request.value);1056    ggml_backend_tensor_memset(tensor, request.value, request.offset, request.size);1057    return true;1058}1059 1060ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor) {1061    // Validate tensor type before using it1062    if (tensor->type >= GGML_TYPE_COUNT) {1063        GGML_LOG_ERROR("[%s] invalid tensor type received: %u\n", __func__, tensor->type);1064        return nullptr;1065    }1066 1067    // Fix: Prevent division by zero if blck_size is 0 (e.g., deprecated types)1068    if (ggml_blck_size((enum ggml_type)tensor->type) == 0) {1069        GGML_LOG_ERROR("[%s] invalid tensor type received (blck_size is 0): %u\n", __func__, tensor->type);1070        return nullptr;1071    }1072 1073    ggml_tensor * result = ggml_new_tensor_4d(ctx, (ggml_type) tensor->type,1074        tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]);1075 1076    // ggml_new_tensor_4d might fail if dimensions are invalid, although less likely to crash than invalid type1077    if (result == nullptr) {1078        GGML_LOG_ERROR("[%s] ggml_new_tensor_4d failed for type %u\n", __func__, tensor->type);1079        return nullptr;1080    }1081 1082    for (uint32_t i = 0; i < GGML_MAX_DIMS; i++) {1083        result->nb[i] = tensor->nb[i];1084    }1085    result->buffer = reinterpret_cast<ggml_backend_buffer_t>(tensor->buffer);1086    if (result->buffer && buffers.find(result->buffer) == buffers.end()) {1087        result->buffer = nullptr;1088    }1089 1090    if (result->buffer) {1091        // require that the tensor data does not go beyond the buffer end1092        uint64_t tensor_size = (uint64_t) ggml_nbytes(result);1093        uint64_t buffer_start = (uint64_t) ggml_backend_buffer_get_base(result->buffer);1094        uint64_t buffer_size = (uint64_t) ggml_backend_buffer_get_size(result->buffer);1095        GGML_ASSERT(tensor->data + tensor_size >= tensor->data); // check for overflow1096        GGML_ASSERT(tensor->data >= buffer_start && tensor->data + tensor_size <= buffer_start + buffer_size);1097    }1098 1099    result->op = (ggml_op) tensor->op;1100    for (uint32_t i = 0; i < GGML_MAX_OP_PARAMS / sizeof(int32_t); i++) {1101        result->op_params[i] = tensor->op_params[i];1102    }1103    result->flags = tensor->flags;1104    result->data = reinterpret_cast<void *>(tensor->data);1105    ggml_set_name(result, tensor->name);1106    return result;1107}1108 1109 1110bool rpc_server::set_tensor(const std::vector<uint8_t> & input) {1111    // serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes) |1112    if (input.size() < sizeof(rpc_tensor) + sizeof(uint64_t)) {1113        return false;1114    }1115    const rpc_tensor * in_tensor = (const rpc_tensor *)input.data();1116    uint64_t offset;1117    memcpy(&offset, input.data() + sizeof(rpc_tensor), sizeof(offset));1118    const size_t size = input.size() - sizeof(rpc_tensor) - sizeof(offset);1119 1120    struct ggml_init_params params {1121        /*.mem_size   =*/ ggml_tensor_overhead(),1122        /*.mem_buffer =*/ NULL,1123        /*.no_alloc   =*/ true,1124    };1125    ggml_context_ptr ctx_ptr { ggml_init(params) };1126    GGML_ASSERT(ctx_ptr != nullptr);1127    ggml_context * ctx = ctx_ptr.get();1128    ggml_tensor * tensor = deserialize_tensor(ctx, in_tensor);1129    if (tensor == nullptr || tensor->buffer == nullptr) {1130        GGML_LOG_ERROR("[%s] error deserializing tensor\n", __func__);1131        return false;1132    }1133    LOG_DBG("[%s] buffer: %p, data: %p, offset: %" PRIu64 ", size: %zu\n", __func__, (void*)tensor->buffer, tensor->data, offset, size);1134 1135    // sanitize tensor->data1136    {1137        const size_t p0 = (size_t) ggml_backend_buffer_get_base(tensor->buffer);1138        const size_t p1 = p0 + ggml_backend_buffer_get_size(tensor->buffer);1139 1140        if (in_tensor->data + offset < p0 || in_tensor->data + offset >= p1 || size > (p1 - in_tensor->data - offset)) {1141            GGML_LOG_ERROR("[%s] tensor data region (data=0x%" PRIx64 ", offset=%" PRIu64 ", size=%zu) out of buffer bounds [0x%zx, 0x%zx)\n",1142                           __func__, in_tensor->data, offset, size, p0, p1);1143            return false;1144        }1145    }1146 1147    const void * data = input.data() + sizeof(rpc_tensor) + sizeof(offset);1148    if (cache_dir && size > HASH_THRESHOLD) {1149        uint64_t hash = fnv_hash((const uint8_t*)data, size);1150        char hash_str[17];1151        snprintf(hash_str, sizeof(hash_str), "%016" PRIx64, hash);1152        // save to cache_dir/hash_str1153        fs::path cache_file = fs::path(cache_dir) / hash_str;1154        std::ofstream ofs(cache_file, std::ios::binary);1155        ofs.write((const char *)data, size);1156        GGML_LOG_INFO("[%s] saved to '%s'\n", __func__, cache_file.string().c_str());1157    }1158    ggml_backend_tensor_set(tensor, data, offset, size);1159    return true;1160}1161 1162bool rpc_server::get_cached_file(uint64_t hash, std::vector<uint8_t> & data) {1163    if (!cache_dir) {1164        return false;1165    }1166    char hash_str[17];1167    snprintf(hash_str, sizeof(hash_str), "%016" PRIx64, hash);1168    fs::path cache_file = fs::path(cache_dir) / hash_str;1169    std::error_code ec;1170    if (!fs::exists(cache_file, ec)) {1171        return false;1172    }1173    std::ifstream ifs(cache_file, std::ios::binary);1174    ifs.seekg(0, std::ios::end);1175    size_t size = ifs.tellg();1176    ifs.seekg(0, std::ios::beg);1177    data.resize(size);1178    ifs.read((char *)data.data(), size);1179    return true;1180}1181 1182bool rpc_server::set_tensor_hash(const rpc_msg_set_tensor_hash_req & request, rpc_msg_set_tensor_hash_rsp & response)1183{1184    std::vector<uint8_t> cached_file;1185    if (!get_cached_file(request.hash, cached_file)) {1186        response.result = 0;1187        return true;1188    }1189    size_t size = cached_file.size();1190    struct ggml_init_params params {1191        /*.mem_size   =*/ ggml_tensor_overhead(),1192        /*.mem_buffer =*/ NULL,1193        /*.no_alloc   =*/ true,1194    };1195    ggml_context_ptr ctx_ptr { ggml_init(params) };1196    GGML_ASSERT(ctx_ptr != nullptr);1197    ggml_context * ctx = ctx_ptr.get();1198    ggml_tensor * tensor = deserialize_tensor(ctx, &request.tensor);1199    if (tensor == nullptr || tensor->buffer == nullptr) {1200        GGML_LOG_ERROR("[%s] error deserializing tensor\n", __func__);

Showing the first 1,200 of 2056 lines. Download the file for the rest.

Brunobkr/llama.cpp_AlgMor24_github · Team Ai