diff --git a/common/arg.cpp b/common/arg.cpp index 52948e045954..ffc1d5604db6 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2784,6 +2784,26 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.lazy_mode = LLAMA_LAZY_MODE_ON; } ).set_env("LLAMA_ARG_NGRAM_ON_DISK")); + add_opt(common_arg( + {"--expert-cache-slots"}, "N", + "DeepSeek V4.1 routed experts resident per layer; requires --expert-cache-mib", + [](common_params & params, int value) { + if (value <= 0) { + throw std::invalid_argument("invalid value"); + } + params.expert_cache_slots = value; + } + ).set_env("LLAMA_ARG_EXPERT_CACHE_SLOTS")); + add_opt(common_arg( + {"--expert-cache-mib"}, "MiB", + "aggregate DeepSeek V4.1 fixed expert slot-tensor capacity; requires --expert-cache-slots", + [](common_params & params, int value) { + if (value <= 0) { + throw std::invalid_argument("invalid value"); + } + params.expert_cache_mib = value; + } + ).set_env("LLAMA_ARG_EXPERT_CACHE_MIB")); add_opt(common_arg( {"-cmoe", "--cpu-moe"}, "keep all Mixture of Experts (MoE) weights in the CPU", diff --git a/common/common.cpp b/common/common.cpp index 9ca280d17bd1..56647ee8bb76 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1693,6 +1693,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.check_tensors = params.check_tensors; mparams.use_extra_bufts = !params.no_extra_bufts; mparams.no_host = params.no_host; + mparams.expert_cache_slots = params.expert_cache_slots; + mparams.expert_cache_bytes = params.expert_cache_mib > 0 ? (size_t) params.expert_cache_mib << 20 : 0; if (params.kv_overrides.empty()) { mparams.kv_overrides = NULL; } else { diff --git a/common/common.h b/common/common.h index 6410b373b3b5..f4626bb5d05c 100644 --- a/common/common.h +++ b/common/common.h @@ -625,6 +625,8 @@ struct common_params { bool no_op_offload = false; // globally disable offload host tensor operations to device bool no_extra_bufts = false; // disable extra buffer types (used for weight repacking) bool no_host = false; // bypass host buffer allowing extra buffers to be used + int32_t expert_cache_slots = 0; // DeepSeek V4.1 routed experts resident per layer + int32_t expert_cache_mib = 0; // aggregate fixed slot-tensor capacity bool single_turn = false; // single turn chat conversation ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index c978ab9d5417..ff400781de50 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -559,6 +559,7 @@ class MODEL_ARCH(IntEnum): DEEPSEEK2OCR = auto() DEEPSEEK32 = auto() DEEPSEEK4 = auto() + DEEPSEEK41 = auto() CHATGLM = auto() GLM4 = auto() GLM4_MOE = auto() @@ -823,6 +824,10 @@ class MODEL_TENSOR(IntEnum): PLE_NORM_QUERY = auto() # qwen4exp PLE_NORM_CONV = auto() # qwen4exp PLE_CONV1D = auto() # qwen4exp + ENGRAM_EMBD = auto() + ENGRAM_Q_NORM = auto() + ENGRAM_K_NORM = auto() + ENGRAM_KV = auto() ATTN_COMPRESSOR_WKV = auto() ATTN_COMPRESSOR_WGATE = auto() ATTN_COMPRESSOR_APE = auto() @@ -1317,6 +1322,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.DEEPSEEK2OCR: "deepseek2-ocr", MODEL_ARCH.DEEPSEEK32: "deepseek32", MODEL_ARCH.DEEPSEEK4: "deepseek4", + MODEL_ARCH.DEEPSEEK41: "deepseek41", MODEL_ARCH.CHATGLM: "chatglm", MODEL_ARCH.GLM4: "glm4", MODEL_ARCH.GLM4_MOE: "glm4moe", @@ -1580,6 +1586,10 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.PLE_NORM_QUERY: "blk.{bid}.ple_norm_query", # qwen4exp MODEL_TENSOR.PLE_NORM_CONV: "blk.{bid}.ple_norm_conv", # qwen4exp MODEL_TENSOR.PLE_CONV1D: "blk.{bid}.ple_conv1d", # qwen4exp + MODEL_TENSOR.ENGRAM_EMBD: "blk.{bid}.engram_embd", + MODEL_TENSOR.ENGRAM_Q_NORM: "blk.{bid}.engram_q_norm", + MODEL_TENSOR.ENGRAM_K_NORM: "blk.{bid}.engram_k_norm", + MODEL_TENSOR.ENGRAM_KV: "blk.{bid}.engram_kv", MODEL_TENSOR.ATTN_COMPRESSOR_WKV: "blk.{bid}.attn_compressor_kv", MODEL_TENSOR.ATTN_COMPRESSOR_WGATE: "blk.{bid}.attn_compressor_gate", MODEL_TENSOR.ATTN_COMPRESSOR_APE: "blk.{bid}.attn_compressor_ape", @@ -3940,6 +3950,47 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.DEEPSEEK41: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_OUT_A, + MODEL_TENSOR.ATTN_OUT_B, + MODEL_TENSOR.HC_ATTN_FN, + MODEL_TENSOR.HC_ATTN_BASE, + MODEL_TENSOR.HC_ATTN_SCALE, + MODEL_TENSOR.HC_FFN_FN, + MODEL_TENSOR.HC_FFN_BASE, + MODEL_TENSOR.HC_FFN_SCALE, + MODEL_TENSOR.ATTN_COMPRESSOR_WKV, + MODEL_TENSOR.ATTN_COMPRESSOR_WGATE, + MODEL_TENSOR.ATTN_COMPRESSOR_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_EXP_PROBS_B_VL, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.ENGRAM_EMBD, + MODEL_TENSOR.ENGRAM_Q_NORM, + MODEL_TENSOR.ENGRAM_K_NORM, + MODEL_TENSOR.ENGRAM_KV, + ], MODEL_ARCH.ERNIE4_5_MOE: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/tests/test_deepseek41_schema.py b/gguf-py/tests/test_deepseek41_schema.py new file mode 100644 index 000000000000..b360e7d02162 --- /dev/null +++ b/gguf-py/tests/test_deepseek41_schema.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 + +import os +import sys +import unittest +from pathlib import Path + +if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent / "gguf-py").exists(): + sys.path.insert(0, str(Path(__file__).parent.parent)) + +from gguf.constants import MODEL_ARCH, MODEL_ARCH_NAMES, MODEL_TENSOR, MODEL_TENSORS, TENSOR_NAMES + + +class TestDeepSeek41Schema(unittest.TestCase): + + def test_architecture_name(self): + self.assertEqual(MODEL_ARCH_NAMES[MODEL_ARCH.DEEPSEEK41], "deepseek41") + + def test_engram_tensor_names(self): + expected = { + MODEL_TENSOR.ENGRAM_EMBD: "blk.14.engram_embd", + MODEL_TENSOR.ENGRAM_Q_NORM: "blk.14.engram_q_norm", + MODEL_TENSOR.ENGRAM_K_NORM: "blk.14.engram_k_norm", + MODEL_TENSOR.ENGRAM_KV: "blk.14.engram_kv", + } + + for tensor, name in expected.items(): + self.assertIn(tensor, MODEL_TENSORS[MODEL_ARCH.DEEPSEEK41]) + self.assertEqual(TENSOR_NAMES[tensor].format(bid=14), name) + + def test_output_head_does_not_require_deepseek4_hc_tensors(self): + tensors = MODEL_TENSORS[MODEL_ARCH.DEEPSEEK41] + + self.assertIn(MODEL_TENSOR.OUTPUT_NORM, tensors) + self.assertIn(MODEL_TENSOR.OUTPUT, tensors) + self.assertNotIn(MODEL_TENSOR.HC_HEAD_FN, tensors) + self.assertNotIn(MODEL_TENSOR.HC_HEAD_BASE, tensors) + self.assertNotIn(MODEL_TENSOR.HC_HEAD_SCALE, tensors) + + def test_kv_a_norm_does_not_use_deepseek4_tensor_kind(self): + tensors = MODEL_TENSORS[MODEL_ARCH.DEEPSEEK41] + + self.assertIn(MODEL_TENSOR.ATTN_KV_A_NORM, tensors) + self.assertNotIn(MODEL_TENSOR.ATTN_KV_NORM, tensors) + self.assertEqual(TENSOR_NAMES[MODEL_TENSOR.ATTN_KV_A_NORM].format(bid=14), "blk.14.attn_kv_a_norm") + + +if __name__ == "__main__": + unittest.main() diff --git a/include/llama.h b/include/llama.h index 8851c113aaed..ff514f0c8363 100644 --- a/include/llama.h +++ b/include/llama.h @@ -345,6 +345,10 @@ extern "C" { // the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE int32_t main_gpu; + // DeepSeek V4.1 routed-expert cache. Both values must be non-zero. + size_t expert_cache_bytes; + int32_t expert_cache_slots; + // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() const float * tensor_split; diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 8dbaee3902c1..c7eb67086811 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -16,6 +16,10 @@ set(LLAMA_CORE_SOURCES llama-chat.cpp llama-context.cpp llama-cparams.cpp + llama-dsv41.cpp + llama-dsv41-engram.cpp + llama-dsv41-expert.cpp + llama-expert-store.cpp llama-grammar.cpp llama-graph.cpp llama-hparams.cpp @@ -28,11 +32,14 @@ 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 llama-memory-recurrent.cpp llama-mmap.cpp + llama-bounded-file.cpp + llama-engram.cpp llama-ple-disk.cpp llama-model-loader.cpp llama-model-saver.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 976f6492ea0c..da12003d3321 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -81,6 +81,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_DEEPSEEK2OCR, "deepseek2-ocr" }, { LLM_ARCH_DEEPSEEK32, "deepseek32" }, { LLM_ARCH_DEEPSEEK4, "deepseek4" }, + { LLM_ARCH_DEEPSEEK41, "deepseek41" }, { LLM_ARCH_CHATGLM, "chatglm" }, { LLM_ARCH_GLM4, "glm4" }, { LLM_ARCH_GLM4_MOE, "glm4moe" }, @@ -310,6 +311,57 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_PLE_EOS_TOKEN_ID, "%s.ple.eos_token_id" }, { LLM_KV_PLE_IMAGE_TOKEN_ID, "%s.ple.image_token_id" }, + { LLM_KV_DSV41_CONFIG, "%s.config" }, + { LLM_KV_DSV41_VOCAB_SIZE, "%s.vocab_size" }, + { LLM_KV_DSV41_HIDDEN_SIZE, "%s.hidden_size" }, + { LLM_KV_DSV41_MOE_INTERMEDIATE_SIZE, "%s.moe_intermediate_size" }, + { LLM_KV_DSV41_NUM_HIDDEN_LAYERS, "%s.num_hidden_layers" }, + { LLM_KV_DSV41_NUM_ATTENTION_HEADS, "%s.num_attention_heads" }, + { LLM_KV_DSV41_NUM_KEY_VALUE_HEADS, "%s.num_key_value_heads" }, + { LLM_KV_DSV41_HEAD_DIM, "%s.head_dim" }, + { LLM_KV_DSV41_QK_ROPE_HEAD_DIM, "%s.qk_rope_head_dim" }, + { LLM_KV_DSV41_Q_LORA_RANK, "%s.q_lora_rank" }, + { LLM_KV_DSV41_O_LORA_RANK, "%s.o_lora_rank" }, + { LLM_KV_DSV41_O_GROUPS, "%s.o_groups" }, + { LLM_KV_DSV41_N_ROUTED_EXPERTS, "%s.n_routed_experts" }, + { LLM_KV_DSV41_N_SHARED_EXPERTS, "%s.n_shared_experts" }, + { LLM_KV_DSV41_NUM_EXPERTS_PER_TOK, "%s.num_experts_per_tok" }, + { LLM_KV_DSV41_MAX_POSITION_EMBEDDINGS, "%s.max_position_embeddings" }, + { LLM_KV_DSV41_SLIDING_WINDOW, "%s.sliding_window" }, + { LLM_KV_DSV41_INDEX_N_HEADS, "%s.index_n_heads" }, + { LLM_KV_DSV41_INDEX_HEAD_DIM, "%s.index_head_dim" }, + { LLM_KV_DSV41_INDEX_TOPK, "%s.index_topk" }, + { LLM_KV_DSV41_CANDIDATE_SOURCE_LAYER_ID, "%s.candidate_source_layer_id" }, + { LLM_KV_DSV41_CANDIDATE_TOPK_BLOCKS, "%s.candidate_topk_blocks" }, + { LLM_KV_DSV41_CANDIDATE_BLOCK_SIZE, "%s.candidate_block_size" }, + { LLM_KV_DSV41_HC_MULT, "%s.hc_mult" }, + { LLM_KV_DSV41_HC_SINKHORN_ITERS, "%s.hc_sinkhorn_iters" }, + { LLM_KV_DSV41_ROPE_THETA, "%s.rope_theta" }, + { LLM_KV_DSV41_COMPRESS_ROPE_THETA, "%s.compress_rope_theta" }, + { LLM_KV_DSV41_RMS_NORM_EPS, "%s.rms_norm_eps" }, + { LLM_KV_DSV41_HC_EPS, "%s.hc_eps" }, + { LLM_KV_DSV41_SWIGLU_LIMIT, "%s.swiglu_limit" }, + { LLM_KV_DSV41_ROUTED_SCALING_FACTOR, "%s.routed_scaling_factor" }, + { LLM_KV_DSV41_SCORING_FUNC, "%s.scoring_func" }, + { LLM_KV_DSV41_HIDDEN_ACT, "%s.hidden_act" }, + { LLM_KV_DSV41_TOPK_METHOD, "%s.topk_method" }, + { LLM_KV_DSV41_NORM_TOPK_PROB, "%s.norm_topk_prob" }, + { LLM_KV_DSV41_COMPRESS_RATIOS, "%s.compress_ratios" }, + { LLM_KV_DSV41_KV_SOURCE_LAYER_IDS, "%s.kv_source_layer_ids" }, + { LLM_KV_DSV41_INDEX_SOURCE_LAYER_IDS, "%s.index_source_layer_ids" }, + { LLM_KV_DSV41_ROPE_SCALING_FACTOR, "%s.rope_scaling.factor" }, + { LLM_KV_DSV41_ROPE_SCALING_BETA_FAST, "%s.rope_scaling.beta_fast" }, + { LLM_KV_DSV41_ROPE_SCALING_BETA_SLOW, "%s.rope_scaling.beta_slow" }, + { LLM_KV_DSV41_ROPE_SCALING_ORIG_CTX_LEN, "%s.rope_scaling.original_max_position_embeddings" }, + { LLM_KV_DSV41_ENGRAM_ENCODING, "%s.engram.encoding" }, + { LLM_KV_DSV41_ENGRAM_LAYER_IDS, "%s.engram.layer_ids" }, + { LLM_KV_DSV41_ENGRAM_ROWS, "%s.engram.rows" }, + { LLM_KV_DSV41_ENGRAM_COMPRESSED_VOCAB_SIZE, "%s.engram.compressed_vocab_size" }, + { LLM_KV_DSV41_ENGRAM_PAD_ID, "%s.engram.pad_id" }, + { LLM_KV_DSV41_ENGRAM_TOKEN_MAP, "%s.engram.token_map" }, + { LLM_KV_DSV41_ENGRAM_PRIMES, "%s.engram.primes" }, + { LLM_KV_DSV41_ENGRAM_MULTIPLIERS, "%s.engram.multipliers" }, + { LLM_KV_HASH_LAYER_COUNT, "%s.hash_layer_count" }, { LLM_KV_ROPE_DIMENSION_COUNT, "%s.rope.dimension_count" }, @@ -546,6 +598,10 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_PLE_NORM_QUERY, "blk.%d.ple_norm_query" }, { LLM_TENSOR_PLE_NORM_CONV, "blk.%d.ple_norm_conv" }, { LLM_TENSOR_PLE_CONV1D, "blk.%d.ple_conv1d" }, + { LLM_TENSOR_ENGRAM_EMBD, "blk.%d.engram_embd" }, + { LLM_TENSOR_ENGRAM_Q_NORM, "blk.%d.engram_q_norm" }, + { LLM_TENSOR_ENGRAM_K_NORM, "blk.%d.engram_k_norm" }, + { LLM_TENSOR_ENGRAM_KV, "blk.%d.engram_kv" }, { LLM_TENSOR_ATTN_COMPRESSOR_WKV, "blk.%d.attn_compressor_kv" }, { LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "blk.%d.attn_compressor_gate" }, { LLM_TENSOR_ATTN_COMPRESSOR_APE, "blk.%d.attn_compressor_ape" }, @@ -777,6 +833,10 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_PLE_NORM_QUERY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_PLE_NORM_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_PLE_CONV1D, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, + {LLM_TENSOR_ENGRAM_EMBD, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_ENGRAM_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_ENGRAM_K_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_ENGRAM_KV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, @@ -1089,6 +1149,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-arch.h b/src/llama-arch.h index 88aaaff41d37..ed04639ba42a 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -86,6 +86,7 @@ enum llm_arch { LLM_ARCH_DEEPSEEK2OCR, LLM_ARCH_DEEPSEEK32, LLM_ARCH_DEEPSEEK4, + LLM_ARCH_DEEPSEEK41, LLM_ARCH_CHATGLM, LLM_ARCH_GLM4, LLM_ARCH_GLM4_MOE, @@ -315,6 +316,57 @@ enum llm_kv { LLM_KV_PLE_EOS_TOKEN_ID, LLM_KV_PLE_IMAGE_TOKEN_ID, + LLM_KV_DSV41_CONFIG, + LLM_KV_DSV41_VOCAB_SIZE, + LLM_KV_DSV41_HIDDEN_SIZE, + LLM_KV_DSV41_MOE_INTERMEDIATE_SIZE, + LLM_KV_DSV41_NUM_HIDDEN_LAYERS, + LLM_KV_DSV41_NUM_ATTENTION_HEADS, + LLM_KV_DSV41_NUM_KEY_VALUE_HEADS, + LLM_KV_DSV41_HEAD_DIM, + LLM_KV_DSV41_QK_ROPE_HEAD_DIM, + LLM_KV_DSV41_Q_LORA_RANK, + LLM_KV_DSV41_O_LORA_RANK, + LLM_KV_DSV41_O_GROUPS, + LLM_KV_DSV41_N_ROUTED_EXPERTS, + LLM_KV_DSV41_N_SHARED_EXPERTS, + LLM_KV_DSV41_NUM_EXPERTS_PER_TOK, + LLM_KV_DSV41_MAX_POSITION_EMBEDDINGS, + LLM_KV_DSV41_SLIDING_WINDOW, + LLM_KV_DSV41_INDEX_N_HEADS, + LLM_KV_DSV41_INDEX_HEAD_DIM, + LLM_KV_DSV41_INDEX_TOPK, + LLM_KV_DSV41_CANDIDATE_SOURCE_LAYER_ID, + LLM_KV_DSV41_CANDIDATE_TOPK_BLOCKS, + LLM_KV_DSV41_CANDIDATE_BLOCK_SIZE, + LLM_KV_DSV41_HC_MULT, + LLM_KV_DSV41_HC_SINKHORN_ITERS, + LLM_KV_DSV41_ROPE_THETA, + LLM_KV_DSV41_COMPRESS_ROPE_THETA, + LLM_KV_DSV41_RMS_NORM_EPS, + LLM_KV_DSV41_HC_EPS, + LLM_KV_DSV41_SWIGLU_LIMIT, + LLM_KV_DSV41_ROUTED_SCALING_FACTOR, + LLM_KV_DSV41_SCORING_FUNC, + LLM_KV_DSV41_HIDDEN_ACT, + LLM_KV_DSV41_TOPK_METHOD, + LLM_KV_DSV41_NORM_TOPK_PROB, + LLM_KV_DSV41_COMPRESS_RATIOS, + LLM_KV_DSV41_KV_SOURCE_LAYER_IDS, + LLM_KV_DSV41_INDEX_SOURCE_LAYER_IDS, + LLM_KV_DSV41_ROPE_SCALING_FACTOR, + LLM_KV_DSV41_ROPE_SCALING_BETA_FAST, + LLM_KV_DSV41_ROPE_SCALING_BETA_SLOW, + LLM_KV_DSV41_ROPE_SCALING_ORIG_CTX_LEN, + LLM_KV_DSV41_ENGRAM_ENCODING, + LLM_KV_DSV41_ENGRAM_LAYER_IDS, + LLM_KV_DSV41_ENGRAM_ROWS, + LLM_KV_DSV41_ENGRAM_COMPRESSED_VOCAB_SIZE, + LLM_KV_DSV41_ENGRAM_PAD_ID, + LLM_KV_DSV41_ENGRAM_TOKEN_MAP, + LLM_KV_DSV41_ENGRAM_PRIMES, + LLM_KV_DSV41_ENGRAM_MULTIPLIERS, + LLM_KV_HASH_LAYER_COUNT, LLM_KV_ROPE_DIMENSION_COUNT, @@ -610,6 +662,10 @@ enum llm_tensor { LLM_TENSOR_PLE_NORM_QUERY, // qwen4exp LLM_TENSOR_PLE_NORM_CONV, // qwen4exp LLM_TENSOR_PLE_CONV1D, // qwen4exp + LLM_TENSOR_ENGRAM_EMBD, + LLM_TENSOR_ENGRAM_Q_NORM, + LLM_TENSOR_ENGRAM_K_NORM, + LLM_TENSOR_ENGRAM_KV, LLM_TENSOR_ATTN_COMPRESSOR_WKV, LLM_TENSOR_ATTN_COMPRESSOR_WGATE, LLM_TENSOR_ATTN_COMPRESSOR_APE, diff --git a/src/llama-bounded-file.cpp b/src/llama-bounded-file.cpp new file mode 100644 index 000000000000..3ae6d5c31a18 --- /dev/null +++ b/src/llama-bounded-file.cpp @@ -0,0 +1,244 @@ +#include "llama-bounded-file.h" + +#include "llama-impl.h" + +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#include +#endif + +struct llama_bounded_file::buffer::impl { + void * ptr = nullptr; + size_t bytes = 0; + + ~impl() { + free(ptr); + } +}; + +llama_bounded_file::buffer::buffer() : pimpl(std::make_unique()) {} +llama_bounded_file::buffer::~buffer() = default; +llama_bounded_file::buffer::buffer(buffer && other) noexcept = default; +llama_bounded_file::buffer & llama_bounded_file::buffer::operator=(buffer && other) noexcept = default; + +bool llama_bounded_file::buffer::empty() const { + return pimpl->ptr == nullptr; +} + +struct llama_bounded_file::impl { + std::string fname; + int fd = -1; + uint64_t file_size = 0; + size_t block = 4096; + bool direct = false; + + impl(const std::string & fname, const params & p) : fname(fname) { +#if defined(_WIN32) + GGML_UNUSED(p); + throw std::runtime_error("llama_bounded_file: not supported on Windows"); +#else + int direct_error = 0; + if (p.direct_io) { +#if defined(O_DIRECT) + fd = open(fname.c_str(), O_RDONLY | O_DIRECT | O_CLOEXEC); + direct = fd >= 0; + direct_error = direct ? 0 : errno; +#else + fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); + direct_error = fd >= 0 ? ENOTSUP : errno; +#if defined(F_NOCACHE) + if (fd >= 0) { + if (fcntl(fd, F_NOCACHE, 1) == 0) { + direct = true; +#if defined(F_RDAHEAD) + if (fcntl(fd, F_RDAHEAD, 0) != 0) { + direct = false; + direct_error = errno; + } +#endif + } else { + direct_error = errno; + } + } +#endif +#endif + if (!direct && p.direct_io_required) { + const int saved = direct_error ? direct_error : ENOTSUP; + if (fd >= 0) { + close(fd); + fd = -1; + } + throw std::runtime_error(format("llama_bounded_file: uncached open of %s failed: %s", + fname.c_str(), strerror(saved))); + } + if (!direct) { + LLAMA_LOG_WARN("%s: uncached open of %s failed (%s); falling back to buffered reads\n", + __func__, fname.c_str(), strerror(direct_error ? direct_error : ENOTSUP)); + } + if (!direct && fd >= 0) { + close(fd); + fd = -1; + } + } + if (fd < 0) { + fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); + if (fd < 0) { + throw std::runtime_error(format("llama_bounded_file: failed to open %s: %s", + fname.c_str(), strerror(errno))); + } + } + + struct stat st = {}; + if (fstat(fd, &st) != 0) { + const int saved = errno; + close(fd); + fd = -1; + throw std::runtime_error(format("llama_bounded_file: fstat of %s failed: %s", + fname.c_str(), strerror(saved))); + } + if (!S_ISREG(st.st_mode) || st.st_size < 0) { + close(fd); + fd = -1; + throw std::runtime_error(format("llama_bounded_file: %s is not a regular file", fname.c_str())); + } + file_size = (uint64_t) st.st_size; +#endif + } + + ~impl() { +#if !defined(_WIN32) + if (fd >= 0) { + close(fd); + } +#endif + } + + size_t read_size(uint64_t offset, size_t size) const { + if (!direct || size == 0) { + return size; + } + const size_t prefix = (size_t) (offset & (block - 1)); + if (size > SIZE_MAX - prefix || size + prefix > SIZE_MAX - (block - 1)) { + throw std::overflow_error("llama_bounded_file: aligned read size overflow"); + } + return (size + prefix + block - 1) & ~(block - 1); + } + +#if !defined(_WIN32) + void pread_full(void * dst, size_t len, uint64_t offset, size_t need) const { + size_t done = 0; + while (done < len) { + const size_t remaining = len - done; + const ssize_t n = pread(fd, (uint8_t *) dst + done, remaining, (off_t) (offset + done)); + if (n < 0) { + if (errno == EINTR) { + continue; + } + throw std::runtime_error(format("llama_bounded_file: pread of %s failed at %llu: %s", + fname.c_str(), (unsigned long long) (offset + done), strerror(errno))); + } + if (n == 0) { + break; + } + done += (size_t) n; + if (direct && done >= need) { + return; + } + if (direct && (size_t) n < remaining) { + break; + } + } + if (done < need) { + throw std::runtime_error(format("llama_bounded_file: short read in %s: %zu of %zu bytes at %llu", + fname.c_str(), done, need, (unsigned long long) offset)); + } + } +#endif +}; + +llama_bounded_file::llama_bounded_file(const std::string & fname, const params & p) + : pimpl(std::make_unique(fname, p)) {} + +llama_bounded_file::~llama_bounded_file() = default; + +llama_bounded_file::buffer llama_bounded_file::make_buffer(size_t size) const { + buffer result; + if (!pimpl->direct || size == 0) { + return result; + } +#if defined(_WIN32) + GGML_UNUSED(size); + throw std::runtime_error("llama_bounded_file: not supported on Windows"); +#else + const size_t bytes = pimpl->read_size(pimpl->block - 1, size); + void * ptr = nullptr; + const int err = posix_memalign(&ptr, pimpl->block, bytes); + if (err != 0) { + throw std::runtime_error(format("llama_bounded_file: posix_memalign failed: %s", strerror(err))); + } + result.pimpl->ptr = ptr; + result.pimpl->bytes = bytes; + return result; +#endif +} + +void llama_bounded_file::read(uint64_t offset, void * dst, size_t size, buffer & scratch) const { +#if defined(_WIN32) + GGML_UNUSED(offset); + GGML_UNUSED(dst); + GGML_UNUSED(size); + GGML_UNUSED(scratch); + throw std::runtime_error("llama_bounded_file: not supported on Windows"); +#else + if (size == 0) { + return; + } + if (dst == nullptr || offset > pimpl->file_size || size > pimpl->file_size - offset || + offset > (uint64_t) INT64_MAX || size > (uint64_t) INT64_MAX - offset) { + throw std::invalid_argument(format("llama_bounded_file: invalid read of %zu bytes at %llu in %s", + size, (unsigned long long) offset, pimpl->fname.c_str())); + } + if (!pimpl->direct) { + pimpl->pread_full(dst, size, offset, size); + return; + } + + const uint64_t aligned_offset = offset & ~(uint64_t) (pimpl->block - 1); + const size_t prefix = (size_t) (offset - aligned_offset); + const size_t aligned_size = pimpl->read_size(offset, size); + if (aligned_size > (uint64_t) INT64_MAX - aligned_offset) { + throw std::overflow_error("llama_bounded_file: aligned read offset overflow"); + } + if (scratch.pimpl->ptr == nullptr || scratch.pimpl->bytes < aligned_size) { + throw std::invalid_argument("llama_bounded_file: aligned scratch buffer is too small"); + } + pimpl->pread_full(scratch.pimpl->ptr, aligned_size, aligned_offset, prefix + size); + memcpy(dst, (uint8_t *) scratch.pimpl->ptr + prefix, size); +#endif +} + +uint64_t llama_bounded_file::size() const { + return pimpl->file_size; +} + +size_t llama_bounded_file::alignment() const { + return pimpl->block; +} + +size_t llama_bounded_file::read_size(uint64_t offset, size_t size) const { + return pimpl->read_size(offset, size); +} + +bool llama_bounded_file::direct_io() const { + return pimpl->direct; +} + +const std::string & llama_bounded_file::name() const { + return pimpl->fname; +} diff --git a/src/llama-bounded-file.h b/src/llama-bounded-file.h new file mode 100644 index 000000000000..3af763580a36 --- /dev/null +++ b/src/llama-bounded-file.h @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#include +#include + +struct llama_bounded_file { + struct params { + bool direct_io = true; + bool direct_io_required = false; // fail instead of using the page cache + }; + + struct buffer { + buffer(); + ~buffer(); + buffer(buffer && other) noexcept; + buffer & operator=(buffer && other) noexcept; + + buffer(const buffer &) = delete; + buffer & operator=(const buffer &) = delete; + + bool empty() const; + + struct impl; + std::unique_ptr pimpl; + }; + + llama_bounded_file(const std::string & fname, const params & p); + ~llama_bounded_file(); + + buffer make_buffer(size_t size) const; + // Read only the requested extent. Direct I/O may read its containing aligned blocks into scratch. + void read(uint64_t offset, void * dst, size_t size, buffer & scratch) const; + + uint64_t size() const; + size_t alignment() const; + size_t read_size(uint64_t offset, size_t size) const; + bool direct_io() const; + const std::string & name() const; + + struct impl; + std::unique_ptr pimpl; +}; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0d289ba03d6b..73e4d7cd3477 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -501,11 +501,15 @@ llama_context::llama_context( sampling.token_ids_full_vocab[i] = i; } } + + model.acquire_runtime_context(); } llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); + model.release_runtime_work(); + model.release_runtime_context(); // when training, ggml_opt allocates extra buffers through the scheduler, so the sizes no longer match the expectation if (!model.hparams.no_alloc && !opt_ctx) { @@ -707,18 +711,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); @@ -1371,11 +1380,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,11 +1474,23 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1); if (status != GGML_STATUS_SUCCESS) { + model.release_runtime_work_after_sync(sched.get()); LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status); ret = status; return nullptr; } + if (model.requires_synchronous_graph()) { + synchronize(); + const std::string error = model.consume_runtime_error(); + if (!error.empty()) { + model.release_runtime_work(); + LLAMA_LOG_ERROR("%s: model runtime failed: %s\n", __func__, error.c_str()); + ret = GGML_STATUS_FAILED; + return nullptr; + } + } + transaction_guard.commit(); ret = GGML_STATUS_SUCCESS; return res; @@ -2371,6 +2415,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 || @@ -3753,7 +3798,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-engram.cpp b/src/llama-dsv41-engram.cpp new file mode 100644 index 000000000000..b2273a5236f2 --- /dev/null +++ b/src/llama-dsv41-engram.cpp @@ -0,0 +1,400 @@ +#include "llama-dsv41-engram.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include + +static llama_dsv41_engram_sequence_state dsv41_engram_initial_state() { + llama_dsv41_engram_sequence_state state; + state.history.reset(); + return state; +} + +void llama_dsv41_validate_engram_extent(const llama_dsv41_engram_extent & extent) { + if (extent.fname.empty() || extent.rows == 0 || extent.columns != LLAMA_ENGRAM_ROW_BYTES || + extent.row_count != extent.rows || extent.type != GGML_TYPE_I8) { + throw std::invalid_argument("DeepSeek V4.1 Engram tensor must be I8 [264, rows]"); + } + const uint64_t bytes = (uint64_t) extent.rows*LLAMA_ENGRAM_ROW_BYTES; + if (extent.offset > (uint64_t) INT64_MAX || bytes > (uint64_t) INT64_MAX - extent.offset) { + throw std::invalid_argument("DeepSeek V4.1 Engram tensor extent overflows"); + } +} + +struct llama_dsv41_engram_transaction::impl { + const llama_dsv41_engram_runtime * owner = nullptr; + uint64_t generation = 0; + bool active = false; + size_t count = 0; + std::vector ids; + std::array, LLAMA_ENGRAM_LAYERS> decoded; + std::vector mask; + std::map next; +}; + +llama_dsv41_engram_transaction::llama_dsv41_engram_transaction() : pimpl(std::make_unique()) {} +llama_dsv41_engram_transaction::~llama_dsv41_engram_transaction() = default; +llama_dsv41_engram_transaction::llama_dsv41_engram_transaction(llama_dsv41_engram_transaction && other) noexcept = default; +llama_dsv41_engram_transaction & llama_dsv41_engram_transaction::operator=(llama_dsv41_engram_transaction && other) noexcept = default; + +size_t llama_dsv41_engram_transaction::token_count() const { + return pimpl->count; +} + +const uint32_t * llama_dsv41_engram_transaction::row_ids(uint32_t layer) const { + if (layer >= LLAMA_ENGRAM_LAYERS || pimpl->count == 0) { + return nullptr; + } + return pimpl->ids.data() + layer*LLAMA_ENGRAM_COLS; +} + +const float * llama_dsv41_engram_transaction::rows(uint32_t layer) const { + if (layer >= LLAMA_ENGRAM_LAYERS || pimpl->decoded[layer].empty()) { + return nullptr; + } + return pimpl->decoded[layer].data(); +} + +const uint8_t * llama_dsv41_engram_transaction::text_mask() const { + return pimpl->mask.empty() ? nullptr : pimpl->mask.data(); +} + +void llama_dsv41_engram_transaction::upload_layer( + uint32_t layer, + size_t token_offset, + size_t token_count, + ggml_tensor * rows_input, + ggml_tensor * text_select_input) const { + if (!pimpl->active) { + throw std::invalid_argument("DeepSeek V4.1 Engram transaction is not active"); + } + if (layer >= LLAMA_ENGRAM_LAYERS || token_offset > pimpl->count || + token_count > pimpl->count - token_offset) { + throw std::invalid_argument("DeepSeek V4.1 Engram upload range is invalid"); + } + const int64_t row_width = LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM; + if (rows_input == nullptr || text_select_input == nullptr || + rows_input->type != GGML_TYPE_F32 || rows_input->ne[0] != row_width || + ggml_nelements(rows_input) != row_width*(int64_t) token_count || + text_select_input->type != GGML_TYPE_I32 || + ggml_nelements(text_select_input) != (int64_t) token_count) { + throw std::invalid_argument("DeepSeek V4.1 Engram input tensor shape mismatch"); + } + + const size_t row_offset = token_offset*row_width; + ggml_backend_tensor_set( + rows_input, + pimpl->decoded[layer].data() + row_offset, + 0, + token_count*row_width*sizeof(float)); + std::vector select(token_count); + for (size_t i = 0; i < token_count; ++i) { + select[i] = pimpl->mask[token_offset + i] != 0 ? (int32_t) (token_count + i) : (int32_t) i; + } + ggml_backend_tensor_set( + text_select_input, + select.data(), + 0, + token_count*sizeof(int32_t)); +} + +struct llama_dsv41_engram_runtime::impl { + llama_engram_hasher hasher; + std::array, LLAMA_ENGRAM_LAYERS> tables; + std::map sequences; + size_t max_tokens; + uint64_t generation = 0; + + impl( + llama_engram_layout layout, + const std::array & extents, + size_t max_tokens) : + hasher(std::move(layout)), + max_tokens(max_tokens) { + if (max_tokens == 0 || max_tokens > (size_t) INT32_MAX/2) { + throw std::invalid_argument("DeepSeek V4.1 Engram token bound is invalid"); + } + for (size_t i = 0; i < LLAMA_ENGRAM_LAYERS; ++i) { + llama_dsv41_validate_engram_extent(extents[i]); + if (extents[i].rows != hasher.layout().rows[i]) { + throw std::invalid_argument("DeepSeek V4.1 Engram tensor rows do not match metadata"); + } + tables[i] = std::make_unique( + extents[i].fname, extents[i].offset, extents[i].rows); + } + } +}; + +llama_dsv41_engram_runtime::llama_dsv41_engram_runtime( + llama_engram_layout layout, + const std::array & extents, + size_t max_tokens) : + pimpl(std::make_unique(std::move(layout), extents, max_tokens)) {} + +llama_dsv41_engram_runtime::~llama_dsv41_engram_runtime() = default; + +llama_dsv41_engram_transaction llama_dsv41_engram_runtime::prepare( + const std::vector & tokens) { + if (tokens.size() > pimpl->max_tokens) { + throw std::invalid_argument("DeepSeek V4.1 Engram batch exceeds the configured bound"); + } + + llama_dsv41_engram_transaction result; + result.pimpl->owner = this; + result.pimpl->generation = pimpl->generation; + result.pimpl->count = tokens.size(); + result.pimpl->ids.resize(tokens.size()*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS); + result.pimpl->mask.resize(tokens.size()); + result.pimpl->next = pimpl->sequences; + + for (size_t i = 0; i < tokens.size(); ++i) { + const llama_dsv41_engram_token & token = tokens[i]; + if (token.seq_ids.empty()) { + throw std::invalid_argument("DeepSeek V4.1 Engram token has no sequence"); + } + + uint32_t expected[LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS] = {}; + bool have_expected = false; + std::set seen; + for (llama_seq_id seq_id : token.seq_ids) { + if (seq_id < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + if (!seen.insert(seq_id).second) { + throw std::invalid_argument("DeepSeek V4.1 Engram token repeats a sequence ID"); + } + auto inserted = result.pimpl->next.emplace(seq_id, dsv41_engram_initial_state()); + llama_dsv41_engram_sequence_state & state = inserted.first->second; + if (token.pos != state.pos + 1) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence position is not contiguous"); + } + + uint32_t current[LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS] = {}; + pimpl->hasher.hash(state.history, &token.token, &token.text, 1, current); + state.pos = token.pos; + if (have_expected && std::memcmp(expected, current, sizeof(expected)) != 0) { + throw std::invalid_argument("DeepSeek V4.1 coupled sequence histories differ"); + } + std::memcpy(expected, current, sizeof(expected)); + have_expected = true; + } + + std::memcpy( + result.pimpl->ids.data() + i*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + expected, + sizeof(expected)); + result.pimpl->mask[i] = token.text != 0; + } + + const size_t row_values = tokens.size()*LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + result.pimpl->decoded[layer].resize(row_values); + pimpl->tables[layer]->read_batch( + result.pimpl->ids.data() + layer*LLAMA_ENGRAM_COLS, + tokens.size(), + LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + result.pimpl->decoded[layer].data()); + } + result.pimpl->active = true; + return result; +} + +void llama_dsv41_engram_runtime::commit(llama_dsv41_engram_transaction & transaction) { + if (!transaction.pimpl->active || transaction.pimpl->owner != this) { + throw std::invalid_argument("DeepSeek V4.1 Engram transaction is not active"); + } + if (transaction.pimpl->generation != pimpl->generation) { + throw std::runtime_error("DeepSeek V4.1 Engram transaction is stale"); + } + auto next = transaction.pimpl->next; + pimpl->sequences.swap(next); + ++pimpl->generation; + transaction.pimpl->active = false; +} + +void llama_dsv41_engram_runtime::rollback(llama_dsv41_engram_transaction & transaction) { + if (transaction.pimpl->owner != this) { + throw std::invalid_argument("DeepSeek V4.1 Engram transaction belongs to another runtime"); + } + transaction.pimpl->active = false; +} + +void llama_dsv41_engram_runtime::seq_reset(llama_seq_id seq_id) { + if (seq_id < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + pimpl->sequences[seq_id] = dsv41_engram_initial_state(); + ++pimpl->generation; +} + +void llama_dsv41_engram_runtime::seq_copy(llama_seq_id seq_id_src, llama_seq_id seq_id_dst) { + if (seq_id_src < 0 || seq_id_dst < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + const auto it = pimpl->sequences.find(seq_id_src); + pimpl->sequences[seq_id_dst] = it == pimpl->sequences.end() ? + dsv41_engram_initial_state() : it->second; + ++pimpl->generation; +} + +void llama_dsv41_engram_runtime::seq_remove(llama_seq_id seq_id) { + if (seq_id < 0) { + throw std::invalid_argument("DeepSeek V4.1 Engram sequence ID is negative"); + } + pimpl->sequences.erase(seq_id); + ++pimpl->generation; +} + +llama_dsv41_engram_snapshot llama_dsv41_engram_runtime::checkpoint() const { + return { pimpl->sequences }; +} + +void llama_dsv41_engram_runtime::restore(const llama_dsv41_engram_snapshot & snapshot) { + auto restored = snapshot.sequences; + for (auto & item : restored) { + if (item.first < 0 || item.second.pos < -1) { + throw std::invalid_argument("DeepSeek V4.1 Engram snapshot is invalid"); + } + pimpl->hasher.hash(item.second.history, nullptr, nullptr, 0, nullptr); + } + pimpl->sequences.swap(restored); + ++pimpl->generation; +} + +llama_dsv41_engram_sequence_state llama_dsv41_engram_runtime::sequence(llama_seq_id seq_id) const { + const auto it = pimpl->sequences.find(seq_id); + return it == pimpl->sequences.end() ? dsv41_engram_initial_state() : it->second; +} + +size_t llama_dsv41_engram_runtime::max_tokens() const { + return pimpl->max_tokens; +} + +static ggml_tensor * dsv41_bf16_f32(ggml_context * ctx, ggml_tensor * tensor) { + return ggml_cast(ctx, ggml_cast(ctx, tensor, GGML_TYPE_BF16), GGML_TYPE_F32); +} + +static void dsv41_engram_gate_f32( + ggml_tensor * dst, + const ggml_tensor * src, + int ith, + int nth, + void *) { + GGML_ASSERT(dst->type == GGML_TYPE_F32 && src->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src)); + const float * input = static_cast(src->data); + float * output = static_cast(dst->data); + const int64_t count = ggml_nelements(src); + for (int64_t i = ith; i < count; i += nth) { + const float signed_root = std::copysign(std::sqrt(std::max(std::abs(input[i]), 1.0e-6f)), input[i]); + output[i] = 1.0f/(1.0f + std::exp(-signed_root)); + } +} + +ggml_tensor * llama_dsv41_build_engram_gate( + ggml_context * ctx, + ggml_tensor * dot, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || dot == nullptr || sched == nullptr || backend_cpu == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 Engram gate input is null"); + } + if (!ggml_backend_is_cpu(backend_cpu)) { + throw std::invalid_argument("DeepSeek V4.1 Engram gate backend is not local CPU"); + } + bool found = false; + for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) { + found = found || ggml_backend_sched_get_backend(sched, i) == backend_cpu; + } + if (!found) { + throw std::invalid_argument("DeepSeek V4.1 Engram gate CPU backend is not in the scheduler"); + } + ggml_tensor * gate = ggml_map_custom1(ctx, dot, dsv41_engram_gate_f32, GGML_N_TASKS_MAX, nullptr); + ggml_backend_sched_set_tensor_backend(sched, gate, backend_cpu); + return gate; +} + +ggml_tensor * llama_dsv41_build_engram_add( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * projected, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_select, + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || residual == nullptr || projected == nullptr || q_norm == nullptr || k_norm == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 Engram graph input is null"); + } + + const int64_t width = residual->ne[0]; + const int64_t streams = residual->ne[1]; + const int64_t tokens = residual->ne[2]; + if (width <= 0 || streams != 4 || tokens <= 0 || tokens > INT32_MAX/2 || + projected->ne[0] != 5*width || projected->ne[1] != tokens || + q_norm->ne[0] != width || q_norm->ne[1] != streams || + k_norm->ne[0] != width || k_norm->ne[1] != streams || + (text_select != nullptr && + (text_select->type != GGML_TYPE_I32 || ggml_nelements(text_select) != tokens))) { + throw std::invalid_argument("DeepSeek V4.1 Engram graph shape mismatch"); + } + + projected = dsv41_bf16_f32(ctx, projected); + ggml_tensor * value = ggml_view_2d( + ctx, projected, width, tokens, projected->nb[1], 4*projected->nb[0]*width); + ggml_tensor * result = nullptr; + for (int64_t stream = 0; stream < streams; ++stream) { + ggml_tensor * hidden = ggml_view_2d( + ctx, residual, width, tokens, residual->nb[2], stream*residual->nb[1]); + ggml_tensor * key = ggml_view_2d( + ctx, projected, width, tokens, projected->nb[1], stream*projected->nb[0]*width); + ggml_tensor * qw = ggml_view_1d(ctx, q_norm, width, stream*q_norm->nb[1]); + ggml_tensor * kw = ggml_view_1d(ctx, k_norm, width, stream*k_norm->nb[1]); + + ggml_tensor * hidden_norm = ggml_rms_norm(ctx, hidden, rms_eps); + ggml_tensor * key_norm = ggml_rms_norm(ctx, key, rms_eps); + ggml_tensor * dot = ggml_mul(ctx, hidden_norm, qw); + dot = ggml_mul(ctx, dot, kw); + dot = ggml_mul(ctx, dot, key_norm); + dot = ggml_scale(ctx, ggml_sum_rows(ctx, dot), 1.0f/std::sqrt((float) width)); + + ggml_tensor * gate = llama_dsv41_build_engram_gate(ctx, dot, sched, backend_cpu); + ggml_tensor * updated = dsv41_bf16_f32(ctx, ggml_add(ctx, hidden, ggml_mul(ctx, value, gate))); + if (text_select != nullptr) { + updated = ggml_get_rows(ctx, ggml_concat(ctx, hidden, updated, 1), text_select); + } + updated = ggml_reshape_3d(ctx, updated, width, 1, tokens); + result = result == nullptr ? updated : ggml_concat(ctx, result, updated, 1); + } + return result; +} + +ggml_tensor * llama_dsv41_build_engram( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * rows, + ggml_tensor * engram_kv, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_select, + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || residual == nullptr || rows == nullptr || engram_kv == nullptr || + rows->ne[0] != LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM || + engram_kv->ne[0] != rows->ne[0] || + engram_kv->ne[1] != 5*residual->ne[0]) { + throw std::invalid_argument("DeepSeek V4.1 Engram projection shape mismatch"); + } + ggml_tensor * projected = ggml_mul_mat(ctx, engram_kv, rows); + return llama_dsv41_build_engram_add( + ctx, residual, projected, q_norm, k_norm, text_select, rms_eps, sched, backend_cpu); +} diff --git a/src/llama-dsv41-engram.h b/src/llama-dsv41-engram.h new file mode 100644 index 000000000000..2edca59c12d0 --- /dev/null +++ b/src/llama-dsv41-engram.h @@ -0,0 +1,123 @@ +#pragma once + +#include "ggml-backend.h" +#include "llama-engram.h" +#include "llama.h" + +#include +#include +#include +#include +#include +#include +#include + +struct ggml_context; +struct ggml_tensor; + +struct llama_dsv41_engram_extent { + std::string fname; + uint64_t offset = 0; + uint32_t rows = 0; + int64_t columns = 0; + int64_t row_count = 0; + int32_t type = 0; +}; + +void llama_dsv41_validate_engram_extent(const llama_dsv41_engram_extent & extent); + +struct llama_dsv41_engram_token { + int32_t token = -1; + llama_pos pos = -1; + std::vector seq_ids; + uint8_t text = 1; +}; + +struct llama_dsv41_engram_sequence_state { + llama_engram_history history = {}; + llama_pos pos = -1; +}; + +struct llama_dsv41_engram_snapshot { + std::map sequences; +}; + +struct llama_dsv41_engram_transaction { + llama_dsv41_engram_transaction(); + ~llama_dsv41_engram_transaction(); + llama_dsv41_engram_transaction(llama_dsv41_engram_transaction && other) noexcept; + llama_dsv41_engram_transaction & operator=(llama_dsv41_engram_transaction && other) noexcept; + + llama_dsv41_engram_transaction(const llama_dsv41_engram_transaction &) = delete; + llama_dsv41_engram_transaction & operator=(const llama_dsv41_engram_transaction &) = delete; + + size_t token_count() const; + const uint32_t * row_ids(uint32_t layer) const; + const float * rows(uint32_t layer) const; + const uint8_t * text_mask() const; + void upload_layer( + uint32_t layer, + size_t token_offset, + size_t token_count, + ggml_tensor * rows_input, + ggml_tensor * text_select_input) const; + + struct impl; + std::unique_ptr pimpl; +}; + +class llama_dsv41_engram_runtime { +public: + llama_dsv41_engram_runtime( + llama_engram_layout layout, + const std::array & extents, + size_t max_tokens); + ~llama_dsv41_engram_runtime(); + + llama_dsv41_engram_transaction prepare(const std::vector & tokens); + void commit(llama_dsv41_engram_transaction & transaction); + void rollback(llama_dsv41_engram_transaction & transaction); + + void seq_reset(llama_seq_id seq_id); + void seq_copy(llama_seq_id seq_id_src, llama_seq_id seq_id_dst); + void seq_remove(llama_seq_id seq_id); + llama_dsv41_engram_snapshot checkpoint() const; + void restore(const llama_dsv41_engram_snapshot & snapshot); + llama_dsv41_engram_sequence_state sequence(llama_seq_id seq_id) const; + + size_t max_tokens() const; + +private: + struct impl; + std::unique_ptr pimpl; +}; + +// text_select row i keeps the original residual; row tokens+i selects the BF16-updated residual. +ggml_tensor * llama_dsv41_build_engram_add( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * projected, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_select, + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); + +ggml_tensor * llama_dsv41_build_engram_gate( + ggml_context * ctx, + ggml_tensor * dot, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); + +ggml_tensor * llama_dsv41_build_engram( + ggml_context * ctx, + ggml_tensor * residual, + ggml_tensor * rows, + ggml_tensor * engram_kv, + ggml_tensor * q_norm, + ggml_tensor * k_norm, + ggml_tensor * text_select, + float rms_eps, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); diff --git a/src/llama-dsv41-expert.cpp b/src/llama-dsv41-expert.cpp new file mode 100644 index 000000000000..3d82f29b4739 --- /dev/null +++ b/src/llama-dsv41-expert.cpp @@ -0,0 +1,552 @@ +#include "llama-dsv41-expert.h" + +#include "llama-dsv41.h" +#include "llama-impl.h" + +#include "ggml-cpp.h" + +#include +#include +#include +#include +#include +#include +#include + +// The remap node follows ggml-org/llama.cpp#25294 commit 4260e4608. +// Lease publication and release follow ggml-org/llama.cpp#27861 commit bccbacdb8. + +namespace { + +size_t checked_add(size_t a, size_t b, const char * message) { + if (b > std::numeric_limits::max() - a) { + throw std::runtime_error(message); + } + return a + b; +} + +size_t checked_mul(size_t a, size_t b, const char * message) { + if (a != 0 && b > std::numeric_limits::max()/a) { + throw std::runtime_error(message); + } + return a*b; +} + +size_t projection_index(llama_expert_projection projection) { + if (projection < LLAMA_EXPERT_PROJECTION_GATE || projection > LLAMA_EXPERT_PROJECTION_DOWN) { + throw std::invalid_argument("DeepSeek V4.1 expert projection is invalid"); + } + return static_cast(projection); +} + +void require_local_cpu(ggml_backend_sched_t sched, ggml_backend_t backend_cpu) { + if (sched == nullptr || backend_cpu == nullptr || !ggml_backend_is_cpu(backend_cpu)) { + throw std::invalid_argument("DeepSeek V4.1 expert callback requires a local CPU backend"); + } + for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) { + if (ggml_backend_sched_get_backend(sched, i) == backend_cpu) { + return; + } + } + throw std::invalid_argument("DeepSeek V4.1 expert CPU backend is not in the scheduler"); +} + +struct dsv41_expert_callback_state { + llama_dsv41_expert_runtime * runtime = nullptr; + int32_t layer = -1; +}; + +} + +struct llama_dsv41_expert_runtime::impl { + struct logical_slot { + int32_t expert_id = -1; + uint32_t pins = 0; + uint64_t last_use = 0; + }; + + struct layer_state { + std::array tensors = {}; + std::vector slots; + std::unique_ptr lease; + std::vector pinned_slots; + std::vector active_ids; + std::vector active_remap; + }; + + llama_dsv41_expert_runtime_params params; + std::unique_ptr store; + std::vector layers; + std::vector callbacks; + std::vector> contexts; + std::vector buffers; + upload_fn upload; + publish_fn before_publish; + size_t bytes_cache = 0; + size_t bytes_staging = 0; + uint64_t use_clock = 0; + std::string error; + bool context_active = false; + mutable std::mutex mutex; + + impl( + llama_dsv41_expert_runtime * owner, + std::vector tensors, + const llama_dsv41_expert_runtime_params & params, + buft_selector select_buft, + upload_fn upload, + publish_fn before_publish) + : params(params), upload(std::move(upload)), before_publish(std::move(before_publish)) { + if (params.cache_bytes == 0 || params.cache_slots == 0) { + throw std::runtime_error("DeepSeek V4.1 expert cache byte and slot capacity must be non-zero"); + } + if (!select_buft) { + throw std::invalid_argument("DeepSeek V4.1 expert cache buffer selector is empty"); + } + if (tensors.size() != LLAMA_DSV41_N_LAYER*3) { + throw std::runtime_error("DeepSeek V4.1 must register 40 gate/up/down routed tensor sets"); + } + + layers.resize(LLAMA_DSV41_N_LAYER); + callbacks.resize(LLAMA_DSV41_N_LAYER); + std::vector layer_plane_bytes(LLAMA_DSV41_N_LAYER, 0); + for (const auto & tensor : tensors) { + llama_expert_store_validate_tensor(tensor); + if (tensor.layer < 0 || tensor.layer >= (int32_t) LLAMA_DSV41_N_LAYER) { + throw std::runtime_error("DeepSeek V4.1 routed tensor layer is invalid"); + } + if (tensor.ne[2] != LLAMA_DSV41_N_EXPERT) { + throw std::runtime_error("DeepSeek V4.1 routed tensor expert count mismatch"); + } + layer_plane_bytes[tensor.layer] = checked_add( + layer_plane_bytes[tensor.layer], + tensor.nb[2], + "DeepSeek V4.1 expert plane byte count overflow"); + bytes_cache = checked_add( + bytes_cache, + checked_mul(tensor.nb[2], params.cache_slots, "DeepSeek V4.1 expert cache byte count overflow"), + "DeepSeek V4.1 expert cache byte count overflow"); + + ggml_backend_buffer_type_t buft = select_buft(tensor); + if (buft == nullptr) { + throw std::runtime_error("DeepSeek V4.1 expert cache has no backend buffer type"); + } + ggml_context * ctx = nullptr; + for (auto & item : contexts) { + if (item.first == buft) { + ctx = item.second.get(); + break; + } + } + if (ctx == nullptr) { + ggml_init_params ctx_params = { + /*.mem_size =*/ ggml_tensor_overhead()*(LLAMA_DSV41_N_LAYER*3 + 1), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx = ggml_init(ctx_params); + if (ctx == nullptr) { + throw std::runtime_error("DeepSeek V4.1 failed to create expert cache tensor context"); + } + contexts.emplace_back(buft, ctx); + } + + ggml_tensor * cache = ggml_new_tensor_3d( + ctx, tensor.type, tensor.ne[0], tensor.ne[1], params.cache_slots); + ggml_format_name(cache, "%s.cache", tensor.name.c_str()); + layers[tensor.layer].tensors[projection_index(tensor.projection)] = cache; + } + if (bytes_cache > params.cache_bytes) { + throw std::runtime_error(format( + "DeepSeek V4.1 expert cache requires %zu bytes for %zu slots per layer, configured %zu", + bytes_cache, params.cache_slots, params.cache_bytes)); + } + + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + for (ggml_tensor * tensor : layers[il].tensors) { + if (tensor == nullptr) { + throw std::runtime_error(format("DeepSeek V4.1 layer %d is missing an expert cache tensor", il)); + } + } + layers[il].slots.resize(params.cache_slots); + callbacks[il] = { owner, il }; + bytes_staging = std::max( + bytes_staging, + checked_mul(layer_plane_bytes[il], params.cache_slots, "DeepSeek V4.1 expert staging byte count overflow")); + } + + for (auto & item : contexts) { + ggml_backend_buffer_t buffer = nullptr; + if (params.no_alloc) { + buffer = ggml_backend_buft_alloc_buffer(item.first, 0); + for (ggml_tensor * tensor = ggml_get_first_tensor(item.second.get()); + tensor != nullptr; + tensor = ggml_get_next_tensor(item.second.get(), tensor)) { + tensor->buffer = buffer; + } + } else { + buffer = ggml_backend_alloc_ctx_tensors_from_buft(item.second.get(), item.first); + } + if (buffer == nullptr) { + throw std::runtime_error(format( + "DeepSeek V4.1 failed to allocate %s expert cache buffer", + ggml_backend_buft_name(item.first))); + } + ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + buffers.emplace_back(buffer); + } + + if (!params.no_alloc) { + llama_expert_store_params store_params; + store_params.cache_bytes = bytes_staging; + store_params.cache_slots = checked_mul(params.cache_slots, 3, "DeepSeek V4.1 expert staging slot count overflow"); + store_params.direct_io = params.direct_io; + store_params.allow_buffered_io = params.allow_buffered_io; + store = std::make_unique(std::move(tensors), store_params); + } + + if (!this->upload) { + this->upload = [](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + ggml_backend_tensor_set(tensor, data, offset, size); + }; + } + } + + const llama_expert_store::payload & payload( + const std::vector & payloads, + int32_t layer, + llama_expert_projection projection, + int32_t expert_id) const { + for (const auto & payload : payloads) { + if (payload.layer == layer && payload.projection == projection && payload.expert_id == expert_id) { + return payload; + } + } + throw std::runtime_error("DeepSeek V4.1 expert store did not return a requested payload"); + } + + std::vector remap(int32_t layer, const std::vector & expert_ids) { + std::lock_guard lock(mutex); + if (store == nullptr) { + throw std::runtime_error("DeepSeek V4.1 expert cache is metadata-only"); + } + if (layer < 0 || layer >= (int32_t) layers.size()) { + throw std::invalid_argument("DeepSeek V4.1 expert layer is invalid"); + } + layer_state & state = layers[layer]; + if (state.lease) { + if (state.active_ids == expert_ids) { + return state.active_remap; + } + throw std::runtime_error(format("DeepSeek V4.1 expert layer %d still has an in-flight lease", layer)); + } + + std::vector unique = expert_ids; + for (int32_t expert_id : unique) { + if (expert_id < 0 || expert_id >= (int32_t) LLAMA_DSV41_N_EXPERT) { + throw std::runtime_error("DeepSeek V4.1 selected expert ID is out of range"); + } + } + std::sort(unique.begin(), unique.end()); + unique.erase(std::unique(unique.begin(), unique.end()), unique.end()); + if (unique.size() > params.cache_slots) { + throw std::runtime_error(format( + "DeepSeek V4.1 selected expert union has %zu entries, cache has %zu slots", + unique.size(), params.cache_slots)); + } + + std::map resident; + std::vector empty; + std::vector victims; + for (uint32_t slot = 0; slot < state.slots.size(); ++slot) { + const logical_slot & entry = state.slots[slot]; + if (entry.expert_id >= 0) { + resident.emplace(entry.expert_id, slot); + if (entry.pins == 0 && !std::binary_search(unique.begin(), unique.end(), entry.expert_id)) { + victims.push_back(slot); + } + } else { + empty.push_back(slot); + } + } + std::sort(victims.begin(), victims.end(), [&](uint32_t a, uint32_t b) { + if (state.slots[a].last_use != state.slots[b].last_use) { + return state.slots[a].last_use < state.slots[b].last_use; + } + return a < b; + }); + + std::vector misses; + for (int32_t expert_id : unique) { + if (resident.count(expert_id) == 0) { + misses.push_back(expert_id); + } + } + if (misses.size() > empty.size() + victims.size()) { + throw std::runtime_error("DeepSeek V4.1 expert cache capacity is pinned"); + } + + std::vector targets = empty; + targets.insert(targets.end(), victims.begin(), victims.end()); + targets.resize(misses.size()); + std::sort(targets.begin(), targets.end()); + + auto lease = std::make_unique(); + if (!misses.empty()) { + *lease = store->acquire({ + { layer, LLAMA_EXPERT_PROJECTION_GATE, misses }, + { layer, LLAMA_EXPERT_PROJECTION_UP, misses }, + { layer, LLAMA_EXPERT_PROJECTION_DOWN, misses }, + }); + const auto payloads = lease->payloads(); + try { + for (size_t i = 0; i < misses.size(); ++i) { + const uint32_t slot = targets[i]; + for (llama_expert_projection projection : { + LLAMA_EXPERT_PROJECTION_GATE, + LLAMA_EXPERT_PROJECTION_UP, + LLAMA_EXPERT_PROJECTION_DOWN }) { + const auto & item = payload(payloads, layer, projection, misses[i]); + ggml_tensor * tensor = state.tensors[projection_index(projection)]; + upload(tensor, slot*tensor->nb[2], item.data, item.size); + } + } + } catch (...) { + for (uint32_t slot : targets) { + state.slots[slot] = {}; + } + throw; + } + + for (size_t i = 0; i < misses.size(); ++i) { + const uint32_t slot = targets[i]; + state.slots[slot].expert_id = misses[i]; + resident[misses[i]] = slot; + } + } + + std::vector pinned_slots; + pinned_slots.reserve(unique.size()); + for (int32_t expert_id : unique) { + const uint32_t slot = resident.at(expert_id); + pinned_slots.push_back(slot); + } + + std::vector result; + result.reserve(expert_ids.size()); + for (int32_t expert_id : expert_ids) { + result.push_back((int32_t) resident.at(expert_id)); + } + std::vector active_ids = expert_ids; + std::vector active_remap = result; + + if (before_publish) { + before_publish(); + } + for (uint32_t slot : pinned_slots) { + logical_slot & entry = state.slots[slot]; + entry.last_use = ++use_clock; + entry.pins++; + } + state.pinned_slots = std::move(pinned_slots); + state.lease = std::move(lease); + state.active_ids = std::move(active_ids); + state.active_remap = std::move(active_remap); + return result; + } + + void release(int32_t layer) { + std::lock_guard lock(mutex); + if (layer < 0 || layer >= (int32_t) layers.size()) { + return; + } + layer_state & state = layers[layer]; + for (uint32_t slot : state.pinned_slots) { + if (state.slots[slot].pins == 0) { + GGML_ABORT("DeepSeek V4.1 expert slot pin underflow"); + } + state.slots[slot].pins--; + } + state.pinned_slots.clear(); + state.lease.reset(); + state.active_ids.clear(); + state.active_remap.clear(); + } +}; + +llama_dsv41_expert_runtime::llama_dsv41_expert_runtime( + std::vector tensors, + const llama_dsv41_expert_runtime_params & params, + buft_selector select_buft, + upload_fn upload, + publish_fn before_publish) + : pimpl(std::make_unique( + this, std::move(tensors), params, std::move(select_buft), std::move(upload), std::move(before_publish))) { +} + +llama_dsv41_expert_runtime::~llama_dsv41_expert_runtime() = default; + +std::vector llama_dsv41_expert_runtime::remap( + int32_t layer, const std::vector & expert_ids) { + return pimpl->remap(layer, expert_ids); +} + +void llama_dsv41_expert_runtime::release(int32_t layer) { + pimpl->release(layer); +} + +void llama_dsv41_expert_runtime::release_all() { + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + pimpl->release(il); + } +} + +void llama_dsv41_expert_runtime::release_all_after_sync(ggml_backend_sched_t sched) { + if (sched != nullptr) { + ggml_backend_sched_synchronize(sched); + } + release_all(); +} + +ggml_tensor * llama_dsv41_expert_runtime::cache_tensor( + int32_t layer, llama_expert_projection projection) const { + if (layer < 0 || layer >= (int32_t) pimpl->layers.size()) { + throw std::invalid_argument("DeepSeek V4.1 expert layer is invalid"); + } + return pimpl->layers[layer].tensors[projection_index(projection)]; +} + +size_t llama_dsv41_expert_runtime::cache_slots() const { + return pimpl->params.cache_slots; +} + +size_t llama_dsv41_expert_runtime::cache_bytes() const { + return pimpl->bytes_cache; +} + +size_t llama_dsv41_expert_runtime::staging_bytes() const { + return pimpl->bytes_staging; +} + +void llama_dsv41_expert_runtime::set_error(const std::string & error) { + std::lock_guard lock(pimpl->mutex); + if (pimpl->error.empty()) { + pimpl->error = error; + } +} + +std::string llama_dsv41_expert_runtime::consume_error() { + std::lock_guard lock(pimpl->mutex); + std::string result; + result.swap(pimpl->error); + return result; +} + +void llama_dsv41_expert_runtime::acquire_context() { + std::lock_guard lock(pimpl->mutex); + if (pimpl->context_active) { + throw std::runtime_error("DeepSeek V4.1 bounded expert runtime supports one context per model"); + } + pimpl->context_active = true; +} + +void llama_dsv41_expert_runtime::release_context() { + std::lock_guard lock(pimpl->mutex); + pimpl->context_active = false; +} + +std::vector llama_dsv41_register_expert_tensors( + const std::function & ne)> & register_tensor) { + if (!register_tensor) { + throw std::invalid_argument("DeepSeek V4.1 expert tensor registrar is empty"); + } + std::vector result; + result.reserve(LLAMA_DSV41_N_LAYER*3); + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + result.push_back(register_tensor( + "blk." + std::to_string(il) + ".ffn_gate_exps.weight", + il, LLAMA_EXPERT_PROJECTION_GATE, + { LLAMA_DSV41_N_EMBD, LLAMA_DSV41_N_FF_EXP, LLAMA_DSV41_N_EXPERT })); + result.push_back(register_tensor( + "blk." + std::to_string(il) + ".ffn_up_exps.weight", + il, LLAMA_EXPERT_PROJECTION_UP, + { LLAMA_DSV41_N_EMBD, LLAMA_DSV41_N_FF_EXP, LLAMA_DSV41_N_EXPERT })); + result.push_back(register_tensor( + "blk." + std::to_string(il) + ".ffn_down_exps.weight", + il, LLAMA_EXPERT_PROJECTION_DOWN, + { LLAMA_DSV41_N_FF_EXP, LLAMA_DSV41_N_EMBD, LLAMA_DSV41_N_EXPERT })); + } + return result; +} + +static void dsv41_expert_remap_callback( + ggml_tensor * dst, + const ggml_tensor * src, + int, + int, + void * userdata) { + auto * state = static_cast(userdata); + GGML_ASSERT(dst->type == GGML_TYPE_I32 && src->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src)); + const int32_t * ids = static_cast(src->data); + const size_t count = ggml_nelements(src); + try { + const std::vector remapped = state->runtime->remap( + state->layer, std::vector(ids, ids + count)); + memcpy(dst->data, remapped.data(), remapped.size()*sizeof(int32_t)); + } catch (const std::exception & error) { + std::fill_n(static_cast(dst->data), count, 0); + state->runtime->set_error(error.what()); + } +} + +static void dsv41_expert_release_callback( + ggml_tensor * dst, + const ggml_tensor * src, + int, + int, + void * userdata) { + auto * state = static_cast(userdata); + memcpy(dst->data, src->data, ggml_nbytes(src)); + state->runtime->release(state->layer); +} + +ggml_tensor * llama_dsv41_build_expert_remap( + ggml_context * ctx, + ggml_tensor * selected_experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || selected_experts == nullptr || selected_experts->type != GGML_TYPE_I32) { + throw std::invalid_argument("DeepSeek V4.1 expert remap input is invalid"); + } + require_local_cpu(sched, backend_cpu); + ggml_tensor * original = ggml_cont(ctx, selected_experts); + ggml_tensor * remapped = ggml_map_custom1( + ctx, original, dsv41_expert_remap_callback, 1, &runtime.pimpl->callbacks.at(layer)); + ggml_backend_sched_set_tensor_backend(sched, remapped, backend_cpu); + return remapped; +} + +ggml_tensor * llama_dsv41_build_expert_release( + ggml_context * ctx, + ggml_tensor * experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || experts == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 expert release input is invalid"); + } + require_local_cpu(sched, backend_cpu); + ggml_tensor * completion = ggml_sum(ctx, experts); + completion = ggml_map_custom1( + ctx, completion, dsv41_expert_release_callback, 1, &runtime.pimpl->callbacks.at(layer)); + ggml_backend_sched_set_tensor_backend(sched, completion, backend_cpu); + return completion; +} diff --git a/src/llama-dsv41-expert.h b/src/llama-dsv41-expert.h new file mode 100644 index 000000000000..82ce5037e8fc --- /dev/null +++ b/src/llama-dsv41-expert.h @@ -0,0 +1,89 @@ +#pragma once + +#include "llama-expert-store.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include + +struct ggml_cgraph; +struct ggml_context; +struct ggml_tensor; + +struct llama_dsv41_expert_runtime_params { + size_t cache_bytes = 0; + size_t cache_slots = 0; + bool direct_io = true; + bool allow_buffered_io = false; + bool no_alloc = false; +}; + +struct llama_dsv41_expert_runtime { + using buft_selector = std::function; + using upload_fn = std::function; + using publish_fn = std::function; + + llama_dsv41_expert_runtime( + std::vector tensors, + const llama_dsv41_expert_runtime_params & params, + buft_selector select_buft, + upload_fn upload = {}, + publish_fn before_publish = {}); + ~llama_dsv41_expert_runtime(); + + llama_dsv41_expert_runtime(const llama_dsv41_expert_runtime &) = delete; + llama_dsv41_expert_runtime & operator=(const llama_dsv41_expert_runtime &) = delete; + + std::vector remap(int32_t layer, const std::vector & expert_ids); + void acquire_context(); + void release_context(); + void release(int32_t layer); + void release_all(); + void release_all_after_sync(ggml_backend_sched_t sched); + + ggml_tensor * cache_tensor(int32_t layer, llama_expert_projection projection) const; + size_t cache_slots() const; + size_t cache_bytes() const; + size_t staging_bytes() const; + + void set_error(const std::string & error); + std::string consume_error(); + +private: + friend ggml_tensor * llama_dsv41_build_expert_remap( + ggml_context *, ggml_tensor *, llama_dsv41_expert_runtime &, int32_t, ggml_backend_sched_t, ggml_backend_t); + friend ggml_tensor * llama_dsv41_build_expert_release( + ggml_context *, ggml_tensor *, llama_dsv41_expert_runtime &, int32_t, ggml_backend_sched_t, ggml_backend_t); + + struct impl; + std::unique_ptr pimpl; +}; + +std::vector llama_dsv41_register_expert_tensors( + const std::function & ne)> & register_tensor); + +ggml_tensor * llama_dsv41_build_expert_remap( + ggml_context * ctx, + ggml_tensor * selected_experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); + +ggml_tensor * llama_dsv41_build_expert_release( + ggml_context * ctx, + ggml_tensor * experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); diff --git a/src/llama-dsv41.cpp b/src/llama-dsv41.cpp new file mode 100644 index 000000000000..2b578afa443c --- /dev/null +++ b/src/llama-dsv41.cpp @@ -0,0 +1,634 @@ +#include "llama-dsv41.h" + +#include "ggml.h" + +#include +#include +#include +#include +#include +#include + +static constexpr int32_t DSV41_KV_SOURCES[] = { 2, 8, 14, 20 }; +static constexpr int32_t DSV41_INDEX_SOURCES[] = { 2, 8, 14, 20, 24, 28, 32, 36 }; +static constexpr uint32_t DSV41_ENGRAM_LAYERS[] = { 1, 14 }; +static constexpr uint32_t DSV41_ENGRAM_ROWS[] = { 384006168, 384016682 }; + +static void dsv41_require(bool condition, const char * message) { + if (!condition) { + throw std::runtime_error(std::string("DeepSeek V4.1 metadata: ") + message); + } +} + +void llama_dsv41_validate_config(const llama_dsv41_config & config) { + dsv41_require(config.n_ctx_train == LLAMA_DSV41_N_CTX, "max_position_embeddings must be 1048576"); + dsv41_require(config.n_embd == LLAMA_DSV41_N_EMBD, "hidden_size must be 5120"); + dsv41_require(config.n_layer == LLAMA_DSV41_N_LAYER, "num_hidden_layers must be 40"); + dsv41_require(config.n_vocab == LLAMA_DSV41_N_VOCAB, "vocab_size must be 129280"); + dsv41_require(config.n_head == LLAMA_DSV41_N_HEAD, "num_attention_heads must be 64"); + dsv41_require(config.n_head_kv == LLAMA_DSV41_N_HEAD_KV, "num_key_value_heads must be 1"); + dsv41_require(config.n_head_dim == LLAMA_DSV41_N_HEAD_DIM, "head_dim must be 512"); + dsv41_require(config.n_rot == LLAMA_DSV41_N_ROT, "qk_rope_head_dim must be 64"); + dsv41_require(config.n_lora_q == LLAMA_DSV41_N_LORA_Q, "q_lora_rank must be 1280"); + dsv41_require(config.n_lora_o == LLAMA_DSV41_N_LORA_O, "o_lora_rank must be 1024"); + dsv41_require(config.n_o_group == LLAMA_DSV41_N_O_GROUP, "o_group_num must be 8"); + dsv41_require(config.n_ff_dense == LLAMA_DSV41_N_FF_DENSE, "intermediate_size must be 18432"); + dsv41_require(config.n_ff_expert == LLAMA_DSV41_N_FF_EXP, "moe_intermediate_size must be 2304"); + dsv41_require(config.n_expert == LLAMA_DSV41_N_EXPERT, "n_routed_experts must be 384"); + dsv41_require(config.n_expert_used == LLAMA_DSV41_N_EXPERT_USED, "num_experts_per_tok must be 6"); + dsv41_require(config.n_expert_shared == LLAMA_DSV41_N_EXPERT_SHARED, "n_shared_experts must be 1"); + dsv41_require(config.indexer_n_head == LLAMA_DSV41_N_INDEX_HEAD, "index_n_heads must be 32"); + dsv41_require(config.indexer_head_size == LLAMA_DSV41_N_INDEX_HEAD_DIM, "index_head_dim must be 128"); + dsv41_require(config.indexer_top_k == LLAMA_DSV41_N_INDEX_TOP_K, "index_topk must be 512"); + dsv41_require(config.hc_count == LLAMA_DSV41_HC_MULT, "hc_num_streams must be 4"); + dsv41_require(config.hc_sinkhorn_iters == LLAMA_DSV41_HC_SINKHORN_ITERS, "hc_sinkhorn_iters must be 20"); + dsv41_require(config.raw_window == LLAMA_DSV41_N_SWA, "sliding_window must be 128"); + dsv41_require(config.candidate_source_layer == LLAMA_DSV41_CANDIDATE_SOURCE_LAYER, "candidate_source_layer must be 20"); + dsv41_require(config.candidate_topk_blocks == LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS, "candidate_topk_blocks must be 2048"); + dsv41_require(config.candidate_block_size == LLAMA_DSV41_CANDIDATE_BLOCK_SIZE, "candidate_block_size must be 8"); + dsv41_require(config.f_norm_rms_eps == 1.0e-20f, "rms_norm_eps must be 1e-20"); + dsv41_require(config.hc_eps == 1.0e-6f, "hc_eps must be 1e-6"); + dsv41_require(config.swiglu_clamp == 10.0f, "swiglu_clamp_limit must be 10"); + dsv41_require(config.routed_scale == 1.5f, "routed_scaling_factor must be 1.5"); + dsv41_require(config.rope_theta == 10000, "rope_theta must be 10000"); + dsv41_require(config.compress_rope_theta == 160000, "compress_rope_theta must be 160000"); + dsv41_require(config.yarn_factor == 16.0f, "rope_scaling.factor must be 16"); + dsv41_require(config.yarn_beta_fast == 32.0f, "rope_scaling.beta_fast must be 32"); + dsv41_require(config.yarn_beta_slow == 1.0f, "rope_scaling.beta_slow must be 1"); + dsv41_require(config.yarn_original_context == 65536.0f, "rope_scaling.original_max_position_embeddings must be 65536"); + dsv41_require(config.expert_weights_norm, "norm_topk_prob must be true"); + dsv41_require(config.hidden_act == "silu", "hidden_act must be silu"); + dsv41_require(config.scoring_func == "sqrtsoftplus", "scoring_func must be sqrtsoftplus"); + dsv41_require(config.topk_method == "noaux_tc", "topk_method must be noaux_tc"); + + dsv41_require(config.compress_ratios.size() == LLAMA_DSV41_N_LAYER, "compress_ratios must have 40 entries"); + for (uint32_t il = 0; il < LLAMA_DSV41_N_LAYER; ++il) { + dsv41_require(config.compress_ratios[il] == llama_dsv41_compress_ratio(il), "compress_ratios has an invalid main-layer value"); + } + dsv41_require(config.kv_sources == std::vector(std::begin(DSV41_KV_SOURCES), std::end(DSV41_KV_SOURCES)), "kv_source_layers must be [2,8,14,20]"); + dsv41_require(config.index_sources == std::vector(std::begin(DSV41_INDEX_SOURCES), std::end(DSV41_INDEX_SOURCES)), "index_source_layers must be [2,8,14,20,24,28,32,36]"); + dsv41_require(config.engram_layers == std::vector(std::begin(DSV41_ENGRAM_LAYERS), std::end(DSV41_ENGRAM_LAYERS)), "engram.layer_ids must be [1,14]"); + dsv41_require(config.engram_rows == std::vector(std::begin(DSV41_ENGRAM_ROWS), std::end(DSV41_ENGRAM_ROWS)), "engram.rows must be [384006168,384016682]"); + dsv41_require(config.engram_encoding == LLAMA_DSV41_ENGRAM_ENCODING, "engram.encoding must be e4m3_e8m0_32_row264"); + dsv41_require(config.engram_compressed_vocab_size == LLAMA_DSV41_ENGRAM_COMPRESSED_VOCAB, "engram.compressed_vocab_size must be 99092"); + dsv41_require(config.engram_pad_id == LLAMA_DSV41_ENGRAM_PAD_ID, "engram.pad_id must be 2"); + dsv41_require(config.engram_token_map_size == LLAMA_DSV41_N_VOCAB, "engram.token_map must contain 129280 entries"); + dsv41_require(config.engram_primes_size == LLAMA_DSV41_ENGRAM_PRIMES_COUNT, "engram.primes must contain 48 entries"); + dsv41_require(config.engram_multipliers_size == LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT, "engram.multipliers must contain 8 entries"); + if (!config.engram_token_map.empty() || !config.engram_primes.empty() || !config.engram_multipliers.empty()) { + llama_dsv41_make_engram_layout(config); + } +} + +llama_engram_layout llama_dsv41_make_engram_layout(const llama_dsv41_config & config) { + dsv41_require(config.engram_token_map.size() == LLAMA_DSV41_N_VOCAB, "engram.token_map data must contain 129280 entries"); + dsv41_require(config.engram_primes.size() == LLAMA_DSV41_ENGRAM_PRIMES_COUNT, "engram.primes data must contain 48 entries"); + dsv41_require(config.engram_multipliers.size() == LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT, "engram.multipliers data must contain 8 entries"); + + llama_engram_layout layout; + layout.encoding = config.engram_encoding; + std::copy(config.engram_layers.begin(), config.engram_layers.end(), layout.layer_ids.begin()); + layout.token_map = config.engram_token_map; + layout.compressed_vocab_size = config.engram_compressed_vocab_size; + layout.pad_id = config.engram_pad_id; + std::copy(config.engram_rows.begin(), config.engram_rows.end(), layout.rows.begin()); + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + std::copy_n( + config.engram_primes.begin() + layer*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS, + layout.primes[layer].begin()); + std::copy_n( + config.engram_multipliers.begin() + layer*LLAMA_ENGRAM_NGRAM, + LLAMA_ENGRAM_NGRAM, + layout.multipliers[layer].begin()); + } + llama_engram_hasher validate(layout); + return layout; +} + + + +static int32_t dsv41_source_layer(const int32_t * sources, size_t n, uint32_t il) { + int32_t result = -1; + for (size_t i = 0; i < n; ++i) { + if ((uint32_t) sources[i] > il) { + break; + } + result = sources[i]; + } + return result; +} + +llama_dsv41_cache_state::llama_dsv41_cache_state(uint32_t compressed_cache_size) : + compressed_cache_size(compressed_cache_size), + raw(LLAMA_DSV41_N_SWA, -1) { + if (compressed_cache_size == 0) { + throw std::runtime_error("DeepSeek V4.1 compressed cache is empty"); + } + for (int32_t source : DSV41_KV_SOURCES) { + const uint32_t ratio = llama_dsv41_compress_ratio(source); + sources.emplace(source, source_state { + ratio, + std::vector(compressed_cache_size, -1), + std::vector(ratio, -1), + }); + } +} + +void llama_dsv41_cache_state::clear() { + pos = -1; + std::fill(raw.begin(), raw.end(), -1); + for (auto & item : sources) { + std::fill(item.second.compressed.begin(), item.second.compressed.end(), -1); + std::fill(item.second.pending.begin(), item.second.pending.end(), -1); + } + candidates.clear(); +} + +void llama_dsv41_cache_state::append(llama_pos next_pos) { + if (next_pos != pos + 1) { + throw std::runtime_error("DeepSeek V4.1 cache requires contiguous single-sequence positions"); + } + + for (const auto & item : sources) { + const source_state & state = item.second; + if ((next_pos + 1)%state.ratio == 0 && (uint64_t) (next_pos/state.ratio) >= compressed_cache_size) { + throw std::runtime_error("DeepSeek V4.1 compressed cache overflow"); + } + } + + raw[next_pos%raw.size()] = next_pos; + for (auto & item : sources) { + source_state & state = item.second; + state.pending[next_pos%state.ratio] = next_pos; + if ((next_pos + 1)%state.ratio != 0) { + continue; + } + + const uint64_t dst = next_pos/state.ratio; + state.compressed[dst] = next_pos + 1 - state.ratio; + } + pos = next_pos; +} + +void llama_dsv41_cache_state::set_candidate_blocks(const std::vector & blocks) { + candidates = blocks; +} + +llama_pos llama_dsv41_cache_state::position() const { + return pos; +} + +const std::vector & llama_dsv41_cache_state::raw_slots() const { + return raw; +} + +const std::vector & llama_dsv41_cache_state::compressed_slots(uint32_t source_layer) const { + const auto it = sources.find(source_layer); + if (it == sources.end()) { + throw std::runtime_error("DeepSeek V4.1 compressed cache source layer is invalid"); + } + return it->second.compressed; +} + +const std::vector & llama_dsv41_cache_state::pending_slots(uint32_t source_layer) const { + const auto it = sources.find(source_layer); + if (it == sources.end()) { + throw std::runtime_error("DeepSeek V4.1 compressor state source layer is invalid"); + } + return it->second.pending; +} + +const std::vector & llama_dsv41_cache_state::candidate_blocks() const { + return candidates; +} + +uint64_t llama_dsv41_memory_accounting::total() const { + return raw_kv + compressed_kv + index_keys + compressor_carry + + candidate_scores + candidate_ids + position_state + graph_workspace; +} + +llama_dsv41_memory_accounting llama_dsv41_account_memory( + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_tokens, + uint32_t kv_element_size, + uint32_t index_element_size, + uint64_t graph_workspace) { + if (n_ctx == 0 || n_seq == 0 || n_tokens == 0 || kv_element_size == 0 || index_element_size == 0) { + throw std::runtime_error("DeepSeek V4.1 memory accounting dimensions must be non-zero"); + } + + llama_dsv41_memory_accounting result; + const uint64_t raw_rows = (uint64_t) LLAMA_DSV41_N_LAYER*LLAMA_DSV41_N_SWA*n_seq; + const uint64_t ratio_2_rows = ((uint64_t) n_ctx + 1)/2; + const uint64_t ratio_1_rows = n_ctx; + const uint64_t compressed_rows = (3*ratio_2_rows + ratio_1_rows)*n_seq; + const uint64_t candidate_blocks = ((uint64_t) n_ctx + LLAMA_DSV41_CANDIDATE_BLOCK_SIZE - 1)/ + LLAMA_DSV41_CANDIDATE_BLOCK_SIZE; + uint64_t pending_rows = 0; + uint64_t gated_pending_rows = 0; + for (int32_t source : DSV41_KV_SOURCES) { + const uint32_t ratio = llama_dsv41_compress_ratio(source); + pending_rows += ratio; + if (ratio == 2) { + gated_pending_rows += ratio; + } + } + + result.raw_kv = raw_rows*LLAMA_DSV41_N_HEAD_DIM*kv_element_size; + result.compressed_kv = compressed_rows*LLAMA_DSV41_N_HEAD_DIM*kv_element_size; + result.index_keys = compressed_rows*LLAMA_DSV41_N_INDEX_HEAD_DIM*index_element_size; + result.compressor_carry = + (pending_rows + gated_pending_rows)*LLAMA_DSV41_N_HEAD_DIM*kv_element_size*n_seq; + result.candidate_scores = candidate_blocks*sizeof(float)*n_tokens; + result.candidate_ids = std::min(candidate_blocks, LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS)* + sizeof(int32_t)*n_tokens; + const uint64_t position_rows = + ((uint64_t) LLAMA_DSV41_N_SWA + compressed_rows/n_seq + pending_rows)*n_seq; + result.position_state = position_rows*sizeof(llama_pos); + result.graph_workspace = graph_workspace; + return result; +} + +int32_t llama_dsv41_kv_source_layer(uint32_t il) { + return dsv41_source_layer(DSV41_KV_SOURCES, sizeof(DSV41_KV_SOURCES)/sizeof(DSV41_KV_SOURCES[0]), il); +} + +int32_t llama_dsv41_index_source_layer(uint32_t il) { + return dsv41_source_layer(DSV41_INDEX_SOURCES, sizeof(DSV41_INDEX_SOURCES)/sizeof(DSV41_INDEX_SOURCES[0]), il); +} + +uint32_t llama_dsv41_compress_ratio(uint32_t il) { + if (il < 2) { + return 0; + } + if (il < 20) { + return 2; + } + if (il < LLAMA_DSV41_N_LAYER) { + return 1; + } + throw std::runtime_error("DeepSeek V4.1 layer index is out of range"); +} + +llama_dsv41_layer_plan llama_dsv41_build_layer_plan( + uint32_t il, + const std::vector & positions, + uint32_t compressed_cache_size) { + if (positions.empty()) { + throw std::runtime_error("DeepSeek V4.1 graph plan requires at least one token"); + } + if (positions.front() < 0) { + throw std::runtime_error("DeepSeek V4.1 graph plan requires non-negative positions"); + } + for (size_t i = 1; i < positions.size(); ++i) { + if (positions[i] != positions[i - 1] + 1) { + throw std::runtime_error("DeepSeek V4.1 graph plan requires one contiguous sequence"); + } + } + + llama_dsv41_layer_plan plan = {}; + plan.layer = il; + plan.ratio = llama_dsv41_compress_ratio(il); + plan.kv_source_layer = llama_dsv41_kv_source_layer(il); + plan.index_source_layer = llama_dsv41_index_source_layer(il); + plan.owns_kv_source = plan.kv_source_layer == (int32_t) il; + plan.owns_index_source = plan.index_source_layer == (int32_t) il; + plan.builds_candidates = il == LLAMA_DSV41_CANDIDATE_SOURCE_LAYER; + plan.uses_candidates = plan.owns_index_source && il > LLAMA_DSV41_CANDIDATE_SOURCE_LAYER; + plan.reuses_index_selection = !plan.owns_index_source && plan.index_source_layer >= 0; + plan.collapses_output = il + 1 == LLAMA_DSV41_N_LAYER; + plan.raw_ring_order = llama_dsv41_raw_ring_order(positions.back(), LLAMA_DSV41_N_SWA); + if (plan.owns_kv_source && plan.ratio != 0) { + plan.compression = llama_dsv41_build_compression_plan(positions, plan.ratio, compressed_cache_size); + } + return plan; +} + +llama_dsv41_compression_plan llama_dsv41_build_compression_plan( + const std::vector & positions, + uint32_t ratio, + uint32_t cache_size) { + if (ratio != 1 && ratio != 2) { + throw std::runtime_error("DeepSeek V4.1 compression ratio must be 1 or 2"); + } + if (cache_size == 0) { + throw std::runtime_error("DeepSeek V4.1 compressed cache is empty"); + } + + llama_dsv41_compression_plan plan; + plan.n_visible.resize(positions.size(), 0); + + std::vector latest_state_src(ratio, -1); + std::vector latest_state_pos(ratio, -1); + + const int32_t scratch_offset = ratio; + for (size_t i = 0; i < positions.size(); ++i) { + const llama_pos pos = positions[i]; + if (pos < 0) { + continue; + } + + plan.n_visible[i] = (int32_t) ((pos + 1)/ratio); + plan.n_kv = std::max(plan.n_kv, plan.n_visible[i]); + + if (ratio == 1) { + if ((uint64_t) pos >= cache_size) { + throw std::runtime_error("DeepSeek V4.1 compressed cache overflow"); + } + plan.write_idxs.push_back(pos); + plan.write_pos.push_back((int32_t) pos); + continue; + } + + const int32_t row = (int32_t) (pos%ratio); + plan.state_pos.push_back(row); + if (latest_state_src[row] < 0 || pos > latest_state_pos[row]) { + latest_state_src[row] = (int32_t) i; + latest_state_pos[row] = pos; + } + + if ((pos + 1)%ratio != 0) { + continue; + } + + const llama_pos source_start = pos + 1 - ratio; + const int64_t write_idx = pos/ratio; + if ((uint64_t) write_idx >= cache_size) { + throw std::runtime_error("DeepSeek V4.1 compressed cache overflow"); + } + + for (uint32_t j = 0; j < ratio; ++j) { + const llama_pos source_pos = source_start + j; + int32_t source_idx = (int32_t) (source_pos%ratio); + for (size_t k = 0; k < positions.size(); ++k) { + if (positions[k] == source_pos) { + source_idx = scratch_offset + (int32_t) k; + break; + } + } + plan.state_read_idxs.push_back(source_idx); + } + + plan.write_idxs.push_back(write_idx); + plan.write_pos.push_back((int32_t) source_start); + } + + if (ratio == 2) { + for (uint32_t row = 0; row < ratio; ++row) { + if (latest_state_src[row] >= 0) { + plan.state_persist_src_idxs.push_back(latest_state_src[row]); + plan.state_persist_dst_idxs.push_back((int32_t) row); + } + } + } + + plan.n_kv = std::min(cache_size, std::max(1, plan.n_kv)); + return plan; +} + +std::vector llama_dsv41_select_candidate_blocks( + const std::vector & scores, + uint32_t n_visible, + uint32_t block_size, + uint32_t top_k_blocks) { + if (block_size == 0) { + throw std::runtime_error("DeepSeek V4.1 candidate block size must be non-zero"); + } + + n_visible = std::min(n_visible, scores.size()); + if (n_visible == 0 || top_k_blocks == 0) { + return {}; + } + + const uint32_t n_blocks = (n_visible + block_size - 1)/block_size; + std::vector block_scores(n_blocks, -std::numeric_limits::infinity()); + for (uint32_t block = 0; block < n_blocks; ++block) { + const uint32_t i0 = block*block_size; + const uint32_t i1 = std::min(n_visible, i0 + block_size); + for (uint32_t i = i0; i < i1; ++i) { + if (std::isnan(scores[i])) { + throw std::runtime_error("DeepSeek V4.1 candidate score is NaN"); + } + block_scores[block] = std::max(block_scores[block], scores[i]); + } + } + + const int32_t final_block = (int32_t) n_blocks - 1; + std::vector blocks(n_blocks - 1); + std::iota(blocks.begin(), blocks.end(), 0); + std::stable_sort(blocks.begin(), blocks.end(), [&](int32_t a, int32_t b) { + if (block_scores[a] != block_scores[b]) { + return block_scores[a] > block_scores[b]; + } + return a < b; + }); + + const uint32_t n_selected = std::min(top_k_blocks, n_blocks); + std::vector selected; + selected.reserve(n_selected); + selected.push_back(final_block); + selected.insert(selected.end(), blocks.begin(), blocks.begin() + n_selected - 1); + return selected; +} + +std::vector llama_dsv41_candidate_rows( + const std::vector & blocks, + uint32_t n_visible, + uint32_t block_size) { + if (block_size == 0) { + throw std::runtime_error("DeepSeek V4.1 candidate block size must be non-zero"); + } + + std::vector rows; + for (int32_t block : blocks) { + if (block < 0) { + throw std::runtime_error("DeepSeek V4.1 candidate block index must be non-negative"); + } + const uint64_t i0 = (uint64_t) block*block_size; + const uint64_t i1 = std::min(n_visible, i0 + block_size); + for (uint64_t i = i0; i < i1; ++i) { + rows.push_back((int32_t) i); + } + } + return rows; +} + +std::vector llama_dsv41_raw_ring_order(llama_pos pos, uint32_t window) { + if (window == 0 || pos < 0) { + return {}; + } + + const uint32_t n_raw = std::min((uint64_t) pos + 1, window); + const uint32_t start = (uint32_t) ((pos + 1 - n_raw)%window); + + std::vector result(n_raw); + for (uint32_t i = 0; i < n_raw; ++i) { + result[i] = (start + i)%window; + } + return result; +} + +std::vector llama_dsv41_output_collapse( + const std::vector & residual, + const std::vector & pre, + uint32_t n_embd, + uint32_t hc_mult) { + if (pre.size() != hc_mult || residual.size() != (size_t) n_embd*hc_mult) { + throw std::runtime_error("DeepSeek V4.1 output collapse shape mismatch"); + } + + std::vector result(n_embd, 0.0f); + for (uint32_t h = 0; h < hc_mult; ++h) { + for (uint32_t d = 0; d < n_embd; ++d) { + result[d] += residual[(size_t) h*n_embd + d]*pre[h]; + } + } + return result; +} + +ggml_tensor * llama_dsv41_build_ratio_pool( + ggml_context * ctx, + ggml_tensor * kv, + ggml_tensor * gate, + uint32_t ratio) { + if (ratio != 1 && ratio != 2) { + throw std::runtime_error("DeepSeek V4.1 graph compression ratio must be 1 or 2"); + } + if (kv->ne[1] != ratio) { + throw std::runtime_error("DeepSeek V4.1 graph compressor input shape mismatch"); + } + if (ratio == 1) { + return ggml_reshape_2d(ctx, kv, kv->ne[0], kv->ne[2]); + } + if (gate == nullptr || !ggml_are_same_shape(kv, gate)) { + throw std::runtime_error("DeepSeek V4.1 ratio-2 graph requires matching KV and gate tensors"); + } + + ggml_tensor * kv_t = ggml_cont(ctx, ggml_permute(ctx, kv, 1, 0, 2, 3)); + ggml_tensor * gate_t = ggml_cont(ctx, ggml_permute(ctx, gate, 1, 0, 2, 3)); + ggml_tensor * weights = ggml_soft_max(ctx, gate_t); + ggml_tensor * pooled = ggml_sum_rows(ctx, ggml_mul(ctx, kv_t, weights)); + return ggml_reshape_2d(ctx, pooled, kv->ne[0], kv->ne[2]); +} + +ggml_tensor * llama_dsv41_build_shared_softmax( + ggml_context * ctx, + ggml_tensor * raw_scores, + ggml_tensor * compressed_scores) { + if (raw_scores == nullptr && compressed_scores == nullptr) { + throw std::runtime_error("DeepSeek V4.1 attention requires raw or compressed scores"); + } + ggml_tensor * scores = raw_scores; + if (scores == nullptr) { + scores = compressed_scores; + } else if (compressed_scores != nullptr) { + if (raw_scores->ne[1] != compressed_scores->ne[1] || + raw_scores->ne[2] != compressed_scores->ne[2] || + raw_scores->ne[3] != compressed_scores->ne[3]) { + throw std::runtime_error("DeepSeek V4.1 raw and compressed score shapes are incompatible"); + } + scores = ggml_concat(ctx, raw_scores, compressed_scores, 0); + } + 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, + ggml_tensor * pre, + uint32_t n_embd, + uint32_t hc_mult, + uint32_t n_tokens) { + if (residual->ne[0] != n_embd || residual->ne[1] != hc_mult || residual->ne[2] != n_tokens || + pre->ne[0] != hc_mult || pre->ne[1] != n_tokens) { + throw std::runtime_error("DeepSeek V4.1 graph output collapse shape mismatch"); + } + + ggml_tensor * residual_t = ggml_cont(ctx, ggml_permute(ctx, residual, 1, 0, 2, 3)); + ggml_tensor * pre_t = ggml_reshape_3d(ctx, pre, hc_mult, 1, n_tokens); + ggml_tensor * collapsed = ggml_sum_rows(ctx, ggml_mul(ctx, residual_t, pre_t)); + collapsed = ggml_reshape_2d(ctx, collapsed, n_embd, n_tokens); + 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, + ggml_tensor * pre, + ggml_tensor * output_norm, + ggml_tensor * output, + float rms_eps, + uint32_t hc_mult) { + if (output_norm->ne[0] != residual->ne[0] || output->ne[0] != residual->ne[0]) { + throw std::runtime_error("DeepSeek V4.1 output tensor shape mismatch"); + } + + 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, llama_dsv41_build_output_norm_input(ctx, collapsed), rms_eps); + normalized = ggml_mul(ctx, normalized, output_norm); + return ggml_mul_mat(ctx, output, normalized); +} diff --git a/src/llama-dsv41.h b/src/llama-dsv41.h new file mode 100644 index 000000000000..86fca484d8de --- /dev/null +++ b/src/llama-dsv41.h @@ -0,0 +1,249 @@ +#pragma once + +#include "llama-engram.h" +#include "llama.h" + +#include +#include +#include +#include +#include + +struct ggml_context; +struct ggml_tensor; + +static constexpr uint32_t LLAMA_DSV41_N_LAYER = 40; +static constexpr uint32_t LLAMA_DSV41_N_EMBD = 5120; +static constexpr uint32_t LLAMA_DSV41_N_VOCAB = 129280; +static constexpr uint32_t LLAMA_DSV41_N_HEAD = 64; +static constexpr uint32_t LLAMA_DSV41_N_HEAD_KV = 1; +static constexpr uint32_t LLAMA_DSV41_N_HEAD_DIM = 512; +static constexpr uint32_t LLAMA_DSV41_N_ROT = 64; +static constexpr uint32_t LLAMA_DSV41_N_LORA_Q = 1280; +static constexpr uint32_t LLAMA_DSV41_N_LORA_O = 1024; +static constexpr uint32_t LLAMA_DSV41_N_O_GROUP = 8; +static constexpr uint32_t LLAMA_DSV41_N_FF_DENSE = 18432; +static constexpr uint32_t LLAMA_DSV41_N_EXPERT = 384; +static constexpr uint32_t LLAMA_DSV41_N_EXPERT_USED = 6; +static constexpr uint32_t LLAMA_DSV41_N_EXPERT_SHARED = 1; +static constexpr uint32_t LLAMA_DSV41_N_FF_EXP = 2304; +static constexpr uint32_t LLAMA_DSV41_N_INDEX_HEAD = 32; +static constexpr uint32_t LLAMA_DSV41_N_INDEX_HEAD_DIM = 128; +static constexpr uint32_t LLAMA_DSV41_N_INDEX_TOP_K = 512; +static constexpr uint32_t LLAMA_DSV41_N_SWA = 128; +static constexpr uint32_t LLAMA_DSV41_N_CTX = 1048576; +static constexpr uint32_t LLAMA_DSV41_HC_MULT = 4; +static constexpr uint32_t LLAMA_DSV41_HC_SINKHORN_ITERS = 20; +static constexpr uint32_t LLAMA_DSV41_CANDIDATE_SOURCE_LAYER = 20; +static constexpr uint32_t LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS = 2048; +static constexpr uint32_t LLAMA_DSV41_CANDIDATE_BLOCK_SIZE = 8; +static constexpr uint32_t LLAMA_DSV41_ENGRAM_COMPRESSED_VOCAB = 99092; +static constexpr uint32_t LLAMA_DSV41_ENGRAM_PAD_ID = 2; +static constexpr uint32_t LLAMA_DSV41_ENGRAM_PRIMES_COUNT = 48; +static constexpr uint32_t LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT = 8; +static constexpr const char * LLAMA_DSV41_ENGRAM_ENCODING = "e4m3_e8m0_32_row264"; + +struct llama_dsv41_config { + uint32_t n_ctx_train; + uint32_t n_embd; + uint32_t n_layer; + uint32_t n_vocab; + uint32_t n_head; + uint32_t n_head_kv; + uint32_t n_head_dim; + uint32_t n_rot; + uint32_t n_lora_q; + uint32_t n_lora_o; + uint32_t n_o_group; + uint32_t n_ff_dense; + uint32_t n_ff_expert; + uint32_t n_expert; + uint32_t n_expert_used; + uint32_t n_expert_shared; + uint32_t indexer_n_head; + uint32_t indexer_head_size; + uint32_t indexer_top_k; + uint32_t hc_count; + uint32_t hc_sinkhorn_iters; + uint32_t raw_window; + uint32_t candidate_source_layer; + uint32_t candidate_topk_blocks; + uint32_t candidate_block_size; + float f_norm_rms_eps; + float hc_eps; + float swiglu_clamp; + float routed_scale; + uint32_t rope_theta; + uint32_t compress_rope_theta; + float yarn_factor; + float yarn_beta_fast; + float yarn_beta_slow; + float yarn_original_context; + bool expert_weights_norm; + std::string hidden_act; + std::string scoring_func; + std::string topk_method; + std::vector compress_ratios; + std::vector kv_sources; + std::vector index_sources; + std::vector engram_layers; + std::vector engram_rows; + std::string engram_encoding; + uint32_t engram_compressed_vocab_size; + uint32_t engram_pad_id; + uint32_t engram_token_map_size; + uint32_t engram_primes_size; + uint32_t engram_multipliers_size; + std::vector engram_token_map; + std::vector engram_primes; + std::vector engram_multipliers; +}; + +void llama_dsv41_validate_config(const llama_dsv41_config & config); +llama_engram_layout llama_dsv41_make_engram_layout(const llama_dsv41_config & config); + +struct llama_dsv41_compression_plan { + std::vector state_pos; + std::vector state_persist_src_idxs; + std::vector state_persist_dst_idxs; + std::vector state_read_idxs; + std::vector write_idxs; + std::vector write_pos; + std::vector n_visible; + int64_t n_kv = 0; +}; + +struct llama_dsv41_layer_plan { + uint32_t layer; + uint32_t ratio; + int32_t kv_source_layer; + int32_t index_source_layer; + bool owns_kv_source; + bool owns_index_source; + bool builds_candidates; + bool uses_candidates; + bool reuses_index_selection; + bool collapses_output; + std::vector raw_ring_order; + llama_dsv41_compression_plan compression; +}; + +class llama_dsv41_cache_state { +public: + explicit llama_dsv41_cache_state(uint32_t compressed_cache_size); + + void clear(); + void append(llama_pos pos); + void set_candidate_blocks(const std::vector & blocks); + + llama_pos position() const; + const std::vector & raw_slots() const; + const std::vector & compressed_slots(uint32_t source_layer) const; + const std::vector & pending_slots(uint32_t source_layer) const; + const std::vector & candidate_blocks() const; + +private: + struct source_state { + uint32_t ratio; + std::vector compressed; + std::vector pending; + }; + + uint32_t compressed_cache_size; + llama_pos pos = -1; + std::vector raw; + std::map sources; + std::vector candidates; +}; + +struct llama_dsv41_memory_accounting { + uint64_t raw_kv = 0; + uint64_t compressed_kv = 0; + uint64_t index_keys = 0; + uint64_t compressor_carry = 0; + uint64_t candidate_scores = 0; + uint64_t candidate_ids = 0; + uint64_t position_state = 0; + uint64_t graph_workspace = 0; + + uint64_t total() const; +}; + +llama_dsv41_memory_accounting llama_dsv41_account_memory( + uint32_t n_ctx, + uint32_t n_seq, + uint32_t n_tokens, + uint32_t kv_element_size, + uint32_t index_element_size, + uint64_t graph_workspace); + +int32_t llama_dsv41_kv_source_layer(uint32_t il); +int32_t llama_dsv41_index_source_layer(uint32_t il); +uint32_t llama_dsv41_compress_ratio(uint32_t il); + +llama_dsv41_layer_plan llama_dsv41_build_layer_plan( + uint32_t il, + const std::vector & positions, + uint32_t compressed_cache_size); + +llama_dsv41_compression_plan llama_dsv41_build_compression_plan( + const std::vector & positions, + uint32_t ratio, + uint32_t cache_size); + +std::vector llama_dsv41_select_candidate_blocks( + const std::vector & scores, + uint32_t n_visible, + uint32_t block_size, + uint32_t top_k_blocks); + +std::vector llama_dsv41_candidate_rows( + const std::vector & blocks, + uint32_t n_visible, + uint32_t block_size); + +std::vector llama_dsv41_raw_ring_order(llama_pos pos, uint32_t window); + +std::vector llama_dsv41_output_collapse( + const std::vector & residual, + const std::vector & pre, + uint32_t n_embd, + uint32_t hc_mult); + +ggml_tensor * llama_dsv41_build_ratio_pool( + ggml_context * ctx, + ggml_tensor * kv, + ggml_tensor * gate, + uint32_t ratio); + +ggml_tensor * llama_dsv41_build_shared_softmax( + ggml_context * ctx, + ggml_tensor * raw_scores, + ggml_tensor * compressed_scores); + +ggml_tensor * llama_dsv41_build_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, + ggml_tensor * pre, + uint32_t n_embd, + 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, + ggml_tensor * pre, + ggml_tensor * output_norm, + ggml_tensor * output, + float rms_eps, + uint32_t hc_mult); diff --git a/src/llama-engram.cpp b/src/llama-engram.cpp new file mode 100644 index 000000000000..336aae39f2d7 --- /dev/null +++ b/src/llama-engram.cpp @@ -0,0 +1,305 @@ +#include "llama-engram.h" + +#include "llama-bounded-file.h" +#include "llama-impl.h" + +#include +#include +#include +#include +#include +#include + +static void llama_engram_validate_layout(const llama_engram_layout & layout) { + if (layout.encoding != "e4m3_e8m0_32_row264") { + throw std::invalid_argument("llama_engram: unsupported row encoding"); + } + if (layout.layer_ids[0] == layout.layer_ids[1]) { + throw std::invalid_argument("llama_engram: layer IDs must be distinct"); + } + if (layout.token_map.empty() || layout.token_map.size() > (size_t) INT32_MAX) { + throw std::invalid_argument("llama_engram: invalid token map size"); + } + if (layout.compressed_vocab_size == 0 || layout.compressed_vocab_size > (uint32_t) INT32_MAX || + layout.pad_id >= layout.compressed_vocab_size) { + throw std::invalid_argument("llama_engram: invalid compressed vocabulary"); + } + for (uint32_t token : layout.token_map) { + if (token >= layout.compressed_vocab_size) { + throw std::invalid_argument("llama_engram: token map entry is out of range"); + } + } + + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + for (size_t i = 0; i < LLAMA_ENGRAM_NGRAM; ++i) { + const uint64_t multiplier = layout.multipliers[layer][i]; + if ((multiplier & 1) == 0 || + multiplier > (uint64_t) INT64_MAX / layout.compressed_vocab_size) { + throw std::invalid_argument("llama_engram: invalid hash multiplier"); + } + } + + uint64_t row_count = 0; + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + const uint32_t prime = layout.primes[layer][col]; + if (prime < 2) { + throw std::invalid_argument("llama_engram: invalid hash prime"); + } + row_count += prime; + } + if (row_count > UINT32_MAX || row_count != layout.rows[layer]) { + throw std::invalid_argument("llama_engram: row extent does not match hash buckets"); + } + } +} + +void llama_engram_history::reset() { + tail.fill(LLAMA_ENGRAM_DEAD); +} + +struct llama_engram_hasher::impl { + llama_engram_layout layout; + + explicit impl(llama_engram_layout layout) : layout(std::move(layout)) { + llama_engram_validate_layout(this->layout); + } + + void validate_history(const llama_engram_history & history) const { + for (int32_t token : history.tail) { + if (token < LLAMA_ENGRAM_DEAD || + (token >= 0 && (uint32_t) token >= layout.compressed_vocab_size)) { + throw std::invalid_argument("llama_engram: invalid token history"); + } + } + } + + void hash( + llama_engram_history & history, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * rows) const { + if (count > SIZE_MAX / (LLAMA_ENGRAM_LAYERS * LLAMA_ENGRAM_COLS * sizeof(*rows)) || + (count != 0 && (tokens == nullptr || rows == nullptr))) { + throw std::invalid_argument("llama_engram: invalid hash buffers"); + } + validate_history(history); + for (size_t i = 0; i < count; ++i) { + if (tokens[i] < 0 || (size_t) tokens[i] >= layout.token_map.size()) { + throw std::invalid_argument("llama_engram: token ID is out of range"); + } + } + + llama_engram_history next = history; + uint32_t * output = rows; + for (size_t pos = 0; pos < count; ++pos) { + const int32_t current = mask != nullptr && mask[pos] == 0 ? + LLAMA_ENGRAM_DEAD : (int32_t) layout.token_map[tokens[pos]]; + uint32_t ids[LLAMA_ENGRAM_NGRAM]; + bool blocked = false; + for (size_t depth = 0; depth < LLAMA_ENGRAM_NGRAM; ++depth) { + const int32_t id = depth == 0 ? current : next.tail[depth - 1]; + blocked = blocked || id == LLAMA_ENGRAM_DEAD; + ids[depth] = blocked ? layout.pad_id : (uint32_t) id; + } + + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + uint64_t hash = (uint64_t) ids[0] * layout.multipliers[layer][0]; + uint64_t offset = 0; + for (size_t depth = 1; depth < LLAMA_ENGRAM_NGRAM; ++depth) { + hash ^= (uint64_t) ids[depth] * layout.multipliers[layer][depth]; + for (size_t head = 0; head < LLAMA_ENGRAM_HEADS; ++head) { + const size_t col = (depth - 1) * LLAMA_ENGRAM_HEADS + head; + const uint32_t prime = layout.primes[layer][col]; + *output++ = (uint32_t) (hash % prime + offset); + offset += prime; + } + } + } + + for (size_t i = next.tail.size() - 1; i > 0; --i) { + next.tail[i] = next.tail[i - 1]; + } + next.tail[0] = current; + } + history = next; + } +}; + +llama_engram_hasher::llama_engram_hasher(llama_engram_layout layout) + : pimpl(std::make_unique(std::move(layout))) {} + +llama_engram_hasher::~llama_engram_hasher() = default; + +void llama_engram_hasher::hash( + llama_engram_history & history, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * rows) const { + pimpl->hash(history, tokens, mask, count, rows); +} + +const llama_engram_layout & llama_engram_hasher::layout() const { + return pimpl->layout; +} + +static float llama_engram_e4m3(uint8_t code) { + const int exponent = (code >> 3) & 15; + const int mantissa = code & 7; + const float value = exponent != 0 ? + std::ldexp((float) (8 + mantissa), exponent - 10) : + std::ldexp((float) mantissa, -9); + return (code & 128) != 0 ? -value : value; +} + +void llama_engram_decode_row(const uint8_t row[LLAMA_ENGRAM_ROW_BYTES], float output[LLAMA_ENGRAM_DIM]) { + if (row == nullptr || output == nullptr) { + throw std::invalid_argument("llama_engram: invalid row decode buffers"); + } + + float decoded[LLAMA_ENGRAM_DIM]; + for (size_t i = 0; i < LLAMA_ENGRAM_DIM; ++i) { + const uint8_t code = row[i]; + const uint8_t scale = row[LLAMA_ENGRAM_DIM + i / 32]; + if ((code & 127) == 127) { + throw std::domain_error("llama_engram: E4M3 NaN encoding"); + } + if (scale == 255) { + throw std::domain_error("llama_engram: E8M0 scale 255"); + } + + float value = std::ldexp(llama_engram_e4m3(code), (int) scale - 127); + uint32_t bits; + memcpy(&bits, &value, sizeof(bits)); + bits = (bits + 0x7fffu + ((bits >> 16) & 1u)) & 0xffff0000u; + memcpy(&value, &bits, sizeof(value)); + if (!std::isfinite(value)) { + throw std::domain_error("llama_engram: decoded value is not finite"); + } + decoded[i] = value; + } + memcpy(output, decoded, sizeof(decoded)); +} + +struct llama_engram_table::impl { + struct request { + uint32_t row; + uint32_t output; + }; + + static constexpr size_t BATCH_TOKENS = 2048; + + llama_bounded_file file; + llama_bounded_file::buffer scratch; + uint64_t offset; + uint32_t rows; + std::mutex mutex; + std::vector requests; + + impl(const std::string & fname, uint64_t offset, uint32_t rows) + : file(fname, { true, true }), + scratch(file.make_buffer(LLAMA_ENGRAM_ROW_BYTES)), + offset(offset), + rows(rows) { + const uint64_t bytes = (uint64_t) rows * LLAMA_ENGRAM_ROW_BYTES; + if (rows == 0 || offset > (uint64_t) INT64_MAX || bytes > (uint64_t) INT64_MAX - offset || + offset > file.size() || bytes > file.size() - offset) { + throw std::invalid_argument("llama_engram: invalid table extent"); + } + requests.reserve(BATCH_TOKENS * LLAMA_ENGRAM_COLS); + } + + void validate_rows(const uint32_t * row_ids, size_t count) const { + if (count > SIZE_MAX / (LLAMA_ENGRAM_DIM * sizeof(float)) || + (count != 0 && row_ids == nullptr)) { + throw std::invalid_argument("llama_engram: invalid row list"); + } + for (size_t i = 0; i < count; ++i) { + if (row_ids[i] >= rows) { + throw std::invalid_argument("llama_engram: row ID is out of range"); + } + } + } + + void read_one(uint32_t row, float * output) { + uint8_t raw[LLAMA_ENGRAM_ROW_BYTES]; + file.read(offset + (uint64_t) row * LLAMA_ENGRAM_ROW_BYTES, raw, sizeof(raw), scratch); + llama_engram_decode_row(raw, output); + } + + void read(const uint32_t * row_ids, size_t count, float * output) { + validate_rows(row_ids, count); + if (count != 0 && output == nullptr) { + throw std::invalid_argument("llama_engram: invalid row output"); + } + + std::lock_guard lock(mutex); + for (size_t i = 0; i < count; ++i) { + read_one(row_ids[i], output + i * LLAMA_ENGRAM_DIM); + } + } + + void read_batch(const uint32_t * row_ids, size_t tokens, size_t stride, float * output) { + if (tokens > SIZE_MAX / (LLAMA_ENGRAM_COLS * LLAMA_ENGRAM_DIM * sizeof(float)) || + (tokens != 0 && (row_ids == nullptr || output == nullptr || stride < LLAMA_ENGRAM_COLS)) || + (tokens != 0 && tokens - 1 > (SIZE_MAX / sizeof(*row_ids) - LLAMA_ENGRAM_COLS) / stride)) { + throw std::invalid_argument("llama_engram: invalid batch buffers"); + } + for (size_t token = 0; token < tokens; ++token) { + validate_rows(row_ids + token * stride, LLAMA_ENGRAM_COLS); + } + if (tokens == 0) { + return; + } + + std::lock_guard lock(mutex); + for (size_t start = 0; start < tokens; start += BATCH_TOKENS) { + const size_t count_tokens = std::min(tokens - start, BATCH_TOKENS); + const size_t count_rows = count_tokens * LLAMA_ENGRAM_COLS; + requests.resize(count_rows); + for (size_t i = 0; i < count_rows; ++i) { + requests[i].row = row_ids[(start + i / LLAMA_ENGRAM_COLS) * stride + i % LLAMA_ENGRAM_COLS]; + requests[i].output = (uint32_t) i; + } + std::sort(requests.begin(), requests.end(), [](const request & a, const request & b) { + return a.row < b.row; + }); + + float * chunk = output + start * LLAMA_ENGRAM_COLS * LLAMA_ENGRAM_DIM; + const float * previous = nullptr; + for (size_t i = 0; i < count_rows; ++i) { + float * dst = chunk + (size_t) requests[i].output * LLAMA_ENGRAM_DIM; + if (i != 0 && requests[i].row == requests[i - 1].row) { + memcpy(dst, previous, LLAMA_ENGRAM_DIM * sizeof(*dst)); + } else { + read_one(requests[i].row, dst); + previous = dst; + } + } + } + } +}; + +llama_engram_table::llama_engram_table(const std::string & fname, uint64_t offset, uint32_t rows) + : pimpl(std::make_unique(fname, offset, rows)) {} + +llama_engram_table::~llama_engram_table() = default; + +void llama_engram_table::read(const uint32_t * rows, size_t count, float * output) { + pimpl->read(rows, count, output); +} + +void llama_engram_table::read_batch(const uint32_t * rows, size_t tokens, size_t stride, float * output) { + pimpl->read_batch(rows, tokens, stride, output); +} + +uint32_t llama_engram_table::n_rows() const { + return pimpl->rows; +} + +std::string llama_engram_table::describe() const { + return format("%s @ %llu: %u rows x %u bytes, uncached aligned reads", + pimpl->file.name().c_str(), (unsigned long long) pimpl->offset, + pimpl->rows, LLAMA_ENGRAM_ROW_BYTES); +} diff --git a/src/llama-engram.h b/src/llama-engram.h new file mode 100644 index 000000000000..6c1c0f78b34c --- /dev/null +++ b/src/llama-engram.h @@ -0,0 +1,73 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +enum { + LLAMA_ENGRAM_LAYERS = 2, + LLAMA_ENGRAM_NGRAM = 4, + LLAMA_ENGRAM_HEADS = 8, + LLAMA_ENGRAM_COLS = 24, + LLAMA_ENGRAM_DIM = 256, + LLAMA_ENGRAM_ROW_BYTES = 264, + LLAMA_ENGRAM_DEAD = -1, +}; + +struct llama_engram_layout { + std::string encoding; + std::array layer_ids = {}; + std::vector token_map; + uint32_t compressed_vocab_size = 0; + uint32_t pad_id = 0; + std::array rows = {}; + std::array, LLAMA_ENGRAM_LAYERS> multipliers = {}; + std::array, LLAMA_ENGRAM_LAYERS> primes = {}; +}; + +struct llama_engram_history { + // Newest compressed token first. LLAMA_ENGRAM_DEAD breaks all n-grams that cross it. + std::array tail = {}; + + void reset(); +}; + +struct llama_engram_hasher { + explicit llama_engram_hasher(llama_engram_layout layout); + ~llama_engram_hasher(); + + // Output is [token][layer][column]. Invalid input does not change history or output. + void hash( + llama_engram_history & history, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * rows) const; + + const llama_engram_layout & layout() const; + + struct impl; + std::unique_ptr pimpl; +}; + +// Decode 256 E4M3 values and eight E8M0 scales to BF16-rounded values in F32 storage. +void llama_engram_decode_row(const uint8_t row[LLAMA_ENGRAM_ROW_BYTES], float output[LLAMA_ENGRAM_DIM]); + +struct llama_engram_table { + // The table is never mapped or cached. Construction fails unless uncached reads are available. + llama_engram_table(const std::string & fname, uint64_t offset, uint32_t rows); + ~llama_engram_table(); + + void read(const uint32_t * rows, size_t count, float * output); + // Read 24 rows per token. Input uses row-ID stride; output is packed [token][column][dimension]. + void read_batch(const uint32_t * rows, size_t tokens, size_t stride, float * output); + + uint32_t n_rows() const; + std::string describe() const; + + struct impl; + std::unique_ptr pimpl; +}; diff --git a/src/llama-expert-store.cpp b/src/llama-expert-store.cpp new file mode 100644 index 000000000000..c8597e52845f --- /dev/null +++ b/src/llama-expert-store.cpp @@ -0,0 +1,646 @@ +#include "llama-expert-store.h" + +#include "llama-impl.h" +#include "llama-mmap.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#endif + +// The positional I/O and reservation model is adapted from ggml-org/llama.cpp#25294. +// Leases add the in-flight publication safety described in ggml-org/llama.cpp#27861. + +namespace { + +bool checked_add_u64(uint64_t a, uint64_t b, uint64_t * result) { + if (b > std::numeric_limits::max() - a) { + return false; + } + *result = a + b; + return true; +} + +bool checked_mul_u64(uint64_t a, uint64_t b, uint64_t * result) { + if (a != 0 && b > std::numeric_limits::max() / a) { + return false; + } + *result = a * b; + return true; +} + +bool is_power_of_two(size_t value) { + return value != 0 && (value & (value - 1)) == 0; +} + +struct expert_key { + int32_t layer; + llama_expert_projection projection; + int32_t expert_id; + + bool operator<(const expert_key & other) const { + if (layer != other.layer) { + return layer < other.layer; + } + if (projection != other.projection) { + return projection < other.projection; + } + return expert_id < other.expert_id; + } + + bool operator==(const expert_key & other) const { + return layer == other.layer && projection == other.projection && expert_id == other.expert_id; + } +}; + +struct tensor_key { + int32_t layer; + llama_expert_projection projection; + + bool operator<(const tensor_key & other) const { + if (layer != other.layer) { + return layer < other.layer; + } + return projection < other.projection; + } +}; + +struct aligned_buffer { + uint8_t * data = nullptr; + size_t size = 0; + + aligned_buffer() = default; + + aligned_buffer(size_t size, size_t alignment) { + reset(size, alignment); + } + + aligned_buffer(aligned_buffer && other) noexcept : data(other.data), size(other.size) { + other.data = nullptr; + other.size = 0; + } + + aligned_buffer & operator=(aligned_buffer && other) noexcept { + if (this != &other) { + clear(); + data = other.data; + size = other.size; + other.data = nullptr; + other.size = 0; + } + return *this; + } + + ~aligned_buffer() { + clear(); + } + + aligned_buffer(const aligned_buffer &) = delete; + aligned_buffer & operator=(const aligned_buffer &) = delete; + + void reset(size_t new_size, size_t alignment) { + clear(); + if (new_size == 0) { + return; + } + alignment = std::max(alignment, alignof(void *)); +#if defined(_WIN32) + data = static_cast(_aligned_malloc(new_size, alignment)); + if (data == nullptr) { + throw std::bad_alloc(); + } +#else + void * ptr = nullptr; + if (posix_memalign(&ptr, alignment, new_size) != 0) { + throw std::bad_alloc(); + } + data = static_cast(ptr); +#endif + size = new_size; + } + + void clear() { +#if defined(_WIN32) + _aligned_free(data); +#else + free(data); +#endif + data = nullptr; + size = 0; + } +}; + +struct expert_file { + std::string fname; + uint64_t size = 0; + bool direct = false; + std::unique_ptr file; + + expert_file(const std::string & fname, bool direct_io, bool allow_buffered_io) : fname(fname) { + reopen(direct_io); + if (direct_io && !direct && !allow_buffered_io) { + throw std::runtime_error(format("llama_expert_store: direct I/O is required but unavailable for %s", fname.c_str())); + } + if (direct_io && !direct) { + LLAMA_LOG_WARN("%s: direct I/O is unavailable for %s; using explicitly allowed buffered reads\n", + __func__, fname.c_str()); + } + } + + expert_file(const expert_file &) = delete; + expert_file & operator=(const expert_file &) = delete; + + void reopen(bool direct_io) { + file = std::make_unique(fname.c_str(), "rb", direct_io); + size = file->size(); + direct = direct_io && file->has_direct_io(); + } + + size_t pread_at_least(void * dst, size_t len, uint64_t offset, size_t need) const { + if (need > len) { + throw std::runtime_error("llama_expert_store: invalid read requirement"); + } + size_t total = 0; + while (total < need) { +#if defined(_WIN32) + file->seek(offset + total, SEEK_SET); + file->read_raw(static_cast(dst) + total, len - total); + const size_t n = len - total; +#else + if (offset + total > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("llama_expert_store: file offset exceeds off_t"); + } + const ssize_t result = pread(file->file_id(), static_cast(dst) + total, len - total, + static_cast(offset + total)); + if (result < 0) { + if (errno == EINTR) { + continue; + } + throw std::runtime_error(format("llama_expert_store: pread failed for %s at %llu: %s", + fname.c_str(), (unsigned long long) (offset + total), strerror(errno))); + } + const size_t n = static_cast(result); +#endif + if (n == 0) { + break; + } + total += n; + } + if (total < need) { + throw std::runtime_error(format("llama_expert_store: short read for %s: %zu bytes, need %zu at %llu", + fname.c_str(), total, need, (unsigned long long) offset)); + } + return total; + } +}; + +} + +llama_expert_store_aligned_read llama_expert_store_align_read( + uint64_t offset, size_t size, size_t alignment, uint64_t file_size) { + if (!is_power_of_two(alignment)) { + throw std::runtime_error("llama_expert_store: I/O alignment must be a power of two"); + } + + uint64_t end; + if (!checked_add_u64(offset, size, &end) || end > file_size) { + throw std::runtime_error("llama_expert_store: read is outside the source file"); + } + + const uint64_t aligned_offset = offset & ~static_cast(alignment - 1); + const size_t prefix = static_cast(offset - aligned_offset); + uint64_t needed; + if (!checked_add_u64(prefix, size, &needed)) { + throw std::runtime_error("llama_expert_store: aligned read size overflow"); + } + uint64_t rounded; + if (!checked_add_u64(needed, alignment - 1, &rounded)) { + throw std::runtime_error("llama_expert_store: aligned read size overflow"); + } + rounded &= ~static_cast(alignment - 1); + llama_expert_store_aligned_read result; + result.offset = aligned_offset; + if (rounded > std::numeric_limits::max()) { + throw std::runtime_error("llama_expert_store: aligned read exceeds addressable memory"); + } + result.size = static_cast(rounded); + result.prefix = prefix; + return result; +} + +void llama_expert_store_validate_tensor(const llama_expert_store_tensor & tensor) { + if (tensor.name.empty() || tensor.fname.empty()) { + throw std::runtime_error("llama_expert_store: tensor name and source file are required"); + } + if (tensor.layer < 0) { + throw std::runtime_error(format("llama_expert_store: tensor %s has an invalid layer", tensor.name.c_str())); + } + if (tensor.projection < LLAMA_EXPERT_PROJECTION_GATE || tensor.projection > LLAMA_EXPERT_PROJECTION_DOWN) { + throw std::runtime_error(format("llama_expert_store: tensor %s has an invalid projection", tensor.name.c_str())); + } + const ggml_type expected_type = tensor.projection == LLAMA_EXPERT_PROJECTION_DOWN ? GGML_TYPE_Q2_K : GGML_TYPE_IQ2_XXS; + if (tensor.type != expected_type) { + throw std::runtime_error(format("llama_expert_store: tensor %s must be %s, got %s", + tensor.name.c_str(), ggml_type_name(expected_type), ggml_type_name(tensor.type))); + } + if (tensor.ne[0] <= 0 || tensor.ne[1] <= 0 || tensor.ne[2] <= 0) { + throw std::runtime_error(format("llama_expert_store: tensor %s has invalid dimensions", tensor.name.c_str())); + } + if (tensor.ne[2] > std::numeric_limits::max()) { + throw std::runtime_error(format("llama_expert_store: tensor %s has too many experts", tensor.name.c_str())); + } + if (tensor.ne[0] % ggml_blck_size(tensor.type) != 0) { + throw std::runtime_error(format("llama_expert_store: tensor %s rows are not whole quantization blocks", tensor.name.c_str())); + } + + const size_t row_size = ggml_row_size(tensor.type, tensor.ne[0]); + uint64_t plane_size; + uint64_t tensor_size; + if (!checked_mul_u64(row_size, static_cast(tensor.ne[1]), &plane_size) || + !checked_mul_u64(plane_size, static_cast(tensor.ne[2]), &tensor_size)) { + throw std::runtime_error(format("llama_expert_store: tensor %s size overflows", tensor.name.c_str())); + } + if (tensor.nb[0] != ggml_type_size(tensor.type) || tensor.nb[1] != row_size || tensor.nb[2] != plane_size) { + throw std::runtime_error(format("llama_expert_store: tensor %s is not a contiguous merged-expert tensor", tensor.name.c_str())); + } + uint64_t tensor_end; + if (!checked_add_u64(tensor.file_offset, tensor_size, &tensor_end) || tensor_end > tensor.file_size) { + throw std::runtime_error(format("llama_expert_store: tensor %s is outside the source file", tensor.name.c_str())); + } +} + +struct llama_expert_store::impl { + struct slot { + bool occupied = false; + expert_key key = {}; + const llama_expert_store_tensor * tensor = nullptr; + aligned_buffer bytes; + uint64_t last_use = 0; + uint32_t pins = 0; + }; + + llama_expert_store_params params; + std::map tensors; + std::map> files; + std::vector slots; + size_t bytes_resident = 0; + uint64_t use_clock = 0; + llama_expert_store_stats counters; + mutable std::mutex mutex; + + impl(std::vector tensors, const llama_expert_store_params & params) : params(params) { + if (params.cache_bytes == 0 || params.cache_slots == 0) { + throw std::runtime_error("llama_expert_store: cache byte and slot budgets must be non-zero"); + } + if (params.cache_slots > std::numeric_limits::max()) { + throw std::runtime_error("llama_expert_store: cache slot budget exceeds the slot ID range"); + } + if (!is_power_of_two(params.io_alignment)) { + throw std::runtime_error("llama_expert_store: I/O alignment must be a power of two"); + } + + for (auto & tensor : tensors) { + llama_expert_store_validate_tensor(tensor); + const tensor_key key = { tensor.layer, tensor.projection }; + if (this->tensors.count(key) != 0) { + throw std::runtime_error(format("llama_expert_store: duplicate tensor for layer %d projection %d", + tensor.layer, static_cast(tensor.projection))); + } + if (tensor.nb[2] > params.cache_bytes) { + throw std::runtime_error(format("llama_expert_store: tensor %s expert plane exceeds the cache byte budget", + tensor.name.c_str())); + } + auto file_it = files.find(tensor.fname); + if (file_it == files.end()) { + file_it = files.emplace(tensor.fname, + std::make_unique(tensor.fname, params.direct_io, params.allow_buffered_io)).first; + } + if (file_it->second->size != tensor.file_size) { + throw std::runtime_error(format("llama_expert_store: source file size changed for %s", tensor.fname.c_str())); + } + this->tensors.emplace(key, std::move(tensor)); + } + + if (this->tensors.empty()) { + throw std::runtime_error("llama_expert_store: no tensors registered"); + } + for (auto it = this->tensors.begin(); it != this->tensors.end();) { + const int32_t layer = it->first.layer; + const auto gate = this->tensors.find({ layer, LLAMA_EXPERT_PROJECTION_GATE }); + const auto up = this->tensors.find({ layer, LLAMA_EXPERT_PROJECTION_UP }); + const auto down = this->tensors.find({ layer, LLAMA_EXPERT_PROJECTION_DOWN }); + if (gate == this->tensors.end() || up == this->tensors.end() || down == this->tensors.end()) { + throw std::runtime_error(format("llama_expert_store: layer %d must register gate, up, and down tensors", layer)); + } + if (gate->second.ne[0] != up->second.ne[0] || + gate->second.ne[1] != up->second.ne[1] || + gate->second.ne[2] != up->second.ne[2] || + down->second.ne[0] != gate->second.ne[1] || + down->second.ne[1] != gate->second.ne[0] || + down->second.ne[2] != gate->second.ne[2]) { + throw std::runtime_error(format("llama_expert_store: layer %d expert tensor dimensions do not match", layer)); + } + it = this->tensors.upper_bound({ layer, LLAMA_EXPERT_PROJECTION_DOWN }); + } + slots.resize(params.cache_slots); + } + + const llama_expert_store_tensor & get_tensor(const expert_key & key) const { + const auto it = tensors.find({ key.layer, key.projection }); + if (it == tensors.end()) { + throw std::runtime_error(format("llama_expert_store: no tensor for layer %d projection %d", + key.layer, static_cast(key.projection))); + } + if (key.expert_id < 0 || key.expert_id >= it->second.ne[2]) { + throw std::runtime_error(format("llama_expert_store: expert ID %d is outside [0, %lld)", + key.expert_id, (long long) it->second.ne[2])); + } + return it->second; + } + + aligned_buffer read_expert(const llama_expert_store_tensor & tensor, int32_t expert_id, uint64_t * bytes_read) const { + uint64_t expert_delta; + uint64_t expert_offset; + if (!checked_mul_u64(static_cast(expert_id), tensor.nb[2], &expert_delta) || + !checked_add_u64(tensor.file_offset, expert_delta, &expert_offset)) { + throw std::runtime_error(format("llama_expert_store: expert offset overflow for %s", tensor.name.c_str())); + } + + aligned_buffer payload(tensor.nb[2], params.io_alignment); + auto & file = *files.at(tensor.fname); + if (file.direct) { + const llama_expert_store_aligned_read read = + llama_expert_store_align_read(expert_offset, tensor.nb[2], params.io_alignment, tensor.file_size); + aligned_buffer bounce(read.size, params.io_alignment); + try { + *bytes_read += file.pread_at_least(bounce.data, read.size, read.offset, read.prefix + tensor.nb[2]); + memcpy(payload.data, bounce.data + read.prefix, tensor.nb[2]); + } catch (const std::runtime_error & e) { + if (!params.allow_buffered_io) { + throw std::runtime_error(format("llama_expert_store: direct I/O failed for %s and buffered fallback is disabled: %s", + tensor.fname.c_str(), e.what())); + } + LLAMA_LOG_WARN("%s: direct read failed for %s; retrying with buffered I/O: %s\n", + __func__, tensor.fname.c_str(), e.what()); + file.reopen(false); + *bytes_read += file.pread_at_least(payload.data, tensor.nb[2], expert_offset, tensor.nb[2]); + } + } else { + *bytes_read += file.pread_at_least(payload.data, tensor.nb[2], expert_offset, tensor.nb[2]); + } + + if (!ggml_validate_row_data(tensor.type, payload.data, tensor.nb[2])) { + throw std::runtime_error(format("llama_expert_store: tensor %s expert %d has invalid payload", + tensor.name.c_str(), expert_id)); + } + return payload; + } + + void unpin(const std::vector & slot_ids) { + std::lock_guard lock(mutex); + for (uint32_t slot_id : slot_ids) { + if (slot_id >= slots.size() || slots[slot_id].pins == 0) { + GGML_ABORT("llama_expert_store: invalid lease slot"); + } + slots[slot_id].pins--; + } + } + + std::vector get_payloads(const std::vector & slot_ids) const { + std::lock_guard lock(mutex); + std::vector result; + result.reserve(slot_ids.size()); + for (uint32_t slot_id : slot_ids) { + const slot & entry = slots.at(slot_id); + GGML_ASSERT(entry.occupied && entry.pins > 0); + result.push_back({ + entry.key.layer, + entry.key.projection, + entry.key.expert_id, + slot_id, + entry.tensor->type, + entry.bytes.data, + entry.bytes.size, + }); + } + return result; + } +}; + +struct llama_expert_store::lease::impl { + std::shared_ptr store; + std::vector pinned_slots; + std::vector> remapped_slots; + + ~impl() { + if (store) { + store->unpin(pinned_slots); + } + } +}; + +llama_expert_store::lease::lease() = default; +llama_expert_store::lease::lease(lease && other) noexcept = default; +llama_expert_store::lease & llama_expert_store::lease::operator=(lease && other) noexcept = default; +llama_expert_store::lease::~lease() = default; + +const std::vector> & llama_expert_store::lease::slot_ids() const { + static const std::vector> empty; + return pimpl ? pimpl->remapped_slots : empty; +} + +std::vector llama_expert_store::lease::payloads() const { + return pimpl ? pimpl->store->get_payloads(pimpl->pinned_slots) : std::vector(); +} + +llama_expert_store::llama_expert_store( + std::vector tensors, const llama_expert_store_params & params) + : pimpl(std::make_shared(std::move(tensors), params)) { +} + +llama_expert_store::~llama_expert_store() = default; + +llama_expert_store::lease llama_expert_store::acquire(const std::vector & requests) { + std::lock_guard lock(pimpl->mutex); + + std::vector> request_keys; + std::vector unique_keys; + request_keys.reserve(requests.size()); + for (const auto & request : requests) { + std::vector keys; + keys.reserve(request.expert_ids.size()); + for (int32_t expert_id : request.expert_ids) { + const expert_key key = { request.layer, request.projection, expert_id }; + pimpl->get_tensor(key); + keys.push_back(key); + unique_keys.push_back(key); + } + request_keys.push_back(std::move(keys)); + } + std::sort(unique_keys.begin(), unique_keys.end()); + unique_keys.erase(std::unique(unique_keys.begin(), unique_keys.end()), unique_keys.end()); + + std::map resident; + for (uint32_t i = 0; i < pimpl->slots.size(); ++i) { + if (pimpl->slots[i].occupied) { + resident.emplace(pimpl->slots[i].key, i); + } + } + + std::vector misses; + std::vector hit_slots; + size_t miss_bytes = 0; + for (const expert_key & key : unique_keys) { + const auto hit = resident.find(key); + if (hit != resident.end()) { + hit_slots.push_back(hit->second); + continue; + } + const auto & tensor = pimpl->get_tensor(key); + if (miss_bytes > pimpl->params.cache_bytes || tensor.nb[2] > pimpl->params.cache_bytes - miss_bytes) { + throw std::runtime_error("llama_expert_store: requested expert union exceeds the cache byte budget"); + } + miss_bytes += tensor.nb[2]; + misses.push_back(key); + } + + std::vector empty_slots; + std::vector candidates; + std::sort(hit_slots.begin(), hit_slots.end()); + for (uint32_t i = 0; i < pimpl->slots.size(); ++i) { + const auto & entry = pimpl->slots[i]; + if (!entry.occupied) { + empty_slots.push_back(i); + } else if (entry.pins == 0 && !std::binary_search(hit_slots.begin(), hit_slots.end(), i)) { + candidates.push_back(i); + } + } + std::sort(candidates.begin(), candidates.end(), [&](uint32_t a, uint32_t b) { + const auto & lhs = pimpl->slots[a]; + const auto & rhs = pimpl->slots[b]; + if (lhs.last_use != rhs.last_use) { + return lhs.last_use < rhs.last_use; + } + return a < b; + }); + + const size_t min_victims = misses.size() > empty_slots.size() ? misses.size() - empty_slots.size() : 0; + std::vector victims; + if (miss_bytes > std::numeric_limits::max() - pimpl->bytes_resident) { + throw std::runtime_error("llama_expert_store: cache byte accounting overflow"); + } + size_t bytes_after = pimpl->bytes_resident + miss_bytes; + for (uint32_t candidate : candidates) { + if (victims.size() >= min_victims && bytes_after <= pimpl->params.cache_bytes) { + break; + } + victims.push_back(candidate); + bytes_after -= pimpl->slots[candidate].bytes.size; + } + if (victims.size() < min_victims || bytes_after > pimpl->params.cache_bytes) { + throw std::runtime_error("llama_expert_store: cache capacity is pinned or too small for the requested expert union"); + } + + std::vector target_slots = empty_slots; + target_slots.insert(target_slots.end(), victims.begin(), victims.end()); + std::sort(target_slots.begin(), target_slots.end()); + + uint64_t bytes_read = 0; + std::vector staged; + staged.reserve(misses.size()); + for (const expert_key & key : misses) { + const auto & tensor = pimpl->get_tensor(key); + staged.push_back(pimpl->read_expert(tensor, key.expert_id, &bytes_read)); + } + + for (uint32_t victim : victims) { + auto & entry = pimpl->slots[victim]; + resident.erase(entry.key); + pimpl->bytes_resident -= entry.bytes.size; + entry.bytes.clear(); + entry.occupied = false; + entry.tensor = nullptr; + entry.last_use = 0; + pimpl->counters.evictions++; + } + + for (size_t i = 0; i < misses.size(); ++i) { + const expert_key & key = misses[i]; + const auto & tensor = pimpl->get_tensor(key); + auto & entry = pimpl->slots[target_slots[i]]; + entry.bytes = std::move(staged[i]); + entry.occupied = true; + entry.key = key; + entry.tensor = &tensor; + entry.pins = 0; + pimpl->bytes_resident += entry.bytes.size; + resident[entry.key] = target_slots[i]; + } + + auto lease_impl = std::make_unique(); + lease_impl->remapped_slots.reserve(request_keys.size()); + for (const auto & keys : request_keys) { + std::vector remapped; + remapped.reserve(keys.size()); + for (const expert_key & key : keys) { + remapped.push_back(resident.at(key)); + } + lease_impl->remapped_slots.push_back(std::move(remapped)); + } + lease_impl->pinned_slots.reserve(unique_keys.size()); + for (const expert_key & key : unique_keys) { + lease_impl->pinned_slots.push_back(resident.at(key)); + } + for (uint32_t slot_id : lease_impl->pinned_slots) { + auto & entry = pimpl->slots[slot_id]; + entry.last_use = ++pimpl->use_clock; + entry.pins++; + } + lease_impl->store = pimpl; + + pimpl->counters.hits += unique_keys.size() - misses.size(); + pimpl->counters.misses += misses.size(); + pimpl->counters.bytes_read += bytes_read; + + lease result; + result.pimpl = std::move(lease_impl); + return result; +} + +llama_expert_store_stats llama_expert_store::stats() const { + std::lock_guard lock(pimpl->mutex); + return pimpl->counters; +} + +size_t llama_expert_store::resident_bytes() const { + std::lock_guard lock(pimpl->mutex); + return pimpl->bytes_resident; +} + +size_t llama_expert_store::resident_entries() const { + std::lock_guard lock(pimpl->mutex); + size_t result = 0; + for (const auto & slot : pimpl->slots) { + result += slot.occupied ? 1 : 0; + } + return result; +} + +bool llama_expert_store::direct_io_active() const { + std::lock_guard lock(pimpl->mutex); + for (const auto & item : pimpl->files) { + if (!item.second->direct) { + return false; + } + } + return true; +} diff --git a/src/llama-expert-store.h b/src/llama-expert-store.h new file mode 100644 index 000000000000..10ce3a1a19d7 --- /dev/null +++ b/src/llama-expert-store.h @@ -0,0 +1,106 @@ +#pragma once + +#include "ggml.h" + +#include +#include +#include +#include +#include + +enum llama_expert_projection { + LLAMA_EXPERT_PROJECTION_GATE = 0, + LLAMA_EXPERT_PROJECTION_UP, + LLAMA_EXPERT_PROJECTION_DOWN, +}; + +struct llama_expert_store_tensor { + std::string name; + std::string fname; + size_t file_index = 0; + int32_t layer = -1; + llama_expert_projection projection = LLAMA_EXPERT_PROJECTION_GATE; + ggml_type type = GGML_TYPE_COUNT; + int64_t ne[3] = {}; + size_t nb[3] = {}; + uint64_t file_offset = 0; + uint64_t file_size = 0; +}; + +struct llama_expert_store_params { + size_t cache_bytes = 0; + size_t cache_slots = 0; + size_t io_alignment = 4096; + bool direct_io = true; + bool allow_buffered_io = false; // opt-in only; page-cache bytes are outside cache_bytes +}; + +struct llama_expert_store_request { + int32_t layer = -1; + llama_expert_projection projection = LLAMA_EXPERT_PROJECTION_GATE; + std::vector expert_ids; +}; + +struct llama_expert_store_stats { + uint64_t hits = 0; + uint64_t misses = 0; + uint64_t bytes_read = 0; + uint64_t evictions = 0; +}; + +struct llama_expert_store_aligned_read { + uint64_t offset = 0; + size_t size = 0; + size_t prefix = 0; +}; + +llama_expert_store_aligned_read llama_expert_store_align_read( + uint64_t offset, size_t size, size_t alignment, uint64_t file_size); + +void llama_expert_store_validate_tensor(const llama_expert_store_tensor & tensor); + +struct llama_expert_store { + struct payload { + int32_t layer = -1; + llama_expert_projection projection = LLAMA_EXPERT_PROJECTION_GATE; + int32_t expert_id = -1; + uint32_t slot_id = 0; + ggml_type type = GGML_TYPE_COUNT; + const uint8_t * data = nullptr; + size_t size = 0; + }; + + struct lease { + lease(); + lease(lease && other) noexcept; + lease & operator=(lease && other) noexcept; + ~lease(); + + lease(const lease &) = delete; + lease & operator=(const lease &) = delete; + + const std::vector> & slot_ids() const; + std::vector payloads() const; + + private: + friend struct llama_expert_store; + + struct impl; + std::unique_ptr pimpl; + }; + + llama_expert_store(std::vector tensors, const llama_expert_store_params & params); + ~llama_expert_store(); + + // The lease pins every unique returned slot. Keep it until the backend upload completes. + lease acquire(const std::vector & requests); + + llama_expert_store_stats stats() const; + size_t resident_bytes() const; + size_t resident_entries() const; + bool direct_io_active() const; + +private: + struct impl; + std::shared_ptr pimpl; +}; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 9d5f5e6241ea..df46edf3e416 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1577,15 +1577,17 @@ ggml_tensor * llm_graph_context::build_lora_mm_id( ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b ggml_tensor * ids, - ggml_tensor * w_s) const { + ggml_tensor * w_s, + ggml_tensor * ids_scale) const { ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids); if (w_s) { + ids_scale = ids_scale ? ids_scale : ids; const int64_t n_expert = w_s->ne[0]; const int64_t n_tokens = cur->ne[2]; ggml_tensor * s = ggml_reshape_3d(ctx0, w_s, 1, n_expert, 1); s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, ids); + s = ggml_get_rows(ctx0, s, ids_scale); res = ggml_mul(ctx0, res, s); } for (const auto & lora : *loras) { @@ -1862,7 +1864,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); @@ -1977,6 +1980,26 @@ ggml_tensor * llm_graph_context::build_ffn( return cur; } +llm_moe_expert_ids llm_build_moe_expert_ids( + ggml_context * ctx, + llm_arch arch, + ggml_tensor * selected, + ggml_tensor * lookup, + int64_t n_expert, + int64_t n_expert_total, + uint32_t n_group_experts) { + if (lookup != nullptr) { + return { selected, lookup }; + } + if (arch == LLM_ARCH_GROVEMOE && n_expert != n_expert_total) { + GGML_ASSERT(n_group_experts > 0); + // TODO: Use scalar div instead when/if implemented + ggml_tensor * f_sel = ggml_cast(ctx, selected, GGML_TYPE_F32); + lookup = ggml_cast(ctx, ggml_scale(ctx, f_sel, 1.0f / float(n_group_experts)), GGML_TYPE_I32); + } + return { selected, lookup ? lookup : selected }; +} + ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * cur, ggml_tensor * gate_inp, @@ -1996,7 +2019,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, - ggml_tensor * selected_experts_in) const { + ggml_tensor * selected_experts_in, + ggml_tensor * selected_experts_lookup_in) const { return build_moe_ffn( cur, gate_inp, /* gate_inp_b */ nullptr, @@ -2017,7 +2041,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( up_exps_s, gate_exps_s, down_exps_s, - selected_experts_in + selected_experts_in, + selected_experts_lookup_in ); } @@ -2045,7 +2070,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, - ggml_tensor * selected_experts_in) const { + ggml_tensor * selected_experts_in, + ggml_tensor * selected_experts_lookup_in) const { const int64_t n_embd = cur->ne[0]; const int64_t n_tokens = cur->ne[1]; const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4; // for llama4, we apply the sigmoid-ed weights before the FFN @@ -2142,10 +2168,13 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } cb(selected_experts, "ffn_moe_topk", il); + const llm_moe_expert_ids expert_ids = llm_build_moe_expert_ids( + ctx0, arch, selected_experts, selected_experts_lookup_in, + n_expert, hparams.n_expert, hparams.n_group_experts); + selected_experts = expert_ids.routing; + ggml_tensor * selected_experts_lookup = expert_ids.lookup; + if (arch == LLM_ARCH_GROVEMOE && n_expert != hparams.n_expert) { - // TODO: Use scalar div instead when/if implemented - ggml_tensor * f_sel = ggml_cast(ctx0, selected_experts, GGML_TYPE_F32); - selected_experts = ggml_cast(ctx0, ggml_scale(ctx0, f_sel, 1.0f / float(hparams.n_group_experts)), GGML_TYPE_I32); probs = ggml_reshape_3d(ctx0, probs, 1, hparams.n_expert, n_tokens); } else { probs = ggml_reshape_3d(ctx0, probs, 1, n_expert, n_tokens); @@ -2199,7 +2228,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( if (gate_up_exps) { // merged gate_up path: one mul_mat_id, then split into gate and up views - ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s); // [n_ff*2, n_expert_used, n_tokens] + ggml_tensor * gate_up = build_lora_mm_id( + gate_up_exps, cur, selected_experts_lookup, up_exps_s, selected_experts); // [n_ff*2, n_expert_used, n_tokens] cb(gate_up, "ffn_moe_gate_up", il); if (up_exps_s) { @@ -2218,7 +2248,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(up, "ffn_moe_up", il); } else { // separate gate and up path - up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s); // [n_ff, n_expert_used, n_tokens] + up = build_lora_mm_id( + up_exps, cur, selected_experts_lookup, up_exps_s, selected_experts); // [n_ff, n_expert_used, n_tokens] cb(up, "ffn_moe_up", il); if (up_exps_s) { @@ -2231,7 +2262,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } if (gate_exps) { - cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s); // [n_ff, n_expert_used, n_tokens] + cur = build_lora_mm_id( + gate_exps, cur, selected_experts_lookup, gate_exps_s, selected_experts); // [n_ff, n_expert_used, n_tokens] cb(cur, "ffn_moe_gate", il); } else { cur = up; @@ -2256,7 +2288,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_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) { + if (arch == LLM_ARCH_MAPLE || 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); @@ -2332,7 +2365,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( GGML_ABORT("fatal error"); } - experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens] + experts = build_lora_mm_id( + down_exps, cur, selected_experts_lookup, down_exps_s, selected_experts); // [n_embd, n_expert_used, n_tokens] cb(experts, "ffn_moe_down", il); if (down_exps_s) { diff --git a/src/llama-graph.h b/src/llama-graph.h index ab9e1664d3ff..f6e3127516b3 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -75,6 +75,20 @@ enum llm_norm_type { LLM_NORM_GROUP, }; +struct llm_moe_expert_ids { + ggml_tensor * routing; + ggml_tensor * lookup; +}; + +llm_moe_expert_ids llm_build_moe_expert_ids( + ggml_context * ctx, + llm_arch arch, + ggml_tensor * selected, + ggml_tensor * lookup, + int64_t n_expert, + int64_t n_expert_total, + uint32_t n_group_experts); + // TODO: tmp - need something better to pass the data from the encoder to the decoder struct llama_cross { // the output embeddings from the encoder as a ggml tensor @@ -1067,7 +1081,8 @@ struct llm_graph_context { ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b ggml_tensor * ids, - ggml_tensor * w_s = nullptr) const; + ggml_tensor * w_s = nullptr, + ggml_tensor * ids_scale = nullptr) const; ggml_tensor * build_norm( ggml_tensor * cur, @@ -1136,7 +1151,8 @@ struct llm_graph_context { ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, ggml_tensor * down_exps_s = nullptr, - ggml_tensor * selected_experts_in = nullptr) const; + ggml_tensor * selected_experts_in = nullptr, + ggml_tensor * selected_experts_lookup_in = nullptr) const; ggml_tensor * build_moe_ffn( ggml_tensor * cur, @@ -1162,7 +1178,8 @@ struct llm_graph_context { ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, ggml_tensor * down_exps_s = nullptr, - ggml_tensor * selected_experts_in = nullptr) const; + ggml_tensor * selected_experts_in = nullptr, + ggml_tensor * selected_experts_lookup_in = nullptr) const; // // inputs diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 34b3c688019a..b3d49ac482f7 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -1,5 +1,6 @@ #include "llama-hparams.h" +#include "llama-dsv41.h" #include "ggml.h" #include @@ -47,6 +48,28 @@ bool llama_hparams::is_swa_any() const { return false; } +int32_t llama_hparams::dsv41_kv_source(uint32_t il) const { + if (il >= n_layer()) { + GGML_ABORT("fatal error"); + } + return dsv41_kv_source_layer[il]; +} + +int32_t llama_hparams::dsv41_index_source(uint32_t il) const { + if (il >= n_layer()) { + GGML_ABORT("fatal error"); + } + return dsv41_index_source_layer[il]; +} + +bool llama_hparams::dsv41_is_kv_source(uint32_t il) const { + return dsv41_kv_source(il) == (int32_t) il; +} + +bool llama_hparams::dsv41_is_index_source(uint32_t il) const { + return dsv41_index_source(il) == (int32_t) il; +} + uint32_t llama_hparams::n_head(uint32_t il) const { if (il < n_layer_all) { return n_head_arr[il]; diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 3afa49ebe861..268e28809494 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -294,6 +294,19 @@ struct llama_hparams { float dsv4_hc_eps = 0.0f; std::array dsv4_compress_ratios; + // DeepSeek-V4.1 + uint32_t dsv41_candidate_source_layer = 0; + uint32_t dsv41_candidate_topk_blocks = 0; + uint32_t dsv41_candidate_block_size = 0; + std::array dsv41_kv_source_layer; + std::array dsv41_index_source_layer; + std::bitset dsv41_engram_layers; + + int32_t dsv41_kv_source(uint32_t il) const; + int32_t dsv41_index_source(uint32_t il) const; + bool dsv41_is_kv_source(uint32_t il) const; + bool dsv41_is_index_source(uint32_t il) const; + // 0 = full rank (DeepSeek-V4) uint32_t hc_low_rank = 0; 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..68bcc1161b98 --- /dev/null +++ b/src/llama-memory-dsv41.cpp @@ -0,0 +1,1562 @@ +#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.attach_no_alloc_buffers = config.no_alloc; + config.engram_enabled = model.hparams.dsv41_engram_layers.any(); + 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 rollback_states; + + explicit impl(llama_dsv41_memory_config config) : config(std::move(config)) { + this->config.engram_enabled = + this->config.engram_enabled || this->config.engram != nullptr; + 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); + ggml_backend_buffer_t buffer = nullptr; + if (this->config.no_alloc && this->config.attach_no_alloc_buffers) { + buffer = ggml_backend_buft_alloc_buffer(group.buft, 0); + for (ggml_tensor * tensor = ggml_get_first_tensor(group.ctx.get()); + tensor != nullptr; + tensor = ggml_get_next_tensor(group.ctx.get(), tensor)) { + tensor->buffer = buffer; + } + } else if (!this->config.no_alloc) { + buffer = + ggml_backend_alloc_ctx_tensors_from_buft(group.ctx.get(), group.buft); + } + if (buffer == nullptr && (!this->config.no_alloc || this->config.attach_no_alloc_buffers)) { + throw std::runtime_error("failed to allocate DeepSeek V4.1 memory buffer"); + } + if (buffer != nullptr) { + group.buffer.reset(buffer); + } + if (buffer != nullptr && !this->config.no_alloc) { + 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->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->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); + } + 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->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->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) { + 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; + 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) { + 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; + } + 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"); + } + 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); + } + } + 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->rollback_states.clear(); + } else { + 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; +} + +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_enabled; +} + +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]); + + auto rollback_slot = mem->pimpl->rollback_states.end(); + bool inserted_rollback = false; + try { + std::tie(rollback_slot, inserted_rollback) = + mem->pimpl->rollback_states.emplace(seq_id, llama_memory_dsv41::impl::rollback_state {}); + if (transaction->engram) { + mem->pimpl->config.engram->commit(*transaction->engram); + } + } catch (...) { + 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); + 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_enabled; + 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..a66a09ba31ea --- /dev/null +++ b/src/llama-memory-dsv41.h @@ -0,0 +1,176 @@ +#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 attach_no_alloc_buffers = false; + bool engram_enabled = 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; + size_t retained_rollback_count() 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-mmap.cpp b/src/llama-mmap.cpp index 715a6e3548e6..141640f9d192 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -172,7 +173,13 @@ struct llama_file::impl { } bool has_direct_io() const { - return true; + // Windows uses cached CRT I/O until FILE_FLAG_NO_BUFFERING support is added. + return false; + } + + void discard_cache(size_t offset, size_t length) const { + GGML_UNUSED(offset); + GGML_UNUSED(length); } ~impl() { @@ -374,6 +381,19 @@ struct llama_file::impl { return fd != -1 && alignment > 1; } + void discard_cache(size_t offset, size_t length) const { +#if defined(POSIX_FADV_DONTNEED) + const int file_id = fd == -1 ? fileno(fp) : fd; + const int result = posix_fadvise(file_id, offset, length, POSIX_FADV_DONTNEED); + if (result != 0) { + LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_DONTNEED) failed: %s\n", strerror(result)); + } +#else + GGML_UNUSED(offset); + GGML_UNUSED(length); +#endif + } + ~impl() { if (fd != -1) { close(fd); @@ -408,6 +428,7 @@ size_t llama_file::size() const { return pimpl->size; } size_t llama_file::read_alignment() const { return pimpl->read_alignment(); } bool llama_file::has_direct_io() const { return pimpl->has_direct_io(); } +void llama_file::discard_cache(size_t offset, size_t length) const { pimpl->discard_cache(offset, length); } int llama_file::file_id() const { #ifdef _WIN32 @@ -439,7 +460,6 @@ void llama_file::write_u32(uint32_t val) const { pimpl->write_u32(val); } // llama_mmap -#if defined(_POSIX_MAPPED_FILES) || defined(_WIN32) // merge `ranges` and return their complement within [0, limit) static llama_mmap::ranges ranges_complement(llama_mmap::ranges ranges, size_t limit) { llama_mmap::ranges res; @@ -457,27 +477,38 @@ static llama_mmap::ranges ranges_complement(llama_mmap::ranges ranges, size_t li if (pos < limit) { res.emplace_back(pos, limit); } - return res; } -#endif struct llama_mmap::impl { #ifdef _POSIX_MAPPED_FILES std::vector> mapped_fragments; - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { size = file->size(); int fd = file->file_id(); int flags = MAP_SHARED; if (numa) { prefetch = 0; } #ifdef __linux__ - if (posix_fadvise(fd, 0, 0, POSIX_FADV_SEQUENTIAL)) { - LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_SEQUENTIAL) failed: %s\n", - strerror(errno)); + const bool sequential = llama_mmap::use_sequential_file_advice(strict_exclusion); + const int file_advice = sequential ? POSIX_FADV_SEQUENTIAL : POSIX_FADV_RANDOM; + const int advice_error = file_advice_override ? + file_advice_override(fd, file_advice) : posix_fadvise(fd, 0, 0, file_advice); + if (advice_error) { + if (strict_exclusion) { + throw std::runtime_error(format( + "posix_fadvise(.., POSIX_FADV_RANDOM) failed for external tensor mapping: %s", + strerror(advice_error))); + } + LLAMA_LOG_WARN("warning: posix_fadvise(.., %s) failed: %s\n", + sequential ? "POSIX_FADV_SEQUENTIAL" : "POSIX_FADV_RANDOM", strerror(advice_error)); } - // MAP_POPULATE would fault in the lazy ranges too - if (prefetch && lazy_ranges.empty()) { flags |= MAP_POPULATE; } + // MAP_POPULATE would fault in excluded ranges too + if (prefetch && excluded_ranges.empty()) { flags |= MAP_POPULATE; } +#else + GGML_UNUSED(file_advice_override); #endif addr = mmap(NULL, file->size(), PROT_READ, flags, fd, 0); if (addr == MAP_FAILED) { @@ -497,12 +528,11 @@ struct llama_mmap::impl { } }; - if (prefetch > 0) { - for (const auto & range : ranges_complement(lazy_ranges, std::min(file->size(), prefetch))) { - advise(range.first, range.second, POSIX_MADV_WILLNEED, "POSIX_MADV_WILLNEED"); - } + for (const auto & range : + llama_mmap::planned_prefetch_ranges(file->size(), prefetch, excluded_ranges, strict_exclusion)) { + advise(range.first, range.second, POSIX_MADV_WILLNEED, "POSIX_MADV_WILLNEED"); } - for (const auto & range : lazy_ranges) { + for (const auto & range : excluded_ranges) { advise(range.first, range.second, POSIX_MADV_RANDOM, "POSIX_MADV_RANDOM"); } if (numa) { @@ -573,8 +603,11 @@ struct llama_mmap::impl { #elif defined(_WIN32) HANDLE hMapping = nullptr; - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { GGML_UNUSED(numa); + GGML_UNUSED(file_advice_override); size = file->size(); @@ -604,7 +637,8 @@ struct llama_mmap::impl { if (pPrefetchVirtualMemory) { std::vector entries; - for (const auto & range : ranges_complement(lazy_ranges, std::min(size, prefetch))) { + for (const auto & range : + llama_mmap::planned_prefetch_ranges(size, prefetch, excluded_ranges, strict_exclusion)) { WIN32_MEMORY_RANGE_ENTRY entry; entry.VirtualAddress = (char *) addr + range.first; entry.NumberOfBytes = (SIZE_T) (range.second - range.first); @@ -642,11 +676,15 @@ struct llama_mmap::impl { } } #else - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { GGML_UNUSED(file); GGML_UNUSED(prefetch); GGML_UNUSED(numa); - GGML_UNUSED(lazy_ranges); + GGML_UNUSED(excluded_ranges); + GGML_UNUSED(strict_exclusion); + GGML_UNUSED(file_advice_override); throw std::runtime_error("mmap not supported"); } @@ -664,9 +702,22 @@ struct llama_mmap::impl { }; llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa, - const ranges & lazy_ranges) : pimpl(std::make_unique(file, prefetch, numa, lazy_ranges)) {} + const ranges & excluded_ranges, bool strict_exclusion, file_advice_override file_advice) : + pimpl(std::make_unique(file, prefetch, numa, excluded_ranges, strict_exclusion, file_advice)) {} llama_mmap::~llama_mmap() = default; +bool llama_mmap::use_sequential_file_advice(bool strict_exclusion) { + return !strict_exclusion; +} + +llama_mmap::ranges llama_mmap::planned_prefetch_ranges( + size_t file_size, size_t prefetch, const ranges & excluded_ranges, bool strict_exclusion) { + if (strict_exclusion || prefetch == 0) { + return {}; + } + return ranges_complement(excluded_ranges, std::min(file_size, prefetch)); +} + size_t llama_mmap::size() const { return pimpl->size; } void * llama_mmap::addr() const { return pimpl->addr; } @@ -782,6 +833,18 @@ struct llama_mlock::impl { impl() : addr(NULL), size(0), failed_already(false) {} + static void align_range(size_t * first, size_t * last) { + const size_t granularity = lock_granularity(); + *first &= ~(granularity - 1); + const size_t remainder = *last & (granularity - 1); + if (remainder != 0) { + if (*last > std::numeric_limits::max() - (granularity - remainder)) { + throw std::runtime_error("mlock range overflow"); + } + *last += granularity - remainder; + } + } + void init(void * ptr) { GGML_ASSERT(addr == NULL && size == 0); addr = ptr; @@ -814,6 +877,7 @@ llama_mlock::~llama_mlock() = default; void llama_mlock::init(void * ptr) { pimpl->init(ptr); } void llama_mlock::grow_to(size_t target_size) { pimpl->grow_to(target_size); } +void llama_mlock::align_range(size_t * first, size_t * last) { impl::align_range(first, last); } #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32) const bool llama_mlock::SUPPORTED = true; diff --git a/src/llama-mmap.h b/src/llama-mmap.h index cc28c8a73fa5..c64c5acfb6a1 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -31,6 +31,8 @@ struct llama_file { void read_aligned_chunk(void * dest, size_t size); uint32_t read_u32(); + void discard_cache(size_t offset, size_t length) const; + void write_raw(const void * ptr, size_t len) const; void write_u32(uint32_t val) const; @@ -44,10 +46,12 @@ struct llama_file { struct llama_mmap { // list of [first, last) byte ranges within a file using ranges = std::vector>; + using file_advice_override = int (*)(int fd, int advice); llama_mmap(const llama_mmap &) = delete; llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false, - const ranges & lazy_ranges = {}); + const ranges & excluded_ranges = {}, bool strict_exclusion = false, + file_advice_override file_advice = nullptr); ~llama_mmap(); size_t size() const; @@ -56,6 +60,9 @@ struct llama_mmap { void unmap_fragment(size_t first, size_t last); static const bool SUPPORTED; + static bool use_sequential_file_advice(bool strict_exclusion); + static ranges planned_prefetch_ranges( + size_t file_size, size_t prefetch, const ranges & excluded_ranges, bool strict_exclusion); private: struct impl; @@ -69,6 +76,7 @@ struct llama_mlock { void init(void * ptr); void grow_to(size_t target_size); + static void align_range(size_t * first, size_t * last); static const bool SUPPORTED; private: diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index fde86b19541f..420656116cf1 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include static const size_t kiB = 1024; @@ -429,6 +430,7 @@ namespace GGUFMeta { template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); @@ -724,6 +726,41 @@ llama_model_loader::llama_model_loader( } else { get_key(llm_kv(LLM_KV_GENERAL_ARCHITECTURE), arch_name, false); llm_kv = LLM_KV(llm_arch_from_string(arch_name)); + if (no_alloc && gguf_get_n_tensors(metadata) > 0) { + const int64_t tensor_count = gguf_get_n_tensors(metadata); + const size_t overhead = ggml_tensor_overhead(); + if ((uint64_t) tensor_count > std::numeric_limits::max()/overhead - 1) { + throw std::runtime_error("no-allocation metadata tensor count overflows the context size"); + } + ggml_init_params params = { + /*.mem_size =*/ overhead*((size_t) tensor_count + 1), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + if (ctx == nullptr) { + throw std::runtime_error("failed to create no-allocation metadata tensor context"); + } + contexts.emplace_back(ctx); + for (int64_t i = 0; i < tensor_count; ++i) { + ggml_tensor * tensor = ggml_new_tensor( + ctx, + gguf_get_tensor_type(metadata, i), + GGML_MAX_DIMS, + gguf_get_tensor_ne(metadata, i)); + ggml_set_name(tensor, gguf_get_tensor_name(metadata, i)); + n_elements += ggml_nelements(tensor); + n_bytes += ggml_nbytes(tensor); + const auto inserted = weights_map.emplace( + ggml_get_name(tensor), + llama_tensor_weight(0, metadata, tensor)); + if (!inserted.second) { + throw std::runtime_error(format( + "invalid model: tensor '%s' is duplicated", + ggml_get_name(tensor))); + } + } + } } n_kv = gguf_get_n_kv(metadata); @@ -1131,6 +1168,14 @@ bool llama_model_loader::lazy_read::add(const std::string & name, const ggml_ten return true; } +void llama_model_loader::external_read::add(const llama_tensor_weight & w) { + const std::string name = ggml_get_name(w.tensor); + if (!tensors.insert(name).second) { + throw std::runtime_error(format("external tensor '%s' is already registered", name.c_str())); + } + ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(w.tensor)); +} + struct ggml_tensor * llama_model_loader::create_tensor( const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { @@ -1197,7 +1242,9 @@ struct ggml_tensor * llama_model_loader::create_tensor( const size_t nbytes = ggml_nbytes(t_meta); LLAMA_LOG_WARN("model has unused tensor %s (size = %zu bytes) -- ignoring\n", tn.str().c_str(), nbytes); - size_data -= nbytes; + if (!files.empty()) { + size_data -= nbytes; + } n_created++; return nullptr; @@ -1311,12 +1358,6 @@ struct ggml_tensor * llama_model_loader::create_tensor( if (flags & TENSOR_SKIP_IF_VIRTUAL) { return nullptr; } - ggml_type type = GGML_TYPE_F32; - const int64_t tid = gguf_find_tensor(metadata, tn.str().c_str()); - if (tid != -1) { - type = gguf_get_tensor_type(metadata, tid); - } - // for tensors that are not required some of the dimensions can be invalid: if (flags & TENSOR_NOT_REQUIRED) { for (size_t dim = 0; dim < ne.size(); dim++) { @@ -1326,6 +1367,15 @@ struct ggml_tensor * llama_model_loader::create_tensor( } } + ggml_type type = GGML_TYPE_F32; + const int64_t tid = gguf_find_tensor(metadata, tn.str().c_str()); + if (tid != -1) { + const ggml_tensor * declared = check_tensor_dims( + tn.str(), ne, true, flags & TENSOR_ALLOW_RESHAPE); + GGML_ASSERT(declared != nullptr); + type = declared->type; + } + ggml_tensor t_meta; memset(&t_meta, 0, sizeof(ggml_tensor)); t_meta.type = type; @@ -1344,10 +1394,15 @@ struct ggml_tensor * llama_model_loader::create_tensor( ggml_set_name(&t_meta, tn.str().c_str()); ggml_backend_buffer_type_t buft = buft_for_tensor(&t_meta); - GGML_ASSERT(buft != nullptr); + if (buft == nullptr) { + return nullptr; + } ggml_context * ctx = ctx_for_buft(buft); ggml_tensor * ret = ggml_dup_tensor(ctx, &t_meta); ggml_set_name(ret, tn.str().c_str()); + if (tid != -1 && !(flags & TENSOR_DUPLICATED)) { + n_created++; + } return ret; } @@ -1407,6 +1462,49 @@ struct ggml_tensor * llama_model_loader::create_tensor( return tensor; } +llama_expert_store_tensor llama_model_loader::register_external_tensor( + const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne) { + const ggml_tensor * tensor = check_tensor_dims(name, ne, true, false); + GGML_ASSERT(tensor != nullptr); + if (tensor->ne[3] != 1) { + throw std::runtime_error(format("external tensor '%s' must have three dimensions", name.c_str())); + } + + const llama_tensor_weight & weight = require_weight(name.c_str()); + if (!no_alloc && (fnames.at(weight.idx).empty() || fnames.at(weight.idx) == "(file*)")) { + throw std::runtime_error(format("external tensor '%s' requires a reopenable source file", name.c_str())); + } + llama_expert_store_tensor result; + result.name = name; + result.fname = no_alloc ? "(no_alloc)" : fnames.at(weight.idx); + result.file_index = weight.idx; + result.layer = layer; + result.projection = projection; + result.type = tensor->type; + for (size_t i = 0; i < 3; ++i) { + result.ne[i] = tensor->ne[i]; + result.nb[i] = tensor->nb[i]; + } + result.file_offset = weight.offs; + if (no_alloc) { + const size_t tensor_size = ggml_nbytes(tensor); + if (result.file_offset > std::numeric_limits::max() - tensor_size) { + throw std::runtime_error(format("external tensor '%s' extent overflows the file size", name.c_str())); + } + result.file_size = result.file_offset + tensor_size; + } else { + result.file_size = files.at(weight.idx)->size(); + } + + llama_expert_store_validate_tensor(result); + external.add(weight); + n_created++; + return result; +} + void llama_model_loader::done_getting_tensors(bool partial) const { if (n_created > n_tensors) { throw std::runtime_error(format("%s: too many tensors created; expected %d, got %d", __func__, n_tensors, n_created)); @@ -1415,7 +1513,7 @@ void llama_model_loader::done_getting_tensors(bool partial) const { if (!partial) { throw std::runtime_error(format("%s: wrong number of tensors; expected %d, got %d", __func__, n_tensors, n_created)); } - LLAMA_LOG_INFO("%s: partial load — used %d of %d tensors in the file (rest belong to a sibling model on the same .gguf)\n", + LLAMA_LOG_INFO("%s: partial load - used %d of %d tensors in the file (rest belong to a sibling model on the same .gguf)\n", __func__, n_created, n_tensors); } if (n_tensors_moved > 0) { @@ -1446,10 +1544,17 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps const size_t prefetch_size = prefetch && use_mmap ? -1 : 0; - std::unique_ptr mapping = std::make_unique(file.get(), prefetch_size, is_numa, - lazy.for_file(idx)); + llama_mmap::ranges excluded = lazy.for_file(idx); + const auto & external_ranges = external.for_file(idx); + excluded.insert(excluded.end(), external_ranges.begin(), external_ranges.end()); + + std::unique_ptr mapping = std::make_unique( + file.get(), prefetch_size, is_numa, excluded, !external_ranges.empty()); + for (const auto & range : external_ranges) { + mapping->unmap_fragment(range.first, range.second); + } mmaps_used.emplace_back(mapping->size(), 0); - if (mlock_mmaps) { + if (mlock_mmaps && external_ranges.empty()) { std::unique_ptr mlock_mmap(new llama_mlock()); mlock_mmap->init(mapping->addr()); mlock_mmaps->emplace_back(std::move(mlock_mmap)); @@ -1460,7 +1565,9 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps // compute the total size of all tensors for progress reporting for (const auto & it : weights_map) { - size_data += ggml_nbytes(it.second.tensor); + if (!external.has(it.second.tensor)) { + size_data += ggml_nbytes(it.second.tensor); + } } } @@ -1511,6 +1618,8 @@ const void * llama_model_loader::load_data_range(const llama_tensor_weight & w, bool llama_model_loader::load_all_data( struct ggml_context * ctx, llama_buf_map & bufs, + bool load_from_mmap, + bool discard_file_cache, llama_mlocks * lmlocks, llama_progress_callback progress_callback, void * progress_callback_user_data) { @@ -1537,12 +1646,44 @@ bool llama_model_loader::load_all_data( // 64MB works well for NVMe drives const size_t buffer_size = alignment != 1 ? 64 * 1024 * 1024 + 2 * alignment : 1 * 1024 * 1024; - std::vector host_buffers; - std::vector events; - std::vector host_ptrs; + struct async_upload_resources { + std::vector host_buffers; + std::vector events; + std::vector host_ptrs; + ggml_backend_t backend = nullptr; + + void reset() { + for (auto * event : events) { + if (backend != nullptr) { + ggml_backend_event_synchronize(event); + } + ggml_backend_event_free(event); + } + events.clear(); + for (auto * buffer : host_buffers) { + ggml_backend_buffer_free(buffer); + } + host_buffers.clear(); + host_ptrs.clear(); + ggml_backend_free(backend); + backend = nullptr; + } + + ~async_upload_resources() { + reset(); + } + } async_upload; + async_upload.host_buffers.reserve(n_buffers); + async_upload.events.reserve(n_buffers); + async_upload.host_ptrs.reserve(n_buffers); + + auto & host_buffers = async_upload.host_buffers; + auto & events = async_upload.events; + auto & host_ptrs = async_upload.host_ptrs; + auto & upload_backend = async_upload.backend; size_t buffer_idx = 0; // buffer to use for async loads - ggml_backend_t upload_backend = [&](const char * func) -> ggml_backend_t { - if (use_mmap || check_tensors) { + upload_backend = [&](const char * func) -> ggml_backend_t { + if (load_from_mmap || check_tensors) { return nullptr; } // When not using mmaped io use async uploads from pinned memory to GPU memory. @@ -1589,6 +1730,7 @@ bool llama_model_loader::load_all_data( if (!buf) { LLAMA_LOG_DEBUG("%s: failed to allocate host buffer for async uploads for device %s\n", func, ggml_backend_dev_name(dev)); + async_upload.reset(); return nullptr; } @@ -1599,6 +1741,7 @@ bool llama_model_loader::load_all_data( if (!event) { LLAMA_LOG_DEBUG("%s: failed to create event for async uploads for device %s\n", func, ggml_backend_dev_name(dev)); + async_upload.reset(); return nullptr; } @@ -1609,6 +1752,7 @@ bool llama_model_loader::load_all_data( if (!backend) { LLAMA_LOG_DEBUG("%s: failed to initialize backend for device %s for async uploads\n", func, ggml_backend_dev_name(dev)); + async_upload.reset(); return nullptr; } @@ -1629,7 +1773,7 @@ bool llama_model_loader::load_all_data( // without mmap, tensors in non-host buffers are staged through a temporary buffer sized like the tensor // load them biggest-first so the largest staging buffer is allocated while the fewest weights are resident - if (!use_mmap) { + if (!load_from_mmap) { std::stable_sort(tensors.begin(), tensors.end(), [](const ggml_tensor * a, const ggml_tensor * b) { const bool staged_a = a->buffer && !ggml_backend_buffer_is_host(a->buffer); const bool staged_b = b->buffer && !ggml_backend_buffer_is_host(b->buffer); @@ -1655,7 +1799,7 @@ bool llama_model_loader::load_all_data( size_t n_size = ggml_nbytes(cur); - const bool from_mapping = use_mmap || lazy.has(cur); + const bool from_mapping = load_from_mmap || lazy.has(cur); if (from_mapping) { const auto & mapping = mappings.at(weight->idx); @@ -1763,20 +1907,15 @@ bool llama_model_loader::load_all_data( } } } + if (discard_file_cache) { + file->discard_cache(weight->offs, n_size); + } } size_done += n_size; } - // free temporary resources used for async uploads - for (auto * event : events) { - ggml_backend_event_synchronize(event); - ggml_backend_event_free(event); - } - for (auto * buf : host_buffers) { - ggml_backend_buffer_free(buf); - } - ggml_backend_free(upload_backend); + async_upload.reset(); // check validation results bool validation_failed = false; diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 72bbd53e7d03..e3e7899b917f 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -4,6 +4,7 @@ #include "llama-impl.h" #include "llama-arch.h" +#include "llama-expert-store.h" #include "llama-hparams.h" #include "llama-mmap.h" @@ -11,6 +12,7 @@ #include #include +#include #include #include #include @@ -48,6 +50,19 @@ struct llama_model_loader { throw std::runtime_error(format("tensor '%s' data is not within the file bounds, model is corrupted or incomplete", ggml_get_name(tensor))); } } + + llama_tensor_weight(uint16_t idx, const struct gguf_context * gguf_ctx, ggml_tensor * tensor) : idx(idx), tensor(tensor) { + const int tensor_idx = gguf_find_tensor(gguf_ctx, ggml_get_name(tensor)); + if (tensor_idx < 0) { + throw std::runtime_error(format("tensor '%s' not found in the model", ggml_get_name(tensor))); + } + const size_t data_offset = gguf_get_data_offset(gguf_ctx); + const size_t tensor_offset = gguf_get_tensor_offset(gguf_ctx, tensor_idx); + if (data_offset > std::numeric_limits::max() - tensor_offset) { + throw std::runtime_error(format("tensor '%s' offset overflows", ggml_get_name(tensor))); + } + offs = data_offset + tensor_offset; + } }; // custom comparator to sort weights more nicely by layer @@ -117,6 +132,38 @@ struct llama_model_loader { std::set tensors; } lazy; + struct external_read { + void add(const llama_tensor_weight & w); + + bool any() const { + return !ranges.empty(); + } + + bool has(const ggml_tensor * t) const { + return tensors.count(ggml_get_name(t)) > 0; + } + + const llama_mmap::ranges & for_file(uint32_t idx) const { + static const llama_mmap::ranges none; + + const auto it = ranges.find(idx); + return it == ranges.end() ? none : it->second; + } + + bool intersects(uint32_t idx, size_t first, size_t last) const { + for (const auto & range : for_file(idx)) { + if (range.first < last && first < range.second) { + return true; + } + } + return false; + } + + private: + std::map ranges; + std::set tensors; + } external; + llama_files files; std::vector fnames; // one per entry of files, for readers that outlive the loader llama_ftype ftype; @@ -239,6 +286,12 @@ struct llama_model_loader { const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags); + llama_expert_store_tensor register_external_tensor( + const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne); + void done_getting_tensors(bool partial = false) const; void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr); @@ -256,6 +309,8 @@ struct llama_model_loader { bool load_all_data( struct ggml_context * ctx, llama_buf_map & bufs, + bool load_from_mmap, + bool discard_file_cache, llama_mlocks * lmlocks, llama_progress_callback progress_callback, void * progress_callback_user_data); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index b546491caead..b7753cb2ad1b 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -6,6 +6,7 @@ #include "llama-impl.h" #include "llama-mmap.h" #include "llama-cparams.h" +#include "llama-dsv41.h" #include "llama-model-loader.h" #include "llama-kv-cache.h" @@ -17,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" @@ -204,6 +206,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_dots3note(params); case LLM_ARCH_DEEPSEEK4: return new llama_model_deepseek4(params); + case LLM_ARCH_DEEPSEEK41: + return new llama_model_deepseek41(params); case LLM_ARCH_GLM_DSA: return new llama_model_glm_dsa(params); case LLM_ARCH_MISTRAL4: @@ -377,6 +381,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const llama_hparams & hparams = ud->model->hparams; const std::string tensor_name = tensor->name; const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 || + ud->model->arch == LLM_ARCH_DEEPSEEK41 || (ud->model->arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0); static const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight"); @@ -388,7 +393,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight"); static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*"); static const std::regex pattern_idx_cache ("cache_idx_(k|v)_l\\d*"); - static const std::regex pattern_dsv4_state ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*"); + static const std::regex pattern_dsv4_state ("dsv4(1)?_(csa|hca|lid|comp|index)_state_(kv|score)_l\\d*"); static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); static const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias"); @@ -1232,6 +1237,33 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { return; } + if (arch == LLM_ARCH_DEEPSEEK41) { + std::fill(hparams.n_head_arr.begin(), hparams.n_head_arr.end(), 0); + std::fill(hparams.n_head_kv_arr.begin(), hparams.n_head_kv_arr.end(), 0); + std::fill(hparams.n_ff_arr.begin(), hparams.n_ff_arr.end(), 0); + std::fill(hparams.n_ff_exp_arr.begin(), hparams.n_ff_exp_arr.end(), 0); + std::fill(hparams.n_expert_used_arr.begin(), hparams.n_expert_used_arr.end(), 0); + std::fill(hparams.rope_sections.begin(), hparams.rope_sections.end(), 0); + std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), 1); + std::fill(hparams.is_swa_impl.begin(), hparams.is_swa_impl.end(), 0); + std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), 0); + std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 0); + std::fill(hparams.dsv41_kv_source_layer.begin(), hparams.dsv41_kv_source_layer.end(), -1); + std::fill(hparams.dsv41_index_source_layer.begin(), hparams.dsv41_index_source_layer.end(), -1); + std::fill(hparams.dsv4_compress_ratios.begin(), hparams.dsv4_compress_ratios.end(), 0); + std::fill(hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.end(), 0.0f); + std::fill(hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.end(), 0.0f); + hparams.dsv41_engram_layers.reset(); + + load_arch_hparams(ml); + + pimpl->n_bytes = ml.n_bytes; + pimpl->desc_str = arch_name() + " " + type_name() + " " + ml.ftype_name(); + pimpl->ftype = ml.ftype; + hparams.rope_type = llama_model_rope_type(this); + return; + } + ml.get_key(LLM_KV_CONTEXT_LENGTH, hparams.n_ctx_train); ml.get_key(LLM_KV_EMBEDDING_LENGTH, hparams.n_embd); ml.get_key(LLM_KV_EMBEDDING_LENGTH_OUT, hparams.n_embd_out_impl, false); @@ -1714,12 +1746,22 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // With the n-gram table read by explicit preads (--lazy-mode on-direct), a populated mapping would pull // the table's third of the file resident for nothing; readahead alone carries the sequential load. - ml.init_mappings(params.lazy_mode != LLAMA_LAZY_MODE_DIRECT, use_mlock ? &pimpl->mlock_mmaps : nullptr); + // Do not prefetch files with disk-owned tensor holes either: unsafe contexts load their + // resident tensors through bounded staging and discard copied source pages. + llama_mlocks * mmap_locks = use_mlock && !ml.external.any() ? &pimpl->mlock_mmaps : nullptr; + ml.init_mappings(!(params.lazy_mode == LLAMA_LAZY_MODE_DIRECT || ml.external.any()), mmap_locks); pimpl->mappings.reserve(ml.mappings.size()); // create the backend buffers - std::vector> ctx_buf_maps; + struct ctx_buf_map { + ggml_context * ctx; + llama_buf_map bufs; + bool load_from_mmap; + bool discard_file_cache; + }; + std::vector ctx_buf_maps; ctx_buf_maps.reserve(ml.ctx_map.size()); + bool keep_mappings = false; // Ensure we have enough capacity for the maximum backend buffer we will potentially create const size_t n_max_backend_buffer = ml.ctx_map.size() * ml.files.size(); @@ -1755,9 +1797,24 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // a lazy context is mapped whatever the load mode, but the memory-fit pass maps nothing const bool is_lazy_mapped = ctx_key.lazy && !ml.no_alloc; + bool context_mmap_safe = true; + if (ml.use_mmap && ml.external.any() && !is_lazy_mapped) { + for (uint32_t idx = 0; idx < ml.files.size(); ++idx) { + void * addr = nullptr; + size_t first; + size_t last; + ml.get_mapping_range(&first, &last, &addr, idx, ctx); + if (first < last && ml.external.intersects(idx, first, last)) { + context_mmap_safe = false; + break; + } + } + } - if ((ml.use_mmap || is_lazy_mapped) && use_mmap_buffer && buffer_from_host_ptr_supported && is_default_buft) { + if ((ml.use_mmap || is_lazy_mapped) && (use_mmap_buffer || is_lazy_mapped) && + context_mmap_safe && buffer_from_host_ptr_supported && is_default_buft) { GGML_ASSERT(!ml.no_alloc); + keep_mappings = true; for (uint32_t idx = 0; idx < ml.files.size(); idx++) { // only the mmap region containing the tensors in the model is mapped to the backend buffer // this is important for metal with apple silicon: if the entire model could be mapped to a metal buffer, @@ -1774,6 +1831,15 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { if (buf == nullptr) { throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); } + if (use_mlock && ml.external.any() && !is_lazy_mapped) { + size_t lock_first = first; + size_t lock_last = last; + llama_mlock::align_range(&lock_first, &lock_last); + pimpl->mlock_mmaps.emplace_back(new llama_mlock); + auto & mlock_mmap = pimpl->mlock_mmaps.back(); + mlock_mmap->init((char *) addr + lock_first); + mlock_mmap->grow_to(lock_last - lock_first); + } bufs.emplace_back(buf); buf_map.emplace(idx, buf); } @@ -1810,7 +1876,8 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { pimpl->ctxs_bufs.emplace_back(std::move(ctx_ptr), std::move(bufs)); - ctx_buf_maps.emplace_back(ctx, buf_map); + const bool load_from_mmap = context_mmap_safe && (ml.use_mmap || is_lazy_mapped); + ctx_buf_maps.push_back({ ctx, std::move(buf_map), load_from_mmap, ml.use_mmap && !load_from_mmap }); } if (llama_supports_gpu_offload()) { @@ -1842,21 +1909,24 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } // without mmap, load non-host buffers first: their tensors go through a staging buffer, which is cheapest while the fewest weights are resident - if (!ml.use_mmap) { + if (!ml.use_mmap || ml.external.any()) { std::stable_partition(ctx_buf_maps.begin(), ctx_buf_maps.end(), [](const auto & ctx_buf_map) { - const auto & buf_map = ctx_buf_map.second; - return !buf_map.empty() && !ggml_backend_buffer_is_host(buf_map.begin()->second); + const auto & buf_map = ctx_buf_map.bufs; + return !ctx_buf_map.load_from_mmap && !buf_map.empty() && + !ggml_backend_buffer_is_host(buf_map.begin()->second); }); } // load tensor data - for (auto & [ctx, buf_map] : ctx_buf_maps) { - if (!ml.load_all_data(ctx, buf_map, use_mlock ? &pimpl->mlock_mmaps : NULL, params.progress_callback, params.progress_callback_user_data)) { + for (auto & ctx_buf_map : ctx_buf_maps) { + if (!ml.load_all_data(ctx_buf_map.ctx, ctx_buf_map.bufs, ctx_buf_map.load_from_mmap, + ctx_buf_map.discard_file_cache, mmap_locks, + params.progress_callback, params.progress_callback_user_data)) { return false; } } - if (use_mmap_buffer) { + if (keep_mappings) { for (auto & mapping : ml.mappings) { pimpl->mappings.emplace_back(std::move(mapping)); } @@ -2215,6 +2285,20 @@ ggml_backend_buffer_type_t llama_model::select_buft(int il) const { }); } +ggml_backend_buffer_type_t llama_model::select_moe_buft( + int il, enum ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2) const { + return ::select_buft( + *pimpl->dev_layer.at(il).buft_list, + [&](ggml_context * ctx) { + const int64_t n_expert_used = hparams.n_expert_used_max(); + GGML_ASSERT(n_expert_used > 0); + ggml_tensor * weight = ggml_new_tensor_3d(ctx, type, ne0, ne1, ne2); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, ne0, n_expert_used, 512); + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512); + return ggml_mul_mat_id(ctx, weight, input, ids); + }); +} + bool llama_model::has_tensor_overrides() const { return pimpl->has_tensor_overrides; } @@ -2613,6 +2697,18 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr); } } break; + case LLM_ARCH_DEEPSEEK41: + { + 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) @@ -2934,6 +3030,8 @@ llama_model_params llama_model_default_params() { /*.load_mode =*/ LLAMA_LOAD_MODE_AUTO, /*.lazy_mode =*/ LLAMA_LAZY_MODE_AUTO, /*.main_gpu =*/ 0, + /*.expert_cache_bytes =*/ 0, + /*.expert_cache_slots =*/ 0, /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, @@ -3000,7 +3098,7 @@ int32_t llama_model_n_head_kv(const llama_model * model) { int32_t llama_model_n_swa(const llama_model * model) { // dsv4 kv-cache has SWA but it cannot be used as a rollback because of // other compression ratios, so we return 0 here - if (model->arch == LLM_ARCH_DEEPSEEK4) { + if (model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) { return 0; } return model->hparams.n_swa; @@ -3086,6 +3184,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_DEEPSEEK2OCR: case LLM_ARCH_DEEPSEEK32: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_MUSE_GLIMMER: case LLM_ARCH_PLM: case LLM_ARCH_CHATGLM: diff --git a/src/llama-model.h b/src/llama-model.h index 4278aa47e584..728d2170beb6 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -564,6 +564,10 @@ struct llama_layer { struct ggml_tensor * indexer_attn_k = nullptr; struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias + struct ggml_tensor * engram_kv = nullptr; + struct ggml_tensor * engram_q_norm = nullptr; + struct ggml_tensor * engram_k_norm = nullptr; + // MSA struct ggml_tensor * index_q_proj = nullptr; struct ggml_tensor * index_k_proj = nullptr; @@ -760,6 +764,8 @@ struct llama_model { ggml_backend_dev_t dev_output() const; ggml_backend_buffer_type_t select_buft(int il) const; + ggml_backend_buffer_type_t select_moe_buft( + int il, enum ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2) const; bool has_tensor_overrides() const; @@ -783,6 +789,13 @@ struct llama_model { virtual void load_vocab (llama_model_loader & ml) = 0; virtual bool load_tensors(llama_model_loader & ml) = 0; // returns false if cancelled by progress_callback + virtual bool requires_synchronous_graph() const { return false; } + virtual std::string consume_runtime_error() const { return {}; } + virtual void release_runtime_work() const {} + virtual void release_runtime_work_after_sync(ggml_backend_sched_t) const { release_runtime_work(); } + virtual void acquire_runtime_context() const {} + virtual void release_runtime_context() const {} + // model must define these virtual void load_arch_hparams(llama_model_loader & ml) = 0; virtual void load_arch_tensors(llama_model_loader & ml) = 0; diff --git a/src/llama-ple-disk.cpp b/src/llama-ple-disk.cpp index da374f9d9df5..f391331a787c 100644 --- a/src/llama-ple-disk.cpp +++ b/src/llama-ple-disk.cpp @@ -1,5 +1,6 @@ #include "llama-ple-disk.h" +#include "llama-bounded-file.h" #include "llama-impl.h" #include @@ -9,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -24,15 +26,15 @@ struct llama_ple_disk::impl { std::string fname; - int fd = -1; bool direct = false; size_t offs = 0; ggml_type type = GGML_TYPE_COUNT; int64_t ne0 = 0; int64_t nrows = 0; size_t rs = 0; // bytes per row - size_t block = 4096; // O_DIRECT alignment for offset, length and buffer ggml_to_float_t to_float = nullptr; + std::unique_ptr file; + int fd = -1; // buffered-mode fd, used only for prefetch() fadvise hints // direct-mapped cache of raw rows: slot = row & mask, tag = row size_t n_slots = 0; @@ -47,7 +49,7 @@ struct llama_ple_disk::impl { std::vector> misses; // (row, index into raw) float * dst = nullptr; // destination of the gather in progress const uint8_t * src_full = nullptr; // preloaded table while a gather reads from it - uint8_t * bounce0 = nullptr; // calling-thread bounce buffer + llama_bounded_file::buffer bounce0; // calling-thread bounce buffer std::mutex mtx; // one gather at a time; a model is shared by every context built on it @@ -64,6 +66,7 @@ struct llama_ple_disk::impl { uint64_t gen = 0; size_t pending = 0; std::atomic next{0}; + std::exception_ptr worker_error; // first error seen by a worker, rethrown on the calling thread bool stop = false; // LLAMA_PLE_PRELOAD=1: the whole table is read once into anonymous (huge-page) RAM by background threads; once @@ -82,6 +85,9 @@ struct llama_ple_disk::impl { #if defined(_WIN32) throw std::runtime_error("llama_ple_disk: not supported on Windows"); #else + if (ne0 <= 0 || nrows <= 0) { + throw std::runtime_error("llama_ple_disk: invalid table shape"); + } if (ggml_is_quantized(type) && ne0 % ggml_blck_size(type) != 0) { throw std::runtime_error(format("llama_ple_disk: row of %lld %s elements is not a whole number of blocks", (long long) ne0, ggml_type_name(type))); @@ -94,35 +100,20 @@ struct llama_ple_disk::impl { } } - direct = p.direct_io; - if (direct) { - // only Linux has O_DIRECT; Darwin turns the page cache off per descriptor - // with F_NOCACHE after the open -#if defined(O_DIRECT) - fd = open(fname.c_str(), O_RDONLY | O_DIRECT | O_CLOEXEC); -#else - fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); -#if defined(F_NOCACHE) - if (fd >= 0 && fcntl(fd, F_NOCACHE, 1) < 0) { - LLAMA_LOG_WARN("%s: F_NOCACHE on %s failed (%s); reads go through the page cache\n", - __func__, fname.c_str(), strerror(errno)); - direct = false; - } -#else - direct = false; -#endif -#endif - if (fd < 0) { - LLAMA_LOG_WARN("%s: direct open of %s failed (%s); falling back to buffered reads\n", - __func__, fname.c_str(), strerror(errno)); - direct = false; - } + llama_bounded_file::params fp; + fp.direct_io = p.direct_io; + file = std::make_unique(fname, fp); + direct = file->direct_io(); + if (rs == 0 || (uint64_t) nrows > UINT64_MAX / rs) { + throw std::runtime_error("llama_ple_disk: table extent overflow"); + } + const uint64_t bytes = (uint64_t) nrows * rs; + if (offs > file->size() || bytes > file->size() - offs) { + throw std::runtime_error("llama_ple_disk: table extent is outside the file"); } - if (fd < 0) { + if (!direct) { + // best-effort fd used only to hint the page cache in prefetch(); failure is not fatal fd = open(fname.c_str(), O_RDONLY | O_CLOEXEC); - if (fd < 0) { - throw std::runtime_error(format("llama_ple_disk: failed to open %s: %s", fname.c_str(), strerror(errno))); - } } n_threads = std::max(1, p.n_threads); @@ -162,7 +153,6 @@ struct llama_ple_disk::impl { for (auto & w : workers) { w.join(); } - free(bounce0); if (fd >= 0) { close(fd); } @@ -253,52 +243,9 @@ struct llama_ple_disk::impl { }); } - size_t bounce_size() const { - return ((rs + block - 1) / block) * block + 2 * block; - } - - uint8_t * alloc_bounce() const { - void * ptr = nullptr; - if (posix_memalign(&ptr, block, bounce_size()) != 0) { - throw std::runtime_error("llama_ple_disk: posix_memalign failed"); - } - return (uint8_t *) ptr; - } - - // read `len` bytes at `off` into `dst`; a short read is only tolerated past `need` - void pread_full(uint8_t * dst, size_t len, off_t off, size_t need) const { - size_t got = 0; - while (got < len) { - const ssize_t r = pread(fd, dst + got, len - got, off + (off_t) got); - if (r < 0) { - if (errno == EINTR) { - continue; - } - GGML_ABORT("llama_ple_disk: pread(%s, %zu @ %lld) failed: %s", - fname.c_str(), len, (long long) off, strerror(errno)); - } - if (r == 0) { - break; // EOF - } - got += (size_t) r; - } - if (got < need) { - GGML_ABORT("llama_ple_disk: short read in %s: %zu of %zu bytes at %lld", - fname.c_str(), got, need, (long long) off); - } - } - - void read_row(int64_t row, uint8_t * dst, uint8_t * bounce) const { - const off_t off = (off_t) offs + (off_t) row * (off_t) rs; - if (!direct) { - pread_full(dst, rs, off, rs); - return; - } - const off_t a0 = off & ~(off_t) (block - 1); - const size_t need = (size_t) (off - a0) + rs; - const size_t len = ((need + block - 1) / block) * block; - pread_full(bounce, len, a0, need); - memcpy(dst, bounce + (off - a0), rs); + void read_row(int64_t row, uint8_t * dst, llama_bounded_file::buffer & bounce) const { + const uint64_t off = (uint64_t) offs + (uint64_t) row * rs; + file->read(off, dst, rs, bounce); } // dequantize distinct row u into the first position that asked for it, then copy it to the others @@ -315,7 +262,7 @@ struct llama_ple_disk::impl { } } - void run_items(uint8_t * bounce) { + void run_items(llama_bounded_file::buffer & bounce) { for (;;) { const size_t i0 = next.fetch_add(grain); if (i0 >= n_items) { @@ -333,8 +280,8 @@ struct llama_ple_disk::impl { } void worker() { - uint8_t * bounce = direct ? alloc_bounce() : nullptr; - uint64_t seen = 0; + llama_bounded_file::buffer bounce; + uint64_t seen = 0; for (;;) { { std::unique_lock lk(pm); @@ -344,15 +291,25 @@ struct llama_ple_disk::impl { } seen = gen; } - run_items(bounce); + std::exception_ptr error; + try { + if (direct && bounce.empty()) { + bounce = file->make_buffer(rs); + } + run_items(bounce); + } catch (...) { + error = std::current_exception(); + } { std::lock_guard lk(pm); + if (error && !worker_error) { + worker_error = error; + } if (--pending == 0) { cv_done.notify_one(); } } } - free(bounce); } // run `kind` over n items; up to serial_max of them stay on the calling thread, @@ -361,8 +318,8 @@ struct llama_ple_disk::impl { job = kind; n_items = n; grain = g; - if (direct && bounce0 == nullptr) { - bounce0 = alloc_bounce(); + if (direct && bounce0.empty()) { + bounce0 = file->make_buffer(rs); } if (n_threads <= 1 || n <= serial_max) { next = 0; @@ -377,14 +334,29 @@ struct llama_ple_disk::impl { } { std::lock_guard lk(pm); - next = 0; - pending = workers.size(); + next = 0; + pending = workers.size(); + worker_error = nullptr; ++gen; } cv_work.notify_all(); - run_items(bounce0); + // the calling thread also helps drain items, so its own read errors join the same + // worker_error slot instead of throwing early and abandoning still-running workers + std::exception_ptr error; + try { + run_items(bounce0); + } catch (...) { + error = std::current_exception(); + } std::unique_lock lk(pm); cv_done.wait(lk, [&] { return pending == 0; }); + if (!error) { + error = worker_error; + } + lk.unlock(); + if (error) { + std::rethrow_exception(error); + } } void gather(const int32_t * idx, size_t n, float * out) { @@ -468,7 +440,9 @@ struct llama_ple_disk::impl { st_uniq += n_uniq; st_hits += n_uniq - misses.size(); st_reads += misses.size(); - st_bytes += misses.size() * (direct ? block : rs); + for (const auto & m : misses) { + st_bytes += file->read_size((uint64_t) offs + (uint64_t) m.first * rs, rs); + } st_ms += std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); } #endif @@ -511,11 +485,13 @@ void llama_ple_disk::prefetch(const int32_t * idx, size_t n) const { } std::sort(uniq.begin(), uniq.end()); uniq.erase(std::unique(uniq.begin(), uniq.end()), uniq.end()); +#if defined(POSIX_FADV_WILLNEED) for (const int32_t row : uniq) { const off_t off = (off_t) pimpl->offs + (off_t) row * (off_t) pimpl->rs; posix_fadvise(pimpl->fd, off, (off_t) pimpl->rs, POSIX_FADV_WILLNEED); } #endif +#endif } bool llama_ple_disk::page_cached() const { return !pimpl->direct; } diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp new file mode 100644 index 000000000000..8f5946ec7900 --- /dev/null +++ b/src/models/deepseek41.cpp @@ -0,0 +1,1582 @@ +#include "llama-dsv41.h" +#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 + +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); +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, + 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 { + if (hparams.no_alloc) { + return nullptr; + } + 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 * candidate_final_block = 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.candidate_final_block = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1, 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.candidate_final_block); + 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); + std::vector final_blocks(n_tokens); + 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; + if (visible_blocks == 0) { + throw std::runtime_error( + "DeepSeek V4.1 candidate source has no visible block"); + } + for (uint32_t block = visible_blocks; block < n_blocks; ++block) { + bias[(size_t) token*n_blocks + block] = + -std::numeric_limits::infinity(); + } + final_blocks[token] = (int32_t) visible_blocks - 1; + } + set_tensor(source.candidate_block_bias, bias); + set_tensor(source.candidate_final_block, final_blocks); + + 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]); + set_tensor( + engram_row_ids[index], + mctx->engram_row_ids(index)); + } + } + + 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) { + 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); +} + +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_cont( + ctx, + 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; + ml.get_key(LLM_KV_DSV41_CONFIG, raw_config); + ml.get_key(LLM_KV_DSV41_MAX_POSITION_EMBEDDINGS, config.n_ctx_train); + ml.get_key(LLM_KV_DSV41_HIDDEN_SIZE, config.n_embd); + ml.get_key(LLM_KV_DSV41_NUM_HIDDEN_LAYERS, config.n_layer); + ml.get_key(LLM_KV_DSV41_VOCAB_SIZE, config.n_vocab); + ml.get_key(LLM_KV_DSV41_NUM_ATTENTION_HEADS, config.n_head); + ml.get_key(LLM_KV_DSV41_NUM_KEY_VALUE_HEADS, config.n_head_kv); + ml.get_key(LLM_KV_DSV41_HEAD_DIM, config.n_head_dim); + ml.get_key(LLM_KV_DSV41_QK_ROPE_HEAD_DIM, config.n_rot); + ml.get_key(LLM_KV_DSV41_Q_LORA_RANK, config.n_lora_q); + ml.get_key(LLM_KV_DSV41_O_LORA_RANK, config.n_lora_o); + ml.get_key(LLM_KV_DSV41_O_GROUPS, config.n_o_group); + ml.get_key(LLM_KV_DSV41_MOE_INTERMEDIATE_SIZE, config.n_ff_expert); + ml.get_key(LLM_KV_DSV41_N_ROUTED_EXPERTS, config.n_expert); + ml.get_key(LLM_KV_DSV41_NUM_EXPERTS_PER_TOK, config.n_expert_used); + ml.get_key(LLM_KV_DSV41_N_SHARED_EXPERTS, config.n_expert_shared); + ml.get_key(LLM_KV_DSV41_INDEX_N_HEADS, config.indexer_n_head); + ml.get_key(LLM_KV_DSV41_INDEX_HEAD_DIM, config.indexer_head_size); + ml.get_key(LLM_KV_DSV41_INDEX_TOPK, config.indexer_top_k); + ml.get_key(LLM_KV_DSV41_HC_MULT, config.hc_count); + ml.get_key(LLM_KV_DSV41_HC_SINKHORN_ITERS, config.hc_sinkhorn_iters); + ml.get_key(LLM_KV_DSV41_SLIDING_WINDOW, config.raw_window); + ml.get_key(LLM_KV_DSV41_CANDIDATE_SOURCE_LAYER_ID, config.candidate_source_layer); + ml.get_key(LLM_KV_DSV41_CANDIDATE_TOPK_BLOCKS, config.candidate_topk_blocks); + ml.get_key(LLM_KV_DSV41_CANDIDATE_BLOCK_SIZE, config.candidate_block_size); + ml.get_key(LLM_KV_DSV41_RMS_NORM_EPS, config.f_norm_rms_eps); + ml.get_key(LLM_KV_DSV41_HC_EPS, config.hc_eps); + ml.get_key(LLM_KV_DSV41_SWIGLU_LIMIT, config.swiglu_clamp); + ml.get_key(LLM_KV_DSV41_ROUTED_SCALING_FACTOR, config.routed_scale); + ml.get_key(LLM_KV_DSV41_ROPE_THETA, config.rope_theta); + ml.get_key(LLM_KV_DSV41_COMPRESS_ROPE_THETA, config.compress_rope_theta); + ml.get_key(LLM_KV_DSV41_ROPE_SCALING_FACTOR, config.yarn_factor); + ml.get_key(LLM_KV_DSV41_ROPE_SCALING_BETA_FAST, config.yarn_beta_fast); + ml.get_key(LLM_KV_DSV41_ROPE_SCALING_BETA_SLOW, config.yarn_beta_slow); + ml.get_key(LLM_KV_DSV41_ROPE_SCALING_ORIG_CTX_LEN, config.yarn_original_context); + ml.get_key(LLM_KV_DSV41_NORM_TOPK_PROB, config.expert_weights_norm); + ml.get_key(LLM_KV_DSV41_HIDDEN_ACT, config.hidden_act); + ml.get_key(LLM_KV_DSV41_SCORING_FUNC, config.scoring_func); + ml.get_key(LLM_KV_DSV41_TOPK_METHOD, config.topk_method); + ml.get_arr(LLM_KV_DSV41_COMPRESS_RATIOS, config.compress_ratios); + ml.get_arr(LLM_KV_DSV41_KV_SOURCE_LAYER_IDS, config.kv_sources); + ml.get_arr(LLM_KV_DSV41_INDEX_SOURCE_LAYER_IDS, config.index_sources); + ml.get_key(LLM_KV_DSV41_ENGRAM_ENCODING, config.engram_encoding); + ml.get_arr(LLM_KV_DSV41_ENGRAM_LAYER_IDS, config.engram_layers); + ml.get_arr(LLM_KV_DSV41_ENGRAM_ROWS, config.engram_rows); + ml.get_key(LLM_KV_DSV41_ENGRAM_COMPRESSED_VOCAB_SIZE, config.engram_compressed_vocab_size); + ml.get_key(LLM_KV_DSV41_ENGRAM_PAD_ID, config.engram_pad_id); + ml.get_arr_n(LLM_KV_DSV41_ENGRAM_TOKEN_MAP, config.engram_token_map_size); + ml.get_arr_n(LLM_KV_DSV41_ENGRAM_PRIMES, config.engram_primes_size); + ml.get_arr_n(LLM_KV_DSV41_ENGRAM_MULTIPLIERS, config.engram_multipliers_size); + ml.get_arr(LLM_KV_DSV41_ENGRAM_TOKEN_MAP, config.engram_token_map); + ml.get_arr(LLM_KV_DSV41_ENGRAM_PRIMES, config.engram_primes); + ml.get_arr(LLM_KV_DSV41_ENGRAM_MULTIPLIERS, config.engram_multipliers); + + config.n_ff_dense = LLAMA_DSV41_N_FF_DENSE; + llama_dsv41_validate_config(config); + engram = std::make_shared(); + engram->layout = llama_dsv41_make_engram_layout(config); + + if (raw_config.empty()) { + throw std::runtime_error("DeepSeek V4.1 metadata: config must not be empty"); + } + hparams.n_ctx_train = config.n_ctx_train; + hparams.n_embd = config.n_embd; + hparams.n_embd_out_impl = config.n_embd; + hparams.n_layer_all = config.n_layer; + hparams.n_layer_nextn = 0; + hparams.n_expert = config.n_expert; + hparams.n_expert_shared = config.n_expert_shared; + hparams.n_lora_q = config.n_lora_q; + hparams.n_ff_shexp = config.n_ff_expert; + hparams.n_embd_head_k_full = config.n_head_dim; + hparams.n_embd_head_v_full = config.n_head_dim; + hparams.n_embd_head_k_swa = config.n_head_dim; + hparams.n_embd_head_v_swa = config.n_head_dim; + hparams.n_rot_full = config.n_rot; + hparams.n_rot_swa = config.n_rot; + hparams.n_swa = config.raw_window; + hparams.indexer_n_head = config.indexer_n_head; + hparams.indexer_head_size = config.indexer_head_size; + hparams.indexer_top_k = config.indexer_top_k; + hparams.dsv4_o_group_count = config.n_o_group; + hparams.dsv4_o_lora_rank = config.n_lora_o; + hparams.dsv4_hc_mult = config.hc_count; + hparams.dsv4_hc_sinkhorn_iters = config.hc_sinkhorn_iters; + hparams.dsv4_compress_rope_base = (float) config.compress_rope_theta; + hparams.dsv4_hc_eps = config.hc_eps; + hparams.dsv41_candidate_source_layer = config.candidate_source_layer; + hparams.dsv41_candidate_topk_blocks = config.candidate_topk_blocks; + hparams.dsv41_candidate_block_size = config.candidate_block_size; + hparams.f_norm_rms_eps = config.f_norm_rms_eps; + hparams.expert_weights_scale = config.routed_scale; + hparams.expert_weights_norm = config.expert_weights_norm; + hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS; + hparams.rope_freq_base_train = (float) config.rope_theta; + hparams.rope_freq_base_train_swa = (float) config.rope_theta; + hparams.rope_freq_scale_train = 1.0f/config.yarn_factor; + hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + hparams.n_ctx_orig_yarn = (uint32_t) config.yarn_original_context; + hparams.yarn_beta_fast = config.yarn_beta_fast; + hparams.yarn_beta_slow = config.yarn_beta_slow; + hparams.yarn_ext_factor = 1.0f; + hparams.rope_attn_factor = dsv41_rope_attn_factor(hparams.rope_freq_scale_train); + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + hparams.causal_attn = true; + + for (uint32_t il = 0; il < config.n_layer; ++il) { + hparams.n_head_arr[il] = config.n_head; + hparams.n_head_kv_arr[il] = config.n_head_kv; + hparams.n_ff_arr[il] = config.n_ff_dense; + hparams.n_ff_exp_arr[il] = config.n_ff_expert; + hparams.n_expert_used_arr[il] = config.n_expert_used; + hparams.swiglu_clamp_exp[il] = config.swiglu_clamp; + hparams.swiglu_clamp_shexp[il] = config.swiglu_clamp; + hparams.dsv4_compress_ratios[il] = config.compress_ratios[il]; + hparams.dsv41_kv_source_layer[il] = llama_dsv41_kv_source_layer(il); + hparams.dsv41_index_source_layer[il] = llama_dsv41_index_source_layer(il); + hparams.is_swa_impl[il] = 1; + } + for (uint32_t il : config.engram_layers) { + hparams.dsv41_engram_layers.set(il); + } + + type = LLM_TYPE_UNKNOWN; +} + +void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t n_ff_exp = hparams.n_ff_exp(); + const int64_t n_expert_shared = hparams.n_expert_shared; + const int64_t n_embd_head = hparams.n_embd_head_k(); + const int64_t o_groups = hparams.dsv4_o_group_count; + const int64_t o_lora_rank = hparams.dsv4_o_lora_rank; + const int64_t hc_mult = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc_mult*n_embd; + const int64_t hc_mix_dim = (2 + hc_mult)*hc_mult; + + const std::vector expert_tensors = + llama_dsv41_register_expert_tensors([&](const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne) { + return ml.register_external_tensor(name, layer, projection, ne); + }); + if (params.expert_cache_bytes == 0 || params.expert_cache_slots <= 0) { + throw std::runtime_error( + "DeepSeek V4.1 requires non-zero expert_cache_bytes and expert_cache_slots before tensor allocation"); + } + llama_dsv41_expert_runtime_params expert_params; + expert_params.cache_bytes = params.expert_cache_bytes; + expert_params.cache_slots = params.expert_cache_slots; + expert_params.direct_io = true; + expert_params.allow_buffered_io = false; + expert_params.no_alloc = ml.no_alloc; + experts = std::make_shared( + expert_tensors, + expert_params, + [this](const llama_expert_store_tensor & tensor) { + return select_moe_buft( + tensor.layer, tensor.type, tensor.ne[0], tensor.ne[1], params.expert_cache_slots); + }); + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, 0); + + for (int32_t il = 0; il < n_layer; ++il) { + auto & layer = layers[il]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, 0); + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", il), { n_head }, 0); + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, 0); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, 0); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head*n_embd_head }, 0); + layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", il), { n_embd, n_embd_head }, 0); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { n_embd_head }, 0); + layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", il), { n_head*n_embd_head/o_groups, o_lora_rank, o_groups }, TENSOR_ALLOW_RESHAPE); + layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", il), { o_groups*o_lora_rank, n_embd }, 0); + + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", il), { hc_dim, hc_mix_dim }, 0); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", il), { hc_mix_dim }, 0); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", il), { 3 }, 0); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", il), { hc_dim, hc_mix_dim }, 0); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", il), { hc_mix_dim }, 0); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", il), { 3 }, 0); + + if (hparams.dsv41_is_kv_source(il)) { + layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", il), { n_embd, n_embd_head }, 0); + layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", il), { n_embd_head }, 0); + if (hparams.dsv4_compress_ratios[il] == 2) { + layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", il), { n_embd, n_embd_head }, 0); + } + layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", il), { n_embd_head, hparams.indexer_head_size }, 0); + layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { hparams.indexer_head_size }, 0); + } + if (hparams.dsv41_is_index_source(il)) { + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", il), { n_embd, hparams.indexer_n_head }, 0); + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", il), { q_lora_rank, hparams.indexer_n_head*hparams.indexer_head_size }, 0); + } + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, 0); + layer.ffn_exp_probs_b_vl = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B_VL, "bias", il), { n_expert }, TENSOR_NOT_REQUIRED); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, 0); + layer.ffn_gate_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_GATE); + layer.ffn_down_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_DOWN); + layer.ffn_up_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_UP); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_exp*n_expert_shared }, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_exp*n_expert_shared, n_embd }, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_exp*n_expert_shared }, 0); + + if (hparams.dsv41_engram_layers.test(il)) { + const size_t index = il == (int32_t) engram->layout.layer_ids[0] ? 0 : 1; + const std::string table_name = tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il).str(); + const auto * table = ml.get_weight(table_name.c_str()); + if (table == nullptr) { + throw std::runtime_error("DeepSeek V4.1 is missing required Engram tensor " + table_name); + } + llama_dsv41_engram_extent & extent = engram->extents[index]; + extent.fname = ml.no_alloc ? "(no_alloc)" : ml.fnames.at(table->idx); + extent.offset = table->offs; + extent.rows = engram->layout.rows[index]; + extent.columns = table->tensor->ne[0]; + extent.row_count = table->tensor->ne[1]; + extent.type = table->tensor->type; + llama_dsv41_validate_engram_extent(extent); + + create_tensor( + tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il), + { LLAMA_ENGRAM_ROW_BYTES, (int64_t) extent.rows }, + TENSOR_SKIP); + layer.engram_q_norm = create_tensor( + tn(LLM_TENSOR_ENGRAM_Q_NORM, "weight", il), + { n_embd, hc_mult }, + 0); + layer.engram_k_norm = create_tensor( + tn(LLM_TENSOR_ENGRAM_K_NORM, "weight", il), + { n_embd, hc_mult }, + 0); + layer.engram_kv = create_tensor( + tn(LLM_TENSOR_ENGRAM_KV, "weight", il), + { LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, (hc_mult + 1)*n_embd }, + 0); + } + } + +} + +bool llama_model_deepseek41::requires_synchronous_graph() const { + return experts != nullptr; +} + +std::string llama_model_deepseek41::consume_runtime_error() const { + return experts ? experts->consume_error() : std::string(); +} + +void llama_model_deepseek41::release_runtime_work() const { + if (experts) { + experts->release_all(); + } +} + +void llama_model_deepseek41::release_runtime_work_after_sync(ggml_backend_sched_t sched) const { + if (experts) { + experts->release_all_after_sync(sched); + } +} + +void llama_model_deepseek41::acquire_runtime_context() const { + if (experts) { + experts->acquire_context(); + } +} + +void llama_model_deepseek41::release_runtime_context() const { + if (experts) { + experts->release_context(); + } +} + +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_concat( + ggml_context * ctx, + ggml_tensor * a, + ggml_tensor * b, + int dim, + const char * label) { + if (a == nullptr || b == nullptr || a->type != b->type) { + throw std::runtime_error(format( + "DeepSeek V4.1 %s concat type mismatch: %s and %s", + label, + a == nullptr ? "null" : ggml_type_name(a->type), + b == nullptr ? "null" : ggml_type_name(b->type))); + } + return ggml_concat(ctx, a, b, dim); +} + +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, 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( + 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, + ggml_reshape_1d( + graph.ctx0, source->read_idxs, + source->read_width*graph.n_tokens)); + index_k = ggml_reshape_3d( + graph.ctx0, index_k, + graph.hparams.indexer_head_size, + source->read_width, graph.n_tokens); + 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 = llama_dsv41_build_candidate_blocks( + graph.ctx0, block_scores, + source->candidate_final_block, 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 = dsv41_concat( + graph.ctx0, raw_store, raw_write, 1, "raw state"); + raw_read = dsv41_append_zero_row(graph.ctx0, raw_read); + raw_read = ggml_get_rows( + graph.ctx0, raw_read, + ggml_reshape_1d( + graph.ctx0, input.raw_read_idxs, + graph.hparams.n_swa*graph.n_tokens)); + 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); + carry_kv_first = dsv41_cast_for_store( + graph.ctx0, carry_kv_first, compressed); + carry_gate_first = dsv41_cast_for_store( + graph.ctx0, carry_gate_first, gate); + ggml_tensor * source_kv = dsv41_concat( + graph.ctx0, carry_kv_first, compressed, 1, "compressor carry KV"); + ggml_tensor * source_gate = dsv41_concat( + graph.ctx0, carry_gate_first, gate, 1, "compressor carry score"); + 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)); + comp_store = dsv41_append_zero_row( + graph.ctx0, comp_store); + ggml_tensor * compressed = ggml_get_rows( + graph.ctx0, comp_store, + ggml_reshape_1d( + graph.ctx0, selected, + selected->ne[0]*graph.n_tokens)); + compressed = ggml_reshape_4d( + graph.ctx0, compressed, n_embd_head, 1, + selected->ne[0], graph.n_tokens); + k_all = dsv41_concat( + graph.ctx0, raw_read, compressed, 2, "raw and compressed attention"); + + 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); + compressed_mask = dsv41_cast_for_store( + graph.ctx0, compressed_mask, input.raw_mask); + mask_all = dsv41_concat( + graph.ctx0, input.raw_mask, compressed_mask, 0, "attention mask"); + 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 }; +} + +} + +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) : + 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 = llama_dsv41_build_output_norm_input(ctx0, cur); + cb(cur, "dsv41_output_collapse_f32", -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 814080059803..6516ac01fac0 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -9,6 +9,8 @@ #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 static inline ggml_tensor * build_gdn_l2_norm(ggml_context * ctx, ggml_tensor * x, float eps) { @@ -1326,6 +1328,31 @@ struct llama_model_deepseek4 : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +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; + std::string consume_runtime_error() const override; + void release_runtime_work() const override; + void release_runtime_work_after_sync(ggml_backend_sched_t sched) const override; + void acquire_runtime_context() const override; + void release_runtime_context() const override; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_deepseek2ocr : public llama_model_base { llama_model_deepseek2ocr(const struct llama_model_params & params) : llama_model_base(params) {} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index ab0d29bddcbf..4cfd9261f4f3 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -196,7 +196,20 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # llama_build_and_test(test-double-float.cpp) # SLOW + llama_build_and_test(test-deepseek41-schema.cpp) + 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) + llama_test( + test-llama-archs + NAME test-deepseek41-no-alloc + LABEL main + ARGS -a deepseek41 -s 1 + ) set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") diff --git a/tests/test-deepseek41-engram.cpp b/tests/test-deepseek41-engram.cpp new file mode 100644 index 000000000000..ee73c48628f2 --- /dev/null +++ b/tests/test-deepseek41-engram.cpp @@ -0,0 +1,529 @@ +#include "../src/llama-dsv41-engram.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#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_engram_layout make_layout() { + llama_engram_layout layout; + layout.encoding = "e4m3_e8m0_32_row264"; + 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] = (uint32_t) 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; +} + +static void fill_row(uint8_t row[LLAMA_ENGRAM_ROW_BYTES], uint32_t id) { + const uint8_t code = (uint8_t) (8 + id%100); + std::fill(row, row + LLAMA_ENGRAM_DIM, code); + std::fill(row + LLAMA_ENGRAM_DIM, row + LLAMA_ENGRAM_ROW_BYTES, 127); +} + +#if !defined(_WIN32) +static void write_full(int fd, const void * data, size_t size, uint64_t offset) { + const ssize_t written = pwrite(fd, data, size, (off_t) offset); + check(written == (ssize_t) size, "failed to write DeepSeek V4.1 Engram test data"); +} + +struct test_file { + std::string path; + int fd = -1; + std::array extents; + + test_file(const llama_engram_layout & layout) { + char name[] = "/tmp/llama-dsv41-engram-XXXXXX"; + fd = mkstemp(name); + check(fd >= 0, "failed to create DeepSeek V4.1 Engram test file"); + path = name; + + 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_id = 0; row_id < layout.rows[layer]; ++row_id) { + uint8_t row[LLAMA_ENGRAM_ROW_BYTES]; + fill_row(row, row_id + 7*layer); + write_full(fd, row, sizeof(row), offset + (uint64_t) row_id*sizeof(row)); + } + offset += (uint64_t) layout.rows[layer]*LLAMA_ENGRAM_ROW_BYTES + 4096; + } + } + + ~test_file() { + if (fd >= 0) { + close(fd); + } + if (!path.empty()) { + unlink(path.c_str()); + } + } +}; + +static void test_extent_validation() { + llama_dsv41_engram_extent extent = { + "/tmp/model.gguf", 4096, 48, LLAMA_ENGRAM_ROW_BYTES, 48, GGML_TYPE_I8, + }; + llama_dsv41_validate_engram_extent(extent); + extent.type = GGML_TYPE_F32; + expect_invalid([&] { llama_dsv41_validate_engram_extent(extent); }, "non-I8 Engram extent was accepted"); + extent.type = GGML_TYPE_I8; + extent.columns = LLAMA_ENGRAM_ROW_BYTES - 1; + expect_invalid([&] { llama_dsv41_validate_engram_extent(extent); }, "short Engram row was accepted"); +} + +static void test_transactions_and_sequences() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime runtime(layout, file.extents, 8); + + std::vector tokens = { + { 3, 0, { 7, 9 }, 1 }, + { 5, 1, { 7, 9 }, 0 }, + { 11, 0, { 12 }, 1 }, + }; + llama_dsv41_engram_transaction transaction = runtime.prepare(tokens); + check(transaction.token_count() == tokens.size(), "Engram transaction token count mismatch"); + check(transaction.row_ids(1) == transaction.row_ids(0) + LLAMA_ENGRAM_COLS, + "Engram layer row selection is not the second 24-ID half"); + llama_engram_hasher hasher(layout); + llama_engram_history expected_history_7; + llama_engram_history expected_history_12; + expected_history_7.reset(); + expected_history_12.reset(); + uint32_t expected[3*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS]; + const int32_t coupled_tokens[] = { 3, 5 }; + const uint8_t coupled_mask[] = { 1, 0 }; + hasher.hash(expected_history_7, coupled_tokens, coupled_mask, 2, expected); + hasher.hash( + expected_history_12, + &tokens[2].token, + &tokens[2].text, + 1, + expected + 2*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS); + for (size_t token = 0; token < tokens.size(); ++token) { + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + check(std::memcmp( + transaction.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + expected + (token*LLAMA_ENGRAM_LAYERS + layer)*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(uint32_t)) == 0, + "DeepSeek V4.1 Engram IDs or 48-ID token stride differ from the core hasher"); + } + } + check(transaction.text_mask()[0] == 1 && transaction.text_mask()[1] == 0, + "Engram text mask changed"); + check(runtime.sequence(7).pos == -1, "Engram prepare committed sequence state early"); + + runtime.commit(transaction); + check(runtime.sequence(7).pos == 1 && runtime.sequence(9).pos == 1, + "Engram commit did not advance coupled sequences"); + check(runtime.sequence(7).history.tail[0] == LLAMA_ENGRAM_DEAD, + "masked Engram token did not break sequence history"); + + runtime.seq_copy(7, 13); + check(runtime.sequence(13).history.tail == runtime.sequence(7).history.tail, + "Engram sequence copy changed history"); + const llama_dsv41_engram_snapshot snapshot = runtime.checkpoint(); + runtime.seq_reset(7); + check(runtime.sequence(7).pos == -1, "Engram sequence reset did not clear position"); + runtime.restore(snapshot); + check(runtime.sequence(7).pos == 1, "Engram checkpoint restore lost position"); + runtime.seq_remove(13); + check(runtime.sequence(13).pos == -1, "Engram sequence remove retained state"); + + llama_dsv41_engram_transaction stale = runtime.prepare({ { 7, 2, { 7 }, 1 } }); + runtime.seq_copy(7, 14); + expect_runtime([&] { runtime.commit(stale); }, "stale Engram transaction was committed"); + runtime.rollback(stale); + + llama_dsv41_engram_snapshot invalid = runtime.checkpoint(); + invalid.sequences[7].history.tail[0] = (int32_t) layout.compressed_vocab_size; + expect_invalid([&] { runtime.restore(invalid); }, "invalid Engram snapshot was restored"); + expect_invalid( + [&] { runtime.prepare({ { 8, 2, { 7, 7 }, 1 } }); }, + "duplicate Engram sequence ID was accepted"); +} + +static void test_scheduler_upload() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime runtime(layout, file.extents, 8); + llama_dsv41_engram_transaction transaction = runtime.prepare({ + { 3, 0, { 0 }, 1 }, + { 5, 1, { 0 }, 0 }, + { 7, 2, { 0 }, 1 }, + }); + + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create DeepSeek V4.1 Engram upload context"); + ggml_tensor * rows = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, transaction.token_count()); + ggml_tensor * select = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, transaction.token_count()); + + ggml_backend_t backend = ggml_backend_cpu_init(); + check(backend != nullptr, "failed to create DeepSeek V4.1 Engram upload backend"); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + check(buffer != nullptr, "failed to allocate DeepSeek V4.1 Engram upload tensors"); + + std::vector actual_rows(ggml_nelements(rows)); + std::vector actual_select(ggml_nelements(select)); + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + transaction.upload_layer(layer, 0, transaction.token_count(), rows, select); + ggml_backend_tensor_get(rows, actual_rows.data(), 0, ggml_nbytes(rows)); + ggml_backend_tensor_get(select, actual_select.data(), 0, ggml_nbytes(select)); + check(std::memcmp( + actual_rows.data(), + transaction.rows(layer), + ggml_nbytes(rows)) == 0, + "scheduler-backed Engram row upload changed the bounded pack"); + check(actual_select == std::vector({ 3, 1, 5 }), + "scheduler-backed Engram text selection upload changed"); + } + + runtime.commit(transaction); + expect_invalid( + [&] { transaction.upload_layer(0, 0, 3, rows, select); }, + "committed Engram transaction uploaded stale input"); + + ggml_backend_buffer_free(buffer); + ggml_backend_free(backend); + ggml_free(ctx); +} + +static void test_chunked_prefill() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime whole_runtime(layout, file.extents, 8); + llama_dsv41_engram_runtime chunked_runtime(layout, file.extents, 8); + const std::vector tokens = { + { 3, 0, { 0 }, 1 }, + { 5, 1, { 0 }, 1 }, + { 7, 2, { 0 }, 0 }, + { 9, 3, { 0 }, 1 }, + }; + + llama_dsv41_engram_transaction whole = whole_runtime.prepare(tokens); + llama_dsv41_engram_transaction first = chunked_runtime.prepare({ tokens[0], tokens[1] }); + for (size_t token = 0; token < 2; ++token) { + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + check(std::memcmp( + whole.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + first.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(uint32_t)) == 0, + "first Engram prefill chunk differs from whole-chunk hashing"); + } + } + chunked_runtime.commit(first); + + llama_dsv41_engram_transaction second = chunked_runtime.prepare({ tokens[2], tokens[3] }); + for (size_t token = 0; token < 2; ++token) { + for (uint32_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + check(std::memcmp( + whole.row_ids(layer) + (token + 2)*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + second.row_ids(layer) + token*LLAMA_ENGRAM_LAYERS*LLAMA_ENGRAM_COLS, + LLAMA_ENGRAM_COLS*sizeof(uint32_t)) == 0, + "later Engram prefill chunk lost committed history"); + } + } +} + +static void test_transactional_read_failure() { + const llama_engram_layout layout = make_layout(); + test_file file(layout); + llama_dsv41_engram_runtime runtime(layout, file.extents, 8); + + llama_dsv41_engram_transaction first = runtime.prepare({ { 1, 0, { 0 }, 1 } }); + runtime.commit(first); + const llama_dsv41_engram_sequence_state before = runtime.sequence(0); + + const uint64_t truncated = file.extents[1].offset + LLAMA_ENGRAM_ROW_BYTES; + check(ftruncate(file.fd, (off_t) truncated) == 0, "failed to truncate Engram transaction test file"); + expect_runtime( + [&] { runtime.prepare({ { 2, 1, { 0 }, 1 } }); }, + "Engram read failure was not surfaced"); + const llama_dsv41_engram_sequence_state after = runtime.sequence(0); + check(after.pos == before.pos && after.history.tail == before.history.tail, + "failed Engram transaction advanced sequence state"); +} +#endif + +static float bf16(float value) { + return ggml_bf16_to_fp32(ggml_fp32_to_bf16(value)); +} + +static void test_graph_gate() { + constexpr int64_t width = 8; + constexpr int64_t streams = 4; + constexpr int64_t tokens = 2; + + ggml_init_params params = { + /*.mem_size =*/ 8*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create DeepSeek V4.1 Engram graph context"); + + ggml_tensor * residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, width, streams, tokens); + ggml_tensor * rows = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, tokens); + ggml_tensor * engram_kv = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, LLAMA_ENGRAM_COLS*LLAMA_ENGRAM_DIM, 5*width); + ggml_tensor * q_norm = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, streams); + ggml_tensor * k_norm = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, streams); + ggml_tensor * select = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens); + + ggml_set_input(residual); + ggml_set_input(rows); + ggml_set_input(engram_kv); + ggml_set_input(q_norm); + ggml_set_input(k_norm); + ggml_set_input(select); + + std::vector residual_data(ggml_nelements(residual)); + std::vector rows_data(ggml_nelements(rows)); + std::vector engram_kv_data(ggml_nelements(engram_kv)); + std::vector q_data(ggml_nelements(q_norm)); + std::vector k_data(ggml_nelements(k_norm)); + for (int64_t i = 0; i < width*streams*tokens; ++i) { + residual_data[i] = i < width*streams ? + bf16(0.125f + (float) (i%13)/16.0f) : + 0.1234567f + (float) (i%13)/17.0f; + } + residual_data[width*streams] = -0.0f; + residual_data[width*streams + 1] = 0.0f; + std::vector projected_data(5*width*tokens); + for (int64_t token = 0; token < tokens; ++token) { + rows_data[token*rows->ne[0]] = 1.0f; + for (int64_t i = 0; i < 5*width; ++i) { + projected_data[token*5*width + i] = 0.0625f + (float) ((token*5*width + i)%11)/32.0f; + engram_kv_data[i*engram_kv->ne[0] + token] = projected_data[token*5*width + i]; + } + } + for (int64_t i = 0; i < width*streams; ++i) { + q_data[i] = 0.5f + (float) (i%5)/8.0f; + k_data[i] = 0.75f - (float) (i%3)/16.0f; + } + const int32_t select_data[] = { (int32_t) tokens, 1 }; + + const std::vector original = residual_data; + ggml_backend_t backend = ggml_backend_cpu_init(); + check(backend != nullptr, "failed to create DeepSeek V4.1 Engram graph backend"); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 512, false, true); + check(sched != nullptr, "failed to create DeepSeek V4.1 Engram graph scheduler"); + ggml_tensor * output = llama_dsv41_build_engram( + ctx, residual, rows, engram_kv, q_norm, k_norm, select, 1.0e-20f, sched, backend); + ggml_set_output(output); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + check(ggml_backend_sched_alloc_graph(sched, graph), "failed to allocate DeepSeek V4.1 Engram graph"); + ggml_backend_tensor_set(residual, residual_data.data(), 0, ggml_nbytes(residual)); + ggml_backend_tensor_set(rows, rows_data.data(), 0, ggml_nbytes(rows)); + ggml_backend_tensor_set(engram_kv, engram_kv_data.data(), 0, ggml_nbytes(engram_kv)); + ggml_backend_tensor_set(q_norm, q_data.data(), 0, ggml_nbytes(q_norm)); + ggml_backend_tensor_set(k_norm, k_data.data(), 0, ggml_nbytes(k_norm)); + ggml_backend_tensor_set(select, select_data, 0, sizeof(select_data)); + check(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS, + "DeepSeek V4.1 Engram graph execution failed"); + + std::vector actual(ggml_nelements(output)); + ggml_backend_tensor_get(output, actual.data(), 0, ggml_nbytes(output)); + for (int64_t stream = 0; stream < streams; ++stream) { + double hidden_sq = 0.0; + double key_sq = 0.0; + double dot = 0.0; + for (int64_t i = 0; i < width; ++i) { + const float hidden = original[stream*width + i]; + const float key = bf16(projected_data[stream*width + i]); + hidden_sq += hidden*hidden; + key_sq += key*key; + dot += hidden*q_data[stream*width + i]*k_data[stream*width + i]*key; + } + dot /= std::sqrt(hidden_sq/width + 1.0e-20); + dot /= std::sqrt(key_sq/width + 1.0e-20); + dot /= std::sqrt((double) width); + const double gate = 1.0/(1.0 + std::exp(-std::copysign(std::sqrt(std::max(std::abs(dot), 1.0e-6)), dot))); + for (int64_t i = 0; i < width; ++i) { + const float value = bf16(projected_data[4*width + i]); + const float expected = bf16(original[stream*width + i] + (float) gate*value); + check(std::abs(actual[stream*width + i] - expected) <= std::max(1.0e-6f, std::abs(expected)/128.0f), + "DeepSeek V4.1 Engram gate differs from scalar reference"); + const size_t masked = width*streams + stream*width + i; + check(std::memcmp(actual.data() + masked, original.data() + masked, sizeof(float)) == 0, + "masked DeepSeek V4.1 Engram row changed"); + } + } + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +static void test_signed_zero_gate() { + constexpr int64_t matrix_size = 32; + + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create DeepSeek V4.1 signed-zero graph context"); + + ggml_tensor * dot = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2); + ggml_tensor * matrix_a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, matrix_size, matrix_size); + ggml_tensor * matrix_b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, matrix_size, matrix_size); + ggml_set_input(dot); + ggml_set_input(matrix_a); + ggml_set_input(matrix_b); + ggml_tensor * accelerated = ggml_mul_mat(ctx, matrix_a, matrix_b); + + ggml_backend_load_all(); + ggml_backend_t backend_accel = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t device = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(device) == GGML_BACKEND_DEVICE_TYPE_ACCEL && + ggml_backend_dev_supports_op(device, accelerated)) { + backend_accel = ggml_backend_dev_init(device, nullptr); + if (backend_accel != nullptr) { + break; + } + } + } + + ggml_backend_t backend_cpu = ggml_backend_cpu_init(); + check(backend_cpu != nullptr, "failed to create DeepSeek V4.1 signed-zero CPU backend"); + ggml_backend_t backends[] = { + backend_accel != nullptr ? backend_accel : backend_cpu, + backend_cpu, + }; + ggml_backend_buffer_type_t bufts[] = { + ggml_backend_get_default_buffer_type(backends[0]), + ggml_backend_cpu_buffer_type(), + }; + const int n_backends = backend_accel != nullptr ? 2 : 1; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, bufts, n_backends, 32, false, true); + check(sched != nullptr, "failed to create DeepSeek V4.1 signed-zero scheduler"); + if (backend_accel != nullptr) { + ggml_backend_sched_set_tensor_backend(sched, accelerated, backend_accel); + expect_invalid( + [&] { llama_dsv41_build_engram_gate(ctx, dot, sched, backend_accel); }, + "DeepSeek V4.1 gate accepted a non-CPU backend"); + } + + ggml_tensor * gate = llama_dsv41_build_engram_gate(ctx, dot, sched, backend_cpu); + ggml_tensor * output = ggml_add(ctx, gate, ggml_repeat(ctx, ggml_sum(ctx, accelerated), gate)); + ggml_set_output(output); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + + check(ggml_backend_sched_alloc_graph(sched, graph), "failed to allocate DeepSeek V4.1 signed-zero graph"); + if (backend_accel != nullptr) { + check(ggml_backend_sched_get_tensor_backend(sched, accelerated) == backend_accel, + "DeepSeek V4.1 gate moved surrounding work off its accelerator"); + } + check(ggml_backend_sched_get_tensor_backend(sched, gate) == backend_cpu, + "DeepSeek V4.1 gate was not assigned to the local CPU backend"); + + const float input[] = { 0.0f, -0.0f }; + std::vector zeros(matrix_size*matrix_size); + ggml_backend_tensor_set(dot, input, 0, sizeof(input)); + ggml_backend_tensor_set(matrix_a, zeros.data(), 0, ggml_nbytes(matrix_a)); + ggml_backend_tensor_set(matrix_b, zeros.data(), 0, ggml_nbytes(matrix_b)); + check(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS, + "DeepSeek V4.1 signed-zero scheduler execution failed"); + const float positive = 1.0f/(1.0f + std::exp(-0.001f)); + const float negative = 1.0f/(1.0f + std::exp(0.001f)); + float actual[2]; + ggml_backend_tensor_get(output, actual, 0, sizeof(actual)); + check(std::abs(actual[0] - positive) < 1.0e-7f && actual[0] > 0.5f, + "DeepSeek V4.1 positive-zero gate lost copysign semantics"); + check(std::abs(actual[1] - negative) < 1.0e-7f && actual[1] < 0.5f, + "DeepSeek V4.1 negative-zero gate lost copysign semantics"); + + ggml_backend_sched_free(sched); + if (backend_accel != nullptr) { + ggml_backend_free(backend_accel); + } + ggml_backend_free(backend_cpu); + ggml_free(ctx); +} + +int main() { +#if !defined(_WIN32) + test_extent_validation(); + test_transactions_and_sequences(); + test_scheduler_upload(); + test_chunked_prefill(); + test_transactional_read_failure(); +#endif + test_graph_gate(); + test_signed_zero_gate(); + std::puts("DeepSeek V4.1 Engram runtime and graph: PASS"); + return 0; +} diff --git a/tests/test-deepseek41-expert.cpp b/tests/test-deepseek41-expert.cpp new file mode 100644 index 000000000000..269922196301 --- /dev/null +++ b/tests/test-deepseek41-expert.cpp @@ -0,0 +1,583 @@ +#include "../src/llama-dsv41-expert.h" +#include "../src/llama-dsv41.h" +#include "../src/llama-graph.h" +#include "../src/llama-model-loader.h" + +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml.h" +#include "gguf.h" +#include "../ggml/src/ggml-backend-impl.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define REQUIRE(cond) do { if (!(cond)) { throw std::runtime_error("requirement failed: " #cond); } } while (0) + +namespace { + +struct temp_file { + std::filesystem::path path; + + temp_file() { + static uint64_t sequence = 0; + const auto stamp = std::chrono::steady_clock::now().time_since_epoch().count(); + path = std::filesystem::temp_directory_path() / + ("llama-dsv41-expert-" + std::to_string(stamp) + "-" + std::to_string(++sequence) + ".gguf"); + } + + ~temp_file() { + std::error_code ec; + std::filesystem::remove(path, ec); + } +}; + +template +void require_throws(F && fn) { + bool threw = false; + try { + fn(); + } catch (const std::exception &) { + threw = true; + } + REQUIRE(threw); +} + +struct fixture { + static constexpr int64_t n_embd = 256; + static constexpr int64_t n_ff = 256; + static constexpr int64_t n_expert = LLAMA_DSV41_N_EXPERT; + + temp_file file; + std::vector tensors; + + fixture() { + const size_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd)*n_ff; + const size_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, n_ff)*n_embd; + const size_t data_size = (2*gate_plane + down_plane)*n_expert; + ggml_init_params params = { + /*.mem_size =*/ data_size + 8*ggml_tensor_overhead() + 4096, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ2_XXS, n_embd, n_ff, n_expert); + ggml_tensor * up = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ2_XXS, n_embd, n_ff, n_expert); + ggml_tensor * down = ggml_new_tensor_3d(ctx, GGML_TYPE_Q2_K, n_ff, n_embd, n_expert); + ggml_set_name(gate, "blk.0.ffn_gate_exps.weight"); + ggml_set_name(up, "blk.0.ffn_up_exps.weight"); + ggml_set_name(down, "blk.0.ffn_down_exps.weight"); + fill(gate, 0x10); + fill(up, 0x20); + fill(down, 0x30); + + gguf_context * gguf = gguf_init_empty(); + REQUIRE(gguf != nullptr); + gguf_set_val_str(gguf, "general.architecture", "deepseek41"); + gguf_add_tensor(gguf, gate); + gguf_add_tensor(gguf, up); + gguf_add_tensor(gguf, down); + REQUIRE(gguf_write_to_file(gguf, file.path.string().c_str(), false)); + gguf_free(gguf); + ggml_free(ctx); + + std::vector splits; + llama_model_loader loader( + nullptr, nullptr, nullptr, file.path.string(), splits, nullptr, + LLAMA_LOAD_MODE_MMAP, false, false, false, nullptr, nullptr); + const auto gate_extent = loader.register_external_tensor( + "blk.0.ffn_gate_exps.weight", 0, LLAMA_EXPERT_PROJECTION_GATE, { n_embd, n_ff, n_expert }); + const auto up_extent = loader.register_external_tensor( + "blk.0.ffn_up_exps.weight", 0, LLAMA_EXPERT_PROJECTION_UP, { n_embd, n_ff, n_expert }); + const auto down_extent = loader.register_external_tensor( + "blk.0.ffn_down_exps.weight", 0, LLAMA_EXPERT_PROJECTION_DOWN, { n_ff, n_embd, n_expert }); + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + for (auto extent : { gate_extent, up_extent, down_extent }) { + extent.layer = il; + extent.name = "blk." + std::to_string(il) + extent.name.substr(5); + tensors.push_back(std::move(extent)); + } + } + } + + static void fill(ggml_tensor * tensor, uint8_t tag) { + memset(tensor->data, 0, ggml_nbytes(tensor)); + for (int64_t expert = 0; expert < tensor->ne[2]; ++expert) { + uint8_t * plane = static_cast(tensor->data) + expert*tensor->nb[2]; + plane[tensor->nb[2] - 1] = tag + expert%32; + } + } + + size_t cache_bytes(size_t slots) const { + size_t result = 0; + for (const auto & tensor : tensors) { + result += tensor.nb[2]*slots; + } + return result; + } + + llama_dsv41_expert_runtime make_runtime( + size_t slots, + llama_dsv41_expert_runtime::upload_fn upload = {}, + llama_dsv41_expert_runtime::publish_fn before_publish = {}) const { + llama_dsv41_expert_runtime_params params; + params.cache_slots = slots; + params.cache_bytes = cache_bytes(slots); + params.direct_io = false; + return llama_dsv41_expert_runtime( + tensors, + params, + [](const llama_expert_store_tensor &) { return ggml_backend_cpu_buffer_type(); }, + std::move(upload), + std::move(before_publish)); + } +}; + +void test_registration() { + std::vector names; + const auto tensors = llama_dsv41_register_expert_tensors( + [&](const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne) { + names.push_back(name); + llama_expert_store_tensor tensor; + tensor.name = name; + tensor.fname = "unused"; + tensor.file_index = layer % 3; + tensor.layer = layer; + tensor.projection = projection; + tensor.type = projection == LLAMA_EXPERT_PROJECTION_DOWN ? GGML_TYPE_Q2_K : GGML_TYPE_IQ2_XXS; + std::copy(ne.begin(), ne.end(), tensor.ne); + tensor.nb[0] = ggml_type_size(tensor.type); + tensor.nb[1] = ggml_row_size(tensor.type, tensor.ne[0]); + tensor.nb[2] = tensor.nb[1]*tensor.ne[1]; + tensor.file_offset = 4096 + (size_t) layer*3*1024 + (size_t) projection*1024; + tensor.file_size = tensor.nb[2]*tensor.ne[2]; + return tensor; + }); + REQUIRE(tensors.size() == LLAMA_DSV41_N_LAYER*3); + REQUIRE(names.front() == "blk.0.ffn_gate_exps.weight"); + REQUIRE(names.back() == "blk.39.ffn_down_exps.weight"); + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + REQUIRE(tensors[3*il + 0].layer == il); + REQUIRE(tensors[3*il + 1].layer == il); + REQUIRE(tensors[3*il + 2].layer == il); + REQUIRE(tensors[3*il + 0].file_index == (size_t) il % 3); + REQUIRE(tensors[3*il + 0].file_offset == 4096 + (size_t) il*3*1024); + REQUIRE(tensors[3*il + 0].ne[2] == LLAMA_DSV41_N_EXPERT); + REQUIRE(tensors[3*il + 0].nb[2] == 3041280); + REQUIRE(tensors[3*il + 1].nb[2] == 3041280); + REQUIRE(tensors[3*il + 2].nb[2] == 3870720); + } + REQUIRE(std::none_of(names.begin(), names.end(), [](const std::string & name) { + return name.find("shexp") != std::string::npos; + })); + size_t one_slot_bytes = 0; + for (const auto & tensor : tensors) { + one_slot_bytes += tensor.nb[2]; + } + REQUIRE(one_slot_bytes == 398131200); +} + +void test_configuration(const fixture & f) { + llama_dsv41_expert_runtime_params params; + params.direct_io = false; + require_throws([&] { + llama_dsv41_expert_runtime runtime( + f.tensors, params, [](const llama_expert_store_tensor &) { return ggml_backend_cpu_buffer_type(); }); + }); + + params.cache_slots = 1; + params.cache_bytes = f.cache_bytes(1) - 1; + require_throws([&] { + llama_dsv41_expert_runtime runtime( + f.tensors, params, [](const llama_expert_store_tensor &) { return ggml_backend_cpu_buffer_type(); }); + }); + + auto runtime = f.make_runtime(1); + runtime.acquire_context(); + require_throws([&] { runtime.acquire_context(); }); + runtime.release_context(); + runtime.acquire_context(); + runtime.release_context(); +} + +void test_remap_upload_and_eviction(const fixture & f) { + struct upload_record { + std::string name; + size_t offset; + std::vector bytes; + }; + std::vector uploads; + auto runtime = f.make_runtime(3, [&](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + uploads.push_back({ tensor->name, offset, std::vector( + static_cast(data), static_cast(data) + size) }); + ggml_backend_tensor_set(tensor, data, offset, size); + }); + + const auto remapped = runtime.remap(0, { 3, 3, 1, 2 }); + REQUIRE(remapped == std::vector({ 2, 2, 0, 1 })); + REQUIRE(uploads.size() == 9); + REQUIRE(uploads[0].bytes.back() == 0x11); + REQUIRE(uploads[1].bytes.back() == 0x21); + REQUIRE(uploads[2].bytes.back() == 0x31); + require_throws([&] { runtime.remap(0, { 0 }); }); + runtime.release(0); + + uploads.clear(); + REQUIRE(runtime.remap(0, { 3 }).front() == 2); + REQUIRE(uploads.empty()); + runtime.release(0); + + REQUIRE(runtime.remap(0, { 4 }).front() == 0); + runtime.release(0); +} + +void test_capacity_and_upload_failure(const fixture & f) { + auto runtime = f.make_runtime(2); + require_throws([&] { runtime.remap(0, { 0, 1, 2 }); }); + REQUIRE(runtime.remap(0, { 0, 0, 1 }).size() == 3); + runtime.release(0); + + size_t calls = 0; + auto failing = f.make_runtime(1, [&](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + if (++calls == 5) { + throw std::runtime_error("synthetic upload failure"); + } + ggml_backend_tensor_set(tensor, data, offset, size); + }); + REQUIRE(failing.remap(0, { 0 }).front() == 0); + failing.release(0); + require_throws([&] { failing.remap(0, { 1 }); }); + const size_t after_failure = calls; + REQUIRE(failing.remap(0, { 0 }).front() == 0); + REQUIRE(calls == after_failure + 3); + failing.release(0); +} + +void test_publication_allocation_failure(const fixture & f) { + size_t attempts = 0; + auto runtime = f.make_runtime(1, {}, [&] { + if (attempts++ == 0) { + throw std::bad_alloc(); + } + }); + + require_throws([&] { runtime.remap(0, { 0 }); }); + REQUIRE(runtime.remap(0, { 1 }).front() == 0); + runtime.release(0); +} + +void test_graph_callbacks(const fixture & f) { + auto runtime = f.make_runtime(2); + ggml_init_params params = { + /*.mem_size =*/ 2*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 3, 1); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, fixture::n_embd, 3, 1); + ggml_set_input(selected); + ggml_set_input(input); + ggml_tensor * remapped = llama_dsv41_build_expert_remap(ctx, selected, runtime, 0, sched, backend); + ggml_tensor * expert_values = ggml_mul_mat_id( + ctx, runtime.cache_tensor(0, LLAMA_EXPERT_PROJECTION_GATE), input, remapped); + ggml_tensor * release = llama_dsv41_build_expert_release(ctx, expert_values, runtime, 0, sched, backend); + ggml_set_output(remapped); + ggml_set_output(release); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, release); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + const int32_t original[] = { 3, 1, 1 }; + std::vector input_data(fixture::n_embd*3, 1.0f); + ggml_backend_tensor_set(selected, original, 0, sizeof(original)); + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size()*sizeof(float)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + const std::string error = runtime.consume_error(); + if (!error.empty()) { + throw std::runtime_error(error); + } + REQUIRE(ggml_backend_sched_get_tensor_backend(sched, remapped) == backend); + REQUIRE(ggml_backend_sched_get_tensor_backend(sched, release) == backend); + REQUIRE(runtime.remap(0, { 0 }).front() >= 0); + runtime.release(0); + runtime.release_all(); + + const int32_t over_capacity[] = { 3, 1, 2 }; + ggml_backend_tensor_set(selected, over_capacity, 0, sizeof(over_capacity)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + REQUIRE(!runtime.consume_error().empty()); + int32_t safe_ids[3] = { -1, -1, -1 }; + ggml_backend_tensor_get(remapped, safe_ids, 0, sizeof(safe_ids)); + REQUIRE(safe_ids[0] == 0 && safe_ids[1] == 0 && safe_ids[2] == 0); + REQUIRE(runtime.remap(0, { 0 }).front() >= 0); + runtime.release(0); + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +void test_graph_upload_failure_sentinel(const fixture & f) { + size_t calls = 0; + auto runtime = f.make_runtime(1, [&](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + if (++calls == 5) { + throw std::runtime_error("synthetic graph upload failure"); + } + ggml_backend_tensor_set(tensor, data, offset, size); + }); + ggml_init_params params = { + /*.mem_size =*/ 2*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1, 1); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, fixture::n_embd, 1, 1); + ggml_set_input(selected); + ggml_set_input(input); + ggml_tensor * remapped = llama_dsv41_build_expert_remap(ctx, selected, runtime, 0, sched, backend); + ggml_tensor * expert_values = ggml_mul_mat_id( + ctx, runtime.cache_tensor(0, LLAMA_EXPERT_PROJECTION_GATE), input, remapped); + ggml_tensor * release = llama_dsv41_build_expert_release(ctx, expert_values, runtime, 0, sched, backend); + ggml_set_output(remapped); + ggml_set_output(release); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, release); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + + std::vector input_data(fixture::n_embd, 1.0f); + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size()*sizeof(float)); + const int32_t first[] = { 0 }; + ggml_backend_tensor_set(selected, first, 0, sizeof(first)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + REQUIRE(runtime.consume_error().empty()); + + const int32_t failed[] = { 1 }; + ggml_backend_tensor_set(selected, failed, 0, sizeof(failed)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + REQUIRE(runtime.consume_error().find("synthetic graph upload failure") != std::string::npos); + int32_t safe_id = -1; + ggml_backend_tensor_get(remapped, &safe_id, 0, sizeof(safe_id)); + REQUIRE(safe_id == 0); + + REQUIRE(runtime.remap(0, { 2 }).front() == 0); + runtime.release(0); + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +void test_grovemoe_lookup_ids() { + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + ggml_tensor * explicit_slots = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + ggml_set_input(selected); + ggml_set_input(explicit_slots); + const llm_moe_expert_ids grovemoe = llm_build_moe_expert_ids( + ctx, LLM_ARCH_GROVEMOE, selected, nullptr, 2, 8, 4); + const llm_moe_expert_ids deepseek = llm_build_moe_expert_ids( + ctx, LLM_ARCH_DEEPSEEK41, selected, explicit_slots, 8, 8, 0); + REQUIRE(grovemoe.routing == selected); + REQUIRE(grovemoe.lookup != selected); + REQUIRE(deepseek.routing == selected); + REQUIRE(deepseek.lookup == explicit_slots); + ggml_set_output(grovemoe.routing); + ggml_set_output(grovemoe.lookup); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, grovemoe.routing); + ggml_build_forward_expand(graph, grovemoe.lookup); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + + const int32_t original[] = { 7, 1 }; + ggml_backend_tensor_set(selected, original, 0, sizeof(original)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + int32_t routing_ids[2] = {}; + int32_t chunk_ids[2] = {}; + ggml_backend_tensor_get(grovemoe.routing, routing_ids, 0, sizeof(routing_ids)); + ggml_backend_tensor_get(grovemoe.lookup, chunk_ids, 0, sizeof(chunk_ids)); + REQUIRE(routing_ids[0] == 7 && routing_ids[1] == 1); + REQUIRE(chunk_ids[0] == 1 && chunk_ids[1] == 0); + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +struct sync_test_context { + llama_dsv41_expert_runtime * runtime = nullptr; + int synchronize_count = 0; + bool saw_pinned = false; +}; + +const char * sync_test_backend_name(ggml_backend_t) { + return "dsv41-sync-test"; +} + +void sync_test_backend_synchronize(ggml_backend_t backend) { + auto * state = static_cast(backend->context); + state->synchronize_count++; + try { + state->runtime->remap(0, { 2 }); + } catch (const std::exception &) { + state->saw_pinned = true; + } +} + +const char * sync_test_device_name(ggml_backend_dev_t) { + return "dsv41-sync-test"; +} + +enum ggml_backend_dev_type sync_test_device_type(ggml_backend_dev_t) { + return GGML_BACKEND_DEVICE_TYPE_CPU; +} + +bool sync_test_device_supports_op(ggml_backend_dev_t, const ggml_tensor *) { + return true; +} + +bool sync_test_device_supports_buft(ggml_backend_dev_t, ggml_backend_buffer_type_t buft) { + return buft == ggml_backend_cpu_buffer_type(); +} + +void test_release_after_sync(const fixture & f) { + auto runtime = f.make_runtime(1); + REQUIRE(runtime.remap(0, { 1 }).front() == 0); + + sync_test_context state = { &runtime }; + ggml_backend_device device = {}; + device.iface.get_name = sync_test_device_name; + device.iface.get_type = sync_test_device_type; + device.iface.supports_op = sync_test_device_supports_op; + device.iface.supports_buft = sync_test_device_supports_buft; + ggml_backend backend = {}; + backend.iface.get_name = sync_test_backend_name; + backend.iface.synchronize = sync_test_backend_synchronize; + backend.device = &device; + backend.context = &state; + ggml_backend_t backends[] = { &backend }; + ggml_backend_buffer_type_t bufts[] = { ggml_backend_cpu_buffer_type() }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, bufts, 1, 16, false, false); + REQUIRE(sched != nullptr); + + runtime.release_all_after_sync(sched); + REQUIRE(state.synchronize_count == 1); + REQUIRE(state.saw_pinned); + REQUIRE(runtime.remap(0, { 2 }).front() == 0); + runtime.release(0); + ggml_backend_sched_free(sched); +} + +void test_original_and_slot_ids() { + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * cache = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, 2); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 2, 1); + ggml_tensor * slots = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + ggml_tensor * probs = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 4, 1); + ggml_tensor * original = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + for (ggml_tensor * tensor : { cache, input, slots, probs, original }) { + ggml_set_input(tensor); + } + + ggml_tensor * expert_values = ggml_mul_mat_id(ctx, cache, input, slots); + ggml_tensor * routing_weights = ggml_get_rows(ctx, probs, original); + ggml_set_output(expert_values); + ggml_set_output(routing_weights); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, expert_values); + ggml_build_forward_expand(graph, routing_weights); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + + const float cache_data[] = { 10.0f, 20.0f }; + const float input_data[] = { 1.0f, 1.0f }; + const int32_t slot_data[] = { 1, 0 }; + const float prob_data[] = { 1.0f, 2.0f, 3.0f, 4.0f }; + const int32_t original_data[] = { 3, 1 }; + ggml_backend_tensor_set(cache, cache_data, 0, sizeof(cache_data)); + ggml_backend_tensor_set(input, input_data, 0, sizeof(input_data)); + ggml_backend_tensor_set(slots, slot_data, 0, sizeof(slot_data)); + ggml_backend_tensor_set(probs, prob_data, 0, sizeof(prob_data)); + ggml_backend_tensor_set(original, original_data, 0, sizeof(original_data)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + + float expert_result[2] = {}; + float weight_result[2] = {}; + ggml_backend_tensor_get(expert_values, expert_result, 0, sizeof(expert_result)); + ggml_backend_tensor_get(routing_weights, weight_result, 0, sizeof(weight_result)); + REQUIRE(expert_result[0] == 20.0f && expert_result[1] == 10.0f); + REQUIRE(weight_result[0] == 4.0f && weight_result[1] == 2.0f); + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +} + +int main() { + try { + test_registration(); + fixture f; + test_configuration(f); + test_remap_upload_and_eviction(f); + test_capacity_and_upload_failure(f); + test_publication_allocation_failure(f); + test_graph_callbacks(f); + test_graph_upload_failure_sentinel(f); + test_grovemoe_lookup_ids(); + test_release_after_sync(f); + test_original_and_slot_ids(); + } catch (const std::exception & error) { + std::fprintf(stderr, "test-deepseek41-expert: %s\n", error.what()); + return 1; + } + return 0; +} diff --git a/tests/test-deepseek41-memory.cpp b/tests/test-deepseek41-memory.cpp new file mode 100644 index 000000000000..057b9e4848f4 --- /dev/null +++ b/tests/test-deepseek41-memory.cpp @@ -0,0 +1,736 @@ +#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; +}; + +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); + + 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"); + + 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() { + 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(2048, 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"); + + 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 + +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 new file mode 100644 index 000000000000..696abea90100 --- /dev/null +++ b/tests/test-deepseek41-runtime.cpp @@ -0,0 +1,657 @@ +#include "../src/llama-dsv41.h" +#include "../src/llama-arch.h" + +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#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) { + std::fprintf(stderr, "%s\n", message.c_str()); + std::exit(1); + } +} + +static void expect_throw(const std::function & fn, const std::string & message) { + try { + fn(); + } catch (const std::runtime_error &) { + return; + } + check(false, message); +} + +static llama_dsv41_config valid_config() { + llama_dsv41_config config = {}; + config.n_ctx_train = LLAMA_DSV41_N_CTX; + config.n_embd = LLAMA_DSV41_N_EMBD; + config.n_layer = LLAMA_DSV41_N_LAYER; + config.n_vocab = LLAMA_DSV41_N_VOCAB; + config.n_head = LLAMA_DSV41_N_HEAD; + config.n_head_kv = LLAMA_DSV41_N_HEAD_KV; + config.n_head_dim = LLAMA_DSV41_N_HEAD_DIM; + config.n_rot = LLAMA_DSV41_N_ROT; + config.n_lora_q = LLAMA_DSV41_N_LORA_Q; + config.n_lora_o = LLAMA_DSV41_N_LORA_O; + config.n_o_group = LLAMA_DSV41_N_O_GROUP; + config.n_ff_dense = LLAMA_DSV41_N_FF_DENSE; + config.n_ff_expert = LLAMA_DSV41_N_FF_EXP; + config.n_expert = LLAMA_DSV41_N_EXPERT; + config.n_expert_used = LLAMA_DSV41_N_EXPERT_USED; + config.n_expert_shared = LLAMA_DSV41_N_EXPERT_SHARED; + config.indexer_n_head = LLAMA_DSV41_N_INDEX_HEAD; + config.indexer_head_size = LLAMA_DSV41_N_INDEX_HEAD_DIM; + config.indexer_top_k = LLAMA_DSV41_N_INDEX_TOP_K; + config.hc_count = LLAMA_DSV41_HC_MULT; + config.hc_sinkhorn_iters = LLAMA_DSV41_HC_SINKHORN_ITERS; + config.raw_window = LLAMA_DSV41_N_SWA; + config.candidate_source_layer = LLAMA_DSV41_CANDIDATE_SOURCE_LAYER; + config.candidate_topk_blocks = LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS; + config.candidate_block_size = LLAMA_DSV41_CANDIDATE_BLOCK_SIZE; + config.f_norm_rms_eps = 1.0e-20f; + config.hc_eps = 1.0e-6f; + config.swiglu_clamp = 10.0f; + config.routed_scale = 1.5f; + config.rope_theta = 10000.0f; + config.compress_rope_theta = 160000.0f; + config.yarn_factor = 16.0f; + config.yarn_beta_fast = 32.0f; + config.yarn_beta_slow = 1.0f; + config.yarn_original_context = 65536; + config.expert_weights_norm = true; + config.hidden_act = "silu"; + config.scoring_func = "sqrtsoftplus"; + config.topk_method = "noaux_tc"; + for (uint32_t il = 0; il < LLAMA_DSV41_N_LAYER; ++il) { + config.compress_ratios.push_back(llama_dsv41_compress_ratio(il)); + } + config.kv_sources = { 2, 8, 14, 20 }; + config.index_sources = { 2, 8, 14, 20, 24, 28, 32, 36 }; + config.engram_layers = { 1, 14 }; + config.engram_rows = { 384006168, 384016682 }; + config.engram_encoding = LLAMA_DSV41_ENGRAM_ENCODING; + config.engram_compressed_vocab_size = LLAMA_DSV41_ENGRAM_COMPRESSED_VOCAB; + config.engram_pad_id = LLAMA_DSV41_ENGRAM_PAD_ID; + config.engram_token_map_size = LLAMA_DSV41_N_VOCAB; + config.engram_primes_size = LLAMA_DSV41_ENGRAM_PRIMES_COUNT; + config.engram_multipliers_size = LLAMA_DSV41_ENGRAM_MULTIPLIERS_COUNT; + return config; +} + +static void test_hparams() { + static_assert(std::is_same_v); + static_assert(std::is_same_v); + static_assert(std::is_same_v); + + llama_dsv41_validate_config(valid_config()); + check(valid_config().engram_rows == std::vector({ 384006168, 384016682 }), "published Engram rows mismatch"); + + llama_dsv41_config config = valid_config(); + config.compress_ratios.insert(config.compress_ratios.end(), 3, 0); + expect_throw([&]() { llama_dsv41_validate_config(config); }, "43-entry source-config compression layout was accepted"); + + config = valid_config(); + config.compress_ratios[20] = 2; + expect_throw([&]() { llama_dsv41_validate_config(config); }, "invalid ratio-1 boundary was accepted"); + + config = valid_config(); + config.kv_sources = { 2, 8, 20 }; + expect_throw([&]() { llama_dsv41_validate_config(config); }, "invalid KV source map was accepted"); + + config = valid_config(); + config.engram_primes_size = 24; + expect_throw([&]() { llama_dsv41_validate_config(config); }, "truncated Engram prime table was accepted"); + + 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() { + check(llama_dsv41_compress_ratio(0) == 0, "layer 0 ratio mismatch"); + check(llama_dsv41_compress_ratio(1) == 0, "layer 1 ratio mismatch"); + check(llama_dsv41_compress_ratio(2) == 2, "layer 2 ratio mismatch"); + check(llama_dsv41_compress_ratio(19) == 2, "layer 19 ratio mismatch"); + check(llama_dsv41_compress_ratio(20) == 1, "layer 20 ratio mismatch"); + + check(llama_dsv41_kv_source_layer(0) == -1, "layer 0 unexpectedly has a KV source"); + check(llama_dsv41_kv_source_layer(2) == 2, "layer 2 KV source mismatch"); + check(llama_dsv41_kv_source_layer(7) == 2, "layer 7 KV source mismatch"); + check(llama_dsv41_kv_source_layer(8) == 8, "layer 8 KV source mismatch"); + check(llama_dsv41_kv_source_layer(19) == 14, "layer 19 KV source mismatch"); + check(llama_dsv41_kv_source_layer(39) == 20, "layer 39 KV source mismatch"); + + check(llama_dsv41_index_source_layer(19) == 14, "layer 19 index source mismatch"); + check(llama_dsv41_index_source_layer(20) == 20, "layer 20 index source mismatch"); + check(llama_dsv41_index_source_layer(23) == 20, "layer 23 index source mismatch"); + check(llama_dsv41_index_source_layer(24) == 24, "layer 24 index source mismatch"); + check(llama_dsv41_index_source_layer(39) == 36, "layer 39 index source mismatch"); +} + +static void test_compression() { + const auto ratio_2 = llama_dsv41_build_compression_plan({ 0, 1, 2 }, 2, 1024); + check(ratio_2.n_visible == std::vector({ 0, 1, 1 }), "ratio-2 visible counts mismatch"); + check(ratio_2.write_idxs == std::vector({ 0 }), "ratio-2 write index mismatch"); + check(ratio_2.write_pos == std::vector({ 0 }), "ratio-2 compressed position mismatch"); + check(ratio_2.state_persist_dst_idxs == std::vector({ 0, 1 }), "ratio-2 state rows mismatch"); + + const auto ratio_1 = llama_dsv41_build_compression_plan({ 19, 20 }, 1, 1024); + check(ratio_1.n_visible == std::vector({ 20, 21 }), "ratio-1 visible counts mismatch"); + check(ratio_1.write_idxs == std::vector({ 19, 20 }), "ratio-1 write indexes mismatch"); + check(ratio_1.write_pos == std::vector({ 19, 20 }), "ratio-1 compressed positions mismatch"); + + const auto layer_0 = llama_dsv41_build_layer_plan(0, { 0 }, 1024); + check(layer_0.ratio == 0 && layer_0.compression.write_idxs.empty(), "layer 0 must use raw attention only"); + const auto layer_2 = llama_dsv41_build_layer_plan(2, { 0, 1 }, 1024); + check(layer_2.ratio == 2 && layer_2.owns_kv_source, "layer 2 compression ownership mismatch"); + check(layer_2.compression.write_idxs == std::vector({ 0 }), "layer 2 graph compression mismatch"); + const auto layer_20 = llama_dsv41_build_layer_plan(20, { 20 }, 1024); + check(layer_20.ratio == 1 && layer_20.owns_kv_source, "layer 20 compression ownership mismatch"); + check(layer_20.builds_candidates && !layer_20.uses_candidates, "layer 20 candidate propagation mismatch"); + const auto layer_21 = llama_dsv41_build_layer_plan(21, { 21 }, 1024); + check(!layer_21.uses_candidates && layer_21.reuses_index_selection, "layer 21 index reuse mismatch"); + const auto layer_24 = llama_dsv41_build_layer_plan(24, { 24 }, 1024); + check(!layer_24.owns_kv_source && layer_24.owns_index_source, "layer 24 source ownership mismatch"); + check(layer_24.uses_candidates, "layer 24 must consume layer-20 candidates"); + check(llama_dsv41_build_layer_plan(39, { 39 }, 1024).collapses_output, "final layer output collapse missing"); + expect_throw([&]() { llama_dsv41_build_layer_plan(20, { 20, 22 }, 1024); }, "non-contiguous graph plan was accepted"); +} + +static void test_state() { + llama_dsv41_cache_state state(1024); + for (llama_pos pos = 0; pos <= 129; ++pos) { + state.append(pos); + } + check(state.position() == 129, "cache position mismatch"); + check(state.raw_slots()[0] == 128 && state.raw_slots()[1] == 129, "raw ring state mismatch"); + check(state.compressed_slots(2)[0] == 0, "ratio-2 first compressed row mismatch"); + check(state.compressed_slots(2)[64] == 128, "ratio-2 boundary row mismatch"); + check(state.pending_slots(2) == std::vector({ 128, 129 }), "ratio-2 pending rows mismatch"); + check(state.compressed_slots(20)[129] == 129, "ratio-1 direct row mismatch"); + state.set_candidate_blocks({ 4, 1 }); + check(state.candidate_blocks() == std::vector({ 4, 1 }), "candidate state mismatch"); + expect_throw([&]() { state.append(131); }, "non-contiguous cache append was accepted"); + state.clear(); + check(state.position() == -1 && state.raw_slots()[0] == -1, "cache clear mismatch"); + + llama_dsv41_cache_state small(1); + small.append(0); + expect_throw([&]() { small.append(1); }, "compressed cache overflow was accepted"); + check(small.position() == 0 && small.raw_slots()[1] == -1, "failed cache append mutated state"); + + const auto bytes = llama_dsv41_account_memory(32768, 1, 8192, 2, 2, 1234); + check(bytes.raw_kv > 0 && bytes.compressed_kv > 0 && bytes.index_keys > 0, "cache memory accounting is incomplete"); + check(bytes.compressor_carry > 0 && bytes.candidate_scores > 0 && bytes.candidate_ids > 0, "state memory accounting is incomplete"); + check(bytes.total() == bytes.raw_kv + bytes.compressed_kv + bytes.index_keys + bytes.compressor_carry + + bytes.candidate_scores + bytes.candidate_ids + bytes.position_state + bytes.graph_workspace, + "memory accounting total mismatch"); +} + +static void test_raw_ring() { + const auto at_127 = llama_dsv41_raw_ring_order(127, 128); + check(at_127.size() == 128 && at_127.front() == 0 && at_127.back() == 127, "raw ring at 127 mismatch"); + const auto at_128 = llama_dsv41_raw_ring_order(128, 128); + check(at_128.front() == 1 && at_128.back() == 0, "raw ring at 128 mismatch"); + const auto at_129 = llama_dsv41_raw_ring_order(129, 128); + check(at_129.front() == 2 && at_129.back() == 1, "raw ring at 129 mismatch"); +} + +static void test_candidates() { + for (uint32_t n_visible : { 1u, 7u, 8u, 9u, 127u, 16385u, 16392u, 17017u }) { + std::vector scores(n_visible); + for (uint32_t i = 0; i < n_visible; ++i) { + scores[i] = -(float) i; + } + const auto blocks = llama_dsv41_select_candidate_blocks(scores, n_visible, 8, 2048); + const int32_t final_block = (int32_t) ((n_visible - 1)/8); + check(std::find(blocks.begin(), blocks.end(), final_block) != blocks.end(), "final visible candidate block was dropped"); + check(blocks.size() == std::min(2048, (n_visible + 7)/8), "candidate block count mismatch"); + const auto rows = llama_dsv41_candidate_rows(blocks, n_visible, 8); + check(std::find(rows.begin(), rows.end(), (int32_t) n_visible - 1) != rows.end(), "final visible row was filtered"); + check(std::all_of(rows.begin(), rows.end(), [&](int32_t row) { return row >= 0 && (uint32_t) row < n_visible; }), "candidate rows crossed causal visibility"); + } + + const auto tie = llama_dsv41_select_candidate_blocks(std::vector(24, 1.0f), 24, 8, 2); + check(tie == std::vector({ 2, 0 }), "candidate tie-break or final block retention mismatch"); + + std::vector partial_scores(9, -100.0f); + partial_scores[0] = 100.0f; + const auto partial = llama_dsv41_select_candidate_blocks(partial_scores, 9, 8, 1); + check(partial == std::vector({ 1 }), "final partial candidate block was not forced"); + + std::vector full_scores(16392, -100.0f); + full_scores[0] = 100.0f; + const auto full = llama_dsv41_select_candidate_blocks(full_scores, 16392, 8, 2048); + check(std::find(full.begin(), full.end(), 2048) != full.end(), "final full candidate block was not forced"); + + const auto inf_tie = llama_dsv41_select_candidate_blocks( + std::vector(24, std::numeric_limits::infinity()), 24, 8, 2); + check(std::find(inf_tie.begin(), inf_tie.end(), 2) != inf_tie.end(), "final candidate block lost an infinity tie"); + + const auto inf_boundary = llama_dsv41_select_candidate_blocks( + std::vector(16392, std::numeric_limits::infinity()), 16392, 8, 2048); + check(std::find(inf_boundary.begin(), inf_boundary.end(), 2048) != inf_boundary.end(), + "final full candidate block lost an infinity tie"); + std::vector unique_blocks = inf_boundary; + std::sort(unique_blocks.begin(), unique_blocks.end()); + check(std::adjacent_find(unique_blocks.begin(), unique_blocks.end()) == unique_blocks.end(), + "candidate selection contains duplicate blocks"); + check(std::all_of(unique_blocks.begin(), unique_blocks.end(), [](int32_t block) { + return block >= 0 && block <= 2048; + }), "candidate selection contains an out-of-range block"); +} + +static void test_output_collapse() { + const std::vector residual = { + 1.0f, 2.0f, + 3.0f, 4.0f, + 5.0f, 6.0f, + 7.0f, 8.0f, + }; + const auto result = llama_dsv41_output_collapse(residual, { 0.1f, 0.2f, 0.3f, 0.4f }, 2, 4); + check(result.size() == 2, "output collapse width mismatch"); + check(std::abs(result[0] - 5.0f) < 1.0e-6f, "output collapse first value mismatch"); + check(std::abs(result[1] - 6.0f) < 1.0e-6f, "output collapse second value mismatch"); +} + +static void test_candidate_graph_selection(uint32_t n_visible, uint32_t n_candidate) { + ggml_init_params params = { + /*.mem_size =*/ 4*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create candidate graph context"); + + ggml_tensor * scores = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_visible, 1); + std::fill_n( + static_cast(scores->data), + ggml_nelements(scores), + std::numeric_limits::infinity()); + ggml_tensor * block_scores = ggml_pool_1d( + ctx, scores, GGML_OP_POOL_MAX, 8, 8, 0); + ggml_tensor * final_blocks = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, 1, 1); + const int32_t final_block = (int32_t) block_scores->ne[0] - 1; + static_cast(final_blocks->data)[0] = final_block; + ggml_tensor * candidate_blocks = llama_dsv41_build_candidate_blocks( + ctx, block_scores, final_blocks, n_candidate); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, candidate_blocks); + check( + ggml_graph_compute_with_ctx(ctx, gf, 1) == GGML_STATUS_SUCCESS, + "candidate selection graph execution failed"); + + const int32_t * ids = static_cast(candidate_blocks->data); + std::vector selected(ids, ids + n_candidate); + check(selected.front() == final_block, "final candidate block is not first"); + check(std::count(selected.begin(), selected.end(), final_block) == 1, + "final candidate block was not retained exactly once"); + std::vector unique = selected; + std::sort(unique.begin(), unique.end()); + check(std::adjacent_find(unique.begin(), unique.end()) == unique.end(), + "candidate graph selected duplicate blocks"); + check(std::all_of(unique.begin(), unique.end(), [&](int32_t block) { + return block >= 0 && block <= final_block; + }), "candidate graph selected an out-of-range block"); + + ggml_free(ctx); +} + +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, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(params); + check(ctx != nullptr, "failed to create graph test context"); + + ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 512, 2, 3); + ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 512, 2, 3); + ggml_tensor * pooled = llama_dsv41_build_ratio_pool(ctx, kv, gate, 2); + check(pooled->ne[0] == 512 && pooled->ne[1] == 3, "ratio-2 graph output shape mismatch"); + + ggml_tensor * direct = llama_dsv41_build_ratio_pool( + ctx, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 512, 1, 3), nullptr, 1); + check(direct->ne[0] == 512 && direct->ne[1] == 3, "ratio-1 graph output shape mismatch"); + + ggml_tensor * compressed_scores = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 32, 64, 2); + ggml_tensor * raw_scores = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 128, 64, 2); + ggml_tensor * probs = llama_dsv41_build_shared_softmax(ctx, raw_scores, compressed_scores); + check(probs->ne[0] == 160 && probs->ne[1] == 64 && probs->ne[2] == 2, "shared-softmax graph shape mismatch"); + + ggml_tensor * raw_order = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2); + ggml_tensor * compressed_order = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + const float raw_values[] = { 1.0f, 2.0f }; + const float compressed_value = 3.0f; + std::memcpy(raw_order->data, raw_values, sizeof(raw_values)); + std::memcpy(compressed_order->data, &compressed_value, sizeof(compressed_value)); + ggml_tensor * ordered_probs = llama_dsv41_build_shared_softmax(ctx, raw_order, compressed_order); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, ordered_probs); + check(ggml_graph_compute_with_ctx(ctx, gf, 1) == GGML_STATUS_SUCCESS, "shared-softmax graph execution failed"); + 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 * 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_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, + "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"); + 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); + ggml_tensor * collapsed = llama_dsv41_build_output_collapse(ctx, residual, pre, 32, 4, 2); + check(collapsed->type == GGML_TYPE_BF16, "output collapse BF16 boundary is missing"); + check(collapsed->ne[0] == 32 && collapsed->ne[1] == 2, "output collapse graph shape mismatch"); + + ggml_tensor * output_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 32); + ggml_tensor * output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 64); + 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_tensor * exec_residual = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 32, 4, 1); + ggml_tensor * exec_pre = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 1); + ggml_tensor * exec_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 32); + ggml_tensor * exec_output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32, 2); + std::fill_n(static_cast(exec_residual->data), ggml_nelements(exec_residual), 1.0f); + std::fill_n(static_cast(exec_pre->data), ggml_nelements(exec_pre), 0.25f); + std::fill_n(static_cast(exec_norm->data), ggml_nelements(exec_norm), 1.0f); + std::fill_n(static_cast(exec_output->data), ggml_nelements(exec_output), 1.0f); + ggml_tensor * exec_collapse = llama_dsv41_build_output_collapse( + ctx, exec_residual, exec_pre, 32, 4, 1); + ggml_tensor * exec_norm_input = llama_dsv41_build_output_norm_input( + ctx, exec_collapse); + check(exec_collapse->type == GGML_TYPE_BF16, + "production output collapse is not BF16"); + check(exec_norm_input->type == GGML_TYPE_F32, + "production output RMSNorm input is not F32"); + ggml_tensor * exec_normalized = ggml_rms_norm( + ctx, exec_norm_input, 1.0e-20f); + check(exec_normalized->src[0] == exec_norm_input, + "production RMSNorm does not consume the F32 collapse"); + exec_normalized = ggml_mul(ctx, exec_normalized, exec_norm); + ggml_tensor * exec_logits = ggml_mul_mat( + ctx, exec_output, exec_normalized); + ggml_cgraph * exec_gf = ggml_new_graph(ctx); + ggml_build_forward_expand(exec_gf, exec_logits); + ggml_backend_t exec_backend = ggml_backend_cpu_init(); + check(exec_backend != nullptr, "failed to create output graph CPU backend"); + for (int i = 0; i < ggml_graph_n_nodes(exec_gf); ++i) { + check(ggml_backend_supports_op(exec_backend, ggml_graph_node(exec_gf, i)), + "CPU backend does not support the production output graph"); + } + ggml_backend_free(exec_backend); + check(ggml_graph_compute_with_ctx(ctx, exec_gf, 1) == GGML_STATUS_SUCCESS, "final output graph execution failed"); + const float * exec_values = static_cast(exec_logits->data); + check(std::isfinite(exec_values[0]) && std::isfinite(exec_values[1]), "final output graph produced non-finite logits"); + check(std::abs(exec_values[0] - 32.0f) < 1.0e-4f && std::abs(exec_values[1] - 32.0f) < 1.0e-4f, + "final output graph numeric mismatch"); + + ggml_free(ctx); +} + +int main() { + test_hparams(); + test_source_maps(); + test_compression(); + test_state(); + test_raw_ring(); + test_candidates(); + test_output_collapse(); + test_candidate_graph_selection(24, 1); + test_candidate_graph_selection(24, 2); + test_candidate_graph_selection(16392, 2048); + test_graph_contract(); + test_graph_construction(); + return 0; +} diff --git a/tests/test-deepseek41-schema.cpp b/tests/test-deepseek41-schema.cpp new file mode 100644 index 000000000000..4ae136d90c1a --- /dev/null +++ b/tests/test-deepseek41-schema.cpp @@ -0,0 +1,86 @@ +#include "../src/llama-arch.h" + +#include +#include +#include +#include +#include + +static void check(bool condition, const std::string & message) { + if (!condition) { + std::fprintf(stderr, "%s\n", message.c_str()); + std::exit(1); + } +} + +int main() { + check(llm_arch_from_string("deepseek41") == LLM_ARCH_DEEPSEEK41, "deepseek41 architecture lookup failed"); + check(std::string(llm_arch_name(LLM_ARCH_DEEPSEEK41)) == "deepseek41", "deepseek41 architecture name failed"); + + const LLM_KV kv(LLM_ARCH_DEEPSEEK41); + const std::vector> keys = { + { LLM_KV_DSV41_CONFIG, "deepseek41.config" }, + { LLM_KV_DSV41_VOCAB_SIZE, "deepseek41.vocab_size" }, + { LLM_KV_DSV41_HIDDEN_SIZE, "deepseek41.hidden_size" }, + { LLM_KV_DSV41_MOE_INTERMEDIATE_SIZE, "deepseek41.moe_intermediate_size" }, + { LLM_KV_DSV41_NUM_HIDDEN_LAYERS, "deepseek41.num_hidden_layers" }, + { LLM_KV_DSV41_NUM_ATTENTION_HEADS, "deepseek41.num_attention_heads" }, + { LLM_KV_DSV41_NUM_KEY_VALUE_HEADS, "deepseek41.num_key_value_heads" }, + { LLM_KV_DSV41_HEAD_DIM, "deepseek41.head_dim" }, + { LLM_KV_DSV41_QK_ROPE_HEAD_DIM, "deepseek41.qk_rope_head_dim" }, + { LLM_KV_DSV41_Q_LORA_RANK, "deepseek41.q_lora_rank" }, + { LLM_KV_DSV41_O_LORA_RANK, "deepseek41.o_lora_rank" }, + { LLM_KV_DSV41_O_GROUPS, "deepseek41.o_groups" }, + { LLM_KV_DSV41_N_ROUTED_EXPERTS, "deepseek41.n_routed_experts" }, + { LLM_KV_DSV41_N_SHARED_EXPERTS, "deepseek41.n_shared_experts" }, + { LLM_KV_DSV41_NUM_EXPERTS_PER_TOK, "deepseek41.num_experts_per_tok" }, + { LLM_KV_DSV41_MAX_POSITION_EMBEDDINGS, "deepseek41.max_position_embeddings" }, + { LLM_KV_DSV41_SLIDING_WINDOW, "deepseek41.sliding_window" }, + { LLM_KV_DSV41_INDEX_N_HEADS, "deepseek41.index_n_heads" }, + { LLM_KV_DSV41_INDEX_HEAD_DIM, "deepseek41.index_head_dim" }, + { LLM_KV_DSV41_INDEX_TOPK, "deepseek41.index_topk" }, + { LLM_KV_DSV41_CANDIDATE_SOURCE_LAYER_ID, "deepseek41.candidate_source_layer_id" }, + { LLM_KV_DSV41_CANDIDATE_TOPK_BLOCKS, "deepseek41.candidate_topk_blocks" }, + { LLM_KV_DSV41_CANDIDATE_BLOCK_SIZE, "deepseek41.candidate_block_size" }, + { LLM_KV_DSV41_HC_MULT, "deepseek41.hc_mult" }, + { LLM_KV_DSV41_HC_SINKHORN_ITERS, "deepseek41.hc_sinkhorn_iters" }, + { LLM_KV_DSV41_ROPE_THETA, "deepseek41.rope_theta" }, + { LLM_KV_DSV41_COMPRESS_ROPE_THETA, "deepseek41.compress_rope_theta" }, + { LLM_KV_DSV41_RMS_NORM_EPS, "deepseek41.rms_norm_eps" }, + { LLM_KV_DSV41_HC_EPS, "deepseek41.hc_eps" }, + { LLM_KV_DSV41_SWIGLU_LIMIT, "deepseek41.swiglu_limit" }, + { LLM_KV_DSV41_ROUTED_SCALING_FACTOR, "deepseek41.routed_scaling_factor" }, + { LLM_KV_DSV41_SCORING_FUNC, "deepseek41.scoring_func" }, + { LLM_KV_DSV41_HIDDEN_ACT, "deepseek41.hidden_act" }, + { LLM_KV_DSV41_TOPK_METHOD, "deepseek41.topk_method" }, + { LLM_KV_DSV41_NORM_TOPK_PROB, "deepseek41.norm_topk_prob" }, + { LLM_KV_DSV41_COMPRESS_RATIOS, "deepseek41.compress_ratios" }, + { LLM_KV_DSV41_KV_SOURCE_LAYER_IDS, "deepseek41.kv_source_layer_ids" }, + { LLM_KV_DSV41_INDEX_SOURCE_LAYER_IDS, "deepseek41.index_source_layer_ids" }, + { LLM_KV_DSV41_ROPE_SCALING_FACTOR, "deepseek41.rope_scaling.factor" }, + { LLM_KV_DSV41_ROPE_SCALING_BETA_FAST, "deepseek41.rope_scaling.beta_fast" }, + { LLM_KV_DSV41_ROPE_SCALING_BETA_SLOW, "deepseek41.rope_scaling.beta_slow" }, + { LLM_KV_DSV41_ROPE_SCALING_ORIG_CTX_LEN, "deepseek41.rope_scaling.original_max_position_embeddings" }, + { LLM_KV_DSV41_ENGRAM_ENCODING, "deepseek41.engram.encoding" }, + { LLM_KV_DSV41_ENGRAM_LAYER_IDS, "deepseek41.engram.layer_ids" }, + { LLM_KV_DSV41_ENGRAM_ROWS, "deepseek41.engram.rows" }, + { LLM_KV_DSV41_ENGRAM_COMPRESSED_VOCAB_SIZE, "deepseek41.engram.compressed_vocab_size" }, + { LLM_KV_DSV41_ENGRAM_PAD_ID, "deepseek41.engram.pad_id" }, + { LLM_KV_DSV41_ENGRAM_TOKEN_MAP, "deepseek41.engram.token_map" }, + { LLM_KV_DSV41_ENGRAM_PRIMES, "deepseek41.engram.primes" }, + { LLM_KV_DSV41_ENGRAM_MULTIPLIERS, "deepseek41.engram.multipliers" }, + }; + + for (const auto & item : keys) { + check(kv(item.first) == item.second, std::string("metadata key mismatch: ") + item.second); + } + + const LLM_TN tn(LLM_ARCH_DEEPSEEK41); + check(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", 2).str() == "blk.2.attn_kv_a_norm.weight", "V4.1 KV A norm tensor name failed"); + check(tn(LLM_TENSOR_ENGRAM_EMBD, "weight", 1).str() == "blk.1.engram_embd.weight", "Engram embedding tensor name failed"); + check(tn(LLM_TENSOR_ENGRAM_Q_NORM, "weight", 1).str() == "blk.1.engram_q_norm.weight", "Engram query norm tensor name failed"); + check(tn(LLM_TENSOR_ENGRAM_K_NORM, "weight", 1).str() == "blk.1.engram_k_norm.weight", "Engram key norm tensor name failed"); + check(tn(LLM_TENSOR_ENGRAM_KV, "weight", 1).str() == "blk.1.engram_kv.weight", "Engram projection tensor name failed"); + + return 0; +} diff --git a/tests/test-engram.cpp b/tests/test-engram.cpp new file mode 100644 index 000000000000..e741f231262b --- /dev/null +++ b/tests/test-engram.cpp @@ -0,0 +1,444 @@ +#include "../src/llama-bounded-file.h" +#include "../src/llama-engram.h" +#include "../src/llama-ple-disk.h" + +#include +#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_domain(const std::function & fn, const char * message) { + try { + fn(); + } catch (const std::domain_error &) { + 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_engram_layout make_layout() { + llama_engram_layout layout; + layout.encoding = "e4m3_e8m0_32_row264"; + layout.layer_ids = { 1, 14 }; + layout.token_map.resize(256); + for (size_t i = 0; i < layout.token_map.size(); ++i) { + layout.token_map[i] = (uint32_t) i / 2; + } + layout.compressed_vocab_size = 128; + layout.pad_id = 1; + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + for (size_t depth = 0; depth < LLAMA_ENGRAM_NGRAM; ++depth) { + layout.multipliers[layer][depth] = 35184372088831ull - 2 * (depth + 4 * layer); + } + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + layout.primes[layer][col] = 16000057; + layout.rows[layer] += layout.primes[layer][col]; + } + } + return layout; +} + +static void reference_hash( + const llama_engram_layout & layout, + const int32_t * tokens, + const uint8_t * mask, + size_t count, + uint32_t * output) { + for (size_t pos = 0; pos < count; ++pos) { + for (size_t layer = 0; layer < LLAMA_ENGRAM_LAYERS; ++layer) { + uint64_t offset = 0; + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + uint64_t hash = 0; + bool blocked = false; + for (size_t shift = 0; shift < col / LLAMA_ENGRAM_HEADS + 2; ++shift) { + const bool before_start = shift > pos; + const size_t source = before_start ? 0 : pos - shift; + blocked = blocked || before_start || (mask != nullptr && mask[source] == 0); + const uint32_t id = blocked ? layout.pad_id : layout.token_map[tokens[source]]; + hash ^= (uint64_t) id * layout.multipliers[layer][shift]; + } + *output++ = (uint32_t) (hash % layout.primes[layer][col] + offset); + offset += layout.primes[layer][col]; + } + } + } +} + +static void test_layout_validation() { + llama_engram_layout layout = make_layout(); + llama_engram_hasher valid(layout); + check(valid.layout().rows == layout.rows, "valid Engram layout changed"); + + llama_engram_layout bad = layout; + bad.encoding = "e4m3"; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad encoding was accepted"); + bad = layout; + bad.layer_ids[1] = bad.layer_ids[0]; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "duplicate layer IDs were accepted"); + bad = layout; + bad.token_map.clear(); + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "empty token map was accepted"); + bad = layout; + bad.compressed_vocab_size = 0; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "empty compressed vocabulary was accepted"); + bad = layout; + bad.pad_id = bad.compressed_vocab_size; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad pad ID was accepted"); + bad = layout; + bad.token_map[10] = bad.compressed_vocab_size; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad token map entry was accepted"); + bad = layout; + bad.multipliers[0][0]--; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "even multiplier was accepted"); + bad = layout; + bad.multipliers[0][0] = UINT64_MAX; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "overflowing multiplier was accepted"); + bad = layout; + bad.primes[0][0] = 1; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad prime was accepted"); + bad = layout; + bad.rows[0]--; + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "bad row extent was accepted"); + bad = layout; + bad.primes[0].fill(UINT32_MAX); + expect_invalid([&] { llama_engram_hasher hasher(bad); }, "overflowing row extent was accepted"); +} + +static void test_hash() { + const llama_engram_layout layout = make_layout(); + const llama_engram_hasher hasher(layout); + constexpr size_t count = 513; + constexpr size_t width = LLAMA_ENGRAM_LAYERS * LLAMA_ENGRAM_COLS; + std::vector tokens(count); + std::vector mask(count); + std::vector expected(count * width); + std::vector actual(count * width); + for (size_t i = 0; i < count; ++i) { + tokens[i] = (int32_t) ((i * 97 + i / 3) % layout.token_map.size()); + mask[i] = i % 17 != 0 && (i < 125 || i > 131); + } + + for (int masked = 0; masked < 2; ++masked) { + const uint8_t * current_mask = masked != 0 ? mask.data() : nullptr; + reference_hash(layout, tokens.data(), current_mask, count, expected.data()); + llama_engram_history expected_history; + expected_history.reset(); + hasher.hash(expected_history, tokens.data(), current_mask, count, actual.data()); + check(actual == expected, "single-chunk Engram hash differs from full-history reference"); + for (size_t chunk = 1; chunk <= count; ++chunk) { + llama_engram_history history; + history.reset(); + for (size_t pos = 0; pos < count; pos += chunk) { + const size_t size = std::min(chunk, count - pos); + hasher.hash(history, tokens.data() + pos, + current_mask != nullptr ? current_mask + pos : nullptr, + size, actual.data() + pos * width); + } + check(actual == expected, "rolling Engram hash differs from full-history reference"); + check(history.tail == expected_history.tail, "chunking changed final Engram history"); + } + } + + llama_engram_history history; + history.reset(); + const llama_engram_history before = history; + uint32_t output[2 * width]; + std::fill(output, output + 2 * width, UINT32_MAX); + const int32_t invalid_high[] = { 0, (int32_t) layout.token_map.size() }; + expect_invalid([&] { hasher.hash(history, invalid_high, nullptr, 2, output); }, + "high invalid token was accepted"); + check(history.tail == before.tail, "invalid token mutated Engram history"); + check(output[0] == UINT32_MAX, "invalid token changed Engram output"); + const int32_t invalid_low[] = { 0, -1 }; + expect_invalid([&] { hasher.hash(history, invalid_low, nullptr, 2, output); }, + "negative token was accepted"); + check(history.tail == before.tail, "negative token mutated Engram history"); + hasher.hash(history, nullptr, nullptr, 0, nullptr); + expect_invalid([&] { hasher.hash(history, invalid_high, nullptr, SIZE_MAX, output); }, + "overflowing hash count was accepted"); + history.tail[0] = (int32_t) layout.compressed_vocab_size; + expect_invalid([&] { hasher.hash(history, invalid_high, nullptr, 1, output); }, + "invalid history was accepted"); +} + +static float decode_reference(uint8_t code, uint8_t scale) { + const int exponent = (code >> 3) & 15; + double value = exponent != 0 ? + (1.0 + (code & 7) / 8.0) * std::pow(2.0, exponent - 7) : + (code & 7) / 512.0; + if ((code & 128) != 0) { + value = -value; + } + float result = (float) (value * std::pow(2.0, (int) scale - 127)); + uint32_t bits; + memcpy(&bits, &result, sizeof(bits)); + bits = (bits + 0x7fffu + ((bits >> 16) & 1u)) & 0xffff0000u; + memcpy(&result, &bits, sizeof(result)); + return result; +} + +static void test_decode() { + uint8_t row[LLAMA_ENGRAM_ROW_BYTES]; + float output[LLAMA_ENGRAM_DIM]; + for (uint32_t code = 0; code < 256; ++code) { + memset(row, (int) code, LLAMA_ENGRAM_DIM); + for (uint32_t scale = 0; scale < 256; ++scale) { + memset(row + LLAMA_ENGRAM_DIM, (int) scale, LLAMA_ENGRAM_ROW_BYTES - LLAMA_ENGRAM_DIM); + const float expected = decode_reference((uint8_t) code, (uint8_t) scale); + const bool valid = (code & 127) != 127 && scale != 255 && std::isfinite(expected); + if (valid) { + llama_engram_decode_row(row, output); + for (float value : output) { + check(memcmp(&value, &expected, sizeof(value)) == 0, "Engram decode differs from BF16 reference"); + } + } else { + output[0] = 123456.0f; + expect_domain([&] { llama_engram_decode_row(row, output); }, "invalid Engram value was accepted"); + check(output[0] == 123456.0f, "invalid Engram value changed output"); + } + } + } + + memset(row, 0, sizeof(row)); + memset(row + LLAMA_ENGRAM_DIM, 127, LLAMA_ENGRAM_ROW_BYTES - LLAMA_ENGRAM_DIM); + row[0] = 128; + llama_engram_decode_row(row, output); + check(output[0] == 0.0f && std::signbit(output[0]), "negative zero was not preserved"); + check(output[1] == 0.0f && !std::signbit(output[1]), "positive zero was not preserved"); +} + +#if !defined(_WIN32) +static void write_full(int fd, const void * data, size_t size, uint64_t offset) { + const ssize_t written = pwrite(fd, data, size, (off_t) offset); + check(written == (ssize_t) size, "failed to write Engram test row"); +} + +static void test_disk_rows() { + char path[] = "/tmp/llama-engram-XXXXXX"; + const int fd = mkstemp(path); + check(fd >= 0, "failed to create Engram test file"); + + const uint64_t offset = (1ull << 33) + 32; + uint8_t raw[3][LLAMA_ENGRAM_ROW_BYTES]; + for (size_t row = 0; row < 3; ++row) { + for (size_t i = 0; i < LLAMA_ENGRAM_DIM; ++i) { + raw[row][i] = (uint8_t) i; + } + raw[row][127] = 0; + raw[row][255] = 128; + for (size_t i = 0; i < LLAMA_ENGRAM_ROW_BYTES - LLAMA_ENGRAM_DIM; ++i) { + raw[row][LLAMA_ENGRAM_DIM + i] = (uint8_t) (126 + row); + } + } + write_full(fd, raw, sizeof(raw), offset); + + llama_engram_table table(path, offset, 3); + check(table.n_rows() == 3, "Engram table row count changed"); + const uint32_t rows[] = { 2, 0, 2, 1 }; + float output[4 * LLAMA_ENGRAM_DIM]; + table.read(rows, 4, output); + for (size_t row = 0; row < 4; ++row) { + float expected[LLAMA_ENGRAM_DIM]; + llama_engram_decode_row(raw[rows[row]], expected); + check(memcmp(output + row * LLAMA_ENGRAM_DIM, expected, sizeof(expected)) == 0, + "Engram row order was not preserved"); + } + + constexpr size_t tokens = 2051; + constexpr size_t stride = LLAMA_ENGRAM_LAYERS * LLAMA_ENGRAM_COLS; + constexpr size_t width = LLAMA_ENGRAM_COLS * LLAMA_ENGRAM_DIM; + std::vector batch_rows(tokens * stride, UINT32_MAX); + std::vector batch(tokens * width + 1); + for (size_t token = 0; token < tokens; ++token) { + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + batch_rows[token * stride + col] = (uint32_t) ((token * 7 + col * 11) % 3); + } + } + const size_t sizes[] = { 1, 2, 31, 65, 257, 2047, 2048, 2049, tokens }; + float decoded[3][LLAMA_ENGRAM_DIM]; + for (size_t row = 0; row < 3; ++row) { + llama_engram_decode_row(raw[row], decoded[row]); + } + for (size_t count : sizes) { + batch[count * width] = 123456.0f; + table.read_batch(batch_rows.data(), count, stride, batch.data()); + for (size_t token = 0; token < count; ++token) { + for (size_t col = 0; col < LLAMA_ENGRAM_COLS; ++col) { + const uint32_t row = batch_rows[token * stride + col]; + check(memcmp(batch.data() + token * width + col * LLAMA_ENGRAM_DIM, + decoded[row], sizeof(decoded[row])) == 0, + "batched Engram row differs from direct decode"); + } + } + check(batch[count * width] == 123456.0f, "batched Engram read overflowed output"); + } + + table.read(nullptr, 0, nullptr); + table.read_batch(nullptr, 0, 0, nullptr); + expect_invalid([&] { table.read_batch(batch_rows.data(), 1, LLAMA_ENGRAM_COLS - 1, batch.data()); }, + "short Engram batch stride was accepted"); + expect_invalid([&] { table.read_batch(batch_rows.data(), 2, SIZE_MAX, batch.data()); }, + "overflowing Engram batch stride was accepted"); + expect_invalid([&] { table.read_batch(batch_rows.data(), SIZE_MAX, stride, batch.data()); }, + "overflowing Engram batch count was accepted"); + + const uint32_t bad_row = 3; + output[0] = 123456.0f; + expect_invalid([&] { table.read(&bad_row, 1, output); }, "invalid Engram row was accepted"); + check(output[0] == 123456.0f, "invalid Engram row changed output"); + batch_rows[(tokens - 1) * stride] = bad_row; + batch[0] = 123456.0f; + expect_invalid([&] { table.read_batch(batch_rows.data(), tokens, stride, batch.data()); }, + "invalid batched Engram row was accepted"); + check(batch[0] == 123456.0f, "invalid batched Engram row changed output"); + batch_rows[(tokens - 1) * stride] = 0; + + uint8_t invalid = 127; + write_full(fd, &invalid, 1, offset); + const uint32_t first_row = 0; + expect_domain([&] { table.read(&first_row, 1, output); }, "E4M3 NaN row was accepted"); + write_full(fd, raw, sizeof(raw), offset); + invalid = 255; + write_full(fd, &invalid, 1, offset + LLAMA_ENGRAM_DIM); + expect_domain([&] { table.read(&first_row, 1, output); }, "E8M0 scale 255 row was accepted"); + write_full(fd, raw, sizeof(raw), offset); + + check(ftruncate(fd, (off_t) (offset + 260)) == 0, "failed to truncate Engram test file"); + expect_runtime([&] { table.read(&first_row, 1, output); }, "short Engram read was accepted"); + close(fd); + + expect_invalid([&] { llama_engram_table invalid_table(path, offset, 1); }, + "truncated Engram extent was accepted"); + expect_invalid([&] { llama_engram_table invalid_table(path, UINT64_MAX - 1, 3); }, + "overflowing Engram extent was accepted"); + expect_invalid([&] { llama_engram_table invalid_table(path, offset, 0); }, + "empty Engram table was accepted"); + check(unlink(path) == 0, "failed to remove Engram test file"); +} + +#if defined(__linux__) +static void test_direct_tail_read() { + char path[] = "/tmp/llama-direct-tail-XXXXXX"; + const int fd = mkstemp(path); + check(fd >= 0, "failed to create direct tail test file"); + + constexpr uint64_t offset = 4096; + const uint8_t expected[] = { 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2 }; + write_full(fd, expected, sizeof(expected), offset); + close(fd); + + llama_bounded_file::params params; + params.direct_io = true; + params.direct_io_required = true; + llama_bounded_file file(path, params); + check(file.direct_io(), "direct tail test did not use O_DIRECT"); + llama_bounded_file::buffer scratch = file.make_buffer(sizeof(expected)); + uint8_t actual[sizeof(expected)] = {}; + file.read(offset, actual, sizeof(actual), scratch); + check(memcmp(actual, expected, sizeof(actual)) == 0, "valid direct tail read failed"); + check(unlink(path) == 0, "failed to remove direct tail test file"); +} +#endif + +static void test_ple_disk_reader() { + char path[] = "/tmp/llama-ple-reader-XXXXXX"; + const int fd = mkstemp(path); + check(fd >= 0, "failed to create PLE test file"); + + constexpr uint64_t offset = 32; + constexpr size_t columns = 4; + constexpr size_t rows = 3; + const float table[rows][columns] = { + { 1.0f, 2.0f, 3.0f, 4.0f }, + { -1.0f, -2.0f, -3.0f, -4.0f }, + { 0.5f, 0.25f, 0.125f, 0.0625f }, + }; + write_full(fd, table, sizeof(table), offset); + close(fd); + + for (int direct = 0; direct < 2; ++direct) { + llama_ple_disk::params params; + params.n_threads = 2; + params.cache_bytes = sizeof(table); + params.direct_io = direct != 0; + llama_ple_disk disk(path, offset, GGML_TYPE_F32, columns, rows, params); + const int32_t ids[] = { 2, 0, 2, 1 }; + float output[4][columns]; + disk.gather(ids, 4, output[0]); + for (size_t i = 0; i < 4; ++i) { + check(memcmp(output[i], table[ids[i]], sizeof(output[i])) == 0, + "shared bounded reader changed PLE row output"); + } + } + + llama_ple_disk::params params; + params.n_threads = 4; + params.cache_bytes = 0; + params.direct_io = false; + llama_ple_disk disk(path, offset, GGML_TYPE_F32, columns, rows, params); + check(truncate(path, (off_t) (offset + sizeof(float))) == 0, "failed to truncate PLE test file"); + const int32_t ids[] = { 0, 1, 2 }; + float output[3][columns]; + expect_runtime([&] { disk.gather(ids, 3, output[0]); }, "threaded PLE read failure did not propagate"); + + const int repair_fd = open(path, O_WRONLY); + check(repair_fd >= 0, "failed to reopen PLE test file"); + write_full(repair_fd, table, sizeof(table), offset); + close(repair_fd); + disk.gather(ids, 3, output[0]); + check(memcmp(output, table, sizeof(table)) == 0, "PLE reader did not recover after worker failure"); + + check(unlink(path) == 0, "failed to remove PLE test file"); +} +#endif + +int main() { + test_layout_validation(); + test_hash(); + test_decode(); +#if !defined(_WIN32) + test_disk_rows(); +#if defined(__linux__) + test_direct_tail_read(); +#endif + test_ple_disk_reader(); +#endif + std::puts("Engram layout, hash and bounded disk rows: PASS"); + return 0; +} diff --git a/tests/test-expert-store.cpp b/tests/test-expert-store.cpp new file mode 100644 index 000000000000..fbf5f0b10954 --- /dev/null +++ b/tests/test-expert-store.cpp @@ -0,0 +1,651 @@ +#include "../src/llama-expert-store.h" +#include "../src/llama-model-loader.h" + +#include "ggml.h" +#include "gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define REQUIRE(cond) do { if (!(cond)) { throw std::runtime_error("requirement failed: " #cond); } } while (0) + +namespace { + +struct temp_file { + std::filesystem::path path; + + explicit temp_file(const char * suffix) { + static uint64_t sequence = 0; + const auto stamp = std::chrono::steady_clock::now().time_since_epoch().count(); + path = std::filesystem::temp_directory_path() / + ("llama-expert-store-" + std::to_string(stamp) + "-" + std::to_string(++sequence) + suffix); + } + + ~temp_file() { + std::error_code ec; + std::filesystem::remove(path, ec); + } +}; + +struct fixture { + static constexpr int64_t n_embd = 512; + static constexpr int64_t n_ff = 256; + static constexpr int64_t n_expert = 4; + + temp_file file { ".gguf" }; + std::vector tensors; + + fixture() { + const size_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd) * n_ff; + const size_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, n_ff) * n_embd; + const size_t data_size = 2 * gate_plane * n_expert + down_plane * n_expert; + + ggml_init_params ggml_params = { + /*.mem_size =*/ data_size + 8 * ggml_tensor_overhead() + 4096, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(ggml_params); + REQUIRE(ctx != nullptr); + + ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ2_XXS, n_embd, n_ff, n_expert); + ggml_tensor * up = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ2_XXS, n_embd, n_ff, n_expert); + ggml_tensor * down = ggml_new_tensor_3d(ctx, GGML_TYPE_Q2_K, n_ff, n_embd, n_expert); + ggml_set_name(gate, "blk.0.ffn_gate_exps.weight"); + ggml_set_name(up, "blk.0.ffn_up_exps.weight"); + ggml_set_name(down, "blk.0.ffn_down_exps.weight"); + + fill_tensor(gate, 0x10); + fill_tensor(up, 0x20); + fill_tensor(down, 0x30); + + gguf_context * gguf = gguf_init_empty(); + REQUIRE(gguf != nullptr); + gguf_set_val_str(gguf, "general.architecture", "deepseek41"); + gguf_add_tensor(gguf, gate); + gguf_add_tensor(gguf, up); + gguf_add_tensor(gguf, down); + REQUIRE(gguf_write_to_file(gguf, file.path.string().c_str(), false)); + gguf_free(gguf); + ggml_free(ctx); + + std::vector splits; + llama_model_loader loader( + nullptr, nullptr, nullptr, file.path.string(), splits, nullptr, + LLAMA_LOAD_MODE_MMAP, false, false, false, nullptr, nullptr); + + REQUIRE(loader.ctx_map.empty()); + tensors.push_back(loader.register_external_tensor( + "blk.0.ffn_gate_exps.weight", 0, LLAMA_EXPERT_PROJECTION_GATE, { n_embd, n_ff, n_expert })); + tensors.push_back(loader.register_external_tensor( + "blk.0.ffn_up_exps.weight", 0, LLAMA_EXPERT_PROJECTION_UP, { n_embd, n_ff, n_expert })); + tensors.push_back(loader.register_external_tensor( + "blk.0.ffn_down_exps.weight", 0, LLAMA_EXPERT_PROJECTION_DOWN, { n_ff, n_embd, n_expert })); + loader.done_getting_tensors(); + + REQUIRE(loader.external.any()); + REQUIRE(loader.ctx_map.empty()); + for (const auto & tensor : tensors) { + REQUIRE(tensor.file_index == 0); + REQUIRE(loader.external.has(loader.require_tensor_meta(tensor.name))); + } + loader.init_mappings(true); + REQUIRE(loader.mappings.size() == 1); + REQUIRE(loader.ctx_map.empty()); + const auto & ranges = loader.external.for_file(0); + REQUIRE(ranges.size() == 3); + for (size_t i = 0; i < ranges.size(); ++i) { + REQUIRE(ranges[i].first == tensors[i].file_offset); + REQUIRE(ranges[i].second == tensors[i].file_offset + tensors[i].nb[2] * tensors[i].ne[2]); + } + } + + static void fill_tensor(ggml_tensor * tensor, uint8_t tag) { + memset(tensor->data, 0, ggml_nbytes(tensor)); + for (int64_t expert = 0; expert < tensor->ne[2]; ++expert) { + uint8_t * plane = static_cast(tensor->data) + expert * tensor->nb[2]; + plane[tensor->nb[2] - 1] = tag + expert; + } + } + + llama_expert_store make_store(size_t slots, size_t bytes) const { + llama_expert_store_params params; + params.cache_slots = slots; + params.cache_bytes = bytes; + params.direct_io = false; + return llama_expert_store(tensors, params); + } + + size_t max_plane_size() const { + size_t result = 0; + for (const auto & tensor : tensors) { + result = std::max(result, tensor.nb[2]); + } + return result; + } + + size_t all_projection_bytes() const { + size_t result = 0; + for (const auto & tensor : tensors) { + result += tensor.nb[2]; + } + return result; + } +}; + +template +void require_throws(F && fn) { + bool threw = false; + try { + fn(); + } catch (const std::exception &) { + threw = true; + } + REQUIRE(threw); +} + +const llama_expert_store::payload & find_payload( + const std::vector & payloads, + llama_expert_projection projection, + int32_t expert_id) { + for (const auto & payload : payloads) { + if (payload.projection == projection && payload.expert_id == expert_id) { + return payload; + } + } + throw std::runtime_error("payload not found"); +} + +void test_layout_and_offsets(const fixture & f) { + const auto & gate = f.tensors[0]; + const auto & up = f.tensors[1]; + const auto & down = f.tensors[2]; + + REQUIRE(gate.type == GGML_TYPE_IQ2_XXS); + REQUIRE(up.type == GGML_TYPE_IQ2_XXS); + REQUIRE(down.type == GGML_TYPE_Q2_K); + REQUIRE(gate.nb[2] == ggml_row_size(GGML_TYPE_IQ2_XXS, fixture::n_embd) * fixture::n_ff); + REQUIRE(down.nb[2] == ggml_row_size(GGML_TYPE_Q2_K, fixture::n_ff) * fixture::n_embd); + REQUIRE(gate.file_offset + 3 * gate.nb[2] > gate.file_offset); + + llama_expert_store store = f.make_store(3, f.all_projection_bytes()); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.resident_bytes() == 0); + auto lease = store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 2 } }, + { 0, LLAMA_EXPERT_PROJECTION_UP, { 2 } }, + { 0, LLAMA_EXPERT_PROJECTION_DOWN, { 2 } }, + }); + const auto payloads = lease.payloads(); + REQUIRE(payloads.size() == 3); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_GATE, 2).data[gate.nb[2] - 1] == 0x12); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_UP, 2).data[up.nb[2] - 1] == 0x22); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_DOWN, 2).data[down.nb[2] - 1] == 0x32); +} + +void test_alignment_and_large_offsets() { + const auto aligned = llama_expert_store_align_read(4097, 5000, 4096, 20000); + REQUIRE(aligned.offset == 4096); + REQUIRE(aligned.prefix == 1); + REQUIRE(aligned.size == 8192); + + const uint64_t large_offset = (uint64_t(1) << 32) + 123; + const auto large = llama_expert_store_align_read(large_offset, 777, 4096, large_offset + 777); + REQUIRE(large.offset > std::numeric_limits::max()); + REQUIRE(large.prefix == 123); + REQUIRE(large.size == 4096); + + require_throws([] { + llama_expert_store_align_read(0, 1, 3000, 1); + }); + require_throws([] { + llama_expert_store_align_read(UINT64_MAX - 4, 8, 4096, UINT64_MAX); + }); + + size_t granularity_first = 1; + size_t granularity_last = 1; + llama_mlock::align_range(&granularity_first, &granularity_last); + const size_t lock_granularity = granularity_last; + REQUIRE(granularity_first == 0); + REQUIRE(lock_granularity > 1); + + size_t lock_first = lock_granularity + 1; + size_t lock_last = 2 * lock_granularity; + llama_mlock::align_range(&lock_first, &lock_last); + REQUIRE(lock_first == lock_granularity); + REQUIRE(lock_last == 2 * lock_granularity); +} + +void test_external_mapping_access_policy() { + const size_t page = 4096; + const llama_mmap::ranges external = { + { 0, page }, + { 2 * page + 1, 4 * page - 1 }, + { 5 * page, 6 * page }, + { 7 * page, 8 * page }, + { 9 * page, 10 * page }, + }; + + REQUIRE(llama_mmap::use_sequential_file_advice(false)); + REQUIRE(!llama_mmap::use_sequential_file_advice(true)); + REQUIRE(llama_mmap::planned_prefetch_ranges(10 * page, 10 * page, external, true).empty()); + + const auto lazy_ranges = llama_mmap::planned_prefetch_ranges(10 * page, 10 * page, external, false); + REQUIRE(lazy_ranges.size() == 4); + REQUIRE(lazy_ranges.front() == std::make_pair(page, 2 * page + 1)); + REQUIRE(lazy_ranges.back() == std::make_pair(8 * page, 9 * page)); +} + +#if defined(__linux__) +int file_advice_calls = 0; + +int fail_file_advice(int, int) { + ++file_advice_calls; + return EIO; +} + +void test_external_mapping_advice_failure() { + temp_file file { ".bin" }; + { + std::ofstream out(file.path, std::ios::binary); + REQUIRE(out.good()); + out.seekp(8191); + out.put('\0'); + } + + llama_file input(file.path.string(), "rb"); + file_advice_calls = 0; + bool continued_after_advice = false; + require_throws([&] { + llama_mmap mapping(&input, 0, false, { { 4096, 8192 } }, true, fail_file_advice); + continued_after_advice = true; + }); + REQUIRE(file_advice_calls == 1); + REQUIRE(!continued_after_advice); + + llama_mmap legacy_mapping(&input, 0, false, {}, false, fail_file_advice); + REQUIRE(file_advice_calls == 2); + REQUIRE(legacy_mapping.addr() != nullptr); +} +#endif + +void test_published_layout_accounting() { + const int64_t n_embd = 7680; + const int64_t n_ff = 1536; + const int64_t n_expert = 384; + const int64_t n_layer = 40; + + const uint64_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd) * n_ff; + const uint64_t up_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd) * n_ff; + const uint64_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, n_ff) * n_embd; + const uint64_t slot_bytes = (gate_plane + up_plane + down_plane) * n_layer; + const uint64_t routed_bytes = slot_bytes * n_expert; + const uint64_t dense_bytes = 10067427328; + const uint64_t engram_bytes = 202758045696; + + REQUIRE(gate_plane == 3041280); + REQUIRE(up_plane == 3041280); + REQUIRE(down_plane == 3870720); + REQUIRE(gate_plane + up_plane + down_plane == 9953280); + REQUIRE(slot_bytes == 398131200); + REQUIRE(routed_bytes == 152882380800); + REQUIRE(224 * slot_bytes == 89181388800); + REQUIRE(256 * slot_bytes == 101921587200); + REQUIRE(dense_bytes + 224 * slot_bytes == 99248816128); + REQUIRE(dense_bytes + 256 * slot_bytes == 111989014528); + REQUIRE(engram_bytes > routed_bytes); +} + +void test_large_offset_read() { + temp_file sparse { ".bin" }; + const int64_t ne0 = 256; + const int64_t ne1 = 256; + const int64_t n_expert = 1; + const size_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, ne0) * ne1; + const size_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, ne0) * ne1; + const uint64_t base = (uint64_t(1) << 32) + 4096; + const uint64_t file_size = base + 2 * gate_plane + down_plane; + + { + std::ofstream out(sparse.path, std::ios::binary | std::ios::trunc); + REQUIRE(out.good()); + out.seekp(static_cast(file_size - 1)); + out.put('\0'); + const std::vector> markers = { + { base + gate_plane - 1, 0x41 }, + { base + 2 * gate_plane - 1, 0x42 }, + { file_size - 1, 0x43 }, + }; + for (const auto & marker : markers) { + out.seekp(static_cast(marker.first)); + out.put(static_cast(marker.second)); + } + } + + auto make_tensor = [&](const char * name, llama_expert_projection projection, ggml_type type, uint64_t offset) { + llama_expert_store_tensor tensor; + tensor.name = name; + tensor.fname = sparse.path.string(); + tensor.layer = 0; + tensor.projection = projection; + tensor.type = type; + tensor.ne[0] = ne0; + tensor.ne[1] = ne1; + tensor.ne[2] = n_expert; + tensor.nb[0] = ggml_type_size(type); + tensor.nb[1] = ggml_row_size(type, ne0); + tensor.nb[2] = tensor.nb[1] * ne1; + tensor.file_offset = offset; + tensor.file_size = file_size; + return tensor; + }; + + std::vector tensors; + tensors.push_back(make_tensor("gate", LLAMA_EXPERT_PROJECTION_GATE, GGML_TYPE_IQ2_XXS, base)); + tensors.push_back(make_tensor("up", LLAMA_EXPERT_PROJECTION_UP, GGML_TYPE_IQ2_XXS, base + gate_plane)); + tensors.push_back(make_tensor("down", LLAMA_EXPERT_PROJECTION_DOWN, GGML_TYPE_Q2_K, base + 2 * gate_plane)); + + llama_expert_store_params params { 2 * gate_plane + down_plane, 3, 4096, false }; + llama_expert_store store(std::move(tensors), params); + auto lease = store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_UP, { 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_DOWN, { 0 } }, + }); + const auto payloads = lease.payloads(); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_GATE, 0).data[gate_plane - 1] == 0x41); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_UP, 0).data[gate_plane - 1] == 0x42); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_DOWN, 0).data[down_plane - 1] == 0x43); +} + +void test_cache_and_remapping(const fixture & f) { + const size_t gate_plane = f.tensors[0].nb[2]; + llama_expert_store store = f.make_store(2, 2 * f.max_plane_size()); + + { + auto lease = store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0, 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1, 0 } }, + }); + REQUIRE(lease.slot_ids().size() == 2); + REQUIRE(lease.slot_ids()[0][0] == lease.slot_ids()[0][1]); + REQUIRE(lease.slot_ids()[1][1] == lease.slot_ids()[0][0]); + REQUIRE(lease.slot_ids()[1][0] != lease.slot_ids()[1][1]); + REQUIRE(store.resident_entries() == 2); + REQUIRE(store.resident_bytes() == 2 * gate_plane); + } + + { + auto hit = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(hit.slot_ids()[0][0] == 0); + } + { + auto miss = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 2 } } }); + REQUIRE(miss.slot_ids()[0][0] == 1); + } + + const llama_expert_store_stats stats = store.stats(); + REQUIRE(stats.hits == 1); + REQUIRE(stats.misses == 3); + REQUIRE(stats.evictions == 1); + REQUIRE(stats.bytes_read == 3 * gate_plane); + + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { -1 } } }); + }); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { fixture::n_expert } } }); + }); +} + +void test_direct_io(const fixture & f) { + llama_expert_store_params params; + params.cache_slots = 1; + params.cache_bytes = f.max_plane_size(); + params.io_alignment = 4096; + params.direct_io = true; + params.allow_buffered_io = true; + + llama_expert_store store(f.tensors, params); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 3 } } }); + const auto payloads = lease.payloads(); + REQUIRE(payloads.size() == 1); + REQUIRE(payloads[0].data[payloads[0].size - 1] == 0x13); + REQUIRE(store.stats().bytes_read >= payloads[0].size); + REQUIRE(store.stats().bytes_read <= payloads[0].size + 2 * params.io_alignment); +} + +#if defined(__linux__) +void test_direct_io_file_tail(const fixture & f) { + const auto & down = f.tensors[2]; + REQUIRE(down.file_offset + down.nb[2] * down.ne[2] == down.file_size); + + llama_expert_store_params params; + params.cache_slots = 1; + params.cache_bytes = f.max_plane_size(); + params.io_alignment = 4096; + params.direct_io = true; + params.allow_buffered_io = false; + + llama_expert_store store(f.tensors, params); + REQUIRE(store.direct_io_active()); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_DOWN, { fixture::n_expert - 1 } } }); + const auto payloads = lease.payloads(); + REQUIRE(payloads.size() == 1); + REQUIRE(payloads[0].data[payloads[0].size - 1] == 0x33); + REQUIRE(store.direct_io_active()); +} +#endif + +#if defined(_WIN32) +void test_windows_direct_io_policy(const fixture & f) { + llama_expert_store_params params; + params.cache_slots = 1; + params.cache_bytes = f.max_plane_size(); + params.direct_io = true; + params.allow_buffered_io = false; + + require_throws([&] { + llama_expert_store store(f.tensors, params); + }); + + params.allow_buffered_io = true; + llama_expert_store store(f.tensors, params); + REQUIRE(!store.direct_io_active()); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(lease.payloads().size() == 1); +} +#endif + +void test_pins_and_atomic_failure(const fixture & f) { + llama_expert_store store = f.make_store(1, f.max_plane_size()); + auto pinned = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + const llama_expert_store_stats before = store.stats(); + + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1 } } }); + }); + REQUIRE(store.resident_entries() == 1); + REQUIRE(store.stats().hits == before.hits); + REQUIRE(store.stats().misses == before.misses); + REQUIRE(pinned.payloads()[0].expert_id == 0); + + pinned = {}; + auto replacement = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1 } } }); + REQUIRE(replacement.payloads()[0].expert_id == 1); + REQUIRE(store.stats().evictions == 1); + + llama_expert_store::lease surviving; + { + llama_expert_store short_lived = f.make_store(1, f.max_plane_size()); + surviving = short_lived.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 2 } } }); + } + REQUIRE(surviving.payloads()[0].expert_id == 2); +} + +void test_limits_and_validation(const fixture & f) { + require_throws([&] { + f.make_store(0, f.max_plane_size()); + }); + require_throws([&] { + f.make_store(1, f.tensors[2].nb[2] - 1); + }); + { + llama_expert_store_params params { f.max_plane_size(), 1, 1, false }; + llama_expert_store store(f.tensors, params); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(lease.payloads().size() == 1); + } + { + llama_expert_store store = f.make_store(3, f.max_plane_size()); + require_throws([&] { + store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_UP, { 0 } }, + }); + }); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.stats().misses == 0); + } + { + llama_expert_store store = f.make_store(1, 2 * f.tensors[0].nb[2]); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0, 1 } } }); + }); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.stats().misses == 0); + } + + auto bad_type = f.tensors; + bad_type[0].type = GGML_TYPE_Q2_K; + require_throws([&] { + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(bad_type), params); + }); + + auto bad_stride = f.tensors; + bad_stride[1].nb[2]++; + require_throws([&] { + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(bad_stride), params); + }); + + auto bad_bounds = f.tensors; + bad_bounds[2].file_size = bad_bounds[2].file_offset + bad_bounds[2].nb[2] - 1; + require_throws([&] { + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(bad_bounds), params); + }); +} + +void test_payload_validation(const fixture & f) { + temp_file copy { ".gguf" }; + std::filesystem::copy_file(f.file.path, copy.path); + auto tensors = f.tensors; + for (auto & tensor : tensors) { + tensor.fname = copy.path.string(); + tensor.file_size = std::filesystem::file_size(copy.path); + } + + { + std::fstream io(copy.path, std::ios::binary | std::ios::in | std::ios::out); + REQUIRE(io.good()); + io.seekp(static_cast(tensors[0].file_offset)); + const uint8_t invalid_scale[2] = { 0x00, 0x7c }; + io.write(reinterpret_cast(invalid_scale), sizeof(invalid_scale)); + } + + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(tensors), params); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + }); + REQUIRE(store.resident_entries() == 0); +} + +void test_truncated_file(const fixture & f) { + temp_file copy { ".gguf" }; + std::filesystem::copy_file(f.file.path, copy.path); + auto tensors = f.tensors; + for (auto & tensor : tensors) { + tensor.fname = copy.path.string(); + tensor.file_size = std::filesystem::file_size(copy.path); + } + + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(tensors), params); + std::filesystem::resize_file(copy.path, f.tensors[0].file_offset + f.tensors[0].nb[2] - 1); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + }); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.stats().misses == 0); +} + +void test_failed_replacement_keeps_resident_entry(const fixture & f) { + temp_file copy { ".gguf" }; + std::filesystem::copy_file(f.file.path, copy.path); + auto tensors = f.tensors; + for (auto & tensor : tensors) { + tensor.fname = copy.path.string(); + tensor.file_size = std::filesystem::file_size(copy.path); + } + + llama_expert_store_params params { f.max_plane_size(), 1, 4096, false }; + llama_expert_store store(std::move(tensors), params); + { + auto resident = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(resident.payloads()[0].expert_id == 0); + } + const llama_expert_store_stats before = store.stats(); + std::filesystem::resize_file(copy.path, f.tensors[0].file_offset + 2*f.tensors[0].nb[2] - 1); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1 } } }); + }); + REQUIRE(store.resident_entries() == 1); + REQUIRE(store.stats().evictions == before.evictions); + auto hit = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(hit.payloads()[0].expert_id == 0); +} + +} + +int main() { + try { + fixture f; + test_layout_and_offsets(f); + test_alignment_and_large_offsets(); + test_external_mapping_access_policy(); +#if defined(__linux__) + test_external_mapping_advice_failure(); +#endif + test_published_layout_accounting(); + test_large_offset_read(); + test_cache_and_remapping(f); + test_direct_io(f); +#if defined(__linux__) + test_direct_io_file_tail(f); +#endif +#if defined(_WIN32) + test_windows_direct_io_policy(f); +#endif + test_pins_and_atomic_failure(f); + test_limits_and_validation(f); + test_payload_validation(f); + test_truncated_file(f); + test_failed_replacement_keeps_resident_entry(f); + } catch (const std::exception & e) { + fprintf(stderr, "test-expert-store: %s\n", e.what()); + return 1; + } + return 0; +} diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 90a6a7162331..71fced38c1dc 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -9,6 +9,8 @@ // TODO: replace with #include "llama-ext.h" in the future #include "../src/llama-arch.h" +#include "../src/llama-dsv41.h" +#include "../src/llama-dsv41-expert.h" #include "../src/llama-model-saver.h" #include @@ -16,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -324,6 +327,107 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); } + if (arch == LLM_ARCH_DEEPSEEK41) { + std::vector compress_ratios; + compress_ratios.reserve(LLAMA_DSV41_N_LAYER); + for (uint32_t il = 0; il < LLAMA_DSV41_N_LAYER; ++il) { + compress_ratios.push_back(llama_dsv41_compress_ratio(il)); + } + std::vector token_map(LLAMA_DSV41_N_VOCAB, 0); + std::vector primes(LLAMA_DSV41_ENGRAM_PRIMES_COUNT, 2); + primes[0] = 384006168 - 2*(LLAMA_ENGRAM_COLS - 1); + primes[LLAMA_ENGRAM_COLS] = 384016682 - 2*(LLAMA_ENGRAM_COLS - 1); + const std::vector multipliers = { + 101, 103, 105, 107, + 109, 111, 113, 115, + }; + + ms.add_kv(LLM_KV_DSV41_CONFIG, "{}"); + ms.add_kv(LLM_KV_DSV41_MAX_POSITION_EMBEDDINGS, LLAMA_DSV41_N_CTX); + ms.add_kv(LLM_KV_DSV41_HIDDEN_SIZE, LLAMA_DSV41_N_EMBD); + ms.add_kv(LLM_KV_DSV41_NUM_HIDDEN_LAYERS, LLAMA_DSV41_N_LAYER); + ms.add_kv(LLM_KV_DSV41_VOCAB_SIZE, LLAMA_DSV41_N_VOCAB); + ms.add_kv(LLM_KV_DSV41_NUM_ATTENTION_HEADS, LLAMA_DSV41_N_HEAD); + ms.add_kv(LLM_KV_DSV41_NUM_KEY_VALUE_HEADS, LLAMA_DSV41_N_HEAD_KV); + ms.add_kv(LLM_KV_DSV41_HEAD_DIM, LLAMA_DSV41_N_HEAD_DIM); + ms.add_kv(LLM_KV_DSV41_QK_ROPE_HEAD_DIM, LLAMA_DSV41_N_ROT); + ms.add_kv(LLM_KV_DSV41_Q_LORA_RANK, LLAMA_DSV41_N_LORA_Q); + ms.add_kv(LLM_KV_DSV41_O_LORA_RANK, LLAMA_DSV41_N_LORA_O); + ms.add_kv(LLM_KV_DSV41_O_GROUPS, LLAMA_DSV41_N_O_GROUP); + ms.add_kv(LLM_KV_DSV41_MOE_INTERMEDIATE_SIZE, LLAMA_DSV41_N_FF_EXP); + ms.add_kv(LLM_KV_DSV41_N_ROUTED_EXPERTS, LLAMA_DSV41_N_EXPERT); + ms.add_kv(LLM_KV_DSV41_NUM_EXPERTS_PER_TOK, LLAMA_DSV41_N_EXPERT_USED); + ms.add_kv(LLM_KV_DSV41_N_SHARED_EXPERTS, LLAMA_DSV41_N_EXPERT_SHARED); + ms.add_kv(LLM_KV_DSV41_INDEX_N_HEADS, LLAMA_DSV41_N_INDEX_HEAD); + ms.add_kv(LLM_KV_DSV41_INDEX_HEAD_DIM, LLAMA_DSV41_N_INDEX_HEAD_DIM); + ms.add_kv(LLM_KV_DSV41_INDEX_TOPK, LLAMA_DSV41_N_INDEX_TOP_K); + ms.add_kv(LLM_KV_DSV41_HC_MULT, LLAMA_DSV41_HC_MULT); + ms.add_kv(LLM_KV_DSV41_HC_SINKHORN_ITERS, LLAMA_DSV41_HC_SINKHORN_ITERS); + ms.add_kv(LLM_KV_DSV41_SLIDING_WINDOW, LLAMA_DSV41_N_SWA); + ms.add_kv(LLM_KV_DSV41_CANDIDATE_SOURCE_LAYER_ID, LLAMA_DSV41_CANDIDATE_SOURCE_LAYER); + ms.add_kv(LLM_KV_DSV41_CANDIDATE_TOPK_BLOCKS, LLAMA_DSV41_CANDIDATE_TOPK_BLOCKS); + ms.add_kv(LLM_KV_DSV41_CANDIDATE_BLOCK_SIZE, LLAMA_DSV41_CANDIDATE_BLOCK_SIZE); + ms.add_kv(LLM_KV_DSV41_RMS_NORM_EPS, 1.0e-20f); + ms.add_kv(LLM_KV_DSV41_HC_EPS, 1.0e-6f); + ms.add_kv(LLM_KV_DSV41_SWIGLU_LIMIT, 10.0f); + ms.add_kv(LLM_KV_DSV41_ROUTED_SCALING_FACTOR, 1.5f); + ms.add_kv(LLM_KV_DSV41_ROPE_THETA, uint32_t(10000)); + ms.add_kv(LLM_KV_DSV41_COMPRESS_ROPE_THETA, uint32_t(160000)); + ms.add_kv(LLM_KV_DSV41_ROPE_SCALING_FACTOR, 16.0f); + ms.add_kv(LLM_KV_DSV41_ROPE_SCALING_BETA_FAST, 32.0f); + ms.add_kv(LLM_KV_DSV41_ROPE_SCALING_BETA_SLOW, 1.0f); + ms.add_kv(LLM_KV_DSV41_ROPE_SCALING_ORIG_CTX_LEN, 65536.0f); + ms.add_kv(LLM_KV_DSV41_NORM_TOPK_PROB, true); + ms.add_kv(LLM_KV_DSV41_HIDDEN_ACT, "silu"); + ms.add_kv(LLM_KV_DSV41_SCORING_FUNC, "sqrtsoftplus"); + ms.add_kv(LLM_KV_DSV41_TOPK_METHOD, "noaux_tc"); + ms.add_kv(LLM_KV_DSV41_COMPRESS_RATIOS, compress_ratios); + ms.add_kv(LLM_KV_DSV41_KV_SOURCE_LAYER_IDS, std::vector({ 2, 8, 14, 20 })); + ms.add_kv(LLM_KV_DSV41_INDEX_SOURCE_LAYER_IDS, std::vector({ 2, 8, 14, 20, 24, 28, 32, 36 })); + ms.add_kv(LLM_KV_DSV41_ENGRAM_ENCODING, LLAMA_DSV41_ENGRAM_ENCODING); + ms.add_kv(LLM_KV_DSV41_ENGRAM_LAYER_IDS, std::vector({ 1, 14 })); + ms.add_kv(LLM_KV_DSV41_ENGRAM_ROWS, std::vector({ 384006168, 384016682 })); + ms.add_kv(LLM_KV_DSV41_ENGRAM_COMPRESSED_VOCAB_SIZE, LLAMA_DSV41_ENGRAM_COMPRESSED_VOCAB); + ms.add_kv(LLM_KV_DSV41_ENGRAM_PAD_ID, LLAMA_DSV41_ENGRAM_PAD_ID); + ms.add_kv(LLM_KV_DSV41_ENGRAM_TOKEN_MAP, token_map); + ms.add_kv(LLM_KV_DSV41_ENGRAM_PRIMES, primes); + ms.add_kv(LLM_KV_DSV41_ENGRAM_MULTIPLIERS, multipliers); + + auto add_tensor = [&](const std::string & name, ggml_type type, const std::initializer_list & ne) { + ggml_tensor tensor = {}; + tensor.type = type; + tensor.ne[0] = tensor.ne[1] = tensor.ne[2] = tensor.ne[3] = 1; + size_t i = 0; + for (int64_t dim : ne) { + tensor.ne[i++] = dim; + } + tensor.nb[0] = ggml_type_size(type); + tensor.nb[1] = ggml_row_size(type, tensor.ne[0]); + tensor.nb[2] = tensor.nb[1]*tensor.ne[1]; + tensor.nb[3] = tensor.nb[2]*tensor.ne[2]; + snprintf(tensor.name, sizeof(tensor.name), "%s", name.c_str()); + gguf_add_tensor(ms.gguf_ctx, &tensor); + }; + llama_dsv41_register_expert_tensors( + [&](const std::string & name, + int32_t, + llama_expert_projection projection, + const std::initializer_list & ne) { + add_tensor( + name, + projection == LLAMA_EXPERT_PROJECTION_DOWN ? + GGML_TYPE_Q2_K : GGML_TYPE_IQ2_XXS, + ne); + return llama_expert_store_tensor(); + }); + const LLM_TN tn(LLM_ARCH_DEEPSEEK41); + for (int32_t il : { 1, 14 }) { + add_tensor( + tn(LLM_TENSOR_ENGRAM_EMBD, "weight", il).str(), + GGML_TYPE_I8, + { LLAMA_ENGRAM_ROW_BYTES, il == 1 ? 384006168 : 384016682 }); + } + } if (arch == LLM_ARCH_MAPLE) { ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 7.0f); @@ -374,18 +478,20 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, 25.0f); ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, -5.0f); - for (uint32_t il = 0; il < n_layer; il++) { - ggml_tensor t; - memset(&t, 0, sizeof(ggml_tensor)); - t.type = GGML_TYPE_F16; - ggml_format_name(&t, "conv%" PRIu32 "d.weight", il); - gguf_add_tensor(ms.gguf_ctx, &t); - ggml_format_name(&t, "posnet.%" PRIu32 ".conv1.weight", il); - gguf_add_tensor(ms.gguf_ctx, &t); - ggml_format_name(&t, "posnet.%" PRIu32 ".conv2.weight", il); - gguf_add_tensor(ms.gguf_ctx, &t); - ggml_format_name(&t, "convnext.%" PRIu32 ".dw.weight", il); - gguf_add_tensor(ms.gguf_ctx, &t); + if (arch != LLM_ARCH_DEEPSEEK41) { + for (uint32_t il = 0; il < n_layer; il++) { + ggml_tensor t; + memset(&t, 0, sizeof(ggml_tensor)); + t.type = GGML_TYPE_F16; + ggml_format_name(&t, "conv%" PRIu32 "d.weight", il); + gguf_add_tensor(ms.gguf_ctx, &t); + ggml_format_name(&t, "posnet.%" PRIu32 ".conv1.weight", il); + gguf_add_tensor(ms.gguf_ctx, &t); + ggml_format_name(&t, "posnet.%" PRIu32 ".conv2.weight", il); + gguf_add_tensor(ms.gguf_ctx, &t); + ggml_format_name(&t, "convnext.%" PRIu32 ".dw.weight", il); + gguf_add_tensor(ms.gguf_ctx, &t); + } } return ret; } @@ -396,7 +502,7 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, - const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { + const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, bool no_alloc = false) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -404,9 +510,15 @@ static std::pair get_model_and_ctx( devs_copy.push_back(nullptr); model_params.devices = devs_copy.data(); model_params.split_mode = split_mode; + model_params.no_alloc = no_alloc; + if (no_alloc) { + model_params.n_gpu_layers = 0; + model_params.expert_cache_bytes = std::numeric_limits::max(); + model_params.expert_cache_slots = LLAMA_DSV41_N_EXPERT_USED; + } llama_context_params ctx_params = llama_context_default_params(); - ctx_params.n_ctx = 0; + ctx_params.n_ctx = no_alloc ? 256 : 0; ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; if (!encode) { @@ -427,6 +539,53 @@ static std::pair get_model_and_ctx( return std::make_pair(std::move(model), std::move(lctx)); } +static int test_deepseek41_no_alloc(const size_t seed) { + gguf_context_ptr malformed = get_gguf_ctx(LLM_ARCH_DEEPSEEK41, true); + ggml_tensor malformed_token_embd = {}; + malformed_token_embd.type = GGML_TYPE_F32; + malformed_token_embd.ne[0] = malformed_token_embd.ne[1] = + malformed_token_embd.ne[2] = malformed_token_embd.ne[3] = 1; + malformed_token_embd.nb[0] = ggml_type_size(malformed_token_embd.type); + malformed_token_embd.nb[1] = ggml_row_size( + malformed_token_embd.type, malformed_token_embd.ne[0]); + malformed_token_embd.nb[2] = + malformed_token_embd.nb[1]*malformed_token_embd.ne[1]; + malformed_token_embd.nb[3] = + malformed_token_embd.nb[2]*malformed_token_embd.ne[2]; + ggml_set_name( + &malformed_token_embd, + LLM_TN(LLM_ARCH_DEEPSEEK41)( + LLM_TENSOR_TOKEN_EMBD, "weight").str().c_str()); + gguf_add_tensor(malformed.get(), &malformed_token_embd); + bool rejected = false; + try { + auto unused = get_model_and_ctx( + malformed.get(), nullptr, seed, {}, + LLAMA_SPLIT_MODE_LAYER, false, true); + GGML_UNUSED(unused); + } catch (const std::runtime_error &) { + rejected = true; + } + if (!rejected) { + throw std::runtime_error( + "DeepSeek V4.1 no-allocation loader accepted a malformed tensor"); + } + + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_DEEPSEEK41, true); + if (gguf_find_tensor(gguf_ctx.get(), "blk.0.ffn_gate_exps.weight") < 0 || + gguf_find_tensor(gguf_ctx.get(), "blk.1.engram_embd.weight") < 0) { + throw std::runtime_error("DeepSeek V4.1 no-allocation fixture is incomplete"); + } + auto model_and_ctx = get_model_and_ctx( + gguf_ctx.get(), nullptr, seed, {}, + LLAMA_SPLIT_MODE_LAYER, false, true); + if (llama_n_ctx(model_and_ctx.second.get()) != 256) { + throw std::runtime_error("DeepSeek V4.1 no-allocation context size mismatch"); + } + printf("deepseek41: no-allocation architecture graph PASS\n"); + return 0; +} + static std::vector get_logits( llama_model * model, llama_context * lctx, const std::vector & tokens, bool encode = false) { const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); @@ -481,6 +640,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_DEEPSEEK32: case LLM_ARCH_DOTS3NOTE: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: case LLM_ARCH_EXAONE_MOE: @@ -568,6 +728,9 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_DEEPSEEK2OCR) { return false; } + if (arch == LLM_ARCH_DEEPSEEK41) { + 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 if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_QWEN4EXP || @@ -878,6 +1041,9 @@ int main(int argc, char ** argv) { if (!out.empty()) { return save_models(arch, seed, verbosity, out); } + if (arch == LLM_ARCH_DEEPSEEK41) { + return test_deepseek41_no_alloc(seed); + } return test_backends(arch, seed, verbosity); } catch (const std::exception & err) { fprintf(stderr, "encountered runtime error: %s\n", err.what());