diff --git a/common/speculative.cpp b/common/speculative.cpp index 3cb0876..cf6f47f 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2274,6 +2274,9 @@ common_speculative_init_result::common_speculative_init_result( if (spec_mtp) { cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + // use the same (possibly fit-reduced) context size as the target + // avoids MTP draft reserving flash attention workspace for full n_ctx + cparams.n_ctx = llama_n_ctx(ctx_tgt); } // note: for small models maybe we can set this to the maximum possible draft from all speculative types diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 744593c..9f15356 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1091,6 +1091,19 @@ private: auto cparams_dft = common_context_params_to_llama(params_dft); if (spec_mtp) { cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + // MTP draft uses the target's (possibly fit-reduced) n_ctx at runtime. + // Estimate with the same n_ctx to avoid over-reserving fit_params_target. + { + auto mparams_tgt = common_model_params_to_llama(params_base); + auto cparams_tgt = common_context_params_to_llama(params_base); + common_fit_params(params_dft.model.path.c_str(), &mparams_tgt, &cparams_tgt, + params_base.tensor_split, + params_base.tensor_buft_overrides.data(), + params_base.fit_params_target.data(), + params_base.fit_params_min_ctx, + GGML_LOG_LEVEL_ERROR); + cparams_dft.n_ctx = cparams_tgt.n_ctx; + } } cparams_dft.n_rs_seq = 0;