mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-18 16:55:05 +02:00
feat: add EAGLE3 speculative decoding support
EAGLE3 is an encoder-decoder based speculative decoding method: - Extracts features from target model at specific layers - Uses feature fusion layer to compress target features - Generates draft tokens with single-layer decoder - Maps draft vocabulary to target vocabulary via d2t tensor Key changes: - Add LLM_ARCH_EAGLE3 architecture - Add EAGLE3 encoder/decoder graph (src/models/eagle3.cpp) - Add feature extraction from target model layers - Add g_embeddings handling for decoder input - Add GGML_TENSOR_FLAG_SYNC for GPU synchronization - Add --eagle3 flag for speculative-simple example - Add EAGLE3 model conversion in convert_hf_to_gguf.py
This commit is contained in:
ruixiangw
@@ -22,6 +22,11 @@ struct common_speculative {
|
||||
llama_tokens prompt_dft;
|
||||
bool vocab_dft_compatible = true; // whether retokenization is needed
|
||||
std::map<std::string, std::string> tgt_dft_replacements = {};
|
||||
|
||||
// EAGLE3 specific
|
||||
struct llama_context * eagle3_encoder = nullptr;
|
||||
struct llama_context * eagle3_decoder = nullptr;
|
||||
int32_t eagle3_n_past = 0; // number of verified positions in decoder KV cache
|
||||
};
|
||||
|
||||
struct common_speculative * common_speculative_init(
|
||||
@@ -74,6 +79,35 @@ struct common_speculative * common_speculative_init(
|
||||
return result;
|
||||
}
|
||||
|
||||
struct common_speculative * common_speculative_init_eagle3(
|
||||
struct llama_context * ctx_tgt,
|
||||
struct llama_context * ctx_encoder,
|
||||
struct llama_context * ctx_decoder) {
|
||||
|
||||
auto * result = new common_speculative {
|
||||
/* .ctx_tgt = */ ctx_tgt,
|
||||
/* .ctx_dft = */ nullptr, // Not used for EAGLE3
|
||||
/* .smpl = */ nullptr,
|
||||
/* .batch = */ llama_batch_init(llama_n_batch(ctx_decoder), 0, 1),
|
||||
/* .prompt_dft = */ {},
|
||||
/* .vocab_dft_compatible = */ true, // EAGLE3 uses same vocab
|
||||
/* .tgt_dft_replacements = */ {},
|
||||
/* .eagle3_encoder = */ ctx_encoder,
|
||||
/* .eagle3_decoder = */ ctx_decoder,
|
||||
};
|
||||
|
||||
// Initialize sampler for EAGLE3 decoder
|
||||
{
|
||||
common_params_sampling params;
|
||||
params.no_perf = false;
|
||||
params.top_k = 10; // set 1 for greedy sampling (argmax) to match vLLM's default behavior but >1 always gets higher acceptance rate for eagle3
|
||||
params.samplers = { COMMON_SAMPLER_TYPE_TOP_K };
|
||||
result->smpl = common_sampler_init(llama_get_model(ctx_decoder), params);
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
void common_speculative_free(struct common_speculative * spec) {
|
||||
if (spec == nullptr) {
|
||||
return;
|
||||
@@ -81,6 +115,14 @@ void common_speculative_free(struct common_speculative * spec) {
|
||||
|
||||
common_sampler_free(spec->smpl);
|
||||
|
||||
// EAGLE3 cleanup
|
||||
if (spec->eagle3_encoder) {
|
||||
llama_free(spec->eagle3_encoder);
|
||||
}
|
||||
if (spec->eagle3_decoder) {
|
||||
llama_free(spec->eagle3_decoder);
|
||||
}
|
||||
|
||||
llama_batch_free(spec->batch);
|
||||
|
||||
delete spec;
|
||||
@@ -181,12 +223,169 @@ static std::string replace_to_tgt(
|
||||
return result;
|
||||
}
|
||||
|
||||
// EAGLE3 Draft Generation with KV Cache Reuse
|
||||
//
|
||||
// ============================================================================
|
||||
// EXAMPLE: Two rounds of speculative decoding
|
||||
// ============================================================================
|
||||
//
|
||||
// ROUND 1 (Initial):
|
||||
// Prompt: [t0, t1, t2, t3, t4], target generates t5
|
||||
// prompt_tgt = [t0, t1, t2, t3, t4], id_last = t5 (GENERATED)
|
||||
// n = 5, n_past = 0, n_new = 5
|
||||
//
|
||||
// Step 1: Encoder
|
||||
// features: [f0, f1, f2, f3, f4] → g_embeddings: [g0, g1, g2, g3, g4]
|
||||
//
|
||||
// Step 2: Decoder batch (positions 0-4)
|
||||
// tokens: [t1, t2, t3, t4, t5] ← prompt[1:] + id_last
|
||||
// g_embd: [g0, g1, g2, g3, g4]
|
||||
// positions: [0, 1, 2, 3, 4 ]
|
||||
// → KV cache: [0, 1, 2, 3, 4]
|
||||
// → sample d1 from logits[4]
|
||||
//
|
||||
// Step 3: Autoregressive (positions 5, 6, ...)
|
||||
// pos 5: token=d1, g_embd=prenorm[4] → KV cache: [0,1,2,3,4,5] → d2
|
||||
// pos 6: token=d2, g_embd=prenorm → KV cache: [0,1,2,3,4,5,6] → d3
|
||||
//
|
||||
// Output: [d1, d2, d3]
|
||||
// Update: n_past = 5 (verified positions from batch decode)
|
||||
//
|
||||
// ROUND 2 (assuming d1 accepted, d2/d3 rejected):
|
||||
// prompt_tgt = [t0, t1, t2, t3, t4, t5, d1], id_last = t6 (new target output)
|
||||
// n = 7, n_past = 5, n_new = 2
|
||||
//
|
||||
// Step 1: Clear KV cache [5, inf) - remove draft positions
|
||||
// KV cache: [0, 1, 2, 3, 4] (reuse from round 1!)
|
||||
//
|
||||
// Step 2: Encoder (only new tokens)
|
||||
// features: [f5, f6] → g_embeddings: [g5, g6]
|
||||
//
|
||||
// Step 3: Decoder batch (only new positions 5-6)
|
||||
// tokens: [d1, t6] (prompt_tgt[6], id_last)
|
||||
// g_embd: [g5, g6]
|
||||
// positions: [5, 6 ]
|
||||
// → KV cache: [0,1,2,3,4] + [5,6] = [0,1,2,3,4,5,6]
|
||||
// → sample d1' from logits[1] (last position in batch)
|
||||
//
|
||||
// Step 4: Autoregressive...
|
||||
//
|
||||
// ============================================================================
|
||||
//
|
||||
// Key insight: Decoder KV cache stores K/V computed from (tok_embd + g_embd).
|
||||
// For verified positions, both tok_embd and g_embd are fixed (encoder output),
|
||||
// so KV cache can be reused. Draft positions use prenorm as g_embd, which
|
||||
// differs from encoder output, so they must be cleared and recomputed.
|
||||
//
|
||||
static llama_tokens gen_eagle3_draft(
|
||||
struct common_speculative * spec,
|
||||
struct common_speculative_params params,
|
||||
const llama_tokens & prompt_tgt,
|
||||
llama_token id_last) {
|
||||
|
||||
auto * ctx_tgt = spec->ctx_tgt;
|
||||
auto * ctx_encoder = spec->eagle3_encoder;
|
||||
auto * ctx_decoder = spec->eagle3_decoder;
|
||||
auto * smpl = spec->smpl;
|
||||
auto & batch = spec->batch;
|
||||
|
||||
const int n_embd = llama_model_n_embd(llama_get_model(ctx_encoder));
|
||||
const int n = (int)prompt_tgt.size();
|
||||
const int n_new = n - spec->eagle3_n_past;
|
||||
|
||||
GGML_ASSERT(n >= 1 && "prompt_tgt is empty");
|
||||
GGML_ASSERT(n_new >= 1 && "must have at least 1 new token");
|
||||
|
||||
// Clear draft positions from decoder KV cache [n_past, inf)
|
||||
llama_memory_seq_rm(llama_get_memory(ctx_decoder), 0, spec->eagle3_n_past, -1);
|
||||
|
||||
// Encoder: features → g_embeddings
|
||||
const float * features = llama_get_eagle3_target_features(ctx_tgt);
|
||||
GGML_ASSERT(features && "no target features");
|
||||
|
||||
llama_batch enc_batch = {
|
||||
/*.n_tokens =*/ n_new,
|
||||
/*.token =*/ nullptr,
|
||||
/*.embd =*/ const_cast<float*>(features),
|
||||
/*.pos =*/ nullptr,
|
||||
/*.n_seq_id =*/ nullptr,
|
||||
/*.seq_id =*/ nullptr,
|
||||
/*.logits =*/ nullptr,
|
||||
};
|
||||
GGML_ASSERT(llama_encode(ctx_encoder, enc_batch) == 0);
|
||||
|
||||
const float * g_embd = llama_get_embeddings(ctx_encoder);
|
||||
GGML_ASSERT(g_embd && "encoder output failed");
|
||||
|
||||
// Decoder batch: process new tokens with KV cache reuse
|
||||
llama_set_eagle3_g_embeddings(ctx_decoder, g_embd, n_embd, n_new);
|
||||
|
||||
common_batch_clear(batch);
|
||||
for (int i = 0; i < n_new; i++) {
|
||||
const int pos = spec->eagle3_n_past + i;
|
||||
const llama_token tok = (pos < n - 1) ? prompt_tgt[pos + 1] : id_last;
|
||||
common_batch_add(batch, tok, pos, {0}, true);
|
||||
}
|
||||
|
||||
GGML_ASSERT(llama_decode(ctx_decoder, batch) == 0);
|
||||
|
||||
spec->eagle3_n_past = n; // update verified positions
|
||||
|
||||
// Sample draft tokens
|
||||
llama_tokens result;
|
||||
common_sampler_reset(smpl);
|
||||
|
||||
// Sample and check probability (consistent with standard speculative decoding)
|
||||
auto sample_and_check = [&](int idx) -> bool {
|
||||
common_sampler_sample(smpl, ctx_decoder, idx);
|
||||
|
||||
const auto * cur_p = common_sampler_get_candidates(smpl, true);
|
||||
const llama_token id = cur_p->data[0].id;
|
||||
|
||||
common_sampler_accept(smpl, id, true);
|
||||
result.push_back(id);
|
||||
|
||||
return cur_p->data[0].p >= params.p_min;
|
||||
};
|
||||
|
||||
// First draft token from batch decode
|
||||
if (!sample_and_check(n_new - 1)) {
|
||||
return result;
|
||||
}
|
||||
|
||||
// Autoregressive: use prenorm as g_embd (-1 = last output)
|
||||
const float * prenorm = llama_get_embeddings_ith(ctx_decoder, -1);
|
||||
|
||||
for (int i = 1; i < params.n_draft; i++) {
|
||||
GGML_ASSERT(prenorm && "prenorm failed");
|
||||
llama_set_eagle3_g_embeddings(ctx_decoder, prenorm, n_embd, 1);
|
||||
|
||||
common_batch_clear(batch);
|
||||
common_batch_add(batch, result.back(), n - 1 + i, {0}, true);
|
||||
GGML_ASSERT(llama_decode(ctx_decoder, batch) == 0);
|
||||
|
||||
prenorm = llama_get_embeddings_ith(ctx_decoder, -1);
|
||||
|
||||
if (!sample_and_check(0)) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
llama_tokens common_speculative_gen_draft(
|
||||
struct common_speculative * spec,
|
||||
struct common_speculative_params params,
|
||||
const llama_tokens & prompt_tgt_main_model, // specified in target model vocab
|
||||
llama_token id_last) {
|
||||
|
||||
// EAGLE3 path
|
||||
if (spec->eagle3_encoder && spec->eagle3_decoder) {
|
||||
return gen_eagle3_draft(spec, params, prompt_tgt_main_model, id_last);
|
||||
}
|
||||
|
||||
// Standard draft model path
|
||||
auto & batch = spec->batch;
|
||||
auto & ctx_tgt = spec->ctx_tgt;
|
||||
auto & ctx_dft = spec->ctx_dft;
|
||||
|
||||
Reference in New Issue
Block a user