From 1794846f7351ba405b3ab0b3654f8c9f9f387d01 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Sat, 15 Aug 2026 19:34:47 +0200 Subject: [PATCH] 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 --- common/speculative.cpp | 21 ++++++++++----------- src/graphs/build_gemma4.cpp | 12 +----------- 2 files changed, 11 insertions(+), 22 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 9783bd9f..b641f3ef 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -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; diff --git a/src/graphs/build_gemma4.cpp b/src/graphs/build_gemma4.cpp index 7268ba49..b7bc00b9 100644 --- a/src/graphs/build_gemma4.cpp +++ b/src/graphs/build_gemma4.cpp @@ -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);