nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix Gemma4 MTP (#2324) 

* Fix Gemma4 MTP

* Committed this change by mistake - reverting

* Fix Gemma4 assistant crash in split mode graph

* Disable some vocabulary compatibility chacks for Gemma4 assistant drafters
This commit is contained in:
Kawrakow 2026-08-15 19:34:47 +02:00 committed by GitHub
parent 7cd62a3eb2
commit 1794846f73
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
common/speculative.cpp
View File

@ -99,18 +99,17 @@ static bool common_speculative_are_compatible(
LOG_DBG("%s: vocab_type dft: %d\n", __func__, vocab_type_dft);
if (vocab_type_tgt != vocab_type_dft) {
LOG_DBG("%s: draft model vocab type must match target model to use speculation but ", __func__);
LOG_DBG("vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt);
LOG_WRN("%s: draft model vocab type must match target model to use speculation but ", __func__);
LOG_WRN("vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt);
return false;
}
if (
llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) ||
if (!llama_model_is_gemma4_mtp_assistant(model_dft) &&
(llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) ||
llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) ||
llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft) ||
llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft)
) {
LOG_DBG("%s: draft model special tokens must match target model to use speculation\n", __func__);
llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft))) {
LOG_WRN("%s: draft model special tokens must match target model to use speculation\n", __func__);
return false;
}
@ -122,8 +121,8 @@ static bool common_speculative_are_compatible(
: n_vocab_dft - n_vocab_tgt;
if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) {
LOG_DBG("%s: draft model vocab must closely match target model to use speculation but ", __func__);
LOG_DBG("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n",
LOG_WRN("%s: draft model vocab must closely match target model to use speculation but ", __func__);
LOG_WRN("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n",
n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE);
return false;
}
@ -133,8 +132,8 @@ static bool common_speculative_are_compatible(
const char * token_text_dft = llama_vocab_get_text(vocab_dft, i);
if (std::strcmp(token_text_tgt, token_text_dft) != 0) {
LOG_DBG("%s: draft model vocab must match target model to use speculation but ", __func__);
LOG_DBG("token %d content differs - target '%s', draft '%s'\n", i,
LOG_WRN("%s: draft model vocab must match target model to use speculation but ", __func__);
LOG_WRN("token %d content differs - target '%s', draft '%s'\n", i,
common_token_to_piece(vocab_tgt, i).c_str(),
common_token_to_piece(vocab_dft, i).c_str());
return false;

12
src/graphs/build_gemma4.cpp
View File

@ -548,10 +548,6 @@ ggml_cgraph * llm_build_context::build_gemma4_mtp() {
ggml_row_size(hidden_state->type, n_backbone), 0);
cb(cur, "mtp_init_hidden_view", -1);
ggml_tensor * mtp_embd = ggml_dup(ctx0, hidden_state);
cb(mtp_embd, "result_mtp_embd", -1);
ggml_build_forward_expand(gf, mtp_embd);
ggml_tensor * logits = build_output(lctx, ctx0, cur, model.output, model.output_norm, cb);
cb(logits, "result_output", -1);
ggml_build_forward_expand(gf, logits);
@ -676,7 +672,7 @@ ggml_cgraph * llm_build_context::build_gemma4_mtp() {
GGML_ASSERT(model.layers[il].attn_q_norm && model.layers[il].attn_q_norm->extra);
Qcur = do_split_norm(ctx0, Qcur, model.layers[il].attn_q_norm, hparams, cb, id, il_cb, false);
cb(Qcur, "Qcur_normed", il_cb);
auto freq_factors = is_sliding ? nullptr : ((const ggml_split_tensor_t *)model.layers[il].rope_freqs->extra)->splits[id];
auto freq_factors = is_sliding || !model.layers[il].rope_freqs ? nullptr : ((const ggml_split_tensor_t *)model.layers[il].rope_freqs->extra)->splits[id];
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, freq_factors, n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(Qcur, "Qcur_rope", il_cb);
@ -1123,12 +1119,6 @@ ggml_cgraph * llm_build_context::build_gemma4() {
cur = inpL;
if (cparams.mtp) {
ggml_tensor * mtp_embd = ggml_dup(ctx0, cur);
cb(mtp_embd, "result_mtp_embd", -1);
ggml_build_forward_expand(gf, mtp_embd);
}
cur = llm_build_norm(ctx0, cur, hparams, model.output_norm, NULL, LLM_NORM_RMS, cb, -1);
cb(cur, "result_norm", -1);