Merge pull request #220 from TheTom/tom/pr-219-full-rebase

dflash: port #219 onto feature (rebased full CL)

Author: TheTom

common/arg.cpp

@@ -3600,6 +3600,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE"));
+    add_opt(common_arg(
+        {"--eagle3"},
+        "use EAGLE3 speculative decoding with the draft model",
+        [](common_params & params) {
+            params.speculative.draft.eagle3 = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--dflash"},
+        "use DFlash speculative decoding with the draft model",
+        [](common_params & params) {
+            params.speculative.draft.dflash = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"-cd", "--ctx-size-draft"}, "N",
+        string_format("size of the prompt context for the draft model (default: %d, 0 = loaded from model)", params.speculative.draft.n_ctx),
+        [](common_params & params, int value) {
+            params.speculative.draft.n_ctx = value;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CTX"));
 
     add_opt(common_arg(
         {"--spec-draft-n-max"}, "N",

common/common.h

@@ -325,6 +325,11 @@ struct common_params_speculative_draft {
     std::vector<ggml_backend_dev_t> devices; // devices to use for offloading
 
     std::vector<llama_model_tensor_buft_override> tensor_buft_overrides;
+
+    bool    eagle3       = false; // use EAGLE3 speculative decoding
+    bool    dflash       = false; // use DFlash speculative decoding
+    int32_t n_ctx        = 0;  // draft context size
+
 };
 
 struct common_params_speculative_ngram_mod {

common/speculative.cpp

@@ -835,6 +835,14 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
 
     // scratch buffer for concatenated target features [n_tokens, n_embd_enc]
     std::vector<float> features_buf;
+    // Stashed encoder outputs for deferred KV cache injection
+    struct StashedG {
+        llama_pos pos;
+        std::vector<float> data; // n_embd_dec floats
+    };
+    std::vector<std::vector<StashedG>> stashed;  // [n_seq][...]
+    static constexpr int32_t MAX_STASH = 64;
+    bool m_use_deferred = true;
 
     common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq)
         : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, n_seq)
@@ -869,7 +877,17 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
         LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);
         LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u\n", __func__, block_size, mask_token_id, target_layer_ids_n);
 
-        // DFlash input is [id_last, <mask> * (block_size-1)], so it can draft at most block_size-1 tokens per step
+        // Detect target model type: Gemma4 models regress with deferred injection (low acceptance rate)
+        {
+            char model_desc[128] = {};
+            llama_model_desc(model_tgt, model_desc, sizeof(model_desc));
+            if (strstr(model_desc, "gemma")) {
+                m_use_deferred = false;
+                LOG_INF("%s: - deferred_kv_injection=0 (disabled for %s)\n", __func__, model_desc);
+            } else {
+                LOG_INF("%s: - deferred_kv_injection=1 (enabled for %s)\n", __func__, model_desc);
+            }
+        }
         if (this->params.n_max > block_size - 1 || this->params.n_min > block_size - 1) {
             LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained DFlash block size %d -- clamping to %d\n",
                     __func__, this->params.n_max, this->params.n_min, block_size, block_size - 1);
@@ -880,6 +898,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
         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);
 
+        stashed.resize(n_seq);
+
         smpls.resize(n_seq);
         for (auto & s : smpls) {
             common_params_sampling sparams;
@@ -889,13 +909,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
             s.reset(common_sampler_init(model_dft, sparams));
         }
 
-        // turn on extraction of the target layers' input embeddings
+        // Enable extraction of target model layer outputs for DFlash encoder
         for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
             llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
         }
 
+        // Enable nextn (encoder) output on the draft context for DFlash decoder K/V injection
         llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
-        llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention
     }
 
     ~common_speculative_impl_draft_dflash() override {
@@ -996,31 +1016,76 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                 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) {
-                    batch_inject.pos[i]       = batch_in.pos[i_batch_beg[seq_id] + offset + i];
-                    batch_inject.n_seq_id[i]  = 1;
-                    batch_inject.seq_id[i][0] = seq_id;
-                    batch_inject.logits[i]    = false;
-                }
-                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);
-                    return false;
+                if (m_use_deferred) {
+                    // stash encoder output for deferred KV cache injection
+                    auto & stash = stashed[seq_id];
+                    for (int32_t i = 0; i < n_chunk; ++i) {
+                        StashedG sg;
+                        sg.pos = batch_in.pos[i_batch_beg[seq_id] + offset + i];
+                        sg.data.resize(n_embd_dec);
+                        std::copy(inp_g + (size_t) i * n_embd_dec, inp_g + (size_t) (i + 1) * n_embd_dec, sg.data.begin());
+                        stash.push_back(std::move(sg));
+                    }
+                    // auto-flush to prevent unbounded accumulation
+                    if ((int32_t) stash.size() >= MAX_STASH) {
+                        flush_injection(seq_id);
+                    }
+                } else {
+                    // immediate injection (original per-chunk behavior)
+                    auto & stash = stashed[seq_id];
+                    for (int32_t i = 0; i < n_chunk; ++i) {
+                        StashedG sg;
+                        sg.pos = batch_in.pos[i_batch_beg[seq_id] + offset + i];
+                        sg.data.resize(n_embd_dec);
+                        std::copy(inp_g + (size_t) i * n_embd_dec, inp_g + (size_t) (i + 1) * n_embd_dec, sg.data.begin());
+                        stash.push_back(std::move(sg));
+                    }
+                    flush_injection(seq_id);
                 }
             }
         }
 
         return true;
     }
 
