Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-11 23:57:07 +02:00
Code Issues Packages Projects Releases Wiki Activity

server: fix speculation after an image (#28715)

Browse Source 
* 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
This commit is contained in:
Jesus Gulfo
2026-09-11 03:33:26 -05:00
committed by GitHub
parent 16378d93f9
commit b0dcb8192b
4 changed files with 13 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/speculative.cpp
+10 -10
View File
@@ -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);
}
common/speculative.h
+1 -1
View File
@@ -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
examples/speculative-simple/speculative-simple.cpp
+1 -1
View File
@@ -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
tools/server/server-context.cpp
+1 -1
View File
@@ -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,