diff --git a/common/arg.cpp b/common/arg.cpp index 94dead3ff3ba..80d23f70b9bd 100644 --- a/common/arg.cpp +++ b/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", diff --git a/common/common.h b/common/common.h index 0ab2a3ffb3d1..458324ec608f 100644 --- a/common/common.h +++ b/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; 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..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 @@ -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 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([&]() {