diff --git a/common/speculative.cpp b/common/speculative.cpp index 058e75b796..4f97d464dd 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -307,7 +307,7 @@ static llama_tokens gen_eagle3_draft( /*.n_tokens =*/ n_new, /*.token =*/ nullptr, /*.embd =*/ const_cast(features), - /*.pos =*/ nullptr, + /*.pos =*/ nullptr, /*.n_seq_id =*/ nullptr, /*.seq_id =*/ nullptr, /*.logits =*/ nullptr, diff --git a/include/llama.h b/include/llama.h index caded2c83e..7d26eebcbd 100644 --- a/include/llama.h +++ b/include/llama.h @@ -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( diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a921112df2..b4f6bb5b99 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -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 eagle3_d2t_map; static thread_local std::vector 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::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 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)); diff --git a/src/llama-graph.h b/src/llama-graph.h index 617ea154c3..69df6b1f4e 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -74,18 +74,18 @@ struct llama_cross { struct llama_eagle3 { // Configuration: which layers to extract from target model std::vector 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 target_features; - + // Encoder output (for decoder input) std::vector g_embeddings; - + // Tensor references for feature extraction from target model std::vector extract_tensors; - + // Clear all stored data void clear() { target_features.clear(); diff --git a/src/models/eagle3.cpp b/src/models/eagle3.cpp index dea887bdd3..629d89d327 100644 --- a/src/models/eagle3.cpp +++ b/src/models/eagle3.cpp @@ -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(); + 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(); - 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); -} \ No newline at end of file + // 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); +} diff --git a/src/models/models.h b/src/models/models.h index 653c962d19..a6d1a2fccf 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -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 {