Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

spec : fuse the DFlash encoder into the KV cache injection (#27310)

Browse Source 
* dflash : fuse the encoder into the KV injection decode

The encoder is a single fc + norm, but running it as a separate
llama_encode forced a device-to-host round trip of its output before the
injection decode could re-upload it, plus a second graph build per
round. Fold the encoder into the decoder's embd branch and feed the
target features directly to one llama_decode.

Assisted-by: Claude Fable

* nit

* Apply batched suggestions from code review

Co-authored-by: Ruixiang Wang <wangruixiang07@outlook.com>

* Fix missing references from renaming

---------

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
Co-authored-by: Ruixiang Wang <wangruixiang07@outlook.com>
This commit is contained in:
王金旭
2026-08-31 11:19:20 +02:00
committed by GitHub
co-authored by Ruixiang Wang Sigbjørn Skjæret
parent 2cdae802e4
commit 662a0b0121
3 changed files with 32 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/speculative.cpp
+6 -46
View File
@@ -941,9 +941,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices
uint32_t target_layer_ids_n = 0;
// scratch buffer for concatenated target features [n_tokens, n_embd_enc]
std::vector<float> features_buf;
common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq,
common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)
: common_speculative_impl(type, n_seq, params.draft.n_max)
@@ -1011,7 +1008,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
this->n_max = this->params.n_max;
batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq);
batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq);
batch_inject = llama_batch_init(llama_n_ubatch(ctx_dft), n_embd_enc, n_seq);
// embd batches on an M-RoPE draft need 4 position rows per token
is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE;
@@ -1137,58 +1134,21 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {
const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);
// gather this chunk's target features, interleaved by extract layer
features_buf.resize((size_t) n_chunk * n_embd_enc);
// gather target features per extract layer; the fused decode encodes and
// injects them into the K/V cache at the target positions
batch_inject.n_tokens = n_chunk;
for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
if (!layer) {
GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);
}
for (int32_t i = 0; i < n_chunk; ++i) {
float * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
float * dst = batch_inject.embd + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
const float * src = layer + (size_t) (i_batch_beg[seq_id] + offset + i) * n_embd_tgt;
std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));
}
}
// fuse extracted features through DFlash encoder
// M-RoPE drafts read 4 position rows per token from embd batches, so pass them explicitly
std::vector<llama_pos> enc_pos;
if (is_mrope) {
enc_pos.resize((size_t) 4 * n_chunk);
for (int32_t i = 0; i < n_chunk; ++i) {
const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i];
enc_pos[0 * n_chunk + i] = p;
enc_pos[1 * n_chunk + i] = p;
enc_pos[2 * n_chunk + i] = p;
enc_pos[3 * n_chunk + i] = 0;
}
}
llama_batch enc_batch = {
/*.n_tokens =*/ n_chunk,
/*.token =*/ nullptr,
/*.embd =*/ features_buf.data(),
/*.pos =*/ is_mrope ? enc_pos.data() : nullptr,
/*.n_seq_id =*/ nullptr,
/*.seq_id =*/ nullptr,
/*.logits =*/ nullptr,
};
int32_t rc = llama_encode(ctx_dft, enc_batch);
if (rc != 0) {
LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
__func__, rc, (int) n_chunk, (int) offset);
return false;
}
const float * inp_g = llama_get_embeddings_nextn(ctx_dft);
GGML_ASSERT(inp_g && "DFlash encoder produced no output.");
// inject the DFlash decoder K/V cache at the tokens' target positions
batch_inject.n_tokens = n_chunk;
std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float));
for (int32_t i = 0; i < n_chunk; ++i) {
const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i];
batch_inject.pos[i] = p;
@@ -1201,7 +1161,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
batch_inject.seq_id[i][0] = seq_id;
batch_inject.logits[i] = false;
}
rc = llama_decode(ctx_dft, batch_inject);
const int32_t rc = llama_decode(ctx_dft, batch_inject);
if (rc != 0) {
LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
__func__, rc, (int) n_chunk, (int) offset);