From a7d4c803b1e3b2e6d55cd6d1d29fc156a8198ee7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 6 Sep 2026 02:25:26 -0700 Subject: [PATCH] rpc: cut the cost of uploading weights to a remote server Loading a 27B layer split spends most of its time in the weight upload, and most of that time is protocol, not wire. Measured on one node with an rpc-server on the RoCE interface (Qwen3.8-27B UD-Q4_K_XL, 15.7 GiB pushed to the remote backend, RDMA active), the upload phase cost 41.5 s and broke down as: hashing 12.6 s every tensor over 10 MiB hashed with FNV-1a at about 1.2 GiB/s staging 4.3 s a fresh zero filled buffer per tensor, then a copy into it wire 4.3 s other 20.4 s client side stalls outside the RPC calls Four changes, all inside ggml/src/ggml-rpc: - The server now says at HELLO whether it keeps a tensor cache, in the byte that used to be padding in the response. Without a cache the answer to SET_TENSOR_HASH is always "not cached", so the hash pass over every large tensor was pure cost. The client only hashes when the server can use it. The message keeps its size, an older server sends a zero byte and an older client ignores it, so both directions interoperate unchanged. - SET_TENSOR is written from the header and the caller's payload directly instead of being copied into one contiguous buffer first. That buffer cost a zero fill and a full copy of every tensor. The bytes on the wire are identical. - The server reads a SET_TENSOR message off the connection instead of into a vector sized to the whole message, and receives the payload straight into the destination when the backend buffer is host memory. Non host backends reuse one staging allocation that is never zero filled. - The RDMA transport keeps up to eight 256 KiB chunks in flight instead of posting one and polling it to completion before posting the next, and drains them at the message boundary that flush() already marks. Receives are now consumed byte by byte from the completed buffer, so a peer that frames a message differently keeps working instead of losing the remainder of a frame. After the change the same upload phase is 23.5 s, with hashing and staging at zero and 4.9 s on the wire. GGML_RPC_LOAD_OPT=0 restores the previous behaviour for A/B. GGML_RPC_LOADPROF=1 turns on a load profiler on both sides: per command counts and times on the client, and the split of the upload into hashing, staging, wire and the gaps between calls. It is off by default and costs one relaxed atomic load per call. --- ggml/src/ggml-rpc/ggml-rpc.cpp | 366 +++++++++++++++++++++++++++++--- ggml/src/ggml-rpc/transport.cpp | 131 +++++++++--- ggml/src/ggml-rpc/transport.h | 4 + 3 files changed, 450 insertions(+), 51 deletions(-) diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 69a8a08ae17..bada5134e97 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -17,6 +17,9 @@ #include #include #include +#include +#include +#include static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG"); @@ -26,6 +29,128 @@ static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG"); namespace fs = std::filesystem; +// --------------------------------------------------------------------------- +// Env-gated load profiler (GGML_RPC_LOADPROF=1). Off by default and zero cost +// when off: one relaxed atomic load per set_tensor. It breaks the weight upload +// down into hashing, host staging, wire time and the gap between calls, which is +// what tells a protocol problem apart from a bandwidth problem. +// --------------------------------------------------------------------------- +struct rpc_load_prof { + bool enabled = false; + const char * tag = "client"; + + std::atomic calls{0}; + std::atomic bytes{0}; + std::atomic ns_hash{0}; + std::atomic ns_stage{0}; + std::atomic ns_wire{0}; + std::atomic ns_gap{0}; + std::atomic hash_calls{0}; + std::atomic hash_hits{0}; + std::atomic hash_bytes{0}; + // size histogram: <4K, <64K, <1M, <10M, >=10M + std::atomic hist[5]; + std::atomic last_end_ns{0}; + std::atomic first_ns{0}; + std::atomic gap_max{0}; + std::atomic gaps_1ms{0}; + + rpc_load_prof(const char * tag) : tag(tag) { + const char * e = std::getenv("GGML_RPC_LOADPROF"); + enabled = e != nullptr && e[0] != '0'; + for (int i = 0; i < 5; i++) hist[i].store(0); + } + + static uint64_t now_ns() { + return (uint64_t) std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()).count(); + } + + void bucket(size_t size) { + int b = size < 4096 ? 0 : size < 65536 ? 1 : size < (1u<<20) ? 2 : size < 10u*(1u<<20) ? 3 : 4; + hist[b].fetch_add(1, std::memory_order_relaxed); + } + + void print() { + if (!enabled || calls.load() == 0) { + return; + } + const double mb = (double) bytes.load() / (1024.0*1024.0); + const double gap_s = ns_gap.load() / 1e9; + const double hash_s = ns_hash.load() / 1e9; + const double stage_s = ns_stage.load() / 1e9; + const double wire_s = ns_wire.load() / 1e9; + const double span_s = (last_end_ns.load() - first_ns.load()) / 1e9; + fprintf(stderr, "RPCLOADPROF[%s] set_tensor calls=%" PRIu64 " bytes=%.1f MiB span=%.2fs\n", + tag, calls.load(), mb, span_s); + fprintf(stderr, "RPCLOADPROF[%s] hash=%.2fs (%" PRIu64 " calls, %" PRIu64 " hits, %.1f MiB) stage=%.2fs wire=%.2fs gap=%.2fs\n", + tag, hash_s, hash_calls.load(), hash_hits.load(), + (double) hash_bytes.load()/(1024.0*1024.0), stage_s, wire_s, gap_s); + fprintf(stderr, "RPCLOADPROF[%s] gap max=%.3fs, %" PRIu64 " gaps over 1 ms\n", + tag, gap_max.load()/1e9, gaps_1ms.load()); + fprintf(stderr, "RPCLOADPROF[%s] sizes <4K=%" PRIu64 " <64K=%" PRIu64 " <1M=%" PRIu64 " <10M=%" PRIu64 " >=10M=%" PRIu64 "\n", + tag, hist[0].load(), hist[1].load(), hist[2].load(), hist[3].load(), hist[4].load()); + if (wire_s > 0) { + fprintf(stderr, "RPCLOADPROF[%s] effective wire rate %.2f MiB/s, end to end %.2f MiB/s\n", + tag, mb/wire_s, span_s > 0 ? mb/span_s : 0.0); + } + } + + ~rpc_load_prof() { print(); } +}; + +// GGML_RPC_LOAD_OPT=0 restores the pre-optimisation upload path: always try SET_TENSOR_HASH, +// stage every tensor into one contiguous buffer before sending, one RDMA chunk in flight. +static bool rpc_load_opt() { + static const bool opt = [] { + const char * e = std::getenv("GGML_RPC_LOAD_OPT"); + return !(e && e[0] == '0'); + }(); + return opt; +} + +static rpc_load_prof g_rpc_loadprof_client("client"); +static rpc_load_prof g_rpc_loadprof_server("server"); + +// Per command client side accounting, so the load can be attributed to a command and not just +// to "somewhere in the RPC backend". Same env gate, same zero cost when off. +static const char * rpc_cmd_name(int cmd); + +struct rpc_cmd_prof { + std::atomic calls[64]; + std::atomic ns[64]; + bool enabled = false; + + rpc_cmd_prof() { + const char * e = std::getenv("GGML_RPC_LOADPROF"); + enabled = e != nullptr && e[0] != '0'; + for (int i = 0; i < 64; i++) { calls[i].store(0); ns[i].store(0); } + } + + void add(int cmd, uint64_t dt) { + if (cmd >= 0 && cmd < 64) { + calls[cmd].fetch_add(1, std::memory_order_relaxed); + ns[cmd].fetch_add(dt, std::memory_order_relaxed); + } + } + + ~rpc_cmd_prof() { + if (!enabled) { + return; + } + for (int i = 0; i < 64; i++) { + if (calls[i].load() > 0) { + // entries at 32 and above are the request plus the wait for the response + fprintf(stderr, "RPCLOADPROF[cmd] %-18s%-5s calls=%8" PRIu64 " time=%8.2fs\n", + rpc_cmd_name(i % 32), i >= 32 ? "+rsp" : "", calls[i].load(), ns[i].load()/1e9); + } + } + } +}; + +static rpc_cmd_prof g_rpc_cmdprof; + + // macro for nicer error messages on server crash #define RPC_STATUS_ASSERT(x) if (!(x)) GGML_ABORT("Remote RPC server crashed or returned malformed response") @@ -77,6 +202,30 @@ enum rpc_cmd { static_assert(RPC_CMD_HELLO == 14, "RPC_CMD_HELLO must be always 14"); +static const char * rpc_cmd_name(int cmd) { + switch (cmd) { + case RPC_CMD_ALLOC_BUFFER: return "ALLOC_BUFFER"; + case RPC_CMD_GET_ALIGNMENT: return "GET_ALIGNMENT"; + case RPC_CMD_GET_MAX_SIZE: return "GET_MAX_SIZE"; + case RPC_CMD_BUFFER_GET_BASE: return "BUFFER_GET_BASE"; + case RPC_CMD_FREE_BUFFER: return "FREE_BUFFER"; + case RPC_CMD_BUFFER_CLEAR: return "BUFFER_CLEAR"; + case RPC_CMD_SET_TENSOR: return "SET_TENSOR"; + case RPC_CMD_SET_TENSOR_HASH: return "SET_TENSOR_HASH"; + case RPC_CMD_GET_TENSOR: return "GET_TENSOR"; + case RPC_CMD_COPY_TENSOR: return "COPY_TENSOR"; + case RPC_CMD_GRAPH_COMPUTE: return "GRAPH_COMPUTE"; + case RPC_CMD_GET_DEVICE_MEMORY: return "GET_DEVICE_MEMORY"; + case RPC_CMD_INIT_TENSOR: return "INIT_TENSOR"; + case RPC_CMD_GET_ALLOC_SIZE: return "GET_ALLOC_SIZE"; + case RPC_CMD_HELLO: return "HELLO"; + case RPC_CMD_DEVICE_COUNT: return "DEVICE_COUNT"; + case RPC_CMD_GRAPH_RECOMPUTE: return "GRAPH_RECOMPUTE"; + case RPC_CMD_MEMSET_TENSOR: return "MEMSET_TENSOR"; + default: return "UNKNOWN"; + } +} + // Try RPC_CMD_SET_TENSOR_HASH first when data size is larger than this threshold const size_t HASH_THRESHOLD = 10 * 1024 * 1024; @@ -88,10 +237,15 @@ struct rpc_msg_hello_rsp { uint8_t major; uint8_t minor; uint8_t patch; - uint8_t padding; + // was a padding byte, always zero. A server that predates this reports no features, which + // is the conservative answer, so the message keeps its size and old and new interoperate. + uint8_t flags; uint8_t conn_caps[RPC_CONN_CAPS_SIZE]; }; +// the server keeps a local tensor cache, so RPC_CMD_SET_TENSOR_HASH can save an upload +#define RPC_SRV_FLAG_HAS_CACHE (1 << 0) + struct rpc_msg_device_count_rsp { uint32_t device_count; }; @@ -298,9 +452,44 @@ static bool parse_endpoint(const std::string & endpoint, std::string & host, int return true; } +// Same wire format as send_rpc_cmd below, with the request written from two buffers instead of +// one. It exists so that a tensor upload does not have to be copied into a staging buffer whose +// only purpose is to make the header and the payload contiguous: for a 27B layer split that copy +// is several gigabytes of pure memory traffic, plus the zero fill of the buffer that receives it. +// The bytes on the wire are identical, so a server built before this change sees no difference. +static bool send_rpc_cmd_hdr_payload(socket_ptr sock, enum rpc_cmd cmd, + const void * hdr, size_t hdr_size, + const void * payload, size_t payload_size) { + const uint64_t t0 = g_rpc_cmdprof.enabled ? rpc_load_prof::now_ns() : 0; + struct prof_guard { + enum rpc_cmd cmd; uint64_t t0; + ~prof_guard() { if (g_rpc_cmdprof.enabled) g_rpc_cmdprof.add(cmd, rpc_load_prof::now_ns() - t0); } + } guard{cmd, t0}; + uint8_t cmd_byte = cmd; + uint64_t input_size = hdr_size + payload_size; + if (!sock->send_data(&cmd_byte, sizeof(cmd_byte))) { + return false; + } + if (!sock->send_data(&input_size, sizeof(input_size))) { + return false; + } + if (!sock->send_data(hdr, hdr_size)) { + return false; + } + if (payload_size > 0 && !sock->send_data(payload, payload_size)) { + return false; + } + return sock->flush(); +} + // RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) | // No response static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) { + const uint64_t t0 = g_rpc_cmdprof.enabled ? rpc_load_prof::now_ns() : 0; + struct prof_guard { + enum rpc_cmd cmd; uint64_t t0; + ~prof_guard() { if (g_rpc_cmdprof.enabled) g_rpc_cmdprof.add(cmd, rpc_load_prof::now_ns() - t0); } + } guard{cmd, t0}; uint8_t cmd_byte = cmd; if (!sock->send_data(&cmd_byte, sizeof(cmd_byte))) { return false; @@ -317,6 +506,12 @@ static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, // RPC request : | rpc_cmd (1 byte) | request_size (8 bytes) | request_data (request_size bytes) | // RPC response: | response_size (8 bytes) | response_data (response_size bytes) | static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size, void * output, size_t output_size) { + const uint64_t t0 = g_rpc_cmdprof.enabled ? rpc_load_prof::now_ns() : 0; + struct prof_guard { + enum rpc_cmd cmd; uint64_t t0; + // the inner call books its own send time, so book only the wait for the response here + ~prof_guard() { if (g_rpc_cmdprof.enabled) g_rpc_cmdprof.add(cmd + 32, rpc_load_prof::now_ns() - t0); } + } guard{cmd, t0}; if (!send_rpc_cmd(sock, cmd, input, input_size)) { return false; } @@ -353,6 +548,8 @@ static bool negotiate_hello(const std::shared_ptr & sock) { return false; } + sock->srv_flags = response.flags; + sock->update_caps(response.conn_caps); return true; } @@ -488,27 +685,82 @@ static void ggml_backend_rpc_buffer_memset_tensor( static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; + rpc_load_prof & prof = g_rpc_loadprof_client; + const bool prof_on = prof.enabled; + uint64_t t_call = 0; + if (prof_on) { + t_call = rpc_load_prof::now_ns(); + uint64_t prev = prof.last_end_ns.load(std::memory_order_relaxed); + if (prev != 0) { + const uint64_t g = t_call - prev; + prof.ns_gap.fetch_add(g, std::memory_order_relaxed); + if (g > 1000000) { + prof.gaps_1ms.fetch_add(1, std::memory_order_relaxed); + } + uint64_t m = prof.gap_max.load(std::memory_order_relaxed); + while (g > m && !prof.gap_max.compare_exchange_weak(m, g, std::memory_order_relaxed)) { } + } else { + prof.first_ns.store(t_call, std::memory_order_relaxed); + } + prof.calls.fetch_add(1, std::memory_order_relaxed); + prof.bytes.fetch_add(size, std::memory_order_relaxed); + prof.bucket(size); + } rpc_tensor rpc_tensor = serialize_tensor(tensor); - if (size > HASH_THRESHOLD) { + // Hashing a tensor only pays off if the server keeps a cache to match it against. Without + // one the reply is always "not cached", so the FNV pass over every large tensor is pure + // cost: on a 27B split it hashes gigabytes at about 1.2 GiB/s and then uploads them anyway. + // The server now says at HELLO whether it has a cache. + const bool try_hash = !rpc_load_opt() || (ctx->sock->srv_flags & RPC_SRV_FLAG_HAS_CACHE); + if (size > HASH_THRESHOLD && try_hash) { rpc_msg_set_tensor_hash_req request; request.tensor = rpc_tensor; request.offset = offset; + uint64_t t0 = prof_on ? rpc_load_prof::now_ns() : 0; request.hash = fnv_hash((const uint8_t*)data, size); + if (prof_on) { + prof.ns_hash.fetch_add(rpc_load_prof::now_ns() - t0, std::memory_order_relaxed); + prof.hash_calls.fetch_add(1, std::memory_order_relaxed); + prof.hash_bytes.fetch_add(size, std::memory_order_relaxed); + } rpc_msg_set_tensor_hash_rsp response; + uint64_t t1 = prof_on ? rpc_load_prof::now_ns() : 0; bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR_HASH, &request, sizeof(request), &response, sizeof(response)); + if (prof_on) { + prof.ns_wire.fetch_add(rpc_load_prof::now_ns() - t1, std::memory_order_relaxed); + } RPC_STATUS_ASSERT(status); if (response.result) { // the server has the same data, no need to send it + if (prof_on) { + prof.hash_hits.fetch_add(1, std::memory_order_relaxed); + prof.last_end_ns.store(rpc_load_prof::now_ns(), std::memory_order_relaxed); + } return; } } // input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes) - size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + size; - std::vector input(input_size, 0); - memcpy(input.data(), &rpc_tensor, sizeof(rpc_tensor)); - memcpy(input.data() + sizeof(rpc_tensor), &offset, sizeof(offset)); - memcpy(input.data() + sizeof(rpc_tensor) + sizeof(offset), data, size); - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR, input.data(), input.size()); + uint64_t t2 = prof_on ? rpc_load_prof::now_ns() : 0; + uint8_t hdr[sizeof(rpc_tensor) + sizeof(uint64_t)]; + memcpy(hdr, &rpc_tensor, sizeof(rpc_tensor)); + memcpy(hdr + sizeof(rpc_tensor), &offset, sizeof(offset)); + std::vector input; + if (!rpc_load_opt()) { + // pre-optimisation path, kept for A/B: one contiguous buffer for header and payload + input.resize(sizeof(hdr) + size, 0); + memcpy(input.data(), hdr, sizeof(hdr)); + memcpy(input.data() + sizeof(hdr), data, size); + } + uint64_t t3 = prof_on ? rpc_load_prof::now_ns() : 0; + bool status = input.empty() + ? send_rpc_cmd_hdr_payload(ctx->sock, RPC_CMD_SET_TENSOR, hdr, sizeof(hdr), data, size) + : send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR, input.data(), input.size()); + if (prof_on) { + uint64_t t4 = rpc_load_prof::now_ns(); + prof.ns_stage.fetch_add(t3 - t2, std::memory_order_relaxed); + prof.ns_wire.fetch_add(t4 - t3, std::memory_order_relaxed); + prof.last_end_ns.store(t4, std::memory_order_relaxed); + } RPC_STATUS_ASSERT(status); } @@ -861,7 +1113,7 @@ class rpc_server { bool free_buffer(const rpc_msg_free_buffer_req & request); bool buffer_clear(const rpc_msg_buffer_clear_req & request); bool memset_tensor(const rpc_msg_memset_tensor_req & request); - bool set_tensor(const std::vector & input); + bool set_tensor_stream(const socket_ptr & sock); bool set_tensor_hash(const rpc_msg_set_tensor_hash_req & request, rpc_msg_set_tensor_hash_rsp & response); bool get_tensor(const rpc_msg_get_tensor_req & request, std::vector & response); bool copy_tensor(const rpc_msg_copy_tensor_req & request, rpc_msg_copy_tensor_rsp & response); @@ -887,6 +1139,9 @@ class rpc_server { std::vector backends; const char * cache_dir; + // reused, never zero filled, and only allocated when the destination is not host memory + std::unique_ptr stage_buf; + size_t stage_cap = 0; std::unordered_set buffers; // store the last computed graph for each backend std::vector stored_graphs; @@ -896,6 +1151,7 @@ void rpc_server::hello(rpc_msg_hello_rsp & response) { response.major = RPC_PROTO_MAJOR_VERSION; response.minor = RPC_PROTO_MINOR_VERSION; response.patch = RPC_PROTO_PATCH_VERSION; + response.flags = cache_dir != nullptr ? RPC_SRV_FLAG_HAS_CACHE : 0; LOG_DBG("[%s] version: %d.%d.%d\n", __func__, response.major, response.minor, response.patch); } @@ -1115,15 +1371,33 @@ ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rp } -bool rpc_server::set_tensor(const std::vector & input) { - // serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes) | - if (input.size() < sizeof(rpc_tensor) + sizeof(uint64_t)) { +// Reads one RPC_CMD_SET_TENSOR message straight off the connection instead of into a +// std::vector sized to the whole message. The vector cost two full passes over every tensor, +// one to zero the fresh buffer and one to copy the payload out of it, and for a host backend +// the payload can be received into its final home with no copy at all. +// The wire format is unchanged: | rpc_tensor | offset (8 bytes) | data (size bytes) |. +bool rpc_server::set_tensor_stream(const socket_ptr & sock) { + constexpr size_t hdr_size = sizeof(rpc_tensor) + sizeof(uint64_t); + + uint64_t msg_size; + if (!sock->recv_data(&msg_size, sizeof(msg_size))) { + return false; + } + if (msg_size < hdr_size) { + GGML_LOG_ERROR("[%s] message too small (%" PRIu64 ")\n", __func__, msg_size); + return false; + } + const size_t size = msg_size - hdr_size; + if (g_rpc_loadprof_server.enabled) { + g_rpc_loadprof_server.bytes.fetch_add(msg_size, std::memory_order_relaxed); + g_rpc_loadprof_server.bucket(msg_size); + } + + rpc_tensor in_tensor; + uint64_t offset; + if (!sock->recv_data(&in_tensor, sizeof(in_tensor)) || !sock->recv_data(&offset, sizeof(offset))) { return false; } - const rpc_tensor * in_tensor = (const rpc_tensor *)input.data(); - uint64_t offset; - memcpy(&offset, input.data() + sizeof(rpc_tensor), sizeof(offset)); - const size_t size = input.size() - sizeof(rpc_tensor) - sizeof(offset); struct ggml_init_params params { /*.mem_size =*/ ggml_tensor_overhead(), @@ -1133,7 +1407,7 @@ bool rpc_server::set_tensor(const std::vector & input) { ggml_context_ptr ctx_ptr { ggml_init(params) }; GGML_ASSERT(ctx_ptr != nullptr); ggml_context * ctx = ctx_ptr.get(); - ggml_tensor * tensor = deserialize_tensor(ctx, in_tensor); + ggml_tensor * tensor = deserialize_tensor(ctx, &in_tensor); if (tensor == nullptr || tensor->buffer == nullptr) { GGML_LOG_ERROR("[%s] error deserializing tensor\n", __func__); return false; @@ -1145,25 +1419,48 @@ bool rpc_server::set_tensor(const std::vector & input) { const size_t p0 = (size_t) ggml_backend_buffer_get_base(tensor->buffer); const size_t p1 = p0 + ggml_backend_buffer_get_size(tensor->buffer); - if (in_tensor->data + offset < p0 || in_tensor->data + offset >= p1 || size > (p1 - in_tensor->data - offset)) { + if (in_tensor.data + offset < p0 || in_tensor.data + offset >= p1 || size > (p1 - in_tensor.data - offset)) { GGML_LOG_ERROR("[%s] tensor data region (data=0x%" PRIx64 ", offset=%" PRIu64 ", size=%zu) out of buffer bounds [0x%zx, 0x%zx)\n", - __func__, in_tensor->data, offset, size, p0, p1); + __func__, in_tensor.data, offset, size, p0, p1); return false; } } - const void * data = input.data() + sizeof(rpc_tensor) + sizeof(offset); + // a host buffer can take the payload directly; anything else, or a run that has to hash the + // payload for the cache, needs it contiguous in host memory first + const bool direct = ggml_backend_buffer_is_host(tensor->buffer) && cache_dir == nullptr; + uint8_t * dst; + if (direct) { + dst = (uint8_t *) tensor->data + offset; + } else { + if (stage_cap < size) { + stage_buf.reset(new (std::nothrow) uint8_t[size]); + if (stage_buf == nullptr) { + GGML_LOG_ERROR("[%s] failed to allocate %zu bytes of staging\n", __func__, size); + stage_cap = 0; + return false; + } + stage_cap = size; + } + dst = stage_buf.get(); + } + if (size > 0 && !sock->recv_data(dst, size)) { + return false; + } + if (cache_dir && size > HASH_THRESHOLD) { - uint64_t hash = fnv_hash((const uint8_t*)data, size); + uint64_t hash = fnv_hash(dst, size); char hash_str[17]; snprintf(hash_str, sizeof(hash_str), "%016" PRIx64, hash); // save to cache_dir/hash_str fs::path cache_file = fs::path(cache_dir) / hash_str; std::ofstream ofs(cache_file, std::ios::binary); - ofs.write((const char *)data, size); + ofs.write((const char *)dst, size); GGML_LOG_INFO("[%s] saved to '%s'\n", __func__, cache_file.string().c_str()); } - ggml_backend_tensor_set(tensor, data, offset, size); + if (!direct) { + ggml_backend_tensor_set(tensor, dst, offset, size); + } return true; } @@ -1679,13 +1976,27 @@ static void rpc_serve_client(const std::vector & backends, const break; } case RPC_CMD_SET_TENSOR: { - std::vector input; - if (!recv_msg(sock, input)) { - return; + rpc_load_prof & prof = g_rpc_loadprof_server; + const bool prof_on = prof.enabled; + uint64_t t0 = 0; + if (prof_on) { + t0 = rpc_load_prof::now_ns(); + uint64_t prev = prof.last_end_ns.load(std::memory_order_relaxed); + if (prev != 0) { + prof.ns_gap.fetch_add(t0 - prev, std::memory_order_relaxed); + } else { + prof.first_ns.store(t0, std::memory_order_relaxed); + } } - if (!server.set_tensor(input)) { + if (!server.set_tensor_stream(sock)) { return; } + if (prof_on) { + uint64_t t2 = rpc_load_prof::now_ns(); + prof.calls.fetch_add(1, std::memory_order_relaxed); + prof.ns_wire.fetch_add(t2 - t0, std::memory_order_relaxed); + prof.last_end_ns.store(t2, std::memory_order_relaxed); + } break; } case RPC_CMD_SET_TENSOR_HASH: { @@ -1849,6 +2160,7 @@ void ggml_backend_rpc_start_server(const char * endpoint, const char * cache_dir printf("Accepted client connection\n"); fflush(stdout); rpc_serve_client(backends, cache_dir, client_socket); + g_rpc_loadprof_server.print(); printf("Client connection closed\n"); fflush(stdout); } diff --git a/ggml/src/ggml-rpc/transport.cpp b/ggml/src/ggml-rpc/transport.cpp index 5ec15dc80c0..060aa37a188 100644 --- a/ggml/src/ggml-rpc/transport.cpp +++ b/ggml/src/ggml-rpc/transport.cpp @@ -53,7 +53,23 @@ using rdma_gid_t = std::array; #if defined(GGML_RPC_RDMA) && !defined(GGML_RPC_RDMA_APPLE) static constexpr size_t RDMA_CHUNK = 256 * 1024; // 256 KiB per send/recv (fits default 8 MiB memlock) -static constexpr int RDMA_RX_DEPTH = 24; // pre-posted recv ring: 24 × 256 KiB = 6 MiB +static constexpr int RDMA_RX_DEPTH = 24; // pre-posted recv ring: 24 x 256 KiB = 6 MiB + +// Sends used to be stop-and-wait: one 256 KiB chunk was posted and then polled to completion +// before the next was posted, so the wire was idle for a full round trip on every chunk and a +// multi-megabyte tensor moved at a fraction of the link rate. The send side now keeps up to +// RDMA_TX_DEPTH chunks in flight against a ring of registered staging buffers; completions are +// drained at every message boundary by flush(). The receive ring is 24 deep, so the sender can +// never outrun it. GGML_RPC_LOAD_OPT=0 restores the one-at-a-time behaviour. +static constexpr int RDMA_TX_DEPTH = 8; // send ring: 8 x 256 KiB = 2 MiB + +static bool rdma_load_opt() { + static const bool opt = [] { + const char * e = std::getenv("GGML_RPC_LOAD_OPT"); + return !(e && e[0] == '0'); + }(); + return opt; +} struct rdma_conn { struct ibv_context * ctx = nullptr; @@ -62,12 +78,22 @@ struct rdma_conn { struct ibv_cq * rcq = nullptr; // recv completions struct ibv_qp * qp = nullptr; - void * tx_buf = nullptr; + void * tx_buf = nullptr; // tx_depth x RDMA_CHUNK contiguous struct ibv_mr * tx_mr = nullptr; + int tx_depth = 1; // registered send slots + int tx_head = 0; // next slot to post from + int tx_inflight = 0; // posted but not yet completed - void * rx_buf = nullptr; // RDMA_RX_DEPTH × RDMA_CHUNK contiguous + void * rx_buf = nullptr; // RDMA_RX_DEPTH x RDMA_CHUNK contiguous struct ibv_mr * rx_mr = nullptr; int rx_head = 0; + // A completed receive is consumed byte by byte, so one posted buffer can serve several + // recv_data() calls. Without this a peer that framed a message differently (an older + // client that sends header and payload in one write) would have the remainder of its + // frame dropped on the floor. + int rx_cur = -1; // slot being consumed, -1 = none + size_t rx_off = 0; // bytes already taken from that slot + size_t rx_len = 0; // bytes the slot holds uint32_t max_inline = 0; @@ -75,6 +101,10 @@ struct rdma_conn { return static_cast(rx_buf) + static_cast(i) * RDMA_CHUNK; } + uint8_t * tx_slot(int i) const { + return static_cast(tx_buf) + static_cast(i) * RDMA_CHUNK; + } + bool post_rx(int i) { struct ibv_sge sge = {}; sge.addr = (uintptr_t)rx_slot(i); @@ -139,6 +169,8 @@ struct socket_t::impl { bool rdma_probe(); bool rdma_send(const void * data, size_t size); bool rdma_recv(void * data, size_t size); + bool rdma_wait_send_one(); + bool rdma_flush_sends(); bool tcp_peer_closed(); bool rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, const uint8_t * remote_gid); bool rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc); @@ -299,7 +331,7 @@ bool socket_t::impl::rdma_probe() { qia.send_cq = rdma->scq; qia.recv_cq = rdma->rcq; qia.qp_type = IBV_QPT_RC; - qia.cap.max_send_wr = 4; + qia.cap.max_send_wr = RDMA_TX_DEPTH + 4; qia.cap.max_recv_wr = RDMA_RX_DEPTH + 4; qia.cap.max_send_sge = 1; qia.cap.max_recv_sge = 1; @@ -309,11 +341,22 @@ bool socket_t::impl::rdma_probe() { if (!rdma->qp) return false; rdma->max_inline = qia.cap.max_inline_data; - rdma->tx_buf = aligned_alloc(4096, RDMA_CHUNK); + // Register as many send slots as the machine's locked memory allows, down to the single + // slot the transport has always used, so a tight memlock limit degrades instead of failing. + for (int d = rdma_load_opt() ? RDMA_TX_DEPTH : 1; d >= 1; d /= 2) { + rdma->tx_buf = aligned_alloc(4096, static_cast(d) * RDMA_CHUNK); + if (!rdma->tx_buf) continue; + rdma->tx_mr = ibv_reg_mr(rdma->pd, rdma->tx_buf, static_cast(d) * RDMA_CHUNK, + IBV_ACCESS_LOCAL_WRITE); + if (rdma->tx_mr) { + rdma->tx_depth = d; + break; + } + free(rdma->tx_buf); + rdma->tx_buf = nullptr; + } rdma->rx_buf = aligned_alloc(4096, static_cast(RDMA_RX_DEPTH) * RDMA_CHUNK); if (!rdma->tx_buf || !rdma->rx_buf) return false; - - rdma->tx_mr = ibv_reg_mr(rdma->pd, rdma->tx_buf, RDMA_CHUNK, IBV_ACCESS_LOCAL_WRITE); rdma->rx_mr = ibv_reg_mr(rdma->pd, rdma->rx_buf, static_cast(RDMA_RX_DEPTH) * RDMA_CHUNK, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); if (!rdma->tx_mr || !rdma->rx_mr) return false; @@ -394,8 +437,8 @@ bool socket_t::impl::rdma_activate(uint32_t remote_qpn, uint32_t remote_psn, con } } - GGML_LOG_INFO("RDMA activated: qpn=%u->%u mtu=%d rx_depth=%d\n", - rdma_local.qpn, remote_qpn, 128 << rdma_local.path_mtu, RDMA_RX_DEPTH); + GGML_LOG_INFO("RDMA activated: qpn=%u->%u mtu=%d rx_depth=%d tx_depth=%d\n", + rdma_local.qpn, remote_qpn, 128 << rdma_local.path_mtu, RDMA_RX_DEPTH, rdma->tx_depth); return true; } @@ -418,6 +461,22 @@ bool socket_t::impl::rdma_poll(struct ibv_cq * cq, struct ibv_wc * wc) { } } +// Retires exactly one posted send. Send completions on an RC queue pair arrive in post order, +// so retiring one completion frees the oldest slot in the ring. +bool socket_t::impl::rdma_wait_send_one() { + struct ibv_wc wc; + if (!rdma_poll(rdma->scq, &wc)) return false; + rdma->tx_inflight--; + return true; +} + +bool socket_t::impl::rdma_flush_sends() { + while (rdma->tx_inflight > 0) { + if (!rdma_wait_send_one()) return false; + } + return true; +} + bool socket_t::impl::rdma_send(const void * data, size_t size) { rdma_conn * c = rdma.get(); const uint8_t * src = (const uint8_t *)data; @@ -425,6 +484,13 @@ bool socket_t::impl::rdma_send(const void * data, size_t size) { while (rem > 0) { size_t chunk = std::min(rem, RDMA_CHUNK); + // Wait only when the ring is full. Because completions retire in post order, an + // in-flight count below the depth means the slot at tx_head is already free. + while (c->tx_inflight >= c->tx_depth) { + if (!rdma_wait_send_one()) return false; + } + const int slot = c->tx_head; + struct ibv_sge sge = {}; struct ibv_send_wr wr = {}, * bad = nullptr; wr.opcode = IBV_WR_SEND; @@ -432,20 +498,22 @@ bool socket_t::impl::rdma_send(const void * data, size_t size) { wr.num_sge = 1; if (chunk <= c->max_inline) { + // an inline send copies into the work request at post time, so the caller's + // buffer is free the moment ibv_post_send returns sge.addr = (uintptr_t)src; sge.length = chunk; wr.send_flags = IBV_SEND_SIGNALED | IBV_SEND_INLINE; } else { - memcpy(c->tx_buf, src, chunk); - sge.addr = (uintptr_t)c->tx_buf; + memcpy(c->tx_slot(slot), src, chunk); + sge.addr = (uintptr_t)c->tx_slot(slot); sge.length = chunk; sge.lkey = c->tx_mr->lkey; wr.send_flags = IBV_SEND_SIGNALED; } if (ibv_post_send(c->qp, &wr, &bad) != 0) return false; - struct ibv_wc wc; - if (!rdma_poll(c->scq, &wc)) return false; + c->tx_inflight++; + c->tx_head = (slot + 1) % c->tx_depth; src += chunk; rem -= chunk; @@ -458,17 +526,27 @@ bool socket_t::impl::rdma_recv(void * data, size_t size) { uint8_t * dst = (uint8_t *)data; size_t rem = size; while (rem > 0) { - struct ibv_wc wc; - if (!rdma_poll(c->rcq, &wc)) return false; - - int slot = (int)wc.wr_id; - size_t got = wc.byte_len; - memcpy(dst, c->rx_slot(slot), got); - - if (!c->post_rx(slot)) return false; - - dst += got; - rem -= got; + if (c->rx_cur < 0) { + struct ibv_wc wc; + if (!rdma_poll(c->rcq, &wc)) return false; + c->rx_cur = (int)wc.wr_id; + c->rx_off = 0; + c->rx_len = wc.byte_len; + } + // Take only what was asked for. A receive that carries more than this call needs is + // kept for the next call instead of being discarded, which is what lets a peer choose + // its own framing (one write per message, or a header write plus a payload write). + const size_t take = std::min(rem, c->rx_len - c->rx_off); + memcpy(dst, c->rx_slot(c->rx_cur) + c->rx_off, take); + c->rx_off += take; + dst += take; + rem -= take; + + if (c->rx_off >= c->rx_len) { + const int slot = c->rx_cur; + c->rx_cur = -1; + if (!c->post_rx(slot)) return false; + } } return true; } @@ -590,6 +668,11 @@ bool socket_t::impl::flush() { if (use_rdma) { return rdma->flush(); } +#elif defined(GGML_RPC_RDMA) + if (use_rdma) { + // a message is only really sent once every chunk of it has completed + return rdma_flush_sends(); + } #endif return true; } diff --git a/ggml/src/ggml-rpc/transport.h b/ggml/src/ggml-rpc/transport.h index 3f747ecffd9..8987b3887d5 100644 --- a/ggml/src/ggml-rpc/transport.h +++ b/ggml/src/ggml-rpc/transport.h @@ -13,6 +13,10 @@ static constexpr size_t RPC_CONN_CAPS_SIZE = 24; struct socket_t { ~socket_t(); + // features the peer advertised in its HELLO response (RPC_SRV_FLAG_* in ggml-rpc.cpp). + // Zero for a server that predates the flag, which is also the conservative value. + uint8_t srv_flags = 0; + bool send_data(const void * data, size_t size); bool recv_data(void * data, size_t size); // Must be called at every message boundary: the RDMA transport coalesces