From c26d21dc4cbe13e0a39ba52034edc2c17c3c0d76 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Fri, 23 Oct 2026 08:19:07 -0300 Subject: [PATCH 1/4] serve: --device hybrid (external prefill provider + HRX decode) The prompt's KV cache comes from a prefill provider named by --prefill-provider or ONEBIT_HYBRID_PREFILL; llama-server restores it into its one slot and decodes on HRX (the CPU in a build without HRX). app/serve_hybrid.h holds the handoff: template and tokenize on the backend, the provider's state file, the slot restore, the request pinned to the slot with cache_prompt on, and the fallback to a plain request when any step fails. docs/hybrid.md states the contract. tests/hybrid_route.sh runs the path with the fake backend and a fake provider; tests/hybrid_e2e.sh runs it on a real model with the decoding server making the state file, and checks the greedy reply equals the plain run's and that the backend prefilled one token. Co-Authored-By: Claude Fable 5.1 --- CMakeLists.txt | 9 ++ README.md | 3 + app/serve.cpp | 123 +++++++++++++-- app/serve_hybrid.h | 271 +++++++++++++++++++++++++++++++++ docs/hybrid.md | 107 +++++++++++++ docs/serve.md | 6 +- tests/fake_backend.py | 25 ++- tests/fake_prefill_provider.py | 67 ++++++++ tests/hybrid_e2e.sh | 95 ++++++++++++ tests/hybrid_route.sh | 97 ++++++++++++ 10 files changed, 791 insertions(+), 12 deletions(-) create mode 100644 app/serve_hybrid.h create mode 100644 docs/hybrid.md create mode 100755 tests/fake_prefill_provider.py create mode 100755 tests/hybrid_e2e.sh create mode 100755 tests/hybrid_route.sh diff --git a/CMakeLists.txt b/CMakeLists.txt index bb8f3300..08b049e9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -137,6 +137,15 @@ add_test(NAME laya_gguf_route COMMAND ${CMAKE_SOURCE_DIR}/tests/laya_gguf_route. add_test(NAME mem_admit_route COMMAND ${CMAKE_SOURCE_DIR}/tests/mem_admit_route.sh $) # engine#315: what --hrx-cwsr-placement puts in an HRX backend's environment (app/hrx_cwsr.h); fake backend, no GPU. add_test(NAME hrx_cwsr_route COMMAND ${CMAKE_SOURCE_DIR}/tests/hrx_cwsr_route.sh $) +# --device hybrid (docs/hybrid.md): the handoff with a fake backend and a fake provider; and, with a model, +# end to end on the CPU llama-server (ONEBIT_SERVE_TEST_LLAMA_SERVER, else the build's own; skips without one) +add_test(NAME hybrid_route COMMAND ${CMAKE_SOURCE_DIR}/tests/hybrid_route.sh $) +set(ONEBIT_SERVE_TEST_LLAMA_SERVER "" CACHE FILEPATH "llama-server for hybrid_e2e (default: the build's own)") +if(ONEBIT_SERVE_TEST_GGUF) + add_test(NAME hybrid_e2e + COMMAND ${CMAKE_SOURCE_DIR}/tests/hybrid_e2e.sh $ ${ONEBIT_SERVE_TEST_GGUF} ${ONEBIT_SERVE_TEST_LLAMA_SERVER}) + set_tests_properties(hybrid_e2e PROPERTIES SKIP_RETURN_CODE 77) +endif() # The NPU route's split into reasoning_content (app/think_split.h); header only, runs in CI. add_executable(think_split_test tests/think_split_test.cpp) target_include_directories(think_split_test PRIVATE app) diff --git a/README.md b/README.md index fdf5f030..8d23d492 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,9 @@ hands to 1bit, it runs the engine as one of its backends, behind an OpenAI-compa same build on the CPU. - **The NPU engine** on the XDNA 2 NPU, in progress: correctness first, speed next ([docs/npu.md](docs/npu.md)). +- **Hybrid**: the prompt on the NPU engine, the tokens on HRX. `1bit serve --device hybrid` hands + the prompt's KV cache from a prefill provider to the HRX llama-server + ([docs/hybrid.md](docs/hybrid.md)). - **DwarfStar**, for DeepSeek V4 Flash, GLM 5.x and Qwen3.8-Flash-Next in its own GGUFs, with SSD expert streaming ([docs/dwarfstar.md](docs/dwarfstar.md)). - **Laya**, opt-in, which picks where each request runs ([docs/laya.md](docs/laya.md)), and diff --git a/app/serve.cpp b/app/serve.cpp index 1a97b94b..80942be6 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -13,8 +13,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -// `1bit serve -m [--port 8000] [--device auto|npu|hrx|cpu|ds4] [--mtp HEAD]` -// [--npu-opt KEY=VALUE ...] +// `1bit serve -m [--port 8000] [--device auto|npu|hrx|hybrid|cpu|ds4] [--mtp HEAD]` +// [--npu-opt KEY=VALUE ...] [--prefill-provider PATH] // // The engine's one front door (docs/serve.md): one model per process behind an // OpenAI-compatible API. It is how the engine runs inside Lemonade: Lemonade's @@ -33,6 +33,9 @@ // -DONEBIT_NPU_PRIVATE (docs/npu.md) // a .gguf, --device hrx -> the HRX build's llama-server on HRX0 // a .gguf, --device cpu -> the same llama-server with no GPU layers +// a .gguf, --device hybrid -> the HRX llama-server decodes; the prompt's KV comes +// from an external prefill provider (the NPU engine, +// ONEBIT_HYBRID_PREFILL; app/serve_hybrid.h, docs/hybrid.md) // a .1bp (1BP v5 package), --device auto|ds4 -> the same ds4-server (our DwarfStar fork reads 1BP) // a DwarfStar .gguf, --device ds4 -> this build's DwarfStar ds4-server (DeepSeek // V4/V4.1 Flash, GLM 5.x, Qwen3.8-Flash-Next in @@ -72,6 +75,7 @@ #include "recipes.h" #include "serve_metrics.h" #include "serve_watchdog.h" +#include "serve_hybrid.h" #include "sysinfo.h" #ifdef ONEBIT_MOE @@ -165,6 +169,7 @@ struct Options { std::string embed, rerank; // RAG: an embedding model and a reranker, served beside the chat model std::string role; // "embedding" or "reranking": the model itself is served in that role std::vector npu_opts; // --npu-opt KEY=VALUE, for a private NPU route + std::string prefill_provider; // --prefill-provider PATH (or ONEBIT_HYBRID_PREFILL): --device hybrid's prefill (docs/hybrid.md) std::string laya_model; // --laya-model DIR: route each request by the Laya scorer (docs/laya.md) bool laya_auto = false; // --laya: the same with the installed checkpoint (laya_checkpoint()) std::string route_policy; // --route-policy FILE: request class -> device (default: built in) @@ -702,12 +707,15 @@ bool wants_stream(const std::string& body) { } // watch, where given, tracks the request for the progress watchdog (app/serve_watchdog.h). +// body_override, where given, is the JSON body to send in place of the request's (hybrid: id_slot and +// cache_prompt added); on_first_byte, where given, runs when the first byte of the reply arrives. void forward(int port, std::shared_ptr held, int draft_max, const std::string& our_id, bool drop_model, const std::string& set_model, bool precise, const httplib::Request& req, httplib::Response& res, - ProgressWatch* watch = nullptr) { + ProgressWatch* watch = nullptr, const json* body_override = nullptr, + std::function on_first_byte = nullptr) { json body; try { - body = json::parse(req.body); + body = body_override ? *body_override : json::parse(req.body); } catch (const std::exception&) { res.status = 400; res.set_content(R"({"error":{"message":"request body is not JSON"}})", "application/json"); @@ -741,6 +749,7 @@ void forward(int port, std::shared_ptr held, int draft_max, const std: c.stop(); } auto r = pending.get(); + if (on_first_byte) on_first_byte(); if (!r) { res.status = 502; res.set_content(R"({"error":{"message":"backend did not answer"}})", "application/json"); @@ -758,12 +767,15 @@ void forward(int port, std::shared_ptr held, int draft_max, const std: } // Streaming: relay the child's SSE bytes as they arrive. const size_t prompt_bytes = req.body.size(); - res.set_chunked_content_provider("text/event-stream", [port, path, payload, held, watch, prompt_bytes](size_t, httplib::DataSink& sink) { + res.set_chunked_content_provider("text/event-stream", [port, path, payload, held, watch, prompt_bytes, on_first_byte](size_t, httplib::DataSink& sink) { const auto tracked = watch ? watch->begin(port, "POST " + path + " (streamed)", prompt_bytes, false) : nullptr; httplib::Client c("127.0.0.1", port); c.set_read_timeout(3600); + bool first = true; c.Post(path, httplib::Headers{}, payload, "application/json", - [&sink, &tracked](const char* data, size_t n) { + [&sink, &tracked, &first, &on_first_byte](const char* data, size_t n) { + if (first && on_first_byte) on_first_byte(); + first = false; if (tracked) ProgressWatch::progress(*tracked, n); return sink.write(data, n); }); @@ -804,6 +816,8 @@ struct Launch { bool mtp = false; std::string ready_path = "/health"; std::vector env_override; // set in place of the inherited value (Child) + bool hybrid = false; // --device hybrid: the prompt's KV comes from the prefill provider (app/serve_hybrid.h) + std::string slot_dir; // hybrid: the backend's --slot-save-path, where the handoff files go }; // A GGUF that carries its own MTP head (.nextn_predict_layers > 0: Unsloth's Qwen3.8-27B @@ -892,6 +906,8 @@ bool delta_net_arch(const std::string& arch) { return arch == "qwen35" || arch == "qwen35moe" || arch == "qwen3next"; } +std::string auto_gguf_device(const Options& o); + Launch launch_for(const Options& o, const std::string& device, int child_port) { Launch l; l.device = device; @@ -900,6 +916,30 @@ Launch launch_for(const Options& o, const std::string& device, int child_port) { std::vector& env = l.env; std::string& set_model = l.set_model; if (o.ssd_streaming && device != "ds4") throw std::runtime_error("--ssd-streaming works with --device ds4"); + if (device == "hybrid") { + // Hybrid (docs/hybrid.md, app/serve_hybrid.h): llama-server decodes on HRX (the CPU in a build + // without HRX, so the path runs anywhere), from a KV cache an external prefill provider computed + // and serve restored into its one slot. The handoff is a sequence-state file in --slot-save-path. + // Nothing else is handed over: no drafter, no images, no second slot. + if (!o.mtp.empty() || !o.dflash.empty() || !o.mmproj.empty() || o.parallel > 1 || !o.role.empty() || o.moe_slots != 0) + throw std::runtime_error("--device hybrid takes no --mtp, --dflash, --mmproj, --moe-slots, --parallel > 1 or a role: " + "one slot, decoding from the provider's KV (docs/hybrid.md)"); +#if defined(_WIN32) + throw std::runtime_error("--device hybrid is a Linux route"); +#else + Options one = o; + one.parallel = 1; + one.no_mtp = true; // the slot's KV comes from the provider; the GGUF's own MTP head would draft over it unverified + Launch l = launch_for(one, auto_gguf_device(o), child_port); + char dir[] = "/tmp/1bit-hybrid-XXXXXX"; + if (!::mkdtemp(dir)) throw std::runtime_error(std::string("--device hybrid: cannot make a slot directory: ") + std::strerror(errno)); + l.slot_dir = dir; + l.argv.insert(l.argv.end(), {"--slot-save-path", l.slot_dir}); + l.device = "hybrid"; + l.hybrid = true; + return l; +#endif + } if (device == "hrx" || device == "cpu") { // --device cpu: llama-server with no GPU layers (what --device auto falls back to where // HRX has no route for a request, and in a build without HRX) @@ -1399,7 +1439,8 @@ int serve_only(const std::string& why) { // in order; "warmup" is the first one, for a lemond that reads only that), and Lemonade's onebit recipe // sends them once the backend is ready, before it marks the model loaded. static bool warmup_wanted(const std::string& device) { - return getenv("ONEBIT_NO_WARMUP") == nullptr && device.substr(0, device.find(' ')) == "hrx"; + const std::string dev = device.substr(0, device.find(' ')); + return getenv("ONEBIT_NO_WARMUP") == nullptr && (dev == "hrx" || dev == "hybrid"); } static const char* const kWarmupPath = "/v1/chat/completions"; @@ -1562,6 +1603,7 @@ int serve_child(const Options& given) { // --print-backend (docs/serve.md, "One hop from Lemonade"): only a lone llama-server can be // started by Lemonade itself; these need serve in the request path if (o.device == "ds4") return serve_only("--device ds4 is not a llama-server"); + if (o.device == "hybrid") return serve_only("--device hybrid hands each prompt's KV to llama-server in serve"); if (o.pm) return serve_only("--pm runs its tool loop in serve"); if (o.laya_auto || !o.laya_model.empty()) return serve_only("--laya routes each request in serve"); if (!o.embed.empty() || !o.rerank.empty()) return serve_only("--embed and --rerank run more servers behind serve"); @@ -1716,6 +1758,30 @@ int serve_child(const Options& given) { const bool drop_model = backends[0].drop_model; const std::string set_model = backends[0].set_model; + // --device hybrid (app/serve_hybrid.h, docs/hybrid.md): the prefill provider, when named; without one + // the backend prefills every prompt itself (said once), so the device still serves. + std::unique_ptr hybrid_prefill; + struct SlotDirGuard { + std::string dir; + ~SlotDirGuard() { if (!dir.empty()) { std::error_code ec; fs::remove_all(dir, ec); } } + } slot_dir_guard{backends[0].hybrid ? backends[0].slot_dir : ""}; + if (backends[0].hybrid) { + if (o.pm) throw std::runtime_error("--device hybrid does not combine with --pm (the tool loop forwards its own requests)"); + std::string provider = o.prefill_provider; + if (provider.empty()) + if (const char* e = std::getenv("ONEBIT_HYBRID_PREFILL"); e && *e) provider = e; + std::string elfs; + if (const char* e = std::getenv("ONEBIT_HYBRID_ELFS"); e && *e) elfs = e; + if (provider.empty()) { + std::fprintf(stderr, "1bit serve: hybrid: no prefill provider (--prefill-provider or ONEBIT_HYBRID_PREFILL): " + "the %s backend prefills every prompt itself (docs/hybrid.md)\n", auto_gguf_device(o).c_str()); + } else { + hybrid_prefill = std::make_unique(provider, elfs, backends[0].slot_dir, o.model); + std::fprintf(stderr, "1bit serve: hybrid: prefill by %s%s%s, decode on %s\n", provider.c_str(), + elfs.empty() ? "" : " --elfs ", elfs.c_str(), auto_gguf_device(o).c_str()); + } + } + std::vector rag; // [0] embedding, then rerank, if given if (!o.embed.empty()) rag.push_back(rag_launch(o, o.embed, "embedding", free_port())); if (!o.rerank.empty()) rag.push_back(rag_launch(o, o.rerank, "reranking", free_port())); @@ -1949,6 +2015,41 @@ int serve_child(const Options& given) { held = std::make_shared(&inflight[pick]); } ++routed[pick]; + if (backends[pick].hybrid && hybrid_prefill) { + // the handoff (app/serve_hybrid.h): template + tokenize on the backend, the provider's prefill, + // restore into slot 0, then the request pinned to that slot. One slot: one request at a time, + // held through its decode so the next handoff cannot replace a KV still in use. + json body; + try { + body = json::parse(q.body); + } catch (const std::exception&) { + r.status = 400; + r.set_content(R"({"error":{"message":"request body is not JSON"}})", "application/json"); + return; + } + // The lock is held until the backend has started on the slot (its first byte): from then + // on llama-server itself defers a restore into a busy slot. httplib runs the handler and + // the streamed reply on one thread, so the unlock is on the locking thread. + auto slot = std::make_shared>(hybrid_prefill->slot_mutex()); + const auto t0 = std::chrono::steady_clock::now(); + const hybrid::Outcome h = hybrid_prefill->prepare(backends[pick].port, q.path, body); + if (!h.handed_over) { + std::fprintf(stderr, "1bit serve: hybrid: %s; the backend prefills this request itself\n", h.why.c_str()); + slot->unlock(); + forward(backends[pick].port, held, -1, id, drop_model, set_model, o.precise, q, r, &watch); + return; + } + const size_t n = h.n_tokens; + const double prefill_ms = h.prefill_ms, restore_ms = h.restore_ms; + forward(backends[pick].port, held, -1, id, drop_model, set_model, o.precise, q, r, &watch, &body, + [t0, n, prefill_ms, restore_ms, slot] { + const double ttft = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + std::fprintf(stderr, "1bit serve: hybrid: %zu prompt tokens: prefill %.0f ms, restore %.0f ms, first byte %.0f ms\n", + n, prefill_ms, restore_ms, ttft); + if (slot->owns_lock()) slot->unlock(); + }); + return; + } forward(backends[pick].port, held, -1, id, drop_model, set_model, o.precise, q, r, &watch); }; srv.Post("/v1/chat/completions", post); @@ -1956,7 +2057,7 @@ int serve_child(const Options& given) { // The rest of llama-server's API, for the devices that run llama-server (Lemonade's // llamacpp backend, which the onebit recipe inherits, forwards these to us). They go to // the first backend; the NPU and DwarfStar answer 501. - const bool llama = device == "hrx" || device == "cpu" || + const bool llama = device == "hrx" || device == "cpu" || device == "hybrid" || (laya && (backends[0].device == "hrx" || backends[0].device == "cpu")); auto unsupported = [&](const httplib::Request& q, httplib::Response& r) { r.status = 501; @@ -2471,8 +2572,11 @@ std::string repack_gguf_lane(const std::string& gguf, const std::string& lane_di void usage(FILE* out) { std::fprintf(out, "usage: 1bit serve -m [--port 8000] [--host 127.0.0.1]\n" - " [--device auto|npu|hrx|cpu|ds4] [--ctx-size N] [--alias NAME]\n" + " [--device auto|npu|hrx|hybrid|cpu|ds4] [--ctx-size N] [--alias NAME]\n" " [--llama-server PATH] [--hrx-libhsa PATH]\n" + " [--prefill-provider PATH] --device hybrid: the program that prefills each prompt\n" + " (the NPU engine) and writes its KV for llama-server, which decodes it on HRX\n" + " (or ONEBIT_HYBRID_PREFILL; ONEBIT_HYBRID_ELFS passes --elfs; docs/hybrid.md)\n" " [--hrx-cwsr-placement=auto|on|off] engine#315: on a CWSR-on amdgpu (auto, the default)\n" " or always (on), HRX backends get HSA_USE_SVM=0 HSA_USERPTR_FOR_PAGED_MEM=0\n" " HRX_ALLOW_NO_SVM=1, which put the CWSR save area in GTT (docs/hrx.md)\n" @@ -2585,6 +2689,7 @@ int run_serve(int argc, char** argv) { else if (a == "--reranking") o.role = "reranking"; else if (a == "--rerank") o.rerank = next(); else if (a == "--npu-opt") o.npu_opts.push_back(next()); + else if (a == "--prefill-provider") o.prefill_provider = next(); else if (a == "--laya-model") o.laya_model = next(); else if (a == "--laya") o.laya_auto = true; else if (a == "--route-policy") o.route_policy = next(); diff --git a/app/serve_hybrid.h b/app/serve_hybrid.h new file mode 100644 index 00000000..1366fa8b --- /dev/null +++ b/app/serve_hybrid.h @@ -0,0 +1,271 @@ +// 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 serve --device hybrid` (docs/hybrid.md): the prompt's KV cache is computed by an external +// prefill provider, and llama-server decodes from it on HRX. The provider is a program named by +// ONEBIT_HYBRID_PREFILL (or --prefill-provider); the engine's public tree knows only its contract: +// +// --tokens --out --model [--elfs ] +// +// ids file the prompt's token ids, one decimal id per line, in order (what llama-server's +// /tokenize returned for the templated prompt) +// state file written by the provider: a llama.cpp sequence-state file for the prompt (the format +// llama_state_seq_save_file writes: the ids, then the KV cells at positions 0..n-1 of +// sequence 0), which llama-server restores into its slot with /slots/0?action=restore +// stdout one JSON line: {"n_tokens": N, "argmax": id, "prefill_ms": x} (prefill_ms optional) +// exit 0 on success; anything else, or a missing state file, falls back to a plain request +// +// The provider also sees ONEBIT_HYBRID_BACKEND_URL=http://127.0.0.1:, the decoding +// llama-server's own address (a test provider makes its state file through it; a real one has no +// use for it). ONEBIT_HYBRID_ELFS, when set, is passed on as --elfs. +// +// Per request: the chat template and tokenizer are the backend's own (/apply-template, /tokenize), +// so the provider sees exactly the ids the backend would have prefilled; the provider runs; the +// state file is restored into slot 0; the request is forwarded with cache_prompt on and id_slot 0, +// so the backend finds its prompt in the slot, re-evaluates the last token for its logits (llama-server +// always does), and decodes. One slot: hybrid requests are serialized here. Any failure on the way +// is said on stderr once per request and the request is forwarded as it came (the backend prefills +// it itself): the device is never unusable because the provider is. + +#pragma once + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#include +#include +#endif + +namespace hybrid { + +using json = nlohmann::json; + +struct Outcome { + bool handed_over = false; // the slot holds the provider's KV; the body carries id_slot/cache_prompt + size_t n_tokens = 0; + double prefill_ms = 0; // the provider's own figure when it gave one, else wall time + double restore_ms = 0; + std::string why; // when not handed over: the reason, for one stderr line +}; + +// The provider's stdout (its last line) and exit code, with a timeout in seconds (SIGTERM, then SIGKILL). +struct Run { + int code = -1; + std::string out; +}; + +inline Run run_provider(const std::vector& argv, const std::vector& env_extra, int timeout_s) { + Run r; +#if defined(_WIN32) + (void) argv; (void) env_extra; (void) timeout_s; + r.code = -1; + return r; +#else + int out[2] = {-1, -1}; + if (::pipe2(out, O_CLOEXEC) != 0) return r; + std::vector args; + for (const auto& a : argv) args.push_back(const_cast(a.c_str())); + args.push_back(nullptr); + std::vector env_store; + for (char** e = environ; *e; ++e) env_store.emplace_back(*e); + for (const auto& e : env_extra) env_store.push_back(e); + std::vector envp; + for (auto& e : env_store) envp.push_back(e.data()); + envp.push_back(nullptr); + const pid_t pid = ::fork(); + if (pid < 0) { ::close(out[0]); ::close(out[1]); return r; } + if (pid == 0) { + ::dup2(out[1], 1); + ::execvpe(args[0], args.data(), envp.data()); + ::_exit(127); + } + ::close(out[1]); + // read stdout to its end (the provider's exit closes it), then reap, with the timeout on the whole + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout_s); + bool killed = false; + char buf[4096]; + for (;;) { + fd_set fds; + FD_ZERO(&fds); + FD_SET(out[0], &fds); + timeval tv{1, 0}; + const int sel = ::select(out[0] + 1, &fds, nullptr, nullptr, &tv); + if (sel > 0) { + const ssize_t n = ::read(out[0], buf, sizeof buf); + if (n < 0 && errno == EINTR) continue; + if (n <= 0) break; + r.out.append(buf, size_t(n)); + continue; + } + if (std::chrono::steady_clock::now() > deadline) { + ::kill(pid, killed ? SIGKILL : SIGTERM); + if (killed) break; + killed = true; + } + } + ::close(out[0]); + int status = 0; + while (::waitpid(pid, &status, 0) < 0) + if (errno != EINTR) break; + r.code = killed ? -2 : (WIFEXITED(status) ? WEXITSTATUS(status) : -1); + return r; +#endif +} + +class Prefill { +public: + // provider: the program; elfs: passed as --elfs when not empty; slot_dir: the backend's + // --slot-save-path (the state files go there, by bare name, as llama-server requires) + Prefill(std::string provider, std::string elfs, std::string slot_dir, std::string model) + : provider_(std::move(provider)), elfs_(std::move(elfs)), slot_dir_(std::move(slot_dir)), model_(std::move(model)) { + if (const char* t = std::getenv("ONEBIT_HYBRID_PREFILL_TIMEOUT_S"); t && *t) timeout_s_ = std::max(1, std::atoi(t)); + } + + const std::string& provider() const { return provider_; } + + // Serializes hybrid requests: one slot, one handoff at a time, held through the decode by the caller. + std::mutex& slot_mutex() { return slot_mu_; } + + // The prompt as the backend sees it: the templated chat (/apply-template) or the completion's prompt. + // Empty when the request has no text prompt (an array prompt, images), in which case there is no handoff. + static std::string prompt_text(int port, const std::string& path, const json& body, std::string* why) { + if (path == "/v1/chat/completions") { + httplib::Client c("127.0.0.1", port); + c.set_read_timeout(60); + auto r = c.Post("/apply-template", body.dump(), "application/json"); + if (!r || r->status != 200) { *why = "apply-template answered " + std::to_string(r ? r->status : -1); return ""; } + try { + return json::parse(r->body).value("prompt", std::string()); + } catch (const std::exception&) { + *why = "apply-template reply is not JSON"; + return ""; + } + } + if (body.contains("prompt") && body["prompt"].is_string()) return body["prompt"].get(); + *why = "no text prompt"; + return ""; + } + + static std::vector tokenize(int port, const std::string& prompt, std::string* why) { + httplib::Client c("127.0.0.1", port); + c.set_read_timeout(60); + const json req = {{"content", prompt}, {"add_special", true}, {"parse_special", true}}; + auto r = c.Post("/tokenize", req.dump(), "application/json"); + if (!r || r->status != 200) { *why = "tokenize answered " + std::to_string(r ? r->status : -1); return {}; } + try { + return json::parse(r->body).at("tokens").get>(); + } catch (const std::exception&) { + *why = "tokenize reply has no tokens"; + return {}; + } + } + + // The whole handoff for one request. On success body carries id_slot 0 and cache_prompt true. + Outcome prepare(int port, const std::string& path, json& body) { + Outcome o; + const std::string prompt = prompt_text(port, path, body, &o.why); + if (prompt.empty()) { if (o.why.empty()) o.why = "empty prompt"; return o; } + const std::vector ids = tokenize(port, prompt, &o.why); + if (ids.size() < 2) { if (o.why.empty()) o.why = "prompt shorter than 2 tokens"; return o; } + o.n_tokens = ids.size(); + const unsigned n = ++seq_; + const std::string stem = "hybrid-" + std::to_string(n); + const std::filesystem::path tokens = std::filesystem::path(slot_dir_) / (stem + ".tokens"); + const std::filesystem::path state = std::filesystem::path(slot_dir_) / (stem + ".bin"); + { + std::ofstream f(tokens); + for (int32_t id : ids) f << id << '\n'; + if (!f) { o.why = "cannot write " + tokens.string(); return o; } + } + std::vector argv = {provider_, "--tokens", tokens.string(), "--out", state.string(), "--model", model_}; + if (!elfs_.empty()) { argv.push_back("--elfs"); argv.push_back(elfs_); } + const auto t0 = std::chrono::steady_clock::now(); + const Run run = run_provider(argv, {"ONEBIT_HYBRID_BACKEND_URL=http://127.0.0.1:" + std::to_string(port)}, timeout_s_); + o.prefill_ms = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + std::error_code ec; + std::filesystem::remove(tokens, ec); + if (run.code != 0) { + o.why = run.code == -2 ? "provider timed out after " + std::to_string(timeout_s_) + " s" + : "provider exited " + std::to_string(run.code); + std::filesystem::remove(state, ec); + return o; + } + if (!std::filesystem::exists(state)) { o.why = "provider wrote no state file"; return o; } + // the provider's own figures, from its last stdout line + try { + std::string last = run.out; + while (!last.empty() && (last.back() == '\n' || last.back() == '\r')) last.pop_back(); + if (const size_t nl = last.rfind('\n'); nl != std::string::npos) last = last.substr(nl + 1); + const json j = json::parse(last); + if (j.contains("prefill_ms") && j["prefill_ms"].is_number()) o.prefill_ms = j["prefill_ms"].get(); + if (j.contains("n_tokens") && j["n_tokens"].is_number_integer() && j["n_tokens"].get() != ids.size()) { + o.why = "provider prefilled " + std::to_string(j["n_tokens"].get()) + " of " + std::to_string(ids.size()) + " tokens"; + std::filesystem::remove(state, ec); + return o; + } + } catch (const std::exception&) { + // no JSON line: the file is the result + } + // restore into slot 0 (the backend reads slot_dir/; its reply counts the cells as n_restored) + httplib::Client c("127.0.0.1", port); + c.set_read_timeout(600); + const auto t1 = std::chrono::steady_clock::now(); + auto r = c.Post("/slots/0?action=restore", json{{"filename", stem + ".bin"}}.dump(), "application/json"); + o.restore_ms = std::chrono::duration(std::chrono::steady_clock::now() - t1).count(); + std::filesystem::remove(state, ec); + if (!r || r->status != 200) { + o.why = "restore answered " + std::to_string(r ? r->status : -1) + (r ? ": " + r->body.substr(0, 200) : ""); + return o; + } + try { + const json j = json::parse(r->body); + const size_t restored = j.value("n_restored", size_t(0)); + if (restored != ids.size()) { + o.why = "restore holds " + std::to_string(restored) + " of " + std::to_string(ids.size()) + " tokens"; + return o; + } + if (j.contains("timings") && j["timings"].contains("restore_ms")) o.restore_ms = j["timings"]["restore_ms"].get(); + } catch (const std::exception&) { + o.why = "restore reply is not JSON"; + return o; + } + body["id_slot"] = 0; + body["cache_prompt"] = true; + o.handed_over = true; + return o; + } + +private: + std::string provider_, elfs_, slot_dir_, model_; + int timeout_s_ = 600; + std::atomic seq_{0}; + std::mutex slot_mu_; +}; + +} // namespace hybrid diff --git a/docs/hybrid.md b/docs/hybrid.md new file mode 100644 index 00000000..a2f06445 --- /dev/null +++ b/docs/hybrid.md @@ -0,0 +1,107 @@ + +# Hybrid: NPU prefill, HRX decode (`--device hybrid`) + +``` +1bit serve -m Qwen3-0.6B-Q4_K_M.gguf --device hybrid --prefill-provider +``` + +Hybrid serves a `.gguf` with the prompt processed on one device and the tokens generated on another: +a **prefill provider** computes the prompt's KV cache, and the HRX llama-server (docs/hrx.md) +decodes from it on `HRX0`. The provider is the engine's NPU route; the public tree knows it only as +a program with the contract below, the way `-DONEBIT_NPU_PRIVATE` keeps the NPU kernels out of the +repository (docs/npu.md). Without a provider, `--device hybrid` is `--device hrx` with one slot. + +In a build without HRX the decoder is the same llama-server on the CPU, so the route runs (and is +tested) on any machine. + +## What happens per request + +1. **Template and tokenize on the backend.** For `/v1/chat/completions` serve asks the llama-server + child for the templated prompt (`/apply-template`) and its ids (`/tokenize`, `add_special` and + `parse_special` on, as the server itself tokenizes); `/v1/completions` tokenizes its `prompt` + string. The provider sees exactly the ids the backend would have prefilled. +2. **The provider prefills.** serve writes the ids to a file in the backend's slot directory and runs + the provider. It writes a llama.cpp sequence-state file for the prompt. +3. **Restore into slot 0.** `POST /slots/0?action=restore` loads the file into the one slot; serve + checks the slot holds all the prompt's tokens. +4. **Decode.** The request is forwarded with `id_slot` 0 and `cache_prompt` true. llama-server finds + the prompt in the slot, re-evaluates its last token (it always does, for the logits), and generates. + The reply, streamed or not, is relayed as for every device. + +One line per request on stderr: `hybrid: N prompt tokens: prefill X ms, restore Y ms, first byte Z ms` +(the provider's own `prefill_ms` when it reports one, else its wall time; the restore from the backend's +timing; first byte measured in serve from the start of the handoff). + +**Fallback.** Any failure on the way (no text prompt, a provider that exits non-zero, times out or +writes no file, a restore that fails or holds fewer tokens) is said on stderr with its reason, and the +request is forwarded as it came: the backend prefills it itself. Without a provider named, serve says +so once at start. The device is never unusable because the provider is. + +**One slot.** Hybrid runs the backend with `-np 1` and `--slot-save-path`, and serializes handoffs: a +request holds the slot from its handoff until the backend has started on it (its first byte), after +which llama-server itself defers the next restore until the slot is free. `--parallel` > 1 is refused; +so are `--mtp`, `--dflash`, `--mmproj`, `--moe-slots`, a role (`--embedding`, `--reranking`) and `--pm`, +and the GGUF's own MTP head is not used: nothing but the prompt's KV is handed over, and a drafter +would draft over a cache it did not compute. + +## The provider contract + +``` + --tokens --out --model [--elfs ] +``` + +| | | +|---|---| +| `--tokens` | the prompt's token ids, one decimal id per line, in order | +| `--out` | where to write the state: a bare filename inside the backend's `--slot-save-path` directory (serve chooses it; the backend restores by that name) | +| `--model` | the model file serve was given (`-m`) | +| `--elfs` | passed when `ONEBIT_HYBRID_ELFS` is set: where the provider's kernels are | +| stdout | one JSON line, last: `{"n_tokens": N, "argmax": id, "prefill_ms": x}`; `n_tokens` must equal the ids given (checked), `argmax` is the next-token argmax after the prompt (for diagnostics), `prefill_ms` is optional | +| exit | 0 on success; anything else, or no file, is the fallback above | +| environment | `ONEBIT_HYBRID_BACKEND_URL=http://127.0.0.1:`, the decoding llama-server's address (a test provider makes its file through it; a real one has no use for it), plus serve's own environment | +| time limit | `ONEBIT_HYBRID_PREFILL_TIMEOUT_S` (default 600): SIGTERM, then SIGKILL | + +**The state file** is what `llama_state_seq_save_file` writes and `/slots/?action=restore` reads +(llama.cpp `src/llama-context.cpp`, `src/llama-kv-cache.cpp` `state_write_meta` / `state_write_data`): +the file magic and version, the token count and ids, then the KV cells of sequence 0 at positions +0..N-1: cell count, each cell's position and sequence ids, then per layer the K type, row size and N +rows, then V the same way (transposed or not, as the server's cache is: `v_trans` 0 with flash +attention, which HRX uses). The K rows are post-RoPE keys, as llama.cpp stores them; the cache type is +the backend's (`f16` by default), so a provider computing in another type converts. A provider must +match the version of the pinned llama.cpp (`third_party/llama.cpp`): the restore refuses another. + +Naming: `--prefill-provider PATH` or `ONEBIT_HYBRID_PREFILL=PATH`; `ONEBIT_HYBRID_ELFS=DIR` adds +`--elfs DIR`. + +## Tests + +- `tests/hybrid_route.sh` (CTest `hybrid_route`, no model, no NPU): `tests/fake_backend.py` stands in + for llama-server and `tests/fake_prefill_provider.py` (`FAKE_PREFILL_COPY=1`) for the NPU engine. + Checks what reaches the backend (`id_slot` 0, `cache_prompt` on, the slot directory), the stderr + lines, the fallbacks (no provider; a provider that fails), the refused flags and `--print-backend`. +- `tests/hybrid_e2e.sh` (CTest `hybrid_e2e` with `ONEBIT_SERVE_TEST_GGUF`; skips without a model or a + llama-server): a real model through the whole path. The fake provider makes a real state file by + having the decoding llama-server itself prefill the ids into slot 0 and save the slot, so the + handoff is exercised end to end with no NPU. Checked: the greedy reply equals the plain `--device + cpu` reply, and the backend's `prompt_n` is 1 (the re-evaluated last token), not the prompt length. + +## Status + +The plumbing (this document) is public. The NPU prefill provider, its kernels and its state-file +writer are the engine's private NPU work; what it measures on Strix Halo is reported there when it +lands. diff --git a/docs/serve.md b/docs/serve.md index 7c2713f9..4a03e204 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -25,8 +25,9 @@ so it also works on its own with any OpenAI client. ```sh 1bit serve -m [--port 8000] [--host 127.0.0.1] - [--device auto|npu|hrx|cpu|ds4] [--ctx-size N] [--alias NAME] + [--device auto|npu|hrx|hybrid|cpu|ds4] [--ctx-size N] [--alias NAME] [--llama-server PATH] [--ds4 PATH] [--ssd-streaming] [--hrx-libhsa PATH] + [--prefill-provider PATH] --device hybrid: NPU prefill, HRX decode (docs/hybrid.md) [--hrx-cwsr-placement=auto|on|off] [--mtp HEAD.gguf | --dflash DRAFT.gguf | --no-mtp] [--mtp-max N] [--mtp-p-min P] [--mmproj MMPROJ.gguf] [--laya | --laya-model DIR] [--route-policy FILE] @@ -252,7 +253,7 @@ two processes (lemond, llama-server) instead of two hops and three. |---|---|---| | 0 | the plan is on stdout | starts it | | 3 | the model does not fit ([above](#will-it-fit)) | fails the load with serve's line in its log | -| 5 | the route needs serve in the request path: an NPU model directory, `--device npu` or `ds4`, `--pm`, `--laya`, `--embed` or `--rerank` | runs `1bit serve` as before | +| 5 | the route needs serve in the request path: an NPU model directory, `--device npu`, `hybrid` or `ds4`, `--pm`, `--laya`, `--embed` or `--rerank` | runs `1bit serve` as before | Any other exit (an older engine without the flag) also falls back to `1bit serve`, and `LEMONADE_ONEBIT_ONE_HOP=0` in lemond's environment turns one hop off. `env` lists only variables @@ -361,6 +362,7 @@ without a GPU. | Qwen3.6-35B-A3B Q4NX directory (`model_type` `qwen3_5_moe`) | `auto`, `npu` | a closed NPU add-on route, in process, in builds with `-DONEBIT_NPU_PRIVATE` ([npu.md](npu.md#add-on-routes)) | | `.gguf` | `auto`, `hrx` | the HRX build's llama-server on `HRX0` with our HIP kernels (docs/hrx.md), `--mtp` included; every node on HRX in all 43 configurations of the 2026-10-20 census (#476, #480, #481). `auto` exceptions: below | | `.gguf` | `cpu` | the same llama-server with no GPU layers; `auto` in a build without HRX | +| `.gguf` | `hybrid` | the HRX llama-server decodes, one slot, from a KV cache a prefill provider computed for the prompt (the NPU engine, named by `--prefill-provider` / `ONEBIT_HYBRID_PREFILL`) and serve restored into the slot; without a provider, HRX prefills too ([hybrid.md](hybrid.md)) | | DwarfStar `.gguf` (DeepSeek V4 Flash, GLM 5.x, Qwen3.8-Flash-Next in its own layouts) | `ds4` | this build's DwarfStar `ds4-server` (HIP; `--ssd-streaming` streams routed experts; docs/dwarfstar.md) | What `--device auto` does with a `.gguf` that HRX has no route for. Until RFC #213 stage 3 these diff --git a/tests/fake_backend.py b/tests/fake_backend.py index 246e0340..b7a96fa5 100755 --- a/tests/fake_backend.py +++ b/tests/fake_backend.py @@ -66,6 +66,26 @@ def client_gone(self): def do_POST(self): req = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + # what `--device hybrid` asks of its backend (app/serve_hybrid.h, tests/hybrid_route.sh): the + # templated prompt, its ids (one per character here), and a slot restore that counts the + # ids in the provider's state file (which the fake provider writes as one id per line) + if self.path == "/apply-template": + self._json(200, {"prompt": "".join(m.get("content", "") for m in req.get("messages", []))}) + return + if self.path == "/tokenize": + self._json(200, {"tokens": [ord(c) for c in req.get("content", "")]}) + return + if self.path.startswith("/slots/"): + slot, _, action = self.path[len("/slots/"):].partition("?action=") + if action != "restore": + self._json(400, {"error": {"message": "Invalid action"}}) + return + path = os.path.join(sys.argv[sys.argv.index("--slot-save-path") + 1], req["filename"]) + with open(path) as f: + n = sum(1 for line in f if line.strip()) + self._json(200, {"id_slot": int(slot), "filename": req["filename"], "n_restored": n, "n_read": 0, + "timings": {"restore_ms": 1.5}}) + return if req.get("fake_device_lost"): # what ggml-hrx logs after a GPU reset; the backend keeps running and fails every request print("graph_compute: launch HRX graph replay: hrx-system/runtime/src/iree/hal/drivers/amdgpu/system_event.c:129: " @@ -103,8 +123,11 @@ def do_POST(self): msg = {"message": {"role": "assistant", "content": text}} if "chat" in self.path else {"text": text} # llama-server's timings, fixed, for tools/bench.py (tests/bench_selftest.sh) timings = {"prompt_per_second": 500.0, "predicted_per_second": 20.0} + # the slot the request was pinned to and whether it may reuse the slot's prompt + # (what the hybrid handoff sets; tests/hybrid_route.sh reads them back) self._json(200, {"id": "x", "model": "fake-backend-id", "choices": [dict(index=0, finish_reason="stop", **msg)], - "timings": timings}) + "timings": timings, "fake_id_slot": req.get("id_slot", -1), + "fake_cache_prompt": req.get("cache_prompt", None)}) # "fake_exit": N: the backend exits with status N shortly after this reply, on its own # (a crashed backend; serve must notice and not wait its SIGTERM grace for a dead pid) if "fake_exit" in req: diff --git a/tests/fake_prefill_provider.py b/tests/fake_prefill_provider.py new file mode 100755 index 00000000..600b5c9a --- /dev/null +++ b/tests/fake_prefill_provider.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +# 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. +# +# fake_prefill_provider.py: a prefill provider for `1bit serve --device hybrid` (docs/hybrid.md) that +# has no NPU. It honours the provider contract (app/serve_hybrid.h): +# +# --tokens --out --model [--elfs DIR] +# stdout: one JSON line {"n_tokens": N, "argmax": id, "prefill_ms": x}; exit 0 +# +# Two ways to make the state file: +# FAKE_PREFILL_COPY=1 copy the ids file (tests/fake_backend.py's restore counts its lines; +# tests/hybrid_route.sh, no model needed) +# default ask the decoding llama-server itself (ONEBIT_HYBRID_BACKEND_URL, which serve +# sets for every provider) to prefill the ids into slot 0 and save the slot: the +# state file is then a real sequence-state file for a real model, made by the +# same model that will decode from it (tests/hybrid_e2e.sh). The file lands in +# the backend's --slot-save-path, which is where --out points. +# +# A real provider computes the KV itself (the NPU engine) and writes the same file format. +import json, os, shutil, sys, time, urllib.request + +args = sys.argv[1:] +opt = {args[i]: args[i + 1] for i in range(0, len(args) - 1, 2)} +tokens_path, out_path = opt["--tokens"], opt["--out"] +with open(tokens_path) as f: + ids = [int(line) for line in f if line.strip()] +t0 = time.time() + +if os.environ.get("FAKE_PREFILL_COPY"): + shutil.copyfile(tokens_path, out_path) + print(json.dumps({"n_tokens": len(ids), "argmax": ids[-1], "prefill_ms": round((time.time() - t0) * 1000, 1)}), flush=True) + sys.exit(0) + +url = os.environ["ONEBIT_HYBRID_BACKEND_URL"] + + +def post(path, body): + req = urllib.request.Request(url + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=600) as r: + return json.loads(r.read()) + + +# the prompt through the backend, no generation: the slot then holds exactly the prompt's cells +reply = post("/completion", {"prompt": ids, "n_predict": 0, "id_slot": 0, "cache_prompt": False, "temperature": 0}) +# the slot's state, by bare name into the backend's --slot-save-path (where --out points) +saved = post("/slots/0?action=save", {"filename": os.path.basename(out_path)}) +n = int(saved.get("n_saved", 0)) +# the slot is emptied again, so the file is the only way the KV reaches the decode +post("/slots/0?action=erase", {}) +if n != len(ids): + print(f"fake_prefill_provider: saved {n} tokens for {len(ids)} ids", file=sys.stderr, flush=True) + sys.exit(1) +print(json.dumps({"n_tokens": n, "argmax": -1, "prefill_ms": round((time.time() - t0) * 1000, 1), + "prompt_n": reply.get("timings", {}).get("prompt_n")}), flush=True) diff --git a/tests/hybrid_e2e.sh b/tests/hybrid_e2e.sh new file mode 100755 index 00000000..b83767e3 --- /dev/null +++ b/tests/hybrid_e2e.sh @@ -0,0 +1,95 @@ +#!/usr/bin/env bash +# 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 serve --device hybrid` (docs/hybrid.md) end to end on a real model, with no NPU: the decoding +# llama-server is the CPU build (or the HRX one, where the build has it), and the prefill provider is +# tests/fake_prefill_provider.py, which makes the state file by having that same server prefill the ids +# into slot 0 and save the slot. So the whole path runs: template + tokenize, the provider, the +# restore, the request pinned to the slot. Checked: the greedy reply equals the plain --device cpu +# reply for the same request, and the backend's timings show it prefilled one token (the last one, +# which llama-server always re-evaluates) rather than the prompt. +# +# usage: tests/hybrid_e2e.sh path/to/1bit [llama-server] +# Skips (exit 77) without a model or a llama-server. +set -uo pipefail + +bin=${1:?usage: hybrid_e2e.sh path/to/1bit [llama-server]} +model=${2:-} +server=${3:-} +here=$(cd "$(dirname "$0")" && pwd) +[ -n "$model" ] && [ -f "$model" ] || { echo "skip: no model"; exit 77; } +if [ -z "$server" ]; then + "$bin" serve -m "$model" --device cpu --print-backend --no-mem-check >/dev/null 2>&1 || { echo "skip: no llama-server in this build"; exit 77; } + server_arg=() +else + [ -x "$server" ] || { echo "skip: no llama-server at $server"; exit 77; } + server_arg=(--llama-server "$server") +fi +scratch=$(mktemp -d) +fail=0 +check() { if eval "$2"; then echo "ok $1"; else echo "FAIL $1"; fail=1; fi; } +pid= +cleanup() { [ -n "$pid" ] && { kill "$pid" 2>/dev/null; wait "$pid" 2>/dev/null; }; rm -rf "$scratch"; } +trap cleanup EXIT + +start() { # start [env...] + local device=$1 log=$2; shift 2 + port=$(python3 -c 'import socket; s=socket.socket(); s.bind(("127.0.0.1", 0)); print(s.getsockname()[1])') + env "$@" "$bin" serve -m "$model" --device "$device" --port "$port" --alias e2e-model --ctx-size 2048 \ + --no-mem-check "${server_arg[@]}" >"$log" 2>&1 & + pid=$! + api="http://127.0.0.1:$port" + local code=000 + for _ in $(seq 1 1200); do + code=$(curl -s -o /dev/null -w '%{http_code}' "$api/health"); [ "$code" = 200 ] && break + kill -0 "$pid" 2>/dev/null || break; sleep 0.5 + done + [ "$code" = 200 ] +} +stop() { kill "$pid" 2>/dev/null; wait "$pid" 2>/dev/null; pid=; } +body='{"model": "e2e-model", "messages": [{"role": "user", "content": "Name three colours of the rainbow, then stop."}], + "temperature": 0, "max_tokens": 24, "chat_template_kwargs": {"enable_thinking": false}}' +content() { python3 -c 'import json,sys; print(json.load(sys.stdin)["choices"][0]["message"]["content"])'; } +field() { python3 -c 'import json,sys; d=json.load(sys.stdin); print(d.get("timings", {}).get(sys.argv[1], ""))' "$1"; } + +check "cpu: serve starts" 'start cpu "$scratch/cpu.log"' +plain=$(curl -s "$api/v1/chat/completions" -H 'Content-Type: application/json' -d "$body") +plain_text=$(echo "$plain" | content) +plain_prompt_n=$(echo "$plain" | field prompt_n) +echo " cpu said: $plain_text (prompt_n $plain_prompt_n)" +check "cpu: answered" '[ -n "$plain_text" ]' +stop + +check "hybrid: serve starts with the fake provider" 'start hybrid "$scratch/hybrid.log" ONEBIT_HYBRID_PREFILL="$here/fake_prefill_provider.py"' +reply=$(curl -s "$api/v1/chat/completions" -H 'Content-Type: application/json' -d "$body") +text=$(echo "$reply" | content) +prompt_n=$(echo "$reply" | field prompt_n) +echo " hybrid said: $text (prompt_n $prompt_n)" +check "hybrid: the greedy reply equals the plain one" '[ "$text" = "$plain_text" ]' +check "hybrid: the backend prefilled only the last token (prompt_n $prompt_n, plain $plain_prompt_n)" '[ "$prompt_n" = 1 ] && [ "$plain_prompt_n" -gt 1 ]' +sleep 0.3 +check "hybrid: stderr has the per-request line" 'grep -q "hybrid: [0-9]* prompt tokens: prefill .* ms, restore .* ms, first byte .* ms" "$scratch/hybrid.log"' +check "hybrid: no fallback happened" '! grep -q "the backend prefills this request itself" "$scratch/hybrid.log"' +# a second turn: the handoff again, same answer +reply2=$(curl -s "$api/v1/chat/completions" -H 'Content-Type: application/json' -d "$body") +check "hybrid: the second request answers the same" '[ "$(echo "$reply2" | content)" = "$plain_text" ]' +chunks=$(curl -sN "$api/v1/chat/completions" -H 'Content-Type: application/json' \ + -d "${body%\}}, \"stream\": true}" | grep -c '^data: {') +check "hybrid: streams ($chunks chunks)" '[ "$chunks" -gt 3 ]' +stop + +if [ $fail -ne 0 ]; then echo "--- hybrid log (tail)"; tail -40 "$scratch/hybrid.log"; echo FAIL; exit 1; fi +echo PASS diff --git a/tests/hybrid_route.sh b/tests/hybrid_route.sh new file mode 100755 index 00000000..58753566 --- /dev/null +++ b/tests/hybrid_route.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +# 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 serve --device hybrid` (docs/hybrid.md) without a model or an NPU: tests/fake_backend.py +# stands in for llama-server (it answers /apply-template, /tokenize and the slot restore) and +# tests/fake_prefill_provider.py for the NPU engine (FAKE_PREFILL_COPY=1: the state file is the ids +# file). Checks the handoff order and what reaches the backend: the request pinned to slot 0 with +# cache_prompt on, the backend's --slot-save-path, the per-request stderr line, and the fallback +# (a provider that fails leaves the request to the backend, which still answers). +# +# usage: tests/hybrid_route.sh path/to/1bit +set -uo pipefail + +bin=${1:?usage: hybrid_route.sh path/to/1bit} +here=$(cd "$(dirname "$0")" && pwd) +scratch=$(mktemp -d) +touch "$scratch/tiny.gguf" +fail=0 +check() { if eval "$2"; then echo "ok $1"; else echo "FAIL $1"; fail=1; fi; } +pid= +cleanup() { [ -n "$pid" ] && { kill "$pid" 2>/dev/null; wait "$pid" 2>/dev/null; }; rm -rf "$scratch"; } +trap cleanup EXIT + +start() { # start [env...]: serve on a free port with the fake backend, env applied to serve + local log=$1; shift + port=$(python3 -c 'import socket; s=socket.socket(); s.bind(("127.0.0.1", 0)); print(s.getsockname()[1])') + env "$@" "$bin" serve -m "$scratch/tiny.gguf" --device hybrid --port "$port" --alias hybrid-model \ + --llama-server "$here/fake_backend.py" --no-mem-check >"$log" 2>&1 & + pid=$! + api="http://127.0.0.1:$port" + local code=000 + for _ in $(seq 1 100); do code=$(curl -s -o /dev/null -w '%{http_code}' "$api/health"); [ "$code" = 200 ] && break; sleep 0.1; done + [ "$code" = 200 ] +} +stop() { kill "$pid" 2>/dev/null; wait "$pid" 2>/dev/null; pid=; } + +# 1. the handoff: provider named, the backend sees id_slot 0 and cache_prompt true +check "serve starts with the fake provider" 'start "$scratch/a.log" ONEBIT_HYBRID_PREFILL="$here/fake_prefill_provider.py" FAKE_PREFILL_COPY=1' +reply=$(curl -s "$api/v1/chat/completions" -H 'Content-Type: application/json' \ + -d '{"model": "hybrid-model", "messages": [{"role": "user", "content": "What is the capital of France?"}]}') +check "chat reply comes through" '[[ "$reply" == *Paris* ]]' +check "request pinned to slot 0 (fake_id_slot)" '[[ "$reply" == *"\"fake_id_slot\":0"* ]]' +check "cache_prompt on (fake_cache_prompt)" '[[ "$reply" == *"\"fake_cache_prompt\":true"* ]]' +sleep 0.3 +check "stderr: the per-request line (prefill, restore, first byte)" 'grep -q "hybrid: 30 prompt tokens: prefill .* ms, restore .* ms, first byte .* ms" "$scratch/a.log"' +check "stderr: the startup line names the provider" 'grep -q "hybrid: prefill by .*fake_prefill_provider.py" "$scratch/a.log"' +reply=$(curl -s "$api/v1/completions" -H 'Content-Type: application/json' -d '{"model": "hybrid-model", "prompt": "Hello hybrid"}') +check "/v1/completions goes through the handoff too" '[[ "$reply" == *"\"fake_id_slot\":0"* ]]' +chunks=$(curl -sN "$api/v1/chat/completions" -H 'Content-Type: application/json' \ + -d '{"model": "hybrid-model", "stream": true, "messages": [{"role": "user", "content": "hi there"}]}' | grep -c '^data: {') +check "streamed reply relayed ($chunks chunks)" '[ "$chunks" = 7 ]' +check "no handoff file left in the slot directory" '[ -z "$(ls /tmp/1bit-hybrid-*/ 2>/dev/null)" ]' +slot_dirs_before=$(ls -d /tmp/1bit-hybrid-* 2>/dev/null | wc -l) +stop +slot_dirs_after=$(ls -d /tmp/1bit-hybrid-* 2>/dev/null | wc -l) +check "the slot directory is removed at exit ($slot_dirs_before -> $slot_dirs_after)" '[ "$slot_dirs_after" -lt "$slot_dirs_before" ]' + +# 2. no provider: the device serves, the backend prefills (said once) +check "serve starts without a provider" 'start "$scratch/b.log"' +reply=$(curl -s "$api/v1/chat/completions" -H 'Content-Type: application/json' \ + -d '{"model": "hybrid-model", "messages": [{"role": "user", "content": "hi"}]}') +check "answers without a provider" '[[ "$reply" == *Paris* ]]' +check "not pinned to a slot without a provider" '[[ "$reply" == *"\"fake_id_slot\":-1"* ]]' +check "stderr: no provider, said once" '[ "$(grep -c "hybrid: no prefill provider" "$scratch/b.log")" = 1 ]' +stop + +# 3. a provider that fails: the request is forwarded as it came +printf '#!/bin/sh\nexit 3\n' > "$scratch/bad_provider.sh"; chmod +x "$scratch/bad_provider.sh" +check "serve starts with a failing provider" 'start "$scratch/c.log" ONEBIT_HYBRID_PREFILL="$scratch/bad_provider.sh"' +reply=$(curl -s "$api/v1/chat/completions" -H 'Content-Type: application/json' \ + -d '{"model": "hybrid-model", "messages": [{"role": "user", "content": "hi"}]}') +check "answers when the provider fails" '[[ "$reply" == *Paris* ]]' +check "not pinned when the provider failed" '[[ "$reply" == *"\"fake_id_slot\":-1"* ]]' +check "stderr: the failure named" 'grep -q "hybrid: provider exited 3; the backend prefills this request itself" "$scratch/c.log"' +stop + +# 4. refused combinations and --print-backend +out=$("$bin" serve -m "$scratch/tiny.gguf" --device hybrid --parallel 2 --llama-server "$here/fake_backend.py" 2>&1) +check "--parallel 2 is refused" '[[ "$out" == *"one slot"* ]]' +"$bin" serve -m "$scratch/tiny.gguf" --device hybrid --print-backend --llama-server "$here/fake_backend.py" >/dev/null 2>&1 +check "--print-backend says the route needs serve (exit 5)" '[ $? = 5 ]' + +if [ $fail -ne 0 ]; then for l in a b c; do echo "--- $l.log"; tail -20 "$scratch/$l.log" 2>/dev/null; done; echo FAIL; exit 1; fi +echo PASS From fe391df36c9a0fa7c7750107eafbd0c1ee839775 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Fri, 23 Oct 2026 10:28:19 -0300 Subject: [PATCH 2/4] tests: fake backend validates the slot-restore file name CodeQL flagged the restore path as user-controlled. The fake backend now accepts only a bare name inside --slot-save-path, as llama-server's fs_validate_filename does, and answers 400 otherwise. Co-Authored-By: Claude Fable 5.1 --- tests/fake_backend.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/fake_backend.py b/tests/fake_backend.py index b7a96fa5..75c4fb45 100755 --- a/tests/fake_backend.py +++ b/tests/fake_backend.py @@ -80,7 +80,14 @@ def do_POST(self): if action != "restore": self._json(400, {"error": {"message": "Invalid action"}}) return - path = os.path.join(sys.argv[sys.argv.index("--slot-save-path") + 1], req["filename"]) + # the file name is a bare name inside --slot-save-path, as llama-server requires + # (fs_validate_filename): no directory part, no "..", nothing outside that directory + name = str(req.get("filename", "")) + slot_dir = os.path.realpath(sys.argv[sys.argv.index("--slot-save-path") + 1]) + path = os.path.realpath(os.path.join(slot_dir, name)) + if not name or name != os.path.basename(name) or name in (".", "..") or os.path.dirname(path) != slot_dir: + self._json(400, {"error": {"message": "Invalid filename"}}) + return with open(path) as f: n = sum(1 for line in f if line.strip()) self._json(200, {"id_slot": int(slot), "filename": req["filename"], "n_restored": n, "n_read": 0, From 999d885980ea05221ad80bd72f3d337171f757b1 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Fri, 23 Oct 2026 10:31:01 -0300 Subject: [PATCH 3/4] tests: fake backend resolves the restore file from the directory listing The restore path is taken from os.listdir of --slot-save-path, keyed by the requested bare name, so no request string reaches a path expression. Co-Authored-By: Claude Fable 5.1 --- tests/fake_backend.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/fake_backend.py b/tests/fake_backend.py index 75c4fb45..b9fad574 100755 --- a/tests/fake_backend.py +++ b/tests/fake_backend.py @@ -83,9 +83,11 @@ def do_POST(self): # the file name is a bare name inside --slot-save-path, as llama-server requires # (fs_validate_filename): no directory part, no "..", nothing outside that directory name = str(req.get("filename", "")) - slot_dir = os.path.realpath(sys.argv[sys.argv.index("--slot-save-path") + 1]) - path = os.path.realpath(os.path.join(slot_dir, name)) - if not name or name != os.path.basename(name) or name in (".", "..") or os.path.dirname(path) != slot_dir: + slot_dir = sys.argv[sys.argv.index("--slot-save-path") + 1] + # the path comes from the directory listing, never from the request string itself + files = {entry: os.path.join(slot_dir, entry) for entry in os.listdir(slot_dir)} + path = files.get(name) if name == os.path.basename(name) else None + if path is None: self._json(400, {"error": {"message": "Invalid filename"}}) return with open(path) as f: From 102cf06b6c6a172c4db5fa01df188d6e6a046b83 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Fri, 23 Oct 2026 10:36:11 -0300 Subject: [PATCH 4/4] tests: fake backend matches the restore name against the listing Co-Authored-By: Claude Fable 5.1 --- tests/fake_backend.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/fake_backend.py b/tests/fake_backend.py index b9fad574..f2b6f5cc 100755 --- a/tests/fake_backend.py +++ b/tests/fake_backend.py @@ -85,8 +85,10 @@ def do_POST(self): name = str(req.get("filename", "")) slot_dir = sys.argv[sys.argv.index("--slot-save-path") + 1] # the path comes from the directory listing, never from the request string itself - files = {entry: os.path.join(slot_dir, entry) for entry in os.listdir(slot_dir)} - path = files.get(name) if name == os.path.basename(name) else None + path = None + for entry in sorted(os.listdir(slot_dir)): + if entry == name: + path = os.path.join(slot_dir, entry) if path is None: self._json(400, {"error": {"message": "Invalid filename"}}) return