speculative : add Step 3.7 MTP support (#2250)
* speculative : add Step 3.7 MTP support * clean unnecesary comments * remove duplicated function and fix split graph mode to work * Simplify speculative MTP compatibility checks
This commit is contained in:
Samuel Oliveira Alves
GitHub
parent
21819724b3
commit
ea791ac569
|
|
@ -4248,7 +4248,7 @@ struct llama_model_params common_model_params_to_llama(const gpt_params & params
|
|||
mparams.validate_quants = params.validate_quants;
|
||||
mparams.merge_qkv = params.merge_qkv;
|
||||
mparams.merge_up_gate_exps = params.merge_up_gate_exps;
|
||||
mparams.mtp = params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP);
|
||||
mparams.mtp = params.has_mtp || params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP);
|
||||
mparams.flash_attn = params.flash_attn;
|
||||
mparams.defer_experts = params.defer_experts;
|
||||
mparams.swa_compress = params.swa_compress;
|
||||
|
|
@ -4345,7 +4345,7 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa
|
|||
cparams.prefetch_experts = params.prefetch_experts;
|
||||
cparams.prefetch_experts_threads = params.prefetch_experts_threads;
|
||||
cparams.max_extra_alloc = params.max_extra_alloc_MiB;
|
||||
cparams.mtp = params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP);
|
||||
cparams.mtp = params.has_mtp || params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP);
|
||||
cparams.mtp_op_type = MTP_OP_NONE;
|
||||
|
||||
cparams.type_k = kv_cache_type_from_str(params.cache_type_k);
|
||||
|
|
|
|||
|
|
@ -143,6 +143,24 @@ static bool common_speculative_are_compatible(
|
|||
return true;
|
||||
}
|
||||
|
||||
static bool common_speculative_target_has_appended_mtp_contract(const llama_model * model) {
|
||||
return llama_model_is_step35(model) || llama_model_is_deepseek4(model);
|
||||
}
|
||||
|
||||
static bool common_speculative_has_recognized_mtp_companion(
|
||||
const llama_model * target,
|
||||
const llama_model * companion) {
|
||||
if (target == nullptr || companion == nullptr) {
|
||||
return false;
|
||||
}
|
||||
if (llama_model_is_gemma4_mtp_assistant(companion)) {
|
||||
return llama_model_arch_string(target) != nullptr &&
|
||||
std::strcmp(llama_model_arch_string(target), "gemma4") == 0;
|
||||
}
|
||||
return common_speculative_target_has_appended_mtp_contract(target) &&
|
||||
llama_model_mtp_package(companion) == LLAMA_MTP_PACKAGE_COMPANION;
|
||||
}
|
||||
|
||||
// state of an implementation of speculative decoding
|
||||
//
|
||||
// each implementation has a unique type and a state that is implementation-specific
|
||||
|
|
@ -1906,6 +1924,9 @@ bool common_speculative_load_draft_model(
|
|||
}
|
||||
|
||||
gpt_params params_dft = params_base;
|
||||
if (params.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP)) {
|
||||
params_dft.has_mtp = true;
|
||||
}
|
||||
params_dft.devices = params.devices.empty() ? params_base.devices : params.devices;
|
||||
params_dft.model = params.model;
|
||||
params_dft.n_gpu_layers = params.n_gpu_layers;
|
||||
|
|
@ -1966,13 +1987,28 @@ bool common_speculative_prepare_mtp_runtime(
|
|||
return false;
|
||||
}
|
||||
|
||||
if (llama_model_n_nextn_layer(model) == 0 && !has_external_mtp) {
|
||||
LOG_WRN("%s: MTP speculative stage requested, but model has 0 NextN layers. Removing MTP from the configured stage chain.\n",
|
||||
const enum llama_mtp_package target_package = llama_model_mtp_package(model);
|
||||
const bool has_embedded_mtp = target_package == LLAMA_MTP_PACKAGE_EMBEDDED;
|
||||
if (target_package == LLAMA_MTP_PACKAGE_COMPANION ||
|
||||
target_package == LLAMA_MTP_PACKAGE_INVALID) {
|
||||
LOG_ERR("%s: the target GGUF is an MTP companion or invalid package, load a complete target model instead\n",
|
||||
__func__);
|
||||
params.remove_stage_type(COMMON_SPECULATIVE_TYPE_MTP);
|
||||
if (!params.needs_dft_model()) {
|
||||
params.clear_dft();
|
||||
return false;
|
||||
}
|
||||
if (target_package == LLAMA_MTP_PACKAGE_TARGET_ONLY && !has_external_mtp) {
|
||||
LOG_ERR("%s: target GGUF contains no MTP tail, provide a matching predictor-only companion with -md\n",
|
||||
__func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (llama_model_is_step35(model) && has_embedded_mtp &&
|
||||
!llama_model_step35_has_nextn_weights(model)) {
|
||||
LOG_ERR("%s: Step target is missing one or more complete MTP heads\n", __func__);
|
||||
return false;
|
||||
}
|
||||
if (!has_external_mtp && !has_embedded_mtp) {
|
||||
LOG_ERR("%s: MTP speculative stage requested, but the target package is not a complete embedded MTP model\n",
|
||||
__func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
|
|
@ -2066,27 +2102,78 @@ bool common_speculative_finalize_startup(
|
|||
return false;
|
||||
}
|
||||
|
||||
if (params.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) &&
|
||||
params.model_dft != nullptr &&
|
||||
llama_model_is_deepseek4(params.model_dft) &&
|
||||
llama_model_n_nextn_layer(params.model_dft) != 1) {
|
||||
LOG_ERR("%s: DeepSeek-V4 MTP draft requires exactly one NextN predictor layer, got %d.\n",
|
||||
__func__, llama_model_n_nextn_layer(params.model_dft));
|
||||
const bool mtp_requested = params.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP);
|
||||
if (mtp_requested && model != nullptr && params.model_dft != nullptr) {
|
||||
const llama_model * companion = params.model_dft;
|
||||
const bool appended_contract = common_speculative_target_has_appended_mtp_contract(model);
|
||||
|
||||
if (appended_contract &&
|
||||
llama_model_mtp_package(companion) != LLAMA_MTP_PACKAGE_COMPANION) {
|
||||
LOG_ERR("%s: -md for an MTP stage must be a predictor-only companion GGUF\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (llama_model_is_step35(model)) {
|
||||
if (!llama_model_is_step35(companion)) {
|
||||
LOG_ERR("%s: Step MTP requires a Step companion\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
const int32_t n_heads = llama_model_n_nextn_layer(companion);
|
||||
if (n_heads != 3) {
|
||||
LOG_ERR("%s: Step MTP companion requires exactly three predictor heads, got %d\n",
|
||||
__func__, n_heads);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!llama_model_step35_has_nextn_weights(companion)) {
|
||||
LOG_ERR("%s: Step MTP companion is missing required predictor tensors\n", __func__);
|
||||
return false;
|
||||
}
|
||||
} else if (llama_model_is_deepseek4(model)) {
|
||||
if (!llama_model_is_deepseek4(companion)) {
|
||||
LOG_ERR("%s: DeepSeek-V4 MTP requires a DeepSeek-V4 companion\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
const int32_t n_heads = llama_model_n_nextn_layer(companion);
|
||||
if (n_heads != 1) {
|
||||
LOG_ERR("%s: DeepSeek-V4 MTP companion requires exactly one predictor layer, got %d\n",
|
||||
__func__, n_heads);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
if (common_speculative_has_recognized_mtp_companion(model, companion)) {
|
||||
if (!common_speculative_are_compatible(model, companion)) {
|
||||
LOG_ERR("%s: MTP target and companion vocabularies are incompatible\n", __func__);
|
||||
return false;
|
||||
}
|
||||
|
||||
const uint32_t target_width = llama_model_mtp_feature_width(model);
|
||||
const uint32_t companion_width = llama_model_mtp_feature_width(companion);
|
||||
|
||||
if (target_width != companion_width) {
|
||||
LOG_ERR("%s: MTP feature-width mismatch: target=%u companion=%u\n",
|
||||
__func__, target_width, companion_width);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
params_base.has_mtp = params.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP);
|
||||
const bool has_external_mtp = params_base.has_mtp && params.model_dft &&
|
||||
(llama_model_is_gemma4_mtp_assistant(params.model_dft) ||
|
||||
(llama_model_is_deepseek4(params.model_dft) &&
|
||||
llama_model_n_nextn_layer(params.model_dft) == 1));
|
||||
common_speculative_has_recognized_mtp_companion(model, params.model_dft);
|
||||
|
||||
params_base.has_mtp = common_speculative_prepare_mtp_runtime(
|
||||
params,
|
||||
params_base,
|
||||
model,
|
||||
has_external_mtp);
|
||||
if (params.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) && !params_base.has_mtp) {
|
||||
return false;
|
||||
}
|
||||
if (params_base.has_mtp) {
|
||||
params_base.pooling_type = LLAMA_POOLING_TYPE_NONE;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -712,7 +712,12 @@ int main(int argc, char ** argv) {
|
|||
// if we run out of context:
|
||||
// - take the n_keep first tokens from the original prompt (via n_past)
|
||||
// - take half of the last (n_ctx - n_keep) tokens and recompute the logits in batches
|
||||
if (n_past + (int) embd.size() + std::max<int>(0, guidance_offset) >= n_ctx) {
|
||||
const bool context_full = n_past + (int) embd.size() + std::max<int>(0, guidance_offset) >= n_ctx;
|
||||
if (context_full && !params.ctx_shift) {
|
||||
LOG_TEE("\n\n%s: context full and context shifting is disabled, stopping\n", __func__);
|
||||
break;
|
||||
}
|
||||
if (context_full) {
|
||||
if (params.n_predict == -2) {
|
||||
LOG_TEE("\n\n%s: context full and n_predict == -%d => stopping\n", __func__, params.n_predict);
|
||||
break;
|
||||
|
|
|
|||
|
|
@ -711,6 +711,8 @@ extern "C" {
|
|||
// Returns true if the model is a Gemma 4 MTP assistant (external frozen-KV speculative drafter)
|
||||
LLAMA_API bool llama_model_is_gemma4_mtp_assistant(const struct llama_model * model);
|
||||
|
||||
LLAMA_API bool llama_model_is_step35(const struct llama_model * model);
|
||||
|
||||
LLAMA_API bool llama_is_gemma4_mtp_file(const char * path);
|
||||
|
||||
LLAMA_API bool llama_model_is_split_mode_graph(const struct llama_model * model);
|
||||
|
|
@ -1609,8 +1611,22 @@ LLAMA_API struct llama_grammar* llama_sampler_init_grammar_lazy_patterns(
|
|||
// MTP
|
||||
//
|
||||
|
||||
enum llama_mtp_package {
|
||||
LLAMA_MTP_PACKAGE_NONE = 0,
|
||||
LLAMA_MTP_PACKAGE_EMBEDDED,
|
||||
LLAMA_MTP_PACKAGE_TARGET_ONLY,
|
||||
LLAMA_MTP_PACKAGE_COMPANION,
|
||||
LLAMA_MTP_PACKAGE_INVALID,
|
||||
};
|
||||
|
||||
LLAMA_API int32_t llama_model_n_nextn_layer(const struct llama_model * model);
|
||||
|
||||
LLAMA_API enum llama_mtp_package llama_model_mtp_package(const struct llama_model * model);
|
||||
|
||||
LLAMA_API uint32_t llama_model_mtp_feature_width(const struct llama_model * model);
|
||||
|
||||
LLAMA_API bool llama_model_step35_has_nextn_weights(const struct llama_model * model);
|
||||
|
||||
// Set which, if any, MTP operation the context will use
|
||||
LLAMA_API void llama_set_mtp_op_type(struct llama_context * ctx, enum llama_mtp_op_type mtp_op_type);
|
||||
|
||||
|
|
|
|||
|
|
@ -5,15 +5,57 @@
|
|||
ggml_cgraph * llm_build_context::build_step35() {
|
||||
ggml_cgraph * gf = new_graph_custom();
|
||||
ggml_tensor * cur;
|
||||
auto inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
|
||||
auto inp_pos = build_inp_pos();
|
||||
|
||||
if (cparams.mtp_op_type != MTP_OP_NONE) {
|
||||
GGML_ASSERT(model.mtp && hparams.nextn_predict_layers > 0);
|
||||
GGML_ASSERT(batch.token && "Step35 MTP requires token batches");
|
||||
|
||||
const int n_layer_base = hparams.n_layer > hparams.nextn_predict_layers
|
||||
? hparams.n_layer - hparams.nextn_predict_layers : hparams.n_layer;
|
||||
const int n_heads_model = (int) hparams.nextn_predict_layers;
|
||||
const int n_heads = lctx.mtp_n_heads > 0
|
||||
? std::max(1, std::min((int) lctx.mtp_n_heads, n_heads_model)) : n_heads_model;
|
||||
const int step = std::max(0, std::min((int) lctx.mtp_step_idx, n_heads - 1));
|
||||
const int il = n_layer_base + step;
|
||||
|
||||
ggml_tensor * hidden_states = build_inp_mtp_states(n_embd);
|
||||
|
||||
const bool step_independent_warmup = model.arch == LLM_ARCH_STEP35 &&
|
||||
(cparams.mtp_op_type == MTP_OP_WARMUP ||
|
||||
cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED) && n_heads > 1;
|
||||
if (step_independent_warmup) {
|
||||
for (int i = n_heads - 1; i >= 0; --i) {
|
||||
const int head_il = n_layer_base + i;
|
||||
const bool is_first = i == 0;
|
||||
const bool emit_logits = is_first && cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED;
|
||||
cur = build_step35_mtp(model.layers[head_il], hidden_states, gf, inp_pos,
|
||||
/*reduce_output=*/is_first, emit_logits);
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
return gf;
|
||||
}
|
||||
|
||||
const bool reduce_mtp_output = cparams.mtp_op_type != MTP_OP_NONE;
|
||||
const bool emit_mtp_logits = cparams.mtp_op_type == MTP_OP_DRAFT_GEN ||
|
||||
cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED;
|
||||
cur = build_step35_mtp(model.layers[il], hidden_states, gf, inp_pos,
|
||||
reduce_mtp_output, emit_mtp_logits);
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
return gf;
|
||||
}
|
||||
|
||||
auto inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
|
||||
auto inp_out_ids = build_inp_out_ids();
|
||||
auto KQ_mask = build_inp_KQ_mask();
|
||||
auto KQ_mask_swa = build_inp_KQ_mask_swa();
|
||||
//const float kq_scale = 1.0f / sqrtf(float(n_rot));
|
||||
const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k));
|
||||
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
const int n_layer_base = hparams.n_layer > hparams.nextn_predict_layers
|
||||
? hparams.n_layer - hparams.nextn_predict_layers : hparams.n_layer;
|
||||
|
||||
for (int il = 0; il < n_layer_base; ++il) {
|
||||
bool is_swa = hparams.swa_layers[il];
|
||||
auto & layer = const_cast<llama_layer&>(model.layers[il]);
|
||||
|
||||
|
|
@ -25,7 +67,7 @@ ggml_cgraph * llm_build_context::build_step35() {
|
|||
auto rope_freqs = layer.rope_freqs;
|
||||
layer.rope_freqs = nullptr;
|
||||
cur = build_std_attention(gf, model.layers[il].attn_norm, inpL,
|
||||
inp_pos, il == n_layer - 1 && n_tokens > 1 ? inp_out_ids : nullptr,
|
||||
inp_pos, il == n_layer_base - 1 && n_tokens > 1 && !cparams.mtp ? inp_out_ids : nullptr,
|
||||
rope_factors, is_swa ? KQ_mask_swa : KQ_mask, nullptr, nullptr, kq_scale, 0.0f, is_swa ? hparams.n_swa : 0,
|
||||
il, true, false, true);
|
||||
layer.rope_freqs = rope_freqs;
|
||||
|
|
@ -65,6 +107,17 @@ ggml_cgraph * llm_build_context::build_step35() {
|
|||
inpL = cur;
|
||||
}
|
||||
|
||||
if (cparams.mtp) {
|
||||
ggml_tensor * mtp_embd = inpL->type == GGML_TYPE_F32 ? inpL : ggml_cast(ctx0, inpL, GGML_TYPE_F32);
|
||||
cb(mtp_embd, "result_mtp_embd", -1);
|
||||
ggml_set_output(mtp_embd);
|
||||
ggml_build_forward_expand(gf, mtp_embd);
|
||||
|
||||
if (inp_out_ids) {
|
||||
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
|
||||
}
|
||||
}
|
||||
|
||||
cur = build_output(lctx, ctx0, inpL, model.output, model.output_norm, cb);
|
||||
cb(cur, "result_output", -1);
|
||||
|
||||
|
|
@ -73,3 +126,93 @@ ggml_cgraph * llm_build_context::build_step35() {
|
|||
return gf;
|
||||
}
|
||||
|
||||
ggml_tensor * llm_build_context::build_step35_mtp(
|
||||
const llama_layer & mtp_layer,
|
||||
ggml_tensor * hidden_states_from_main_model,
|
||||
ggml_cgraph * gf,
|
||||
ggml_tensor * inp_pos,
|
||||
bool reduce_output,
|
||||
bool emit_logits,
|
||||
ggml_tensor ** hidden_out) {
|
||||
const int il = (int) (&mtp_layer - model.layers.data());
|
||||
|
||||
GGML_ASSERT(mtp_layer.nextn.eh_proj && mtp_layer.nextn.enorm && mtp_layer.nextn.hnorm);
|
||||
GGML_ASSERT(mtp_layer.wq && mtp_layer.wk && mtp_layer.wv && mtp_layer.wo);
|
||||
|
||||
ggml_tensor * inp_out_ids = (n_tokens > 1 && n_outputs < n_tokens) ? build_inp_out_ids() : nullptr;
|
||||
ggml_tensor * tok_embd_w = mtp_layer.nextn.embed_tokens ? mtp_layer.nextn.embed_tokens : model.tok_embd;
|
||||
ggml_tensor * tok_embd = build_inp_embd_mtp(tok_embd_w);
|
||||
ggml_tensor * cur = build_mtp_input(mtp_layer, hidden_states_from_main_model,
|
||||
tok_embd, il, "mtp_eh_proj");
|
||||
|
||||
const bool is_swa = hparams.swa_layers[il];
|
||||
ggml_tensor * rope_factors = nullptr;
|
||||
const uint32_t apply_mask = hparams.rope_scaling_apply_mask;
|
||||
if ((is_swa && (apply_mask & 0x2)) || (!is_swa && (apply_mask & 0x1))) {
|
||||
rope_factors = build_rope_factors(il);
|
||||
}
|
||||
auto KQ_mask = is_swa ? build_inp_KQ_mask_swa() : build_inp_KQ_mask();
|
||||
const float kq_scale = 1.0f / sqrtf(float(hparams.n_embd_head_k(il)));
|
||||
|
||||
cur = build_std_attention(gf, mtp_layer.attn_norm, cur, inp_pos, nullptr,
|
||||
rope_factors, KQ_mask, nullptr, nullptr, kq_scale, 0.0f,
|
||||
is_swa ? hparams.n_swa : 0, il, true, false, true, false, false, nullptr, il);
|
||||
|
||||
if (mtp_layer.ffn_gate_inp == nullptr) {
|
||||
cur = llm_build_ffn(ctx0, lctx, mtp_layer.ffn_norm, cur,
|
||||
mtp_layer.ffn_up, nullptr, nullptr,
|
||||
mtp_layer.ffn_gate, nullptr, nullptr,
|
||||
mtp_layer.ffn_down, nullptr, nullptr,
|
||||
nullptr, LLM_FFN_SILU, LLM_FFN_PAR, cb, il, gf, true);
|
||||
} else {
|
||||
cur = llm_build_std_moe_ffn(ctx0, lctx, mtp_layer.ffn_norm, cur,
|
||||
mtp_layer.ffn_gate_inp, nullptr,
|
||||
mtp_layer.ffn_up_exps, nullptr,
|
||||
mtp_layer.ffn_gate_exps, nullptr,
|
||||
mtp_layer.ffn_down_exps, nullptr,
|
||||
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, hparams.expert_weights_scale != 0.0f,
|
||||
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, "mtp_post_ffn", il);
|
||||
if (hidden_out) {
|
||||
*hidden_out = cur;
|
||||
}
|
||||
|
||||
ggml_tensor * output_hidden = cur;
|
||||
if (reduce_output) {
|
||||
if (cparams.mtp_op_type != MTP_OP_NONE && n_tokens > 1) {
|
||||
output_hidden = ggml_view_2d(ctx0, cur, n_embd, 1,
|
||||
cur->nb[1], (size_t) (n_tokens - 1) * cur->nb[1]);
|
||||
} else if (inp_out_ids) {
|
||||
output_hidden = ggml_get_rows(ctx0, cur, inp_out_ids);
|
||||
}
|
||||
}
|
||||
if (reduce_output) {
|
||||
ggml_tensor * mtp_embd = output_hidden->type == GGML_TYPE_F32 ? output_hidden : ggml_cast(ctx0, output_hidden, GGML_TYPE_F32);
|
||||
cb(mtp_embd, "result_mtp_embd", -1);
|
||||
ggml_set_output(mtp_embd);
|
||||
ggml_build_forward_expand(gf, mtp_embd);
|
||||
}
|
||||
|
||||
if (!emit_logits) {
|
||||
return output_hidden;
|
||||
}
|
||||
|
||||
ggml_tensor * head_norm = mtp_layer.nextn.shared_head_norm
|
||||
? mtp_layer.nextn.shared_head_norm : model.output_norm;
|
||||
ggml_tensor * head = mtp_layer.nextn.shared_head_head
|
||||
? mtp_layer.nextn.shared_head_head : model.output;
|
||||
GGML_ASSERT(head_norm && head);
|
||||
cur = llm_build_context::build_output(lctx, ctx0, output_hidden, head, head_norm, cb);
|
||||
cb(cur, "result_output", -1);
|
||||
return cur;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -171,7 +171,8 @@ ggml_cgraph * llm_build_context::build_k_shift() {
|
|||
cb(lctx.inp_K_shift, "K_shift", -1);
|
||||
ggml_set_input(lctx.inp_K_shift);
|
||||
|
||||
for (int il = 0; il < n_layer; ++il) {
|
||||
const int n_kv_layers = model.mtp ? hparams.n_layer : hparams.n_layer - hparams.nextn_predict_layers;
|
||||
for (int il = 0; il < n_kv_layers; ++il) {
|
||||
if (llm_arch_is_hybrid(model.arch) && hparams.is_recurrent(il)) {
|
||||
continue;
|
||||
}
|
||||
|
|
@ -748,6 +749,15 @@ ggml_tensor * llm_build_context::build_inp_s_seq() {
|
|||
ggml_cgraph * llm_build_context::append_pooling(struct ggml_cgraph * gf) {
|
||||
// find result_norm tensor for input
|
||||
struct ggml_tensor * inp = nullptr;
|
||||
if (lctx.cparams.mtp) {
|
||||
for (int i = gf->n_nodes - 1; i >= 0; --i) {
|
||||
if (strcmp(gf->nodes[i]->name, "result_mtp_embd") == 0) {
|
||||
inp = gf->nodes[i];
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!inp) {
|
||||
for (int i = gf->n_nodes - 1; i >= 0; --i) {
|
||||
inp = gf->nodes[i];
|
||||
|
||||
|
|
@ -758,7 +768,9 @@ ggml_cgraph * llm_build_context::append_pooling(struct ggml_cgraph * gf) {
|
|||
}
|
||||
inp = nullptr;
|
||||
}
|
||||
}
|
||||
GGML_ASSERT(inp != nullptr && "missing result_norm/result_embd tensor");
|
||||
const bool is_mtp_hidden = strcmp(inp->name, "result_mtp_embd") == 0;
|
||||
|
||||
struct ggml_tensor * cur;
|
||||
|
||||
|
|
@ -784,7 +796,9 @@ ggml_cgraph * llm_build_context::append_pooling(struct ggml_cgraph * gf) {
|
|||
}
|
||||
}
|
||||
|
||||
if (!is_mtp_hidden || pooling_type != LLAMA_POOLING_TYPE_NONE) {
|
||||
cb(cur, "result_embd_pooled", -1);
|
||||
}
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
|
||||
|
|
@ -2957,7 +2971,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
|
|||
ggml_tensor * input, ggml_tensor * inp_pos, ggml_tensor * inp_out_ids, ggml_tensor * rope_factors_in,
|
||||
ggml_tensor * KQ_mask, ggml_tensor * sinks, ggml_tensor * inp_attn_scale, float KQ_scale, float f_attn_scale,
|
||||
int n_swa, int il, bool do_rope, bool add_graph_split, bool add_input, bool is_norm, bool is_multi,
|
||||
ggml_tensor * post_norm) {
|
||||
ggml_tensor * post_norm, int kv_il) {
|
||||
|
||||
float freq_base_l = n_swa > 0 ? hparams.rope_freq_base_train_swa : cparams.rope_freq_base;
|
||||
float freq_scale_l = n_swa > 0 ? hparams.rope_freq_scale_train_swa : hparams.rope_freq_scale_train;
|
||||
|
|
@ -3325,7 +3339,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
|
|||
if (auto wqkv_gate = model.layers[il].wqkv_gate; wqkv_gate != nullptr) {
|
||||
cur = llm_build_kv(ctx0, lctx, kv_self, gf,
|
||||
nullptr, nullptr,
|
||||
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa);
|
||||
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il);
|
||||
cb(cur, "wqkv", il);
|
||||
auto gate = llm_build_lora_mm(lctx, ctx0, wqkv_gate, input_normed);
|
||||
if (model.arch == LLM_ARCH_LAGUNA) {
|
||||
|
|
@ -3359,7 +3373,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
|
|||
} else {
|
||||
if (gate) {
|
||||
cur = llm_build_kv(ctx0, lctx, kv_self, gf, nullptr, nullptr,
|
||||
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa);
|
||||
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il);
|
||||
if (false && cur->ne[1] == 1) { // we need to add GGML_UNARY_OP_SIGMOID to the ops supported by ggml_fused_mul_unary
|
||||
cur = ggml_fused_mul_unary(ctx0, cur, gate, GGML_UNARY_OP_SIGMOID);
|
||||
} else {
|
||||
|
|
@ -3376,7 +3390,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
|
|||
} else {
|
||||
cur = llm_build_kv(ctx0, lctx, kv_self, gf,
|
||||
model.layers[il].wo, model.layers[il].bo,
|
||||
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa);
|
||||
Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, KQ_scale, cb, il, sinks, n_swa, kv_il);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -3403,7 +3417,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
|
|||
}
|
||||
|
||||
int32_t llama_model_n_nextn_layer(const llama_model * model) {
|
||||
return model->hparams.nextn_predict_layers;
|
||||
return model ? model->hparams.nextn_predict_layers : 0;
|
||||
}
|
||||
|
||||
ggml_cgraph * llm_build_context::new_graph_custom() {
|
||||
|
|
|
|||
|
|
@ -431,6 +431,15 @@ struct llm_build_context {
|
|||
|
||||
ggml_cgraph * build_step35();
|
||||
|
||||
ggml_tensor * build_step35_mtp(
|
||||
const llama_layer & mtp_layer,
|
||||
ggml_tensor * hidden_states_from_main_model,
|
||||
ggml_cgraph * gf,
|
||||
ggml_tensor * inp_pos,
|
||||
bool reduce_output = true,
|
||||
bool emit_logits = true,
|
||||
ggml_tensor ** hidden_out = nullptr);
|
||||
|
||||
//
|
||||
static ggml_tensor * llm_build_lora_mm(llama_context & lctx, ggml_context * ctx0,
|
||||
ggml_tensor * w, ggml_tensor * cur);
|
||||
|
|
@ -579,7 +588,7 @@ llm_expert_gating_func_type gating_op,
|
|||
ggml_tensor * inp_pos, ggml_tensor * inp_out_ids, ggml_tensor * rope_factors,
|
||||
ggml_tensor * KQ_mask, ggml_tensor * sinks, ggml_tensor * inp_attn_scale, float KQ_scale, float f_attn_scale,
|
||||
int n_swa, int il, bool do_rope = true, bool add_graph_split = false, bool add_input = false, bool is_norm = false,
|
||||
bool is_multi = false, ggml_tensor * post_norm = nullptr);
|
||||
bool is_multi = false, ggml_tensor * post_norm = nullptr, int kv_il = -1);
|
||||
|
||||
static ggml_tensor * build_output(llama_context & lctx, ggml_context * ctx, ggml_tensor * cur, ggml_tensor * output, const llm_build_cb & cb);
|
||||
|
||||
|
|
|
|||
|
|
@ -1564,6 +1564,12 @@ void llm_load_hparams(
|
|||
case LLM_ARCH_STEP35:
|
||||
{
|
||||
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
|
||||
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
|
||||
if (hparams.nextn_predict_layers > hparams.n_layer) {
|
||||
throw std::runtime_error(format("step35.nextn_predict_layers (%u) exceeds block_count (%u)",
|
||||
hparams.nextn_predict_layers, hparams.n_layer));
|
||||
}
|
||||
hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers;
|
||||
//hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
|
||||
// MoE + SWA parameters
|
||||
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
|
||||
|
|
|
|||
|
|
@ -1190,6 +1190,12 @@ bool create_tensors_helper::create_step35_tensors(const LLM_TN & tn) {
|
|||
// output
|
||||
model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
|
||||
model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
|
||||
|
||||
// If, for any reason, an MTP-only GGUF reports a count of physical blocks in addition to the MTP blk, skip them
|
||||
const bool mtp_only = hparams.nextn_predict_layers > 0 &&
|
||||
ml.get_tensor_meta("blk.0.attn_norm.weight") == nullptr;
|
||||
const int trunk_flags = mtp_only
|
||||
? llama_model_loader::TENSOR_SKIP | llama_model_loader::TENSOR_NOT_REQUIRED : 0;
|
||||
// STEP35 supports per-layer partial RoPE dims; rope factors are stored as a single shared tensor
|
||||
// ("rope_freqs.weight") and ggml uses only the first (n_rot_l/2) entries per layer.
|
||||
uint32_t n_rot_max = 0;
|
||||
|
|
@ -1202,22 +1208,38 @@ bool create_tensors_helper::create_step35_tensors(const LLM_TN & tn) {
|
|||
for (int i = 0; i < n_layer; ++i) {
|
||||
ggml_context * ctx_split = ctx_for_layer_split(i);
|
||||
auto & layer = model.layers[i];
|
||||
const bool is_mtp_layer = hparams.nextn_predict_layers > 0 &&
|
||||
static_cast<uint32_t>(i) >= hparams.n_layer - hparams.nextn_predict_layers;
|
||||
const std::string mtp_probe_name = tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i);
|
||||
const bool mtp_layer_present = !is_mtp_layer || ml.get_tensor_meta(mtp_probe_name.c_str()) != nullptr;
|
||||
const int layer_flags = mtp_only && !is_mtp_layer
|
||||
? trunk_flags
|
||||
: (is_mtp_layer && !mtp_layer_present
|
||||
? llama_model_loader::TENSOR_SKIP | llama_model_loader::TENSOR_NOT_REQUIRED : 0);
|
||||
const int optional_layer_flags = layer_flags | llama_model_loader::TENSOR_NOT_REQUIRED;
|
||||
const int nextn_required_flags = mtp_layer_present ? 0 : llama_model_loader::TENSOR_NOT_REQUIRED;
|
||||
const uint32_t n_head_l = hparams.n_head(i);
|
||||
layer.attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
|
||||
layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
layer.attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, layer_flags);
|
||||
layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, optional_layer_flags);
|
||||
layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, optional_layer_flags);
|
||||
// optional rope factors (llama3) / longrope tensors
|
||||
const int rope_flags = mtp_only
|
||||
? llama_model_loader::TENSOR_NOT_REQUIRED
|
||||
: optional_layer_flags;
|
||||
if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) {
|
||||
layer.rope_long = create_tensor(ctx_split, tn(LLM_TENSOR_ROPE_FACTORS_LONG, "weight", i), {n_rot_max/2}, llama_model_loader::TENSOR_NOT_REQUIRED | (i != 0 ? llama_model_loader::TENSOR_DUPLICATED : 0));
|
||||
layer.rope_short = create_tensor(ctx_split, tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot_max/2}, llama_model_loader::TENSOR_NOT_REQUIRED | (i != 0 ? llama_model_loader::TENSOR_DUPLICATED : 0));
|
||||
layer.rope_long = create_tensor(ctx_split, tn(LLM_TENSOR_ROPE_FACTORS_LONG, "weight", i), {n_rot_max/2}, rope_flags | (i != 0 ? llama_model_loader::TENSOR_DUPLICATED : 0));
|
||||
layer.rope_short = create_tensor(ctx_split, tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot_max/2}, rope_flags | (i != 0 ? llama_model_loader::TENSOR_DUPLICATED : 0));
|
||||
} else {
|
||||
layer.rope_freqs = create_tensor(ctx_split, tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot_max/2}, llama_model_loader::TENSOR_NOT_REQUIRED | (i != 0 ? llama_model_loader::TENSOR_DUPLICATED : 0));
|
||||
layer.rope_freqs = create_tensor(ctx_split, tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot_max/2}, rope_flags | (i != 0 ? llama_model_loader::TENSOR_DUPLICATED : 0));
|
||||
}
|
||||
if (layer_flags == 0) {
|
||||
use_mmap_buffer &= !merge_qkv(tn, i, 0);
|
||||
//layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head_l}, 0);
|
||||
//layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, 0);
|
||||
//layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, 0);
|
||||
layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v * n_head_l, n_embd}, 0);
|
||||
} else {
|
||||
layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head_l}, layer_flags);
|
||||
layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, layer_flags);
|
||||
layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, layer_flags);
|
||||
}
|
||||
layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v * n_head_l, n_embd}, layer_flags);
|
||||
const std::string attn_gate_name = tn(LLM_TENSOR_ATTN_GATE, "weight", i);
|
||||
int64_t n_attn_gate = n_head_l;
|
||||
if (model.arch == LLM_ARCH_LAGUNA) {
|
||||
|
|
@ -1229,18 +1251,18 @@ bool create_tensors_helper::create_step35_tensors(const LLM_TN & tn) {
|
|||
n_attn_gate = n_embd_head_v * n_head_l;
|
||||
}
|
||||
}
|
||||
layer.wqkv_gate = create_tensor(ctx_split, attn_gate_name, {n_embd, n_attn_gate}, llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
|
||||
layer.wqkv_gate = create_tensor(ctx_split, attn_gate_name, {n_embd, n_attn_gate}, optional_layer_flags);
|
||||
layer.ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, layer_flags);
|
||||
// dense MLP (leading dense blocks)
|
||||
layer.ffn_gate = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_gate = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, optional_layer_flags);
|
||||
layer.ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, optional_layer_flags);
|
||||
layer.ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, optional_layer_flags);
|
||||
// MoE routed experts + selection bias (router_bias)
|
||||
const int64_t n_ff_exp = hparams.n_ff_exp;
|
||||
if (!layer.ffn_gate) {
|
||||
layer.ffn_gate_inp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert},
|
||||
llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, i, n_ff_exp);
|
||||
optional_layer_flags);
|
||||
use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, i, optional_layer_flags);
|
||||
//layer.ffn_gate_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert},
|
||||
// llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
//layer.ffn_down_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert},
|
||||
|
|
@ -1248,14 +1270,29 @@ bool create_tensors_helper::create_step35_tensors(const LLM_TN & tn) {
|
|||
//layer.ffn_up_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert},
|
||||
// llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
layer.ffn_exp_probs_b = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert},
|
||||
llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
optional_layer_flags);
|
||||
// shared expert MLP
|
||||
layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp},
|
||||
llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
optional_layer_flags);
|
||||
layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp},
|
||||
llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
optional_layer_flags);
|
||||
layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd},
|
||||
llama_model_loader::TENSOR_NOT_REQUIRED);
|
||||
optional_layer_flags);
|
||||
}
|
||||
|
||||
if (is_mtp_layer) {
|
||||
layer.nextn.eh_proj = create_tensor(ctx_split,
|
||||
tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, nextn_required_flags);
|
||||
layer.nextn.embed_tokens = create_tensor(ctx_split,
|
||||
tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, optional_layer_flags);
|
||||
layer.nextn.enorm = create_tensor(ctx_split,
|
||||
tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, nextn_required_flags);
|
||||
layer.nextn.hnorm = create_tensor(ctx_split,
|
||||
tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, nextn_required_flags);
|
||||
layer.nextn.shared_head_head = create_tensor(ctx_split,
|
||||
tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, optional_layer_flags);
|
||||
layer.nextn.shared_head_norm = create_tensor(ctx_split,
|
||||
tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, optional_layer_flags);
|
||||
}
|
||||
}
|
||||
return use_mmap_buffer;
|
||||
|
|
|
|||
|
|
@ -1760,6 +1760,12 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
|
|||
{ LLM_TENSOR_FFN_DOWN_SHEXP, "blk.%d.ffn_down_shexp" },
|
||||
{ LLM_TENSOR_FFN_UP_SHEXP, "blk.%d.ffn_up_shexp" },
|
||||
{ LLM_TENSOR_FFN_EXP_PROBS_B, "blk.%d.exp_probs_b" },
|
||||
{ LLM_TENSOR_NEXTN_EH_PROJ, "blk.%d.nextn.eh_proj" },
|
||||
{ LLM_TENSOR_NEXTN_EMBED_TOKENS, "blk.%d.nextn.embed_tokens" },
|
||||
{ LLM_TENSOR_NEXTN_ENORM, "blk.%d.nextn.enorm" },
|
||||
{ LLM_TENSOR_NEXTN_HNORM, "blk.%d.nextn.hnorm" },
|
||||
{ LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "blk.%d.nextn.shared_head_head" },
|
||||
{ LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "blk.%d.nextn.shared_head_norm" },
|
||||
},
|
||||
},
|
||||
{
|
||||
|
|
@ -2193,6 +2199,90 @@ bool llama_model_is_gemma4_mtp_assistant(const llama_model * model) {
|
|||
return model && (model->arch == LLM_ARCH_GEMMA4_MTP || model->arch == LLM_ARCH_GEMMA4_ASSISTANT);
|
||||
}
|
||||
|
||||
bool llama_model_is_step35(const llama_model * model) {
|
||||
return model && model->arch == LLM_ARCH_STEP35;
|
||||
}
|
||||
|
||||
enum llama_mtp_package llama_model_mtp_package(const llama_model * model) {
|
||||
if (!model) {
|
||||
return LLAMA_MTP_PACKAGE_INVALID;
|
||||
}
|
||||
|
||||
if (llama_model_is_gemma4_mtp_assistant(model) && model->hparams.nextn_predict_layers == 0) {
|
||||
return LLAMA_MTP_PACKAGE_COMPANION;
|
||||
}
|
||||
|
||||
const size_t n_nextn = model->hparams.nextn_predict_layers;
|
||||
const bool has_common_package_contract =
|
||||
llama_model_is_step35(model) || llama_model_is_deepseek4(model) ||
|
||||
llama_model_is_gemma4_mtp_assistant(model);
|
||||
if (!has_common_package_contract) {
|
||||
return n_nextn > 0 ? LLAMA_MTP_PACKAGE_EMBEDDED : LLAMA_MTP_PACKAGE_NONE;
|
||||
}
|
||||
|
||||
if (n_nextn == 0) {
|
||||
if (llama_model_is_step35(model) || llama_model_is_deepseek4(model)) {
|
||||
for (const auto & layer : model->layers) {
|
||||
if (layer.attn_norm != nullptr) {
|
||||
return LLAMA_MTP_PACKAGE_TARGET_ONLY;
|
||||
}
|
||||
}
|
||||
}
|
||||
return LLAMA_MTP_PACKAGE_NONE;
|
||||
}
|
||||
|
||||
const size_t n_layers = model->layers.size();
|
||||
if (n_nextn > n_layers) {
|
||||
return LLAMA_MTP_PACKAGE_INVALID;
|
||||
}
|
||||
|
||||
const size_t first = n_layers - n_nextn;
|
||||
const bool has_tail = model->layers[first].nextn.eh_proj != nullptr;
|
||||
|
||||
bool has_trunk = false;
|
||||
for (size_t il = 0; il < first; ++il) {
|
||||
if (model->layers[il].attn_norm != nullptr) {
|
||||
has_trunk = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (has_trunk && has_tail) {
|
||||
return LLAMA_MTP_PACKAGE_EMBEDDED;
|
||||
}
|
||||
if (has_trunk && !has_tail) {
|
||||
return LLAMA_MTP_PACKAGE_TARGET_ONLY;
|
||||
}
|
||||
if (!has_trunk && has_tail) {
|
||||
return LLAMA_MTP_PACKAGE_COMPANION;
|
||||
}
|
||||
return LLAMA_MTP_PACKAGE_INVALID;
|
||||
}
|
||||
|
||||
bool llama_model_step35_has_nextn_weights(const llama_model * model) {
|
||||
if (!model || !llama_model_is_step35(model) || model->hparams.nextn_predict_layers == 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const size_t n_layers = model->layers.size();
|
||||
const size_t n_nextn = model->hparams.nextn_predict_layers;
|
||||
if (n_nextn > n_layers) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const size_t first = n_layers - n_nextn;
|
||||
for (size_t il = first; il < n_layers; ++il) {
|
||||
const llama_layer & layer = model->layers[il];
|
||||
const llama_layer_nextn & nextn = layer.nextn;
|
||||
if (!nextn.eh_proj || !nextn.enorm || !nextn.hnorm ||
|
||||
!layer.attn_norm || !layer.wq || !layer.wk || !layer.wv || !layer.wo ||
|
||||
!layer.ffn_norm || !layer.ffn_gate || !layer.ffn_up || !layer.ffn_down) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool llama_is_gemma4_mtp_file(const char * path) {
|
||||
if (!path || !*path) return false;
|
||||
struct gguf_init_params params = { /*.no_alloc =*/ true, /*.ctx =*/ nullptr };
|
||||
|
|
|
|||
|
|
@ -10,15 +10,26 @@ uint32_t llama_mtp_state_n_embd(const struct llama_context * ctx) {
|
|||
return 0;
|
||||
}
|
||||
|
||||
const auto & hparams = ctx->model.hparams;
|
||||
if (ctx->cparams.mtp && (ctx->model.arch == LLM_ARCH_GEMMA4_MTP || ctx->model.arch == LLM_ARCH_GEMMA4_ASSISTANT) && hparams.mtp_backbone_n_embd > 0) {
|
||||
if (!ctx->cparams.mtp) {
|
||||
return ctx->model.hparams.n_embd;
|
||||
}
|
||||
|
||||
return llama_model_mtp_feature_width(&ctx->model);
|
||||
}
|
||||
|
||||
uint32_t llama_model_mtp_feature_width(const struct llama_model * model) {
|
||||
if (model == nullptr) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const auto & hparams = model->hparams;
|
||||
if ((model->arch == LLM_ARCH_GEMMA4_MTP || model->arch == LLM_ARCH_GEMMA4_ASSISTANT) &&
|
||||
hparams.mtp_backbone_n_embd > 0) {
|
||||
return hparams.mtp_backbone_n_embd;
|
||||
}
|
||||
|
||||
if (ctx->cparams.mtp && ctx->model.arch == LLM_ARCH_DEEPSEEK4 && hparams.n_embd_out > hparams.n_embd) {
|
||||
if (model->arch == LLM_ARCH_DEEPSEEK4 && hparams.n_embd_out > hparams.n_embd) {
|
||||
return hparams.n_embd_out;
|
||||
}
|
||||
|
||||
return hparams.n_embd;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1106,6 +1106,13 @@ static inline bool llama_kv_qnext_seq_id_in_range(const llama_kv_cache & cache,
|
|||
return n_slots > 0 && seq_id >= 0 && (uint32_t) seq_id < n_slots;
|
||||
}
|
||||
|
||||
static bool llama_mtp_tail_uses_layer_cache(const llama_model & model) {
|
||||
return model.hparams.nextn_predict_layers > 0 &&
|
||||
(model.arch == LLM_ARCH_GLM_DSA ||
|
||||
model.arch == LLM_ARCH_QWEN35MOE ||
|
||||
model.arch == LLM_ARCH_STEP35);
|
||||
}
|
||||
|
||||
static bool llama_kv_cache_init(
|
||||
struct llama_kv_cache & cache,
|
||||
const llama_context * ctx,
|
||||
|
|
@ -1215,9 +1222,7 @@ 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 is_mtp = (model.arch == LLM_ARCH_GLM_DSA ||
|
||||
//model.arch == LLM_ARCH_QWEN35 ||
|
||||
model.arch == LLM_ARCH_QWEN35MOE) && hparams.nextn_predict_layers > 0;
|
||||
const bool is_mtp = llama_mtp_tail_uses_layer_cache(model);
|
||||
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 = is_mtp && i >= n_mtp_first;
|
||||
|
|
@ -1322,10 +1327,7 @@ static bool llama_kv_cache_init(
|
|||
const uint32_t n_head_kv = hparams.n_head_kv(i);
|
||||
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_GLM_DSA) &&
|
||||
hparams.nextn_predict_layers > 0 && i >= n_mtp_first_layer;
|
||||
const bool is_mtp_tail_layer = llama_mtp_tail_uses_layer_cache(model) && i >= 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();
|
||||
ggml_tensor * k = nullptr;
|
||||
|
|
@ -4930,6 +4932,14 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam
|
|||
return -2;
|
||||
}
|
||||
|
||||
const enum llama_mtp_package mtp_package = llama_model_mtp_package(&model);
|
||||
if (mtp_package == LLAMA_MTP_PACKAGE_INVALID) {
|
||||
throw std::runtime_error("invalid MTP package: missing target trunk and predictor tail tensors");
|
||||
}
|
||||
if (mtp_package == LLAMA_MTP_PACKAGE_COMPANION && !params.mtp) {
|
||||
throw std::runtime_error("MTP companion GGUF cannot be used as the target model, pass it with -md/--draft");
|
||||
}
|
||||
|
||||
// ---- populate reload registry ONLY when hot-swap is requested ----
|
||||
if (std::getenv("LLAMA_HOTSWAP_ENABLED") != nullptr) {
|
||||
model.reload = std::make_unique<reload_info>(ml);
|
||||
|
|
@ -5839,6 +5849,7 @@ static uint32_t llama_output_embd_width(const llama_context & lctx) {
|
|||
static bool llama_context_has_mtp_outputs(const llama_context & lctx) {
|
||||
return lctx.cparams.mtp && (
|
||||
lctx.model.hparams.nextn_predict_layers > 0 ||
|
||||
llama_model_mtp_package(&lctx.model) == LLAMA_MTP_PACKAGE_TARGET_ONLY ||
|
||||
lctx.model.arch == LLM_ARCH_GEMMA4 ||
|
||||
lctx.model.arch == LLM_ARCH_GEMMA4_MTP ||
|
||||
lctx.model.arch == LLM_ARCH_GEMMA4_ASSISTANT ||
|
||||
|
|
@ -6396,13 +6407,8 @@ static int llama_decode_internal(
|
|||
}
|
||||
else {
|
||||
const bool has_mtp = llama_context_has_mtp_outputs(lctx);
|
||||
const bool use_raw_mtp_embd = has_mtp && (lctx.model.arch == LLM_ARCH_GEMMA4 ||
|
||||
lctx.model.arch == LLM_ARCH_GEMMA4_MTP||
|
||||
lctx.model.arch == LLM_ARCH_GEMMA4_ASSISTANT ||
|
||||
lctx.model.arch == LLM_ARCH_DEEPSEEK4);
|
||||
// For DSV4 we want to extract the 16,384-dim embedding first
|
||||
if (cparams.embeddings || has_mtp) {
|
||||
if (use_raw_mtp_embd) {
|
||||
if (has_mtp) {
|
||||
for (int i = gf->n_nodes - 1; i >= 0; --i) {
|
||||
if (strcmp(gf->nodes[i]->name, "result_mtp_embd") == 0) {
|
||||
embd = gf->nodes[i];
|
||||
|
|
@ -8141,6 +8147,7 @@ struct llama_context * llama_init_from_model(
|
|||
model->arch != LLM_ARCH_QWEN35MOE && model->arch != LLM_ARCH_GEMMA4 &&
|
||||
model->arch != LLM_ARCH_GEMMA4_MTP && model->arch != LLM_ARCH_GLM_DSA &&
|
||||
model->arch != LLM_ARCH_DEEPSEEK4 &&
|
||||
model->arch != LLM_ARCH_STEP35 &&
|
||||
model->arch != LLM_ARCH_GEMMA4_ASSISTANT &&
|
||||
model->arch != LLM_ARCH_OPENPANGU &&
|
||||
cparams.mtp != 0) {
|
||||
|
|
|
|||
Loading…
Reference in New Issue