serve --long-model: short conversations on -m's route, long ones on W4A4

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 $<TARGET_FILE:onebit>)
 add_test(NAME hadamard_route COMMAND ${CMAKE_SOURCE_DIR}/tests/hadamard_route.sh $<TARGET_FILE:onebit>)
+add_test(NAME long_route COMMAND ${CMAKE_SOURCE_DIR}/tests/long_route.sh $<TARGET_FILE:onebit>)
 # The NPU route's <think> 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)

app/serve.cpp

@@ -56,6 +56,9 @@
 //                                                  drafts a block per step (docs/serve.md)
 //   --mmproj <mmproj.gguf> on any llama.cpp route -> images: the vision encoder runs beside the
 //                                                  model, and chat messages may carry image parts
+//   --long-model <rotated.gguf>                 -> 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<std::string>();
+    } 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<std::atomic<int>> inflight(backends.size());
     std::vector<std::atomic<long>> routed(backends.size());
     std::mutex route_mu;
+    // --long-model: conversation -> (backend, prompt tokens of its first request)
+    std::mutex long_mu;
+    std::unordered_map<size_t, std::pair<size_t, long>> long_cache;
+    std::deque<size_t> long_order;
 #ifdef ONEBIT_LAYA
     std::mutex scorer_mu, class_mu;
     std::unordered_map<size_t, onebit::laya::RequestClass> 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<size_t, long> side{0, -1};
+            bool fresh = true;
+            {
+                std::lock_guard<std::mutex> 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<std::mutex> 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>(&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<std::mutex> 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";

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

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 <tokens> <device>`.
+
+```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`,

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("<Q", len(b)) + b
+    kv.append(s("general.architecture") + struct.pack("<I", 8) + s("llama"))
+    if stamp:
+        kv.append(s("onebit.hadamard_q4_0") + struct.pack("<I", 5) + struct.pack("<i", 32))
+    open(path, "wb").write(b"GGUF" + struct.pack("<IQQ", 3, 0, len(kv)) + b"".join(kv))
+gguf(sys.argv[1] + "/h32.gguf", True)
+gguf(sys.argv[1] + "/plain.gguf", False)
+PY
+
+cat > "$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() {  # <path> <json body>: 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