model : deepseek41 quantized K cache, shared prompt tokens, Engram host placement

src/llama-context.cpp

@@ -1323,8 +1323,9 @@ void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void
 void llama_context::set_embeddings(bool value) {
     LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
 
-    // a pooled context computes the logits only without embeddings, so the toggle changes the graph shape
-    if (cparams.embeddings != value && cparams.pooling_type != LLAMA_POOLING_TYPE_NONE) {
+    // the toggle changes the graph shape (pooled: logits only; pooling NONE: the reserve graph outputs all tokens,
+    // which crashed ggml_gallocr on the first /completion), so re-reserve on every toggle
+    if (cparams.embeddings != value) {
         sched_need_reserve = true;
     }
 

src/llama-hparams.h

@@ -339,6 +339,7 @@ struct llama_hparams {
     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;

src/llama-kv-cache-dsv4.cpp

@@ -1376,12 +1376,76 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
     clear_compressed(-1, true);
 }
 
+// Engram hashes the n-gram ending at each token, so a token shared by several sequences (llama-batched prompt)
+// is well-defined only if all of them agree on its predecessors; returns false (and logs) otherwise
+static bool dsv41_shared_history_ok(const llama_batch & batch, const llama_kv_cache * kv, uint32_t n_hist) {
+    if (n_hist == 0 || !batch.token || !batch.n_seq_id) {
+        return true;
+    }
+
+    bool any_shared = false;
+    for (int32_t i = 0; i < batch.n_tokens && !any_shared; ++i) {
+        any_shared = batch.n_seq_id[i] > 1;
+    }
+    if (!any_shared) {
+        return true;
+    }
+
+    std::map<std::pair<llama_seq_id, llama_pos>, llama_token> in_batch;
+    for (int32_t i = 0; i < batch.n_tokens; ++i) {
+        for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {
+            in_batch[{ batch.seq_id[i][s], batch.pos[i] }] = batch.token[i];
+        }
+    }
+
+    const auto tok_at = [&](llama_seq_id seq_id, llama_pos p) -> llama_token {
+        // batch first: the batch continues the cache
+        auto it = in_batch.upper_bound({ seq_id, p });
+        if (it != in_batch.begin()) {
+            --it;
+            if (it->first.first == seq_id) {
+                return it->second;
+            }
+        }
+        return kv ? kv->seq_pos_tok_le(seq_id, p) : LLAMA_TOKEN_NULL;
+    };
+
+    for (int32_t i = 0; i < batch.n_tokens; ++i) {
+        if (batch.n_seq_id[i] <= 1) {
+            continue;
+        }
+        for (uint32_t d = 1; d <= n_hist; ++d) {
+            const llama_pos p = batch.pos[i] - (llama_pos) d;
+            if (p < 0) {
+                break;
+            }
+            const llama_token t0 = tok_at(batch.seq_id[i][0], p);
+            for (int32_t s = 1; s < batch.n_seq_id[i]; ++s) {
+                const llama_token ts = tok_at(batch.seq_id[i][s], p);
+                if (ts != t0) {
+                    LLAMA_LOG_ERROR("%s: token %d at pos %d is shared by seq %d and seq %d, whose Engram n-gram histories "
+                            "differ at pos %d (%d vs %d); decode such tokens per sequence\n", __func__,
+                            i, batch.pos[i], batch.seq_id[i][0], batch.seq_id[i][s], p, t0, ts);
+                    return false;
+                }
+            }
+        }
+    }
+
+    return true;
+}
+
 llama_memory_context_ptr llama_kv_cache_dsv4::init_batch(
             llama_batch_allocr & balloc,
             uint32_t n_ubatch,
             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)) {
+        return std::make_unique<llama_kv_cache_dsv4_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
+    }
+
     const bool raw_per_seq  = kv_raw->get_base()->get_n_stream() != 1;
     const bool comp_per_seq = csa_state->get_n_stream() > 1;
     const bool has_coupled = dsv4_batch_has_coupled(balloc.get_batch());

src/llama-kv-cache.cpp

@@ -1150,7 +1150,8 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
 
             cells.pos_set(idx, ubatch.pos[i]);
 
-            if (ubatch.is_pos_2d() || ubatch.token || hparams.ple_n_heads > 0) {
+            // 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) {
                 llama_kv_cell_ext ext;
 
                 if (ubatch.is_pos_2d()) {
@@ -1909,6 +1910,14 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st
     }
 }
 
+llama_token llama_kv_cache::seq_pos_tok_le(llama_seq_id seq_id, llama_pos p) const {
+    if (!has_cell_ext() || seq_id < 0 || (size_t) seq_id >= seq_to_stream.size() || p < 0) {
+        return LLAMA_TOKEN_NULL;
+    }
+
+    return v_cells[seq_to_stream[seq_id]].seq_pos_tok_le(seq_id, p);
+}
+
 size_t llama_kv_cache::total_size() const {
     size_t size = 0;
 

src/llama-kv-cache.h

@@ -250,6 +250,9 @@ class llama_kv_cache : public llama_memory_i {
     // note: used by n-gram input embeddings
     void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;
 
+    // token of seq_id at the largest cached position <= p, LLAMA_TOKEN_NULL if none or no cell ext
+    llama_token seq_pos_tok_le(llama_seq_id seq_id, llama_pos p) const;
+
 private:
     const llama_model & model;
     const llama_hparams & hparams;

src/llama-model-loader.cpp

@@ -1269,6 +1269,14 @@ struct ggml_tensor * llama_model_loader::create_tensor(
             }
         }
 
+        // the deepseek41 Engram table (~200 GB) is a CPU get_rows gather table: never put it in device memory
+        if (get_arch() == LLM_ARCH_DEEPSEEK41 && tn.tensor == LLM_TENSOR_PER_LAYER_TOKEN_EMBD && !ggml_backend_buft_is_host(buft)) {
+            LLAMA_LOG_WARN("llama_model_loader: %s (%zu MiB %s) must stay in host memory; ignoring placement on %s, using %s\n",
+                    tn.str().c_str(), ggml_nbytes(t_meta) / 1024 / 1024, ggml_type_name(t_meta->type),
+                    ggml_backend_buft_name(buft), ggml_backend_buft_name(ggml_backend_cpu_buffer_type()));
+            buft = ggml_backend_cpu_buffer_type();
+        }
+
         // avoid using a host buffer when using mmap
         auto * buft_dev = ggml_backend_buft_get_device(buft);
         if (use_mmap && buft_dev && buft == ggml_backend_dev_host_buffer_type(buft_dev)) {

src/llama-quant.cpp

@@ -316,6 +316,9 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param
     // do not quantize the i32 token-id -> expert-id routing table (DeepSeek-V4)
     quantize &= name.find("ffn_gate_tid2eid.weight") == std::string::npos;
 
+    // Engram gate is an elementwise weight (ggml_mul), not a matmul operand
+    quantize &= name.find("engram_gate.weight") == std::string::npos;
+
     // these are very small (e.g. 4x4)
     quantize &= name.find("altup")  == std::string::npos;
     quantize &= name.find("laurel") == std::string::npos;

src/models/deepseek41.cpp

@@ -107,6 +107,7 @@ void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) {
         }
         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;
     }
 }
 
