Prompt lookup that pays, sampling fixes, shallow-path setup fix (engine 0.1.7)

- Prompt lookup on by default: the suffix drafter's windows are taken only where a learned policy (measured
  acceptance per match length, measured window cost) expects more tokens per ms than the MTP window; lookup windows
  up to --spec + 2. Code edits +6-11%, other text unchanged, verified output identical. --suffix-draft 0 = off.
- Per-request sampling (#19) keeps greedy requests on the in-graph sampler (no extra launch per window), and draws
  row t of a window with Philox(seed, pos0 + t), so a seed replays the same text whatever the drafts were.
- tools: a shallow install (G:\Strata) no longer crashes the MTP packer and the pack tools (IndexError, #20).
- setup: says at step 1 when the PC has less RAM than the smallest model needs; requires engine 0.1.7.

CMakeLists.txt

@@ -424,7 +424,7 @@ endif()
 #     worse opt-out.  `cpu_require_expert_support()` therefore refuses at startup on a CPU without the
 #     features, and the fallback exists for tests, exactly as P2.S3 specifies.
 # Plan v0.3 P6: speculation components that need no GPU.
-add_library(strata_spec STATIC src/spec/suffix_drafter.cpp src/spec/controller.cpp)
+add_library(strata_spec STATIC src/spec/suffix_drafter.cpp src/spec/controller.cpp src/spec/draft_policy.cpp)
 target_include_directories(strata_spec PUBLIC include)
 if(STRATA_BUILD_TESTS)
   add_executable(suffix_drafter_test src/spec/suffix_drafter_test.cpp)
@@ -433,6 +433,9 @@ if(STRATA_BUILD_TESTS)
   add_executable(controller_test src/spec/controller_test.cpp)
   target_link_libraries(controller_test PRIVATE strata_spec)
   add_test(NAME controller_test COMMAND controller_test)
+  add_executable(draft_policy_test src/spec/draft_policy_test.cpp)
+  target_link_libraries(draft_policy_test PRIVATE strata_spec)
+  add_test(NAME draft_policy_test COMMAND draft_policy_test)
 endif()
 
 add_library(strata_kernels_cpu STATIC src/kernels/cpu/expert.cpp src/kernels/cpu/pool.cpp

docs/DETAILS.md

@@ -215,8 +215,15 @@ assistant turn and every 16K prompt tokens). A checkpoint is used only when the
 and pictures. Engine options: `--prompt-cache N` (0 = off), `--prompt-cache-every N`, `--turn-token ID`.
 
 **Current limits (v1):** one request at a time, and one conversation cached at a time (switching between two chats
-re-reads the other one); greedy decoding (temperature is ignored); images only when set up with them (below); no
-video.
+re-reads the other one); images only when set up with them (below); no video. **Temperature / top_p / top_k / min_p /
+seed** are honored per request (OpenAI and Anthropic fields); with the default adaptive expert tier a sampled result
+is not reproducible run to run - for seed-reproducible output add `--adapt-every 100000` (static residency) to the
+engine arguments. The run config's optional `sampling` block sets the defaults for requests that leave the fields out
+(`"sampling": {"temperature": 1.0, "top_p": 0.95, "top_k": 20}`); a request's own fields always win, and with no
+block at all a request without sampling keys decodes greedy. The penalties (`presence_penalty`, `frequency_penalty`,
+`repetition_penalty`, with `penalty_last_n` capping how many recent tokens they count over, default 64 when any
+penalty is set) ride the same path; they count the tokens the request has consumed, so a repetition penalty
+suppresses what the model itself just said, not the prompt alone.
 
 ---
 
@@ -325,7 +332,9 @@ test images.
   as the GPU works on the cached ones (AVX-512 / AVX2 kernels, ggml's for the i-quants).
 - **SSD:** the 28.8 GB n-gram table, read a few rows per token through the OS cache.
 - **Speculation:** the model's own MTP layer drafts up to 3 tokens; one pass over all 48 layers checks them. 2.4-3.2
-  tokens per pass on average.
+  tokens per pass on average. When the reply repeats the context (code edits, quoted text), **prompt lookup** (engine
+  0.1.7) drafts up to 5 tokens from the earlier copy, but only where its measured acceptance and cost say it pays:
+  code edits 6-11% faster, other text unchanged. The drafts are checked like the MTP's, so the output is the same.
 - **Prompts** are processed in 2,048-token chunks with the experts streamed to the GPU over PCIe.
 
 The full story, with measurements, bottlenecks and what comes next: **[docs/paper/Strata-Paper.pdf](paper/Strata-Paper.pdf)**.

include/strata/core/layer.hpp

@@ -256,7 +256,8 @@ struct QsaState {
 uint64_t qsa_state_bytes(const ModelGeometry& g, int64_t max_cells, bool with_rope = true, int64_t ring_cells = 0);
 /// KV streaming: keep `cells` cells of each QSA layer in VRAM and the rest in pinned host memory (0: all in VRAM,
 /// the default). Set before sizing and initializing the session; a context that fits in `cells` is not streamed.
-/// Also puts the MTP drafter's K/V in a ring of its window (`ring_cells` of qsa_state_bytes/init).
+/// Also puts the MTP drafter's K/V in a ring of its window (`ring_cells` of qsa_state_bytes/init; -1 forces a fully
+/// resident state).
 void qsa_set_kv_resident(int64_t cells);
 int64_t qsa_kv_resident();
 /// The fewest resident cells a streamed layer may have: one verify window's selections (8 queries x 2,051 cells

include/strata/core/verify.hpp

@@ -24,6 +24,7 @@
 #include "strata/core/expert_source.hpp"
 #include "strata/core/layer.hpp"
 #include "strata/core/session.hpp"
+#include "strata/kernels/sampler.hpp"
 
 #include <cuda_runtime.h>
 
@@ -58,6 +59,21 @@ public:
     /// One window: `tokens[0..T)` at positions pos0.., the pool served per layer; `out[t]` = argmax after token t.
     /// The PLE rows are gathered here from `ss.ple_prev` and the tokens.  Captures the T-token graph on first use.
     bool run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool, void* user, int32_t* out, std::string& err);
+    /// The sampling the verify window's head applies (temperature / top_p / top_k / seed).  Set per
+    /// request; greedy by default.  The sampling itself runs OUTSIDE the captured graph - its
+    /// parameters would otherwise be baked forever - so this can change between requests freely.
+    void set_sampling(const strata::kernels::SamplerParams& sp) {
+        sampling_ = sp;   // row t of a window at pos0 draws Philox(seed, pos0 + t): see run()
+    }
+
+    /// The penalty-history row for `sampling_.penalty_last_n`: ONE row of `history_len` int32 slots, most
+    /// recent token LAST, unused front slots -1 (the kernel reads only the tail window).  Null disables the
+    /// penalties entirely - the neutral run's sampling call is byte-for-byte what it was.  The engine
+    /// re-uploads the request's tail before every window; the pointer must stay alive across the request.
+    void set_history(const int32_t* history, int history_len) {
+        hist_d_ = history;
+        hist_len_ = history_len;
+    }
 
     /// Keep the first `n_keep` (1..T) tokens of the last window; advances `ss.ple_prev` by them.
     bool commit(int n_keep, std::string& err);
@@ -83,6 +99,14 @@ public:
 
 private:
     bool capture(int T, std::string& err);
+    strata::kernels::SamplerParams sampling_ = [] {
+        strata::kernels::SamplerParams s;
+        s.greedy = true;
+        s.temperature = 0.0f;
+        return s;
+    }();   ///< greedy by default; per-request via set_sampling
+    const int32_t* hist_d_ = nullptr;   ///< penalty-history row (set_history); null = no penalties apply
+    int hist_len_ = 0;
     bool capture_commit(std::string& err);
     bool record_window(int T, cudaStream_t cs, std::string& err);
 

include/strata/kernels/sampler.hpp

@@ -1,12 +1,10 @@
 // include/strata/kernels/sampler.hpp - the sampler chain, host-callable (P2.S2).
 //
-//     penalties  ->  top_k  ->  top_p  ->  temperature  ->  pick
+//     penalties -> top_k -> min_p -> top_p -> temperature -> penalties -> pick
 //
-// The ORDER is specified in docs/sampling.md, transcribed from llama.cpp's own chain.  Two facts there are
-// easy to get backwards: temperature comes AFTER the truncation filters, and penalties come AFTER temperature.
-//
-// PENALTIES ARE NOT IMPLEMENTED YET: they need the token history, which is not in this signature, so a caller
-// passing a non-zero `penalty_last_n` gets a loud failure rather than a silently unpenalised sample.
+// The ORDER is specified in docs/sampling.md, transcribed from llama.cpp's own chain.  Three facts there are
+// easy to get backwards: temperature comes AFTER the truncation filters, penalties come AFTER temperature as
+// well as before the selection, and min_p cuts BEFORE top_p (both on the pre-temperature logits).
 #pragma once
 
 #include <cstdint>
@@ -16,6 +14,7 @@ namespace strata::kernels {
 struct SamplerParams {
     int top_k = 20;              // 0 disables the filter
     float top_p = 0.95f;         // 1.0 disables the filter
+    float min_p = 0.0f;          // 0 disables; keeps tokens with p >= min_p * p_max (a prefix of the top_k list)
     float temperature = 1.0f;    // <= 0 means greedy
     int min_keep = 1;            // top_p keeps at least this many
     int penalty_last_n = 0;      // 0 disables; the window over `history` to count occurrences in

include/strata/spec/draft_policy.hpp

@@ -0,0 +1,52 @@
+// include/strata/spec/draft_policy.hpp - per verify round: the MTP's window, or a lookup (suffix) window?
+//
+// The suffix drafter (prompt lookup) proposes the tokens that followed an earlier repeat of the context. Taken
+// whenever it proposes more than the MTP, it lost 2-8% on ordinary text: a long lookup window costs much more to
+// verify than the MTP's usual 3-4 tokens, and it is only worth that when enough of it is accepted. llama.cpp's
+// lookup decoding answers the same problem with confidence thresholds on its n-gram statistics; here the policy
+// learns both sides online and compares expected committed tokens per millisecond:
+//
+//   MTP     E = the mean tokens a window of that size has committed (EMA), at the measured cost of that size
+//   lookup  E(k) = 1 + q + q^2 + ... + q^k for k <= the proposal, q = the acceptance rate of lookup drafts whose
+//           match was about as long (4 buckets of match length, decayed counts), at the measured cost of k + 1
+//
+// and takes the lookup window only when its best E/cost beats the MTP's by `margin`. Costs are the measured round
+// times per window size (EMA; sizes not seen yet are scaled from seen ones by a prior shape), so the policy adapts
+// to the machine and the context length. It only chooses which drafts to verify: the output is unchanged.
+#pragma once
+
+#include <array>
+
+namespace strata::spec {
+
+class DraftPolicy {
+public:
+    static constexpr int kMaxT = 8;
+    static constexpr int kBuckets = 4;
+
+    explicit DraftPolicy(int max_t, double margin = 0.03);
+
+    struct Pick {
+        bool lookup = false;
+        int t = 1;                      // window size (1 + drafts)
+    };
+    /// `t_mtp`: the MTP's window; `lookup_k`: the lookup proposal's length (0 = none); `match`: its match length.
+    Pick choose(int t_mtp, int lookup_k, int match) const;
+    /// After the round: the window it used, the drafts accepted, and the round's time (verify + commit + draft).
+    void observe(bool lookup, int t, int accepted, int match, double round_ms);
+
+    double lookup_rate(int match) const;   // current q for a match length
+    double cost_ms(int t) const;           // measured or scaled round time of a window of t tokens
+
+private:
+    static int bucket(int match);
+    double mtp_tokens(int t) const;
+
+    int max_t_;
+    double margin_;
+    std::array<double, kMaxT + 1> cost_{}, cost_n_{};      // round ms by window size
+    std::array<double, kMaxT + 1> mtp_tok_{}, mtp_n_{};    // tokens committed by MTP windows of that size
+    std::array<double, kBuckets> ok_{}, bad_{};            // lookup drafts accepted / windows cut short, decayed
+};
+
+}  // namespace strata::spec

serve/server.py

@@ -10,6 +10,7 @@ data, http(s) URLs or local file paths) go through `strata-vision` (the model's
 embeddings (`GENI`).  JPEG/PNG/BMP/GIF go straight in; WebP, TIFF, AVIF, ... (agents like omp send WebP) are
 converted to PNG first with Pillow.
 Requests whose prompt plus max tokens exceed the engine's context are REJECTED with 400, never truncated.
+An unset (or 0, or -1) max tokens means "unlimited": whatever the prompt leaves of the context.
 
 The engine boundary is `Engine.generate(prompt_ids, max_new, sampling, cancel) -> iterator of token ids`.
 `StrataEngine` keeps one `strata --serve` process resident (weights, expert arena and VRAM tier load once) and
@@ -42,6 +43,7 @@ from serve.frontend import (ChatTemplate, Event, OutputParser, anthropic_to_mess
 
 IM_END = "<|im_end|>"
 IMAGE_PAD = "<|image_pad|>"
+CTX_SLACK = 8               # `strata --serve` rejects prompt + max_new + 8 > context: keep the same margin here
 
 
 # ------------------------------------------------------------------------------------------------ engines
@@ -71,7 +73,12 @@ class MockEngine:
 
 class StrataEngine:
     """The resident engine: `strata --serve` reads `GEN <max_new> <ids>` lines and streams `T <id>` lines, then
-    `DONE ...`.  Requests are serialized by the service's FIFO, so one pipe is enough."""
+    `DONE ...`.  Requests are serialized by the service's FIFO, so one pipe is enough.
+
+    Per-request sampling rides the same line as engine-side keys between max_new and the ids
+    (`temperature=F top_p=F top_k=N seed=N`, the engine's own spelling).  An absent temperature keeps the
+    engine's default, which is greedy; `temperature=0` means the same thing, so it is not forwarded.
+    """
 
     def __init__(self, exe: str, args: list[str], cwd: str | None = None, log: str | None = None,
                  env: dict | None = None):
@@ -103,11 +110,51 @@ class StrataEngine:
         self.last = {"generated": int(f[1]), "prompt_tokens": int(f[2]), "prompt_ms": float(f[3]),
                      "decode_ms": float(f[4]), "finish": f[5]}
 
+    @staticmethod
+    def sampling_keys(sampling: dict) -> str:
+        keys = ""
+        t = sampling.get("temperature")
+        if isinstance(t, (int, float)) and float(t) > 0.0:
+            keys += f" temperature={float(t)!r}"
+        tp = sampling.get("top_p")
+        if isinstance(tp, (int, float)) and float(tp) < 1.0:
+            keys += f" top_p={float(tp)!r}"
+        tk = sampling.get("top_k")
+        if isinstance(tk, int) and 1 <= tk <= 64:
+            keys += f" top_k={tk}"        # the engine's sampled path takes 1..64; outside it keeps its 20
+        mp = sampling.get("min_p")
+        if isinstance(mp, (int, float)) and 0.0 < float(mp) <= 1.0:
+            keys += f" min_p={float(mp)!r}"
+        rp = sampling.get("repetition_penalty")
+        rp_on = isinstance(rp, (int, float)) and float(rp) != 1.0
+        pf = sampling.get("frequency_penalty")
+        pf_on = isinstance(pf, (int, float)) and float(pf) != 0.0
+        pp = sampling.get("presence_penalty")
+        pp_on = isinstance(pp, (int, float)) and float(pp) != 0.0
+        if rp_on:
+            keys += f" penalty_repeat={float(rp)!r}"
+        if pf_on:
+            keys += f" penalty_freq={float(pf)!r}"
+        if pp_on:
+            keys += f" penalty_present={float(pp)!r}"
+        if rp_on or pf_on or pp_on:
+            # a penalty without a window counts over nothing: the engine's default is the last 64 tokens
+            pln = sampling.get("penalty_last_n")
+            if isinstance(pln, int) and not isinstance(pln, bool) and pln > 0:
+                keys += f" penalty_last_n={pln}"
+            else:
+                keys += " penalty_last_n=64"
+        seed = sampling.get("seed")
+        if isinstance(seed, int) and seed > 0:
+            keys += f" seed={seed}"
+        return keys
+
     def generate(self, ids, max_new, sampling, cancel, embeddings=None):
         """Yields token ids, and None as a heartbeat every 10 s while the engine is quiet (reading a long prompt):
         the HTTP layer turns it into an SSE comment, which keeps clients' watchdogs calm and notices a client that
         has gone.  A consumer that stops early (or `cancel`) makes the engine STOP, so it does not run to max_new."""
-        head = f"GENI {int(max_new)} {embeddings}" if embeddings else f"GEN {int(max_new)}"
+        head = f"GENI {int(max_new)} {embeddings}" if embeddings else \
+            f"GEN {int(max_new)}{self.sampling_keys(sampling or {}) if not embeddings else ''}"
         self.proc.stdin.write(f"{head} {','.join(str(int(t)) for t in ids)}\n")
         self.proc.stdin.flush()
         done = False
@@ -311,8 +358,9 @@ class Detokenizer:
 
 class Service:
     def __init__(self, engine: Engine, tokenizer, template: ChatTemplate, model_name: str = "qwen3.8-flash-next",
-                 vision: Vision | None = None):
+                 vision: Vision | None = None, sampling_defaults: dict | None = None):
         self.engine, self.tok, self.template, self.model, self.vision = engine, tokenizer, template, model_name, vision
+        self.sampling_defaults = dict(sampling_defaults or {})   # the run config's `sampling` block
         self.fifo = threading.Lock()
         self.embeddings = threading.local()           # the current request's image embeddings file (GENI)
         self.api_key = ""                              # when set, /v1/* needs it (Bearer or x-api-key)
@@ -321,7 +369,9 @@ class Service:
         self.stop_ids = set(tokenizer.encode(IM_END, parse_special=True) +
                             tokenizer.encode("<|endoftext|>", parse_special=True))
 
-    def prepare(self, messages, tools, kwargs, max_new):
+    def prepare(self, messages, tools, kwargs, max_new=None):
+        """-> (ids, thinking, max_new). An unset or non-positive max_new (some clients send -1) means "unlimited":
+        the rest of the context."""
         prompt = self.template.render(messages, tools=tools, **kwargs)
         ids = self.tok.encode(prompt, parse_special=True)
         self.embeddings.path = None
@@ -352,10 +402,16 @@ class Service:
                 for path, _ in encoded:
                     f.write(path.read_bytes())
             self.embeddings.path = combined
-        if len(ids) + max_new > self.engine.max_context:
+        room = self.engine.max_context - CTX_SLACK - len(ids)
+        if max_new is None or max_new <= 0:
+            if room < 1:
+                raise ValueError(f"prompt ({len(ids)} tokens) leaves no room to answer in the context "
+                                 f"({self.engine.max_context}); requests are never truncated")
+            max_new = room
+        elif max_new > room:
             raise ValueError(f"prompt ({len(ids)} tokens) + max tokens ({max_new}) exceeds the context "
                              f"({self.engine.max_context}); requests are never truncated")
-        return ids, kwargs.get("enable_thinking", True) is not False
+        return ids, kwargs.get("enable_thinking", True) is not False, max_new
 
     def _note(self, n, evs):
         with self.status_lock:
@@ -392,6 +448,9 @@ class Service:
 
     def run(self, ids, thinking, tools, max_new, sampling, cancel) -> Iterator[tuple[str, object]]:
         """Yields ("event", Event) as text arrives, then ("done", {"finish": .., "completion_tokens": ..})."""
+        if self.sampling_defaults:     # config defaults; the request's own fields win (explicit 0 stays greedy)
+            req_values = {k: v for k, v in (sampling or {}).items() if v is not None}
+            sampling = {**self.sampling_defaults, **req_values}
         parser = OutputParser(thinking=thinking, tools=tools, stream_tools=True)
         detok, n, finish = Detokenizer(self.tok), 0, "length"
         raw_ids = []                                    # every generated id (STRATA_DEBUG: dump raw model text)
@@ -703,10 +762,8 @@ def make_handler(svc: Service):
 
         def _openai(self, req):
             messages, tools, kw = openai_to_messages(req)
-            max_new = int(req.get("max_completion_tokens") or req.get("max_tokens") or 1024)
-            if max_new <= 0:                             # a non-positive budget (some clients send -1) means "unset"
-                max_new = 1024
-            ids, thinking = svc.prepare(messages, tools, kw, max_new)
+            max_new = int(req.get("max_completion_tokens") or req.get("max_tokens") or 0)   # 0/-1: the rest of the context
+            ids, thinking, max_new = svc.prepare(messages, tools, kw, max_new)
             _debug_req("openai", req, messages, tools, max_new, thinking, len(ids))
             cancel = threading.Event()
             chunks = openai_chunks(svc, req, ids, thinking, tools, max_new, cancel)
@@ -727,10 +784,8 @@ def make_handler(svc: Service):
 
         def _anthropic(self, req):
             messages, tools, kw = anthropic_to_messages(req)
-            max_new = int(req.get("max_tokens") or 1024)
-            if max_new <= 0:                             # a non-positive budget (some clients send -1) means "unset"
-                max_new = 1024
-            ids, thinking = svc.prepare(messages, tools, kw, max_new)
+            max_new = int(req.get("max_tokens") or 0)                  # 0/-1: the rest of the context
+            ids, thinking, max_new = svc.prepare(messages, tools, kw, max_new)
             _debug_req("anthropic", req, messages, tools, max_new, thinking, len(ids))
             cancel = threading.Event()
             events = anthropic_events(svc, req, ids, thinking, tools, max_new, cancel)
@@ -765,6 +820,59 @@ def serve(svc: Service, host="127.0.0.1", port=8095) -> ThreadingHTTPServer:
     return httpd
 
 
+def sampling_defaults_from_config(cfg: dict) -> dict:
+    """The run config's optional `sampling` block: defaults for the sampling fields a request leaves out, so
+    a plain client gets configured sampling instead of greedy.  Supported: temperature, top_p, top_k, min_p,
+    presence_penalty, repetition_penalty, frequency_penalty, penalty_last_n, seed.  The request's own fields
+    always win - an explicit temperature=0 still means greedy, a field set to null falls back to the default.
+    A bad value refuses to start the server (a typo'd config should not quietly change sampling); unknown keys
+    are named at startup and ignored."""
+    out = {}
+    for key, value in (cfg.get("sampling") or {}).items():
+        if value is None:
+            continue
+        number = isinstance(value, (int, float)) and not isinstance(value, bool)
+        if key == "temperature":
+            if not number or value < 0:
+                raise SystemExit(f"[strata] config sampling.temperature={value!r}: expected a number >= 0 (0 = greedy)")
+            out[key] = float(value)
+        elif key == "top_p":
+            if not number or not 0 < value <= 1:
+                raise SystemExit(f"[strata] config sampling.top_p={value!r}: expected 0 < top_p <= 1")
+            out[key] = float(value)
+        elif key == "min_p":
+            if not number or not 0 <= value <= 1:
+                raise SystemExit(f"[strata] config sampling.min_p={value!r}: expected 0 <= min_p <= 1")
+            out[key] = float(value)
+        elif key == "top_k":
+            if not number or value != int(value) or not 1 <= value <= 64:
+                raise SystemExit(f"[strata] config sampling.top_k={value!r}: the sampled path takes an integer 1..64")
+            out[key] = int(value)
+        elif key == "presence_penalty":
+            if not number or value < 0:
+                raise SystemExit(f"[strata] config sampling.presence_penalty={value!r}: expected a number >= 0")
+            out[key] = float(value)
+        elif key == "frequency_penalty":
+            if not number or value < 0:
+                raise SystemExit(f"[strata] config sampling.frequency_penalty={value!r}: expected a number >= 0")
+            out[key] = float(value)
+        elif key == "repetition_penalty":
+            if not number or value <= 0:
+                raise SystemExit(f"[strata] config sampling.repetition_penalty={value!r}: expected a number > 0 (1 = off)")
+            out[key] = float(value)
+        elif key == "penalty_last_n":
+            if not number or value != int(value) or value < 0:
+                raise SystemExit(f"[strata] config sampling.penalty_last_n={value!r}: expected a non-negative integer")
+            out[key] = int(value)
+        elif key == "seed":
+            if not number or value != int(value) or value <= 0:
+                raise SystemExit(f"[strata] config sampling.seed={value!r}: expected a positive integer")
+            out[key] = int(value)
+        else:
+            print(f"[strata] config sampling.{key}={value!r}: unknown key, ignored", flush=True)
+    return out
+
+
 def main() -> int:
     ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
     ap.add_argument("--engine", choices=["mock", "strata"], default="mock")
@@ -805,6 +913,10 @@ def main() -> int:
             ap.error("--engine strata needs --config")
         vision = None
         env = child_env(cfg)
+        sampling_defaults = sampling_defaults_from_config(cfg)
+        if sampling_defaults:
+            pretty = ", ".join(f"{k}={v}" for k, v in sampling_defaults.items())
+            print(f"[strata] sampling defaults from the config: {pretty}", flush=True)
         if cfg.get("vision"):
             print("loading the vision encoder ...", flush=True)
             vision = Vision(cfg["vision"], log=open(cfg["log"], "a", encoding="utf-8") if cfg.get("log") else None,
@@ -812,11 +924,12 @@ def main() -> int:
         print("loading the model (the first start takes a minute or two) ...", flush=True)
         engine = StrataEngine(cfg["exe"], cfg["args"], cwd=cfg.get("cwd"), log=cfg.get("log"), env=env)
     else:
-        engine, vision = MockEngine(tok, a.script), None
+        engine, vision, sampling_defaults = MockEngine(tok, a.script), None, {}
     # the model's own chat template (exported with its tokenizer), else the original model's
     tpl = tpath / "chat_template.jinja"
     svc = Service(engine, tok, ChatTemplate(tpl if tpl.exists() else ROOT / "serve/chat_template.jinja"),
-                  model_name=cfg.get("model_name", "qwen3.8-flash-next"), vision=vision)
+                  model_name=cfg.get("model_name", "qwen3.8-flash-next"), vision=vision,
+                  sampling_defaults=sampling_defaults)
     svc.api_key = a.api_key or cfg.get("api_key", "")
     httpd = serve(svc, host=a.host, port=a.port)
     print(f"ready: http://{a.host}:{a.port}/v1  (OpenAI: /v1/chat/completions, Anthropic: /v1/messages, "

serve/test_server.py

@@ -0,0 +1,130 @@
+"""serve/test_server.py - the max tokens budget over both APIs, against the mock engine (no GPU, no pack).
+
+    python -m unittest serve.test_server -v
+"""
+from __future__ import annotations
+
+import json
+import os
+import sys
+import unittest
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from serve.frontend import ChatTemplate  # noqa: E402
+from serve.server import CTX_SLACK, ByteTokenizer, MockEngine, Service, serve  # noqa: E402
+
+ROOT = Path(__file__).resolve().parents[1]
+CTX = 4096
+ANSWER = "x" * 2000                              # longer than the old 1024 fallback: one token per byte
+
+
+class RecordingEngine(MockEngine):
+    def generate(self, ids, max_new, sampling, cancel, embeddings=None):
+        self.last_max_new = max_new
+        yield from super().generate(ids, max_new, sampling, cancel, embeddings)
+
+
+class MaxTokens(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        tok = ByteTokenizer()
+        cls.engine = RecordingEngine(tok, "</think>\n\n" + ANSWER, max_context=CTX)
+        cls.svc = Service(cls.engine, tok, ChatTemplate(ROOT / "serve/chat_template.jinja"))
+        cls.httpd = serve(cls.svc, port=0)
+        cls.base = f"http://127.0.0.1:{cls.httpd.server_address[1]}"
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.httpd.shutdown()
+        cls.httpd.server_close()
+
+    def post(self, path, body):
+        req = urllib.request.Request(self.base + path, data=json.dumps(body).encode(),
+                                     headers={"Content-Type": "application/json"})
+        try:
+            with urllib.request.urlopen(req, timeout=30) as r:
+                return r.status, json.loads(r.read())
+        except urllib.error.HTTPError as e:
+            with e:
+                return e.code, json.loads(e.read())
+
+    def call(self, api, text="hi", **budget):
+        """-> (status, body, prompt tokens, completion tokens); `budget` is merged into the request as given."""
+        msgs = [{"role": "user", "content": text}]
+        if api == "openai":
+            s, b = self.post("/v1/chat/completions", {"model": "m", "messages": msgs, **budget})
+            u = b.get("usage", {})
+            return s, b, u.get("prompt_tokens"), u.get("completion_tokens")
+        s, b = self.post("/v1/messages", {"model": "m", "messages": msgs, **budget})
+        u = b.get("usage", {})
+        return s, b, u.get("input_tokens"), u.get("output_tokens")
+
+    def test_unset_budget_is_the_rest_of_the_context(self):
+        cases = {"openai": [{"max_tokens": -1}, {"max_tokens": 0}, {}, {"max_tokens": None},
+                            {"max_completion_tokens": -1}, {"max_completion_tokens": None, "max_tokens": None}],
+                 "anthropic": [{"max_tokens": -1}, {"max_tokens": 0}, {}, {"max_tokens": None}]}
+        for api, budgets in cases.items():
+            for budget in budgets:
+                with self.subTest(api=api, budget=budget):
+                    s, b, pt, ct = self.call(api, **budget)
+                    self.assertEqual(s, 200, b)
+                    self.assertEqual(self.engine.last_max_new, CTX - CTX_SLACK - pt)
+                    self.assertGreater(ct, 1024)          # the whole answer, not cut at the old 1024 fallback
+
+    def test_explicit_budget_is_honoured(self):
+        for api, budget in [("openai", {"max_tokens": 50}), ("openai", {"max_completion_tokens": 50}),
+                            ("openai", {"max_completion_tokens": 50, "max_tokens": 9}),
+                            ("anthropic", {"max_tokens": 50}), ("openai", {"max_tokens": 1500}),
+                            ("anthropic", {"max_tokens": 1500})]:
+            with self.subTest(api=api, budget=budget):
+                want = budget.get("max_completion_tokens") or budget["max_tokens"]
+                s, b, _, ct = self.call(api, **budget)
+                self.assertEqual(s, 200, b)
+                self.assertEqual(self.engine.last_max_new, want)
+                self.assertEqual(ct, want)
+
+    def test_explicit_budget_over_the_context_is_rejected(self):
+        for api in ("openai", "anthropic"):
+            with self.subTest(api=api):
+                s, b, _, _ = self.call(api, max_tokens=CTX)
+                self.assertEqual(s, 400)
+                self.assertIn("exceeds the context", b["error"]["message"])
+
+    def test_unset_budget_with_a_near_full_prompt(self):
+        _, _, pt0, _ = self.call("openai", max_tokens=1)
+        overhead = pt0 - len("hi")                  # the template's tokens around the user text
+        for api in ("openai", "anthropic"):
+            _, _, pa, _ = self.call(api, max_tokens=1)
+            over = pa - pt0                          # the Anthropic template may differ slightly
+            with self.subTest(api=api, room=5):     # a few tokens left: the budget is exactly those
+                text = "y" * (CTX - CTX_SLACK - overhead - over - 5)
+                s, b, pt, ct = self.call(api, text=text, max_tokens=-1)
+                self.assertEqual(s, 200, b)
+                self.assertEqual(self.engine.last_max_new, 5)
+                self.assertEqual(ct, 5)
+            with self.subTest(api=api, room=0):     # nothing left: rejected, not truncated
+                text = "y" * (CTX - CTX_SLACK - overhead - over)
+                s, b, _, _ = self.call(api, text=text)
+                self.assertEqual(s, 400, b)
+                self.assertIn("no room to answer", b["error"]["message"])
+
+    def test_debug_log_shows_the_resolved_budget(self):
+        import contextlib
+        import io
+        os.environ["STRATA_DEBUG"] = "1"
+        try:
+            for api in ("openai", "anthropic"):
+                with self.subTest(api=api):
+                    out = io.StringIO()
+                    with contextlib.redirect_stdout(out):
+                        _, _, pt, _ = self.call(api, max_tokens=-1)
+                    self.assertIn(f"max_new={CTX - CTX_SLACK - pt} ", out.getvalue())
+        finally:
+            del os.environ["STRATA_DEBUG"]
+
+
+if __name__ == "__main__":
+    unittest.main()

setup.py

@@ -54,7 +54,7 @@ PREBUILT_ASSET = "strata-windows-x64.zip" if WIN else "strata-linux-x64.zip"
 # the CUDA libraries the ready-made engine loads (the same CUDA 13.0 it is built with), from NVIDIA's pip packages
 CUDA_WHEELS = ["nvidia-cublas==13.0.2.14", "nvidia-cuda-runtime==13.0.96"]
 MIN_DRIVER = 580                       # CUDA 13.0
-MIN_ENGINE = (0, 1, 5)                 # KV streaming (--kv-resident), v0.1.5
+MIN_ENGINE = (0, 1, 7)                 # per-request sampling on the GEN line (the server sends it), v0.1.7
 PY_PACKAGES = ["numpy", "jinja2", "regex", "pyyaml", "tqdm", "requests", "cmake", "ninja", "pillow"]
 
 MODELS = {
@@ -696,7 +696,14 @@ def main() -> int:
         warn("less than 12 GB of VRAM: Strata will run, but most experts stay on the CPU and it will be slow")
     ram = ram_gb()
     cpu, avx2, avx512 = cpu_info()
-    ok(f"RAM: {ram:.0f} GB")
+    need = min(d["ram_gb"] for d in MODELS.values())
+    if ram < need - 4 and not a.check:
+        # every model keeps ALL its experts in RAM (34+ GB); VRAM only holds a copy of the most-used ones, so a
+        # bigger GPU does not lower this
+        fail(f"RAM: {ram:.0f} GB - the smallest model (Q2_0 / IQ2_XS) needs about {need} GB",
+             "Strata keeps all of the model's experts in RAM (34-50 GB, whatever the GPU) and the GPU holds a copy "
+             "of the most-used ones: it needs 48 GB of RAM or more")
+    ok(f"RAM: {ram:.0f} GB" if ram >= need - 4 else f"RAM: {ram:.0f} GB (less than the {need} GB the smallest model needs)")
     ok(f"CPU: {cpu} ({'AVX-512' if avx512 else 'AVX2' if avx2 else 'no AVX2'})")
     if not avx2:
         fail("this CPU has no AVX2; Strata needs at least AVX2")

src/core/layer.cpp

@@ -492,7 +492,7 @@ KvPlan kv_plan(const QsaShapes& s, int64_t max_cells, int64_t ring_cells) {
     p.pages = (max_cells + s.page_size - 1) / s.page_size;
     p.slots = p.pages;
     p.pooled_rows = max_cells / s.idx_block + 2;
-    if (g_kv_resident <= 0) return p;
+    if (g_kv_resident <= 0 || ring_cells < 0) return p;   // ring_cells < 0: always fully resident
     // A/B only: STRATA_KV_RING_OFF keeps the drafter fully resident, STRATA_KV_MAIN_OFF the main layers
     static const bool ring_off = std::getenv("STRATA_KV_RING_OFF") != nullptr;
     static const bool main_off = std::getenv("STRATA_KV_MAIN_OFF") != nullptr;
@@ -599,8 +599,11 @@ uint64_t qsa_state_init(const ModelGeometry& g, int64_t max_cells, void* base, Q
         uint8_t* h = nullptr;
         uint8_t* d = nullptr;
         if (cudaHostAlloc((void**) &h, bytes, cudaHostAllocMapped | cudaHostAllocPortable) != cudaSuccess ||
-            cudaHostGetDevicePointer((void**) &d, h, 0) != cudaSuccess)
+            cudaHostGetDevicePointer((void**) &d, h, 0) != cudaSuccess) {
+            if (p.mode == 1) std::fprintf(stderr, "strata: KV streaming: cannot pin %.2f GiB of RAM for a layer's KV copy "
+                                 "(lower the context, or run without --kv-resident)\n", (double) bytes / 1073741824.0);
             return 0;
+        }
         g_kv_host_bytes += bytes;
         Cursor hc{d};
         if (g_kv_int8) {

src/core/mtp.cpp

@@ -158,11 +158,22 @@ bool MtpDrafter::load(const std::string& rt_dir, const ModelGeometry& g, Session
     const strata::kernels::QsaShapes s = shapes_of(g);
     const int64_t max_cells = ss.qsa_states[0].max_cells;
     // KV streaming: the drafter only reads its last `window` cells, so with streaming on its K/V is a ring of the
-    // window (plus the cells a round writes ahead of its queries) over a host copy, refilled on a resume
-    const int64_t ring = (window > 0 && window < max_cells) ? window + 4 * (int64_t) max_t + 64 : 0;
-    const uint64_t sb = qsa_state_bytes(g, max_cells, false, ring);
+    // window (plus the cells a round writes ahead of its queries) over a host copy, refilled on a resume. The host copy
+    // is pinned after the expert arena has pinned what it could: if it does not fit, the K/V stays whole in VRAM.
+    int64_t ring = (window > 0 && window < max_cells) ? window + 4 * (int64_t) max_t + 64 : 0;
+    uint64_t sb = qsa_state_bytes(g, max_cells, false, ring);
     if (cudaMalloc(&state_arena_, sb) != cudaSuccess) { err = "mtp: the K/V state does not fit"; return false; }
-    if (qsa_state_init(g, max_cells, state_arena_, st_, &ss.qsa_states[0], ring) == 0) { err = "mtp: state init failed"; return false; }
+    if (qsa_state_init(g, max_cells, state_arena_, st_, &ss.qsa_states[0], ring) == 0) {
+        if (st_.kv_mode == 0) { err = "mtp: state init failed"; return false; }
+        std::fprintf(stderr, "strata mtp: no pinned RAM left for the draft layer's K/V copy; keeping it in VRAM\n");
+        cudaGetLastError();
+        cudaFree(state_arena_);
+        st_ = QsaState{};
+        ring = -1;   // fully resident
+        sb = qsa_state_bytes(g, max_cells, false, ring);
+        if (cudaMalloc(&state_arena_, sb) != cudaSuccess) { err = "mtp: the K/V state does not fit"; return false; }
+        if (qsa_state_init(g, max_cells, state_arena_, st_, &ss.qsa_states[0], ring) == 0) { err = "mtp: state init failed"; return false; }
+    }
     qsa_state_zero(st_, g, nullptr);
     cudaDeviceSynchronize();
     vram_ += sb;

src/core/verify.cpp

@@ -111,6 +111,8 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
     hits_ = hits;
     head_ = head;
     max_t_ = max_t;
+    sampling_.greedy = true;      // a fresh verifier samples greedily until set_sampling says otherwise
+    sampling_.temperature = 0.0f;
     if (max_t < 2 || max_t > strata::kernels::kVerifyMaxT || max_t > strata::kernels::cpu::MAXT) {
         err = "verify: the window must hold 2.." + std::to_string(strata::kernels::kVerifyMaxT) + " tokens";
         return false;
@@ -594,6 +596,9 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
                 return false;
             }
         }
+        // Greedy, the default, is recorded here as before (no extra launch or sync per window). A request that
+        // samples or penalizes is sampled again host-side after the replay (run()) with its own parameters and a
+        // fresh draw counter: a captured sampler would bake them in and replay the same draws forever.
         SamplerParams sp;
         sp.greedy = true;
         sp.temperature = 0.0f;
@@ -793,6 +798,20 @@ bool Verifier::run(int T, const int32_t* tokens, int64_t pos0, PoolMultiFn pool,
     const cudaError_t se = cudaStreamSynchronize(cs_);
     if (se != cudaSuccess) { err = std::string("verify: ") + cudaGetErrorString(se); return false; }
     cudaStreamSynchronize(copy_);   // no host function of this window may raise flag B in the next one
+    // ---- a sampled or penalized request: the head's sampling again, host-side so its parameters are this call's
+    // own (a captured kernel would replay the same draws forever).  Row t's draw is Philox(seed, pos0 + t): tied to
+    // the POSITION it samples, not to how the text was cut into windows, so a seed replays the same text whatever
+    // the drafts were. Exact: a rejected row's draw is discarded, and no kept decision depends on a reused draw.
+    const bool sampled = !sampling_.greedy && sampling_.temperature > 0.0f;
+    if (sampled || hist_d_ != nullptr) {
+        SamplerParams sp = sampling_;
+        sp.counter = (uint64_t) pos0;
+        sample_tokens(head_logits_, T, (int) n_vocab_, hist_d_, hist_len_, sp, m_out_, cs_);
+        if (cudaStreamSynchronize(cs_) != cudaSuccess) {   // m_out_ is the mapped h_out_: synced, it is readable
+            err = "verify: the head sampling failed";
+            return false;
+        }
+    }
     for (int t = 0; t < T; ++t) out[t] = ((volatile int32_t*) h_out_)[t];
     VDBG("window done\n");
     ++windows;

src/kernels/cuda/sampler.cu

@@ -1,6 +1,6 @@
 // src/kernels/cuda/sampler.cu - P2.S2: the sampler chain, in the order docs/sampling.md settles.
 //
-//     penalties  ->  top_k  ->  top_p  ->  temperature  ->  pick
+//     penalties -> top_k -> min_p -> top_p -> temperature -> penalties -> pick
 //
 // THE ORDER IS THE WHOLE CONTENT OF THIS FILE.  `docs/sampling.md` transcribes it from llama.cpp's own chain
 // (`common/sampling.cpp` L357/360/375/381/399) and the two facts that are easy to get backwards are that
@@ -9,10 +9,9 @@
 // valid token, so only a comparison at the distribution level can tell them apart; the parity test does that
 // explicitly by running the wrong order and requiring it to differ.
 //
-// Everything happens on the logits of ONE token in one thread.  The vocab is 248,320, which is far too large
-// to sort per token on a naive path, so `top_k` uses `k` passes of a maximum scan - 20 x 248,320 = 5.0M
-// comparisons per token.  Phase 3's note ("sort-free sampler: top-k = 20 makes this easy") is the optimisation;
-// this is the correctness version.
+// Both kernels put ONE BLOCK per token over the vocabulary: `sampler_greedy_kernel` is the plain argmax,
+// `sampler_kernel` runs the sampled chain as `top_k` block-argmax rounds followed by the top_p / temperature /
+// draw chain (its header says why the selection must be parallel and why the tie rule keeps the semantics).
 #include "strata/kernels/sampler.hpp"
 
 #include <cuda_runtime.h>
@@ -102,12 +101,32 @@ __global__ void sampler_greedy_kernel(const float* __restrict__ logits, int n_vo
         hrow += history_len - hlen;          // the window is the TAIL
     }
 
+    // PENALTY MEMBERSHIP AS A BITMAP.  The history touches at most `hlen` tokens of a quarter-million
+    // vocabulary, but the naive `history_count` per candidate per argmax round costs O(k x n_vocab x hlen)
+    // integer compares (~318 M per token at k=20, hlen=64 - measured 45 -> 31 tok/s on a real workload).
+    // A shared bitmap gives an O(1) membership test, and only the (at most hlen) hits pay the count scan;
+    // the counts - and therefore every sampled value - are exactly what the per-candidate scan produced.
+    extern __shared__ unsigned int penal_bits[];
+    const int bits_words = (int) ((n_vocab + 31) / 32);
+    const bool use_bits = hrow != nullptr && bits_words > 0;
+    if (use_bits) {
+        for (int w = threadIdx.x; w < bits_words; w += blockDim.x) penal_bits[w] = 0u;
+        __syncthreads();
+        for (int i = threadIdx.x; i < hlen; i += blockDim.x)
+            if (hrow[i] >= 0) atomicOr(&penal_bits[hrow[i] >> 5], 1u << (hrow[i] & 31));
+        __syncthreads();
+    }
+    auto hit_count = [&](int v) -> int {
+        if (!use_bits || !(penal_bits[v >> 5] & (1u << (v & 31)))) return 0;
+        return history_count(hrow, hlen, v);
+    };
+
     // `n_vocab` is the "no candidate" index: it loses every comparison to a real one, so a thread with no
     // elements contributes nothing rather than contributing a bogus zero.
     float bv = __int_as_float(0xff800000);   // -inf
     int best = n_vocab;
     for (int v = threadIdx.x; v < n_vocab; v += blockDim.x) {
-        const float s = apply_penalties(l[v], hrow ? history_count(hrow, hlen, v) : 0, p);
+        const float s = apply_penalties(l[v], hit_count(v), p);
         if (s > bv) { bv = s; best = v; }
     }
     for (int off = 16; off > 0; off >>= 1) {
@@ -134,47 +153,32 @@ __global__ void sampler_greedy_kernel(const float* __restrict__ logits, int n_vo
     }
 }
 
+/// **THE SAMPLED PATH, ONE BLOCK PER TOKEN.**  The kernel below replaced a version that ran the whole chain
+/// in ONE THREAD per token (`<<<ceil(T/64), 64>>>`, so a 4-token window fielded four threads): `top_k` alone
+/// was `k` sequential scans of the vocabulary with an inner sweep over the already-taken list - 20 x 248,320
+/// iterations of dependent work on one SM - and a verify window measured **1.6 s in the sampler**, which made
+/// every temperature-bearing request ~30x slower than a greedy one.  The selection is `k` argmax rounds, and
+/// an argmax over the vocabulary parallelises exactly like `sampler_greedy_kernel` (block over the vocab), so
+/// the rounds run back to back inside a block-per-token launch: the per-token cost falls to
+/// `k x n_vocab / 1024` plus `k` block reductions.
+///
+/// THE SEMANTICS ARE THE SERIAL ONES, EXACTLY.  Each round's argmax resolves ties to the LOWEST index (the
+/// serial scan's strict `>` keeps the first maximum it meets), so the kept sequence - both its set and its
+/// order - is unchanged; `top_p`'s cut reads that order in double arithmetic as before; temperature and the
+/// Philox draw apply after the cut.  `sampler_parity` pins all of it against the host reference.
 __global__ void sampler_kernel(const float* __restrict__ logits, int n_vocab, int n_tokens,
                                const int* __restrict__ history, int history_len, const SamplerParams p,
                                int* __restrict__ out) {
-    const int t = blockIdx.x * blockDim.x + threadIdx.x;
+    const int t = blockIdx.x;
     if (t >= n_tokens) return;
-
-    // Work on a private copy: the chain mutates the distribution and the caller's logits are shared with the
-    // next stage.  A 248,320-float allocation per thread is far too much for shared memory, so this uses a
-    // two-pass approach instead - first find the keep set, then normalise over it - and never materialises a
-    // second full-size buffer.
     const float* l = logits + (size_t) t * n_vocab;
 
-    // Temperature is needed by BOTH paths below, so it is computed here; the chain still APPLIES it after
+    // Temperature is needed by BOTH stages below, so it is computed here; the chain still APPLIES it after
     // the truncation filters - the survivors are chosen on the raw logits and only then scaled.
     const float inv_t = p.temperature > 0.0f ? 1.0f / p.temperature : 0.0f;
 
-    // ---- penalties: NOT IMPLEMENTED, and the ORDER they must take is transcribed here so the next person
-    // does not have to guess it.  `llama_sampler_penalties_apply` (src/llama-sampler.cpp) runs AFTER
-    // temperature in the chain (L381 against L375), and for every token seen in the last `penalty_last_n`:
-    //
-    //     if (logit <= 0) logit *= penalty_repeat;   else logit /= penalty_repeat;
-    //     logit -= float(count) * penalty_freq + float(count > 0) * penalty_present;
-    //
-    // TWO details there are counter-intuitive.  The repeat penalty MULTIPLIES for non-positive logits and
-    // DIVIDES for positive ones - the source notes the paper only ever divided, "but that would cause tokens
-    // with negative logits to become more likely, which is obviously wrong".  And the presence penalty rides
-    // on `float(count > 0)`, a BOOLEAN cast to float, not the count - so it applies once however many times
-    // the token appeared.  Both are the kind of thing a plausible implementation gets wrong while still
-    // producing valid tokens.
-    //
-    // It needs the token history, which is not in this signature, so a caller passing a non-zero
-    // `penalty_last_n` gets a loud failure rather than a silently unpenalised sample.
-    // ---- top_k: k passes of a maximum scan, skipping the already-taken
-    // ---- top_p: cumulative mass over the survivors in descending order
-    // The two are done together in one selection pass so the survivors and their mass are consistent.
-    // GREEDY IS THE GLOBAL ARGMAX, WHATEVER THE FILTERS SAY - so it is computed directly and the filter chain
-    // is skipped entirely.  No filter can remove the maximum: top_k keeps the k largest, top_p keeps a prefix
-    // of the descending order that always contains the largest (min_keep >= 1), and temperature is monotonic so
-    // it cannot reorder anything.  That also makes greedy INDEPENDENT OF THE ORDER, which is why the ordering
-    // test cannot use it - see sampler_parity.cpp.
-    // The penalty window is the last `penalty_last_n` entries of this row's history.
+    // The penalty window is the last `penalty_last_n` entries of this row's history (disabled at this
+    // launch: `sample_tokens` refuses a non-zero `penalty_last_n` without a history buffer).
     const int* hrow = history ? history + (size_t) t * history_len : nullptr;
     int hlen = 0;
     if (hrow) {
@@ -183,109 +187,117 @@ __global__ void sampler_kernel(const float* __restrict__ logits, int n_vocab, in
         hrow += history_len - hlen;          // the window is the TAIL
     }
 
-    if (p.greedy || p.temperature <= 0.0f) {
-        // Greedy is the global argmax ONLY when no penalty is active: a penalty can lower a token enough to
-        // change the maximum, so the argmax has to be taken over the penalised logits, not the raw ones.
-        //
-        // *** TEMPERATURE MUST NOT BE APPLIED HERE.  THIS WAS A REAL BUG. ***
-        // This branch used to read `apply_penalties(l[v] * inv_t, ...)`.  `inv_t` is 0.0f whenever
-        // temperature <= 0 (see above), so at temperature 0 EVERY logit became 0.0f and the argmax returned
-        // index 0 - the sampler emitted token 0 forever, whatever the model predicted.  OpenAI clients send
-        // `temperature: 0` for greedy decoding, so this was reachable from any ordinary client.
-        // It is also pointless even when inv_t is non-zero: scaling by a positive constant is monotonic and
-        // cannot reorder the argmax, which is exactly what the note above this block already says.
-        int best = 0;
-        float bv = apply_penalties(l[0], hrow ? history_count(hrow, hlen, 0) : 0, p);
-        for (int v = 1; v < n_vocab; ++v) {
-            const float s = apply_penalties(l[v], hrow ? history_count(hrow, hlen, v) : 0, p);
-            if (s > bv) { bv = s; best = v; }
-        }
-        out[t] = best;
-        return;
+    // the membership bitmap, as in `sampler_greedy_kernel` - see the cost note there
+    extern __shared__ unsigned int penal_bits[];
+    const int bits_words = (int) ((n_vocab + 31) / 32);
+    const bool use_bits = hrow != nullptr && bits_words > 0;
+    if (use_bits) {
+        for (int w = threadIdx.x; w < bits_words; w += blockDim.x) penal_bits[w] = 0u;
+        __syncthreads();
+        for (int i = threadIdx.x; i < hlen; i += blockDim.x)
+            if (hrow[i] >= 0) atomicOr(&penal_bits[hrow[i] >> 5], 1u << (hrow[i] & 31));
+        __syncthreads();
     }
+    auto hit_count = [&](int v) -> int {
+        if (!use_bits || !(penal_bits[v >> 5] & (1u << (v & 31)))) return 0;
+        return history_count(hrow, hlen, v);
+    };
 
     const int KMAX = 64;
-    int keep_ids[KMAX];
-    float keep_logit[KMAX];
-    int k = p.top_k > 0 ? (p.top_k < KMAX ? p.top_k : KMAX) : KMAX;
-    if (p.top_k <= 0) k = 0;                      // 0 means "no top-k filter" -> handled by the mask below
-
+    int k = p.top_k > 0 ? (p.top_k < KMAX ? p.top_k : KMAX) : 0;
     if (k <= 0) {
-        std::printf("sampler: the sampled path needs top_k in 1..%d (got %d); greedy needs no filters\n", KMAX, p.top_k);
+        if (threadIdx.x == 0)
+            std::printf("sampler: the sampled path needs top_k in 1..%d (got %d); greedy needs no filters\n",
+                        KMAX, p.top_k);
         return;   // leave out[t] unwritten rather than returning an uninitialised token
     }
-    int n_keep = 0;
-    {
-        for (int i = 0; i < k; ++i) {
-            int best = -1;
-            float bv = 0.0f;
-            for (int v = 0; v < n_vocab; ++v) {
-                bool taken = false;
-                for (int j = 0; j < i; ++j) if (keep_ids[j] == v) { taken = true; break; }
-                if (taken) continue;
-                if (best < 0 || l[v] > bv) { best = v; bv = l[v]; }
+
+    // ---- top_k: k rounds of a block argmax over the not-yet-taken.  `sel_*` holds the kept ids and their
+    // raw logits in selection order: descending by value, ties to the lower index, which is the order the
+    // top_p cut below is defined over.
+    __shared__ int sel_ids[KMAX];
+    __shared__ float sel_logit[KMAX];
+    __shared__ float sv[32];
+    __shared__ int si[32];
+    for (int i = 0; i < k; ++i) {
+        // `n_vocab` is the "no candidate" index: it loses every comparison to a real one (same convention as
+        // the greedy kernel, whose tie rule this reduction shares).
+        float bv = __int_as_float(0xff800000);   // -inf
+        int best = n_vocab;
+        for (int v = threadIdx.x; v < n_vocab; v += blockDim.x) {
+            bool taken = false;
+            for (int j = 0; j < i; ++j) if (sel_ids[j] == v) { taken = true; break; }
+            if (taken) continue;
+            const float s = apply_penalties(l[v], hit_count(v), p);
+            if (s > bv) { bv = s; best = v; }
+        }
+        for (int off = 16; off > 0; off >>= 1) {
+            const float ov = __shfl_down_sync(0xFFFFFFFFu, bv, off);
+            const int oi = __shfl_down_sync(0xFFFFFFFFu, best, off);
+            if (ov > bv || (ov == bv && oi < best)) { bv = ov; best = oi; }
+        }
+        const int warp = (int) (threadIdx.x >> 5), lane = (int) (threadIdx.x & 31);
+        if (lane == 0) { sv[warp] = bv; si[warp] = best; }
+        __syncthreads();
+        if (warp == 0) {
+            const int nw = (int) ((blockDim.x + 31) >> 5);
+            float wv = lane < nw ? sv[lane] : __int_as_float(0xff800000);
+            int wi = lane < nw ? si[lane] : n_vocab;
+            for (int off = 16; off > 0; off >>= 1) {
+                const float ov = __shfl_down_sync(0xFFFFFFFFu, wv, off);
+                const int oi = __shfl_down_sync(0xFFFFFFFFu, wi, off);
+                if (ov > wv || (ov == wv && oi < wi)) { wv = ov; wi = oi; }
             }
-            keep_ids[i] = best;
-            keep_logit[i] = bv;
-        }
-        n_keep = k;
-
-        // ---- top_p over the survivors, in descending order (which the selection above already produced)
-        if (p.top_p < 1.0f) {
-            // softmax over the keep set for the cumulative mass
-            float mx = keep_logit[0];
-            for (int i = 1; i < n_keep; ++i) mx = fmaxf(mx, keep_logit[i]);
-            double sum = 0.0;
-            for (int i = 0; i < n_keep; ++i) sum += exp((double) keep_logit[i] - (double) mx);
-            double cum = 0.0;
-            int cut = n_keep;
-            for (int i = 0; i < n_keep; ++i) {
-                cum += exp((double) keep_logit[i] - (double) mx) / sum;
-                if (cum >= (double) p.top_p) { cut = i + 1; break; }
-            }
-            if (cut < p.min_keep) cut = p.min_keep < n_keep ? p.min_keep : n_keep;
-            n_keep = cut;
+            if (lane == 0) { sel_ids[i] = (wi < n_vocab) ? wi : 0; sel_logit[i] = wv; }
         }
+        __syncthreads();
     }
 
+    // ---- min_p: keep the descending prefix whose probability is at least `min_p` of the top token's.  The
+    // kept list is in selection order (descending), so the survivors are a PREFIX and the cut composes with
+    // top_p's below.  In logit space the threshold is `sel_logit[0] + logf(min_p)` - equivalent to
+    // `p >= min_p * p_max` without the overflow an exp of raw logits risks.  0 disables, and the head itself
+    // always survives (`expf(0) == 1 >= min_p` for min_p in 0..1), so the count never reaches zero.
+    int n_minp = k;
+    if (p.min_p > 0.0f) {
+        const float thresh = sel_logit[0] + logf(p.min_p);
+        for (int i = 0; i < k; ++i)
+            if (sel_logit[i] < thresh) { n_minp = i; break; }
+    }
 
-    if (p.greedy || p.temperature <= 0.0f) {
-        // DEAD CODE AS WRITTEN: the identical condition at the top of this function already returned, so this
-        // block cannot be reached.  It is kept and corrected rather than deleted because if that early return
-        // is ever narrowed, this is the path that would run - and it carried the same `* inv_t` bug (every
-        // logit 0.0f at temperature 0, so `best` would always be `keep_ids[0]`).  Temperature is monotonic and
-        // cannot reorder a maximum, so it must not appear here at all.
-        // argmax over the SURVIVORS, ties by smallest id (the same convention as the router and top_k)
-        int best = keep_ids[0];
-        float bv = keep_logit[0];
-        for (int i = 1; i < n_keep; ++i) {
-            const float v = keep_logit[i];
-            if (v > bv) { bv = v; best = keep_ids[i]; }
+    // ---- top_p over the survivors, in descending order (which the selection produced), then temperature and
+    // one Philox draw.  Every thread computes the same chain redundantly over `sel_*` - the arithmetic is the
+    // serial kernel's, instruction for instruction - so they agree on `pick` and thread 0 writes it.
+    int n_keep = n_minp;
+    float mx = sel_logit[0];
+    for (int i = 1; i < n_minp; ++i) mx = fmaxf(mx, sel_logit[i]);
+    if (p.top_p < 1.0f) {
+        double sum = 0.0;
+        for (int i = 0; i < n_minp; ++i) sum += exp((double) sel_logit[i] - (double) mx);
+        double cum = 0.0;
+        int cut = n_minp;
+        for (int i = 0; i < n_minp; ++i) {
+            cum += exp((double) sel_logit[i] - (double) mx) / sum;
+            if (cum >= (double) p.top_p) { cut = i + 1; break; }
         }
-        out[t] = best;
-        return;
+        if (cut < p.min_keep) cut = p.min_keep < n_minp ? p.min_keep : n_minp;
+        n_keep = cut;
     }
-
-    // ---- pick: temperature, then penalties, then softmax over the survivors, then one uniform draw.
-    // The ORDER is the chain's: the survivors were chosen on the RAW logits, temperature scales them, and the
-    // penalty is applied to the scaled value - which is why the `logit <= 0` branch inside the penalty sees a
-    // temperature-scaled logit and not the model's own.
-    for (int i = 0; i < n_keep; ++i) {
-        keep_logit[i] = apply_penalties(keep_logit[i] * inv_t,
-                                        hrow ? history_count(hrow, hlen, keep_ids[i]) : 0, p);
-    }
-    float mx = keep_logit[0];
-    for (int i = 1; i < n_keep; ++i) mx = fmaxf(mx, keep_logit[i]);
+    auto scaled = [&](int i) {
+        return apply_penalties(sel_logit[i] * inv_t, hit_count(sel_ids[i]), p);
+    };
+    float smx = scaled(0);
+    for (int i = 1; i < n_keep; ++i) smx = fmaxf(smx, scaled(i));
     double sum = 0.0;
-    for (int i = 0; i < n_keep; ++i) sum += exp((double) keep_logit[i] - (double) mx);
+    for (int i = 0; i < n_keep; ++i) sum += exp((double) scaled(i) - (double) smx);
     const float u = philox_uniform(p.seed, p.counter + (uint64_t) t);
     double cum = 0.0;
-    int pick = keep_ids[n_keep - 1];
+    int pick = sel_ids[n_keep - 1];
     for (int i = 0; i < n_keep; ++i) {
-        cum += exp((double) keep_logit[i] - (double) mx) / sum;
-        if ((double) u < cum) { pick = keep_ids[i]; break; }
+        cum += exp((double) scaled(i) - (double) smx) / sum;
+        if ((double) u < cum) { pick = sel_ids[i]; break; }
     }
-    out[t] = pick;
+    if (threadIdx.x == 0) out[t] = pick;
 }
 
 }  // namespace
@@ -298,16 +310,19 @@ void sample_tokens(const float* logits, int n_tokens, int n_vocab, const int* hi
                      p.penalty_last_n, (const void*) history, history_len);
         std::exit(1);
     }
-    const int threads = 64;
-    const unsigned grid = (unsigned) ((n_tokens + threads - 1) / threads);
+    const unsigned shmem = (history != nullptr && history_len > 0 && p.penalty_last_n > 0)
+                               ? (unsigned) ((n_vocab + 31) / 32) * sizeof(unsigned)   // the penalty bitmap
+                               : 0;
     if (p.greedy || p.temperature <= 0.0f) {
         // One block per token, 1,024 threads over the vocabulary.  See `sampler_greedy_kernel`.
         const int gthreads = 1024;
-        sampler_greedy_kernel<<<(unsigned) n_tokens, gthreads, 0, (cudaStream_t) stream>>>(
+        sampler_greedy_kernel<<<(unsigned) n_tokens, gthreads, shmem, (cudaStream_t) stream>>>(
             logits, n_vocab, history, history_len, p, p.penalty_last_n, p.penalty_last_n, out);
     } else {
-        sampler_kernel<<<grid, threads, 0, (cudaStream_t) stream>>>(logits, n_vocab, n_tokens, history,
-                                                                   history_len, p, out);
+        // The same block-per-token shape: the selection's k argmax rounds reduce inside the block.  See
+        // `sampler_kernel`'s header for what the old one-thread-per-token launch cost.
+        sampler_kernel<<<(unsigned) n_tokens, 1024, shmem, (cudaStream_t) stream>>>(
+            logits, n_vocab, n_tokens, history, history_len, p, out);
     }
     const cudaError_t e = cudaGetLastError();
     if (e != cudaSuccess) {

src/kernels/sampler_parity.cpp

@@ -12,6 +12,7 @@
 
 #include <cuda_runtime.h>
 
+#include <algorithm>
 #include <cmath>
 #include <cstdio>
 #include <cstdlib>
@@ -126,6 +127,139 @@ int run(const char* name, const std::vector<float>& logits, int n_tokens, const
     return bad;
 }
 
+// The Philox draw, host side - a transcription of the kernel's `philox_uniform` so the SAMPLED pick (not
+// just the greedy argmax) can be pinned against a reference.  `__umulhi(a, b)` is the high half of a 32x32
+// multiply, spelled `(uint32_t)(((uint64_t) a * b) >> 32)` here.
+struct PhiloxRound {
+    uint32_t& c0; uint32_t& c1; uint32_t& c2; uint32_t& c3;
+    void step(uint32_t k0, uint32_t k1) const {
+        const uint32_t hi0 = (uint32_t) (((uint64_t) 0x9E3779B9u * c0) >> 32);
+        const uint32_t hi1 = (uint32_t) (((uint64_t) 0xBB67AE85u * c2) >> 32);
+        const uint32_t lo0 = 0x9E3779B9u * c0;
+        const uint32_t lo1 = 0xBB67AE85u * c2;
+        const uint32_t n0 = hi1 ^ c1 ^ k0;
+        const uint32_t n1 = lo1;
+        const uint32_t n2 = hi0 ^ c3 ^ k1;
+        const uint32_t n3 = lo0;
+        c0 = n0; c1 = n1; c2 = n2; c3 = n3;
+    }
+};
+
+float host_philox_uniform(uint64_t seed, uint64_t counter) {
+    uint32_t c0 = (uint32_t) counter, c1 = (uint32_t) (counter >> 32);
+    uint32_t c2 = (uint32_t) seed, c3 = (uint32_t) (seed >> 32);
+    PhiloxRound r{c0, c1, c2, c3};
+    for (int i = 0; i < 10; ++i) r.step((uint32_t) i, 0u);
+    return (float) (c0 >> 8) * (1.0f / 16777216.0f);
+}
+
+// The full SAMPLED chain, host side - the kernel's `sampler_kernel` in serial form: penalties on the raw
+// logits during the top_k selection (ties to the lowest index), the min_p prefix cut, top_p's cut in double,
+// then penalties AGAIN on the temperature-scaled survivors (docs/sampling.md's after-temperature pass), and
+// one Philox draw at (seed, counter + row).
+int sampled_reference(const std::vector<float>& l, const std::vector<int>& hist,
+                      const strata::kernels::SamplerParams& p, int row) {
+    auto penal = [&](float logit, int count) {
+        if (count <= 0) return logit;
+        if (logit <= 0.0f) logit *= p.penalty_repeat; else logit /= p.penalty_repeat;
+        logit -= (float) count * p.penalty_freq + (count > 0 ? 1.0f : 0.0f) * p.penalty_present;
+        return logit;
+    };
+    auto count = [&](int v) { int c = 0; for (int h : hist) if (h == v) ++c; return c; };
+    const int nv = (int) l.size();
+    const int KMAX = 64;
+    const int k = p.top_k > 0 ? (p.top_k < KMAX ? p.top_k : KMAX) : 0;
+    std::vector<int> sel_ids;
+    std::vector<float> sel_logit;
+    std::vector<char> taken((size_t) nv, 0);
+    for (int i = 0; i < k; ++i) {
+        int best = -1; float bv = 0;
+        for (int v = 0; v < nv; ++v) {
+            if (taken[(size_t) v]) continue;
+            const float s = penal(l[(size_t) v], count(v));
+            if (best < 0 || s > bv) { best = v; bv = s; }
+        }
+        taken[(size_t) best] = 1;
+        sel_ids.push_back(best); sel_logit.push_back(bv);
+    }
+    int n_minp = (int) sel_ids.size();
+    if (p.min_p > 0.0f) {
+        const float thresh = sel_logit[0] + std::log(p.min_p);
+        for (int i = 0; i < (int) sel_ids.size(); ++i)
+            if (sel_logit[(size_t) i] < thresh) { n_minp = i; break; }
+    }
+    int n_keep = n_minp;
+    if (p.top_p < 1.0f) {
+        double sum = 0.0;
+        for (int i = 0; i < n_minp; ++i) sum += std::exp((double) sel_logit[(size_t) i] - (double) sel_logit[0]);
+        double cum = 0.0;
+        int cut = n_minp;
+        for (int i = 0; i < n_minp; ++i) {
+            cum += std::exp((double) sel_logit[(size_t) i] - (double) sel_logit[0]) / sum;
+            if (cum >= (double) p.top_p) { cut = i + 1; break; }
+        }
+        if (cut < p.min_keep) cut = p.min_keep < n_minp ? p.min_keep : n_minp;
+        n_keep = cut;
+    }
+    const float inv_t = p.temperature > 0.0f ? 1.0f / p.temperature : 0.0f;
+    auto scaled = [&](int i) { return penal(sel_logit[(size_t) i] * inv_t, count(sel_ids[(size_t) i])); };
+    float smx = scaled(0);
+    for (int i = 1; i < n_keep; ++i) smx = std::fmax(smx, scaled(i));
+    double sum = 0.0;
+    for (int i = 0; i < n_keep; ++i) sum += std::exp((double) scaled(i) - (double) smx);
+    const float u = host_philox_uniform(p.seed, p.counter + (uint64_t) row);
+    double cum = 0.0;
+    int pick = sel_ids[(size_t) (n_keep - 1)];
+    for (int i = 0; i < n_keep; ++i) {
+        cum += std::exp((double) scaled(i) - (double) smx) / sum;
+        if ((double) u < cum) { pick = sel_ids[(size_t) i]; break; }
+    }
+    return pick;
+}
+
+// Survivors after the min_p + top_p cuts, in the sampled chain - the quantity an order or a threshold
+// actually changes, used to assert a fixture can SEE the feature before asserting the kernel matches.
+int sampled_cut(const std::vector<float>& l, const std::vector<int>& hist, const strata::kernels::SamplerParams& p) {
+    auto penal = [&](float logit, int count) {
+        if (count <= 0) return logit;
+        if (logit <= 0.0f) logit *= p.penalty_repeat; else logit /= p.penalty_repeat;
+        logit -= (float) count * p.penalty_freq + (count > 0 ? 1.0f : 0.0f) * p.penalty_present;
+        return logit;
+    };
+    auto count = [&](int v) { int c = 0; for (int h : hist) if (h == v) ++c; return c; };
+    const int nv = (int) l.size();
+    const int KMAX = 64;
+    const int k = p.top_k > 0 ? (p.top_k < KMAX ? p.top_k : KMAX) : 0;
+    std::vector<float> sel;
+    std::vector<char> taken((size_t) nv, 0);
+    for (int i = 0; i < k; ++i) {
+        int best = -1; float bv = 0;
+        for (int v = 0; v < nv; ++v) {
+            if (taken[(size_t) v]) continue;
+            const float s = penal(l[(size_t) v], count(v));
+            if (best < 0 || s > bv) { best = v; bv = s; }
+        }
+        taken[(size_t) best] = 1; sel.push_back(bv);
+    }
+    int n_minp = (int) sel.size();
+    if (p.min_p > 0.0f) {
+        const float thresh = sel[0] + std::log(p.min_p);
+        for (int i = 0; i < (int) sel.size(); ++i)
+            if (sel[(size_t) i] < thresh) { n_minp = i; break; }
+    }
+    if (p.top_p >= 1.0f) return n_minp;
+    double sum = 0.0;
+    for (int i = 0; i < n_minp; ++i) sum += std::exp((double) sel[(size_t) i] - (double) sel[0]);
+    double cum = 0.0;
+    int cut = n_minp;
+    for (int i = 0; i < n_minp; ++i) {
+        cum += std::exp((double) sel[(size_t) i] - (double) sel[0]) / sum;
+        if (cum >= (double) p.top_p) { cut = i + 1; break; }
+    }
+    if (cut < p.min_keep) cut = p.min_keep < n_minp ? p.min_keep : n_minp;
+    return cut;
+}
+
 }  // namespace
 
 int main(int argc, char** argv) {
@@ -326,6 +460,116 @@ int main(int argc, char** argv) {
         bad += run("T=1 greedy=true  is the argmax", l, NT3, p2, want);
     }
 
+    // ---- fixture 6: PENALTIES IN THE SAMPLED CHAIN.  Fixture 4 pins the greedy (argmax) path; the sampled
+    // chain applies penalties a SECOND time - after the temperature scaling, on the survivors - so it gets its
+    // own reference (the full chain with the host Philox) and its own observability check: with the penalties
+    // on, the history row's favourite must LOSE a pick it would win penalty-free.
+    {
+        const int NV2 = 8, NT2 = 2;
+        strata::kernels::SamplerParams p;
+        p.top_k = 5; p.top_p = 0.9f; p.temperature = 0.8f; p.seed = 9; p.counter = 0;
+        p.penalty_last_n = 4; p.penalty_repeat = 3.0f; p.penalty_freq = 0.2f; p.penalty_present = 0.6f;
+
+        std::vector<float> l((size_t) NV2 * NT2, -8.0f);
+        std::vector<int> hist((size_t) NT2 * 4, -1);
+        for (int t = 0; t < NT2; ++t) {
+            float* row = l.data() + (size_t) t * NV2;
+            row[0] = 9.0f; row[1] = 4.5f; row[2] = 4.4f; row[3] = 4.3f;   // token 0 leads clean (9 vs 4.5)
+            hist[(size_t) t * 4 + 0] = 0;                                  // and falls to 2.2/2.0 penalised
+            hist[(size_t) t * 4 + 1] = t == 1 ? 0 : -1;                    // (repeat 3, freq, presence)
+        }
+        std::vector<int> want((size_t) NT2), clean((size_t) NT2);
+        strata::kernels::SamplerParams clean_p = p;
+        clean_p.penalty_last_n = 0; clean_p.penalty_repeat = 1.0f;
+        clean_p.penalty_freq = 0.0f; clean_p.penalty_present = 0.0f;
+        for (int t = 0; t < NT2; ++t) {
+            const std::vector<float> row(l.begin() + (size_t) t * NV2, l.begin() + (size_t) (t + 1) * NV2);
+            const std::vector<int> h(hist.begin() + (size_t) t * 4, hist.begin() + (size_t) (t + 1) * 4);
+            want[(size_t) t] = sampled_reference(row, h, p, t);
+            clean[(size_t) t] = sampled_reference(row, h, clean_p, t);
+        }
+        const bool visible = want[0] != clean[0] || want[1] != clean[1];
+        std::printf("  %-34s %s (penalised picks %d/%d, clean %d/%d)\n",
+                    "sampled penalties are observable", visible ? "yes" : "*** NO ***", want[0], want[1],
+                    clean[0], clean[1]);
+        if (!visible) ++bad;
+        else bad += run("sampled chain: penalties + top_k/p", l, NT2, p, want, hist, 4);
+    }
+
+    // ---- fixture 7: MIN_P.  The cut is a PREFIX of the descending top_k list (logit >= max + log(min_p)),
+    // so the fixture asserts the survivor count moves with the threshold (the observability half) and that
+    // the kernel's pick equals the reference's through the full sampled chain (the correctness half).
+    {
+        const int NV3 = 8, NT3 = 2;
+        std::vector<float> l((size_t) NV3 * NT3, -8.0f);
+        for (int t = 0; t < NT3; ++t) {
+            float* row = l.data() + (size_t) t * NV3;
+            row[0] = 4.0f; row[1] = 3.5f; row[2] = 3.2f; row[3] = 3.1f;   // gaps keep the cut off the
+            row[4] = 2.0f;                                                // logf/rounding knife edge
+        }
+        strata::kernels::SamplerParams base;
+        base.top_k = 6; base.top_p = 1.0f; base.temperature = 0.9f; base.seed = 77;
+
+        int c0 = 0, c05 = 0, c09 = 0;
+        for (int t = 0; t < NT3; ++t) {
+            const std::vector<float> row(l.begin() + (size_t) t * NV3, l.begin() + (size_t) (t + 1) * NV3);
+            strata::kernels::SamplerParams q = base; q.min_p = 0.0f;
+            c0 += sampled_cut(row, {}, q);
+            q.min_p = 0.5f; c05 += sampled_cut(row, {}, q);
+            q.min_p = 0.9f; c09 += sampled_cut(row, {}, q);
+        }
+        const bool visible = c0 > c05 && c05 > c09 && c09 >= NT3;
+        std::printf("  %-34s %s (survivors: min_p 0 -> %d, 0.5 -> %d, 0.9 -> %d)\n",
+                    "min_p cut is observable", visible ? "yes" : "*** NO ***", c0, c05, c09);
+        if (!visible) ++bad;
+
+        for (float mp : {0.0f, 0.5f, 0.9f}) {
+            strata::kernels::SamplerParams q = base; q.min_p = mp;
+            std::vector<int> want((size_t) NT3);
+            for (int t = 0; t < NT3; ++t) {
+                const std::vector<float> row(l.begin() + (size_t) t * NV3, l.begin() + (size_t) (t + 1) * NV3);
+                want[(size_t) t] = sampled_reference(row, {}, q, t);
+            }
+            char name[64];
+            std::snprintf(name, sizeof name, "sampled chain: min_p=%.1f", (double) mp);
+            bad += run(name, l, NT3, q, want);
+        }
+    }
+
+    // ---- fixture 8: THE PENALTY WINDOW IS THE TAIL.  With an 8-entry history and penalty_last_n = 4, only
+    // the LAST four entries count: a token punished in the old half must come back to full strength, and one
+    // punished in the tail half stays down.  The reference counts the same tail; the observability check runs
+    // the reference once more WITHOUT the clamp (counting all 8) and requires the picks to differ.
+    {
+        const int NV4 = 8;
+        strata::kernels::SamplerParams p;
+        p.top_k = 0; p.top_p = 1.0f; p.temperature = 1.0f; p.greedy = true;
+        p.penalty_last_n = 4; p.penalty_repeat = 3.0f; p.penalty_freq = 0.3f; p.penalty_present = 0.5f;
+
+        std::vector<float> l((size_t) NV4, -8.0f);
+        l[0] = 6.0f; l[3] = 6.5f;                       // token 0 leads clean; token 3 is the tail offender
+        std::vector<int> hist = {0, 0, 0, 0, 3, 3, 3, 3};   // token 0 old (out), token 3 in the tail
+
+        auto pick_clamped = [&](bool clamp) {
+            int best = 0; float bv = 0; bool first = true;
+            for (int v = 0; v < NV4; ++v) {
+                int c = 0;
+                for (int i = 0; i < (clamp ? 4 : 8); ++i) if (hist[(size_t) (8 - (clamp ? 4 : 8) + i)] == v) ++c;
+                float logit = l[(size_t) v];
+                if (c > 0) { logit = logit <= 0.0f ? logit * p.penalty_repeat : logit / p.penalty_repeat;
+                             logit -= (float) c * p.penalty_freq + p.penalty_present; }
+                if (first || logit > bv) { bv = logit; best = v; first = false; }
+            }
+            return best;
+        };
+        const int want = pick_clamped(true), unclamped = pick_clamped(false);
+        const bool visible = want != unclamped;
+        std::printf("  %-34s %s (clamped pick %d, full-history pick %d)\n",
+                    "penalty window clamp is observable", visible ? "yes" : "*** NO ***", want, unclamped);
+        if (!visible) ++bad;
+        else bad += run("penalty window: tail only", {l.begin(), l.end()}, 1, p, {want}, hist, 8);
+    }
+
     // A continuous stream and individual decode calls consume the same draw counters.
     {
         constexpr int count = 32, vocab = 16;

src/program/generate.cpp

@@ -41,6 +41,7 @@
 #include "strata/prefill/prefill.hpp"
 #include "strata/core/native_dense.hpp"
 #include "strata/program/logits_selection.hpp"
+#include "strata/spec/draft_policy.hpp"
 #include "strata/spec/suffix_drafter.hpp"
 
 #include <cuda_runtime.h>
@@ -250,10 +251,11 @@ struct Options {
     /// result) goes through the verify windows, S tokens at a time, instead of the batched prompt path (0 = always
     /// the batched path)
     int64_t short_read = 64;
-    /// The suffix drafter: when the text being written repeats an earlier stretch of the context (code edits, quoted
-    /// input, tool-call JSON) by at least this many tokens, the window is filled with what followed it there instead
-    /// of the MTP's drafts, up to --spec - 1 of them, as long as the MTP's own first guess agrees.  0 = MTP only.
-    int suffix_draft = 0;
+    /// The suffix drafter (prompt lookup): when the text being written repeats an earlier stretch of the context (code
+    /// edits, quoted input, tool-call JSON) by at least this many tokens, the window may be filled with what followed
+    /// it there instead of the MTP's drafts, where the MTP's own first guess agrees and the draft policy expects it to
+    /// pay (strata/spec/draft_policy.hpp).  On by default; 0 = MTP only.
+    int suffix_draft = 3;
     /// The MTP's own window cap (0 = --spec): with --spec 6 --mtp-max-t 4 the long windows come from suffix matches.
     int mtp_max_t = 0;
 };
@@ -322,7 +324,8 @@ void usage() {
                  "  --turn-token ID      --serve: the token that opens a chat turn (default 248045, <|im_start|>)\n"
                  "  --short-read N       --serve: read at most N fresh text tokens through the decode windows instead\n"
                  "                       of the batched prompt path (default 64, 0 = off)\n"
-                 "  --suffix-draft N     draft from an earlier repeat of the last N+ tokens of context (0 = MTP only)\n"
+                 "  --suffix-draft N     prompt lookup: draft from an earlier repeat of the last N+ tokens of context\n"
+                 "                       when it pays (default 3; 0 = MTP only)\n"
                  "  --mtp-max-t M        cap the MTP's windows at M tokens (0 = --spec; longer ones come from suffixes)\n"
                  "  --no-token-graph     A/B: two graphs per layer (the host launches each) instead of one per token\n"
                  "  --no-fused-gr        A/B: the six-kernel hyper-connection read and a separate write (native)\n"
@@ -732,6 +735,13 @@ int main(int argc, char** argv) {
         return 2;
     }
     strata::core::qsa_set_kv_resident(o.kv_resident);
+    // Prompt lookup (the suffix drafter, on by default): the MTP keeps its --spec windows and a lookup window may be
+    // up to 2 tokens longer; the draft policy (strata/spec/draft_policy.hpp) takes one only where it pays. Code
+    // edits +6-11%, ordinary text unchanged (bench/results/2026-09-27-spec). --suffix-draft 0 turns it off.
+    if (o.suffix_draft > 0 && o.spec >= 2 && o.mtp_max_t == 0) {
+        o.mtp_max_t = o.spec;
+        o.spec = std::min(o.spec + 2, 8);   // kVerifyMaxT
+    }
     strata::core::layer_set_shared_early(!o.shared_late);
     if (!o.native_preset.empty()) {
         if (o.no_ple || o.ple_gguf.empty()) {
@@ -1991,6 +2001,16 @@ int main(int argc, char** argv) {
             return 1;
         }
         mem_mark("the head and the prompt path");
+        // the penalty-history buffer: ONE row, the request's last `penalty_last_n` tokens, -1 padded in
+        // front.  Allocated once at the cap; a request without penalties gets a null row and takes the
+        // byte-for-byte neutral path (no upload, no buffer handed to the sampler).
+        constexpr int kPenaltyWindowCap = 4096;
+        int32_t* d_hist = nullptr;
+        std::vector<int32_t> hist_stage((size_t) kPenaltyWindowCap, -1);
+        if (cudaMalloc(&d_hist, (size_t) kPenaltyWindowCap * sizeof(int32_t)) != cudaSuccess) {
+            std::fprintf(stderr, "strata serve: the penalty-history allocation failed\n");
+            return 1;
+        }
         strata::core::Verifier ver;
         strata::core::VerifyHits vh;
         vh.d_res = thits.d_res;
@@ -2164,6 +2184,7 @@ int main(int argc, char** argv) {
         const int S_mtp = o.mtp_max_t > 0 ? std::min(o.mtp_max_t, S) : S;   // the MTP's windows; suffixes go up to S
         if (S_mtp < S) mtp.set_max_drafts(S_mtp - 1);
         strata::spec::SuffixDrafter sfx(std::max(1, o.suffix_draft), 64, (size_t) o.max_context + 4096);
+        strata::spec::DraftPolicy policy(S);   // MTP or lookup window, learned over the whole process
         // The vision path (--vision): GENI <max_new> <embeddings file> <id,id,...> carries images.  The file is one
         // or more strata-vision records (int32 'SVE1', n, nx, ny, n_embd, then n x n_embd floats) in prompt order;
         // each image's rows go to its run of <|image_pad|> tokens, whose M-RoPE positions are mtmd's: t = p,
@@ -2182,6 +2203,37 @@ int main(int argc, char** argv) {
             }
             char* endp = nullptr;
             const long long max_new = std::strtoll(line.c_str() + (geni ? 5 : 4), &endp, 10);
+            // optional sampling keys between max_new and the ids: temperature=F, top_p=F, top_k=N, min_p=F,
+            // penalty_last_n=N, penalty_repeat=F, penalty_freq=F, penalty_present=F, seed=N (text requests
+            // only).  Absent keys keep today's behavior: greedy, no penalties.
+            float req_temperature = 0.0f, req_top_p = 1.0f;
+            int req_top_k = 20;   // the sampler's own default; the sampled path REQUIRES top_k in 1..64
+            unsigned long long req_seed = 0;
+            float req_min_p = 0.0f, req_penalty_repeat = 1.0f, req_penalty_freq = 0.0f, req_penalty_present = 0.0f;
+            int req_penalty_last_n = 0;
+            if (!geni && endp != nullptr) {
+                for (;;) {
+                    while (*endp == ' ') ++endp;
+                    const char* start = endp;
+                    while (*endp != '\0' && *endp != ' ') ++endp;
+                    if (endp == start) break;
+                    const std::string tok(start, (size_t) (endp - start));
+                    const size_t eq = tok.find('=');
+                    if (eq == std::string::npos) { endp = const_cast<char*>(start); break; }
+                    const std::string key = tok.substr(0, eq);
+                    const float fv = std::strtof(tok.c_str() + eq + 1, nullptr);
+                    if (key == "temperature") req_temperature = fv;
+                    else if (key == "top_p") req_top_p = fv;
+                    else if (key == "top_k") req_top_k = std::atoi(tok.c_str() + eq + 1);
+                    else if (key == "min_p") req_min_p = fv;
+                    else if (key == "penalty_last_n") req_penalty_last_n = std::atoi(tok.c_str() + eq + 1);
+                    else if (key == "penalty_repeat") req_penalty_repeat = fv;
+                    else if (key == "penalty_freq") req_penalty_freq = fv;
+                    else if (key == "penalty_present") req_penalty_present = fv;
+                    else if (key == "seed") req_seed = std::strtoull(tok.c_str() + eq + 1, nullptr, 10);
+                    // unknown keys are skipped: the ids start at the first token without '='
+                }
+            }
             std::string emb_path;
             if (geni && endp != nullptr) {
                 while (*endp == ' ') ++endp;
@@ -2417,6 +2469,23 @@ int main(int argc, char** argv) {
                 return true;
             };
             apply_pending(true);
+            // per-request sampling for the verify window's head (greedy when temperature is absent)
+            strata::kernels::SamplerParams req_sp;
+            req_sp.greedy = req_temperature <= 0.0f;
+            req_sp.temperature = req_temperature;
+            req_sp.top_p = req_top_p;
+            req_sp.top_k = req_top_k;
+            req_sp.seed = req_seed ? req_seed
+                                   : (unsigned long long) std::chrono::steady_clock::now().time_since_epoch().count();
+            req_sp.min_p = std::clamp(req_min_p, 0.0f, 1.0f);
+            req_sp.penalty_last_n = std::max(req_penalty_last_n, 0);
+            req_sp.penalty_repeat = req_penalty_repeat;
+            req_sp.penalty_freq = req_penalty_freq;
+            req_sp.penalty_present = req_penalty_present;
+            req_sp.counter = 0;
+            ver.set_sampling(req_sp);
+            const int hist_n = std::min(req_sp.penalty_last_n, kPenaltyWindowCap);
+            ver.set_history(hist_n > 0 ? d_hist : nullptr, hist_n);
             bool cancelled = false;
             tr("prompt start", n - 1);
             // The prompt is read in two parts when it has a turn boundary past `resume`: up to the last <|im_start|>
@@ -2494,13 +2563,20 @@ int main(int argc, char** argv) {
                     while (T < S_mtp && dprob[(size_t) T - 1] >= (float) o.spec_min_p) ++T;
                 }
                 if (first_window) T = 1;
-                // a repeat of earlier context proposes more than the MTP is confident about: take it, but only where the
-                // MTP's own first guess agrees (a repeat the model is not following costs a whole window)
+                // a repeat of earlier context (prompt lookup) where the MTP's own first guess agrees: the policy takes it
+                // when its expected tokens per ms, from the measured acceptance and window costs, beat the MTP window's
                 bool from_sfx = false;
+                int sfx_match = 0;
                 if (o.suffix_draft > 0 && !first_window) {
                     const int k = sfx.propose(S - 1, sbuf.data());
-                    if (k + 1 > T && sbuf[0] == drafts[0]) { T = k + 1; from_sfx = true; }
+                    sfx_match = sfx.last_match();
+                    if (k > 0 && sbuf[0] == drafts[0]) {
+                        const strata::spec::DraftPolicy::Pick pk = policy.choose(T, k, sfx_match);
+                        if (pk.lookup) { T = pk.t; from_sfx = true; }
+                    }
                 }
+                const bool timed_round = !first_window;
+                const Clock::time_point round0 = Clock::now();
                 if (p + T > o.max_context) break;
                 window[0] = x;
                 for (int i = 1; i < T; ++i) window[(size_t) i] = from_sfx ? sbuf[(size_t) i - 1] : drafts[(size_t) i - 1];
@@ -2508,6 +2584,19 @@ int main(int argc, char** argv) {
                 drive.d.experts = 0;
                 drive.d.failed = false;
                 apply_pending(false);
+                if (hist_n > 0) {
+                    // the tail the penalties count over: the tokens the state has consumed plus the fed-back
+                    // head `x` (it joins `consumed` only after this window commits).  Most recent LAST,
+                    // -1 in the unused front slots - the sampler's row layout.
+                    const int64_t avail = (int64_t) consumed.size() + 1;
+                    const int take = (int) std::min<int64_t>(hist_n, avail);
+                    std::fill(hist_stage.begin(), hist_stage.begin() + hist_n, -1);
+                    for (int j = 0; j < take - 1; ++j)
+                        hist_stage[(size_t) (hist_n - take + j)] =
+                            consumed[(size_t) ((int64_t) consumed.size() - (take - 1) + j)];
+                    hist_stage[(size_t) (hist_n - 1)] = (int32_t) x;
+                    cudaMemcpy(d_hist, hist_stage.data(), (size_t) hist_n * sizeof(int32_t), cudaMemcpyHostToDevice);
+                }
                 tr("window", p, T);
                 if (!ver.run(T, window.data(), p, &drive_pool_multi, &drive, outv.data(), err) || drive.d.failed) {
                     std::printf("ERR %s\n", drive.d.failed && drive.d.fail ? drive.d.fail : err.c_str());
@@ -2550,6 +2639,9 @@ int main(int argc, char** argv) {
                     std::printf("ERR %s\n", err.c_str());
                     return 1;
                 }
+                if (timed_round && !eos)
+                    policy.observe(from_sfx, T, a, sfx_match,
+                                   std::chrono::duration<double, std::milli>(Clock::now() - round0).count());
                 if (eos) { finish = "stop"; break; }
                 if (stop_req.load()) { finish = "cancel"; break; }
                 x = outv[(size_t) a];
@@ -3110,6 +3202,7 @@ int main(int argc, char** argv) {
         const int S_mtp = o.mtp_max_t > 0 ? std::min(o.mtp_max_t, o.spec) : o.spec;
         if (use_mtp && S_mtp < o.spec) mtp.set_max_drafts(S_mtp - 1);
         strata::spec::SuffixDrafter sfx(std::max(1, o.suffix_draft), 64, (size_t) o.max_context + 4096);
+        strata::spec::DraftPolicy policy(o.spec);   // MTP or lookup window (see draft_policy.hpp)
         std::vector<int32_t> sbuf((size_t) o.spec, 0);
         int64_t sfx_windows = 0, sfx_drafts = 0, sfx_ok = 0;
         if (o.suffix_draft > 0) {
@@ -3127,10 +3220,16 @@ int main(int argc, char** argv) {
             }
             if (first_window) T = 1;
             bool from_sfx = false;
+            int sfx_match = 0;
             if (o.suffix_draft > 0 && !first_window) {
                 const int k = sfx.propose(o.spec - 1, sbuf.data());
-                if (k + 1 > T && (!use_mtp || sbuf[0] == drafts[0])) { T = k + 1; from_sfx = true; }
+                sfx_match = sfx.last_match();
+                if (k > 0 && (!use_mtp || sbuf[0] == drafts[0])) {
+                    const strata::spec::DraftPolicy::Pick pk = policy.choose(T, k, sfx_match);
+                    if (pk.lookup) { T = pk.t; from_sfx = true; }
+                }
             }
+            const bool timed_round = !first_window;
             ++window_hist[(size_t) T];
             if (p + T > o.max_context) {
                 std::fprintf(stderr, "strata generate: ran out of context at position %lld\n", (long long) p);
@@ -3201,7 +3300,9 @@ int main(int argc, char** argv) {
             }
             x = outv[(size_t) a];
             p += a + 1;
-            total_ms += std::chrono::duration<double, std::milli>(Clock::now() - t0).count();
+            const double round_ms = std::chrono::duration<double, std::milli>(Clock::now() - t0).count();
+            total_ms += round_ms;
+            if (timed_round) policy.observe(from_sfx, T, a, sfx_match, round_ms);
             if (rounds % 64 == 0)
                 std::fprintf(stderr, "strata generate: position %lld, %lld tokens, %lld rounds\n", (long long) p,
                              (long long) produced.size(), (long long) rounds);

src/spec/draft_policy.cpp

@@ -0,0 +1,101 @@
+// src/spec/draft_policy.cpp - see include/strata/spec/draft_policy.hpp.
+#include "strata/spec/draft_policy.hpp"
+
+#include <algorithm>
+
+namespace strata::spec {
+namespace {
+
+// The shape of a round's cost by window size, relative to one token, used only for sizes not measured yet (the
+// measured round times replace it). Between the RTX 5070's measured curves: +10 ms per token with every missed expert
+// on the CPU (bench/results/2026-09-27-spec/window-cost), flatter with the default CPU/DMA split.
+constexpr double kShape[DraftPolicy::kMaxT + 1] = {0.0, 1.0, 1.35, 1.7, 2.05, 2.45, 2.85, 3.25, 3.6};
+constexpr double kCostAlpha = 0.1;    // EMA weight of a new round time
+constexpr double kTokAlpha = 0.05;    // EMA weight of a new MTP window outcome
+constexpr double kDecay = 0.97;       // lookup counts: older windows fade
+// Before a bucket has data: the longer the match, the likelier its continuation (llama.cpp's lookup decoding gates
+// on the same thing); worth 4 observations, so a few real windows override it
+constexpr double kPriorQ[DraftPolicy::kBuckets] = {0.75, 0.88, 0.93, 0.96};
+constexpr double kPriorN = 4.0;
+constexpr int kProbes = 3;       // a lookup window size is tried this often before its guessed cost can veto it
+
+}  // namespace
+
+DraftPolicy::DraftPolicy(int max_t, double margin)
+    : max_t_(std::clamp(max_t, 1, kMaxT)), margin_(margin) {}
+
+int DraftPolicy::bucket(int match) {
+    return match < 6 ? 0 : match < 12 ? 1 : match < 24 ? 2 : 3;
+}
+
+double DraftPolicy::lookup_rate(int match) const {
+    const int b = bucket(match);
+    return (ok_[b] + kPriorN * kPriorQ[b]) / (ok_[b] + bad_[b] + kPriorN);
+}
+
+double DraftPolicy::cost_ms(int t) const {
+    t = std::clamp(t, 1, kMaxT);
+    if (cost_n_[t] > 0) return cost_[t];
+    // scale from the measured sizes, weighting each by how often it was seen
+    double num = 0.0, den = 0.0;
+    for (int u = 1; u <= kMaxT; ++u)
+        if (cost_n_[u] > 0) {
+            const double w = std::min(cost_n_[u], 20.0);
+            num += w * cost_[u] * kShape[t] / kShape[u];
+            den += w;
+        }
+    return den > 0 ? num / den : kShape[t];
+}
+
+double DraftPolicy::mtp_tokens(int t) const {
+    if (mtp_n_[t] > 0) return mtp_tok_[t];
+    return 1.0 + 0.7 * (t - 1);       // before any MTP window of this size: a typical acceptance
+}
+
+DraftPolicy::Pick DraftPolicy::choose(int t_mtp, int lookup_k, int match) const {
+    Pick p;
+    p.t = std::clamp(t_mtp, 1, max_t_);
+    if (lookup_k <= 0) return p;
+    const double base = mtp_tokens(p.t) / cost_ms(p.t);
+    const double q = lookup_rate(match);
+    double e = 1.0, qi = 1.0, best = 0.0;
+    int best_t = 0;
+    for (int k = 1; k <= std::min(lookup_k, max_t_ - 1); ++k) {
+        qi *= q;
+        e += qi;
+        const double r = e / cost_ms(k + 1);
+        if (r > best) { best = r; best_t = k + 1; }
+    }
+    if (best_t > 0 && best > base * (1.0 + margin_)) {
+        p.lookup = true;
+        p.t = best_t;
+        return p;
+    }
+    // a guessed cost can keep the policy from ever measuring a size: the first few times a confident lookup would
+    // need a size not measured yet, it is tried (verification keeps the output; only the one round's speed is at stake)
+    const int t_full = std::min(lookup_k, max_t_ - 1) + 1;
+    if (t_full > p.t && cost_n_[t_full] < kProbes && q >= 0.85) {
+        p.lookup = true;
+        p.t = t_full;
+    }
+    return p;
+}
+
+void DraftPolicy::observe(bool lookup, int t, int accepted, int match, double round_ms) {
+    t = std::clamp(t, 1, kMaxT);
+    if (round_ms > 0) {
+        cost_[t] = cost_n_[t] > 0 ? (1.0 - kCostAlpha) * cost_[t] + kCostAlpha * round_ms : round_ms;
+        cost_n_[t] += 1.0;
+    }
+    if (lookup) {
+        const int b = bucket(match);
+        ok_[b] = kDecay * ok_[b] + accepted;
+        bad_[b] = kDecay * bad_[b] + (accepted < t - 1 ? 1.0 : 0.0);
+    } else {
+        const double got = accepted + 1.0;
+        mtp_tok_[t] = mtp_n_[t] > 0 ? (1.0 - kTokAlpha) * mtp_tok_[t] + kTokAlpha * got : got;
+        mtp_n_[t] += 1.0;
+    }
+}
+
+}  // namespace strata::spec

src/spec/draft_policy_test.cpp

@@ -0,0 +1,63 @@
+// src/spec/draft_policy_test.cpp - DraftPolicy: when does a lookup window beat the MTP's?
+//
+// Simulated rounds with the costs measured on the RTX 5070 (window-cost: ~10 ms more per token) check that
+//   1. with no lookup proposal the MTP window is kept;
+//   2. lookup drafts that are mostly rejected stop being taken (their bucket's rate falls);
+//   3. lookup drafts that are always accepted are taken, and the window grows with them;
+//   4. match-length buckets learn separately (short matches failing does not stop long ones);
+//   5. the policy never proposes a window beyond its cap.
+#include "strata/spec/draft_policy.hpp"
+
+#include <cstdio>
+
+using strata::spec::DraftPolicy;
+
+namespace {
+int g_fail = 0;
+void check(bool ok, const char* what) {
+    std::printf("  %-66s %s\n", what, ok ? "ok" : "FAIL");
+    if (!ok) ++g_fail;
+}
+double cost(int t) { return 19.0 + 10.5 * (t - 1); }   // ms per round, measured shape
+}  // namespace
+
+int main() {
+    std::printf("draft_policy_test\n");
+    {
+        DraftPolicy p(6);
+        for (int i = 0; i < 50; ++i) p.observe(false, 4, 2, 0, cost(4));   // MTP windows of 4: 3 tokens each
+        const DraftPolicy::Pick k = p.choose(4, 0, 0);
+        check(!k.lookup && k.t == 4, "no proposal: the MTP window");
+    }
+    {
+        DraftPolicy p(6);
+        for (int i = 0; i < 50; ++i) p.observe(false, 4, 2, 0, cost(4));
+        for (int t = 2; t <= 6; ++t) p.observe(false, t, 0, 0, cost(t));
+        for (int i = 0; i < 40; ++i) p.observe(true, 6, 0, 4, cost(6));      // short matches, all rejected
+        check(p.lookup_rate(4) < 0.15, "rejected short-match drafts: their rate falls below 0.15");
+        check(!p.choose(4, 5, 4).lookup, "rejected short-match drafts: no longer taken");
+        for (int i = 0; i < 40; ++i) p.observe(true, 6, 5, 30, cost(6));     // long matches, all accepted
+        check(p.lookup_rate(30) > 0.9, "accepted long-match drafts: their rate rises above 0.9");
+        const DraftPolicy::Pick k = p.choose(4, 5, 30);
+        check(k.lookup && k.t == 6, "accepted long matches: the full lookup window is taken");
+        check(!p.choose(4, 5, 4).lookup, "buckets are separate: short matches still not taken");
+        check(p.choose(4, 20, 30).t <= 6, "never beyond the window cap");
+    }
+    {
+        DraftPolicy p(8);
+        for (int i = 0; i < 50; ++i) p.observe(false, 3, 2, 0, cost(3));      // a very good MTP: 3 of 3 tokens
+        for (int t = 2; t <= 8; ++t) p.observe(false, t, t - 1, 0, cost(t));
+        for (int i = 0; i < 40; ++i) p.observe(true, 4, 2, 8, cost(4));       // lookup at q ~ 0.67
+        check(!p.choose(3, 7, 8).lookup, "a mediocre lookup does not replace a strong MTP window");
+    }
+    {
+        DraftPolicy p(6);
+        for (int i = 0; i < 50; ++i) p.observe(false, 4, 3, 0, cost(4));      // a near-perfect MTP, only size 4 seen
+        const DraftPolicy::Pick k = p.choose(4, 5, 40);
+        check(k.lookup && k.t == 6, "an unmeasured size is probed for a confident lookup");
+        for (int i = 0; i < 3; ++i) p.observe(true, 6, 5, 40, 3.0 * cost(6));   // it turns out very expensive
+        check(!p.choose(4, 5, 40).lookup, "after the probes, the measured cost decides");
+    }
+    std::printf(g_fail ? "FAIL\n" : "PASS\n");
+    return g_fail ? 1 : 0;
+}

tools/_paths.py

@@ -14,7 +14,7 @@ ENGINE = Path(__file__).resolve().parents[1]
 
 def gguf_py() -> str:
     cands = [os.environ.get("STRATA_GGUF_PY"), ENGINE / "third_party" / "llama.cpp" / "gguf-py",
-             ENGINE.parents[1] / ".ref" / "llama.cpp" / "gguf-py"]
+             ENGINE.parents[min(1, len(ENGINE.parents) - 1)] / ".ref" / "llama.cpp" / "gguf-py"]   # the dev tree; a shallow install (G:\Strata) has no grandparent
     for c in cands:
         if c and (Path(c) / "gguf").is_dir():
             return str(c)

tools/mtp_rt.py

@@ -19,7 +19,6 @@ from pathlib import Path
 
 import numpy as np
 
-STRATA = Path(__file__).resolve().parents[3]
 sys.path.insert(0, str(Path(__file__).resolve().parent))
 from _paths import add_gguf_py  # noqa: E402
 add_gguf_py()

tools/pack_layer.py

@@ -34,7 +34,7 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
 from gguf_reader import GGUFFile  # noqa: E402
 
 REPO = pathlib.Path(__file__).resolve().parents[1]
-STRATA = REPO.parents[1]          # the development layout's default shards (setup.py always passes --gguf)
+STRATA = REPO.parents[min(1, len(REPO.parents) - 1)]   # the dev layout's default shards (setup.py passes --gguf)
 SHARD1 = STRATA / "Q2_0" / "Qwen3.8-Flash-Next-GSQ-RCO-Q2_0-00001-of-00002.gguf"
 
 H, FF, NE = 2560, 640, 512

tools/strata_pack.py

@@ -41,7 +41,8 @@ from canonical_xcheck import (MAPPINGS, KV_IQ4NL, open_shard,             # noqa
 
 FORMAT_VERSION = 1
 ALIGN = 64                       # every plane starts on a 64-byte boundary (P1.S7)
-STRATA = pathlib.Path(__file__).resolve().parents[3]   # the development layout; build takes the shard from --gguf
+_HERE = pathlib.Path(__file__).resolve()
+STRATA = _HERE.parents[min(3, len(_HERE.parents) - 1)]   # the development layout; build takes the shard from --gguf
 SHARD2 = STRATA / "Q2_0" / "Qwen3.8-Flash-Next-GSQ-RCO-Q2_0-00002-of-00002.gguf"
 
 EXPERT_RE = ("ffn_gate_exps.weight", "ffn_up_exps.weight", "ffn_down_exps.weight")