diff --git a/common/arg.cpp b/common/arg.cpp index 12e33698dc27..8ada9a1718ff 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4156,6 +4156,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", diff --git a/common/common.h b/common/common.h index fd2dde72ff99..9a46edb8e18e 100644 --- a/common/common.h +++ b/common/common.h @@ -333,12 +333,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; diff --git a/common/speculative.cpp b/common/speculative.cpp index ed482143e115..519bce59cbbd 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -3023,7 +3023,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { result.push_back(id); - if (std::min((int) params.n_max, common_lab_knob("LAB_SPEC_DRAFT_N_MAX", params.n_max)) <= (int) result.size()) { + // the global cap (with LAB_SPEC_DRAFT_N_MAX) and the caller's per-draft cap (the server's depth-dependent draft size) + if (std::min((int) params.n_max, common_lab_knob("LAB_SPEC_DRAFT_N_MAX", params.n_max)) <= (int) result.size() || + (dp.n_max > 0 && dp.n_max <= (int) result.size())) { drafting[seq_id] = false; n_drafting--; continue; @@ -3811,8 +3813,29 @@ 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 + // of each sequence before each batch), or, with one sequence, rows up to --spec-draft-depth-max (the server + // feeds a batch only when it starts at or below that depth, and asserts that it ends below depth + n_batch). + // With several sequences one batch can mix a shallow and a deep sequence, so depth-max does not bound the + // rows of each sequence there and the draft context keeps the full size. Sizing it for the bound 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 : (cparams.n_seq_max <= 1 ? depth_max : 0); + 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 5c76f4b20973..d7314ac2eb81 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -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; @@ -3682,6 +3684,38 @@ struct server_context_impl { } } + // with one slot and no window, common_speculative_init sizes the draft context for depth-max: a view that + // reaches it starts at or below the cutoff and holds at most n_batch tokens + if (spec_process && params_base.speculative.draft.n_depth_max > 0 && params_base.speculative.draft.n_window == 0 && + params_base.n_parallel == 1) { + GGML_ASSERT(batch_view.pos[batch_view.n_tokens - 1] < params_base.speculative.draft.n_depth_max + params_base.n_batch); + } + + // --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 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([&]() { diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index 871e5abaf93c..dab67f2ab37b 100644 --- a/tools/server/tests/unit/test_speculative.py +++ b/tools/server/tests/unit/test_speculative.py @@ -106,6 +106,31 @@ def test_different_draft_min_draft_max(): last_content = res.body["content"] +def test_draft_window_keeps_output(): + # a 64-row draft window with a longer prompt: every batch trims the draft context (the trim path), + # drafting goes on, and the greedy output is the same as without the window + global server + request = { + "prompt": "Once upon a time there was a little girl named Lily. " * 12, + "temperature": 0.0, + "top_k": 1, + "n_predict": 32, + "return_tokens": True, + } + server.start() + res = server.make_request("POST", "/completion", data=request) + assert res.status_code == 200 + tokens_full = res.body["tokens"] + server.stop() + + server.spec_draft_window = 64 + server.start() + res = server.make_request("POST", "/completion", data=request) + assert res.status_code == 200 + assert res.body["timings"]["draft_n"] > 0 + assert res.body["tokens"] == tokens_full + + def test_slot_ctx_not_exceeded(): global server server.n_ctx = 256 diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index a0d2dfa3c591..496e7388b957 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -99,6 +99,7 @@ class ServerProcess: spec_type: str | None = None spec_draft_n_min: int | None = None spec_draft_n_max: int | None = None + spec_draft_window: int | None = None no_ui: bool | None = None jinja: bool | None = None reasoning_format: Literal['deepseek', 'none', 'nothink'] | None = None @@ -245,6 +246,8 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: server_args.extend(["--spec-draft-n-max", self.spec_draft_n_max]) if self.spec_draft_n_min: server_args.extend(["--spec-draft-n-min", self.spec_draft_n_min]) + if self.spec_draft_window: + server_args.extend(["--spec-draft-window", self.spec_draft_window]) if self.no_ui: server_args.append("--no-ui") if self.no_models_autoload: