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.
This commit is contained in:
ruanrms
2026-08-10 00:15:04 +02:00
committed by Pedro Cuenca
parent f00ae69308
commit c8ee4945b1
+7 -1
View File
@@ -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;
}