From 289ecb05c9423e25d5e856c2fc68b94c9bd03d40 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 10 Jul 2026 10:41:19 +0300 Subject: [PATCH] talk-llama : sync llama.cpp --- examples/talk-llama/CMakeLists.txt | 1 + examples/talk-llama/llama-arch.cpp | 57 + examples/talk-llama/llama-arch.h | 34 + examples/talk-llama/llama-batch.cpp | 76 +- examples/talk-llama/llama-batch.h | 3 +- examples/talk-llama/llama-context.cpp | 223 +- examples/talk-llama/llama-context.h | 4 + examples/talk-llama/llama-graph.cpp | 463 ++++- examples/talk-llama/llama-graph.h | 100 +- examples/talk-llama/llama-hparams.h | 11 + examples/talk-llama/llama-impl.h | 24 +- examples/talk-llama/llama-kv-cache-dsa.cpp | 2 +- examples/talk-llama/llama-kv-cache-dsv4.cpp | 1807 +++++++++++++++++ examples/talk-llama/llama-kv-cache-dsv4.h | 362 ++++ examples/talk-llama/llama-kv-cache-iswa.cpp | 25 +- examples/talk-llama/llama-kv-cache-iswa.h | 20 +- examples/talk-llama/llama-kv-cache.cpp | 54 +- examples/talk-llama/llama-kv-cache.h | 3 + .../talk-llama/llama-memory-hybrid-iswa.cpp | 18 +- examples/talk-llama/llama-memory-hybrid.cpp | 18 +- .../talk-llama/llama-memory-recurrent.cpp | 15 +- examples/talk-llama/llama-model-loader.cpp | 95 +- examples/talk-llama/llama-model.cpp | 57 +- examples/talk-llama/llama-model.h | 27 + examples/talk-llama/llama-quant.cpp | 4 +- examples/talk-llama/llama-vocab.cpp | 26 +- examples/talk-llama/llama.h | 7 + examples/talk-llama/models/deepseek4.cpp | 1165 +++++++++++ examples/talk-llama/models/delta-net-base.cpp | 12 +- examples/talk-llama/models/dflash.cpp | 276 +++ examples/talk-llama/models/models.h | 131 ++ examples/talk-llama/models/qwen3next.cpp | 2 + 32 files changed, 4809 insertions(+), 313 deletions(-) create mode 100644 examples/talk-llama/llama-kv-cache-dsv4.cpp create mode 100644 examples/talk-llama/llama-kv-cache-dsv4.h create mode 100644 examples/talk-llama/models/deepseek4.cpp create mode 100644 examples/talk-llama/models/dflash.cpp diff --git a/examples/talk-llama/CMakeLists.txt b/examples/talk-llama/CMakeLists.txt index 13b284ed0..59643c390 100644 --- a/examples/talk-llama/CMakeLists.txt +++ b/examples/talk-llama/CMakeLists.txt @@ -21,6 +21,7 @@ if (WHISPER_SDL2) llama-kv-cache.cpp llama-kv-cache-iswa.cpp llama-kv-cache-dsa.cpp + llama-kv-cache-dsv4.cpp llama-memory-recurrent.cpp llama-memory-hybrid.cpp llama-memory-hybrid-iswa.cpp diff --git a/examples/talk-llama/llama-arch.cpp b/examples/talk-llama/llama-arch.cpp index 4a52d9772..b890e66fc 100644 --- a/examples/talk-llama/llama-arch.cpp +++ b/examples/talk-llama/llama-arch.cpp @@ -77,6 +77,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_DEEPSEEK2, "deepseek2" }, { LLM_ARCH_DEEPSEEK2OCR, "deepseek2-ocr" }, { LLM_ARCH_DEEPSEEK32, "deepseek32" }, + { LLM_ARCH_DEEPSEEK4, "deepseek4" }, { LLM_ARCH_CHATGLM, "chatglm" }, { LLM_ARCH_GLM4, "glm4" }, { LLM_ARCH_GLM4_MOE, "glm4moe" }, @@ -129,6 +130,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_PANGU_EMBED, "pangu-embedded" }, { LLM_ARCH_MISTRAL3, "mistral3" }, { LLM_ARCH_EAGLE3, "eagle3" }, + { LLM_ARCH_DFLASH, "dflash" }, { LLM_ARCH_MISTRAL4, "mistral4" }, { LLM_ARCH_PADDLEOCR, "paddleocr" }, { LLM_ARCH_MIMO2, "mimo2" }, @@ -249,9 +251,19 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" }, { LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" }, { LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" }, + { LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" }, + { LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" }, + { LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" }, + { LLM_KV_ATTENTION_COMPRESS_RATIOS, "%s.attention.compress_ratios" }, { LLM_KV_ATTENTION_SHARED_KV_LAYERS, "%s.attention.shared_kv_layers" }, { LLM_KV_ATTENTION_RECURRENT_LAYERS, "%s.attention.recurrent_layers" }, + { LLM_KV_HYPER_CONNECTION_COUNT, "%s.hyper_connection.count" }, + { LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, "%s.hyper_connection.sinkhorn_iterations" }, + { LLM_KV_HYPER_CONNECTION_EPSILON, "%s.hyper_connection.epsilon" }, + + { LLM_KV_HASH_LAYER_COUNT, "%s.hash_layer_count" }, + { LLM_KV_ROPE_DIMENSION_COUNT, "%s.rope.dimension_count" }, { LLM_KV_ROPE_DIMENSION_COUNT_SWA, "%s.rope.dimension_count_swa" }, { LLM_KV_ROPE_DIMENSION_SECTIONS, "%s.rope.dimension_sections" }, @@ -439,6 +451,23 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_ATTN_Q_B, "blk.%d.attn_q_b" }, { LLM_TENSOR_ATTN_KV_A_MQA, "blk.%d.attn_kv_a_mqa" }, { LLM_TENSOR_ATTN_KV_B, "blk.%d.attn_kv_b" }, + { LLM_TENSOR_ATTN_KV, "blk.%d.attn_kv" }, + { LLM_TENSOR_ATTN_KV_NORM, "blk.%d.attn_kv_a_norm" }, + { LLM_TENSOR_ATTN_OUT_A, "blk.%d.attn_output_a" }, + { LLM_TENSOR_ATTN_OUT_B, "blk.%d.attn_output_b" }, + { LLM_TENSOR_HC_HEAD_FN, "output_hc_fn" }, + { LLM_TENSOR_HC_HEAD_BASE, "output_hc_base" }, + { LLM_TENSOR_HC_HEAD_SCALE, "output_hc_scale" }, + { LLM_TENSOR_HC_ATTN_FN, "blk.%d.hc_attn_fn" }, + { LLM_TENSOR_HC_ATTN_BASE, "blk.%d.hc_attn_base" }, + { LLM_TENSOR_HC_ATTN_SCALE, "blk.%d.hc_attn_scale" }, + { LLM_TENSOR_HC_FFN_FN, "blk.%d.hc_ffn_fn" }, + { LLM_TENSOR_HC_FFN_BASE, "blk.%d.hc_ffn_base" }, + { LLM_TENSOR_HC_FFN_SCALE, "blk.%d.hc_ffn_scale" }, + { 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" }, + { LLM_TENSOR_ATTN_COMPRESSOR_NORM, "blk.%d.attn_compressor_norm" }, { LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "per_layer_token_embd" }, { LLM_TENSOR_PER_LAYER_MODEL_PROJ, "per_layer_model_proj" }, { LLM_TENSOR_PER_LAYER_PROJ_NORM, "per_layer_proj_norm" }, @@ -565,6 +594,11 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" }, { LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" }, { LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" }, + { LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "blk.%d.indexer_compressor_kv" }, + { LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "blk.%d.indexer_compressor_gate" }, + { LLM_TENSOR_INDEXER_COMPRESSOR_APE, "blk.%d.indexer_compressor_ape" }, + { LLM_TENSOR_INDEXER_COMPRESSOR_NORM, "blk.%d.indexer_compressor_norm" }, + { LLM_TENSOR_FFN_GATE_TID2EID, "blk.%d.ffn_gate_tid2eid" }, { LLM_TENSOR_MASKED_EMBD_CENTROIDS, "masked_embd_centroids" }, { LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" }, { LLM_TENSOR_FC, "fc" }, @@ -615,6 +649,23 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_KV_A_MQA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_KV_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_KV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_KV_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_ATTN_OUT_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_OUT_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_HEAD_FN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_HEAD_BASE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}}, + {LLM_TENSOR_HC_HEAD_SCALE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_HC_ATTN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_ATTN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, + {LLM_TENSOR_HC_ATTN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_FFN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, + {LLM_TENSOR_HC_FFN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {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_ADD}}, + {LLM_TENSOR_ATTN_COMPRESSOR_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_ATTN_K_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_V_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_SINKS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SCALE}}, @@ -778,6 +829,11 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_INDEXER_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_INDEXER_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, + {LLM_TENSOR_INDEXER_COMPRESSOR_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_GATE_TID2EID, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, {LLM_TENSOR_NEXTN_PROJ_PRE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_NEXTN_PROJ_POST, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, // NextN/MTP tensors are stored per-block (blk.%d.nextn.*) even though only the @@ -932,6 +988,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_OLMOE: case LLM_ARCH_DEEPSEEK2: case LLM_ARCH_DEEPSEEK32: + case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_GLM_DSA: case LLM_ARCH_BITNET: case LLM_ARCH_T5: diff --git a/examples/talk-llama/llama-arch.h b/examples/talk-llama/llama-arch.h index 989da06d8..a4f5091e7 100644 --- a/examples/talk-llama/llama-arch.h +++ b/examples/talk-llama/llama-arch.h @@ -82,6 +82,7 @@ enum llm_arch { LLM_ARCH_DEEPSEEK2, LLM_ARCH_DEEPSEEK2OCR, LLM_ARCH_DEEPSEEK32, + LLM_ARCH_DEEPSEEK4, LLM_ARCH_CHATGLM, LLM_ARCH_GLM4, LLM_ARCH_GLM4_MOE, @@ -143,6 +144,7 @@ enum llm_arch { LLM_ARCH_TALKIE, LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, + LLM_ARCH_DFLASH, LLM_ARCH_UNKNOWN, }; @@ -254,9 +256,19 @@ enum llm_kv { LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, LLM_KV_ATTENTION_INDEXER_TOP_K, + LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, + LLM_KV_ATTENTION_OUTPUT_LORA_RANK, + LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, + LLM_KV_ATTENTION_COMPRESS_RATIOS, LLM_KV_ATTENTION_SHARED_KV_LAYERS, LLM_KV_ATTENTION_RECURRENT_LAYERS, + LLM_KV_HYPER_CONNECTION_COUNT, + LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, + LLM_KV_HYPER_CONNECTION_EPSILON, + + LLM_KV_HASH_LAYER_COUNT, + LLM_KV_ROPE_DIMENSION_COUNT, LLM_KV_ROPE_DIMENSION_COUNT_SWA, LLM_KV_ROPE_DIMENSION_SECTIONS, @@ -500,10 +512,27 @@ enum llm_tensor { LLM_TENSOR_ATTN_Q_B, LLM_TENSOR_ATTN_KV_A_MQA, LLM_TENSOR_ATTN_KV_B, + LLM_TENSOR_ATTN_KV, + LLM_TENSOR_ATTN_KV_NORM, + LLM_TENSOR_ATTN_OUT_A, + LLM_TENSOR_ATTN_OUT_B, LLM_TENSOR_ATTN_K_B, LLM_TENSOR_ATTN_V_B, LLM_TENSOR_ATTN_Q_A_NORM, LLM_TENSOR_ATTN_KV_A_NORM, + LLM_TENSOR_HC_HEAD_FN, + LLM_TENSOR_HC_HEAD_BASE, + LLM_TENSOR_HC_HEAD_SCALE, + LLM_TENSOR_HC_ATTN_FN, + LLM_TENSOR_HC_ATTN_BASE, + LLM_TENSOR_HC_ATTN_SCALE, + LLM_TENSOR_HC_FFN_FN, + LLM_TENSOR_HC_FFN_BASE, + LLM_TENSOR_HC_FFN_SCALE, + LLM_TENSOR_ATTN_COMPRESSOR_WKV, + LLM_TENSOR_ATTN_COMPRESSOR_WGATE, + LLM_TENSOR_ATTN_COMPRESSOR_APE, + LLM_TENSOR_ATTN_COMPRESSOR_NORM, LLM_TENSOR_ATTN_SUB_NORM, LLM_TENSOR_FFN_SUB_NORM, LLM_TENSOR_DEC_ATTN_NORM, @@ -565,6 +594,11 @@ enum llm_tensor { LLM_TENSOR_INDEXER_PROJ, LLM_TENSOR_INDEXER_ATTN_K, LLM_TENSOR_INDEXER_ATTN_Q_B, + LLM_TENSOR_INDEXER_COMPRESSOR_WKV, + LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, + LLM_TENSOR_INDEXER_COMPRESSOR_APE, + LLM_TENSOR_INDEXER_COMPRESSOR_NORM, + LLM_TENSOR_FFN_GATE_TID2EID, LLM_TENSOR_NEXTN_PROJ_PRE, LLM_TENSOR_NEXTN_PROJ_POST, LLM_TENSOR_NEXTN_EH_PROJ, diff --git a/examples/talk-llama/llama-batch.cpp b/examples/talk-llama/llama-batch.cpp index 6bf76939c..2b98a552f 100644 --- a/examples/talk-llama/llama-batch.cpp +++ b/examples/talk-llama/llama-batch.cpp @@ -379,6 +379,8 @@ bool llama_batch_allocr::init( LLAMA_LOG_ERROR("%s: sequence %d positions are decreasing (not allowed)\n", __func__, seq_id); return false; } + + cur_seq_pos[seq_id] = pos; } } } @@ -505,7 +507,7 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) { return ubatch_add(idxs, idxs.size(), false); } -llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) { +llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail) { if (sequential && has_cpl) { LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__); @@ -548,7 +550,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) } } - const uint32_t n_seqs = cur_seq_set.size(); + uint32_t n_seqs = cur_seq_set.size(); // we are done if (n_seqs == 0) { @@ -569,7 +571,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) std::vector idxs_per_seq(n_seqs); while (true) { - // we can only add new n_seq_tokens tokens if all the sequence sets have at least one more unused token and + // we can only add new n_seq_tokens tokens if all the sequence sets have at least 1 more unused tokens and // if we haven't reached n_ubatch bool can_expand = true; @@ -600,6 +602,72 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) } } + // if n_keep_tail > 0, keep only the seqs that either finish in this ubatch or have at least + // n_keep_tail tokens remaining for a future ubatch, so that the trailing n_keep_tail tokens + // of each seq are never split across ubatches + if (n_keep_tail > 0) { + GGML_ASSERT(n_ubatch > n_keep_tail); + + auto n_remaining = [&](uint32_t s) { + return (uint32_t) (seq_set_map[cur_seq_set[s]].size() - cur_idx[s]); + }; + + // keep the longest prefix of seqs that satisfy the constraint, to preserve sequential seq ids + uint32_t n_keep = 0; + while (n_keep < n_seqs) { + const uint32_t remaining = n_remaining(n_keep); + + if (remaining != 0 && remaining < n_keep_tail) { + break; + } + + n_keep++; + } + + // all seqs violate the constraint - resolve the first one directly and emit it alone + if (n_keep == 0) { + auto & idxs = idxs_per_seq[0]; + + const auto & seq_idxs = seq_set_map[cur_seq_set[0]]; + + if (idxs.size() + n_remaining(0) <= n_ubatch) { + // extend the seq to completion + while (n_remaining(0) > 0) { + const int32_t idx = seq_idxs[cur_idx[0]]; + + idxs.push_back(idx); + + used[idx] = true; + ++n_used; + + ++cur_idx[0]; + } + } else { + // truncate the seq so that at least n_keep_tail tokens remain + while (n_remaining(0) < n_keep_tail) { + used[idxs.back()] = false; + --n_used; + + idxs.pop_back(); + + --cur_idx[0]; + } + } + + n_keep = 1; + } + + // return the tokens of the deferred seqs back to the pool + for (uint32_t s = n_keep; s < n_seqs; ++s) { + for (const int32_t idx : idxs_per_seq[s]) { + used[idx] = false; + --n_used; + } + } + + n_seqs = n_keep; + } + // concat the per-sequence-set lists std::vector idxs; @@ -814,7 +882,7 @@ void llama_batch_allocr::ubatch_print(const llama_ubatch & ubatch, int debug) { LLAMA_LOG_DEBUG("%s: output = %p\n", __func__, (void *) ubatch.output); LLAMA_LOG_DEBUG("%s: n_outputs = %d\n", __func__, n_outputs); - if (debug > 1) { + if (debug > 0) { int seq_id_max = 0; for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { for (int s = 0; s < ubatch.n_seq_id[i]; ++s) { diff --git a/examples/talk-llama/llama-batch.h b/examples/talk-llama/llama-batch.h index f77520e86..a3d1889d4 100644 --- a/examples/talk-llama/llama-batch.h +++ b/examples/talk-llama/llama-batch.h @@ -104,7 +104,8 @@ public: // make ubatches of equal-length sequences sets // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids - llama_ubatch split_equal(uint32_t n_ubatch, bool sequential); + // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch + llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail); // sequence-set-wise split - each ubatch contains a single sequence-set llama_ubatch split_seq(uint32_t n_ubatch); diff --git a/examples/talk-llama/llama-context.cpp b/examples/talk-llama/llama-context.cpp index 220240ea9..5edfc85ab 100644 --- a/examples/talk-llama/llama-context.cpp +++ b/examples/talk-llama/llama-context.cpp @@ -17,6 +17,7 @@ #include #include #include +#include // // llama_context @@ -30,6 +31,30 @@ static llm_graph_type ctx_type_to_graph_type(llama_context_type ctx_type) { throw std::runtime_error("Unsupported ctx type"); } +struct llm_fused_op_probe { + llm_fused_op op; + const char * name; + uint32_t n_tokens_per_seq; +}; + +static const llm_fused_op_probe llm_fused_op_flash_attn_probe = { + /*.op =*/ LLM_FUSED_OP_FLASH_ATTN, + /*.name =*/ "Flash Attention", + /*.n_tokens_per_seq =*/ 1, +}; + +static const llm_fused_op_probe llm_fused_op_gdn_ar_probe = { + /*.op =*/ LLM_FUSED_OP_GDN_AR, + /*.name =*/ "fused Gated Delta Net (autoregressive)", + /*.n_tokens_per_seq =*/ 1, +}; + +static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = { + /*.op =*/ LLM_FUSED_OP_GDN_CH, + /*.name =*/ "fused Gated Delta Net (chunked)", + /*.n_tokens_per_seq =*/ 16, +}; + llama_context::llama_context( const llama_model & model, llama_context_params params) : @@ -100,10 +125,10 @@ llama_context::llama_context( cparams.ctx_other = params.ctx_other; } - if (model.arch == LLM_ARCH_EAGLE3) { + if (model.arch == LLM_ARCH_EAGLE3 || model.arch == LLM_ARCH_DFLASH) { if (model.tok_embd == nullptr || model.output == nullptr) { if (params.ctx_other == nullptr) { - throw std::runtime_error("EAGLE3 requires ctx_other to be set (this warning is normal during memory fitting)"); + throw std::runtime_error(model.arch_name() + " requires ctx_other to be set (this warning is normal during memory fitting)"); } cparams.ctx_other = params.ctx_other; } @@ -256,7 +281,7 @@ llama_context::llama_context( LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max); if (cparams.n_ctx_seq < hparams.n_ctx_train) { - LLAMA_LOG_WARN("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n", + LLAMA_LOG_INFO("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n", __func__, cparams.n_ctx_seq, hparams.n_ctx_train); } @@ -436,6 +461,69 @@ llama_context::~llama_context() { ggml_opt_free(opt_ctx); } +void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs) { + const char * func = __func__; + auto resolve = [&](const llm_fused_op_probe & probe, bool & enabled) { + if (!enabled) { + return; + } + + const uint32_t n_tokens_probe = probe.n_tokens_per_seq*n_seqs; + + auto * gf = graph_reserve(n_tokens_probe, n_seqs, n_tokens_probe, mctx, true); + if (!gf) { + throw std::runtime_error(std::string("failed to reserve graph for ") + probe.name + " check"); + } + + bool device_mismatch = false; + for (const auto & node : get_gf_res_reserve()->get_fused_nodes()) { + if (node.op != probe.op) { + continue; + } + + GGML_ASSERT(node.il >= 0); + + ggml_backend_t backend_fused = ggml_backend_sched_get_tensor_backend(sched.get(), node.tensor); + ggml_backend_dev_t device_fused = backend_fused ? ggml_backend_get_device(backend_fused) : nullptr; + + // TODO: make this descriptor-specific; model.dev_layer() preserves the current behavior, + // but is still wrong for cases like --no-kv-offload. + ggml_backend_dev_t device_layer = model.dev_layer(node.il); + + if (device_fused != device_layer) { + LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but %s " + "is assigned to device %s (usually due to missing support)\n", + func, node.il, + device_layer ? ggml_backend_dev_name(device_layer) : "none", + probe.name, + device_fused ? ggml_backend_dev_name(device_fused) : "none"); + device_mismatch = true; + break; + } + } + + if (device_mismatch) { + enabled = false; + LLAMA_LOG_WARN("%s: %s not supported, set to disabled\n", func, probe.name); + } else { + enabled = true; + LLAMA_LOG_INFO("%s: %s enabled\n", func, probe.name); + } + }; + + if (cparams.auto_fa) { + resolve(llm_fused_op_flash_attn_probe, cparams.flash_attn); + cparams.auto_fa = false; + } + + if (cparams.auto_fgdn) { + LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", func); + resolve(llm_fused_op_gdn_ar_probe, cparams.fused_gdn_ar); + resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch); + cparams.auto_fgdn = false; + } +} + void llama_context::sched_reserve() { if (!sched_need_reserve) { return; @@ -475,128 +563,7 @@ void llama_context::sched_reserve() { LLAMA_LOG_DEBUG("%s: worst-case: n_tokens = %d, n_seqs = %d, n_outputs = %d\n", __func__, n_tokens, n_seqs, n_outputs); - // resolve automatic Flash Attention use - if (cparams.auto_fa) { - auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true); - if (!gf) { - throw std::runtime_error("failed to reserve graph for Flash Attention check"); - } - - const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FATTN) + 1; - bool fa_device_mismatch = false; - for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { - ggml_tensor * n = ggml_graph_node(gf, i); - if (n->op != GGML_OP_FLASH_ATTN_EXT) { - continue; - } - ggml_backend_dev_t device_fa = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n)); - - // TODO: instead of the tensor names, use a map to keep track of which (FA) tensors belong to which layer - GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FATTN "-", prefix_len) == 0); - const int il = std::stoi(n->name + prefix_len); - ggml_backend_dev_t device_kv = model.dev_layer(il); - if (device_fa != device_kv) { - LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the Flash Attention tensor " - "is assigned to device %s (usually due to missing support)\n", - __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_fa)); - // FIXME: fa_device_mismatch logic is wrong for --no-kv-offload, but this is broken anyways - fa_device_mismatch = true; - break; - } - } - - if (fa_device_mismatch) { - cparams.flash_attn = false; - LLAMA_LOG_WARN("%s: Flash Attention was auto, set to disabled\n", __func__); - } else { - cparams.flash_attn = true; - LLAMA_LOG_INFO("%s: Flash Attention was auto, set to enabled\n", __func__); - } - - cparams.auto_fa = false; - } - - if (cparams.auto_fgdn) { - LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", __func__); - - if (cparams.fused_gdn_ar) { - auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true); - if (!gf) { - throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (autoregressive)"); - } - - const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FGDN_AR) + 1; - bool gdn_device_mismatch = false; - for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { - ggml_tensor * n = ggml_graph_node(gf, i); - if (n->op != GGML_OP_GATED_DELTA_NET) { - continue; - } - ggml_backend_dev_t device_gdn = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n)); - - GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FGDN_AR "-", prefix_len) == 0); - const int il = std::stoi(n->name + prefix_len); - ggml_backend_dev_t device_kv = model.dev_layer(il); - if (device_gdn != device_kv) { - LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the fused Gated Delta Net tensor " - "is assigned to device %s (usually due to missing support)\n", - __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_gdn)); - gdn_device_mismatch = true; - break; - } - } - - if (gdn_device_mismatch) { - cparams.fused_gdn_ar = false; - LLAMA_LOG_WARN("%s: fused Gated Delta Net (autoregressive) not supported, set to disabled\n", __func__); - } else { - LLAMA_LOG_INFO("%s: fused Gated Delta Net (autoregressive) enabled\n", __func__); - } - } - - if (cparams.fused_gdn_ch) { - // more than one token in the batch per sequence in order to take the chunked path - // note: n_outputs must match n_tokens for embedding models with mean/rank pooling, - // because build_pooling creates inp_mean with shape [n_tokens, n_seqs] and multiplies - // it with t_embd which is reduced to [n_outputs, ...] via out_ids. if n_outputs != n_tokens, - // the ggml_mul_mat assertion fails. - const uint32_t n_tokens_ch = 16*n_seqs; - auto * gf = graph_reserve(n_tokens_ch, n_seqs, n_tokens_ch, mctx.get(), true); - if (!gf) { - throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (chunked)"); - } - - const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FGDN_CH) + 1; - bool gdn_device_mismatch = false; - for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { - ggml_tensor * n = ggml_graph_node(gf, i); - if (n->op != GGML_OP_GATED_DELTA_NET) { - continue; - } - ggml_backend_dev_t device_gdn = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n)); - - GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FGDN_CH "-", prefix_len) == 0); - const int il = std::stoi(n->name + prefix_len); - ggml_backend_dev_t device_kv = model.dev_layer(il); - if (device_gdn != device_kv) { - LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the fused Gated Delta Net tensor " - "is assigned to device %s (usually due to missing support)\n", - __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_gdn)); - gdn_device_mismatch = true; - break; - } - } - - if (gdn_device_mismatch) { - cparams.fused_gdn_ch = false; - LLAMA_LOG_WARN("%s: fused Gated Delta Net (chunked) not supported, set to disabled\n", __func__); - } else { - LLAMA_LOG_INFO("%s: fused Gated Delta Net (chunked) enabled\n", __func__); - } - } - - cparams.auto_fgdn = false; - } + resolve_fused_ops(mctx.get(), n_seqs); // reserve worst-case graph int n_splits_pp = -1; @@ -2321,7 +2288,11 @@ void llama_context::output_reorder() { // uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { - if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE) { + if (model.arch == LLM_ARCH_QWEN3NEXT || + model.arch == LLM_ARCH_KIMI_LINEAR || + model.arch == LLM_ARCH_QWEN35 || + model.arch == LLM_ARCH_QWEN35MOE || + model.arch == LLM_ARCH_DEEPSEEK4) { return std::max(n_tokens * 40, 32u * model.n_tensors()); } uint32_t res = std::max(1024u, 8u*model.n_tensors()); diff --git a/examples/talk-llama/llama-context.h b/examples/talk-llama/llama-context.h index f8b780587..bf91daa8b 100644 --- a/examples/talk-llama/llama-context.h +++ b/examples/talk-llama/llama-context.h @@ -262,6 +262,10 @@ private: llm_graph_cb graph_get_cb() const; + // disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device + // that differs from the layer it belongs to (usually due to missing backend support) + void resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs); + // TODO: read/write lora adapters and cvec size_t state_write_data(llama_io_write_i & io); size_t state_read_data (llama_io_read_i & io); diff --git a/examples/talk-llama/llama-graph.cpp b/examples/talk-llama/llama-graph.cpp index 68c9e606c..a8fd11ebc 100644 --- a/examples/talk-llama/llama-graph.cpp +++ b/examples/talk-llama/llama-graph.cpp @@ -8,6 +8,7 @@ #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" #include "llama-kv-cache-dsa.h" +#include "llama-kv-cache-dsv4.h" #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-recurrent.h" @@ -17,6 +18,7 @@ #include #include #include +#include #include // dedup helpers @@ -61,26 +63,6 @@ static bool can_reuse_kq_mask( // impl -static ggml_tensor * ggml_mul_mat_aux( - ggml_context * ctx, - ggml_tensor * cur, - ggml_tensor * rot) { - const auto n = rot->ne[0]; - - ggml_tensor * res; - - if (!ggml_is_contiguous(cur)) { - res = ggml_cont_2d (ctx, cur, n, ggml_nelements(cur)/n); - } else { - res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n); - } - res = ggml_mul_mat (ctx, rot, res); - ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD); - res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]); - - return res; -} - void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) { if (ubatch->token) { const int64_t n_tokens = ubatch->n_tokens; @@ -486,13 +468,17 @@ void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { mctx->set_input_k_idxs(self_k_idxs, ubatch); mctx->set_input_v_idxs(self_v_idxs, ubatch); - mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + // the mask is left unallocated when the graph only stores K/V without attending + // (e.g. DFlash's KV-injection pass) + if (self_kq_mask && self_kq_mask->buffer) { + mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + } - if (self_k_rot) { + if (self_k_rot && self_k_rot->buffer) { mctx->set_input_k_rot(self_k_rot); } - if (self_v_rot) { + if (self_v_rot && self_v_rot->buffer) { mctx->set_input_v_rot(self_v_rot); } } @@ -564,7 +550,9 @@ void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) { // base tensors may not be allocated if there are no non-SWA attention layers if (self_k_idxs && self_k_idxs->buffer) { mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch); - mctx->get_base()->set_input_v_idxs(self_v_idxs, ubatch); + if (self_v_idxs) { + mctx->get_base()->set_input_v_idxs(self_v_idxs, ubatch); + } } // the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live @@ -575,26 +563,28 @@ void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) { // swa tensors may not be allocated if there are no SWA attention layers if (self_k_idxs_swa && self_k_idxs_swa->buffer) { mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch); - mctx->get_swa()->set_input_v_idxs(self_v_idxs_swa, ubatch); + if (self_v_idxs_swa) { + mctx->get_swa()->set_input_v_idxs(self_v_idxs_swa, ubatch); + } } if (self_kq_mask_swa && self_kq_mask_swa->buffer) { mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn); } - if (self_k_rot) { + if (self_k_rot && self_k_rot->buffer) { mctx->get_base()->set_input_k_rot(self_k_rot); } - if (self_v_rot) { + if (self_v_rot && self_v_rot->buffer) { mctx->get_base()->set_input_v_rot(self_v_rot); } - if (self_k_rot_swa) { + if (self_k_rot_swa && self_k_rot_swa->buffer) { mctx->get_swa()->set_input_k_rot(self_k_rot_swa); } - if (self_v_rot_swa) { + if (self_v_rot_swa && self_v_rot_swa->buffer) { mctx->get_swa()->set_input_v_rot(self_v_rot_swa); } } @@ -629,6 +619,305 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) { return res; } +static void dsv4_set_i64(ggml_tensor * dst, const std::vector & src) { + if (!dst || !dst->buffer) { + return; + } + + GGML_ASSERT(dst->ne[0] == (int64_t) src.size()); + ggml_backend_tensor_set(dst, src.data(), 0, src.size()*ggml_element_size(dst)); +} + +static void dsv4_set_i32(ggml_tensor * dst, const std::vector & src) { + if (!dst || !dst->buffer) { + return; + } + + GGML_ASSERT(dst->ne[0] == (int64_t) src.size()); + ggml_backend_tensor_set(dst, src.data(), 0, src.size()*ggml_element_size(dst)); +} + +static void dsv4_set_kq_mask( + ggml_tensor * dst, + const llama_kv_cache_dsv4_context::comp_plan & plan, + uint32_t n_tokens, + int64_t n_stream) { + if (!dst || !dst->buffer) { + return; + } + + GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); + GGML_ASSERT(n_stream > 0); + GGML_ASSERT(n_tokens%n_stream == 0); + GGML_ASSERT(dst->ne[0] == plan.n_kv); + GGML_ASSERT(dst->ne[1] == (int64_t) n_tokens/n_stream); + GGML_ASSERT(dst->ne[2] == 1); + GGML_ASSERT(dst->ne[3] == n_stream); + GGML_ASSERT((int64_t) plan.n_visible.size() == (int64_t) n_tokens); + GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); + + if (dst->type == GGML_TYPE_F32) { + float * data = (float *) dst->data; + + for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { + const int32_t n_visible = plan.n_visible[i]; + + for (int64_t j = 0; j < dst->ne[0]; ++j) { + data[i*dst->ne[0] + j] = j < n_visible ? 0.0f : -INFINITY; + } + } + } else if (dst->type == GGML_TYPE_F16) { + ggml_fp16_t * data = (ggml_fp16_t *) dst->data; + const ggml_fp16_t fp16_ninf = llama_cast(-INFINITY); + const ggml_fp16_t fp16_zero = llama_cast(0.0f); + + for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { + const int32_t n_visible = plan.n_visible[i]; + + for (int64_t j = 0; j < dst->ne[0]; ++j) { + data[i*dst->ne[0] + j] = j < n_visible ? fp16_zero : fp16_ninf; + } + } + } +} + +static ggml_tensor * dsv4_build_raw_kq_mask( + ggml_context * ctx, + const llama_kv_cache_dsv4_raw_context * mctx, + const llama_ubatch & ubatch, + const llama_cparams & cparams, + int64_t n_stream) { + const auto n_kv = mctx->get_n_kv(); + const auto n_tokens = ubatch.n_tokens; + + GGML_ASSERT(n_stream > 0); + GGML_ASSERT(n_tokens%n_stream == 0); + + const auto type = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32; + + ggml_tensor * res = ggml_new_tensor_4d(ctx, type, n_kv, n_tokens/n_stream, 1, n_stream); + ggml_set_input(res); + ggml_set_name(res, "attn_inp_kq_mask"); + + return res; +} + +static bool dsv4_can_reuse_raw_kq_mask( + ggml_tensor * kq_mask, + const llama_kv_cache_dsv4_raw_context * mctx, + const llama_ubatch & ubatch, + int64_t n_stream) { + const auto n_kv = mctx->get_n_kv(); + const auto n_tokens = ubatch.n_tokens; + + GGML_ASSERT(n_stream > 0); + + bool res = true; + + res &= (kq_mask->ne[0] == n_kv); + res &= (kq_mask->ne[1] == n_tokens/n_stream); + res &= (kq_mask->ne[2] == 1); + res &= (kq_mask->ne[3] == n_stream); + + return res; +} + +static std::string dsv4_plan_positions(const std::vector & values) { + std::ostringstream ss; + ss << "["; + for (size_t i = 0; i < values.size(); ++i) { + if (i > 0) { + ss << ", "; + } + ss << values[i]; + } + ss << "]"; + return ss.str(); +} + +static bool dsv4_compress_debug() { + static const bool debug = []() { + const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG"); + return env && atoi(env) > 0; + }(); + + return debug; +} + +static void dsv4_set_comp_inputs( + const llm_graph_input_dsv4::comp_input & inp, + const llama_kv_cache_dsv4_context::comp_plan & plan, + const char * name, + bool debug, + uint32_t n_tokens, + int64_t n_stream) { + dsv4_set_i32(inp.state_pos, plan.state_pos); + dsv4_set_i32(inp.state_persist_src_idxs, plan.state_persist_src_idxs); + dsv4_set_i32(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs); + dsv4_set_i32(inp.state_read_idxs, plan.state_read_idxs); + dsv4_set_i64(inp.state_write_idxs, plan.state_write_idxs); + dsv4_set_i32(inp.state_write_pos, plan.state_write_pos); + dsv4_set_kq_mask(inp.kq_mask, plan, n_tokens, n_stream); + + if (debug || dsv4_compress_debug()) { + LLAMA_LOG_INFO("%s: %s n_tokens=%u, n_stream=%d, state_persist_dst=%s, state_write_pos=%s\n", + __func__, name, n_tokens, (int) n_stream, + dsv4_plan_positions(plan.state_persist_dst_idxs).c_str(), + dsv4_plan_positions(plan.state_write_pos).c_str()); + } +} + +static bool dsv4_can_reuse_tensor_1d(ggml_tensor * t, int64_t ne0) { + return (t == nullptr && ne0 == 0) || (t != nullptr && t->ne[0] == ne0); +} + +static bool dsv4_can_reuse_kq_mask( + ggml_tensor * t, + const llama_kv_cache_dsv4_context::comp_plan & plan, + uint32_t n_tokens, + int64_t n_stream) { + if (plan.n_kv == 0) { + return t == nullptr; + } + + GGML_ASSERT(n_stream > 0); + + return t != nullptr && + t->ne[0] == plan.n_kv && + t->ne[1] == (int64_t) n_tokens/n_stream && + t->ne[2] == 1 && + t->ne[3] == n_stream; +} + +static bool dsv4_can_reuse_comp_input( + const llm_graph_input_dsv4::comp_input & inp, + const llama_kv_cache_dsv4_context::comp_plan & plan, + uint32_t n_tokens, + int64_t n_stream) { + bool res = true; + res &= dsv4_can_reuse_tensor_1d(inp.state_pos, plan.state_pos.size()); + res &= dsv4_can_reuse_tensor_1d(inp.state_persist_src_idxs, plan.state_persist_src_idxs.size()); + res &= dsv4_can_reuse_tensor_1d(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs.size()); + res &= dsv4_can_reuse_tensor_1d(inp.state_read_idxs, plan.state_read_idxs.size()); + res &= dsv4_can_reuse_tensor_1d(inp.state_write_idxs, plan.state_write_idxs.size()); + res &= dsv4_can_reuse_tensor_1d(inp.state_write_pos, plan.state_write_pos.size()); + res &= dsv4_can_reuse_kq_mask(inp.kq_mask, plan, n_tokens, n_stream); + + return res; +} + +static ggml_tensor * dsv4_build_input_1d( + ggml_context * ctx, + ggml_type type, + int64_t ne0, + const std::string & name) { + if (ne0 == 0) { + return nullptr; + } + + ggml_tensor * res = ggml_new_tensor_1d(ctx, type, ne0); + ggml_set_input(res); + ggml_set_name(res, name.c_str()); + + return res; +} + +static void dsv4_build_comp_inputs( + ggml_context * ctx, + llm_graph_input_dsv4::comp_input & inp, + const llama_kv_cache_dsv4_context::comp_plan & plan, + const char * name, + const llama_cparams & cparams, + int64_t n_stream) { + inp.state_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_pos.size(), std::string("dsv4_") + name + "_state_pos"); + inp.state_persist_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_src_idxs.size(), std::string("dsv4_") + name + "_state_persist_src_idxs"); + inp.state_persist_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_dst_idxs.size(), std::string("dsv4_") + name + "_state_persist_dst_idxs"); + inp.state_read_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_read_idxs.size(), std::string("dsv4_") + name + "_state_read_idxs"); + inp.state_write_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I64, plan.state_write_idxs.size(), std::string("dsv4_") + name + "_state_write_idxs"); + inp.state_write_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_write_pos.size(), std::string("dsv4_") + name + "_state_write_pos"); + + if (plan.n_kv > 0) { + const int64_t n_tokens = (int64_t) plan.n_visible.size(); + + GGML_ASSERT(n_stream > 0); + GGML_ASSERT(n_tokens%n_stream == 0); + + inp.kq_mask = ggml_new_tensor_4d(ctx, cparams.flash_attn && strcmp(name, "lid") != 0 ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream); + ggml_set_input(inp.kq_mask); + ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str()); + } +} + +void llm_graph_input_dsv4_raw::set_input(const llama_ubatch * ubatch) { + if (self_k_idxs && self_k_idxs->buffer) { + mctx->set_input_k_idxs(self_k_idxs); + } + + if (self_kq_mask && self_kq_mask->buffer) { + mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + } + + if (self_k_rot) { + mctx->set_input_k_rot(self_k_rot); + } +} + +void llm_graph_input_dsv4::set_input(const llama_ubatch * ubatch) { + const auto & plan_csa = mctx->get_csa_plan(*ubatch); + const auto & plan_hca = mctx->get_hca_plan(*ubatch); + const auto & plan_lid = mctx->get_lid_plan(*ubatch); + const int64_t n_stream = plan_csa.n_stream; + + inp_raw->mctx = mctx->get_raw(); + inp_raw->set_input(ubatch); + + dsv4_set_comp_inputs(inp_csa, plan_csa, "csa", debug > 0, ubatch->n_tokens, n_stream); + dsv4_set_comp_inputs(inp_hca, plan_hca, "hca", debug > 0, ubatch->n_tokens, n_stream); + dsv4_set_comp_inputs(inp_lid, plan_lid, "lid", debug > 0, ubatch->n_tokens, n_stream); + + if (inp_csa.k_rot && inp_csa.k_rot->buffer) { + mctx->get_csa()->set_input_k_rot(inp_csa.k_rot); + } + + if (inp_hca.k_rot && inp_hca.k_rot->buffer) { + mctx->get_hca()->set_input_k_rot(inp_hca.k_rot); + } + + if (inp_lid.k_rot && inp_lid.k_rot->buffer) { + mctx->get_lid()->set_input_k_rot(inp_lid.k_rot); + } +} + +bool llm_graph_input_dsv4::can_reuse(const llm_graph_params & params) { + const auto * mctx = static_cast(params.mctx); + + this->mctx = mctx; + inp_raw->mctx = mctx->get_raw(); + + bool res = true; + + const auto & plan_csa = mctx->get_csa_plan(params.ubatch); + const auto & plan_hca = mctx->get_hca_plan(params.ubatch); + const auto & plan_lid = mctx->get_lid_plan(params.ubatch); + const int64_t n_stream = plan_csa.n_stream; + + const auto * raw_ctx = mctx->get_raw(); + inp_raw->mctx = raw_ctx; + + if (inp_raw->self_k_idxs && inp_raw->self_k_idxs->buffer) { + res &= inp_raw->self_k_idxs->ne[0] == raw_ctx->get_n_write(); + } + if (inp_raw->self_kq_mask && inp_raw->self_kq_mask->buffer) { + res &= dsv4_can_reuse_raw_kq_mask(inp_raw->self_kq_mask, raw_ctx, params.ubatch, n_stream); + } + + res &= dsv4_can_reuse_comp_input(inp_csa, plan_csa, params.ubatch.n_tokens, n_stream); + res &= dsv4_can_reuse_comp_input(inp_hca, plan_hca, params.ubatch.n_tokens, n_stream); + res &= dsv4_can_reuse_comp_input(inp_lid, plan_lid, params.ubatch.n_tokens, n_stream); + + return res; +} + void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) { GGML_ASSERT(cross_kq_mask); @@ -904,6 +1193,7 @@ void llm_graph_result::reset() { t_logits = nullptr; t_embd = nullptr; t_embd_pooled = nullptr; + t_h_nextn = nullptr; t_layer_inp.resize(LLAMA_MAX_LAYERS); std::fill(t_layer_inp.begin(), t_layer_inp.end(), nullptr); @@ -916,6 +1206,7 @@ void llm_graph_result::reset() { params = {}; inputs.clear(); + fused_nodes.clear(); buf_compute_meta.resize(ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false)); @@ -1017,6 +1308,10 @@ llm_graph_input_i * llm_graph_result::add_input(llm_graph_input_ptr input) { return inputs.back().get(); } +void llm_graph_result::add_fused_node(llm_graph_fused_node result) { + fused_nodes.push_back(result); +} + void llm_graph_result::set_params(const llm_graph_params & params) { this->params = params; } @@ -1076,6 +1371,8 @@ void llm_graph_context::cb(ggml_tensor * cur, const char * name, int il) const { } } + + ggml_tensor * llm_graph_context::build_cvec( ggml_tensor * cur, int il) const { @@ -1346,20 +1643,24 @@ ggml_tensor * llm_graph_context::build_ffn( switch (type_op) { case LLM_FFN_SILU: if (gate && type_gate == LLM_FFN_PAR) { - // Step35: HF clamps gate (after SiLU) and up before multiplication - if (arch == LLM_ARCH_STEP35 && il >= 0) { + if (il >= 0) { const float limit = hparams.swiglu_clamp_shexp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - ggml_tensor * gate_act = ggml_silu(ctx0, cur); - cb(gate_act, "ffn_silu", il); - gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit); - cb(gate_act, "ffn_silu_clamped", il); - tmp = ggml_clamp(ctx0, tmp, -limit, limit); cb(tmp, "ffn_up_clamped", il); - cur = ggml_mul(ctx0, gate_act, tmp); + if (arch == LLM_ARCH_DEEPSEEK4) { + cur = ggml_clamp(ctx0, cur, -INFINITY, limit); + cb(cur, "ffn_gate_clamped", il); + cur = ggml_swiglu_split(ctx0, cur, tmp); + } else { + ggml_tensor * gate_act = ggml_silu(ctx0, cur); + cb(gate_act, "ffn_silu", il); + gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit); + cb(gate_act, "ffn_silu_clamped", il); + cur = ggml_mul(ctx0, gate_act, tmp); + } cb(cur, "ffn_swiglu_limited", il); type_gate = LLM_FFN_SEQ; break; @@ -1469,7 +1770,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * gate_up_exps, ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, - ggml_tensor * down_exps_s) const { + ggml_tensor * down_exps_s, + ggml_tensor * selected_experts_in) const { return build_moe_ffn( cur, gate_inp, /* gate_inp_b */ nullptr, @@ -1489,7 +1791,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( /* gate_up_exps_b */ nullptr, up_exps_s, gate_exps_s, - down_exps_s + down_exps_s, + selected_experts_in ); } @@ -1516,7 +1819,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * gate_up_exps_b, ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, - ggml_tensor * down_exps_s) const { + ggml_tensor * down_exps_s, + ggml_tensor * selected_experts_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 @@ -1525,6 +1829,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn( if (probs_in == nullptr) { logits = build_lora_mm(gate_inp, cur); // [n_expert, n_tokens] + if (gating_op == LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) { + ggml_mul_mat_set_prec(logits, GGML_PREC_F32); + } cb(logits, "ffn_moe_logits", il); } else { logits = probs_in; @@ -1549,6 +1856,10 @@ ggml_tensor * llm_graph_context::build_moe_ffn( { probs = logits; // [n_expert, n_tokens] } break; + case LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS: + { + probs = ggml_sqrt(ctx0, ggml_softplus(ctx0, logits)); // [n_expert, n_tokens] + } break; default: GGML_ABORT("fatal error"); } @@ -1599,8 +1910,11 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } // select experts - ggml_tensor * selected_experts = ggml_argsort_top_k(ctx0, selection_probs, n_expert_used); // [n_expert_used, n_tokens] - cb(selected_experts->src[0], "ffn_moe_argsort", il); + ggml_tensor * selected_experts = selected_experts_in; + if (selected_experts == nullptr) { + selected_experts = ggml_argsort_top_k(ctx0, selection_probs, n_expert_used); // [n_expert_used, n_tokens] + cb(selected_experts->src[0], "ffn_moe_argsort", il); + } cb(selected_experts, "ffn_moe_topk", il); if (arch == LLM_ARCH_GROVEMOE && n_expert != hparams.n_expert) { @@ -1713,20 +2027,24 @@ ggml_tensor * llm_graph_context::build_moe_ffn( switch (type_op) { case LLM_FFN_SILU: if (gate_exps) { - // Step35: per-layer clamp for routed experts - if (arch == LLM_ARCH_STEP35 && il >= 0) { + if (il >= 0) { const float limit = hparams.swiglu_clamp_exp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - ggml_tensor * gate_act = ggml_silu(ctx0, cur); - cb(gate_act, "ffn_moe_silu", il); - gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit); - cb(gate_act, "ffn_moe_silu_clamped", il); - up = ggml_clamp(ctx0, up, -limit, limit); cb(up, "ffn_moe_up_clamped", il); - cur = ggml_mul(ctx0, gate_act, up); + if (arch == LLM_ARCH_DEEPSEEK4) { + cur = ggml_clamp(ctx0, cur, -INFINITY, limit); + cb(cur, "ffn_moe_gate_clamped", il); + cur = ggml_swiglu_split(ctx0, cur, up); + } else { + ggml_tensor * gate_act = ggml_silu(ctx0, cur); + cb(gate_act, "ffn_moe_silu", il); + gate_act = ggml_clamp(ctx0, gate_act, -INFINITY, limit); + cb(gate_act, "ffn_moe_silu_clamped", il); + cur = ggml_mul(ctx0, gate_act, up); + } cb(cur, "ffn_moe_swiglu_limited", il); break; } @@ -2105,7 +2423,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias, hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); - cb(cur, LLAMA_TENSOR_NAME_FATTN, il); + res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); ggml_flash_attn_ext_add_sinks(cur, sinks); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); @@ -2324,12 +2642,12 @@ ggml_tensor * llm_graph_context::build_attn( GGML_ASSERT(v_mla == nullptr); if (inp->self_k_rot) { - q_cur = ggml_mul_mat_aux(ctx0, q_cur, inp->self_k_rot); - k_cur = ggml_mul_mat_aux(ctx0, k_cur, inp->self_k_rot); + q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot); + k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot); } if (inp->self_v_rot) { - v_cur = ggml_mul_mat_aux(ctx0, v_cur, inp->self_v_rot); + v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot); } // these nodes are added to the graph together so that they are not reordered @@ -2360,7 +2678,7 @@ ggml_tensor * llm_graph_context::build_attn( cb(cur, "kqv_out", il); if (inp->self_v_rot) { - cur = ggml_mul_mat_aux(ctx0, cur, inp->self_v_rot); + cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); } if (wo) { @@ -2565,14 +2883,14 @@ ggml_tensor * llm_graph_context::build_attn( auto * v_rot = is_swa ? inp->self_v_rot_swa : inp->self_v_rot; if (k_rot) { - q_cur = ggml_mul_mat_aux(ctx0, q_cur, k_rot); + q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot); if (k_cur) { - k_cur = ggml_mul_mat_aux(ctx0, k_cur, k_rot); + k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot); } } if (v_rot) { if (v_cur) { - v_cur = ggml_mul_mat_aux(ctx0, v_cur, v_rot); + v_cur = llama_mul_mat_hadamard(ctx0, v_cur, v_rot); } } @@ -2615,7 +2933,7 @@ ggml_tensor * llm_graph_context::build_attn( cb(cur, "kqv_out", il); if (v_rot) { - cur = ggml_mul_mat_aux(ctx0, cur, v_rot); + cur = llama_mul_mat_hadamard(ctx0, cur, v_rot); } if (wo) { @@ -2755,6 +3073,33 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const return (llm_graph_input_attn_kv_iswa *) res->add_input(std::move(inp)); } +llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const { + const auto * mctx_cur = static_cast(mctx); + const auto * raw_ctx = mctx_cur->get_raw(); + + auto inp_raw = std::make_unique(cparams, raw_ctx); + + const int64_t n_stream = mctx_cur->get_csa_plan(ubatch).n_stream; + + GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "DSV4 expects SWA raw cache"); + + inp_raw->self_k_idxs = raw_ctx->build_input_k_idxs(ctx0, ubatch); + inp_raw->self_kq_mask = dsv4_build_raw_kq_mask(ctx0, raw_ctx, ubatch, cparams, n_stream); + inp_raw->self_kq_mask_cnv = inp_raw->self_kq_mask; + + inp_raw->self_k_rot = raw_ctx->build_input_k_rot(ctx0); + auto inp = std::make_unique(cparams, std::move(inp_raw), mctx_cur); + + dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(ubatch), "csa", cparams, n_stream); + dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(ubatch), "hca", cparams, n_stream); + dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(ubatch), "lid", cparams, n_stream); + inp->inp_csa.k_rot = mctx_cur->get_csa()->build_input_k_rot(ctx0); + inp->inp_hca.k_rot = mctx_cur->get_hca()->build_input_k_rot(ctx0); + inp->inp_lid.k_rot = mctx_cur->get_lid()->build_input_k_rot(ctx0); + + return (llm_graph_input_dsv4 *) res->add_input(std::move(inp)); +} + ggml_tensor * llm_graph_context::build_rs( ggml_tensor * s, ggml_tensor * state_copy_main, diff --git a/examples/talk-llama/llama-graph.h b/examples/talk-llama/llama-graph.h index a6e8c3985..97141ef93 100644 --- a/examples/talk-llama/llama-graph.h +++ b/examples/talk-llama/llama-graph.h @@ -23,6 +23,8 @@ struct llama_memory_context_i; class llama_kv_cache_context; class llama_kv_cache_dsa_context; +class llama_kv_cache_dsv4_raw_context; +class llama_kv_cache_dsv4_context; class llama_kv_cache_iswa_context; class llama_memory_recurrent_context; class llama_memory_hybrid_context; @@ -36,6 +38,12 @@ enum llm_graph_type { LLM_GRAPH_TYPE_DECODER_MTP, }; +enum llm_fused_op { + LLM_FUSED_OP_FLASH_ATTN, + LLM_FUSED_OP_GDN_AR, + LLM_FUSED_OP_GDN_CH, +}; + enum llm_ffn_op_type : int { LLM_FFN_NONE = 0, // sentinel: unset; archs must assign before use LLM_FFN_SILU, @@ -459,6 +467,79 @@ public: const llama_kv_cache_iswa_context * mctx; }; +// DSV4 raw graph inputs are SWA-only, but their mask may be stream-shaped +// so raw K can be concatenated with DSV4 compressed K in one attention op. +class llm_graph_input_dsv4_raw { +public: + llm_graph_input_dsv4_raw( + const llama_cparams & cparams, + const llama_kv_cache_dsv4_raw_context * mctx) : + cparams(cparams), + mctx(mctx) { + } + + void set_input(const llama_ubatch * ubatch); + + ggml_tensor * get_k_idxs() const { return self_k_idxs; } + ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; } + + ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch] + + ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + + ggml_tensor * self_k_rot = nullptr; + + const llama_cparams cparams; + + const llama_kv_cache_dsv4_raw_context * mctx; +}; + +class llm_graph_input_dsv4 : public llm_graph_input_i { +public: + struct comp_input { + ggml_tensor * state_pos = nullptr; // I32 [n_state] + ggml_tensor * state_persist_src_idxs = nullptr; // I32 [n_state_persist] + ggml_tensor * state_persist_dst_idxs = nullptr; // I32 [n_state_persist] + ggml_tensor * state_read_idxs = nullptr; // I32 [ratio*n_state_write] + ggml_tensor * state_write_idxs = nullptr; // I64 [n_state_write] + ggml_tensor * state_write_pos = nullptr; // I32 [n_state_write] + + ggml_tensor * kq_mask = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream] + + ggml_tensor * k_rot = nullptr; + }; + + llm_graph_input_dsv4( + const llama_cparams & cparams, + std::unique_ptr inp_raw, + const llama_kv_cache_dsv4_context * mctx) : + inp_raw(std::move(inp_raw)), + cparams(cparams), + mctx(mctx) { + } + ~llm_graph_input_dsv4() = default; + + void set_input(const llama_ubatch * ubatch) override; + + bool can_reuse(const llm_graph_params & params) override; + + llm_graph_input_dsv4_raw * get_raw() const { return inp_raw.get(); } + const comp_input & get_csa() const { return inp_csa; } + const comp_input & get_hca() const { return inp_hca; } + const comp_input & get_lid() const { return inp_lid; } + + std::unique_ptr inp_raw; + + comp_input inp_csa; + comp_input inp_hca; + comp_input inp_lid; + + const llama_cparams cparams; + + const llama_kv_cache_dsv4_context * mctx; +}; + class llm_graph_input_attn_cross : public llm_graph_input_i { public: llm_graph_input_attn_cross(const llama_cross * cross) : cross(cross) {} @@ -700,6 +781,12 @@ struct llm_graph_params { } }; +struct llm_graph_fused_node { + llm_fused_op op; + ggml_tensor * tensor; + int il; +}; + class llm_graph_result { public: llm_graph_result(int64_t max_nodes); @@ -733,6 +820,10 @@ public: llm_graph_input_i * add_input(llm_graph_input_ptr input); + void add_fused_node(llm_graph_fused_node result); + + const std::vector & get_fused_nodes() const { return fused_nodes; } + void set_params(const llm_graph_params & params); // important graph nodes @@ -751,6 +842,7 @@ public: std::map t_sampled_probs; std::vector inputs; + std::vector fused_nodes; ggml_context_ptr ctx_compute; @@ -920,7 +1012,8 @@ struct llm_graph_context { ggml_tensor * gate_up_exps = nullptr, ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, - ggml_tensor * down_exps_s = nullptr) const; + ggml_tensor * down_exps_s = nullptr, + ggml_tensor * selected_experts_in = nullptr) const; ggml_tensor * build_moe_ffn( ggml_tensor * cur, @@ -945,7 +1038,8 @@ struct llm_graph_context { ggml_tensor * gate_up_exps_b = nullptr, ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, - ggml_tensor * down_exps_s = nullptr) const; + ggml_tensor * down_exps_s = nullptr, + ggml_tensor * selected_experts_in = nullptr) const; // // inputs @@ -1045,6 +1139,8 @@ struct llm_graph_context { llm_graph_input_attn_kv_iswa * build_attn_inp_kv_iswa() const; + llm_graph_input_dsv4 * build_inp_dsv4() const; + // note: if k_cur or v_cur are not provided, they will not be stored in the memory ggml_tensor * build_attn( llm_graph_input_attn_kv_iswa * inp, diff --git a/examples/talk-llama/llama-hparams.h b/examples/talk-llama/llama-hparams.h index 2eadeb214..8be5f28f3 100644 --- a/examples/talk-llama/llama-hparams.h +++ b/examples/talk-llama/llama-hparams.h @@ -14,6 +14,7 @@ enum llama_expert_gating_func_type { LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX = 1, LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID = 2, LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT = 3, // applied to the router weights instead of the logits + LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS = 4, }; enum llama_swa_type { @@ -226,6 +227,16 @@ struct llama_hparams { uint32_t indexer_head_size = 0; uint32_t indexer_top_k = 0; + // DeepSeek-V4 + uint32_t dsv4_o_group_count = 0; + uint32_t dsv4_o_lora_rank = 0; + uint32_t dsv4_hc_mult = 0; + uint32_t dsv4_hc_sinkhorn_iters = 0; + uint32_t dsv4_hash_layer_count = 0; + float dsv4_compress_rope_base = 0.0f; + float dsv4_hc_eps = 0.0f; + std::array dsv4_compress_ratios; + // qwen3vl deepstack // When parsed from GGUF, this implies the first N layers consume the first // N deepstack embeddings. Use deepstack_mapping_arr if you need a more diff --git a/examples/talk-llama/llama-impl.h b/examples/talk-llama/llama-impl.h index 7923c3f7e..4988b06d2 100644 --- a/examples/talk-llama/llama-impl.h +++ b/examples/talk-llama/llama-impl.h @@ -54,6 +54,26 @@ static inline dst_t llama_cast(src_t v) { } } +static inline ggml_tensor * llama_mul_mat_hadamard( + ggml_context * ctx, + ggml_tensor * cur, + ggml_tensor * rot) { + const auto n = rot->ne[0]; + + ggml_tensor * res; + + if (!ggml_is_contiguous(cur)) { + res = ggml_cont_2d(ctx, cur, n, ggml_nelements(cur)/n); + } else { + res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n); + } + res = ggml_mul_mat(ctx, rot, res); + ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD); + res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]); + + return res; +} + struct time_meas { time_meas(int64_t & t_acc, bool disable = false); ~time_meas(); @@ -83,7 +103,3 @@ std::string llama_format_tensor_shape(const std::vector & ne); std::string llama_format_tensor_shape(const struct ggml_tensor * t); std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i); - -#define LLAMA_TENSOR_NAME_FATTN "__fattn__" -#define LLAMA_TENSOR_NAME_FGDN_AR "__fgdn_ar__" -#define LLAMA_TENSOR_NAME_FGDN_CH "__fgdn_ch__" diff --git a/examples/talk-llama/llama-kv-cache-dsa.cpp b/examples/talk-llama/llama-kv-cache-dsa.cpp index 916ab6537..241c50365 100644 --- a/examples/talk-llama/llama-kv-cache-dsa.cpp +++ b/examples/talk-llama/llama-kv-cache-dsa.cpp @@ -113,7 +113,7 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_batch( std::vector ubatches; while (true) { - auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true); + auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0); if (ubatch.n_tokens == 0) { break; diff --git a/examples/talk-llama/llama-kv-cache-dsv4.cpp b/examples/talk-llama/llama-kv-cache-dsv4.cpp new file mode 100644 index 000000000..9fccf347e --- /dev/null +++ b/examples/talk-llama/llama-kv-cache-dsv4.cpp @@ -0,0 +1,1807 @@ +#include "llama-kv-cache-dsv4.h" + +#include "ggml-backend.h" +#include "llama-impl.h" +#include "llama-batch.h" +#include "llama-io.h" +#include "llama-model.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +static constexpr uint32_t DSV4_CSA_RATIO = 4; +static constexpr uint32_t DSV4_HCA_RATIO = 128; + +static constexpr uint32_t DSV4_STATE_MAGIC = 0x34565344; // DSV4 +static constexpr uint32_t DSV4_STATE_VERSION = 1; +static constexpr uint32_t DSV4_STATE_MODE_FULL = 0; +static constexpr uint32_t DSV4_STATE_MODE_PARTIAL = 1; +static constexpr uint32_t DSV4_K_CACHE_STATE_VER = 1; +static constexpr uint32_t DSV4_COMP_STATE_VER = 1; + +static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) { + return std::max(1, (kv_size + ratio - 1)/ratio); +} + +static int64_t dsv4_stream_offset(uint32_t n_stream, llama_seq_id seq_id, uint32_t size) { + if (n_stream <= 1) { + return 0; + } + if (seq_id < 0 || (uint32_t) seq_id >= n_stream) { + throw std::runtime_error("DSV4 sequence id out of stream range"); + } + + return (int64_t) seq_id*size; +} + +static bool dsv4_ubatch_has_coupled(const llama_ubatch & ubatch) { + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + if (ubatch.n_seq_id[i] > 1) { + return true; + } + } + + return false; +} + +static bool dsv4_token_has_seq(const llama_ubatch & ubatch, uint32_t i, llama_seq_id seq_id) { + for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) { + if (ubatch.seq_id[i][s] == seq_id) { + return true; + } + } + + return false; +} + +static llama_ubatch dsv4_build_raw_write_ubatch(const llama_ubatch & ubatch) { + if (!dsv4_ubatch_has_coupled(ubatch)) { + return ubatch; + } + if (ubatch.embd) { + throw std::runtime_error("DSV4 coupled embedding ubatches are not supported"); + } + + std::vector counts(ubatch.n_seqs_unq, 0); + uint32_t n_tokens = 0; + for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { + const llama_seq_id seq_id = ubatch.seq_id_unq[s]; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + if (dsv4_token_has_seq(ubatch, i, seq_id)) { + ++counts[s]; + ++n_tokens; + } + } + } + + if (n_tokens == 0) { + return ubatch; + } + + const uint32_t n_seq_tokens = counts[0]; + for (uint32_t s = 1; s < counts.size(); ++s) { + if (counts[s] != n_seq_tokens) { + throw std::runtime_error("DSV4 coupled raw writes require equal sequence lengths"); + } + } + + auto data = std::make_shared(); + data->pos.resize((size_t) n_tokens*ubatch.n_pos); + data->n_seq_id.reserve(n_tokens); + data->seq_id.reserve(n_tokens); + data->seq_id_data.reserve(n_tokens); + data->seq_id_unq.assign(ubatch.seq_id_unq, ubatch.seq_id_unq + ubatch.n_seqs_unq); + data->seq_idx.assign(LLAMA_MAX_SEQ, -1); + data->output.assign(n_tokens, 0); + if (ubatch.token) { + data->token.reserve(n_tokens); + } + + for (uint32_t s = 0; s < data->seq_id_unq.size(); ++s) { + data->seq_idx[data->seq_id_unq[s]] = s; + } + + for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { + const llama_seq_id seq_id = ubatch.seq_id_unq[s]; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + if (!dsv4_token_has_seq(ubatch, i, seq_id)) { + continue; + } + + const uint32_t dst = data->n_seq_id.size(); + if (ubatch.token) { + data->token.push_back(ubatch.token[i]); + } + for (uint32_t p = 0; p < ubatch.n_pos; ++p) { + data->pos[(size_t) p*n_tokens + dst] = ubatch.pos[(size_t) p*ubatch.n_tokens + i]; + } + data->n_seq_id.push_back(1); + data->seq_id_data.push_back(seq_id); + } + } + + for (uint32_t i = 0; i < n_tokens; ++i) { + data->seq_id.push_back(&data->seq_id_data[i]); + } + + llama_ubatch res { + /*.b_equal_seqs =*/ true, + /*.n_tokens =*/ n_tokens, + /*.n_seq_tokens =*/ n_seq_tokens, + /*.n_seqs =*/ ubatch.n_seqs_unq, + /*.n_seqs_unq =*/ ubatch.n_seqs_unq, + /*.n_pos =*/ ubatch.n_pos, + /*.token =*/ data->token.empty() ? nullptr : data->token.data(), + /*.embd =*/ nullptr, + /*.pos =*/ data->pos.data(), + /*.n_seq_id =*/ data->n_seq_id.data(), + /*.seq_id =*/ data->seq_id.data(), + /*.seq_id_unq =*/ data->seq_id_unq.data(), + /*.seq_idx =*/ data->seq_idx.data(), + /*.output =*/ data->output.data(), + /*.data =*/ data, + }; + + return res; +} + +static std::vector dsv4_build_raw_write_ubatches(const std::vector & ubatches) { + std::vector res; + res.reserve(ubatches.size()); + for (const llama_ubatch & ubatch : ubatches) { + res.push_back(dsv4_build_raw_write_ubatch(ubatch)); + } + return res; +} + +static bool dsv4_batch_has_coupled(const llama_batch & batch) { + if (!batch.n_seq_id) { + return false; + } + + for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.n_seq_id[i] > 1) { + return true; + } + } + + return false; +} + +static int64_t dsv4_comp_graph_n_stream(const llama_ubatch & ubatch, uint32_t n_stream) { + // Coupled sequence sets must stay in one graph stream because their + // compressed state is shared. Independent per-seq state can fan out. + if (n_stream <= 1 || ubatch.n_seqs_unq <= 1 || dsv4_ubatch_has_coupled(ubatch)) { + return 1; + } + + return ubatch.n_seqs_unq; +} + +static void dsv4_state_src_stream_range( + uint32_t n_stream, + llama_seq_id seq_id, + uint32_t & s0, + uint32_t & ns) { + if (seq_id >= 0 && n_stream > 1) { + if ((uint32_t) seq_id >= n_stream) { + throw std::runtime_error("DSV4 state sequence id out of stream range"); + } + + s0 = (uint32_t) seq_id; + ns = 1; + return; + } + + s0 = 0; + ns = seq_id >= 0 ? 1 : n_stream; +} + +static void dsv4_state_dst_stream_range( + uint32_t n_stream, + llama_seq_id seq_id, + uint32_t ns, + uint32_t & s0) { + if (seq_id >= 0) { + if (ns != 1) { + throw std::runtime_error("DSV4 sequence state stream count mismatch"); + } + if (n_stream > 1 && (uint32_t) seq_id >= n_stream) { + throw std::runtime_error("DSV4 state sequence id out of stream range"); + } + + s0 = n_stream > 1 ? (uint32_t) seq_id : 0; + return; + } + + if (ns != n_stream) { + throw std::runtime_error("DSV4 full state stream count mismatch"); + } + + s0 = 0; +} + +static void dsv4_state_write_tensor_streams( + llama_io_write_i & io, + ggml_tensor * tensor, + uint32_t n_rows, + uint32_t s0, + uint32_t ns) { + const int32_t type_i = (int32_t) tensor->type; + const uint64_t ne0 = tensor->ne[0]; + const uint64_t rows = n_rows; + const uint64_t row_size = ggml_row_size(tensor->type, tensor->ne[0]); + + io.write(&type_i, sizeof(type_i)); + io.write(&ne0, sizeof(ne0)); + io.write(&rows, sizeof(rows)); + io.write(&row_size, sizeof(row_size)); + + const size_t offset = (size_t) s0*n_rows*row_size; + const size_t size = (size_t) ns*n_rows*row_size; + + io.write_tensor(tensor, offset, size); +} + +static void dsv4_state_read_tensor_streams( + llama_io_read_i & io, + ggml_tensor * tensor, + uint32_t n_rows, + uint32_t s0, + uint32_t ns) { + int32_t type_i_ref; + uint64_t ne0_ref; + uint64_t rows_ref; + uint64_t row_size_ref; + + io.read(&type_i_ref, sizeof(type_i_ref)); + io.read(&ne0_ref, sizeof(ne0_ref)); + io.read(&rows_ref, sizeof(rows_ref)); + io.read(&row_size_ref, sizeof(row_size_ref)); + + const int32_t type_i = (int32_t) tensor->type; + const uint64_t ne0 = tensor->ne[0]; + const uint64_t rows = n_rows; + const uint64_t row_size = ggml_row_size(tensor->type, tensor->ne[0]); + + if (type_i != type_i_ref || ne0 != ne0_ref || rows != rows_ref || row_size != row_size_ref) { + throw std::runtime_error("DSV4 state tensor metadata mismatch"); + } + + const size_t offset = (size_t) s0*n_rows*row_size; + const size_t size = (size_t) ns*n_rows*row_size; + + io.read_tensor(tensor, offset, size); +} + +static void dsv4_state_write_k_cache( + llama_io_write_i & io, + const llama_kv_cache * kv, + llama_seq_id seq_id, + llama_state_seq_flags flags) { + GGML_UNUSED(flags); + + uint32_t s0; + uint32_t ns; + dsv4_state_src_stream_range(kv->get_n_stream(), seq_id, s0, ns); + + const uint32_t version = DSV4_K_CACHE_STATE_VER; + const uint32_t kv_size = kv->get_size(); + const auto layer_ids = kv->get_layer_ids(); + const uint32_t n_layer = layer_ids.size(); + + io.write(&version, sizeof(version)); + io.write(&kv_size, sizeof(kv_size)); + io.write(&ns, sizeof(ns)); + io.write(&n_layer, sizeof(n_layer)); + + for (uint32_t il : layer_ids) { + io.write(&il, sizeof(il)); + dsv4_state_write_tensor_streams(io, kv->get_k_storage(il), kv_size, s0, ns); + } +} + +static void dsv4_state_read_k_cache( + llama_io_read_i & io, + llama_kv_cache * kv, + llama_seq_id seq_id, + llama_state_seq_flags flags) { + GGML_UNUSED(flags); + + uint32_t version; + uint32_t kv_size_ref; + uint32_t ns; + uint32_t n_layer_ref; + + io.read(&version, sizeof(version)); + io.read(&kv_size_ref, sizeof(kv_size_ref)); + io.read(&ns, sizeof(ns)); + io.read(&n_layer_ref, sizeof(n_layer_ref)); + + if (version != DSV4_K_CACHE_STATE_VER) { + throw std::runtime_error("DSV4 K-cache state version mismatch"); + } + if (kv_size_ref != kv->get_size()) { + throw std::runtime_error("DSV4 K-cache state size mismatch"); + } + + uint32_t s0; + dsv4_state_dst_stream_range(kv->get_n_stream(), seq_id, ns, s0); + + const auto layer_ids = kv->get_layer_ids(); + if (n_layer_ref != layer_ids.size()) { + throw std::runtime_error("DSV4 K-cache layer count mismatch"); + } + + for (uint32_t il : layer_ids) { + uint32_t il_ref; + io.read(&il_ref, sizeof(il_ref)); + if (il_ref != il) { + throw std::runtime_error("DSV4 K-cache layer id mismatch"); + } + + dsv4_state_read_tensor_streams(io, kv->get_k_storage(il), kv->get_size(), s0, ns); + } +} + +static std::string dsv4_plan_positions(const std::vector & values) { + std::ostringstream ss; + ss << "["; + for (size_t i = 0; i < values.size(); ++i) { + if (i > 0) { + ss << ", "; + } + ss << values[i]; + } + ss << "]"; + return ss.str(); +} + +static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( + const llama_ubatch & ubatch, + uint32_t ratio, + bool overlap, + uint32_t state_size, + uint32_t kv_size, + uint32_t n_stream) { + llama_kv_cache_dsv4_context::comp_plan plan; + plan.n_visible.resize(ubatch.n_tokens); + plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream); + + // n_stream is the persistent cache/state layout; plan.n_stream is the + // graph view for this ubatch and can be a subset of those streams. + if (n_stream <= 1 && ubatch.n_seqs_unq > 1) { + throw std::runtime_error("DSV4 single compressed stream cannot serve multiple sequences"); + } + + const int64_t state_rows = (int64_t) state_size*n_stream; + + struct persist_row { + int32_t dst; + int32_t src; + llama_pos pos; + }; + + std::vector persist_rows; + + // For the overlap compressor, build_overlap_compressed_kv_from_state() consumes + // state_read_idxs as two contiguous halves: the first ratio*n_blocks entries are + // the "previous-window" gather indices for every block, followed by the + // "current-window" indices for every block. Collect them separately here and + // append cur after prev once the loop has visited all completed blocks + std::vector overlap_prev_reads; + std::vector overlap_cur_reads; + + std::map, int64_t> curr_token_idx_map; + + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) { + curr_token_idx_map[std::make_pair(ubatch.seq_id[i][s], ubatch.pos[i])] = i; + } + } + + const auto state_source_idx = [&](llama_seq_id seq_id, llama_pos pos) -> int32_t { + if (pos < 0) { + // The overlap compressor needs a zero/-inf source for the first + // block's previous half. The graph appends that row after the + // current-ubatch scratch rows. + return (int32_t) (state_rows + ubatch.n_tokens); + } + + const auto key = std::make_pair(seq_id, pos); + if (curr_token_idx_map.find(key) != curr_token_idx_map.end()) { + return (int32_t) (state_rows + curr_token_idx_map.at(key)); + } + + const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size); + return (int32_t) (stream_off + pos%state_size); + }; + + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + const llama_pos pos = ubatch.pos[i]; + + if (pos < 0) { + continue; + } + + plan.state_pos.push_back((int32_t) (pos%ratio)); + + const int64_t n_visible = (int64_t) (pos + 1)/ratio; + plan.n_visible[i] = (int32_t) n_visible; + plan.n_kv = std::max(plan.n_kv, n_visible); + + for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) { + const llama_seq_id seq_id = ubatch.seq_id[i][s]; + const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size); + const int32_t state_idx = (int32_t) (stream_off + pos%state_size); + + const auto it = std::find_if(persist_rows.begin(), persist_rows.end(), + [state_idx](const persist_row & row) { + return row.dst == state_idx; + }); + if (it == persist_rows.end()) { + persist_rows.push_back({ state_idx, (int32_t) i, pos }); + } else if (pos > it->pos) { + it->src = (int32_t) i; + it->pos = pos; + } + + if ((pos + 1) % ratio != 0) { + continue; + } + + const llama_pos source_start = pos + 1 - ratio; + const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size); + + plan.state_write_idxs.push_back(cache_off + pos/ratio); + plan.state_write_pos.push_back((int32_t) source_start); + + if (overlap) { + const llama_pos prev_start = source_start - ratio; + + for (uint32_t j = 0; j < ratio; ++j) { + overlap_prev_reads.push_back(state_source_idx(seq_id, prev_start + j)); + } + for (uint32_t j = 0; j < ratio; ++j) { + overlap_cur_reads.push_back(state_source_idx(seq_id, source_start + j)); + } + } else { + for (uint32_t j = 0; j < ratio; ++j) { + plan.state_read_idxs.push_back(state_source_idx(seq_id, source_start + j)); + } + } + } + } + + if (ratio == DSV4_CSA_RATIO && plan.state_write_idxs.empty() && !plan.state_pos.empty()) { + // Non-boundary CSA steps still need a write op so their graph matches + // boundary steps. Use a padded scratch row that is masked from attention. + assert(kv_size > 0); + + uint32_t i = 0; + while (i < ubatch.n_tokens && ubatch.pos[i] < 0) { + ++i; + } + assert(i < ubatch.n_tokens); + + const llama_pos pos = ubatch.pos[i]; + const llama_seq_id seq_id = ubatch.seq_id[i][0]; + const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size); + const int32_t source_idx = state_source_idx(seq_id, pos); + + plan.state_write_idxs.push_back(cache_off + kv_size - 1); + plan.state_write_pos .push_back(0); + + if (overlap) { + for (uint32_t j = 0; j < ratio; ++j) { + overlap_prev_reads.push_back(source_idx); + overlap_cur_reads .push_back(source_idx); + } + } else { + for (uint32_t j = 0; j < ratio; ++j) { + plan.state_read_idxs.push_back(source_idx); + } + } + } + + if (overlap) { + // [ all blocks' prev-window indices | all blocks' cur-window indices ] + plan.state_read_idxs.reserve(overlap_prev_reads.size() + overlap_cur_reads.size()); + plan.state_read_idxs.insert(plan.state_read_idxs.end(), + overlap_prev_reads.begin(), overlap_prev_reads.end()); + plan.state_read_idxs.insert(plan.state_read_idxs.end(), + overlap_cur_reads.begin(), overlap_cur_reads.end()); + } + + plan.n_kv = GGML_PAD(plan.n_kv, 256u); + + std::sort(persist_rows.begin(), persist_rows.end(), + [](const persist_row & a, const persist_row & b) { + return a.dst < b.dst; + }); + + for (const persist_row & row : persist_rows) { + plan.state_persist_src_idxs.push_back(row.src); + plan.state_persist_dst_idxs.push_back(row.dst); + } + + static const bool debug = []() { + const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG"); + return env && atoi(env) > 0; + }(); + + if (debug) { + LLAMA_LOG_INFO("%s: ratio=%u, n_tokens=%u, state_persist_dst=%s, state_write_pos=%s\n", + __func__, ratio, ubatch.n_tokens, + dsv4_plan_positions(plan.state_persist_dst_idxs).c_str(), + dsv4_plan_positions(plan.state_write_pos).c_str()); + } + + return plan; +} + +static std::vector dsv4_build_comp_plans( + const std::vector & ubatches, + uint32_t ratio, + bool overlap, + uint32_t state_size, + uint32_t kv_size, + uint32_t n_stream) { + std::vector plans; + plans.reserve(ubatches.size()); + + for (const llama_ubatch & ubatch : ubatches) { + plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream)); + } + + return plans; +} + +static llama_kv_cache::slot_info_vec_t dsv4_build_comp_sinfos( + const std::vector & ubatches, + uint32_t n_stream) { + llama_kv_cache::slot_info_vec_t sinfos; + sinfos.reserve(ubatches.size()); + + for (const llama_ubatch & ubatch : ubatches) { + if (n_stream <= 1 && ubatch.n_seqs_unq > 1) { + throw std::runtime_error("DSV4 single compressed stream cannot serve multiple sequences"); + } + + const uint32_t ns = (uint32_t) dsv4_comp_graph_n_stream(ubatch, n_stream); + llama_kv_cache::slot_info sinfo; + sinfo.s0 = n_stream > 1 ? LLAMA_MAX_SEQ : 0; + sinfo.s1 = 0; + sinfo.resize(ns); + + for (uint32_t s = 0; s < ns; ++s) { + const llama_seq_id seq_id = n_stream > 1 ? ubatch.seq_id_unq[s] : 0; + const uint32_t strm = (uint32_t) dsv4_stream_offset(n_stream, seq_id, 1); + + sinfo.s0 = std::min(sinfo.s0, strm); + sinfo.s1 = std::max(sinfo.s1, strm); + sinfo.strm[s] = strm; + sinfo.idxs[s].resize(1, 0); + } + + if (n_stream > 1 && sinfo.s1 - sinfo.s0 + 1 != ns) { + throw std::runtime_error("DSV4 compressed streams are not contiguous in ubatch"); + } + + sinfos.push_back(std::move(sinfo)); + } + + return sinfos; +} + +static llama_kv_cache::slot_info_vec_t dsv4_build_raw_read_sinfos( + const llama_kv_cache::slot_info_vec_t & sinfos_write, + const std::vector & ubatches) { + llama_kv_cache::slot_info_vec_t sinfos; + sinfos.reserve(ubatches.size()); + + for (size_t i = 0; i < ubatches.size(); ++i) { + const llama_ubatch & ubatch = ubatches[i]; + const auto & sinfo_write = sinfos_write[i]; + + if (!dsv4_ubatch_has_coupled(ubatch)) { + sinfos.push_back(sinfo_write); + continue; + } + + const llama_seq_id seq_id = ubatch.seq_id[0][0]; + uint32_t i_stream = 0; + for (; i_stream < sinfo_write.n_stream(); ++i_stream) { + if (sinfo_write.strm[i_stream] == seq_id) { + break; + } + } + if (i_stream == sinfo_write.n_stream()) { + throw std::runtime_error("DSV4 raw write stream not found for coupled read"); + } + + llama_kv_cache::slot_info sinfo; + sinfo.s0 = sinfo_write.strm[i_stream]; + sinfo.s1 = sinfo_write.strm[i_stream]; + sinfo.resize(1); + sinfo.strm[0] = sinfo_write.strm[i_stream]; + sinfo.idxs[0] = sinfo_write.idxs[i_stream]; + sinfos.push_back(std::move(sinfo)); + } + + return sinfos; +} + +static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan( + const llama_ubatch & ubatch, + uint32_t ratio, + bool overlap, + uint32_t state_size, + uint32_t kv_size, + uint32_t n_stream) { + llama_kv_cache_dsv4_context::comp_plan plan; + plan.n_visible.resize(ubatch.n_tokens); + plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream); + plan.n_kv = kv_size; + + if (ubatch.n_tokens == 0) { + return plan; + } + + const uint32_t n_seqs = std::max(1, ubatch.n_seqs); + const uint32_t n_seq_tokens = std::max(1, ubatch.n_seq_tokens); + const uint64_t n_blocks_u64 = (uint64_t) n_seqs*((n_seq_tokens + ratio - 1)/ratio); + const size_t n_blocks = (size_t) std::max(1, n_blocks_u64); + GGML_ASSERT((uint64_t) n_blocks == std::max(1, n_blocks_u64)); + + const uint64_t state_rows = (uint64_t) state_size*n_stream; + const size_t n_persist = (size_t) std::min(ubatch.n_tokens, state_rows); + + plan.state_pos .resize(ubatch.n_tokens); + plan.state_persist_src_idxs.resize(n_persist); + plan.state_persist_dst_idxs.resize(n_persist); + plan.state_read_idxs .resize((overlap ? 2u : 1u)*ratio*n_blocks); + plan.state_write_idxs.resize(n_blocks); + plan.state_write_pos .resize(n_blocks); + + return plan; +} + +static void dsv4_make_k_only(llama_hparams & hparams) { + // llama_kv_cache uses hparams.is_mla() to allocate K-only storage. + hparams.n_embd_head_k_mla_impl = hparams.n_embd_head_k(); + hparams.n_embd_head_v_mla_impl = hparams.n_embd_head_k(); +} + +// +// llama_dsv4_comp_state +// + +llama_dsv4_comp_state::llama_dsv4_comp_state( + const llama_model & model, + bool offload, + bool unified, + uint32_t n_seq_max, + uint32_t ratio, + uint32_t state_size, + uint32_t n_embd_state, + const char * name, + const llama_memory_i::layer_filter_cb & filter) : + ratio(ratio), + state_size(state_size), + n_embd_state(n_embd_state), + n_stream(unified ? 1 : n_seq_max) { + const llama_hparams & hparams = model.hparams; + + struct ggml_backend_buft_comparator { + bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const { + return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0; + } + }; + + std::map ctx_map; + + auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * { + auto it = ctx_map.find(buft); + if (it == ctx_map.end()) { + ggml_init_params params = { + /*.mem_size =*/ size_t(2u*hparams.n_layer()*ggml_tensor_overhead()), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + + ggml_context * ctx = ggml_init(params); + if (!ctx) { + return nullptr; + } + + ctx_map.emplace(buft, ctx); + + return ctx; + } + + return it->second.get(); + }; + + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + if (filter && !filter(il)) { + continue; + } + + const char * dev_name = "CPU"; + + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + + if (offload) { + auto * dev = model.dev_layer(il); + buft = ggml_backend_dev_buffer_type(dev); + + dev_name = ggml_backend_dev_name(dev); + } + + LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name); + + ggml_context * ctx = ctx_for_buft(buft); + if (!ctx) { + throw std::runtime_error("failed to create ggml context for DSV4 compressor state"); + } + + ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream); + ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream); + + ggml_format_name(kv, "dsv4_%s_state_kv_l%d", name, il); + ggml_format_name(score, "dsv4_%s_state_score_l%d", name, il); + + map_layer_ids[il] = layers.size(); + + layers.push_back({ il, kv, score }); + } + + for (auto & [buft, ctx] : ctx_map) { + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + if (!buf) { + throw std::runtime_error("failed to allocate buffer for DSV4 compressor state"); + } + + ggml_backend_buffer_clear(buf, 0); + + LLAMA_LOG_INFO("%s: %10s DSV4 %s state buffer size = %8.2f MiB\n", + __func__, ggml_backend_buffer_name(buf), name, ggml_backend_buffer_get_size(buf)/1024.0/1024.0); + + ctxs_bufs.emplace_back(std::move(ctx), buf); + } + + LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, layers = %zu, size = %7.2f MiB\n", + __func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0); +} + +void llama_dsv4_comp_state::clear(bool data) { + if (!data) { + return; + } + + for (auto & [_, buf] : ctxs_bufs) { + ggml_backend_buffer_clear(buf.get(), 0); + } +} + +uint32_t llama_dsv4_comp_state::get_ratio() const { + return ratio; +} + +uint32_t llama_dsv4_comp_state::get_state_size() const { + return state_size; +} + +uint32_t llama_dsv4_comp_state::get_n_stream() const { + return n_stream; +} + +std::map llama_dsv4_comp_state::memory_breakdown() const { + std::map ret; + for (const auto & [_, buf] : ctxs_bufs) { + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get()); + ret[buft] += ggml_backend_buffer_get_size(buf.get()); + } + return ret; +} + +void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { + GGML_UNUSED(flags); + + uint32_t s0; + uint32_t ns; + dsv4_state_src_stream_range(n_stream, seq_id, s0, ns); + + const uint32_t version = DSV4_COMP_STATE_VER; + const uint32_t n_layer = layers.size(); + + io.write(&version, sizeof(version)); + io.write(&ratio, sizeof(ratio)); + io.write(&state_size, sizeof(state_size)); + io.write(&n_embd_state, sizeof(n_embd_state)); + io.write(&ns, sizeof(ns)); + io.write(&n_layer, sizeof(n_layer)); + + for (const auto & layer : layers) { + io.write(&layer.il, sizeof(layer.il)); + + dsv4_state_write_tensor_streams(io, layer.kv, state_size, s0, ns); + dsv4_state_write_tensor_streams(io, layer.score, state_size, s0, ns); + } +} + +void llama_dsv4_comp_state::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + GGML_UNUSED(flags); + + uint32_t version; + uint32_t ratio_ref; + uint32_t state_size_ref; + uint32_t n_embd_state_ref; + uint32_t ns; + uint32_t n_layer_ref; + + io.read(&version, sizeof(version)); + io.read(&ratio_ref, sizeof(ratio_ref)); + io.read(&state_size_ref, sizeof(state_size_ref)); + io.read(&n_embd_state_ref, sizeof(n_embd_state_ref)); + io.read(&ns, sizeof(ns)); + io.read(&n_layer_ref, sizeof(n_layer_ref)); + + if (version != DSV4_COMP_STATE_VER) { + throw std::runtime_error("DSV4 compressor state version mismatch"); + } + if (ratio_ref != ratio || state_size_ref != state_size || n_embd_state_ref != n_embd_state) { + throw std::runtime_error("DSV4 compressor state metadata mismatch"); + } + if (n_layer_ref != layers.size()) { + throw std::runtime_error("DSV4 compressor state layer count mismatch"); + } + + uint32_t s0; + dsv4_state_dst_stream_range(n_stream, seq_id, ns, s0); + + for (const auto & layer : layers) { + uint32_t il_ref; + io.read(&il_ref, sizeof(il_ref)); + if (il_ref != layer.il) { + throw std::runtime_error("DSV4 compressor state layer id mismatch"); + } + + dsv4_state_read_tensor_streams(io, layer.kv, state_size, s0, ns); + dsv4_state_read_tensor_streams(io, layer.score, state_size, s0, ns); + } +} + +ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const { + const int32_t ids = map_layer_ids.at(il); + + ggml_tensor * state = layers[ids].kv; + + return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]); +} + +ggml_tensor * llama_dsv4_comp_state::get_score(ggml_context * ctx, int32_t il) const { + const int32_t ids = map_layer_ids.at(il); + + ggml_tensor * state = layers[ids].score; + + return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]); +} + +ggml_tensor * llama_dsv4_comp_state::cpy_kv(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const { + return ggml_set_rows(ctx, get_kv(ctx, il), cur, idxs); +} + +ggml_tensor * llama_dsv4_comp_state::cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const { + return ggml_set_rows(ctx, get_score(ctx, il), cur, idxs); +} + +size_t llama_dsv4_comp_state::total_size() const { + size_t size = 0; + + for (const auto & [_, buf] : ctxs_bufs) { + size += ggml_backend_buffer_get_size(buf.get()); + } + + return size; +} + +// +// llama_kv_cache_dsv4 +// + +llama_kv_cache_dsv4::llama_kv_cache_dsv4( + const llama_model & model, + ggml_type type_k, + ggml_type type_v, + bool v_trans, + bool offload, + bool swa_full, + bool unified, + uint32_t kv_size, + uint32_t n_seq_max, + uint32_t n_ubatch, + uint32_t n_pad, + const layer_filter_cb & filter, + const layer_reuse_cb & reuse) : + hparams_raw(model.hparams), + hparams_csa(model.hparams), + hparams_hca(model.hparams), + hparams_lid(model.hparams), + n_seq_max(n_seq_max) { + + const layer_filter_cb filter_raw = [&](int32_t il) { + if (filter && !filter(il)) { + return false; + } + + return true; + }; + + GGML_UNUSED(unified); + + // Keep DSV4 KV/state streams per sequence even when public KV mode is unified. + const bool unified_raw = false; + + LLAMA_LOG_INFO("%s: creating DSV4 raw KV cache\n", __func__); + + dsv4_make_k_only(hparams_raw); + + kv_raw = std::make_unique( + model, hparams_raw, type_k, type_v, + v_trans, offload, swa_full, unified_raw, kv_size, n_seq_max, n_ubatch, n_pad, + nullptr, filter_raw, reuse, nullptr); + + dsv4_make_k_only(hparams_csa); + dsv4_make_k_only(hparams_hca); + + std::fill(hparams_lid.n_head_kv_arr.begin(), hparams_lid.n_head_kv_arr.end(), 1); + hparams_lid.n_embd_head_k_full = model.hparams.indexer_head_size; + hparams_lid.n_embd_head_v_full = model.hparams.indexer_head_size; + hparams_lid.n_embd_head_k_swa = model.hparams.indexer_head_size; + hparams_lid.n_embd_head_v_swa = model.hparams.indexer_head_size; + hparams_lid.rope_type = LLAMA_ROPE_TYPE_NEOX; + dsv4_make_k_only(hparams_lid); + + const layer_filter_cb filter_csa = [&](int32_t il) { + if (filter && !filter(il)) { + return false; + } + + return model.hparams.dsv4_compress_ratios[il] == DSV4_CSA_RATIO; + }; + + const layer_filter_cb filter_hca = [&](int32_t il) { + if (filter && !filter(il)) { + return false; + } + + return model.hparams.dsv4_compress_ratios[il] == DSV4_HCA_RATIO; + }; + + const bool unified_compressed = false; + + LLAMA_LOG_INFO("%s: creating DSV4 CSA compressed KV cache, size = %u cells\n", + __func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO)); + + kv_csa = std::make_unique( + model, hparams_csa, type_k, type_v, + v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad, + 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr); + + LLAMA_LOG_INFO("%s: creating DSV4 HCA compressed KV cache, size = %u cells\n", + __func__, dsv4_comp_size(kv_size, DSV4_HCA_RATIO)); + + kv_hca = std::make_unique( + model, hparams_hca, type_k, type_v, + v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_HCA_RATIO), 256u), n_seq_max, n_pad, + 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_hca, nullptr, nullptr); + + LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer KV cache, size = %u cells\n", + __func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO)); + + kv_lid = std::make_unique( + model, hparams_lid, type_k, type_v, + v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad, + 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr); + + LLAMA_LOG_INFO("%s: creating DSV4 CSA compressor state\n", __func__); + + csa_state = std::make_unique( + model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, + 2*model.hparams.n_embd_head_k(), "csa", filter_csa); + + LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__); + + hca_state = std::make_unique( + model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO, + model.hparams.n_embd_head_k(), "hca", filter_hca); + + LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer compressor state\n", __func__); + + lid_state = std::make_unique( + model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, + 2*model.hparams.indexer_head_size, "lid", filter_csa); + + // DSV4 attention reads compressed-K / compressor-state rows that the current + // graph does not necessarily overwrite; uninitialized buffer contents would + // otherwise leak in (instance-specific garbage) and corrupt recall. Zero all + // compressed buffers up front so reads of un-written rows are deterministic. + clear_compressed(true); +} + +llama_memory_context_ptr llama_kv_cache_dsv4::init_batch( + llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool embd_all) { + GGML_UNUSED(embd_all); + + const bool raw_per_seq = kv_raw->get_base()->get_n_stream() != 1; + const bool comp_per_seq = csa_state->get_n_stream() > 1; + const bool has_coupled = dsv4_batch_has_coupled(balloc.get_batch()); + + const auto make_context = [&](std::vector ubatches) -> llama_memory_context_ptr { + auto ubatches_raw = dsv4_build_raw_write_ubatches(ubatches); + + auto sinfos_raw_base_write = kv_raw->get_base()->prepare(ubatches_raw); + if (sinfos_raw_base_write.empty()) { + return nullptr; + } + + auto sinfos_raw_swa_write = kv_raw->get_swa()->prepare(ubatches_raw); + if (sinfos_raw_swa_write.empty()) { + return nullptr; + } + + auto sinfos_raw_swa_read = dsv4_build_raw_read_sinfos(sinfos_raw_swa_write, ubatches); + + return std::make_unique( + this, + std::move(sinfos_raw_base_write), + std::move(sinfos_raw_swa_write), + std::move(sinfos_raw_swa_read), + std::move(ubatches), + std::move(ubatches_raw)); + }; + + // Match llama_kv_cache_iswa splitting when DSV4 compressed state does not + // require per-sequence graph layout. + do { + if (raw_per_seq || comp_per_seq) { + break; + } + + balloc.split_reset(); + + std::vector ubatches; + while (true) { + auto ubatch = balloc.split_simple(n_ubatch); + if (ubatch.n_tokens == 0) { + break; + } + ubatches.push_back(std::move(ubatch)); // NOLINT + } + + if (balloc.get_n_used() < balloc.get_n_tokens()) { + break; + } + + if (auto ctx = make_context(std::move(ubatches))) { + return ctx; + } + } while (false); + + // When raw or compressed state is per-sequence, independent sequences can + // share an equal-length ubatch. Coupled sequence sets still serialize until + // DSV4 has explicit shared-state handling for compressed streams. + do { + balloc.split_reset(); + + std::vector ubatches; + while (true) { + llama_ubatch ubatch; + if (has_coupled) { + ubatch = balloc.split_seq(n_ubatch); + } else { + ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq, 0); + } + + if (ubatch.n_tokens == 0) { + break; + } + ubatches.push_back(std::move(ubatch)); // NOLINT + } + + if (balloc.get_n_used() < balloc.get_n_tokens()) { + break; + } + + if (auto ctx = make_context(std::move(ubatches))) { + return ctx; + } + } while (false); + + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); +} + +llama_memory_context_ptr llama_kv_cache_dsv4::init_full() { + return std::make_unique(this); +} + +llama_memory_context_ptr llama_kv_cache_dsv4::init_update(llama_context * lctx, bool optimize) { + return std::make_unique(this, lctx, optimize); +} + +bool llama_kv_cache_dsv4::get_can_shift() const { + // Compressed row metadata uses block-derived positions. Keep shifting + // disabled until DSV4 compressed-cache shift semantics are wired. + return false; +} + +void llama_kv_cache_dsv4::clear(bool data) { + kv_raw->clear(data); + clear_compressed(true); // DSV4 compressed buffers must never expose stale/uninit rows +} + +bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { + if (p1 >= 0) { + return false; + } + + if (p0 > 0) { + // DSV4 compressed cache rows are derived from running compressor state, + // so arbitrary rollback is not reconstructible from the raw cache alone. + // Allow the common prompt-cache cleanup no-op: remove [end, infinity). + if (seq_id >= 0 && p0 > kv_raw->seq_pos_max(seq_id)) { + return true; + } + + return false; + } + + const bool res = kv_raw->seq_rm(seq_id, p0, p1); + + if (res) { + clear_compressed(true); + } + + return res; +} + +void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + kv_raw->seq_cp(seq_id_src, seq_id_dst, p0, p1); + clear_compressed(true); +} + +void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) { + kv_raw->seq_keep(seq_id); + clear_compressed(true); +} + +void llama_kv_cache_dsv4::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + kv_raw->seq_add(seq_id, p0, p1, shift); + clear_compressed(true); +} + +void llama_kv_cache_dsv4::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + kv_raw->seq_div(seq_id, p0, p1, d); + clear_compressed(true); +} + +llama_pos llama_kv_cache_dsv4::seq_pos_min(llama_seq_id seq_id) const { + if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) { + return -1; + } + + // The raw SWA cache may contain a wider window, but the compressed DSV4 + // state cannot be rolled back within that window. Report only the current + // boundary so server-context uses checkpoints for rollback. + return kv_raw->seq_pos_max(seq_id); +} + +llama_pos llama_kv_cache_dsv4::seq_pos_max(llama_seq_id seq_id) const { + if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) { + return -1; + } + + return kv_raw->seq_pos_max(seq_id); +} + +std::map llama_kv_cache_dsv4::memory_breakdown() const { + std::map mb = kv_raw->memory_breakdown(); + for (const auto & buft_size : kv_csa->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + for (const auto & buft_size : kv_hca->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + for (const auto & buft_size : kv_lid->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + for (const auto & buft_size : csa_state->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + for (const auto & buft_size : hca_state->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + for (const auto & buft_size : lid_state->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + return mb; +} + +void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { + const bool partial_only = flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; + + const uint32_t magic = DSV4_STATE_MAGIC; + const uint32_t version = DSV4_STATE_VERSION; + const uint32_t mode = partial_only ? DSV4_STATE_MODE_PARTIAL : DSV4_STATE_MODE_FULL; + + io.write(&magic, sizeof(magic)); + io.write(&version, sizeof(version)); + io.write(&mode, sizeof(mode)); + + kv_raw->state_write(io, seq_id, flags); + + if (!partial_only) { + dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags); + dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags); + dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags); + } + + csa_state->state_write(io, seq_id, flags); + hca_state->state_write(io, seq_id, flags); + lid_state->state_write(io, seq_id, flags); +} + +void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + uint32_t magic; + uint32_t version; + uint32_t mode = DSV4_STATE_MODE_FULL; + + io.read(&magic, sizeof(magic)); + io.read(&version, sizeof(version)); + + if (magic != DSV4_STATE_MAGIC) { + throw std::runtime_error("DSV4 state magic mismatch"); + } + if (version != DSV4_STATE_VERSION) { + throw std::runtime_error("DSV4 state version mismatch"); + } + + io.read(&mode, sizeof(mode)); + if (mode != DSV4_STATE_MODE_FULL && mode != DSV4_STATE_MODE_PARTIAL) { + throw std::runtime_error("DSV4 state mode mismatch"); + } + + const bool partial_only = mode == DSV4_STATE_MODE_PARTIAL; + if (partial_only != !!(flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)) { + throw std::runtime_error("DSV4 state flags mismatch"); + } + + kv_raw->state_read(io, seq_id, flags); + + if (!partial_only) { + dsv4_state_read_k_cache(io, kv_csa.get(), seq_id, flags); + dsv4_state_read_k_cache(io, kv_hca.get(), seq_id, flags); + dsv4_state_read_k_cache(io, kv_lid.get(), seq_id, flags); + } + + csa_state->state_read(io, seq_id, flags); + hca_state->state_read(io, seq_id, flags); + lid_state->state_read(io, seq_id, flags); + +} + +llama_kv_cache_iswa * llama_kv_cache_dsv4::get_raw() const { + return kv_raw.get(); +} + +llama_kv_cache * llama_kv_cache_dsv4::get_csa() const { + return kv_csa.get(); +} + +llama_kv_cache * llama_kv_cache_dsv4::get_hca() const { + return kv_hca.get(); +} + +llama_kv_cache * llama_kv_cache_dsv4::get_lid() const { + return kv_lid.get(); +} + +llama_dsv4_comp_state * llama_kv_cache_dsv4::get_csa_state() const { + return csa_state.get(); +} + +llama_dsv4_comp_state * llama_kv_cache_dsv4::get_hca_state() const { + return hca_state.get(); +} + +llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const { + return lid_state.get(); +} + +void llama_kv_cache_dsv4::clear_compressed(bool data) { + kv_csa->clear(data); + kv_hca->clear(data); + kv_lid->clear(data); + csa_state->clear(data); + hca_state->clear(data); + lid_state->clear(data); +} + +// +// llama_kv_cache_dsv4_raw_context +// + +static llama_kv_cache::slot_info dsv4_build_full_sinfo(const llama_kv_cache * kv) { + const uint32_t n_stream = kv->get_n_stream(); + + llama_kv_cache::slot_info sinfo; + sinfo.s0 = 0; + sinfo.s1 = n_stream - 1; + sinfo.resize(n_stream); + for (uint32_t s = 0; s < n_stream; ++s) { + sinfo.strm[s] = s; + sinfo.idxs[s].resize(1, 0); + } + + return sinfo; +} + +llama_kv_cache_dsv4_raw_context::llama_kv_cache_dsv4_raw_context(llama_kv_cache_iswa * kv) : + kv_swa(kv->get_swa()), + ctx_base_mem(nullptr), + ctx_swa_mem(nullptr), + n_kv(kv_swa->get_size()), + status(LLAMA_MEMORY_STATUS_SUCCESS) { + sinfos_read.push_back(dsv4_build_full_sinfo(kv_swa)); + sinfos_write = sinfos_read; +} + +llama_kv_cache_dsv4_raw_context::llama_kv_cache_dsv4_raw_context( + llama_kv_cache_iswa * kv, + llama_context * lctx, + bool optimize) : + kv_swa(kv->get_swa()), + ctx_base_mem(kv->get_base()->init_update(lctx, optimize)), + ctx_swa_mem(kv->get_swa()->init_update(lctx, optimize)), + n_kv(kv_swa->get_size()), + status(llama_memory_status_combine(ctx_base_mem->get_status(), ctx_swa_mem->get_status())) { +} + +llama_kv_cache_dsv4_raw_context::llama_kv_cache_dsv4_raw_context( + llama_kv_cache_iswa * kv, + slot_info_vec_t sinfos_base_write, + slot_info_vec_t sinfos_swa_write, + slot_info_vec_t sinfos_swa_read, + std::vector ubatches, + std::vector ubatches_write) : + kv_swa(kv->get_swa()), + sinfos_write(std::move(sinfos_swa_write)), + sinfos_read(std::move(sinfos_swa_read)), + ubatches(std::move(ubatches)), + ubatches_write(std::move(ubatches_write)), + ctx_base_mem(std::make_unique( + kv->get_base(), std::move(sinfos_base_write), this->ubatches_write)), + ctx_swa_mem(nullptr), + n_kv(kv_swa->get_size()), + status(LLAMA_MEMORY_STATUS_SUCCESS) { +} + +bool llama_kv_cache_dsv4_raw_context::next() { + if (ubatches.empty()) { + return true; + } + + if (ctx_base_mem) { + ctx_base_mem->next(); + } + + if (++i_next >= ubatches.size()) { + return false; + } + + return true; +} + +bool llama_kv_cache_dsv4_raw_context::apply() { + bool res = true; + + if (ctx_base_mem) { + res = res & ctx_base_mem->apply(); + } + if (ctx_swa_mem) { + res = res & ctx_swa_mem->apply(); + } + if (!ubatches_write.empty()) { + kv_swa->apply_ubatch(sinfos_write[i_next], ubatches_write[i_next]); + n_kv = kv_swa->get_n_kv(sinfos_read[i_next]); + } + + return res; +} + +llama_memory_status llama_kv_cache_dsv4_raw_context::get_status() const { + return status; +} + +const llama_ubatch & llama_kv_cache_dsv4_raw_context::get_ubatch() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return ubatches[i_next]; +} + +uint32_t llama_kv_cache_dsv4_raw_context::get_n_kv() const { + return n_kv; +} + +uint32_t llama_kv_cache_dsv4_raw_context::get_n_write() const { + if (ubatches_write.empty()) { + return 0; + } + + return ubatches_write[i_next].n_tokens; +} + +ggml_tensor * llama_kv_cache_dsv4_raw_context::get_k(ggml_context * ctx, int32_t il) const { + return kv_swa->get_k(ctx, il, n_kv, sinfos_read[i_next]); +} + +ggml_tensor * llama_kv_cache_dsv4_raw_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const { + const auto & sinfo = sinfos_write[i_next]; + + if (k_cur->ne[2] == k_idxs->ne[0]) { + return kv_swa->cpy_k(ctx, k_cur, k_idxs, il, sinfo); + } + + // k_idxs may be expanded to one block per stream while k_cur is only + // the token block. Keep zero deps on all copies so each write executes. + const int64_t n_fanout = (int64_t) sinfo.size()*sinfo.n_stream(); + + GGML_ASSERT(sinfo.n_stream() > 1); + GGML_ASSERT(k_cur->ne[2] == (int64_t) sinfo.size()); + GGML_ASSERT(k_idxs->ne[0] == n_fanout); + + ggml_tensor * res = nullptr; + for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { + ggml_tensor * k_idxs_s = ggml_view_1d(ctx, k_idxs, sinfo.size(), s*sinfo.size()*ggml_element_size(k_idxs)); + ggml_tensor * cur = kv_swa->cpy_k(ctx, k_cur, k_idxs_s, il, sinfo); + if (res == nullptr) { + res = cur; + } else { + res = ggml_add(ctx, res, ggml_sub(ctx, cur, cur)); + } + } + + return res; +} + +ggml_tensor * llama_kv_cache_dsv4_raw_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { + const uint32_t n_tokens = ubatches_write.empty() ? ubatch.n_tokens : ubatches_write[i_next].n_tokens; + + ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens); + ggml_set_input(k_idxs); + + return k_idxs; +} + +ggml_tensor * llama_kv_cache_dsv4_raw_context::build_input_k_rot(ggml_context * ctx) const { + return kv_swa->build_input_k_rot(ctx); +} + +void llama_kv_cache_dsv4_raw_context::set_input_k_idxs(ggml_tensor * dst) const { + kv_swa->set_input_k_idxs(dst, &ubatches_write[i_next], sinfos_write[i_next]); +} + +void llama_kv_cache_dsv4_raw_context::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const { + kv_swa->set_input_kq_mask(dst, ubatch, causal_attn); +} + +void llama_kv_cache_dsv4_raw_context::set_input_k_rot(ggml_tensor * dst) const { + kv_swa->set_input_k_rot(dst); +} + +// +// llama_kv_cache_dsv4_comp_context +// + +llama_kv_cache_dsv4_comp_context::llama_kv_cache_dsv4_comp_context(llama_kv_cache * kv) : kv(kv), n_kv(kv->get_size()) { + const uint32_t n_stream = kv->get_n_stream(); + + sinfos.resize(1); + sinfos[0].s0 = 0; + sinfos[0].s1 = n_stream - 1; + sinfos[0].idxs.resize(n_stream); + for (uint32_t s = 0; s < n_stream; ++s) { + sinfos[0].strm.push_back(s); + sinfos[0].idxs[s].resize(1, 0); + } +} + +llama_kv_cache_dsv4_comp_context::llama_kv_cache_dsv4_comp_context( + llama_kv_cache * kv, + slot_info_vec_t sinfos, + std::vector ubatches) : + kv(kv), + sinfos(std::move(sinfos)), + ubatches(std::move(ubatches)), + n_kv(kv->get_size()) { +} + +bool llama_kv_cache_dsv4_comp_context::next() { + if (ubatches.empty()) { + return true; + } + + if (++i_cur >= ubatches.size()) { + return false; + } + + return true; +} + +uint32_t llama_kv_cache_dsv4_comp_context::get_n_kv() const { + return n_kv; +} + +ggml_tensor * llama_kv_cache_dsv4_comp_context::get_k(ggml_context * ctx, int32_t il) const { + return kv->get_k(ctx, il, n_kv, sinfos[i_cur]); +} + +ggml_tensor * llama_kv_cache_dsv4_comp_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const { + return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]); +} + +ggml_tensor * llama_kv_cache_dsv4_comp_context::build_input_k_rot(ggml_context * ctx) const { + return kv->build_input_k_rot(ctx); +} + +void llama_kv_cache_dsv4_comp_context::set_input_k_rot(ggml_tensor * dst) const { + kv->set_input_k_rot(dst); +} + +// +// llama_kv_cache_dsv4_context +// + +llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(llama_memory_status status) : status(status) {} + +llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( + llama_kv_cache_dsv4 * kv) : + ctx_raw(std::make_unique(kv->get_raw())), + ctx_csa_mem(kv->get_csa()->init_full()), + ctx_hca_mem(kv->get_hca()->init_full()), + ctx_lid_mem(kv->get_lid()->init_full()), + ctx_csa(std::make_unique(kv->get_csa())), + ctx_hca(std::make_unique(kv->get_hca())), + ctx_lid(std::make_unique(kv->get_lid())), + csa_state(kv->get_csa_state()), + hca_state(kv->get_hca_state()), + lid_state(kv->get_lid_state()), + reserve_plans(true), + status(llama_memory_status_combine( + llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()), + llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status()))) { +} + +llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( + llama_kv_cache_dsv4 * kv, + llama_context * lctx, + bool optimize) : + ctx_raw(std::make_unique(kv->get_raw(), lctx, optimize)), + ctx_csa_mem(kv->get_csa()->init_update(lctx, optimize)), + ctx_hca_mem(kv->get_hca()->init_update(lctx, optimize)), + ctx_lid_mem(kv->get_lid()->init_update(lctx, optimize)), + ctx_csa(std::make_unique(kv->get_csa())), + ctx_hca(std::make_unique(kv->get_hca())), + ctx_lid(std::make_unique(kv->get_lid())), + csa_state(kv->get_csa_state()), + hca_state(kv->get_hca_state()), + lid_state(kv->get_lid_state()), + status(llama_memory_status_combine( + llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()), + llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status()))) { +} + +llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( + llama_kv_cache_dsv4 * kv, + slot_info_vec_t sinfos_raw_base_write, + slot_info_vec_t sinfos_raw_swa_write, + slot_info_vec_t sinfos_raw_swa_read, + std::vector ubatches, + std::vector ubatches_raw) : + ubatches(std::move(ubatches)), + plans_csa(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true, + kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream())), + plans_hca(dsv4_build_comp_plans(this->ubatches, DSV4_HCA_RATIO, false, + kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream())), + plans_lid(plans_csa), + ctx_raw(std::make_unique( + kv->get_raw(), + std::move(sinfos_raw_base_write), + std::move(sinfos_raw_swa_write), + std::move(sinfos_raw_swa_read), + this->ubatches, + std::move(ubatches_raw))), + ctx_csa_mem(nullptr), + ctx_hca_mem(nullptr), + ctx_lid_mem(nullptr), + ctx_csa(std::make_unique( + kv->get_csa(), + dsv4_build_comp_sinfos(this->ubatches, kv->get_csa()->get_n_stream()), + this->ubatches)), + ctx_hca(std::make_unique( + kv->get_hca(), + dsv4_build_comp_sinfos(this->ubatches, kv->get_hca()->get_n_stream()), + this->ubatches)), + ctx_lid(std::make_unique( + kv->get_lid(), + dsv4_build_comp_sinfos(this->ubatches, kv->get_lid()->get_n_stream()), + this->ubatches)), + csa_state(kv->get_csa_state()), + hca_state(kv->get_hca_state()), + lid_state(kv->get_lid_state()), + status(ctx_raw->get_status()) { +} + +llama_kv_cache_dsv4_context::~llama_kv_cache_dsv4_context() = default; + +bool llama_kv_cache_dsv4_context::next() { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + ctx_raw->next(); + ctx_csa->next(); + ctx_hca->next(); + ctx_lid->next(); + + if (++i_next >= ubatches.size()) { + return false; + } + + return true; +} + +bool llama_kv_cache_dsv4_context::apply() { + assert(!llama_memory_status_is_fail(status)); + + bool res = true; + + res = res & ctx_raw->apply(); + + return res; +} + +llama_memory_status llama_kv_cache_dsv4_context::get_status() const { + return status; +} + +const llama_ubatch & llama_kv_cache_dsv4_context::get_ubatch() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return ubatches[i_next]; +} + +const llama_kv_cache_dsv4_raw_context * llama_kv_cache_dsv4_context::get_raw() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return ctx_raw.get(); +} + +const llama_kv_cache_dsv4_comp_context * llama_kv_cache_dsv4_context::get_csa() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return ctx_csa.get(); +} + +const llama_kv_cache_dsv4_comp_context * llama_kv_cache_dsv4_context::get_hca() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return ctx_hca.get(); +} + +const llama_kv_cache_dsv4_comp_context * llama_kv_cache_dsv4_context::get_lid() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return ctx_lid.get(); +} + +const llama_dsv4_comp_state * llama_kv_cache_dsv4_context::get_csa_state() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return csa_state; +} + +const llama_dsv4_comp_state * llama_kv_cache_dsv4_context::get_hca_state() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return hca_state; +} + +const llama_dsv4_comp_state * llama_kv_cache_dsv4_context::get_lid_state() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + return lid_state; +} + +const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_csa_plan() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + static const comp_plan empty; + if (plans_csa.empty()) { + return empty; + } + + return plans_csa[i_next]; +} + +const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_hca_plan() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + static const comp_plan empty; + if (plans_hca.empty()) { + return empty; + } + + return plans_hca[i_next]; +} + +const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_lid_plan() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + static const comp_plan empty; + if (plans_lid.empty()) { + return empty; + } + + return plans_lid[i_next]; +} + +const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_csa_plan(const llama_ubatch & ubatch) const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + if (!reserve_plans) { + return get_csa_plan(); + } + + reserve_plan_csa = dsv4_build_reserve_comp_plan( + ubatch, DSV4_CSA_RATIO, true, + csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream()); + + return reserve_plan_csa; +} + +const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_hca_plan(const llama_ubatch & ubatch) const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + if (!reserve_plans) { + return get_hca_plan(); + } + + reserve_plan_hca = dsv4_build_reserve_comp_plan( + ubatch, DSV4_HCA_RATIO, false, + hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream()); + + return reserve_plan_hca; +} + +const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_lid_plan(const llama_ubatch & ubatch) const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + + if (!reserve_plans) { + return get_lid_plan(); + } + + reserve_plan_lid = dsv4_build_reserve_comp_plan( + ubatch, DSV4_CSA_RATIO, true, + lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream()); + + return reserve_plan_lid; +} diff --git a/examples/talk-llama/llama-kv-cache-dsv4.h b/examples/talk-llama/llama-kv-cache-dsv4.h new file mode 100644 index 000000000..772b428cd --- /dev/null +++ b/examples/talk-llama/llama-kv-cache-dsv4.h @@ -0,0 +1,362 @@ +#pragma once + +#include "llama-kv-cache.h" +#include "llama-kv-cache-iswa.h" + +#include +#include +#include +#include + +class llama_dsv4_comp_state { +public: + llama_dsv4_comp_state( + const llama_model & model, + bool offload, + bool unified, + uint32_t n_seq_max, + uint32_t ratio, + uint32_t state_size, + uint32_t n_embd_state, + const char * name, + const llama_memory_i::layer_filter_cb & filter); + + void clear(bool data); + + uint32_t get_ratio() const; + uint32_t get_state_size() const; + uint32_t get_n_stream() const; + + std::map memory_breakdown() const; + + void state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const; + void state_read (llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags); + + ggml_tensor * get_kv (ggml_context * ctx, int32_t il) const; + ggml_tensor * get_score(ggml_context * ctx, int32_t il) const; + + ggml_tensor * cpy_kv (ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const; + ggml_tensor * cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const; + +private: + struct layer { + uint32_t il; + + ggml_tensor * kv; + ggml_tensor * score; + }; + + const uint32_t ratio; + const uint32_t state_size; + const uint32_t n_embd_state; + const uint32_t n_stream; + + std::vector> ctxs_bufs; + + std::vector layers; + + std::unordered_map map_layer_ids; + + size_t total_size() const; +}; + +// +// llama_kv_cache_dsv4 +// + +// DSV4 uses a normal raw/SWA token cache plus compressed K-only block caches. +// The compressed caches are storage only; DSV4-specific visibility and block +// planning are handled by llama_kv_cache_dsv4_context / llm_graph_input_dsv4. + +class llama_kv_cache_dsv4 : public llama_memory_i { +public: + llama_kv_cache_dsv4( + const llama_model & model, + ggml_type type_k, + ggml_type type_v, + bool v_trans, + bool offload, + bool swa_full, + bool unified, + uint32_t kv_size, + uint32_t n_seq_max, + uint32_t n_ubatch, + uint32_t n_pad, + const layer_filter_cb & filter, + const layer_reuse_cb & reuse); + + ~llama_kv_cache_dsv4() = default; + + // + // llama_memory_i + // + + 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; + void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override; + 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 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; + + // + // llama_kv_cache_dsv4 specific API + // + + llama_kv_cache_iswa * get_raw() const; + llama_kv_cache * get_csa() const; + llama_kv_cache * get_hca() const; + llama_kv_cache * get_lid() const; + llama_dsv4_comp_state * get_csa_state() const; + llama_dsv4_comp_state * get_hca_state() const; + llama_dsv4_comp_state * get_lid_state() const; + +private: + llama_hparams hparams_raw; + llama_hparams hparams_csa; + llama_hparams hparams_hca; + llama_hparams hparams_lid; + + const uint32_t n_seq_max; + + std::unique_ptr kv_raw; + std::unique_ptr kv_csa; + std::unique_ptr kv_hca; + std::unique_ptr kv_lid; + std::unique_ptr csa_state; + std::unique_ptr hca_state; + std::unique_ptr lid_state; + + void clear_compressed(bool data); +}; + +// DSV4 raw attention only uses the SWA half of kv_raw. The base half is kept +// for generic ISWA bookkeeping, but it has no DSV4 layers to expose here. +class llama_kv_cache_dsv4_raw_context : public llama_memory_context_i { +public: + using slot_info_vec_t = llama_kv_cache::slot_info_vec_t; + + llama_kv_cache_dsv4_raw_context(llama_kv_cache_iswa * kv); + + llama_kv_cache_dsv4_raw_context( + llama_kv_cache_iswa * kv, + llama_context * lctx, + bool optimize); + + llama_kv_cache_dsv4_raw_context( + llama_kv_cache_iswa * kv, + slot_info_vec_t sinfos_base_write, + slot_info_vec_t sinfos_swa_write, + slot_info_vec_t sinfos_swa_read, + std::vector ubatches, + std::vector ubatches_write); + + bool next() override; + bool apply() override; + + llama_memory_status get_status() const override; + const llama_ubatch & get_ubatch() const override; + + uint32_t get_n_kv() const; + uint32_t get_n_write() const; + + ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; + ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const; + + ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const; + ggml_tensor * build_input_k_rot(ggml_context * ctx) const; + + void set_input_k_idxs(ggml_tensor * dst) const; + void set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; + void set_input_k_rot(ggml_tensor * dst) const; + +private: + size_t i_next = 0; + + llama_kv_cache * kv_swa = nullptr; + + slot_info_vec_t sinfos_write; + slot_info_vec_t sinfos_read; + std::vector ubatches; + std::vector ubatches_write; + + const llama_memory_context_ptr ctx_base_mem; + const llama_memory_context_ptr ctx_swa_mem; + + uint32_t n_kv = 0; + + const llama_memory_status status; +}; + +// DSV4 compressed KV rows are graph outputs, not normal token KV writes. +// Keep a small context that exposes K tensors without generic apply() semantics. +class llama_kv_cache_dsv4_comp_context { +public: + using slot_info_vec_t = llama_kv_cache::slot_info_vec_t; + + llama_kv_cache_dsv4_comp_context(llama_kv_cache * kv); + + llama_kv_cache_dsv4_comp_context( + llama_kv_cache * kv, + slot_info_vec_t sinfos, + std::vector ubatches); + + bool next(); + + uint32_t get_n_kv() const; + + ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; + ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const; + + ggml_tensor * build_input_k_rot(ggml_context * ctx) const; + void set_input_k_rot(ggml_tensor * dst) const; + +private: + llama_kv_cache * kv; + + size_t i_cur = 0; + slot_info_vec_t sinfos; + std::vector ubatches; + + uint32_t n_kv; +}; + +class llama_kv_cache_dsv4_context : public llama_memory_context_i { +public: + using slot_info_vec_t = llama_kv_cache::slot_info_vec_t; + + struct comp_plan { + // Per-ubatch recipe for updating compressor state, committing completed + // compressed rows, and masking the compressed attention source. + + // APE row ids, i.e. pos % ratio, for the compressor-state updates. + std::vector state_pos; + + // Current-ubatch source row ids and unique persistent-state + // destination row ids for deterministic ring-state updates. + std::vector state_persist_src_idxs; + std::vector state_persist_dst_idxs; + + // Flattened source row ids used for state-backed commits. Source rows + // index the graph-local [persistent_state | current_ubatch_scratch] + // tensor. For overlapped compression the first half is previous rows + // and the second half is current rows; a final synthetic zero/-inf row + // may be addressed for the first block's previous half. + std::vector state_read_idxs; + + // Final compressed-cache row ids written by state-backed commits. + // A non-boundary CSA/LID decode step can target a masked scratch row. + std::vector state_write_idxs; + + // RoPE positions for state-backed commits. + std::vector state_write_pos; + + // Number of completed compressed rows visible for each query token. + std::vector n_visible; + + // Number of streams used by the attention graph for this ubatch. + int64_t n_stream = 1; + + // Graph-width for compressed rows. This can be larger than n_visible + // so masked padding rows do not force a new graph at every CSA block. + int64_t n_kv = 0; + }; + + llama_kv_cache_dsv4_context(llama_memory_status status); + + llama_kv_cache_dsv4_context( + llama_kv_cache_dsv4 * kv); + + llama_kv_cache_dsv4_context( + llama_kv_cache_dsv4 * kv, + llama_context * lctx, + bool optimize); + + llama_kv_cache_dsv4_context( + llama_kv_cache_dsv4 * kv, + slot_info_vec_t sinfos_raw_base_write, + slot_info_vec_t sinfos_raw_swa_write, + slot_info_vec_t sinfos_raw_swa_read, + std::vector ubatches, + std::vector ubatches_raw); + + virtual ~llama_kv_cache_dsv4_context(); + + // + // llama_memory_context_i + // + + bool next() override; + bool apply() override; + + llama_memory_status get_status() const override; + const llama_ubatch & get_ubatch() const override; + + // + // llama_kv_cache_dsv4_context specific API + // + + const llama_kv_cache_dsv4_raw_context * get_raw() const; + const llama_kv_cache_dsv4_comp_context * get_csa() const; + const llama_kv_cache_dsv4_comp_context * get_hca() const; + const llama_kv_cache_dsv4_comp_context * get_lid() const; + const llama_dsv4_comp_state * get_csa_state() const; + const llama_dsv4_comp_state * get_hca_state() const; + const llama_dsv4_comp_state * get_lid_state() const; + + const comp_plan & get_csa_plan() const; + const comp_plan & get_hca_plan() const; + const comp_plan & get_lid_plan() const; + + const comp_plan & get_csa_plan(const llama_ubatch & ubatch) const; + const comp_plan & get_hca_plan(const llama_ubatch & ubatch) const; + const comp_plan & get_lid_plan(const llama_ubatch & ubatch) const; + +private: + size_t i_next = 0; + + std::vector ubatches; + + std::vector plans_csa; + std::vector plans_hca; + std::vector plans_lid; + + const std::unique_ptr ctx_raw; + const llama_memory_context_ptr ctx_csa_mem; + const llama_memory_context_ptr ctx_hca_mem; + const llama_memory_context_ptr ctx_lid_mem; + + const std::unique_ptr ctx_csa; + const std::unique_ptr ctx_hca; + const std::unique_ptr ctx_lid; + + const llama_dsv4_comp_state * csa_state = nullptr; + const llama_dsv4_comp_state * hca_state = nullptr; + const llama_dsv4_comp_state * lid_state = nullptr; + + bool reserve_plans = false; + mutable comp_plan reserve_plan_csa; + mutable comp_plan reserve_plan_hca; + mutable comp_plan reserve_plan_lid; + + const llama_memory_status status; +}; diff --git a/examples/talk-llama/llama-kv-cache-iswa.cpp b/examples/talk-llama/llama-kv-cache-iswa.cpp index aa1b1b72e..e91866469 100644 --- a/examples/talk-llama/llama-kv-cache-iswa.cpp +++ b/examples/talk-llama/llama-kv-cache-iswa.cpp @@ -26,7 +26,28 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( llama_memory_t mem_other, const layer_filter_cb & filter, const layer_reuse_cb & reuse, - const layer_share_cb & share) : hparams(model.hparams), unified(unified) { + const layer_share_cb & share) : + llama_kv_cache_iswa(model, model.hparams, type_k, type_v, v_trans, offload, swa_full, unified, + kv_size, n_seq_max, n_ubatch, n_pad, mem_other, filter, reuse, share) { +} + +llama_kv_cache_iswa::llama_kv_cache_iswa( + const llama_model & model, + const llama_hparams & hparams, + ggml_type type_k, + ggml_type type_v, + bool v_trans, + bool offload, + bool swa_full, + bool unified, + uint32_t kv_size, + uint32_t n_seq_max, + uint32_t n_ubatch, + uint32_t n_pad, + llama_memory_t mem_other, + const layer_filter_cb & filter, + const layer_reuse_cb & reuse, + const layer_share_cb & share) : unified(unified) { // chain filters const layer_filter_cb filter_base = [&](int32_t il) { @@ -185,7 +206,7 @@ llama_memory_context_ptr llama_kv_cache_iswa::init_batch(llama_batch_allocr & ba std::vector ubatches; while (true) { - auto ubatch = balloc.split_equal(n_ubatch, !unified); + auto ubatch = balloc.split_equal(n_ubatch, !unified, 0); if (ubatch.n_tokens == 0) { break; diff --git a/examples/talk-llama/llama-kv-cache-iswa.h b/examples/talk-llama/llama-kv-cache-iswa.h index dfafc1ef5..7dab6eaa8 100644 --- a/examples/talk-llama/llama-kv-cache-iswa.h +++ b/examples/talk-llama/llama-kv-cache-iswa.h @@ -30,6 +30,24 @@ public: const layer_reuse_cb & reuse, const layer_share_cb & share); + llama_kv_cache_iswa( + const llama_model & model, + const llama_hparams & hparams, + ggml_type type_k, + ggml_type type_v, + bool v_trans, + bool offload, + bool swa_full, + bool unified, + uint32_t kv_size, + uint32_t n_seq_max, + uint32_t n_ubatch, + uint32_t n_pad, + llama_memory_t mem_other, + const layer_filter_cb & filter, + const layer_reuse_cb & reuse, + const layer_share_cb & share); + ~llama_kv_cache_iswa() = default; // @@ -73,8 +91,6 @@ public: llama_kv_cache * get_swa () const; private: - const llama_hparams & hparams; - const bool unified; std::unique_ptr kv_base; diff --git a/examples/talk-llama/llama-kv-cache.cpp b/examples/talk-llama/llama-kv-cache.cpp index 2802103bd..e70583e64 100644 --- a/examples/talk-llama/llama-kv-cache.cpp +++ b/examples/talk-llama/llama-kv-cache.cpp @@ -57,22 +57,6 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) { } } -static ggml_tensor * ggml_mul_mat_aux( - ggml_context * ctx, - ggml_tensor * cur, - ggml_tensor * rot) { - const auto n = rot->ne[0]; - - ggml_tensor * res; - - res = ggml_reshape_2d(ctx, cur, n, ggml_nelements(cur)/n); - res = ggml_mul_mat (ctx, rot, res); - ggml_mul_mat_set_hint(res, GGML_HINT_SRC0_IS_HADAMARD); - res = ggml_reshape_4d(ctx, res, cur->ne[0], cur->ne[1], cur->ne[2], cur->ne[3]); - - return res; -} - // // llama_kv_cache // @@ -211,10 +195,12 @@ llama_kv_cache::llama_kv_cache( n_embd_head_k_all = -1; } - if (n_embd_head_v_all == 0) { - n_embd_head_v_all = (int32_t) hparams.n_embd_head_v(il); - } else if (n_embd_head_v_all > 0 && n_embd_head_v_all != (int32_t) hparams.n_embd_head_v(il)) { - n_embd_head_v_all = -1; + if (!is_mla) { + if (n_embd_head_v_all == 0) { + n_embd_head_v_all = (int32_t) hparams.n_embd_head_v(il); + } else if (n_embd_head_v_all > 0 && n_embd_head_v_all != (int32_t) hparams.n_embd_head_v(il)) { + n_embd_head_v_all = -1; + } } // [TAG_V_CACHE_VARIABLE] @@ -336,8 +322,9 @@ llama_kv_cache::llama_kv_cache( ggml_is_quantized(type_k) && hparams.n_embd_head_k() % 64 == 0; - // always create Hadamard rotation tensors for DeepSeek V3.2 DSA lightning indexer - if (model.arch == LLM_ARCH_DEEPSEEK32 && hparams.n_embd_head_k_full == hparams.indexer_head_size) { + // always create Hadamard rotation tensors for DeepSeek lightning indexers + if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4) && + hparams.n_embd_head_k_full == hparams.indexer_head_size) { attn_rot_k = true; } @@ -719,7 +706,7 @@ llama_memory_context_ptr llama_kv_cache::init_batch( std::vector ubatches; while (true) { - auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true); + auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0); if (ubatch.n_tokens == 0) { break; @@ -1220,6 +1207,23 @@ ggml_type llama_kv_cache::type_v() const { return layers[0].v->type; } +std::vector llama_kv_cache::get_layer_ids() const { + std::vector res; + res.reserve(layers.size()); + + for (const auto & layer : layers) { + res.push_back(layer.il); + } + + return res; +} + +ggml_tensor * llama_kv_cache::get_k_storage(int32_t il) const { + const int32_t ikv = map_layer_ids.at(il); + + return layers[ikv].k; +} + uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { uint32_t result = 0; @@ -1855,14 +1859,14 @@ ggml_tensor * llama_kv_cache::build_rope_shift( tmp = ggml_cast(ctx, cur, GGML_TYPE_F32); // rotate back - tmp = ggml_mul_mat_aux(ctx, tmp, rot); + tmp = llama_mul_mat_hadamard(ctx, tmp, rot); tmp = ggml_rope_ext(ctx, tmp, shift, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, yarn_ext_factor, yarn_attn_factor, yarn_beta_fast, yarn_beta_slow); // rotate fwd - tmp = ggml_mul_mat_aux(ctx, tmp, rot); + tmp = llama_mul_mat_hadamard(ctx, tmp, rot); tmp = ggml_cpy(ctx, tmp, cur); } else { diff --git a/examples/talk-llama/llama-kv-cache.h b/examples/talk-llama/llama-kv-cache.h index 3d68f98c1..531d99dbd 100644 --- a/examples/talk-llama/llama-kv-cache.h +++ b/examples/talk-llama/llama-kv-cache.h @@ -161,6 +161,9 @@ public: ggml_type type_k() const; ggml_type type_v() const; + std::vector get_layer_ids() const; + ggml_tensor * get_k_storage(int32_t il) const; + // // graph_build API // diff --git a/examples/talk-llama/llama-memory-hybrid-iswa.cpp b/examples/talk-llama/llama-memory-hybrid-iswa.cpp index c7d4bcd41..06f7fd542 100644 --- a/examples/talk-llama/llama-memory-hybrid-iswa.cpp +++ b/examples/talk-llama/llama-memory-hybrid-iswa.cpp @@ -77,15 +77,15 @@ llama_memory_context_ptr llama_memory_hybrid_iswa::init_batch(llama_batch_allocr // if all tokens are output, split by sequence ubatch = balloc.split_seq(n_ubatch); } else { - if (mem_recr->n_rs_seq > 0) { - // [TAG_RECURRENT_ROLLBACK_SPLITS] - // TODO: recurrent state rollback does not support equal splits - ubatch = balloc.split_seq(n_ubatch); - } else { - // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice) - const bool unified = (mem_attn->get_base()->get_n_stream() == 1); - ubatch = balloc.split_equal(n_ubatch, !unified); - } + // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice) + const bool unified = (mem_attn->get_base()->get_n_stream() == 1); + + // [TAG_RECURRENT_ROLLBACK_SPLITS] + // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch + // so that the rollback snapshots remain valid + const uint32_t n_rs_seq = mem_recr->n_rs_seq; + + ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0); } if (ubatch.n_tokens == 0) { diff --git a/examples/talk-llama/llama-memory-hybrid.cpp b/examples/talk-llama/llama-memory-hybrid.cpp index f2d49cbce..42c7381a9 100644 --- a/examples/talk-llama/llama-memory-hybrid.cpp +++ b/examples/talk-llama/llama-memory-hybrid.cpp @@ -78,15 +78,15 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba // if all tokens are output, split by sequence ubatch = balloc.split_seq(n_ubatch); } else { - if (mem_recr->n_rs_seq > 0) { - // [TAG_RECURRENT_ROLLBACK_SPLITS] - // TODO: recurrent state rollback does not support equal splits - ubatch = balloc.split_seq(n_ubatch); - } else { - // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice) - const bool unified = (mem_attn->get_n_stream() == 1); - ubatch = balloc.split_equal(n_ubatch, !unified); - } + // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice) + const bool unified = (mem_attn->get_n_stream() == 1); + + // [TAG_RECURRENT_ROLLBACK_SPLITS] + // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch + // so that the rollback snapshots remain valid + const uint32_t n_rs_seq = mem_recr->n_rs_seq; + + ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0); } if (ubatch.n_tokens == 0) { diff --git a/examples/talk-llama/llama-memory-recurrent.cpp b/examples/talk-llama/llama-memory-recurrent.cpp index 6a4892fb4..3d6c6db87 100644 --- a/examples/talk-llama/llama-memory-recurrent.cpp +++ b/examples/talk-llama/llama-memory-recurrent.cpp @@ -416,15 +416,12 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & // if all tokens are output, split by sequence ubatch = balloc.split_seq(n_ubatch); } else { - if (n_rs_seq > 0) { - // [TAG_RECURRENT_ROLLBACK_SPLITS] - // TODO: recurrent state rollback does not support equal splits - ubatch = balloc.split_seq(n_ubatch); - } else { - // TODO: non-sequential equal split can be done if using unified KV cache - // for simplicity, we always use sequential equal split for now - ubatch = balloc.split_equal(n_ubatch, true); - } + // TODO: non-sequential equal split can be done if using unified KV cache + // for simplicity, we always use sequential equal split for now + // [TAG_RECURRENT_ROLLBACK_SPLITS] + // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch + // so that the rollback snapshots remain valid + ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0); } if (ubatch.n_tokens == 0) { diff --git a/examples/talk-llama/llama-model-loader.cpp b/examples/talk-llama/llama-model-loader.cpp index 474cabdfc..28f8bb793 100644 --- a/examples/talk-llama/llama-model-loader.cpp +++ b/examples/talk-llama/llama-model-loader.cpp @@ -27,52 +27,55 @@ const char * llama_file_version_name(llama_fver version) { return "unknown"; } -static std::string llama_model_ftype_name(llama_ftype ftype) { - if (ftype & LLAMA_FTYPE_GUESSED) { - return llama_model_ftype_name((enum llama_ftype) (ftype & ~LLAMA_FTYPE_GUESSED)) + " (guessed)"; - } +#define LLAMA_FTYPE_PREFIX "(guessed) " - switch (ftype) { - case LLAMA_FTYPE_ALL_F32: return "all F32"; - case LLAMA_FTYPE_MOSTLY_F16: return "F16"; - case LLAMA_FTYPE_MOSTLY_BF16: return "BF16"; - case LLAMA_FTYPE_MOSTLY_Q1_0: return "Q1_0"; - case LLAMA_FTYPE_MOSTLY_Q4_0: return "Q4_0"; - case LLAMA_FTYPE_MOSTLY_Q4_1: return "Q4_1"; - case LLAMA_FTYPE_MOSTLY_Q5_0: return "Q5_0"; - case LLAMA_FTYPE_MOSTLY_Q5_1: return "Q5_1"; - case LLAMA_FTYPE_MOSTLY_Q8_0: return "Q8_0"; - case LLAMA_FTYPE_MOSTLY_MXFP4_MOE: return "MXFP4 MoE"; - case LLAMA_FTYPE_MOSTLY_NVFP4: return "NVFP4"; - case LLAMA_FTYPE_MOSTLY_Q2_K: return "Q2_K - Medium"; - case LLAMA_FTYPE_MOSTLY_Q2_K_S: return "Q2_K - Small"; - case LLAMA_FTYPE_MOSTLY_Q3_K_S: return "Q3_K - Small"; - case LLAMA_FTYPE_MOSTLY_Q3_K_M: return "Q3_K - Medium"; - case LLAMA_FTYPE_MOSTLY_Q3_K_L: return "Q3_K - Large"; - case LLAMA_FTYPE_MOSTLY_Q4_K_S: return "Q4_K - Small"; - case LLAMA_FTYPE_MOSTLY_Q4_K_M: return "Q4_K - Medium"; - case LLAMA_FTYPE_MOSTLY_Q5_K_S: return "Q5_K - Small"; - case LLAMA_FTYPE_MOSTLY_Q5_K_M: return "Q5_K - Medium"; - case LLAMA_FTYPE_MOSTLY_Q6_K: return "Q6_K"; - case LLAMA_FTYPE_MOSTLY_TQ1_0: return "TQ1_0 - 1.69 bpw ternary"; - case LLAMA_FTYPE_MOSTLY_TQ2_0: return "TQ2_0 - 2.06 bpw ternary"; - case LLAMA_FTYPE_MOSTLY_IQ2_XXS: return "IQ2_XXS - 2.0625 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ2_XS: return "IQ2_XS - 2.3125 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ2_S: return "IQ2_S - 2.5 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ2_M: return "IQ2_M - 2.7 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ3_XS: return "IQ3_XS - 3.3 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ3_XXS: return "IQ3_XXS - 3.0625 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ1_S: return "IQ1_S - 1.5625 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ1_M: return "IQ1_M - 1.75 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ4_NL: return "IQ4_NL - 4.5 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ4_XS: return "IQ4_XS - 4.25 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ3_S: return "IQ3_S - 3.4375 bpw"; - case LLAMA_FTYPE_MOSTLY_IQ3_M: return "IQ3_S mix - 3.66 bpw"; - - default: return "unknown, may not work"; +const char * llama_ftype_name(llama_ftype ftype) { + static constexpr size_t guessed_prefix_len = sizeof(LLAMA_FTYPE_PREFIX) - 1; + const char * name; + switch ((enum llama_ftype) (ftype & ~LLAMA_FTYPE_GUESSED)) { + case LLAMA_FTYPE_ALL_F32: name = LLAMA_FTYPE_PREFIX "all F32"; break; + case LLAMA_FTYPE_MOSTLY_F16: name = LLAMA_FTYPE_PREFIX "F16"; break; + case LLAMA_FTYPE_MOSTLY_BF16: name = LLAMA_FTYPE_PREFIX "BF16"; break; + case LLAMA_FTYPE_MOSTLY_Q1_0: name = LLAMA_FTYPE_PREFIX "Q1_0"; break; + case LLAMA_FTYPE_MOSTLY_Q2_0: name = LLAMA_FTYPE_PREFIX "Q2_0"; break; + case LLAMA_FTYPE_MOSTLY_Q4_0: name = LLAMA_FTYPE_PREFIX "Q4_0"; break; + case LLAMA_FTYPE_MOSTLY_Q4_1: name = LLAMA_FTYPE_PREFIX "Q4_1"; break; + case LLAMA_FTYPE_MOSTLY_Q5_0: name = LLAMA_FTYPE_PREFIX "Q5_0"; break; + case LLAMA_FTYPE_MOSTLY_Q5_1: name = LLAMA_FTYPE_PREFIX "Q5_1"; break; + case LLAMA_FTYPE_MOSTLY_Q8_0: name = LLAMA_FTYPE_PREFIX "Q8_0"; break; + case LLAMA_FTYPE_MOSTLY_MXFP4_MOE: name = LLAMA_FTYPE_PREFIX "MXFP4 MoE"; break; + case LLAMA_FTYPE_MOSTLY_NVFP4: name = LLAMA_FTYPE_PREFIX "NVFP4"; break; + case LLAMA_FTYPE_MOSTLY_Q2_K: name = LLAMA_FTYPE_PREFIX "Q2_K - Medium"; break; + case LLAMA_FTYPE_MOSTLY_Q2_K_S: name = LLAMA_FTYPE_PREFIX "Q2_K - Small"; break; + case LLAMA_FTYPE_MOSTLY_Q3_K_S: name = LLAMA_FTYPE_PREFIX "Q3_K - Small"; break; + case LLAMA_FTYPE_MOSTLY_Q3_K_M: name = LLAMA_FTYPE_PREFIX "Q3_K - Medium"; break; + case LLAMA_FTYPE_MOSTLY_Q3_K_L: name = LLAMA_FTYPE_PREFIX "Q3_K - Large"; break; + case LLAMA_FTYPE_MOSTLY_Q4_K_S: name = LLAMA_FTYPE_PREFIX "Q4_K - Small"; break; + case LLAMA_FTYPE_MOSTLY_Q4_K_M: name = LLAMA_FTYPE_PREFIX "Q4_K - Medium"; break; + case LLAMA_FTYPE_MOSTLY_Q5_K_S: name = LLAMA_FTYPE_PREFIX "Q5_K - Small"; break; + case LLAMA_FTYPE_MOSTLY_Q5_K_M: name = LLAMA_FTYPE_PREFIX "Q5_K - Medium"; break; + case LLAMA_FTYPE_MOSTLY_Q6_K: name = LLAMA_FTYPE_PREFIX "Q6_K"; break; + case LLAMA_FTYPE_MOSTLY_TQ1_0: name = LLAMA_FTYPE_PREFIX "TQ1_0 - 1.69 bpw ternary"; break; + case LLAMA_FTYPE_MOSTLY_TQ2_0: name = LLAMA_FTYPE_PREFIX "TQ2_0 - 2.06 bpw ternary"; break; + case LLAMA_FTYPE_MOSTLY_IQ2_XXS: name = LLAMA_FTYPE_PREFIX "IQ2_XXS - 2.0625 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ2_XS: name = LLAMA_FTYPE_PREFIX "IQ2_XS - 2.3125 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ2_S: name = LLAMA_FTYPE_PREFIX "IQ2_S - 2.5 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ2_M: name = LLAMA_FTYPE_PREFIX "IQ2_M - 2.7 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ3_XS: name = LLAMA_FTYPE_PREFIX "IQ3_XS - 3.3 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ3_XXS: name = LLAMA_FTYPE_PREFIX "IQ3_XXS - 3.0625 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ1_S: name = LLAMA_FTYPE_PREFIX "IQ1_S - 1.5625 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ1_M: name = LLAMA_FTYPE_PREFIX "IQ1_M - 1.75 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ4_NL: name = LLAMA_FTYPE_PREFIX "IQ4_NL - 4.5 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ4_XS: name = LLAMA_FTYPE_PREFIX "IQ4_XS - 4.25 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ3_S: name = LLAMA_FTYPE_PREFIX "IQ3_S - 3.4375 bpw"; break; + case LLAMA_FTYPE_MOSTLY_IQ3_M: name = LLAMA_FTYPE_PREFIX "IQ3_S mix - 3.66 bpw"; break; + default: name = LLAMA_FTYPE_PREFIX "unknown, may not work"; break; } + return (ftype & LLAMA_FTYPE_GUESSED) ? name : name + guessed_prefix_len; } +#undef LLAMA_FTYPE_PREFIX + // return a list of splits for a given path // for example, given "-00002-of-00004.gguf", returns list of all 4 splits static std::vector llama_get_list_splits(const std::string & path, const int idx, const int n_split) { @@ -294,6 +297,8 @@ namespace GGUFMeta { } template bool llama_model_loader::get_arr_n(enum llm_kv kid, uint32_t & result, bool required); + template std::enable_if::value, bool>::type + llama_model_loader::get_arr_n(const std::string & key, uint32_t & result, bool required); template bool llama_model_loader::get_arr(const std::string & key, std::vector & result, bool required) { @@ -395,6 +400,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::array & result, bool required); template bool llama_model_loader::get_key(const std::string & key, T & result, bool required) { @@ -762,6 +768,7 @@ llama_model_loader::llama_model_loader( case GGML_TYPE_IQ3_S: ftype = LLAMA_FTYPE_MOSTLY_IQ3_S; break; case GGML_TYPE_NVFP4: ftype = LLAMA_FTYPE_MOSTLY_NVFP4; break; case GGML_TYPE_Q1_0: ftype = LLAMA_FTYPE_MOSTLY_Q1_0; break; + case GGML_TYPE_Q2_0: ftype = LLAMA_FTYPE_MOSTLY_Q2_0; break; default: { LLAMA_LOG_WARN("%s: unknown type %s\n", __func__, ggml_type_name(type_max)); @@ -1690,12 +1697,12 @@ bool llama_model_loader::load_all_data( } std::string llama_model_loader::ftype_name() const { - return llama_model_ftype_name(ftype); + return llama_ftype_name(ftype); } void llama_model_loader::print_info() const { LLAMA_LOG_INFO("%s: file format = %s\n", __func__, llama_file_version_name(fver)); - LLAMA_LOG_INFO("%s: file type = %s\n", __func__, llama_model_ftype_name(ftype).c_str()); + LLAMA_LOG_INFO("%s: file type = %s\n", __func__, llama_ftype_name(ftype)); if (n_bytes < GiB) { LLAMA_LOG_INFO("%s: file size = %.2f MiB (%.2f BPW) \n", __func__, n_bytes/1024.0/1024.0, n_bytes*8.0/n_elements); } else { diff --git a/examples/talk-llama/llama-model.cpp b/examples/talk-llama/llama-model.cpp index 6cb0ec379..adacf702d 100644 --- a/examples/talk-llama/llama-model.cpp +++ b/examples/talk-llama/llama-model.cpp @@ -11,6 +11,7 @@ #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" #include "llama-kv-cache-dsa.h" +#include "llama-kv-cache-dsv4.h" #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-recurrent.h" @@ -181,6 +182,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_deepseek2ocr(params); case LLM_ARCH_DEEPSEEK32: return new llama_model_deepseek32(params); + case LLM_ARCH_DEEPSEEK4: + return new llama_model_deepseek4(params); case LLM_ARCH_GLM_DSA: return new llama_model_glm_dsa(params); case LLM_ARCH_MISTRAL4: @@ -291,6 +294,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_mistral3(params); case LLM_ARCH_EAGLE3: return new llama_model_eagle3(params); + case LLM_ARCH_DFLASH: + return new llama_model_dflash(params); case LLM_ARCH_MIMO2: return new llama_model_mimo2(params); case LLM_ARCH_KIMI_LINEAR: @@ -815,6 +820,7 @@ static const char * llama_expert_gating_func_name(llama_expert_gating_func_type switch (type) { case LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX: return "softmax"; case LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID: return "sigmoid"; + case LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS: return "sqrtsoftplus"; default: return "unknown"; } } @@ -947,6 +953,8 @@ static buft_list_t make_gpu_buft_list(ggml_backend_dev_t dev, llama_split_mode s if (buft != nullptr) { buft_list.emplace_back(dev, buft); } + } else { + throw std::runtime_error(format("device %s does not support split buffers", ggml_backend_dev_name(dev))); } } @@ -981,6 +989,8 @@ struct llama_model::impl { std::string desc_str; + llama_ftype ftype = LLAMA_FTYPE_ALL_F32; + // model memory mapped files llama_mmaps mappings; @@ -1004,9 +1014,17 @@ struct llama_model::impl { std::vector dev_layer; bool has_tensor_overrides; + + std::vector tensor_split_owned; }; llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique()) { + if (params.tensor_split != nullptr) { + // llama_model_params stores tensor_split as a borrowed pointer, but the model + // may need it later for tensor-parallel KV-cache split metadata. + pimpl->tensor_split_owned.assign(params.tensor_split, params.tensor_split + llama_max_devices()); + this->params.tensor_split = pimpl->tensor_split_owned.data(); + } pimpl->has_tensor_overrides = params.tensor_buft_overrides && params.tensor_buft_overrides[0].pattern; } @@ -1194,6 +1212,8 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { pimpl->desc_str = arch_name() + " " + type_name() + " " + ml.ftype_name(); + pimpl->ftype = ml.ftype; + if (hparams.f_max_alibi_bias > 0.0f) { hparams.use_alibi = true; } @@ -1640,6 +1660,10 @@ std::string llama_model::desc() const { return pimpl->desc_str; } +llama_ftype llama_model::ftype() const { + return pimpl->ftype; +} + size_t llama_model::size() const { return pimpl->n_bytes; } @@ -2154,7 +2178,24 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } } - if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { + if (arch == LLM_ARCH_DEEPSEEK4) { + GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); + + res = new llama_kv_cache_dsv4( + *this, + params.type_k, + params.type_v, + !cparams.flash_attn, + cparams.offload_kqv, + params.swa_full, + cparams.kv_unified, + cparams.n_ctx_seq, + cparams.n_seq_max, + cparams.n_ubatch, + 1, + filter, + reuse); + } else if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { GGML_ASSERT(hparams.is_swa_any()); if (arch == LLM_ARCH_GEMMA4_ASSISTANT) { @@ -2326,6 +2367,11 @@ 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) { + return 0; + } return model->hparams.n_swa; } @@ -2407,6 +2453,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_DEEPSEEK2: case LLM_ARCH_DEEPSEEK2OCR: case LLM_ARCH_DEEPSEEK32: + case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_PLM: case LLM_ARCH_CHATGLM: case LLM_ARCH_GRANITE: @@ -2494,6 +2541,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_STEP35: case LLM_ARCH_TALKIE: case LLM_ARCH_MELLUM: + case LLM_ARCH_DFLASH: return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_QWEN2VL: @@ -2586,6 +2634,10 @@ int32_t llama_model_desc(const llama_model * model, char * buf, size_t buf_size) return snprintf(buf, buf_size, "%s", model->desc().c_str()); } +llama_ftype llama_model_ftype(const llama_model * model) { + return model->ftype(); +} + uint64_t llama_model_size(const llama_model * model) { return model->size(); } @@ -2617,7 +2669,8 @@ bool llama_model_has_encoder(const llama_model * model) { switch (model->arch) { case LLM_ARCH_T5: case LLM_ARCH_T5ENCODER: - case LLM_ARCH_EAGLE3: return true; + case LLM_ARCH_EAGLE3: + case LLM_ARCH_DFLASH: return true; default: return false; } } diff --git a/examples/talk-llama/llama-model.h b/examples/talk-llama/llama-model.h index 77d8d3b62..45b054ced 100644 --- a/examples/talk-llama/llama-model.h +++ b/examples/talk-llama/llama-model.h @@ -255,9 +255,11 @@ struct llama_layer { struct ggml_tensor * wq_b = nullptr; struct ggml_tensor * wkv_a_mqa = nullptr; struct ggml_tensor * wkv_b = nullptr; + struct ggml_tensor * wkv = nullptr; struct ggml_tensor * wk_b = nullptr; struct ggml_tensor * wv_b = nullptr; struct ggml_tensor * wqkv_b = nullptr; + struct ggml_tensor * wo_a = nullptr; struct ggml_tensor * wo_b = nullptr; struct ggml_tensor * wq_cross = nullptr; struct ggml_tensor * wk_cross = nullptr; @@ -333,6 +335,7 @@ struct llama_layer { struct ggml_tensor * ffn_up_b = nullptr; // b3 struct ggml_tensor * ffn_act = nullptr; struct ggml_tensor * ffn_exp_probs_b = nullptr; + struct ggml_tensor * ffn_gate_tid2eid = nullptr; // mamba proj struct ggml_tensor * ssm_in = nullptr; @@ -463,6 +466,23 @@ struct llama_layer { // openai-moe struct ggml_tensor * attn_sinks = nullptr; + // DeepSeek-V4 + struct ggml_tensor * attn_kv_norm = nullptr; + struct ggml_tensor * hc_attn_fn = nullptr; + struct ggml_tensor * hc_attn_base = nullptr; + struct ggml_tensor * hc_attn_scale = nullptr; + struct ggml_tensor * hc_ffn_fn = nullptr; + struct ggml_tensor * hc_ffn_base = nullptr; + struct ggml_tensor * hc_ffn_scale = nullptr; + struct ggml_tensor * attn_comp_wkv = nullptr; + struct ggml_tensor * attn_comp_wgate = nullptr; + struct ggml_tensor * attn_comp_ape = nullptr; + struct ggml_tensor * attn_comp_norm = nullptr; + struct ggml_tensor * indexer_comp_wkv = nullptr; + struct ggml_tensor * indexer_comp_wgate = nullptr; + struct ggml_tensor * indexer_comp_ape = nullptr; + struct ggml_tensor * indexer_comp_norm = nullptr; + // cogvlm struct ggml_tensor * visexp_attn_wqkv = nullptr; struct ggml_tensor * visexp_attn_wo = nullptr; @@ -553,6 +573,11 @@ struct llama_model { struct ggml_tensor * nextn_proj_pre = nullptr; struct ggml_tensor * nextn_proj_post = nullptr; + // DeepSeek-V4 + struct ggml_tensor * hc_head_fn = nullptr; + struct ggml_tensor * hc_head_base = nullptr; + struct ggml_tensor * hc_head_scale = nullptr; + // classifier struct ggml_tensor * cls = nullptr; struct ggml_tensor * cls_b = nullptr; @@ -612,6 +637,8 @@ struct llama_model { std::string desc() const; + llama_ftype ftype() const; + size_t size() const; // file size size_t n_tensors() const; size_t n_devices() const; diff --git a/examples/talk-llama/llama-quant.cpp b/examples/talk-llama/llama-quant.cpp index 847e79f46..aebbc1ffb 100644 --- a/examples/talk-llama/llama-quant.cpp +++ b/examples/talk-llama/llama-quant.cpp @@ -380,6 +380,7 @@ static ggml_type tensor_type_fallback(quantize_state_impl & qs, const ggml_tenso case GGML_TYPE_IQ3_XXS: case GGML_TYPE_IQ3_S: // types on the right: block size 32 case GGML_TYPE_IQ4_XS: return_type = GGML_TYPE_IQ4_NL; break; + case GGML_TYPE_Q2_0: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_TQ1_0: @@ -480,7 +481,7 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type else if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS) { new_type = GGML_TYPE_IQ3_S; } - else if (ftype == LLAMA_FTYPE_MOSTLY_TQ1_0 || ftype == LLAMA_FTYPE_MOSTLY_TQ2_0) { + else if (ftype == LLAMA_FTYPE_MOSTLY_TQ1_0 || ftype == LLAMA_FTYPE_MOSTLY_TQ2_0 || ftype == LLAMA_FTYPE_MOSTLY_Q2_0) { new_type = GGML_TYPE_Q4_K; } } @@ -800,6 +801,7 @@ ggml_type llama_ftype_get_default_type(llama_ftype ftype) { case LLAMA_FTYPE_MOSTLY_BF16: return GGML_TYPE_BF16; case LLAMA_FTYPE_ALL_F32: return GGML_TYPE_F32; case LLAMA_FTYPE_MOSTLY_Q1_0: return GGML_TYPE_Q1_0; + case LLAMA_FTYPE_MOSTLY_Q2_0: return GGML_TYPE_Q2_0; case LLAMA_FTYPE_MOSTLY_MXFP4_MOE: return GGML_TYPE_MXFP4; diff --git a/examples/talk-llama/llama-vocab.cpp b/examples/talk-llama/llama-vocab.cpp index 6e78a3f6c..fdd447147 100644 --- a/examples/talk-llama/llama-vocab.cpp +++ b/examples/talk-llama/llama-vocab.cpp @@ -887,9 +887,6 @@ struct llm_tokenizer_ugm : llm_tokenizer { // blob containing XOR-compressed compact double array (XCDA) entries uint32_t xcda_blob_size = *(const uint32_t *) &precompiled_charsmap[0]; charsmap_offset += sizeof(xcda_blob_size); - if (xcda_blob_size + charsmap_offset >= precompiled_charsmap.size()) { - throw std::runtime_error("Index out of array bounds in precompiled charsmap!"); - } // Next xcda_blob_size bytes contain entries of XOR-compressed compact // double array (XCDA). Each entry is bit-packed into a 32-bit integer. @@ -1205,7 +1202,15 @@ private: throw std::runtime_error("Index out of array bounds in precompiled charsmap!"); } const char * prefix_replacement = &(tokenizer.prefix_replacements)[longest_prefix_offset]; - return { prefix_replacement, strlen(prefix_replacement), longest_prefix_length }; + size_t max_len = tokenizer.prefix_replacements_size - longest_prefix_offset; + size_t repl_len = 0; + while (repl_len < max_len && prefix_replacement[repl_len] != '\0') { + repl_len++; + } + if (repl_len == max_len) { + throw std::runtime_error("Unterminated string in precompiled charsmap!"); + } + return { prefix_replacement, repl_len, longest_prefix_length }; } // check if the input prefix contains a valid sequence of UTF-8 code units @@ -2018,11 +2023,18 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { const size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx); const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx); precompiled_charsmap.assign(pc, pc + n_precompiled_charsmap); + if (precompiled_charsmap.size() < sizeof(uint32_t)) { + throw std::runtime_error("precompiled_charsmap too small for xcda_blob_size header!"); + } + uint32_t * xcda_blob_size = (uint32_t *) &precompiled_charsmap[0]; +#if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ + *xcda_blob_size = __builtin_bswap32(*xcda_blob_size); +#endif + if (*xcda_blob_size + sizeof(uint32_t) >= precompiled_charsmap.size()) { + throw std::runtime_error("Index out of array bounds in precompiled charsmap!"); + } #if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ // correct endianness of data in precompiled_charsmap binary blob - uint32_t * xcda_blob_size = (uint32_t *) &precompiled_charsmap[0]; - *xcda_blob_size = __builtin_bswap32(*xcda_blob_size); - assert(*xcda_blob_size + sizeof(uint32_t) < n_precompiled_charsmap); size_t xcda_array_size = *xcda_blob_size / sizeof(uint32_t); uint32_t * xcda_array = (uint32_t *) &precompiled_charsmap[sizeof(uint32_t)]; for (size_t i = 0; i < xcda_array_size; ++i) { diff --git a/examples/talk-llama/llama.h b/examples/talk-llama/llama.h index f723c9f60..a311ac202 100644 --- a/examples/talk-llama/llama.h +++ b/examples/talk-llama/llama.h @@ -155,10 +155,14 @@ extern "C" { LLAMA_FTYPE_MOSTLY_MXFP4_MOE = 38, // except 1d tensors LLAMA_FTYPE_MOSTLY_NVFP4 = 39, // except 1d tensors LLAMA_FTYPE_MOSTLY_Q1_0 = 40, // except 1d tensors + LLAMA_FTYPE_MOSTLY_Q2_0 = 41, // except 1d tensors LLAMA_FTYPE_GUESSED = 1024, // not specified in the model file }; + // Get the model file type (quantization) as a string, e.g. "Q8_0" or "Q4_K - Medium" + LLAMA_API const char * llama_ftype_name(enum llama_ftype ftype); + enum llama_rope_scaling_type { LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED = -1, LLAMA_ROPE_SCALING_TYPE_NONE = 0, @@ -606,6 +610,9 @@ extern "C" { // Get a string describing the model type LLAMA_API int32_t llama_model_desc(const struct llama_model * model, char * buf, size_t buf_size); + // Get the model file type (quantization), e.g. LLAMA_FTYPE_MOSTLY_Q8_0 + LLAMA_API enum llama_ftype llama_model_ftype(const struct llama_model * model); + // Returns the total size of all the tensors in the model in bytes LLAMA_API uint64_t llama_model_size(const struct llama_model * model); diff --git a/examples/talk-llama/models/deepseek4.cpp b/examples/talk-llama/models/deepseek4.cpp new file mode 100644 index 000000000..07aa477e1 --- /dev/null +++ b/examples/talk-llama/models/deepseek4.cpp @@ -0,0 +1,1165 @@ +#include "models.h" + +#include "llama-kv-cache-dsv4.h" + +#include +#include +#include +#include + +static float dsv4_rope_attn_factor(float freq_scale, float ext_factor) { + if (ext_factor == 0.0f) { + return 1.0f; + } + + return 1.0f / (1.0f + 0.1f*logf(1.0f/freq_scale)); +} + +void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer()); + if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer(), 0)) { + hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp; + } + + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + + ml.get_key(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count); + ml.get_key(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank); + ml.get_key(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, hparams.dsv4_compress_rope_base); + ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); + ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); + ml.get_key(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count); + + uint32_t n_compress_ratios = 0; + ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_compress_ratios); + if (n_compress_ratios < hparams.n_layer()) { + throw std::runtime_error("DeepSeek-V4 compress_ratios is shorter than block_count"); + } + ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios); + + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) { + throw std::runtime_error("DeepSeek-V4 loader currently expects sqrtsoftplus MoE scoring"); + } + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + hparams.set_swa_pattern(0); + + switch (hparams.n_layer()) { + case 43: type = LLM_TYPE_UNKNOWN; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) { + 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; + + 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); + + hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc_dim, hc_mult}, 0); + hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0); + hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0); + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0); + layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0); + layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0); + layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank * o_groups}, 0); + layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0); + + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0); + + const int64_t ratio = hparams.dsv4_compress_ratios[i]; + if (ratio != 0) { + const int64_t coff = ratio == 4 ? 2 : 1; + + layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", i), {n_embd, coff * n_embd_head}, 0); + layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, coff * n_embd_head}, 0); + layer.attn_comp_ape = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_APE, "weight", i), {coff * n_embd_head, ratio}, 0); + layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", i), {n_embd_head}, 0); + + if (ratio == 4) { + const int64_t n_embd_indexer = hparams.indexer_head_size; + + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, 0); + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, 0); + + layer.indexer_comp_wkv = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "weight", i), {n_embd, 2 * n_embd_indexer}, 0); + layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, 2 * n_embd_indexer}, 0); + layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", i), {2 * n_embd_indexer, ratio}, 0); + layer.indexer_comp_norm = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_NORM, "weight", i), {n_embd_indexer}, 0); + } else if (ratio != 128) { + throw std::runtime_error("DeepSeek-V4 loader only supports compression ratios 0, 4, and 128"); + } + } + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + if ((uint32_t) i < hparams.dsv4_hash_layer_count) { + layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, "weight", i), {n_expert_used, n_vocab}, 0); + } else { + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + } + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0); + } +} + +std::unique_ptr llama_model_deepseek4::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +static size_t dsv4_elem_offset(const ggml_tensor * t, int64_t i) { + return ggml_row_size(t->type, i); +} + +static ggml_tensor * dsv4_view_1d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t i0) { + return ggml_view_1d(ctx, t, ne0, dsv4_elem_offset(t, i0)); +} + +static ggml_tensor * dsv4_view_2d( + ggml_context * ctx, + ggml_tensor * t, + int64_t ne0, + int64_t ne1, + int64_t i0) { + return ggml_view_2d(ctx, t, ne0, ne1, t->nb[1], dsv4_elem_offset(t, i0)); +} + +static ggml_tensor * dsv4_append_zero_row(ggml_context * ctx, ggml_tensor * t, bool neg_inf) { + ggml_tensor * row = ggml_view_1d(ctx, t, t->ne[0], 0); + row = neg_inf ? ggml_scale_bias(ctx, row, 0.0f, -INFINITY) : ggml_scale(ctx, row, 0.0f); + row = ggml_reshape_2d(ctx, row, t->ne[0], 1); + + return ggml_concat(ctx, t, row, 1); +} + +static ggml_tensor * dsv4_with_zero_dep(ggml_context * ctx, ggml_tensor * t, ggml_tensor * dep) { + if (dep == nullptr) { + return t; + } + + ggml_tensor * zero = ggml_scale(ctx, ggml_sum(ctx, dep), 0.0f); + return ggml_add(ctx, t, zero); +} + +static constexpr int64_t DSV4_CSA_RATIO = 4; +static constexpr int64_t DSV4_HCA_RATIO = 128; + +static ggml_tensor * dsv4_hc_affine( + ggml_context * ctx, + ggml_tensor * x, + ggml_tensor * scale, + ggml_tensor * base) { + x = ggml_mul(ctx, x, scale); + x = ggml_add(ctx, x, base); + return x; +} + +ggml_tensor * llama_model_deepseek4::graph::build_hc_weighted_sum( + ggml_tensor * x, + ggml_tensor * weights) const { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = x->ne[2]; + + ggml_tensor * acc = nullptr; + for (int64_t ih = 0; ih < hc; ++ih) { + ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih*x->nb[1]); + ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih*weights->nb[0]); + + ggml_tensor * cur = ggml_mul(ctx0, xh, wh); + acc = acc ? ggml_add(ctx0, acc, cur) : cur; + } + + return acc; +} + +ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn( + ggml_tensor * comb, + int il) const { + GGML_UNUSED(il); + + // comb is [dst_hc, src_hc, n_tokens]. Sinkhorn follows the reference: + // row softmax over dst, one column normalization, then repeated row/column normalization. + comb = ggml_soft_max(ctx0, comb); + + ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1); + eps = ggml_fill(ctx0, eps, hparams.dsv4_hc_eps); + + comb = ggml_add(ctx0, comb, eps); + + auto norm_cols = [&]() { + ggml_tensor * comb_src_dst = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3)); + ggml_tensor * col_sum = ggml_sum_rows(ctx0, comb_src_dst); + col_sum = ggml_add(ctx0, col_sum, eps); + col_sum = ggml_permute(ctx0, col_sum, 1, 0, 2, 3); + comb = ggml_div(ctx0, comb, col_sum); + }; + + auto norm_rows = [&]() { + ggml_tensor * row_sum = ggml_sum_rows(ctx0, comb); + row_sum = ggml_add(ctx0, row_sum, eps); + comb = ggml_div(ctx0, comb, row_sum); + }; + + norm_cols(); + for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) { + norm_rows(); + norm_cols(); + } + + return comb; +} + +ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + ggml_tensor ** post, + ggml_tensor ** comb, + int il) const { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc*n_embd; + const int64_t hc_mix_dim = (2 + hc)*hc; + const int64_t nt = x->ne[2]; + + GGML_ASSERT(hc == 4); + GGML_ASSERT(hc_fn->ne[1] == hc_mix_dim); + + ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt); + ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps); + ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); + cb(mixes, "hc_mixes", il); + + ggml_tensor * scale_pre = dsv4_view_1d(ctx0, hc_scale, 1, 0); + ggml_tensor * scale_post = dsv4_view_1d(ctx0, hc_scale, 1, 1); + ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2); + + ggml_tensor * base_pre = dsv4_view_1d(ctx0, hc_base, hc, 0); + ggml_tensor * base_post = dsv4_view_1d(ctx0, hc_base, hc, hc); + ggml_tensor * base_comb = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc); + + ggml_tensor * pre = dsv4_view_2d(ctx0, mixes, hc, nt, 0); + pre = dsv4_hc_affine(ctx0, pre, scale_pre, base_pre); + pre = ggml_sigmoid(ctx0, pre); + pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps); + cb(pre, "hc_pre", il); + + *post = dsv4_view_2d(ctx0, mixes, hc, nt, hc); + *post = dsv4_hc_affine(ctx0, *post, scale_post, base_post); + *post = ggml_sigmoid(ctx0, *post); + *post = ggml_scale(ctx0, *post, 2.0f); + cb(*post, "hc_post", il); + + *comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc); + *comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb); + *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt); + *comb = build_hc_sinkhorn(*comb, il); + cb(*comb, "hc_comb", il); + + return build_hc_weighted_sum(x, pre); +} + +ggml_tensor * llama_model_deepseek4::graph::build_hc_post( + ggml_tensor * x, + ggml_tensor * residual, + ggml_tensor * post, + ggml_tensor * comb, + int il) const { + GGML_UNUSED(il); + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = x->ne[1]; + + ggml_tensor * out = nullptr; + for (int64_t dst = 0; dst < hc; ++dst) { + ggml_tensor * post_dst = ggml_view_2d(ctx0, post, 1, nt, post->nb[1], dst*post->nb[0]); + ggml_tensor * cur = ggml_mul(ctx0, x, post_dst); + + for (int64_t src = 0; src < hc; ++src) { + ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src*residual->nb[1]); + ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2], dst*comb->nb[0] + src*comb->nb[1]); + cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst)); + } + + cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt); + out = out ? ggml_concat(ctx0, out, cur, 1) : cur; + } + + return out; +} + +ggml_tensor * llama_model_deepseek4::graph::build_hc_head( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base) const { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc*n_embd; + const int64_t nt = x->ne[2]; + + ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt); + ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps); + ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); + cb(mixes, "hc_head_mixes", -1); + + ggml_tensor * pre = dsv4_hc_affine(ctx0, mixes, hc_scale, hc_base); + pre = ggml_sigmoid(ctx0, pre); + pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps); + cb(pre, "hc_head_pre", -1); + + return build_hc_weighted_sum(x, pre); +} + +ggml_tensor * llama_model_deepseek4::graph::build_hca_compressed_kv_from_state( + ggml_tensor * kv_state, + ggml_tensor * score_state, + ggml_tensor * state_read_idxs, + ggml_tensor * comp_pos, + ggml_tensor * norm, + int64_t n_embd_head, + const char * name, + int il) const { + const int64_t n_embd_head_rope = hparams.n_rot(); + const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope; + const int64_t n_blocks = comp_pos ? comp_pos->ne[0] : 0; + + GGML_ASSERT(n_blocks > 0); + GGML_ASSERT(state_read_idxs); + GGML_ASSERT(state_read_idxs->ne[0] == DSV4_HCA_RATIO*n_blocks); + GGML_ASSERT(n_embd_head >= n_embd_head_rope); + + ggml_tensor * kv = ggml_get_rows(ctx0, kv_state, state_read_idxs); + kv = ggml_reshape_3d(ctx0, kv, n_embd_head, DSV4_HCA_RATIO, n_blocks); + cb(kv, name, il); + + ggml_tensor * score = ggml_get_rows(ctx0, score_state, state_read_idxs); + score = ggml_reshape_3d(ctx0, score, n_embd_head, DSV4_HCA_RATIO, n_blocks); + cb(score, name, il); + + ggml_tensor * values = ggml_cont(ctx0, ggml_permute(ctx0, kv, 1, 0, 2, 3)); + ggml_tensor * scores = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); + + ggml_tensor * weights = ggml_soft_max(ctx0, scores); + ggml_tensor * comp = ggml_mul(ctx0, values, weights); + comp = ggml_sum_rows(ctx0, comp); + comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); + cb(comp, name, il); + + comp = build_norm(comp, norm, nullptr, LLM_NORM_RMS, il); + cb(comp, name, il); + + ggml_tensor * comp_nope = ggml_view_3d(ctx0, comp, n_embd_head_nope, 1, n_blocks, + ggml_row_size(comp->type, n_embd_head), + ggml_row_size(comp->type, n_embd_head), + 0); + ggml_tensor * comp_pe = ggml_view_3d(ctx0, comp, n_embd_head_rope, 1, n_blocks, + ggml_row_size(comp->type, n_embd_head), + ggml_row_size(comp->type, n_embd_head), + ggml_row_size(comp->type, n_embd_head_nope)); + + comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig, + hparams.dsv4_compress_rope_base, freq_scale, ext_factor, + dsv4_rope_attn_factor(freq_scale, ext_factor), beta_fast, beta_slow); + cb(comp_pe, name, il); + + comp = ggml_concat(ctx0, comp_nope, comp_pe, 0); + cb(comp, name, il); + + return comp; +} + +ggml_tensor * llama_model_deepseek4::graph::build_overlap_compressed_kv_from_state( + ggml_tensor * kv_state, + ggml_tensor * score_state, + ggml_tensor * state_read_idxs, + ggml_tensor * comp_pos, + ggml_tensor * norm, + int64_t ratio, + int64_t n_embd_head, + const char * name, + int il) const { + const int64_t n_embd_head_rope = hparams.n_rot(); + const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope; + const int64_t n_blocks = comp_pos ? comp_pos->ne[0] : 0; + + GGML_ASSERT(n_blocks > 0); + GGML_ASSERT(state_read_idxs); + GGML_ASSERT(state_read_idxs->ne[0] == 2*ratio*n_blocks); + GGML_ASSERT(kv_state->ne[0] == 2*n_embd_head); + GGML_ASSERT(score_state->ne[0] == 2*n_embd_head); + GGML_ASSERT(n_embd_head >= n_embd_head_rope); + + kv_state = dsv4_append_zero_row(ctx0, kv_state, false); + score_state = dsv4_append_zero_row(ctx0, score_state, true); + + ggml_tensor * prev_idxs = dsv4_view_1d(ctx0, state_read_idxs, ratio*n_blocks, 0); + ggml_tensor * cur_idxs = dsv4_view_1d(ctx0, state_read_idxs, ratio*n_blocks, ratio*n_blocks); + + ggml_tensor * kv_prev = ggml_get_rows(ctx0, kv_state, prev_idxs); + kv_prev = ggml_cont(ctx0, ggml_view_2d(ctx0, kv_prev, n_embd_head, ratio*n_blocks, kv_prev->nb[1], 0)); + kv_prev = ggml_reshape_3d(ctx0, kv_prev, n_embd_head, ratio, n_blocks); + cb(kv_prev, name, il); + + ggml_tensor * score_prev = ggml_get_rows(ctx0, score_state, prev_idxs); + score_prev = ggml_cont(ctx0, ggml_view_2d(ctx0, score_prev, n_embd_head, ratio*n_blocks, score_prev->nb[1], 0)); + score_prev = ggml_reshape_3d(ctx0, score_prev, n_embd_head, ratio, n_blocks); + cb(score_prev, name, il); + + ggml_tensor * kv_cur = ggml_get_rows(ctx0, kv_state, cur_idxs); + kv_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, kv_cur, n_embd_head, ratio*n_blocks, kv_cur->nb[1], + ggml_row_size(kv_cur->type, n_embd_head))); + kv_cur = ggml_reshape_3d(ctx0, kv_cur, n_embd_head, ratio, n_blocks); + + ggml_tensor * score_cur = ggml_get_rows(ctx0, score_state, cur_idxs); + score_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, score_cur, n_embd_head, ratio*n_blocks, score_cur->nb[1], + ggml_row_size(score_cur->type, n_embd_head))); + score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks); + + ggml_tensor * values = ggml_concat(ctx0, kv_prev, kv_cur, 1); + ggml_tensor * scores = ggml_concat(ctx0, score_prev, score_cur, 1); + + values = ggml_cont(ctx0, ggml_permute(ctx0, values, 1, 0, 2, 3)); + scores = ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3)); + + ggml_tensor * weights = ggml_soft_max(ctx0, scores); + ggml_tensor * comp = ggml_mul(ctx0, values, weights); + comp = ggml_sum_rows(ctx0, comp); + comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); + cb(comp, name, il); + + comp = build_norm(comp, norm, nullptr, LLM_NORM_RMS, il); + cb(comp, name, il); + + ggml_tensor * comp_nope = ggml_view_3d(ctx0, comp, n_embd_head_nope, 1, n_blocks, + ggml_row_size(comp->type, n_embd_head), + ggml_row_size(comp->type, n_embd_head), + 0); + ggml_tensor * comp_pe = ggml_view_3d(ctx0, comp, n_embd_head_rope, 1, n_blocks, + ggml_row_size(comp->type, n_embd_head), + ggml_row_size(comp->type, n_embd_head), + ggml_row_size(comp->type, n_embd_head_nope)); + + comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig, + hparams.dsv4_compress_rope_base, freq_scale, ext_factor, + dsv4_rope_attn_factor(freq_scale, ext_factor), beta_fast, beta_slow); + cb(comp_pe, name, il); + + comp = ggml_concat(ctx0, comp_nope, comp_pe, 0); + cb(comp, name, il); + + return comp; +} + +ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k( + const llama_model & model, + llm_graph_input_dsv4 * inp_dsv4, + ggml_tensor * qr, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int il) const { + const auto & layer = model.layers[il]; + const auto & inp_lid = inp_dsv4->get_lid(); + const int64_t n_embd_indexer_head = hparams.indexer_head_size; + const int64_t n_embd_indexer_head_rope = hparams.n_rot(); + const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope; + const int64_t n_indexer_head = hparams.indexer_n_head; + const int64_t nt = cur->ne[1]; + + GGML_ASSERT(inp_lid.kq_mask); + GGML_ASSERT(inp_lid.k_rot); + GGML_ASSERT(n_embd_indexer_head >= n_embd_indexer_head_rope); + + ggml_tensor * indexer_q = build_lora_mm(layer.indexer_attn_q_b, qr); + indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, nt); + cb(indexer_q, "lid_q", il); + + ggml_tensor * indexer_q_nope = ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, nt, + ggml_row_size(indexer_q->type, n_embd_indexer_head), + ggml_row_size(indexer_q->type, n_embd_indexer_head)*n_indexer_head, + 0); + ggml_tensor * indexer_q_pe = ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, nt, + ggml_row_size(indexer_q->type, n_embd_indexer_head), + ggml_row_size(indexer_q->type, n_embd_indexer_head)*n_indexer_head, + ggml_row_size(indexer_q->type, n_embd_indexer_head_nope)); + + indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_embd_indexer_head_rope, + rope_type, n_ctx_orig, hparams.dsv4_compress_rope_base, freq_scale, + ext_factor, dsv4_rope_attn_factor(freq_scale, ext_factor), beta_fast, beta_slow); + cb(indexer_q_pe, "lid_q_pe", il); + + indexer_q = ggml_concat(ctx0, indexer_q_nope, indexer_q_pe, 0); + indexer_q = llama_mul_mat_hadamard(ctx0, indexer_q, inp_lid.k_rot); + cb(indexer_q, "lid_q_rot", il); + + ggml_tensor * indexer_weights = build_lora_mm(layer.indexer_proj, cur); + indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f/sqrtf(float(n_embd_indexer_head*n_indexer_head))); + cb(indexer_weights, "lid_weights", il); + + ggml_tensor * indexer_k = inp_dsv4->mctx->get_lid()->get_k(ctx0, il); + const int64_t n_lid = inp_lid.kq_mask->ne[0]; + GGML_ASSERT(n_lid > 0); + GGML_ASSERT(n_lid <= indexer_k->ne[2]); + + indexer_k = ggml_view_4d(ctx0, indexer_k, + indexer_k->ne[0], indexer_k->ne[1], n_lid, indexer_k->ne[3], + indexer_k->nb[1], indexer_k->nb[2], indexer_k->nb[3], 0); + cb(indexer_k, "lid_k", il); + + const int64_t n_stream = indexer_k->ne[3]; + indexer_q = ggml_view_4d(ctx0, indexer_q, + indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, + indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0); + indexer_weights = ggml_view_4d(ctx0, indexer_weights, + indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, + indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0); + + indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); + cb(indexer_q, "lid_q", il); + indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3); + cb(indexer_k, "lid_k", il); + + ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); + cb(indexer_kq, "lid_kq", il); + + indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); + cb(indexer_kq, "lid_kq", il); + + ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq); + indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); + indexer_score = ggml_sum_rows(ctx0, indexer_score); + indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); + cb(indexer_score, "lid_score", il); + + indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask); + cb(indexer_score, "lid_score_masked", il); + + const uint32_t n_top_k = indexer_score->ne[0] < hparams.indexer_top_k ? indexer_score->ne[0] : hparams.indexer_top_k; + ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k)); + cb(top_k, "lid_top_k", il); + + return top_k; +} + +ggml_tensor * llama_model_deepseek4::graph::build_top_k_mask( + ggml_tensor * kq_mask, + ggml_tensor * top_k, + const char * name, + int il) const { + GGML_ASSERT(kq_mask); + GGML_ASSERT(top_k); + + ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); + kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], + kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); + + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, + top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); + + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); + + ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); + kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, + kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], + kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); + + kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); + cb(kq_mask_top_k, name, il); + + return kq_mask_top_k; +} + +ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( + const llama_model & model, + llm_graph_input_dsv4 * inp_dsv4, + llm_graph_input_dsv4_raw * inp_attn, + ggml_tensor * q, + ggml_tensor * kv, + ggml_tensor * qr, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * sinks, + float kq_scale, + int il) const { + const auto & inp_csa = inp_dsv4->get_csa(); + GGML_ASSERT(inp_csa.kq_mask); + + ggml_tensor * top_k = build_lid_top_k(model, inp_dsv4, qr, cur, inp_pos, il); + + ggml_tensor * k_rot = inp_attn->self_k_rot; + if (k_rot) { + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + kv = llama_mul_mat_hadamard(ctx0, kv, k_rot); + } + + ggml_build_forward_expand(gf, q); + ggml_build_forward_expand(gf, kv); + + const llama_kv_cache_dsv4_raw_context * mctx_raw = inp_attn->mctx; + + ggml_build_forward_expand(gf, mctx_raw->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il)); + + ggml_tensor * raw_k = mctx_raw->get_k(ctx0, il); + cb(raw_k, "csa_raw_k", il); + + ggml_tensor * csa_k = inp_dsv4->mctx->get_csa()->get_k(ctx0, il); + const int64_t n_csa = inp_csa.kq_mask->ne[0]; + GGML_ASSERT(n_csa > 0); + GGML_ASSERT(n_csa <= csa_k->ne[2]); + + csa_k = ggml_view_4d(ctx0, csa_k, + csa_k->ne[0], csa_k->ne[1], n_csa, csa_k->ne[3], + csa_k->nb[1], csa_k->nb[2], csa_k->nb[3], 0); + cb(csa_k, "csa_comp_k", il); + + ggml_tensor * k_all = ggml_concat(ctx0, raw_k, csa_k, 2); + cb(k_all, "csa_k_all", il); + + ggml_tensor * raw_mask = inp_attn->get_kq_mask(); + ggml_tensor * csa_mask = build_top_k_mask(inp_csa.kq_mask, top_k, "csa_top_k_mask", il); + + ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); + cb(kq_mask, "csa_lid_kq_mask", il); + + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } + cb(out, "attn_csa_lid", il); + + return out; +} + +ggml_tensor * llama_model_deepseek4::graph::build_hca_attention( + llm_graph_input_dsv4 * inp_dsv4, + llm_graph_input_dsv4_raw * inp_attn, + ggml_tensor * q, + ggml_tensor * kv, + ggml_tensor * sinks, + float kq_scale, + int il) const { + const auto & inp_hca = inp_dsv4->get_hca(); + GGML_ASSERT(inp_hca.kq_mask); + + ggml_tensor * k_rot = inp_attn->self_k_rot; + if (k_rot) { + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + kv = llama_mul_mat_hadamard(ctx0, kv, k_rot); + } + + ggml_build_forward_expand(gf, q); + ggml_build_forward_expand(gf, kv); + + const llama_kv_cache_dsv4_raw_context * mctx_raw = inp_attn->mctx; + + ggml_build_forward_expand(gf, mctx_raw->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il)); + + ggml_tensor * raw_k = mctx_raw->get_k(ctx0, il); + cb(raw_k, "hca_raw_k", il); + + ggml_tensor * hca_k = inp_dsv4->mctx->get_hca()->get_k(ctx0, il); + const int64_t n_hca = inp_hca.kq_mask->ne[0]; + GGML_ASSERT(n_hca > 0); + GGML_ASSERT(n_hca <= hca_k->ne[2]); + + hca_k = ggml_view_4d(ctx0, hca_k, + hca_k->ne[0], hca_k->ne[1], n_hca, hca_k->ne[3], + hca_k->nb[1], hca_k->nb[2], hca_k->nb[3], 0); + cb(hca_k, "hca_comp_k", il); + + ggml_tensor * k_all = ggml_concat(ctx0, raw_k, hca_k, 2); + cb(k_all, "hca_k_all", il); + + ggml_tensor * raw_mask = inp_attn->get_kq_mask(); + ggml_tensor * hca_mask = inp_hca.kq_mask; + + ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0); + cb(kq_mask, "hca_kq_mask", il); + + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } + cb(out, "attn_hca", il); + + return out; +} + +ggml_tensor * llama_model_deepseek4::graph::build_raw_attention( + llm_graph_input_dsv4_raw * inp_attn, + ggml_tensor * q, + ggml_tensor * kv, + ggml_tensor * sinks, + float kq_scale, + int il) const { + GGML_ASSERT(hparams.is_swa(il)); + + ggml_tensor * k_rot = inp_attn->self_k_rot; + + if (k_rot) { + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + kv = llama_mul_mat_hadamard(ctx0, kv, k_rot); + } + + ggml_build_forward_expand(gf, q); + ggml_build_forward_expand(gf, kv); + + const llama_kv_cache_dsv4_raw_context * mctx_cur = inp_attn->mctx; + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il)); + + ggml_tensor * kq_mask = inp_attn->get_kq_mask(); + + ggml_tensor * k = mctx_cur->get_k(ctx0, il); + + ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } + cb(out, "attn_raw", il); + + return out; +} + +ggml_tensor * llama_model_deepseek4::graph::build_attention( + const llama_model & model, + llm_graph_input_dsv4 * inp_dsv4, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int il) const { + const auto & layer = model.layers[il]; + llm_graph_input_dsv4_raw * inp_attn = inp_dsv4->get_raw(); + + const int64_t n_embd_head = hparams.n_embd_head_k(); + const int64_t n_embd_head_rope = hparams.n_rot(); + const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope; + const int64_t n_groups = hparams.dsv4_o_group_count; + const int64_t n_heads_group = n_head / n_groups; + const int64_t o_lora_rank = hparams.dsv4_o_lora_rank; + const int64_t o_group_dim = n_heads_group*n_embd_head; + const int64_t nt = cur->ne[1]; + + GGML_ASSERT(n_embd_head == n_embd_head_v); + GGML_ASSERT(n_head % n_groups == 0); + + const bool use_compress_rope = hparams.dsv4_compress_ratios[il] != 0; + const float freq_base_l = use_compress_rope ? hparams.dsv4_compress_rope_base : freq_base; + const float freq_scale_l = use_compress_rope ? freq_scale : 1.0f; + const float ext_factor_l = use_compress_rope ? ext_factor : 0.0f; + const float attn_factor_l = dsv4_rope_attn_factor(freq_scale_l, ext_factor_l); + const float beta_fast_l = use_compress_rope ? beta_fast : 0.0f; + const float beta_slow_l = use_compress_rope ? beta_slow : 0.0f; + const int32_t n_ctx_orig_l = use_compress_rope ? n_ctx_orig : 0; + + ggml_tensor * qr = build_lora_mm(layer.wq_a, cur); + cb(qr, "qr", il); + + qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(qr, "qr_norm", il); + + ggml_tensor * q = build_lora_mm(layer.wq_b, qr); + q = ggml_reshape_3d(ctx0, q, n_embd_head, n_head, nt); + q = ggml_rms_norm(ctx0, q, norm_rms_eps); + cb(q, "q_norm", il); + + ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_nope, n_head, nt, + ggml_row_size(q->type, n_embd_head), + ggml_row_size(q->type, n_embd_head)*n_head, + 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_embd_head_rope, n_head, nt, + ggml_row_size(q->type, n_embd_head), + ggml_row_size(q->type, n_embd_head)*n_head, + ggml_row_size(q->type, n_embd_head_nope)); + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, + freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l); + cb(q_pe, "q_pe", il); + q = ggml_concat(ctx0, q_nope, q_pe, 0); + cb(q, "q", il); + + ggml_tensor * kv = build_lora_mm(layer.wkv, cur); + kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il); + kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, nt); + cb(kv, "kv_norm", il); + + ggml_tensor * kv_nope = ggml_view_3d(ctx0, kv, n_embd_head_nope, 1, nt, + ggml_row_size(kv->type, n_embd_head), + ggml_row_size(kv->type, n_embd_head), + 0); + ggml_tensor * kv_pe = ggml_view_3d(ctx0, kv, n_embd_head_rope, 1, nt, + ggml_row_size(kv->type, n_embd_head), + ggml_row_size(kv->type, n_embd_head), + ggml_row_size(kv->type, n_embd_head_nope)); + kv_pe = ggml_rope_ext(ctx0, kv_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, + freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l); + cb(kv_pe, "kv_pe", il); + kv = ggml_concat(ctx0, kv_nope, kv_pe, 0); + cb(kv, "kv", il); + + const int64_t ratio = hparams.dsv4_compress_ratios[il]; + + ggml_tensor * hca_state_kv = nullptr; + ggml_tensor * hca_state_score = nullptr; + if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) { + hca_state_kv = build_lora_mm(layer.attn_comp_wkv, cur); + cb(hca_state_kv, "hca_state_kv", il); + + hca_state_score = build_lora_mm(layer.attn_comp_wgate, cur); + cb(hca_state_score, "hca_state_score", il); + + ggml_tensor * ape = layer.attn_comp_ape; + + ggml_tensor * ape_rows = ggml_get_rows(ctx0, ape, inp_dsv4->get_hca().state_pos); + hca_state_score = ggml_add(ctx0, hca_state_score, ape_rows); + cb(hca_state_score, "hca_state_score_ape", il); + + } + + if (ratio == DSV4_CSA_RATIO && inp_dsv4->get_csa().state_pos) { + ggml_tensor * csa_state_kv = build_lora_mm(layer.attn_comp_wkv, cur); + cb(csa_state_kv, "csa_state_kv", il); + + ggml_tensor * csa_state_score = build_lora_mm(layer.attn_comp_wgate, cur); + cb(csa_state_score, "csa_state_score", il); + + ggml_tensor * csa_ape = layer.attn_comp_ape; + + ggml_tensor * csa_ape_rows = ggml_get_rows(ctx0, csa_ape, inp_dsv4->get_csa().state_pos); + csa_state_score = ggml_add(ctx0, csa_state_score, csa_ape_rows); + cb(csa_state_score, "csa_state_score_ape", il); + + GGML_ASSERT(inp_dsv4->get_csa().state_write_idxs); + + ggml_tensor * csa_source_kv = ggml_concat(ctx0, + inp_dsv4->mctx->get_csa_state()->get_kv(ctx0, il), csa_state_kv, 1); + ggml_tensor * csa_source_score = ggml_concat(ctx0, + inp_dsv4->mctx->get_csa_state()->get_score(ctx0, il), csa_state_score, 1); + + ggml_tensor * kv_comp_csa_state = build_overlap_compressed_kv_from_state( + csa_source_kv, + csa_source_score, + inp_dsv4->get_csa().state_read_idxs, + inp_dsv4->get_csa().state_write_pos, + layer.attn_comp_norm, + DSV4_CSA_RATIO, + n_embd_head, + "csa_state_compress", + il); + + if (inp_dsv4->get_csa().k_rot) { + kv_comp_csa_state = llama_mul_mat_hadamard(ctx0, kv_comp_csa_state, inp_dsv4->get_csa().k_rot); + cb(kv_comp_csa_state, "csa_state_compress_rot", il); + } + + ggml_build_forward_expand(gf, inp_dsv4->mctx->get_csa()->cpy_k(ctx0, + kv_comp_csa_state, inp_dsv4->get_csa().state_write_idxs, il)); + + csa_state_kv = dsv4_with_zero_dep(ctx0, csa_state_kv, kv_comp_csa_state); + csa_state_score = dsv4_with_zero_dep(ctx0, csa_state_score, kv_comp_csa_state); + + ggml_tensor * csa_persist_kv = ggml_get_rows(ctx0, csa_state_kv, inp_dsv4->get_csa().state_persist_src_idxs); + ggml_tensor * csa_persist_score = ggml_get_rows(ctx0, csa_state_score, inp_dsv4->get_csa().state_persist_src_idxs); + + csa_state_kv = inp_dsv4->mctx->get_csa_state()->cpy_kv(ctx0, + csa_persist_kv, inp_dsv4->get_csa().state_persist_dst_idxs, il); + csa_state_score = inp_dsv4->mctx->get_csa_state()->cpy_score(ctx0, + csa_persist_score, inp_dsv4->get_csa().state_persist_dst_idxs, il); + + ggml_build_forward_expand(gf, csa_state_kv); + ggml_build_forward_expand(gf, csa_state_score); + + ggml_tensor * lid_state_kv = build_lora_mm(layer.indexer_comp_wkv, cur); + cb(lid_state_kv, "lid_state_kv", il); + + ggml_tensor * lid_state_score = build_lora_mm(layer.indexer_comp_wgate, cur); + cb(lid_state_score, "lid_state_score", il); + + ggml_tensor * lid_ape = layer.indexer_comp_ape; + + ggml_tensor * lid_ape_rows = ggml_get_rows(ctx0, lid_ape, inp_dsv4->get_lid().state_pos); + lid_state_score = ggml_add(ctx0, lid_state_score, lid_ape_rows); + cb(lid_state_score, "lid_state_score_ape", il); + + GGML_ASSERT(inp_dsv4->get_lid().state_write_idxs); + + ggml_tensor * lid_source_kv = ggml_concat(ctx0, + inp_dsv4->mctx->get_lid_state()->get_kv(ctx0, il), lid_state_kv, 1); + ggml_tensor * lid_source_score = ggml_concat(ctx0, + inp_dsv4->mctx->get_lid_state()->get_score(ctx0, il), lid_state_score, 1); + + ggml_tensor * kv_comp_lid_state = build_overlap_compressed_kv_from_state( + lid_source_kv, + lid_source_score, + inp_dsv4->get_lid().state_read_idxs, + inp_dsv4->get_lid().state_write_pos, + layer.indexer_comp_norm, + DSV4_CSA_RATIO, + hparams.indexer_head_size, + "lid_state_compress", + il); + + if (inp_dsv4->get_lid().k_rot) { + kv_comp_lid_state = llama_mul_mat_hadamard(ctx0, kv_comp_lid_state, inp_dsv4->get_lid().k_rot); + cb(kv_comp_lid_state, "lid_state_compress_rot", il); + } + + ggml_build_forward_expand(gf, inp_dsv4->mctx->get_lid()->cpy_k(ctx0, + kv_comp_lid_state, inp_dsv4->get_lid().state_write_idxs, il)); + + lid_state_kv = dsv4_with_zero_dep(ctx0, lid_state_kv, kv_comp_lid_state); + lid_state_score = dsv4_with_zero_dep(ctx0, lid_state_score, kv_comp_lid_state); + + ggml_tensor * lid_persist_kv = ggml_get_rows(ctx0, lid_state_kv, inp_dsv4->get_lid().state_persist_src_idxs); + ggml_tensor * lid_persist_score = ggml_get_rows(ctx0, lid_state_score, inp_dsv4->get_lid().state_persist_src_idxs); + + lid_state_kv = inp_dsv4->mctx->get_lid_state()->cpy_kv(ctx0, + lid_persist_kv, inp_dsv4->get_lid().state_persist_dst_idxs, il); + lid_state_score = inp_dsv4->mctx->get_lid_state()->cpy_score(ctx0, + lid_persist_score, inp_dsv4->get_lid().state_persist_dst_idxs, il); + + ggml_build_forward_expand(gf, lid_state_kv); + ggml_build_forward_expand(gf, lid_state_score); + } + + ggml_tensor * hca_state_dep = nullptr; + if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_write_idxs) { + GGML_ASSERT(hca_state_kv); + GGML_ASSERT(hca_state_score); + + ggml_tensor * hca_source_kv = ggml_concat(ctx0, + inp_dsv4->mctx->get_hca_state()->get_kv(ctx0, il), hca_state_kv, 1); + ggml_tensor * hca_source_score = ggml_concat(ctx0, + inp_dsv4->mctx->get_hca_state()->get_score(ctx0, il), hca_state_score, 1); + + ggml_tensor * kv_comp_hca = build_hca_compressed_kv_from_state( + hca_source_kv, + hca_source_score, + inp_dsv4->get_hca().state_read_idxs, + inp_dsv4->get_hca().state_write_pos, + layer.attn_comp_norm, + n_embd_head, + "hca_state_compress", + il); + + if (inp_dsv4->get_hca().k_rot) { + kv_comp_hca = llama_mul_mat_hadamard(ctx0, kv_comp_hca, inp_dsv4->get_hca().k_rot); + cb(kv_comp_hca, "hca_state_compress_rot", il); + } + + ggml_build_forward_expand(gf, inp_dsv4->mctx->get_hca()->cpy_k(ctx0, + kv_comp_hca, inp_dsv4->get_hca().state_write_idxs, il)); + hca_state_dep = kv_comp_hca; + } + + if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) { + GGML_ASSERT(hca_state_kv); + GGML_ASSERT(hca_state_score); + + hca_state_kv = dsv4_with_zero_dep(ctx0, hca_state_kv, hca_state_dep); + hca_state_score = dsv4_with_zero_dep(ctx0, hca_state_score, hca_state_dep); + + ggml_tensor * hca_persist_kv = ggml_get_rows(ctx0, hca_state_kv, inp_dsv4->get_hca().state_persist_src_idxs); + ggml_tensor * hca_persist_score = ggml_get_rows(ctx0, hca_state_score, inp_dsv4->get_hca().state_persist_src_idxs); + + hca_state_kv = inp_dsv4->mctx->get_hca_state()->cpy_kv(ctx0, + hca_persist_kv, inp_dsv4->get_hca().state_persist_dst_idxs, il); + hca_state_score = inp_dsv4->mctx->get_hca_state()->cpy_score(ctx0, + hca_persist_score, inp_dsv4->get_hca().state_persist_dst_idxs, il); + + ggml_build_forward_expand(gf, hca_state_kv); + ggml_build_forward_expand(gf, hca_state_score); + } + + ggml_tensor * out = nullptr; + if (ratio == DSV4_CSA_RATIO && + inp_dsv4->get_csa().kq_mask && + inp_dsv4->get_lid().kq_mask && + inp_dsv4->get_lid().k_rot) { + out = build_csa_lid_attention(model, inp_dsv4, inp_attn, q, kv, qr, cur, inp_pos, layer.attn_sinks, + 1.0f/sqrtf(float(n_embd_head)), il); + } else if (ratio == DSV4_HCA_RATIO && + inp_dsv4->get_hca().kq_mask) { + out = build_hca_attention(inp_dsv4, inp_attn, q, kv, layer.attn_sinks, + 1.0f/sqrtf(float(n_embd_head)), il); + } else { + out = build_raw_attention(inp_attn, q, kv, layer.attn_sinks, + 1.0f/sqrtf(float(n_embd_head)), il); + } + + out = ggml_reshape_3d(ctx0, out, n_embd_head, n_head, nt); + ggml_tensor * out_nope = ggml_view_3d(ctx0, out, n_embd_head_nope, n_head, nt, + ggml_row_size(out->type, n_embd_head), + ggml_row_size(out->type, n_embd_head)*n_head, + 0); + ggml_tensor * out_pe = ggml_view_3d(ctx0, out, n_embd_head_rope, n_head, nt, + ggml_row_size(out->type, n_embd_head), + ggml_row_size(out->type, n_embd_head)*n_head, + ggml_row_size(out->type, n_embd_head_nope)); + out_pe = ggml_rope_ext_back(ctx0, out_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, + freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l); + out = ggml_concat(ctx0, out_nope, out_pe, 0); + cb(out, "attn_derope", il); + + out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, nt); + out = ggml_permute(ctx0, out, 0, 2, 1, 3); + ggml_tensor * oa = ggml_mul_mat(ctx0, ggml_reshape_3d(ctx0, layer.wo_a, layer.wo_a->ne[0], o_lora_rank, n_groups), out); + cb(oa, "attn_wo_a", il); + oa = ggml_permute(ctx0, oa, 0, 2, 1, 3); + oa = ggml_cont_2d(ctx0, oa, o_lora_rank*n_groups, nt); + + out = build_lora_mm(layer.wo_b, oa); + cb(out, "attn_out", il); + + return out; +} + +llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + ggml_tensor * cur; + + 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(); + llm_graph_input_dsv4 * inp_dsv4 = build_inp_dsv4(); + llm_graph_input_dsv4_raw * inp_attn = inp_dsv4->get_raw(); + ggml_build_forward_expand(gf, inp_attn->self_kq_mask); + + const int64_t hc = hparams.dsv4_hc_mult; + ggml_tensor * inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens); + inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1); + cb(inpL, "hc_init", -1); + + for (int il = 0; il < n_layer; ++il) { + ggml_tensor * residual = inpL; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + + cur = build_hc_pre(inpL, + model.layers[il].hc_attn_fn, + model.layers[il].hc_attn_scale, + model.layers[il].hc_attn_base, + &post, &comb, il); + cb(cur, "hc_attn_pre", il); + + cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + cur = build_attention(model, inp_dsv4, cur, inp_pos, il); + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "hc_attn_post", il); + + residual = inpL; + cur = build_hc_pre(inpL, + model.layers[il].hc_ffn_fn, + model.layers[il].hc_ffn_scale, + model.layers[il].hc_ffn_base, + &post, &comb, il); + cb(cur, "hc_ffn_pre", il); + + cur = build_norm(cur, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + const auto & layer = model.layers[il]; + ggml_tensor * selected_experts = nullptr; + ggml_tensor * exp_probs_b = layer.ffn_exp_probs_b; + if ((uint32_t) il < hparams.dsv4_hash_layer_count) { + selected_experts = ggml_get_rows(ctx0, layer.ffn_gate_tid2eid, res->t_inp_tokens); + exp_probs_b = nullptr; + } + + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + 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, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + selected_experts); + cb(moe_out, "ffn_moe_out", il); + + ggml_tensor * ffn_shexp = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + + inpL = build_hc_post(cur, residual, post, comb, il); + inpL = build_cvec(inpL, il); + cb(inpL, "l_out", il); + } + + if (inp_out_ids) { + ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); + flat = ggml_get_rows(ctx0, flat, inp_out_ids); + inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs); + } + + cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base); + cb(cur, "hc_head", -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); +} diff --git a/examples/talk-llama/models/delta-net-base.cpp b/examples/talk-llama/models/delta-net-base.cpp index ad9ce7714..ad6612647 100644 --- a/examples/talk-llama/models/delta-net-base.cpp +++ b/examples/talk-llama/models/delta-net-base.cpp @@ -401,9 +401,9 @@ std::pair llm_build_delta_net_base::build_delta_ne // K=1: output carries the final state only. state s is 4D [S_v, S_v, H_v, n_seqs]. ggml_tensor * result = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, /*K=*/1); if (n_tokens == 1) { - cb(result, LLAMA_TENSOR_NAME_FGDN_AR, il); + res->add_fused_node({LLM_FUSED_OP_GDN_AR, result, il}); } else { - cb(result, LLAMA_TENSOR_NAME_FGDN_CH, il); + res->add_fused_node({LLM_FUSED_OP_GDN_CH, result, il}); } ggml_tensor * output = ggml_view_4d(ctx0, result, @@ -496,8 +496,8 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update)); } else { // [TAG_RECURRENT_ROLLBACK_SPLITS] - // TODO: this logic incorrectly assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are - // inside the same ubatch. currently with `split_equal()` this is not correct + // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside + // the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument const int64_t K = (int64_t) cparams.n_rs_seq + 1; @@ -566,9 +566,9 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( // state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output. ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K); if (n_seq_tokens > 1) { - cb(gdn_out, LLAMA_TENSOR_NAME_FGDN_CH, il); + res->add_fused_node({LLM_FUSED_OP_GDN_CH, gdn_out, il}); } else { - cb(gdn_out, LLAMA_TENSOR_NAME_FGDN_AR, il); + res->add_fused_node({LLM_FUSED_OP_GDN_AR, gdn_out, il}); } const int64_t attn_score_elems = S_v * H_v * n_seq_tokens * n_seqs; diff --git a/examples/talk-llama/models/dflash.cpp b/examples/talk-llama/models/dflash.cpp new file mode 100644 index 000000000..a7b4f4435 --- /dev/null +++ b/examples/talk-llama/models/dflash.cpp @@ -0,0 +1,276 @@ +#include "models.h" + +#include "llama-kv-cache.h" +#include "llama-kv-cache-iswa.h" + +void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { + + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) { + throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata"); + } + + hparams.n_embd_inp_enc_impl = (uint32_t) target_layer_ids.size() * hparams.n_embd; + + LLAMA_LOG_INFO("%s: DFlash extract_layers = [", __func__); + for (size_t i = 0; i < target_layer_ids.size(); ++i) { + LLAMA_LOG_INFO("%d%s", target_layer_ids[i], i + 1 < target_layer_ids.size() ? ", " : ""); + } + LLAMA_LOG_INFO("]\n"); + + // optional interleaved sliding-window attention with per-layer pattern array. + // DFlash has a single rope, so the SWA rope == main rope. + if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) { + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer()); + hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; + hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + } + + type = LLM_TYPE_UNKNOWN; +} + +void llama_model_dflash::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t n_embd_inp = hparams.n_embd_inp_enc(); + + fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); + output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0); + + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head }, 0); + layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), { n_embd, n_embd_k_gqa }, 0); + layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0); + + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0); + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0); + } +} + +std::unique_ptr llama_model_dflash::build_arch_graph(const llm_graph_params & params) const { + switch (params.gtype) { + case LLM_GRAPH_TYPE_ENCODER: + return std::make_unique>(*this, params); + case LLM_GRAPH_TYPE_DEFAULT: + case LLM_GRAPH_TYPE_DECODER: + return std::make_unique>(*this, params); + default: + GGML_ABORT("invalid graph type"); + }; +} + +template <> +ggml_tensor * llama_model_dflash::graph::build_inp_embd_enc() const { + auto inp_target = std::make_unique(hparams.n_embd_inp_enc()); + + inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp_enc(), n_tokens); + ggml_set_input(inp_target->embd); + + ggml_tensor * cur = inp_target->embd; + cb(cur, "inp_embd", -1); + + res->add_input(std::move(inp_target)); + + return cur; +} + +// DFlash Encoder: processes target model features through feature fusion layer +template <> +llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + ggml_tensor * cur = build_inp_embd_enc(); + + cur = build_lora_mm(model.fc, cur); + cb(cur, "fc_out", -1); + + cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1); + cb(cur, "enc_norm_out", -1); + + ggml_set_output(cur); + res->t_h_nextn = cur; + + ggml_build_forward_expand(gf, cur); +} + +// DFlash decoder, dual-mode by batch type: +// * embd batch -> fused target features: project + inject K/V into the cache. +// * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens +template <> +llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + ggml_tensor * inp_pos = build_inp_pos(); + + // optional iSWA: pick the matching attention input + const bool use_iswa = hparams.swa_type != LLAMA_SWA_TYPE_NONE; + + llm_graph_input_attn_kv * inp_attn = nullptr; + llm_graph_input_attn_kv_iswa * inp_attn_iswa = nullptr; + if (use_iswa) { + inp_attn_iswa = build_attn_inp_kv_iswa(); + } else { + inp_attn = build_attn_inp_kv(); + } + + const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + + // KV cache injection + if (ubatch.embd) { + auto inp = std::make_unique(n_embd); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * inp_g = inp->embd; + cb(inp_g, "inp_g_embeddings", -1); + + res->add_input(std::move(inp)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * Kcur = build_lora_mm(layer.wk, inp_g); + ggml_tensor * Vcur = build_lora_mm(layer.wv, inp_g); + + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + cb(Kcur, "Kcur_injected", il); + cb(Vcur, "Vcur_injected", il); + + if (use_iswa) { + // route each layer's K/V to its sub-cache: SWA layers -> sliding cache, full -> dense + const bool is_swa = hparams.is_swa(il); + const auto * kv = is_swa ? inp_attn_iswa->mctx->get_swa() : inp_attn_iswa->mctx->get_base(); + ggml_tensor * k_idxs = is_swa ? inp_attn_iswa->get_k_idxs_swa() : inp_attn_iswa->get_k_idxs(); + ggml_tensor * v_idxs = is_swa ? inp_attn_iswa->get_v_idxs_swa() : inp_attn_iswa->get_v_idxs(); + ggml_build_forward_expand(gf, kv->cpy_k(ctx0, Kcur, k_idxs, il)); + ggml_build_forward_expand(gf, kv->cpy_v(ctx0, Vcur, v_idxs, il)); + } else { + ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il)); + ggml_build_forward_expand(gf, inp_attn->mctx->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il)); + } + } + + res->t_embd = inp_g; + + ggml_build_forward_expand(gf, inp_g); + return; + } + + // tok_embd from the target model (shared via ctx_other) + auto * tok_embd = model.tok_embd; + if (tok_embd == nullptr) { + GGML_ASSERT(cparams.ctx_other != nullptr); + const auto * model_other = llama_get_model(cparams.ctx_other); + + GGML_ASSERT(model_other->tok_embd != nullptr && "DFlash decoder requires the target model's token embeddings"); + tok_embd = model_other->tok_embd; + } + + auto inp = std::make_unique(n_embd); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens); + cb(inpL, "inp_noise_embd", -1); + + res->add_input(std::move(inp)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(noise_norm, "noise_norm", il); + + ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm); + ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm); + ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm); + + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + + Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + // cache-aware, non-causal attention + ggml_tensor * cur = use_iswa + ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il) + : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + layer.ffn_up, NULL, NULL, + layer.ffn_gate, NULL, NULL, + layer.ffn_down, NULL, NULL, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "l_out", il); + + inpL = cur; + } + + ggml_tensor * cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + + res->t_embd = cur; + + // lm_head from the target model (shared via ctx_other) + auto * output = model.output; + if (output == nullptr) { + GGML_ASSERT(cparams.ctx_other != nullptr); + const auto * model_other = llama_get_model(cparams.ctx_other); + GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection"); + output = model_other->output; + } + + cur = build_lora_mm(output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/examples/talk-llama/models/models.h b/examples/talk-llama/models/models.h index 2ac8415a3..7a52e7bc1 100644 --- a/examples/talk-llama/models/models.h +++ b/examples/talk-llama/models/models.h @@ -1085,6 +1085,121 @@ struct llama_model_deepseek32 : public llama_model_base { }; +struct llama_model_deepseek4 : public llama_model_base { + llama_model_deepseek4(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + + ggml_tensor * build_hc_pre( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + ggml_tensor ** post, + ggml_tensor ** comb, + int il) const; + + ggml_tensor * build_hc_post( + ggml_tensor * x, + ggml_tensor * residual, + ggml_tensor * post, + ggml_tensor * comb, + int il) const; + + ggml_tensor * build_hc_head( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base) const; + + ggml_tensor * build_attention( + const llama_model & model, + llm_graph_input_dsv4 * inp_dsv4, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int il) const; + + ggml_tensor * build_hca_compressed_kv_from_state( + ggml_tensor * kv_state, + ggml_tensor * score_state, + ggml_tensor * state_read_idxs, + ggml_tensor * comp_pos, + ggml_tensor * norm, + int64_t n_embd_head, + const char * name, + int il) const; + + ggml_tensor * build_overlap_compressed_kv_from_state( + ggml_tensor * kv_state, + ggml_tensor * score_state, + ggml_tensor * state_read_idxs, + ggml_tensor * comp_pos, + ggml_tensor * norm, + int64_t ratio, + int64_t n_embd_head, + const char * name, + int il) const; + + ggml_tensor * build_lid_top_k( + const llama_model & model, + llm_graph_input_dsv4 * inp_dsv4, + ggml_tensor * qr, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int il) const; + + ggml_tensor * build_top_k_mask( + ggml_tensor * kq_mask, + ggml_tensor * top_k, + const char * name, + int il) const; + + ggml_tensor * build_csa_lid_attention( + const llama_model & model, + llm_graph_input_dsv4 * inp_dsv4, + llm_graph_input_dsv4_raw * inp_attn, + ggml_tensor * q, + ggml_tensor * kv, + ggml_tensor * qr, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * sinks, + float kq_scale, + int il) const; + + ggml_tensor * build_hca_attention( + llm_graph_input_dsv4 * inp_dsv4, + llm_graph_input_dsv4_raw * inp_attn, + ggml_tensor * q, + ggml_tensor * kv, + ggml_tensor * sinks, + float kq_scale, + int il) const; + + ggml_tensor * build_raw_attention( + llm_graph_input_dsv4_raw * inp_attn, + ggml_tensor * q, + ggml_tensor * kv, + ggml_tensor * sinks, + float kq_scale, + int il) const; + + ggml_tensor * build_hc_weighted_sum( + ggml_tensor * x, + ggml_tensor * weights) const; + + ggml_tensor * build_hc_sinkhorn( + ggml_tensor * comb, + int il) const; + }; + + 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) {} void load_arch_hparams(llama_model_loader & ml) override; @@ -1122,6 +1237,22 @@ struct llama_model_eagle3 : public llama_model_base { }; +struct llama_model_dflash : public llama_model_base { + llama_model_dflash(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + template + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + + ggml_tensor * build_inp_embd_enc() const; + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_mistral4 : public llama_model_deepseek2 { llama_model_mistral4(const struct llama_model_params & params) : llama_model_deepseek2(params) {} // reuse load_arch_hparams and load_arch_tensors from llama_model_deepseek2 diff --git a/examples/talk-llama/models/qwen3next.cpp b/examples/talk-llama/models/qwen3next.cpp index 97200a440..09b66423d 100644 --- a/examples/talk-llama/models/qwen3next.cpp +++ b/examples/talk-llama/models/qwen3next.cpp @@ -121,6 +121,8 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p ggml_tensor * inp_out_ids = build_inp_out_ids(); for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; + ggml_tensor * inpSA = inpL; cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);