diff --git a/app/forward_serve.cpp b/app/forward_serve.cpp index c70f989e0..effed11be 100644 --- a/app/forward_serve.cpp +++ b/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 prefix; std::vector 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 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 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 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 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(); diff --git a/app/serve.cpp b/app/serve.cpp index 9556af962..164d6623f 100644 --- a/app/serve.cpp +++ b/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 fav; for (auto& a : fargs) fav.push_back(a.data()); return run_forward_serve(int(fav.size()), fav.data()); diff --git a/docs/npu.md b/docs/npu.md index 2f7fdf31c..a5a47c2a0 100644 --- a/docs/npu.md +++ b/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 diff --git a/tests/npu_sampling_test.cpp b/tests/npu_sampling_test.cpp index 94e6efdf8..56851d2b4 100644 --- a/tests/npu_sampling_test.cpp +++ b/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 a1 = prompt(0, 5000, 200), a2 = prompt(0, 6000, 200); + const std::vector 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 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");