diff --git a/src/llama-context.cpp b/src/llama-context.cpp index d8060b8f1ed0..facd0556cdfc 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -268,6 +268,8 @@ llama_context::llama_context( cparams.op_offload = params.op_offload; cparams.kv_unified = params.kv_unified; + cparams.type_k = params.type_k; + cparams.type_v = params.type_v; // initialized later cparams.pipeline_parallel = false; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 5018170ed85e..0b5e65281702 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -55,6 +55,11 @@ struct llama_cparams { std::vector embeddings_layer_inp; // [n_layer()] extract input embeddings for layer + // [1bit] the KV cache types this context asked for, so a shared-KV provider can be told what to build + // (llama_kv_share_params, llama-kv-share.cpp); kept here because llama_memory_params only lives in init. + ggml_type type_k = GGML_TYPE_F16; + ggml_type type_v = GGML_TYPE_F16; + enum llama_context_type ctx_type; enum llama_pooling_type pooling_type; diff --git a/src/llama-ext.h b/src/llama-ext.h index 43360b80d0ab..0e604fb329c8 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -48,6 +48,20 @@ LLAMA_API bool llama_kv_share_import(struct llama_context * ctx, const struct ll // Copies no data: this is the file route's ?action=restore, without the file. LLAMA_API bool llama_kv_cells_adopt(struct llama_context * ctx, uint32_t n_tokens, llama_seq_id seq_id); +// [1bit] Target 1: the context params a provider has to reproduce for its cache to lay out like the owner's +// (the layout hash only says yes or no; this says what to build). +struct llama_kv_share_cparams { + uint32_t n_ctx; + uint32_t n_batch; + uint32_t n_ubatch; + uint32_t n_seq_max; + int32_t flash_attn; // llama_flash_attn_type + int32_t type_k; // ggml_type + int32_t type_v; +}; + +LLAMA_API bool llama_kv_share_params(const struct llama_context * ctx, struct llama_kv_share_cparams * out); + // Get the default ggml_type for a given ftype. LLAMA_API ggml_type llama_ftype_get_default_type(llama_ftype ftype); diff --git a/src/llama-kv-share.cpp b/src/llama-kv-share.cpp index 0219019c5528..c94bdeccafe8 100644 --- a/src/llama-kv-share.cpp +++ b/src/llama-kv-share.cpp @@ -322,3 +322,19 @@ bool llama_kv_cells_adopt(llama_context * ctx, uint32_t n_tokens, llama_seq_id s return d.cells_adopt(n_tokens, seq_id); }); } + +// [1bit] what a provider in another process has to reproduce for its own cache to lay out like this one's. +bool llama_kv_share_params(const llama_context * ctx, struct llama_kv_share_cparams * out) { + if (ctx == nullptr || out == nullptr) { + return false; + } + const llama_cparams & cp = ctx->get_cparams(); + out->n_ctx = cp.n_ctx; + out->n_batch = cp.n_batch; + out->n_ubatch = cp.n_ubatch; + out->n_seq_max = cp.n_seq_max; + out->flash_attn = cp.flash_attn ? 1 : 0; + out->type_k = (int32_t) cp.type_k; + out->type_v = (int32_t) cp.type_v; + return true; +} diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 780df3266132..d6f4e31b43b5 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -27,6 +27,7 @@ else() add_subdirectory(server) endif() add_subdirectory(tokenize) + add_subdirectory(kvshare-prefill) # [1bit] target 1: reference provider for a shared KV region add_subdirectory(parser) add_subdirectory(tts) add_subdirectory(mtmd) diff --git a/tools/kvshare-prefill/CMakeLists.txt b/tools/kvshare-prefill/CMakeLists.txt new file mode 100644 index 000000000000..e45bd457b676 --- /dev/null +++ b/tools/kvshare-prefill/CMakeLists.txt @@ -0,0 +1,11 @@ +# [1bit] TZEROCOPY target 1: the reference prefill provider for a shared KV region (tools/kvshare-prefill). +set(TARGET llama-kvshare-prefill) +add_executable(${TARGET} kvshare-prefill.cpp) +# nlohmann/json is the vendored header the rest of the tree uses (vendor/nlohmann) +target_include_directories(${TARGET} PRIVATE ${PROJECT_SOURCE_DIR}/vendor) +target_link_libraries(${TARGET} PRIVATE llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) + +if(LLAMA_TOOLS_INSTALL) + install(TARGETS ${TARGET} RUNTIME) +endif() diff --git a/tools/kvshare-prefill/kvshare-prefill.cpp b/tools/kvshare-prefill/kvshare-prefill.cpp new file mode 100644 index 000000000000..42ded08583cc --- /dev/null +++ b/tools/kvshare-prefill/kvshare-prefill.cpp @@ -0,0 +1,230 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// [1bit] TZEROCOPY target 1: the reference prefill provider for a shared KV region. +// +// The decoder (llama-server with ONEBIT_KV_SHARE=1) exports its KV cache as dma-buf regions and hands them to +// this process as inherited fds + a JSON header (`POST /slots/?action=share {"socket": ...}`). This program +// maps those regions into its own context's KV tensors -- no copy, no state file -- prefills the prompt into +// them, and exits. The decoder then takes the cells over with `POST /slots/?action=adopt`. +// +// kvshare-prefill --model --tokens --layout
--fds 3,4 [--gpu-layers N] +// kvshare-prefill --model --tokens --out (the file route) +// +// The two modes exist so one provider can serve both transports: `--fds` prefills into the decoder's own +// KV memory, `--out` prefills into its own and writes the sequence state the decoder restores. Same math, +// same model, same engine -- which is what makes fd-vs-file a controlled comparison. +// +// stdout: one JSON line {"n_tokens": N, "prefill_ms": x, "argmax": id}. exit 0 on success. +// +// The real provider on this box is the NPU runtime (npu/forward), which imports the same descriptors as XDNA BOs +// instead of mmap'ing them. This one exists so the fd plumbing can be gated end to end with the CPU, and so the +// contract between the engine and any provider has a written-down implementation. + +#include "llama.h" +#include "../../src/llama-ext.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +using json = nlohmann::json; + +static void usage(const char * argv0) { + std::fprintf(stderr, + "usage: %s --model --tokens --layout
--fds \n" + " [--gpu-layers N] [--device NAME]\n", argv0); +} + +int main(int argc, char ** argv) { + std::string model_path; + std::string tokens_path; + std::string layout_path; + std::string fds_arg; + std::string out_path; + std::string device; + int32_t n_gpu_layers = 0; + + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + auto next = [&](const char * what) -> std::string { + if (i + 1 >= argc) { std::fprintf(stderr, "%s: %s needs a value\n", argv[0], what); std::exit(2); } + return argv[++i]; + }; + if (arg == "--model") { model_path = next("--model"); } + else if (arg == "--tokens") { tokens_path = next("--tokens"); } + else if (arg == "--layout") { layout_path = next("--layout"); } + else if (arg == "--fds") { fds_arg = next("--fds"); } + else if (arg == "--out") { out_path = next("--out"); } + else if (arg == "--device") { device = next("--device"); } + else if (arg == "--gpu-layers") { n_gpu_layers = std::atoi(next("--gpu-layers").c_str()); } + else if (arg == "-h" || arg == "--help") { usage(argv[0]); return 0; } + else { std::fprintf(stderr, "%s: unknown argument %s\n", argv[0], arg.c_str()); usage(argv[0]); return 2; } + } + if (model_path.empty() || tokens_path.empty()) { + usage(argv[0]); + return 2; + } + const bool shared = !fds_arg.empty() && !layout_path.empty(); + if (!shared && out_path.empty()) { + usage(argv[0]); + return 2; + } + + // the prompt the decoder tokenized, so both sides prefill exactly the same ids + std::vector tokens; + { + std::ifstream f(tokens_path); + if (!f) { std::fprintf(stderr, "kvshare-prefill: cannot read %s\n", tokens_path.c_str()); return 1; } + long v; + while (f >> v) tokens.push_back((llama_token) v); + } + if (tokens.empty()) { std::fprintf(stderr, "kvshare-prefill: %s has no tokens\n", tokens_path.c_str()); return 1; } + + // the header the decoder sent: region descriptors plus the context params its cache was built with + json header = json{{"params", json::object()}, {"regions", json::array()}}; + if (shared) try { + std::ifstream f(layout_path); + if (!f) throw std::runtime_error("cannot read " + layout_path); + f >> header; + } catch (const std::exception & e) { + std::fprintf(stderr, "kvshare-prefill: %s\n", e.what()); + return 1; + } + const json & params = header.at("params"); + const json & jregions = header.at("regions"); + + // the fds: inherited from the engine, which received them over the unix socket from the decoder + std::vector fds; + for (size_t pos = 0; pos < fds_arg.size();) { + const size_t comma = fds_arg.find(',', pos); + const std::string one = fds_arg.substr(pos, comma == std::string::npos ? std::string::npos : comma - pos); + fds.push_back(std::atoi(one.c_str())); + if (comma == std::string::npos) break; + pos = comma + 1; + } + if (shared && fds.size() != jregions.size()) { + std::fprintf(stderr, "kvshare-prefill: %zu fd(s) for %zu region(s)\n", fds.size(), jregions.size()); + return 1; + } + + std::vector regions(jregions.size()); + for (size_t i = 0; i < regions.size(); ++i) { + regions[i].fd = fds[i]; + regions[i].offset = jregions[i].at("offset").get(); + regions[i].size = jregions[i].at("size").get(); + regions[i].layout = jregions[i].at("layout").get(); + } + + // the model, and a context built the way the decoder's was: the region's layout hash is checked against this + const auto t_load0 = std::chrono::steady_clock::now(); + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = n_gpu_layers; + llama_model * model = llama_model_load_from_file(model_path.c_str(), mparams); + if (model == nullptr) { std::fprintf(stderr, "kvshare-prefill: cannot load %s\n", model_path.c_str()); return 1; } + + llama_context_params cparams = llama_context_default_params(); + if (shared) { + cparams.n_ctx = params.at("n_ctx").get(); + cparams.n_batch = params.at("n_batch").get(); + cparams.n_ubatch = params.at("n_ubatch").get(); + cparams.n_seq_max = params.at("n_seq_max").get(); + cparams.flash_attn_type = params.at("flash_attn").get() != 0 + ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + cparams.type_k = (ggml_type) params.at("type_k").get(); + cparams.type_v = (ggml_type) params.at("type_v").get(); + } else { + // file route: the state the decoder restores must match its cache, and the decoder runs -fa on + cparams.n_ctx = (uint32_t) std::max(tokens.size() + 8, 256); + cparams.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + } + cparams.n_threads = 8; + cparams.n_threads_batch = 8; + + llama_context * ctx = llama_init_from_model(model, cparams); + if (ctx == nullptr) { std::fprintf(stderr, "kvshare-prefill: cannot create a context\n"); return 1; } + + // shared route: map the decoder's KV memory as this context's KV tensors (zero copy: the tensors ARE it) + const auto t_import0 = std::chrono::steady_clock::now(); + if (shared && !llama_kv_share_import(ctx, regions.data(), regions.size())) { + std::fprintf(stderr, "kvshare-prefill: cannot map the shared KV region (layout differs: same model, -fa, " + "-c, -b/-ub and KV types as the decoder are required)\n"); + return 1; + } + const double import_ms = std::chrono::duration(std::chrono::steady_clock::now() - t_import0).count(); + const double load_ms = std::chrono::duration(std::chrono::steady_clock::now() - t_load0).count(); + + // prefill into it, in chunks of n_ubatch, exactly as a normal prompt prefill would + const int32_t n_batch_max = (int32_t) std::min(cparams.n_batch, tokens.size()); + llama_batch batch = llama_batch_init(n_batch_max, 0, 1); + const auto t_pre0 = std::chrono::steady_clock::now(); + for (size_t done = 0; done < tokens.size();) { + const size_t k = std::min(cparams.n_ubatch, tokens.size() - done); + batch.n_tokens = (int32_t) k; + for (size_t i = 0; i < k; ++i) { + batch.token[i] = tokens[done + i]; + batch.pos[i] = (llama_pos) (done + i); + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + batch.logits[i] = (done + i == tokens.size() - 1) ? 1 : 0; + } + if (llama_decode(ctx, batch) != 0) { + std::fprintf(stderr, "kvshare-prefill: decode failed at token %zu\n", done); + return 1; + } + done += k; + } + llama_synchronize(ctx); + const double prefill_ms = std::chrono::duration(std::chrono::steady_clock::now() - t_pre0).count(); + + // the last token's argmax, so the engine can see that the provider and the decoder agree on the prefix + int32_t argmax = -1; + if (const float * logits = llama_get_logits_ith(ctx, batch.n_tokens - 1)) { + const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + float best = logits[0]; + argmax = 0; + for (int i = 1; i < n_vocab; ++i) { + if (logits[i] > best) { best = logits[i]; argmax = i; } + } + } + + // file route: the same cells, written where the decoder's ?action=restore will read them + if (!shared) { + const size_t written = llama_state_seq_save_file(ctx, out_path.c_str(), 0, tokens.data(), tokens.size()); + if (written == 0) { + std::fprintf(stderr, "kvshare-prefill: cannot write %s\n", out_path.c_str()); + return 1; + } + } + + std::printf("{\"n_tokens\": %zu, \"prefill_ms\": %.1f, \"import_ms\": %.3f, \"load_ms\": %.1f, \"argmax\": %d}\n", + tokens.size(), prefill_ms, import_ms, load_ms, argmax); + std::fflush(stdout); + + llama_batch_free(batch); + llama_free(ctx); + llama_model_free(model); + return 0; +} diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index fd168020167f..8e6e1f2d9696 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -30,6 +30,15 @@ #include #include +#include +#include + +#ifndef _WIN32 +#include +#include +#include +#endif + // fix problem with std::min and std::max #if defined(_WIN32) #define WIN32_LEAN_AND_MEAN @@ -976,6 +985,12 @@ struct server_context_impl { // slots / clients std::vector slots; + // [1bit] target 1: the last KV-region share issued per slot. An adopt has to quote it back, so cells that + // nobody exported (stale, forgotten, or a client inventing ids) are refused instead of silently becoming a + // prompt's cache. Single use: adopting erases it. + std::map kv_share_issued; + uint64_t kv_share_seq = 0; + int trace = 0; int slots_debug = 0; int n_empty_consecutive = 0; @@ -2679,6 +2694,108 @@ struct server_context_impl { res->t_ms = t_restore_ms; queue_results.send(std::move(res)); } break; + case SERVER_TASK_TYPE_SLOT_SHARE: + { + // [1bit] target 1: hand the KV region itself to a prefill provider outside this process, as + // dma-buf fds. Nothing is copied here: the region already exists and the provider maps it. + const int id_slot = task.slot_action.id_slot; + server_slot * slot = get_slot_by_id(id_slot); + if (slot == nullptr) { + send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST); + break; + } + if (slot->is_processing()) { + SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id); + queue_tasks.defer(std::move(task)); + break; + } + + const int64_t t_start = ggml_time_us(); + + // the count first (n_out = 0 fills nothing), then the regions + const size_t n_regions = llama_kv_share_regions(ctx_tgt, nullptr, 0); + if (n_regions == 0) { + send_error(task, "This server's KV cache has no shareable region: either it was not created with " + "ONEBIT_KV_SHARE=1, or the model is hybrid and its recurrent state is not " + "shareable (target 1 shares llama_kv_cache only)", + ERROR_TYPE_NOT_SUPPORTED); + break; + } + std::vector regions(n_regions); + llama_kv_share_cparams share_params{}; + if (llama_kv_share_regions(ctx_tgt, regions.data(), regions.size()) != n_regions || + !llama_kv_share_params(ctx_tgt, &share_params)) { + send_error(task, "Cannot describe this server's KV region", ERROR_TYPE_SERVER); + break; + } + + auto res = std::make_unique(); + for (const auto & r : regions) { + res->regions.push_back({ r.fd, r.offset, r.size, r.layout }); + } + res->id = task.id; + res->id_slot = id_slot; + res->share_id = ++kv_share_seq; + kv_share_issued[id_slot] = res->share_id; + res->n_ctx = share_params.n_ctx; + res->n_batch = share_params.n_batch; + res->n_ubatch = share_params.n_ubatch; + res->n_seq_max = share_params.n_seq_max; + res->flash_attn = share_params.flash_attn; + res->type_k = share_params.type_k; + res->type_v = share_params.type_v; + res->t_ms = (ggml_time_us() - t_start) / 1000.0; + queue_results.send(std::move(res)); + } break; + case SERVER_TASK_TYPE_SLOT_ADOPT: + { + // [1bit] target 1: the provider has filled the shared region; this slot's cache already holds + // those bytes, so only the cell metadata is written (llama_kv_cells_adopt). + const int id_slot = task.slot_action.id_slot; + server_slot * slot = get_slot_by_id(id_slot); + if (slot == nullptr) { + send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST); + break; + } + if (slot->is_processing()) { + SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id); + queue_tasks.defer(std::move(task)); + break; + } + + const llama_tokens & tokens = task.slot_action.tokens; + if (tokens.empty()) { + send_error(task, "adopt needs the prompt tokens the provider prefilled", ERROR_TYPE_INVALID_REQUEST); + break; + } + // the cells are only meaningful for the export they were written into + const auto issued = kv_share_issued.find(id_slot); + if (task.slot_action.share_id == 0 || issued == kv_share_issued.end() || + issued->second != task.slot_action.share_id) { + send_error(task, "adopt must quote the share_id of the export these cells came from " + "(POST /slots/?action=share first)", ERROR_TYPE_INVALID_REQUEST); + break; + } + + const int64_t t_start = ggml_time_us(); + + if (!llama_kv_cells_adopt(ctx_tgt, (uint32_t) tokens.size(), slot->id)) { + slot->prompt.clear(); // the cache did not take these cells (another sequence holds them) + send_error(task, "Cannot adopt the shared cells: this slot's cache does not hold the provider's prompt", ERROR_TYPE_SERVER); + break; + } + slot->prompt.clear(); + slot->prompt.tokens.insert(tokens); + kv_share_issued.erase(issued); // single use + + auto res = std::make_unique(); + res->id = task.id; + res->id_slot = id_slot; + res->share_id = task.slot_action.share_id; + res->n_tokens = tokens.size(); + res->t_ms = (ggml_time_us() - t_start) / 1000.0; + queue_results.send(std::move(res)); + } break; case SERVER_TASK_TYPE_SLOT_ERASE: { const int id_slot = task.slot_action.id_slot; @@ -4618,6 +4735,12 @@ void server_routes::init_routes() { if (action == "erase") { return handle_slots_erase(req, id_slot); } + if (action == "share") { + return handle_slots_share(req, id_slot); + } + if (action == "adopt") { + return handle_slots_adopt(req, id_slot); + } res->error(format_error_response("Invalid action", ERROR_TYPE_INVALID_REQUEST)); return res; @@ -5299,6 +5422,184 @@ std::unique_ptr server_routes::handle_slots_erase(const se return res; } +#ifndef _WIN32 + +// [1bit] target 1: one message to the provider: [uint32 length][JSON header], the region fds in the same sendmsg +// as SCM_RIGHTS. Linux has no way to reopen a dma-buf from /proc//fd (measured), so this is the only path. +static bool kv_share_send(const std::string & socket_path, const server_task_result_slot_share & r, std::string * why) { + nlohmann::json header = { + { "n_regions", r.regions.size() }, + { "params", { + { "n_ctx", r.n_ctx }, + { "n_batch", r.n_batch }, + { "n_ubatch", r.n_ubatch }, + { "n_seq_max", r.n_seq_max }, + { "flash_attn", r.flash_attn }, + { "type_k", r.type_k }, + { "type_v", r.type_v }, + }}, + }; + nlohmann::json regions = nlohmann::json::array(); + for (const auto & reg : r.regions) { + regions.push_back({ { "offset", reg.offset }, { "size", reg.size }, { "layout", reg.layout } }); + } + header["regions"] = regions; + const std::string body = header.dump(); + + const int fd = ::socket(AF_UNIX, SOCK_STREAM, 0); + if (fd < 0) { *why = std::string("socket: ") + std::strerror(errno); return false; } + sockaddr_un addr{}; + addr.sun_family = AF_UNIX; + if (socket_path.size() >= sizeof(addr.sun_path)) { + ::close(fd); + *why = "socket path too long"; + return false; + } + std::memcpy(addr.sun_path, socket_path.c_str(), socket_path.size() + 1); + if (::connect(fd, reinterpret_cast(&addr), sizeof(addr)) != 0) { + *why = "connect " + socket_path + ": " + std::strerror(errno); + ::close(fd); + return false; + } + + const uint32_t len = (uint32_t) body.size(); + std::vector buf(4 + body.size()); + std::memcpy(buf.data(), &len, 4); + std::memcpy(buf.data() + 4, body.data(), body.size()); + + std::vector control(CMSG_SPACE(r.regions.size() * sizeof(int))); + msghdr msg{}; + iovec iov{ buf.data(), buf.size() }; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + msg.msg_control = control.data(); + msg.msg_controllen = control.size(); + cmsghdr * cmsg = CMSG_FIRSTHDR(&msg); + cmsg->cmsg_level = SOL_SOCKET; + cmsg->cmsg_type = SCM_RIGHTS; + cmsg->cmsg_len = CMSG_LEN(r.regions.size() * sizeof(int)); + for (size_t i = 0; i < r.regions.size(); ++i) { + std::memcpy(CMSG_DATA(cmsg) + i * sizeof(int), &r.regions[i].fd, sizeof(int)); + } + + bool ok = true; + size_t sent = 0; + while (sent < iov.iov_len) { + const ssize_t n = ::sendmsg(fd, &msg, MSG_NOSIGNAL); + if (n < 0 && errno == EINTR) { + continue; + } + if (n <= 0) { + *why = std::string("sendmsg: ") + std::strerror(errno); + ok = false; + break; + } + sent += (size_t) n; + // a partial send cannot carry the control message a second time + if (sent < iov.iov_len) { + *why = "partial send to the provider"; + ok = false; + break; + } + } + ::close(fd); + return ok; +} + +#else + +static bool kv_share_send(const std::string &, const server_task_result_slot_share &, std::string * why) { + *why = "shared KV regions need unix sockets (POSIX only)"; + return false; +} + +#endif + +std::unique_ptr server_routes::handle_slots_share(const server_http_req & req, int id_slot) { + auto res = create_response(); + std::string socket_path; + try { + socket_path = json::parse(req.body).at("socket").get(); + } catch (const std::exception &) { + res->error(format_error_response("share needs {\"socket\": \"\"}", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + + auto & rd = res->rd; + { + server_task task(SERVER_TASK_TYPE_SLOT_SHARE); + task.id = rd.get_new_id(); + task.slot_action.id_slot = id_slot; + task.slot_action.socket = socket_path; + rd.post_task(std::move(task)); + } + + auto result = rd.next(req.should_stop); + if (!result) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + auto * shared = dynamic_cast(result.get()); + GGML_ASSERT(shared != nullptr); + + std::string why; + if (!kv_share_send(socket_path, *shared, &why)) { + res->error(format_error_response("cannot hand the KV region to the provider: " + why, ERROR_TYPE_SERVER)); + return res; + } + SRV_INF("shared KV region: %zu fd(s) handed to the provider on %s (%.3f ms)\n", + shared->regions.size(), socket_path.c_str(), shared->t_ms); + res->ok(shared->to_json()); + return res; +} + +std::unique_ptr server_routes::handle_slots_adopt(const server_http_req & req, int id_slot) { + auto res = create_response(); + json body; + try { + body = json::parse(req.body); + } catch (const std::exception &) { + res->error(format_error_response("adopt needs {\"tokens\": [...]} (the ids the provider prefilled)", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + if (!body.contains("tokens") || !body.at("tokens").is_array()) { + res->error(format_error_response("adopt needs {\"tokens\": [...], \"share_id\": N} (the ids the provider prefilled, " + "and the export they were written into)", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + if (!body.contains("share_id") || !body.at("share_id").is_number_unsigned()) { + res->error(format_error_response("adopt needs the \"share_id\" that ?action=share returned", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + + auto & rd = res->rd; + { + server_task task(SERVER_TASK_TYPE_SLOT_ADOPT); + task.id = rd.get_new_id(); + task.slot_action.id_slot = id_slot; + task.slot_action.tokens = body.at("tokens").get(); + task.slot_action.share_id = body.at("share_id").get(); + rd.post_task(std::move(task)); + } + + auto result = rd.next(req.should_stop); + if (!result) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + GGML_ASSERT(dynamic_cast(result.get()) != nullptr); + res->ok(result->to_json()); + return res; +} + std::unique_ptr server_routes::handle_embeddings_impl(const server_http_req & req, task_response_type res_type) { auto res = create_response(); if (!params.embedding) { diff --git a/tools/server/server-context.h b/tools/server/server-context.h index f9ab1132b197..5ce9f3d55960 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -167,6 +167,9 @@ struct server_routes { std::unique_ptr handle_slots_save(const server_http_req & req, int id_slot); std::unique_ptr handle_slots_restore(const server_http_req & req, int id_slot); std::unique_ptr handle_slots_erase(const server_http_req &, int id_slot); + // [1bit] target 1: the shared-region handoff to a prefill provider in another process + std::unique_ptr handle_slots_share (const server_http_req & req, int id_slot); + std::unique_ptr handle_slots_adopt (const server_http_req & req, int id_slot); std::unique_ptr handle_embeddings_impl(const server_http_req & req, task_response_type res_type); std::unique_ptr handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const server_http_req & req, task_response_type res_type); diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 070f1ade241c..b30b70026707 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1602,9 +1602,56 @@ json server_task_result_slot_erase::to_json() { } // -// server_task_result_get_lora +// server_task_result_slot_share [1bit] target 1 +// +json server_task_result_slot_share::to_json() { + json out_regions = json::array(); + for (const auto & r : regions) { + // no fd here on purpose: the reply goes over HTTP, the fds went over the socket + out_regions.push_back(json { + { "offset", r.offset }, + { "size", r.size }, + { "layout", r.layout }, + }); + } + return json { + { "id_slot", id_slot }, + { "share_id", share_id }, + { "n_regions", regions.size() }, + { "regions", out_regions }, + { "params", { + { "n_ctx", n_ctx }, + { "n_batch", n_batch }, + { "n_ubatch", n_ubatch }, + { "n_seq_max", n_seq_max }, + { "flash_attn", flash_attn }, + { "type_k", type_k }, + { "type_v", type_v }, + }}, + { "timings", { + { "share_ms", t_ms } + }}, + }; +} + +// +// server_task_result_slot_adopt [1bit] target 1 // +json server_task_result_slot_adopt::to_json() { + return json { + { "id_slot", id_slot }, + { "share_id", share_id }, + { "n_adopted", n_tokens }, + { "n_read", 0 }, + { "timings", { + { "adopt_ms", t_ms } + }}, + }; +} +// +// server_task_result_get_lora +// json server_task_result_get_lora::to_json() { json result = json::array(); for (size_t i = 0; i < loras.size(); ++i) { diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 411d91807955..a0ed05a7980d 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -25,6 +25,8 @@ enum server_task_type { SERVER_TASK_TYPE_SLOT_SAVE, SERVER_TASK_TYPE_SLOT_RESTORE, SERVER_TASK_TYPE_SLOT_ERASE, + SERVER_TASK_TYPE_SLOT_SHARE, // [1bit] target 1: export the KV region to a prefill provider (no copy) + SERVER_TASK_TYPE_SLOT_ADOPT, // [1bit] target 1: take the provider's cells into this slot (no copy) SERVER_TASK_TYPE_GET_LORA, SERVER_TASK_TYPE_SET_LORA, }; @@ -161,11 +163,15 @@ struct server_task { server_task_type type; - // used by SERVER_TASK_TYPE_SLOT_SAVE, SERVER_TASK_TYPE_SLOT_RESTORE, SERVER_TASK_TYPE_SLOT_ERASE + // used by SERVER_TASK_TYPE_SLOT_SAVE, SERVER_TASK_TYPE_SLOT_RESTORE, SERVER_TASK_TYPE_SLOT_ERASE, + // SERVER_TASK_TYPE_SLOT_SHARE, SERVER_TASK_TYPE_SLOT_ADOPT struct slot_action { int id_slot; std::string filename; std::string filepath; + std::string socket; // SHARE: where the provider listens for the regions + std::vector tokens; // ADOPT: the prompt tokens the provider prefilled into the region + uint64_t share_id = 0; // SHARE out / ADOPT in: which export these cells belong to }; slot_action slot_action; @@ -556,6 +562,44 @@ struct server_task_result_slot_erase : server_task_result { virtual json to_json() override; }; +// [1bit] target 1: the KV region a slot's decoder exported, plus the context params a provider has to reproduce +// for its own cache to lay out the same way. The fds are deliberately not in to_json(): they only ever travel to +// the provider process over the unix socket, as SCM_RIGHTS. +struct server_task_result_slot_share : server_task_result { + struct region { + int fd = -1; + size_t offset = 0; + size_t size = 0; + uint64_t layout = 0; + }; + + int id_slot = 0; + uint64_t share_id = 0; // adopt must quote this back: cells nobody exported are not adopted + std::vector regions; + + uint32_t n_ctx = 0; + uint32_t n_batch = 0; + uint32_t n_ubatch = 0; + uint32_t n_seq_max = 0; + int32_t flash_attn = 0; + int32_t type_k = 0; + int32_t type_v = 0; + + double t_ms = 0; + + virtual json to_json() override; +}; + +// [1bit] target 1: the provider's cells are this slot's cache now; this is a restore with no file and no copy +struct server_task_result_slot_adopt : server_task_result { + int id_slot = 0; + uint64_t share_id = 0; + size_t n_tokens = 0; + double t_ms = 0; + + virtual json to_json() override; +}; + struct server_task_result_control : server_task_result { bool success = false; std::string message; // optional detail when success is false