Skip to content

speculative: fix MTP draft crash on vision inputs #25144

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

Draft
ServeurpersoCom wants to merge 4 commits into
base: master
Choose a base branch
from
+16 −6
Draft

speculative: fix MTP draft crash on vision inputs #25144

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6d37ca2
speculative: fix MTP draft crash on vision inputs
ServeurpersoCom Jun 29, 2026
04f5fd8
speculative: renumber MTP draft positions across vision gaps
ServeurpersoCom Jun 30, 2026
6f45650
restore the comment
ServeurpersoCom Jun 30, 2026
1e408b7
speculative: add FIXME for proper draft vision support
ServeurpersoCom Jul 7, 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
22 changes: 16 additions & 6 deletions common/speculative.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1380,8 +1380,18 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
if (!is_mem_shared) {
common_batch_clear(batch);

// renumber the draft KV draft-local so a vision gap leaves no position jump:
// a 1-axis RoPE draft keeps consecutive positions (Y = X + 1) across the gap
// FIXME: fix this properly with vision support in the draft context
auto * mem_dft = llama_get_memory(ctx_dft);
std::vector<llama_pos> dbase(n_seq, 0);
for (llama_seq_id s = 0; s < (llama_seq_id) n_seq; ++s) {
dbase[s] = llama_memory_seq_pos_max(mem_dft, s) + 1;
}
std::vector<llama_pos> dcur = dbase;
for (int k = 0; k < n_tokens; ++k) {
common_batch_add(batch, batch_in.token[k], batch_in.pos[k], { batch_in.seq_id[k][0] }, 0);
const llama_seq_id s = batch_in.seq_id[k][0];
common_batch_add(batch, batch_in.token[k], dcur[s]++, { s }, 0);
}

// shift the tgt embeddings to the right by one position
Expand All @@ -1407,8 +1417,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
set_h(i_batch_beg[seq_id], pending_h[seq_id].data());
}

auto * mem_dft = llama_get_memory(ctx_dft);

bool ok = true;
for (int head = 0; head < n_mtp_layers; ++head) {
if (chain_heads) {
Expand All @@ -1417,7 +1425,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
if (i_batch_beg[seq_id] < 0) {
continue;
}
llama_memory_seq_rm(mem_dft, seq_id, batch_in.pos[i_batch_beg[seq_id]], -1);
llama_memory_seq_rm(mem_dft, seq_id, dbase[seq_id], -1);
}
llama_set_nextn_layer_offset(ctx_dft, head);
}
Expand Down Expand Up @@ -1468,6 +1476,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
// keep track of which sequences are still drafting
int n_drafting = 0;
std::vector<bool> drafting(n_seq);
std::vector<llama_pos> ddbase(n_seq, 0);

const size_t row_bytes = (size_t) n_embd * sizeof(float);

Expand All @@ -1482,7 +1491,8 @@ 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);
ddbase[seq_id] = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id) + 1;
common_batch_add(batch, dp.id_last, ddbase[seq_id], { 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;
Expand All @@ -1505,7 +1515,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
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, ddbase[seq_id], -1);
}
}
llama_set_nextn_layer_offset(ctx_dft, i);
Expand Down
Loading