From f5525f7e7a7e7cbecd386144299493ea40499bd3 Mon Sep 17 00:00:00 2001 From: Alex <18387287+wadealexc@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:20:42 -0400 Subject: [PATCH] server : fix draft model fit vs load inconsistency (#25056) * fix: draft model fit vs load inconsistency * refactor(server): unify draft/mtp parameter initialization, model, and context load - moves speculative init to speculative.cpp - changes server_context_impl model_dft and ctx_dft to use raw pointers - fix: don't throttle progress callback when loading draft model - refactor: rename draft model/ctx load method * fix: valign --- common/speculative.cpp | 107 +++++++++++++++++++++ common/speculative.h | 18 ++++ tools/server/server-context.cpp | 158 +++++++++++--------------------- 3 files changed, 176 insertions(+), 107 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 5b26597f62..1d2977e424 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2221,6 +2221,113 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) { return n_max; } +common_params common_base_params_to_speculative(const common_params & params) { + const bool has_draft = params.speculative.has_dft(); + + const auto & params_spec = params.speculative.draft; + common_params result = params; + + if (has_draft) { + result.devices = params_spec.devices; + result.model = params_spec.mparams; + result.n_gpu_layers = params_spec.n_gpu_layers; + result.tensor_buft_overrides = params_spec.tensor_buft_overrides; + + if (params_spec.cpuparams.n_threads > 0) { + result.cpuparams.n_threads = params_spec.cpuparams.n_threads; + result.cpuparams_batch.n_threads = params_spec.cpuparams_batch.n_threads; + } + } + + result.cache_type_k = params_spec.cache_type_k; + result.cache_type_v = params_spec.cache_type_v; + result.n_outputs_max = params.n_parallel; + + return result; +} + +struct common_init_speculative_result::impl { + impl() = default; + ~impl() = default; + + // note: the order in which model, context, etc. are declared matters because their destructors will be called bottom-to-top + llama_model_ptr model; + llama_context_ptr context; +}; + +common_init_speculative_result::common_init_speculative_result( + common_params & params, + llama_model * model_tgt, + llama_context * ctx_tgt) : + pimpl(new impl{}) { + const bool has_draft = params.speculative.has_dft(); + const bool spec_mtp = std::find(params.speculative.types.begin(), + params.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); + GGML_ASSERT(has_draft || spec_mtp); + + auto mparams = common_model_params_to_llama(params); + auto cparams = common_context_params_to_llama(params); + + if (spec_mtp) { + cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + } + + // note: for small models maybe we can set this to the maximum possible draft from all speculative types + // the extra memory for small models is likely negligible? + cparams.n_rs_seq = 0; + cparams.ctx_other = ctx_tgt; + + std::string model_path; + if (has_draft) { + model_path = params.speculative.draft.mparams.path; + LOG_TRC("%s: loading draft model '%s'\n", __func__, model_path.c_str()); + + llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams); + if (model_dft == NULL) { + LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str()); + return; + } + + pimpl->model.reset(model_dft); + + llama_context * ctx_dft = llama_init_from_model(model_dft, cparams); + if (ctx_dft == nullptr) { + LOG_ERR("%s: failed to create MTP context\n", __func__); + return; + } + + pimpl->context.reset(ctx_dft); + } else if (spec_mtp) { + model_path = params.model.path; + + LOG_TRC("%s: creating MTP draft context against the target model '%s'\n", + __func__, model_path.c_str()); + + llama_context * ctx_dft = llama_init_from_model(model_tgt, cparams); + if (ctx_dft == nullptr) { + LOG_ERR("%s: failed to create MTP context\n", __func__); + return; + } + + pimpl->context.reset(ctx_dft); + } +} + +common_init_speculative_result::~common_init_speculative_result() = default; + +llama_model * common_init_speculative_result::model() { + return pimpl->model.get(); +} + +llama_context * common_init_speculative_result::context() { + return pimpl->context.get(); +} + +common_init_speculative_result_ptr common_init_speculative_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) { + return std::make_unique(params, model_tgt, ctx_tgt); +} + // initialization of the speculative decoding system // common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) { diff --git a/common/speculative.h b/common/speculative.h index c58fac3cc6..2f6d225b6f 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -23,6 +23,8 @@ std::string common_speculative_type_to_str(enum common_speculative_type type); // return the max number of draft tokens based on the speculative parameters int32_t common_speculative_n_max(const common_params_speculative * spec); +common_params common_base_params_to_speculative(const common_params & params); + common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq); void common_speculative_free(common_speculative * spec); @@ -80,3 +82,19 @@ struct common_speculative_deleter { }; typedef std::unique_ptr common_speculative_ptr; + +struct common_init_speculative_result { + common_init_speculative_result(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); + ~common_init_speculative_result(); + + llama_model * model(); + llama_context * context(); + +private: + struct impl; + std::unique_ptr pimpl; +}; + +using common_init_speculative_result_ptr = std::unique_ptr; + +common_init_speculative_result_ptr common_init_speculative_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index bb3b91ab5e..3890fe5f05 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -897,8 +897,10 @@ private: server_batch batch; - llama_model_ptr model_dft; - llama_context_ptr ctx_dft; + llama_model * model_dft = nullptr; + llama_context * ctx_dft = nullptr; + + common_init_speculative_result_ptr spec_init; 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; @@ -939,8 +941,10 @@ private: void destroy() { spec.reset(); - ctx_dft.reset(); - model_dft.reset(); + spec_init.reset(); + + ctx_dft = nullptr; + model_dft = nullptr; llama_init.reset(); @@ -1084,30 +1088,15 @@ private: // optionally reserve VRAM for the draft / MTP context before fitting the target model if (params_base.fit_params) { if (has_spec) { - common_params params_dft = params_base; - bool measure_model_bytes = true; + // MTP draft context lives on the target model, only context+compute are new + bool measure_model_bytes = has_draft; - if (has_draft) { - const auto & params_spec = params_base.speculative.draft; - params_dft.devices = params_spec.devices; - params_dft.model = params_spec.mparams; - params_dft.n_gpu_layers = params_spec.n_gpu_layers; - params_dft.cache_type_k = params_spec.cache_type_k; - params_dft.cache_type_v = params_spec.cache_type_v; - params_dft.tensor_buft_overrides = params_spec.tensor_buft_overrides; - } else { - // MTP draft context lives on the target model, only context+compute are new - measure_model_bytes = false; - } - - params_dft.n_outputs_max = params_base.n_parallel; + common_params params_dft = common_base_params_to_speculative(params_base); auto mparams_dft = common_model_params_to_llama(params_dft); auto cparams_dft = common_context_params_to_llama(params_dft); if (spec_mtp) { cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; - cparams_dft.type_k = params_base.speculative.draft.cache_type_k; - cparams_dft.type_v = params_base.speculative.draft.cache_type_v; } cparams_dft.n_rs_seq = 0; @@ -1175,82 +1164,35 @@ private: add_bos_token = llama_vocab_get_add_bos(vocab); - if (has_draft) { - // TODO speculative: move to common/speculative.cpp? - const auto & params_spec = params_base.speculative.draft; - - SRV_TRC("loading draft model '%s'\n", params_spec.mparams.path.c_str()); - - auto params_dft = params_base; - - params_dft.devices = params_spec.devices; - params_dft.model = params_spec.mparams; - params_dft.n_gpu_layers = params_spec.n_gpu_layers; - params_dft.cache_type_k = params_spec.cache_type_k; - params_dft.cache_type_v = params_spec.cache_type_v; - - if (params_spec.cpuparams.n_threads > 0) { - params_dft.cpuparams.n_threads = params_spec.cpuparams.n_threads; - params_dft.cpuparams_batch.n_threads = params_spec.cpuparams_batch.n_threads; - } - - params_dft.tensor_buft_overrides = params_spec.tensor_buft_overrides; - - auto mparams_dft = common_model_params_to_llama(params_dft); - - // progress callback - mparams_dft.progress_callback = load_progress_callback; - mparams_dft.progress_callback_user_data = &load_progress_spec; - - model_dft.reset(llama_model_load_from_file(params_dft.model.path.c_str(), mparams_dft)); - if (model_dft == nullptr) { - SRV_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str()); - return false; - } - - auto cparams = common_context_params_to_llama(params_dft); - - if (spec_mtp) { - cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; - } - - // note: for small models maybe we can set this to the maximum possible draft from all speculative types - // the extra memory for small models is likely negligible? - cparams.n_rs_seq = 0; - cparams.ctx_other = ctx_tgt; - - ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams)); - if (ctx_dft == nullptr) { - SRV_ERR("%s", "failed to create draft context\n"); - return false; - } - - params_base.speculative.draft.ctx_tgt = ctx_tgt; - params_base.speculative.draft.ctx_dft = ctx_dft.get(); - } else if (spec_mtp) { - // no new model load, so we simply report 0.0 and 1.0 progress + if (has_spec) { + // spec_mtp doesn't use load a model internally, so we report 0.0 and 1.0 manually load_progress_callback(0.0f, &load_progress_spec); + load_progress_spec.t_last_load_progress_ms = 0; // reset so internal cbs aren't delayed - SRV_TRC("creating MTP draft context against the target model '%s'\n", - params_base.model.path.c_str()); + { + common_params params_dft = common_base_params_to_speculative(params_base); + // progress callback + params_dft.load_progress_callback = load_progress_callback; + params_dft.load_progress_callback_user_data = &load_progress_spec; - auto cparams_mtp = common_context_params_to_llama(params_base); - cparams_mtp.ctx_type = LLAMA_CONTEXT_TYPE_MTP; - cparams_mtp.type_k = params_base.speculative.draft.cache_type_k; - cparams_mtp.type_v = params_base.speculative.draft.cache_type_v; - cparams_mtp.n_rs_seq = 0; - cparams_mtp.n_outputs_max = params_base.n_parallel; - cparams_mtp.ctx_other = ctx_tgt; + spec_init = common_init_speculative_from_params(params_dft, model_tgt, ctx_tgt); + model_dft = spec_init->model(); + ctx_dft = spec_init->context(); - ctx_dft.reset(llama_init_from_model(model_tgt, cparams_mtp)); - if (ctx_dft == nullptr) { - SRV_ERR("%s", "failed to create MTP context\n"); - return false; + if (has_draft && model_dft == nullptr) { + SRV_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str()); + return false; + } + + if (ctx_dft == nullptr) { + SRV_ERR("%s", "failed to create MTP context\n"); + return false; + } + + params_base.speculative.draft.ctx_tgt = ctx_tgt; + params_base.speculative.draft.ctx_dft = ctx_dft; } - params_base.speculative.draft.ctx_tgt = ctx_tgt; - params_base.speculative.draft.ctx_dft = ctx_dft.get(); - load_progress_callback(1.0f, &load_progress_spec); } @@ -1343,13 +1285,15 @@ private: } if (ctx_dft) { - ctx_dft_seq_rm_type = common_context_can_seq_rm(ctx_dft.get()); + ctx_dft_seq_rm_type = common_context_can_seq_rm(ctx_dft); } if (spec) { SRV_TRC("%s", "speculative decoding context initialized\n"); } else { - ctx_dft.reset(); + spec_init.reset(); + ctx_dft = nullptr; + model_dft = nullptr; } for (int i = 0; i < params_base.n_parallel; i++) { @@ -1357,7 +1301,7 @@ private: slot.id = i; slot.ctx_tgt = ctx_tgt; - slot.ctx_dft = ctx_dft.get(); + slot.ctx_dft = ctx_dft; slot.spec = spec.get(); slot.n_ctx = n_ctx_slot; @@ -2363,7 +2307,7 @@ private: cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max); cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - cur.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + cur.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // stash the draft's speculative state with the checkpoint common_speculative_get_state(spec.get(), slot.id, cur.data_spec); @@ -2899,8 +2843,8 @@ private: common_context_seq_add(ctx_tgt, slot.id, n_keep + n_discard, slot.prompt.n_tokens(), -n_discard); if (ctx_dft) { - common_context_seq_rm (ctx_dft.get(), slot.id, n_keep , n_keep + n_discard); - common_context_seq_add(ctx_dft.get(), slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); + common_context_seq_rm (ctx_dft, slot.id, n_keep , n_keep + n_discard); + common_context_seq_add(ctx_dft, slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); } // add generated tokens to cache @@ -2972,7 +2916,7 @@ private: llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id)); if (use_ckpt_dft) { - slot.spec_ckpt.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } slot.spec_prompt = slot.prompt.tokens.get_text_tokens(); @@ -3009,10 +2953,10 @@ private: if (ctx_dft) { if (use_ckpt_dft) { - ckpt.load_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } - common_context_seq_rm(ctx_dft.get(), slot.id, ckpt.pos_max + 1, -1); + common_context_seq_rm(ctx_dft, slot.id, ckpt.pos_max + 1, -1); } if (!draft.empty()) { @@ -3021,7 +2965,7 @@ private: (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_tgt)); const bool use_ckpt_dft = - (ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft.get())); + (ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft)); if (use_ckpt_tgt) { //const int64_t t_start = ggml_time_us(); @@ -3038,7 +2982,7 @@ private: } if (use_ckpt_dft) { - ckpt.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } } }); @@ -3219,8 +3163,8 @@ private: common_context_seq_add(ctx_tgt, slot.id, head_c, head_c + n_match, kv_shift); if (ctx_dft) { - common_context_seq_rm (ctx_dft.get(), slot.id, head_p, head_c); - common_context_seq_add(ctx_dft.get(), slot.id, head_c, head_c + n_match, kv_shift); + common_context_seq_rm (ctx_dft, slot.id, head_p, head_c); + common_context_seq_add(ctx_dft, slot.id, head_c, head_c + n_match, kv_shift); } for (size_t i = 0; i < n_match; i++) { @@ -3321,7 +3265,7 @@ private: if (!do_reset) { // restore the context checkpoint it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - it->load_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // restore the draft's speculative state common_speculative_set_state(spec.get(), slot.id, it->data_spec); @@ -3395,7 +3339,7 @@ private: common_context_seq_rm(ctx_tgt, slot.id, p0, -1); if (ctx_dft) { - common_context_seq_rm(ctx_dft.get(), slot.id, p0, -1); + common_context_seq_rm(ctx_dft, slot.id, p0, -1); } // If using an alora, there may be uncached tokens that come