spec : update speculative-simple (#26904)

* spec : update speculative-simple

* cont : simplify

* cont : clean-up

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<uint8_t> & /*data*/) const { return false; }
     virtual void set_state(llama_seq_id /*seq_id*/, const std::vector<uint8_t> & /*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);

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);
 

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 \
+        -hf  ggml-org/Qwen3-8B-Base-GGUF:Q8_0 \
-    -md ../models/qwen2.5-1.5b-coder-instruct/ggml-model-q4_0.gguf \
+        -hfd ggml-org/Qwen3-0.6B-Base-GGUF \
-    -f test.txt -c 0 -ngl 99 --color on \
+        -p "Here is a quick sort implementation in C++. Just code, no comments:\n\n#include" \
-    --sampling-seq k --top-k 1 -fa on --temp 0.0 \
+        --spec-type draft-simple --spec-draft-n-max 7 -ngld 99 --color on \
-    -ngld 99 --spec-draft-n-max 16 --spec-draft-n-draft-min 5 --draft-p-min 0.9
+        -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 <s> 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 "<s>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
 ```

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
+    // load the draft model (if any) - this also creates the MTP draft context when MTP speculation is enabled
-    llama_model_ptr model_dft;
+    common_speculative_init_result_ptr spec_init;
-    llama_context_ptr ctx_dft;
 
-    // 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;
+        spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt);
-
-        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));
 
         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_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_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
+    // init the speculator
-    llama_decode(ctx_tgt,       llama_batch_get_one(inp.data(), inp.size() - 1));
+    const auto & params_spec = params.speculative;
-    llama_decode(ctx_dft.get(), llama_batch_get_one(inp.data(), inp.size() - 1));
+
+    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) {
                 }
             }
 
-            {
+            // reset the draft context to the checkpoint before verification
-                ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+            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
+        // feed the batch to the speculative implementation(s) - this drives the draft model, MTP, Eagle3, etc.
-        {
+        if (!common_speculative_process(spec, batch_tgt)) {
-            // TODO: extend to support MTP, Eagle, etc. See server code for reference
+            LOG_ERR("%s", "failed to process speculative batch\n");
-            llama_decode(ctx_dft.get(), batch_tgt);
+            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()) {
+        if (use_ckpt_tgt && ids.size() - 1 < n_draft) {
-            LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, draft.size());
+            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);
             }
 
-            {
+            if (ctx_dft) {
-                ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+                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);
@@ -330,7 +330,10 @@ 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);
+
+            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");

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));
+        return task->need_embd();
-    }
-
-    bool need_embd_nextn() const {
-        GGML_ASSERT(task);
-        return spec && common_speculative_need_embd_nextn(spec);
     }
 
     // if the context does not have a memory module then all embeddings have to be computed within a single ubatch