forward-serve: --prefix-pin, the pinned shared prefix on the device path, measured (Refs #1785)

app/forward_serve.cpp

@@ -109,6 +109,10 @@ struct ForwardEngine {
     // slot's attention KV and GDN state.  Only ever touched under mu.
     std::unique_ptr<npu::PrefixCache> prefix;
     std::vector<Q4nxNpuForward::StateSnapshot> prefix_state;
+    // --prefix-pin (engine#1785): keep the state at the point where a request stops agreeing with the
+    // previous one, so single-turn requests under one system prompt resume from it (PrefixCache::pin_point).
+    bool prefix_pin = false;
+    std::vector<int> prev_ids;  // the previous request's prompt, what pin_point() compares against
     // The d3 decision head (#653), loaded when decision_config.json and decision.q4nx sit next to the
     // model: it is what turns this server into a decision endpoint as well as a chat one.
     std::unique_ptr<npu::DecisionHead> head;
@@ -332,11 +336,42 @@ void handle_generate(ForwardEngine& e, const httplib::Request& req, httplib::Res
             std::fprintf(stderr, "[prefix] slot=%d cached=%d of %zu reason=%s\n", stats.prefix.slot,
                          stats.prefix.cached_tokens, ids.size(), stats.prefix.reason.c_str());
         }
+        // (engine#1785) Where this request should leave a pinned snapshot, as a position: the state is saved
+        // one token short of the shared part for the same reason the whole-prompt snapshot below is saved one
+        // token short of the prompt. -1 = no pin for this request. pin_point() is 0 when a slot (the pin
+        // itself, on every request after the one that made it) already covers the shared part, and it is
+        // always less than the prompt's length, so this cut is never the last-token cut below.
+        int pin_cut = -1;
+        size_t pin_tokens = 0;
+        if (e.prefix_pin && e.prefix && e.prefix->enabled()) {
+            pin_tokens = e.prefix->pin_point(e.prev_ids, ids);
+            if (pin_tokens > 0 && int(pin_tokens) - 1 >= resume) pin_cut = int(pin_tokens) - 1;
+        }
         {
             const int chunk_max = std::max(1, e.fwd->max_batch_rows());
             const int n = int(ids.size());
             for (int off = resume; off < n;) {
-                const int take = std::min(chunk_max, n - off);
+                int take = std::min(chunk_max, n - off);
+                if (pin_cut >= 0 && off == pin_cut) {
+                    Q4nxNpuForward::StateSnapshot snap;
+                    if (e.fwd->snapshot_state(off, snap)) {
+                        std::string why;
+                        const std::vector<int> shared(ids.begin(), ids.begin() + long(pin_tokens));
+                        const int slot = e.prefix->pin(shared, snap.bytes(), &why);
+                        if (slot >= 0 && size_t(slot) < e.prefix_state.size()) e.prefix_state[size_t(slot)] = std::move(snap);
+                        std::fprintf(stderr, "[prefix] pin slot=%d tokens=%zu state_at=%d %s\n", slot, pin_tokens, off,
+                                     why.c_str());
+                    }
+                    pin_cut = -1;
+                } else if (pin_cut > off && off + take > pin_cut) {
+                    // Stop this chunk at the pin position; the next pass saves the state there. One more launch.
+                    take = pin_cut - off;
+                    const std::vector<int> head(ids.begin() + off, ids.begin() + off + take);
+                    if (!e.fwd->step_batch(head, off, logits))
+                        throw std::runtime_error(std::string("forward step failed: ") + e.fwd->last_error());
+                    off += take;
+                    continue;
+                }
                 // The state the prefix cache keeps is taken one token short of the prompt, because the answer is
                 // drawn from the last position's logits and a snapshot holds state, not logits -- so the snapshot
                 // is saved just before the last token and that token is run again on a resume.  On a GDN model
@@ -368,6 +403,7 @@ void handle_generate(ForwardEngine& e, const httplib::Request& req, httplib::Res
                 off += take;
             }
         }
+        if (e.prefix_pin) e.prev_ids = ids;
 
         // Prompt processing ends here: every prompt token has been through the model and the first
         // token's distribution is in `logits`.
@@ -520,7 +556,8 @@ void handle_generate(ForwardEngine& e, const httplib::Request& req, httplib::Res
                     {"prompt_tokens_details", {{"cached_tokens", st.prefix.cached_tokens}}},
                     {"flm_prefix_cache", {{"slot", st.prefix.slot},
                                           {"cached_tokens", st.prefix.cached_tokens},
-                                          {"reason", st.prefix.reason}}}};
+                                          {"reason", st.prefix.reason},
+                                          {"pinned", st.prefix.pinned}}}};
     };
 
     // OpenAI's logprobs.content entry: the token's text, its log-probability, and the most likely
