speculative: fix failed to decode mtmd chunk with DFlash

[jesdga95]

Overview

When using DFlash w/ vision models, the drafter memory fails to allocate new tokens because images report a fixed offset and errors out (Tested on Qwen3.8 27B Q5_S, RTX 5090 CUDA 13.3):

E process: llama_decode(ctx_dft) failed rc=1 (n_tokens=1024, offset=1024)
E post-decode callback failed
E slot operator(): failed to decode mtmd chunk, idx = 8409, res = 1
E slot operator(): failed to process mtmd chunk, res = -1
E srv  send_error: error: failed to process mtmd chunk

Stop copying them to allow the drafter to continue.

Additional information

Repro: llama-server with a target + --mmproj + a DFlash draft, use an image bigger than whatever -ub is set at.

Requirements

• I have read and agree with the contributing guidelines: Yes
• AI usage disclosure: Yes, used Fable 5.1 to understand the surrounding code and possible solutions.

[ggml-gh-bot]

Hi @jesdga95, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 2 open PRs.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[pwilkin]

/bot review

[ggml-gh-bot]

Automated code review

Review of common/speculative.cpp change (DFlash + M-RoPE media fix). I verified the diagnosis against the mtmd position code before judging the fix, and the root cause analysis is correct: for MTMD_POS_TYPE_MROPE, every image token gets the same section-0 position (pos.t = pos_0 in mtmd_image_tokens_get_decoder_pos, tools/mtmd/mtmd.cpp), so injecting a large image puts thousands of draft KV rows at one pinned position and the draft cache cannot reclaim cells while the position stays fixed. Skipping the injection is a reasonable minimal fix, and returning true (success) is the right contract. No blocking findings. Points that will come up in review:

Will slow the review

(point 1) The guard keys on the wrong model. is_mrope is computed from model_dft (line 1014), but the pinned positions come from the target's M-RoPE media encoding (mtmd-helper.cpp builds the batch from the target's decoder rope type). If the target is M-RoPE and the DFlash draft is not (NORM/NEOX rope), batch_in.pos[k] is still constant across the image, process() still injects all rows at one position, and the exact same cache exhaustion occurs - the guard never fires. Conversely, a M-RoPE draft with a non-M-RoPE target gets skipped even though injection would work. Consider storing the target's rope type as a member and gating on that, or gating on whether batch_in.pos actually advances within the batch (which is the true invariant you care about).

(point 2) The skip is broader than the comment claims. if (has_embeddings && is_mrope) skips all embedding batches, but audio chunks on M-RoPE targets use set_position_mrope_1d (tools/mtmd/mtmd-helper-common.h), where all four position rows advance sequentially - those inject fine and are now silently dropped, losing draft conditioning on audio. Either narrow the condition (image-only, or non-advancing positions only) or make the comment say "media" instead of "image".

(point 3) The new comment contradicts the comment six lines above it, which states that skipping embedding batches "leaves a hole in the draft's cache and the next injection fails to initialize". The new code now deliberately skips them and asserts "the draft can jump over the gap". If the hole is genuinely tolerable for DFlash, please update the older comment to carve out this case (or reference it); as written a reader gets two opposing claims about the same situation and cannot tell which is true.

Nits

(point 4) Comment precision, common/speculative.cpp:1105-1106: only the section-0 (text) position is shared across image tokens; the h/w sections vary per token. Also "never frees cells" is slightly overstated - the cells do free once the sequence position advances past the window; the failure is that one image can exceed the entire draft cache while the position is pinned. A one-line rewording would make the invariant exact.

(point 5) Follow-up worth checking after this lands: dp.n_past is set from slot.prompt.n_tokens() (token count, tools/server/server-context.cpp:3015) while the draft KV is seeded at rope positions; after media these two diverge (an image costs n_tokens cache rows but only max(nx, ny) positions). With the skip making media prompts reachable for DFlash, the noise blocks are now decoded at positions ahead of the seeded rows, and begin()'s pos_max < N-1 warning will fire on every media prompt. This is a draft-quality concern, not a crash, but it affects how well "the draft can jump over the gap" holds in practice - worth measuring accept rates after images before/after.

Otherwise the change is minimal, single-purpose, ASCII-clean, and in the right spot; disclosure is filled in.

This review was generated automatically by pi coding agent using zai-org/GLM-5.3. It may contain mistakes. Maintainers make the final call.

[ruixiang63]

