diff --git a/common/speculative.cpp b/common/speculative.cpp index 2ee1e6b848..0082e5fc5d 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -171,12 +171,6 @@ struct common_speculative_impl { // (optional) serialize/restore per-seq internal state (e.g. eagle3's deferred boundary). virtual bool get_state(llama_seq_id /*seq_id*/, std::vector & /*data*/) const { return false; } virtual void set_state(llama_seq_id /*seq_id*/, const std::vector & /*data*/) {} - - // true if this implementation requires the target context to extract post-norm embeddings - virtual bool need_embd() const = 0; - - // true if this implementation requires the target context to extract pre-norm embeddings - virtual bool need_embd_nextn() const { return false; } }; struct common_speculative_impl_draft_simple : public common_speculative_impl { @@ -193,6 +187,10 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { auto * ctx_dft = this->params.ctx_dft; auto * ctx_tgt = this->params.ctx_tgt; + if (!ctx_dft) { + throw std::runtime_error("draft-simple requires a draft context"); + } + SPC_TRC("%s", "adding speculative implementation 'draft-simple'\n"); SPC_TRC("- n_max=%d, n_min=%d, p_min=%f\n", this->params.n_max, this->params.n_min, this->params.p_min); SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", @@ -385,10 +383,6 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override { // noop } - - bool need_embd() const override { - return false; - } }; @@ -907,10 +901,6 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { pending_g_last[seq_id].resize(n_embd_dec); std::memcpy(pending_g_last[seq_id].data(), data.data() + sizeof(llama_pos), (size_t) n_embd_dec * sizeof(float)); } - - bool need_embd() const override { - return false; - } }; // DFlash: block-diffusion drafting with a draft-side KV cache injection @@ -1247,10 +1237,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override { // noop } - - bool need_embd() const override { - return false; - } }; struct common_speculative_impl_draft_mtp : public common_speculative_impl { @@ -1689,14 +1675,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const size_t row_bytes = (size_t) n_embd * sizeof(float); std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes); } - - bool need_embd() const override { - return false; - } - - bool need_embd_nextn() const override { - return true; - } }; // state of self-speculation (simple implementation, not ngram-map) @@ -1743,10 +1721,6 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl { void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override { // noop } - - bool need_embd() const override { - return false; - } }; struct common_speculative_impl_ngram_map_k : public common_speculative_impl { @@ -1801,10 +1775,6 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl { common_ngram_map_accept(config[seq_id], n_accepted); } - - bool need_embd() const override { - return false; - } }; struct common_speculative_impl_ngram_mod : public common_speculative_impl { @@ -1980,10 +1950,6 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { } } } - - bool need_embd() const override { - return false; - } }; struct common_speculative_impl_ngram_cache : public common_speculative_impl { @@ -2123,10 +2089,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override { // noop } - - bool need_embd() const override { - return false; - } }; struct common_speculative { @@ -2322,7 +2284,6 @@ common_speculative_init_result::common_speculative_init_result( 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); @@ -2560,34 +2521,6 @@ bool common_speculative_process(common_speculative * spec, const llama_batch & b return result; } -bool common_speculative_need_embd(common_speculative * spec) { - if (spec == nullptr) { - return false; - } - - for (auto & impl : spec->impls) { - if (impl->need_embd()) { - return true; - } - } - - return false; -} - -bool common_speculative_need_embd_nextn(common_speculative * spec) { - if (spec == nullptr) { - return false; - } - - for (auto & impl : spec->impls) { - if (impl->need_embd_nextn()) { - return true; - } - } - - return false; -} - void common_speculative_draft(common_speculative * spec) { if (spec == nullptr) { return; @@ -2672,7 +2605,10 @@ void common_speculative_draft(common_speculative * spec) { void common_speculative_accept(common_speculative * spec, llama_seq_id seq_id, uint16_t n_accepted) { common_speculative_impl * impl = spec->impl_last[seq_id]; - GGML_ASSERT(impl); + if (impl == nullptr) { + GGML_ASSERT(n_accepted == 0); + return; + } { common_time_meas tm(impl->t_accept_us, !impl->gen_perf); diff --git a/common/speculative.h b/common/speculative.h index c6986affda..06b0992ed6 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -67,12 +67,6 @@ void common_speculative_begin(common_speculative * spec, llama_seq_id seq_id, co // process the batch and update the internal state of the speculative context bool common_speculative_process(common_speculative * spec, const llama_batch & batch); -// true if any implementation requires target post-norm embeddings to be extracted -bool common_speculative_need_embd(common_speculative * spec); - -// true if any implementation requires target nextn embeddings to be extracted -bool common_speculative_need_embd_nextn(common_speculative * spec); - // generate drafts for the sequences specified with `common_speculative_get_draft_params` void common_speculative_draft(common_speculative * spec); diff --git a/examples/speculative-simple/README.md b/examples/speculative-simple/README.md index f72129b3f9..b81583f00b 100644 --- a/examples/speculative-simple/README.md +++ b/examples/speculative-simple/README.md @@ -3,10 +3,47 @@ Demonstration of basic greedy speculative decoding ```bash +# spec-type draft-simple ./bin/llama-speculative-simple \ - -m ../models/qwen2.5-32b-coder-instruct/ggml-model-q8_0.gguf \ - -md ../models/qwen2.5-1.5b-coder-instruct/ggml-model-q4_0.gguf \ - -f test.txt -c 0 -ngl 99 --color on \ - --sampling-seq k --top-k 1 -fa on --temp 0.0 \ - -ngld 99 --spec-draft-n-max 16 --spec-draft-n-draft-min 5 --draft-p-min 0.9 + -hf ggml-org/Qwen3-8B-Base-GGUF:Q8_0 \ + -hfd ggml-org/Qwen3-0.6B-Base-GGUF \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-simple --spec-draft-n-max 7 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 + +# spec-type draft-mtp +./bin/llama-speculative-simple \ + -hf ggml-org/Qwen3.6-27B-GGUF:Q8_0 \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-mtp --spec-draft-n-max 3 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 + +# spec-type draft-mtp (with shared KV cache) +# note: this model needs a token at the start to somewhat work without the chat template +./bin/llama-speculative-simple \ + -hf ggml-org/Gemma-4-31B-it-GGUF:Q8_0 \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-mtp --spec-draft-n-max 3 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 + +# spec-type draft-eagle3 +./bin/llama-speculative-simple \ + -hf ggml-org/gpt-oss-20b-GGUF \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-eagle3 --spec-draft-n-max 3 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 + +# spec-type draft-dflash +./bin/llama-speculative-simple \ + -hf ggml-org/Qwen3-8B-GGUF \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-dflash --spec-draft-n-max 7 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 + +# spec-type draft-dspark +./bin/llama-speculative-simple \ + -hf ggml-org/Qwen3-8B-GGUF \ + -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \ + --spec-type draft-dspark --spec-draft-n-max 7 -ngld 99 --color on \ + -n 256 --temp 0 --top-k 1 --seed 42 -ngl 99 -lv 4 ``` diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index c727e8139d..487ae03abf 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -51,48 +51,23 @@ int main(int argc, char ** argv) { const llama_vocab * vocab = llama_model_get_vocab(model_tgt); - // load the draft model - llama_model_ptr model_dft; - llama_context_ptr ctx_dft; + // load the draft model (if any) - this also creates the MTP draft context when MTP speculation is enabled + common_speculative_init_result_ptr spec_init; - // TODO: simplify this logic { - const auto & params_spec = params.speculative.draft; + common_params params_dft = common_base_params_to_speculative(params); - auto params_dft = params; - - params_dft.n_outputs_max = params.n_parallel; - params_dft.n_outputs_max_per_seq = 1; - - params_dft.devices = params_spec.devices; - params_dft.model = params_spec.mparams; - params_dft.n_gpu_layers = params_spec.n_gpu_layers; - - if (params_spec.cpuparams.n_threads > 0) { - params_dft.cpuparams.n_threads = params.speculative.draft.cpuparams.n_threads; - params_dft.cpuparams_batch.n_threads = params.speculative.draft.cpuparams_batch.n_threads; - } - - params_dft.tensor_buft_overrides = params.speculative.draft.tensor_buft_overrides; - - auto mparams_dft = common_model_params_to_llama(params_dft); - - model_dft.reset(llama_model_load_from_file(params_dft.model.path.c_str(), mparams_dft)); - if (model_dft == nullptr) { - LOG_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str()); - return 1; - } - - auto cparams = common_context_params_to_llama(params_dft); - ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams)); + spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt); params.speculative.draft.ctx_tgt = ctx_tgt; - params.speculative.draft.ctx_dft = ctx_dft.get(); + params.speculative.draft.ctx_dft = spec_init->context(); } + llama_context * ctx_dft = params.speculative.draft.ctx_dft; + // check if the context supports partial sequence removal - const bool use_ckpt_tgt = (common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL); - const bool use_ckpt_dft = (common_context_can_seq_rm(ctx_dft.get()) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL); + const bool use_ckpt_tgt = common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; + const bool use_ckpt_dft = common_context_can_seq_rm(ctx_dft) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; if (use_ckpt_tgt) { LOG_INF("speculative decoding will use checkpoints (context does not support partial sequence removal)\n"); @@ -138,9 +113,30 @@ int main(int argc, char ** argv) { // target model sampling context common_sampler_ptr smpl(common_sampler_init(model_tgt, params.sampling)); - // eval the prompt - llama_decode(ctx_tgt, llama_batch_get_one(inp.data(), inp.size() - 1)); - llama_decode(ctx_dft.get(), llama_batch_get_one(inp.data(), inp.size() - 1)); + // init the speculator + const auto & params_spec = params.speculative; + + struct common_speculative * spec = common_speculative_init(params.speculative, 1); + + if (spec == nullptr) { + LOG_ERR("%s", "failed to initialize speculative decoding\n"); + return 1; + } + + // eval the prompt on the target and feed it to the speculative implementation(s) + { + llama_batch batch_prompt = llama_batch_init(inp.size(), 0, 1); + for (size_t i = 0; i < inp.size() - 1; ++i) { + common_batch_add(batch_prompt, inp[i], i, { seq_id }, false); + } + + llama_decode(ctx_tgt, batch_prompt); + + if (!common_speculative_process(spec, batch_prompt)) { + LOG_ERR("%s", "failed to process speculative prompt\n"); + return 1; + } + } // note: keep the last token separate! llama_token id_last = inp.back(); @@ -151,18 +147,12 @@ int main(int argc, char ** argv) { int n_past = inp.size() - 1; - // init the speculator - const auto & params_spec = params.speculative; - - struct common_speculative * spec = common_speculative_init(params.speculative, 1); - common_speculative_begin(spec, seq_id, prompt_tgt); llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1); - size_t n_draft = 0; - llama_tokens draft; + common_prompt_checkpoint ckpt; const auto t_enc_end = ggml_time_us(); @@ -184,13 +174,20 @@ int main(int argc, char ** argv) { llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), seq_id)); if (use_ckpt_dft) { - ckpt.update_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } + // determine the max draft that fits the remaining context and generation budget + int n_draft_max = (int) llama_n_ctx(ctx_tgt) - n_past - 2; + if (params.n_predict >= 0) { + n_draft_max = std::min(n_draft_max, params.n_predict - n_predict - 1); + } + n_draft_max = std::max(n_draft_max, 0); + // generate a new draft common_speculative_get_draft_params(spec, seq_id) = { /* .drafting = */ true, - /* .n_max = */ -1, + /* .n_max = */ n_draft_max, /* .n_past = */ n_past, /* .id_last = */ id_last, /* .prompt = */ &prompt_tgt, @@ -198,9 +195,6 @@ int main(int argc, char ** argv) { }; common_speculative_draft(spec); - // save the original draft size - n_draft = draft.size(); - // save a checkpoint of the target context before evaluating the draft // this allows us to restore the state if partial draft acceptance occurs if (!draft.empty()) { @@ -209,10 +203,13 @@ int main(int argc, char ** argv) { } } - { - ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + // reset the draft context to the checkpoint before verification + if (ctx_dft) { + if (use_ckpt_dft) { + ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } - llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), seq_id, ckpt.pos_max + 1, -1); + llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1); } } else { // we have a previous (partial) draft to reuse from checkpoint restoration @@ -236,10 +233,10 @@ int main(int argc, char ** argv) { llama_decode(ctx_tgt, batch_tgt); } - // evaluate the same batch with the draft model - { - // TODO: extend to support MTP, Eagle, etc. See server code for reference - llama_decode(ctx_dft.get(), batch_tgt); + // feed the batch to the speculative implementation(s) - this drives the draft model, MTP, Eagle3, etc. + if (!common_speculative_process(spec, batch_tgt)) { + LOG_ERR("%s", "failed to process speculative batch\n"); + break; } // only save the sampler sampler state if we use checkpoints @@ -248,6 +245,9 @@ int main(int argc, char ** argv) { smpl_save.reset(common_sampler_clone(smpl.get())); } + // save the size of the draft being verified + const size_t n_draft = draft.size(); + // sample from the full target batch and return the accepted tokens based on the target sampler // // for each token to be accepted, the sampler would have to sample that same token @@ -264,8 +264,8 @@ int main(int argc, char ** argv) { // check for partial draft acceptance: // if the context doesn't support partial sequence removal, restore the checkpoint // and make the accepted tokens the new partial draft for the next iteration - if (use_ckpt_tgt && ids.size() - 1 < draft.size()) { - LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, draft.size()); + if (use_ckpt_tgt && ids.size() - 1 < n_draft) { + LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, n_draft); draft = std::move(ids); @@ -275,10 +275,10 @@ int main(int argc, char ** argv) { llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, ckpt.pos_max + 1, -1); } - { - ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + if (ctx_dft) { + ckpt.load_dft(ctx_dft, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), seq_id, ckpt.pos_max + 1, -1); + llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, ckpt.pos_max + 1, -1); } prompt_tgt.resize(ckpt.n_tokens); @@ -329,8 +329,11 @@ int main(int argc, char ** argv) { { LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past); - llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, n_past, -1); - llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), seq_id, n_past, -1); + llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, n_past, -1); + + if (ctx_dft) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, n_past, -1); + } } if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) { @@ -356,6 +359,7 @@ int main(int argc, char ** argv) { LOG_INF("\n"); LOG_INF("draft:\n\n"); + common_speculative_print_stats(spec); LOG_INF("\n"); LOG_INF("target:\n\n"); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 3b5f6a1218..d75c6856df 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -397,12 +397,7 @@ struct server_slot { bool need_embd() const { GGML_ASSERT(task); - return task->need_embd() || (spec && common_speculative_need_embd(spec)); - } - - bool need_embd_nextn() const { - GGML_ASSERT(task); - return spec && common_speculative_need_embd_nextn(spec); + return task->need_embd(); } // if the context does not have a memory module then all embeddings have to be computed within a single ubatch