diff --git a/.gitignore b/.gitignore index 9b589615a402..66a4bed61e9d 100644 --- a/.gitignore +++ b/.gitignore @@ -153,3 +153,4 @@ a.out.* AGENTS.local.md .pi/SYSTEM.md +bench/ diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 40cea024c3d2..5b68cce0d2c5 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -8,7 +8,10 @@ extern "C" { #endif - #define GGML_BACKEND_API_VERSION 2 + // 3: ggml_backend_i gained cpy_tensor_from_async. A dynamically loaded backend compiled + // against version 2 supplies a struct one member short, so the version has to move with it + // for the load-time check to reject that pairing instead of reading past the end. + #define GGML_BACKEND_API_VERSION 3 // // Backend buffer type @@ -138,6 +141,15 @@ extern "C" { // (optional) sort/optimize the nodes in the graph void (*graph_optimize) (ggml_backend_t backend, struct ggml_cgraph * cgraph); + + // (optional) copy a tensor of this backend into another backend's tensor, driven from the + // source side. Kept separate from cpy_tensor_async because every existing implementation + // of that is written as a destination side handler: it casts backend_dst to its own + // context before deciding anything, so dispatching a source side call there would + // reinterpret a foreign backend_dst, or dereference a null one. Implement this only if + // the source role is genuinely supported. Appended last so backends that list the + // members positionally are unaffected. + bool (*cpy_tensor_from_async)(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst); }; struct ggml_backend { diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index e519bdf50a1b..f0a97a11bff8 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -499,6 +499,27 @@ void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor } } +// src is only asked when its implementation differs: same-type backends share one implementation that already declined. +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; + } + } + + // backend_dst must be a real backend before anything is delegated to the source. An + // implementation of the source role has to inspect the destination to decide whether it can + // help, so handing it a null destination pushes that dereference into every implementer. + // Guarding here keeps the invariant in one place instead of depending on all of them. + if (backend_src != NULL && backend_dst != NULL && backend_src->iface.cpy_tensor_from_async != NULL) { + if (backend_src->iface.cpy_tensor_from_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 +528,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 +1747,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 69a8a08ae172..d33cf62b716e 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, }; @@ -84,11 +87,20 @@ struct rpc_msg_hello_req { uint8_t conn_caps[RPC_CONN_CAPS_SIZE]; }; +// Server feature flags, carried in the byte that used to be pure padding in the HELLO response. +// Advertised here rather than by bumping the protocol minor: a client rejects any server whose +// minor exceeds its own, so a bump locks out every already-deployed older client even when the new +// command is purely additive and such a client would never send it. This byte is fixed size, +// already on the wire, and read as padding by existing clients, which see zero. +enum rpc_srv_flag { + RPC_SRV_FLAG_BATCHED_GET = 1 << 0, // supports RPC_CMD_GET_TENSORS +}; + struct rpc_msg_hello_rsp { uint8_t major; uint8_t minor; uint8_t patch; - uint8_t padding; + uint8_t srv_flags; uint8_t conn_caps[RPC_CONN_CAPS_SIZE]; }; @@ -176,6 +188,21 @@ struct rpc_msg_get_tensor_req { uint64_t size; }; +// Ceiling on one batched GET_TENSORS response. Per-entry validation already bounds each region by +// a really allocated buffer, but nothing bounds the number of entries naming the same large buffer, +// so a handful of valid entries could still ask for a terabyte. This is far above any legitimate +// batch: the deferred gets this batches are activations and single tensor regions, orders of +// magnitude smaller, so the limit costs nothing real while keeping the sum in a range where the +// checked addition below cannot wrap. +static constexpr size_t MAX_GET_TENSORS_RESPONSE = 4ull * 1024 * 1024 * 1024; // 4 GiB + +// GET_TENSORS request: | n_entries (4 bytes) | n_entries x entry |, response: regions in entry order +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 +239,6 @@ struct ggml_backend_rpc_device_context { uint32_t device; std::string name; std::string description; - uint64_t last_graph_uid; }; struct ggml_backend_rpc_buffer_type_context { @@ -298,9 +324,90 @@ static bool parse_endpoint(const std::string & endpoint, std::string & host, int return true; } +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 ops must keep their wire order, so every send flushes first; the flush sends too, hence the 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) { +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 +421,58 @@ 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); +} + +// The server answers a connection in request order, so the n-th response belongs to the n-th +// response-bearing request. Construct with mtx_send held; always released, so a failed send +// cannot strand the later waiters. +struct rpc_response_ticket { + rpc_conn_state & conn; + uint64_t seq; + + 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, or a later waiter is woken with a response that is not theirs + failed = true; + } + } + + if (failed) { + ticket->wait(); return false; } + + ticket->wait(); + uint64_t out_size; if (!sock->recv_data(&out_size, sizeof(out_size))) { return false; @@ -353,17 +506,32 @@ static bool negotiate_hello(const std::shared_ptr & sock) { return false; } + sock->conn.server_minor = response.minor; + sock->conn.server_flags = response.srv_flags; + sock->update_caps(response.conn_caps); return true; } +// The server serves one connection of a client at a time, so opening a second connection to a live +// endpoint blocks until the first closes: 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 +554,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 +624,248 @@ static rpc_tensor serialize_tensor(const ggml_tensor * tensor) { return result; } + +struct rpc_staging { + ggml_backend_buffer_t buffer = nullptr; + uint8_t * base = nullptr; + size_t capacity = 0; + size_t used = 0; + + // Every event recorded for a copy still reading this arena. This was a single slot, which + // silently assumed that one arena is written from one backend: the last event recorded would + // then be ordered after all the earlier ones, so waiting on it alone was enough. The arena is + // keyed by socket, meaning one per endpoint, and with --pipeline-groups > 1 each llama_context + // drives its own destination backend and therefore its own stream. Two groups copying over the + // same endpoint then overwrote each other's event here, and the wrap below waited only on the + // survivor before recycling the whole arena, while the other group's copy could still be + // reading its region. Events on different streams have no ordering between them, so the fix is + // to keep all of them and wait for all of them. + std::vector outstanding; + + std::vector events; + size_t events_used = 0; + + // Regions handed out by rpc_staging_alloc() that are not yet represented in `outstanding`. + // The allocation and the event that makes the region trackable cannot happen under one lock, + // because the copy in between blocks and holding the mutex across it would serialize the + // pipeline groups this arena exists to keep concurrent. So the region is reserved instead: a + // wrap waits for every reservation to be handed over before it resets `used`, otherwise it + // could give the same bytes to another group while the first is still reading into them or + // its destination stream is still consuming them. + size_t in_flight = 0; + std::condition_variable cv_reserved; + + rpc_staging() = default; + + // The arena is pinned host memory and the events are driver objects. Nothing used to free + // either: the map is keyed by a raw socket_t * and kept its entry after the socket expired + // from the weak cache, so a process that opens and closes RPC connections over its lifetime + // retained every connection's pinned allocation until it exited. + ~rpc_staging() { + for (ggml_backend_event_t ev : outstanding) { + ggml_backend_event_synchronize(ev); + } + for (ggml_backend_event_t ev : events) { + ggml_backend_event_free(ev); + } + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + rpc_staging(const rpc_staging &) = delete; + rpc_staging & operator=(const rpc_staging &) = delete; +}; + +static std::mutex rpc_staging_mutex; +static std::unordered_map rpc_staging_map; + +// called from socket_t's destructor, which is where a raw-pointer key stops being valid. Without +// this the entry outlives the socket, leaks its pinned buffer and events, and a later socket +// allocated at the same address would inherit a stale arena. +void rpc_staging_drop(socket_t * sock) { + std::lock_guard lock(rpc_staging_mutex); + rpc_staging_map.erase(sock); +} + +// caller must hold rpc_staging_mutex through `lock`. On success the returned region is reserved +// and the caller must hand it over with rpc_staging_commit() once it is safe to recycle. +static uint8_t * rpc_staging_alloc(std::unique_lock & lock, rpc_staging & st, + ggml_backend_buffer_type_t host_buft, size_t size) { + if (st.used + size > st.capacity) { + // Wait for regions that have been handed out but are not yet in `outstanding`. Without + // this the synchronize below sees an incomplete picture and the reset hands live bytes to + // the next caller. The wait releases the mutex, and every reservation is released without + // needing anything from this thread, so it cannot deadlock against us. + st.cv_reserved.wait(lock, [&st] { return st.in_flight == 0; }); + + for (ggml_backend_event_t ev : st.outstanding) { + ggml_backend_event_synchronize(ev); + } + st.outstanding.clear(); + st.used = 0; + // recycle the event pool here too, not only in rpc_flush_deferred(): that reset is behind + // rpc_flush_deferred_guarded()'s empty-queue early return, and the RPC-to-local copy path + // blocks in ggml_backend_tensor_get() with nothing deferred, so it never ran and every + // copy allocated a new backend event that was never reused. Recycling is safe because the + // loop above waited on every outstanding event, not merely the most recent one. + st.events_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; + st.in_flight++; + return ptr; +} + +// caller must hold rpc_staging_mutex. Ends the reservation taken by rpc_staging_alloc(), either +// because the region's event is now in `outstanding` or because the copy has been synchronized. +static void rpc_staging_commit(rpc_staging & st) { + GGML_ASSERT(st.in_flight > 0); + st.in_flight--; + st.cv_reserved.notify_all(); +} + +// 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; +} + +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; + } + + // an RDMA completion carries exactly one send: reading it back in pieces overruns and hangs + 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)); + + 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 +1084,208 @@ 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; + // find_socket, not get_socket: nothing connected means nothing queued + auto sock = find_socket(rpc_ctx->endpoint); + if (sock != nullptr) { + rpc_flush_deferred_guarded(sock); + } +} + +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_flags & RPC_SRV_FLAG_BATCHED_GET); +} + +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); +} + +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; + } + + // ggml_backend_is_rpc() tolerates null, but nothing below does: a null destination against an + // RPC source passes the differing-kind test and then reaches ggml_backend_get_device(other). + if (backend_src == nullptr || backend_dst == nullptr) { + return false; + } + + const bool src_is_rpc = ggml_backend_is_rpc(backend_src); + const bool dst_is_rpc = ggml_backend_is_rpc(backend_dst); + + 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); + + // staging must 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; + } + + // a backend's async entry points only accept tensors in its own default buffer type + 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) { + 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::unique_lock lock(rpc_staging_mutex); + rpc_staging & st = rpc_staging_map[sock.get()]; + staging = rpc_staging_alloc(lock, st, host_buft, size); + if (staging == nullptr) { + return false; + } + event = rpc_staging_event(st, other_dev); + if (event == nullptr) { + // the region was reserved, so it has to be handed back or a later wrap waits forever + rpc_staging_commit(st); + } + } + if (event == nullptr) { + return false; + } + + ggml_backend_tensor_get_async(backend_src, src, staging, 0, size); + ggml_backend_event_record(event, backend_src); + + { + // Make the region trackable before releasing it. This event was previously known only + // to the deferred op, so a wrap driven from the other staging path reset `used` without + // waiting for this read to land. It is still the flush on wrap above that keeps the + // bytes alive until the deferred send has copied them out; this only ensures a reset + // never happens while the fill is still in flight. + std::lock_guard lock(rpc_staging_mutex); + rpc_staging & st = rpc_staging_map[sock.get()]; + st.outstanding.push_back(event); + rpc_staging_commit(st); + } + + 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; + } + + if (backend_dst->iface.set_tensor_async == nullptr) { + return false; + } + + // a wrap here has to send whatever is already queued before it can recycle the arena, and that + // cannot be done from inside rpc_staging_alloc() because rpc_flush_deferred() takes + // rpc_staging_mutex itself. Predict it the same way the other staging path does. + { + 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; + { + std::unique_lock lock(rpc_staging_mutex); + rpc_staging & st = rpc_staging_map[sock.get()]; + staging = rpc_staging_alloc(lock, 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.outstanding.push_back(event); + } else { + ggml_backend_synchronize(backend_dst); + } + // the destination stream's progress is now either tracked by the event or already complete + rpc_staging_commit(st); + } + return true; } static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::vector & tensors, std::unordered_set & visited) { @@ -731,21 +1341,31 @@ 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); + + // a connection is shared across llama_contexts, so the uid check must be under the same lock + // as the send, or RECOMPUTE re-runs a 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 +1373,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, @@ -766,6 +1386,9 @@ static ggml_backend_i ggml_backend_rpc_interface = { /* .event_record = */ NULL, /* .event_wait = */ NULL, /* .graph_optimize = */ NULL, + // safe in the source role: it checks ggml_backend_is_rpc on both sides and declines unless + // exactly one of them is an RPC backend, so it never reinterprets a foreign backend_dst + /* .cpy_tensor_from_async = */ ggml_backend_rpc_cpy_tensor_async, }; ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, uint32_t device) { @@ -864,6 +1487,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); @@ -896,7 +1520,9 @@ 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; - LOG_DBG("[%s] version: %d.%d.%d\n", __func__, response.major, response.minor, response.patch); + response.srv_flags = RPC_SRV_FLAG_BATCHED_GET; + LOG_DBG("[%s] version: %d.%d.%d flags: 0x%02x\n", __func__, + response.major, response.minor, response.patch, response.srv_flags); } bool rpc_server::get_alloc_size(const rpc_msg_get_alloc_size_req & request, rpc_msg_get_alloc_size_rsp & response) { @@ -1299,6 +1925,90 @@ bool rpc_server::get_tensor(const rpc_msg_get_tensor_req & request, std::vector< return true; } + +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)); + + // Validate every entry before allocating anything. The sizes come straight off the wire, and + // the per-entry bounds check further down only constrains a region against its own source + // buffer, never against the response, so summing first and allocating on that sum was wrong + // two ways. An unchecked sum can exceed anything allocatable and throw an uncaught bad_alloc, + // which terminates the server. Worse, the sum is accumulated into size_t from uint64_t sizes, + // so it can wrap: a small total then allocates a small response while the copy loop below + // still writes entries[i].size bytes at out_offset, running off the end of the heap block. + size_t total = 0; + for (uint32_t i = 0; i < n_entries; i++) { + struct ggml_init_params vparams { + /*.mem_size =*/ ggml_tensor_overhead(), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + ggml_context_ptr vctx { ggml_init(vparams) }; + GGML_ASSERT(vctx != nullptr); + ggml_tensor * t = deserialize_tensor(vctx.get(), &entries[i].tensor); + if (t == nullptr || t->buffer == nullptr) { + GGML_LOG_ERROR("[%s] error deserializing tensor %u\n", __func__, i); + return false; + } + const size_t p0 = (size_t) ggml_backend_buffer_get_base(t->buffer); + const size_t p1 = p0 + ggml_backend_buffer_get_size(t->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; + } + if (entries[i].size > MAX_GET_TENSORS_RESPONSE - total) { // checked add, no wrap + GGML_LOG_ERROR("[%s] batched read of %" PRIu64 " bytes exceeds the response limit\n", + __func__, entries[i].size); + return false; + } + total += (size_t) 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 +2439,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 +2768,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.cpp b/ggml/src/ggml-rpc/transport.cpp index 5ec15dc80c0c..4c3345b1c80a 100644 --- a/ggml/src/ggml-rpc/transport.cpp +++ b/ggml/src/ggml-rpc/transport.cpp @@ -598,7 +598,13 @@ bool socket_t::impl::flush() { socket_t::socket_t(std::unique_ptr p) : pimpl(std::move(p)) {} -socket_t::~socket_t() = default; +// defined in ggml-rpc.cpp, which owns the staging arenas keyed by socket +void rpc_staging_drop(socket_t * sock); + +socket_t::~socket_t() { + // the arena is keyed by this pointer, so this is the last moment the entry can be found + rpc_staging_drop(this); +} bool socket_t::send_data(const void * data, size_t size) { return pimpl->send_data(data, size); diff --git a/ggml/src/ggml-rpc/transport.h b/ggml/src/ggml-rpc/transport.h index 3f747ecffd97..aae0d9202912 100644 --- a/ggml/src/ggml-rpc/transport.h +++ b/ggml/src/ggml-rpc/transport.h @@ -1,18 +1,58 @@ #pragma once +#include #include #include #include +#include +#include +#include struct socket_t; typedef std::shared_ptr socket_ptr; +struct rpc_deferred_op { + enum kind_t { GET, SET } kind; + + // serialized at queue time: the owning graph result is reset once per ubatch, so a pointer would dangle + std::vector tensor_bytes; + + void * data = nullptr; // GET: host destination; SET: host staging source + uint64_t offset = 0; + uint64_t size = 0; + + // SET only: must complete before the staging buffer holds the data. void to keep ggml-backend out. + void * event = nullptr; +}; + static constexpr size_t MAX_CHUNK_SIZE = 1024ull * 1024ull * 1024ull; // 1 GiB static constexpr size_t RPC_CONN_CAPS_SIZE = 24; +// One connection is shared by every backend of an endpoint, across llama_contexts, so: mtx_send +// keeps a message atomic on the wire; seq_* hands responses out in request order without holding +// mtx_send; last_graph_uid stops RECOMPUTE re-running another context's graph. +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; + + uint32_t server_minor = 0; + uint8_t server_flags = 0; + + // lock order: mtx_defer before mtx_send, never the reverse + std::mutex mtx_defer; + std::vector deferred; +}; + struct socket_t { ~socket_t(); + 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 93736c3edfa9..bef9088f39c4 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 a9edbd7be8b4..9a78e0310459 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -35,6 +35,11 @@ #include #endif +#include +#include +#include +#include + constexpr int HTTP_POLLING_SECONDS = 1; static common_speculative_output_limits server_output_limits(const common_params & params) { @@ -69,6 +74,7 @@ struct server_batch { struct token { int32_t id_slot; + int32_t seq_id; llama_token token; llama_pos pos; bool output; @@ -108,22 +114,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 +165,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 +200,10 @@ struct server_batch { struct server_slot { int id; + int id_group = 0; + + int seq_id = 0; + llama_context * ctx_tgt = nullptr; llama_context * ctx_dft = nullptr; @@ -255,8 +265,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 +278,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 +298,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 +361,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 +473,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 +493,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 +686,8 @@ 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); + 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 +790,37 @@ static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch return try_decode(); } +// One llama_context with its own batch, decode loop and contiguous range of slots: with N > 1 one +// group drives one stage of a layer split while another drives the other. N == 1 is the old path. +struct server_group { + int id = 0; + + llama_context * ctx = nullptr; + + // groups run concurrently, so they cannot share one thread pool: common_init_from_params only + // builds and attaches a pool for group 0's context. without this every extra group falls back + // to a disposable pool per graph, which ignores the configured cpu mask, priority and strict + // placement. one pool per concurrently executing context. + common_threadpools threadpools; + + server_batch batch; + + std::vector slots; + + // llama_decode() is async, so these are only valid after a sync. not in server_metrics, which + // is copied as-is into the task result. + int64_t t_decode_start = 0; + int64_t t_prompt_start = 0; + 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 +845,9 @@ struct server_context_impl { server_state_callback_t callback_state = [](server_state, json) -> void {}; + // 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 +879,14 @@ struct server_context_impl { llama_context * ctx_tgt = nullptr; - server_batch batch; + int n_groups = 1; + std::vector> groups; + + int n_seq_per_group = 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 +913,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 +937,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 +1100,49 @@ struct server_context_impl { params_base.load_progress_callback_user_data = &load_progress_text; } - llama_init = common_init_from_params(params_base); + 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) { + // 1/N of the sequences and of the context each: per-slot context and total KV are unchanged + params_ctx.n_parallel = n_seq_per_group; + params_ctx.n_ctx = params_base.n_ctx / n_groups; + } + + // common_init_from_params() fits the model to device memory for the single context it + // creates, but this server then builds n_groups - 1 more contexts from the same model. + // Their KV and compute buffers are invisible to the fit, so without this the fit would + // happily offload enough layers to fill the devices and a later group would fail to + // allocate. Reserve the aggregate memory of the extra contexts up front, the same way + // the mmproj estimate above does. + if (n_groups > 1 && params_ctx.fit_params) { + reserve_memory_for_extra_groups(params_ctx); + } + + llama_init = common_init_from_params(params_ctx); model_tgt = llama_init->model(); ctx_tgt = llama_init->context(); + if (n_groups > 1) { + 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 +1155,44 @@ struct server_context_impl { vocab = llama_model_get_vocab(model_tgt); - n_ctx = llama_n_ctx(ctx_tgt); + { + 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); + + // the extra contexts must be 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)); + + // group 0 already has the pool common_init_from_params attached to ctx_tgt + groups[g]->threadpools.init(groups[g]->ctx, params_ctx); + } + } + + n_ctx = llama_n_ctx(ctx_tgt) * n_groups; add_bos_token = llama_vocab_get_add_bos(vocab); @@ -1182,7 +1309,7 @@ struct server_context_impl { } // try speculative decoding - if (ctx_tgt_seq_rm_type != COMMON_CONTEXT_SEQ_RM_TYPE_NO) { + 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 +1332,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 +1397,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 +1451,120 @@ struct server_context_impl { return true; } + // Account for the pipeline group contexts that are created after common_init_from_params(). + // The fit only ever sees one context, so measure what a single group context costs per device + // and add (n_groups - 1) times that to the per-device margins the fit has to leave free. + void reserve_memory_for_extra_groups(common_params & params_ctx) const { + GGML_ASSERT(n_groups > 1); + + const int n_extra = n_groups - 1; + + auto mparams = common_model_params_to_llama(params_ctx); + auto cparams = common_context_params_to_llama(params_ctx); + + std::vector devs; + uint32_t hp_ngl = 0; + uint32_t hp_n_ctx_train = 0; + uint32_t hp_n_expert = 0; + + common_device_memory_data_vec dmd; + + const int64_t t_start = ggml_time_us(); + + try { + dmd = common_get_device_memory_data(params_ctx.model.path.c_str(), &mparams, &cparams, + devs, hp_ngl, hp_n_ctx_train, hp_n_expert, + params_ctx.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR); + } catch (const std::exception & e) { + SRV_WRN("failed to estimate the memory of the extra pipeline group contexts (%s), " + "the fit will only account for one of the %d groups\n", e.what(), n_groups); + return; + } + + const int64_t t_elapsed = ggml_time_us() - t_start; + + // dmd is indexed like devs, with the host buffers in the extra last entry + if (dmd.size() != devs.size() + 1) { + SRV_WRN("%s", "unexpected memory breakdown size, skipping pipeline group memory reservation\n"); + return; + } + + size_t total_per_ctx = 0; + + // the fit assumes host memory is unlimited, so it is only reported, never reserved + const size_t host_per_ctx = dmd.back().context + dmd.back().compute; + + // note: common_fit_params() indexes the margins by the model device order, which is + // exactly the order of devs / dmd here, not by the global ggml_backend_dev_get() order + for (size_t i = 0; i < devs.size() && i < params_ctx.fit_params_target.size(); i++) { + // model weights are shared by every group, only the KV cache and the compute buffers + // are paid for again by each additional context + const size_t per_ctx = dmd[i].context + dmd[i].compute; + + total_per_ctx += per_ctx; + + if (per_ctx == 0) { + continue; + } + + SRV_DBG("reserving %.2f MiB (%d x %.2f MiB) on device %s for the extra pipeline group contexts\n", + n_extra * per_ctx / (1024.0 * 1024.0), n_extra, per_ctx / (1024.0 * 1024.0), + ggml_backend_dev_name(devs[i])); + + params_ctx.fit_params_target[i] += n_extra * per_ctx; + } + + SRV_INF("fitting %d pipeline groups: one context of n_ctx = %d needs %.2f MiB of device memory " + "(+ %.2f MiB host), reserving %.2f MiB of device memory for the %d extra group(s) (took %.2f ms)\n", + n_groups, params_ctx.n_ctx, total_per_ctx / (1024.0 * 1024.0), host_per_ctx / (1024.0 * 1024.0), + n_extra * total_per_ctx / (1024.0 * 1024.0), n_extra, t_elapsed / 1000.0); + } + + 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 it only to the context it creates + if (!params.control_vectors.empty()) { + return refuse("--control-vector"); + } + + // sleeping 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 +1577,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 group runs its own update loop on its own thread + return; + } + if (groups.empty()) { + return; + } + update_slots(*groups[0]); }); queue_tasks.on_sleeping_state([this](bool sleeping) { handle_sleeping_state(sleeping); @@ -1424,6 +1681,120 @@ struct server_context_impl { return true; } + // Holding mtx_engine keeps every group out of a new iteration, so the slot state is stable as + // soon as the guard exists; touching a llama_context additionally needs wait_for(its group). + 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); + + // stop every group at its next iteration: wait_for() releases the lock meanwhile + for (auto & grp : srv->groups) { + grp->n_pause_req++; + } + } + + 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; }); + } + + 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; + }; + + void group_loop(server_group & grp) { + while (true) { + { + std::unique_lock lk(mtx_engine); + if (groups_stop) { + return; + } + } + + if (update_slots(grp)) { + continue; + } + + 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(); @@ -1451,11 +1822,30 @@ struct server_context_impl { return nullptr; } - server_slot * get_available_slot(const server_task & task) { + // n_slots_needed is the total number of slots the task occupies at once, i.e. n_cmpl: + // the parent slot plus one slot per child task. The children take their KV from the parent, + // so all of them have to come from the parent's pipeline group and only groups with that + // many idle slots are eligible - otherwise a group whose siblings are busy would be picked + // by LRU / LCP and the task deferred while another group sits completely idle. + server_slot * get_available_slot(const server_task & task, bool * out_update_cache, size_t n_slots_needed = 1) { server_slot * ret = nullptr; bool update_cache = false; + std::vector n_free_per_group(n_groups, 0); + + if (n_slots_needed > 1) { + for (const server_slot & slot : slots) { + if (!slot.is_processing()) { + n_free_per_group[slot.id_group]++; + } + } + } + + auto group_has_room = [&](const server_slot & slot) { + return n_slots_needed <= 1 || (size_t) n_free_per_group[slot.id_group] >= n_slots_needed; + }; + // if a specific slot is requested, use it (still goes through cache update logic below) if (task.id_slot != -1) { ret = get_slot_by_id(task.id_slot); @@ -1479,6 +1869,13 @@ struct server_context_impl { continue; } + // skip the slot if its group cannot host all the children as well + if (!group_has_room(slot)) { + SLT_TRC(slot, " - skipping, group %d has %d free slots < %zu needed\n", + slot.id_group, n_free_per_group[slot.id_group], n_slots_needed); + continue; + } + const auto & tokens = slot.prompt.tokens; // skip the slot if it does not contains cached tokens @@ -1526,6 +1923,13 @@ struct server_context_impl { continue; } + // skip the slot if its group cannot host all the children as well + if (!group_has_room(slot)) { + SLT_TRC(slot, " - skipping, group %d has %d free slots < %zu needed\n", + slot.id_group, n_free_per_group[slot.id_group], n_slots_needed); + continue; + } + // select the current slot if the criteria match if (!ret || slot.t_last_used <= t_last) { t_last = slot.t_last_used; @@ -1546,24 +1950,32 @@ struct server_context_impl { // cache prompts only for completion tasks update_cache = update_cache && task.type == SERVER_TASK_TYPE_COMPLETION; - if (update_cache) { - SRV_TRC("%s", "updating prompt cache\n"); + // the update itself is left to update_prompt_cache(), which the caller runs after it + // has waited for this slot's group. prompt_save() and prompt_load() call + // llama_state_seq_* on the slot's context, and with more than one pipeline group + // another slot of that group can still be inside llama_decode() at this point. + // Selecting a slot must not touch its context. + *out_update_cache = update_cache; + } - const int64_t t_start = ggml_time_us(); + return ret; + } - ret->prompt_save(*prompt_cache); + // must be called with the slot's group waited for + void update_prompt_cache(server_slot & slot, const server_task & task) { + SRV_TRC("%s", "updating prompt cache\n"); - if (!ret->prompt_load(*prompt_cache, task.tokens)) { - ret->prompt_clear(); - } + const int64_t t_start = ggml_time_us(); - prompt_cache->update(); + slot.prompt_save(*prompt_cache); - SRV_TRC("prompt cache update took %.2f ms\n", (ggml_time_us() - t_start) / 1000.0); - } + if (!slot.prompt_load(*prompt_cache, task.tokens)) { + slot.prompt_clear(); } - return ret; + prompt_cache->update(); + + SRV_TRC("prompt cache update took %.2f ms\n", (ggml_time_us() - t_start) / 1000.0); } // return true if at least one slot has been cleared @@ -1571,14 +1983,16 @@ 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) { + for (auto * slot_ptr : grp.slots) { + auto & slot = *slot_ptr; + if (slot.is_processing()) { continue; } @@ -1703,9 +2117,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 +2138,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 +2302,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 +2325,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 +2422,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 +2445,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 +2475,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 +2515,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 +2561,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 +2662,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 +2681,8 @@ struct server_context_impl { return false; } + engine_guard guard(this); + switch (task.type) { case SERVER_TASK_TYPE_COMPLETION: case SERVER_TASK_TYPE_INFILL: @@ -2279,7 +2699,20 @@ struct server_context_impl { const int id_task = task.id; - server_slot * slot = get_available_slot(task); + // the children take their KV from the parent, so they must fit its group. + // check this before selecting a slot, otherwise an impossible request would + // find no eligible group and be deferred forever instead of being rejected. + const size_t n_slots_needed = task.child_tasks.size() + 1; + + if (task.is_parent() && (int) n_slots_needed > 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; + } + + bool update_cache = false; + server_slot * slot = get_available_slot(task, &update_cache, n_slots_needed); // // slot scheduling logic @@ -2299,10 +2732,18 @@ struct server_context_impl { break; } + guard.wait_for(slot->id_group); + + // only now is every slot of this group out of llama_decode(), so the + // llama_state_seq_* calls behind the prompt cache can touch its context + if (update_cache) { + update_prompt_cache(*slot, task); + } + 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); + 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 +2759,8 @@ struct server_context_impl { } if (params_base.cache_idle_slots) { + 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 +2783,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 +2804,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 +2888,8 @@ struct server_context_impl { break; } + guard.wait_for(slot->id_group); + const int64_t t_start = ggml_time_us(); std::string filename = task.slot_action.filename; @@ -2456,7 +2905,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 +2940,8 @@ struct server_context_impl { break; } + guard.wait_for(slot->id_group); + const int64_t t_start = ggml_time_us(); std::string filename = task.slot_action.filename; @@ -2500,10 +2951,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 +2967,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 +3007,8 @@ struct server_context_impl { break; } + guard.wait_for(slot->id_group); + // Erase token cache const size_t n_erased = slot->prompt.tokens.size(); @@ -2595,6 +3048,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 +3091,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 +3130,21 @@ struct server_context_impl { }; #endif - void update_slots() { + 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; + + 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 +3162,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 +3172,30 @@ 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; - } 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, no spinning needed } 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 +3229,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 +3246,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 +3309,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 +3377,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 +3420,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 +3439,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 +3451,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 +3627,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 +3658,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 +3727,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 +3802,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 +3848,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 +3864,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 +3900,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 +3970,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 +4005,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 +4040,39 @@ 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 across the compute, so another group can drive the other stage. + // RAII: a throwing decode must not leave the group busy, nor skip re-taking the lock. + 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 so metrics tasks are still handled while decoding; the sync is inside so the wait is covered + 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 +4098,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 +4115,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 +4124,13 @@ 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: several groups refuse speculative decoding, so yield_to_queue() is safe here if (spec) { bool ok = true; queue_tasks.yield_to_queue([&]() { @@ -3640,12 +4146,13 @@ struct server_context_impl { } // handle `n_cmpl > 1` tasks - when the main prompt is processed, activate all child tasks too - for (auto & slot : 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 +4172,11 @@ 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) { + 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 +4332,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 +4385,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 +4426,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 +4482,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 +4502,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 +4548,16 @@ bool server_context::load_model(common_params & params) { void server_context::start_loop() { auto & params = impl->params_base; + + 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 5d464b8e8cb7..c489bb2ac84b 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -110,6 +110,9 @@ struct server_context { // note: must be set before load_model() is called void set_state_callback(server_state_callback_t callback); + + // 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 5fe2729ba1b2..002675a67553 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,53 @@ static std::function shutdown_handler; static std::atomic_flag is_terminating = ATOMIC_FLAG_INIT; +// parsed here rather than in common/arg.cpp: everything --pipeline-groups changes is 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; + + // the router strips this flag from argv before handing argv to server_models, so a child + // spawned by the router would otherwise always run at one group no matter what the operator + // asked for. the router re-exports the resolved value and the child picks it up here, which + // is the same channel LLAMA_ARG_HF_REPO and LLAMA_SERVER_ROUTER_PORT already use. + // an explicit flag on the command line still wins over the inherited value. + if (const char * env = getenv("LLAMA_ARG_PIPELINE_GROUPS")) { + g_pipeline_groups = std::atoi(env); + } + + 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 +146,8 @@ int llama_server(int argc, char ** argv) { // own arguments required by this example common_params params; + 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 +220,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)) { @@ -186,6 +239,22 @@ int llama_server(common_params & params, int argc, char ** argv) { std::optional models_routes{}; if (is_router_server) { + // the router itself never loads a model, so it cannot check the group settings the way a + // normal server does. children can only satisfy --pipeline-groups > 1 if they are also + // given a context size, and the router only renders --ctx-size into the child args when the + // operator passed one. without this check every model request would fail at load time with + // an error from a subprocess, long after the mistake was made. + if (g_pipeline_groups > 1 && params.n_ctx <= 0) { + SRV_ERR("%s", "--pipeline-groups > 1 in router mode requires an explicit context size, pass -c N\n"); + return 1; + } + + // server_models snapshots the environment at construction and passes that snapshot to every + // child it spawns. --pipeline-groups has already been stripped from argv by this point, so + // exporting it here is what actually carries the operator's setting through to the children + // that do the loading. must happen before the emplace below, which takes the snapshot. + common_set_env("LLAMA_ARG_PIPELINE_GROUPS", std::to_string(g_pipeline_groups)); + // setup server instances manager try { models_routes.emplace(params, argc, argv);