diff --git a/common/speculative.cpp b/common/speculative.cpp index d34d1c9c59..851a47b9a5 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -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 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 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); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9aed801332..a6d2f264fc 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -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; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 037723a902..036bcc14aa 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -257,9 +257,10 @@ std::unique_ptr llama_model_dflash::build_arch_graph(const ll template <> ggml_tensor * llama_model_dflash::graph::build_inp_embd_enc() const { - auto inp_target = std::make_unique(hparams.n_embd_inp_enc()); + const int64_t n_embd_inp = hparams.n_embd_inp_enc(); + auto inp_target = std::make_unique(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::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::graph(const llama_model & model, const llm_gra // KV cache injection if (ubatch.embd) { - auto inp = std::make_unique(n_embd); + auto inp = std::make_unique(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::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(n_embd); + auto inp = std::make_unique(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];