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
{ 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" },
{ 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" },
{ 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" },
{ 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" },
{ 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" },
{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}},
{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
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:
LLM_ARCH_DEEPSEEK2,
LLM_ARCH_DEEPSEEK2OCR,
LLM_ARCH_DEEPSEEK32,
+ LLM_ARCH_DEEPSEEK4,
LLM_ARCH_CHATGLM,
LLM_ARCH_GLM4,
LLM_ARCH_GLM4_MOE,
LLM_ARCH_TALKIE,
LLM_ARCH_MELLUM,
LLM_ARCH_EAGLE3,
+ LLM_ARCH_DFLASH,
LLM_ARCH_UNKNOWN,
};
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,
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,
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,
LLAMA_LOG_ERROR("%s: sequence %d positions are decreasing (not allowed)\n", __func__, seq_id);
return false;
}
+
+ cur_seq_pos[seq_id] = pos;
}
}
}
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__);
}
}
- const uint32_t n_seqs = cur_seq_set.size();
+ uint32_t n_seqs = cur_seq_set.size();
// we are done
if (n_seqs == 0) {
std::vector<idx_vec_t> 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;
}
}
+ // 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<int32_t> idxs;
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) {
// 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);
#include <cstring>
#include <limits>
#include <stdexcept>
+#include <string>
//
// llama_context
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) :
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;
}
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);
}
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;
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;
//
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<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
}
uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
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);
#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"
#include <cstring>
#include <numeric>
#include <sstream>
+#include <string>
#include <unordered_set>
// dedup helpers
// 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;
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);
}
}
// 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
// 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);
}
}
return res;
}
+static void dsv4_set_i64(ggml_tensor * dst, const std::vector<int64_t> & 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<int32_t> & 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<ggml_fp16_t>(-INFINITY);
+ const ggml_fp16_t fp16_zero = llama_cast<ggml_fp16_t>(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<int32_t> & 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<const llama_kv_cache_dsv4_context *>(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);
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);
params = {};
inputs.clear();
+ fused_nodes.clear();
buf_compute_meta.resize(ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false));
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;
}
}
}
+
+
ggml_tensor * llm_graph_context::build_cvec(
ggml_tensor * cur,
int il) const {
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;
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,
/* gate_up_exps_b */ nullptr,
up_exps_s,
gate_exps_s,
- down_exps_s
+ down_exps_s,
+ selected_experts_in
);
}
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
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;
{
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");
}
}
// 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) {
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;
}
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);
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
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) {
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);
}
}
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) {
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<const llama_kv_cache_dsv4_context *>(mctx);
+ const auto * raw_ctx = mctx_cur->get_raw();
+
+ auto inp_raw = std::make_unique<llm_graph_input_dsv4_raw>(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<llm_graph_input_dsv4>(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,
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;
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,
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<llm_graph_input_dsv4_raw> 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<llm_graph_input_dsv4_raw> 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) {}
}
};
+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);
llm_graph_input_i * add_input(llm_graph_input_ptr input);
+ void add_fused_node(llm_graph_fused_node result);
+
+ const std::vector<llm_graph_fused_node> & get_fused_nodes() const { return fused_nodes; }
+
void set_params(const llm_graph_params & params);
// important graph nodes
std::map<llama_seq_id, ggml_tensor *> t_sampled_probs;
std::vector<llm_graph_input_ptr> inputs;
+ std::vector<llm_graph_fused_node> fused_nodes;
ggml_context_ptr ctx_compute;
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,
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
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,
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 {
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<uint32_t, LLAMA_MAX_LAYERS> 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
}
}
+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();
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__"
std::vector<llama_ubatch> 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;
--- /dev/null
+#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 <algorithm>
+#include <cassert>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <map>
+#include <sstream>
+#include <stdexcept>
+
+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<uint32_t>(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<uint32_t> 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<llama_ubatch::data_t>();
+ 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<llama_ubatch> dsv4_build_raw_write_ubatches(const std::vector<llama_ubatch> & ubatches) {
+ std::vector<llama_ubatch> 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<int32_t> & 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_row> 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<int32_t> overlap_prev_reads;
+ std::vector<int32_t> overlap_cur_reads;
+
+ std::map<std::pair<llama_seq_id, llama_pos>, 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<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans(
+ const std::vector<llama_ubatch> & ubatches,
+ uint32_t ratio,
+ bool overlap,
+ uint32_t state_size,
+ uint32_t kv_size,
+ uint32_t n_stream) {
+ std::vector<llama_kv_cache_dsv4_context::comp_plan> 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<llama_ubatch> & 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<llama_ubatch> & 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<uint32_t>(1, ubatch.n_seqs);
+ const uint32_t n_seq_tokens = std::max<uint32_t>(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<uint64_t>(1, n_blocks_u64);
+ GGML_ASSERT((uint64_t) n_blocks == std::max<uint64_t>(1, n_blocks_u64));
+
+ const uint64_t state_rows = (uint64_t) state_size*n_stream;
+ const size_t n_persist = (size_t) std::min<uint64_t>(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<ggml_backend_buffer_type_t, ggml_context_ptr, ggml_backend_buft_comparator> 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<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_breakdown() const {
+ std::map<ggml_backend_buffer_type_t, size_t> 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<llama_kv_cache_iswa>(
+ 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<llama_kv_cache>(
+ 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<llama_kv_cache>(
+ 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<llama_kv_cache>(
+ 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<llama_dsv4_comp_state>(
+ 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<llama_dsv4_comp_state>(
+ 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<llama_dsv4_comp_state>(
+ 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<llama_ubatch> 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<llama_kv_cache_dsv4_context>(
+ 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<llama_ubatch> 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<llama_ubatch> 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_kv_cache_dsv4_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
+}
+
+llama_memory_context_ptr llama_kv_cache_dsv4::init_full() {
+ return std::make_unique<llama_kv_cache_dsv4_context>(this);
+}
+
+llama_memory_context_ptr llama_kv_cache_dsv4::init_update(llama_context * lctx, bool optimize) {
+ return std::make_unique<llama_kv_cache_dsv4_context>(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<ggml_backend_buffer_type_t, size_t> llama_kv_cache_dsv4::memory_breakdown() const {
+ std::map<ggml_backend_buffer_type_t, size_t> 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<llama_ubatch> ubatches,
+ std::vector<llama_ubatch> 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<llama_kv_cache_context>(
+ 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<llama_ubatch> 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<llama_kv_cache_dsv4_raw_context>(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<llama_kv_cache_dsv4_comp_context>(kv->get_csa())),
+ ctx_hca(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_hca())),
+ ctx_lid(std::make_unique<llama_kv_cache_dsv4_comp_context>(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<llama_kv_cache_dsv4_raw_context>(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<llama_kv_cache_dsv4_comp_context>(kv->get_csa())),
+ ctx_hca(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_hca())),
+ ctx_lid(std::make_unique<llama_kv_cache_dsv4_comp_context>(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<llama_ubatch> ubatches,
+ std::vector<llama_ubatch> 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<llama_kv_cache_dsv4_raw_context>(
+ 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<llama_kv_cache_dsv4_comp_context>(
+ kv->get_csa(),
+ dsv4_build_comp_sinfos(this->ubatches, kv->get_csa()->get_n_stream()),
+ this->ubatches)),
+ ctx_hca(std::make_unique<llama_kv_cache_dsv4_comp_context>(
+ kv->get_hca(),
+ dsv4_build_comp_sinfos(this->ubatches, kv->get_hca()->get_n_stream()),
+ this->ubatches)),
+ ctx_lid(std::make_unique<llama_kv_cache_dsv4_comp_context>(
+ 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;
+}
--- /dev/null
+#pragma once
+
+#include "llama-kv-cache.h"
+#include "llama-kv-cache-iswa.h"
+
+#include <map>
+#include <memory>
+#include <unordered_map>
+#include <vector>
+
+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<ggml_backend_buffer_type_t, size_t> 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<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
+
+ std::vector<layer> layers;
+
+ std::unordered_map<int32_t, int32_t> 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<ggml_backend_buffer_type_t, size_t> 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<llama_kv_cache_iswa> kv_raw;
+ std::unique_ptr<llama_kv_cache> kv_csa;
+ std::unique_ptr<llama_kv_cache> kv_hca;
+ std::unique_ptr<llama_kv_cache> kv_lid;
+ std::unique_ptr<llama_dsv4_comp_state> csa_state;
+ std::unique_ptr<llama_dsv4_comp_state> hca_state;
+ std::unique_ptr<llama_dsv4_comp_state> 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<llama_ubatch> ubatches,
+ std::vector<llama_ubatch> 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<llama_ubatch> ubatches;
+ std::vector<llama_ubatch> 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<llama_ubatch> 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<llama_ubatch> 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<int32_t> state_pos;
+
+ // Current-ubatch source row ids and unique persistent-state
+ // destination row ids for deterministic ring-state updates.
+ std::vector<int32_t> state_persist_src_idxs;
+ std::vector<int32_t> 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<int32_t> 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<int64_t> state_write_idxs;
+
+ // RoPE positions for state-backed commits.
+ std::vector<int32_t> state_write_pos;
+
+ // Number of completed compressed rows visible for each query token.
+ std::vector<int32_t> 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<llama_ubatch> ubatches,
+ std::vector<llama_ubatch> 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<llama_ubatch> ubatches;
+
+ std::vector<comp_plan> plans_csa;
+ std::vector<comp_plan> plans_hca;
+ std::vector<comp_plan> plans_lid;
+
+ const std::unique_ptr<llama_kv_cache_dsv4_raw_context> 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<llama_kv_cache_dsv4_comp_context> ctx_csa;
+ const std::unique_ptr<llama_kv_cache_dsv4_comp_context> ctx_hca;
+ const std::unique_ptr<llama_kv_cache_dsv4_comp_context> 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;
+};
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) {
std::vector<llama_ubatch> 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;
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;
//
llama_kv_cache * get_swa () const;
private:
- const llama_hparams & hparams;
-
const bool unified;
std::unique_ptr<llama_kv_cache> kv_base;
}
}
-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
//
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]
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;
}
std::vector<llama_ubatch> 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;
return layers[0].v->type;
}
+std::vector<uint32_t> llama_kv_cache::get_layer_ids() const {
+ std::vector<uint32_t> 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;
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 {
ggml_type type_k() const;
ggml_type type_v() const;
+ std::vector<uint32_t> get_layer_ids() const;
+ ggml_tensor * get_k_storage(int32_t il) const;
+
//
// graph_build API
//
// 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) {
// 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) {
// 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) {
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)";
- }
-
- 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";
+#define LLAMA_FTYPE_PREFIX "(guessed) "
+
+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 "<name>-00002-of-00004.gguf", returns list of all 4 splits
static std::vector<std::string> llama_get_list_splits(const std::string & path, const int idx, const int n_split) {
}
template bool llama_model_loader::get_arr_n(enum llm_kv kid, uint32_t & result, bool required);
+ template std::enable_if<std::is_integral<uint32_t>::value, bool>::type
+ llama_model_loader::get_arr_n<uint32_t>(const std::string & key, uint32_t & result, bool required);
template<typename T>
bool llama_model_loader::get_arr(const std::string & key, std::vector<T> & result, bool required) {
template bool llama_model_loader::get_arr<std::vector<std::string>>(enum llm_kv kid, std::vector<std::string> & result, bool required);
template bool llama_model_loader::get_arr<std::array<int32_t, 512>>(enum llm_kv kid, std::array<int32_t, 512> & result, bool required);
template bool llama_model_loader::get_arr<std::vector<int32_t>>(enum llm_kv kid, std::vector<int32_t> & result, bool required);
+ template bool llama_model_loader::get_arr<std::array<uint32_t, LLAMA_MAX_LAYERS>>(enum llm_kv kid, std::array<uint32_t, LLAMA_MAX_LAYERS> & result, bool required);
template<typename T>
bool llama_model_loader::get_key(const std::string & key, T & result, bool required) {
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));
}
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 {
#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"
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:
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:
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";
}
}
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)));
}
}
std::string desc_str;
+ llama_ftype ftype = LLAMA_FTYPE_ALL_F32;
+
// model memory mapped files
llama_mmaps mappings;
std::vector<layer_dev> dev_layer;
bool has_tensor_overrides;
+
+ std::vector<float> tensor_split_owned;
};
llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique<impl>()) {
+ 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;
}
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;
}
return pimpl->desc_str;
}
+llama_ftype llama_model::ftype() const {
+ return pimpl->ftype;
+}
+
size_t llama_model::size() const {
return pimpl->n_bytes;
}
}
}
- 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) {
}
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;
}
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:
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:
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();
}
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;
}
}
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;
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;
// 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;
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;
std::string desc() const;
+ llama_ftype ftype() const;
+
size_t size() const; // file size
size_t n_tensors() const;
size_t n_devices() const;
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:
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;
}
}
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;
// 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.
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
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 defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
- // correct endianness of data in precompiled_charsmap binary blob
+ 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);
- assert(*xcda_blob_size + sizeof(uint32_t) < n_precompiled_charsmap);
+#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
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) {
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,
// 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);
--- /dev/null
+#include "models.h"
+
+#include "llama-kv-cache-dsv4.h"
+
+#include <algorithm>
+#include <cmath>
+#include <stdexcept>
+#include <string>
+
+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<llm_graph_context> llama_model_deepseek4::build_arch_graph(const llm_graph_params & params) const {
+ return std::make_unique<graph>(*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);
+}
// 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,
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;
// 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;
--- /dev/null
+#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<llm_graph_context> llama_model_dflash::build_arch_graph(const llm_graph_params & params) const {
+ switch (params.gtype) {
+ case LLM_GRAPH_TYPE_ENCODER:
+ return std::make_unique<graph<true>>(*this, params);
+ case LLM_GRAPH_TYPE_DEFAULT:
+ case LLM_GRAPH_TYPE_DECODER:
+ return std::make_unique<graph<false>>(*this, params);
+ default:
+ GGML_ABORT("invalid graph type");
+ };
+}
+
+template <>
+ggml_tensor * llama_model_dflash::graph<true>::build_inp_embd_enc() const {
+ auto inp_target = std::make_unique<llm_graph_input_embd>(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<true>::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<false>::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<llm_graph_input_embd>(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<llm_graph_input_embd>(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);
+}
};
+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<llm_graph_context> 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;
};
+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 <bool is_enc>
+ 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<llm_graph_context> 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
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);
overview / pkg / ggml / sources / whisper.cpp / commitdiff