From ccf844f69c9f2d72f3d9492f507da7d10c56804b Mon Sep 17 00:00:00 2001 From: Pranav Dhinakar Date: Wed, 10 Jun 2026 11:28:49 -0700 Subject: [PATCH 1/3] speculative-simple : support MTP speculative decoding --- .../speculative-simple/speculative-simple.cpp | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index d87ba48beb14..7693cf4f3f09 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -75,6 +75,16 @@ int main(int argc, char ** argv) { } auto cparams = common_context_params_to_llama(params_dft); + + const bool spec_mtp = std::find(params.speculative.types.begin(), + params.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); + if (spec_mtp) { + cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + } + cparams.n_rs_seq = 0; + cparams.ctx_other = ctx_tgt; // ctx_tgt already loaded above (line 44) + ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams)); params.speculative.draft.ctx_tgt = ctx_tgt; @@ -228,9 +238,9 @@ int main(int argc, char ** argv) { } // 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); + if (!common_speculative_process(spec, batch_tgt)) { + LOG_ERR("%s: failed to process speculative batch\n", __func__); + return 1; } // only save the sampler sampler state if we use checkpoints From 746ec878e667e9450b3f4e00d9e8bf525a0e8561 Mon Sep 17 00:00:00 2001 From: Pranav Dhinakar Date: Fri, 17 Jul 2026 12:22:00 -0700 Subject: [PATCH 2/3] speculative-simple : now uses common helpers for draft model loading --- .../speculative-simple/speculative-simple.cpp | 62 +++++++------------ 1 file changed, 21 insertions(+), 41 deletions(-) diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index 7693cf4f3f09..95ce3f353817 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -45,55 +45,35 @@ 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 using the shared common helpers + common_speculative_init_result_ptr spec_init; + llama_model * model_dft = nullptr; + llama_context * ctx_dft = nullptr; - // 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); + model_dft = spec_init->model(); + ctx_dft = spec_init->context(); - 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) { + if (params.speculative.has_dft() && 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); - - const bool spec_mtp = std::find(params.speculative.types.begin(), - params.speculative.types.end(), - COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); - if (spec_mtp) { - cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + if (ctx_dft == nullptr) { + LOG_ERR("%s: failed to create draft/MTP context\n", __func__); + return 1; } - cparams.n_rs_seq = 0; - cparams.ctx_other = ctx_tgt; // ctx_tgt already loaded above (line 44) - - 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 = 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_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"); @@ -141,7 +121,7 @@ int main(int argc, char ** argv) { // 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)); + llama_decode(ctx_dft, llama_batch_get_one(inp.data(), inp.size() - 1)); // note: keep the last token separate! llama_token id_last = inp.back(); @@ -185,7 +165,7 @@ 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); } // generate a new draft @@ -211,9 +191,9 @@ int main(int argc, char ** argv) { } { - 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); } } else { // we have a previous (partial) draft to reuse from checkpoint restoration @@ -277,9 +257,9 @@ int main(int argc, char ** argv) { } { - 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); @@ -331,7 +311,7 @@ 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_dft), seq_id, n_past, -1); } if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) { From 2016af956f56caabdc768cfb8211bd7db6f84e94 Mon Sep 17 00:00:00 2001 From: Pranav Dhinakar Date: Fri, 31 Jul 2026 14:00:50 -0700 Subject: [PATCH 3/3] speculative-simple : fix MTP prompt prefill corrupting shared KV cache --- .../speculative-simple/speculative-simple.cpp | 50 ++++++++++++++----- 1 file changed, 37 insertions(+), 13 deletions(-) diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index 95ce3f353817..ab5832342614 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -75,6 +75,10 @@ int main(int argc, char ** argv) { 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); + // some implementations (e.g. Gemma 4 MTP) share the target's memory with the draft context + // in that case any ctx_dft memory operation also affects the target, so they have to be skipped + const bool is_mem_shared = llama_get_memory(ctx_dft) == llama_get_memory(ctx_tgt); + if (use_ckpt_tgt) { LOG_INF("speculative decoding will use checkpoints (context does not support partial sequence removal)\n"); } @@ -119,10 +123,6 @@ 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, llama_batch_get_one(inp.data(), inp.size() - 1)); - // note: keep the last token separate! llama_token id_last = inp.back(); @@ -137,10 +137,30 @@ int main(int argc, char ** argv) { 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); + // eval the prompt + if (inp.size() > 1) { + common_batch_clear(batch_tgt); + + for (size_t i = 0; i + 1 < inp.size(); ++i) { + // request the outputs at every position - the speculator may need them + common_batch_add(batch_tgt, inp[i], (llama_pos) i, { seq_id }, true); + } + + if (llama_decode(ctx_tgt, batch_tgt) != 0) { + LOG_ERR("%s: failed to eval the prompt\n", __func__); + return 1; + } + + if (!common_speculative_process(spec, batch_tgt)) { + LOG_ERR("%s: failed to process the prompt with the speculator\n", __func__); + return 1; + } + } + + common_speculative_begin(spec, seq_id, prompt_tgt); + size_t n_draft = 0; llama_tokens draft; @@ -190,11 +210,11 @@ int main(int argc, char ** argv) { } } - { + 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), 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 if (use_ckpt_tgt) { @@ -256,12 +276,12 @@ int main(int argc, char ** argv) { llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, ckpt.pos_max + 1, -1); } - { + 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), 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); smpl = std::move(smpl_save); @@ -311,7 +331,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), seq_id, n_past, -1); + + if (!is_mem_shared) { + 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) { @@ -352,3 +375,4 @@ int main(int argc, char ** argv) { return 0; } +