From c8ee4945b1d9a1f1230a618dc111b4063aa9115f Mon Sep 17 00:00:00 2001 From: ruanrms Date: Sun, 9 Aug 2026 14:16:51 -0700 Subject: [PATCH] DFlash: seed the draft KV cache from multimodal embedding batches `common_speculative_impl_draft_dflash::process()` returned early on any batch carrying embeddings, so an image prefill never had its target-layer features fused through the DFlash encoder and injected into the draft's KV cache. That left a hole spanning the image's positions, and the next injection at a post-image position failed to initialize its batch: ``` decoding image batch 1/1, n_tokens_batch = 256 decode: failed to initialize batch llama_decode: failed to decode, ret = -1 process: llama_decode(ctx_dft) failed rc=-1 (n_tokens=17, offset=0) srv decode: failed to process speculative batch ``` Every image request with `--spec-type draft-dflash` failed with HTTP 500. Text-only was unaffected, since those batches carry token ids and were let through. Restore the earlier condition, which admits a batch that is either tokens or embeddings and skips only the degenerate neither/both cases. The rest of `process()` is already layout-agnostic -- it gathers features via `llama_get_embeddings_layer_inp()` and indexes `batch_in.pos[]` / `batch_in.seq_id[]`, none of which assume token ids -- so this is the whole fix. Validated against `muse-glimmer-30B-bf16.gguf` + `mmproj-muse-glimmer-30B-bf16.gguf` + a DFlash draft head, on an image describe-the-shapes request: - before: HTTP 500, `failed to process speculative batch` - after: HTTP 200, draft acceptance 0.34012 (167 accepted / 491 generated), mean len 3.04 Output equivalence holds, which is the property that matters: at temperature 0 the drafted response is byte-identical to the same request served with no draft attached (1213/1213 chars), so the draft is drafting correctly through the image context rather than merely not crashing. --- common/speculative.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 5653a90b88..4d1f6c25ca 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1032,7 +1032,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { return true; } - if (batch_in.token == nullptr || batch_in.embd != nullptr) { + // Target prefill may contain token IDs or multimodal embeddings. Both + // produce the target-layer features used to seed the draft KV cache, so + // skipping the embedding batches leaves a hole in the draft's cache and + // the next injection fails to initialize. + const bool has_tokens = batch_in.token != nullptr; + const bool has_embeddings = batch_in.embd != nullptr; + if (has_tokens == has_embeddings) { return true; }