From 9147b60536f41c26567b3eb700b271a398129650 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 21:03:10 -0700 Subject: [PATCH 1/5] deepseek41 : activate scheduler-backed full graph Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/CMakeLists.txt | 1 + src/llama-arch.cpp | 1 + src/llama-context.cpp | 33 +- src/llama-dsv41.cpp | 4 +- src/llama-dsv41.h | 1 - src/llama-graph.cpp | 6 +- src/llama-kv-cache.cpp | 1 + src/llama-memory-dsv41.cpp | 1540 +++++++++++++++++++++++++++++ src/llama-memory-dsv41.h | 173 ++++ src/llama-memory.h | 5 + src/llama-model.cpp | 13 +- src/models/deepseek41.cpp | 1205 +++++++++++++++++++++- src/models/models.h | 9 +- tests/CMakeLists.txt | 1 + tests/test-deepseek41-memory.cpp | 642 ++++++++++++ tests/test-deepseek41-runtime.cpp | 220 ++++- tests/test-llama-archs.cpp | 2 +- 17 files changed, 3843 insertions(+), 14 deletions(-) create mode 100644 src/llama-memory-dsv41.cpp create mode 100644 src/llama-memory-dsv41.h create mode 100644 tests/test-deepseek41-memory.cpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7c403a37e0d9..971d75f5706a 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -32,6 +32,7 @@ set(LLAMA_CORE_SOURCES llama-kv-cache-msa.cpp llama-kv-cache-dsv4.cpp llama-memory.cpp + llama-memory-dsv41.cpp llama-memory-hybrid.cpp llama-memory-hybrid-iswa.cpp llama-memory-hybrid-idx.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 996f18973127..dd22cfbc17a9 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1142,6 +1142,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_QWEN35MOE: case LLM_ARCH_QWEN4EXP: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_MINIMAX_01: return true; default: diff --git a/src/llama-context.cpp b/src/llama-context.cpp index bf77844264c0..ad8f71bac9be 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -691,18 +691,23 @@ void llama_context::sched_reserve() { } } + size_t graph_workspace_size = 0; for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; ggml_backend_buffer_type_t buft = backend_buft[i]; if (!model.hparams.no_alloc) { backend_buf_exp_size[i] = ggml_backend_sched_get_buffer_size(sched.get(), backend); } + graph_workspace_size += backend_buf_exp_size[i]; if (backend_buf_exp_size[i] > 1) { LLAMA_LOG_INFO("%s: %10s compute buffer size = %8.2f MiB\n", __func__, ggml_backend_buft_name(buft), backend_buf_exp_size[i] / 1024.0 / 1024.0); } } + if (memory) { + memory->set_graph_workspace_size(graph_workspace_size); + } if (n_nodes_pp == n_nodes_tg) { LLAMA_LOG_INFO("%s: graph nodes = %d\n", __func__, n_nodes_pp); @@ -1355,11 +1360,34 @@ bool llama_context::set_adapter_cvec( llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, llm_graph_type gtype, llama_memory_context_i * mctx, ggml_status & ret) { if (mctx && !mctx->apply()) { + mctx->rollback(); LLAMA_LOG_ERROR("%s: failed to apply memory context\n", __func__); ret = GGML_STATUS_FAILED; return nullptr; } + struct memory_transaction_guard { + llama_memory_context_i * context; + bool active; + + ~memory_transaction_guard() { + if (active) { + try { + context->rollback(); + } catch (const std::exception & error) { + LLAMA_LOG_ERROR("%s: memory rollback failed: %s\n", __func__, error.what()); + } + } + } + + void commit() { + if (active) { + context->commit(); + active = false; + } + } + } transaction_guard { mctx, mctx != nullptr }; + auto * res = gf_res_prev.get(); auto * gf = res->get_gf(); @@ -1442,6 +1470,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll } } + transaction_guard.commit(); ret = GGML_STATUS_SUCCESS; return res; @@ -2358,6 +2387,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_QWEN4EXP || model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_DEEPSEEK41 || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_01 || @@ -3737,7 +3767,8 @@ llama_context * llama_init_from_model( } } - if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && params.type_k != params.type_v) { + if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) && + params.type_k != params.type_v) { LLAMA_LOG_ERROR("%s: model does not support different K (%s) and V (%s) cache types\n", __func__, ggml_type_name(params.type_k), ggml_type_name(params.type_v)); return nullptr; } diff --git a/src/llama-dsv41.cpp b/src/llama-dsv41.cpp index 5de36da9da8f..5faf063ac12a 100644 --- a/src/llama-dsv41.cpp +++ b/src/llama-dsv41.cpp @@ -106,9 +106,7 @@ llama_engram_layout llama_dsv41_make_engram_layout(const llama_dsv41_config & co return layout; } -const char * llama_dsv41_runtime_dependency_error() { - return "DeepSeek V4.1 execution requires routed-expert streaming support"; -} + static int32_t dsv41_source_layer(const int32_t * sources, size_t n, uint32_t il) { int32_t result = -1; diff --git a/src/llama-dsv41.h b/src/llama-dsv41.h index 19d73a2f04dd..8cd8c55f22fa 100644 --- a/src/llama-dsv41.h +++ b/src/llama-dsv41.h @@ -101,7 +101,6 @@ struct llama_dsv41_config { void llama_dsv41_validate_config(const llama_dsv41_config & config); llama_engram_layout llama_dsv41_make_engram_layout(const llama_dsv41_config & config); -const char * llama_dsv41_runtime_dependency_error(); struct llama_dsv41_compression_plan { std::vector state_pos; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 9a7219184874..3c9debeb5965 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1834,7 +1834,8 @@ ggml_tensor * llm_graph_context::build_ffn( const float limit = hparams.swiglu_clamp_shexp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || + (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_swiglu_clamp(ctx0, cur, tmp, limit); } else { tmp = ggml_clamp(ctx0, tmp, -limit, limit); @@ -2257,7 +2258,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( const float limit = hparams.swiglu_clamp_exp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || + (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) { cur = ggml_swiglu_clamp(ctx0, cur, up, limit); } else { up = ggml_clamp(ctx0, up, -limit, limit); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index a5c0970e9bb3..1a284eac0440 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -341,6 +341,7 @@ llama_kv_cache::llama_kv_cache( // always create Hadamard rotation tensors for DeepSeek lightning indexers if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_DEEPSEEK41 || model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_DOTS3NOTE) && hparams.n_embd_head_k_full == hparams.indexer_head_size) { attn_rot_k = true; diff --git a/src/llama-memory-dsv41.cpp b/src/llama-memory-dsv41.cpp new file mode 100644 index 000000000000..42f3ed8c7215 --- /dev/null +++ b/src/llama-memory-dsv41.cpp @@ -0,0 +1,1540 @@ +#include "llama-memory-dsv41.h" + +#include "ggml-backend.h" +#include "ggml-cpp.h" +#include "llama-batch.h" +#include "llama-impl.h" +#include "llama-io.h" +#include "llama-model.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr uint64_t DSV41_STATE_MAGIC = 0x314d454d31345644ULL; +constexpr uint32_t DSV41_STATE_VERSION = 1; + +uint64_t hash_mix(uint64_t hash, uint64_t value) { + hash ^= value; + return hash*1099511628211ULL; +} + +uint64_t hash_string(uint64_t hash, const char * value) { + while (*value != '\0') { + hash = hash_mix(hash, (uint8_t) *value++); + } + return hash; +} + +uint32_t source_capacity(const llama_dsv41_memory_config & config, uint32_t source) { + const uint32_t ratio = config.ratios.at(source); + return (config.n_ctx + ratio - 1)/ratio; +} + +llama_dsv41_memory_config make_model_config( + const llama_model & model, + ggml_type type_k, + bool offload, + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_ubatch, + std::unique_ptr engram) { + llama_dsv41_memory_config config; + config.n_ctx = n_ctx; + config.n_seq = n_seq; + config.n_ubatch = n_ubatch; + config.n_layer = model.hparams.n_layer(); + config.raw_window = model.hparams.n_swa; + config.kv_width = model.hparams.n_embd_head_k(); + config.index_width = model.hparams.indexer_head_size; + config.candidate_topk_blocks = model.hparams.dsv41_candidate_topk_blocks; + config.candidate_block_size = model.hparams.dsv41_candidate_block_size; + config.candidate_source_layer = model.hparams.dsv41_candidate_source_layer; + config.type_k = type_k; + config.type_index = type_k; + config.no_alloc = model.hparams.no_alloc; + config.expert_enabled = model.requires_synchronous_graph(); + config.ratios.resize(config.n_layer); + config.kv_sources.clear(); + config.index_sources.clear(); + for (uint32_t il = 0; il < config.n_layer; ++il) { + config.ratios[il] = model.hparams.dsv4_compress_ratios[il]; + if (model.hparams.dsv41_is_kv_source(il)) { + config.kv_sources.push_back(il); + } + if (model.hparams.dsv41_is_index_source(il)) { + config.index_sources.push_back(il); + } + } + config.buft_for_layer = [&model, offload](int32_t il) { + return offload ? model.select_buft(il) : ggml_backend_cpu_buffer_type(); + }; + config.engram = std::move(engram); + return config; +} + +} + +bool llama_dsv41_graph_topology::same_topology(const llama_dsv41_graph_topology & other) const { + return n_tokens == other.n_tokens && + n_seqs == other.n_seqs && + n_outputs == other.n_outputs && + visible_raw_widths == other.visible_raw_widths && + visible_compressed_widths == other.visible_compressed_widths && + source_ratios == other.source_ratios && + source_carry_counts == other.source_carry_counts && + candidate_width == other.candidate_width && + backend_layout == other.backend_layout && + engram_enabled == other.engram_enabled && + expert_enabled == other.expert_enabled; +} + +struct llama_memory_dsv41::impl { + struct sequence_state { + llama_pos pos = -1; + std::vector candidates; + }; + + struct tensor_snapshot { + ggml_tensor * tensor = nullptr; + size_t offset = 0; + std::vector data; + }; + + struct rollback_state { + llama_pos start_pos = -1; + sequence_state sequence; + std::vector snapshots; + llama_dsv41_engram_sequence_state engram; + bool has_engram = false; + }; + + struct source_storage { + uint32_t layer = 0; + uint32_t ratio = 0; + uint32_t capacity = 0; + ggml_tensor * kv = nullptr; + ggml_tensor * index = nullptr; + ggml_tensor * carry_kv = nullptr; + ggml_tensor * carry_score = nullptr; + }; + + struct buffer_group { + ggml_backend_buffer_type_t buft = nullptr; + ggml_context_ptr ctx; + ggml_backend_buffer_ptr buffer; + }; + + llama_dsv41_memory_config config; + std::vector sequences; + std::vector raw; + std::map sources; + std::vector groups; + ggml_tensor * candidate_scores = nullptr; + ggml_tensor * candidate_ids = nullptr; + ggml_tensor * committed_candidate_ids = nullptr; + ggml_tensor * positions = nullptr; + uint32_t position_rows = 0; + uint32_t candidate_blocks = 0; + uint32_t candidate_width = 0; + uint64_t generation = 0; + uint64_t backend_layout = 1469598103934665603ULL; + uint64_t graph_workspace = 0; + bool transaction_active = false; + std::map> engram_boundaries; + std::map rollback_states; + + explicit impl(llama_dsv41_memory_config config) : config(std::move(config)) { + if (this->config.n_ctx == 0 || this->config.n_seq == 0 || this->config.n_ubatch == 0 || + this->config.n_layer == 0 || this->config.raw_window == 0 || + this->config.kv_width == 0 || this->config.index_width == 0 || + this->config.candidate_block_size == 0 || this->config.candidate_topk_blocks == 0) { + throw std::invalid_argument("DeepSeek V4.1 memory dimensions must be non-zero"); + } + if (this->config.ratios.empty()) { + this->config.ratios.resize(this->config.n_layer); + for (uint32_t il = 0; il < this->config.n_layer; ++il) { + this->config.ratios[il] = llama_dsv41_compress_ratio(il); + } + } + if (this->config.ratios.size() != this->config.n_layer) { + throw std::invalid_argument("DeepSeek V4.1 memory ratio map has the wrong size"); + } + if (this->config.candidate_source_layer >= this->config.n_layer) { + throw std::invalid_argument("DeepSeek V4.1 candidate source layer is out of range"); + } + if (!this->config.buft_for_layer) { + this->config.buft_for_layer = [](int32_t) { return ggml_backend_cpu_buffer_type(); }; + } + for (uint32_t source : this->config.kv_sources) { + if (source >= this->config.n_layer || (this->config.ratios[source] != 1 && this->config.ratios[source] != 2)) { + throw std::invalid_argument("DeepSeek V4.1 memory KV source is invalid"); + } + } + for (uint32_t source : this->config.index_sources) { + if (source >= this->config.n_layer) { + throw std::invalid_argument("DeepSeek V4.1 memory index source is invalid"); + } + } + if (this->config.engram && this->config.engram->max_tokens() < this->config.n_ubatch) { + throw std::invalid_argument("DeepSeek V4.1 Engram runtime is smaller than n_ubatch"); + } + + sequences.resize(this->config.n_seq); + raw.resize(this->config.n_layer); + candidate_blocks = (this->config.n_ctx + this->config.candidate_block_size - 1)/ + this->config.candidate_block_size; + candidate_width = std::min(candidate_blocks, this->config.candidate_topk_blocks); + + std::map tensor_counts; + for (uint32_t il = 0; il < this->config.n_layer; ++il) { + tensor_counts[this->config.buft_for_layer(il)]++; + } + for (uint32_t source : this->config.kv_sources) { + tensor_counts[this->config.buft_for_layer(source)] += this->config.ratios[source] == 2 ? 4 : 2; + } + ggml_backend_buffer_type_t candidate_buft = + this->config.buft_for_layer(this->config.candidate_source_layer); + tensor_counts[candidate_buft] += 3; + tensor_counts[ggml_backend_cpu_buffer_type()]++; + + for (const auto & entry : tensor_counts) { + ggml_init_params params = { + /*.mem_size =*/ (entry.second + 4)*ggml_tensor_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + if (ctx == nullptr) { + throw std::runtime_error("failed to create DeepSeek V4.1 memory tensor context"); + } + groups.push_back({ entry.first, ggml_context_ptr(ctx), nullptr }); + } + + auto context_for = [&](ggml_backend_buffer_type_t buft) -> ggml_context * { + for (auto & group : groups) { + if (group.buft == buft) { + return group.ctx.get(); + } + } + throw std::runtime_error("DeepSeek V4.1 memory buffer type is missing"); + }; + + for (uint32_t il = 0; il < this->config.n_layer; ++il) { + ggml_context * ctx = context_for(this->config.buft_for_layer(il)); + raw[il] = ggml_new_tensor_3d( + ctx, this->config.type_k, this->config.kv_width, this->config.raw_window, this->config.n_seq); + ggml_format_name(raw[il], "dsv41_raw_k_l%u", il); + } + + for (uint32_t source : this->config.kv_sources) { + source_storage storage; + storage.layer = source; + storage.ratio = this->config.ratios[source]; + storage.capacity = source_capacity(this->config, source); + ggml_context * ctx = context_for(this->config.buft_for_layer(source)); + storage.kv = ggml_new_tensor_3d( + ctx, this->config.type_k, this->config.kv_width, storage.capacity, this->config.n_seq); + storage.index = ggml_new_tensor_3d( + ctx, this->config.type_index, this->config.index_width, storage.capacity, this->config.n_seq); + ggml_format_name(storage.kv, "dsv41_comp_kv_l%u", source); + ggml_format_name(storage.index, "dsv41_index_k_l%u", source); + if (storage.ratio == 2) { + storage.carry_kv = ggml_new_tensor_3d( + ctx, this->config.type_k, this->config.kv_width, storage.ratio, this->config.n_seq); + storage.carry_score = ggml_new_tensor_3d( + ctx, this->config.type_k, this->config.kv_width, storage.ratio, this->config.n_seq); + ggml_format_name(storage.carry_kv, "dsv41_carry_kv_l%u", source); + ggml_format_name(storage.carry_score, "dsv41_carry_score_l%u", source); + } + sources.emplace(source, storage); + } + + { + ggml_context * ctx = context_for(candidate_buft); + candidate_scores = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, candidate_blocks, this->config.n_ubatch); + candidate_ids = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, candidate_width, this->config.n_ubatch); + committed_candidate_ids = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, candidate_width, this->config.n_seq); + ggml_set_name(candidate_scores, "dsv41_candidate_scores"); + ggml_set_name(candidate_ids, "dsv41_candidate_ids"); + ggml_set_name(committed_candidate_ids, "dsv41_committed_candidate_ids"); + } + + position_rows = 1 + this->config.raw_window; + for (uint32_t source : this->config.kv_sources) { + position_rows += source_capacity(this->config, source); + if (this->config.ratios[source] == 2) { + position_rows += 2; + } + } + { + ggml_context * ctx = context_for(ggml_backend_cpu_buffer_type()); + positions = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, position_rows, this->config.n_seq); + ggml_set_name(positions, "dsv41_position_state"); + } + + for (auto & group : groups) { + const size_t buffer_size = + ggml_backend_alloc_ctx_tensors_from_buft_size(group.ctx.get(), group.buft); + if (!this->config.no_alloc) { + ggml_backend_buffer_t buffer = + ggml_backend_alloc_ctx_tensors_from_buft(group.ctx.get(), group.buft); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate DeepSeek V4.1 memory buffer"); + } + group.buffer.reset(buffer); + ggml_backend_buffer_clear(buffer, 0); + } + backend_layout = hash_string(backend_layout, ggml_backend_buft_name(group.buft)); + backend_layout = hash_mix(backend_layout, buffer_size); + } + for (ggml_tensor * tensor : all_tensors()) { + backend_layout = hash_mix(backend_layout, tensor->type); + backend_layout = hash_mix(backend_layout, tensor->ne[0]); + backend_layout = hash_mix(backend_layout, tensor->ne[1]); + backend_layout = hash_mix(backend_layout, tensor->ne[2]); + } + for (uint32_t seq = 0; seq < this->config.n_seq; ++seq) { + update_position_state(seq); + } + } + + std::vector all_tensors() const { + std::vector result; + result.insert(result.end(), raw.begin(), raw.end()); + for (const auto & entry : sources) { + result.push_back(entry.second.kv); + result.push_back(entry.second.index); + if (entry.second.carry_kv != nullptr) { + result.push_back(entry.second.carry_kv); + result.push_back(entry.second.carry_score); + } + } + result.push_back(candidate_scores); + result.push_back(candidate_ids); + result.push_back(committed_candidate_ids); + result.push_back(positions); + return result; + } + + std::vector sequence_tensors() const { + std::vector result; + result.insert(result.end(), raw.begin(), raw.end()); + for (const auto & entry : sources) { + result.push_back(entry.second.kv); + result.push_back(entry.second.index); + if (entry.second.carry_kv != nullptr) { + result.push_back(entry.second.carry_kv); + result.push_back(entry.second.carry_score); + } + } + result.push_back(committed_candidate_ids); + result.push_back(positions); + return result; + } + + size_t sequence_plane(const ggml_tensor * tensor) const { + if (tensor->ne[2] == (int64_t) config.n_seq) { + return tensor->nb[2]; + } + if (tensor->ne[1] == (int64_t) config.n_seq) { + return tensor->nb[1]; + } + throw std::runtime_error("DeepSeek V4.1 tensor has no sequence plane"); + } + + bool valid_seq(llama_seq_id seq_id) const { + return seq_id >= 0 && (uint32_t) seq_id < config.n_seq; + } + + void require_idle() const { + if (transaction_active) { + throw std::runtime_error("DeepSeek V4.1 memory operation is not allowed during a transaction"); + } + } + + void update_position_state(llama_seq_id seq_id) { + const std::vector values = position_state_values(sequences[seq_id]); + if (!config.no_alloc) { + ggml_backend_tensor_set( + positions, values.data(), (size_t) seq_id*positions->nb[1], values.size()*sizeof(int32_t)); + } + } + + std::vector position_state_values(const sequence_state & sequence) const { + std::vector values(position_rows, -1); + const llama_pos pos = sequence.pos; + values[0] = pos; + if (pos >= 0) { + const llama_pos first_raw = std::max(0, pos + 1 - config.raw_window); + for (llama_pos current = first_raw; current <= pos; ++current) { + values[1 + current%config.raw_window] = current; + } + } + uint32_t offset = 1 + config.raw_window; + for (uint32_t source : config.kv_sources) { + const uint32_t ratio = config.ratios[source]; + const uint32_t capacity = source_capacity(config, source); + const uint32_t visible = pos < 0 ? 0 : (uint32_t) (pos + 1)/ratio; + for (uint32_t row = 0; row < visible; ++row) { + values[offset + row] = row*ratio; + } + offset += capacity; + if (ratio == 2) { + for (uint32_t row = 0; row < 2; ++row) { + if (pos >= (llama_pos) row) { + values[offset + row] = pos - ((pos - row)%2); + } + } + offset += 2; + } + } + return values; + } + + void update_candidate_state(llama_seq_id seq_id) { + const std::vector values = candidate_state_values(sequences[seq_id]); + if (!config.no_alloc) { + ggml_backend_tensor_set( + committed_candidate_ids, + values.data(), + (size_t) seq_id*committed_candidate_ids->nb[1], + values.size()*sizeof(int32_t)); + } + } + + std::vector candidate_state_values(const sequence_state & sequence) const { + std::vector values(candidate_width); + const auto & candidates = sequence.candidates; + std::copy_n(candidates.begin(), std::min(candidates.size(), values.size()), values.begin()); + return values; + } + + void clear_sequence_data(llama_seq_id seq_id) { + if (config.no_alloc) { + return; + } + for (ggml_tensor * tensor : sequence_tensors()) { + const size_t plane = sequence_plane(tensor); + std::vector zeros(plane); + ggml_backend_tensor_set(tensor, zeros.data(), (size_t) seq_id*plane, plane); + } + } + + void copy_sequence_data(llama_seq_id src, llama_seq_id dst) { + if (config.no_alloc || src == dst) { + return; + } + for (ggml_tensor * tensor : sequence_tensors()) { + const size_t plane = sequence_plane(tensor); + std::vector data(plane); + ggml_backend_tensor_get(tensor, data.data(), (size_t) src*plane, plane); + ggml_backend_tensor_set(tensor, data.data(), (size_t) dst*plane, plane); + } + } +}; + +struct llama_memory_dsv41_context::transaction_state { + llama_dsv41_memory_plan plan; + std::vector next_sequences; + std::vector snapshots; + std::unique_ptr engram; + std::map engram_before; + std::vector seq_ids; + std::vector start_positions; +}; + +llama_memory_dsv41::llama_memory_dsv41(llama_dsv41_memory_config config) : + pimpl(std::make_unique(std::move(config))) { +} + +llama_memory_dsv41::llama_memory_dsv41( + const llama_model & model, + ggml_type type_k, + bool offload, + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_ubatch, + std::unique_ptr engram) : + llama_memory_dsv41(make_model_config( + model, type_k, offload, n_ctx, n_seq, n_ubatch, std::move(engram))) { +} + +llama_memory_dsv41::~llama_memory_dsv41() = default; + +llama_memory_context_ptr llama_memory_dsv41::init_batch( + llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool embd_all) { + GGML_UNUSED(embd_all); + if (n_ubatch == 0 || n_ubatch > pimpl->config.n_ubatch || pimpl->transaction_active) { + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + + balloc.split_reset(); + std::vector ubatches; + while (true) { + llama_ubatch ubatch = balloc.split_seq(n_ubatch); + if (ubatch.n_tokens == 0) { + break; + } + ubatches.push_back(std::move(ubatch)); + } + if (balloc.get_n_used() != balloc.get_n_tokens() || ubatches.empty()) { + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + return std::make_unique(this, std::move(ubatches)); +} + +llama_memory_context_ptr llama_memory_dsv41::init_full() { + return std::make_unique(this, true); +} + +llama_memory_context_ptr llama_memory_dsv41::init_update(llama_context * lctx, bool optimize) { + GGML_UNUSED(lctx); + GGML_UNUSED(optimize); + return std::make_unique(LLAMA_MEMORY_STATUS_NO_UPDATE); +} + +bool llama_memory_dsv41::get_can_shift() const { + return false; +} + +void llama_memory_dsv41::clear(bool data) { + pimpl->require_idle(); + for (auto & sequence : pimpl->sequences) { + sequence = {}; + } + pimpl->engram_boundaries.clear(); + pimpl->rollback_states.clear(); + if (pimpl->config.engram) { + for (uint32_t seq = 0; seq < pimpl->config.n_seq; ++seq) { + pimpl->config.engram->seq_remove(seq); + } + } + if (data && !pimpl->config.no_alloc) { + for (auto & group : pimpl->groups) { + ggml_backend_buffer_clear(group.buffer.get(), 0); + } + } + for (uint32_t seq = 0; seq < pimpl->config.n_seq; ++seq) { + pimpl->update_position_state(seq); + pimpl->update_candidate_state(seq); + } + ++pimpl->generation; +} + +bool llama_memory_dsv41::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { + try { + pimpl->require_idle(); + } catch (const std::exception & error) { + LLAMA_LOG_ERROR("%s: %s\n", __func__, error.what()); + return false; + } + if (seq_id < 0) { + if (p0 == -1 && p1 == -1) { + clear(true); + return true; + } + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 only supports wildcard removal for the full memory\n", __func__); + return false; + } + if (!pimpl->valid_seq(seq_id)) { + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 sequence ID %d is out of range\n", __func__, seq_id); + return false; + } + + const llama_pos max_pos = pimpl->sequences[seq_id].pos; + const llama_pos begin = p0 < 0 ? 0 : p0; + const llama_pos end = p1 < 0 ? std::numeric_limits::max() : p1; + if (begin >= end || begin > max_pos) { + return true; + } + if (end <= max_pos) { + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 only supports full removal or suffix rollback\n", __func__); + return false; + } + + auto rollback = pimpl->rollback_states.end(); + if (begin > 0) { + rollback = pimpl->rollback_states.find(seq_id); + if (rollback == pimpl->rollback_states.end() || rollback->second.start_pos != begin) { + LLAMA_LOG_ERROR("%s: DeepSeek V4.1 only supports rollback of the most recent committed ubatch\n", __func__); + return false; + } + } + + if (begin == 0) { + pimpl->sequences[seq_id] = {}; + pimpl->clear_sequence_data(seq_id); + pimpl->update_position_state(seq_id); + pimpl->update_candidate_state(seq_id); + if (pimpl->config.engram) { + pimpl->config.engram->seq_remove(seq_id); + } + pimpl->engram_boundaries.erase(seq_id); + pimpl->rollback_states.erase(seq_id); + } else { + for (const auto & snapshot : rollback->second.snapshots) { + ggml_backend_tensor_set( + snapshot.tensor, snapshot.data.data(), + snapshot.offset, snapshot.data.size()); + } + pimpl->sequences[seq_id] = rollback->second.sequence; + if (rollback->second.has_engram) { + llama_dsv41_engram_snapshot snapshot = pimpl->config.engram->checkpoint(); + snapshot.sequences[seq_id] = rollback->second.engram; + pimpl->config.engram->restore(snapshot); + auto & boundaries = pimpl->engram_boundaries[seq_id]; + boundaries.erase(boundaries.lower_bound(begin), boundaries.end()); + } + pimpl->update_position_state(seq_id); + pimpl->update_candidate_state(seq_id); + pimpl->rollback_states.erase(rollback); + } + ++pimpl->generation; + return true; +} + +void llama_memory_dsv41::seq_cp( + llama_seq_id seq_id_src, + llama_seq_id seq_id_dst, + llama_pos p0, + llama_pos p1) { + pimpl->require_idle(); + if (!pimpl->valid_seq(seq_id_src) || !pimpl->valid_seq(seq_id_dst)) { + throw std::invalid_argument("DeepSeek V4.1 sequence copy ID is out of range"); + } + const llama_pos max_pos = pimpl->sequences[seq_id_src].pos; + if ((p0 > 0) || (p1 >= 0 && p1 <= max_pos)) { + throw std::invalid_argument("DeepSeek V4.1 only supports full sequence copy"); + } + pimpl->copy_sequence_data(seq_id_src, seq_id_dst); + pimpl->sequences[seq_id_dst] = pimpl->sequences[seq_id_src]; + pimpl->rollback_states.erase(seq_id_dst); + if (pimpl->config.engram) { + pimpl->config.engram->seq_copy(seq_id_src, seq_id_dst); + pimpl->engram_boundaries[seq_id_dst] = pimpl->engram_boundaries[seq_id_src]; + } + ++pimpl->generation; +} + +void llama_memory_dsv41::seq_keep(llama_seq_id seq_id) { + pimpl->require_idle(); + if (!pimpl->valid_seq(seq_id)) { + throw std::invalid_argument("DeepSeek V4.1 sequence keep ID is out of range"); + } + for (uint32_t current = 0; current < pimpl->config.n_seq; ++current) { + if ((llama_seq_id) current == seq_id) { + continue; + } + pimpl->sequences[current] = {}; + pimpl->clear_sequence_data(current); + pimpl->update_position_state(current); + pimpl->update_candidate_state(current); + if (pimpl->config.engram) { + pimpl->config.engram->seq_remove(current); + } + pimpl->engram_boundaries.erase(current); + pimpl->rollback_states.erase(current); + } + ++pimpl->generation; +} + +[[noreturn]] void llama_memory_dsv41::seq_add( + llama_seq_id, + llama_pos, + llama_pos, + llama_pos) { + throw std::invalid_argument("DeepSeek V4.1 memory does not support position shifts"); +} + +[[noreturn]] void llama_memory_dsv41::seq_div( + llama_seq_id, + llama_pos, + llama_pos, + int) { + throw std::invalid_argument("DeepSeek V4.1 memory does not support position division"); +} + +llama_pos llama_memory_dsv41::seq_pos_min(llama_seq_id seq_id) const { + if (!pimpl->valid_seq(seq_id)) { + return -1; + } + return pimpl->sequences[seq_id].pos < 0 ? -1 : 0; +} + +llama_pos llama_memory_dsv41::seq_pos_max(llama_seq_id seq_id) const { + if (!pimpl->valid_seq(seq_id)) { + return -1; + } + return pimpl->sequences[seq_id].pos; +} + +std::map llama_memory_dsv41::memory_breakdown() const { + std::map result; + for (const auto & group : pimpl->groups) { + const size_t size = pimpl->config.no_alloc ? + ggml_backend_alloc_ctx_tensors_from_buft_size(group.ctx.get(), group.buft) : + ggml_backend_buffer_get_size(group.buffer.get()); + result[group.buft] += size; + } + return result; +} + +void llama_memory_dsv41::set_graph_workspace_size(size_t size) { + pimpl->graph_workspace = size; +} + +void llama_memory_dsv41::state_write( + llama_io_write_i & io, + llama_seq_id seq_id, + llama_state_seq_flags flags) const { + GGML_UNUSED(flags); + pimpl->require_idle(); + if (seq_id != -1 && !pimpl->valid_seq(seq_id)) { + throw std::invalid_argument("DeepSeek V4.1 state sequence ID is out of range"); + } + + io.write(&DSV41_STATE_MAGIC, sizeof(DSV41_STATE_MAGIC)); + io.write(&DSV41_STATE_VERSION, sizeof(DSV41_STATE_VERSION)); + const uint32_t count = seq_id == -1 ? pimpl->config.n_seq : 1; + io.write(&count, sizeof(count)); + const auto tensors = pimpl->sequence_tensors(); + const uint32_t tensor_count = tensors.size(); + io.write(&tensor_count, sizeof(tensor_count)); + + const llama_dsv41_engram_snapshot engram_snapshot = + pimpl->config.engram ? pimpl->config.engram->checkpoint() : llama_dsv41_engram_snapshot {}; + for (uint32_t i = 0; i < count; ++i) { + const llama_seq_id current = seq_id == -1 ? (llama_seq_id) i : seq_id; + io.write(¤t, sizeof(current)); + const auto & state = pimpl->sequences[current]; + io.write(&state.pos, sizeof(state.pos)); + const uint32_t n_candidates = state.candidates.size(); + io.write(&n_candidates, sizeof(n_candidates)); + if (n_candidates > 0) { + io.write(state.candidates.data(), n_candidates*sizeof(int32_t)); + } + const uint8_t has_engram = pimpl->config.engram ? 1 : 0; + io.write(&has_engram, sizeof(has_engram)); + if (has_engram) { + const auto found = engram_snapshot.sequences.find(current); + llama_dsv41_engram_sequence_state engram_state; + engram_state.history.reset(); + if (found != engram_snapshot.sequences.end()) { + engram_state = found->second; + } + io.write(&engram_state.pos, sizeof(engram_state.pos)); + io.write(engram_state.history.tail.data(), sizeof(engram_state.history.tail)); + } + for (ggml_tensor * tensor : tensors) { + const uint64_t plane = pimpl->sequence_plane(tensor); + io.write(&plane, sizeof(plane)); + io.write_tensor(tensor, (size_t) current*plane, plane); + } + } +} + +void llama_memory_dsv41::state_read( + llama_io_read_i & io, + llama_seq_id seq_id, + llama_state_seq_flags flags) { + GGML_UNUSED(flags); + pimpl->require_idle(); + if (seq_id != -1 && !pimpl->valid_seq(seq_id)) { + throw std::invalid_argument("DeepSeek V4.1 state destination sequence ID is out of range"); + } + + uint64_t magic = 0; + uint32_t version = 0; + uint32_t count = 0; + uint32_t tensor_count = 0; + io.read(&magic, sizeof(magic)); + io.read(&version, sizeof(version)); + io.read(&count, sizeof(count)); + io.read(&tensor_count, sizeof(tensor_count)); + const auto tensors = pimpl->sequence_tensors(); + if (magic != DSV41_STATE_MAGIC || version != DSV41_STATE_VERSION || + tensor_count != tensors.size() || count == 0 || + (seq_id != -1 && count != 1) || (seq_id == -1 && count != pimpl->config.n_seq)) { + throw std::runtime_error("DeepSeek V4.1 state header is incompatible"); + } + + const auto sequence_is_empty = [&](llama_seq_id current) { + if (pimpl->sequences[current].pos >= 0 || + !pimpl->sequences[current].candidates.empty() || + pimpl->rollback_states.count(current) != 0 || + pimpl->engram_boundaries.count(current) != 0) { + return false; + } + return !pimpl->config.engram || pimpl->config.engram->sequence(current).pos < 0; + }; + if (seq_id == -1) { + for (uint32_t current = 0; current < pimpl->config.n_seq; ++current) { + if (!sequence_is_empty(current)) { + throw std::runtime_error("DeepSeek V4.1 full state restore requires empty memory"); + } + } + } else if (!sequence_is_empty(seq_id)) { + throw std::runtime_error("DeepSeek V4.1 sequence state restore requires an empty destination"); + } + + llama_dsv41_engram_snapshot engram_snapshot = + pimpl->config.engram && seq_id != -1 ? + pimpl->config.engram->checkpoint() : llama_dsv41_engram_snapshot {}; + std::set restored; + std::vector tensor_data; + + try { + for (uint32_t i = 0; i < count; ++i) { + llama_seq_id stored = -1; + llama_pos pos = -1; + uint32_t n_candidates = 0; + uint8_t has_engram = 0; + io.read(&stored, sizeof(stored)); + io.read(&pos, sizeof(pos)); + io.read(&n_candidates, sizeof(n_candidates)); + const llama_seq_id target = seq_id == -1 ? stored : seq_id; + if (!pimpl->valid_seq(stored) || !pimpl->valid_seq(target) || + (seq_id == -1 && !restored.insert(target).second) || + pos < -1 || pos >= (llama_pos) pimpl->config.n_ctx || + n_candidates > pimpl->candidate_width) { + throw std::runtime_error("DeepSeek V4.1 state metadata is invalid"); + } + std::vector candidates(n_candidates); + if (n_candidates > 0) { + io.read(candidates.data(), n_candidates*sizeof(int32_t)); + } + io.read(&has_engram, sizeof(has_engram)); + if ((has_engram != 0) != (pimpl->config.engram != nullptr)) { + throw std::runtime_error("DeepSeek V4.1 Engram state availability differs"); + } + if (has_engram) { + llama_dsv41_engram_sequence_state engram_state; + io.read(&engram_state.pos, sizeof(engram_state.pos)); + io.read(engram_state.history.tail.data(), sizeof(engram_state.history.tail)); + engram_snapshot.sequences[target] = engram_state; + } + for (ggml_tensor * tensor : tensors) { + uint64_t plane = 0; + io.read(&plane, sizeof(plane)); + if (plane != pimpl->sequence_plane(tensor)) { + throw std::runtime_error("DeepSeek V4.1 state tensor layout differs"); + } + tensor_data.resize(plane); + io.read(tensor_data.data(), plane); + ggml_backend_tensor_set(tensor, tensor_data.data(), (size_t) target*plane, plane); + } + pimpl->sequences[target].pos = pos; + pimpl->sequences[target].candidates = std::move(candidates); + pimpl->update_position_state(target); + pimpl->update_candidate_state(target); + } + if (pimpl->config.engram) { + pimpl->config.engram->restore(engram_snapshot); + } + if (seq_id == -1) { + pimpl->engram_boundaries.clear(); + pimpl->rollback_states.clear(); + } else { + pimpl->engram_boundaries.erase(seq_id); + pimpl->rollback_states.erase(seq_id); + } + ++pimpl->generation; + } catch (...) { + if (seq_id == -1) { + clear(true); + } else { + seq_rm(seq_id, -1, -1); + } + throw; + } +} + +ggml_tensor * llama_memory_dsv41::raw_k(uint32_t layer) const { + return layer < pimpl->raw.size() ? pimpl->raw[layer] : nullptr; +} + +ggml_tensor * llama_memory_dsv41::compressed_kv(uint32_t source_layer) const { + const auto found = pimpl->sources.find(source_layer); + return found == pimpl->sources.end() ? nullptr : found->second.kv; +} + +ggml_tensor * llama_memory_dsv41::index_keys(uint32_t source_layer) const { + const auto found = pimpl->sources.find(source_layer); + return found == pimpl->sources.end() ? nullptr : found->second.index; +} + +ggml_tensor * llama_memory_dsv41::compressor_carry_kv(uint32_t source_layer) const { + const auto found = pimpl->sources.find(source_layer); + return found == pimpl->sources.end() ? nullptr : found->second.carry_kv; +} + +ggml_tensor * llama_memory_dsv41::compressor_carry_score(uint32_t source_layer) const { + const auto found = pimpl->sources.find(source_layer); + return found == pimpl->sources.end() ? nullptr : found->second.carry_score; +} + +ggml_tensor * llama_memory_dsv41::candidate_scores() const { + return pimpl->candidate_scores; +} + +ggml_tensor * llama_memory_dsv41::candidate_ids() const { + return pimpl->candidate_ids; +} + +ggml_tensor * llama_memory_dsv41::committed_candidate_ids() const { + return pimpl->committed_candidate_ids; +} + +ggml_tensor * llama_memory_dsv41::position_state() const { + return pimpl->positions; +} + +const llama_dsv41_memory_config & llama_memory_dsv41::config() const { + return pimpl->config; +} + +llama_dsv41_memory_accounting llama_memory_dsv41::accounting() const { + llama_dsv41_memory_accounting result; + for (ggml_tensor * tensor : pimpl->raw) { + result.raw_kv += ggml_nbytes(tensor); + } + for (const auto & entry : pimpl->sources) { + result.compressed_kv += ggml_nbytes(entry.second.kv); + result.index_keys += ggml_nbytes(entry.second.index); + if (entry.second.carry_kv != nullptr) { + result.compressor_carry += ggml_nbytes(entry.second.carry_kv); + result.compressor_carry += ggml_nbytes(entry.second.carry_score); + } + } + result.candidate_scores = ggml_nbytes(pimpl->candidate_scores); + result.candidate_ids = ggml_nbytes(pimpl->candidate_ids) + ggml_nbytes(pimpl->committed_candidate_ids); + result.position_state = ggml_nbytes(pimpl->positions); + result.graph_workspace = pimpl->graph_workspace; + return result; +} + +std::vector llama_memory_dsv41::sequence_candidate_ids(llama_seq_id seq_id) const { + if (!pimpl->valid_seq(seq_id)) { + throw std::invalid_argument("DeepSeek V4.1 candidate sequence ID is out of range"); + } + return pimpl->sequences[seq_id].candidates; +} + +bool llama_memory_dsv41::engram_enabled() const { + return pimpl->config.engram != nullptr; +} + +llama_memory_dsv41_context::llama_memory_dsv41_context(llama_memory_status status) : + status(status) { +} + +llama_memory_dsv41_context::llama_memory_dsv41_context( + llama_memory_dsv41 * memory, + bool full) : + status(LLAMA_MEMORY_STATUS_SUCCESS), + mem(memory), + full(full) { + if (!memory || !full) { + return; + } + + llama_batch_allocr allocator(1); + llama_ubatch ubatch = allocator.ubatch_reserve(memory->pimpl->config.n_ubatch, 1); + const uint32_t n_tokens = ubatch.n_tokens; + const uint32_t n_streams = 1; + ubatch.data->seq_id_data.resize((size_t) n_tokens*n_streams); + ubatch.data->seq_id_unq.resize(n_streams); + ubatch.seq_id_unq = ubatch.data->seq_id_unq.data(); + ubatch.n_seqs_unq = n_streams; + for (uint32_t seq = 0; seq < n_streams; ++seq) { + ubatch.data->seq_id_unq[seq] = seq; + ubatch.seq_idx[seq] = seq; + } + for (uint32_t token = 0; token < n_tokens; ++token) { + ubatch.token[token] = 0; + ubatch.pos[token] = token; + ubatch.n_seq_id[token] = n_streams; + ubatch.seq_id[token] = ubatch.data->seq_id_data.data() + (size_t) token*n_streams; + for (uint32_t seq = 0; seq < n_streams; ++seq) { + ubatch.seq_id[token][seq] = seq; + } + } + ubatches.push_back(std::move(ubatch)); + + transaction = std::make_unique(); + transaction->plan.generation = memory->pimpl->generation; + transaction->next_sequences = memory->pimpl->sequences; + for (uint32_t seq = 0; seq < n_streams; ++seq) { + transaction->seq_ids.push_back(seq); + transaction->start_positions.push_back(0); + transaction->next_sequences[seq].pos = n_tokens - 1; + } + const uint32_t raw_width = memory->pimpl->config.raw_window; + const uint32_t persist_first = n_tokens > raw_width ? n_tokens - raw_width : 0; + for (uint32_t token = 0; token < n_tokens; ++token) { + std::vector token_sequences(n_streams); + for (uint32_t seq = 0; seq < n_streams; ++seq) { + token_sequences[seq] = seq; + if (token >= persist_first) { + transaction->plan.raw.persist_src_idxs.push_back(token); + transaction->plan.raw.write_idxs.push_back( + (int64_t) seq*raw_width + token%raw_width); + } + } + transaction->plan.token_seq_ids.push_back(std::move(token_sequences)); + transaction->plan.positions.push_back(token); + const uint32_t visible = std::min(raw_width, token + 1); + transaction->plan.raw.n_visible.push_back(visible); + const auto order = llama_dsv41_raw_ring_order(token, raw_width); + for (uint32_t row = 0; row < raw_width; ++row) { + transaction->plan.raw.read_idxs.push_back(row < visible ? order[row] : -1); + transaction->plan.raw.mask.push_back( + row < visible ? 0.0f : -std::numeric_limits::infinity()); + } + } + for (uint32_t source : memory->pimpl->config.kv_sources) { + llama_dsv41_source_plan source_plan; + source_plan.source_layer = source; + source_plan.ratio = memory->pimpl->config.ratios[source]; + source_plan.capacity = source_capacity(memory->pimpl->config, source); + source_plan.compression = llama_dsv41_build_compression_plan( + transaction->plan.positions, source_plan.ratio, source_plan.capacity); + const uint32_t max_visible = source_plan.compression.n_visible.empty() ? 0 : + *std::max_element( + source_plan.compression.n_visible.begin(), + source_plan.compression.n_visible.end()); + for (int32_t visible : source_plan.compression.n_visible) { + for (uint32_t row = 0; row < max_visible; ++row) { + source_plan.read_idxs.push_back(row < (uint32_t) visible ? row : -1); + } + } + transaction->plan.sources.push_back(std::move(source_plan)); + } + const uint32_t visible_blocks = + (n_tokens + memory->pimpl->config.candidate_block_size - 1)/ + memory->pimpl->config.candidate_block_size; + transaction->plan.candidate_width = std::min( + visible_blocks, memory->pimpl->config.candidate_topk_blocks); +} + +llama_memory_dsv41_context::llama_memory_dsv41_context( + llama_memory_dsv41 * memory, + std::vector ubatches) : + status(LLAMA_MEMORY_STATUS_SUCCESS), + mem(memory), + ubatches(std::move(ubatches)) { +} + +llama_memory_dsv41_context::~llama_memory_dsv41_context() { + rollback(); +} + +bool llama_memory_dsv41_context::next() { + if (status != LLAMA_MEMORY_STATUS_SUCCESS || transaction) { + return false; + } + if (++i_next >= ubatches.size()) { + return false; + } + return true; +} + +bool llama_memory_dsv41_context::apply() { + if (status == LLAMA_MEMORY_STATUS_NO_UPDATE) { + return true; + } + if (status != LLAMA_MEMORY_STATUS_SUCCESS || full || mem == nullptr || transaction || + i_next >= ubatches.size() || mem->pimpl->transaction_active) { + status = LLAMA_MEMORY_STATUS_FAILED_PREPARE; + return false; + } + + const llama_ubatch & ubatch = ubatches[i_next]; + auto next = std::make_unique(); + try { + if (ubatch.n_tokens == 0 || ubatch.n_tokens > mem->pimpl->config.n_ubatch || + ubatch.n_pos != 1 || ubatch.pos == nullptr || ubatch.n_seq_id == nullptr || + ubatch.seq_id == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 ubatch shape is invalid"); + } + + next->next_sequences = mem->pimpl->sequences; + std::vector topology; + for (uint32_t token = 0; token < ubatch.n_tokens; ++token) { + if (ubatch.n_seq_id[token] <= 0 || ubatch.seq_id[token] == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 token has no sequence"); + } + std::vector token_sequences( + ubatch.seq_id[token], ubatch.seq_id[token] + ubatch.n_seq_id[token]); + std::sort(token_sequences.begin(), token_sequences.end()); + if (std::adjacent_find(token_sequences.begin(), token_sequences.end()) != token_sequences.end()) { + throw std::invalid_argument("DeepSeek V4.1 token repeats a sequence ID"); + } + if (token == 0) { + topology = token_sequences; + } else if (topology != token_sequences) { + throw std::invalid_argument("DeepSeek V4.1 ubatch changes sequence topology"); + } + next->plan.token_seq_ids.push_back(token_sequences); + next->plan.positions.push_back(ubatch.pos[token]); + } + if (topology.empty()) { + throw std::invalid_argument("DeepSeek V4.1 ubatch has no sequence topology"); + } + if (topology.size() != 1 || ubatch.n_seqs != 1) { + throw std::invalid_argument("DeepSeek V4.1 memory supports one sequence per ubatch"); + } + next->seq_ids = topology; + for (llama_seq_id seq_id : topology) { + if (!mem->pimpl->valid_seq(seq_id)) { + throw std::invalid_argument("DeepSeek V4.1 sequence ID is out of range"); + } + next->start_positions.push_back(mem->pimpl->sequences[seq_id].pos + 1); + } + if (!std::all_of(next->start_positions.begin(), next->start_positions.end(), + [&](llama_pos pos) { return pos == next->start_positions.front(); })) { + throw std::invalid_argument("DeepSeek V4.1 coupled sequences have different positions"); + } + for (uint32_t token = 0; token < ubatch.n_tokens; ++token) { + const llama_pos expected = next->start_positions.front() + token; + if (ubatch.pos[token] != expected || expected < 0 || expected >= (llama_pos) mem->pimpl->config.n_ctx) { + throw std::invalid_argument("DeepSeek V4.1 positions must be contiguous and within n_ctx"); + } + } + + next->plan.generation = mem->pimpl->generation + 1; + const uint32_t raw_width = mem->pimpl->config.raw_window; + const uint32_t persist_first = + ubatch.n_tokens > raw_width ? ubatch.n_tokens - raw_width : 0; + next->plan.raw.n_visible.resize(ubatch.n_tokens); + for (uint32_t token = 0; token < ubatch.n_tokens; ++token) { + const llama_pos pos = ubatch.pos[token]; + for (llama_seq_id seq_id : topology) { + if (token >= persist_first) { + next->plan.raw.persist_src_idxs.push_back(token); + next->plan.raw.write_idxs.push_back( + (int64_t) seq_id*raw_width + pos%raw_width); + } + next->next_sequences[seq_id].pos = pos; + } + const uint32_t visible = std::min(raw_width, pos + 1); + next->plan.raw.n_visible[token] = visible; + const auto order = llama_dsv41_raw_ring_order(pos, raw_width); + for (uint32_t row = 0; row < raw_width; ++row) { + next->plan.raw.read_idxs.push_back( + row < visible ? (int32_t) ((uint32_t) topology.front()*raw_width + order[row]) : -1); + next->plan.raw.mask.push_back(row < visible ? 0.0f : -std::numeric_limits::infinity()); + } + } + + for (uint32_t source : mem->pimpl->config.kv_sources) { + llama_dsv41_source_plan source_plan; + source_plan.source_layer = source; + source_plan.ratio = mem->pimpl->config.ratios[source]; + source_plan.capacity = source_capacity(mem->pimpl->config, source); + source_plan.compression = llama_dsv41_build_compression_plan( + next->plan.positions, source_plan.ratio, source_plan.capacity); + const uint32_t max_visible = source_plan.compression.n_visible.empty() ? 0 : + *std::max_element(source_plan.compression.n_visible.begin(), source_plan.compression.n_visible.end()); + for (int32_t visible : source_plan.compression.n_visible) { + for (uint32_t row = 0; row < max_visible; ++row) { + source_plan.read_idxs.push_back( + row < (uint32_t) visible ? + (int32_t) ((uint32_t) topology.front()*source_plan.capacity + row) : -1); + } + } + next->plan.sources.push_back(std::move(source_plan)); + } + const llama_pos final_pos = next->plan.positions.back(); + const uint32_t candidate_visible = final_pos + 1; + const uint32_t visible_blocks = + (candidate_visible + mem->pimpl->config.candidate_block_size - 1)/ + mem->pimpl->config.candidate_block_size; + next->plan.candidate_width = std::min( + visible_blocks, mem->pimpl->config.candidate_topk_blocks); + + if (mem->pimpl->config.engram) { + if (ubatch.token == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 Engram requires token inputs"); + } + std::vector tokens; + tokens.reserve(ubatch.n_tokens); + for (uint32_t token = 0; token < ubatch.n_tokens; ++token) { + tokens.push_back({ + ubatch.token[token], + ubatch.pos[token], + next->plan.token_seq_ids[token], + 1, + }); + } + for (llama_seq_id seq_id : topology) { + next->engram_before[seq_id] = mem->pimpl->config.engram->sequence(seq_id); + } + next->engram = std::make_unique( + mem->pimpl->config.engram->prepare(tokens)); + } + + std::set> snapshot_keys; + auto snapshot = [&](ggml_tensor * tensor, size_t offset, size_t size) { + if (tensor == nullptr || mem->pimpl->config.no_alloc || + !snapshot_keys.emplace(tensor, offset, size).second) { + return; + } + llama_memory_dsv41::impl::tensor_snapshot entry; + entry.tensor = tensor; + entry.offset = offset; + entry.data.resize(size); + ggml_backend_tensor_get(tensor, entry.data.data(), offset, size); + next->snapshots.push_back(std::move(entry)); + }; + for (llama_seq_id seq_id : topology) { + for (ggml_tensor * raw : mem->pimpl->raw) { + const size_t row_bytes = raw->nb[1]; + for (llama_pos pos : next->plan.positions) { + const size_t row = (size_t) seq_id*mem->pimpl->config.raw_window + + (uint32_t) pos%mem->pimpl->config.raw_window; + snapshot(raw, row*row_bytes, row_bytes); + } + } + for (const auto & source_plan : next->plan.sources) { + const auto & storage = mem->pimpl->sources.at(source_plan.source_layer); + for (int64_t row : source_plan.compression.write_idxs) { + const size_t physical = (size_t) seq_id*storage.capacity + row; + snapshot(storage.kv, physical*storage.kv->nb[1], storage.kv->nb[1]); + snapshot(storage.index, physical*storage.index->nb[1], storage.index->nb[1]); + } + if (storage.carry_kv != nullptr) { + snapshot( + storage.carry_kv, + (size_t) seq_id*storage.carry_kv->nb[2], + storage.carry_kv->nb[2]); + snapshot( + storage.carry_score, + (size_t) seq_id*storage.carry_score->nb[2], + storage.carry_score->nb[2]); + } + } + } + + mem->pimpl->transaction_active = true; + transaction = std::move(next); + return true; + } catch (const std::exception & error) { + if (next->engram) { + mem->pimpl->config.engram->rollback(*next->engram); + } + LLAMA_LOG_ERROR("%s: %s\n", __func__, error.what()); + status = LLAMA_MEMORY_STATUS_FAILED_PREPARE; + return false; + } +} + +void llama_memory_dsv41_context::commit() { + if (!transaction || mem == nullptr) { + return; + } + GGML_ASSERT(transaction->seq_ids.size() == 1); + const llama_seq_id seq_id = transaction->seq_ids.front(); + llama_memory_dsv41::impl::rollback_state rollback; + rollback.start_pos = transaction->start_positions.front(); + rollback.sequence = mem->pimpl->sequences[seq_id]; + if (transaction->engram) { + rollback.engram = transaction->engram_before.at(seq_id); + rollback.has_engram = true; + } + + std::vector committed_candidates; + if (!mem->pimpl->config.no_alloc && transaction->plan.candidate_width > 0) { + const size_t row_bytes = mem->pimpl->candidate_ids->nb[1]; + std::vector ids(transaction->plan.candidate_width); + uint32_t token = 0; + for (uint32_t current = 0; current < transaction->plan.token_seq_ids.size(); ++current) { + if (std::find( + transaction->plan.token_seq_ids[current].begin(), + transaction->plan.token_seq_ids[current].end(), + seq_id) != transaction->plan.token_seq_ids[current].end()) { + token = current; + } + } + ggml_backend_tensor_get( + mem->pimpl->candidate_ids, + ids.data(), + (size_t) token*row_bytes, + ids.size()*sizeof(int32_t)); + transaction->next_sequences[seq_id].candidates = std::move(ids); + } + + const std::vector position_values = + mem->pimpl->position_state_values(transaction->next_sequences[seq_id]); + committed_candidates = + mem->pimpl->candidate_state_values(transaction->next_sequences[seq_id]); + std::map next_boundaries; + if (transaction->engram) { + const auto found = mem->pimpl->engram_boundaries.find(seq_id); + if (found != mem->pimpl->engram_boundaries.end()) { + next_boundaries = found->second; + } + next_boundaries[transaction->start_positions.front()] = + transaction->engram_before.at(seq_id); + } + + auto rollback_slot = mem->pimpl->rollback_states.end(); + bool inserted_rollback = false; + auto boundary_slot = mem->pimpl->engram_boundaries.end(); + bool inserted_boundary = false; + try { + std::tie(rollback_slot, inserted_rollback) = + mem->pimpl->rollback_states.emplace(seq_id, llama_memory_dsv41::impl::rollback_state {}); + if (transaction->engram) { + std::tie(boundary_slot, inserted_boundary) = + mem->pimpl->engram_boundaries.emplace( + seq_id, std::map {}); + } + if (transaction->engram) { + mem->pimpl->config.engram->commit(*transaction->engram); + } + } catch (...) { + if (inserted_boundary) { + mem->pimpl->engram_boundaries.erase(boundary_slot); + } + if (inserted_rollback) { + mem->pimpl->rollback_states.erase(rollback_slot); + } + throw; + } + + if (!mem->pimpl->config.no_alloc) { + ggml_backend_tensor_set( + mem->pimpl->committed_candidate_ids, + committed_candidates.data(), + (size_t) seq_id*mem->pimpl->committed_candidate_ids->nb[1], + committed_candidates.size()*sizeof(int32_t)); + ggml_backend_tensor_set( + mem->pimpl->positions, + position_values.data(), + (size_t) seq_id*mem->pimpl->positions->nb[1], + position_values.size()*sizeof(int32_t)); + } + mem->pimpl->sequences.swap(transaction->next_sequences); + if (transaction->engram) { + boundary_slot->second.swap(next_boundaries); + } + rollback.snapshots.swap(transaction->snapshots); + rollback_slot->second = std::move(rollback); + mem->pimpl->generation = transaction->plan.generation; + mem->pimpl->transaction_active = false; + transaction.reset(); +} + +void llama_memory_dsv41_context::rollback() { + if (!transaction || mem == nullptr) { + return; + } + if (full) { + transaction.reset(); + return; + } + if (transaction->engram) { + mem->pimpl->config.engram->rollback(*transaction->engram); + } + if (!mem->pimpl->config.no_alloc) { + for (const auto & snapshot : transaction->snapshots) { + ggml_backend_tensor_set( + snapshot.tensor, snapshot.data.data(), snapshot.offset, snapshot.data.size()); + } + } + mem->pimpl->transaction_active = false; + transaction.reset(); +} + +llama_memory_status llama_memory_dsv41_context::get_status() const { + return status; +} + +const llama_ubatch & llama_memory_dsv41_context::get_ubatch() const { + if (status != LLAMA_MEMORY_STATUS_SUCCESS || i_next >= ubatches.size()) { + throw std::runtime_error("DeepSeek V4.1 memory context has no current ubatch"); + } + return ubatches[i_next]; +} + +const llama_dsv41_memory_plan & llama_memory_dsv41_context::plan() const { + if (!transaction) { + throw std::runtime_error("DeepSeek V4.1 memory transaction is not prepared"); + } + return transaction->plan; +} + +const llama_dsv41_memory_plan & llama_memory_dsv41_context::graph_plan( + const llama_ubatch & ubatch) const { + if (!full) { + return plan(); + } + if (mem == nullptr || ubatch.n_tokens == 0 || + ubatch.n_tokens > mem->pimpl->config.n_ubatch) { + throw std::invalid_argument("DeepSeek V4.1 full graph ubatch is invalid"); + } + if (transaction && transaction->plan.positions.size() == ubatch.n_tokens) { + return transaction->plan; + } + + auto next = std::make_unique(); + next->plan.generation = mem->pimpl->generation; + next->next_sequences = mem->pimpl->sequences; + next->seq_ids = { 0 }; + next->start_positions = { 0 }; + next->next_sequences[0].pos = ubatch.n_tokens - 1; + + const uint32_t raw_width = mem->pimpl->config.raw_window; + const uint32_t persist_first = + ubatch.n_tokens > raw_width ? ubatch.n_tokens - raw_width : 0; + for (uint32_t token = 0; token < ubatch.n_tokens; ++token) { + next->plan.token_seq_ids.push_back({ 0 }); + next->plan.positions.push_back(token); + if (token >= persist_first) { + next->plan.raw.persist_src_idxs.push_back(token); + next->plan.raw.write_idxs.push_back(token%raw_width); + } + const uint32_t visible = std::min(raw_width, token + 1); + next->plan.raw.n_visible.push_back(visible); + const auto order = llama_dsv41_raw_ring_order(token, raw_width); + for (uint32_t row = 0; row < raw_width; ++row) { + next->plan.raw.read_idxs.push_back(row < visible ? order[row] : -1); + next->plan.raw.mask.push_back( + row < visible ? 0.0f : -std::numeric_limits::infinity()); + } + } + for (uint32_t source : mem->pimpl->config.kv_sources) { + llama_dsv41_source_plan source_plan; + source_plan.source_layer = source; + source_plan.ratio = mem->pimpl->config.ratios[source]; + source_plan.capacity = source_capacity(mem->pimpl->config, source); + source_plan.compression = llama_dsv41_build_compression_plan( + next->plan.positions, source_plan.ratio, source_plan.capacity); + const uint32_t max_visible = source_plan.compression.n_visible.empty() ? 0 : + *std::max_element( + source_plan.compression.n_visible.begin(), + source_plan.compression.n_visible.end()); + for (int32_t visible : source_plan.compression.n_visible) { + for (uint32_t row = 0; row < max_visible; ++row) { + source_plan.read_idxs.push_back(row < (uint32_t) visible ? row : -1); + } + } + next->plan.sources.push_back(std::move(source_plan)); + } + const uint32_t visible_blocks = + (ubatch.n_tokens + mem->pimpl->config.candidate_block_size - 1)/ + mem->pimpl->config.candidate_block_size; + next->plan.candidate_width = std::min( + visible_blocks, mem->pimpl->config.candidate_topk_blocks); + transaction = std::move(next); + return transaction->plan; +} + +llama_dsv41_graph_topology llama_memory_dsv41_context::topology( + const llama_ubatch & ubatch, + uint32_t n_outputs) const { + if (mem == nullptr) { + throw std::runtime_error("DeepSeek V4.1 memory context has no memory"); + } + if (full) { + graph_plan(ubatch); + } + llama_dsv41_graph_topology result; + result.n_tokens = ubatch.n_tokens; + result.n_seqs = transaction ? transaction->seq_ids.size() : ubatch.n_seqs; + result.n_outputs = n_outputs; + result.backend_layout = mem->pimpl->backend_layout; + result.engram_enabled = mem->pimpl->config.engram != nullptr; + result.expert_enabled = mem->pimpl->config.expert_enabled; + result.transaction_generation = transaction ? transaction->plan.generation : mem->pimpl->generation; + if (transaction) { + result.seq_ids = transaction->seq_ids; + result.start_positions = transaction->start_positions; + for (llama_seq_id seq_id : transaction->seq_ids) { + const llama_pos final_pos = transaction->next_sequences[seq_id].pos; + result.visible_raw_widths.push_back(std::min( + mem->pimpl->config.raw_window, final_pos + 1)); + } + for (const auto & source : transaction->plan.sources) { + result.source_ratios.push_back(source.ratio); + result.visible_compressed_widths.push_back( + source.compression.n_visible.empty() ? 0 : + *std::max_element(source.compression.n_visible.begin(), source.compression.n_visible.end())); + result.source_carry_counts.push_back( + source.ratio == 2 ? (transaction->plan.positions.back() + 1)%source.ratio : 0); + } + result.candidate_width = transaction->plan.candidate_width; + } else { + const uint32_t n_seqs = std::min(ubatch.n_seqs, mem->pimpl->config.n_seq); + for (uint32_t seq = 0; seq < n_seqs; ++seq) { + result.seq_ids.push_back(seq); + result.start_positions.push_back(0); + result.visible_raw_widths.push_back(mem->pimpl->config.raw_window); + } + for (uint32_t source : mem->pimpl->config.kv_sources) { + result.source_ratios.push_back(mem->pimpl->config.ratios[source]); + result.visible_compressed_widths.push_back(source_capacity(mem->pimpl->config, source)); + result.source_carry_counts.push_back(mem->pimpl->config.ratios[source] == 2 ? 1 : 0); + } + result.candidate_width = mem->pimpl->candidate_width; + } + return result; +} + +const llama_dsv41_engram_transaction * llama_memory_dsv41_context::engram_transaction() const { + return transaction && transaction->engram ? transaction->engram.get() : nullptr; +} + +std::vector llama_memory_dsv41_context::engram_row_ids(uint32_t engram_layer) const { + const auto * prepared = engram_transaction(); + if (prepared == nullptr) { + return {}; + } + if (engram_layer >= LLAMA_ENGRAM_LAYERS) { + throw std::invalid_argument("DeepSeek V4.1 Engram layer index is out of range"); + } + + const uint32_t * source = prepared->row_ids(engram_layer); + std::vector result(prepared->token_count()*LLAMA_ENGRAM_COLS); + for (size_t token = 0; token < prepared->token_count(); ++token) { + for (uint32_t column = 0; column < LLAMA_ENGRAM_COLS; ++column) { + const uint32_t row = + source[token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS + column]; + if (row > (uint32_t) INT32_MAX) { + throw std::runtime_error("DeepSeek V4.1 Engram row ID exceeds I32"); + } + result[token*LLAMA_ENGRAM_COLS + column] = row; + } + } + return result; +} + +void llama_memory_dsv41_context::stage_candidate_ids( + uint32_t token, + const std::vector & ids) { + if (!transaction || mem == nullptr || token >= transaction->plan.positions.size() || + ids.size() != transaction->plan.candidate_width) { + throw std::invalid_argument("DeepSeek V4.1 candidate staging shape is invalid"); + } + if (!mem->pimpl->config.no_alloc && !ids.empty()) { + ggml_backend_tensor_set( + mem->pimpl->candidate_ids, + ids.data(), + (size_t) token*mem->pimpl->candidate_ids->nb[1], + ids.size()*sizeof(int32_t)); + } +} diff --git a/src/llama-memory-dsv41.h b/src/llama-memory-dsv41.h new file mode 100644 index 000000000000..0eb074c12fe5 --- /dev/null +++ b/src/llama-memory-dsv41.h @@ -0,0 +1,173 @@ +#pragma once + +#include "llama-dsv41.h" +#include "llama-dsv41-engram.h" +#include "llama-memory.h" + +#include +#include +#include +#include +#include +#include + +struct ggml_tensor; +struct llama_model; + +struct llama_dsv41_memory_config { + uint32_t n_ctx = 0; + uint32_t n_seq = 0; + uint32_t n_ubatch = 0; + uint32_t n_layer = LLAMA_DSV41_N_LAYER; + uint32_t raw_window = LLAMA_DSV41_N_SWA; + uint32_t kv_width = LLAMA_DSV41_N_HEAD_DIM; + uint32_t index_width = LLAMA_DSV41_N_INDEX_HEAD_DIM; + uint32_t candidate_topk_blocks = LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS; + uint32_t candidate_block_size = LLAMA_DSV41_CANDIDATE_BLOCK_SIZE; + uint32_t candidate_source_layer = LLAMA_DSV41_CANDIDATE_SOURCE_LAYER; + ggml_type type_k = GGML_TYPE_F16; + ggml_type type_index = GGML_TYPE_F16; + bool no_alloc = false; + bool expert_enabled = false; + std::vector kv_sources = { 2, 8, 14, 20 }; + std::vector index_sources = { 2, 8, 14, 20, 24, 28, 32, 36 }; + std::vector ratios; + std::function buft_for_layer; + std::unique_ptr engram; +}; + +struct llama_dsv41_raw_plan { + std::vector persist_src_idxs; + std::vector write_idxs; + std::vector read_idxs; + std::vector mask; + std::vector n_visible; +}; + +struct llama_dsv41_source_plan { + uint32_t source_layer = 0; + uint32_t ratio = 0; + uint32_t capacity = 0; + llama_dsv41_compression_plan compression; + std::vector read_idxs; +}; + +struct llama_dsv41_graph_topology { + uint32_t n_tokens = 0; + uint32_t n_seqs = 0; + uint32_t n_outputs = 0; + std::vector seq_ids; + std::vector start_positions; + std::vector visible_raw_widths; + std::vector visible_compressed_widths; + std::vector source_ratios; + std::vector source_carry_counts; + uint32_t candidate_width = 0; + uint64_t backend_layout = 0; + uint64_t transaction_generation = 0; + bool engram_enabled = false; + bool expert_enabled = false; + + bool same_topology(const llama_dsv41_graph_topology & other) const; +}; + +struct llama_dsv41_memory_plan { + uint64_t generation = 0; + std::vector positions; + std::vector> token_seq_ids; + llama_dsv41_raw_plan raw; + std::vector sources; + uint32_t candidate_width = 0; +}; + +class llama_memory_dsv41; + +class llama_memory_dsv41_context : public llama_memory_context_i { +public: + llama_memory_dsv41_context(llama_memory_status status); + llama_memory_dsv41_context(llama_memory_dsv41 * memory, bool full); + llama_memory_dsv41_context(llama_memory_dsv41 * memory, std::vector ubatches); + ~llama_memory_dsv41_context() override; + + bool next() override; + bool apply() override; + void commit() override; + void rollback() override; + + llama_memory_status get_status() const override; + const llama_ubatch & get_ubatch() const override; + + const llama_dsv41_memory_plan & plan() const; + const llama_dsv41_memory_plan & graph_plan(const llama_ubatch & ubatch) const; + llama_dsv41_graph_topology topology(const llama_ubatch & ubatch, uint32_t n_outputs) const; + const llama_dsv41_engram_transaction * engram_transaction() const; + std::vector engram_row_ids(uint32_t engram_layer) const; + const llama_memory_dsv41 * memory() const { return mem; } + + void stage_candidate_ids(uint32_t token, const std::vector & ids); + +private: + struct transaction_state; + + llama_memory_status status; + llama_memory_dsv41 * mem = nullptr; + bool full = false; + size_t i_next = 0; + std::vector ubatches; + mutable std::unique_ptr transaction; +}; + +class llama_memory_dsv41 : public llama_memory_i { +public: + explicit llama_memory_dsv41(llama_dsv41_memory_config config); + llama_memory_dsv41( + const llama_model & model, + ggml_type type_k, + bool offload, + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_ubatch, + std::unique_ptr engram); + ~llama_memory_dsv41() override; + + llama_memory_context_ptr init_batch( + llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool embd_all) override; + llama_memory_context_ptr init_full() override; + llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override; + + bool get_can_shift() const override; + void clear(bool data) override; + bool seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; + void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override; + void seq_keep(llama_seq_id seq_id) override; + [[noreturn]] void seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override; + [[noreturn]] void seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) override; + llama_pos seq_pos_min(llama_seq_id seq_id) const override; + llama_pos seq_pos_max(llama_seq_id seq_id) const override; + std::map memory_breakdown() const override; + void set_graph_workspace_size(size_t size) override; + void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; + void state_read(llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override; + + ggml_tensor * raw_k(uint32_t layer) const; + ggml_tensor * compressed_kv(uint32_t source_layer) const; + ggml_tensor * index_keys(uint32_t source_layer) const; + ggml_tensor * compressor_carry_kv(uint32_t source_layer) const; + ggml_tensor * compressor_carry_score(uint32_t source_layer) const; + ggml_tensor * candidate_scores() const; + ggml_tensor * candidate_ids() const; + ggml_tensor * committed_candidate_ids() const; + ggml_tensor * position_state() const; + + const llama_dsv41_memory_config & config() const; + llama_dsv41_memory_accounting accounting() const; + std::vector sequence_candidate_ids(llama_seq_id seq_id) const; + bool engram_enabled() const; + +private: + friend class llama_memory_dsv41_context; + struct impl; + std::unique_ptr pimpl; +}; diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645e..9ab61a480355 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -59,6 +59,10 @@ struct llama_memory_context_i { // return false on failure virtual bool apply() = 0; + // publish or discard state prepared by apply() + virtual void commit() {} + virtual void rollback() {} + // get the current ubatch virtual const llama_ubatch & get_ubatch() const = 0; @@ -117,6 +121,7 @@ struct llama_memory_i { virtual llama_pos seq_pos_max(llama_seq_id seq_id) const = 0; virtual std::map memory_breakdown() const = 0; + virtual void set_graph_workspace_size(size_t size) { GGML_UNUSED(size); } // // state write/read diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 554a374708b5..895c29be0489 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -18,6 +18,7 @@ #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-hybrid-idx.h" +#include "llama-memory-dsv41.h" #include "llama-memory-recurrent.h" #include "llama.h" @@ -2543,7 +2544,17 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } } break; case LLM_ARCH_DEEPSEEK41: - throw std::runtime_error(llama_dsv41_runtime_dependency_error()); + { + const auto & model_dsv41 = static_cast(*this); + res = new llama_memory_dsv41( + *this, + params.type_k, + cparams.offload_kqv, + cparams.n_ctx_seq, + cparams.n_seq_max, + cparams.n_ubatch, + model_dsv41.create_memory_engram_runtime(cparams.n_ubatch)); + } break; case LLM_ARCH_DFLASH: { // DSV4 DSpark stages store a single MLA-style K per position (window = the draft ring) diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index ba7923663f53..4d6b82d5881d 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -2,11 +2,14 @@ #include "llama-dsv41-engram.h" #include "llama-dsv41-expert.h" #include "llama-hparams.h" +#include "llama-memory-dsv41.h" #include "models.h" #include #include #include +#include +#include #include #include #include @@ -15,11 +18,438 @@ static float dsv41_rope_attn_factor(float freq_scale) { return 1.0f/(1.0f + 0.1f*logf(1.0f/freq_scale)); } +std::string llama_dsv41_graph_trace_name(const char * trace, uint32_t layer); + +std::string llama_dsv41_graph_trace_name( + const char * trace, + uint32_t layer) { + return "dsv41.trace." + std::string(trace) + ".l" + + std::to_string(layer); +} + struct llama_model_deepseek41::engram_model { llama_engram_layout layout; std::array extents; }; +std::unique_ptr llama_model_deepseek41::create_memory_engram_runtime( + size_t max_tokens) const { + return engram ? + std::make_unique(engram->layout, engram->extents, max_tokens) : + nullptr; +} + +namespace { + +struct dsv41_graph_source_input { + uint32_t layer = 0; + uint32_t ratio = 0; + uint32_t capacity = 0; + uint32_t read_width = 0; + + ggml_tensor * read_idxs = nullptr; + ggml_tensor * mask = nullptr; + ggml_tensor * carry_read_idxs = nullptr; + ggml_tensor * state_read_idxs = nullptr; + ggml_tensor * state_persist_src_idxs = nullptr; + ggml_tensor * state_persist_dst_idxs = nullptr; + ggml_tensor * write_idxs = nullptr; + ggml_tensor * write_pos = nullptr; + ggml_tensor * candidate_pad_mask = nullptr; + ggml_tensor * candidate_block_bias = nullptr; + ggml_tensor * row_blocks = nullptr; +}; + +class dsv41_graph_input final : public llm_graph_input_i { +public: + dsv41_graph_input( + ggml_context * ctx, + const llama_cparams & cparams, + const llama_hparams & hparams, + const llama_memory_dsv41_context * mctx, + const llama_ubatch & ubatch, + uint32_t n_outputs) : + cparams(cparams), + mctx(mctx), + topology(mctx->topology(ubatch, n_outputs)) { + const llama_memory_dsv41 * memory = mctx->memory(); + const uint32_t n_tokens = ubatch.n_tokens; + const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32; + const llama_dsv41_memory_plan & plan = mctx->graph_plan(ubatch); + + initial_pre = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.dsv4_hc_mult, n_tokens); + ggml_set_input(initial_pre); + ggml_set_name(initial_pre, "dsv41.inp.carried_pre"); + + raw_persist_src_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, plan.raw.persist_src_idxs.size()); + raw_write_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, plan.raw.write_idxs.size()); + raw_read_idxs = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, hparams.n_swa, n_tokens); + raw_mask = ggml_new_tensor_4d(ctx, type_mask, hparams.n_swa, 1, 1, n_tokens); + ggml_set_input(raw_persist_src_idxs); + ggml_set_input(raw_write_idxs); + ggml_set_input(raw_read_idxs); + ggml_set_input(raw_mask); + ggml_set_name(raw_persist_src_idxs, "dsv41.inp.raw.persist_src_idxs"); + ggml_set_name(raw_write_idxs, "dsv41.inp.raw.write_idxs"); + ggml_set_name(raw_read_idxs, "dsv41.inp.raw.read_idxs"); + ggml_set_name(raw_mask, "dsv41.inp.raw.mask"); + + for (const llama_dsv41_source_plan & source_plan : plan.sources) { + dsv41_graph_source_input source; + source.layer = source_plan.source_layer; + source.ratio = source_plan.ratio; + source.capacity = source_plan.capacity; + source.read_width = source_plan.read_idxs.empty() ? 0 : + source_plan.read_idxs.size()/n_tokens; + if (source.read_width > 0) { + source.read_idxs = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, source.read_width, n_tokens); + source.mask = ggml_new_tensor_4d( + ctx, type_mask, source.read_width, 1, 1, n_tokens); + ggml_set_input(source.read_idxs); + ggml_set_input(source.mask); + } + if (!source_plan.compression.state_read_idxs.empty()) { + source.carry_read_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, source.ratio); + source.state_read_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, source_plan.compression.state_read_idxs.size()); + ggml_set_input(source.carry_read_idxs); + ggml_set_input(source.state_read_idxs); + } + if (!source_plan.compression.state_persist_src_idxs.empty()) { + source.state_persist_src_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, source_plan.compression.state_persist_src_idxs.size()); + source.state_persist_dst_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, + source_plan.compression.state_persist_dst_idxs.size()*topology.n_seqs); + ggml_set_input(source.state_persist_src_idxs); + ggml_set_input(source.state_persist_dst_idxs); + } + if (!source_plan.compression.write_idxs.empty()) { + source.write_idxs = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, + source_plan.compression.write_idxs.size()*topology.n_seqs); + source.write_pos = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, source_plan.compression.write_pos.size()); + ggml_set_input(source.write_idxs); + ggml_set_input(source.write_pos); + } + if (source.layer == hparams.dsv41_candidate_source_layer && source.read_width > 0) { + const uint32_t block_size = hparams.dsv41_candidate_block_size; + const uint32_t n_blocks = (source.read_width + block_size - 1)/block_size; + const uint32_t padded = n_blocks*block_size; + source.candidate_pad_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, padded, n_tokens); + source.candidate_block_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_blocks, n_tokens); + source.row_blocks = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, source.read_width); + ggml_set_input(source.candidate_pad_mask); + ggml_set_input(source.candidate_block_bias); + ggml_set_input(source.row_blocks); + } + sources.emplace(source.layer, source); + } + + if (topology.engram_enabled) { + for (uint32_t index = 0; index < LLAMA_ENGRAM_LAYERS; ++index) { + engram_rows[index] = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, n_tokens); + engram_select[index] = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); + engram_row_ids[index] = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, LLAMA_ENGRAM_COLS, n_tokens); + ggml_set_input(engram_rows[index]); + ggml_set_input(engram_select[index]); + ggml_set_input(engram_row_ids[index]); + const uint32_t layer = index == 0 ? 1 : 14; + ggml_format_name( + engram_row_ids[index], + "dsv41.inp.engram.row_ids.l%u", layer); + } + } + + GGML_UNUSED(memory); + } + + void set_input(const llama_ubatch * ubatch) override { + const llama_dsv41_memory_plan & plan = mctx->plan(); + const llama_memory_dsv41 * memory = mctx->memory(); + const uint32_t n_tokens = ubatch->n_tokens; + const uint32_t n_streams = plan.token_seq_ids.empty() ? 1 : plan.token_seq_ids.front().size(); + if (n_streams != 1) { + throw std::runtime_error( + "DeepSeek V4.1 graph supports one active sequence per ubatch"); + } + + std::vector pre((size_t) LLAMA_DSV41_HC_MULT*n_tokens, 0.0f); + for (uint32_t token = 0; token < n_tokens; ++token) { + pre[(size_t) token*LLAMA_DSV41_HC_MULT] = 1.0f; + } + set_tensor(initial_pre, pre); + set_tensor(raw_persist_src_idxs, plan.raw.persist_src_idxs); + set_tensor(raw_write_idxs, plan.raw.write_idxs); + + std::vector raw_idxs = plan.raw.read_idxs; + const int32_t raw_cache_rows = + (int32_t) (memory->config().raw_window*memory->config().n_seq); + const int32_t raw_sentinel = raw_cache_rows + n_tokens; + const llama_pos batch_start = plan.positions.front(); + for (uint32_t token = 0; token < n_tokens; ++token) { + const uint32_t visible = plan.raw.n_visible[token]; + const llama_pos first = plan.positions[token] + 1 - visible; + for (uint32_t row = 0; row < visible; ++row) { + const llama_pos pos = first + row; + if (pos >= batch_start) { + raw_idxs[(size_t) token*memory->config().raw_window + row] = + raw_cache_rows + pos - batch_start; + } + } + } + std::replace(raw_idxs.begin(), raw_idxs.end(), -1, raw_sentinel); + set_tensor(raw_read_idxs, raw_idxs); + set_mask(raw_mask, plan.raw.mask); + + for (const llama_dsv41_source_plan & source_plan : plan.sources) { + dsv41_graph_source_input & source = sources.at(source_plan.source_layer); + if (source.read_idxs != nullptr) { + std::vector read_idxs = source_plan.read_idxs; + const int32_t sentinel = + (int32_t) (source_plan.capacity*memory->config().n_seq); + std::replace(read_idxs.begin(), read_idxs.end(), -1, sentinel); + set_tensor(source.read_idxs, read_idxs); + std::vector mask(read_idxs.size()); + std::transform( + source_plan.read_idxs.begin(), + source_plan.read_idxs.end(), + mask.begin(), + [](int32_t idx) { + return idx >= 0 ? 0.0f : -std::numeric_limits::infinity(); + }); + set_mask(source.mask, mask); + } + if (source.state_read_idxs != nullptr) { + std::vector carry_read_idxs(source.ratio); + const llama_seq_id seq_id = plan.token_seq_ids.front().front(); + for (uint32_t row = 0; row < source.ratio; ++row) { + carry_read_idxs[row] = seq_id*source.ratio + row; + } + set_tensor(source.carry_read_idxs, carry_read_idxs); + set_tensor(source.state_read_idxs, source_plan.compression.state_read_idxs); + } + if (source.state_persist_src_idxs != nullptr) { + set_tensor(source.state_persist_src_idxs, source_plan.compression.state_persist_src_idxs); + std::vector dst; + dst.reserve(source_plan.compression.state_persist_dst_idxs.size()*n_streams); + for (int32_t row : source_plan.compression.state_persist_dst_idxs) { + for (llama_seq_id seq_id : plan.token_seq_ids.front()) { + dst.push_back((int64_t) seq_id*source.ratio + row); + } + } + set_tensor(source.state_persist_dst_idxs, dst); + } + if (source.write_idxs != nullptr) { + std::vector write_idxs; + write_idxs.reserve(source_plan.compression.write_idxs.size()*n_streams); + for (int64_t row : source_plan.compression.write_idxs) { + for (llama_seq_id seq_id : plan.token_seq_ids.front()) { + write_idxs.push_back((int64_t) seq_id*source.capacity + row); + } + } + set_tensor(source.write_idxs, write_idxs); + set_tensor(source.write_pos, source_plan.compression.write_pos); + } + if (source.candidate_pad_mask != nullptr) { + const uint32_t block_size = memory->config().candidate_block_size; + const uint32_t n_blocks = + (source.read_width + block_size - 1)/block_size; + const uint32_t padded = n_blocks*block_size; + std::vector pad((size_t) padded*n_tokens, 0.0f); + for (uint32_t token = 0; token < n_tokens; ++token) { + for (uint32_t row = source.read_width; row < padded; ++row) { + pad[(size_t) token*padded + row] = + -std::numeric_limits::infinity(); + } + } + set_tensor(source.candidate_pad_mask, pad); + + std::vector bias((size_t) n_blocks*n_tokens, 0.0f); + for (uint32_t token = 0; token < n_tokens; ++token) { + const uint32_t visible = + source_plan.compression.n_visible[token]; + const uint32_t visible_blocks = + (visible + block_size - 1)/block_size; + for (uint32_t block = visible_blocks; block < n_blocks; ++block) { + bias[(size_t) token*n_blocks + block] = + -std::numeric_limits::infinity(); + } + if (visible_blocks > 0) { + bias[(size_t) token*n_blocks + visible_blocks - 1] = + std::numeric_limits::infinity(); + } + } + set_tensor(source.candidate_block_bias, bias); + + std::vector row_blocks(source.read_width); + for (uint32_t row = 0; row < source.read_width; ++row) { + row_blocks[row] = row/block_size; + } + set_tensor(source.row_blocks, row_blocks); + } + } + + const llama_dsv41_engram_transaction * transaction = mctx->engram_transaction(); + for (uint32_t index = 0; index < LLAMA_ENGRAM_LAYERS; ++index) { + if (engram_rows[index] == nullptr) { + continue; + } + if (transaction == nullptr) { + std::vector zeros( + (size_t) LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM*n_tokens); + std::vector select(n_tokens); + std::vector ids((size_t) LLAMA_ENGRAM_COLS*n_tokens); + set_tensor(engram_rows[index], zeros); + set_tensor(engram_select[index], select); + set_tensor(engram_row_ids[index], ids); + continue; + } + transaction->upload_layer( + index, 0, n_tokens, engram_rows[index], engram_select[index]); + std::vector ids((size_t) LLAMA_ENGRAM_COLS*n_tokens); + const uint32_t * source = transaction->row_ids(index); + for (uint32_t token = 0; token < n_tokens; ++token) { + std::copy_n( + source + (size_t) token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS, + ids.begin() + (size_t) token*LLAMA_ENGRAM_COLS); + } + set_tensor(engram_row_ids[index], ids); + } + } + + bool can_reuse(const llm_graph_params & params) override { + const auto * next = static_cast(params.mctx); + const llama_dsv41_graph_topology next_topology = + next->topology(params.ubatch, params.n_outputs); + if (!topology.same_topology(next_topology)) { + return false; + } + mctx = next; + topology = next_topology; + return true; + } + + dsv41_graph_source_input * source(uint32_t layer) { + const auto found = sources.find(layer); + return found == sources.end() ? nullptr : &found->second; + } + + const llama_cparams cparams; + const llama_memory_dsv41_context * mctx; + llama_dsv41_graph_topology topology; + ggml_tensor * initial_pre = nullptr; + ggml_tensor * raw_persist_src_idxs = nullptr; + ggml_tensor * raw_write_idxs = nullptr; + ggml_tensor * raw_read_idxs = nullptr; + ggml_tensor * raw_mask = nullptr; + std::map sources; + std::array engram_rows = {}; + std::array engram_select = {}; + std::array engram_row_ids = {}; + +private: + template + static void set_tensor(ggml_tensor * tensor, const std::vector & values) { + if (tensor == nullptr || tensor->buffer == nullptr) { + return; + } + if (ggml_nelements(tensor) != (int64_t) values.size()) { + throw std::runtime_error("DeepSeek V4.1 graph input shape changed"); + } + ggml_backend_tensor_set(tensor, values.data(), 0, values.size()*sizeof(T)); + } + + static void set_mask(ggml_tensor * tensor, const std::vector & values) { + if (tensor == nullptr || tensor->buffer == nullptr) { + return; + } + if (ggml_nelements(tensor) != (int64_t) values.size()) { + throw std::runtime_error("DeepSeek V4.1 graph mask shape changed"); + } + if (tensor->type == GGML_TYPE_F16) { + std::vector converted(values.size()); + ggml_fp32_to_fp16_row(values.data(), converted.data(), values.size()); + ggml_backend_tensor_set( + tensor, converted.data(), 0, converted.size()*sizeof(ggml_fp16_t)); + } else { + ggml_backend_tensor_set( + tensor, values.data(), 0, values.size()*sizeof(float)); + } + } +}; + +static ggml_tensor * dsv41_flatten_memory( + ggml_context * ctx, + ggml_tensor * tensor) { + return ggml_reshape_2d( + ctx, tensor, tensor->ne[0], tensor->ne[1]*tensor->ne[2]); +} + +static ggml_tensor * dsv41_append_zero_row( + ggml_context * ctx, + ggml_tensor * tensor) { + ggml_tensor * row = ggml_view_2d( + ctx, tensor, tensor->ne[0], 1, tensor->nb[1], 0); + row = ggml_scale(ctx, row, 0.0f); + return ggml_concat(ctx, tensor, row, 1); +} + +static ggml_tensor * dsv41_hc_mean( + ggml_context * ctx, + ggml_tensor * streams) { + ggml_tensor * result = ggml_view_2d( + ctx, streams, streams->ne[0], streams->ne[2], streams->nb[2], 0); + for (int64_t stream = 1; stream < streams->ne[1]; ++stream) { + result = ggml_add(ctx, result, ggml_view_2d( + ctx, streams, streams->ne[0], streams->ne[2], + streams->nb[2], stream*streams->nb[1])); + } + return ggml_scale(ctx, result, 1.0f/streams->ne[1]); +} + +static ggml_tensor * dsv41_view_1d( + ggml_context * ctx, + ggml_tensor * tensor, + int64_t ne0, + int64_t offset) { + return ggml_view_1d(ctx, tensor, ne0, ggml_row_size(tensor->type, offset)); +} + +static ggml_tensor * dsv41_view_2d( + ggml_context * ctx, + ggml_tensor * tensor, + int64_t ne0, + int64_t ne1, + int64_t offset) { + return ggml_view_2d( + ctx, tensor, ne0, ne1, tensor->nb[1], + ggml_row_size(tensor->type, offset)); +} + +static ggml_tensor * dsv41_sort_row_ids( + ggml_context * ctx, + ggml_tensor * ids) { + ggml_tensor * order = ggml_argsort( + ctx, ggml_cast(ctx, ids, GGML_TYPE_F32), + GGML_SORT_ORDER_ASC); + ggml_tensor * sorted = ggml_get_rows( + ctx, + ggml_reshape_3d(ctx, ids, 1, ids->ne[0], ids->ne[1]), + order); + return ggml_reshape_2d(ctx, sorted, ids->ne[0], ids->ne[1]); +} + +} + void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { llama_dsv41_config config = {}; std::string raw_config; @@ -305,6 +735,777 @@ void llama_model_deepseek41::release_runtime_context() const { } } -[[noreturn]] std::unique_ptr llama_model_deepseek41::build_arch_graph(const llm_graph_params &) const { - throw std::runtime_error(llama_dsv41_runtime_dependency_error()); +namespace { + +struct dsv41_hc_mix { + ggml_tensor * pre; + ggml_tensor * post; + ggml_tensor * comb; +}; + +static dsv41_hc_mix dsv41_build_hc_mix( + const llama_model_deepseek41::graph & graph, + ggml_tensor * streams, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + int il) { + const int64_t hc = graph.hparams.dsv4_hc_mult; + const int64_t nt = streams->ne[2]; + ggml_tensor * flat = ggml_reshape_2d(graph.ctx0, streams, graph.n_embd*hc, nt); + flat = ggml_rms_norm(graph.ctx0, flat, graph.norm_rms_eps); + ggml_tensor * mixes = ggml_mul_mat(graph.ctx0, hc_fn, flat); + graph.cb(mixes, "hc_mixes", il); + + ggml_tensor * scale_pre = dsv41_view_1d(graph.ctx0, hc_scale, 1, 0); + ggml_tensor * scale_post = dsv41_view_1d(graph.ctx0, hc_scale, 1, 1); + ggml_tensor * scale_comb = dsv41_view_1d(graph.ctx0, hc_scale, 1, 2); + ggml_tensor * base_pre = dsv41_view_1d(graph.ctx0, hc_base, hc, 0); + ggml_tensor * base_post = dsv41_view_1d(graph.ctx0, hc_base, hc, hc); + ggml_tensor * base_comb = dsv41_view_1d(graph.ctx0, hc_base, hc*hc, 2*hc); + + ggml_tensor * pre = dsv41_view_2d(graph.ctx0, mixes, hc, nt, 0); + pre = ggml_add(graph.ctx0, ggml_mul(graph.ctx0, pre, scale_pre), base_pre); + pre = ggml_scale_bias( + graph.ctx0, ggml_sigmoid(graph.ctx0, pre), 1.0f, + graph.hparams.dsv4_hc_eps); + graph.cb(pre, "hc_pre", il); + + ggml_tensor * post = dsv41_view_2d(graph.ctx0, mixes, hc, nt, hc); + post = ggml_add(graph.ctx0, ggml_mul(graph.ctx0, post, scale_post), base_post); + post = ggml_scale(graph.ctx0, ggml_sigmoid(graph.ctx0, post), 2.0f); + graph.cb(post, "hc_post", il); + + ggml_tensor * comb = nullptr; + if (graph.cparams.fused_dsv4_hc_comb) { + comb = ggml_dsv4_hc_comb( + graph.ctx0, mixes, hc_scale, hc_base, + graph.hparams.dsv4_hc_eps, + (int32_t) graph.hparams.dsv4_hc_sinkhorn_iters); + graph.res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, comb, il}); + } else { + comb = dsv41_view_2d(graph.ctx0, mixes, hc*hc, nt, 2*hc); + comb = ggml_add(graph.ctx0, ggml_mul(graph.ctx0, comb, scale_comb), base_comb); + comb = ggml_reshape_3d(graph.ctx0, comb, hc, hc, nt); + comb = graph.build_hc_sinkhorn(comb, il); + } + graph.cb(comb, "hc_comb", il); + return { pre, post, comb }; +} + +static ggml_tensor * dsv41_cast_for_store( + ggml_context * ctx, + ggml_tensor * source, + const ggml_tensor * destination) { + return source->type == destination->type ? + source : ggml_cast(ctx, source, destination->type); +} + +static ggml_tensor * dsv41_completion_zero( + ggml_context * ctx, + ggml_tensor * dependency, + ggml_type type) { + ggml_tensor * marker = ggml_view_1d(ctx, dependency, 1, 0); + marker = ggml_argsort_top_k(ctx, marker, 1); + marker = ggml_cast(ctx, marker, type); + return ggml_scale(ctx, marker, 0.0f); +} + +static ggml_tensor * dsv41_build_candidate_mask( + const llama_model_deepseek41::graph & graph, + ggml_tensor * candidate_blocks, + const dsv41_graph_source_input & source) { + const int64_t n_blocks = source.candidate_block_bias->ne[0]; + const int64_t n_tokens = candidate_blocks->ne[1]; + ggml_tensor * all = ggml_new_tensor_3d( + graph.ctx0, GGML_TYPE_F32, 1, n_blocks, n_tokens); + all = ggml_fill(graph.ctx0, all, -INFINITY); + ggml_tensor * zeros = ggml_scale( + graph.ctx0, ggml_cast(graph.ctx0, candidate_blocks, GGML_TYPE_F32), 0.0f); + zeros = ggml_reshape_3d( + graph.ctx0, zeros, 1, candidate_blocks->ne[0], n_tokens); + ggml_tensor * blocks = ggml_set_rows(graph.ctx0, all, zeros, candidate_blocks); + blocks = ggml_reshape_2d(graph.ctx0, blocks, n_blocks, n_tokens); + blocks = ggml_cont(graph.ctx0, ggml_transpose(graph.ctx0, blocks)); + ggml_tensor * rows = ggml_get_rows(graph.ctx0, blocks, source.row_blocks); + return ggml_cont(graph.ctx0, ggml_transpose(graph.ctx0, rows)); +} + +static ggml_tensor * dsv41_build_index_selection( + const llama_model_deepseek41::graph & graph, + const llama_model & model, + dsv41_graph_input & input, + ggml_tensor * qr, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * & candidate_blocks, + ggml_tensor * & selected_local, + int il) { + const auto & layer = model.layers[il]; + const uint32_t source_layer = graph.hparams.dsv41_kv_source_layer[il]; + dsv41_graph_source_input * source = input.source(source_layer); + if (source == nullptr || source->read_width == 0) { + return nullptr; + } + + const llama_memory_dsv41 * memory = input.mctx->memory(); + ggml_tensor * index_cache = dsv41_flatten_memory( + graph.ctx0, memory->index_keys(source_layer)); + index_cache = dsv41_append_zero_row(graph.ctx0, index_cache); + ggml_tensor * index_k = ggml_get_rows(graph.ctx0, index_cache, source->read_idxs); + graph.cb(index_k, "dsv41_index_k", il); + + ggml_tensor * index_q = graph.build_lora_mm(layer.indexer_attn_q_b, qr); + index_q = ggml_reshape_3d( + graph.ctx0, index_q, + graph.hparams.indexer_head_size, + graph.hparams.indexer_n_head, + graph.n_tokens); + index_q = ggml_rope_ext( + graph.ctx0, index_q, inp_pos, nullptr, graph.hparams.n_rot(), + graph.rope_type, graph.n_ctx_orig, + graph.hparams.dsv4_compress_rope_base, + graph.freq_scale, graph.ext_factor, + dsv41_rope_attn_factor(graph.freq_scale), + graph.beta_fast, graph.beta_slow); + index_q = ggml_rope_set_offset( + index_q, graph.hparams.indexer_head_size - graph.hparams.n_rot()); + graph.cb(index_q, "dsv41_index_q", il); + + ggml_tensor * weights = graph.build_lora_mm(layer.indexer_proj, cur); + weights = ggml_scale( + graph.ctx0, weights, + 1.0f/std::sqrt( + (float) (graph.hparams.indexer_head_size*graph.hparams.indexer_n_head))); + weights = ggml_reshape_3d( + graph.ctx0, weights, graph.hparams.indexer_n_head, 1, graph.n_tokens); + + ggml_tensor * scores = ggml_mul_mat(graph.ctx0, index_k, index_q); + ggml_prec_set_acc(scores, GGML_PREC_F32); + scores = ggml_relu(graph.ctx0, scores); + scores = ggml_cont(graph.ctx0, ggml_permute(graph.ctx0, scores, 1, 0, 2, 3)); + scores = ggml_mul(graph.ctx0, scores, weights); + scores = ggml_sum_rows(graph.ctx0, scores); + scores = ggml_reshape_2d(graph.ctx0, scores, source->read_width, graph.n_tokens); + scores = ggml_add( + graph.ctx0, scores, + ggml_reshape_2d( + graph.ctx0, source->mask, source->read_width, graph.n_tokens)); + + if (candidate_blocks != nullptr) { + scores = ggml_add( + graph.ctx0, scores, + dsv41_build_candidate_mask(graph, candidate_blocks, *source)); + } + graph.cb(scores, "dsv41_index_scores", il); + + if ((uint32_t) il == graph.hparams.dsv41_candidate_source_layer) { + const uint32_t block_size = graph.hparams.dsv41_candidate_block_size; + const int64_t padded = source->candidate_pad_mask->ne[0]; + ggml_tensor * padded_scores = ggml_pad( + graph.ctx0, scores, padded - source->read_width, 0, 0, 0); + padded_scores = ggml_add( + graph.ctx0, padded_scores, source->candidate_pad_mask); + ggml_tensor * block_scores = ggml_pool_1d( + graph.ctx0, padded_scores, GGML_OP_POOL_MAX, + block_size, block_size, 0); + block_scores = ggml_add( + graph.ctx0, block_scores, source->candidate_block_bias); + graph.cb(block_scores, "dsv41_candidate_scores", il); + + const uint32_t n_candidate = input.topology.candidate_width; + candidate_blocks = ggml_cont( + graph.ctx0, + ggml_argsort_top_k( + graph.ctx0, block_scores, n_candidate)); + ggml_set_name( + candidate_blocks, + llama_dsv41_graph_trace_name( + "attn.candidate_blocks", il).c_str()); + ggml_build_forward_expand(graph.gf, candidate_blocks); + + ggml_tensor * score_store = memory->candidate_scores(); + ggml_tensor * score_view = ggml_view_2d( + graph.ctx0, score_store, block_scores->ne[0], graph.n_tokens, + score_store->nb[1], 0); + ggml_build_forward_expand( + graph.gf, ggml_cpy(graph.ctx0, block_scores, score_view)); + ggml_tensor * id_store = memory->candidate_ids(); + ggml_tensor * id_view = ggml_view_2d( + graph.ctx0, id_store, candidate_blocks->ne[0], graph.n_tokens, + id_store->nb[1], 0); + ggml_build_forward_expand( + graph.gf, ggml_cpy(graph.ctx0, candidate_blocks, id_view)); + } + + const uint32_t n_top_k = std::min( + source->read_width, graph.hparams.indexer_top_k); + selected_local = ggml_cont( + graph.ctx0, + ggml_argsort_top_k(graph.ctx0, scores, n_top_k)); + selected_local = dsv41_sort_row_ids( + graph.ctx0, selected_local); + ggml_tensor * selected = ggml_get_rows( + graph.ctx0, + ggml_reshape_3d( + graph.ctx0, source->read_idxs, 1, + source->read_width, graph.n_tokens), + selected_local); + selected = ggml_cont( + graph.ctx0, + ggml_reshape_2d( + graph.ctx0, selected, + n_top_k, graph.n_tokens)); + if ((uint32_t) il > graph.hparams.dsv41_candidate_source_layer) { + ggml_set_name( + selected, + llama_dsv41_graph_trace_name( + "attn.candidates", il).c_str()); + ggml_build_forward_expand(graph.gf, selected); + } + return selected; +} + +static ggml_tensor * dsv41_build_attention( + const llama_model_deepseek41::graph & graph, + const llama_model & model, + dsv41_graph_input & input, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * & selected, + ggml_tensor * & selected_local, + ggml_tensor * & candidate_blocks, + int il) { + const auto & layer = model.layers[il]; + const int64_t n_embd_head = graph.hparams.n_embd_head_k(); + const int64_t n_rot = graph.hparams.n_rot(); + const int64_t n_nope = n_embd_head - n_rot; + const int64_t ratio = graph.hparams.dsv4_compress_ratios[il]; + const int64_t n_groups = graph.hparams.dsv4_o_group_count; + const int64_t heads_per_group = graph.n_head/n_groups; + const int64_t group_width = heads_per_group*n_embd_head; + const int64_t o_lora_rank = graph.hparams.dsv4_o_lora_rank; + const float rope_base = ratio == 0 ? + graph.freq_base : graph.hparams.dsv4_compress_rope_base; + const float rope_scale = ratio == 0 ? 1.0f : graph.freq_scale; + const float rope_ext = ratio == 0 ? 0.0f : graph.ext_factor; + const float rope_attn = ratio == 0 ? + 1.0f : dsv41_rope_attn_factor(graph.freq_scale); + const float rope_beta_fast = ratio == 0 ? 0.0f : graph.beta_fast; + const float rope_beta_slow = ratio == 0 ? 0.0f : graph.beta_slow; + const int32_t rope_ctx = ratio == 0 ? 0 : graph.n_ctx_orig; + + ggml_tensor * qr = graph.build_lora_mm(layer.wq_a, cur); + qr = graph.build_norm( + qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + graph.cb(qr, "dsv41_qr", il); + + ggml_tensor * q = graph.build_lora_mm(layer.wq_b, qr); + q = ggml_reshape_3d( + graph.ctx0, q, n_embd_head, graph.n_head, graph.n_tokens); + q = ggml_rms_norm(graph.ctx0, q, graph.norm_rms_eps); + q = ggml_rope_ext( + graph.ctx0, q, inp_pos, nullptr, n_rot, graph.rope_type, + rope_ctx, rope_base, rope_scale, rope_ext, rope_attn, + rope_beta_fast, rope_beta_slow); + q = ggml_rope_set_offset(q, n_nope); + graph.cb(q, "dsv41_q", il); + + ggml_tensor * kv = graph.build_lora_mm(layer.wkv, cur); + kv = graph.build_norm( + kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + kv = ggml_reshape_3d( + graph.ctx0, kv, n_embd_head, 1, graph.n_tokens); + ggml_tensor * raw_kv = ggml_rope_ext( + graph.ctx0, kv, inp_pos, nullptr, n_rot, graph.rope_type, + rope_ctx, rope_base, rope_scale, rope_ext, rope_attn, + rope_beta_fast, rope_beta_slow); + raw_kv = ggml_rope_set_offset(raw_kv, n_nope); + graph.cb(raw_kv, "dsv41_raw_kv", il); + + const llama_memory_dsv41 * memory = input.mctx->memory(); + ggml_tensor * raw_store = dsv41_flatten_memory( + graph.ctx0, memory->raw_k(il)); + ggml_tensor * raw_write = ggml_reshape_2d( + graph.ctx0, raw_kv, n_embd_head, graph.n_tokens); + ggml_tensor * raw_persist = ggml_get_rows( + graph.ctx0, raw_write, input.raw_persist_src_idxs); + raw_write = dsv41_cast_for_store( + graph.ctx0, raw_write, raw_store); + + ggml_tensor * raw_read = ggml_concat( + graph.ctx0, raw_store, raw_write, 1); + raw_read = dsv41_append_zero_row(graph.ctx0, raw_read); + raw_read = ggml_get_rows( + graph.ctx0, raw_read, input.raw_read_idxs); + raw_read = ggml_reshape_4d( + graph.ctx0, raw_read, n_embd_head, 1, + graph.hparams.n_swa, graph.n_tokens); + + const int32_t kv_source = graph.hparams.dsv41_kv_source_layer[il]; + dsv41_graph_source_input * source = + kv_source >= 0 ? input.source(kv_source) : nullptr; + if (source != nullptr && source->layer == (uint32_t) il) { + ggml_tensor * compressed = graph.build_lora_mm( + layer.attn_comp_wkv, cur); + ggml_tensor * gate = ratio == 2 ? + graph.build_lora_mm(layer.attn_comp_wgate, cur) : nullptr; + ggml_tensor * compressed_current = compressed; + ggml_tensor * gate_current = gate; + ggml_tensor * carry_kv = nullptr; + ggml_tensor * carry_gate = nullptr; + ggml_tensor * persist_kv = nullptr; + ggml_tensor * persist_gate = nullptr; + + if (ratio == 2) { + carry_kv = dsv41_flatten_memory( + graph.ctx0, memory->compressor_carry_kv(source->layer)); + carry_gate = dsv41_flatten_memory( + graph.ctx0, memory->compressor_carry_score(source->layer)); + if (source->write_idxs != nullptr) { + ggml_tensor * carry_kv_first = ggml_get_rows( + graph.ctx0, carry_kv, source->carry_read_idxs); + ggml_tensor * carry_gate_first = ggml_get_rows( + graph.ctx0, carry_gate, source->carry_read_idxs); + ggml_tensor * source_kv = ggml_concat( + graph.ctx0, carry_kv_first, compressed, 1); + ggml_tensor * source_gate = ggml_concat( + graph.ctx0, carry_gate_first, gate, 1); + source_kv = ggml_get_rows( + graph.ctx0, source_kv, source->state_read_idxs); + source_gate = ggml_get_rows( + graph.ctx0, source_gate, source->state_read_idxs); + const int64_t n_write = source->write_pos->ne[0]; + source_kv = ggml_reshape_3d( + graph.ctx0, source_kv, n_embd_head, ratio, n_write); + source_gate = ggml_reshape_3d( + graph.ctx0, source_gate, n_embd_head, ratio, n_write); + compressed = llama_dsv41_build_ratio_pool( + graph.ctx0, source_kv, source_gate, ratio); + } else { + compressed = nullptr; + } + + persist_kv = ggml_get_rows( + graph.ctx0, compressed_current, + source->state_persist_src_idxs); + persist_gate = ggml_get_rows( + graph.ctx0, gate_current, + source->state_persist_src_idxs); + } else { + compressed = llama_dsv41_build_ratio_pool( + graph.ctx0, + ggml_reshape_3d( + graph.ctx0, compressed, n_embd_head, 1, + graph.n_tokens), + nullptr, 1); + } + + if (compressed != nullptr) { + compressed = graph.build_norm( + compressed, layer.attn_comp_norm, nullptr, + LLM_NORM_RMS, il); + graph.cb(compressed, "dsv41_compressed_unrotated", il); + + ggml_tensor * index_k = graph.build_lora_mm( + layer.indexer_attn_k, compressed); + index_k = graph.build_norm( + index_k, layer.indexer_k_norm, nullptr, + LLM_NORM_RMS, il); + index_k = ggml_reshape_3d( + graph.ctx0, index_k, + graph.hparams.indexer_head_size, 1, index_k->ne[1]); + index_k = ggml_rope_ext( + graph.ctx0, index_k, source->write_pos, nullptr, + graph.hparams.n_rot(), graph.rope_type, graph.n_ctx_orig, + graph.hparams.dsv4_compress_rope_base, + graph.freq_scale, graph.ext_factor, + dsv41_rope_attn_factor(graph.freq_scale), + graph.beta_fast, graph.beta_slow); + index_k = ggml_rope_set_offset( + index_k, + graph.hparams.indexer_head_size - graph.hparams.n_rot()); + + ggml_tensor * compressed_rope = ggml_reshape_3d( + graph.ctx0, compressed, n_embd_head, 1, compressed->ne[1]); + compressed_rope = ggml_rope_ext( + graph.ctx0, compressed_rope, source->write_pos, nullptr, + n_rot, graph.rope_type, graph.n_ctx_orig, + graph.hparams.dsv4_compress_rope_base, + graph.freq_scale, graph.ext_factor, + dsv41_rope_attn_factor(graph.freq_scale), + graph.beta_fast, graph.beta_slow); + compressed_rope = ggml_rope_set_offset(compressed_rope, n_nope); + + ggml_tensor * comp_store = dsv41_flatten_memory( + graph.ctx0, memory->compressed_kv(source->layer)); + ggml_tensor * index_store = dsv41_flatten_memory( + graph.ctx0, memory->index_keys(source->layer)); + ggml_tensor * comp_write = ggml_reshape_2d( + graph.ctx0, compressed_rope, n_embd_head, + compressed_rope->ne[2]); + ggml_tensor * index_write = ggml_reshape_2d( + graph.ctx0, index_k, graph.hparams.indexer_head_size, + index_k->ne[2]); + comp_write = dsv41_cast_for_store( + graph.ctx0, comp_write, comp_store); + index_write = dsv41_cast_for_store( + graph.ctx0, index_write, index_store); + ggml_build_forward_expand( + graph.gf, ggml_set_rows( + graph.ctx0, comp_store, comp_write, + source->write_idxs)); + ggml_build_forward_expand( + graph.gf, ggml_set_rows( + graph.ctx0, index_store, index_write, + source->write_idxs)); + } + + if (persist_kv != nullptr) { + if (compressed != nullptr) { + ggml_tensor * completion = dsv41_completion_zero( + graph.ctx0, compressed, persist_kv->type); + persist_kv = ggml_add( + graph.ctx0, persist_kv, completion); + persist_gate = ggml_add( + graph.ctx0, persist_gate, completion); + } + persist_kv = dsv41_cast_for_store( + graph.ctx0, persist_kv, carry_kv); + persist_gate = dsv41_cast_for_store( + graph.ctx0, persist_gate, carry_gate); + ggml_build_forward_expand( + graph.gf, ggml_set_rows( + graph.ctx0, carry_kv, persist_kv, + source->state_persist_dst_idxs)); + ggml_build_forward_expand( + graph.gf, ggml_set_rows( + graph.ctx0, carry_gate, persist_gate, + source->state_persist_dst_idxs)); + } + } + + if (graph.hparams.dsv41_is_index_source(il)) { + selected = dsv41_build_index_selection( + graph, model, input, qr, cur, inp_pos, + candidate_blocks, selected_local, + il); + } + + ggml_tensor * source_trace = nullptr; + if (selected != nullptr) { + source_trace = ggml_cont(graph.ctx0, selected); + ggml_set_name( + source_trace, + llama_dsv41_graph_trace_name( + "attn.source", il).c_str()); + } else { + source_trace = ggml_cont(graph.ctx0, input.raw_read_idxs); + ggml_set_name( + source_trace, + llama_dsv41_graph_trace_name( + "attn.source", il).c_str()); + } + ggml_build_forward_expand(graph.gf, source_trace); + + ggml_tensor * k_all = raw_read; + ggml_tensor * mask_all = input.raw_mask; + int64_t n_kv_max = graph.hparams.n_swa; + if (selected != nullptr) { + ggml_tensor * comp_store = dsv41_flatten_memory( + graph.ctx0, memory->compressed_kv(kv_source)); + ggml_tensor * compressed = ggml_get_rows( + graph.ctx0, comp_store, selected); + compressed = ggml_reshape_4d( + graph.ctx0, compressed, n_embd_head, 1, + selected->ne[0], graph.n_tokens); + k_all = ggml_concat(graph.ctx0, raw_read, compressed, 2); + + ggml_tensor * source_mask = ggml_reshape_3d( + graph.ctx0, source->mask, 1, + source->read_width, graph.n_tokens); + ggml_tensor * compressed_mask = ggml_get_rows( + graph.ctx0, source_mask, selected_local); + compressed_mask = ggml_cont( + graph.ctx0, + ggml_permute( + graph.ctx0, compressed_mask, 1, 0, 2, 3)); + compressed_mask = ggml_reshape_4d( + graph.ctx0, compressed_mask, + selected->ne[0], 1, 1, graph.n_tokens); + mask_all = ggml_concat( + graph.ctx0, input.raw_mask, compressed_mask, 0); + n_kv_max += selected->ne[0]; + } + + ggml_tensor * out = graph.build_attn_mha( + q, k_all, k_all, nullptr, mask_all, + layer.attn_sinks, nullptr, n_kv_max, + 1.0f/std::sqrt((float) n_embd_head), il); + + // Keep the prior ring intact until attention has consumed it. + ggml_tensor * completion = dsv41_completion_zero( + graph.ctx0, out, raw_persist->type); + raw_persist = ggml_add( + graph.ctx0, raw_persist, completion); + raw_persist = dsv41_cast_for_store( + graph.ctx0, raw_persist, raw_store); + ggml_tensor * raw_update = ggml_set_rows( + graph.ctx0, raw_store, raw_persist, + input.raw_write_idxs); + ggml_build_forward_expand(graph.gf, raw_update); + out = ggml_reshape_3d( + graph.ctx0, out, n_embd_head, graph.n_head, + graph.n_tokens); + out = ggml_rope_ext_back( + graph.ctx0, out, inp_pos, nullptr, n_rot, + graph.rope_type, rope_ctx, rope_base, rope_scale, + rope_ext, rope_attn, rope_beta_fast, rope_beta_slow); + out = ggml_rope_set_offset(out, n_nope); + graph.cb(out, "dsv41_attn_derope", il); + + out = ggml_reshape_3d( + graph.ctx0, out, group_width, n_groups, + graph.n_tokens); + out = ggml_permute(graph.ctx0, out, 0, 2, 1, 3); + if (graph.n_tokens > 1 && graph.n_tokens <= 8) { + out = ggml_cont(graph.ctx0, out); + } + ggml_tensor * oa = ggml_mul_mat(graph.ctx0, layer.wo_a, out); + oa = ggml_permute(graph.ctx0, oa, 0, 2, 1, 3); + oa = ggml_cont_2d( + graph.ctx0, oa, o_lora_rank*n_groups, + graph.n_tokens); + out = graph.build_lora_mm(layer.wo_b, oa); + graph.cb(out, "dsv41_attn_out", il); + return out; +} + +static std::pair dsv41_build_router( + const llama_model_deepseek41::graph & graph, + const llama_layer & layer, + ggml_tensor * cur, + int il) { + ggml_tensor * logits = graph.build_lora_mm(layer.ffn_gate_inp, cur); + ggml_prec_set_acc(logits, GGML_PREC_F32); + ggml_tensor * probs = ggml_sqrt( + graph.ctx0, ggml_softplus(graph.ctx0, logits)); + ggml_tensor * selection = ggml_add( + graph.ctx0, probs, layer.ffn_exp_probs_b); + ggml_tensor * ids = ggml_cont( + graph.ctx0, + ggml_argsort_top_k( + graph.ctx0, selection, + graph.hparams.n_expert_used())); + ids = dsv41_sort_row_ids(graph.ctx0, ids); + ggml_set_name( + ids, + llama_dsv41_graph_trace_name( + "expert.ids", il).c_str()); + + ggml_tensor * weights = ggml_get_rows( + graph.ctx0, + ggml_reshape_3d( + graph.ctx0, probs, 1, graph.n_expert, + graph.n_tokens), + ids); + weights = ggml_reshape_2d( + graph.ctx0, weights, + graph.hparams.n_expert_used(), graph.n_tokens); + ggml_tensor * sum = ggml_clamp( + graph.ctx0, ggml_sum_rows(graph.ctx0, weights), + 6.103515625e-5, INFINITY); + weights = ggml_div(graph.ctx0, weights, sum); + weights = ggml_scale( + graph.ctx0, weights, + graph.hparams.expert_weights_scale); + weights = ggml_cont(graph.ctx0, weights); + ggml_set_name( + weights, + llama_dsv41_graph_trace_name( + "expert.weights", il).c_str()); + ggml_build_forward_expand(graph.gf, weights); + return { logits, ids }; +} + +} + +llama_model_deepseek41::graph::graph( + const llama_model & model, + const llm_graph_params & params) : + llama_model_deepseek4::graph(params) { + GGML_ASSERT(n_layer == LLAMA_DSV41_N_LAYER); + GGML_ASSERT(hparams.dsv4_hc_mult == LLAMA_DSV41_HC_MULT); + GGML_ASSERT(mctx != nullptr); + + const auto * dsv41_mctx = + static_cast(mctx); + auto input_owner = std::make_unique( + ctx0, cparams, hparams, dsv41_mctx, + ubatch, n_outputs); + auto * input = static_cast( + res->add_input(std::move(input_owner))); + if (input->topology.n_seqs != 1) { + throw std::invalid_argument( + "DeepSeek V4.1 graph supports one active sequence per ubatch"); + } + + ggml_tensor * inp = build_inp_embd(model.tok_embd); + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + ggml_tensor * streams = ggml_reshape_3d( + ctx0, inp, n_embd, 1, n_tokens); + streams = ggml_repeat_4d( + ctx0, streams, n_embd, + hparams.dsv4_hc_mult, n_tokens, 1); + ggml_tensor * carried_pre = input->initial_pre; + ggml_tensor * selected = nullptr; + ggml_tensor * selected_local = nullptr; + ggml_tensor * candidate_blocks = nullptr; + cb(streams, "dsv41_hc_init", -1); + + for (int il = 0; il < n_layer; ++il) { + if ((size_t) il < cparams.embeddings_layer_inp.size() && + cparams.embeddings_layer_inp[il]) { + res->t_layer_inp[il] = dsv41_hc_mean(ctx0, streams); + cb(res->t_layer_inp[il], "layer_inp", il); + ggml_build_forward_expand(gf, res->t_layer_inp[il]); + } + + if (hparams.dsv41_engram_layers.test(il)) { + const uint32_t index = il == 1 ? 0 : 1; + ggml_tensor * row_ids = ggml_cont( + ctx0, input->engram_row_ids[index]); + ggml_set_name( + row_ids, + llama_dsv41_graph_trace_name( + "engram.row_ids", il).c_str()); + ggml_build_forward_expand(gf, row_ids); + streams = llama_dsv41_build_engram( + ctx0, streams, + input->engram_rows[index], + model.layers[il].engram_kv, + model.layers[il].engram_q_norm, + model.layers[il].engram_k_norm, + input->engram_select[index], + norm_rms_eps, sched, backend_cpu); + cb(streams, "dsv41_engram", il); + } + + ggml_tensor * residual = streams; + const dsv41_hc_mix attn_mix = dsv41_build_hc_mix( + *this, streams, + model.layers[il].hc_attn_fn, + model.layers[il].hc_attn_scale, + model.layers[il].hc_attn_base, + il); + ggml_tensor * cur = build_hc_pre( + streams, carried_pre, il); + cb(cur, "dsv41_hc_attn_carried_pre", il); + cur = build_norm( + cur, model.layers[il].attn_norm, + nullptr, LLM_NORM_RMS, il); + cur = dsv41_build_attention( + *this, model, *input, cur, inp_pos, + selected, selected_local, candidate_blocks, il); + streams = build_hc_post( + cur, residual, attn_mix.post, attn_mix.comb, il); + cb(streams, "dsv41_hc_attn_post", il); + + residual = streams; + const dsv41_hc_mix ffn_mix = dsv41_build_hc_mix( + *this, streams, + model.layers[il].hc_ffn_fn, + model.layers[il].hc_ffn_scale, + model.layers[il].hc_ffn_base, + il); + cur = build_hc_pre(streams, attn_mix.pre, il); + cb(cur, "dsv41_hc_ffn_attn_pre", il); + cur = build_norm( + cur, model.layers[il].ffn_norm, + nullptr, LLM_NORM_RMS, il); + + const auto [router_logits, original_ids] = + dsv41_build_router(*this, model.layers[il], cur, il); + ggml_tensor * slot_ids = llama_dsv41_build_expert_remap( + ctx0, original_ids, + *static_cast(model).experts, + il, sched, backend_cpu); + ggml_tensor * moe_out = build_moe_ffn( + cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + model.layers[il].ffn_exp_probs_b, + n_expert, hparams.n_expert_used(), + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + router_logits, + nullptr, nullptr, nullptr, nullptr, + original_ids, + slot_ids); + cb(moe_out, "dsv41_ffn_moe", il); + + ggml_tensor * shared = build_ffn( + cur, + model.layers[il].ffn_up_shexp, nullptr, nullptr, + model.layers[il].ffn_gate_shexp, nullptr, nullptr, + model.layers[il].ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(shared, "dsv41_ffn_shared", il); + cur = ggml_add(ctx0, moe_out, shared); + cb(cur, "dsv41_ffn_out", il); + + ggml_tensor * release = llama_dsv41_build_expert_release( + ctx0, moe_out, + *static_cast(model).experts, + il, sched, backend_cpu); + ggml_build_forward_expand(gf, release); + + streams = build_hc_post( + cur, residual, ffn_mix.post, ffn_mix.comb, il); + streams = build_cvec(streams, il); + carried_pre = ffn_mix.pre; + cb(streams, "dsv41_layer_out", il); + } + + if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && + cparams.embeddings_layer_inp[n_layer]) { + res->t_layer_inp[n_layer] = dsv41_hc_mean(ctx0, streams); + cb(res->t_layer_inp[n_layer], "layer_inp", n_layer); + ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]); + } + + if (inp_out_ids != nullptr) { + ggml_tensor * flat = ggml_reshape_2d( + ctx0, streams, + n_embd*hparams.dsv4_hc_mult, n_tokens); + flat = ggml_get_rows(ctx0, flat, inp_out_ids); + streams = ggml_reshape_3d( + ctx0, flat, n_embd, + hparams.dsv4_hc_mult, n_outputs); + carried_pre = ggml_get_rows(ctx0, carried_pre, inp_out_ids); + } + + ggml_tensor * cur = llama_dsv41_build_output_collapse( + ctx0, streams, carried_pre, n_embd, + hparams.dsv4_hc_mult, + inp_out_ids ? n_outputs : n_tokens); + cb(cur, "dsv41_output_collapse", -1); + cur = build_norm( + cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} + +std::unique_ptr llama_model_deepseek41::build_arch_graph( + const llm_graph_params & params) const { + return std::make_unique(*this, params); } diff --git a/src/models/models.h b/src/models/models.h index ef05debe660c..4acf4d90e1ad 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -9,6 +9,7 @@ #include class llama_memory_hybrid_idx_context; +class llama_dsv41_engram_runtime; struct llama_dsv41_expert_runtime; // ref: https://github.com/ggml-org/llama.cpp/pull/28068 @@ -1317,10 +1318,16 @@ struct llama_model_deepseek4 : public llama_model_base { struct llama_model_deepseek41 : public llama_model_deepseek4 { llama_model_deepseek41(const struct llama_model_params & params) : llama_model_deepseek4(params) {} + struct graph : public llama_model_deepseek4::graph { + graph(const llama_model & model, const llm_graph_params & params); + }; + struct engram_model; std::shared_ptr engram; std::shared_ptr experts; + std::unique_ptr create_memory_engram_runtime(size_t max_tokens) const; + void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; bool requires_synchronous_graph() const override; @@ -1330,7 +1337,7 @@ struct llama_model_deepseek41 : public llama_model_deepseek4 { void acquire_runtime_context() const override; void release_runtime_context() const override; - [[noreturn]] std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2c4242b749c9..98e9cfce4b7a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -200,6 +200,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-deepseek41-engram.cpp) llama_build_and_test(test-deepseek41-expert.cpp) llama_build_and_test(test-expert-store.cpp) + llama_build_and_test(test-deepseek41-memory.cpp) llama_build_and_test(test-deepseek41-runtime.cpp) llama_build_and_test(test-engram.cpp) llama_build(test-llama-archs.cpp) diff --git a/tests/test-deepseek41-memory.cpp b/tests/test-deepseek41-memory.cpp new file mode 100644 index 000000000000..e4f6302f4ba1 --- /dev/null +++ b/tests/test-deepseek41-memory.cpp @@ -0,0 +1,642 @@ +#include "../src/llama-batch.h" +#include "../src/llama-io.h" +#include "../src/llama-memory-dsv41.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#endif + +static void check(bool condition, const char * message) { + if (!condition) { + std::fprintf(stderr, "%s\n", message); + std::exit(1); + } +} + +static void expect_invalid(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::invalid_argument &) { + return; + } + check(false, message); +} + +static void expect_runtime(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::runtime_error &) { + return; + } + check(false, message); +} + +static llama_dsv41_memory_config small_config( + uint32_t n_ctx = 256, + uint32_t n_seq = 3, + uint32_t n_ubatch = 128) { + llama_dsv41_memory_config config; + config.n_ctx = n_ctx; + config.n_seq = n_seq; + config.n_ubatch = n_ubatch; + config.kv_width = 8; + config.index_width = 4; + config.candidate_topk_blocks = 4; + config.candidate_block_size = 8; + config.ratios.resize(LLAMA_DSV41_N_LAYER); + for (uint32_t il = 0; il < LLAMA_DSV41_N_LAYER; ++il) { + config.ratios[il] = llama_dsv41_compress_ratio(il); + } + config.buft_for_layer = [](int32_t) { return ggml_backend_cpu_buffer_type(); }; + return config; +} + +static llama_ubatch make_ubatch(llama_pos start, uint32_t count, llama_seq_id seq_id) { + llama_batch_allocr allocator(1); + llama_ubatch ubatch = allocator.ubatch_reserve(count, 1); + ubatch.data->seq_id_data.resize(count); + ubatch.data->seq_id_unq = { seq_id }; + ubatch.seq_id_unq = ubatch.data->seq_id_unq.data(); + ubatch.seq_idx[seq_id] = 0; + for (uint32_t i = 0; i < count; ++i) { + ubatch.token[i] = 10 + i; + ubatch.pos[i] = start + i; + ubatch.n_seq_id[i] = 1; + ubatch.data->seq_id_data[i] = seq_id; + ubatch.seq_id[i] = &ubatch.data->seq_id_data[i]; + ubatch.output[i] = 1; + } + return ubatch; +} + +static llama_ubatch make_coupled_ubatch( + llama_pos start, + uint32_t count, + llama_seq_id first, + llama_seq_id second) { + llama_ubatch ubatch = make_ubatch(start, count, first); + ubatch.data->seq_id_data.resize((size_t) count*2); + ubatch.data->seq_id_unq = { first, second }; + ubatch.seq_id_unq = ubatch.data->seq_id_unq.data(); + ubatch.n_seqs_unq = 2; + ubatch.seq_idx[first] = 0; + ubatch.seq_idx[second] = 1; + for (uint32_t i = 0; i < count; ++i) { + ubatch.n_seq_id[i] = 2; + ubatch.data->seq_id_data[2*i] = first; + ubatch.data->seq_id_data[2*i + 1] = second; + ubatch.seq_id[i] = ubatch.data->seq_id_data.data() + 2*i; + } + return ubatch; +} + +static size_t state_sequence_tensor_bytes(const llama_memory_dsv41 & memory) { + size_t result = 0; + const auto add = [&](const ggml_tensor * tensor) { + result += sizeof(uint64_t); + result += tensor->ne[2] == (int64_t) memory.config().n_seq ? + tensor->nb[2] : tensor->nb[1]; + }; + for (uint32_t il = 0; il < memory.config().n_layer; ++il) { + add(memory.raw_k(il)); + } + for (uint32_t source : memory.config().kv_sources) { + add(memory.compressed_kv(source)); + add(memory.index_keys(source)); + if (memory.config().ratios[source] == 2) { + add(memory.compressor_carry_kv(source)); + add(memory.compressor_carry_score(source)); + } + } + add(memory.committed_candidate_ids()); + add(memory.position_state()); + return result; +} + +class vector_writer : public llama_io_write_i { +public: + void write(const void * src, size_t size) override { + const uint8_t * bytes = static_cast(src); + data.insert(data.end(), bytes, bytes + size); + } + + void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { + const size_t old_size = data.size(); + data.resize(old_size + size); + ggml_backend_tensor_get(tensor, data.data() + old_size, offset, size); + } + + size_t n_bytes() override { + return data.size(); + } + + std::vector data; +}; + +class vector_reader : public llama_io_read_i { +public: + explicit vector_reader(const std::vector & data) : data(data) {} + + void read(void * dst, size_t size) override { + if (offset > data.size() || size > data.size() - offset) { + throw std::runtime_error("test state buffer is truncated"); + } + std::memcpy(dst, data.data() + offset, size); + offset += size; + } + + void read_tensor(ggml_tensor * tensor, size_t tensor_offset, size_t size) override { + (void) tensor; + (void) tensor_offset; + (void) size; + throw std::runtime_error("DeepSeek V4.1 state restore must read tensor bytes before publication"); + } + + size_t n_bytes() override { + return offset; + } + +private: + const std::vector & data; + size_t offset = 0; +}; + +static void test_transaction_commit_rollback() { + llama_memory_dsv41 memory(small_config()); + llama_ubatch ubatch = make_ubatch(0, 1, 0); + + auto full_base = memory.init_full(); + auto * full = dynamic_cast(full_base.get()); + check(full != nullptr && full->get_ubatch().n_tokens == memory.config().n_ubatch, + "full memory context did not expose a bounded reserve plan"); + check(full->plan().sources.size() == memory.config().kv_sources.size(), + "full memory context source plan mismatch"); + llama_ubatch reserve_decode = make_ubatch(0, 1, 0); + check(full->topology(reserve_decode, 1).n_tokens == 1 && + full->plan().positions.size() == 1, + "full memory context did not resize its synthetic decode plan"); + llama_ubatch reserve_prefill = make_ubatch(0, 7, 0); + check(full->graph_plan(reserve_prefill).positions.size() == 7 && + full->topology(reserve_prefill, 3).n_outputs == 3, + "full memory context did not resize its synthetic prefill plan"); + + const size_t row_bytes = memory.raw_k(0)->nb[1]; + std::vector original(row_bytes, 0x31); + std::vector changed(row_bytes, 0x72); + std::vector actual(row_bytes); + ggml_backend_tensor_set(memory.raw_k(0), original.data(), 0, row_bytes); + + llama_memory_dsv41_context rollback_context(&memory, std::vector { ubatch }); + check(rollback_context.apply(), "transaction prepare failed"); + check(memory.seq_pos_max(0) == -1, "prepare published the position early"); + ggml_backend_tensor_set(memory.raw_k(0), changed.data(), 0, row_bytes); + rollback_context.stage_candidate_ids(0, { 7 }); + rollback_context.rollback(); + ggml_backend_tensor_get(memory.raw_k(0), actual.data(), 0, row_bytes); + check(actual == original, "rollback did not restore a graph-written raw row"); + check(memory.seq_pos_max(0) == -1, "rollback changed the committed position"); + check(memory.sequence_candidate_ids(0).empty(), "rollback published candidate IDs"); + + llama_memory_dsv41_context commit_context(&memory, std::vector { ubatch }); + check(commit_context.apply(), "commit prepare failed"); + commit_context.stage_candidate_ids(0, { 3 }); + const auto topology = commit_context.topology(ubatch, 1); + check(topology.n_tokens == 1 && topology.n_seqs == 1 && topology.n_outputs == 1, + "graph topology batch identity mismatch"); + check(topology.start_positions == std::vector({ 0 }), + "graph topology start position mismatch"); + auto reusable_topology = topology; + reusable_topology.seq_ids = { 2 }; + reusable_topology.start_positions = { 17 }; + reusable_topology.transaction_generation++; + check(topology.same_topology(reusable_topology), + "graph reuse rejected refreshed sequence and position inputs"); + reusable_topology.candidate_width++; + check(!topology.same_topology(reusable_topology), + "graph reuse accepted a changed candidate workspace shape"); + commit_context.commit(); + check(memory.seq_pos_max(0) == 0, "commit did not publish the position"); + check(memory.sequence_candidate_ids(0) == std::vector({ 3 }), + "commit did not publish candidate IDs"); + + llama_memory_dsv41_context candidate_rollback( + &memory, std::vector { make_ubatch(1, 1, 0) }); + check(candidate_rollback.apply(), "candidate rollback prepare failed"); + candidate_rollback.stage_candidate_ids(0, { 9 }); + candidate_rollback.rollback(); + check(memory.sequence_candidate_ids(0) == std::vector({ 3 }), + "candidate rollback changed committed IDs"); + + llama_memory_dsv41_context over_capacity( + &memory, + std::vector { + make_ubatch(memory.config().n_ctx, 1, 1), + }); + check(!over_capacity.apply(), "over-capacity transaction was accepted"); + check(memory.seq_pos_max(1) == -1, "over-capacity prepare mutated sequence state"); + + llama_memory_dsv41_context coupled( + &memory, + std::vector { + make_coupled_ubatch(0, 1, 1, 2), + }); + check(!coupled.apply(), "coupled sequence evaluation was accepted"); + check(memory.seq_pos_max(1) == -1 && memory.seq_pos_max(2) == -1, + "rejected coupled sequence evaluation mutated state"); +} + +static void test_window_compression_and_sequences() { + llama_memory_dsv41 memory(small_config()); + llama_ubatch prefill = make_ubatch(0, 128, 0); + llama_memory_dsv41_context prefill_context( + &memory, std::vector { prefill }); + check(prefill_context.apply(), "128-token prefill prepare failed"); + const auto & prefill_plan = prefill_context.plan(); + check(prefill_plan.raw.n_visible.back() == 128, "raw window width at 127 mismatch"); + check(prefill_plan.raw.read_idxs[127*128] == 0 && + prefill_plan.raw.read_idxs[127*128 + 127] == 127, + "raw window order at 127 mismatch"); + check(prefill_plan.sources[0].ratio == 2 && + prefill_plan.sources[0].compression.n_visible.back() == 64, + "ratio-2 visibility at 127 mismatch"); + prefill_context.stage_candidate_ids(127, { 4, 3, 2, 1 }); + prefill_context.commit(); + check(memory.seq_pos_max(0) == 127, "prefill position mismatch"); + + const size_t raw_row_bytes = memory.raw_k(0)->nb[1]; + const size_t carry_plane_bytes = memory.compressor_carry_kv(2)->nb[2]; + std::vector raw_before(raw_row_bytes, 0x21); + std::vector raw_after(raw_row_bytes, 0x72); + std::vector carry_before(carry_plane_bytes, 0x32); + std::vector carry_after(carry_plane_bytes, 0x83); + std::vector actual; + ggml_backend_tensor_set(memory.raw_k(0), raw_before.data(), 0, raw_row_bytes); + ggml_backend_tensor_set(memory.compressor_carry_kv(2), carry_before.data(), 0, carry_plane_bytes); + + llama_ubatch decode = make_ubatch(128, 1, 0); + llama_memory_dsv41_context decode_context( + &memory, std::vector { decode }); + check(decode_context.apply(), "boundary decode prepare failed"); + const auto & decode_plan = decode_context.plan(); + check(decode_plan.raw.write_idxs == std::vector({ 0 }), + "raw ring write did not wrap at 128"); + check(decode_plan.raw.read_idxs.front() == 1 && + decode_plan.raw.read_idxs[127] == 0, + "raw ring order at 128 mismatch"); + check(decode_plan.sources[0].compression.n_visible == std::vector({ 64 }), + "ratio-2 carry visibility at 128 mismatch"); + check(decode_plan.sources[0].compression.state_persist_dst_idxs == std::vector({ 0 }), + "ratio-2 carry publication row mismatch"); + ggml_backend_tensor_set(memory.raw_k(0), raw_after.data(), 0, raw_row_bytes); + ggml_backend_tensor_set(memory.compressor_carry_kv(2), carry_after.data(), 0, carry_plane_bytes); + decode_context.stage_candidate_ids(0, { 8, 7, 6, 5 }); + decode_context.commit(); + check(memory.seq_pos_max(0) == 128, "decode position mismatch"); + check(memory.seq_rm(0, 128, -1), "immediate rollback failed"); + check(memory.seq_pos_max(0) == 127, "immediate rollback position mismatch"); + check(memory.sequence_candidate_ids(0) == std::vector({ 4, 3, 2, 1 }), + "immediate rollback did not restore candidate state"); + actual.resize(raw_row_bytes); + ggml_backend_tensor_get(memory.raw_k(0), actual.data(), 0, raw_row_bytes); + check(actual == raw_before, "immediate rollback did not restore the raw ring row"); + actual.resize(carry_plane_bytes); + ggml_backend_tensor_get(memory.compressor_carry_kv(2), actual.data(), 0, carry_plane_bytes); + check(actual == carry_before, "immediate rollback did not restore compressor carry"); + + llama_memory_dsv41_context recommit_context( + &memory, std::vector { decode }); + check(recommit_context.apply(), "decode after rollback prepare failed"); + recommit_context.stage_candidate_ids(0, { 8, 7, 6, 5 }); + recommit_context.commit(); + memory.seq_cp(0, 1, -1, -1); + check(memory.seq_pos_max(1) == 128, "full sequence copy lost the position"); + check(memory.sequence_candidate_ids(1) == std::vector({ 8, 7, 6, 5 }), + "full sequence copy lost candidate state"); + memory.seq_keep(1); + check(memory.seq_pos_max(0) == -1 && memory.seq_pos_max(1) == 128, + "sequence keep retained another sequence"); + check(!memory.seq_rm(1, 128, -1), "copied sequence rollback was accepted without a retained snapshot"); + check(memory.seq_pos_max(1) == 128, "rejected copied-sequence rollback mutated the position"); + check(!memory.seq_rm(1, 64, -1), "non-immediate suffix rollback was accepted"); + check(memory.seq_pos_max(1) == 128, "rejected suffix rollback mutated the position"); + check(!memory.seq_rm(1, 64, 96), "interior range removal was accepted"); + check(!memory.seq_rm(-1, 0, -1), "partial wildcard removal was accepted"); + check(memory.seq_pos_max(1) == 128, "rejected wildcard removal mutated the position"); + expect_invalid( + [&] { memory.seq_cp(1, 2, 64, -1); }, + "partial sequence copy was accepted"); + check(memory.seq_rm(-1, -1, -1), "wildcard full-memory removal failed"); + check(memory.seq_pos_max(0) == -1 && + memory.seq_pos_max(1) == -1 && + memory.seq_pos_max(2) == -1, + "wildcard full-memory removal retained sequence state"); + llama_memory_dsv41_context negative_wildcard_context( + &memory, std::vector { make_ubatch(0, 1, 2) }); + check(negative_wildcard_context.apply(), "negative wildcard setup failed"); + negative_wildcard_context.commit(); + check(memory.seq_rm(-2, -1, -1), "negative wildcard full-memory removal failed"); + check(memory.seq_pos_max(2) == -1, "negative wildcard removal retained sequence state"); +} + +static void test_long_prefill_raw_publication() { + llama_memory_dsv41 memory(small_config(256, 1, 256)); + llama_memory_dsv41_context context( + &memory, + std::vector { make_ubatch(0, 256, 0) }); + check(context.apply(), "long prefill prepare failed"); + const auto & raw = context.plan().raw; + check(raw.persist_src_idxs.size() == LLAMA_DSV41_N_SWA, + "long prefill retained duplicate raw ring writes"); + check(raw.persist_src_idxs.front() == 128 && + raw.persist_src_idxs.back() == 255, + "long prefill did not retain the final raw window"); + std::vector unique = raw.write_idxs; + std::sort(unique.begin(), unique.end()); + check(std::adjacent_find(unique.begin(), unique.end()) == unique.end(), + "long prefill raw ring write indexes are not unique"); + context.rollback(); +} + +static void test_state_save_load() { + llama_memory_dsv41 memory(small_config()); + llama_ubatch ubatch = make_ubatch(0, 3, 0); + llama_memory_dsv41_context context(&memory, std::vector { ubatch }); + check(context.apply(), "state test prepare failed"); + context.stage_candidate_ids(2, { 11 }); + context.commit(); + + const size_t row_bytes = memory.raw_k(0)->nb[1]; + std::vector expected(row_bytes, 0x5a); + std::vector actual(row_bytes); + ggml_backend_tensor_set(memory.raw_k(0), expected.data(), 2*row_bytes, row_bytes); + + vector_writer writer; + memory.state_write(writer, 0); + expect_runtime( + [&] { + vector_reader nonempty_reader(writer.data); + memory.state_read(nonempty_reader, 0); + }, + "state restore overwrote a non-empty destination"); + memory.clear(true); + vector_reader reader(writer.data); + memory.state_read(reader, 2); + check(memory.seq_pos_max(2) == 2, "state restore lost the position"); + check(memory.sequence_candidate_ids(2) == std::vector({ 11 }), + "state restore lost candidate IDs"); + ggml_backend_tensor_get( + memory.raw_k(0), + actual.data(), + ((size_t) 2*memory.config().raw_window + 2)*row_bytes, + row_bytes); + check(actual == expected, "state restore lost raw cache data"); + + std::vector truncated = writer.data; + truncated.pop_back(); + expect_runtime( + [&] { + vector_reader truncated_reader(truncated); + memory.state_read(truncated_reader, 1); + }, + "truncated state restore was accepted"); + check(memory.seq_pos_max(1) == -1 && memory.seq_pos_max(2) == 2, + "failed state restore changed committed sequence state"); + + vector_writer full_writer; + memory.state_write(full_writer, -1); + memory.clear(true); + std::vector incomplete = full_writer.data; + const uint32_t incomplete_count = memory.config().n_seq - 1; + std::memcpy( + incomplete.data() + sizeof(uint64_t) + sizeof(uint32_t), + &incomplete_count, + sizeof(incomplete_count)); + expect_runtime( + [&] { + vector_reader incomplete_reader(incomplete); + memory.state_read(incomplete_reader, -1); + }, + "incomplete full state restore was accepted"); + for (uint32_t seq = 0; seq < memory.config().n_seq; ++seq) { + check(memory.seq_pos_max(seq) == -1, "incomplete full restore changed memory"); + } + + std::vector duplicate = full_writer.data; + const size_t header_size = + sizeof(uint64_t) + 3*sizeof(uint32_t); + const size_t first_record_size = + sizeof(llama_seq_id) + sizeof(llama_pos) + sizeof(uint32_t) + + sizeof(uint8_t) + state_sequence_tensor_bytes(memory); + const llama_seq_id duplicate_id = 0; + std::memcpy( + duplicate.data() + header_size + first_record_size, + &duplicate_id, + sizeof(duplicate_id)); + expect_runtime( + [&] { + vector_reader duplicate_reader(duplicate); + memory.state_read(duplicate_reader, -1); + }, + "duplicate full state sequence was accepted"); + for (uint32_t seq = 0; seq < memory.config().n_seq; ++seq) { + check(memory.seq_pos_max(seq) == -1, "duplicate full restore changed memory"); + } + + vector_reader full_reader(full_writer.data); + memory.state_read(full_reader, -1); + check(memory.seq_pos_max(2) == 2, "full state restore lost sequence state"); +} + +static void test_accounting() { + const auto config = small_config(); + llama_memory_dsv41 memory(small_config()); + memory.set_graph_workspace_size(1234); + const auto bytes = memory.accounting(); + const uint64_t raw_expected = + (uint64_t) config.n_layer*config.kv_width*config.raw_window*config.n_seq*sizeof(uint16_t); + const uint64_t compressed_rows = + (uint64_t) 3*((config.n_ctx + 1)/2) + config.n_ctx; + const uint64_t compressed_expected = + compressed_rows*config.kv_width*config.n_seq*sizeof(uint16_t); + const uint64_t index_expected = + compressed_rows*config.index_width*config.n_seq*sizeof(uint16_t); + const uint64_t carry_expected = + (uint64_t) 3*2*config.kv_width*config.n_seq*sizeof(uint16_t)*2; + const uint64_t candidate_scores_expected = + (uint64_t) (config.n_ctx/config.candidate_block_size)*config.n_ubatch*sizeof(float); + const uint64_t candidate_ids_expected = + (uint64_t) config.candidate_topk_blocks*(config.n_ubatch + config.n_seq)*sizeof(int32_t); + const uint64_t position_rows = + 1 + config.raw_window + 3*((config.n_ctx + 1)/2 + 2) + config.n_ctx; + const uint64_t position_expected = position_rows*config.n_seq*sizeof(int32_t); + + check(bytes.raw_kv == raw_expected, "raw cache accounting mismatch"); + check(bytes.compressed_kv == compressed_expected, "compressed cache accounting mismatch"); + check(bytes.index_keys == index_expected, "index key accounting mismatch"); + check(bytes.compressor_carry == carry_expected, "compressor carry accounting mismatch"); + check(bytes.candidate_scores == candidate_scores_expected, "candidate score accounting mismatch"); + check(bytes.candidate_ids == candidate_ids_expected, "candidate ID accounting mismatch"); + check(bytes.position_state == position_expected, "position accounting mismatch"); + check(bytes.graph_workspace == 1234, "graph workspace accounting mismatch"); + const auto breakdown = memory.memory_breakdown(); + check(breakdown.size() == 1 && + breakdown.begin()->second >= bytes.total() - bytes.graph_workspace, + "memory breakdown does not include all allocated state"); + + auto no_alloc_config = small_config(); + no_alloc_config.no_alloc = true; + llama_memory_dsv41 no_alloc_memory(std::move(no_alloc_config)); + check(no_alloc_memory.raw_k(0)->buffer == nullptr, + "no-allocation memory probe allocated cache storage"); + const auto no_alloc_bytes = no_alloc_memory.accounting(); + check(no_alloc_bytes.total() == bytes.total() - bytes.graph_workspace, + "no-allocation memory accounting differs from allocated state"); +} + +#if !defined(_WIN32) +static llama_engram_layout make_engram_layout() { + llama_engram_layout layout; + layout.encoding = LLAMA_DSV41_ENGRAM_ENCODING; + layout.layer_ids = { 1, 14 }; + layout.token_map.resize(32); + for (size_t i = 0; i < layout.token_map.size(); ++i) { + layout.token_map[i] = i; + } + layout.compressed_vocab_size = 32; + layout.pad_id = 2; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + for (size_t i = 0; i < LLAMA_ENGRAM_NGRAM; ++i) { + layout.multipliers[layer][i] = 101 + 8*layer + 2*i; + } + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + layout.primes[layer][col] = 2; + layout.rows[layer] += 2; + } + } + return layout; +} + +struct engram_test_file { + std::string path = "test-dsv41-memory-engram.bin"; + int fd = -1; + std::array extents; + + explicit engram_test_file(const llama_engram_layout & layout) { + fd = open(path.c_str(), O_CREAT | O_TRUNC | O_RDWR, 0600); + check(fd >= 0, "failed to create Engram memory test file"); + uint64_t offset = 4096; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + extents[layer] = { + path, + offset, + layout.rows[layer], + LLAMA_ENGRAM_ROW_BYTES, + layout.rows[layer], + GGML_TYPE_I8, + }; + for (uint32_t row = 0; row < layout.rows[layer]; ++row) { + uint8_t data[LLAMA_ENGRAM_ROW_BYTES]; + std::fill(data, data + LLAMA_ENGRAM_ROW_BYTES, (uint8_t) (8 + row)); + check( + pwrite(fd, data, sizeof(data), offset + (uint64_t) row*sizeof(data)) == + (ssize_t) sizeof(data), + "failed to write Engram memory test row"); + } + offset += (uint64_t) layout.rows[layer]*LLAMA_ENGRAM_ROW_BYTES + 4096; + } + } + + ~engram_test_file() { + if (fd >= 0) { + close(fd); + } + unlink(path.c_str()); + } +}; + +static void test_engram_transaction() { + const llama_engram_layout layout = make_engram_layout(); + engram_test_file file(layout); + auto config = small_config(8, 1, 4); + auto runtime = std::make_unique( + layout, file.extents, config.n_ubatch); + llama_dsv41_engram_runtime * runtime_ptr = runtime.get(); + config.engram = std::move(runtime); + llama_memory_dsv41 memory(std::move(config)); + + const size_t row_bytes = memory.raw_k(0)->nb[1]; + std::vector changed(row_bytes, 0x6b); + std::vector actual(row_bytes); + llama_memory_dsv41_context stale( + &memory, std::vector { make_ubatch(0, 1, 0) }); + check(stale.apply(), "stale Engram transaction prepare failed"); + ggml_backend_tensor_set(memory.raw_k(0), changed.data(), 0, row_bytes); + stale.stage_candidate_ids(0, { 9 }); + runtime_ptr->seq_remove(0); + expect_runtime([&] { stale.commit(); }, "stale Engram transaction commit was accepted"); + stale.rollback(); + ggml_backend_tensor_get(memory.raw_k(0), actual.data(), 0, row_bytes); + check(std::all_of(actual.begin(), actual.end(), [](uint8_t value) { return value == 0; }), + "failed commit did not restore graph-written state"); + check(memory.seq_pos_max(0) == -1 && memory.sequence_candidate_ids(0).empty(), + "failed commit published sequence state"); + + llama_memory_dsv41_context first( + &memory, std::vector { make_ubatch(0, 2, 0) }); + check(first.apply() && first.engram_transaction() != nullptr, + "Engram transaction was not exposed"); + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + const auto packed = first.engram_row_ids(layer); + const uint32_t * strided = first.engram_transaction()->row_ids(layer); + check(packed.size() == 2*LLAMA_ENGRAM_COLS, "packed Engram row ID shape mismatch"); + for (uint32_t token = 0; token < 2; ++token) { + check(std::memcmp( + packed.data() + token*LLAMA_ENGRAM_COLS, + strided + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(int32_t)) == 0, + "packed Engram row IDs changed the token stride"); + } + } + first.stage_candidate_ids(1, { 1 }); + first.commit(); + + llama_memory_dsv41_context rolled_back( + &memory, std::vector { make_ubatch(2, 1, 0) }); + check(rolled_back.apply(), "Engram rollback prepare failed"); + rolled_back.rollback(); + + llama_memory_dsv41_context second( + &memory, std::vector { make_ubatch(2, 1, 0) }); + check(second.apply(), "Engram state advanced during rollback"); + second.stage_candidate_ids(0, { 2 }); + second.commit(); + check(memory.seq_pos_max(0) == 2, "Engram commit position mismatch"); + check(memory.seq_rm(0, 2, -1), "Engram suffix rollback at transaction boundary failed"); + check(memory.seq_pos_max(0) == 1, "Engram suffix rollback did not restore position"); +} +#endif + +int main() { + test_transaction_commit_rollback(); + test_window_compression_and_sequences(); + test_long_prefill_raw_publication(); + test_state_save_load(); + test_accounting(); +#if !defined(_WIN32) + test_engram_transaction(); +#endif + return 0; +} diff --git a/tests/test-deepseek41-runtime.cpp b/tests/test-deepseek41-runtime.cpp index 2d5945b407fb..a03743a8551d 100644 --- a/tests/test-deepseek41-runtime.cpp +++ b/tests/test-deepseek41-runtime.cpp @@ -1,4 +1,5 @@ #include "../src/llama-dsv41.h" +#include "../src/llama-arch.h" #include "ggml.h" @@ -8,11 +9,14 @@ #include #include #include +#include #include #include #include #include +std::string llama_dsv41_graph_trace_name(const char * trace, uint32_t layer); + static void check(bool condition, const std::string & message) { if (!condition) { std::fprintf(stderr, "%s\n", message.c_str()); @@ -110,8 +114,9 @@ static void test_hparams() { config.engram_primes_size = 24; expect_throw([&]() { llama_dsv41_validate_config(config); }, "truncated Engram prime table was accepted"); - const std::string dependency_error = llama_dsv41_runtime_dependency_error(); - check(dependency_error.find("routed-expert streaming") != std::string::npos, "dependency error omits expert streaming"); + check(llm_arch_is_hybrid(LLM_ARCH_DEEPSEEK41), "DeepSeek V4.1 must use hybrid context handling"); + check(!llm_arch_supports_rs_rollback(LLM_ARCH_DEEPSEEK41), + "DeepSeek V4.1 must not advertise partial rollback support"); } static void test_source_maps() { @@ -245,6 +250,102 @@ static void test_output_collapse() { check(std::abs(result[1] - 6.0f) < 1.0e-6f, "output collapse second value mismatch"); } +static void test_graph_contract() { + enum stage { + STAGE_ENGRAM, + STAGE_ATTN_HC, + STAGE_CARRIED_PRE, + STAGE_ATTN, + STAGE_ATTN_POST, + STAGE_FFN_HC, + STAGE_ATTN_PRE, + STAGE_ROUTED_EXPERTS, + STAGE_SHARED_EXPERT, + STAGE_FFN_POST, + STAGE_CARRY_FFN_PRE, + }; + + const std::vector common = { + STAGE_ATTN_HC, + STAGE_CARRIED_PRE, + STAGE_ATTN, + STAGE_ATTN_POST, + STAGE_FFN_HC, + STAGE_ATTN_PRE, + STAGE_ROUTED_EXPERTS, + STAGE_SHARED_EXPERT, + STAGE_FFN_POST, + STAGE_CARRY_FFN_PRE, + }; + uint32_t ratio_count[3] = {}; + uint32_t kv_sources = 0; + uint32_t index_sources = 0; + uint32_t candidate_sources = 0; + std::vector candidate_trace_layers; + for (uint32_t il = 0; il < LLAMA_DSV41_N_LAYER; ++il) { + std::vector stages = common; + if (il == 1 || il == 14) { + stages.insert(stages.begin(), STAGE_ENGRAM); + check(stages[0] == STAGE_ENGRAM && stages[1] == STAGE_ATTN_HC, + "Engram must precede attention HC"); + } + check(stages[stages.size() - 1] == STAGE_CARRY_FFN_PRE, + "FFN pre must be carried to the next layer"); + check(stages[1 + (il == 1 || il == 14)] == STAGE_CARRIED_PRE, + "attention must collapse with carried pre"); + check(stages[5 + (il == 1 || il == 14)] == STAGE_ATTN_PRE, + "FFN must collapse with current attention pre"); + check(stages[7 + (il == 1 || il == 14)] == STAGE_SHARED_EXPERT, + "shared expert must be added after routed experts"); + + const uint32_t ratio = llama_dsv41_compress_ratio(il); + ratio_count[ratio]++; + kv_sources += llama_dsv41_kv_source_layer(il) == (int32_t) il; + index_sources += llama_dsv41_index_source_layer(il) == (int32_t) il; + candidate_sources += il == LLAMA_DSV41_CANDIDATE_SOURCE_LAYER; + + check( + llama_dsv41_graph_trace_name("expert.ids", il) == + "dsv41.trace.expert.ids.l" + std::to_string(il), + "expert ID trace name is unstable"); + check( + llama_dsv41_graph_trace_name("expert.weights", il) == + "dsv41.trace.expert.weights.l" + std::to_string(il), + "expert weight trace name is unstable"); + check( + llama_dsv41_graph_trace_name("attn.source", il) == + "dsv41.trace.attn.source.l" + std::to_string(il), + "attention source trace name is unstable"); + if (il > LLAMA_DSV41_CANDIDATE_SOURCE_LAYER && + llama_dsv41_index_source_layer(il) == (int32_t) il) { + candidate_trace_layers.push_back(il); + check( + llama_dsv41_graph_trace_name("attn.candidates", il) == + "dsv41.trace.attn.candidates.l" + std::to_string(il), + "attention candidate trace name is unstable"); + } + } + check(ratio_count[0] == 2 && ratio_count[1] == 20 && ratio_count[2] == 18, + "ratio 0/1/2 layer counts mismatch"); + check(kv_sources == 4, "KV source ownership count mismatch"); + check(index_sources == 8, "index source ownership count mismatch"); + check(candidate_sources == 1, "candidate source ownership count mismatch"); + check(candidate_trace_layers == std::vector({ 24, 28, 32, 36 }), + "attention candidate trace layer coverage mismatch"); + check( + llama_dsv41_graph_trace_name("attn.candidate_blocks", 20) == + "dsv41.trace.attn.candidate_blocks.l20", + "candidate block trace name is unstable"); + check( + llama_dsv41_graph_trace_name("engram.row_ids", 1) == + "dsv41.trace.engram.row_ids.l1" && + llama_dsv41_graph_trace_name("engram.row_ids", 14) == + "dsv41.trace.engram.row_ids.l14", + "Engram row trace names are unstable"); + check(llama_dsv41_build_layer_plan(39, { 39 }, 1024).collapses_output, + "final layer must preserve streams for carried-pre output collapse"); +} + static void test_graph_construction() { ggml_init_params params = { /*.mem_size =*/ 4*1024*1024, @@ -281,6 +382,108 @@ static void test_graph_construction() { const float * ordered = static_cast(ordered_probs->data); check(ordered[0] < ordered[1] && ordered[1] < ordered[2], "shared-softmax segment order mismatch"); + ggml_tensor * candidate_scores = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 24, 2); + ggml_tensor * candidate_bias = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 3, 2); + float * score_data = static_cast(candidate_scores->data); + float * bias_data = static_cast(candidate_bias->data); + std::fill_n(score_data, 48, -100.0f); + std::fill_n(bias_data, 6, 0.0f); + score_data[0] = 100.0f; + score_data[24] = 100.0f; + std::fill(score_data + 9, score_data + 24, -std::numeric_limits::infinity()); + std::fill(score_data + 24 + 17, score_data + 48, -std::numeric_limits::infinity()); + bias_data[1] = std::numeric_limits::infinity(); + bias_data[2] = -std::numeric_limits::infinity(); + bias_data[3 + 2] = std::numeric_limits::infinity(); + ggml_tensor * block_scores = ggml_pool_1d( + ctx, candidate_scores, GGML_OP_POOL_MAX, 8, 8, 0); + block_scores = ggml_add(ctx, block_scores, candidate_bias); + ggml_tensor * candidate_ids = ggml_cont( + ctx, ggml_argsort_top_k(ctx, block_scores, 2)); + ggml_set_name( + candidate_ids, + llama_dsv41_graph_trace_name( + "attn.candidate_blocks", 20).c_str()); + ggml_cgraph * candidate_gf = ggml_new_graph(ctx); + ggml_build_forward_expand(candidate_gf, candidate_ids); + check( + ggml_graph_compute_with_ctx(ctx, candidate_gf, 1) == + GGML_STATUS_SUCCESS, + "candidate block graph execution failed"); + const int32_t * ids = static_cast(candidate_ids->data); + check( + ids[0] == 1 && ids[1] == 0 && + ids[2] == 2 && ids[3] == 0, + "candidate blocks are not pinned and sorted"); + check( + ggml_is_contiguous(candidate_ids) && + std::string(ggml_get_name(candidate_ids)) == + "dsv41.trace.attn.candidate_blocks.l20", + "candidate block trace is not stable and contiguous"); + + ggml_tensor * prior_ring = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 1, 2); + ggml_tensor * current_k = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 1, 1); + ggml_tensor * read_idx = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, 1); + ggml_tensor * write_idx = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, 1); + static_cast(prior_ring->data)[0] = 10.0f; + static_cast(prior_ring->data)[1] = 20.0f; + static_cast(current_k->data)[0] = 30.0f; + static_cast(read_idx->data)[0] = 0; + static_cast(write_idx->data)[0] = 0; + ggml_tensor * prior_read = ggml_get_rows( + ctx, prior_ring, read_idx); + ggml_tensor * attention = ggml_add( + ctx, prior_read, current_k); + ggml_tensor * completion = ggml_argsort_top_k( + ctx, ggml_view_1d(ctx, attention, 1, 0), 1); + completion = ggml_scale( + ctx, ggml_cast(ctx, completion, GGML_TYPE_F32), 0.0f); + ggml_tensor * delayed_k = ggml_add( + ctx, current_k, completion); + ggml_tensor * ring_update = ggml_set_rows( + ctx, prior_ring, delayed_k, write_idx); + ggml_cgraph * ring_gf = ggml_new_graph(ctx); + ggml_build_forward_expand(ring_gf, ring_update); + check( + ggml_graph_compute_with_ctx(ctx, ring_gf, 1) == + GGML_STATUS_SUCCESS, + "ordered raw-ring update graph execution failed"); + check( + static_cast(attention->data)[0] == 40.0f && + static_cast(prior_ring->data)[0] == 30.0f, + "raw-ring write did not wait for the prior-ring read"); + + ggml_tensor * selected_ids = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, 3, 2); + const int32_t selected_values[] = { 5, 1, 3, 4, 0, 2 }; + std::memcpy(selected_ids->data, selected_values, sizeof(selected_values)); + ggml_tensor * selected_order = ggml_argsort( + ctx, ggml_cast(ctx, selected_ids, GGML_TYPE_F32), + GGML_SORT_ORDER_ASC); + ggml_tensor * selected_sorted = ggml_get_rows( + ctx, ggml_reshape_3d(ctx, selected_ids, 1, 3, 2), + selected_order); + selected_sorted = ggml_cont( + ctx, ggml_reshape_2d(ctx, selected_sorted, 3, 2)); + ggml_cgraph * selected_gf = ggml_new_graph(ctx); + ggml_build_forward_expand(selected_gf, selected_sorted); + check( + ggml_graph_compute_with_ctx(ctx, selected_gf, 1) == + GGML_STATUS_SUCCESS, + "selected ID ordering graph execution failed"); + const int32_t selected_expected[] = { 1, 3, 5, 0, 2, 4 }; + check( + std::memcmp( + selected_sorted->data, selected_expected, + sizeof(selected_expected)) == 0, + "selected IDs are not accumulated in original ID order"); + ggml_tensor * residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 32, 4, 2); ggml_tensor * pre = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 2); ggml_tensor * collapsed = llama_dsv41_build_output_collapse(ctx, residual, pre, 32, 4, 2); @@ -292,6 +495,18 @@ static void test_graph_construction() { ggml_tensor * logits = llama_dsv41_build_output(ctx, residual, pre, output_norm, output, 1.0e-20f, 4); check(logits->ne[0] == 64 && logits->ne[1] == 2, "final output graph shape mismatch"); + ggml_tensor * original_ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 6, 2); + ggml_tensor * slot_ids = ggml_cont(ctx, original_ids); + check(ggml_is_contiguous(original_ids), "original expert IDs must be contiguous"); + check(ggml_is_contiguous(slot_ids) && slot_ids != original_ids, + "slot IDs must be a distinct contiguous remap"); + + ggml_tensor * routed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 2); + ggml_tensor * shared = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 2); + ggml_tensor * combined = ggml_add(ctx, routed, shared); + check(combined->src[0] == routed && combined->src[1] == shared, + "shared expert output is not added to routed output"); + ggml_free(ctx); } @@ -303,6 +518,7 @@ int main() { test_raw_ring(); test_candidates(); test_output_collapse(); + test_graph_contract(); test_graph_construction(); return 0; } diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 917d3214a630..68f26ccd8911 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -563,7 +563,7 @@ static bool arch_supported(const llm_arch arch) { return false; } if (arch == LLM_ARCH_DEEPSEEK41) { - return false; // Published-model construction is gated until Engram and expert streaming are available. + return false; // The fixed published geometry is too large for the compact generated-model fixture. } // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU From 8cbc9e66e3331b5962e56e6b7effaba14ef7943f Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 21:17:49 -0700 Subject: [PATCH 2/5] deepseek41 : fix graph sentinel and state restore Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-memory-dsv41.cpp | 48 +++++++++++++++--- src/models/deepseek41.cpp | 36 ++++++++++++-- tests/test-deepseek41-memory.cpp | 81 +++++++++++++++++++++++++++++++ tests/test-deepseek41-runtime.cpp | 71 +++++++++++++++++++++++++++ 4 files changed, 226 insertions(+), 10 deletions(-) diff --git a/src/llama-memory-dsv41.cpp b/src/llama-memory-dsv41.cpp index 42f3ed8c7215..2fd088f5a962 100644 --- a/src/llama-memory-dsv41.cpp +++ b/src/llama-memory-dsv41.cpp @@ -795,6 +795,16 @@ void llama_memory_dsv41::state_read( pimpl->config.engram->checkpoint() : llama_dsv41_engram_snapshot {}; std::set restored; std::vector tensor_data; + struct staged_sequence { + llama_seq_id target = -1; + llama_pos pos = -1; + std::vector candidates; + std::vector planes; + }; + std::vector staged; + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + staged.reserve(count); + } try { for (uint32_t i = 0; i < count; ++i) { @@ -826,24 +836,48 @@ void llama_memory_dsv41::state_read( io.read(engram_state.history.tail.data(), sizeof(engram_state.history.tail)); engram_snapshot.sequences[target] = engram_state; } + std::vector planes; + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + planes.reserve(tensors.size()); + } for (ggml_tensor * tensor : tensors) { uint64_t plane = 0; io.read(&plane, sizeof(plane)); if (plane != pimpl->sequence_plane(tensor)) { throw std::runtime_error("DeepSeek V4.1 state tensor layout differs"); } - tensor_data.resize(plane); - io.read(tensor_data.data(), plane); - ggml_backend_tensor_set(tensor, tensor_data.data(), (size_t) target*plane, plane); + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + planes.push_back(plane); + } else { + tensor_data.resize(plane); + io.read(tensor_data.data(), plane); + ggml_backend_tensor_set(tensor, tensor_data.data(), (size_t) target*plane, plane); + } + } + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + staged.push_back({ target, pos, std::move(candidates), std::move(planes) }); + } else { + pimpl->sequences[target].pos = pos; + pimpl->sequences[target].candidates = std::move(candidates); + pimpl->update_position_state(target); + pimpl->update_candidate_state(target); } - pimpl->sequences[target].pos = pos; - pimpl->sequences[target].candidates = std::move(candidates); - pimpl->update_position_state(target); - pimpl->update_candidate_state(target); } if (pimpl->config.engram) { pimpl->config.engram->restore(engram_snapshot); } + for (auto & sequence : staged) { + for (size_t i = 0; i < tensors.size(); ++i) { + io.read_tensor( + tensors[i], + (size_t) sequence.target*sequence.planes[i], + sequence.planes[i]); + } + pimpl->sequences[sequence.target].pos = sequence.pos; + pimpl->sequences[sequence.target].candidates = std::move(sequence.candidates); + pimpl->update_position_state(sequence.target); + pimpl->update_candidate_state(sequence.target); + } if (seq_id == -1) { pimpl->engram_boundaries.clear(); pimpl->rollback_states.clear(); diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index 4d6b82d5881d..6de976159211 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -19,6 +19,11 @@ static float dsv41_rope_attn_factor(float freq_scale) { } std::string llama_dsv41_graph_trace_name(const char * trace, uint32_t layer); +ggml_tensor * llama_dsv41_graph_append_zero_row(ggml_context * ctx, ggml_tensor * tensor); +ggml_tensor * llama_dsv41_graph_completion_zero( + ggml_context * ctx, + ggml_tensor * dependency, + ggml_type type); std::string llama_dsv41_graph_trace_name( const char * trace, @@ -397,9 +402,16 @@ static ggml_tensor * dsv41_flatten_memory( static ggml_tensor * dsv41_append_zero_row( ggml_context * ctx, ggml_tensor * tensor) { + const ggml_type type = tensor->type; ggml_tensor * row = ggml_view_2d( ctx, tensor, tensor->ne[0], 1, tensor->nb[1], 0); + if (row->type != GGML_TYPE_F32) { + row = ggml_cast(ctx, row, GGML_TYPE_F32); + } row = ggml_scale(ctx, row, 0.0f); + if (type != GGML_TYPE_F32) { + row = ggml_cast(ctx, row, type); + } return ggml_concat(ctx, tensor, row, 1); } @@ -445,7 +457,9 @@ static ggml_tensor * dsv41_sort_row_ids( ctx, ggml_reshape_3d(ctx, ids, 1, ids->ne[0], ids->ne[1]), order); - return ggml_reshape_2d(ctx, sorted, ids->ne[0], ids->ne[1]); + return ggml_cont( + ctx, + ggml_reshape_2d(ctx, sorted, ids->ne[0], ids->ne[1])); } } @@ -807,8 +821,9 @@ static ggml_tensor * dsv41_completion_zero( ggml_type type) { ggml_tensor * marker = ggml_view_1d(ctx, dependency, 1, 0); marker = ggml_argsort_top_k(ctx, marker, 1); - marker = ggml_cast(ctx, marker, type); - return ggml_scale(ctx, marker, 0.0f); + marker = ggml_cast(ctx, marker, GGML_TYPE_F32); + marker = ggml_scale(ctx, marker, 0.0f); + return type == GGML_TYPE_F32 ? marker : ggml_cast(ctx, marker, type); } static ggml_tensor * dsv41_build_candidate_mask( @@ -1214,6 +1229,8 @@ static ggml_tensor * dsv41_build_attention( if (selected != nullptr) { ggml_tensor * comp_store = dsv41_flatten_memory( graph.ctx0, memory->compressed_kv(kv_source)); + comp_store = dsv41_append_zero_row( + graph.ctx0, comp_store); ggml_tensor * compressed = ggml_get_rows( graph.ctx0, comp_store, selected); compressed = ggml_reshape_4d( @@ -1330,6 +1347,19 @@ static std::pair dsv41_build_router( } +ggml_tensor * llama_dsv41_graph_append_zero_row( + ggml_context * ctx, + ggml_tensor * tensor) { + return dsv41_append_zero_row(ctx, tensor); +} + +ggml_tensor * llama_dsv41_graph_completion_zero( + ggml_context * ctx, + ggml_tensor * dependency, + ggml_type type) { + return dsv41_completion_zero(ctx, dependency, type); +} + llama_model_deepseek41::graph::graph( const llama_model & model, const llm_graph_params & params) : diff --git a/tests/test-deepseek41-memory.cpp b/tests/test-deepseek41-memory.cpp index e4f6302f4ba1..b5a3898e2e5c 100644 --- a/tests/test-deepseek41-memory.cpp +++ b/tests/test-deepseek41-memory.cpp @@ -174,6 +174,66 @@ class vector_reader : public llama_io_read_i { size_t offset = 0; }; +class device_writer : public llama_io_write_i { +public: + void write(const void * src, size_t size) override { + const uint8_t * bytes = static_cast(src); + metadata.insert(metadata.end(), bytes, bytes + size); + } + + void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { + tensors.emplace_back(size); + ggml_backend_tensor_get(tensor, tensors.back().data(), offset, size); + tensor_bytes += size; + } + + size_t n_bytes() override { + return metadata.size(); + } + + std::vector metadata; + std::vector> tensors; + size_t tensor_bytes = 0; +}; + +class device_reader : public llama_io_read_i { +public: + explicit device_reader(const device_writer & writer) : + metadata(writer.metadata), + tensors(writer.tensors) { + } + + void read(void * dst, size_t size) override { + if (offset > metadata.size() || size > metadata.size() - offset) { + throw std::runtime_error("test on-device metadata is truncated"); + } + std::memcpy(dst, metadata.data() + offset, size); + offset += size; + } + + void read_tensor(ggml_tensor * tensor, size_t tensor_offset, size_t size) override { + if (i_tensor >= tensors.size() || tensors[i_tensor].size() != size) { + throw std::runtime_error("test on-device tensor layout differs"); + } + ggml_backend_tensor_set(tensor, tensors[i_tensor].data(), tensor_offset, size); + ++i_tensor; + } + + size_t n_bytes() override { + return offset; + } + + size_t tensor_reads() const { + return i_tensor; + } + +private: + const std::vector & metadata; + const std::vector> & tensors; + size_t offset = 0; + size_t i_tensor = 0; +}; + static void test_transaction_commit_rollback() { llama_memory_dsv41 memory(small_config()); llama_ubatch ubatch = make_ubatch(0, 1, 0); @@ -458,6 +518,27 @@ static void test_state_save_load() { vector_reader full_reader(full_writer.data); memory.state_read(full_reader, -1); check(memory.seq_pos_max(2) == 2, "full state restore lost sequence state"); + + device_writer on_device_writer; + memory.state_write(on_device_writer, 2, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); + vector_writer host_writer; + memory.state_write(host_writer, 2); + check(on_device_writer.metadata.size() + on_device_writer.tensor_bytes == host_writer.data.size(), + "on-device state metadata includes tensor payload bytes"); + memory.clear(true); + device_reader on_device_reader(on_device_writer); + memory.state_read(on_device_reader, 1, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); + check(on_device_reader.tensor_reads() == on_device_writer.tensors.size(), + "on-device state restore did not consume every tensor"); + check(memory.seq_pos_max(1) == 2 && + memory.sequence_candidate_ids(1) == std::vector({ 11 }), + "on-device state restore lost sequence metadata"); + ggml_backend_tensor_get( + memory.raw_k(0), + actual.data(), + ((size_t) memory.config().raw_window + 2)*row_bytes, + row_bytes); + check(actual == expected, "on-device state restore lost raw cache data"); } static void test_accounting() { diff --git a/tests/test-deepseek41-runtime.cpp b/tests/test-deepseek41-runtime.cpp index a03743a8551d..fed172a333e8 100644 --- a/tests/test-deepseek41-runtime.cpp +++ b/tests/test-deepseek41-runtime.cpp @@ -1,6 +1,8 @@ #include "../src/llama-dsv41.h" #include "../src/llama-arch.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" #include "ggml.h" #include @@ -16,6 +18,11 @@ #include std::string llama_dsv41_graph_trace_name(const char * trace, uint32_t layer); +ggml_tensor * llama_dsv41_graph_append_zero_row(ggml_context * ctx, ggml_tensor * tensor); +ggml_tensor * llama_dsv41_graph_completion_zero( + ggml_context * ctx, + ggml_tensor * dependency, + ggml_type type); static void check(bool condition, const std::string & message) { if (!condition) { @@ -382,6 +389,46 @@ static void test_graph_construction() { const float * ordered = static_cast(ordered_probs->data); check(ordered[0] < ordered[1] && ordered[1] < ordered[2], "shared-softmax segment order mismatch"); + ggml_tensor * f16_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 32, 1); + std::vector cache_values(32, 1.0f); + ggml_fp32_to_fp16_row( + cache_values.data(), + static_cast(f16_cache->data), + cache_values.size()); + ggml_tensor * cache_with_sentinel = + llama_dsv41_graph_append_zero_row(ctx, f16_cache); + ggml_tensor * sentinel_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 2); + static_cast(sentinel_ids->data)[0] = 1; + static_cast(sentinel_ids->data)[1] = 0; + ggml_tensor * sentinel_rows = ggml_get_rows(ctx, cache_with_sentinel, sentinel_ids); + ggml_tensor * completion_zero = + llama_dsv41_graph_completion_zero(ctx, sentinel_rows, GGML_TYPE_F16); + ggml_cgraph * support_gf = ggml_new_graph(ctx); + ggml_build_forward_expand(support_gf, sentinel_rows); + ggml_build_forward_expand(support_gf, completion_zero); + ggml_backend_t backend = ggml_backend_cpu_init(); + check(backend != nullptr, "failed to create graph support backend"); + for (int i = 0; i < ggml_graph_n_nodes(support_gf); ++i) { + ggml_tensor * node = ggml_graph_node(support_gf, i); + if (node->op == GGML_OP_SCALE) { + check(node->src[0]->type == GGML_TYPE_F32, + "DeepSeek V4.1 graph contains a non-F32 SCALE input"); + } + check(ggml_backend_supports_op(backend, node), + "CPU backend does not support a DeepSeek V4.1 dependency node"); + } + ggml_backend_free(backend); + check(cache_with_sentinel->ne[1] == 2, + "compressed cache sentinel row was not allocated"); + check(ggml_graph_compute_with_ctx(ctx, support_gf, 1) == GGML_STATUS_SUCCESS, + "compressed sentinel graph execution failed"); + for (uint32_t i = 0; i < 32; ++i) { + check(ggml_get_f32_1d(sentinel_rows, i) == 0.0f, + "compressed sentinel row is not zero"); + check(ggml_get_f32_1d(sentinel_rows, 32 + i) == 1.0f, + "compressed real row changed"); + } + ggml_tensor * candidate_scores = ggml_new_tensor_2d( ctx, GGML_TYPE_F32, 24, 2); ggml_tensor * candidate_bias = ggml_new_tensor_2d( @@ -471,8 +518,24 @@ static void test_graph_construction() { selected_order); selected_sorted = ggml_cont( ctx, ggml_reshape_2d(ctx, selected_sorted, 3, 2)); + ggml_tensor * routing_probs = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 6, 2); + const float routing_values[] = { + 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, + 20.0f, 21.0f, 22.0f, 23.0f, 24.0f, 25.0f, + }; + std::memcpy( + routing_probs->data, routing_values, + sizeof(routing_values)); + ggml_tensor * selected_weights = ggml_get_rows( + ctx, + ggml_reshape_3d(ctx, routing_probs, 1, 6, 2), + selected_sorted); + selected_weights = ggml_cont( + ctx, ggml_reshape_2d(ctx, selected_weights, 3, 2)); ggml_cgraph * selected_gf = ggml_new_graph(ctx); ggml_build_forward_expand(selected_gf, selected_sorted); + ggml_build_forward_expand(selected_gf, selected_weights); check( ggml_graph_compute_with_ctx(ctx, selected_gf, 1) == GGML_STATUS_SUCCESS, @@ -483,6 +546,14 @@ static void test_graph_construction() { selected_sorted->data, selected_expected, sizeof(selected_expected)) == 0, "selected IDs are not accumulated in original ID order"); + const float selected_weight_expected[] = { + 11.0f, 13.0f, 15.0f, 20.0f, 22.0f, 24.0f, + }; + check( + std::memcmp( + selected_weights->data, selected_weight_expected, + sizeof(selected_weight_expected)) == 0, + "routing weights are not paired with sorted original IDs"); ggml_tensor * residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 32, 4, 2); ggml_tensor * pre = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 2); From 081c549451b92c01fe832ffb0bdd9ad90370e21a Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 21:30:29 -0700 Subject: [PATCH 3/5] deepseek41 : reuse packed Engram row IDs Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/models/deepseek41.cpp | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index 6de976159211..a7e1094f07fc 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -319,15 +319,9 @@ class dsv41_graph_input final : public llm_graph_input_i { } transaction->upload_layer( index, 0, n_tokens, engram_rows[index], engram_select[index]); - std::vector ids((size_t) LLAMA_ENGRAM_COLS*n_tokens); - const uint32_t * source = transaction->row_ids(index); - for (uint32_t token = 0; token < n_tokens; ++token) { - std::copy_n( - source + (size_t) token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, - LLAMA_ENGRAM_COLS, - ids.begin() + (size_t) token*LLAMA_ENGRAM_COLS); - } - set_tensor(engram_row_ids[index], ids); + set_tensor( + engram_row_ids[index], + mctx->engram_row_ids(index)); } } From 4f1a4ebd74b58f1b22580b5c7f5ab976fdb99506 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Mon, 14 Sep 2026 00:27:29 -0700 Subject: [PATCH 4/5] deepseek41 : bound Engram rollback state Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-memory-dsv41.cpp | 38 +++++--------------------------- src/llama-memory-dsv41.h | 1 + tests/test-deepseek41-memory.cpp | 15 ++++++++++++- 3 files changed, 20 insertions(+), 34 deletions(-) diff --git a/src/llama-memory-dsv41.cpp b/src/llama-memory-dsv41.cpp index 2fd088f5a962..18aa10a8c564 100644 --- a/src/llama-memory-dsv41.cpp +++ b/src/llama-memory-dsv41.cpp @@ -148,7 +148,6 @@ struct llama_memory_dsv41::impl { uint64_t backend_layout = 1469598103934665603ULL; uint64_t graph_workspace = 0; bool transaction_active = false; - std::map> engram_boundaries; std::map rollback_states; explicit impl(llama_dsv41_memory_config config) : config(std::move(config)) { @@ -515,7 +514,6 @@ void llama_memory_dsv41::clear(bool data) { for (auto & sequence : pimpl->sequences) { sequence = {}; } - pimpl->engram_boundaries.clear(); pimpl->rollback_states.clear(); if (pimpl->config.engram) { for (uint32_t seq = 0; seq < pimpl->config.n_seq; ++seq) { @@ -582,7 +580,6 @@ bool llama_memory_dsv41::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) if (pimpl->config.engram) { pimpl->config.engram->seq_remove(seq_id); } - pimpl->engram_boundaries.erase(seq_id); pimpl->rollback_states.erase(seq_id); } else { for (const auto & snapshot : rollback->second.snapshots) { @@ -595,8 +592,6 @@ bool llama_memory_dsv41::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) llama_dsv41_engram_snapshot snapshot = pimpl->config.engram->checkpoint(); snapshot.sequences[seq_id] = rollback->second.engram; pimpl->config.engram->restore(snapshot); - auto & boundaries = pimpl->engram_boundaries[seq_id]; - boundaries.erase(boundaries.lower_bound(begin), boundaries.end()); } pimpl->update_position_state(seq_id); pimpl->update_candidate_state(seq_id); @@ -624,7 +619,6 @@ void llama_memory_dsv41::seq_cp( pimpl->rollback_states.erase(seq_id_dst); if (pimpl->config.engram) { pimpl->config.engram->seq_copy(seq_id_src, seq_id_dst); - pimpl->engram_boundaries[seq_id_dst] = pimpl->engram_boundaries[seq_id_src]; } ++pimpl->generation; } @@ -645,7 +639,6 @@ void llama_memory_dsv41::seq_keep(llama_seq_id seq_id) { if (pimpl->config.engram) { pimpl->config.engram->seq_remove(current); } - pimpl->engram_boundaries.erase(current); pimpl->rollback_states.erase(current); } ++pimpl->generation; @@ -774,8 +767,7 @@ void llama_memory_dsv41::state_read( const auto sequence_is_empty = [&](llama_seq_id current) { if (pimpl->sequences[current].pos >= 0 || !pimpl->sequences[current].candidates.empty() || - pimpl->rollback_states.count(current) != 0 || - pimpl->engram_boundaries.count(current) != 0) { + pimpl->rollback_states.count(current) != 0) { return false; } return !pimpl->config.engram || pimpl->config.engram->sequence(current).pos < 0; @@ -879,10 +871,8 @@ void llama_memory_dsv41::state_read( pimpl->update_candidate_state(sequence.target); } if (seq_id == -1) { - pimpl->engram_boundaries.clear(); pimpl->rollback_states.clear(); } else { - pimpl->engram_boundaries.erase(seq_id); pimpl->rollback_states.erase(seq_id); } ++pimpl->generation; @@ -967,6 +957,10 @@ std::vector llama_memory_dsv41::sequence_candidate_ids(llama_seq_id seq return pimpl->sequences[seq_id].candidates; } +size_t llama_memory_dsv41::retained_rollback_count() const { + return pimpl->rollback_states.size(); +} + bool llama_memory_dsv41::engram_enabled() const { return pimpl->config.engram != nullptr; } @@ -1315,35 +1309,16 @@ void llama_memory_dsv41_context::commit() { mem->pimpl->position_state_values(transaction->next_sequences[seq_id]); committed_candidates = mem->pimpl->candidate_state_values(transaction->next_sequences[seq_id]); - std::map next_boundaries; - if (transaction->engram) { - const auto found = mem->pimpl->engram_boundaries.find(seq_id); - if (found != mem->pimpl->engram_boundaries.end()) { - next_boundaries = found->second; - } - next_boundaries[transaction->start_positions.front()] = - transaction->engram_before.at(seq_id); - } auto rollback_slot = mem->pimpl->rollback_states.end(); bool inserted_rollback = false; - auto boundary_slot = mem->pimpl->engram_boundaries.end(); - bool inserted_boundary = false; try { std::tie(rollback_slot, inserted_rollback) = mem->pimpl->rollback_states.emplace(seq_id, llama_memory_dsv41::impl::rollback_state {}); - if (transaction->engram) { - std::tie(boundary_slot, inserted_boundary) = - mem->pimpl->engram_boundaries.emplace( - seq_id, std::map {}); - } if (transaction->engram) { mem->pimpl->config.engram->commit(*transaction->engram); } } catch (...) { - if (inserted_boundary) { - mem->pimpl->engram_boundaries.erase(boundary_slot); - } if (inserted_rollback) { mem->pimpl->rollback_states.erase(rollback_slot); } @@ -1363,9 +1338,6 @@ void llama_memory_dsv41_context::commit() { position_values.size()*sizeof(int32_t)); } mem->pimpl->sequences.swap(transaction->next_sequences); - if (transaction->engram) { - boundary_slot->second.swap(next_boundaries); - } rollback.snapshots.swap(transaction->snapshots); rollback_slot->second = std::move(rollback); mem->pimpl->generation = transaction->plan.generation; diff --git a/src/llama-memory-dsv41.h b/src/llama-memory-dsv41.h index 0eb074c12fe5..53b82e0d0596 100644 --- a/src/llama-memory-dsv41.h +++ b/src/llama-memory-dsv41.h @@ -164,6 +164,7 @@ class llama_memory_dsv41 : public llama_memory_i { const llama_dsv41_memory_config & config() const; llama_dsv41_memory_accounting accounting() const; std::vector sequence_candidate_ids(llama_seq_id seq_id) const; + size_t retained_rollback_count() const; bool engram_enabled() const; private: diff --git a/tests/test-deepseek41-memory.cpp b/tests/test-deepseek41-memory.cpp index b5a3898e2e5c..057b9e4848f4 100644 --- a/tests/test-deepseek41-memory.cpp +++ b/tests/test-deepseek41-memory.cpp @@ -651,7 +651,7 @@ struct engram_test_file { static void test_engram_transaction() { const llama_engram_layout layout = make_engram_layout(); engram_test_file file(layout); - auto config = small_config(8, 1, 4); + auto config = small_config(2048, 1, 4); auto runtime = std::make_unique( layout, file.extents, config.n_ubatch); llama_dsv41_engram_runtime * runtime_ptr = runtime.get(); @@ -707,6 +707,19 @@ static void test_engram_transaction() { check(memory.seq_pos_max(0) == 2, "Engram commit position mismatch"); check(memory.seq_rm(0, 2, -1), "Engram suffix rollback at transaction boundary failed"); check(memory.seq_pos_max(0) == 1, "Engram suffix rollback did not restore position"); + + for (llama_pos pos = 2; pos < 1024; ++pos) { + llama_memory_dsv41_context decode( + &memory, std::vector { make_ubatch(pos, 1, 0) }); + check(decode.apply(), "long Engram decode prepare failed"); + decode.commit(); + check(memory.retained_rollback_count() == 1, + "Engram decode retained more than the immediate rollback state"); + } + check(memory.seq_rm(0, 1023, -1), "long Engram decode immediate rollback failed"); + check(memory.seq_pos_max(0) == 1022, "long Engram decode rollback restored the wrong position"); + check(memory.retained_rollback_count() == 0, + "Engram rollback retained an obsolete rollback state"); } #endif From dd2fd630d03edf6e1bf42f82ed4f034fd89b5d26 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Mon, 14 Sep 2026 02:59:39 -0700 Subject: [PATCH 5/5] deepseek41 : bind full graph provenance Assisted-by: GPT-5.6 Sol Copilot-Session: 30a08176-cebe-482f-a2f6-eebfac0a3d34 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>