This is a wrong fix, vision input works well from my testing.

[jesdga95]

@ruixiang63 I am able to reproduce this every single time in the checkpoint I posted, if you don't mind, what is your checkpoint? large ubatches might make this is a non issue (try with -ub 64, although I could reproduce with -ub 1024). I'm addressing the comments from the review either way.

There is also this second error that's the same root issue when -ub is big enough:

Sep 08 13:27:00 server llama-server[1712899]: [56291] 1.03.302.447 W decode: failed to find a memory slot for batch of size 512
Sep 08 13:27:00 server llama-server[1712899]: [56291] 1.03.302.450 E process: llama_decode(ctx_dft) failed rc=1 (n_tokens=512, offset=0)
Sep 08 13:27:00 server llama-server[1712899]: [56291] 1.03.302.476 E srv        decode: failed to process speculative batch
Sep 08 13:27:00 server llama-server[1712899]: [56291] 1.03.302.495 E srv  update_slots: decode() failed: failed to process speculative batch
Sep 08 13:27:00 server llama-server[1712899]: [56291] 1.03.302.499 E srv    send_error: task id = 590, error: decode() failed: failed to process speculative batch
Sep 08 13:27:00 server llama-server[1712899]: [56291] 1.03.302.502 I slot      release: id  0 | task 590 | stop processing: n_tokens = 3076, truncated = 0
Sep 08 13:27:00 server llama-server[1712899]: [56291] 1.03.302.520 W srv          stop: cancel task, id_task = 590

[ruixiang63]

Did you build the current master branch to test?

[jesdga95]

@ruixiang63 rebuilt on current master (5d806aa), same errors:

Small image, fails on the text prefill right after the image:

W decode: failed to find a memory slot for batch of size 512
E process: llama_decode(ctx_dft) failed rc=1 (n_tokens=512, offset=512)
E srv        decode: failed to process speculative batch

Larger image, fails while injecting the image itself, third chunk:

W decode: failed to find a memory slot for batch of size 512
E process: llama_decode(ctx_dft) failed rc=1 (n_tokens=512, offset=1024)
E post-decode callback failed
E slot   operator(): id  0 | task 11 | failed to decode mtmd chunk, idx = 1504, res = 1

Here's the image if it matters:

---

Setup: Qwen3.8-27B-UD-Q5_K_S + mmproj-qwen3.8-27b-F16 + Qwen3.8-27B-DFlash2-Q8_0, -ub 512:

[Qwen3.8-27B Q5]
device = CUDA0
model = /mnt/nvme/models/Qwen3.8-27B-UD-Q5_K_S.gguf
model-draft = /mnt/nvme/models/DFLASH2/Qwen3.8-27B-DFlash2-Q8_0.gguf
spec-type = draft-dflash
spec-draft-n-max = 7
spec-draft-ngl = all
spec-draft-device = CUDA0
fit-target = 2000
mmproj = /mnt/nvme/models/mmproj-qwen3.8-27b-F16.gguf
mmproj-device = CUDA1
n-gpu-layers = 99
jinja = true
flash-attn = on
parallel = 1
threads-http = 5
ctx-size = 262144
cache-type-k = q8_0
cache-type-v = q8_0
batch-size = 2048
ubatch-size = 512
temp = 1.0
top-p = 0.95
top-k = 20
min-p = 0.0
presence-penalty = 0.0

[ruixiang63]

I can't reproduce that. cc @ngxson to take a look.

[jesdga95]

@ruixiang63 I have a theory on why you might not be able to repro it, the z-lab DFlash2 for Qwen 3.8 27B has dflash.attention.sliding_window = 2048.

Tested on a new clean build from master (5d806aa):

• full-res image, 46 tokens of text: 500
• same image resized to 1000x1500 (~1457 tokens), 46 tokens of text: 200 - works!
• resized image behind ~3k tokens of text: 500 again, at offset 512

