The repeat, frequency and presence penalties count the tokens the reply has written, not the last tokens of the prompt; STRATA_PENALTY_PROMPT=1 brings the old window back (#4590)

With the model-card presence penalty 1.5 over the last 64 tokens, the prompt's own tokens (the digits of a new task, a quoted JSON answer) were penalised, and a conversation whose last turn asks for something new could get its previous answer back. The window now starts at the end of the prompt: the fed-back last prompt token and the reply's tokens. Earlier assistant turns inside the prompt are not counted. The history is built from the request's prompt ids on every request, so a conversation restored from the conversation cache or a saved slot gets the same window as a cold read. The STRATA_ADMIT_MERGE_FIRST merged first window is not used for a request with a penalty window, since its prompt rows would join the reply window.

Measured with the opt-in switch (618cd7a4b) on a P100, IQ2_XS, the reporter's saved 4-message conversation, 32 seeds per arm: old JSON returned 13/32 as shipped, 4/32 reply-only, 3/32 with presence 0. Qwen IQ3_S on an R9700, 16 seeds: old JSON 8/16 as shipped, 0/16 reply-only. 61-task quality027, greedy, one run per arm: 56 and 55 passed as shipped (two runs), 55 reply-only, 54 with presence 0.

Added penalty_window_test (header-only, CPU) and a line in docs/DETAILS.md.

CMakeLists.txt

@@ -1821,6 +1821,10 @@ if(STRATA_BUILD_TESTS OR STRATA_BUILD_CONVERSATION_TESTS)
   add_executable(coupled_draft_test src/core/coupled_draft_test.cpp)
   target_include_directories(coupled_draft_test PRIVATE include)
   add_test(NAME coupled_draft_test COMMAND coupled_draft_test)
+  # #4590: the penalty window is the current reply by default, STRATA_PENALTY_PROMPT=1 the old prompt + reply window (header-only, CPU)
+  add_executable(penalty_window_test src/kernels/penalty_window_test.cpp)
+  target_include_directories(penalty_window_test PRIVATE include)
+  add_test(NAME penalty_window_test COMMAND penalty_window_test)
   # probabilistic draft acceptance (STRATA_SPEC_PROB): rejection sampling preserves p - chi-square over 2M trials per case (CPU)
   add_executable(spec_prob_test src/core/spec_prob_test.cpp)
   target_include_directories(spec_prob_test PRIVATE include)

docs/DETAILS.md

@@ -535,6 +535,10 @@ applies at the next start. Network, key, MCP and program keys are not editable t
   `top_k` keeps at most 64 candidates. Penalties (`presence_penalty`, `frequency_penalty`, `repetition_penalty`,
   `penalty_last_n`; llama.cpp's `repeat_penalty` and `repeat_last_n` work too) apply to every token the speculative check
   verifies, as if decoded one at a time, and make a request 1-11% slower because more drafts are rejected.
+  `penalty_last_n` (64 when a penalty is set and the request gives none) is a window over the tokens of the current answer: the
+  prompt, including earlier assistant turns in it, is not counted, so the numbers or a quoted earlier answer in a new task are not
+  penalised. `STRATA_PENALTY_PROMPT=1` counts the last `penalty_last_n` tokens of prompt and answer together, the old window.
+  This is the Qwen3.8 engine path; the Qwen3.6-35B-A3B path still counts the prompt.
 - Sampling defaults, in order of precedence: the request, the web page's shared Chat settings, the config's `sampling` block,
   then the model card. Without a block, thinking requests get temperature 1.0, top_p 0.95, top_k 20, and non-thinking
   requests get 0.7, 0.8, 20 with presence_penalty 1.5 (Qwen3.6-35B-A3B uses presence_penalty 1.5 in both). Greedy is only for

include/strata/kernels/sampler.hpp

@@ -91,6 +91,18 @@ bool sample_tokens_proposal(const float* logits, int n_tokens, int n_vocab, cons
                         const SamplerParams& p, const int32_t* dtok, const int32_t* qbuf, int n_q, int* out,
                         void* stream);
 
+// #4590: which tokens the repeat / frequency / presence penalties count.  By default only the current reply: the
+// tokens the answer has written (the window starts at the last prompt token, the one the first decode step feeds in).
+// The prompt, including earlier assistant turns in it, is not counted, so a number or a quoted answer in the new
+// task is not penalised just for being there.  STRATA_PENALTY_PROMPT=1 restores the old window, the last
+// penalty_last_n tokens of prompt + reply.  Any other value (or none) is the default.
+inline bool penalty_counts_prompt(const char* env) { return env != nullptr && env[0] == '1'; }
+// The first index of the session's consumed tokens (prompt, then reply) the penalty window may read, given how
+// many were consumed when the request's decode began (the prompt without its last token).
+inline size_t penalty_window_from(size_t consumed_at_decode_start, bool counts_prompt) {
+    return counts_prompt ? 0 : consumed_at_decode_start;
+}
+
 // The penalty-history rows of a verify window, on the host: row t of `out` (T rows of `h` slots) is the last `h`
 // tokens of `tail[0..n_tail)` followed by `window[0..t]`, most recent LAST, -1 in the unused front slots.
 // `tail` is what the state consumed before the window, `window[0]` the fed-back token and `window[1..]` the

src/kernels/penalty_window_test.cpp

@@ -0,0 +1,125 @@
+// #4590: which tokens the repeat / frequency / presence penalties count.  Header-only, CPU, no model.
+//
+// Default: the current reply only.  The prompt (the new task's digits, a quoted earlier answer, earlier assistant turns)
+// is not in the window; the last prompt token, which the first decode step feeds in, is.  STRATA_PENALTY_PROMPT=1 is the
+// old window: the last penalty_last_n tokens of prompt + reply.
+#include "strata/kernels/sampler.hpp"
+
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+using strata::kernels::penalty_counts_prompt;
+using strata::kernels::penalty_rows;
+using strata::kernels::penalty_window_from;
+
+namespace {
+int checks = 0;
+void check(bool ok, const char* label) {
+    ++checks;
+    if (!ok) { std::fprintf(stderr, "FAIL: %s\n", label); std::exit(1); }
+}
+
+constexpr int kH = 64;           // penalty_last_n 64 (the server's default window)
+constexpr int32_t kTaskTok = 12; // a token of the new task ("12"), in the last 64 prompt tokens
+constexpr int32_t kOldAnswer = 7777;   // a token of an earlier assistant turn, quoted near the end of the prompt
+
+// A request's prompt: filler, an earlier assistant answer (kOldAnswer tokens), then the new task (kTaskTok), then the
+// assistant header whose last token is the fed-back head x.
+std::vector<int32_t> prompt() {
+    std::vector<int32_t> ids;
+    for (int i = 0; i < 300; ++i) ids.push_back(1000 + i);
+    for (int i = 0; i < 20; ++i) ids.push_back(kOldAnswer + (i % 3));
+    for (int i = 0; i < 10; ++i) ids.push_back(kTaskTok);
+    for (int i = 0; i < 4; ++i) ids.push_back(50 + i);
+    return ids;
+}
+
+// The row the engine stages for the pick after `reply` has written its tokens: the tokens `consumed` holds when decode
+// began (the prompt without its last token) followed by the reply, read from penalty_window_from().
+std::vector<int32_t> row_after(const std::vector<int32_t>& ids, const std::vector<int32_t>& reply, bool counts_prompt) {
+    std::vector<int32_t> consumed(ids.begin(), ids.end() - 1);
+    const size_t from = penalty_window_from(consumed.size(), counts_prompt);
+    const int32_t x = ids.back();
+    // window = [x, reply...]: the single row of the last one is "the pick after the last reply token"
+    std::vector<int32_t> window = {x};
+    window.insert(window.end(), reply.begin(), reply.end());
+    const int T = (int) window.size();
+    std::vector<int32_t> rows((size_t) T * kH);
+    penalty_rows(consumed.data() + from, (int64_t) (consumed.size() - from), window.data(), T, kH, rows.data());
+    return std::vector<int32_t>(rows.end() - kH, rows.end());   // the last row
+}
+
+int count_of(const std::vector<int32_t>& row, int32_t tok) {
+    int c = 0;
+    for (int32_t t : row) c += t == tok;
+    return c;
+}
+int filled(const std::vector<int32_t>& row) {
+    int c = 0;
+    for (int32_t t : row) c += t >= 0;
+    return c;
+}
+}  // namespace
+
+int main() {
+    // ---- the switch: unset, empty, 0 and anything but "1" is the default (reply only)
+    check(!penalty_counts_prompt(nullptr), "env unset: reply only");
+    check(!penalty_counts_prompt(""), "env empty: reply only");
+    check(!penalty_counts_prompt("0"), "env 0 (the old opt-in value): reply only");
+    check(!penalty_counts_prompt("2"), "env 2: reply only");
+    check(penalty_counts_prompt("1"), "env 1: the old window");
+
+    const std::vector<int32_t> ids = prompt();
+    const int32_t x = ids.back();
+
+    // ---- first pick (nothing written yet)
+    {
+        const std::vector<int32_t> r = row_after(ids, {}, false);
+        check(filled(r) == 1 && r[kH - 1] == x, "default: the first pick counts only the fed-back head x");
+        check(count_of(r, kTaskTok) == 0 && count_of(r, kOldAnswer) == 0, "default: no prompt task / earlier-answer token");
+        const std::vector<int32_t> o = row_after(ids, {}, true);
+        check(filled(o) == kH && o[kH - 1] == x, "opt-out: a full window of prompt tokens ends in x");
+        check(count_of(o, kTaskTok) == 10, "opt-out: the new task's token is in the window (the old behaviour)");
+    }
+
+    // ---- mid reply: the reply's own tokens are counted in both modes, a loop still gets penalised
+    {
+        const std::vector<int32_t> reply = {900, 901, 900, 902, 900};
+        const std::vector<int32_t> r = row_after(ids, reply, false);
+        check(filled(r) == 1 + (int) reply.size(), "default: the window is x + the reply");
+        check(count_of(r, 900) == 3 && count_of(r, 901) == 1 && count_of(r, 902) == 1, "default: reply tokens counted");
+        check(count_of(r, kTaskTok) == 0 && count_of(r, kOldAnswer) == 0, "default: still no prompt token mid reply");
+        const std::vector<int32_t> o = row_after(ids, reply, true);
+        check(count_of(o, 900) == 3 && count_of(o, kTaskTok) > 0, "opt-out: reply and prompt tokens both counted");
+    }
+
+    // ---- a long reply: the window slides over the reply alone, the prompt never comes back
+    {
+        std::vector<int32_t> reply;
+        for (int i = 0; i < 200; ++i) reply.push_back(2000 + (i % 7));
+        const std::vector<int32_t> r = row_after(ids, reply, false);
+        check(filled(r) == kH, "default: a reply longer than the window fills it");
+        bool prompt_token = false;
+        for (int32_t t : r) prompt_token |= (t >= 0 && t < 2000);
+        check(!prompt_token, "default: after 63 reply tokens no prompt token (not even x) is in the window");
+        const std::vector<int32_t> o = row_after(ids, reply, true);
+        check(filled(o) == kH && count_of(o, kTaskTok) == 0, "opt-out: the same long reply has pushed the prompt out too");
+    }
+
+    // ---- a restored conversation: the session comes back from the conversation cache or a saved slot with its state, but
+    // the history is built from the request's ids and the tokens written since, so it is the same as a cold read of the
+    // same ids.  A second turn that quotes the first answer in the prompt does not carry the first reply's window over.
+    {
+        const std::vector<int32_t> cold = row_after(ids, {900, 901}, false);
+        const std::vector<int32_t> restored = row_after(ids, {900, 901}, false);
+        check(cold == restored, "restored session: same window as a cold read of the same ids");
+        std::vector<int32_t> next = ids;   // turn 2: the prompt is turn 1's prompt + its reply + a new task
+        next.insert(next.end(), {900, 901, 902, 99});
+        const std::vector<int32_t> r2 = row_after(next, {}, false);
+        check(filled(r2) == 1 && r2[kH - 1] == 99, "turn 2 default: turn 1's reply tokens in the prompt are not counted");
+    }
+
+    std::printf("penalty_window_test: %d checks passed\n", checks);
+    return 0;
+}

src/program/generate.cpp

@@ -15861,6 +15861,13 @@ int main(int argc, char** argv) {
                                                ? drive.d.pcie_num : pcie_num_of(stages[(size_t) st - 1]->pcie_frac);
             if (pipe) for (int st = 0; st < split_drive.n; ++st) split_drive_b.pcie_num[st] = split_drive.pcie_num[st];
             const int hist_n = std::min(req_sp.penalty_last_n, kPenaltyWindowCap);
+            // #4590: the repeat / frequency / presence penalties count the tokens this reply has written, not the prompt's last
+            // tokens (the prompt's own numbers and quoted answers were penalised, and a conversation whose last turn asks for
+            // something new could get its previous answer back).  Earlier assistant turns inside the prompt are not counted
+            // either.  STRATA_PENALTY_PROMPT=1 brings the old window back (the last penalty_last_n tokens of prompt + reply).
+            // The history is rebuilt from the request's prompt ids every time, so a conversation restored from the cache or a
+            // saved slot seeds it exactly as a cold read does.
+            static const bool pen_prompt = strata::kernels::penalty_counts_prompt(std::getenv("STRATA_PENALTY_PROMPT"));
             ver.set_history(hist_n > 0 ? d_hist : nullptr, hist_n);
             if (pipe) ver_b.set_history(hist_n > 0 ? d_hist : nullptr, hist_n);
             bool cancelled = false;
@@ -15940,6 +15947,7 @@ int main(int argc, char** argv) {
             int64_t merge_at = -1;
             auto merge_from = [&](int64_t a0, int64_t to) -> int64_t {
                 if (!admit_merge_first_env() || admit_slot < 0 || max_new != 1 || batch_mtp || pipe) return -1;
+                if (hist_n > 0 && !pen_prompt) return -1;   // #4590: the merged window's prompt rows would join the reply-only penalty window
                 if (to != n - 1 || to == turn_at || to == root_at || to == message_at || to == pin_at) return -1;
                 if (!req_imgs.empty() || kvg.on || peer.valid() || remote_caches || remote_opt || fs_swap) return -1;
                 if (std::getenv("STRATA_LOGPOS") != nullptr) return -1;
@@ -16095,10 +16103,9 @@ int main(int argc, char** argv) {
             std::vector<int32_t> consumed;
             consumed.reserve((size_t) (n + max_new + S));
             for (int64_t i = 0; i < p; ++i) consumed.push_back((int32_t) ids[(size_t) i]);
-            // #4590 (opt-in, STRATA_PENALTY_PROMPT=0): the repeat / frequency / presence penalties count the tokens this reply has
-            // written, not the prompt's last tokens.  Default unchanged: the window is the last penalty_last_n tokens of prompt + reply.
-            static const bool pen_gen_only = [] { const char* e = std::getenv("STRATA_PENALTY_PROMPT"); return e != nullptr && e[0] == '0'; }();
-            const size_t pen_gen0 = pen_gen_only ? consumed.size() : 0;
+            // #4590: where the penalty window may start in `consumed` (see pen_prompt above): the end of the prompt without its
+            // last token, which joins the window as the fed-back head `x`.
+            const size_t pen_gen0 = strata::kernels::penalty_window_from(consumed.size(), pen_prompt);
             const char* finish = "length";
             const Clock::time_point d0 = Clock::now();
             // STRATA_DECODE_TIMING=1: where a request's decode time goes (one line per request)