furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 10:15:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make MTP work with split mode graph

Browse Source
This commit is contained in:
Kawrakow
2026-05-28 04:12:34 +00:00
parent de4257303f
commit 7cf668f797
2 changed files with 12 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/graphs/build_deepseek2.cpp
+3
View File
@@ -32,6 +32,9 @@ ggml_tensor * llm_build_context::build_deepseek2_tp_attention(
const uint32_t n_embd_head_v = hparams.n_embd_head_v(il);
auto cache_repl = (const ggml_split_tensor_t *)kv_self.k_l[il]->extra;
if (!cache_repl) {
LLAMA_LOG_ERROR("%s: no cache split for layer %d?\n", __func__, il);
}
GGML_ASSERT(cache_repl);
GGML_ASSERT(cache_repl->n_device == n_device);
src/llama.cpp
+9 -7
View File
@@ -770,7 +770,7 @@ static bool llama_kv_cache_init(
const struct llama_hparams & hparams = model.hparams;
const int64_t n_layer = model.mtp ? hparams.n_layer
: hparams.n_layer - hparams.nextn_predict_layers;
: hparams.n_layer - hparams.nextn_predict_layers;
cache.has_shift = false;
@@ -818,11 +818,12 @@ static bool llama_kv_cache_init(
// count used buffer types
std::map<ggml_backend_buffer_type_t, int> buft_layer_count;
if (offload) {
const bool qwen_mtp = (model.arch == LLM_ARCH_QWEN35 ||
model.arch == LLM_ARCH_QWEN35MOE) && hparams.nextn_predict_layers > 0;
const int64_t n_mtp_first = n_layer - hparams.nextn_predict_layers;
const bool is_mtp = (model.arch == LLM_ARCH_GLM_DSA ||
model.arch == LLM_ARCH_QWEN35 ||
model.arch == LLM_ARCH_QWEN35MOE) && hparams.nextn_predict_layers > 0;
const int64_t n_mtp_first = hparams.n_layer - hparams.nextn_predict_layers;
for (int64_t i = 0; i < n_layer; ++i) {
const bool is_mtp_tail = qwen_mtp && i >= n_mtp_first;
const bool is_mtp_tail = is_mtp && i >= n_mtp_first;
if ((split_cache || replicate_mla) && !is_mtp_tail) {
buft_layer_count[model.buft_layer[i].buft_matrix]++;
if (model.buft_layer[i].buft != model.buft_layer[i].buft_matrix) {
@@ -898,7 +899,7 @@ static bool llama_kv_cache_init(
int n_mla = 0;
int n_kv_active_layers = 0;
const int64_t n_mtp_first_layer = n_layer - hparams.nextn_predict_layers;
const int64_t n_mtp_first_layer = hparams.n_layer - hparams.nextn_predict_layers;
for (int i = 0; i < (int) n_layer; i++) {
// For MTP-only context, skip KV allocation for non-MTP layers
if (cparams.mtp_op_type != MTP_OP_NONE && i < (int)n_mtp_first_layer) {
@@ -915,7 +916,8 @@ static bool llama_kv_cache_init(
const uint32_t n_embd_head_k= hparams.n_embd_head_k(i);
const bool is_mtp_tail_layer = (model.arch == LLM_ARCH_QWEN35 ||
model.arch == LLM_ARCH_QWEN35MOE) &&
model.arch == LLM_ARCH_QWEN35MOE ||
model.arch == LLM_ARCH_GLM_DSA) &&
hparams.nextn_predict_layers > 0 && i >= (int)n_mtp_first_layer;
//struct ggml_context * ctx = split_cache && !qnext_recurrent ? ctx_map.at(model.buft_layer[i].buft_matrix) : offload ? ctx_map.at(model.buft_layer[i].buft) : cache.ctxs.front();
struct ggml_context * ctx = ((split_cache || replicate_mla) && !is_mtp_tail_layer) ? ctx_map.at(model.buft_layer[i].buft_matrix) : offload ? ctx_map.at(model.buft_layer[i].buft) : cache.ctxs.front();