So full repro: use a high resolution image and send it as the first message with this dflash file, and -ub 512 OR add some text to the prompt (I'm using a harness so system instructions + tools). Hopefully this info helps.

[ruixiang63]

@ruixiang63 I have a theory on why you might not be able to repro it, the z-lab DFlash2 for Qwen 3.8 27B has dflash.attention.sliding_window = 2048.

Tested on a new clean build from master (5d806aa):

• full-res image, 46 tokens of text: 500
• same image resized to 1000x1500 (~1457 tokens), 46 tokens of text: 200 - works!
• resized image behind ~3k tokens of text: 500 again, at offset 512

So full repro: use a high resolution image and send it as the first message with this dflash file, and -ub 512 OR add some text to the prompt (I'm using a harness so system instructions + tools). Hopefully this info helps.

Thanks! Will give it a try.

[nouveau-nvc0]

I faced the same crash on 3× gfx1030 GPUs with ROCm. This fix prevents the crash for me, but DFlash performance appears to degrade significantly after vision input - TG drops to ~11 t/s, which is about the same as without DFlash enabled.

[jesdga95]

@nouveau-nvc0 thanks for testing. The slowdown after an image is a different bug which I was also able to reproduce.

One-line fix here: jesdga95@0c5dff206 (use pos_next() instead of n_tokens). I can't open a second PR until this one is resolved, but if you can try it on your ROCm setup that would be a useful checkpoint. It also affects MTP and Eagle3, so I need to do some further testing.

[ruixiang63]

Thanks for identifying this and proposing a fix. Vision inputs are not well supported by current speculative decoding logic, and the lack of vision-specific training can also reduce the acceptance rate for DFlash/DFlash2.

That said, I would prefer to avoid a hotfix, since the current solution in dflash is also a workaround. The proper long-term fix appears to be #24669. I suggest waiting for that PR to be merged and then verifying whether it resolves the issue definitively.

[nouveau-nvc0]

@jesdga95

I can't open a second PR until this one is resolved, but if you can try it on your ROCm setup that would be a useful checkpoint.

This fix worked. Reasoning performance is back to my normal 20–26 t/s (1 slot)

[ngxson]

/bot review

[ggml-gh-bot]

Automated code review

This follow-up review covers the changes since commit 7146cdc (the address PR feedback commit), which reworked the skip guard in common/speculative_impl_draft_dflash::process().

Previous review points - all addressed

(point 1) The wrong-model gating (is_mrope from model_dft) is gone. The guard now keys on the actual batch invariant (pos_pinned via batch_in.pos[beg] == batch_in.pos[end]), which I verified against the mtmd position code: for MTMD_POS_TYPE_MROPE images, section 0 is laid out first in the flat pos array and pos.t = pos_0 for every image token (mtmd_image_tokens_get_decoder_pos, tools/mtmd/mtmd.cpp:2471), so the check fires regardless of which side (target or draft) uses M-RoPE. A M-RoPE target with a non-M-RoPE draft is now covered, and a M-RoPE draft with a normal-positioned batch is no longer over-skipped. This also stays correct for chunked images: decode_embd_batch::get_view keeps the section-major layout, so each ubatch of a large image is independently detected as pinned.

(point 2) Audio and other 1d-media embedding batches (whose positions advance via set_position_mrope_1d) are no longer skipped - only genuinely pinned ranges are. The n_rows > 1 term keeps single-row pinned batches injectable, which is harmless for the cache.

(point 3) The contradictory comment is resolved - the old text about "leaves a hole ... and the next injection fails to initialize" is replaced with "embeddings are injected too, except the pinned ones skipped below", which now matches the code.

Nits

(point 4) common/speculative.cpp:1107-1110 (begin()): the pos_max < N-1 warning ("process() did not run on every prefill ubatch") will now fire on every prompt containing media, since skipped image rows are expected and deliberate. It reads as an error condition to users. Consider either leaving it (it is technically true - drafts do degrade) or making it clear media prompts are a known case, whichever fits the pos_next() follow-up discussed in the PR thread.

(point 5) Known follow-up, not for this PR: dp.n_past is token-count based while the draft KV is seeded at rope positions, so after a media prompt the noise blocks decode at positions ahead of the seeded rows. That is the post-image slowdown reported by @nouveau-nvc0, and the author already has a one-line pos_next() fix queued as a separate PR. Fine to keep this PR single-purpose; just noting it so a reviewer does not mistake the residual slowdown for a failure of this change.

The change remains minimal, single-purpose, ASCII-clean, and the new comment is concise and accurate (the injected rows really are all written at one draft position: the is_mrope branch sets all four sections to p, so "pins all its rows to one position" is exact). No blocking findings.

This review was generated automatically by pi coding agent using zai-org/GLM-5.3. It may contain mistakes. Maintainers make the final call.

[ngxson]

no strong opinion about this fix tbh, but I think it's fine as it's only at libcommon level

I'll revisit it after #24669 , expected to finalize later this month