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
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];