Superminaren/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-04 02:37:27 +02:00
Code Issues 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 <[email protected]>

* Fix missing references from renaming

---------

Co-authored-by: Sigbjørn Skjæret <[email protected]>
Co-authored-by: Ruixiang Wang <[email protected]>
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);
src/llama-context.cpp
+3 -1
View File
@@ -1660,7 +1660,9 @@ int llama_context::decode(const llama_batch & batch_inp) {
const int64_t n_vocab = vocab.n_tokens();
const bool mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && batch_inp.embd;
const int64_t n_embd = mtp_embd ? hparams.n_embd_out() : hparams.n_embd_inp();
// DFlash embd batches carry the fused target features at the encoder input width
const bool dflash_embd = model.arch == LLM_ARCH_DFLASH && batch_inp.embd;
const int64_t n_embd = mtp_embd ? hparams.n_embd_out() : dflash_embd ? hparams.n_embd_inp_enc() : hparams.n_embd_inp();
// when computing embeddings, all tokens are output
const bool output_all = cparams.embeddings;
src/models/dflash.cpp
+23 -10
View File
@@ -257,9 +257,10 @@ std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const ll
template <>
ggml_tensor * llama_model_dflash::graph<true>::build_inp_embd_enc() const {
auto inp_target = std::make_unique<llm_graph_input_embd>(hparams.n_embd_inp_enc());
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
auto inp_target = std::make_unique<llm_graph_input_embd>(n_embd_inp);
inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp_enc(), n_tokens);
inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
ggml_set_input(inp_target->embd);
ggml_tensor * cur = inp_target->embd;
@@ -567,6 +568,7 @@ static void build_dflash2_selector(llm_graph_context & g, const llama_model & mo
// * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens
template <>
llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
@@ -602,16 +604,21 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
// KV cache injection
if (ubatch.embd) {
auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
ggml_set_input(inp->embd);
ggml_tensor * inp_g = inp->embd;
cb(inp_g, "inp_g_embeddings", -1);
ggml_tensor * inp_target = inp->embd;
cb(inp_target, "inp_target_features", -1);
res->add_input(std::move(inp));
// fuse the target features through the encoder
ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
inp_g = build_norm(inp_g, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
cb(inp_g, "inp_g_embeddings", -1);
for (int il = 0; il < n_layer; ++il) {
const auto & layer = model.layers[il];
@@ -823,6 +830,7 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
// * token batch -> noise block through 3 full DSV4 stages (hc + MLA + MoE), markov + confidence heads
llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_graph_params & params) :
llama_model_deepseek4::graph(params) {
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
const int64_t n_embd_head = hparams.n_embd_head_k();
const int64_t n_embd_head_rope = hparams.n_rot();
const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
@@ -833,16 +841,21 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
// KV cache injection: fused target features from the encoder
if (ubatch.embd) {
auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
ggml_set_input(inp->embd);
ggml_tensor * inp_g = inp->embd;
cb(inp_g, "inp_g_embeddings", -1);
ggml_tensor * inp_target = inp->embd;
cb(inp_target, "inp_target_features", -1);
res->add_input(std::move(inp));
// fuse the target features through the encoder
ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
inp_g = build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1);
cb(inp_g, "inp_g_embeddings", -1);
for (int il = 0; il < n_layer; ++il) {
const auto & layer = model.layers[il];