+    // Batch-inject all stashed encoder outputs into the draft decoder's KV cache
+    void flush_injection(llama_seq_id seq_id) {
+        auto & stash = stashed[seq_id];
+        if (stash.empty()) return;
+
+        const int32_t n = (int32_t) stash.size();
+        batch_inject.n_tokens = n;
+        for (int32_t i = 0; i < n; ++i) {
+            std::memcpy(batch_inject.embd + (size_t) i * n_embd_dec,
+                        stash[i].data.data(), (size_t) n_embd_dec * sizeof(float));
+            batch_inject.pos[i]       = stash[i].pos;
+            batch_inject.n_seq_id[i]  = 1;
+            batch_inject.seq_id[i][0] = seq_id;
+            batch_inject.logits[i]    = false;
+        }
+
+        auto * ctx_dft = params.ctx_dft;
+        int32_t rc = llama_decode(ctx_dft, batch_inject);
+        if (rc != 0) {
+            LOG_ERR("%s: flush_injection llama_decode(ctx_dft) failed rc=%d (n_tokens=%d)\n",
+                    __func__, rc, n);
+        }
+
+        stash.clear();
+    }
+
     void draft(common_speculative_draft_params_vec & dparams) override {
         auto & ctx_dft = params.ctx_dft;
 
+        if (m_use_deferred) {
+            // flush all stashed KV injections before drafting
+            for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+                if (dparams[seq_id].drafting) {
+                    flush_injection(seq_id);
+                }
+            }
+        }
+
         common_batch_clear(batch);
 
         // build one batch holding every drafting sequence's noise block into a single decode)
@@ -2082,6 +2147,13 @@ common_speculative * common_speculative_init(common_params_speculative & params,
         bool has_mtp          = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_MTP))    && params.draft.ctx_dft != nullptr;
         bool has_draft_dflash = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)) && params.draft.ctx_dft != nullptr;
 
+        // If --dflash or --eagle3 flags are set, enable the corresponding type
+        if (!has_draft_dflash && params.draft.dflash && params.draft.ctx_dft != nullptr) {
+            has_draft_dflash = true;
+        }
+        if (!has_draft_eagle3 && params.draft.eagle3 && params.draft.ctx_dft != nullptr) {
+            has_draft_eagle3 = true;
+        }
 
 
         bool has_ngram_cache   = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_NGRAM_CACHE));

examples/speculative-simple/speculative-simple.cpp

@@ -75,6 +75,8 @@ int main(int argc, char ** argv) {
         }
 
         auto cparams = common_context_params_to_llama(params_dft);
+        // DFlash/EAGLE3 draft models may reuse the target model's embeddings/output
+        cparams.ctx_other = ctx_tgt;
         ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams));
 
         params.speculative.draft.ctx_tgt = ctx_tgt;