diff --git a/src/llama-hparams.h b/src/llama-hparams.h index a093ae78a2..2b92f5d608 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -10,8 +10,8 @@ // bump if necessary #define LLAMA_MAX_LAYERS 512 #define LLAMA_MAX_EXPERTS 1024 // Kimi K3 -#define LLAMA_MAX_PLE_NGRAM 8 // qwen4exp -#define LLAMA_MAX_PLE_HEADS 64 // qwen4exp +#define LLAMA_MAX_PLE_NGRAM 8 // qwen4exp, deepseek41 (all hash modules) +#define LLAMA_MAX_PLE_HEADS 64 // qwen4exp, deepseek41 (all hash modules) enum llama_expert_gating_func_type { LLAMA_EXPERT_GATING_FUNC_TYPE_NONE = 0, @@ -339,8 +339,6 @@ struct llama_hparams { int32_t dsv41_cand_src = -1; // -1 = no candidate block selection uint32_t dsv41_cand_topk_blocks = 0; uint32_t dsv41_cand_block_size = 0; - bool dsv41_tok_hist = false; // KV cells keep token ids for the Engram n-gram history - uint32_t dsv41_hist_len = 0; // Engram: predecessor tokens an n-gram hash reads (max_ngram - 1) // 0 = full rank (DeepSeek-V4) uint32_t hc_low_rank = 0; @@ -356,6 +354,10 @@ struct llama_hparams { uint32_t ple_eos_token_id = 0; // the id the PLE hash stands in at image positions; 0 makes the loader fall back to EOS uint32_t ple_image_token_id = 0; + // the id the KV cells record at embedding (image) positions for the n-gram history; LLAMA_TOKEN_NULL = none + llama_token ple_embd_token = LLAMA_TOKEN_NULL; + // hash constant sets, stored back to back: 1 shared by all PLE layers (qwen4exp), or 1 per PLE layer in layer order (deepseek41) + uint32_t ple_n_modules = 1; // the file lists PLE layer indices, so this is never a per-layer gguf array and can hold one bit per layer std::bitset is_ple_impl; // the hash multipliers reach ~2e13 and have to stay 64-bit diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index cdbabeb0bb..976897b790 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -1522,8 +1522,8 @@ llama_memory_context_ptr llama_kv_cache_dsv4::init_batch( bool embd_all) { GGML_UNUSED(embd_all); - if (hparams_raw.dsv41_tok_hist && - !dsv41_shared_history_ok(balloc.get_batch(), kv_raw->get_swa(), hparams_raw.dsv41_hist_len)) { + if (hparams_raw.ple_n_heads > 0 && + !dsv41_shared_history_ok(balloc.get_batch(), kv_raw->get_swa(), hparams_raw.ple_ngram_size - 1)) { return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); } diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 501cb24b61..4ab280a619 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1153,8 +1153,7 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & cells.pos_set(idx, ubatch.pos[i]); - // dsv41_tok_hist: embd cells must record LLAMA_TOKEN_NULL (DEAD for Engram), never a stale id - if (ubatch.is_pos_2d() || ubatch.token || hparams.ple_n_heads > 0 || hparams.dsv41_tok_hist) { + if (ubatch.is_pos_2d() || ubatch.token || hparams.ple_n_heads > 0) { llama_kv_cell_ext ext; if (ubatch.is_pos_2d()) { @@ -1169,9 +1168,7 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & } else if (hparams.ple_n_heads > 0) { // embd batch (multimodal input) has no token ids, need to pad it with the correct ID for PLE layers // TODO @ngxson : check if we can do the same as gemma 3n / gemma 4 - ext.tok = hparams.ple_image_token_id != 0 - ? (llama_token) hparams.ple_image_token_id - : (llama_token) hparams.ple_eos_token_id; + ext.tok = hparams.ple_embd_token; } cells.ext_set(idx, ext); @@ -1851,7 +1848,7 @@ void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const { bool llama_kv_cache::has_cell_ext() const { // M-RoPE needs the 2D position, the PLE / Engram n-gram hashes need the token id - return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0 || hparams.dsv41_tok_hist; + return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0; } void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const { diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 4564dbfb11..752e6361fe 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -361,16 +361,9 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, hparams.dsv41_cand_topk_blocks); add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE, hparams.dsv41_cand_block_size); } - if (!m41->engram_layers.empty()) { - add_kv(LLM_KV_PLE_LAYERS, std::vector(m41->engram_layers.begin(), m41->engram_layers.end())); - add_kv(LLM_KV_PLE_NGRAM_SIZE, m41->engram_max_ngram); - add_kv(LLM_KV_PLE_HEADS_PER_NGRAM, m41->engram_n_head); - add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, m41->engram_head_dim); - add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, m41->engram_mult); - add_kv(LLM_KV_PLE_HEAD_OFFSETS, m41->engram_offsets); - add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES, m41->engram_primes); - add_kv(LLM_KV_ENGRAM_PAD_TOKEN_ID, m41->engram_pad_id); - add_kv(LLM_KV_ENGRAM_TOKEN_MAP, m41->engram_token_map); + if (hparams.ple_n_heads > 0) { + add_kv(LLM_KV_ENGRAM_PAD_TOKEN_ID, m41->engram_pad_id); + add_kv(LLM_KV_ENGRAM_TOKEN_MAP, m41->engram_token_map); } } @@ -390,13 +383,13 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.ple_head_dim); add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, std::vector( hparams.ple_layer_multipliers.begin(), - hparams.ple_layer_multipliers.begin() + hparams.ple_ngram_size)); + hparams.ple_layer_multipliers.begin() + hparams.ple_n_modules*hparams.ple_ngram_size)); add_kv(LLM_KV_PLE_HEAD_OFFSETS, std::vector( hparams.ple_head_offsets.begin(), - hparams.ple_head_offsets.begin() + hparams.ple_n_heads)); + hparams.ple_head_offsets.begin() + hparams.ple_n_modules*hparams.ple_n_heads)); add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES, std::vector( hparams.ple_head_vocab_sizes.begin(), - hparams.ple_head_vocab_sizes.begin() + hparams.ple_n_heads)); + hparams.ple_head_vocab_sizes.begin() + hparams.ple_n_modules*hparams.ple_n_heads)); } const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train; diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index a29a7ddf0d..03e45f9b9f 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -83,39 +83,49 @@ void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { } } + // Engram: the PLE n-gram hash of qwen4exp with one set of hash constants per Engram layer, stored back to back std::vector e_layers; ml.get_arr(LLM_KV_PLE_LAYERS, e_layers, false); if (!e_layers.empty()) { - ml.get_key(LLM_KV_PLE_NGRAM_SIZE, engram_max_ngram); - ml.get_key(LLM_KV_PLE_HEADS_PER_NGRAM, engram_n_head); - ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, engram_head_dim); - ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, engram_mult); - ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS, engram_offsets); - ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES, engram_primes); + std::vector mult; + std::vector offsets; + std::vector primes; + ml.get_key(LLM_KV_PLE_NGRAM_SIZE, hparams.ple_ngram_size); + ml.get_key(LLM_KV_PLE_HEADS_PER_NGRAM, hparams.ple_heads_per_ngram); + ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.ple_head_dim); + ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, mult); + ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS, offsets); + ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES, primes); ml.get_key(LLM_KV_ENGRAM_PAD_TOKEN_ID, engram_pad_id); ml.get_arr(LLM_KV_ENGRAM_TOKEN_MAP, engram_token_map); const size_t n_e = e_layers.size(); - const size_t n_cols = engram_n_cols(); - if (engram_max_ngram < 2 || engram_n_head == 0 || engram_head_dim == 0 || - engram_mult.size() != n_e*engram_max_ngram || engram_primes.size() != n_e*n_cols || - engram_offsets.size() != n_e*n_cols || engram_pad_id >= engram_token_map.size()) { + const size_t n_gram = hparams.ple_ngram_size; + const size_t n_cols = n_gram >= 2 ? (n_gram - 1)*hparams.ple_heads_per_ngram : 0; + if (n_cols == 0 || hparams.ple_head_dim == 0 || + mult.size() != n_e*n_gram || primes.size() != n_e*n_cols || offsets.size() != n_e*n_cols || + mult.size() > LLAMA_MAX_PLE_NGRAM || primes.size() > LLAMA_MAX_PLE_HEADS || + engram_pad_id >= engram_token_map.size()) { throw std::runtime_error("DeepSeek-V4.1 engram metadata is inconsistent"); } for (size_t e = 0; e < n_e; ++e) { - if (e_layers[e] >= n_layer) { - throw std::runtime_error(format("DeepSeek-V4.1 engram layer %u is out of range", e_layers[e])); + // module e belongs to the e-th Engram layer + if (e_layers[e] >= n_layer || (e > 0 && e_layers[e] <= e_layers[e - 1])) { + throw std::runtime_error(format("DeepSeek-V4.1 engram layer %u is out of range or out of order", e_layers[e])); } + hparams.is_ple_impl.set(e_layers[e]); } for (size_t c = 0; c < n_e*n_cols; ++c) { // a bucket count of 0 traps in the hash; row ids are int32 - if (engram_primes[c] < 2 || engram_offsets[c] + engram_primes[c] > (uint64_t) INT32_MAX) { + if (primes[c] < 2 || offsets[c] + primes[c] > (uint64_t) INT32_MAX) { throw std::runtime_error(format("DeepSeek-V4.1 engram head %zu: bucket range out of range", c)); } + hparams.ple_head_offsets[c] = (uint32_t) offsets[c]; + hparams.ple_head_vocab_sizes[c] = (uint32_t) primes[c]; } - engram_layers.assign(e_layers.begin(), e_layers.end()); - hparams.dsv41_tok_hist = true; - hparams.dsv41_hist_len = engram_max_ngram > 0 ? (uint32_t) engram_max_ngram - 1 : 0; + std::copy(mult.begin(), mult.end(), hparams.ple_layer_multipliers.begin()); + hparams.ple_n_heads = (uint32_t) n_cols; + hparams.ple_n_modules = (uint32_t) n_e; } } @@ -139,10 +149,10 @@ void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); // all Engram tables in one lazily read tensor, like the qwen4exp n-gram table - if (!engram_layers.empty()) { + if (hparams.ple_n_heads > 0) { int64_t n_rows = 0; - for (size_t c = 0; c < engram_offsets.size(); ++c) { - n_rows = std::max(n_rows, (int64_t) (engram_offsets[c] + engram_primes[c])); + for (size_t c = 0; c < (size_t) hparams.ple_n_modules*hparams.ple_n_heads; ++c) { + n_rows = std::max(n_rows, (int64_t) hparams.ple_head_offsets[c] + hparams.ple_head_vocab_sizes[c]); } const std::string name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str(); if (const auto * w = ml.get_weight(name.c_str())) { @@ -152,7 +162,7 @@ void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { } n_rows = w->tensor->ne[1]; } - per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {(int64_t) engram_head_dim, n_rows}, TENSOR_READ_LAZY); + per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {(int64_t) hparams.ple_head_dim, n_rows}, TENSOR_READ_LAZY); } for (int i = 0; i < n_layer; ++i) { @@ -190,7 +200,7 @@ void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { } if (engram_index(i) >= 0) { - layer.engram_kv = create_tensor(tn(LLM_TENSOR_ENGRAM_KV, "weight", i), {(int64_t) engram_n_cols() * engram_head_dim, n_embd * (hc_mult + 1)}, 0); + layer.engram_kv = create_tensor(tn(LLM_TENSOR_ENGRAM_KV, "weight", i), {(int64_t) hparams.ple_n_heads * hparams.ple_head_dim, n_embd * (hc_mult + 1)}, 0); layer.engram_gate = create_tensor(tn(LLM_TENSOR_ENGRAM_GATE, "weight", i), {n_embd, hc_mult}, 0); } @@ -219,20 +229,26 @@ std::unique_ptr llama_model_deepseek41::build_arch_graph(cons // Engram n-gram hash (engram.py NgramHashState.forward) int llama_model_deepseek41::engram_index(int il) const { - const auto it = std::find(engram_layers.begin(), engram_layers.end(), il); - return it == engram_layers.end() ? -1 : (int) (it - engram_layers.begin()); + if (hparams.ple_n_heads == 0 || !hparams.is_ple(il)) { + return -1; + } + int e = 0; + for (int j = 0; j < il; ++j) { + e += hparams.is_ple(j); + } + return e; } void llama_model_deepseek41::engram_hash(int e, const llama_ubatch & ubatch, const std::vector & prev, std::vector & out) const { - const int64_t n_gram = engram_max_ngram; + const int64_t n_gram = hparams.ple_ngram_size; const int64_t n_prev = n_gram - 1; - const int64_t n_head = engram_n_head; - const int64_t n_cols = engram_n_cols(); + const int64_t n_head = hparams.ple_heads_per_ngram; + const int64_t n_cols = hparams.ple_n_heads; - const uint64_t * mult = engram_mult.data() + (size_t) e*n_gram; - const uint64_t * primes = engram_primes.data() + (size_t) e*n_cols; + const uint64_t * mult = hparams.ple_layer_multipliers.data() + (size_t) e*n_gram; + const uint32_t * primes = hparams.ple_head_vocab_sizes.data() + (size_t) e*n_cols; const int64_t pad = engram_token_map[engram_pad_id]; - const uint64_t * offs = engram_offsets.data() + (size_t) e*n_cols; + const uint32_t * offs = hparams.ple_head_offsets.data() + (size_t) e*n_cols; const auto cid = [&](llama_token t) -> int64_t { return t < 0 || t >= (llama_token) engram_token_map.size() ? DSV41_ENGRAM_DEAD : engram_token_map[t]; @@ -291,7 +307,7 @@ void llm_graph_input_dsv41::set_input(const llama_ubatch * ubatch) { if (!engram_ids.empty()) { // a token shared by several sequences has the same history in each, so its first sequence resolves it std::vector prev; - mctx->get_raw()->get_prev_tokens(*ubatch, model.engram_max_ngram - 1, prev); + mctx->get_raw()->get_prev_tokens(*ubatch, model.hparams.ple_ngram_size - 1, prev); // hash all Engram layers first so a lazily read table gets one batched readahead before the first gather std::vector> ids(engram_ids.size()); std::vector all; @@ -341,7 +357,7 @@ bool llm_graph_input_dsv41::can_reuse(const llm_graph_params & params) { bool res = true; for (auto * t : engram_ids) { - res &= t->ne[0] == (int64_t) model.engram_n_cols()*params.ubatch.n_tokens; + res &= t->ne[0] == (int64_t) model.hparams.ple_n_heads*params.ubatch.n_tokens; } res &= (engram_keep != nullptr) == (!engram_ids.empty() && params.ubatch.embd != nullptr); if (engram_keep) { @@ -366,7 +382,7 @@ ggml_tensor * llama_model_deepseek41::graph::build_engram( const int64_t hc = hparams.dsv4_hc_mult; const int64_t nt = x->ne[2]; - GGML_ASSERT(rows->ne[0] == (int64_t) (model.engram_head_dim*model.engram_n_cols()) && rows->ne[1] == nt); + GGML_ASSERT(rows->ne[0] == (int64_t) (hparams.ple_head_dim*hparams.ple_n_heads) && rows->ne[1] == nt); ggml_tensor * kv = build_lora_mm(layer.engram_kv, rows); cb(kv, "engram_kv", il); @@ -773,21 +789,23 @@ llama_model_deepseek41::graph::graph(const llama_model & model_base, const llm_g // embedding input has no token ids: it skips Engram, like the reference's masked (image) positions // the gathers go right after the token embeddings, so they share the first graph split (the table stays on the host) std::vector engram_rows; - for (size_t e = 0; ubatch.token && e < model.engram_layers.size(); ++e) { - const int32_t il = model.engram_layers[e]; - ggml_tensor * ids = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, (int64_t) model.engram_n_cols()*n_tokens); + for (int il = 0; ubatch.token && hparams.ple_n_heads > 0 && il < n_layer; ++il) { + if (!hparams.is_ple(il)) { + continue; + } + ggml_tensor * ids = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, (int64_t) hparams.ple_n_heads*n_tokens); ggml_set_name(ids, ("engram_ids-" + std::to_string(il)).c_str()); ggml_set_input(ids); inp41->engram_ids.push_back(ids); ggml_tensor * rows = ggml_get_rows(ctx0, model.per_layer_tok_embd, ids); - rows = ggml_reshape_2d(ctx0, rows, model.engram_head_dim*model.engram_n_cols(), n_tokens); + rows = ggml_reshape_2d(ctx0, rows, hparams.ple_head_dim*hparams.ple_n_heads, n_tokens); cb(rows, "engram_rows", il); ggml_build_forward_expand(gf, rows); engram_rows.push_back(rows); } ggml_tensor * engram_keep = nullptr; - if (!model.engram_layers.empty() && ubatch.token && ubatch.embd) { + if (hparams.ple_n_heads > 0 && ubatch.token && ubatch.embd) { engram_keep = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens); ggml_set_input(engram_keep); ggml_set_name(engram_keep, "engram_keep"); diff --git a/src/models/models.h b/src/models/models.h index fc83a7fe91..4eeb361eab 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1498,19 +1498,12 @@ struct llama_model_deepseek41 : public llama_model_deepseek4 { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; - // Engram hash constants, hashed on the host (ggml has no int64 mul/xor/mod) - std::vector engram_layers; - uint32_t engram_max_ngram = 0; - uint32_t engram_n_head = 0; - uint32_t engram_head_dim = 0; - uint32_t engram_pad_id = 0; - std::vector engram_token_map; - std::vector engram_mult; - std::vector engram_primes; - std::vector engram_offsets; // first row of each (layer, order, head) in per_layer_token_embd + // Engram: the hash constants are hparams.ple_* (one module per Engram layer), hashed on the host (ggml has no + // int64 mul/xor/mod); the token map and its pad id are V4.1 only + uint32_t engram_pad_id = 0; + std::vector engram_token_map; int engram_index(int il) const; - uint32_t engram_n_cols() const { return (engram_max_ngram - 1)*engram_n_head; } void engram_hash(int e, const llama_ubatch & ubatch, const std::vector & prev, std::vector & out) const; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index fc9acd43e3..e20f20591f 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -115,6 +115,9 @@ void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram; hparams.ple_head_dim = hparams.n_embd_per_layer; + hparams.ple_embd_token = hparams.ple_image_token_id != 0 + ? (llama_token) hparams.ple_image_token_id + : (llama_token) hparams.ple_eos_token_id; if (hparams.ple_ngram_size < 2 || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) { throw std::runtime_error(format("PLE n-gram size %u is out of range", hparams.ple_ngram_size)); }