diff --git a/common/speculative.cpp b/common/speculative.cpp index d32db0fb..5acef4b4 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -169,7 +169,7 @@ struct common_speculative_state_mtp : public common_speculative_state { if (ctx_mtp) { LOG_INF("%s: created MTP context (n_ctx=%d)\n", __func__, llama_n_ctx(ctx_mtp)); } else { - LOG_ERR("%s: failed to create MTP context, falling back to shared context\n", __func__); + LOG_ERR("%s: failed to create MTP context\n", __func__); } } @@ -193,14 +193,12 @@ struct common_speculative_state_mtp : public common_speculative_state { int32_t n_past = (int32_t)prompt_tgt.size(); llama_seq_id seq_id = 0; - if (ctx_mtp) { - llama_pos mtp_pos_max = llama_kv_cache_seq_pos_max(ctx_mtp, seq_id); - if (mtp_pos_max >= n_past) { - llama_kv_cache_seq_rm(ctx_mtp, seq_id, n_past, -1); - } + llama_pos mtp_pos_max = llama_kv_cache_seq_pos_max(ctx_mtp, seq_id); + if (mtp_pos_max >= n_past) { + llama_kv_cache_seq_rm(ctx_mtp, seq_id, n_past, -1); } - llama_context * ctx = ctx_mtp ? ctx_mtp : ctx_tgt; + llama_context * ctx = ctx_mtp; result = mtp_speculative_gen_draft( smpl, @@ -974,10 +972,15 @@ common_speculative * common_speculative_init( break; } case COMMON_SPECULATIVE_TYPE_MTP: { - impls.push_back(std::make_unique(config.type, + auto mtp_state = std::make_unique(config.type, /* .ctx_tgt = */ ctx_tgt, /* .mtp_cparams = */ params.cparams_dft - )); + ); + if (!mtp_state->ctx_mtp) { + LOG_ERR("%s: failed to create MTP context\n", __func__); + return nullptr; + } + impls.push_back(std::move(mtp_state)); break; } case COMMON_SPECULATIVE_TYPE_EAGLE3: { @@ -1280,7 +1283,7 @@ void mtp_update_kv_cache(struct llama_context * ctx, const llama_batch& batch, b return; } - llama_seq_id seq_id = batch.seq_id[0][0]; + llama_seq_id seq_id = batch.seq_id[0][0]; llama_pos start_pos = batch.pos[0]; if (llama_kv_cache_seq_pos_max(ctx, seq_id) >= start_pos) { diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index 81422dc5..6bb46172 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -255,7 +255,8 @@ void server_context::init() { SLT_INF(slot, "%s", "speculative decoding context initialized\n"); } else { if (slot.has_mtp) { - SRV_ERR("%s", "failed to initialize MTP speculative context\n"); + SRV_ERR("%s", "failed to initialize MTP speculative context, aborting\n"); + GGML_ABORT("MTP context creation failed"); } else { SLT_INF(slot, "%s", "speculative decoding context not initialized\n"); } diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 82239652..317308e3 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -8183,59 +8183,57 @@ struct ggml_tensor * llm_build_context::build_mtp_tail( cb(combined, "mtp_concat", il); ggml_tensor* cur = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, combined); - struct ggml_tensor * inpSA = cur; - - cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.attn_norm, NULL, LLM_NORM_RMS, cb, il); - cb(cur, "attn_norm", il); - // Self-Attention - { + const float kq_scale = 1.0f / sqrtf(float(n_embd_head)); + ggml_tensor * ffn_inp; + if (rope_cache == nullptr) { + cur = build_std_attention(gf, mtp_layer.attn_norm, cur, + inp_pos, nullptr, nullptr, + KQ_mask, nullptr, nullptr, + kq_scale, 0.0f, 0, il, true, false, true, false, false, nullptr); + ffn_inp = cur; + } else { + struct ggml_tensor * inpSA = cur; + cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.attn_norm, NULL, LLM_NORM_RMS, cb, il); + cb(cur, "attn_norm", il); auto [Qcur, Kcur, Vcur] = llm_build_mul_mat_qkv(gf, cur, - nullptr, nullptr, // wqkv, bqkv (not used in GLM usually?) - nullptr, nullptr, // wqk, bqk + nullptr, nullptr, + nullptr, nullptr, mtp_layer.wq, mtp_layer.bq, mtp_layer.wk, mtp_layer.bk, mtp_layer.wv, mtp_layer.bv, - mtp_layer.attn_q_norm, mtp_layer.attn_k_norm, + mtp_layer.attn_q_norm, mtp_layer.attn_k_norm, 0.f, il); - - // RoPE - if (rope_cache) { - Qcur = ggml_rope_fast(ctx0, Qcur, rope_cache); - Kcur = ggml_rope_fast(ctx0, Kcur, rope_cache); - } else { - 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); - } - + Qcur = ggml_rope_fast(ctx0, Qcur, rope_cache); + Kcur = ggml_rope_fast(ctx0, Kcur, rope_cache); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il); - - // KV Cache & Attention cur = llm_build_kv(ctx0, lctx, kv_self, gf, - model.layers[il].wo, NULL, + mtp_layer.wo, NULL, Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, - 1.0f/sqrtf(float(n_embd_head)), cb, il); + kq_scale, cb, il); + ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "mtp_ffn_inp", il); } - ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); - cb(ffn_inp, "mtp_ffn_inp", il); - + // FFN cur = llm_build_std_moe_ffn(ctx0, lctx, mtp_layer.ffn_norm, ffn_inp, - mtp_layer.ffn_gate_inp, NULL, - mtp_layer.ffn_up_exps, NULL, - mtp_layer.ffn_gate_exps, NULL, - mtp_layer.ffn_down_exps, NULL, - mtp_layer.ffn_exp_probs_b, - mtp_layer.ffn_up_shexp, nullptr, - mtp_layer.ffn_gate_shexp, nullptr, - mtp_layer.ffn_down_shexp, nullptr, - n_expert, n_expert_used, - LLM_FFN_SILU, hparams.expert_weights_norm, true, hparams.expert_weights_scale, - (llm_expert_gating_func_type) hparams.expert_gating_func, - LLM_FFN_SILU, cb, il, gf, true, mtp_layer.ffn_up_gate_exps); + mtp_layer.ffn_gate_inp, NULL, + mtp_layer.ffn_up_exps, NULL, + mtp_layer.ffn_gate_exps, NULL, + mtp_layer.ffn_down_exps, NULL, + mtp_layer.ffn_exp_probs_b, + mtp_layer.ffn_up_shexp, nullptr, + mtp_layer.ffn_gate_shexp, nullptr, + mtp_layer.ffn_down_shexp, nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, true, hparams.expert_weights_scale, + (llm_expert_gating_func_type) hparams.expert_gating_func, + LLM_FFN_SILU, cb, il, gf, true, mtp_layer.ffn_up_gate_exps); + + cur = lctx.cvec.apply_to(ctx0, cur, il); cb(cur, "ffn_out", il); cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.nextn.shared_head_norm, NULL, LLM_NORM_RMS, cb, il); diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 1ac10081..a65e7f20 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -309,6 +309,14 @@ create_tensors_helper::create_tensors_helper(llama_model_loader & _ml, llama_mod } } + // Split MTP layer's to graph + if ((model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN) && + model.hparams.nextn_predict_layers > 0 && model.splits.size() > 1) { + int mtp_first = n_layer - model.hparams.nextn_predict_layers; + LLAMA_LOG_DEBUG("%s: MTP layer(s) %d-%d: split attention+FFN, nextn on per-device CUDA\n", + __func__, mtp_first, n_layer - 1); + } + auto n_tensors = ml.n_tensors; if (ml.merge_qkv) n_tensors += n_layer; if (ml.merge_up_gate_exps) n_tensors += n_layer; @@ -2616,35 +2624,32 @@ bool create_tensors_helper::create_glm4_moe_tensors(const LLM_TN & tn) { layer.ffn_down = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, flags); layer.ffn_up = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, flags); } - // --- NextN / MTP tensors (preserved but unused), on the final layer --- + // --- NextN / MTP tensors on the final layer --- if (hparams.nextn_predict_layers > 0 && static_cast(i) >= n_layer - hparams.nextn_predict_layers) { const int final_layer = n_layer - 1; - // EH_PROJ: [2*embd, embd] - layer.nextn.eh_proj = create_tensor(ctx_for_layer(final_layer), + auto nextn_ctx = ctx_for_layer(final_layer); + auto nextn_host_ctx = ctx_input; + layer.nextn.eh_proj = create_tensor(nextn_ctx, tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", final_layer), { 2*n_embd, n_embd }, flags); - // EMBED_TOKENS: [embd, vocab] - layer.nextn.embed_tokens = create_tensor(ctx_for_layer(final_layer), + layer.nextn.embed_tokens = create_tensor(nextn_host_ctx, tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", final_layer), { n_embd, n_vocab }, flags | llama_model_loader::TENSOR_NOT_REQUIRED); - // ENORM, HNORM: [embd] - layer.nextn.enorm = create_tensor(ctx_for_layer(final_layer), + layer.nextn.enorm = create_tensor(nextn_ctx, tn(LLM_TENSOR_NEXTN_ENORM, "weight", final_layer), { n_embd }, flags); - layer.nextn.hnorm = create_tensor(ctx_for_layer(final_layer), + layer.nextn.hnorm = create_tensor(nextn_ctx, tn(LLM_TENSOR_NEXTN_HNORM, "weight", final_layer), { n_embd }, flags); - // SHARED_HEAD_HEAD: [embd, vocab] - layer.nextn.shared_head_head = create_tensor(ctx_for_layer(final_layer), + layer.nextn.shared_head_head = create_tensor(nextn_host_ctx, tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", final_layer), { n_embd, n_vocab }, flags | llama_model_loader::TENSOR_NOT_REQUIRED); - // SHARED_HEAD_NORM: [embd] - layer.nextn.shared_head_norm = create_tensor(ctx_for_layer(final_layer), + layer.nextn.shared_head_norm = create_tensor(nextn_ctx, tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", final_layer), { n_embd }, flags | llama_model_loader::TENSOR_NOT_REQUIRED);