speculative: --spec-draft-window (MTP drafting at every depth)

common/arg.cpp

@@ -4128,6 +4128,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.speculative.draft.n_depth_max = value;
         }
     ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SPEC_DRAFT_DEPTH_MAX"));
+    add_opt(common_arg(
+        {"--spec-draft-window"}, "N",
+        string_format("draft (MTP) context keeps only the last N rows, so its cache stays small and a draft pass costs "
+                      "the same at any depth; 0 = full history (default: %d)", params.speculative.draft.n_window),
+        [](common_params & params, int value) {
+            if (value < 0) {
+                throw std::invalid_argument("invalid value");
+            }
+            params.speculative.draft.n_window = value;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SPEC_DRAFT_WINDOW"));
 
     add_opt(common_arg(
         {"--spec-draft-p-split", "--draft-p-split"}, "P",

common/common.h

@@ -329,12 +329,18 @@ struct common_params_speculative_draft {
     float p_split = 0.1f; // speculative decoding split probability
     float p_min   = 0.0f; // minimum speculative decoding probability (greedy)
 
-    // stop drafting once the sequence is this long (0 = never). Deep in the context a step is bound
-    // by reading the KV cache, and the draft passes plus the multi-column verify add to that read
-    // without shortening it: on a 4070 with Bonsai 2 27B the draft is +85% at zero depth, breaks
-    // even near 24k tokens and costs 30% at 64k. Past the cutoff the slot decodes one token per step.
+    // stop drafting once the sequence is this long (0 = never). Measured when a deep verify batch ran on the
+    // vector FA kernel and the draft context grew with the conversation (Bonsai 2 27B on a 4070: +85% at zero
+    // depth, even near 24k, -30% at 64k). With quantized-KV decode on the MMA FA kernel and n_window below,
+    // drafting pays at every depth (64k: 48 -> 90 tok/s), so the server default is 0.
     int32_t n_depth_max = 0;
 
+    // draft (MTP) context window (0 = full history): it keeps only the last n_window rows and is sized for
+    // them, so its cells are reused, its cache stays small and on the device, and a draft pass costs the same
+    // at any depth. The MTP head predicts the next few tokens from recent context: 16k rows accept as many
+    // drafts as the full history at 131k.
+    int32_t n_window = 0;
+
     bool backend_sampling = true; // offload draft sampling to the backend (default: on)
 
     common_params_model mparams;

common/speculative.cpp

@@ -2783,7 +2783,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
 
                 result.push_back(id);
 
-                if (params.n_max <= (int) result.size()) {
+                // the caller's per-draft cap (the server's depth-dependent draft size) as well as the global one
+                if (params.n_max <= (int) result.size() || (dp.n_max > 0 && dp.n_max <= (int) result.size())) {
                     drafting[seq_id] = false;
                     n_drafting--;
                     continue;
@@ -3535,8 +3536,26 @@ common_speculative_init_result::common_speculative_init_result(
         cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
     }
 
-    // the draft context holds as many tokens per sequence as the target context
+    // the draft context holds as many tokens per sequence as the target context, unless the MTP draft only
+    // ever sees a bounded part of the sequence: the last --spec-draft-window rows (the server drops older ones
+    // before each batch), or rows up to --spec-draft-depth-max (never fed past it). Sizing it for that keeps
+    // its cache small; with the window its cells are reused and its attention span stays that short.
     cparams.n_ctx = llama_n_ctx(ctx_tgt);
+    if (spec_mtp) {
+        const int32_t window    = params.speculative.draft.n_window;
+        const int32_t depth_max = params.speculative.draft.n_depth_max;
+        const int32_t span      = window > 0 ? window : depth_max;
+        if (span > 0) {
+            const uint32_t n_seq  = std::max<uint32_t>(1, cparams.n_seq_max);
+            const uint32_t need   = (uint32_t) span + 2*std::max<uint32_t>(cparams.n_batch, cparams.n_ubatch) + 256;
+            const uint32_t capped = GGML_PAD(need, 256) * n_seq;
+            if (capped < cparams.n_ctx) {
+                LOG_INF("%s: draft context sized to %u tokens (%s %d, target %u)\n", __func__, capped,
+                        window > 0 ? "spec-draft-window" : "spec-draft-depth-max", span, cparams.n_ctx);
+                cparams.n_ctx = capped;
+            }
+        }
+    }
 
     // n_rs_seq stays as common_context_params_to_llama set it: the draft context needs the same rollback window as the target, with n_rs_seq == 0 its seq_rm fails silently on partial acceptance and keeps stale positions
     cparams.ctx_other = ctx_tgt;

tools/server/server-context.cpp

@@ -22,6 +22,7 @@
 #include <cinttypes>
 #include <exception>
 #include <iterator>
+#include <map>
 #include <memory>
 #include <filesystem>
 #include <utility>
@@ -972,6 +973,7 @@ struct server_context_impl {
         const bool is_resume = sleeping;
 
         params_base = params;
+
         const auto output_limits = server_output_limits(params_base);
         params_base.n_outputs_max = output_limits.total;
         params_base.n_outputs_max_per_seq = output_limits.per_seq;
@@ -3675,6 +3677,31 @@ struct server_context_impl {
             }
         }
 
+        // --spec-draft-window: before feeding this view to the draft context, drop the rows of each sequence in the
+        // view that are older than the window, measured from that sequence's first position in the view. Sequences
+        // that are not in the view keep their rows. The draft context is sized for the window
+        // (common_speculative_init), so freed cells are reused and its attention span stays short.
+        if (spec_process && ctx_dft && params_base.speculative.draft.n_window > 0) {
+            std::map<llama_seq_id, llama_pos> seq_pos_min;
+            for (int i = 0; i < batch_view.n_tokens; ++i) {
+                for (int j = 0; j < batch_view.n_seq_id[i]; ++j) {
+                    const llama_seq_id sid = batch_view.seq_id[i][j];
+                    auto it = seq_pos_min.find(sid);
+                    if (it == seq_pos_min.end()) {
+                        seq_pos_min.emplace(sid, batch_view.pos[i]);
+                    } else {
+                        it->second = std::min(it->second, batch_view.pos[i]);
+                    }
+                }
+            }
+            for (const auto & [sid, pos_min] : seq_pos_min) {
+                const llama_pos hi = pos_min - params_base.speculative.draft.n_window;
+                if (hi > 0) {
+                    llama_memory_seq_rm(llama_get_memory(ctx_dft), sid, 0, hi);
+                }
+            }
+        }
+
         if (spec_process) {
             bool ok = true;
             queue_tasks.yield_to_queue([&]() {