deepseek41 : add compressed attention runtime

src/CMakeLists.txt

@@ -16,6 +16,7 @@ set(LLAMA_CORE_SOURCES
     llama-chat.cpp
     llama-context.cpp
     llama-cparams.cpp
+    llama-dsv41.cpp
     llama-grammar.cpp
     llama-graph.cpp
     llama-hparams.cpp

src/llama-dsv41.h

@@ -0,0 +1,234 @@
+#pragma once
+
+#include "llama.h"
+
+#include <cstdint>
+#include <map>
+#include <string>
+#include <vector>
+
+struct ggml_context;
+struct ggml_tensor;
+
+static constexpr uint32_t LLAMA_DSV41_N_LAYER                 = 40;
+static constexpr uint32_t LLAMA_DSV41_N_EMBD                  = 5120;
+static constexpr uint32_t LLAMA_DSV41_N_VOCAB                 = 129280;
+static constexpr uint32_t LLAMA_DSV41_N_HEAD                  = 64;
+static constexpr uint32_t LLAMA_DSV41_N_HEAD_KV               = 1;
+static constexpr uint32_t LLAMA_DSV41_N_HEAD_DIM              = 512;
+static constexpr uint32_t LLAMA_DSV41_N_ROT                   = 64;
+static constexpr uint32_t LLAMA_DSV41_N_LORA_Q                = 1280;
+static constexpr uint32_t LLAMA_DSV41_N_LORA_O                = 1024;
+static constexpr uint32_t LLAMA_DSV41_N_O_GROUP               = 8;
+static constexpr uint32_t LLAMA_DSV41_N_FF_DENSE              = 18432;
+static constexpr uint32_t LLAMA_DSV41_N_EXPERT                = 384;
+static constexpr uint32_t LLAMA_DSV41_N_EXPERT_USED           = 6;
+static constexpr uint32_t LLAMA_DSV41_N_EXPERT_SHARED         = 1;
+static constexpr uint32_t LLAMA_DSV41_N_FF_EXP                = 2304;
+static constexpr uint32_t LLAMA_DSV41_N_INDEX_HEAD            = 32;
+static constexpr uint32_t LLAMA_DSV41_N_INDEX_HEAD_DIM        = 128;
+static constexpr uint32_t LLAMA_DSV41_N_INDEX_TOP_K           = 512;
+static constexpr uint32_t LLAMA_DSV41_N_SWA                   = 128;
+static constexpr uint32_t LLAMA_DSV41_N_CTX                   = 1048576;
+static constexpr uint32_t LLAMA_DSV41_HC_MULT                 = 4;
+static constexpr uint32_t LLAMA_DSV41_HC_SINKHORN_ITERS       = 20;
+static constexpr uint32_t LLAMA_DSV41_CANDIDATE_SOURCE_LAYER  = 20;
+static constexpr uint32_t LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS   = 2048;
+static constexpr uint32_t LLAMA_DSV41_CANDIDATE_BLOCK_SIZE    = 8;
+static constexpr uint32_t LLAMA_DSV41_ENGRAM_COMPRESSED_VOCAB = 99092;
+static constexpr uint32_t LLAMA_DSV41_ENGRAM_PAD_ID            = 2;
+static constexpr uint32_t LLAMA_DSV41_ENGRAM_PRIMES_COUNT     = 48;
+static constexpr uint32_t LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT = 8;
+static constexpr const char * LLAMA_DSV41_ENGRAM_ENCODING = "e4m3_e8m0_32_row264";
+
+struct llama_dsv41_config {
+    uint32_t n_ctx_train;
+    uint32_t n_embd;
+    uint32_t n_layer;
+    uint32_t n_vocab;
+    uint32_t n_head;
+    uint32_t n_head_kv;
+    uint32_t n_head_dim;
+    uint32_t n_rot;
+    uint32_t n_lora_q;
+    uint32_t n_lora_o;
+    uint32_t n_o_group;
+    uint32_t n_ff_dense;
+    uint32_t n_ff_expert;
+    uint32_t n_expert;
+    uint32_t n_expert_used;
+    uint32_t n_expert_shared;
+    uint32_t indexer_n_head;
+    uint32_t indexer_head_size;
+    uint32_t indexer_top_k;
+    uint32_t hc_count;
+    uint32_t hc_sinkhorn_iters;
+    uint32_t raw_window;
+    uint32_t candidate_source_layer;
+    uint32_t candidate_topk_blocks;
+    uint32_t candidate_block_size;
+    float f_norm_rms_eps;
+    float hc_eps;
+    float swiglu_clamp;
+    float routed_scale;
+    uint32_t rope_theta;
+    uint32_t compress_rope_theta;
+    float yarn_factor;
+    float yarn_beta_fast;
+    float yarn_beta_slow;
+    float yarn_original_context;
+    bool expert_weights_norm;
+    std::string hidden_act;
+    std::string scoring_func;
+    std::string topk_method;
+    std::vector<uint32_t> compress_ratios;
+    std::vector<uint32_t> kv_sources;
+    std::vector<uint32_t> index_sources;
+    std::vector<uint32_t> engram_layers;
+    std::vector<uint32_t> engram_rows;
+    std::string engram_encoding;
+    uint32_t engram_compressed_vocab_size;
+    uint32_t engram_pad_id;
+    uint32_t engram_token_map_size;
+    uint32_t engram_primes_size;
+    uint32_t engram_multipliers_size;
+};
+
+void llama_dsv41_validate_config(const llama_dsv41_config & config);
+const char * llama_dsv41_runtime_dependency_error();
+
+struct llama_dsv41_compression_plan {
+    std::vector<int32_t> state_pos;
+    std::vector<int32_t> state_persist_src_idxs;
+    std::vector<int32_t> state_persist_dst_idxs;
+    std::vector<int32_t> state_read_idxs;
+    std::vector<int64_t> write_idxs;
+    std::vector<int32_t> write_pos;
+    std::vector<int32_t> n_visible;
+    int64_t n_kv = 0;
+};
+
+struct llama_dsv41_layer_plan {
+    uint32_t layer;
+    uint32_t ratio;
+    int32_t kv_source_layer;
+    int32_t index_source_layer;
+    bool owns_kv_source;
+    bool owns_index_source;
+    bool builds_candidates;
+    bool uses_candidates;
+    bool reuses_index_selection;
+    bool collapses_output;
+    std::vector<uint32_t> raw_ring_order;
+    llama_dsv41_compression_plan compression;
+};
+
+class llama_dsv41_cache_state {
+public:
+    explicit llama_dsv41_cache_state(uint32_t compressed_cache_size);
+
+    void clear();
+    void append(llama_pos pos);
+    void set_candidate_blocks(const std::vector<int32_t> & blocks);
+
+    llama_pos position() const;
+    const std::vector<llama_pos> & raw_slots() const;
+    const std::vector<llama_pos> & compressed_slots(uint32_t source_layer) const;
+    const std::vector<llama_pos> & pending_slots(uint32_t source_layer) const;
+    const std::vector<int32_t> & candidate_blocks() const;
+
+private:
+    struct source_state {
+        uint32_t ratio;
+        std::vector<llama_pos> compressed;
+        std::vector<llama_pos> pending;
+    };
+
+    uint32_t compressed_cache_size;
+    llama_pos pos = -1;
+    std::vector<llama_pos> raw;
+    std::map<uint32_t, source_state> sources;
+    std::vector<int32_t> candidates;
+};
+
+struct llama_dsv41_memory_accounting {
+    uint64_t raw_kv = 0;
+    uint64_t compressed_kv = 0;
+    uint64_t index_keys = 0;
+    uint64_t compressor_carry = 0;
+    uint64_t candidate_scores = 0;
+    uint64_t candidate_ids = 0;
+    uint64_t position_state = 0;
+    uint64_t graph_workspace = 0;
+
+    uint64_t total() const;
+};
+
+llama_dsv41_memory_accounting llama_dsv41_account_memory(
+        uint32_t n_ctx,
+        uint32_t n_seq,
+        uint32_t n_tokens,
+        uint32_t kv_element_size,
+        uint32_t index_element_size,
+        uint64_t graph_workspace);
+
+int32_t llama_dsv41_kv_source_layer(uint32_t il);
+int32_t llama_dsv41_index_source_layer(uint32_t il);
+uint32_t llama_dsv41_compress_ratio(uint32_t il);
+
+llama_dsv41_layer_plan llama_dsv41_build_layer_plan(
+        uint32_t il,
+        const std::vector<llama_pos> & positions,
+        uint32_t compressed_cache_size);
+
+llama_dsv41_compression_plan llama_dsv41_build_compression_plan(
+        const std::vector<llama_pos> & positions,
+        uint32_t ratio,
+        uint32_t cache_size);
+
+std::vector<int32_t> llama_dsv41_select_candidate_blocks(
+        const std::vector<float> & scores,
+        uint32_t n_visible,
+        uint32_t block_size,
+        uint32_t top_k_blocks);
+
+std::vector<int32_t> llama_dsv41_candidate_rows(
+        const std::vector<int32_t> & blocks,
+        uint32_t n_visible,
+        uint32_t block_size);
+
+std::vector<uint32_t> llama_dsv41_raw_ring_order(llama_pos pos, uint32_t window);
+
+std::vector<float> llama_dsv41_output_collapse(
+        const std::vector<float> & residual,
+        const std::vector<float> & pre,
+        uint32_t n_embd,
+        uint32_t hc_mult);
+
+ggml_tensor * llama_dsv41_build_ratio_pool(
+        ggml_context * ctx,
+        ggml_tensor * kv,
+        ggml_tensor * gate,
+        uint32_t ratio);
+
+ggml_tensor * llama_dsv41_build_shared_softmax(
+        ggml_context * ctx,
+        ggml_tensor * raw_scores,
+        ggml_tensor * compressed_scores);
+
+ggml_tensor * llama_dsv41_build_output_collapse(
+        ggml_context * ctx,
+        ggml_tensor * residual,
+        ggml_tensor * pre,
+        uint32_t n_embd,
+        uint32_t hc_mult,
+        uint32_t n_tokens);
+
+ggml_tensor * llama_dsv41_build_output(
+        ggml_context * ctx,
+        ggml_tensor * residual,
+        ggml_tensor * pre,
+        ggml_tensor * output_norm,
+        ggml_tensor * output,
+        float rms_eps,
+        uint32_t hc_mult);

src/llama-hparams.cpp

@@ -1,5 +1,6 @@
 #include "llama-hparams.h"
 
+#include "llama-dsv41.h"
 #include "ggml.h"
 
 #include <algorithm>
@@ -47,6 +48,28 @@ bool llama_hparams::is_swa_any() const {
     return false;
 }
 
+int32_t llama_hparams::dsv41_kv_source(uint32_t il) const {
+    if (il >= n_layer()) {
+        GGML_ABORT("fatal error");
+    }
+    return dsv41_kv_source_layer[il];
+}
+
+int32_t llama_hparams::dsv41_index_source(uint32_t il) const {
+    if (il >= n_layer()) {
+        GGML_ABORT("fatal error");
+    }
+    return dsv41_index_source_layer[il];
+}
+
+bool llama_hparams::dsv41_is_kv_source(uint32_t il) const {
+    return dsv41_kv_source(il) == (int32_t) il;
+}
+
+bool llama_hparams::dsv41_is_index_source(uint32_t il) const {
+    return dsv41_index_source(il) == (int32_t) il;
+}
+
 uint32_t llama_hparams::n_head(uint32_t il) const {
     if (il < n_layer_all) {
         return n_head_arr[il];

src/llama-hparams.h

@@ -294,6 +294,19 @@ struct llama_hparams {
     float    dsv4_hc_eps               = 0.0f;
     std::array<uint32_t, LLAMA_MAX_LAYERS> dsv4_compress_ratios;
 
+    // DeepSeek-V4.1
+    uint32_t dsv41_candidate_source_layer = 0;
+    uint32_t dsv41_candidate_topk_blocks  = 0;
+    uint32_t dsv41_candidate_block_size   = 0;
+    std::array<int32_t, LLAMA_MAX_LAYERS> dsv41_kv_source_layer;
+    std::array<int32_t, LLAMA_MAX_LAYERS> dsv41_index_source_layer;
+    std::bitset<LLAMA_MAX_LAYERS> dsv41_engram_layers;
+
+    int32_t dsv41_kv_source(uint32_t il) const;
+    int32_t dsv41_index_source(uint32_t il) const;
+    bool dsv41_is_kv_source(uint32_t il) const;
+    bool dsv41_is_index_source(uint32_t il) const;
+
     // 0 = full rank (DeepSeek-V4)
     uint32_t hc_low_rank = 0;
 

src/llama-model.cpp

@@ -6,6 +6,7 @@
 #include "llama-impl.h"
 #include "llama-mmap.h"
 #include "llama-cparams.h"
+#include "llama-dsv41.h"
 #include "llama-model-loader.h"
 
 #include "llama-kv-cache.h"
@@ -200,6 +201,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
             return new llama_model_dots3note(params);
         case LLM_ARCH_DEEPSEEK4:
             return new llama_model_deepseek4(params);
+        case LLM_ARCH_DEEPSEEK41:
+            return new llama_model_deepseek41(params);
         case LLM_ARCH_GLM_DSA:
             return new llama_model_glm_dsa(params);
         case LLM_ARCH_MISTRAL4:
@@ -373,6 +376,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     const llama_hparams & hparams = ud->model->hparams;
     const std::string tensor_name = tensor->name;
     const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 ||
+        ud->model->arch == LLM_ARCH_DEEPSEEK41 ||
         (ud->model->arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0);
 
     static const std::regex pattern_q_weight        ("blk\\.\\d*\\.attn_q.weight");
@@ -384,7 +388,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     static const std::regex pattern_qk_norm         ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight");
     static const std::regex pattern_kv_cache        ("cache_(k|v)_l\\d*");
     static const std::regex pattern_idx_cache       ("cache_idx_(k|v)_l\\d*");
-    static const std::regex pattern_dsv4_state      ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*");
+    static const std::regex pattern_dsv4_state      ("dsv4(1)?_(csa|hca|lid|comp|index)_state_(kv|score)_l\\d*");
     static const std::regex pattern_attn_sinks      ("blk\\.\\d*\\.attn_sinks.weight");
     static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight");
     static const std::regex pattern_attn_out_bias   ("blk\\.\\d*\\.attn_output.bias");
@@ -1224,6 +1228,33 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
         return;
     }
 
+    if (arch == LLM_ARCH_DEEPSEEK41) {
+        std::fill(hparams.n_head_arr.begin(),       hparams.n_head_arr.end(),       0);
+        std::fill(hparams.n_head_kv_arr.begin(),    hparams.n_head_kv_arr.end(),    0);
+        std::fill(hparams.n_ff_arr.begin(),         hparams.n_ff_arr.end(),         0);
+        std::fill(hparams.n_ff_exp_arr.begin(),     hparams.n_ff_exp_arr.end(),     0);
+        std::fill(hparams.n_expert_used_arr.begin(), hparams.n_expert_used_arr.end(), 0);
+        std::fill(hparams.rope_sections.begin(),    hparams.rope_sections.end(),    0);
+        std::fill(hparams.rope_pattern.begin(),     hparams.rope_pattern.end(),     1);
+        std::fill(hparams.is_swa_impl.begin(),      hparams.is_swa_impl.end(),      0);
+        std::fill(hparams.is_recr_impl.begin(),     hparams.is_recr_impl.end(),     0);
+        std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 0);
+        std::fill(hparams.dsv41_kv_source_layer.begin(), hparams.dsv41_kv_source_layer.end(), -1);
+        std::fill(hparams.dsv41_index_source_layer.begin(), hparams.dsv41_index_source_layer.end(), -1);
+        std::fill(hparams.dsv4_compress_ratios.begin(), hparams.dsv4_compress_ratios.end(), 0);
+        std::fill(hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.end(), 0.0f);
+        std::fill(hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.end(), 0.0f);
+        hparams.dsv41_engram_layers.reset();
+
+        load_arch_hparams(ml);
+
+        pimpl->n_bytes = ml.n_bytes;
+        pimpl->desc_str = arch_name() + " " + type_name() + " " + ml.ftype_name();
+        pimpl->ftype = ml.ftype;
+        hparams.rope_type = llama_model_rope_type(this);
+        return;
+    }
+
     ml.get_key(LLM_KV_CONTEXT_LENGTH,          hparams.n_ctx_train);
     ml.get_key(LLM_KV_EMBEDDING_LENGTH,        hparams.n_embd);
     ml.get_key(LLM_KV_EMBEDDING_LENGTH_OUT,    hparams.n_embd_out_impl, false);
@@ -2480,6 +2511,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                             nullptr);
                 }
             } break;
+        case LLM_ARCH_DEEPSEEK41:
+            throw std::runtime_error(llama_dsv41_runtime_dependency_error());
         case LLM_ARCH_DFLASH:
             {
                 // DSV4 DSpark stages store a single MLA-style K per position (window = the draft ring)
@@ -2838,7 +2871,7 @@ int32_t llama_model_n_head_kv(const llama_model * model) {
 int32_t llama_model_n_swa(const llama_model * model) {
     // dsv4 kv-cache has SWA but it cannot be used as a rollback because of
     // other compression ratios, so we return 0 here
-    if (model->arch == LLM_ARCH_DEEPSEEK4) {
+    if (model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) {
         return 0;
     }
     return model->hparams.n_swa;