deepseek41 : activate scheduler-backed full graph

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

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:

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;
     }

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;
@@ -536,6 +534,57 @@ ggml_tensor * llama_dsv41_build_shared_softmax(
     return ggml_soft_max(ctx, scores);
 }
 
+ggml_tensor * llama_dsv41_build_candidate_blocks(
+        ggml_context * ctx,
+        ggml_tensor * block_scores,
+        ggml_tensor * final_blocks,
+        uint32_t n_candidate) {
+    if (block_scores == nullptr || final_blocks == nullptr ||
+            block_scores->type != GGML_TYPE_F32 ||
+            final_blocks->type != GGML_TYPE_I32 ||
+            final_blocks->ne[0] != 1 ||
+            final_blocks->ne[1] != block_scores->ne[1]) {
+        throw std::runtime_error("DeepSeek V4.1 candidate graph shape mismatch");
+    }
+
+    const int64_t n_blocks = block_scores->ne[0];
+    const int64_t n_tokens = block_scores->ne[1];
+    if (n_candidate == 0 || n_candidate > (uint32_t) n_blocks) {
+        throw std::runtime_error("DeepSeek V4.1 candidate graph width is invalid");
+    }
+    if (n_candidate == 1) {
+        return ggml_cont(ctx, final_blocks);
+    }
+
+    ggml_tensor * local = ggml_reshape_2d(
+            ctx, ggml_arange(ctx, 0.0f, (float) (n_blocks - 1), 1.0f),
+            n_blocks - 1, 1);
+    local = ggml_repeat_4d(
+            ctx, local, n_blocks - 1, n_tokens, 1, 1);
+    ggml_tensor * final_f32 = ggml_cast(ctx, final_blocks, GGML_TYPE_F32);
+    ggml_tensor * shift = ggml_step(
+            ctx, ggml_scale_bias(
+                ctx, ggml_sub(ctx, local, final_f32), 1.0f, 0.5f));
+    ggml_tensor * ordinary_ids = ggml_cast(
+            ctx, ggml_add(ctx, local, shift), GGML_TYPE_I32);
+
+    ggml_tensor * ordinary_scores = ggml_get_rows(
+            ctx,
+            ggml_reshape_3d(ctx, block_scores, 1, n_blocks, n_tokens),
+            ordinary_ids);
+    ordinary_scores = ggml_reshape_2d(
+            ctx, ordinary_scores, n_blocks - 1, n_tokens);
+    ggml_tensor * selected_local = ggml_argsort_top_k(
+            ctx, ordinary_scores, n_candidate - 1);
+    ggml_tensor * selected = ggml_get_rows(
+            ctx,
+            ggml_reshape_3d(ctx, ordinary_ids, 1, n_blocks - 1, n_tokens),
+            selected_local);
+    selected = ggml_reshape_2d(
+            ctx, selected, n_candidate - 1, n_tokens);
+    return ggml_cont(ctx, ggml_concat(ctx, final_blocks, selected, 0));
+}
+
 ggml_tensor * llama_dsv41_build_output_collapse(
         ggml_context * ctx,
         ggml_tensor * residual,
@@ -555,6 +604,15 @@ ggml_tensor * llama_dsv41_build_output_collapse(
     return ggml_cast(ctx, collapsed, GGML_TYPE_BF16);
 }
 
+ggml_tensor * llama_dsv41_build_output_norm_input(
+        ggml_context * ctx,
+        ggml_tensor * collapsed) {
+    if (collapsed == nullptr || collapsed->type != GGML_TYPE_BF16) {
+        throw std::runtime_error("DeepSeek V4.1 output collapse must be BF16");
+    }
+    return ggml_cast(ctx, collapsed, GGML_TYPE_F32);
+}
+
 ggml_tensor * llama_dsv41_build_output(
         ggml_context * ctx,
         ggml_tensor * residual,
@@ -569,7 +627,8 @@ ggml_tensor * llama_dsv41_build_output(
 
     ggml_tensor * collapsed = llama_dsv41_build_output_collapse(
             ctx, residual, pre, residual->ne[0], hc_mult, residual->ne[2]);
-    ggml_tensor * normalized = ggml_rms_norm(ctx, ggml_cast(ctx, collapsed, GGML_TYPE_F32), rms_eps);
+    ggml_tensor * normalized = ggml_rms_norm(
+            ctx, llama_dsv41_build_output_norm_input(ctx, collapsed), rms_eps);
     normalized = ggml_mul(ctx, normalized, output_norm);
     return ggml_mul_mat(ctx, output, normalized);
 }

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<int32_t> state_pos;
@@ -222,6 +221,12 @@ ggml_tensor * llama_dsv41_build_shared_softmax(
         ggml_tensor * raw_scores,
         ggml_tensor * compressed_scores);
 
+ggml_tensor * llama_dsv41_build_candidate_blocks(
+        ggml_context * ctx,
+        ggml_tensor * block_scores,
+        ggml_tensor * final_blocks,
+        uint32_t n_candidate);
+
 ggml_tensor * llama_dsv41_build_output_collapse(
         ggml_context * ctx,
         ggml_tensor * residual,
@@ -230,6 +235,10 @@ ggml_tensor * llama_dsv41_build_output_collapse(
         uint32_t hc_mult,
         uint32_t n_tokens);
 
+ggml_tensor * llama_dsv41_build_output_norm_input(
+        ggml_context * ctx,
+        ggml_tensor * collapsed);
+
 ggml_tensor * llama_dsv41_build_output(
         ggml_context * ctx,
         ggml_tensor * residual,

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);

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;