diff --git a/CMakeLists.txt b/CMakeLists.txt index 271eb2f7b..1a2073004 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -67,6 +67,7 @@ target_compile_definitions(onebit PRIVATE ONEBIT_BUILD_DIR="${CMAKE_BINARY_DIR}" # Runs everywhere, CI included: serve's proxy against a fake backend. add_test(NAME smoke_serve COMMAND ${CMAKE_SOURCE_DIR}/tests/smoke_serve.sh $) add_test(NAME hadamard_route COMMAND ${CMAKE_SOURCE_DIR}/tests/hadamard_route.sh $) +add_test(NAME long_route COMMAND ${CMAKE_SOURCE_DIR}/tests/long_route.sh $) # 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/app/serve.cpp b/app/serve.cpp index e1b1e198f..ce6ea959c 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -56,6 +56,9 @@ // drafts a block per step (docs/serve.md) // --mmproj on any llama.cpp route -> images: the vision encoder runs beside the // model, and chat messages may carry image parts +// --long-model -> conversations that start with --long-from tokens +// (default 2048) go to that Hadamard-rotated Q4_0 +// on ROCm (W4A4 prefill), the rest to -m's route // a Hugging Face id, --device mlx (macOS) -> lemon-mlx-engine's server // For a .gguf, the engine starts that server as a private child on a loopback // port and forwards the OpenAI routes to it, streaming included. `auto` picks @@ -158,6 +161,8 @@ struct Options { 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) + std::string long_model; // --long-model FILE: a Hadamard-rotated Q4_0 for long prompts (W4A4 on ROCm) + int long_from = 2048; // --long-from N: a conversation whose first prompt has N tokens or more }; // HRX dlopens the HSA runtime, and a distro libhsa rejects gfx1151's @@ -846,6 +851,38 @@ std::string request_state(const httplib::Request& req) { return out; } +// The prompt length of a request, counted by a llama-server backend: /v1/completions' prompt as +// is, a chat through the model's template (/apply-template, which takes the request body as it +// comes, tools included). If the backend cannot tell, about four bytes a token. +long prompt_tokens(int port, const httplib::Request& q) { + json body; + try { + body = json::parse(q.body); + } catch (const std::exception&) { + return long(q.body.size() / 4); + } + httplib::Client c("127.0.0.1", port); + c.set_read_timeout(std::chrono::seconds(30)); + std::string prompt; + if (body.contains("prompt") && body["prompt"].is_string()) { + prompt = body["prompt"].get(); + } else if (auto t = c.Post("/apply-template", q.body, "application/json"); t && t->status == 200) { + try { + prompt = json::parse(t->body).value("prompt", ""); + } catch (const std::exception&) { + } + } + if (prompt.empty()) return long(request_state(q).size() / 4); + if (auto t = c.Post("/tokenize", json{{"content", prompt}}.dump(), "application/json"); t && t->status == 200) { + try { + const json tok = json::parse(t->body); + if (tok.contains("tokens") && tok["tokens"].is_array()) return long(tok["tokens"].size()); + } catch (const std::exception&) { + } + } + return long(prompt.size() / 4); +} + int serve_child(const Options& given) { Options o = given; if (hadamard_q4_0(o.model)) { @@ -937,6 +974,34 @@ int serve_child(const Options& given) { } else { backends.push_back(launch_for(o, o.device == "auto" ? "vulkan" : o.device, free_port())); } + // --long-model: the same model as a Hadamard-rotated Q4_0 on the ROCm W4A4 route, for + // conversations that start with a long prompt. What the user waits for is the first token + // plus the answer: with a short prompt the faster decode wins (Vulkan), with a long one the + // faster prefill does (W4A4 from about 2,000 tokens; docs/lean.md, "End to end"). Both + // backends load at the start; a conversation stays where its first request went, so its + // prompt cache stays useful as it grows. + const bool long_route = !o.long_model.empty(); + if (long_route) { + if (laya || o.adaptive || !o.role.empty() || !o.mmproj.empty()) + throw std::runtime_error("--long-model does not combine with --laya, --adaptive, --embedding/--reranking or --mmproj"); + if (!hadamard_q4_0(o.long_model)) + throw std::runtime_error("--long-model takes a Hadamard-rotated Q4_0 (tools/hadamard_q4_0.py, docs/lean.md): " + o.long_model); + if (hadamard_q4_0(o.model)) + throw std::runtime_error(o.model + " is already on the W4A4 route; -m is the model for short prompts"); + const std::string& d = backends[0].device; + if (d != "vulkan" && d != "hrx" && d != "rocm") + throw std::runtime_error("--long-model needs -m on a llama-server route (vulkan, hrx or rocm), not " + d); + if (o.long_from < 1) throw std::runtime_error("--long-from takes a positive token count"); + Options lo = o; + lo.model = o.long_model; + lo.mtp.clear(); + lo.dflash.clear(); + lo.llama_server.clear(); + lo.prefill_device.clear(); + lo.moe_slots = 0; + lo.lean = false; + backends.push_back(launch_for(lo, "rocm", free_port())); + } const std::string device = laya ? "laya" : o.adaptive ? "vulkan+rocm" : backends[0].device; const bool drop_model = backends[0].drop_model; const std::string set_model = backends[0].set_model; @@ -981,6 +1046,10 @@ int serve_child(const Options& given) { std::vector> inflight(backends.size()); std::vector> routed(backends.size()); std::mutex route_mu; + // --long-model: conversation -> (backend, prompt tokens of its first request) + std::mutex long_mu; + std::unordered_map> long_cache; + std::deque long_order; #ifdef ONEBIT_LAYA std::mutex scorer_mu, class_mu; std::unordered_map class_cache; // conversation -> class @@ -1047,6 +1116,33 @@ int serve_child(const Options& given) { return; } #endif + if (long_route) { + const size_t key = conversation_key(q); + std::pair side{0, -1}; + bool fresh = true; + { + std::lock_guard lock(long_mu); + if (const auto c = long_cache.find(key); c != long_cache.end()) side = c->second, fresh = false; + } + if (fresh) { + side.second = prompt_tokens(backends[0].port, q); + side.first = side.second >= o.long_from ? 1 : 0; + std::lock_guard lock(long_mu); + if (long_cache.size() >= 4096) { // oldest conversation first + long_cache.erase(long_order.front()); + long_order.pop_front(); + } + if (long_cache.emplace(key, side).second) long_order.push_back(key); + } + const size_t i = side.first; + char route_hdr[96]; + std::snprintf(route_hdr, sizeof route_hdr, "%s %ld %s", i ? "long" : "short", side.second, backends[i].device.c_str()); + held = std::make_shared(&inflight[i]); + ++routed[i]; + forward(backends[i].port, held, -1, id, backends[i].drop_model, backends[i].set_model, q, r); + r.set_header("X-1bit-Route", route_hdr); // short or long, the first prompt's tokens, device + return; + } { std::lock_guard lock(route_mu); if (backends.size() > 1 && inflight[0] >= o.adaptive_at) { @@ -1352,6 +1448,8 @@ void usage(FILE* out) { " [--mmproj MMPROJ.gguf] images in chat messages (vulkan, hrx, rocm)\n" " [--parallel N] N requests decoded together (continuous batching)\n" " [--adaptive] [--adaptive-at N] Vulkan (+MTP) for N in flight (default 1), ROCm batches the rest\n" + " [--long-model ROTATED.gguf] [--long-from N] conversations starting with N or more prompt\n" + " tokens (default 2048) go to that Hadamard-rotated Q4_0 on ROCm (W4A4)\n" " [--embed MODEL.gguf] [--rerank MODEL.gguf] RAG: /v1/embeddings and /v1/rerank\n" " [--embedding | --reranking] serve the model itself as an embedding or reranking model\n" " [--npu-opt KEY=VALUE ...] an option for a private NPU route (docs/npu.md)\n" @@ -1402,6 +1500,8 @@ int run_serve(int argc, char** argv) { else if (a == "--parallel") o.parallel = std::stoi(next()); else if (a == "--adaptive") o.adaptive = true; else if (a == "--adaptive-at") o.adaptive_at = std::stoi(next()); + else if (a == "--long-model") o.long_model = next(); + else if (a == "--long-from") o.long_from = std::stoi(next()); else if (a == "--embed") o.embed = next(); else if (a == "--embedding" || a == "--embeddings") o.role = "embedding"; else if (a == "--reranking") o.role = "reranking"; diff --git a/docs/lean.md b/docs/lean.md index 54466dbfb..40b6a3e4f 100644 --- a/docs/lean.md +++ b/docs/lean.md @@ -247,6 +247,9 @@ effective speed is +56%. Agents and coding tools send long prompts (files, tool history), so that is where this route is for. The two files are not the same quantization (Q4_K_M against rotated Q4_0), so the table compares routes as they ship, not kernels. +`1bit serve --long-model` takes both at once: short conversations on Vulkan, long ones on +this route ([serve.md](serve.md#short-and-long-prompts---long-model)). + The rotation removes a third of the 4-bit error. Against the exact path the rotated file costs +1.2% perplexity for +28% prompt speed. The last row shows the rotation itself is exact. Through `1bit serve -m Qwen3.8-27B-Q4_0-H32.gguf` (no device flag: the stamp picks the route), a diff --git a/docs/serve.md b/docs/serve.md index bd4b6f7a4..00d190d0a 100644 --- a/docs/serve.md +++ b/docs/serve.md @@ -281,6 +281,36 @@ It needs both builds (`ONEBIT_VULKAN`, and `ONEBIT_LEAN` + `ONEBIT_LEAN_ROCM`) a two copies of the model; `--ctx-size` applies to each backend and is split across its slots. On exit, serve logs how many requests each backend took. +## Short and long prompts (`--long-model`) + +What a user waits for is the first token plus the answer. With a short prompt the faster +decode wins (Vulkan); with a long one the faster prefill does (the W4A4 route, from about +2,000 prompt tokens; [lean.md](lean.md#end-to-end-time-to-first-token-and-effective-speed)). +`--long-model` serves the same model both ways: `-m` on its usual route, and a Hadamard-rotated +Q4_0 of it (`tools/hadamard_q4_0.py`) on ROCm with W4A4. A conversation whose first request has +`--long-from` prompt tokens or more (default 2048; chats are counted through the model's +template) goes to the rotated file, the rest to `-m`. A conversation stays where its first +request went, so its prompt cache stays useful as it grows. Each response carries +`X-1bit-Route: short|long `. + +```sh +1bit serve -m Qwen3-Coder-30B-A3B-Instruct-Q4_K_M.gguf --long-model Qwen3-Coder-30B-A3B-Q4_0-H32.gguf --ctx-size 20480 +``` + +Qwen3-Coder-30B-A3B, 256 output tokens, effective tok/s (output tokens over the whole wait; +`tools/e2e_bench.py`, 2026-09-28): + +| Prompt tokens | 512 | 2,048 | 8,192 | 16,384 | +|---|---|---|---|---| +| Vulkan, Q4_K_M | **70.7** | 52.1 | 20.3 | 8.6 | +| ROCm W4A4, rotated Q4_0 | 66.6 | **55.0** | **27.0** | **13.4** | +| **`--long-model`** | **71.3** | 53.1 | 25.5 | 13.1 | + +The `--long-model` row ran while another job used the box's memory bus, which cost its long +side 2-5% against the W4A4 row. Both files stay loaded (here 17 + 16 GiB). It needs the Vulkan +build and `ONEBIT_LEAN` + `ONEBIT_LEAN_ROCM`, and `-m` on a llama-server route (vulkan, hrx or +rocm); it does not combine with `--laya`, `--adaptive` or `--mmproj`. + ## RAG (`--embed`, `--rerank`) `--embed MODEL.gguf` serves `/v1/embeddings` and `--rerank MODEL.gguf` serves `/v1/rerank`, diff --git a/tests/long_route.sh b/tests/long_route.sh new file mode 100755 index 000000000..454cfb0bb --- /dev/null +++ b/tests/long_route.sh @@ -0,0 +1,115 @@ +#!/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. +# +# How `1bit serve --long-model` routes by prompt length (docs/lean.md, "End to end"), without a +# GPU. Fake llama-servers stand in for Vulkan0 (-m) and ROCm0 (--long-model); each answers +# /apply-template and /tokenize (one token per word) and names its device in every reply: +# - a short first prompt goes to Vulkan0, a long one (>= --long-from tokens) to ROCm0, +# - chats are counted through /apply-template, completions from their prompt, +# - a conversation stays where its first request went, even once it is long, +# - the ROCm backend runs with the Hadamard W4A4 environment, +# - an unrotated --long-model and a rotated -m are refused. +# +# usage: tests/long_route.sh path/to/1bit +set -uo pipefail + +bin=${1:?usage: long_route.sh path/to/1bit} +scratch=$(mktemp -d) +trap 'kill -9 $pid 2>/dev/null; wait $pid 2>/dev/null; rm -rf "$scratch"' EXIT +pid= +fail=0 +check() { if eval "$2"; then echo "ok $1"; else echo "FAIL $1"; fail=1; fi; } + +python3 - "$scratch" <<'PY' +import struct, sys +def gguf(path, stamp): + kv = [] + def s(x): b = x.encode(); return struct.pack(" "$scratch/backend.py" <<'PY' +#!/usr/bin/env python3 +import json, os, sys +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +argv = sys.argv[1:] +port = int(argv[argv.index("--port") + 1]) +dev = argv[argv.index("--device") + 1] +env = {k: v for k, v in os.environ.items() if k.startswith("GGML_")} +open(os.path.join(os.environ["RECORDS"], dev + ".json"), "w").write(json.dumps({"argv": argv, "env": env})) +class H(BaseHTTPRequestHandler): + def log_message(self, *a): pass + def _json(self, obj): + b = json.dumps(obj).encode() + self.send_response(200); self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(b))); self.end_headers(); self.wfile.write(b) + def do_GET(self): self._json({"status": "ok"}) + def do_POST(self): + req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))) or b"{}") + if self.path == "/apply-template": + return self._json({"prompt": " ".join("<%s> %s" % (m["role"], m["content"]) for m in req["messages"])}) + if self.path == "/tokenize": + return self._json({"tokens": list(range(len(req["content"].split())))}) + text = "device:" + dev + msg = {"message": {"role": "assistant", "content": text}} if "chat" in self.path else {"text": text} + self._json({"id": "x", "model": "fake", "choices": [dict(index=0, finish_reason="stop", **msg)]}) +ThreadingHTTPServer(("127.0.0.1", port), H).serve_forever() +PY +chmod +x "$scratch/backend.py" + +bad=$("$bin" serve -m "$scratch/plain.gguf" --long-model "$scratch/plain.gguf" --port 1 --llama-server "$scratch/backend.py" 2>&1) +check "an unrotated --long-model is refused" '[[ "$bad" == *"Hadamard-rotated"* ]]' +bad=$("$bin" serve -m "$scratch/h32.gguf" --long-model "$scratch/h32.gguf" --port 1 --llama-server "$scratch/backend.py" 2>&1) +check "a rotated -m is refused" '[[ "$bad" == *"already on the W4A4 route"* ]]' + +port=$(python3 -c 'import socket; s=socket.socket(); s.bind(("127.0.0.1", 0)); print(s.getsockname()[1])') +mkdir -p "$scratch/rec" +env -u GGML_Q4_0_HADAMARD -u GGML_W4A4_TENSORS RECORDS="$scratch/rec" ONEBIT_LEAN_SERVER="$scratch/backend.py" \ + "$bin" serve -m "$scratch/plain.gguf" --long-model "$scratch/h32.gguf" --long-from 100 --port "$port" \ + --llama-server "$scratch/backend.py" >"$scratch/serve.log" 2>&1 & +pid=$! +for _ in $(seq 1 100); do curl -sf "127.0.0.1:$port/health" >/dev/null && break; sleep 0.1; done + +words() { python3 -c 'import sys; print(" ".join(["w"] * int(sys.argv[1])))' "$1"; } +ask() { # : the device that answered, then the X-1bit-Route header + curl -s -D "$scratch/h" "127.0.0.1:$port$1" -H 'Content-Type: application/json' -d "$2" | + python3 -c 'import json,sys; c=json.load(sys.stdin)["choices"][0]; print((c.get("message") or {}).get("content") or c.get("text"))' + tr -d '\r' < "$scratch/h" | sed -n 's/^X-1bit-Route: //Ip' +} +chat() { python3 -c 'import json,sys; print(json.dumps({"messages": [{"role": r, "content": c} for r, c in zip(sys.argv[1::2], sys.argv[2::2])]}))' "$@"; } + +out=$(ask /v1/chat/completions "$(chat user "hello there")") +check "a short chat goes to Vulkan0" '[[ "$out" == "device:Vulkan0"*"short"*"vulkan"* ]]' +out=$(ask /v1/chat/completions "$(chat user "$(words 150)")") +check "a long chat goes to ROCm0" '[[ "$out" == "device:ROCm0"*"long"*"rocm"* ]]' +out=$(ask /v1/completions "{\"prompt\": \"$(words 40)\"}") +check "a short completion goes to Vulkan0" '[[ "$out" == "device:Vulkan0"* ]]' +out=$(ask /v1/completions "{\"prompt\": \"$(words 120)\"}") +check "a long completion goes to ROCm0" '[[ "$out" == "device:ROCm0"* ]]' +out=$(ask /v1/chat/completions "$(chat user "hello there" assistant "hi" user "$(words 300)")") +check "a conversation that grows long stays on Vulkan0" '[[ "$out" == "device:Vulkan0"*"short"* ]]' +check "the ROCm backend has the Hadamard W4A4 environment" \ + 'python3 -c "import json,sys; e=json.load(open(sys.argv[1]))[\"env\"]; sys.exit(not (e.get(\"GGML_Q4_0_HADAMARD\")==\"1\" and e.get(\"GGML_W4A4_TENSORS\")==\"all\"))" "$scratch/rec/ROCm0.json"' +check "the Vulkan backend does not" \ + 'python3 -c "import json,sys; sys.exit(bool(json.load(open(sys.argv[1]))[\"env\"]))" "$scratch/rec/Vulkan0.json"' + +if [ $fail -ne 0 ]; then cat "$scratch/serve.log"; echo FAIL; exit 1; fi +echo PASS