eagle3 : improve naming

This commit is contained in:
Georgi Gerganov 2025-12-17 15:49:03 +02:00
parent 3e7f376b53
commit 5a79c1900f
No known key found for this signature in database
GPG Key ID: 449E073F9DC10735
6 changed files with 132 additions and 130 deletions

View File

@ -307,7 +307,7 @@ static llama_tokens gen_eagle3_draft(
/*.n_tokens =*/ n_new,
/*.token =*/ nullptr,
/*.embd =*/ const_cast<float*>(features),
/*.pos =*/ nullptr,
/*.pos =*/ nullptr,
/*.n_seq_id =*/ nullptr,
/*.seq_id =*/ nullptr,
/*.logits =*/ nullptr,

View File

@ -364,7 +364,7 @@ extern "C" {
bool kv_unified; // use a unified buffer across the input sequences when computing the attention
// try to disable when n_seq_max > 1 for improved performance when the sequences do not share a large prefix
// ref: https://github.com/ggml-org/llama.cpp/pull/14363
// EAGLE3 extraction configuration
// When eagle3_model is set, layer extraction is automatically enabled
const struct llama_model * eagle3_model; // EAGLE3 model to read extract_layers configuration from
@ -876,7 +876,7 @@ extern "C" {
// Returns NULL if no features are available
// Format: [3*n_embd, n_tokens] - use model.hparams.n_embd and batch.n_tokens for dimensions
LLAMA_API const float * llama_get_eagle3_target_features(struct llama_context * ctx);
// Set g_embeddings from EAGLE3 encoder output for decoder input
// g_embd: pointer to encoder output embeddings
LLAMA_API void llama_set_eagle3_g_embeddings(

View File

@ -353,7 +353,7 @@ llama_context::llama_context(
// Allocate tensors array for extraction
eagle3.extract_tensors.resize(eagle3.extract_layer_indices.size(), nullptr);
LLAMA_LOG_INFO("%s: EAGLE3 extraction enabled for layers [%d, %d, %d]\n", __func__,
eagle3.extract_layer_indices[0],
eagle3.extract_layer_indices[1],
@ -879,7 +879,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll
//const auto t_start_us = ggml_time_us();
res->set_inputs(&ubatch);
// EAGLE3: Fill g_embeddings for decoder input
if (model.arch == LLM_ARCH_EAGLE3 && gtype == LLM_GRAPH_TYPE_DECODER && !eagle3.g_embeddings.empty()) {
ggml_tensor * g_embd = ggml_graph_get_tensor(gf, "inp_g_embeddings");
@ -1265,32 +1265,32 @@ int llama_context::decode(const llama_batch & batch_inp) {
if (n_outputs) {
GGML_ASSERT( n_outputs_prev + n_outputs <= n_outputs_all);
GGML_ASSERT((n_outputs_prev + n_outputs)*n_vocab <= (int64_t) logits_size);
// EAGLE3: Map draft vocab to target vocab
if (model.arch == LLM_ARCH_EAGLE3 && model.d2t) {
static thread_local std::vector<int64_t> eagle3_d2t_map;
static thread_local std::vector<float> eagle3_draft_logits;
const int64_t draft_vocab_size = t_logits->ne[0];
const uint32_t last_idx = n_outputs - 1;
// Load d2t mapping once (on first call)
if (eagle3_d2t_map.empty()) {
eagle3_d2t_map.resize(model.d2t->ne[0]);
ggml_backend_tensor_get(model.d2t, eagle3_d2t_map.data(), 0, eagle3_d2t_map.size() * sizeof(int64_t));
}
// Read only the last token's draft logits
eagle3_draft_logits.resize(draft_vocab_size);
const size_t last_offset = last_idx * draft_vocab_size * sizeof(float);
ggml_backend_tensor_get_async(backend_res, t_logits, eagle3_draft_logits.data(), last_offset, draft_vocab_size * sizeof(float));
synchronize();
// Map only the last token's draft logits to target vocab
float * last_logits_out = logits_out + last_idx * n_vocab;
std::fill(last_logits_out, last_logits_out + n_vocab, -std::numeric_limits<float>::infinity());
for (int64_t j = 0; j < draft_vocab_size; j++) {
const int64_t target_id = j + eagle3_d2t_map[j];
GGML_ASSERT(target_id >= 0 && target_id < n_vocab);
@ -1656,7 +1656,7 @@ llm_graph_cb llama_context::graph_get_cb() const {
if (cparams.eagle3_extract_enabled) {
static constexpr const char * prefix = "eagle3_extract_";
static constexpr size_t prefix_len = 15; // strlen("eagle3_extract_")
if (strncmp(name, prefix, prefix_len) == 0) {
// Parse the extraction index from the name (e.g., "eagle3_extract_0" -> 0)
size_t extract_idx = 0;
@ -1667,7 +1667,7 @@ llm_graph_cb llama_context::graph_get_cb() const {
eagle3.extract_tensors[extract_idx] = cur;
LLAMA_LOG_DEBUG("%s: EAGLE3 stored tensor reference for extraction: "
"index=%zu, layer=%d, target_layer=%d, tensor=%s\n",
__func__, extract_idx, il,
__func__, extract_idx, il,
eagle3.extract_layer_indices[extract_idx], name);
}
}
@ -1702,36 +1702,36 @@ void llama_context::extract_eagle3_features(const llama_ubatch & ubatch) {
const int64_t n_tokens = ubatch.n_tokens;
const int64_t n_embd = model.hparams.n_embd;
const size_t n_layers = eagle3.extract_tensors.size();
// Allocate storage for concatenated features
const int64_t n_embd_concat = n_embd * n_layers;
eagle3.target_features.resize(n_embd_concat * n_tokens);
// Temporary buffer to hold layer features before transposing
static thread_local std::vector<float> temp_layer_features;
temp_layer_features.resize(n_embd * n_tokens);
LLAMA_LOG_DEBUG("%s: Start to extract EAGLE3 features: %zu layers, %lld tokens, %lld embd\n",
__func__, n_layers, (long long)n_tokens, (long long)n_embd);
// Extract each layer's features and interleave into token-major layout
for (size_t layer_idx = 0; layer_idx < n_layers; ++layer_idx) {
ggml_tensor * tensor = eagle3.extract_tensors[layer_idx];
GGML_ASSERT(tensor != nullptr && "EAGLE3 extraction tensor is null");
// Get the backend where this tensor is stored
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched.get(), tensor);
GGML_ASSERT(backend != nullptr && "EAGLE3 tensor has no backend");
// Verify tensor shape: should be [n_embd, n_tokens]
GGML_ASSERT(tensor->ne[0] == n_embd && tensor->ne[1] == n_tokens &&
"EAGLE3 extraction tensor has unexpected shape");
// Get layer features to temp buffer
const size_t size_bytes = n_embd * n_tokens * sizeof(float);
ggml_backend_tensor_get_async(backend, tensor, temp_layer_features.data(), 0, size_bytes);
ggml_backend_sched_synchronize(sched.get());
// Then copy to correct position in target_features
// target_features layout: [token_0_all_layers, token_1_all_layers, ...]
// Each token has [layer_0_embd, layer_1_embd, layer_2_embd]
@ -1743,7 +1743,7 @@ void llama_context::extract_eagle3_features(const llama_ubatch & ubatch) {
std::memcpy(dest, src, n_embd * sizeof(float));
}
}
}
//
@ -3235,7 +3235,7 @@ const float * llama_context::get_eagle3_target_features() const {
void llama_context::set_eagle3_g_embeddings(const float * g_embd, int32_t n_embd, int32_t n_tokens) {
GGML_ASSERT(g_embd != nullptr && "g_embeddings cannot be null");
GGML_ASSERT(n_embd > 0 && n_tokens > 0 && "invalid dimensions");
const size_t size = n_embd * n_tokens;
eagle3.g_embeddings.resize(size);
std::memcpy(eagle3.g_embeddings.data(), g_embd, size * sizeof(float));

View File

@ -74,18 +74,18 @@ struct llama_cross {
struct llama_eagle3 {
// Configuration: which layers to extract from target model
std::vector<int> extract_layer_indices;
// Extracted features from target model (for encoder input)
// Concatenated [layer_l, layer_m, layer_h] embeddings
// Shape: [n_layers * n_embd, n_tokens] where n_layers = extract_layer_indices.size()
std::vector<float> target_features;
// Encoder output (for decoder input)
std::vector<float> g_embeddings;
// Tensor references for feature extraction from target model
std::vector<ggml_tensor *> extract_tensors;
// Clear all stored data
void clear() {
target_features.clear();

View File

@ -1,103 +1,109 @@
#include "models.h"
ggml_tensor * llm_build_eagle3_encode::build_inp_embd() const {
const int64_t n_embd_target_features = 3 * hparams.eagle3_target_hidden_size;
ggml_tensor * cur = nullptr;
// Input: Target model features (3 layers concatenated: low, mid, high)
// Data will be provided via ubatch->embd in encode_eagle3_features()
auto inp_target = std::make_unique<llm_graph_input_embd>();
inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_target_features, n_tokens);
ggml_set_input(inp_target->embd);
cur = inp_target->embd;
cb(cur, "inp_embd", -1);
res->add_input(std::move(inp_target));
return cur;
}
// EAGLE3 Encoder: processes target model features through feature fusion layer
// Input: target_features e.g. [12288, n_tokens] from target model layers low, middle, high
// Output: g_embeddings e.g. [4096, n_tokens] stored in context
llm_build_eagle3_encode::llm_build_eagle3_encode(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
ggml_tensor * cur = nullptr;
const int64_t n_embd_target_features = 3 * hparams.eagle3_target_hidden_size;
cur = build_inp_embd();
ggml_tensor * cur;
// Feature fusion layer
cur = build_lora_mm(model.fc, cur);
cb(cur, "fc_out", -1);
// Input: Target model features (3 layers concatenated: low, mid, high)
// Data will be provided via ubatch->embd in encode_eagle3_features()
auto inp_target = std::make_unique<llm_graph_input_embd>();
inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_target_features, n_tokens);
ggml_set_input(inp_target->embd);
ggml_tensor * target_features = inp_target->embd;
res->add_input(std::move(inp_target));
cb(target_features, "inp_target_features", -1);
// Output: g_embeddings e.g. [4096, n_tokens]
res->t_embd = cur;
// Feature fusion layer
ggml_tensor * fused_target = build_lora_mm(model.fc, target_features);
cb(fused_target, "fc_out", -1);
// Output: g_embeddings e.g. [4096, n_tokens]
cur = fused_target;
res->t_embd = cur;
ggml_build_forward_expand(gf, cur);
ggml_build_forward_expand(gf, cur);
}
// EAGLE3 Decoder: processes draft tokens using g_embeddings from encoder
// Input: draft tokens + g_embeddings from encoder
// Output: draft logits
llm_build_eagle3_decode::llm_build_eagle3_decode(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_layer == 1); // EAGLE-3 has only one decoder layer
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_layer == 1); // EAGLE-3 has only one decoder layer
ggml_tensor * cur;
ggml_tensor * inpL;
ggml_tensor * cur;
ggml_tensor * inpL;
// EAGLE3 Decoder receives:
// 1. Token embeddings (e.g.from EAGLE3's own tok_embd for Llama 3.3 70B, or target model for Llama 3.1 8B)
// 2. g_embeddings from encoder
// Choose token_embd_eagle3: prefer EAGLE3's own if available (Llama 3.3 70B), else use target's (Llama 3.1 8B)
ggml_tensor * token_embd_eagle3 = (model.tok_embd != nullptr) ? model.tok_embd : model.target_tok_embd;
GGML_ASSERT(token_embd_eagle3 != nullptr && "EAGLE3 decoder requires token embeddings (own or from target model)");
ggml_tensor * input_embeds = build_inp_embd(token_embd_eagle3);
cb(input_embeds, "token_embd_eagle3", -1);
ggml_tensor * g_embeddings = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
ggml_set_input(g_embeddings);
ggml_set_name(g_embeddings, "inp_g_embeddings");
cb(g_embeddings, "inp_g_embeddings", -1);
// EAGLE3 Decoder receives:
// 1. Token embeddings (e.g.from EAGLE3's own tok_embd for Llama 3.3 70B, or target model for Llama 3.1 8B)
// 2. g_embeddings from encoder
// Choose token_embd_eagle3: prefer EAGLE3's own if available (Llama 3.3 70B), else use target's (Llama 3.1 8B)
ggml_tensor * token_embd_eagle3 = (model.tok_embd != nullptr) ? model.tok_embd : model.target_tok_embd;
GGML_ASSERT(token_embd_eagle3 != nullptr && "EAGLE3 decoder requires token embeddings (own or from target model)");
ggml_tensor * inp_embd = build_inp_embd(token_embd_eagle3);
cb(inp_embd, "inp_embd", -1);
// Store raw g_embeddings as residual
ggml_tensor * residual = g_embeddings;
// TODO: refactor into llm_graph_input
ggml_tensor * inp_g = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
ggml_set_input(inp_g);
cb(inp_g, "inp_g_embeddings", -1); // TODO: do not change the name! refactor into llm_graph_input
// Apply input_layernorm to the token embeddings
ggml_tensor * input_embeds_normed = build_norm(input_embeds,
model.layers[0].attn_norm, NULL,
LLM_NORM_RMS, 0);
cb(input_embeds_normed, "input_layernorm", -1);
inpL = inp_g;
// Apply hidden_norm to g_embeddings
ggml_tensor * g_embeddings_normed = build_norm(g_embeddings,
model.layers[0].eagle3_hidden_norm, NULL,
LLM_NORM_RMS, -1);
cb(g_embeddings_normed, "g_embeddings_normed", -1);
// inp_pos - contains the positions
ggml_tensor * inp_pos = build_inp_pos();
// Concatenate normalized input_embeds and normalized g_embeddings
cur = ggml_concat(ctx0, input_embeds_normed, g_embeddings_normed, 0);
cb(cur, "concat_embeds_g", -1);
inpL = cur;
auto * inp_attn = build_attn_inp_kv();
// inp_pos - contains the positions
ggml_tensor * inp_pos = build_inp_pos();
const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * inp_out_ids = build_inp_out_ids();
ggml_tensor * inp_out_ids = build_inp_out_ids();
const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
// Single decoder layer (il = 0)
const int il = 0;
{
// inpL is the concatenated input (normalized input_embeds + normalized g_embeddings)
// Single decoder layer (il = 0)
const int il = 0;
{
// inpL is the concatenated input (normalized inp_embd + normalized inp_g)
ggml_tensor * inpSA = inpL;
// Apply input_layernorm to the token embeddings
ggml_tensor * embd_norm = build_norm(inp_embd,
model.layers[il].attn_norm, NULL,
LLM_NORM_RMS, il);
cb(embd_norm, "embd_norm", il);
// Apply hidden_norm to inp_g
ggml_tensor * g_norm = build_norm(inp_g,
model.layers[il].eagle3_hidden_norm, NULL,
LLM_NORM_RMS, -1);
cb(g_norm, "g_norm", il);
// Concatenate normalized inp_embd and normalized inp_g
cur = ggml_concat(ctx0, embd_norm, g_norm, il);
cb(cur, "concat_embd", il);
// Self-attention with concatenated input
ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, inpL);
ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur);
cb(Qcur, "Qcur", il);
ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, inpL);
ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
cb(Kcur, "Kcur", il);
ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, inpL);
ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
cb(Vcur, "Vcur", il);
Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
@ -127,25 +133,19 @@ llm_build_eagle3_decode::llm_build_eagle3_decode(const llama_model & model, cons
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
if (inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
residual = ggml_get_rows(ctx0, residual, inp_out_ids);
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
// Add residual and update it
ggml_tensor * attn_with_residual = ggml_add(ctx0, cur, residual);
cb(attn_with_residual, "attn_with_residual", il);
// Update residual
residual = attn_with_residual;
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cb(ffn_inp, "ffn_inp", il);
// Apply FFN norm to the sum
ggml_tensor * ffn_inp = build_norm(attn_with_residual,
cur = build_norm(ffn_inp,
model.layers[il].ffn_norm, NULL,
LLM_NORM_RMS, il);
cb(ffn_inp, "post_attn_norm", il);
cur = ffn_inp;
cb(cur, "post_attn_norm", il);
cur = build_ffn(cur,
model.layers[il].ffn_up, NULL, NULL,
@ -154,30 +154,30 @@ llm_build_eagle3_decode::llm_build_eagle3_decode(const llama_model & model, cons
NULL,
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(cur, "ffn_out", il);
inpL = cur;
}
cur = inpL;
// Output norm with residual
ggml_tensor * final_with_residual = ggml_add(ctx0, cur, residual);
cb(final_with_residual, "eagle3_prenorm", -1);
// Output prenorm state (for next token's g_embeddings in autoregressive generation)
ggml_set_output(final_with_residual);
res->t_embd = final_with_residual;
cur = build_norm(final_with_residual,
model.output_norm, NULL,
LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
cur = ggml_add(ctx0, cur, ffn_inp);
cb(cur, "eagle3_prenorm", il);
// lm_head - projects to draft vocabulary
cur = build_lora_mm(model.output, cur);
inpL = cur;
}
cb(cur, "result_output", -1);
res->t_logits = cur;
cur = inpL;
ggml_build_forward_expand(gf, cur);
}
// Output prenorm state (for next token's g_embeddings in autoregressive generation)
ggml_set_output(cur);
res->t_embd = cur;
cur = build_norm(cur,
model.output_norm, NULL,
LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
// lm_head - projects to draft vocabulary
cur = build_lora_mm(model.output, cur);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}

View File

@ -152,6 +152,8 @@ struct llm_build_dream : public llm_graph_context {
struct llm_build_eagle3_encode : public llm_graph_context {
llm_build_eagle3_encode(const llama_model & model, const llm_graph_params & params);
private:
ggml_tensor * build_inp_embd() const;
};
struct llm_build_eagle3_decode : public llm_graph_context {