@@ -230,7 +231,8 @@ void llama_model_deepseek41::engram_hash(int e, const llama_ubatch & ubatch, con
     std::vector<uint64_t> toks(n_gram);
     for (int64_t i = 0; i < ubatch.n_tokens; ++i) {
         // shift 0 is the token itself; once a shift is blocked (sequence start or dead token) every longer one is too
-        int64_t c = cid(ubatch.token[i]);
+        const bool is_embd = !ubatch.token || (ubatch.is_mixed() && ubatch.type[i]);
+        int64_t c = is_embd ? DSV41_ENGRAM_DEAD : cid(ubatch.token[i]);
         bool blocked = c == DSV41_ENGRAM_DEAD;
         toks[0] = blocked ? pad : c;
         for (int64_t s = 1; s < n_gram; ++s) {
@@ -261,6 +263,7 @@ class llm_graph_input_dsv41 : public llm_graph_input_i {
     bool can_reuse(const llm_graph_params & params) override;
 
     std::vector<ggml_tensor *> engram_ids;
+    ggml_tensor * engram_keep = nullptr;    // 0 for embedding (image) tokens; only in mixed ubatches
     ggml_tensor * cand_pin = nullptr;
     bool cand_csa = false;
 
@@ -295,6 +298,16 @@ void llm_graph_input_dsv41::set_input(const llama_ubatch * ubatch) {
         }
     }
 
+    // embedding (image) tokens are hashed DEAD and get no Engram contribution (gate 0)
+    if (engram_keep && engram_keep->buffer) {
+        std::vector<float> keep(ubatch->n_tokens);
+        for (uint32_t i = 0; i < ubatch->n_tokens; ++i) {
+            const bool is_embd = !ubatch->token || (ubatch->is_mixed() && ubatch->type[i]);
+            keep[i] = is_embd ? 0.0f : 1.0f;
+        }
+        ggml_backend_tensor_set(engram_keep, keep.data(), 0, keep.size()*sizeof(float));
+    }
+
     if (cand_pin && cand_pin->buffer) {
         const auto & plan = cand_plan(*ubatch);
         const int64_t n_blk = cand_pin->ne[0];
@@ -319,6 +332,10 @@ bool llm_graph_input_dsv41::can_reuse(const llm_graph_params & params) {
     for (auto * t : engram_ids) {
         res &= t->ne[0] == (int64_t) model.engram_n_cols()*params.ubatch.n_tokens;
     }
+    res &= (engram_keep != nullptr) == (!engram_ids.empty() && params.ubatch.embd != nullptr);
+    if (engram_keep) {
+        res &= engram_keep->ne[2] == (int64_t) params.ubatch.n_tokens;
+    }
     if (cand_pin) {
         const auto & plan = cand_plan(params.ubatch);
         const int64_t bs = model.hparams.dsv41_cand_block_size;
@@ -332,6 +349,7 @@ ggml_tensor * llama_model_deepseek41::graph::build_engram(
         const llama_model_deepseek41 & model,
         ggml_tensor * x,
         ggml_tensor * ids,
+        ggml_tensor * keep,
         int il) const {
     const auto & layer = model.layers[il];
     const int64_t hc = hparams.dsv4_hc_mult;
@@ -348,14 +366,19 @@ ggml_tensor * llama_model_deepseek41::graph::build_engram(
     ggml_tensor * value = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd, 1,  nt, ggml_row_size(kv->type, n_embd), kv->nb[1],
             ggml_row_size(kv->type, hc*n_embd)));
 
-    ggml_tensor * dot = ggml_mul(ctx0, ggml_mul(ctx0, ggml_rms_norm(ctx0, x, norm_rms_eps), ggml_rms_norm(ctx0, key, norm_rms_eps)), layer.engram_gate);
+    // quantized copy has no binary-op kernel
+    ggml_tensor * egate = ggml_is_quantized(layer.engram_gate->type) ? ggml_cast(ctx0, layer.engram_gate, GGML_TYPE_F32) : layer.engram_gate;
+    ggml_tensor * dot = ggml_mul(ctx0, ggml_mul(ctx0, ggml_rms_norm(ctx0, x, norm_rms_eps), ggml_rms_norm(ctx0, key, norm_rms_eps)), egate);
     dot = ggml_scale(ctx0, ggml_sum_rows(ctx0, dot), 1.0f/sqrtf((float) n_embd));
     cb(dot, "engram_dot", il);
 
     // sigmoid(copysign(sqrt(max(|dot|, 1e-6)), dot)); sign(0) = +1 as in copysign
     ggml_tensor * mag  = ggml_sqrt(ctx0, ggml_clamp(ctx0, ggml_abs(ctx0, dot), 1e-6f, INFINITY));
     ggml_tensor * sign = ggml_scale_bias(ctx0, ggml_step(ctx0, ggml_neg(ctx0, dot)), -2.0f, 1.0f);
     ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_mul(ctx0, mag, sign));
+    if (keep) {
+        gate = ggml_mul(ctx0, gate, keep);
+    }
     cb(gate, "engram_gate", il);
 
     return ggml_add(ctx0, x, ggml_mul(ctx0, ggml_repeat_4d(ctx0, value, n_embd, hc, nt, 1), gate));
@@ -476,8 +499,10 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
 
         ggml_tensor * top_k = topk_of[hparams.dsv41_topk_src[il]];
         if (hparams.dsv41_topk_src[il] == il) {
-            ggml_tensor * ik = ggml_cont(ctx0, ggml_view_4d(ctx0, comp_all, n_idx_embd, 1, n_comp, n_stream,
-                    comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], ggml_row_size(comp_all->type, n_embd_head)));
+            ggml_tensor * ik = ggml_view_4d(ctx0, comp_all, n_idx_embd, 1, n_comp, n_stream,
+                    comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], ggml_row_size(comp_all->type, n_embd_head));
+            // CUDA has no same-type copy of quantized views
+            ik = ggml_is_quantized(ik->type) ? ggml_cast(ctx0, ik, GGML_TYPE_F32) : ggml_cont(ctx0, ik);
 
             ggml_tensor * iq = ggml_reshape_3d(ctx0, build_lora_mm(layer.indexer_attn_q_b, qr), n_idx_embd, n_idx_head, nt);
             iq = rope(iq, inp_pos, n_idx_embd - n_embd_head_rope);
@@ -566,9 +591,11 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
         }
         GGML_ASSERT(top_k && "DeepSeek-V4.1: top-k source layer produced no top-k");
 
-        if (inp_attn->self_k_rot) {
-            // thrown at the first graph reserve, so context creation fails cleanly
-            throw std::runtime_error("DeepSeek-V4.1: a quantized (rotated) KV cache is not supported yet");
+        // a quantized K cache stores the window Hadamard-rotated (as build_raw_attention)
+        ggml_tensor * k_rot = inp_attn->self_k_rot;
+        if (k_rot) {
+            q  = llama_mul_mat_hadamard(ctx0, q, k_rot);
+            kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);
         }
         ggml_build_forward_expand(gf, q);
         ggml_build_forward_expand(gf, kv);
@@ -577,14 +604,30 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
         ggml_tensor * raw_k  = inp_attn->mctx->get_k(ctx0, il);
         ggml_tensor * comp_k = ggml_view_4d(ctx0, comp_all, n_embd_head, 1, n_comp, n_stream,
                 comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], 0);
-        comp_k = comp_k->type != raw_k->type ? ggml_cast(ctx0, comp_k, raw_k->type) : ggml_cont(ctx0, comp_k);
+        if (k_rot) {
+            // compressed rows are stored unrotated: rotate in F32, then cast to the cache type to concat with the window
+            comp_k = comp_k->type != GGML_TYPE_F32 ? ggml_cast(ctx0, comp_k, GGML_TYPE_F32) : ggml_cont(ctx0, comp_k);
+            comp_k = llama_mul_mat_hadamard(ctx0, comp_k, k_rot);
+            if (comp_k->type != raw_k->type) {
+                comp_k = ggml_cast(ctx0, comp_k, raw_k->type);
+            }
+        } else if (ggml_is_quantized(comp_k->type)) {
+            // LLAMA_ATTN_ROT_DISABLE=1: go through F32 for the same reason
+            comp_k = ggml_cast(ctx0, ggml_cast(ctx0, comp_k, GGML_TYPE_F32), raw_k->type);
+        } else {
+            comp_k = comp_k->type != raw_k->type ? ggml_cast(ctx0, comp_k, raw_k->type) : ggml_cont(ctx0, comp_k);
+        }
+        cb(comp_k, "comp_k", il);
         ggml_tensor * k_all = ggml_concat(ctx0, raw_k, comp_k, 2);
 
         ggml_tensor * sel_mask = build_top_k_mask(ci.kq_mask, top_k, "comp_top_k_mask", il);
         ggml_tensor * kq_mask  = ggml_concat(ctx0, inp_attn->get_kq_mask(), sel_mask, 0);
         cb(kq_mask, "comp_kq_mask", il);
 
         out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, layer.attn_sinks, nullptr, 0, kq_scale, il);
+        if (k_rot) {
+            out = llama_mul_mat_hadamard(ctx0, out, k_rot);
+        }
         cb(out, "attn_comp", il);
     }
 
