Skip to content

speculative-simple : support MTP speculative decoding #25679

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
pdhinaka wants to merge 3 commits into from
Closed

speculative-simple : support MTP speculative decoding #25679

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ccf844f
speculative-simple : support MTP speculative decoding
pdhinaka Jun 10, 2026
746ec87
speculative-simple : now uses common helpers for draft model loading
pdhinaka Jul 17, 2026
2016af9
speculative-simple : fix MTP prompt prefill corrupting shared KV cache
pdhinaka Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 58 additions & 44 deletions examples/speculative-simple/speculative-simple.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,45 +45,39 @@ 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);
ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams));
if (ctx_dft == nullptr) {
LOG_ERR("%s: failed to create draft/MTP context\n", __func__);
return 1;
}

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

@ggerganov ggerganov Jul 28, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fix the vertical alignment here


// 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");
Expand Down Expand Up @@ -129,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.get(), llama_batch_get_one(inp.data(), inp.size() - 1));

// note: keep the last token separate!
llama_token id_last = inp.back();

Expand All @@ -147,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;
Expand All @@ -175,7 +185,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
Expand All @@ -200,11 +210,11 @@ int main(int argc, char ** argv) {
}
}

{
ckpt.load_dft(ctx_dft.get(), 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);
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);
} else {
// we have a previous (partial) draft to reuse from checkpoint restoration
if (use_ckpt_tgt) {
Expand All @@ -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
Expand Down Expand Up @@ -266,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);
}

{
ckpt.load_dft(ctx_dft.get(), 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);
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);

prompt_tgt.resize(ckpt.n_tokens);
smpl = std::move(smpl_save);

Expand Down Expand Up @@ -321,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.get()), 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) {
Expand Down Expand Up @@ -362,3 +375,4 @@ int main(int argc, char ** argv) {

return 0;
}