From b0dcb8192b201e402ec3eff524e55450f8070e3e Mon Sep 17 00:00:00 2001 From: Jesus Gulfo Date: Fri, 11 Sep 2026 03:33:26 -0500 Subject: [PATCH] server: fix speculation after an image (#28715) * server: fix speculation after an image Pass the actual position to the drafter after an image, instead of the token count. Affects every drafter, not just DFlash. * rename draft n_past to pos0 n_past is used to denote number of tokens and this parameter is meant to be a position --- common/speculative.cpp | 20 +++++++++---------- common/speculative.h | 2 +- .../speculative-simple/speculative-simple.cpp | 2 +- tools/server/server-context.cpp | 2 +- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index b7811b853e..7c8a06365c 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -296,7 +296,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { drafting[seq_id] = true; common_sampler_reset(smpls[seq_id].get()); - common_batch_add(batch, dp.id_last, dp.n_past, { seq_id }, true); + common_batch_add(batch, dp.id_last, dp.pos0, { seq_id }, true); } int ret = llama_decode(ctx_dft, batch); @@ -355,7 +355,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { continue; } - common_batch_add(batch, id, dp.n_past + i + 1, { seq_id }, true); + common_batch_add(batch, id, dp.pos0 + i + 1, { seq_id }, true); } if (batch.n_tokens == 0) { @@ -1197,7 +1197,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { common_sampler_reset(smpls[seq_id].get()); - const int32_t n = (int32_t) dp.n_past; + const int32_t n = (int32_t) dp.pos0; const int32_t n_draft = params.n_max; @@ -1621,7 +1621,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { drafting[seq_id] = true; common_sampler_reset(smpls[seq_id].get()); - common_batch_add(batch, dp.id_last, dp.n_past, { seq_id }, true); + common_batch_add(batch, dp.id_last, dp.pos0, { seq_id }, true); std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, pending_h[seq_id].data(), row_bytes); i_last[seq_id] = batch.n_tokens - 1; @@ -1635,16 +1635,16 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { while (n_drafting > 0) { // each step decodes under a different head, i.e. a different decoder layer, and - // KV is per layer. process() filled this layer's KV only for positions < n_past + // KV is per layer. process() filled this layer's KV only for positions < pos0 // (prompt + accepted prefix) — nothing in the draft region yet. so reset the - // draft region (the seq_rm lower bound is n_past, leaving the prompt KV intact) + // draft region (the seq_rm lower bound is pos0, leaving the prompt KV intact) // and select head i so it rebuilds its own layer's KV there; decoding just the // latest token would leave its attention reading cells only another head wrote. if (chain_heads) { auto * mem_dft = llama_get_memory(ctx_dft); for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { if (drafting[seq_id]) { - llama_memory_seq_rm(mem_dft, seq_id, dparams[seq_id].n_past, -1); + llama_memory_seq_rm(mem_dft, seq_id, dparams[seq_id].pos0, -1); } } llama_set_nextn_layer_offset(ctx_dft, i); @@ -1710,17 +1710,17 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const int n_rows = (int) result.size() + 1; // id_last + tokens drafted so far for (int t = 0; t < n_rows; ++t) { const llama_token tok = (t == 0) ? dp.id_last : result[t - 1]; - common_batch_add(batch, tok, dp.n_past + t, { seq_id }, t == n_rows - 1); + common_batch_add(batch, tok, dp.pos0 + t, { seq_id }, t == n_rows - 1); std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, chain_h[seq_id].data() + (size_t) t * n_embd, row_bytes); } } else if (is_mem_shared) { // note: with shared memory (e.g. Gemma4 assistants) we use the same position for all draft tokens // ref: https://github.com/huggingface/transformers/blob/effde20942e3f82a1b97449f60b3a48c5ff96145/docs/source/en/model_doc/gemma4_assistant.md?plain=1#L36-L37 - common_batch_add(batch, id, dp.n_past, { seq_id }, true); + common_batch_add(batch, id, dp.pos0, { seq_id }, true); std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, h_row, row_bytes); } else { - common_batch_add(batch, id, dp.n_past + i + 1, { seq_id }, true); + common_batch_add(batch, id, dp.pos0 + i + 1, { seq_id }, true); std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, h_row, row_bytes); } diff --git a/common/speculative.h b/common/speculative.h index 22505891f7..c968750e2d 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -61,7 +61,7 @@ struct common_speculative_draft_params { // can be used to constraint the max draft based on the remaining context size int32_t n_max = -1; - llama_pos n_past; + llama_pos pos0; llama_token id_last; // TODO: remove in the future by keeping track of the prompt from the _begin() call and the consecutive accept calls diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index 487ae03abf..863af5a2c7 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -188,7 +188,7 @@ int main(int argc, char ** argv) { common_speculative_get_draft_params(spec, seq_id) = { /* .drafting = */ true, /* .n_max = */ n_draft_max, - /* .n_past = */ n_past, + /* .pos0 = */ n_past, /* .id_last = */ id_last, /* .prompt = */ &prompt_tgt, /* .result = */ &draft, // output diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index fe068d3e91..b6835e4345 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3028,7 +3028,7 @@ private: common_speculative_get_draft_params(spec.get(), slot.id) = { /* .drafting = */ true, /* .n_max = */ n_draft_max, - /* .n_past = */ slot.prompt.n_tokens(), + /* .pos0 = */ slot.prompt.tokens.pos_next(), /* .id_last = */ slot.sampled, /* .prompt = */ &slot.spec_prompt, /* .result = */ &slot.spec_draft,