speculative: fix MTP draft crash on vision inputs

[ServeurpersoCom]

Overview

Some models that use an MTP draft for speculative decoding abort the request with failed to process speculative batch when an image is sent (e.g. Step-3.7-Flash with its MTP draft). Text-only works; the failure only shows up once an image enters the conversation. Not all MTP models are affected, which is what made it confusing (see the three families below).
Why: the draft can't "see" images (there is no token for image data), so while the main model reads through the picture, the draft stays at the position just before it. On the next token the draft is asked to continue from far ahead of where it sits, and a 1-axis RoPE draft refuses, since its positions must stay consecutive.

The fix gives the draft its own position space: instead of mirroring the target's positions, the draft KV is numbered draft-local and consecutive (position = its own last position + 1). The draft never numbers the image, so no hole ever opens: it simply continues, the 1-axis RoPE stays valid (Y = X + 1), and speculation keeps running across the image instead of erroring out. Because the draft never uses target positions, the server's KV bookkeeping (rollback on rejection, checkpoints, prompt cache, parallel slots, request cancellation) stays correct.

This is strictly better than dropping speculation on the affected turn: the speedup is kept even with an image in context.

Additional information

Validated across the three MTP families, all with images in context (output is unchanged, since speculative decoding verifies every drafted token against the target):

• separate-cache + 1-axis RoPE (Step-3.7-Flash + its MTP draft): crashed before, now drafts correctly across the image and keeps a high acceptance;
• separate-cache + M-RoPE (Qwen3.6): drafts across images and is faster than non-MTP;
• shared-cache (Gemma-4): the draft shares the target KV, so the catch-up is skipped and no gap can form; unaffected by this change, and still faster than non-MTP with images.

Fixes #25129

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES

[ServeurpersoCom]

Tested whether the draft can speculate across the image instead of just disabling it to try to achieve maximum performance

On Step-3.7-Flash:

Relaxing the position check (let it jump the gap): no crash but 0 accepted drafts, the hole in the draft KV makes it useless.
renumbering the draft positions so they stay consecutive (hidden states still come from the target, which saw the image): 0.75 acceptance, so it works.

So it's doable, but it's a feature not a tweak: needs a persistent per-seq offset at every draft-position site, plus multi-image / rollback / checkpoints to handle.

It's more complex than just disabling the draft, as this PR does.

Edit: went ahead with the renumber (commit 2) and it's simpler than I feared, no persistent offset needed. Numbering the draft KV stateless (draft pos = its own pos_max + 1) means the draft never sees a hole, so there's nothing to carry across the gap.
What made me call it a "feature" was the server touching the draft KV with target positions (rollback, checkpoints, prompt cache, parallel slots, cancellation). But the draft stays in its own consecutive space and never uses target positions, so those paths stay correct, I checked them all with images in context. Acceptance holds up too, better than the 0.75 from my first quick probe.
So commit 2 keeps the speedup on image turns instead of dropping speculation; commit 1 stays as the conservative first step.

[ngxson]

I still don't quite get the logic added via this PR: the m-rope temporal position (first axis) can jump non-linear, for example: 1, 2, 3, 3, 3, 3, 7

So will it still be the case after this PR?

[ServeurpersoCom]

The non-linear positions only exist in the target KV, untouched by this PR. The draft never decodes image embeddings, so its stream is text-only, now numbered in its own consecutive space.

This holds for the three MTP families:

• separate-cache + M-RoPE broadcasts the single position across all rope sections and only requires X < Y;
• separate-cache + 1-axis RoPE is the crashing case this fixes (Y = X + 1 now holds);
• shared-cache skips the catch-up entirely, so nothing changes there.

[ServeurpersoCom]

Concretely, with your example:

target (temporal axis): 0 1 2 [3 3 3 3] 7 8 9 ...
draft  (own space):     0 1 2           3 4 5 ...

The image tokens never enter the draft KV (no token ids), and the post-image text just continues the draft's own numbering. The non-linear pattern stays target-side only.

[ngxson]

in such case, I think this is a temporary patch, as it will likely degrade the draft accept rate for vision input

(note: currently, draft context doesn't accept vision input, so this won't change much; but the vision support is another FIXME)

probably better to add a FIXME and fix it properly in the future. cc @am17an for visibility too

[ngxson]

@ServeurpersoCom can you add a FIXME so that we can merge this?

[ServeurpersoCom]

You're right, and I think a proper support is doable: an MTP draft already lives in the target's hidden space, so vision support means reinjecting the target-side image embeddings into the draft KV during the catch-up, with draft-local positions. No second mmproj needed. This is MTP-specific: independent drafts (draft-simple, shared vocab only) already get a compacted text-only prompt server-side (get_text_tokens), which stays the correct behavior for them since they can't consume target-space embeddings anyway.

[ggerganov]

The proper way to fix this is to extend the batches to be able to provide both vision and target embeddings at the same time. Currently, we cannot handle that, so we have this logic which skips the vision embeddings from being processed by the draft context:

llama.cpp/common/speculative.cpp

Lines 1359 to 1362 in 33ca0dc

	// TODO: how to make it work with vision tokens?
	if (batch_in.token == nullptr || batch_in.embd != nullptr) {
	return true;
	}

This sort of works for M-RoPE models (e.g. Qwen3.6) but does not work for non-M-RoPE models (like Step 3.7) because we allow position jumps only for the M-Rope case:

llama.cpp/src/llama-batch.cpp

Lines 255 to 289 in 33ca0dc

	if (n_pos_per_embd > 1) {
	// M-RoPE case: allow position to "jump" forward only (non-continuous positions are allowed)
	for (uint32_t s = 0; s < n_seq_max; ++s) {
	if (seq_pos[s].empty()) {
	continue;
	}
	const llama_pos p0 = memory ? memory->seq_pos_max(s) : -1;
	if (batch.token) {
	if (p0 >= 0 && p0 >= seq_pos_min(s)) {
	LLAMA_LOG_ERROR(
	"%s: the tokens of sequence %d in the input batch have inconsistent sequence positions:\n"
	" - the last position stored in the memory module of the context (i.e. the KV cache) for sequence %d is X = %d\n"
	" - the tokens for sequence %d in the input batch have a starting position of Y = %d\n"
	" for M-RoPE, it is required that the position satisfies: X < Y\n",
	__func__, s, s, p0, s, seq_pos_min(s));
	return false;
	}
	} else {
	// embedding inputs can have overlapping positions
	if (p0 >= 0 && p0 > seq_pos_min(s)) {
	LLAMA_LOG_ERROR(
	"%s: the tokens of sequence %d in the input batch have inconsistent sequence positions:\n"
	" - the last position stored in the memory module of the context (i.e. the KV cache) for sequence %d is X = %d\n"
	" - the tokens for sequence %d in the input batch have a starting position of Y = %d\n"
	" for M-RoPE, it is required that the position satisfies: X <= Y\n",
	__func__, s, s, p0, s, seq_pos_min(s));
	return false;
	}
	}
	}
	} else {

My suggestion is instead of merging the changes from this PR, to focus on refactoring the llama_batch (#24669) and supporting this properly for all cases.