diff --git a/.gitignore b/.gitignore index 9b589615a40..66a4bed61e9 100644 --- a/.gitignore +++ b/.gitignore @@ -153,3 +153,4 @@ a.out.* AGENTS.local.md .pi/SYSTEM.md +bench/ diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 059e4496269..5d403012824 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -7,7 +7,7 @@ extern "C" { #endif #define RPC_PROTO_MAJOR_VERSION 5 -#define RPC_PROTO_MINOR_VERSION 1 +#define RPC_PROTO_MINOR_VERSION 2 #define RPC_PROTO_PATCH_VERSION 0 #ifdef __cplusplus diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index e519bdf50a1..436d44d1d7f 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -499,6 +499,31 @@ void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor } } +// Asks the destination backend to take the copy, then the source backend. Only the destination +// used to be asked, which left a backend that can accelerate reads out of itself (the RPC +// backend, which stages through pinned memory on the peer device) on the synchronous fallback. +// +// The source is only asked when the two backends have DIFFERENT implementations. Two backends +// of the same type (two CUDA devices, say) share one implementation that already saw the pair +// and declined, so asking it again would be a second call with the same answer. That keeps +// every single-type setup, including one GPU and multi GPU, on exactly the old path. +static bool ggml_backend_cpy_tensor_async_impl(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst) { + if (backend_dst != NULL && backend_dst->iface.cpy_tensor_async != NULL) { + if (backend_dst->iface.cpy_tensor_async(backend_src, backend_dst, src, dst)) { + return true; + } + } + + if (backend_src != NULL && backend_src->iface.cpy_tensor_async != NULL && + (backend_dst == NULL || backend_src->iface.cpy_tensor_async != backend_dst->iface.cpy_tensor_async)) { + if (backend_src->iface.cpy_tensor_async(backend_src, backend_dst, src, dst)) { + return true; + } + } + + return false; +} + void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst) { GGML_ASSERT(ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts"); @@ -507,10 +532,8 @@ void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t b } GGML_ASSERT(backend_dst); - if (backend_dst->iface.cpy_tensor_async != NULL) { - if (backend_dst->iface.cpy_tensor_async(backend_src, backend_dst, src, dst)) { - return; - } + if (ggml_backend_cpy_tensor_async_impl(backend_src, backend_dst, src, dst)) { + return; } // an async copy would normally happen after all the queued operations on both backends are completed @@ -1728,7 +1751,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } else { // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface - if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) { + if (!ggml_backend_cpy_tensor_async_impl(input_backend, split_backend, input, input_cpy)) { ggml_backend_synchronize(input_backend); if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 69a8a08ae17..694dfd22f2e 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -17,6 +17,8 @@ #include #include #include +#include +#include static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG"); @@ -72,6 +74,7 @@ enum rpc_cmd { RPC_CMD_DEVICE_COUNT, RPC_CMD_GRAPH_RECOMPUTE, RPC_CMD_MEMSET_TENSOR, + RPC_CMD_GET_TENSORS, RPC_CMD_COUNT, }; @@ -176,6 +179,15 @@ struct rpc_msg_get_tensor_req { uint64_t size; }; +// RPC_CMD_GET_TENSORS reads several tensor regions in one round trip. The request is +// | n_entries (4 bytes) | n_entries x rpc_msg_get_tensors_entry |, the response is the +// requested regions concatenated in the order of the entries. +struct rpc_msg_get_tensors_entry { + rpc_tensor tensor; + uint64_t offset; + uint64_t size; +}; + struct rpc_msg_copy_tensor_req { rpc_tensor src; rpc_tensor dst; @@ -212,7 +224,7 @@ struct ggml_backend_rpc_device_context { uint32_t device; std::string name; std::string description; - uint64_t last_graph_uid; + // note: the uid of the last graph stored on the server is tracked per connection, see socket_t }; struct ggml_backend_rpc_buffer_type_context { @@ -298,9 +310,96 @@ static bool parse_endpoint(const std::string & endpoint, std::string & host, int return true; } +// Client side command counters, enabled with GGML_RPC_STATS=1. GGML_RPC_STATS_MS sets how often +// the running totals are printed (default 5000 ms). Used to count the RPC commands per decode +// step, which is what the backend sampling and async copy work is measured against. +static const char * RPC_STATS = std::getenv("GGML_RPC_STATS"); +static const int RPC_STATS_MS = std::getenv("GGML_RPC_STATS_MS") ? atoi(std::getenv("GGML_RPC_STATS_MS")) : 5000; + +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"; + case RPC_CMD_GET_TENSORS: return "GET_TENSORS"; + default: return "?"; + } +} + +static std::atomic rpc_stats_count[RPC_CMD_COUNT]; +static std::atomic rpc_stats_bytes[RPC_CMD_COUNT]; + +static void rpc_stats_record(enum rpc_cmd cmd, size_t bytes) { + rpc_stats_count[cmd].fetch_add(1, std::memory_order_relaxed); + rpc_stats_bytes[cmd].fetch_add(bytes, std::memory_order_relaxed); + + static std::mutex mtx; + static auto last = std::chrono::steady_clock::now(); + + std::unique_lock lock(mtx, std::try_to_lock); + if (!lock.owns_lock()) { + return; + } + const auto now = std::chrono::steady_clock::now(); + if (std::chrono::duration_cast(now - last).count() < RPC_STATS_MS) { + return; + } + last = now; + + std::string line = "RPCSTATS"; + for (int i = 0; i < RPC_CMD_COUNT; i++) { + const uint64_t n = rpc_stats_count[i].load(std::memory_order_relaxed); + if (n == 0) { + continue; + } + line += " " + std::string(rpc_cmd_name(i)) + "=" + std::to_string(n) + + "/" + std::to_string(rpc_stats_bytes[i].load(std::memory_order_relaxed)) + "B"; + } + fprintf(stderr, "%s\n", line.c_str()); +} + +// Deferred data movements (see rpc_deferred_op). Any command that is written to the socket has +// to keep its place in the wire order, so every entry point that sends flushes the queue first. +// The flush itself sends, hence the recursion guard. +static void rpc_flush_deferred(const socket_ptr & sock); + +static thread_local bool rpc_in_flush = false; + +static void rpc_flush_deferred_guarded(const socket_ptr & sock) { + if (rpc_in_flush) { + return; + } + { + std::lock_guard lock(sock->conn.mtx_defer); + if (sock->conn.deferred.empty()) { + return; + } + } + rpc_flush_deferred(sock); +} + // 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) { +// writes one whole message; the caller must hold sock->conn.mtx_send +static bool send_rpc_cmd_locked(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) { + if (RPC_STATS) { + rpc_stats_record(cmd, input_size); + } uint8_t cmd_byte = cmd; if (!sock->send_data(&cmd_byte, sizeof(cmd_byte))) { return false; @@ -314,12 +413,63 @@ static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, return sock->flush(); } +static bool send_rpc_cmd(socket_ptr sock, enum rpc_cmd cmd, const void * input, size_t input_size) { + rpc_flush_deferred_guarded(sock); + std::lock_guard lock(sock->conn.mtx_send); + return send_rpc_cmd_locked(sock, cmd, input, input_size); +} + +// Reserves this thread's place in the response order of a connection. The server answers the +// commands of one connection strictly in the order it received them, so the n-th response +// belongs to the n-th response-bearing request that was written to the socket. The ticket is +// taken while mtx_send is still held by the sender, and always released, so a failed send +// cannot leave the later waiters stuck. +struct rpc_response_ticket { + rpc_conn_state & conn; + uint64_t seq; + + // must be constructed with conn.mtx_send held + explicit rpc_response_ticket(rpc_conn_state & conn) : conn(conn) { + std::lock_guard lock(conn.mtx_seq); + seq = conn.seq_next++; + } + + void wait() { + std::unique_lock lock(conn.mtx_seq); + conn.cv_seq.wait(lock, [this] { return conn.seq_serving == seq; }); + } + + ~rpc_response_ticket() { + std::lock_guard lock(conn.mtx_seq); + conn.seq_serving = seq + 1; + conn.cv_seq.notify_all(); + } +}; + // 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) { - if (!send_rpc_cmd(sock, cmd, input, input_size)) { + rpc_flush_deferred_guarded(sock); + std::unique_ptr ticket; + bool failed = false; + { + std::lock_guard lock(sock->conn.mtx_send); + ticket.reset(new rpc_response_ticket(sock->conn)); + if (!send_rpc_cmd_locked(sock, cmd, input, input_size)) { + // still take our turn, so the ticket is released in order and no later waiter is + // woken with a response that is not theirs + failed = true; + } + } + + if (failed) { + ticket->wait(); return false; } + + // the response is read outside mtx_send, so the other threads can keep submitting + ticket->wait(); + uint64_t out_size; if (!sock->recv_data(&out_size, sizeof(out_size))) { return false; @@ -353,17 +503,33 @@ static bool negotiate_hello(const std::shared_ptr & sock) { return false; } + sock->conn.server_minor = response.minor; + sock->update_caps(response.conn_caps); return true; } +// The connections of an endpoint, looked up by endpoint. The server serves the connections of a +// client one at a time, so opening a second connection to an endpoint that already has a live one +// would block until the first closes. Anything that only wants the CURRENT connection must use +// find_socket, which never opens one. +static std::mutex g_sockets_mutex; +static std::unordered_map> g_sockets; + +static std::shared_ptr find_socket(const std::string & endpoint) { + std::lock_guard lock(g_sockets_mutex); + auto it = g_sockets.find(endpoint); + if (it != g_sockets.end()) { + return it->second.lock(); + } + return nullptr; +} + static std::shared_ptr get_socket(const std::string & endpoint) { - static std::mutex mutex; - std::lock_guard lock(mutex); - static std::unordered_map> sockets; + std::lock_guard lock(g_sockets_mutex); - auto it = sockets.find(endpoint); - if (it != sockets.end()) { + auto it = g_sockets.find(endpoint); + if (it != g_sockets.end()) { if (auto sock = it->second.lock()) { return sock; } @@ -386,7 +552,7 @@ static std::shared_ptr get_socket(const std::string & endpoint) { return nullptr; } LOG_DBG("[%s] connected to %s\n", __func__, endpoint.c_str()); - sockets[endpoint] = sock; + g_sockets[endpoint] = sock; return sock; } @@ -456,6 +622,201 @@ static rpc_tensor serialize_tensor(const ggml_tensor * tensor) { return result; } + +// +// asynchronous data movement over RPC +// +// The RPC backend used to be fully synchronous: every tensor read was a round trip and every +// hidden state that crossed the split was staged through a host malloc after a full +// ggml_backend_synchronize() of the producing device. Two things are added here: +// +// * get_tensor_async queues the read instead of performing it, and the queue is drained as a +// single RPC_CMD_GET_TENSORS at the next flush point. A decode step with backend sampling +// reads four small tensors per sequence, which used to be one round trip each. +// * cpy_tensor_async takes the device to device copies. For device -> RPC the producing +// backend copies into pinned staging on its own stream and records an event; the dispatcher +// sends from that staging once the event has completed, so the producing device is never +// fully synchronized and the host thread can serialize the graph meanwhile. +// + +struct rpc_staging { + ggml_backend_buffer_t buffer = nullptr; + uint8_t * base = nullptr; + size_t capacity = 0; + size_t used = 0; + + // an event on the consuming backend that still reads from this arena (RPC -> device) + ggml_backend_event_t inflight = nullptr; + + // events recorded on producing backends, reused across steps + std::vector events; + size_t events_used = 0; +}; + +static std::mutex rpc_staging_mutex; +static std::unordered_map rpc_staging_map; + +// caller must hold rpc_staging_mutex +static uint8_t * rpc_staging_alloc(rpc_staging & st, ggml_backend_buffer_type_t host_buft, size_t size) { + if (st.used + size > st.capacity) { + if (st.inflight != nullptr) { + ggml_backend_event_synchronize(st.inflight); + st.inflight = nullptr; + } + st.used = 0; + + if (size > st.capacity) { + const size_t want = std::max(size * 4, 1024 * 1024); + ggml_backend_buffer_t buf = ggml_backend_buft_alloc_buffer(host_buft, want); + if (buf == nullptr) { + return nullptr; + } + if (st.buffer != nullptr) { + ggml_backend_buffer_free(st.buffer); + } + st.buffer = buf; + st.base = (uint8_t *) ggml_backend_buffer_get_base(buf); + st.capacity = want; + } + } + + uint8_t * ptr = st.base + st.used; + st.used += size; + return ptr; +} + +// caller must hold rpc_staging_mutex +static ggml_backend_event_t rpc_staging_event(rpc_staging & st, ggml_backend_dev_t dev) { + if (st.events_used < st.events.size()) { + return st.events[st.events_used++]; + } + ggml_backend_event_t ev = ggml_backend_event_new(dev); + if (ev == nullptr) { + return nullptr; + } + st.events.push_back(ev); + st.events_used++; + return ev; +} + +// | n_entries (4 bytes) | n_entries x rpc_msg_get_tensors_entry |, response scattered into the +// destination pointers of the entries +static bool send_get_tensors(const socket_ptr & sock, const std::vector & gets) { + const uint32_t n = (uint32_t) gets.size(); + + std::vector input(sizeof(uint32_t) + n * sizeof(rpc_msg_get_tensors_entry)); + memcpy(input.data(), &n, sizeof(n)); + + auto * entries = (rpc_msg_get_tensors_entry *) (input.data() + sizeof(uint32_t)); + uint64_t total = 0; + for (uint32_t i = 0; i < n; i++) { + GGML_ASSERT(gets[i]->tensor_bytes.size() == sizeof(rpc_tensor)); + memcpy(&entries[i].tensor, gets[i]->tensor_bytes.data(), sizeof(rpc_tensor)); + entries[i].offset = gets[i]->offset; + entries[i].size = gets[i]->size; + total += gets[i]->size; + } + + std::unique_ptr ticket; + bool failed = false; + { + std::lock_guard lock(sock->conn.mtx_send); + ticket.reset(new rpc_response_ticket(sock->conn)); + if (!send_rpc_cmd_locked(sock, RPC_CMD_GET_TENSORS, input.data(), input.size())) { + failed = true; + } + } + ticket->wait(); + if (failed) { + return false; + } + + uint64_t out_size; + if (!sock->recv_data(&out_size, sizeof(out_size))) { + return false; + } + if (out_size != total) { + return false; + } + + // One receive for the whole response, then scatter. The RDMA transport is not a byte stream: + // a receive completion carries exactly one send, and recv_data copies all of it, so reading a + // single sent message back in several pieces overruns the first destination and then blocks + // for a completion that never comes. + std::vector response(total); + if (total > 0 && !sock->recv_data(response.data(), total)) { + return false; + } + size_t off = 0; + for (uint32_t i = 0; i < n; i++) { + memcpy(gets[i]->data, response.data() + off, gets[i]->size); + off += gets[i]->size; + } + return true; +} + +static void rpc_flush_deferred(const socket_ptr & sock) { + std::lock_guard lock_defer(sock->conn.mtx_defer); + + if (sock->conn.deferred.empty()) { + return; + } + + rpc_in_flush = true; + + std::vector ops; + ops.swap(sock->conn.deferred); + + // a run of consecutive reads goes out as one command; writes keep their place in the order + std::vector gets; + auto flush_gets = [&]() { + if (gets.empty()) { + return; + } + bool status = send_get_tensors(sock, gets); + RPC_STATUS_ASSERT(status); + gets.clear(); + }; + + for (auto & op : ops) { + if (op.kind == rpc_deferred_op::GET) { + gets.push_back(&op); + continue; + } + + flush_gets(); + + if (op.event != nullptr) { + ggml_backend_event_synchronize((ggml_backend_event_t) op.event); + } + + GGML_ASSERT(op.tensor_bytes.size() == sizeof(rpc_tensor)); + rpc_tensor rpc_dst; + memcpy(&rpc_dst, op.tensor_bytes.data(), sizeof(rpc_tensor)); + + // input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes) + std::vector input(sizeof(rpc_dst) + sizeof(uint64_t) + op.size); + memcpy(input.data(), &rpc_dst, sizeof(rpc_dst)); + memcpy(input.data() + sizeof(rpc_dst), &op.offset, sizeof(op.offset)); + memcpy(input.data() + sizeof(rpc_dst) + sizeof(op.offset), op.data, op.size); + + std::lock_guard lock_send(sock->conn.mtx_send); + bool status = send_rpc_cmd_locked(sock, RPC_CMD_SET_TENSOR, input.data(), input.size()); + RPC_STATUS_ASSERT(status); + } + flush_gets(); + + { + std::lock_guard lock(rpc_staging_mutex); + auto it = rpc_staging_map.find(sock.get()); + if (it != rpc_staging_map.end()) { + it->second.events_used = 0; + } + } + + rpc_in_flush = false; +} + static enum ggml_status ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; @@ -674,8 +1035,180 @@ static void ggml_backend_rpc_free(ggml_backend_t backend) { } static void ggml_backend_rpc_synchronize(ggml_backend_t backend) { + ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context; + // never open a connection here: with nothing connected there is nothing queued either + auto sock = find_socket(rpc_ctx->endpoint); + if (sock != nullptr) { + rpc_flush_deferred_guarded(sock); + } +} + +// true when the server understands RPC_CMD_GET_TENSORS +static bool rpc_supports_batched_get(const socket_ptr & sock) { + static const bool disabled = std::getenv("GGML_RPC_NO_BATCHED_GET") != nullptr; + return !disabled && sock->conn.server_minor >= 2; +} + +// The connection a tensor already lives on. Never opens one. +static socket_ptr tensor_socket(const ggml_tensor * tensor) { + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + if (buf == nullptr || !ggml_backend_buffer_is_rpc(buf)) { + return nullptr; + } + return ((ggml_backend_rpc_buffer_context *) buf->context)->sock; +} + +static void ggml_backend_rpc_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { GGML_UNUSED(backend); - // this is no-op because we don't have any async operations + auto sock = tensor_socket(tensor); + + if (sock == nullptr || !rpc_supports_batched_get(sock)) { + ggml_backend_tensor_get(tensor, data, offset, size); + return; + } + + const rpc_tensor rpc_src = serialize_tensor(tensor); + + rpc_deferred_op op; + op.kind = rpc_deferred_op::GET; + op.tensor_bytes.assign((const uint8_t *) &rpc_src, (const uint8_t *) &rpc_src + sizeof(rpc_src)); + op.data = data; + op.offset = offset; + op.size = size; + + std::lock_guard lock(sock->conn.mtx_defer); + sock->conn.deferred.push_back(op); +} + +// Device to device movement across the split, without a full synchronize of the producing +// device and without the host malloc that ggml_backend_tensor_copy would do. +static bool ggml_backend_rpc_cpy_tensor_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, + const ggml_tensor * src, ggml_tensor * dst) { + static const bool disabled = std::getenv("GGML_RPC_NO_ASYNC_COPY") != nullptr; + if (disabled) { + return false; + } + + const bool src_is_rpc = ggml_backend_is_rpc(backend_src); + const bool dst_is_rpc = ggml_backend_is_rpc(backend_dst); + + // RPC to RPC on the same server is handled by the buffer level copy + if (src_is_rpc == dst_is_rpc) { + return false; + } + + const size_t size = ggml_nbytes(src); + if (size != ggml_nbytes(dst)) { + return false; + } + + ggml_backend_t other = src_is_rpc ? backend_dst : backend_src; + ggml_backend_dev_t other_dev = ggml_backend_get_device(other); + + // the staging buffer has to be pinned on the other device, or its copies are not async + ggml_backend_buffer_type_t host_buft = other_dev != nullptr ? ggml_backend_dev_host_buffer_type(other_dev) : nullptr; + if (host_buft == nullptr) { + return false; + } + + auto sock = tensor_socket(src_is_rpc ? src : dst); + if (sock == nullptr) { + return false; + } + + // the asynchronous entry points of a backend only accept tensors that live in its own + // default buffer type; anything else (a host buffer that the backend can also reach) stays + // on the synchronous path + const ggml_tensor * other_t = src_is_rpc ? dst : src; + ggml_backend_buffer_t other_buf = other_t->view_src ? other_t->view_src->buffer : other_t->buffer; + if (other_buf == nullptr || ggml_backend_buffer_get_type(other_buf) != ggml_backend_get_default_buffer_type(other)) { + return false; + } + + if (!src_is_rpc) { + // device -> RPC: pinned D2H on the producing stream, an event to tell when it landed, + // and the send happens at the next flush point + if (backend_src->iface.get_tensor_async == nullptr) { + return false; + } + + // never wrap the arena over staging that a queued send has not read yet + { + bool wraps = false; + { + std::lock_guard lock(rpc_staging_mutex); + rpc_staging & st = rpc_staging_map[sock.get()]; + wraps = st.used + size > st.capacity; + } + if (wraps) { + rpc_flush_deferred_guarded(sock); + } + } + + uint8_t * staging = nullptr; + ggml_backend_event_t event = nullptr; + { + std::lock_guard lock(rpc_staging_mutex); + rpc_staging & st = rpc_staging_map[sock.get()]; + staging = rpc_staging_alloc(st, host_buft, size); + if (staging == nullptr) { + return false; + } + event = rpc_staging_event(st, other_dev); + } + if (event == nullptr) { + return false; + } + + ggml_backend_tensor_get_async(backend_src, src, staging, 0, size); + ggml_backend_event_record(event, backend_src); + + const rpc_tensor rpc_dst = serialize_tensor(dst); + + rpc_deferred_op op; + op.kind = rpc_deferred_op::SET; + op.tensor_bytes.assign((const uint8_t *) &rpc_dst, (const uint8_t *) &rpc_dst + sizeof(rpc_dst)); + op.data = staging; + op.offset = 0; + op.size = size; + op.event = event; + + std::lock_guard lock(sock->conn.mtx_defer); + sock->conn.deferred.push_back(op); + return true; + } + + // RPC -> device: read into pinned staging, then an async H2D on the consuming stream, so + // the consuming device is not synchronized and the caller can queue its graph right after + if (backend_dst->iface.set_tensor_async == nullptr) { + return false; + } + + uint8_t * staging = nullptr; + { + std::lock_guard lock(rpc_staging_mutex); + rpc_staging & st = rpc_staging_map[sock.get()]; + staging = rpc_staging_alloc(st, host_buft, size); + } + if (staging == nullptr) { + return false; + } + + ggml_backend_tensor_get(src, staging, 0, size); + ggml_backend_tensor_set_async(backend_dst, dst, staging, 0, size); + + { + std::lock_guard lock(rpc_staging_mutex); + rpc_staging & st = rpc_staging_map[sock.get()]; + ggml_backend_event_t event = rpc_staging_event(st, other_dev); + if (event != nullptr) { + ggml_backend_event_record(event, backend_dst); + st.inflight = event; + } else { + ggml_backend_synchronize(backend_dst); + } + } + return true; } static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::vector & tensors, std::unordered_set & visited) { @@ -731,21 +1264,34 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g ggml_backend_rpc_device_context * rpc_dev_ctx = (ggml_backend_rpc_device_context *)rpc_dev->context; GGML_ASSERT(cgraph->n_nodes > 0); - bool reuse = cgraph->uid != 0 && rpc_dev_ctx->last_graph_uid == cgraph->uid; - if (reuse) { + GGML_UNUSED(rpc_dev_ctx); + + auto sock = get_socket(rpc_ctx->endpoint); + + // the queued inputs of this graph have to be on the wire before the compute command + rpc_flush_deferred_guarded(sock); + + // The graph stored by RPC_CMD_GRAPH_COMPUTE lives on the server per connection and device, + // and one connection is shared by every backend of this endpoint - including the backends of + // other llama_contexts. So the uid of the last graph sent has to be tracked per connection, + // and the check has to happen under the same lock as the send, or a RECOMPUTE could re-run + // the graph another context stored in between. + std::unique_lock lock(sock->conn.mtx_send); + + auto & last_uid = sock->conn.last_graph_uid[rpc_ctx->device]; + if (cgraph->uid != 0 && last_uid == cgraph->uid) { rpc_msg_graph_recompute_req request; request.device = rpc_ctx->device; - auto sock = get_socket(rpc_ctx->endpoint); - bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request)); - RPC_STATUS_ASSERT(status); - } else { - rpc_dev_ctx->last_graph_uid = cgraph->uid; - std::vector input; - serialize_graph(rpc_ctx->device, cgraph, input); - auto sock = get_socket(rpc_ctx->endpoint); - bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size()); + bool status = send_rpc_cmd_locked(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request)); RPC_STATUS_ASSERT(status); + return GGML_STATUS_SUCCESS; } + + last_uid = cgraph->uid; + std::vector input; + serialize_graph(rpc_ctx->device, cgraph, input); + bool status = send_rpc_cmd_locked(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size()); + RPC_STATUS_ASSERT(status); return GGML_STATUS_SUCCESS; } @@ -753,10 +1299,10 @@ static ggml_backend_i ggml_backend_rpc_interface = { /* .get_name = */ ggml_backend_rpc_name, /* .free = */ ggml_backend_rpc_free, /* .set_tensor_async = */ NULL, - /* .get_tensor_async = */ NULL, + /* .get_tensor_async = */ ggml_backend_rpc_get_tensor_async, /* .set_tensor_2d_async = */ NULL, /* .get_tensor_2d_async = */ NULL, - /* .cpy_tensor_async = */ NULL, + /* .cpy_tensor_async = */ ggml_backend_rpc_cpy_tensor_async, /* .synchronize = */ ggml_backend_rpc_synchronize, /* .graph_plan_create = */ NULL, /* .graph_plan_free = */ NULL, @@ -864,6 +1410,7 @@ class rpc_server { bool set_tensor(const std::vector & input); 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 get_tensors(const std::vector & input, std::vector & response); bool copy_tensor(const rpc_msg_copy_tensor_req & request, rpc_msg_copy_tensor_rsp & response); bool graph_compute(const std::vector & input); bool graph_recompute(const rpc_msg_graph_recompute_req & request); @@ -1299,6 +1846,61 @@ bool rpc_server::get_tensor(const rpc_msg_get_tensor_req & request, std::vector< return true; } + +// Reads several tensor regions in one command. The regions are concatenated into the response in +// request order, so one decode step of a backend sampled batch is one round trip instead of one +// per sequence and per sampler output. +bool rpc_server::get_tensors(const std::vector & input, std::vector & response) { + if (input.size() < sizeof(uint32_t)) { + return false; + } + uint32_t n_entries; + memcpy(&n_entries, input.data(), sizeof(n_entries)); + if (input.size() != sizeof(uint32_t) + (size_t) n_entries * sizeof(rpc_msg_get_tensors_entry)) { + return false; + } + const auto * entries = (const rpc_msg_get_tensors_entry *) (input.data() + sizeof(uint32_t)); + + size_t total = 0; + for (uint32_t i = 0; i < n_entries; i++) { + total += entries[i].size; + } + response.resize(total, 0); + + size_t out_offset = 0; + for (uint32_t i = 0; i < n_entries; i++) { + struct ggml_init_params params { + /*.mem_size =*/ ggml_tensor_overhead(), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + ggml_context_ptr ctx_ptr { ggml_init(params) }; + GGML_ASSERT(ctx_ptr != nullptr); + ggml_tensor * tensor = deserialize_tensor(ctx_ptr.get(), &entries[i].tensor); + if (tensor == nullptr || tensor->buffer == nullptr) { + GGML_LOG_ERROR("[%s] error deserializing tensor %u\n", __func__, i); + return false; + } + + // sanitize tensor->data + { + 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 (entries[i].tensor.data + entries[i].offset < p0 || + entries[i].tensor.data + entries[i].offset >= p1 || + entries[i].size > (p1 - entries[i].tensor.data - entries[i].offset)) { + GGML_LOG_ERROR("[%s] requested tensor region out of buffer bounds\n", __func__); + return false; + } + } + + ggml_backend_tensor_get(tensor, response.data() + out_offset, entries[i].offset, entries[i].size); + out_offset += entries[i].size; + } + return true; +} + bool rpc_server::copy_tensor(const rpc_msg_copy_tensor_req & request, rpc_msg_copy_tensor_rsp & response) { struct ggml_init_params params { /*.mem_size =*/ 2*ggml_tensor_overhead(), @@ -1729,6 +2331,20 @@ static void rpc_serve_client(const std::vector & backends, const } break; } + case RPC_CMD_GET_TENSORS: { + std::vector input; + if (!recv_msg(sock, input)) { + return; + } + std::vector response; + if (!server.get_tensors(input, response)) { + return; + } + if (!send_msg(sock, response.data(), response.size())) { + return; + } + break; + } case RPC_CMD_COPY_TENSOR: { rpc_msg_copy_tensor_req request; if (!recv_msg(sock, &request, sizeof(request))) { @@ -2044,7 +2660,6 @@ ggml_backend_reg_t ggml_backend_rpc_add_server(const char * endpoint) { /* .device = */ ind, /* .name = */ dev_name, /* .description = */ dev_desc, - /* .last_graph_uid = */ 0, }; ggml_backend_dev_t dev = new ggml_backend_device { diff --git a/ggml/src/ggml-rpc/transport.h b/ggml/src/ggml-rpc/transport.h index 3f747ecffd9..78d708dfa41 100644 --- a/ggml/src/ggml-rpc/transport.h +++ b/ggml/src/ggml-rpc/transport.h @@ -1,18 +1,71 @@ #pragma once +#include #include #include #include +#include +#include +#include struct socket_t; typedef std::shared_ptr socket_ptr; +// One deferred RPC data movement, queued by the asynchronous tensor entry points of the RPC +// backend and executed on the connection at the next flush point (a synchronize, or any other +// command that has to keep its place in the wire order). +struct rpc_deferred_op { + enum kind_t { GET, SET } kind; + + // The tensor is serialized when the operation is queued, not when it is flushed: the graph + // result that owns it can be reset before the next flush point (one llama_decode allocates + // and resets a graph per ubatch), so keeping the pointer would leave a dangling read. + std::vector tensor_bytes; + + void * data = nullptr; // GET: host destination; SET: host staging source + uint64_t offset = 0; + uint64_t size = 0; + + // SET only: an event on the producing backend that must complete before the staging + // buffer holds the data. Opaque here so the transport stays free of ggml-backend types. + void * event = nullptr; +}; + static constexpr size_t MAX_CHUNK_SIZE = 1024ull * 1024ull * 1024ull; // 1 GiB static constexpr size_t RPC_CONN_CAPS_SIZE = 24; +// State shared by every client thread that uses one connection. A connection is looked up by +// endpoint and is therefore shared by all backends of that endpoint, including the backends of +// different llama_contexts, so all of it has to be serialised: +// - mtx_send makes a whole RPC message atomic on the wire +// - seq_* hands the responses out in request order (the server answers strictly in order), +// without holding mtx_send while waiting, so another thread can keep submitting work +// - last_graph_uid mirrors the server's per-connection stored graph for a device, so that +// RPC_CMD_GRAPH_RECOMPUTE can never re-run a graph submitted by another context +struct rpc_conn_state { + std::mutex mtx_send; + std::mutex mtx_seq; + std::condition_variable cv_seq; + uint64_t seq_next = 0; + uint64_t seq_serving = 0; + + std::unordered_map last_graph_uid; + + // minor version reported by the server in HELLO, used to gate optional commands + uint32_t server_minor = 0; + + // deferred data movements, see rpc_deferred_op. Guarded by mtx_defer, which is always + // taken before mtx_send and never while holding it. + std::mutex mtx_defer; + std::vector deferred; +}; + struct socket_t { ~socket_t(); + // guarded by conn.mtx_send / conn.mtx_seq, see rpc_conn_state + rpc_conn_state conn; + 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 diff --git a/tools/server/README.md b/tools/server/README.md index 93736c3edfa..bef9088f39c 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -2074,6 +2074,40 @@ Note that the following endpoints are exempt from being considered as incoming t - `GET /models` - `GET /metrics` +## Pipeline groups + +`--pipeline-groups N` (default `1`) runs the server's slots over `N` independent `llama_context` +objects created from the same model. Each group has its own batch, its own sampling and its own +decode thread; the model weights, the task queue, the results queue and the HTTP layer are shared. + +This is meant for a layer split across two machines, e.g. + +```sh +llama-server -m model.gguf -c 32768 --parallel 16 \ + --rpc peer:50052 --device CUDA0,RPC0 -sm layer -ngl 99 \ + --pipeline-groups 2 +``` + +With one context, a layer split is a two-stage pipeline that is fed one batch at a time, so each +stage is idle while the other one computes. With two groups there are two batches in flight, so +while group A is being computed on the second stage, group B is being computed on the first one. + +Details: + +- The slots are partitioned contiguously: with `--parallel P` and `--pipeline-groups N`, group `g` + owns slots `[g*P/N, (g+1)*P/N)`. `--parallel` must be a positive multiple of `--pipeline-groups`. +- Each context is created with `n_seq_max = P/N` and `n_ctx = C/N`, so the per-slot context and the + total KV memory over all groups are the same as with a single context. `-c` must be given + explicitly and must be a multiple of `N`. +- Slot selection for an incoming request still runs over *all* slots, so prompt cache similarity and + the slot save / restore endpoints work exactly as before: a returning conversation lands on the + slot that still holds its prefix, whichever group that slot belongs to. +- Task processing briefly pauses the decode loops, so `/slots`, `/metrics` and cancellations are + answered after the in-flight decode of each group finishes rather than during it. +- `N > 1` is refused at startup together with speculative decoding (`--model-draft`, MTP), + multimodal (`--mmproj`) and `--sleep-idle-seconds`. +- With `N = 1` nothing changes: one context, one batch and one update loop on the main thread. + ## More examples ### Interactive mode diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b..2b1f2e32157 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -35,6 +35,12 @@ #include #endif +// used by the --pipeline-groups decode threads +#include +#include +#include +#include + constexpr int HTTP_POLLING_SECONDS = 1; static common_speculative_output_limits server_output_limits(const common_params & params) { @@ -68,7 +74,8 @@ struct server_batch { bool batch_rendered = false; struct token { - int32_t id_slot; + int32_t id_slot; // global slot id, used to attribute stats + int32_t seq_id; // sequence id inside the context of the slot's pipeline group llama_token token; llama_pos pos; bool output; @@ -108,22 +115,22 @@ struct server_batch { tokens.reserve(n_tokens_alloc); } - bool add(int32_t id_slot, llama_token token, llama_pos pos, bool output, bool is_prompt) { + bool add(int32_t id_slot, int32_t seq_id, llama_token token, llama_pos pos, bool output, bool is_prompt) { GGML_ASSERT(!has_embd); // cannot mix tokens + embd in same batch GGML_ASSERT(batch.pos != nullptr); if ((int32_t)tokens.size() >= n_tokens_alloc) { return false; } - tokens.push_back({ id_slot, token, pos, output, is_prompt }); + tokens.push_back({ id_slot, seq_id, token, pos, output, is_prompt }); return true; } - bool add(int32_t id_slot, const std::vector & embd_in, llama_pos pos, bool output, bool is_prompt) { + bool add(int32_t id_slot, int32_t seq_id, const std::vector & embd_in, llama_pos pos, bool output, bool is_prompt) { GGML_ASSERT(batch.pos != nullptr); if ((int32_t)tokens.size() >= n_tokens_alloc) { return false; } - tokens.push_back({ id_slot, LLAMA_TOKEN_NULL, pos, output, is_prompt }); + tokens.push_back({ id_slot, seq_id, LLAMA_TOKEN_NULL, pos, output, is_prompt }); has_embd = true; embd.insert(embd.end(), embd_in.begin(), embd_in.end()); return true; @@ -159,7 +166,7 @@ struct server_batch { common_batch_clear(batch); for (int32_t i = 0; i < size(); i++) { const auto & t = tokens[i]; - common_batch_add(batch, t.token, t.pos, { t.id_slot }, t.output); + common_batch_add(batch, t.token, t.pos, { t.seq_id }, t.output); } if (has_embd) { batch.token = nullptr; // will be restored on clear() @@ -194,6 +201,14 @@ struct server_batch { struct server_slot { int id; + // pipeline group that owns this slot, i.e. the index of ctx_tgt in server_context_impl::groups + // always 0 unless --pipeline-groups > 1 + int id_group = 0; + + // sequence id of this slot inside ctx_tgt / ctx_dft + // equal to id unless --pipeline-groups > 1, where each context only holds n_parallel/N sequences + int seq_id = 0; + llama_context * ctx_tgt = nullptr; llama_context * ctx_dft = nullptr; @@ -255,8 +270,8 @@ struct server_slot { return false; } - const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); - const size_t cur_size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; + const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE); + const size_t cur_size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0; const size_t cur_size = cur_size_tgt + cur_size_dft; @@ -268,16 +283,16 @@ struct server_slot { return false; } - llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); + llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE); if (ctx_dft) { - llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE); + llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_NONE); } return true; } bool prompt_load(server_prompt_cache & prompt_cache, const server_tokens & tokens) { - bool res = prompt_cache.load(prompt, tokens, ctx_tgt, ctx_dft, id); + bool res = prompt_cache.load(prompt, tokens, ctx_tgt, ctx_dft, seq_id); if (!res) { SLT_WRN(*this, "%s", "failed to load prompt from cache\n"); } @@ -288,7 +303,7 @@ struct server_slot { void prompt_clear() { SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); - mem.seq_rm(id, -1, -1); + mem.seq_rm(seq_id, -1, -1); prompt.clear(); } @@ -351,7 +366,7 @@ struct server_slot { n_predict_max = -1; - llama_set_sampler(ctx_tgt, id, nullptr); + llama_set_sampler(ctx_tgt, seq_id, nullptr); // clear alora start alora_invocation_start = -1; @@ -463,9 +478,9 @@ struct server_slot { i_batch = batch.size(); if (!inp_embd.empty()) { - add_ok &= batch.add(id, inp_embd, prompt.tokens.pos_next(), true, false); + add_ok &= batch.add(id, seq_id, inp_embd, prompt.tokens.pos_next(), true, false); } else { - add_ok &= batch.add(id, sampled, prompt.tokens.pos_next(), true, false); + add_ok &= batch.add(id, seq_id, sampled, prompt.tokens.pos_next(), true, false); } SLT_DBG(*this, "slot decode token, id=%d, n_ctx = %d, n_tokens = %d, truncated = %d\n", @@ -483,9 +498,9 @@ struct server_slot { auto pos0 = prompt.tokens.pos_next(); - add_ok &= batch.add(id, sampled, pos0++, true, false); + add_ok &= batch.add(id, seq_id, sampled, pos0++, true, false); for (auto token : spec_draft) { - add_ok &= batch.add(this->id, token, pos0++, true, false); + add_ok &= batch.add(this->id, seq_id, token, pos0++, true, false); } } @@ -676,8 +691,9 @@ struct server_slot { void copy_state_to(server_slot & other) const { GGML_ASSERT(state == SLOT_STATE_DONE_PROMPT); - mem.seq_rm(other.id, -1, -1); - mem.seq_cp(id, other.id, -1, -1); + // note: parent and child slots are always in the same pipeline group, see get_free_slots() + mem.seq_rm(other.seq_id, -1, -1); + mem.seq_cp(seq_id, other.seq_id, -1, -1); other.i_batch = i_batch; @@ -780,6 +796,37 @@ static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch return try_decode(); } +// A pipeline group is one llama_context with its own batch, its own decode loop and its own +// contiguous range of slots. With --pipeline-groups 1 (the default) there is exactly one group: +// it owns ctx_tgt and every slot, and its update loop runs on the main thread, as before. +// +// With N > 1 the point is that while group A's batch is being computed on the second stage of a +// layer split (the RPC peer), group B's batch can be computed on the first stage (the local GPU), +// so both devices are busy instead of each idling half of every decode step. +struct server_group { + int id = 0; + + llama_context * ctx = nullptr; + + server_batch batch; + + // slots owned by this group, in slot id order (slots are partitioned contiguously) + std::vector slots; + + // queued prompt stats - llama_decode() is async, so the timing is only valid after a sync + // note: kept out of server_metrics, which is copied as-is into the task result + int64_t t_decode_start = 0; // start of the last submitted decode of this group + int64_t t_prompt_start = 0; // start of the oldest queued prompt decode of this group + uint64_t n_prompt_queued = 0; + + int n_empty_consecutive = 0; + + // only used when n_groups > 1, all guarded by server_context_impl::mtx_engine + std::thread thread; + bool busy = false; // a decode is in flight, no one may touch ctx + int n_pause_req = 0; // someone wants the engine stopped, do not start a new iteration +}; + // // server_context_impl (private implementation) // @@ -804,6 +851,9 @@ struct server_context_impl { server_state_callback_t callback_state = [](server_state, json) -> void {}; + // number of pipeline groups requested via --pipeline-groups, must be set before load_model() + int n_pipeline_groups_req = 1; + server_context_impl() { mtmd_helper_log_set(common_log_default_callback, nullptr); } @@ -835,7 +885,18 @@ struct server_context_impl { llama_context * ctx_tgt = nullptr; - server_batch batch; + // pipeline groups, see --pipeline-groups and struct server_group + // groups[0]->ctx is always ctx_tgt; n_groups == 1 unless the user asked for more + int n_groups = 1; + std::vector> groups; + + // number of sequences per context, == params_base.n_parallel when n_groups == 1 + int n_seq_per_group = 1; + + // the following are only ever touched when n_groups > 1 + std::mutex mtx_engine; + std::condition_variable cv_engine; + bool groups_stop = false; llama_model * model_dft = nullptr; llama_context * ctx_dft = nullptr; @@ -862,18 +923,10 @@ struct server_context_impl { int slots_debug = 0; // env: LLAMA_SERVER_SLOTS_DEBUG int slots_n_diff = 0; // env: LLAMA_SERVER_SLOTS_N_DIFF - int n_empty_consecutive = 0; - std::unique_ptr prompt_cache; server_metrics metrics; - // queued prompt stats - llama_decode() is async, so the timing is only valid after a sync - // note: kept out of server_metrics, which is copied as-is into the task result - int64_t t_decode_start = 0; // start of the last submitted decode - int64_t t_prompt_start = 0; // start of the oldest queued prompt decode - uint64_t n_prompt_queued = 0; - json json_ui_settings = json::object(); // Necessary similarity of prompt for slot selection @@ -894,6 +947,15 @@ struct server_context_impl { ctx_dft = nullptr; model_dft = nullptr; + // groups[0]->ctx is owned by llama_init, the rest were created by llama_init_from_model() + for (size_t g = 1; g < groups.size(); ++g) { + if (groups[g]->ctx != nullptr) { + llama_free(groups[g]->ctx); + groups[g]->ctx = nullptr; + } + } + groups.clear(); + llama_init.reset(); ctx_tgt = nullptr; @@ -1048,11 +1110,44 @@ struct server_context_impl { params_base.load_progress_callback_user_data = &load_progress_text; } - llama_init = common_init_from_params(params_base); + // --pipeline-groups: run the slots over N independent contexts of one model, so that the + // stages of a layer split can be busy at the same time. N == 1 is the default and keeps + // every code path below exactly as it was. + n_groups = std::max(1, n_pipeline_groups_req); + + if (n_groups > 1 && !validate_pipeline_groups(params_base, has_spec, has_mmproj)) { + return false; + } + + n_seq_per_group = params_base.n_parallel / n_groups; + + // note: with a single group this reference IS params_base, so nothing changes + common_params params_grp = n_groups > 1 ? params_base : common_params{}; + common_params & params_ctx = n_groups > 1 ? params_grp : params_base; + + if (n_groups > 1) { + // each context gets 1/N of the sequences and 1/N of the total context, so the per-slot + // context (n_ctx / n_seq_max) and the total KV memory over all contexts are unchanged + params_ctx.n_parallel = n_seq_per_group; + params_ctx.n_ctx = params_base.n_ctx / n_groups; + } + + llama_init = common_init_from_params(params_ctx); model_tgt = llama_init->model(); ctx_tgt = llama_init->context(); + if (n_groups > 1) { + // pick up whatever the parameter fitting resolved, but keep the totals the user asked for + const int n_parallel_total = params_base.n_parallel; + const int n_ctx_total = params_base.n_ctx; + + params_base = params_ctx; + + params_base.n_parallel = n_parallel_total; + params_base.n_ctx = n_ctx_total; + } + if (model_tgt == nullptr) { SRV_ERR("failed to load model, '%s'\n", params_base.model.path.c_str()); return false; @@ -1065,7 +1160,43 @@ struct server_context_impl { vocab = llama_model_get_vocab(model_tgt); - n_ctx = llama_n_ctx(ctx_tgt); + // the remaining contexts of the pipeline are created from the same model + { + groups.clear(); + groups.reserve(n_groups); + + for (int g = 0; g < n_groups; ++g) { + groups.emplace_back(new server_group()); + groups[g]->id = g; + } + + groups[0]->ctx = ctx_tgt; + + for (int g = 1; g < n_groups; ++g) { + llama_context_params cparams = common_context_params_to_llama(params_ctx); + + // make sure the extra contexts are identical to the one common_init_from_params made + cparams.n_ctx = llama_n_ctx(ctx_tgt); + cparams.n_seq_max = llama_n_seq_max(ctx_tgt); + + groups[g]->ctx = llama_init_from_model(model_tgt, cparams); + if (groups[g]->ctx == nullptr) { + SRV_ERR("failed to create llama_context for pipeline group %d\n", g); + for (int j = 1; j < g; ++j) { + llama_free(groups[j]->ctx); + groups[j]->ctx = nullptr; + } + groups.clear(); + return false; + } + + SRV_INF("created llama_context for pipeline group %d, n_ctx = %d, n_seq_max = %d\n", + g, (int) llama_n_ctx(groups[g]->ctx), (int) llama_n_seq_max(groups[g]->ctx)); + } + } + + // the total context over all groups, as requested by the user + n_ctx = llama_n_ctx(ctx_tgt) * n_groups; add_bos_token = llama_vocab_get_add_bos(vocab); @@ -1182,7 +1313,9 @@ struct server_context_impl { } // try speculative decoding - if (ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO) { + // note: a common_speculative is bound to one target context, and its code paths yield to + // the task queue, which only one thread may do - so it is off with several groups + if (ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO && n_groups == 1) { try { spec.reset(common_speculative_init(params_base.speculative, params_base.n_parallel)); } catch (const std::exception & e) { @@ -1205,10 +1338,17 @@ struct server_context_impl { for (int i = 0; i < params_base.n_parallel; i++) { server_slot & slot = slots[i]; - slot.id = i; - slot.ctx_tgt = ctx_tgt; - slot.ctx_dft = ctx_dft; - slot.mem.init(ctx_tgt, ctx_dft); + // slots are partitioned contiguously: group g owns slots [g*S, (g+1)*S) + server_group & grp = *groups[i / n_seq_per_group]; + + slot.id = i; + slot.id_group = grp.id; + slot.seq_id = i % n_seq_per_group; + slot.ctx_tgt = grp.ctx; + slot.ctx_dft = ctx_dft; + slot.mem.init(grp.ctx, ctx_dft); + + grp.slots.push_back(&slot); slot.spec = spec.get(); slot.n_ctx = n_ctx_slot; @@ -1263,7 +1403,9 @@ struct server_context_impl { { const int32_t n_batch = llama_n_batch(ctx_tgt); const int32_t n_embd = llama_model_n_embd_inp(model_tgt); - batch.init(std::max(n_batch, params_base.n_parallel), n_embd); + for (auto & grp : groups) { + grp->batch.init(std::max(n_batch, n_seq_per_group), n_embd); + } } if (params_base.cache_ram_mib != 0) { @@ -1315,6 +1457,54 @@ struct server_context_impl { return true; } + // refuse everything we cannot make safe with more than one context, rather than half-support it + bool validate_pipeline_groups(const common_params & params, bool has_spec, bool has_mmproj) const { + auto refuse = [](const char * what) { + SRV_ERR("--pipeline-groups > 1 is not supported together with %s\n", what); + return false; + }; + + if (params.n_parallel < n_groups || params.n_parallel % n_groups != 0) { + SRV_ERR("--parallel (%d) must be a positive multiple of --pipeline-groups (%d)\n", + params.n_parallel, n_groups); + return false; + } + + if (params.n_ctx <= 0) { + SRV_ERR("%s", "--pipeline-groups > 1 requires an explicit context size, pass -c N\n"); + return false; + } + + if (params.n_ctx % n_groups != 0) { + SRV_ERR("--ctx-size (%d) must be a multiple of --pipeline-groups (%d)\n", params.n_ctx, n_groups); + return false; + } + + // a common_speculative and its draft context are bound to one target context + if (has_spec) { + return refuse("speculative decoding (--model-draft / MTP)"); + } + + // mtmd_context is bound to one llama_context + if (has_mmproj) { + return refuse("multimodal (--mmproj)"); + } + + // common_init_from_params() applies the control vector to the context it creates and only + // to that one, so the extra contexts would silently run without it + if (!params.control_vectors.empty()) { + return refuse("--control-vector"); + } + + // entering / leaving the sleeping state destroys and rebuilds the contexts under the + // running group threads + if (params.sleep_idle_seconds >= 0) { + return refuse("--sleep-idle"); + } + + return true; + } + // unlike load_model(), this is only called once during initialization bool init() { GGML_ASSERT(ctx_tgt != nullptr); @@ -1327,7 +1517,14 @@ struct server_context_impl { return process_single_task(std::move(task), is_yielding); }); queue_tasks.on_update_slots([this]() { - update_slots(); + if (n_groups > 1) { + // each pipeline group runs its own update loop on its own thread + return; + } + if (groups.empty()) { + return; // no model loaded + } + update_slots(*groups[0]); }); queue_tasks.on_sleeping_state([this](bool sleeping) { handle_sleeping_state(sleeping); @@ -1424,6 +1621,132 @@ struct server_context_impl { return true; } + // Holds the engine so that the caller can look at the slot state safely, and, on request, + // waits for the in-flight decode of the groups whose context the caller is going to touch. + // Constructing this is a no-op when there is a single group: the single update loop and the + // task processing then run on the same thread, exactly as before. + // + // Taking mtx_engine already keeps every group out of a new iteration, so the slot state is + // stable as soon as the guard exists. Only touching a llama_context needs more than that, + // and only for the group that owns it: wait_for() blocks until that group's decode is done + // while the other groups keep computing. Tasks that touch no context wait for nobody. + struct engine_guard { + server_context_impl * srv = nullptr; + std::unique_lock lk; + + explicit engine_guard(server_context_impl * srv_) { + if (srv_->n_groups <= 1) { + return; + } + + srv = srv_; + lk = std::unique_lock(srv->mtx_engine); + + // ask every group to stop at the start of its next iteration, so that the slot state + // cannot change under us while wait_for() releases the lock + for (auto & grp : srv->groups) { + grp->n_pause_req++; + } + } + + // wait until this group is not inside llama_decode, so its context can be touched + void wait_for(int id_group) { + if (srv == nullptr) { + return; + } + + GGML_ASSERT(id_group >= 0 && id_group < (int) srv->groups.size()); + server_group * grp = srv->groups[id_group].get(); + + srv->cv_engine.wait(lk, [&] { return !grp->busy; }); + } + + // wait for every group, for tasks that are not tied to one slot + void wait_for_all() { + if (srv == nullptr) { + return; + } + + srv->cv_engine.wait(lk, [&] { + for (auto & grp : srv->groups) { + if (grp->busy) { + return false; + } + } + return true; + }); + } + + ~engine_guard() { + if (srv == nullptr) { + return; + } + + for (auto & grp : srv->groups) { + grp->n_pause_req--; + } + + lk.unlock(); + srv->cv_engine.notify_all(); + } + + engine_guard(const engine_guard &) = delete; + engine_guard & operator=(const engine_guard &) = delete; + }; + + // the decode loop of one pipeline group, only used when n_groups > 1 + void group_loop(server_group & grp) { + while (true) { + { + std::unique_lock lk(mtx_engine); + if (groups_stop) { + return; + } + } + + if (update_slots(grp)) { + continue; + } + + // nothing to do for this group, wait for a task to be assigned to one of its slots + std::unique_lock lk(mtx_engine); + cv_engine.wait_for(lk, std::chrono::milliseconds(5), [&] { return groups_stop; }); + } + } + + void start_groups() { + if (n_groups <= 1) { + return; + } + + groups_stop = false; + + for (auto & grp : groups) { + server_group * g = grp.get(); + g->thread = std::thread([this, g]() { group_loop(*g); }); + } + + SRV_INF("started %d pipeline group decode threads\n", n_groups); + } + + void stop_groups() { + if (n_groups <= 1) { + return; + } + + { + std::unique_lock lk(mtx_engine); + groups_stop = true; + } + cv_engine.notify_all(); + + for (auto & grp : groups) { + if (grp->thread.joinable()) { + grp->thread.join(); + } + } + } + server_slot * get_slot_by_id(int id_slot) { // note: allow id_slot to be out of bounds (wrap around) id_slot = id_slot % slots.size(); @@ -1571,14 +1894,17 @@ struct server_context_impl { // - smarter decision which slot to clear (LRU or longest prompt?) // - move slot to level 2 cache instead of removing? // - instead of purging, try to store and resume later? - bool try_clear_idle_slots() { + bool try_clear_idle_slots(server_group & grp) { bool res = false; if (!params_base.kv_unified) { return res; } - for (auto & slot : slots) { + // only slots of this group, their KV lives in this group's context + for (auto * slot_ptr : grp.slots) { + auto & slot = *slot_ptr; + if (slot.is_processing()) { continue; } @@ -1703,9 +2029,9 @@ struct server_context_impl { // TODO: tmp until backend sampling is fully implemented if (use_backend_sampling) { - llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get())); + llama_set_sampler(slot.ctx_tgt, slot.seq_id, common_sampler_get(slot.smpl.get())); } else { - llama_set_sampler(ctx_tgt, slot.id, nullptr); + llama_set_sampler(slot.ctx_tgt, slot.seq_id, nullptr); } SLT_TRC(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl.get()).c_str()); @@ -1724,7 +2050,7 @@ struct server_context_impl { : SLOT_STATE_STARTED; // reset server kill-switch counter - n_empty_consecutive = 0; + groups[slot.id_group]->n_empty_consecutive = 0; SLT_INF(slot, "processing task, is_child = %d\n", slot.task->is_child()); return true; @@ -1888,12 +2214,12 @@ struct server_context_impl { result.probs.push_back({ cur_p->data[i].id, - common_token_to_piece(ctx_tgt, cur_p->data[i].id, special), + common_token_to_piece(slot.ctx_tgt, cur_p->data[i].id, special), cur_p->data[i].p }); } } else { - std::vector cur = get_token_probabilities(ctx_tgt, idx, n_probs_request); + std::vector cur = get_token_probabilities(slot.ctx_tgt, idx, n_probs_request); const size_t max_probs = cur.size(); const size_t n_probs = std::min(max_probs, n_probs_request); @@ -1911,7 +2237,7 @@ struct server_context_impl { for (size_t i = 0; i < n_probs; i++) { result.probs.push_back({ cur[i].id, - common_token_to_piece(ctx_tgt, cur[i].id, special), + common_token_to_piece(slot.ctx_tgt, cur[i].id, special), cur[i].p }); } @@ -2008,7 +2334,7 @@ struct server_context_impl { res->tokens = std::move(slot.generated_tokens); } res->stats = slot.stats; - res->prompt = slot.task->tokens.detokenize(ctx_tgt, true); + res->prompt = slot.task->tokens.detokenize(slot.ctx_tgt, true); res->response_fields = std::move(slot.task->params.response_fields); res->truncated = slot.truncated; @@ -2031,7 +2357,7 @@ struct server_context_impl { // populate res.probs_output if (slot.task->params.sampling.n_probs > 0) { if (!slot.task->params.stream && slot.stop == STOP_TYPE_WORD) { - const llama_tokens stop_word_toks = common_tokenize(ctx_tgt, slot.stopping_word, false); + const llama_tokens stop_word_toks = common_tokenize(slot.ctx_tgt, slot.stopping_word, false); size_t safe_offset = std::min(slot.generated_token_probs.size(), stop_word_toks.size()); res->probs_output = std::vector( @@ -2061,7 +2387,7 @@ struct server_context_impl { std::vector embd_res(n_embd_out, 0.0f); for (int i = 0; i < batch.n_tokens; ++i) { - if (!batch.logits[i] || batch.seq_id[i][0] != slot.id) { + if (!batch.logits[i] || batch.seq_id[i][0] != slot.seq_id) { continue; } @@ -2101,13 +2427,13 @@ struct server_context_impl { res->n_tokens = slot.task->n_tokens(); for (int i = 0; i < batch.n_tokens; ++i) { - if (!batch.logits[i] || batch.seq_id[i][0] != slot.id) { + if (!batch.logits[i] || batch.seq_id[i][0] != slot.seq_id) { continue; } - const float * embd = llama_get_embeddings_seq(ctx_tgt, batch.seq_id[i][0]); + const float * embd = llama_get_embeddings_seq(slot.ctx_tgt, batch.seq_id[i][0]); if (embd == NULL) { - embd = llama_get_embeddings_ith(ctx_tgt, i); + embd = llama_get_embeddings_ith(slot.ctx_tgt, i); } if (embd == NULL) { @@ -2147,9 +2473,13 @@ struct server_context_impl { return true; } - std::vector get_free_slots(size_t n_slots_needed, int exclude_id_slot) { + std::vector get_free_slots(size_t n_slots_needed, int exclude_id_slot, int id_group) { std::vector free_slots; for (auto & slot : slots) { + // the parent copies its KV into the children, so they must live in the same context + if (slot.id_group != id_group) { + continue; + } if (!slot.is_processing() && slot.id != exclude_id_slot) { free_slots.push_back(&slot); } @@ -2244,8 +2574,8 @@ struct server_context_impl { // this is not true for SWA models: https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225 cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max); - cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - cur.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + cur.update_tgt(slot.ctx_tgt, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + cur.update_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // stash the draft's speculative state with the checkpoint common_speculative_get_state(spec.get(), slot.id, cur.data_spec); @@ -2263,6 +2593,12 @@ struct server_context_impl { return false; } + // with more than one group the update loops run on their own threads. Holding the engine + // is enough to look at and modify the slot state; the cases below additionally wait for + // the in-flight decode of the group whose context they touch, and only for that group. + // no-op with a single group. + engine_guard guard(this); + switch (task.type) { case SERVER_TASK_TYPE_COMPLETION: case SERVER_TASK_TYPE_INFILL: @@ -2299,10 +2635,23 @@ struct server_context_impl { break; } + // from here on the slot's context is touched (prompt cache, KV), so the group + // that owns it has to finish its decode. the other groups keep computing. + guard.wait_for(slot->id_group); + if (task.is_parent()) { // try getting free slots for all child tasks size_t n_child_tasks = task.child_tasks.size(); - std::vector child_slots = get_free_slots(n_child_tasks, slot->id); + // the children take their KV from the parent, so they must fit in the + // parent's group. with a single group this is the limit the request + // schema already enforces, so nothing changes there. + if ((int) n_child_tasks + 1 > n_seq_per_group) { + send_error(task, string_format( + "n_cmpl must not exceed the number of slots per pipeline group (%d)", n_seq_per_group), + ERROR_TYPE_INVALID_REQUEST); + break; + } + std::vector child_slots = get_free_slots(n_child_tasks, slot->id, slot->id_group); if (child_slots.size() < n_child_tasks) { SRV_DBG("not enough free slots for child tasks, n_free = %zu, n_children = %zu, defer task, id_task = %d\n", child_slots.size(), n_child_tasks, id_task); queue_tasks.defer(std::move(task)); @@ -2318,6 +2667,9 @@ struct server_context_impl { } if (params_base.cache_idle_slots) { + // this walks every slot of every group + guard.wait_for_all(); + for (auto & slot : slots) { if (!slot.is_processing()) { SLT_TRC(slot, "%s", "saving idle slot to prompt cache\n"); @@ -2340,6 +2692,7 @@ struct server_context_impl { // release slot linked with the task id for (auto & slot : slots) { if (slot.task && slot.task->id == task.id_target) { + guard.wait_for(slot.id_group); slot.release(); break; } @@ -2360,6 +2713,9 @@ struct server_context_impl { break; } + // the sampler of this slot is used by its group between decodes + guard.wait_for(slot->id_group); + if (task.params.control_action == "reasoning_end") { // the budget sampler only exists when reasoning control was armed if (!slot->task->params.sampling.reasoning_control) { @@ -2441,6 +2797,9 @@ struct server_context_impl { break; } + // reads this slot's KV out of its context + guard.wait_for(slot->id_group); + const int64_t t_start = ggml_time_us(); std::string filename = task.slot_action.filename; @@ -2456,7 +2815,7 @@ struct server_context_impl { GGML_ASSERT(packed.size() % sizeof(llama_token) == 0); const size_t nwrite = llama_state_seq_save_file( - ctx_tgt, filepath.c_str(), slot->id, + slot->ctx_tgt, filepath.c_str(), slot->seq_id, reinterpret_cast(packed.data()), packed.size() / sizeof(llama_token)); if (nwrite == 0) { send_error(task, "Unable to save slot", ERROR_TYPE_SERVER); @@ -2491,6 +2850,9 @@ struct server_context_impl { break; } + // writes this slot's KV into its context + guard.wait_for(slot->id_group); + const int64_t t_start = ggml_time_us(); std::string filename = task.slot_action.filename; @@ -2500,10 +2862,10 @@ struct server_context_impl { try { size_t n_packed = 0; llama_tokens packed; - nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, nullptr, 0, &n_packed); + nread = llama_state_seq_load_file(slot->ctx_tgt, filepath.c_str(), slot->seq_id, nullptr, 0, &n_packed); if (nread != 0) { packed.resize(std::max(1, n_packed)); - nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, packed.data(), packed.size(), &n_packed); + nread = llama_state_seq_load_file(slot->ctx_tgt, filepath.c_str(), slot->seq_id, packed.data(), packed.size(), &n_packed); } if (nread == 0) { throw std::runtime_error("No available space in KV cache or invalid slot save file"); @@ -2516,7 +2878,7 @@ struct server_context_impl { throw std::runtime_error("Restored prompt does not fit in the slot context"); } - if (!restored.validate(ctx_tgt)) { + if (!restored.validate(slot->ctx_tgt)) { throw std::runtime_error("Invalid tokens in slot save file"); } @@ -2556,6 +2918,9 @@ struct server_context_impl { break; } + // prompt_clear() drops this slot's KV from its context + guard.wait_for(slot->id_group); + // Erase token cache const size_t n_erased = slot->prompt.tokens.size(); @@ -2595,6 +2960,9 @@ struct server_context_impl { } break; case SERVER_TASK_TYPE_SET_LORA: { + // the adapters are applied to every context + guard.wait_for_all(); + auto new_loras = construct_lora_list(task.set_lora); // logging for (size_t i = 0; i < new_loras.size(); ++i) { @@ -2635,11 +3003,11 @@ struct server_context_impl { } } - void abort_all_slots(const std::string & reason) { - for (auto & slot : slots) { - if (slot.is_processing()) { - send_error(slot, reason, ERROR_TYPE_SERVER); - slot.release(); + void abort_all_slots(server_group & grp, const std::string & reason) { + for (auto * slot : grp.slots) { + if (slot->is_processing()) { + send_error(*slot, reason, ERROR_TYPE_SERVER); + slot->release(); } } } @@ -2674,7 +3042,24 @@ struct server_context_impl { }; #endif - void update_slots() { + // runs one iteration of the decode loop of a single pipeline group + // returns true if the group had work to do + bool update_slots(server_group & grp) { + // shadow the single-context members - everything below operates on this group only + auto * ctx_tgt = grp.ctx; + auto & batch = grp.batch; + auto & slots = grp.slots; + + // when there is only one group there is only one thread and this lock is never engaged + std::unique_lock lk; + if (n_groups > 1) { + lk = std::unique_lock(mtx_engine); + cv_engine.wait(lk, [&]{ return groups_stop || grp.n_pause_req == 0; }); + if (groups_stop) { + return false; + } + } + #ifdef DEBUG_TIMINGS static int64_t t_prev = 0; int64_t t_start = ggml_time_us(); @@ -2692,8 +3077,8 @@ struct server_context_impl { { bool all_idle = true; - for (auto & slot : slots) { - if (slot.is_processing()) { + for (auto * slot : slots) { + if (slot->is_processing()) { all_idle = false; break; } @@ -2702,29 +3087,31 @@ struct server_context_impl { if (all_idle) { SRV_TRC("%s", "all slots are idle\n"); - metrics_flush_idle(); + metrics_flush_idle(grp); - return; // skip further processing + return false; // skip further processing - } else { + } else if (n_groups == 1) { SRV_DBG("%s", "posting NEXT_RESPONSE\n"); server_task task(SERVER_TASK_TYPE_NEXT_RESPONSE); task.id = queue_tasks.get_new_id(); queue_tasks.post(std::move(task)); } + // note: with more than one group each group drives its own loop, so there is no need + // to keep the shared task loop spinning } try { scoped_timer t(t_pre_decode, n_pre_decode); - pre_decode(); + pre_decode(grp); batch.render(); } catch (const std::exception & e) { SRV_ERR("pre_decode() failed: %s\n", e.what()); - abort_all_slots("pre_decode() failed: " + std::string(e.what())); + abort_all_slots(grp, "pre_decode() failed: " + std::string(e.what())); // the batch is half-built and not rendered, skip now to avoid UB - return; + return true; } GGML_ASSERT(batch.slot_batched || batch.size() == 0); @@ -2758,7 +3145,7 @@ struct server_context_impl { // TODO @ngxson : maybe handle n_batch == 1 here instead of inside decode() batch_view = batch.get_view(off, n_tokens); - bool ok = decode(n_batch, off, batch_view); + bool ok = decode(grp, lk, n_batch, off, batch_view); #ifdef DEBUG_TIMINGS llama_synchronize(ctx_tgt); #endif @@ -2775,22 +3162,28 @@ struct server_context_impl { } } catch (const std::exception & e) { SRV_ERR("decode() failed: %s\n", e.what()); - abort_all_slots("decode() failed: " + std::string(e.what())); + abort_all_slots(grp, "decode() failed: " + std::string(e.what())); break; // stop any further processing } try { scoped_timer t(t_post_decode, n_post_decode); - post_decode(n_tokens, off, batch_view); + post_decode(grp, n_tokens, off, batch_view); } catch (const std::exception & e) { SRV_ERR("post_decode() failed: %s\n", e.what()); - abort_all_slots("post_decode() failed: " + std::string(e.what())); + abort_all_slots(grp, "post_decode() failed: " + std::string(e.what())); break; // stop any further processing } } + + return true; } - void pre_decode() { + void pre_decode(server_group & grp) { + auto * ctx_tgt = grp.ctx; + auto & batch = grp.batch; + auto & slots = grp.slots; + (void) ctx_tgt; // apply context-shift if needed // TODO: simplify and improve iterate(slots, [&](server_slot & slot) { @@ -2832,8 +3225,8 @@ struct server_context_impl { SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard); - slot.mem.seq_rm (slot.id, n_keep , n_keep + n_discard); - slot.mem.seq_add(slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); + slot.mem.seq_rm (slot.seq_id, n_keep , n_keep + n_discard); + slot.mem.seq_add(slot.seq_id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); // add generated tokens to cache // ref: https://github.com/ggml-org/llama.cpp/pull/16818#discussion_r2473269481 @@ -2900,11 +3293,11 @@ struct server_context_impl { slot.spec_ckpt.update_pos( slot.prompt.n_tokens(), - llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id), - llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id)); + llama_memory_seq_pos_min(llama_get_memory(slot.ctx_tgt), slot.seq_id), + llama_memory_seq_pos_max(llama_get_memory(slot.ctx_tgt), slot.seq_id)); if (use_ckpt_dft) { - slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + slot.spec_ckpt.update_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } slot.spec_prompt = slot.prompt.tokens.get_text_tokens(); @@ -2943,11 +3336,11 @@ struct server_context_impl { if (ctx_dft) { if (use_ckpt_dft) { - ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } - if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) { - GGML_ABORT("failed to remove sequence %d\n", slot.id); + if (!llama_memory_seq_rm(llama_get_memory(slot.ctx_dft), slot.seq_id, ckpt.pos_max + 1, -1)) { + GGML_ABORT("failed to remove sequence %d\n", slot.seq_id); } } @@ -2962,7 +3355,7 @@ struct server_context_impl { if (use_ckpt_tgt) { //const int64_t t_start = ggml_time_us(); - ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_tgt(slot.ctx_tgt, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); //const int64_t t_total = ggml_time_us() - t_start; //printf("checkpoint total: %f ms\n", t_total / 1000.0); @@ -2974,7 +3367,7 @@ struct server_context_impl { } if (use_ckpt_dft) { - ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } } }); @@ -3150,8 +3543,8 @@ struct server_context_impl { const int64_t kv_shift = (int64_t) head_p - (int64_t) head_c; - slot.mem.seq_rm (slot.id, head_p, head_c); - slot.mem.seq_add(slot.id, head_c, head_c + n_match, kv_shift); + slot.mem.seq_rm (slot.seq_id, head_p, head_c); + slot.mem.seq_add(slot.seq_id, head_c, head_c + n_match, kv_shift); for (size_t i = 0; i < n_match; i++) { slot.prompt.tokens.set_token(head_p + i, slot.prompt.tokens[head_c + i]); @@ -3181,9 +3574,9 @@ struct server_context_impl { const auto pos_min_thold = std::max(0, pos_next - n_swa - (has_new_tokens ? 0 : 1)); if (n_past > 0 && n_past <= slot.prompt.n_tokens()) { - const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id); + const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(slot.ctx_tgt), slot.seq_id); if (pos_min == -1) { - SLT_ERR(slot, "n_past = %d, slot.prompt.tokens.size() = %d, seq_id = %d, pos_min = %d\n", n_past, (int) slot.prompt.tokens.size(), slot.id, pos_min); + SLT_ERR(slot, "n_past = %d, slot.prompt.tokens.size() = %d, seq_id = %d, pos_min = %d\n", n_past, (int) slot.prompt.tokens.size(), slot.seq_id, pos_min); GGML_ABORT("pos_min == -1, but n_past > 0 - should not happen: https://github.com/ggml-org/llama.cpp/pull/13833#discussion_r2116181237"); } @@ -3250,8 +3643,8 @@ struct server_context_impl { if (!do_reset) { // restore the context checkpoint - it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + it->load_tgt(slot.ctx_tgt, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + it->load_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // restore the draft's speculative state common_speculative_set_state(spec.get(), slot.id, it->data_spec); @@ -3325,7 +3718,7 @@ struct server_context_impl { SLT_TRC(slot, "cached n_tokens = %d, memory_seq_rm [%d, end)\n", slot.prompt.n_tokens(), p0); - slot.mem.seq_rm(slot.id, p0, -1); + slot.mem.seq_rm(slot.seq_id, p0, -1); // If using an alora, there may be uncached tokens that come // before the invocation sequence. When this happens, the @@ -3371,7 +3764,7 @@ struct server_context_impl { // process the mtmd chunk // note: it submits its own decode, potentially be async // so the timing is queued and flushed on the next sync - metrics_pre_decode(); + metrics_pre_decode(grp); // encode on the worker thread, so we can still handle metrics tasks size_t n_tokens_out = 0; @@ -3387,7 +3780,7 @@ struct server_context_impl { return; // the slot is done, skip it entirely } - metrics_queue_prompt(n_tokens_out); + metrics_queue_prompt(grp, n_tokens_out); slot.stats.n_prompt_processed += n_tokens_out; slot.stats.update_prompt_last(); @@ -3423,7 +3816,7 @@ struct server_context_impl { // embedding requires all tokens in the batch to be output; // MTP also wants logits at every prompt position so the // streaming hook can mirror t_h_nextn into ctx_dft. - add_ok &= batch.add(slot.id, + add_ok &= batch.add(slot.id, slot.seq_id, cur_tok, /* pos = */ slot.prompt.tokens.pos_next(), /* output = */ slot.need_embd(), @@ -3493,8 +3886,8 @@ struct server_context_impl { } } - const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id); - const auto pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id); + const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(slot.ctx_tgt), slot.seq_id); + const auto pos_max = llama_memory_seq_pos_max(llama_get_memory(slot.ctx_tgt), slot.seq_id); // nothing to checkpoint yet // TODO: is this check needed? @@ -3528,21 +3921,25 @@ struct server_context_impl { // returns true = success ; false = retry with smaller batch size // throw std::runtime_error on fatal error - bool decode(int32_t & n_batch, int32_t off, llama_batch & batch_view) { + bool decode(server_group & grp, std::unique_lock & lk, int32_t & n_batch, int32_t off, llama_batch & batch_view) { + auto * ctx_tgt = grp.ctx; + auto & batch = grp.batch; + auto & slots = grp.slots; + SRV_DBG("n_batch (effective) = %d, off = %d\n", n_batch, off); - metrics_pre_decode(); + metrics_pre_decode(grp); if (batch.size() == 0) { SRV_WRN("%s", "no tokens to decode\n"); - if (++n_empty_consecutive > 3) { + if (++grp.n_empty_consecutive > 3) { GGML_ABORT("fatal error - please provide logs and repro in %s\n", "https://github.com/ggml-org/llama.cpp/pull/20277"); } return true; // nothing to decode } else { - n_empty_consecutive = 0; + grp.n_empty_consecutive = 0; } // TODO @ngxson : dft model may have different n_embd than the tgt model, so we check & reject if that's the case @@ -3559,15 +3956,43 @@ struct server_context_impl { has_output |= batch.tokens[i].output; } - // yield to the queue, so we can still handle metrics tasks while decoding - // note: the sync is done here too, so that the wait is also covered by the yield int ret = 0; - queue_tasks.yield_to_queue([&]() { + if (n_groups > 1) { + // release the engine for the duration of the compute - this is the whole point of the + // feature: while this group is on one stage of the layer split, the other group can + // run its own pre_decode / post_decode and submit its batch to the other stage + // note: RAII, so a throwing decode cannot leave the group marked busy forever, nor + // return to the caller's error handling without the engine lock held + struct decode_window { + server_context_impl * srv; + server_group * grp; + std::unique_lock * lk; + decode_window(server_context_impl * srv, server_group * grp, std::unique_lock * lk) + : srv(srv), grp(grp), lk(lk) { + grp->busy = true; + lk->unlock(); + } + ~decode_window() { + lk->lock(); + grp->busy = false; + srv->cv_engine.notify_all(); + } + } window(this, &grp, &lk); + ret = llama_decode(ctx_tgt, batch_view); if (ret == 0 && has_output) { llama_synchronize(ctx_tgt); } - }); + } else { + // yield to the queue, so we can still handle metrics tasks while decoding + // note: the sync is done here too, so that the wait is also covered by the yield + queue_tasks.yield_to_queue([&]() { + ret = llama_decode(ctx_tgt, batch_view); + if (ret == 0 && has_output) { + llama_synchronize(ctx_tgt); + } + }); + } if (ret != 0) { { @@ -3593,14 +4018,14 @@ struct server_context_impl { if (!err.empty()) { SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret); - for (auto & slot : slots) { - if (slot.is_processing()) { - send_error(slot, err); - slot.release(); + for (auto * slot : slots) { + if (slot->is_processing()) { + send_error(*slot, err); + slot->release(); // note: it's complicated to keep track of how much of the current batch has been // processed before the error occurred, so we simply clear the entire context - slot.prompt_clear(); + slot->prompt_clear(); } } @@ -3610,7 +4035,7 @@ struct server_context_impl { } // retry with half the batch size to try to find a free slot in the KV cache - if (!try_clear_idle_slots()) { + if (!try_clear_idle_slots(grp)) { n_batch /= 2; } @@ -3619,12 +4044,14 @@ struct server_context_impl { return false; // retry with the updated n_batch } else { // success, apply batch metrics - metrics_post_decode(off, batch_view.n_tokens, has_output); + metrics_post_decode(grp, off, batch_view.n_tokens, has_output); } // TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL] // for now, always re-evaluate for simplicity // ref: https://github.com/ggml-org/llama.cpp/pull/22728#issuecomment-4400925384 + // note: speculative decoding is refused with more than one group, so this always runs on + // the main thread and yield_to_queue() is safe here if (spec) { bool ok = true; queue_tasks.yield_to_queue([&]() { @@ -3640,12 +4067,14 @@ struct server_context_impl { } // handle `n_cmpl > 1` tasks - when the main prompt is processed, activate all child tasks too - for (auto & slot : slots) { + // note: children are always in the same group as the parent, see get_free_slots() + for (auto * slot_ptr : slots) { + auto & slot = *slot_ptr; if (slot.state == SLOT_STATE_DONE_PROMPT && slot.task->is_parent()) { std::vector children; - for (auto & other : slots) { - if (other.state == SLOT_STATE_WAIT_OTHER && slot.task->id == other.task->id_parent) { - children.push_back(&other); + for (auto * other : slots) { + if (other->state == SLOT_STATE_WAIT_OTHER && slot.task->id == other->task->id_parent) { + children.push_back(other); } } @@ -3665,7 +4094,12 @@ struct server_context_impl { return true; } - void post_decode(int32_t n_batch_tokens, int32_t off, llama_batch & batch_view) { + void post_decode(server_group & grp, int32_t n_batch_tokens, int32_t off, llama_batch & batch_view) { + // shadow the single-context members, as update_slots() does + auto * ctx_tgt = grp.ctx; + auto & slots = grp.slots; + (void) ctx_tgt; + // for checking if a given batch index is inside batch_view auto is_inside_view = [&](int32_t idx) { return idx >= off && idx < off + n_batch_tokens; @@ -3821,13 +4255,13 @@ struct server_context_impl { SLT_DBG(slot, "restoring speculative checkpoint (pos_min = %d, pos_max = %d, size = %zu)\n", ckpt.pos_min, ckpt.pos_max, ckpt.size()); - ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_tgt(slot.ctx_tgt, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); if (slot.ctx_dft) { - ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_dft(slot.ctx_dft, slot.seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } - slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1); + slot.mem.seq_rm(slot.seq_id, ckpt.pos_max + 1, -1); slot.prompt.tokens.keep_first(ckpt.n_tokens); common_sampler_copy(smpl_save.get(), slot.smpl.get()); @@ -3874,7 +4308,7 @@ struct server_context_impl { slot.sampled = ids.back(); // last accepted token SLT_DBG(slot, "add accepted tokens: sampled=%d, ids.size=%zu, n_draft=%zu\n", slot.sampled, ids.size(), n_draft); - slot.mem.seq_rm(slot.id, slot.prompt.tokens.pos_next(), -1); + slot.mem.seq_rm(slot.seq_id, slot.prompt.tokens.pos_next(), -1); for (size_t i = 0; i < ids.size(); ++i) { completion_token_output result; @@ -3915,32 +4349,34 @@ struct server_context_impl { // // call before submitting a decode, so that the queued prompt stats can be timed - void metrics_pre_decode() { - t_decode_start = ggml_time_us(); + void metrics_pre_decode(server_group & grp) { + grp.t_decode_start = ggml_time_us(); } // the batch is submitted, but its compute may not be done yet - void metrics_queue_prompt(uint64_t n_tokens) { + void metrics_queue_prompt(server_group & grp, uint64_t n_tokens) { if (n_tokens == 0) { return; } - if (n_prompt_queued == 0) { - t_prompt_start = t_decode_start; + if (grp.n_prompt_queued == 0) { + grp.t_prompt_start = grp.t_decode_start; } - n_prompt_queued += n_tokens; + grp.n_prompt_queued += n_tokens; } // call only after the context is synchronized, otherwise the time is meaningless - void metrics_flush_prompt() { - if (n_prompt_queued == 0) { + void metrics_flush_prompt(server_group & grp) { + if (grp.n_prompt_queued == 0) { return; } - metrics.add_prompt(n_prompt_queued, ggml_time_us() - t_prompt_start); - n_prompt_queued = 0; + metrics.add_prompt(grp.n_prompt_queued, ggml_time_us() - grp.t_prompt_start); + grp.n_prompt_queued = 0; } // has_output is computed by the caller, which also already synchronized the context if it is set - void metrics_post_decode(int32_t off, int32_t n_tokens, bool has_output) { + void metrics_post_decode(server_group & grp, int32_t off, int32_t n_tokens, bool has_output) { + auto & batch = grp.batch; + metrics.n_decode++; for (const auto & slot : slots) { if (slot.is_processing()) { @@ -3969,11 +4405,11 @@ struct server_context_impl { } } - metrics_queue_prompt(n_prompt_tokens); + metrics_queue_prompt(grp, n_prompt_tokens); if (has_output) { // the context is already synchronized, so the timings are correct - metrics_flush_prompt(); + metrics_flush_prompt(grp); } // advance the prompt timing of the slots that had tokens in this batch @@ -3989,13 +4425,13 @@ struct server_context_impl { } // flush any queued prompt metrics if all slots are now idle - void metrics_flush_idle() { - if (n_prompt_queued == 0) { + void metrics_flush_idle(server_group & grp) { + if (grp.n_prompt_queued == 0) { return; } - llama_synchronize(ctx_tgt); - metrics_flush_prompt(); + llama_synchronize(grp.ctx); + metrics_flush_prompt(grp); } void metrics_on_prediction(const server_slot & slot) { @@ -4035,7 +4471,17 @@ bool server_context::load_model(common_params & params) { void server_context::start_loop() { auto & params = impl->params_base; + + // no-op unless --pipeline-groups > 1 + impl->start_groups(); + impl->queue_tasks.start_loop(params.sleep_idle_seconds * 1000); + + impl->stop_groups(); +} + +void server_context::set_pipeline_groups(int n_groups) { + impl->n_pipeline_groups_req = n_groups; } void server_context::terminate() { diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 5d464b8e8cb..315c96faa85 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -110,6 +110,11 @@ struct server_context { // note: must be set before load_model() is called void set_state_callback(server_state_callback_t callback); + + // number of pipeline groups, i.e. independent llama_contexts over the one model, each with its + // own slots, batch and decode thread (--pipeline-groups, default 1) + // note: must be set before load_model() is called + void set_pipeline_groups(int n_groups); }; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 5fe2729ba1b..45c8d7f7005 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -14,8 +14,11 @@ #include #include +#include +#include #include #include +#include #include // for std::thread::hardware_concurrency #if defined(_WIN32) @@ -25,6 +28,48 @@ static std::function shutdown_handler; static std::atomic_flag is_terminating = ATOMIC_FLAG_INIT; +// --pipeline-groups N: run the slots over N independent llama_contexts of the same model, each +// with its own batch and decode thread. Useful with a layer split over two nodes (--rpc), where a +// single context leaves each stage idle for half of every decode step. +// The option is parsed here instead of in common/arg.cpp because it only means anything for the +// server; everything it changes lives under tools/server. +static int g_pipeline_groups = 1; + +static void server_take_pipeline_groups(int & argc, char ** argv) { + static const char * opt = "--pipeline-groups"; + const size_t opt_len = strlen(opt); + + int n_kept = 1; + + for (int i = 1; i < argc; i++) { + const std::string arg = argv[i]; + + if (arg == opt) { + if (i + 1 >= argc) { + fprintf(stderr, "error: %s requires a value\n", opt); + exit(1); + } + g_pipeline_groups = std::atoi(argv[++i]); + continue; + } + + if (arg.size() > opt_len + 1 && arg.compare(0, opt_len, opt) == 0 && arg[opt_len] == '=') { + g_pipeline_groups = std::atoi(arg.c_str() + opt_len + 1); + continue; + } + + argv[n_kept++] = argv[i]; + } + + argc = n_kept; + argv[n_kept] = nullptr; + + if (g_pipeline_groups < 1) { + fprintf(stderr, "error: %s must be >= 1\n", opt); + exit(1); + } +} + static inline void signal_handler(int signal) { if (is_terminating.test_and_set()) { // in case it hangs, we can force terminate the server by hitting Ctrl+C twice @@ -96,6 +141,9 @@ int llama_server(int argc, char ** argv) { // own arguments required by this example common_params params; + // strip the server-only --pipeline-groups before the common parser sees it + server_take_pipeline_groups(argc, argv); + common_init(); // start the stream session manager GC right after common init, before any HTTP route can @@ -168,6 +216,7 @@ int llama_server(common_params & params, int argc, char ** argv) { // struct that contains llama context and inference server_context ctx_server; + ctx_server.set_pipeline_groups(g_pipeline_groups); server_http_context ctx_http; if (!ctx_http.init(params)) {