@@ -868,6 +905,7 @@ int run_forward_serve(int argc, char** argv) {
     // instead of being ignored and served by whatever target happened to be linked.
     std::string npu_target = ::npu::target::xdna::XdnaTarget::kName;
     npu::PrefixOptions prefix_opts;  // --prefix-slots / --prefix-cache-mb / --prefix-min-gain (#650)
+    bool prefix_pin = false;         // --prefix-pin (engine#1785)
     for (int i = 0; i < argc; i++) {
         const std::string a = argv[i];
         auto next = [&]() -> std::string {
@@ -882,6 +920,7 @@ int run_forward_serve(int argc, char** argv) {
         else if (a == "--prefix-slots") prefix_opts.slots = std::stoi(next());
         else if (a == "--prefix-cache-mb") prefix_opts.cache_mb = std::stoul(next());
         else if (a == "--prefix-min-gain") prefix_opts.min_gain = std::stoi(next());
+        else if (a == "--prefix-pin") prefix_pin = true;
         else if (a == "--npu-streams") streams = std::stoi(next());
         else if (a == "--npu-target") npu_target = next();
         else if (a == "--dump-logits") dump_logits = next();
@@ -979,6 +1018,8 @@ int run_forward_serve(int argc, char** argv) {
         std::fprintf(stderr,
                      "1bit forward-serve: prefix cache on: %d slots, %zu MB, min gain %d tokens\n",
                      prefix_opts.slots, prefix_opts.cache_mb, prefix_opts.min_gain);
+        e.prefix_pin = prefix_pin;
+        if (prefix_pin) std::fprintf(stderr, "1bit forward-serve: prefix pin on: the part a request shares with the previous one keeps a slot\n");
     }
 
     e.fwd = std::make_unique<Q4nxNpuForward>();

app/serve.cpp

@@ -185,6 +185,7 @@ struct Options {
     int prefix_slots    = -1;   // --prefix-slots N (0 = off)
     int prefix_cache_mb = -1;   // --prefix-cache-mb N
     int prefix_min_gain = -1;   // --prefix-min-gain N: tokens a prompt must add to be worth keeping
+    bool prefix_pin = false;    // --prefix-pin: keep the part single-turn requests share (docs/npu.md)
     std::string npu_target;  // --npu-target NAME: which NPU target serves the per-op route (empty = the build's default)
     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)
@@ -3171,6 +3172,7 @@ void usage(FILE* out) {
                  "                  [--npu-streams N]   with --device npu: this serve is one of N sharing the NPU\n"
                  "                  [--prefix-slots N] [--prefix-cache-mb N] [--prefix-min-gain N]   the prompt prefix cache\n"
                  "                  (docs/npu.md; 0 slots turns it off). Unset keeps the forward's defaults: 8 / 4096 / 128\n"
+                 "                  [--prefix-pin]   keep the part a request shares with the previous one (one system prompt, many questions)\n"
                  "                  (default 1, the device exclusive; docs/npu.md)\n"
                  "                  [--load-queue-ms N]   a load refused because another client holds the NPU array\n"
                  "                  waits at most N ms for it to free before the driver is called dead\n"
@@ -3286,6 +3288,7 @@ int run_serve(int argc, char** argv) {
         else if (a == "--prefix-slots")    o.prefix_slots    = std::stoi(next());
         else if (a == "--prefix-cache-mb") o.prefix_cache_mb = std::stoi(next());
         else if (a == "--prefix-min-gain") o.prefix_min_gain = std::stoi(next());
+        else if (a == "--prefix-pin") o.prefix_pin = true;
         else if (a == "--npu-target") o.npu_target = next();
         else if (a == "--prefill-provider") o.prefill_provider = next();
         else if (a == "--laya-model") o.laya_model = next();
@@ -3452,6 +3455,7 @@ int run_serve(int argc, char** argv) {
                 if (o.prefix_slots    >= 0) { fargs.push_back("--prefix-slots");    fargs.push_back(std::to_string(o.prefix_slots)); }
                 if (o.prefix_cache_mb >= 0) { fargs.push_back("--prefix-cache-mb"); fargs.push_back(std::to_string(o.prefix_cache_mb)); }
                 if (o.prefix_min_gain >= 0) { fargs.push_back("--prefix-min-gain"); fargs.push_back(std::to_string(o.prefix_min_gain)); }
+                if (o.prefix_pin) fargs.push_back("--prefix-pin");
                 std::vector<char*> fav;
                 for (auto& a : fargs) fav.push_back(a.data());
                 return run_forward_serve(int(fav.size()), fav.data());

docs/npu.md

@@ -1610,10 +1610,28 @@ The slot table (`app/npu_prefix_cache.h`) has a slot kind for the part such requ
 
 It reads token ids only, so it does not depend on the template or on which message is the system prompt.
 
-**State of this, plainly:** the table and its tests are in; `forward-serve` does not call `pin_point()`/`pin()`
-yet, so no request is faster because of it today. The forward half is one more cut in the prefill loop (stop one
-token short of the pin point, `snapshot_state`, `pin`), and it needs a device run to be claimed.
-
+**`--prefix-pin`** (on `1bit serve --device npu` and on `forward-serve`; off by default) turns it on. The
+forward keeps the previous request's prompt ids, and when `pin_point()` is non-zero it stops the prefill one
+token short of that point, takes a snapshot and pins it, the same cut it already makes before the last token.
+The cost is one more launch on the request that makes the pin. Every decision is logged
+(`[prefix] pin slot=1 tokens=342 state_at=341 pinned`), and `usage.flm_prefix_cache.pinned` says when a hit
+came from the pinned slot.
+
+Measured 2026-11-10 on `bonsai2-t2s2-perrole876` with its `kernels-m16`, `--ctx-size 1024`, one `npu-hold`
+ticket, one server per arm, four requests: the same 14-rule system prompt, then three different questions and
+the second one again. 16 tokens generated, temperature 0.
+
+| request | cache off | default cache | `--prefix-pin` |
+|---|---|---|---|
+| 1 (question A) | cached 0, 58.6 s | cached 0 | cached 0, 59.5 s |
+| 2 (question B) | cached 0, 59.4 s | cached 0 | cached 0, 62.1 s (makes the pin: 342 tokens) |
+| 3 (question C) | cached 0, 61.5 s | cached 0 | **cached 341 of 355, pinned, 23.4 s** |
+| 4 (question B again) | cached 0, 60.0 s | not reached | cached 354 of 355 (the whole-prompt slot), 20.4 s |
+
+The generated text of all four requests is identical between the cache-off arm and the `--prefix-pin` arm.
+The default-cache arm's times are not in the table: the ticket ended during its third request, so only its
+three cache decisions (from the server log) are known. A request's time here is the whole HTTP call, prefill
+plus 16 generated tokens, not TTFT.
 
 ### Sampling and structured output
 

tests/npu_sampling_test.cpp

@@ -500,6 +500,28 @@ void test_prefix_pin() {
         check(c.pin_point(prompt(100, 5000, 40), prompt(100, 6000, 40)) == 0,
               "100 shared tokens do not beat prefix-min-gain 128");
         check(c.pin_point(prompt(900, 5000, 40), prompt(0, 6000, 40)) == 0, "prompts that share nothing have no pin point");
+        {   // only the leading tokens count: the same 900 tokens at the end of both prompts are not a prefix
+            std::vector<int> a1 = prompt(0, 5000, 200), a2 = prompt(0, 6000, 200);
+            const std::vector<int> tail = prompt(900, 0, 0);
+            a1.insert(a1.end(), tail.begin(), tail.end());
+            a2.insert(a2.end(), tail.begin(), tail.end());
+            check(shared_prefix_len(a1, a2) == 0 && c.pin_point(a1, a2) == 0,
+                  "a shared ending under different openings is not a pin point");
+            a2[0] = a1[0];  // one shared leading token, then they part: still nothing worth a slot
+            check(shared_prefix_len(a1, a2) == 1 && c.pin_point(a1, a2) == 0,
+                  "and one shared leading token before that ending is not one either");
+        }
+        // the contract the forward relies on: a pin point is a real common prefix, never longer than either
+        for (int sys : {0, 1, 127, 128, 129, 500, 900}) {
+            for (int turn : {1, 40}) {
+                const std::vector<int> x = prompt(sys, 5000, turn), y = prompt(sys, 6000, turn + 7);
+                const size_t n = c.pin_point(x, y);
+                check(n <= x.size() && n <= y.size() && std::equal(x.begin(), x.begin() + long(n), y.begin()),
+                      "a pin point never reaches past the tokens both prompts start with");
+                check(n == (sys >= 128 ? size_t(sys) : 0u),
+                      "it is the whole shared opening when that is worth a slot, and 0 (no pin) when it is not");
+            }
+        }
         PrefixOptions off = o;
         off.slots = 0;
         check(PrefixCache(off).pin_point(q1, q2) == 0, "an off cache pins nothing");