eagle3: add recurrent verification state support

Author: wudixiaoniu

convert_hf_to_gguf.py

@@ -2748,6 +2748,7 @@ def prepare_tensors(self):
     "LlavaForConditionalGeneration",
     "VoxtralForConditionalGeneration",
     "LlamaForCausalLMEagle3",
+    "Eagle3LlamaForCausalLM",
     "Eagle3Speculator",
     "Eagle3DraftModel",
     "IQuestCoderForCausalLM",
@@ -2786,7 +2787,9 @@ def __init__(self, *args, **kwargs):
                 target_config = json.load(f)
 
             # EAGLE3 extract_layers
-            target_num_layers = target_config["num_hidden_layers"]
+            # Support nested config (e.g. VL/MoE models with text_config)
+            tc = target_config.get("text_config", target_config)
+            target_num_layers = tc["num_hidden_layers"]
             extract_layers = [2, target_num_layers // 2, target_num_layers - 3]
             logger.info(f"EAGLE3: extract_layers = {extract_layers} (target model has {target_num_layers} layers)")
             self.gguf_writer.add_array(f"{self.gguf_writer.arch}.extract_layers", extract_layers)
@@ -2796,7 +2799,7 @@ def __init__(self, *args, **kwargs):
                 target_hidden_size = eagle3_raw_config["target_hidden_size"]
                 logger.info(f"EAGLE3: target_hidden_size = {target_hidden_size} (from EAGLE3 config)")
             else:
-                target_hidden_size = target_config["hidden_size"]
+                target_hidden_size = tc["hidden_size"]
                 logger.info(f"EAGLE3: target_hidden_size = {target_hidden_size} (from target model config)")
             self.gguf_writer.add_uint32(f"{self.gguf_writer.arch}.target_hidden_size", target_hidden_size)
 

examples/speculative-simple/speculative-simple.cpp

@@ -6,6 +6,7 @@
 #include "llama.h"
 #include "chat.h"
 
+#include <algorithm>
 #include <clocale>
 #include <cstdio>
 #include <cstring>
@@ -33,6 +34,9 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
+    params.kv_unified = true;
+    params.n_parallel = std::max(params.n_parallel, params.speculative.n_max + 2);
+
     // init llama.cpp
     llama_backend_init();
     llama_numa_init(params.numa);
@@ -111,7 +115,7 @@ int main(int argc, char ** argv) {
     int n_predict = 0;
     int n_drafted = 0;
     int n_accept  = 0;
-
+    int draft_len_slot[10] = {0};
     // used to determine end of generation
     bool has_eos = false;
 
@@ -150,6 +154,7 @@ int main(int argc, char ** argv) {
     llama_token id_last;
     llama_tokens prompt_tgt;
     int n_past;
+    bool pending_first_token_print = false;
 
     // TODO: simplify
     if (params.speculative.eagle3) {
@@ -158,7 +163,7 @@ int main(int argc, char ** argv) {
 
         id_last = common_sampler_sample(smpl, ctx_tgt, -1);
         common_sampler_accept(smpl, id_last, true);
-        LOG("%s", common_token_to_piece(ctx_tgt, id_last).c_str());
+        pending_first_token_print = true;
         n_predict++;
 
         // all tokens currently in the target context
@@ -179,14 +184,18 @@ int main(int argc, char ** argv) {
         n_past = inp.size() - 1;
     }
 
-    // init the speculator
     const auto & params_spec = params.speculative;
-
     struct common_speculative * spec = common_speculative_init(params.speculative, ctx_tgt);
 
     common_speculative_begin(spec, prompt_tgt);
 
+    if (pending_first_token_print) {
+        LOG("%s", common_token_to_piece(ctx_tgt, id_last).c_str());
+        pending_first_token_print = false;
+    }
+
     llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1);
+    llama_memory_t mem_tgt = llama_get_memory(ctx_tgt);
 
     const auto t_enc_end = ggml_time_us();
 
@@ -200,37 +209,34 @@ int main(int argc, char ** argv) {
         // offloaded to a remote device. it doesn't even have to be based on an LLM. instead, it can provide tokens
         // from a cache or lookup tables.
         //
+
+
         llama_tokens draft = common_speculative_draft(spec, params_spec, prompt_tgt, id_last);
 
         //LOG_DBG("draft: %s\n", string_from(ctx_dft, draft).c_str());
 
-        // always have a token to evaluate from before - id_last
-        common_batch_clear(batch_tgt);
-        common_batch_add  (batch_tgt, id_last, n_past++, { 0 }, true);
+        // do not waste time on small drafts
+        if (draft.size() < (size_t) params_spec.n_min) {
+            draft.clear();
+        }
 
-        // evaluate the target model on [id_last, draft0, draft1, ..., draftN-1]
-        {
-            // do not waste time on small drafts
-            if (draft.size() < (size_t) params_spec.n_min) {
-                draft.clear();
-            }
+        const int n_verify = (int) draft.size() + 1;
 
-            for (size_t i = 0; i < draft.size(); ++i) {
-                common_batch_add(batch_tgt, draft[i], n_past + i, { 0 }, true);
-            }
+        common_batch_clear(batch_tgt);
 
-            //LOG_DBG("target batch: %s\n", string_from(ctx_tgt, batch_tgt).c_str());
+        if (!llama_memory_eagle3_recurrent_begin(mem_tgt, 0, n_verify, n_past)) {
+            LOG_ERR("%s: failed to reserve EAGLE3 recurrent verification slots for depth %d\n", __func__, n_verify);
+            return 1;
+        }
 
-            llama_decode(ctx_tgt, batch_tgt);
+        common_batch_add(batch_tgt, id_last, n_past++, { 0 }, true);
+
+        for (size_t i = 0; i < draft.size(); ++i) {
+            common_batch_add(batch_tgt, draft[i], n_past + i, { 0 }, true);
         }
 
-        // sample from the full target batch and return the accepted tokens based on the target sampler
-        //
-        // for each token to be accepted, the sampler would have to sample that same token
-        // in such cases, instead of decoding the sampled token as we normally do, we simply continue with the
-        // available logits from the batch and sample the next token until we run out of logits or the sampler
-        // disagrees with the draft
-        //
+        llama_decode(ctx_tgt, batch_tgt);
+
         const auto ids = common_sampler_sample_and_accept_n(smpl, ctx_tgt, draft);
 
         //LOG_DBG("ids: %s\n", string_from(ctx_tgt, ids).c_str());
@@ -240,8 +246,11 @@ int main(int argc, char ** argv) {
         n_past    += ids.size() - 1;
         n_drafted += draft.size(); // note: we ignore the discarded small drafts
         n_accept  += ids.size() - 1;
+        draft_len_slot[ids.size()-1] += 1;
         n_predict += ids.size();
 
+        common_speculative_accept(spec, ids.size() - 1);
+
         // process the accepted tokens and update contexts
         //
         // this is the standard token post-processing that we normally do
@@ -268,10 +277,11 @@ int main(int argc, char ** argv) {
 
         LOG_DBG("accepted %d/%d draft tokens, the last target token is: (%d)\n", (int) ids.size() - 1, (int) draft.size(), id_last);
 
+        GGML_ASSERT(llama_memory_eagle3_recurrent_promote(mem_tgt, 0, ids.size()));
+
         {
             LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past);
-
-            llama_memory_seq_rm(llama_get_memory(ctx_tgt), 0, n_past, -1);
+            llama_memory_seq_rm(mem_tgt, 0, n_past, -1);
         }
 
         if ((params.n_predict >= 0 && n_predict > params.n_predict) || has_eos) {
@@ -293,6 +303,13 @@ int main(int argc, char ** argv) {
     LOG_INF("n_predict = %d\n", n_predict);
     LOG_INF("n_drafted = %d\n", n_drafted);
     LOG_INF("n_accept  = %d\n", n_accept);
+
+    
+    for(int j=0;j<10;j++){
+        LOG_INF("draft_len_slot[%d]  = %d\n",j,draft_len_slot[j]);
+        
+    }
+    
     LOG_INF("accept    = %.3f%%\n", 100.0f * n_accept / n_drafted);
 
     LOG_INF("\n");

examples/speculative/speculative.cpp

@@ -188,6 +188,7 @@ int main(int argc, char ** argv) {
     int n_predict = 0;
     int n_drafted = 0;
     int n_accept  = 0;
+    int n_seq_accept[10] = {0};
 
     int n_past_tgt = inp.size();
     int n_past_dft = inp.size();
@@ -381,13 +382,14 @@ int main(int argc, char ** argv) {
                     token_str = common_token_to_piece(ctx_tgt, token_id);
 
                     for (int s = 0; s < n_seq_dft; ++s) {
+                        // n_seq_accept
                         if (!drafts[s].active) {
                             continue;
                         }
 
                         if (i_dft < (int) drafts[s].tokens.size() && token_id == drafts[s].tokens[i_dft]) {
                             LOG_DBG("the sampled target token matches the %dth drafted token of sequence %d (%d, '%s') - accepted\n", i_dft, s, token_id, token_str.c_str());
-
+                            n_seq_accept[i_dft] += 1;
                             s_keep = s;
                             accept = true;
                         } else {
@@ -625,6 +627,13 @@ int main(int argc, char ** argv) {
     LOG_INF("n_predict = %d\n", n_predict);
     LOG_INF("n_drafted = %d\n", n_drafted);
     LOG_INF("n_accept  = %d\n", n_accept);
+    LOG_INF("n_accept  = %d\n", n_accept);
+    for (int i = 0; i < 10; i++)
+    {
+        LOG_INF("n_seq_accept[%d]-accept_len-%d\n", i,n_seq_accept[i]);
+        
+    }
+    
     LOG_INF("accept    = %.3f%%\n", 100.0f * n_accept / n_drafted);
 
     LOG_INF("\n");

include/llama.h

@@ -777,6 +777,26 @@ extern "C" {
     // Check if the memory supports shifting
     LLAMA_API bool llama_memory_can_shift(llama_memory_t mem);
 
+    // Begin a transient EAGLE3 recurrent verification round for a live sequence.
+    // The round reserves one recurrent state cell per verification depth, starting at position p0.
+    LLAMA_API bool llama_memory_eagle3_recurrent_begin(
+            llama_memory_t mem,
+              llama_seq_id live_seq_id,
+                 uint32_t  n_depth,
+                llama_pos  p0);
+
+    // Promote the accepted verification depth back to the live sequence.
+    // Returns false if no active EAGLE3 recurrent round exists for the sequence.
+    LLAMA_API bool llama_memory_eagle3_recurrent_promote(
+            llama_memory_t mem,
+              llama_seq_id live_seq_id,
+                 uint32_t  depth);
+
+    // Clear any transient EAGLE3 recurrent verification state for the live sequence.
+    LLAMA_API void llama_memory_eagle3_recurrent_clear(
+            llama_memory_t mem,
+              llama_seq_id live_seq_id);
+
     //
     // State / sessions
     //

src/llama-context.cpp

@@ -2193,11 +2193,20 @@ ggml_cgraph * llama_context::graph_reserve(
 
     // EAGLE3: auto-detect encoder (embeddings+no target_model) or decoder (has target_model)
     llm_graph_type gtype = LLM_GRAPH_TYPE_DEFAULT;
+    // if (model.arch == LLM_ARCH_EAGLE3) {
+    //     if (cparams.embeddings && model.target_tok_embd == nullptr) {
+    //         gtype = LLM_GRAPH_TYPE_ENCODER;
+    //     } else if (model.target_tok_embd != nullptr) {
+    //         gtype = LLM_GRAPH_TYPE_DECODER;
+    //     }
+    // }
+
+    // 修改后：
     if (model.arch == LLM_ARCH_EAGLE3) {
-        if (cparams.embeddings && model.target_tok_embd == nullptr) {
-            gtype = LLM_GRAPH_TYPE_ENCODER;
-        } else if (model.target_tok_embd != nullptr) {
+        if (model.target_tok_embd != nullptr) {
             gtype = LLM_GRAPH_TYPE_DECODER;
+        } else {
+            gtype = LLM_GRAPH_TYPE_ENCODER;
         }
     }
     const auto gparams = graph_params(res, ubatch, mctx, gtype);
@@ -3468,6 +3477,39 @@ bool llama_memory_can_shift(llama_memory_t mem) {
     return mem->get_can_shift();
 }
 
+bool llama_memory_eagle3_recurrent_begin(
+        llama_memory_t mem,
+          llama_seq_id live_seq_id,
+             uint32_t  n_depth,
+            llama_pos  p0) {
+    if (!mem) {
+        return false;
+    }
+
+    return mem->eagle3_recurrent_round_begin(live_seq_id, n_depth, p0);
+}
+
+bool llama_memory_eagle3_recurrent_promote(
+        llama_memory_t mem,
+          llama_seq_id live_seq_id,
+             uint32_t  depth) {
+    if (!mem) {
+        return false;
+    }
+
+    return mem->eagle3_recurrent_round_promote(live_seq_id, depth);
+}
+
+void llama_memory_eagle3_recurrent_clear(
+        llama_memory_t mem,
+          llama_seq_id live_seq_id) {
+    if (!mem) {
+        return;
+    }
+
+    mem->eagle3_recurrent_round_clear(live_seq_id);
+}
+
 // llama state API
 
 // deprecated

src/llama-memory-hybrid.cpp

@@ -176,6 +176,18 @@ std::map<ggml_backend_buffer_type_t, size_t> llama_memory_hybrid::memory_breakdo
     return mb;
 }
 
+bool llama_memory_hybrid::eagle3_recurrent_round_begin(llama_seq_id live_seq_id, uint32_t n_depth, llama_pos p0) {
+    return mem_recr->eagle3_recurrent_round_begin(live_seq_id, n_depth, p0);
+}
+
+bool llama_memory_hybrid::eagle3_recurrent_round_promote(llama_seq_id live_seq_id, uint32_t depth) {
+    return mem_recr->eagle3_recurrent_round_promote(live_seq_id, depth);
+}
+
+void llama_memory_hybrid::eagle3_recurrent_round_clear(llama_seq_id live_seq_id) {
+    mem_recr->eagle3_recurrent_round_clear(live_seq_id);
+}
+
 void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
     if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
         mem_attn->state_write(io, seq_id, flags);

src/llama-memory-hybrid.h

@@ -70,6 +70,17 @@ class llama_memory_hybrid : public llama_memory_i {
 
     std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const override;
 
+    bool eagle3_recurrent_round_begin(
+            llama_seq_id live_seq_id,
+               uint32_t  n_depth,
+              llama_pos  p0) override;
+
+    bool eagle3_recurrent_round_promote(
+            llama_seq_id live_seq_id,
+               uint32_t  depth) override;
+
+    void eagle3_recurrent_round_clear(llama_seq_id live_seq_id) override;
+
     // state write/load
 
     void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;