From 971dcaddc72a951abbee6fd5e6f27e86f6b9ec14 Mon Sep 17 00:00:00 2001 From: Cary Palmer <24235924+professorpalmer@users.noreply.github.com> Date: Sun, 27 Sep 2026 02:34:55 -0500 Subject: [PATCH 1/2] speculative: --spec-draft-window and --spec-draft-n-max-tail (MTP drafting at every depth) --spec-draft-window N: the draft (MTP) context keeps only the last N rows. The server drops older rows before feeding each batch and the draft context is sized for the window, 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: a 16k window accepts as many drafts as the full history at 131k. Also sizes the draft context for --spec-draft-depth-max when no window is set. --spec-draft-n-max-tail N: draft size once the sequence reaches --kv-vram-cells. Past the tiered-KV line a step is bound by reading the host tail over PCIe, and a wider verify reads it once for all columns. The drafter and the output limits are built for max(n_max, n_max_tail); each slot caps a draft by depth, and the MTP draft loop now honours that per-draft cap. Bonsai 2 27B on an RTX 4070 with the MMA decode route: drafting pays at every depth, so the --spec-draft-depth-max cutoff is no longer needed (32k 54.5 -> 103.6 tok/s, 64k 48.0 -> 90.1, same acceptance; window alone +17-18% at 131k; tail draft 4 vs 2: +26% code / +15% prose at 180k). Co-Authored-By: Claude Opus 5.5 (cherry picked from commit d2e29649b20916fbf5a03e759bc1f4fda4296c34) --- common/arg.cpp | 22 +++++++++++++++++++ common/common.h | 18 +++++++++++---- common/speculative.cpp | 23 +++++++++++++++++-- tools/server/server-context.cpp | 39 ++++++++++++++++++++++++++++++++- 4 files changed, 95 insertions(+), 7 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 94dead3ff3ba..96c239301c56 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4128,6 +4128,28 @@ 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\n" + "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-n-max-tail"}, "N", + string_format("draft size once the sequence reaches --kv-vram-cells, where a wider verify reads the host tail\n" + "once for all columns; 0 = --spec-draft-n-max (default: %d)", params.speculative.draft.n_max_tail), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + params.speculative.draft.n_max_tail = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MAX_TAIL")); add_opt(common_arg( {"--spec-draft-p-split", "--draft-p-split"}, "P", diff --git a/common/common.h b/common/common.h index 0ab2a3ffb3d1..551af6ec61a9 100644 --- a/common/common.h +++ b/common/common.h @@ -329,12 +329,22 @@ 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; + + // draft size once the sequence reaches the tiered-KV line (--kv-vram-cells; 0 = n_max). Past it a step + // is bound by reading the host tail over PCIe and a wider verify batch reads it once for all columns. + int32_t n_max_tail = 0; + bool backend_sampling = true; // offload draft sampling to the backend (default: on) common_params_model mparams; diff --git a/common/speculative.cpp b/common/speculative.cpp index 7b53a51257dd..f523f0c3008a 100644 --- a/common/speculative.cpp +++ b/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(1, cparams.n_seq_max); + const uint32_t need = (uint32_t) span + 2*std::max(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; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 988d9f365190..5ee7316bb617 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -207,6 +207,9 @@ struct server_slot { // speculative decoding common_speculative * spec; int32_t spec_depth_max = 0; // no drafting once the sequence is longer than this (0 = always draft) + int32_t spec_n_max = 0; // draft size (--spec-draft-n-max) below spec_tail_depth + int32_t spec_n_max_tail = 0; // draft size from spec_tail_depth on (--spec-draft-n-max-tail, 0 = spec_n_max) + int32_t spec_tail_depth = 0; // start of the tiered-KV host tail (--kv-vram-cells, 0 = none) llama_tokens spec_draft; llama_tokens spec_prompt; @@ -459,6 +462,14 @@ struct server_slot { n_draft_max = std::min(n_draft_max, n_remaining() - 1); } + // past the tiered-KV line a step is bound by reading the host tail over PCIe, and a wider verify + // batch reads it once for all columns: draft --spec-draft-n-max-tail there, --spec-draft-n-max below + const bool tail = spec_tail_depth > 0 && spec_n_max_tail > 0 && prompt.n_tokens() >= spec_tail_depth; + const int32_t cap = tail ? spec_n_max_tail : spec_n_max; + if (cap > 0) { + n_draft_max = std::min(n_draft_max, (int) cap); + } + SLT_DBG(*this, "max possible draft: %d\n", n_draft_max); return n_draft_max; @@ -850,6 +861,7 @@ struct server_context_impl { llama_context * ctx_dft = nullptr; common_speculative_init_result_ptr spec_init; + int32_t spec_n_max = 0; // --spec-draft-n-max as given (the drafter is built for max(n_max, n_max_tail)) common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; @@ -972,6 +984,12 @@ struct server_context_impl { const bool is_resume = sleeping; params_base = params; + + // the drafter and the output limits are built for the larger of the two draft sizes; each slot caps a + // draft to --spec-draft-n-max below the tiered-KV line (see server_slot::get_n_draft_max) + spec_n_max = params_base.speculative.draft.n_max; + params_base.speculative.draft.n_max = std::max(params_base.speculative.draft.n_max, params_base.speculative.draft.n_max_tail); + 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; @@ -1245,7 +1263,10 @@ struct server_context_impl { slot.ctx_dft = ctx_dft; slot.mem.init(ctx_tgt, ctx_dft); slot.spec = spec.get(); - slot.spec_depth_max = params_base.speculative.draft.n_depth_max; + slot.spec_depth_max = params_base.speculative.draft.n_depth_max; + slot.spec_n_max = spec_n_max; + slot.spec_n_max_tail = params_base.speculative.draft.n_max_tail; + slot.spec_tail_depth = params_base.n_kv_vram_cells; slot.n_ctx = n_ctx_slot; slot.stats.speculative = slot.can_speculate(); @@ -3675,6 +3696,22 @@ struct server_context_impl { } } + // --spec-draft-window: before feeding this view to the draft context, drop its rows older than the window. + // The draft context is sized for the window (common_speculative_init), so freed cells are reused and its + // attention span stays short; the MTP head predicts the next few tokens from recent context. + if (spec_process && ctx_dft && params_base.speculative.draft.n_window > 0) { + llama_pos pos_min = std::numeric_limits::max(); + for (int i = 0; i < batch_view.n_tokens; ++i) { + pos_min = std::min(pos_min, batch_view.pos[i]); + } + const llama_pos hi = pos_min - params_base.speculative.draft.n_window; + if (hi > 0) { + for (auto & slot : slots) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, 0, hi); + } + } + } + if (spec_process) { bool ok = true; queue_tasks.yield_to_queue([&]() { From 93fcb95a3e8d37c40f514f3e6cb54dd701434f83 Mon Sep 17 00:00:00 2001 From: Cary Palmer <24235924+professorpalmer@users.noreply.github.com> Date: Thu, 8 Oct 2026 13:01:11 -0500 Subject: [PATCH 2/2] speculative: draft window trims each batched sequence by its own cutoff; drop --spec-draft-n-max-tail Review of #320: - The window cutoff was the minimum position of the whole batch and was applied to every slot, so a deep slot could remove the rows of a shallow slot that was not in the batch, and a shallow slot in the batch left extra rows in a deep one. Now each sequence in the batch gets its own cutoff, and sequences that are not in the batch keep their rows. - --spec-draft-n-max-tail needed n_kv_vram_cells from the tiered-KV change (#319), so this branch did not compile on its own. The tail flag, its per-slot fields and the n_max widening at load are removed; the tail returns stacked on the tiered-KV change, with the requested draft limits kept apart from the allocation size (the widening was also recaptured as the normal limit after sleep/wake). - --spec-draft-window help text: one sentence, no forced line break. --- common/arg.cpp | 13 +-------- common/common.h | 4 --- tools/server/server-context.cpp | 50 +++++++++++++-------------------- 3 files changed, 21 insertions(+), 46 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 96c239301c56..80d23f70b9bd 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4130,7 +4130,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).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\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) { @@ -4139,17 +4139,6 @@ common_params_context common_params_parser_init(common_params & params, llama_ex 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-n-max-tail"}, "N", - string_format("draft size once the sequence reaches --kv-vram-cells, where a wider verify reads the host tail\n" - "once for all columns; 0 = --spec-draft-n-max (default: %d)", params.speculative.draft.n_max_tail), - [](common_params & params, int value) { - if (value < 0) { - throw std::invalid_argument("invalid value"); - } - params.speculative.draft.n_max_tail = value; - } - ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MAX_TAIL")); add_opt(common_arg( {"--spec-draft-p-split", "--draft-p-split"}, "P", diff --git a/common/common.h b/common/common.h index 551af6ec61a9..458324ec608f 100644 --- a/common/common.h +++ b/common/common.h @@ -341,10 +341,6 @@ struct common_params_speculative_draft { // drafts as the full history at 131k. int32_t n_window = 0; - // draft size once the sequence reaches the tiered-KV line (--kv-vram-cells; 0 = n_max). Past it a step - // is bound by reading the host tail over PCIe and a wider verify batch reads it once for all columns. - int32_t n_max_tail = 0; - bool backend_sampling = true; // offload draft sampling to the backend (default: on) common_params_model mparams; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5ee7316bb617..f7571ed31b6a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -207,9 +208,6 @@ struct server_slot { // speculative decoding common_speculative * spec; int32_t spec_depth_max = 0; // no drafting once the sequence is longer than this (0 = always draft) - int32_t spec_n_max = 0; // draft size (--spec-draft-n-max) below spec_tail_depth - int32_t spec_n_max_tail = 0; // draft size from spec_tail_depth on (--spec-draft-n-max-tail, 0 = spec_n_max) - int32_t spec_tail_depth = 0; // start of the tiered-KV host tail (--kv-vram-cells, 0 = none) llama_tokens spec_draft; llama_tokens spec_prompt; @@ -462,14 +460,6 @@ struct server_slot { n_draft_max = std::min(n_draft_max, n_remaining() - 1); } - // past the tiered-KV line a step is bound by reading the host tail over PCIe, and a wider verify - // batch reads it once for all columns: draft --spec-draft-n-max-tail there, --spec-draft-n-max below - const bool tail = spec_tail_depth > 0 && spec_n_max_tail > 0 && prompt.n_tokens() >= spec_tail_depth; - const int32_t cap = tail ? spec_n_max_tail : spec_n_max; - if (cap > 0) { - n_draft_max = std::min(n_draft_max, (int) cap); - } - SLT_DBG(*this, "max possible draft: %d\n", n_draft_max); return n_draft_max; @@ -861,7 +851,6 @@ struct server_context_impl { llama_context * ctx_dft = nullptr; common_speculative_init_result_ptr spec_init; - int32_t spec_n_max = 0; // --spec-draft-n-max as given (the drafter is built for max(n_max, n_max_tail)) common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; @@ -985,11 +974,6 @@ struct server_context_impl { params_base = params; - // the drafter and the output limits are built for the larger of the two draft sizes; each slot caps a - // draft to --spec-draft-n-max below the tiered-KV line (see server_slot::get_n_draft_max) - spec_n_max = params_base.speculative.draft.n_max; - params_base.speculative.draft.n_max = std::max(params_base.speculative.draft.n_max, params_base.speculative.draft.n_max_tail); - 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; @@ -1263,10 +1247,7 @@ struct server_context_impl { slot.ctx_dft = ctx_dft; slot.mem.init(ctx_tgt, ctx_dft); slot.spec = spec.get(); - slot.spec_depth_max = params_base.speculative.draft.n_depth_max; - slot.spec_n_max = spec_n_max; - slot.spec_n_max_tail = params_base.speculative.draft.n_max_tail; - slot.spec_tail_depth = params_base.n_kv_vram_cells; + slot.spec_depth_max = params_base.speculative.draft.n_depth_max; slot.n_ctx = n_ctx_slot; slot.stats.speculative = slot.can_speculate(); @@ -3696,18 +3677,27 @@ struct server_context_impl { } } - // --spec-draft-window: before feeding this view to the draft context, drop its rows older than the window. - // The draft context is sized for the window (common_speculative_init), so freed cells are reused and its - // attention span stays short; the MTP head predicts the next few tokens from recent context. + // --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) { - llama_pos pos_min = std::numeric_limits::max(); + std::map seq_pos_min; for (int i = 0; i < batch_view.n_tokens; ++i) { - pos_min = std::min(pos_min, batch_view.pos[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]); + } + } } - const llama_pos hi = pos_min - params_base.speculative.draft.n_window; - if (hi > 0) { - for (auto & slot : slots) { - llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, 0, hi); + 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); } } }