@@ -627,6 +670,13 @@ llama_model_deepseek41::graph::graph(const llama_model & model_base, const llm_g
         ggml_set_name(ids, ("engram_ids-" + std::to_string(model.engram_layers[e])).c_str());
         inp41->engram_ids.push_back(ids);
     }
+    ggml_tensor * engram_keep = nullptr;
+    if (!model.engram_layers.empty() && 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");
+        inp41->engram_keep = engram_keep;
+    }
     ggml_tensor * cand_pin = nullptr;
     if (hparams.dsv41_cand_src >= 0) {
         inp41->cand_csa = hparams.dsv4_compress_ratios[hparams.dsv41_cand_src] == 2;
@@ -654,7 +704,7 @@ llama_model_deepseek41::graph::graph(const llama_model & model_base, const llm_g
 
         const int e = model.engram_index(il);
         if (e >= 0 && !inp41_ptr->engram_ids.empty()) {
-            inpL = build_engram(model, inpL, inp41_ptr->engram_ids[e], il);
+            inpL = build_engram(model, inpL, inp41_ptr->engram_ids[e], engram_keep, il);
             cb(inpL, "engram_out", il);
         }
 

src/models/models.h

@@ -1473,6 +1473,7 @@ struct llama_model_deepseek41 : public llama_model_deepseek4 {
                 const llama_model_deepseek41 & model,
                 ggml_tensor * x,
                 ggml_tensor * ids,
+                ggml_tensor * keep,
                 int il) const;
 
         ggml_tensor * build_attention41(