From 86e33fd6f44632b6e2354f736dab5000bb18f4b0 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Mon, 6 Apr 2026 10:01:08 +0200 Subject: [PATCH] Initial Gemma4 support (#1581) * Gemma4: WIP * Gemma4: WIP - runs with totally wrong results * Gemma4: WIP - add CPU 512, 512 FA * Gemma4: WIP It gives a meaningful response in llama-cli, but PPL is still much too high. Is this due to tokenizer issues? * Gemma4: this works I had forgotten the softcap on the final output. * Remove log * Gemma4: WIP E4B/E2B * Gemma4: Q4B/E2B appear to work now * gemma4: tokenizer fixes --- ggml/src/CMakeLists.txt | 1 + ggml/src/ggml-backend.cpp | 4 + ggml/src/ggml-cuda/fattn-new-mma.cu | 67 ++- ggml/src/ggml-cuda/fattn.cu | 5 +- ggml/src/iqk/fa/iqk_fa_512_512.cpp | 124 +++++ ggml/src/iqk/fa/iqk_fa_templates.h | 1 + ggml/src/iqk/iqk_mul_mat.cpp | 4 + src/llama-arch.cpp | 7 +- src/llama-arch.h | 17 + src/llama-build-context.cpp | 700 ++++++++++++++++++++-------- src/llama-build-context.h | 11 +- src/llama-hparams.cpp | 78 +++- src/llama-hparams.h | 30 +- src/llama-load-tensors.cpp | 131 +++++- src/llama-model.cpp | 37 +- src/llama-model.h | 16 + src/llama-vocab.cpp | 402 ++++++++++++---- src/llama-vocab.h | 12 +- src/llama.cpp | 22 +- src/unicode.cpp | 241 ++++++---- src/unicode.h | 44 +- 21 files changed, 1513 insertions(+), 441 deletions(-) create mode 100644 ggml/src/iqk/fa/iqk_fa_512_512.cpp diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 4cf521b1..078876dd 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -264,6 +264,7 @@ if (GGML_IQK_MUL_MAT) set(GGML_SOURCES_IQK_MM iqk/iqk_mul_mat.cpp iqk/iqk_flash_attn.cpp iqk/fa/iqk_fa_576_512.cpp + iqk/fa/iqk_fa_512_512.cpp iqk/fa/iqk_fa_320_256.cpp iqk/fa/iqk_fa_192_128.cpp iqk/fa/iqk_fa_192_192.cpp diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 17f57123..97c7481a 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1824,6 +1824,10 @@ static void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct gg SET_CAUSE(tensor_copy, "4.cpy"); } int n_inputs = split->n_inputs++; + if (n_inputs >= GGML_SCHED_MAX_SPLIT_INPUTS) { + printf("======================== Oops, too many inputs (%d)\n", n_inputs+1); + for (int k = 0; k < n_inputs; ++k) printf("Input %2d: %s\n", k, split->inputs[k]->name); + } GGML_ASSERT(n_inputs < GGML_SCHED_MAX_SPLIT_INPUTS); split->inputs[n_inputs] = src; } diff --git a/ggml/src/ggml-cuda/fattn-new-mma.cu b/ggml/src/ggml-cuda/fattn-new-mma.cu index 3d6e4665..3fe8633a 100644 --- a/ggml/src/ggml-cuda/fattn-new-mma.cu +++ b/ggml/src/ggml-cuda/fattn-new-mma.cu @@ -332,6 +332,38 @@ struct fattn_mma_f16_config<192, 192> { } }; +template <> +struct fattn_mma_f16_config<512, 512> { + static constexpr int nbatch_fa = 64; + static constexpr int nwarps_max = 4; + static constexpr bool Q_in_reg = false; + static constexpr int nstages_target = 1; + + static int get_nbatch_K2_host([[maybe_unused]] const int cc, [[maybe_unused]] const int ncols) { + return 256; + } + + static constexpr __device__ int get_nbatch_K2_device([[maybe_unused]] int ncols) { + return 256; + } + + static int get_nbatch_V2_host([[maybe_unused]] const int cc, [[maybe_unused]] const int ncols) { + return 256; + } + + static constexpr __device__ int get_nbatch_V2_device([[maybe_unused]] int ncols) { + return 256; + } + + static int get_nbatch_combine_host(const int /*cc*/, const int /*ncols*/) { + return 256; + } + + static constexpr __device__ int get_nbatch_combine_device(int /*ncols*/) { + return 256; + } +}; + template <> struct fattn_mma_f16_config<576, 512> { static constexpr int nbatch_fa = 32; @@ -2025,7 +2057,8 @@ static void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ct const int nstages = cp_async_available(cc) ? c::nstages_target : 0; constexpr int ncols = ncols1 * ncols2; - constexpr int ntiles = ncols <= 8 && DKQ < 576 ? 1 : 2; // Number of tiles per warp. + //constexpr int ntiles = DKQ == 512 ? 2 : ncols <= 8 && DKQ < 576 ? 1 : 2; // Number of tiles per warp. + constexpr int ntiles = ncols <= 8 && DKQ < 512 ? 1 : 2; // Number of tiles per warp. constexpr int cols_per_warp = ntiles * tile_B::I; constexpr int nwarps_max_x = (ncols + cols_per_warp - 1) / cols_per_warp; constexpr int nwarps_max_y = c::nbatch_fa / tile_A::I; @@ -2089,15 +2122,24 @@ template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * Q = dst->src[0]; - if constexpr (DKQ == 576 && ncols2 <= 4) { + if constexpr ((DKQ == 576 || DKQ == 512) && ncols2 <= 4) { ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); } else { if constexpr (ncols2 <= 8) { + //if constexpr (DKQ == 512) { + // ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + // return; + //} else { if (Q->ne[1] <= 8/ncols2) { - ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + if constexpr (DKQ == 512) { + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + } return; } + //} } if (Q->ne[1] <= 16/ncols2) { @@ -2109,8 +2151,11 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); return; } - - ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + if constexpr (DKQ == 512) { + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + } else { + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + } } } @@ -2209,6 +2254,18 @@ void ggml_cuda_flash_attn_ext_mma_new(ggml_backend_cuda_context & ctx, ggml_tens ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<320, 256, 16>(ctx, dst); return; } + if (Q->ne[0] == 512 && K->ne[0] == 512 && V->ne[0] == 512) { + if (gqa_ratio == 8) { + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<512, 512, 8>(ctx, dst); + } + else if (gqa_ratio == 4) { + ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<512, 512, 4>(ctx, dst); + } + else { + GGML_ABORT("Fatal error"); + } + return; + } GGML_ASSERT(Q->ne[0] == 576 && K->ne[0] == 576 && V->ne[0] == 512); if (gqa_ratio == 20 && Q->ne[1] <= 4 && K->ne[1] >= 2048) { if (ggml_cuda_info().devices[ctx.device].cc >= CC_ADA_LOVELACE) { diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index b5347194..4eff602e 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -122,6 +122,7 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst if (new_mma_available(cc) && ((K->ne[0] == 576 && V->ne[0] == 512) || (K->ne[0] == 320 && V->ne[0] == 256) || + (K->ne[0] == 512 && V->ne[0] == 512) || (K->ne[0] == 192 && V->ne[0] == 128 && mma_better_than_turing(cc)))) { //printf("Using ggml_cuda_flash_attn_ext_mma_new\n"); ggml_cuda_flash_attn_ext_mma_new(ctx, dst); @@ -194,8 +195,8 @@ bool ggml_cuda_fattn_is_supported(ggml_backend_cuda_context & ctx, const ggml_te } if (new_mma_available(cc) && - (Q->ne[0] == 576 || Q->ne[0] == 320 || (K->ne[0] == 192 && V->ne[0] == 128 && mma_better_than_turing(cc)))) { - if (Q->ne[0] == 576 || Q->ne[0] == 320) { + (Q->ne[0] == 576 || Q->ne[0] == 320 || Q->ne[0] == 512 || (K->ne[0] == 192 && V->ne[0] == 128 && mma_better_than_turing(cc)))) { + if (Q->ne[0] == 576 || Q->ne[0] == 512 || Q->ne[0] == 320) { int gqa_ratio = Q->ne[2]/K->ne[2]; return (gqa_ratio % 4) == 0; } diff --git a/ggml/src/iqk/fa/iqk_fa_512_512.cpp b/ggml/src/iqk/fa/iqk_fa_512_512.cpp new file mode 100644 index 00000000..6e2e7935 --- /dev/null +++ b/ggml/src/iqk/fa/iqk_fa_512_512.cpp @@ -0,0 +1,124 @@ +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include "iqk/fa/iqk_fa_templates.h" + +namespace { + +template +inline void iqk_deepseek_helper(KHelper& kh, VHelper& vh, + int nq1, int nk1, int stride_q, int stride_m, int stride_qkv, + const float * q, const char * mask, float scale, float softcap, float * qkv, + const float * sinkf, float * M, float * S) { + auto update = [&nq1, &mask, &q, &qkv, &M, &S, stride_q, stride_m, stride_qkv] (int n) { + nq1 -= n; + if (nq1 == 0) return true; + q += n*stride_q; + mask += n*stride_m; + qkv += n*stride_qkv; + if (M && S) { M += n; S += n; } + return false; + }; + if (nq1 >= 16) { + int n_step = nq1/16; + FlashAttn<512, 512, 16, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 16*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + if (update(16*n_step)) return; + } + if (nq1 >= 8) { + int n_step = nq1/8; + FlashAttn<512, 512, 8, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 8*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + if (update(8*n_step)) return; + } + if (nq1 >= 4) { + int n_step = nq1/4; + FlashAttn<512, 512, 4, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 4*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + if (update(4*n_step)) return; + } + if (nq1 == 3) { + FlashAttn<512, 512, 3, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 3, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + } + else if (nq1 == 2) { + FlashAttn<512, 512, 2, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 2, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + } else { + FlashAttn<512, 512, 1, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 1, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + } +} + +template +inline bool iqk_deepseek_helper(ggml_type type_k, + int nq1, int nk1, int stride_q, int stride_k, int stride_v, int stride_m, int stride_qkv, + const float * q, const char * k, const char * v, const char * mask, + float scale, float softcap, float * qkv, const float * sinkf, float * M, float * S) { + if (type_k == GGML_TYPE_Q8_0) { + HelperQ80 kh((const char *)k, stride_k); + HelperQ80 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_Q8_0_R8) { + HelperQ80R8<512> kh((const char *)k, stride_k); + HelperQ80 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_Q6_0) { + HelperQ60 kh((const char *)k, stride_k); + HelperQ60 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } +#if GGML_IQK_FA_ALL_QUANTS + if (type_k == GGML_TYPE_Q8_KV) { + HelperQ8KV<512> kh((const char *)k, stride_k); + HelperQ8KV<512> vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + if (type_k == GGML_TYPE_F16) { + HelperF16 kh((const char *)k, stride_k); + HelperF16 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + HelperBF16<512, step_k> kh((const char *)k, stride_k); + HelperBF16<512, step_k> vh((const char *)v, stride_v); + if (nq1 % 8 == 0) { + FlashAttnBF16<512, 512, 8, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + } else { + FlashAttnBF16<512, 512, 1, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + } + return true; + } +#endif + return false; +} + +} + +IQK_FA_CASE(iqk_fa_512_512) { + + auto type_k = ggml_type(int_type_k); + auto type_v = ggml_type(int_type_v); + + if (!(type_k == type_v || (type_k == GGML_TYPE_Q8_0_R8 && type_v == GGML_TYPE_Q8_0))) { + return false; + } + stride_q /= sizeof(float); // q stride as float + return iqk_deepseek_helper<32>(type_k, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, (const char *)k, (const char *)v, (const char *)mask, scale, softcap, qkv, sinkf, M, S); + +} + +#endif diff --git a/ggml/src/iqk/fa/iqk_fa_templates.h b/ggml/src/iqk/fa/iqk_fa_templates.h index 2cd0c18f..3aaff08c 100644 --- a/ggml/src/iqk/fa/iqk_fa_templates.h +++ b/ggml/src/iqk/fa/iqk_fa_templates.h @@ -2239,6 +2239,7 @@ inline bool iqk_flash_helper_T(ggml_type type_k, ggml_type type_v, float * qkv, const float * sinkf, float * M, float * S) IQK_FA_CASE(iqk_fa_576_512); +IQK_FA_CASE(iqk_fa_512_512); IQK_FA_CASE(iqk_fa_320_256); IQK_FA_CASE(iqk_fa_192_128); IQK_FA_CASE(iqk_fa_192_192); diff --git a/ggml/src/iqk/iqk_mul_mat.cpp b/ggml/src/iqk/iqk_mul_mat.cpp index 9d98081a..5a3800e8 100644 --- a/ggml/src/iqk/iqk_mul_mat.cpp +++ b/ggml/src/iqk/iqk_mul_mat.cpp @@ -1389,6 +1389,10 @@ bool iqk_flash_attn_impl(int int_type_k, // type of k return iqk_fa_576_512(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, q, k, v, mask, scale, softcap, qkv, sinksf, M, S); } + if (Dk == 512 && Dv == 512) { + return iqk_fa_512_512(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, k, v, mask, scale, softcap, qkv, sinksf, M, S); + } if (Dk == 320 && Dv == 256) { return iqk_fa_320_256(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, q, k, v, mask, scale, softcap, qkv, sinksf, M, S); diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 18a29815..ff6fe0be 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -77,6 +77,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_STEP35, "step35" }, { LLM_ARCH_GLM_DSA, "glm-dsa" }, { LLM_ARCH_MISTRAL4, "mistral4" }, + { LLM_ARCH_GEMMA4, "gemma4" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; @@ -107,6 +108,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_VOCAB_SIZE, "%s.vocab_size" }, { LLM_KV_CONTEXT_LENGTH, "%s.context_length" }, { LLM_KV_EMBEDDING_LENGTH, "%s.embedding_length" }, + { LLM_KV_EMBEDDING_LENGTH_PER_LAYER, "%s.embedding_length_per_layer_input" }, { LLM_KV_BLOCK_COUNT, "%s.block_count" }, { LLM_KV_LEADING_DENSE_BLOCK_COUNT, "%s.leading_dense_block_count" }, { LLM_KV_FEED_FORWARD_LENGTH, "%s.feed_forward_length" }, @@ -163,9 +165,12 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" }, { LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" }, { LLM_KV_FULL_ATTENTION_INTERVAL, "%s.full_attention_interval" }, - + { LLM_KV_ATTENTION_SHARED_KV_LAYERS, "%s.attention.shared_kv_layers" }, + { LLM_KV_ATTENTION_KEY_LENGTH_SWA, "%s.attention.key_length_swa" }, + { LLM_KV_ATTENTION_VALUE_LENGTH_SWA, "%s.attention.value_length_swa" }, { LLM_KV_ROPE_DIMENSION_COUNT, "%s.rope.dimension_count" }, + { LLM_KV_ROPE_DIMENSION_COUNT_SWA, "%s.rope.dimension_count_swa" }, { LLM_KV_ROPE_DIMENSION_COUNT_PER_LAYER,"%s.rope.dimension_count_per_layer" }, { LLM_KV_ROPE_DIMENSION_SECTIONS, "%s.rope.dimension_sections" }, { LLM_KV_ROPE_FREQ_BASE, "%s.rope.freq_base" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 3f621174..20d6085c 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -76,6 +76,7 @@ enum llm_arch { LLM_ARCH_STEP35, LLM_ARCH_GLM_DSA, LLM_ARCH_MISTRAL4, + LLM_ARCH_GEMMA4, LLM_ARCH_UNKNOWN, }; @@ -131,6 +132,7 @@ enum llm_kv { LLM_KV_SWIGLU_LIMITS_SHARED, LLM_KV_SWIGLU_CLAMP_EXP, LLM_KV_SWIGLU_CLAMP_SHEXP, + LLM_KV_EMBEDDING_LENGTH_PER_LAYER, LLM_KV_ATTENTION_HEAD_COUNT, LLM_KV_ATTENTION_HEAD_COUNT_KV, @@ -156,8 +158,12 @@ enum llm_kv { LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, LLM_KV_ATTENTION_INDEXER_TOP_K, LLM_KV_FULL_ATTENTION_INTERVAL, + LLM_KV_ATTENTION_SHARED_KV_LAYERS, + LLM_KV_ATTENTION_KEY_LENGTH_SWA, + LLM_KV_ATTENTION_VALUE_LENGTH_SWA, LLM_KV_ROPE_DIMENSION_COUNT, + LLM_KV_ROPE_DIMENSION_COUNT_SWA, LLM_KV_ROPE_DIMENSION_COUNT_PER_LAYER, LLM_KV_ROPE_DIMENSION_SECTIONS, LLM_KV_ROPE_FREQ_BASE, @@ -341,6 +347,17 @@ enum llm_tensor { LLM_TENSOR_INDEXER_ATTN_K, LLM_TENSOR_INDEXER_ATTN_Q_B, // 97 + LLM_TENSOR_PER_LAYER_TOKEN_EMBD, + LLM_TENSOR_PER_LAYER_MODEL_PROJ, + LLM_TENSOR_PER_LAYER_INP_GATE, // 100 + LLM_TENSOR_PER_LAYER_PROJ, + LLM_TENSOR_PER_LAYER_PROJ_NORM, + LLM_TENSOR_PER_LAYER_POST_NORM, + LLM_TENSOR_LAYER_OUT_SCALE, + LLM_TENSOR_FFN_PRE_NORM_2, // 105 + LLM_TENSOR_FFN_POST_NORM_1, + LLM_TENSOR_FFN_POST_NORM_2, + LLM_TENSOR_UNKNOWN, }; diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 6338a391..2b02f6b5 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -48,9 +48,9 @@ llm_build_context::llm_build_context( n_ctx (cparams.n_ctx), n_head (hparams.n_head()), n_head_kv (hparams.n_head_kv()), - n_embd_head_k (hparams.n_embd_head_k), + n_embd_head_k (hparams.n_embd_head_k(0)), n_embd_k_gqa (hparams.n_embd_k_gqa()), - n_embd_head_v (hparams.n_embd_head_v), + n_embd_head_v (hparams.n_embd_head_v(0)), n_embd_v_gqa (hparams.n_embd_v_gqa()), n_expert (hparams.n_expert), n_expert_used (warmup ? hparams.n_expert : hparams.n_expert_used), @@ -560,7 +560,7 @@ void llm_build_context::llm_build_kv_store( const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); const int64_t n_head_kv = hparams.n_head_kv(il); - const int64_t n_embd_head_k = hparams.n_embd_head_k; + const int64_t n_embd_head_k = hparams.n_embd_head_k(il); GGML_ASSERT(kv.size == n_ctx); @@ -568,36 +568,39 @@ void llm_build_context::llm_build_kv_store( // (ggml_row_size(kv.k_l[il]->type, n_embd_k_gqa))*kv_head); //cb(k_cache_view, "k_cache_view", il); - GGML_ASSERT(2*il+1 < (int)lctx.cache_copies.size()); - auto k_row_size = ggml_row_size(kv.k_l[il]->type, n_embd_head_k); - ggml_tensor * k_cache_view = ggml_view_2d(ctx, kv.k_l[il], n_embd_head_k, n_tokens*n_head_kv, - k_row_size, k_row_size*n_head_kv*kv_head); + if (k_cur) { + GGML_ASSERT(2*il+1 < (int)lctx.cache_copies.size()); + auto k_row_size = ggml_row_size(kv.k_l[il]->type, n_embd_head_k); + ggml_tensor * k_cache_view = ggml_view_2d(ctx, kv.k_l[il], n_embd_head_k, n_tokens*n_head_kv, + k_row_size, k_row_size*n_head_kv*kv_head); - lctx.cache_copies[2*il+0].cpy = ggml_cpy(ctx, k_cur, k_cache_view); - lctx.cache_copies[2*il+0].step = k_row_size*n_head_kv; + lctx.cache_copies[2*il+0].cpy = ggml_cpy(ctx, k_cur, k_cache_view); + lctx.cache_copies[2*il+0].step = k_row_size*n_head_kv; - // note: storing RoPE-ed version of K in the KV cache - ggml_build_forward_expand(graph, lctx.cache_copies[2*il+0].cpy); - - struct ggml_tensor * v_cache_view = nullptr; - - if (!kv.v_trans) { - v_cache_view = ggml_view_1d(ctx, kv.v_l[il], n_tokens*n_embd_v_gqa, - (kv_head)*ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa)); - lctx.cache_copies[2*il+1].step = ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa); - } else { - // note: the V cache is transposed for legacy non-FA layouts - v_cache_view = ggml_view_2d(ctx, kv.v_l[il], n_tokens, n_embd_v_gqa, - ( n_ctx)*ggml_element_size(kv.v_l[il]), - (kv_head)*ggml_element_size(kv.v_l[il])); - lctx.cache_copies[2*il+1].step = ggml_element_size(kv.v_l[il]); - - v_cur = ggml_transpose(ctx, v_cur); + // note: storing RoPE-ed version of K in the KV cache + ggml_build_forward_expand(graph, lctx.cache_copies[2*il+0].cpy); } - cb(v_cache_view, "v_cache_view", il); - lctx.cache_copies[2*il+1].cpy = ggml_cpy(ctx, v_cur, v_cache_view); - ggml_build_forward_expand(graph, lctx.cache_copies[2*il+1].cpy); + if (v_cur) { + ggml_tensor * v_cache_view = nullptr; + if (!kv.v_trans) { + v_cache_view = ggml_view_1d(ctx, kv.v_l[il], n_tokens*n_embd_v_gqa, + (kv_head)*ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa)); + lctx.cache_copies[2*il+1].step = ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa); + } else { + // note: the V cache is transposed for legacy non-FA layouts + v_cache_view = ggml_view_2d(ctx, kv.v_l[il], n_tokens, n_embd_v_gqa, + ( n_ctx)*ggml_element_size(kv.v_l[il]), + (kv_head)*ggml_element_size(kv.v_l[il])); + lctx.cache_copies[2*il+1].step = ggml_element_size(kv.v_l[il]); + + v_cur = ggml_transpose(ctx, v_cur); + } + cb(v_cache_view, "v_cache_view", il); + + lctx.cache_copies[2*il+1].cpy = ggml_cpy(ctx, v_cur, v_cache_view); + ggml_build_forward_expand(graph, lctx.cache_copies[2*il+1].cpy); + } } ggml_tensor * llm_build_context::llm_build_lora_mm( @@ -1016,7 +1019,10 @@ ggml_tensor * llm_build_context::llm_build_moe_ffn( float w_scale, llm_expert_gating_func_type gating_op, const llm_build_cb & cb, int il, ggml_cgraph * graph, bool add_input, - ggml_tensor * up_gate_exps, ggml_tensor * up_gate_exps_b) { + ggml_tensor * up_gate_exps, ggml_tensor * up_gate_exps_b, + ggml_tensor * input_logits, ggml_tensor * down_exps_s) { + + GGML_ASSERT(gate_inp || input_logits); auto input = cur; @@ -1024,7 +1030,7 @@ llm_expert_gating_func_type gating_op, int64_t n_tokens = cur->ne[1]; bool weight_before_ffn = lctx.model.arch == LLM_ARCH_LLAMA4; // for llama4, we apply the sigmoid-ed weights before the FFN - ggml_tensor * logits = llm_build_lora_mm(lctx, ctx, gate_inp, cur); // [n_expert, n_tokens] + ggml_tensor * logits = gate_inp ? llm_build_lora_mm(lctx, ctx, gate_inp, cur) : input_logits; // [n_expert, n_tokens] cb(logits, "ffn_moe_logits", il); if (gate_inp_b) { @@ -1209,6 +1215,14 @@ llm_expert_gating_func_type gating_op, cb(experts, "ffn_moe_down_biased", il); } + if (down_exps_s) { + ggml_tensor * s = ggml_reshape_3d(ctx, down_exps_s, 1, n_expert, 1); + s = ggml_repeat_4d(ctx, s, 1, n_expert, n_tokens, 1); + s = ggml_get_rows(ctx, s, selected_experts); // [1, n_expert_used, n_tokens] + experts = ggml_mul(ctx, experts, s); + cb(experts, "ffn_moe_down_scaled", il); + } + if (!weight_before_ffn) { if (lctx.cparams.fused_mmad) { experts = ggml_mul_multi_add(ctx, experts, weights); @@ -1517,19 +1531,24 @@ static ggml_tensor * llm_build_kqv( const int64_t n_ctx = cparams.n_ctx; const int64_t n_head = hparams.n_head(il); const int64_t n_head_kv = hparams.n_head_kv(il); - const int64_t n_embd_head_k = hparams.n_embd_head_k; + const int64_t n_embd_head_k = hparams.n_embd_head_k(il); //const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il); - const int64_t n_embd_head_v = hparams.n_embd_head_v; + const int64_t n_embd_head_v = hparams.n_embd_head_v(il); const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); struct ggml_tensor * q = ggml_permute(ctx, q_cur, 0, 2, 1, 3); cb(q, "q", il); + auto k_cache = lctx.model.hparams.has_kv(il) ? kv.k_l[il] + : lctx.model.hparams.swa_layers[il] ? kv.k_l[hparams.n_layer_kv_from_start-2] : kv.k_l[hparams.n_layer_kv_from_start-1]; + auto v_cache = lctx.model.hparams.has_kv(il) ? kv.v_l[il] + : lctx.model.hparams.swa_layers[il] ? kv.v_l[hparams.n_layer_kv_from_start-2] : kv.v_l[hparams.n_layer_kv_from_start-1]; + struct ggml_tensor * k = - ggml_view_3d(ctx, kv.k_l[il], + ggml_view_3d(ctx, k_cache, n_embd_head_k, n_kv, n_head_kv, - ggml_row_size(kv.k_l[il]->type, n_embd_head_k)*n_head_kv, //n_embd_k_gqa), - ggml_row_size(kv.k_l[il]->type, n_embd_head_k), + ggml_row_size(k_cache->type, n_embd_head_k)*n_head_kv, //n_embd_k_gqa), + ggml_row_size(k_cache->type, n_embd_head_k), 0); cb(k, "k", il); @@ -1559,10 +1578,10 @@ static ggml_tensor * llm_build_kqv( // split cached v into n_head heads (not transposed) struct ggml_tensor * v = - ggml_view_3d(ctx, kv.v_l[il], + ggml_view_3d(ctx, v_cache, n_embd_head_v, n_kv, n_head_kv, - ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa), - ggml_row_size(kv.v_l[il]->type, n_embd_head_v), + ggml_row_size(v_cache->type, n_embd_v_gqa), + ggml_row_size(v_cache->type, n_embd_head_v), 0); cb(v, "v", il); @@ -1592,16 +1611,16 @@ static ggml_tensor * llm_build_kqv( // split cached v into n_head heads struct ggml_tensor * v; if (kv.v_trans) { - v = ggml_view_3d(ctx, kv.v_l[il], + v = ggml_view_3d(ctx, v_cache, n_kv, n_embd_head_v, n_head_kv, - ggml_element_size(kv.v_l[il])*n_ctx, - ggml_element_size(kv.v_l[il])*n_ctx*n_embd_head_v, + ggml_element_size(v_cache)*n_ctx, + ggml_element_size(v_cache)*n_ctx*n_embd_head_v, 0); } else { - v = ggml_view_3d(ctx, kv.v_l[il], + v = ggml_view_3d(ctx, v_cache, n_embd_head_v, n_kv, n_head_kv, - ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa), - ggml_row_size(kv.v_l[il]->type, n_embd_head_v), + ggml_row_size(v_cache->type, n_embd_v_gqa), + ggml_row_size(v_cache->type, n_embd_head_v), 0); v = ggml_cont(ctx, ggml_transpose(ctx, v)); } @@ -1744,27 +1763,32 @@ ggml_tensor * llm_build_context::llm_build_kv( const llama_cparams & cparams = lctx.cparams; if (cparams.k_cache_hadamard) { - q_cur = ggml_hadamard(ctx, q_cur, hparams.n_embd_head_k); - k_cur = ggml_hadamard(ctx, k_cur, hparams.n_embd_head_k); + q_cur = ggml_hadamard(ctx, q_cur, hparams.n_embd_head_k(il)); + if (k_cur) { + k_cur = ggml_hadamard(ctx, k_cur, hparams.n_embd_head_k(il)); + cb(k_cur, "Kcur_hadamard", il); + } cb(q_cur, "Qcur_hadamard", il); - cb(k_cur, "Kcur_hadamard", il); } - if (cparams.v_cache_hadamard) { - v_cur = ggml_hadamard(ctx, v_cur, hparams.n_embd_head_v); + if (cparams.v_cache_hadamard && v_cur) { + v_cur = ggml_hadamard(ctx, v_cur, hparams.n_embd_head_v(il)); } // these nodes are added to the graph together so that they are not reordered // by doing so, the number of splits in the graph is reduced ggml_build_forward_expand(graph, q_cur); - ggml_build_forward_expand(graph, k_cur); - ggml_build_forward_expand(graph, v_cur); + if (k_cur) { + ggml_build_forward_expand(graph, k_cur); + } + if (v_cur) { + ggml_build_forward_expand(graph, v_cur); + } - llm_build_kv_store(lctx, ctx, hparams, cparams, kv, graph, k_cur, v_cur, n_tokens, kv_head, cb, il); + if (k_cur || v_cur) { + llm_build_kv_store(lctx, ctx, hparams, cparams, kv, graph, k_cur, v_cur, n_tokens, kv_head, cb, il); + } - struct ggml_tensor * cur; - - cur = llm_build_kqv(ctx, lctx, kv, graph, wo, wo_b, - q_cur, kq_mask, n_tokens, n_kv, kq_scale, cb, il, sinks, n_swa); + auto cur = llm_build_kqv(ctx, lctx, kv, graph, wo, wo_b, q_cur, kq_mask, n_tokens, n_kv, kq_scale, cb, il, sinks, n_swa); cb(cur, "kqv_out", il); return cur; @@ -1868,7 +1892,7 @@ std::tuple llm_build_context::llm_buil ggml_tensor * q_norm, ggml_tensor * k_norm, float attention_scale, int il, bool add_graph_split) const { int n_head = hparams.n_head(il); int n_head_kv = hparams.n_head_kv(il); - const int64_t n_embd_head_k = hparams.n_embd_head_k; + const int64_t n_embd_head_k = hparams.n_embd_head_k(il); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(il); if (wqkv) { auto qkv = llm_build_lora_mm(lctx, ctx0, wqkv, cur); @@ -2036,8 +2060,8 @@ ggml_cgraph * llm_build_context::build_llama() { // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); ggml_tensor * cur; @@ -2248,9 +2272,9 @@ ggml_cgraph * llm_build_context::build_llama() { ggml_cgraph * llm_build_context::build_mistral3() { auto gf = ggml_new_graph_custom(ctx0, model.max_nodes(n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); ggml_tensor * cur; @@ -2343,8 +2367,8 @@ ggml_cgraph * llm_build_context::build_deci() { // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -2477,8 +2501,8 @@ ggml_cgraph * llm_build_context::build_deci() { ggml_cgraph * llm_build_context::build_baichuan() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -2579,8 +2603,8 @@ static inline size_t llama_model_max_nodes(const llama_model & model, int n_toke ggml_cgraph * llm_build_context::build_xverse() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -2669,9 +2693,9 @@ ggml_cgraph * llm_build_context::build_xverse() { ggml_cgraph * llm_build_context::build_falcon() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -2783,8 +2807,8 @@ ggml_cgraph * llm_build_context::build_grok() { // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -2923,9 +2947,9 @@ ggml_cgraph * llm_build_context::build_dbrx() { // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -3042,9 +3066,9 @@ ggml_cgraph * llm_build_context::build_dbrx() { ggml_cgraph * llm_build_context::build_starcoder() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -3134,8 +3158,8 @@ ggml_cgraph * llm_build_context::build_starcoder() { ggml_cgraph * llm_build_context::build_refact() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -3214,10 +3238,10 @@ ggml_cgraph * llm_build_context::build_refact() { ggml_cgraph * llm_build_context::build_bert() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -3402,9 +3426,9 @@ ggml_cgraph * llm_build_context::build_bert() { ggml_cgraph * llm_build_context::build_bloom() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -3488,9 +3512,9 @@ ggml_cgraph * llm_build_context::build_bloom() { ggml_cgraph * llm_build_context::build_mpt() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * pos; @@ -3611,8 +3635,8 @@ ggml_cgraph * llm_build_context::build_mpt() { ggml_cgraph * llm_build_context::build_stablelm() { struct ggml_cgraph * gf = ggml_new_graph(ctx0); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -3728,8 +3752,8 @@ ggml_cgraph * llm_build_context::build_stablelm() { ggml_cgraph * llm_build_context::build_seedoss() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); ggml_tensor * cur; @@ -3849,8 +3873,8 @@ ggml_cgraph * llm_build_context::build_step35() { ggml_cgraph * llm_build_context::build_qwen() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -3952,8 +3976,8 @@ ggml_cgraph * llm_build_context::build_qwen() { ggml_cgraph * llm_build_context::build_qwen2() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -4043,9 +4067,9 @@ ggml_cgraph * llm_build_context::build_qwen2() { ggml_cgraph * llm_build_context::build_qwen2vl() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); ggml_tensor * cur; @@ -4147,8 +4171,8 @@ ggml_cgraph * llm_build_context::build_qwen2moe() { // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -4274,8 +4298,8 @@ ggml_cgraph * llm_build_context::build_qwen2moe() { ggml_cgraph * llm_build_context::build_qwen3() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -4375,8 +4399,8 @@ ggml_cgraph * llm_build_context::build_qwen3() { ggml_cgraph * llm_build_context::build_qwen3moe() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -4441,8 +4465,8 @@ ggml_cgraph * llm_build_context::build_qwen3next() { delta_net delta(lctx, batch); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); ggml_tensor * inp_pos = build_inp_pos(); @@ -4526,8 +4550,8 @@ ggml_cgraph * llm_build_context::build_qwen35moe() { delta_net delta(lctx, batch); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); ggml_tensor * inp_pos = build_inp_pos(); @@ -4585,8 +4609,8 @@ ggml_cgraph * llm_build_context::build_qwen35() { delta_net delta(lctx, batch); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); ggml_tensor * inp_pos = build_inp_pos(); @@ -4637,9 +4661,9 @@ ggml_cgraph * llm_build_context::build_qwen3vl() { const int64_t n_embd_full = hparams.n_embd; // main embd + deepstack embds const size_t n_deepstack_layers = hparams.n_deepstack_layers; const int64_t n_embd = n_embd_full / (n_deepstack_layers + 1); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -4709,9 +4733,9 @@ ggml_cgraph * llm_build_context::build_qwen3vlmoe() { const int64_t n_embd_full = hparams.n_embd; // main embd + deepstack embds const size_t n_deepstack_layers = hparams.n_deepstack_layers; const int64_t n_embd = n_embd_full / (n_deepstack_layers + 1); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -4849,9 +4873,9 @@ ggml_cgraph * llm_build_context::build_qwen3vlmoe() { ggml_cgraph * llm_build_context::build_phi2() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * attn_norm_output; @@ -4964,9 +4988,9 @@ ggml_cgraph * llm_build_context::build_phi2() { ggml_cgraph * llm_build_context::build_phi3() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -5086,8 +5110,8 @@ ggml_cgraph * llm_build_context::build_phi3() { ggml_cgraph * llm_build_context::build_plamo() { struct ggml_cgraph * gf = ggml_new_graph(ctx0); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -5179,9 +5203,9 @@ ggml_cgraph * llm_build_context::build_plamo() { ggml_cgraph * llm_build_context::build_gpt2() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * pos; @@ -5272,9 +5296,9 @@ ggml_cgraph * llm_build_context::build_gpt2() { ggml_cgraph * llm_build_context::build_codeshell() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -5371,8 +5395,8 @@ ggml_cgraph * llm_build_context::build_codeshell() { ggml_cgraph * llm_build_context::build_orion() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -5461,8 +5485,8 @@ ggml_cgraph * llm_build_context::build_orion() { ggml_cgraph * llm_build_context::build_internlm2() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -5554,8 +5578,8 @@ ggml_cgraph * llm_build_context::build_internlm2() { ggml_cgraph * llm_build_context::build_minicpm() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); const int64_t n_embd = hparams.n_embd; @@ -5671,7 +5695,7 @@ ggml_cgraph * llm_build_context::build_minicpm() { ggml_cgraph * llm_build_context::build_gemma() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head_k = hparams.n_embd_head_k; + const int64_t n_embd_head_k = hparams.n_embd_head_k(0); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -5763,7 +5787,7 @@ ggml_cgraph * llm_build_context::build_gemma() { ggml_cgraph * llm_build_context::build_gemma2() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head_k = hparams.n_embd_head_k; + const int64_t n_embd_head_k = hparams.n_embd_head_k_full; struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -5997,11 +6021,277 @@ ggml_cgraph * llm_build_context::build_gemma3() { return gf; } +static ggml_tensor * gemma4_project_per_layer_inputs(ggml_context * ctx0, const llama_model & model, const llm_build_cb & cb, + int n_embd, int n_embd_per_layer, int n_layer, int n_tokens, + ggml_tensor * inputs_embeds, ggml_tensor * inp_per_layer) { + const float per_layer_projection_scale = 1.0f / sqrtf((float) n_embd); + const float per_layer_input_scale = 1.0f / sqrtf(2.0f); + + ggml_tensor * per_layer_proj = ggml_mul_mat(ctx0, model.per_layer_model_proj, inputs_embeds); + per_layer_proj = ggml_scale(ctx0, per_layer_proj, per_layer_projection_scale); + per_layer_proj = ggml_reshape_3d(ctx0, per_layer_proj, n_embd_per_layer, n_layer, n_tokens); + per_layer_proj = llm_build_context::llm_build_norm(ctx0, per_layer_proj, model.hparams, + model.per_layer_proj_norm, nullptr, LLM_NORM_RMS, cb, -1); // [n_embd_per_layer, n_layer, n_tokens] + cb(per_layer_proj, "per_layer_proj", -1); + + inp_per_layer = ggml_add(ctx0, per_layer_proj, inp_per_layer); + inp_per_layer = ggml_scale(ctx0, inp_per_layer, per_layer_input_scale); + cb(inp_per_layer, "inp_per_layer", -1); + + // permute to shape: [n_embd_per_layer, n_tokens, n_layer] + inp_per_layer = ggml_cont(ctx0, ggml_permute(ctx0, inp_per_layer, 0, 2, 1, 3)); + return inp_per_layer; +} + +ggml_cgraph * llm_build_context::build_gemma4() { + struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); + + struct ggml_tensor * cur; + struct ggml_tensor * inpL; + + inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); + + // important: do not normalize weights for raw embeddings input (i.e. encoded image emdeddings) + if (batch.token) { + inpL = ggml_scale(ctx0, inpL, sqrtf(n_embd)); + cb(inpL, "inp_scaled", -1); + } + + // inp_pos - contains the positions + struct ggml_tensor * inp_pos = build_inp_pos(); + + // KQ_mask (mask for 1 head, it will be broadcasted to all heads) + // gemma3 requires different mask for layers using sliding window (SWA) + struct ggml_tensor * KQ_mask = build_inp_KQ_mask(true); + struct ggml_tensor * KQ_mask_swa = build_inp_KQ_mask_swa(true); + + auto inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; + + ggml_tensor * inp_per_layer = nullptr; + if (model.tok_embd_per_layer) { + if (batch.token) { + inp_per_layer = ggml_get_rows(ctx0, model.tok_embd_per_layer, lctx.inp_tokens); + inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, hparams.n_embd_per_layer, n_layer, n_tokens); + inp_per_layer = ggml_scale(ctx0, inp_per_layer, sqrtf((float) hparams.n_embd_per_layer)); + cb(inp_per_layer, "inp_per_layer_selected", -1); + } else { + // Vision embedding path: use padding token (ID=0) embedding + // TODO: verify if this is the correct behavior in transformers implementation + const int64_t embd_size = model.tok_embd_per_layer->ne[0]; // n_embd_per_layer * n_layer + + // Extract and dequantize padding token embedding (row 0) + auto padding = ggml_view_1d(ctx0, model.tok_embd_per_layer, embd_size, 0); + inp_per_layer = ggml_cast(ctx0, padding, GGML_TYPE_F32); + + // Reshape to [n_embd_per_layer, n_layer, 1] + inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, hparams.n_embd_per_layer, n_layer, 1); + cb(inp_per_layer, "inp_per_layer_vision", -1); + } + inp_per_layer = gemma4_project_per_layer_inputs(ctx0, model, cb, n_embd, + model.hparams.n_embd_per_layer, n_layer, n_tokens, inpL, inp_per_layer); + + } + + // "5-to-1 interleaved attention" + // 5 layers of local attention followed by 1 layer of global attention + // we do this via swa_layers now + // static const int sliding_window_pattern = 6; + + for (int il = 0; il < n_layer; ++il) { + const bool is_sliding = hparams.swa_layers[il] ? true : false; + const float freq_base_l = is_sliding ? hparams.rope_freq_base_train_swa : cparams.rope_freq_base; + const float freq_scale_l = is_sliding ? hparams.rope_freq_scale_train_swa : cparams.rope_freq_scale; + const int n_rot_l = is_sliding ? hparams.n_rot_swa : hparams.n_rot; + const int n_swa = is_sliding ? hparams.n_swa : 0; + const int n_embd_head = hparams.n_embd_head_k(il); + const int n_head = hparams.n_head(il); + const int n_head_kv = hparams.n_head_kv(il); + + struct ggml_tensor * KQ_mask_l = is_sliding ? KQ_mask_swa : KQ_mask; + + // norm + cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il); + cb(cur, "attn_norm", il); + + // self-attention + { + auto freq_factors = !is_sliding ? model.layers[il].rope_freqs : nullptr; + + ggml_tensor *Qcur, *Kcur = nullptr, *Vcur = nullptr; + Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wq, cur); + cb(Qcur, "Qcur", il); + if (hparams.has_kv(il)) { + Kcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, cur); + cb(Kcur, "Kcur", il); + if (model.layers[il].wv) { + Vcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, cur); + cb(Vcur, "Vcur", il); + } else { + Vcur = Kcur; + } + } + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); + Qcur = llm_build_norm(ctx0, Qcur, hparams, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, cb, il); + if (hparams.has_kv(il)) { + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + Kcur = llm_build_norm(ctx0, Kcur, hparams, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, cb, il); + Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); + cb(Kcur, "Kcur_normed", il); + cb(Vcur, "Vcur_normed", il); + } + 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); + if (hparams.has_kv(il)) { + Kcur = ggml_rope_ext(ctx0, Kcur, 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(Kcur, "Kcur_rope", il); + } + cur = llm_build_kv(ctx0, lctx, kv_self, gf, model.layers[il].wo, model.layers[il].bo, + Kcur, Vcur, Qcur, KQ_mask_l, n_tokens, kv_head, n_kv, hparams.f_attention_scale, cb, il, nullptr, n_swa); + + } + + if (il == n_layer - 1 && inp_out_ids) { + // skip computing output for unused tokens + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpL = ggml_get_rows(ctx0, inpL, inp_out_ids); + } + + cur = llm_build_norm(ctx0, cur, hparams, model.layers[il].attn_post_norm, NULL, LLM_NORM_RMS, cb, il); + cb(cur, "attn_post_norm", il); + + auto attn_out = ggml_add(ctx0, cur, inpL); + cb(attn_out, "attn_out", il); + + if (model.layers[il].ffn_gate_inp) { + + auto cur_mlp = llm_build_norm(ctx0, attn_out, hparams, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, cb, il); + cb(cur_mlp, "ffn_norm_1", il); + + cur_mlp = llm_build_ffn(ctx0, lctx, nullptr, cur_mlp, + model.layers[il].ffn_up, nullptr, nullptr, + model.layers[il].ffn_gate, nullptr, nullptr, + model.layers[il].ffn_down, nullptr, nullptr, + nullptr, + LLM_FFN_GELU, LLM_FFN_PAR, cb, il, gf); + + cur_mlp = llm_build_norm(ctx0, cur_mlp, hparams, model.layers[il].ffn_post_norm_1, nullptr, LLM_NORM_RMS, cb, il); + cb(cur_mlp, "ffn_mlp", il); + + // Expert FFN + auto cur_moe = llm_build_norm(ctx0, attn_out, hparams, model.layers[il].ffn_pre_norm_2, nullptr, LLM_NORM_RMS, cb, il); + cb(cur_moe, "ffn_norm_2", il); + + // custom MoE logits calculation (router operates on attn_out, not cur) + auto tmp = ggml_rms_norm(ctx0, attn_out, hparams.f_norm_rms_eps); + tmp = ggml_scale(ctx0, tmp, 1.0f / sqrtf((float) n_embd)); + tmp = ggml_mul(ctx0, tmp, model.layers[il].ffn_gate_inp_s); + auto logits = llm_build_lora_mm(lctx, ctx0, model.layers[il].ffn_gate_inp, tmp); // [n_expert, n_tokens] + cb(logits, "ffn_moe_logits", il); + + // TODO: we need to pass in the above logits + cur_moe = llm_build_moe_ffn(ctx0, lctx, cur_moe, + nullptr, // gate_inp + nullptr, // up_exps + nullptr, // gate_exps + model.layers[il].ffn_down_exps, + nullptr, // exp_probs_b (not used for gemma4) + n_expert, n_expert_used, + LLM_FFN_GELU, true, false, 0.0f, + LLM_EXPERT_GATING_FUNC_SOFTMAX, + cb, il, gf, + false, + model.layers[il].ffn_up_gate_exps, + nullptr, logits, model.layers[il].ffn_down_exps_s); + + cur_moe = llm_build_norm(ctx0, cur_moe, hparams, model.layers[il].ffn_post_norm_2, nullptr, LLM_NORM_RMS, cb, il); + cb(cur_moe, "ffn_moe", il); + + cur = ggml_add(ctx0, cur_mlp, cur_moe); + cb(cur, "ffn_moe_combined", il); + + } else { + cur = llm_build_norm(ctx0, attn_out, hparams, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, cb, il); + cb(cur, "ffn_norm", il); + + cur = llm_build_ffn(ctx0, lctx, nullptr, cur, + model.layers[il].ffn_up, nullptr, nullptr, + model.layers[il].ffn_gate, nullptr, nullptr, + model.layers[il].ffn_down, nullptr, nullptr, + nullptr, + LLM_FFN_GELU, LLM_FFN_PAR, cb, il, gf); + cb(cur, "ffn_out", il); + + } + + cur = llm_build_norm(ctx0, cur, hparams, model.layers[il].ffn_post_norm, NULL, LLM_NORM_RMS, cb, -1); + cb(cur, "ffn_post_norm", -1); + + cur = ggml_add(ctx0, cur, attn_out); + + if (inp_per_layer) { + ggml_tensor * pe_in = cur; + cb(cur, "pe_in", il); + + cur = llm_build_lora_mm(lctx, ctx0, model.layers[il].per_layer_inp_gate, cur); // [n_embd_per_layer, n_tokens] + cur = ggml_gelu(ctx0, cur); + ggml_tensor * inp_this_layer = ggml_view_2d(ctx0, inp_per_layer, inp_per_layer->ne[0], inp_per_layer->ne[1], + ggml_row_size(inp_per_layer->type, inp_per_layer->ne[0]), + il*inp_per_layer->ne[0]*inp_per_layer->ne[1]*ggml_element_size(inp_per_layer)); // [n_embd_per_layer, n_tokens] + + // TODO @ngxson : improve this + if (il == n_layer - 1 && inp_out_ids) { + inp_this_layer = ggml_get_rows(ctx0, inp_this_layer, inp_out_ids); + } + + cur = ggml_mul(ctx0, cur, inp_this_layer); + cur = llm_build_lora_mm(lctx, ctx0, model.layers[il].per_layer_proj, cur); // [n_embd, n_tokens] + cur = llm_build_norm(ctx0, cur, hparams, model.layers[il].per_layer_post_norm, nullptr, LLM_NORM_RMS, cb, il); + cb(cur, "per_layer_embd_out", il); + + // residual connection + cur = ggml_add(ctx0, pe_in, cur); + } + + // layer_scalar + if (model.layers[il].out_scale) { + cur = ggml_mul(ctx0, cur, model.layers[il].out_scale); + cb(cur, "out_scaled", il); + } + + cur = lctx.cvec.apply_to(ctx0, cur, il); + cb(cur, "l_out", il); + + // input for next layer + inpL = cur; + } + + cur = inpL; + + cur = llm_build_norm(ctx0, cur, hparams, model.output_norm, NULL, LLM_NORM_RMS, cb, -1); + cb(cur, "result_norm", -1); + + // lm_head + cur = llm_build_lora_mm(lctx, ctx0, model.output, cur); + + if (hparams.f_final_logit_softcapping > 0) { + cur = ggml_softcap(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping, hparams.f_final_logit_softcapping); + } + + cb(cur, "result_output", -1); + + ggml_build_forward_expand(gf, cur); + + return gf; +} + ggml_cgraph * llm_build_context::build_starcoder2() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -6236,8 +6526,8 @@ ggml_cgraph * llm_build_context::build_command_r() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); const float f_logit_scale = hparams.f_logit_scale; struct ggml_tensor * cur; @@ -6318,8 +6608,8 @@ ggml_cgraph * llm_build_context::build_olmo() { // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -6433,8 +6723,8 @@ ggml_cgraph * llm_build_context::build_olmo() { ggml_cgraph * llm_build_context::build_openelm() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -6545,9 +6835,9 @@ ggml_cgraph * llm_build_context::build_openelm() { ggml_cgraph * llm_build_context::build_gptneox() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -6678,8 +6968,8 @@ ggml_cgraph * llm_build_context::build_arctic() { // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -6805,11 +7095,11 @@ ggml_cgraph * llm_build_context::build_deepseek2() { // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly. // See https://github.com/ggerganov/llama.cpp/discussions/7416 for detailed explanation. const float mscale = attn_factor * (1.0f + hparams.rope_yarn_log_mul * logf(1.0f / freq_scale)); - const float kq_scale = 1.0f*mscale*mscale/sqrtf(float(hparams.n_embd_head_k)); + const float kq_scale = 1.0f*mscale*mscale/sqrtf(float(hparams.n_embd_head_k(0))); const float attn_factor_scaled = 1.0f / (1.0f + 0.1f * logf(1.0f / freq_scale)); const uint32_t n_embd_head_qk_rope = hparams.n_rot; - const uint32_t n_embd_head_qk_nope = hparams.n_embd_head_k - hparams.n_rot; + const uint32_t n_embd_head_qk_nope = hparams.n_embd_head_k(0) - hparams.n_rot; const uint32_t kv_lora_rank = hparams.n_lora_kv; const uint32_t q_lora_rank = hparams.n_lora_q; @@ -6868,9 +7158,9 @@ ggml_cgraph * llm_build_context::build_deepseek2() { qnb1 = mqa->nb[1]; } q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, - ggml_row_size(q->type, hparams.n_embd_head_k), qnb1, 0); + ggml_row_size(q->type, hparams.n_embd_head_k(il)), qnb1, 0); q_rope = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, - ggml_row_size(q->type, hparams.n_embd_head_k), qnb1, ggml_row_size(q->type, n_embd_head_qk_nope)); + ggml_row_size(q->type, hparams.n_embd_head_k(il)), qnb1, ggml_row_size(q->type, n_embd_head_qk_nope)); k_rope = ggml_view_3d(ctx0, kv_rope_compressed, n_embd_head_qk_rope, 1, n_tokens, mqa->nb[1], mqa->nb[1], ggml_row_size(kv_rope_compressed->type, kv_lora_rank)); kv_compressed = ggml_view_2d(ctx0, kv_rope_compressed, kv_lora_rank, n_tokens, mqa->nb[1], 0); @@ -6903,12 +7193,12 @@ ggml_cgraph * llm_build_context::build_deepseek2() { } q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, - ggml_row_size(q->type, hparams.n_embd_head_k), - ggml_row_size(q->type, hparams.n_embd_head_k * n_head), 0); + ggml_row_size(q->type, hparams.n_embd_head_k(il)), + ggml_row_size(q->type, hparams.n_embd_head_k(il) * n_head), 0); q_rope = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, - ggml_row_size(q->type, hparams.n_embd_head_k), - ggml_row_size(q->type, hparams.n_embd_head_k * n_head), + ggml_row_size(q->type, hparams.n_embd_head_k(il)), + ggml_row_size(q->type, hparams.n_embd_head_k(il) * n_head), ggml_row_size(q->type, n_embd_head_qk_nope)); k_rope = ggml_view_3d(ctx0, kv_rope_compressed, n_embd_head_qk_rope, 1, n_tokens, @@ -7036,15 +7326,15 @@ ggml_cgraph * llm_build_context::build_deepseek2() { auto kv_f32 = ggml_mul_mat(ctx0, wkv_b, kv_cache_nope); cb(kv_f32, "kv_f32", il); - auto v_f32 = ggml_view_3d(ctx0, kv_f32, hparams.n_embd_head_v, n_kv, n_max_head, - ggml_row_size(kv_f32->type, n_max_head * (n_embd_head_qk_nope + hparams.n_embd_head_v)), - ggml_row_size(kv_f32->type, n_embd_head_qk_nope + hparams.n_embd_head_v), + auto v_f32 = ggml_view_3d(ctx0, kv_f32, hparams.n_embd_head_v_full, n_kv, n_max_head, + ggml_row_size(kv_f32->type, n_max_head * (n_embd_head_qk_nope + hparams.n_embd_head_v_full)), + ggml_row_size(kv_f32->type, n_embd_head_qk_nope + hparams.n_embd_head_v_full), ggml_row_size(kv_f32->type, n_embd_head_qk_nope)); cb(v_f32, "v_f32", il); auto k_nope_f32 = ggml_view_3d(ctx0, kv_f32, n_embd_head_qk_nope, n_kv, n_max_head, - ggml_row_size(kv_f32->type, n_max_head * (n_embd_head_qk_nope + hparams.n_embd_head_v)), - ggml_row_size(kv_f32->type, n_embd_head_qk_nope + hparams.n_embd_head_v), 0); + ggml_row_size(kv_f32->type, n_max_head * (n_embd_head_qk_nope + hparams.n_embd_head_v_full)), + ggml_row_size(kv_f32->type, n_embd_head_qk_nope + hparams.n_embd_head_v_full), 0); cb(k_nope_f32, "k_nope_f32", il); auto v = ggml_cast(ctx0, v_f32, kv_type); @@ -7224,23 +7514,23 @@ ggml_cgraph * llm_build_context::build_deepseek2() { // split into {n_head * n_embd_head_qk_nope, n_tokens} struct ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens, - ggml_row_size(kv->type, n_embd_head_qk_nope + hparams.n_embd_head_v), - ggml_row_size(kv->type, n_head * (n_embd_head_qk_nope + hparams.n_embd_head_v)), + ggml_row_size(kv->type, n_embd_head_qk_nope + hparams.n_embd_head_v_full), + ggml_row_size(kv->type, n_head * (n_embd_head_qk_nope + hparams.n_embd_head_v_full)), 0); cb(k_nope, "k_nope", il); // and {n_head * n_embd_head_v, n_tokens} - struct ggml_tensor * v_states = ggml_view_3d(ctx0, kv, hparams.n_embd_head_v, n_head, n_tokens, - ggml_row_size(kv->type, (n_embd_head_qk_nope + hparams.n_embd_head_v)), - ggml_row_size(kv->type, (n_embd_head_qk_nope + hparams.n_embd_head_v)*n_head), + struct ggml_tensor * v_states = ggml_view_3d(ctx0, kv, hparams.n_embd_head_v_full, n_head, n_tokens, + ggml_row_size(kv->type, (n_embd_head_qk_nope + hparams.n_embd_head_v_full)), + ggml_row_size(kv->type, (n_embd_head_qk_nope + hparams.n_embd_head_v_full)*n_head), ggml_row_size(kv->type, (n_embd_head_qk_nope))); cb(v_states, "v_states", il); v_states = ggml_cont(ctx0, v_states); cb(v_states, "v_states", il); - v_states = ggml_view_2d(ctx0, v_states, hparams.n_embd_head_v * n_head, n_tokens, - ggml_row_size(kv->type, hparams.n_embd_head_v * n_head), + v_states = ggml_view_2d(ctx0, v_states, hparams.n_embd_head_v_full * n_head, n_tokens, + ggml_row_size(kv->type, hparams.n_embd_head_v_full * n_head), 0); cb(v_states, "v_states", il); @@ -7338,8 +7628,8 @@ ggml_cgraph * llm_build_context::build_glm4_moe() { // create a new graph struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); ggml_tensor * cur; @@ -7594,8 +7884,8 @@ struct ggml_tensor * llm_build_context::build_mtp_tail( ggml_cgraph * llm_build_context::build_bitnet() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -7742,8 +8032,8 @@ ggml_cgraph * llm_build_context::build_bitnet_158() { // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -7857,8 +8147,8 @@ ggml_cgraph * llm_build_context::build_bitnet_158() { ggml_cgraph * llm_build_context::build_cohere2() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); const float f_logit_scale = hparams.f_logit_scale; struct ggml_tensor * cur; @@ -7937,9 +8227,9 @@ ggml_cgraph * llm_build_context::build_t5_encoder() { // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -8055,9 +8345,9 @@ ggml_cgraph * llm_build_context::build_t5_decoder() { // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -8241,9 +8531,9 @@ ggml_cgraph * llm_build_context::build_t5_decoder() { ggml_cgraph * llm_build_context::build_jais() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -8321,9 +8611,9 @@ ggml_cgraph * llm_build_context::build_jais() { ggml_cgraph * llm_build_context::build_chatglm() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -8423,10 +8713,10 @@ ggml_cgraph * llm_build_context::build_chatglm() { ggml_cgraph * llm_build_context::build_glm4() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); struct ggml_tensor * cur; struct ggml_tensor * inpL; @@ -8553,9 +8843,9 @@ ggml_cgraph * llm_build_context::build_glm4() { ggml_cgraph * llm_build_context::build_dots1() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); ggml_tensor * cur; @@ -8693,9 +8983,9 @@ ggml_cgraph * llm_build_context::build_dots1() { ggml_cgraph * llm_build_context::build_ernie4_5() { struct ggml_cgraph* gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); ggml_tensor * cur; @@ -8809,9 +9099,9 @@ ggml_cgraph * llm_build_context::build_ernie4_5() { ggml_cgraph * llm_build_context::build_ernie4_5_moe() { struct ggml_cgraph* gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); ggml_tensor * cur; @@ -8879,9 +9169,9 @@ ggml_cgraph * llm_build_context::build_ernie4_5_moe() { ggml_cgraph * llm_build_context::build_hunyuan_moe() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_rot); ggml_tensor * cur; @@ -8936,8 +9226,8 @@ ggml_cgraph * llm_build_context::build_hunyuan_moe() { ggml_cgraph * llm_build_context::build_mimo2() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - //const int64_t n_embd_head = hparams.n_embd_head_v; - //GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + //const int64_t n_embd_head = hparams.n_embd_head_v(0); + //GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); //GGML_ASSERT(n_embd_head == hparams.n_rot); struct ggml_tensor * cur; @@ -9008,8 +9298,8 @@ ggml_cgraph * llm_build_context::build_mimo2() { ggml_cgraph * llm_build_context::build_openai_moe() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); ggml_tensor * cur; ggml_tensor * inpL; @@ -9069,10 +9359,10 @@ ggml_cgraph * llm_build_context::build_openai_moe() { ggml_cgraph * llm_build_context::build_bailingmoe2() { ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; + const int64_t n_embd_head = hparams.n_embd_head_v(0); //const int64_t n_embd_gqa = hparams.n_embd_v_gqa(); - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); ggml_tensor * cur; ggml_tensor * inpL; @@ -9198,8 +9488,8 @@ ggml_cgraph * llm_build_context::build_bailingmoe2() { ggml_cgraph* llm_build_context::build_minimaxm2() { ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); // GGML_ASSERT(n_embd_head == hparams.n_rot); this is wrong in case of minimax, head_dim = 128, n_rot = 64 ggml_tensor * cur; @@ -9297,15 +9587,15 @@ ggml_cgraph* llm_build_context::build_minimaxm2() { cb(Kcur, "Kcur_roped", il_id); if (cparams.k_cache_hadamard) { - Qcur = ggml_hadamard(ctx0, Qcur, hparams.n_embd_head_k); - Kcur = ggml_hadamard(ctx0, Kcur, hparams.n_embd_head_k); + Qcur = ggml_hadamard(ctx0, Qcur, n_embd_head_k); + Kcur = ggml_hadamard(ctx0, Kcur, n_embd_head_k); cb(Qcur, "Qcur_hadamard", il_id); cb(Kcur, "Kcur_hadamard", il_id); } ggml_build_forward_expand(gf, Qcur); ggml_build_forward_expand(gf, Kcur); if (cparams.v_cache_hadamard) { - Vcur = ggml_hadamard(ctx0, Vcur, hparams.n_embd_head_v); + Vcur = ggml_hadamard(ctx0, Vcur, n_embd_head_v); cb(Vcur, "Vcur_hadamard", il_id); ggml_build_forward_expand(gf, Vcur); } @@ -9455,8 +9745,8 @@ ggml_cgraph* llm_build_context::build_minimaxm2() { ggml_cgraph* llm_build_context::build_smollm3() { ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model, n_tokens), false); - const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + const int64_t n_embd_head = hparams.n_embd_head_v(0); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); // GGML_ASSERT(n_embd_head == hparams.n_rot); this is wrong in case of minimax, head_dim = 128, n_rot = 64 ggml_tensor * cur; @@ -9798,6 +10088,10 @@ ggml_cgraph * llm_build_context::llama_build_graph( { result = llm.build_gemma3(); } break; + case LLM_ARCH_GEMMA4: + { + result = llm.build_gemma4(); + } break; case LLM_ARCH_STARCODER2: { result = llm.build_starcoder2(); @@ -10076,20 +10370,20 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens cb(Qcur, "Qcur_temp_scaled", il_cb); } if (cparams.k_cache_hadamard) { - Qcur = ggml_hadamard(ctx0, Qcur, hparams.n_embd_head_k); - Kcur = ggml_hadamard(ctx0, Kcur, hparams.n_embd_head_k); + Qcur = ggml_hadamard(ctx0, Qcur, hparams.n_embd_head_k(il)); + Kcur = ggml_hadamard(ctx0, Kcur, hparams.n_embd_head_k(il)); cb(Qcur, "Qcur_hadamard", il_cb); cb(Kcur, "Kcur_hadamard", il_cb); } if (cparams.v_cache_hadamard) { - Vcur = ggml_hadamard(ctx0, Vcur, hparams.n_embd_head_v); + Vcur = ggml_hadamard(ctx0, Vcur, hparams.n_embd_head_v(il)); cb(Vcur, "Vcur_hadamard", il_cb); } ggml_build_forward_expand(gf, Qcur); ggml_build_forward_expand(gf, Kcur); ggml_build_forward_expand(gf, Vcur); - const int64_t n_embd_head_k = hparams.n_embd_head_k; + const int64_t n_embd_head_k = hparams.n_embd_head_k(il); const int64_t n_head_kv = split_wk->ne[1] / n_embd_head_k; GGML_ASSERT(kv_self.size == cparams.n_ctx); diff --git a/src/llama-build-context.h b/src/llama-build-context.h index 29edf2d7..58d68f90 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -238,6 +238,8 @@ struct llm_build_context { ggml_cgraph * build_gemma3(); + ggml_cgraph * build_gemma4(); + ggml_cgraph * build_starcoder2(); ggml_cgraph * build_mamba(); @@ -371,7 +373,8 @@ struct llm_build_context { float w_scale, llm_expert_gating_func_type gating_op, const llm_build_cb & cb, int il, ggml_cgraph * graph = nullptr, bool add_input = false, - ggml_tensor * up_gate_exps = nullptr, ggml_tensor * up_gate_exps_b = nullptr); + ggml_tensor * up_gate_exps = nullptr, ggml_tensor * up_gate_exps_b = nullptr, + ggml_tensor * input_logits = nullptr, ggml_tensor * down_exps_s = nullptr); static ggml_tensor * llm_build_moe_ffn(ggml_context * ctx, llama_context & lctx, ggml_tensor * cur, @@ -388,7 +391,8 @@ llm_expert_gating_func_type gating_op, float w_scale, llm_expert_gating_func_type gating_op, const llm_build_cb & cb, int il, ggml_cgraph * graph = nullptr, bool add_input = false, - ggml_tensor * up_gate_exps = nullptr, ggml_tensor * up_gate_exps_b = nullptr) { + ggml_tensor * up_gate_exps = nullptr, ggml_tensor * up_gate_exps_b = nullptr, + ggml_tensor * input_logits = nullptr, ggml_tensor * down_exps_s = nullptr) { return llm_build_moe_ffn(ctx, lctx, cur, gate_inp, nullptr, up_exps, nullptr, @@ -397,7 +401,8 @@ llm_expert_gating_func_type gating_op, exp_probs_b, n_expert, n_expert_used, type_op, norm_w, scale_w, w_scale, - gating_op, cb, il, graph, add_input, up_gate_exps, up_gate_exps_b); + gating_op, cb, il, graph, add_input, up_gate_exps, up_gate_exps_b, + input_logits, down_exps_s); } static ggml_tensor * llm_build_std_moe_ffn(ggml_context * ctx, llama_context & lctx, diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 7c093400..ad13fe7b 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -129,26 +129,36 @@ void llm_load_hparams( // gpt-neox n_rot = rotary_pct * (n_embd / n_head) // gpt-j n_rot = rotary_dim - hparams.n_embd_head_k = hparams.n_embd / hparams.n_head(); - ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH, hparams.n_embd_head_k, false); + hparams.n_embd_head_k_full = hparams.n_embd / hparams.n_head(); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH, hparams.n_embd_head_k_full, false); - hparams.n_embd_head_v = hparams.n_embd / hparams.n_head(); - ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH, hparams.n_embd_head_v, false); + hparams.n_embd_head_v_full = hparams.n_embd / hparams.n_head(); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH, hparams.n_embd_head_v_full, false); // sanity check for n_rot (optional) - hparams.n_rot = hparams.n_embd_head_k; + hparams.n_rot = hparams.n_embd_head_k_full; ml.get_key(LLM_KV_ROPE_DIMENSION_COUNT, hparams.n_rot, false); if (model.arch == LLM_ARCH_LLAMA || model.arch == LLM_ARCH_FALCON || model.arch == LLM_ARCH_BITNET_25 || model.arch == LLM_ARCH_BITNET_B158 || model.arch == LLM_ARCH_DECI) { - if (hparams.n_rot != hparams.n_embd_head_k) { - throw std::runtime_error(format("invalid n_rot: %u, expected %u", hparams.n_rot, hparams.n_embd_head_k)); + if (hparams.n_rot != hparams.n_embd_head_k_full) { + throw std::runtime_error(format("invalid n_rot: %u, expected %u", hparams.n_rot, hparams.n_embd_head_k_full)); } } } else { hparams.n_rot = 0; - hparams.n_embd_head_k = 0; - hparams.n_embd_head_v = 0; + hparams.n_embd_head_k_full = 0; + hparams.n_embd_head_v_full = 0; + } + + { + hparams.n_embd_head_k_swa = hparams.n_embd_head_k_full; + hparams.n_embd_head_v_swa = hparams.n_embd_head_v_full; + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_SWA, hparams.n_embd_head_k_swa, false); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, hparams.n_embd_head_v_swa, false); + + hparams.n_rot_swa = hparams.n_rot; + ml.get_key(LLM_KV_ROPE_DIMENSION_COUNT_SWA, hparams.n_rot_swa, false); } // arch-specific KVs @@ -681,8 +691,36 @@ void llm_load_hparams( hparams.f_attention_scale = model.type == e_model::MODEL_27B ? 1.0f / std::sqrt(float(hparams.n_embd / hparams.n_head(0))) - : 1.0f / std::sqrt(float(hparams.n_embd_head_k)); + : 1.0f / std::sqrt(float(hparams.n_embd_head_k_full)); } break; + case LLM_ARCH_GEMMA4: + { + //hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.swa_layers, hparams.n_layer); + + uint32_t n_kv_shared_layers = 0; + ml.get_key(LLM_KV_ATTENTION_SHARED_KV_LAYERS, n_kv_shared_layers, false); + + hparams.n_layer_kv_from_start = hparams.n_layer - (int32_t)n_kv_shared_layers; + hparams.f_attention_scale = 1.0f; // Gemma4 uses self.scaling = 1.0 (no pre-attn scaling) + hparams.f_final_logit_softcapping = 30.0f; + + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_SWA, hparams.n_embd_head_k_swa); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, hparams.n_embd_head_v_swa); + ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); + + switch (hparams.n_layer) { + case 35: model.type = e_model::MODEL_2B; break; + case 42: model.type = e_model::MODEL_4B; break; // to confirm: E4B or E5B? + default: model.type = e_model::MODEL_UNKNOWN; + } + } break; + case LLM_ARCH_STARCODER2: { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); @@ -844,7 +882,7 @@ void llm_load_hparams( int expected_head_size_v = model.arch == LLM_ARCH_DEEPSEEK2 ? 512 : 256; if (hparams.n_head_kv() == 1) { int n_nead_kv = hparams.n_gqa(); - if (n_nead_kv%4 != 0 || hparams.n_embd_head_k != expected_head_size_k || hparams.n_embd_head_v != expected_head_size_v || + if (n_nead_kv%4 != 0 || hparams.n_embd_head_k(0) != expected_head_size_k || hparams.n_embd_head_v(0) != expected_head_size_v || hparams.n_rot != 64) { printf("==========================================================================\n"); printf("Detected incompatible DeepSeek model without a known way to fix it.\n"); @@ -856,10 +894,10 @@ void llm_load_hparams( } printf("================= Adjusted mainline llama.cpp MLA tensors to ik_llama.cpp\n"); for (auto& item : hparams.n_head_kv_arr) item = n_nead_kv; - hparams.n_embd_head_k = 192; - hparams.n_embd_head_v = 128; - ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k); - ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v); + hparams.n_embd_head_k_full = 192; + hparams.n_embd_head_v_full = 128; + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_full); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_full); } bool is_lite = (hparams.n_layer == 27 || hparams.n_layer == 26) || (hparams.n_layer == 48 && hparams.n_vocab == 128256); ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); @@ -1295,7 +1333,7 @@ void llm_load_hparams( } if (hparams.n_head_kv() == 1) { int n_nead_kv = hparams.n_gqa(); - if (n_nead_kv%4 != 0 || hparams.n_embd_head_k != 576 || hparams.n_embd_head_v != 512 || + if (n_nead_kv%4 != 0 || hparams.n_embd_head_k_full != 576 || hparams.n_embd_head_v_full != 512 || hparams.n_rot != 64) { printf("==========================================================================\n"); printf("Detected incompatible DeepSeek model without a known way to fix it.\n"); @@ -1305,10 +1343,10 @@ void llm_load_hparams( } printf("================= Adjusted mainline llama.cpp MLA tensors to ik_llama.cpp\n"); for (auto& item : hparams.n_head_kv_arr) item = n_nead_kv; - hparams.n_embd_head_k = 192; - hparams.n_embd_head_v = 128; - ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k); - ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v); + hparams.n_embd_head_k_full = 192; + hparams.n_embd_head_v_full = 128; + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_full); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_full); } } break; default: (void)0; diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 1f79557a..b5f481c0 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -26,10 +26,13 @@ struct llama_hparams { uint32_t n_layer; int32_t n_layer_kv_from_start = -1; // if non-negative, the first n_layer_kv_from_start layers have KV cache uint32_t n_rot; + uint32_t n_rot_swa; uint32_t n_swa = 0; // sliding window attention (SWA) uint32_t n_swa_pattern = 1; // by default, all layers use non-sliding-window attention - uint32_t n_embd_head_k; // dimension of keys (d_k). d_q is assumed to be the same, but there are n_head q heads, and only n_head_kv k-v heads - uint32_t n_embd_head_v; // dimension of values (d_v) aka n_embd_head + uint32_t n_embd_head_k_full; // dimension of keys (d_k). d_q is assumed to be the same, but there are n_head q heads, and only n_head_kv k-v heads + uint32_t n_embd_head_v_full; // dimension of values (d_v) aka n_embd_head + uint32_t n_embd_head_k_swa; + uint32_t n_embd_head_v_swa; uint32_t n_expert = 0; uint32_t n_expert_used = 0; uint32_t n_vocab_type = 0; // for BERT-style token types @@ -127,6 +130,9 @@ struct llama_hparams { // qwen3vl deepstack uint32_t n_deepstack_layers = 0; + // gemma4 per-layer embedding + uint32_t n_embd_per_layer = 0; + // needed by encoder-decoder models (e.g. T5, FLAN-T5) // ref: https://github.com/ggerganov/llama.cpp/pull/8141 llama_token dec_start_token_id = -1; @@ -149,8 +155,8 @@ struct llama_hparams { if (this->n_rot != other.n_rot) return true; if (this->n_swa != other.n_swa) return true; if (this->n_swa_pattern != other.n_swa_pattern) return false; - if (this->n_embd_head_k != other.n_embd_head_k) return true; - if (this->n_embd_head_v != other.n_embd_head_v) return true; + if (this->n_embd_head_k_full != other.n_embd_head_k_full) return true; + if (this->n_embd_head_v_full != other.n_embd_head_v_full) return true; if (this->n_expert != other.n_expert) return true; if (this->n_expert_used != other.n_expert_used) return true; @@ -194,6 +200,10 @@ struct llama_hparams { return false; } + bool has_kv(uint32_t il) const { + return n_layer_kv_from_start > 0 ? il < n_layer_kv_from_start : true; + } + uint32_t n_head(uint32_t il = 0) const { if (il < n_layer) { return n_head_arr[il]; @@ -239,16 +249,16 @@ struct llama_hparams { return n_head/n_head_kv; } - uint32_t n_embd_k_gqa(uint32_t il = 0) const { // dimension of key embeddings across all k-v heads - const uint32_t n_head_kv = this->n_head_kv(il); + uint32_t n_embd_head_k(int il) const { return swa_layers[il] ? n_embd_head_k_swa : n_embd_head_k_full; } - return n_embd_head_k * n_head_kv; + uint32_t n_embd_head_v(int il) const { return swa_layers[il] ? n_embd_head_v_swa : n_embd_head_v_full; } + + uint32_t n_embd_k_gqa(uint32_t il = 0) const { // dimension of key embeddings across all k-v heads + return n_head_kv(il) * n_embd_head_k(il); } uint32_t n_embd_v_gqa(uint32_t il = 0) const { // dimension of value embeddings across all k-v heads - const uint32_t n_head_kv = this->n_head_kv(il); - - return n_embd_head_v * n_head_kv; + return n_head_kv(il) * n_embd_head_v(il); } uint32_t n_embd_k_s() const { // dimension of the rolling state embeddings diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index fa12fc53..2adbee0f 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -93,6 +93,8 @@ struct create_tensors_helper : public create_tensors_helper_interface { bool create_gemma_tensors(const LLM_TN & tn, int version); + bool create_gemma4_tensors(const LLM_TN & tn); + bool create_starcoder2_tensors(const LLM_TN & tn); bool create_mamba_tensors(const LLM_TN & tn); @@ -448,8 +450,8 @@ ggml_tensor * create_tensors_helper::create_tensor(ggml_context * ctx, const std [[maybe_unused]] const int64_t n_embd = hparams.n_embd / (hparams.n_deepstack_layers + 1); /* For Qwen3-VL we need to divide by the number of deepstack layers + 1, for other models n_deepstack_layers value is 0 by default */ \ [[maybe_unused]] const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(); \ [[maybe_unused]] const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(); \ - [[maybe_unused]] const int64_t n_embd_head_k = hparams.n_embd_head_k; \ - [[maybe_unused]] const int64_t n_embd_head_v = hparams.n_embd_head_v; \ + [[maybe_unused]] const int64_t n_embd_head_k = hparams.n_embd_head_k(0); \ + [[maybe_unused]] const int64_t n_embd_head_v = hparams.n_embd_head_v(0); \ [[maybe_unused]] const int64_t n_ff = hparams.n_ff(); \ [[maybe_unused]] const int64_t n_embd_gqa = n_embd_v_gqa; \ [[maybe_unused]] const int64_t n_vocab = hparams.n_vocab; \ @@ -1927,6 +1929,105 @@ bool create_tensors_helper::create_gemma_tensors(const LLM_TN & tn, int version) return use_mmap_buffer; } +bool create_tensors_helper::create_gemma4_tensors(const LLM_TN & tn) { + LOADING_PRELUDE + + const uint32_t n_embd_per_layer = hparams.n_embd_per_layer; + const int64_t n_ff_exp = hparams.n_ff_exp; + + if (n_embd_head_k != n_embd_head_v) { + throw std::runtime_error("Gemma 4 requires n_embd_head_k == n_embd_head_v"); + } + if (hparams.n_embd_head_k_swa != hparams.n_embd_head_v_swa) { + throw std::runtime_error("Gemma 4 requires n_embd_head_k_swa == n_embd_head_v_swa"); + } + + model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED); + // if output is NULL, init from the input tok embed + if (!model.output) { + model.output = create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED); + } + + model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + if (n_embd_per_layer > 0) { + model.tok_embd_per_layer = create_tensor(ctx_output, tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, 0); + model.per_layer_model_proj = create_tensor(ctx_output, tn(LLM_TENSOR_PER_LAYER_MODEL_PROJ, "weight"), {n_embd, n_embd_per_layer * n_layer}, 0); + model.per_layer_proj_norm = create_tensor(ctx_output, tn(LLM_TENSOR_PER_LAYER_PROJ_NORM, "weight"), {n_embd_per_layer}, 0); + } + + model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + + int rope_freqs_flag = 0; + + for (int i = 0; i < n_layer; ++i) { + ggml_context * ctx_split = ctx_for_layer_split(i); + auto & layer = model.layers[i]; + const int64_t n_head = hparams.n_head(i); + const int64_t n_embd_head = hparams.n_embd_head_k(i); + const int64_t n_embd_k = hparams.n_embd_k_gqa(i); + const int64_t n_embd_v = hparams.n_embd_v_gqa(i); + + layer.attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + // note: use_alternative_attention (v_proj is optional, if it's not present, use k_proj) + layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head * n_head}, 0); + layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k}, 0); + layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v}, llama_model_loader::TENSOR_NOT_REQUIRED); + layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head * n_head, n_embd}, 0); + + layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head}, 0); + layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head}, 0); + layer.attn_post_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0); + + layer.out_scale = create_tensor(ctx_split, tn(LLM_TENSOR_LAYER_OUT_SCALE, "weight", i), {1u}, llama_model_loader::TENSOR_NOT_REQUIRED); + + if (!hparams.swa_layers[i]) { + // full_attention layers use rope_freqs for proportional rope + layer.rope_freqs = create_tensor(ctx_split, tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_embd_head/2}, rope_freqs_flag); + rope_freqs_flag = llama_model_loader::TENSOR_DUPLICATED; + } + + // handle use_double_wide_mlp + int64_t n_ff_cur = hparams.n_ff(i); + + // for expert layers, we use normal FFN as shared expert (same as python code) + layer.ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_gate = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff_cur}, 0); + layer.ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff_cur}, 0); + layer.ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff_cur, n_embd}, 0); + layer.ffn_post_norm = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0); + + // MoE router + 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); + bool has_expert = layer.ffn_gate_inp != nullptr; + + // norm + if (has_expert) { + layer.ffn_gate_inp_s = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP, "scale", i), {n_embd}, 0); + + layer.ffn_pre_norm_2 = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_PRE_NORM_2, "weight", i), {n_embd}, 0); + layer.ffn_post_norm_1 = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_POST_NORM_1, "weight", i), {n_embd}, 0); + layer.ffn_post_norm_2 = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_POST_NORM_2, "weight", i), {n_embd}, 0); + + // MoE FFN + layer.ffn_up_gate_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_UP_EXPS, "weight", i), {n_embd, n_ff_exp * 2, n_expert}, 0); + layer.ffn_down_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); + layer.ffn_down_exps_s = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_EXPS, "scale", i), {n_expert}, llama_model_loader::TENSOR_NOT_REQUIRED); + + // per-expert scale will be loaded as down_exps_s at the end of the current switch case + } + + // per-layer embeddings + if (n_embd_per_layer > 0) { + layer.per_layer_inp_gate = create_tensor(ctx_split, tn(LLM_TENSOR_PER_LAYER_INP_GATE, "weight", i), {n_embd, n_embd_per_layer}, 0); + layer.per_layer_proj = create_tensor(ctx_split, tn(LLM_TENSOR_PER_LAYER_PROJ, "weight", i), {n_embd_per_layer, n_embd}, 0); + layer.per_layer_post_norm = create_tensor(ctx_split, tn(LLM_TENSOR_PER_LAYER_POST_NORM, "weight", i), {n_embd}, 0); + } + } + return use_mmap_buffer; +} + bool create_tensors_helper::create_starcoder2_tensors(const LLM_TN & tn) { LOADING_PRELUDE @@ -2213,7 +2314,7 @@ bool create_tensors_helper::create_deepseek2_tensors(const LLM_TN & tn) { const bool is_lite = (hparams.n_layer == 27 || hparams.n_layer == 26); const int64_t n_embd_head_qk_rope = hparams.n_rot; - const int64_t n_embd_head_qk_nope = hparams.n_embd_head_k - hparams.n_rot; + const int64_t n_embd_head_qk_nope = hparams.n_embd_head_k(0) - hparams.n_rot; const int64_t q_lora_rank = hparams.n_lora_q; const int64_t kv_lora_rank = hparams.n_lora_kv; @@ -2321,7 +2422,7 @@ bool create_tensors_helper::create_glm_dsa_tensors(const LLM_TN & tn) { LOADING_PRELUDE const int64_t n_embd_head_qk_rope = hparams.n_rot; - const int64_t n_embd_head_qk_nope = hparams.n_embd_head_k - hparams.n_rot; + const int64_t n_embd_head_qk_nope = hparams.n_embd_head_k(0) - hparams.n_rot; const int64_t q_lora_rank = hparams.n_lora_q; const int64_t kv_lora_rank = hparams.n_lora_kv; @@ -2838,8 +2939,8 @@ bool create_tensors_helper::create_chatglm_tensors(const LLM_TN & tn) { layer.attn_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}); - layer.wqkv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd + (hparams.n_embd_head_k << 2)}); - layer.bqkv = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_QKV, "bias", i), {n_embd + (hparams.n_embd_head_k << 2)}); + layer.wqkv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd + (hparams.n_embd_head_k(i) << 2)}); + layer.bqkv = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_QKV, "bias", i), {n_embd + (hparams.n_embd_head_k(i) << 2)}); layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}); @@ -3340,7 +3441,7 @@ bool create_tensors_helper::merge_qkv(const LLM_TN & tn, int i, int bias, bool i const int64_t n_head_kv = hparams.n_head_kv(i); const int64_t n_embd = hparams.n_embd / (hparams.n_deepstack_layers + 1); // For Qwen3-VL we need to divide by the number of deepstack layers + 1, for other models n_deepstack_layers value is 0 by default const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(i); - const int64_t n_embd_head_k = hparams.n_embd_head_k; + const int64_t n_embd_head_k = hparams.n_embd_head_k(i); const int64_t n_embd_gqa = n_embd_v_gqa; ggml_context * ctx_layer = ctx_for_layer(i); @@ -3880,6 +3981,8 @@ bool create_tensors_helper::create_tensors() { use_mmap_buffer = create_gemma_tensors(tn, 2); break; case LLM_ARCH_GEMMA3: use_mmap_buffer = create_gemma_tensors(tn, 3); break; + case LLM_ARCH_GEMMA4: + use_mmap_buffer = create_gemma4_tensors(tn); break; case LLM_ARCH_STARCODER2: use_mmap_buffer = create_starcoder2_tensors(tn); break; case LLM_ARCH_MAMBA: @@ -4018,16 +4121,16 @@ bool create_tensors_helper::create_tensors() { split_recurrent_tensors(hparams, layer, cur_splits, mem_used, ctx_split, il); //, model.arch == LLM_ARCH_QWEN3NEXT ? 0 : 1); } else if (layer.wo && layer.wq && layer.wk && layer.wv) { - auto granularity_kq = hparams.n_embd_head_k * gqa_ratio; + auto granularity_kq = hparams.n_embd_head_k(il) * gqa_ratio; int wq_ne1 = layer.wq->ne[1]; if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_QWEN35) { granularity_kq *= 2; wq_ne1 /= 2; } - auto granularity_vo = hparams.n_embd_head_v * gqa_ratio; + auto granularity_vo = hparams.n_embd_head_v(il) * gqa_ratio; if (ggml_is_quantized(layer.wo->type)) { auto tt = ggml_internal_get_type_traits(layer.wo->type); if (tt.blck_size > granularity_vo) granularity_vo = tt.blck_size; - GGML_ASSERT(granularity_vo % hparams.n_embd_head_v == 0); + GGML_ASSERT(granularity_vo % hparams.n_embd_head_v(il) == 0); // Command-R: align KQ split to wo's block size so wq row // counts remain valid after splitting. if (model.arch == LLM_ARCH_COMMAND_R) { @@ -4059,7 +4162,7 @@ bool create_tensors_helper::create_tensors() { if (layer.attn_q_norm->ne[1] > 1) { // 2D per-head norm (e.g., Command-R+): split along the Q-head dimension auto split_q_heads = split_kq; - for (auto & s : split_q_heads) s /= hparams.n_embd_head_k; + for (auto & s : split_q_heads) s /= hparams.n_embd_head_k(il); prepare_split_tensors(1, ctx_split, layer.attn_q_norm, layer.split_q_norm, split_q_heads, mem_used); } else { prepare_split_tensors(-1, ctx_split, layer.attn_q_norm, layer.split_q_norm, split_kq, mem_used); @@ -4076,7 +4179,7 @@ bool create_tensors_helper::create_tensors() { if (layer.attn_sinks) { auto split_sinks = split_kq; for (auto & s : split_sinks) { - s /= hparams.n_embd_head_k; + s /= hparams.n_embd_head_k(il); } prepare_split_tensors(0, ctx_split, layer.attn_sinks, layer.split_sinks, split_sinks, mem_used); } @@ -4084,7 +4187,7 @@ bool create_tensors_helper::create_tensors() { auto wqkv_gate_split = split_kq; LLAMA_LOG_DEBUG("=================== wqkv_gate_split:"); for (auto & s : wqkv_gate_split) { - s /= hparams.n_embd_head_k; + s /= hparams.n_embd_head_k(il); LLAMA_LOG_DEBUG(" %d", s); } LLAMA_LOG_DEBUG("\n"); @@ -4117,7 +4220,7 @@ bool create_tensors_helper::create_tensors() { // split_kq has already been divided by gqa_ratio, so values are in // (n_embd_head_k * n_head_kv) units; divide again to get head units auto split_k_heads = split_kq; - for (auto & s : split_k_heads) s /= hparams.n_embd_head_k; + for (auto & s : split_k_heads) s /= hparams.n_embd_head_k(il); prepare_split_tensors(1, ctx_split, layer.attn_k_norm, layer.split_k_norm, split_k_heads, mem_used); } else { prepare_split_tensors(-1, ctx_split, layer.attn_k_norm, layer.split_k_norm, split_kq, mem_used); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 2d5f483e..94f266c2 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -760,6 +760,41 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" }, }, }, + { + LLM_ARCH_GEMMA4, + { + { LLM_TENSOR_TOKEN_EMBD, "token_embd" }, + { LLM_TENSOR_OUTPUT_NORM, "output_norm" }, + { LLM_TENSOR_OUTPUT, "output" }, + { LLM_TENSOR_ROPE_FREQS, "rope_freqs" }, + { LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "per_layer_token_embd" }, + { LLM_TENSOR_PER_LAYER_MODEL_PROJ, "per_layer_model_proj" }, + { LLM_TENSOR_PER_LAYER_PROJ_NORM, "per_layer_proj_norm" }, + { LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" }, + { LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" }, + { LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" }, + { LLM_TENSOR_ATTN_K, "blk.%d.attn_k" }, + { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, + { LLM_TENSOR_ATTN_V, "blk.%d.attn_v" }, + { LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" }, + { LLM_TENSOR_ATTN_POST_NORM, "blk.%d.post_attention_norm" }, + { LLM_TENSOR_FFN_NORM, "blk.%d.ffn_norm" }, + { LLM_TENSOR_FFN_GATE, "blk.%d.ffn_gate" }, + { LLM_TENSOR_FFN_DOWN, "blk.%d.ffn_down" }, + { LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" }, + { LLM_TENSOR_FFN_GATE_UP_EXPS, "blk.%d.ffn_gate_up_exps" }, + { LLM_TENSOR_FFN_DOWN_EXPS, "blk.%d.ffn_down_exps" }, + { LLM_TENSOR_FFN_GATE_INP, "blk.%d.ffn_gate_inp" }, + { LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" }, + { LLM_TENSOR_FFN_POST_NORM_1, "blk.%d.post_ffw_norm_1" }, + { LLM_TENSOR_FFN_POST_NORM_2, "blk.%d.post_ffw_norm_2" }, + { LLM_TENSOR_FFN_PRE_NORM_2, "blk.%d.pre_ffw_norm_2" }, + { LLM_TENSOR_LAYER_OUT_SCALE, "blk.%d.layer_output_scale" }, + { LLM_TENSOR_PER_LAYER_INP_GATE, "blk.%d.inp_gate" }, + { LLM_TENSOR_PER_LAYER_PROJ, "blk.%d.proj" }, + { LLM_TENSOR_PER_LAYER_POST_NORM, "blk.%d.post_norm" }, + }, + }, { LLM_ARCH_STARCODER2, { @@ -1890,7 +1925,7 @@ size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, uint3 return size; } auto n_head_kv = hparams.n_head_kv(il); - auto k_size = ggml_row_size(type_k, hparams.n_embd_head_k) * n_head_kv*kv_size; + auto k_size = ggml_row_size(type_k, hparams.n_embd_head_k(il)) * n_head_kv*kv_size; auto v_size = ggml_row_size(type_v, hparams.n_embd_v_gqa(il)) * kv_size; return k_size + v_size; } diff --git a/src/llama-model.h b/src/llama-model.h index a89aac59..b48e1f86 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -157,6 +157,11 @@ struct llama_layer { struct ggml_tensor * attn_norm_cross = nullptr; struct ggml_tensor * attn_norm_enc = nullptr; struct ggml_tensor * wqkv_gate = nullptr; + struct ggml_tensor * ffn_post_norm_1 = nullptr; // gemma4 + struct ggml_tensor * ffn_post_norm_2 = nullptr; // gemma4 + struct ggml_tensor * ffn_pre_norm_2 = nullptr; // gemma4 + struct ggml_tensor * ffn_gate_inp_s = nullptr; // gemma4 + // attention struct ggml_tensor * wq = nullptr; @@ -275,6 +280,7 @@ struct llama_layer { struct ggml_tensor * ffn_gate_exps_b_dup = nullptr; struct ggml_tensor * ffn_down_exps_b_dup = nullptr; struct ggml_tensor * ffn_up_exps_b_dup = nullptr; + struct ggml_tensor * ffn_down_exps_s = nullptr; // ff shared expert (shexp) struct ggml_tensor * ffn_gate_inp_shexp = nullptr; @@ -306,6 +312,11 @@ struct llama_layer { llama_split_tensor split_ffn_act; llama_split_tensor split_ffn_exp_probs_b; + // misc + struct ggml_tensor * per_layer_inp_gate; + struct ggml_tensor * per_layer_proj; + struct ggml_tensor * per_layer_post_norm; + // mamba proj struct ggml_tensor * ssm_in = nullptr; struct ggml_tensor * ssm_x = nullptr; @@ -347,6 +358,7 @@ struct llama_layer { struct ggml_tensor * ffn_gate_scale = nullptr; struct ggml_tensor * ffn_up_scale = nullptr; struct ggml_tensor * ffn_down_scale = nullptr; + struct ggml_tensor * out_scale = nullptr; // gemma4 layer output scale struct llama_layer_nextn nextn; @@ -380,6 +392,10 @@ struct llama_model { struct ggml_tensor * tok_norm; struct ggml_tensor * tok_norm_b; + struct ggml_tensor * tok_embd_per_layer = nullptr; + struct ggml_tensor * per_layer_model_proj = nullptr; + struct ggml_tensor * per_layer_proj_norm = nullptr; + struct ggml_tensor * output_norm; struct ggml_tensor * output_norm_b; struct ggml_tensor * output; diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 918a8ea4..28633339 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -288,6 +288,15 @@ struct llm_tokenizer_bpe : llm_tokenizer { "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", }; break; + case LLAMA_VOCAB_PRE_TYPE_JAIS2: + regex_exprs = { + // original regex from tokenizer.json + //"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s{512}(?!\\S)|\\s{256}(?!\\S)|\\s{128}(?!\\S)|\\s{64}(?!\\S)|\\s{32}(?!\\S)|\\s{16}(?!\\S)|\\s{8}(?!\\S)|\\s{4}(?!\\S)|\\s{1,2}(?!\\S)|\\s{1}", + + // adapted: same as llama3 but with cascading whitespace pattern + "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s{512}(?!\\S)|\\s{256}(?!\\S)|\\s{128}(?!\\S)|\\s{64}(?!\\S)|\\s{32}(?!\\S)|\\s{16}(?!\\S)|\\s{8}(?!\\S)|\\s{4}(?!\\S)|\\s{1,2}(?!\\S)|\\s{1}", + }; + break; case LLAMA_VOCAB_PRE_TYPE_DBRX: case LLAMA_VOCAB_PRE_TYPE_SMAUG: regex_exprs = { @@ -307,12 +316,19 @@ struct llm_tokenizer_bpe : llm_tokenizer { break; case LLAMA_VOCAB_PRE_TYPE_DEEPSEEK3_LLM: case LLAMA_VOCAB_PRE_TYPE_HUNYUAN_DENSE: + case LLAMA_VOCAB_PRE_TYPE_JOYAI_LLM: regex_exprs = { "\\p{N}{1,3}", "[一-龥぀-ゟ゠-ヿ]+", "[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\r\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\r\n]*|\\s*[\r\n]+|\\s+(?!\\S)|\\s+", }; break; + case LLAMA_VOCAB_PRE_TYPE_YOUTU: + regex_exprs = { + "[가-힣ㄱ-ㆎ]+|[！…“”‘’—：；，、-〿︰-﹏]+|[ㄅ-ㄯ]+|[一-龥぀-ゟ゠-ヿ]+", + "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", + }; + break; case LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_CODER: regex_exprs = { "[\r\n]", @@ -354,6 +370,7 @@ struct llm_tokenizer_bpe : llm_tokenizer { case LLAMA_VOCAB_PRE_TYPE_STABLELM2: case LLAMA_VOCAB_PRE_TYPE_QWEN2: case LLAMA_VOCAB_PRE_TYPE_HUNYUAN: + case LLAMA_VOCAB_PRE_TYPE_SOLAR_OPEN: regex_exprs = { // original regex from tokenizer.json // "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" @@ -414,6 +431,14 @@ struct llm_tokenizer_bpe : llm_tokenizer { "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))*((?=[\\p{L}])([^A-Z]))+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))+((?=[\\p{L}])([^A-Z]))*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", }; break; + case LLAMA_VOCAB_PRE_TYPE_TINY_AYA: + regex_exprs = { + // original regex from tokenizer.json: "\\d{1,3}(?=(?:\\d{3})*\\b)" + "\\d{1,3}(?=(?:\\d{3})*\\b)", + // original regex from tokenizer.json: "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" + "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", + }; + break; case LLAMA_VOCAB_PRE_TYPE_KIMI_K2: regex_exprs = { // K2 trigger pattern - this will activate the custom K2 handler in unicode.cpp @@ -449,6 +474,34 @@ struct llm_tokenizer_bpe : llm_tokenizer { "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", }; break; + case LLAMA_VOCAB_PRE_TYPE_AFMOE: + regex_exprs = { + // Digit handling - uses custom implementation in unicode.cpp + // Groups digits with leading 1-2 based on total length modulo 3 + "\\p{AFMoE_digits}", + // CJK and Asian scripts (using direct Unicode literals) + "[一-鿿㐀-䶿豈-﫿぀-ゟ゠-ヿ･-ﾟ⼀-⿟เ-๿຀-໿ក-៿က-႟ꩠ-ꩿꧠ-꧿가-힯ᄀ-ᇿ]+", + // Main BPE pattern + "[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\\r\\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", + }; + break; + case LLAMA_VOCAB_PRE_TYPE_EXAONE_MOE: + regex_exprs = { + // original regex from tokenizer.json + // "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?(?:\\p{L}\\p{M}*(?: \\p{L}\\p{M}*)*)+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]?|\\s*[\\r\\n]|\\s+(?!\\S)|\\s+" + "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?(?:\\p{L}\\p{M}*(?: \\p{L}\\p{M}*)*)+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]?|\\s*[\\r\\n]|\\s+(?!\\S)|\\s+", + }; + break; + case LLAMA_VOCAB_PRE_TYPE_GEMMA4: + // Gemma4 uses SPM-style BPE: spaces are replaced with ▁ by the + // normalizer, then BPE merges run on the whole text without + // word-level pre-splitting. We only need to split on newlines + // since BPE merge lookup asserts no newlines in tokens. + regex_exprs = { + "[^\\n]+|[\\n]+", + }; + byte_encode = false; // uses raw UTF-8, not GPT-2 byte encoding + break; default: // default regex for BPE tokenization pre-processing regex_exprs = { @@ -462,6 +515,7 @@ struct llm_tokenizer_bpe : llm_tokenizer { } std::vector regex_exprs; + bool byte_encode = true; // GPT-2 byte encoding; false for SPM-style BPE (raw UTF-8) }; struct llm_tokenizer_bpe_session { @@ -506,9 +560,10 @@ struct llm_tokenizer_bpe_session { void tokenize(const std::string & text, std::vector & output) { int final_prev_index = -1; - const auto word_collection = unicode_regex_split(text, tokenizer.regex_exprs); + const auto word_collection = unicode_regex_split(text, tokenizer.regex_exprs, tokenizer.byte_encode); symbols_final.clear(); + auto tok_pre = vocab.get_pre_type(); for (const auto & word : word_collection) { work_queue = llm_bigram_bpe::queue(); @@ -521,6 +576,13 @@ struct llm_tokenizer_bpe_session { if (vocab.get_ignore_merges() && vocab.text_to_token(word) != LLAMA_TOKEN_NULL) { symbols.emplace_back(llm_symbol{-1, -1, word.c_str(), word.size()}); offset = word.size(); + } else if (tok_pre == LLAMA_VOCAB_PRE_TYPE_GEMMA4 && word.find_first_not_of('\n') == std::string::npos) { + // fix for gemma 4, ref: https://github.com/ggml-org/llama.cpp/pull/21343 + auto tok = vocab.text_to_token(word); + if (tok != LLAMA_TOKEN_NULL) { + symbols.emplace_back(llm_symbol{-1, -1, word.c_str(), word.size()}); + offset = word.size(); + } } while (offset < word.size()) { @@ -1019,7 +1081,7 @@ private: } private: uint32_t get_node(size_t index) { - if (index > xcda_array_size) { + if (index >= xcda_array_size) { throw std::runtime_error("Index out of array bounds in XCDA array!"); } return xcda_array[index]; @@ -1276,6 +1338,7 @@ struct llm_tokenizer_plamo2 : llm_tokenizer { // Build suffix list in lexicographical order of reversed strings std::vector suffixes; + suffixes.reserve(suffix_to_score.size() + 1); for (const auto & pair : suffix_to_score) { suffixes.push_back(pair.first); } @@ -1583,7 +1646,7 @@ struct llama_vocab::impl { std::vector id_to_token; std::vector cache_special_tokens; - std::vector cache_token_to_piece; // common_token_to_piece(special = true); + std::vector cache_token_to_piece; // llama_token_to_piece(special = true); struct pair_hash { size_t operator()(const std::pair & p) const { return std::hash{}(p.first) ^ //create some hash for pair @@ -1732,26 +1795,33 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { // read bpe merges and populate bpe ranks const int merges_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_MERGES).c_str()); + // Kimi-K2 uses custom tokenization without traditional BPE merges + const bool is_kimi_k2 = (tokenizer_pre == "kimi-k2"); + if (merges_keyidx == -1) { - throw std::runtime_error("cannot find tokenizer merges in model file\n"); - } - - const int n_merges = gguf_get_arr_n(ctx, merges_keyidx); - for (int i = 0; i < n_merges; i++) { - const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i); - //GGML_ASSERT(unicode_cpts_from_utf8(word).size() > 0); - - std::string first; - std::string second; - - const size_t pos = word.find(' ', 1); - - if (pos != std::string::npos) { - first = word.substr(0, pos); - second = word.substr(pos + 1); + if (!is_kimi_k2) { + throw std::runtime_error("cannot find tokenizer merges in model file\n"); } + // Kimi-K2 doesn't need merges, skip + LLAMA_LOG_INFO("%s: Kimi-K2 tokenizer detected, skipping BPE merges\n", __func__); + } else { + const int n_merges = gguf_get_arr_n(ctx, merges_keyidx); + for (int i = 0; i < n_merges; i++) { + const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i); + //GGML_ASSERT(unicode_cpts_from_utf8(word).size() > 0); - bpe_ranks.emplace(std::make_pair(first, second), i); + std::string first; + std::string second; + + const size_t pos = word.find(' ', 1); + + if (pos != std::string::npos) { + first = word.substr(0, pos); + second = word.substr(pos + 1); + } + + bpe_ranks.emplace(std::make_pair(first, second), i); + } } // default special tokens @@ -1780,8 +1850,8 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { const size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx); const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx); precompiled_charsmap.assign(pc, pc + n_precompiled_charsmap); -#ifdef IS_BIG_ENDIAN - // correct endiannes of data in precompiled_charsmap binary blob +#if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ + // correct endianness of data in precompiled_charsmap binary blob uint32_t * xcda_blob_size = (uint32_t *) &precompiled_charsmap[0]; *xcda_blob_size = __builtin_bswap32(*xcda_blob_size); assert(*xcda_blob_size + sizeof(uint32_t) < n_precompiled_charsmap); @@ -1811,6 +1881,42 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { special_sep_id = LLAMA_TOKEN_NULL; special_pad_id = 3; // <|plamo:pad|> special_mask_id = LLAMA_TOKEN_NULL; + } else if (tokenizer_model == "gemma4") { + type = LLAMA_VOCAB_TYPE_BPE; + + // read bpe merges and populate bpe ranks + const int merges_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_MERGES).c_str()); + if (merges_keyidx == -1) { + throw std::runtime_error("cannot find tokenizer merges in model file\n"); + } + { + const int n_merges = gguf_get_arr_n(ctx, merges_keyidx); + for (int i = 0; i < n_merges; i++) { + const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i); + + std::string first; + std::string second; + + const size_t pos = word.find(' ', 1); + + if (pos != std::string::npos) { + first = word.substr(0, pos); + second = word.substr(pos + 1); + } + + bpe_ranks.emplace(std::make_pair(first, second), i); + } + } + + // default special tokens (to be read from GGUF) + special_bos_id = LLAMA_TOKEN_NULL; + special_eos_id = LLAMA_TOKEN_NULL; + special_unk_id = LLAMA_TOKEN_NULL; + special_sep_id = LLAMA_TOKEN_NULL; + special_pad_id = LLAMA_TOKEN_NULL; + special_mask_id = LLAMA_TOKEN_NULL; + + tokenizer_pre = "gemma4"; } else { throw std::runtime_error(format("unknown tokenizer: '%s'", tokenizer_model.c_str())); } @@ -1818,6 +1924,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { // for now, only BPE models have pre-tokenizers if (type == LLAMA_VOCAB_TYPE_BPE) { add_space_prefix = false; + escape_whitespaces = false; clean_spaces = true; if (tokenizer_pre.empty()) { LLAMA_LOG_WARN("%s: missing pre-tokenizer type, using: 'default'\n", __func__); @@ -1838,7 +1945,8 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "falcon-h1" || tokenizer_pre == "pixtral" || tokenizer_pre == "midm-2.0" || - tokenizer_pre == "lfm2") { + tokenizer_pre == "lfm2" || + tokenizer_pre == "jina-v5-nano") { pre_type = LLAMA_VOCAB_PRE_TYPE_LLAMA3; ignore_merges = true; add_bos = true; @@ -1854,6 +1962,11 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "deepseek-v3") { pre_type = LLAMA_VOCAB_PRE_TYPE_DEEPSEEK3_LLM; clean_spaces = false; + } else if ( + tokenizer_pre == "youtu") { + pre_type = LLAMA_VOCAB_PRE_TYPE_YOUTU; + clean_spaces = false; + ignore_merges = true; } else if ( tokenizer_pre == "falcon") { pre_type = LLAMA_VOCAB_PRE_TYPE_FALCON; @@ -1872,8 +1985,16 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "jina-v2-es" || tokenizer_pre == "jina-v2-de" || tokenizer_pre == "a.x-4.0" || - tokenizer_pre == "mellum") { + tokenizer_pre == "mellum" || + tokenizer_pre == "modern-bert") { pre_type = LLAMA_VOCAB_PRE_TYPE_GPT2; + } else if ( + tokenizer_pre == "jais-2") { + pre_type = LLAMA_VOCAB_PRE_TYPE_JAIS2; + } else if ( + tokenizer_pre == "gemma4") { + pre_type = LLAMA_VOCAB_PRE_TYPE_GEMMA4; + escape_whitespaces = true; } else if ( tokenizer_pre == "jina-v1-en" || tokenizer_pre == "jina-v2-code" || @@ -1889,11 +2010,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { clean_spaces = false; } else if ( tokenizer_pre == "qwen2" || - tokenizer_pre == "deepseek-r1-qwen") { + tokenizer_pre == "deepseek-r1-qwen" || + tokenizer_pre == "kormo" || + tokenizer_pre == "f2llmv2") { pre_type = LLAMA_VOCAB_PRE_TYPE_QWEN2; clean_spaces = false; } else if ( - tokenizer_pre == "qwen35") { + tokenizer_pre == "qwen35") { pre_type = LLAMA_VOCAB_PRE_TYPE_QWEN35; clean_spaces = false; } else if ( @@ -1949,6 +2072,9 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { } else if ( tokenizer_pre == "exaone4") { pre_type = LLAMA_VOCAB_PRE_TYPE_GPT2; + } else if ( + tokenizer_pre == "exaone-moe") { + pre_type = LLAMA_VOCAB_PRE_TYPE_EXAONE_MOE; } else if ( tokenizer_pre == "chameleon") { pre_type = LLAMA_VOCAB_PRE_TYPE_CHAMELEON; @@ -1961,10 +2087,15 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "megrez") { pre_type = LLAMA_VOCAB_PRE_TYPE_QWEN2; } else if ( - tokenizer_pre == "gpt-4o" || - tokenizer_pre == "llama4") { + tokenizer_pre == "gpt-4o" || + tokenizer_pre == "llama4" || + tokenizer_pre == "kanana2") { pre_type = LLAMA_VOCAB_PRE_TYPE_GPT4O; clean_spaces = false; + } else if ( + tokenizer_pre == "tiny_aya") { + pre_type = LLAMA_VOCAB_PRE_TYPE_TINY_AYA; + clean_spaces = false; } else if ( tokenizer_pre == "superbpe") { pre_type = LLAMA_VOCAB_PRE_TYPE_SUPERBPE; @@ -1979,7 +2110,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { clean_spaces = false; } else if ( tokenizer_pre == "bailingmoe" || - tokenizer_pre == "bailingmoe2"|| + tokenizer_pre == "bailingmoe2" || tokenizer_pre == "llada-moe") { pre_type = LLAMA_VOCAB_PRE_TYPE_BAILINGMOE; clean_spaces = false; @@ -1995,19 +2126,30 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "hunyuan-dense") { pre_type = LLAMA_VOCAB_PRE_TYPE_HUNYUAN_DENSE; clean_spaces = false; + } else if ( + tokenizer_pre == "joyai-llm") { + pre_type = LLAMA_VOCAB_PRE_TYPE_JOYAI_LLM; + clean_spaces = false; } else if ( tokenizer_pre == "kimi-k2") { pre_type = LLAMA_VOCAB_PRE_TYPE_KIMI_K2; clean_spaces = false; - } - else if ( + } else if ( tokenizer_pre == "grok-2") { pre_type = LLAMA_VOCAB_PRE_TYPE_GROK_2; clean_spaces = false; + } else if ( + tokenizer_pre == "afmoe") { + pre_type = LLAMA_VOCAB_PRE_TYPE_AFMOE; + clean_spaces = false; } else if ( tokenizer_pre == "minimax-m2") { pre_type = LLAMA_VOCAB_PRE_TYPE_MINIMAX_M2; clean_spaces = false; + } else if ( + tokenizer_pre == "solar-open") { + pre_type = LLAMA_VOCAB_PRE_TYPE_SOLAR_OPEN; + clean_spaces = false; } else { throw std::runtime_error(format("unknown pre-tokenizer type: '%s'", tokenizer_pre.c_str())); } @@ -2047,19 +2189,28 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { throw std::runtime_error("cannot find tokenizer vocab in model file\n"); } + const uint32_t n_tokens = gguf_get_arr_n(ctx, token_idx); + const float * scores = nullptr; const int score_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SCORES).c_str()); if (score_idx != -1) { + const uint32_t n_scores = gguf_get_arr_n(ctx, score_idx); + if (n_scores < n_tokens) { + throw std::runtime_error("Index out of array bounds for scores (" + std::to_string(n_scores) + " < " + std::to_string(n_tokens) + ")\n"); + } scores = (const float * ) gguf_get_arr_data(ctx, score_idx); } const int * toktypes = nullptr; const int toktype_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str()); if (toktype_idx != -1) { + const uint32_t n_toktypes = gguf_get_arr_n(ctx, toktype_idx); + if (n_toktypes < n_tokens) { + throw std::runtime_error("Index out of array bounds for toktypes (" + std::to_string(n_toktypes) + " < " + std::to_string(n_tokens) + ")\n"); + } toktypes = (const int * ) gguf_get_arr_data(ctx, toktype_idx); } - uint32_t n_tokens = gguf_get_arr_n(ctx, token_idx); id_to_token.resize(n_tokens); for (uint32_t i = 0; i < n_tokens; i++) { @@ -2180,6 +2331,8 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { // for now, we apply this workaround to find the tokens based on their text for (const auto & t : token_to_id) { + auto & attr = id_to_token[t.second].attr; + // find EOT token: "<|eot_id|>", "<|im_end|>", "", etc. if (special_eot_id == LLAMA_TOKEN_NULL) { if (false @@ -2188,16 +2341,18 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { || t.first == "<|end|>" || t.first == "" || t.first == "<|endoftext|>" + || t.first == "<|end_of_text|>" // granite || t.first == "" || t.first == "_" + || t.first == "[EOT]" // Kimi-K2 || t.first == "<｜end▁of▁sentence｜>" // DeepSeek || t.first == "" // smoldocling ) { special_eot_id = t.second; - if ((id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) { + if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) { LLAMA_LOG_WARN("%s: control-looking token: %6d '%s' was not control-type; this is probably a bug in the model. its type will be overridden\n", __func__, t.second, t.first.c_str()); - id_to_token[t.second].attr = LLAMA_TOKEN_ATTR_CONTROL; + attr = (llama_token_attr) (attr | LLAMA_TOKEN_ATTR_CONTROL); } } } @@ -2208,10 +2363,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { || t.first == "<|eom_id|>" ) { special_eom_id = t.second; - if ((id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) { + if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) { LLAMA_LOG_WARN("%s: control-looking token: %6d '%s' was not control-type; this is probably a bug in the model. its type will be overridden\n", __func__, t.second, t.first.c_str()); - id_to_token[t.second].attr = LLAMA_TOKEN_ATTR_CONTROL; + attr = (llama_token_attr) (attr | LLAMA_TOKEN_ATTR_CONTROL); } } } @@ -2226,12 +2381,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { || t.first == "
"
                         || t.first == "▁
"          // CodeLlama
                         || t.first == "<|code_prefix|>" // GLM-4.5
+                        || t.first == "<|prefix|>"      // Falcon-H1-Tiny-Coder
                         ) {
                     special_fim_pre_id = t.second;
-                    if ((id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
+                    if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
                         LLAMA_LOG_WARN("%s: control-looking token: %6d '%s' was not control-type; this is probably a bug in the model. its type will be overridden\n",
                                 __func__, t.second, t.first.c_str());
-                        id_to_token[t.second].attr = LLAMA_TOKEN_ATTR_CONTROL;
+                        attr = (llama_token_attr) (attr | LLAMA_TOKEN_ATTR_CONTROL);
                     }
                 }
             }
@@ -2246,12 +2402,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                         || t.first == ""
                         || t.first == "▁"         // CodeLlama
                         || t.first == "<|code_suffix|>" // GLM-4.5
+                        || t.first == "<|suffix|>"      // Falcon-H1-Tiny-Coder
                         ) {
                     special_fim_suf_id = t.second;
-                    if ((id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
+                    if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
                         LLAMA_LOG_WARN("%s: control-looking token: %6d '%s' was not control-type; this is probably a bug in the model. its type will be overridden\n",
                                 __func__, t.second, t.first.c_str());
-                        id_to_token[t.second].attr = LLAMA_TOKEN_ATTR_CONTROL;
+                        attr = (llama_token_attr) (attr | LLAMA_TOKEN_ATTR_CONTROL);
                     }
                 }
             }
@@ -2266,12 +2423,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                         || t.first == ""
                         || t.first == "▁"         // CodeLlama
                         || t.first == "<|code_middle|>" // GLM-4.5
+                        || t.first == "<|middle|>"      // Falcon-H1-Tiny-Coder
                         ) {
                     special_fim_mid_id = t.second;
-                    if ((id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
+                    if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
                         LLAMA_LOG_WARN("%s: control-looking token: %6d '%s' was not control-type; this is probably a bug in the model. its type will be overridden\n",
                                 __func__, t.second, t.first.c_str());
-                        id_to_token[t.second].attr = LLAMA_TOKEN_ATTR_CONTROL;
+                        attr = (llama_token_attr) (attr | LLAMA_TOKEN_ATTR_CONTROL);
                     }
                 }
             }
@@ -2283,12 +2441,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                         || t.first == ""
                         || t.first == ""   // Granite
                         || t.first == ""
+                        || t.first == "[PAD]" // Kimi-K2
                         ) {
                     special_fim_pad_id = t.second;
-                    if ((id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
+                    if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
                         LLAMA_LOG_WARN("%s: control-looking token: %6d '%s' was not control-type; this is probably a bug in the model. its type will be overridden\n",
                                 __func__, t.second, t.first.c_str());
-                        id_to_token[t.second].attr = LLAMA_TOKEN_ATTR_CONTROL;
+                        attr = (llama_token_attr) (attr | LLAMA_TOKEN_ATTR_CONTROL);
                     }
                 }
             }
@@ -2303,10 +2462,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                         || t.first == ""    // Granite
                         ) {
                     special_fim_rep_id = t.second;
-                    if ((id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
+                    if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
                         LLAMA_LOG_WARN("%s: control-looking token: %6d '%s' was not control-type; this is probably a bug in the model. its type will be overridden\n",
                                 __func__, t.second, t.first.c_str());
-                        id_to_token[t.second].attr = LLAMA_TOKEN_ATTR_CONTROL;
+                        attr = (llama_token_attr) (attr | LLAMA_TOKEN_ATTR_CONTROL);
                     }
                 }
             }
@@ -2317,15 +2476,41 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                         || t.first == "<|file_sep|>" // Qwen
                         ) {
                     special_fim_sep_id = t.second;
-                    if ((id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
+                    if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
                         LLAMA_LOG_WARN("%s: control-looking token: %6d '%s' was not control-type; this is probably a bug in the model. its type will be overridden\n",
                                 __func__, t.second, t.first.c_str());
-                        id_to_token[t.second].attr = LLAMA_TOKEN_ATTR_CONTROL;
+                        attr = (llama_token_attr) (attr | LLAMA_TOKEN_ATTR_CONTROL);
                     }
                 }
             }
         }
 
+        // auto-detect unused tokens: e.g. control tokens with the word "unused"
+        // ideally, these tokens should be marked as unused during conversion
+        {
+            uint32_t n_unused = 0;
+
+            for (const auto & t : token_to_id) {
+                auto & attr = id_to_token[t.second].attr;
+
+                if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
+                    continue;
+                }
+
+                if ((attr & LLAMA_TOKEN_ATTR_UNUSED) == 0) {
+                    if (strstr(t.first.c_str(), "unused") != NULL) {
+                        attr = (llama_token_attr) (attr | LLAMA_TOKEN_ATTR_UNUSED);
+                    }
+                }
+
+                if (attr & LLAMA_TOKEN_ATTR_UNUSED) {
+                    n_unused++;
+                }
+            }
+
+            LLAMA_LOG_INFO("%s: %u unused tokens\n", __func__, n_unused);
+        }
+
         // maintain a list of tokens that cause end-of-generation
         // this is currently determined based on the token text, which is obviously not ideal
         // ref: https://github.com/ggerganov/llama.cpp/issues/9606
@@ -2344,39 +2529,56 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
         }
 
         for (const auto & t : token_to_id) {
+            auto & attr = id_to_token[t.second].attr;
+
             if (false
                     || t.first == "<|eot_id|>"
                     || t.first == "<|im_end|>"
                     || t.first == "<|end|>"
                     || t.first == "<|return|>" // o200k_harmony
                     || t.first == "<|call|>"   // o200k_harmony
+                    || t.first == "<|flush|>"  // solar-open
+                    || t.first == "<|calls|>"  // solar-open
                     || t.first == ""
                     || t.first == "<|endoftext|>"
+                    || t.first == ""      // paddleocr
                     || t.first == "<|eom_id|>"
                     || t.first == ""
                     || t.first == "_"
+                    || t.first == "[EOT]" // Kimi-K2
+                    || t.first == "[EOS]" // Kimi-K2
                     || t.first == "<|end_of_text|>"
                     || t.first == "" // smoldocling
+                    || t.first == "" // gemma4
+                    || t.first == "<|tool_response>" // gemma4
+                    || t.first == "<｜end▁of▁sentence｜>" // deepseek-ocr
                ) {
                 special_eog_ids.insert(t.second);
-                if ((id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
+                if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
                     LLAMA_LOG_WARN("%s: control-looking token: %6d '%s' was not control-type; this is probably a bug in the model. its type will be overridden\n",
                             __func__, t.second, t.first.c_str());
-                    id_to_token[t.second].attr = LLAMA_TOKEN_ATTR_CONTROL;
+                    attr = (llama_token_attr) (attr | LLAMA_TOKEN_ATTR_CONTROL);
                 }
             } else {
-                // token is control, but not marked as EOG -> print a debug log
-                //if (id_to_token[t.second].attr & LLAMA_TOKEN_ATTR_CONTROL && special_eog_ids.count(t.second) == 0) {
-                //    LLAMA_LOG_DEBUG("%s: control token: %6d '%s' is not marked as EOG\n",
-                //            __func__, t.second, t.first.c_str());
-                //}
+                if (attr & LLAMA_TOKEN_ATTR_CONTROL && !(attr & LLAMA_TOKEN_ATTR_UNUSED)) {
+                    // token is control, but not marked as EOG -> print a debug log
+                    if (special_eog_ids.count(t.second) == 0) {
+                        LLAMA_LOG_DEBUG("%s: control token: %6d '%s' is not marked as EOG\n",
+                                __func__, t.second, t.first.c_str());
+                    }
+                }
             }
         }
 
         // @ngxson : quick hack for gpt-oss, always render these tokens
         for (const auto & t : token_to_id) {
+            auto & attr = id_to_token[t.second].attr;
+
             if (t.first == "<|channel|>" || t.first == "<|message|>" || t.first == "<|start|>" || t.first == "<|constrain|>") {
-                id_to_token[t.second].attr = LLAMA_TOKEN_ATTR_USER_DEFINED;
+                LLAMA_LOG_WARN("%s: setting token '%s' (%d) attribute to USER_DEFINED (%u), old attributes: %u\n",
+                        __func__, t.first.c_str(), t.second, LLAMA_TOKEN_ATTR_USER_DEFINED, attr);
+
+                attr = LLAMA_TOKEN_ATTR_USER_DEFINED;
             }
         }
 
@@ -2396,34 +2598,42 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
             LLAMA_LOG_WARN("%s: special_eom_id is not in special_eog_ids - the tokenizer config may be incorrect\n", __func__);
         }
 
-        // TODO: workaround for o200k_harmony tokenizer: the "<|end|>" token should not be EOG
-        //       we don't have a good way to detect this, so for now, if we have "<|return|>" and "<|call|>" tokens,
+        // TODO: workaround for o200k_harmony and solar-open tokenizer: the "<|end|>" token should not be EOG
+        //       we don't have a good way to detect this, so for now, if we have "<|return|>" and "<|call|>" tokens ("<|calls|>" and "<|flush|>" for solar-open),
         //       we remove the "<|end|>" token from the EOG list
         {
             bool has_return = false;
             bool has_call   = false;
             bool has_end    = false;
+            bool has_flush  = false;
 
             llama_token end_id = LLAMA_TOKEN_NULL;
 
             LLAMA_LOG_INFO("%s: printing all EOG tokens:\n", __func__);
             for (auto tid : special_eog_ids) {
-                LLAMA_LOG_INFO("%s:   - %d ('%s')\n", __func__, tid, id_to_token[tid].text.c_str());
+                auto & text = id_to_token[tid].text;
 
-                if (id_to_token[tid].text == "<|return|>") {
+                LLAMA_LOG_INFO("%s:   - %d ('%s')\n", __func__, tid, text.c_str());
+
+                if (text == "<|return|>") {
                     has_return = true;
-                } else if (id_to_token[tid].text == "<|call|>") {
+                } else if (text == "<|call|>" || text == "<|calls|>") {
                     has_call = true;
-                } else if (id_to_token[tid].text == "<|end|>") {
+                } else if (text == "<|flush|>") {
+                    has_flush = true;
+                } else if (text == "<|end|>") {
                     has_end = true;
                     end_id = tid;
                 }
             }
 
-            if (has_return && has_call && has_end) {
+            if ((has_return && has_call && has_end) || (has_call && has_flush && has_end)) {
                 special_eog_ids.erase(end_id);
-                id_to_token[end_id].attr = LLAMA_TOKEN_ATTR_USER_DEFINED;
-                LLAMA_LOG_WARN("%s: special_eog_ids contains both '<|return|>' and '<|call|>' tokens, removing '<|end|>' token from EOG list\n", __func__);
+
+                auto & attr = id_to_token[end_id].attr;
+                attr = LLAMA_TOKEN_ATTR_USER_DEFINED;
+
+                LLAMA_LOG_WARN("%s: special_eog_ids contains both '<|return|>' and '<|call|>', or '<|calls|>' and '<|flush|>' tokens, removing '<|end|>' token from EOG list\n", __func__);
             }
         }
     }
@@ -2521,6 +2731,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
             for (const auto * token : {"", "", "<|endoftext|>"}) {
                 _set_token_attr(token, LLAMA_TOKEN_ATTR_RSTRIP, false);
             }
+        } else if (_contains_any(model_name, {"modern-bert"})) {
+            if (token_to_id.count("[MASK]") == 0 ) {
+                LLAMA_LOG_WARN("%s: Mask token missing in vocab!\n", __func__);
+            }
+            else {
+                _set_token_attr("[MASK]", LLAMA_TOKEN_ATTR_LSTRIP, true);
+            }
         }
     }
 }
@@ -2876,6 +3093,10 @@ std::vector llama_vocab::impl::tokenize(
                     if (fragment.type == FRAGMENT_BUFFER_VARIANT_TYPE_RAW_TEXT) {
                         std::string text = fragment.raw_text.substr(fragment.offset, fragment.length);
 
+                        if (escape_whitespaces) {
+                            llama_escape_whitespace(text);
+                        }
+
 #ifdef PRETOKENIZERDEBUG
                         LLAMA_LOG_WARN("TT: (%ld %ld %ld) '%s'\n", text.length(), fragment.offset, fragment.length, text.c_str());
 #endif
@@ -3055,6 +3276,12 @@ int32_t llama_vocab::impl::token_to_piece(llama_token token, char * buf, int32_t
                     return _try_copy(token_text.data(), token_text.size());
                 }
                 if (attr & LLAMA_TOKEN_ATTR_NORMAL) {
+                    if (escape_whitespaces) {
+                        // SPM-style BPE: tokens contain ▁ for spaces
+                        std::string result = token_text;
+                        llama_unescape_whitespace(result);
+                        return _try_copy(result.data(), result.size());
+                    }
                     std::string result = llama_decode_text(token_text);
                     return _try_copy(result.data(), result.size());
                 }
@@ -3214,41 +3441,40 @@ int32_t llama_vocab::impl::detokenize(
 }
 
 void llama_vocab::impl::print_info() const {
-    LLAMA_LOG_INFO("%s: vocab type       = %s\n",     __func__, type_name().c_str());
-    LLAMA_LOG_INFO("%s: n_vocab          = %u\n",     __func__, vocab.n_tokens());
-    LLAMA_LOG_INFO("%s: n_merges         = %u\n",     __func__, (uint32_t) bpe_ranks.size());
+    LLAMA_LOG_INFO("%s: vocab type            = %s\n",     __func__, type_name().c_str());
+    LLAMA_LOG_INFO("%s: n_vocab               = %u\n",     __func__, vocab.n_tokens());
+    LLAMA_LOG_INFO("%s: n_merges              = %u\n",     __func__, (uint32_t) bpe_ranks.size());
 
     // special tokens
-    if (special_bos_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: BOS token        = %d '%s'\n", __func__, special_bos_id,     id_to_token.at(special_bos_id).text.c_str() );  }
-    if (special_eos_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: EOS token        = %d '%s'\n", __func__, special_eos_id,     id_to_token.at(special_eos_id).text.c_str() );  }
-    if (special_eot_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: EOT token        = %d '%s'\n", __func__, special_eot_id,     id_to_token.at(special_eot_id).text.c_str() );  }
-    if (special_eom_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: EOM token        = %d '%s'\n", __func__, special_eom_id,     id_to_token.at(special_eom_id).text.c_str() );  }
-    if (special_unk_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: UNK token        = %d '%s'\n", __func__, special_unk_id,     id_to_token.at(special_unk_id).text.c_str() );  }
-    if (special_sep_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: SEP token        = %d '%s'\n", __func__, special_sep_id,     id_to_token.at(special_sep_id).text.c_str() );  }
-    if (special_pad_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: PAD token        = %d '%s'\n", __func__, special_pad_id,     id_to_token.at(special_pad_id).text.c_str() );  }
-    if (special_mask_id != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: MASK token       = %d '%s'\n", __func__, special_mask_id,    id_to_token.at(special_mask_id).text.c_str() ); }
+    if (special_bos_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: BOS token             = %d '%s'\n", __func__, special_bos_id,     id_to_token.at(special_bos_id).text.c_str() );  }
+    if (special_eos_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: EOS token             = %d '%s'\n", __func__, special_eos_id,     id_to_token.at(special_eos_id).text.c_str() );  }
+    if (special_eot_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: EOT token             = %d '%s'\n", __func__, special_eot_id,     id_to_token.at(special_eot_id).text.c_str() );  }
+    if (special_eom_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: EOM token             = %d '%s'\n", __func__, special_eom_id,     id_to_token.at(special_eom_id).text.c_str() );  }
+    if (special_unk_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: UNK token             = %d '%s'\n", __func__, special_unk_id,     id_to_token.at(special_unk_id).text.c_str() );  }
+    if (special_sep_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: SEP token             = %d '%s'\n", __func__, special_sep_id,     id_to_token.at(special_sep_id).text.c_str() );  }
+    if (special_pad_id  != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: PAD token             = %d '%s'\n", __func__, special_pad_id,     id_to_token.at(special_pad_id).text.c_str() );  }
+    if (special_mask_id != LLAMA_TOKEN_NULL)    { LLAMA_LOG_INFO( "%s: MASK token            = %d '%s'\n", __func__, special_mask_id,    id_to_token.at(special_mask_id).text.c_str() ); }
 
-    if (linefeed_id != LLAMA_TOKEN_NULL)        { LLAMA_LOG_INFO( "%s: LF token         = %d '%s'\n", __func__, linefeed_id,        id_to_token.at(linefeed_id).text.c_str() ); }
+    if (linefeed_id != LLAMA_TOKEN_NULL)        { LLAMA_LOG_INFO( "%s: LF token              = %d '%s'\n", __func__, linefeed_id,        id_to_token.at(linefeed_id).text.c_str() ); }
 
-    if (special_fim_pre_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM PRE token    = %d '%s'\n", __func__, special_fim_pre_id, id_to_token.at(special_fim_pre_id).text.c_str() ); }
-    if (special_fim_suf_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM SUF token    = %d '%s'\n", __func__, special_fim_suf_id, id_to_token.at(special_fim_suf_id).text.c_str() ); }
-    if (special_fim_mid_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM MID token    = %d '%s'\n", __func__, special_fim_mid_id, id_to_token.at(special_fim_mid_id).text.c_str() ); }
-    if (special_fim_pad_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM PAD token    = %d '%s'\n", __func__, special_fim_pad_id, id_to_token.at(special_fim_pad_id).text.c_str() ); }
-    if (special_fim_rep_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM REP token    = %d '%s'\n", __func__, special_fim_rep_id, id_to_token.at(special_fim_rep_id).text.c_str() ); }
-    if (special_fim_sep_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM SEP token    = %d '%s'\n", __func__, special_fim_sep_id, id_to_token.at(special_fim_sep_id).text.c_str() ); }
+    if (special_fim_pre_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM PRE token         = %d '%s'\n", __func__, special_fim_pre_id, id_to_token.at(special_fim_pre_id).text.c_str() ); }
+    if (special_fim_suf_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM SUF token         = %d '%s'\n", __func__, special_fim_suf_id, id_to_token.at(special_fim_suf_id).text.c_str() ); }
+    if (special_fim_mid_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM MID token         = %d '%s'\n", __func__, special_fim_mid_id, id_to_token.at(special_fim_mid_id).text.c_str() ); }
+    if (special_fim_pad_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM PAD token         = %d '%s'\n", __func__, special_fim_pad_id, id_to_token.at(special_fim_pad_id).text.c_str() ); }
+    if (special_fim_rep_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM REP token         = %d '%s'\n", __func__, special_fim_rep_id, id_to_token.at(special_fim_rep_id).text.c_str() ); }
+    if (special_fim_sep_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: FIM SEP token         = %d '%s'\n", __func__, special_fim_sep_id, id_to_token.at(special_fim_sep_id).text.c_str() ); }
 
     for (const auto & id : special_eog_ids) {
-        LLAMA_LOG_INFO( "%s: EOG token        = %d '%s'\n", __func__, id, id_to_token.at(id).text.c_str() );
+        LLAMA_LOG_INFO( "%s: EOG token             = %d '%s'\n", __func__, id, id_to_token.at(id).text.c_str() );
     }
 
-    LLAMA_LOG_INFO("%s: max token length = %d\n", __func__, max_token_len);
+    LLAMA_LOG_INFO("%s: max token length      = %d\n", __func__, max_token_len);
 }
 
 llama_vocab::llama_vocab() : pimpl(new impl(*this)) {
 }
 
-llama_vocab::~llama_vocab() {
-}
+llama_vocab::~llama_vocab() = default;
 
 void llama_vocab::load(llama_model_loader & ml, const LLM_KV & kv) {
     pimpl->load(ml, kv);
@@ -3486,9 +3712,7 @@ int llama_vocab::max_token_len() const {
 
 int llama_vocab::find_bpe_rank(const std::string & token_left, const std::string & token_right) const {
     GGML_ASSERT(token_left.find(' ')   == std::string::npos);
-    GGML_ASSERT(token_left.find('\n')  == std::string::npos);
     GGML_ASSERT(token_right.find(' ')  == std::string::npos);
-    GGML_ASSERT(token_right.find('\n') == std::string::npos);
 
     auto it = pimpl->bpe_ranks.find(std::make_pair(token_left, token_right));
     if (it == pimpl->bpe_ranks.end()) {
diff --git a/src/llama-vocab.h b/src/llama-vocab.h
index 6f064bb0..829be07d 100644
--- a/src/llama-vocab.h
+++ b/src/llama-vocab.h
@@ -48,9 +48,17 @@ enum llama_vocab_pre_type {
     LLAMA_VOCAB_PRE_TYPE_KIMI_K2         = 37,
     LLAMA_VOCAB_PRE_TYPE_HUNYUAN_DENSE   = 38,
     LLAMA_VOCAB_PRE_TYPE_GROK_2          = 39,
-    LLAMA_VOCAB_PRE_TYPE_MINIMAX_M2      = 40,
-    LLAMA_VOCAB_PRE_TYPE_GRANITE_DOCLING = 41,
+    LLAMA_VOCAB_PRE_TYPE_GRANITE_DOCLING = 40,
+    LLAMA_VOCAB_PRE_TYPE_MINIMAX_M2      = 41,
+    LLAMA_VOCAB_PRE_TYPE_AFMOE           = 42,
+    LLAMA_VOCAB_PRE_TYPE_SOLAR_OPEN      = 43,
+    LLAMA_VOCAB_PRE_TYPE_YOUTU           = 44,
+    LLAMA_VOCAB_PRE_TYPE_EXAONE_MOE      = 45,
     LLAMA_VOCAB_PRE_TYPE_QWEN35          = 46,
+    LLAMA_VOCAB_PRE_TYPE_TINY_AYA        = 47,
+    LLAMA_VOCAB_PRE_TYPE_JOYAI_LLM       = 48,
+    LLAMA_VOCAB_PRE_TYPE_JAIS2           = 49,
+    LLAMA_VOCAB_PRE_TYPE_GEMMA4          = 50,
 };
 
 struct LLM_KV;
diff --git a/src/llama.cpp b/src/llama.cpp
index ee92f55a..31be5479 100644
--- a/src/llama.cpp
+++ b/src/llama.cpp
@@ -846,7 +846,7 @@ static bool llama_kv_cache_init(
         const bool qnext_recurrent = llama_is_recurrent_layer(hparams, i);
         const uint32_t n_embd_v_row = llama_kv_v_row_embd(model, hparams, i);
         const uint32_t n_head_kv    = hparams.n_head_kv(i);
-        const uint32_t n_embd_head_k= hparams.n_embd_head_k;
+        const uint32_t n_embd_head_k= hparams.n_embd_head_k(i);
 
         //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 ? ctx_map.at(model.buft_layer[i].buft_matrix) : offload ? ctx_map.at(model.buft_layer[i].buft) : cache.ctxs.front();
@@ -876,6 +876,11 @@ static bool llama_kv_cache_init(
             n_mla++;
         }
         else {
+            if (!hparams.has_kv(i)) {
+                cache.k_l.push_back(nullptr);
+                cache.v_l.push_back(nullptr);
+                continue;
+            }
             if (qnext_recurrent) {
                 s = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.n_embd_v_s(), qnext_state_slots);
                 auto s_name = std::string{"cache_s_l"} + std::to_string(i);
@@ -917,7 +922,7 @@ static bool llama_kv_cache_init(
                 ctx = offload ? ctx_map.at(model.buft_layer[i].buft) : cache.ctxs.front();
                 split_cache_i = false;
             }
-            int n_embd_head_v = hparams.n_embd_head_v;
+            int n_embd_head_v = hparams.n_embd_head_v(i);
             k = ggml_new_tensor_2d(ctx, type_k, n_embd_head_k, n_head_kv*kv_size);
 
             int64_t v_ne = int64_t(n_embd_v_row)*kv_size;
@@ -1537,8 +1542,8 @@ static void llm_load_print_meta(llama_model_loader & ml, llama_model & model) {
         LLAMA_LOG_INFO("%s: n_rot            = %u\n",     __func__, hparams.n_rot);
         LLAMA_LOG_INFO("%s: n_swa            = %u\n",     __func__, hparams.n_swa);
         LLAMA_LOG_INFO("%s: n_swa_pattern    = %u\n",     __func__, hparams.n_swa_pattern);
-        LLAMA_LOG_INFO("%s: n_embd_head_k    = %u\n",     __func__, hparams.n_embd_head_k);
-        LLAMA_LOG_INFO("%s: n_embd_head_v    = %u\n",     __func__, hparams.n_embd_head_v);
+        LLAMA_LOG_INFO("%s: n_embd_head_k    = %u\n",     __func__, hparams.n_embd_head_k(0));
+        LLAMA_LOG_INFO("%s: n_embd_head_v    = %u\n",     __func__, hparams.n_embd_head_v(0));
         LLAMA_LOG_INFO("%s: n_gqa            = %s\n",     __func__, print_f([&](uint32_t il) { return hparams.n_gqa(il);        }, hparams.n_layer).c_str());
         LLAMA_LOG_INFO("%s: n_embd_k_gqa     = %s\n",     __func__, print_f([&](uint32_t il) { return hparams.n_embd_k_gqa(il); }, hparams.n_layer).c_str());
         LLAMA_LOG_INFO("%s: n_embd_v_gqa     = %s\n",     __func__, print_f([&](uint32_t il) { return hparams.n_embd_v_gqa(il); }, hparams.n_layer).c_str());
@@ -1676,9 +1681,9 @@ static void llm_prepare_mla(llama_model & model, int mla) {
         // is no MLA + FA for prompt processing, and MLA + FA for token generation, it would be useful
         // to change the MLA setting on the fly, depending on context. In that case, having prepared
         // the MLA tensors here is the right ting to do^TM.
-        const uint32_t n_embd_head_qk_nope = hparams.n_embd_head_k - hparams.n_rot;
+        const uint32_t n_embd_head_qk_nope = hparams.n_embd_head_k(0) - hparams.n_rot;
         const uint32_t kv_lora_rank = hparams.n_lora_kv;
-        const int32_t n_embd_head_v = hparams.n_embd_head_v;
+        const int32_t n_embd_head_v = hparams.n_embd_head_v(0);
         const int32_t n_head        = hparams.n_head(0);
         std::vector work_data;
         LLAMA_LOG_INFO("============ %s: need to compute %d wk_b/wv_b tensors\n", __func__, n_to_compute);
@@ -5489,8 +5494,8 @@ struct llama_context * llama_init_from_model(
         type_v = GGML_TYPE_F32; // required by ggml_ssm_scan for Mamba's ssm_states
     }
 
-    GGML_ASSERT(hparams.n_embd_head_k % ggml_blck_size(type_k) == 0);
-    GGML_ASSERT(hparams.n_embd_head_v % ggml_blck_size(type_v) == 0);
+    GGML_ASSERT(hparams.n_embd_head_k(0) % ggml_blck_size(type_k) == 0);
+    GGML_ASSERT(hparams.n_embd_head_v(0) % ggml_blck_size(type_v) == 0);
     if (!hparams.vocab_only) {
         // initialize backends
 #if defined(GGML_USE_METAL)
@@ -5939,6 +5944,7 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) {
         case LLM_ARCH_MIMO2:
         case LLM_ARCH_SEED_OSS:
         case LLM_ARCH_STEP35:
+        case LLM_ARCH_GEMMA4:
             return LLAMA_ROPE_TYPE_NEOX;
 
         case LLM_ARCH_QWEN2VL:
diff --git a/src/unicode.cpp b/src/unicode.cpp
index 19d8d409..50654016 100644
--- a/src/unicode.cpp
+++ b/src/unicode.cpp
@@ -1,16 +1,10 @@
-#if defined(_MSC_VER)
-#define _SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING
-#endif
-
 #include "unicode.h"
 #include "unicode-data.h"
 
 #include 
 #include 
-#include 
 #include 
 #include 
-#include 
 #include 
 #include 
 #include 
@@ -199,27 +193,6 @@ static std::unordered_map unicode_utf8_to_byte_map() {
     return map;
 }
 
-static inline std::wstring unicode_wstring_from_utf8(const std::string & s) {
-#if defined(__clang__)
-    // disable C++17 deprecation warning for std::codecvt_utf8
-#    pragma clang diagnostic push
-#    pragma clang diagnostic ignored "-Wdeprecated-declarations"
-#elif defined(__GNUC__)
-#    pragma GCC diagnostic push
-#    pragma GCC diagnostic ignored "-Wdeprecated-declarations"
-#endif
-
-    std::wstring_convert> conv;
-
-#if defined(__clang__)
-#    pragma clang diagnostic pop
-#elif defined(__GNUC__)
-#    pragma GCC diagnostic pop
-#endif
-
-    return conv.from_bytes(s);
-}
-
 static std::vector unicode_byte_encoding_process(const std::vector & bpe_words) {
     std::vector bpe_encoded_words;
     for (const auto & word : bpe_words) {
@@ -497,49 +470,26 @@ static std::vector unicode_regex_split_custom_llama3(const std::string &
     return bpe_offsets;
 }
 
-// use std::wregex to split the text
-static std::vector unicode_regex_split_stl(const std::wstring & wtext, const std::wstring & regex_expr, const std::vector & offsets) {
-    std::wregex expr(regex_expr);
+template 
+static std::vector unicode_regex_split_stl(const std::basic_string & text, const std::basic_string & regex, const std::vector & offsets) {
+    using BidirIt = typename std::basic_string::const_iterator;
+#ifdef _MSC_VER
+    // Bypass bug in MSVC: https://github.com/ggml-org/llama.cpp/issues/17830
+    constexpr auto regex_flags = std::regex_constants::ECMAScript;
+#else
+    constexpr auto regex_flags = std::regex_constants::optimize | std::regex_constants::nosubs;
+#endif
+    std::basic_regex expr(regex, regex_flags);
     std::vector bpe_offsets; // store the offset of each word
     bpe_offsets.reserve(offsets.size()); // Reserve memory for the approximate size
     size_t start = 0;
     for (auto offset : offsets) {
-        std::wcregex_iterator it(wtext.data() + start, wtext.data() + start + offset, expr);
-        std::wcregex_iterator end;
+        std::regex_iterator it(text.begin() + start, text.begin() + start + offset, expr);
+        std::regex_iterator end;
 
         int64_t start_idx = 0;
         while (it != end) {
-            std::wcmatch match = *it;
-            if (match.position() > start_idx) {
-                bpe_offsets.emplace_back(match.position() - start_idx);
-            }
-            bpe_offsets.emplace_back(match.length());
-            start_idx = match.position() + match.length();
-            ++it;
-        }
-
-        if (start_idx < (int64_t) offset) {
-            bpe_offsets.emplace_back(offset - start_idx);
-        }
-        start += offset;
-    }
-
-    return bpe_offsets;
-}
-
-// use std::regex to split the text
-static std::vector unicode_regex_split_stl(const std::string & text, const std::string & regex_expr, const std::vector & offsets) {
-    std::regex expr(regex_expr);
-    std::vector bpe_offsets; // store the offset of each word
-    bpe_offsets.reserve(offsets.size()); // Reserve memory for the approximate size
-    size_t start = 0;
-    for (auto offset : offsets) {
-        std::cregex_iterator it(text.data() + start, text.data() + start + offset, expr);
-        std::cregex_iterator end;
-
-        int64_t start_idx = 0;
-        while (it != end) {
-            std::cmatch match = *it;
+            std::match_results match = *it;
             if (match.position() > start_idx) {
                 bpe_offsets.emplace_back(match.position() - start_idx);
             }
@@ -729,6 +679,109 @@ static std::vector unicode_regex_split_custom_kimi_k2(const std::string
     return bpe_offsets;
 }
 
+// AFMOE digit handling: splits digits with leading 1-2 based on total length modulo 3
+static std::vector unicode_regex_split_custom_afmoe(const std::string & text, const std::vector & offsets) {
+    std::vector bpe_offsets;
+    bpe_offsets.reserve(offsets.size());
+
+    const auto cpts = unicode_cpts_from_utf8(text);
+
+    size_t start = 0;
+    for (auto offset : offsets) {
+        const size_t offset_ini = start;
+        const size_t offset_end = start + offset;
+        assert(offset_end <= cpts.size());
+        start = offset_end;
+
+        auto _get_flags = [&] (const size_t pos) -> unicode_cpt_flags {
+            return (offset_ini <= pos && pos < offset_end) ? unicode_cpt_flags_from_cpt(cpts[pos]) : unicode_cpt_flags{};
+        };
+
+        size_t _prev_end = offset_ini;
+        auto _add_token = [&] (const size_t end) -> size_t {
+            assert(_prev_end <= end && end <= offset_end);
+            size_t len = end - _prev_end;
+            if (len > 0) {
+                bpe_offsets.push_back(len);
+            }
+            _prev_end = end;
+            return len;
+        };
+
+        for (size_t pos = offset_ini; pos < offset_end; ) {
+            const auto flags = _get_flags(pos);
+
+            // Handle digit sequences with special splitting logic
+            if (flags.is_number) {
+                size_t digit_start = pos;
+                size_t digit_count = 0;
+
+                // Count consecutive digits
+                while (_get_flags(pos).is_number && pos < offset_end) {
+                    digit_count++;
+                    pos++;
+                }
+
+                // Split based on total length modulo 3
+                size_t remainder = digit_count % 3;
+                size_t current = digit_start;
+
+                // Emit leading 1-2 digits if needed
+                if (remainder > 0) {
+                    _add_token(current + remainder);
+                    current += remainder;
+                }
+
+                // Emit groups of 3
+                while (current < digit_start + digit_count) {
+                    _add_token(current + 3);
+                    current += 3;
+                }
+                continue;
+            }
+
+            // For non-digits, just move forward
+            pos++;
+        }
+
+        // Add any remaining content
+        if (_prev_end < offset_end) {
+            _add_token(offset_end);
+        }
+    }
+
+    return bpe_offsets;
+}
+
+// regex: [^\n]+|[\n]+
+// splits text into runs of non-newline characters and runs of newline characters
+static std::vector unicode_regex_split_custom_newlines(const std::string & text, const std::vector & offsets) {
+    std::vector bpe_offsets;
+    bpe_offsets.reserve(offsets.size());
+
+    const auto cpts = unicode_cpts_from_utf8(text);
+
+    size_t start = 0;
+    for (auto offset : offsets) {
+        const size_t offset_ini = start;
+        const size_t offset_end = start + offset;
+        assert(offset_end <= cpts.size());
+        start = offset_end;
+
+        size_t pos = offset_ini;
+        while (pos < offset_end) {
+            const bool is_newline = (cpts[pos] == '\n');
+            const size_t run_start = pos;
+            while (pos < offset_end && (cpts[pos] == '\n') == is_newline) {
+                pos++;
+            }
+            bpe_offsets.push_back(pos - run_start);
+        }
+    }
+
+    return bpe_offsets;
+}
+
 static std::vector unicode_regex_split_custom(const std::string & text, const std::string & regex_expr, const std::vector & offsets) {
     std::vector bpe_offsets;
 
@@ -742,6 +795,17 @@ static std::vector unicode_regex_split_custom(const std::string & text,
     } else if (regex_expr == "\\p{Han}+") {
         // K2's first pattern - handle all K2 patterns together
         bpe_offsets = unicode_regex_split_custom_kimi_k2(text, offsets);
+    } else if (regex_expr == "\\p{AFMoE_digits}") {
+        // AFMOE digit pattern - use custom implementation for proper splitting
+        bpe_offsets = unicode_regex_split_custom_afmoe(text, offsets);
+    } else if (regex_expr == "[^\\n]+|[\\n]+") {
+        bpe_offsets = unicode_regex_split_custom_newlines(text, offsets);
+    } else if (regex_expr == "\\d{1,3}(?=(?:\\d{3})*\\b)") {
+        // tiny_aya digit grouping pattern from tokenizer.json:
+        //   {"type": "Split", "pattern": {"Regex": "\\d{1,3}(?=(?:\\d{3})*\\b)"}, "behavior": "Isolated"}
+        // Splits digits into groups of 3 from the right (e.g., 1234567 -> 1, 234, 567)
+        // TODO: Revisit this regex, in case there are any subtle tokenization differences with the original regex.
+        bpe_offsets = unicode_regex_split_custom_afmoe(text, offsets);
     }
 
     return bpe_offsets;
@@ -802,7 +866,7 @@ std::vector unicode_cpts_from_utf8(const std::string & utf8) {
             result.push_back(unicode_cpt_from_utf8(utf8, offset));
         }
         catch (const std::invalid_argument & /*ex*/) {
-            // Silently ignore invalid UTF-8 input to avoid leaking the exception beyond common_tokenize
+            // Silently ignore invalid UTF-8 input to avoid leaking the exception beyond llama_tokenize
             ++offset;
             result.emplace_back(0xFFFD); // replacement character
         }
@@ -879,7 +943,7 @@ bool unicode_cpt_is_han(uint32_t cpt) {
     return false;
 }
 
-std::vector unicode_regex_split(const std::string & text, const std::vector & regex_exprs) {
+std::vector unicode_regex_split(const std::string & text, const std::vector & regex_exprs, bool byte_encode) {
     // unicode categories
     static const std::map k_ucat_enum = {
         { "\\p{N}", unicode_cpt_flags::NUMBER },
@@ -887,6 +951,11 @@ std::vector unicode_regex_split(const std::string & text, const std
         { "\\p{P}", unicode_cpt_flags::PUNCTUATION },
         { "\\p{M}", unicode_cpt_flags::ACCENT_MARK },
         { "\\p{S}", unicode_cpt_flags::SYMBOL },
+        { "\\p{Lu}", unicode_cpt_flags::LETTER }, // Uppercase letter
+        { "\\p{Ll}", unicode_cpt_flags::LETTER }, // Lowercase letter
+        { "\\p{Lt}", unicode_cpt_flags::LETTER }, // Titlecase letter
+        { "\\p{Lm}", unicode_cpt_flags::LETTER }, // Modifier letter
+        { "\\p{Lo}", unicode_cpt_flags::LETTER }, // Other letter
     };
 
     static const std::map k_ucat_cpt = {
@@ -969,10 +1038,10 @@ std::vector unicode_regex_split(const std::string & text, const std
                     break;
                 }
             }
+            const auto cpts_regex = unicode_cpts_from_utf8(regex_expr);
 
             if (use_collapsed) {
                 // sanity-check that the original regex does not contain any non-ASCII characters
-                const auto cpts_regex = unicode_cpts_from_utf8(regex_expr);
                 for (size_t i = 0; i < cpts_regex.size(); ++i) {
                     if (cpts_regex[i] >= 128) {
                         throw std::runtime_error("Regex includes both unicode categories and non-ASCII characters - not supported");
@@ -997,22 +1066,26 @@ std::vector unicode_regex_split(const std::string & text, const std
                         continue;
                     }
 
-                    if (regex_expr[i + 0] == '\\' && i + 4 < regex_expr.size() &&
+                    // Match \p{...} Unicode properties of varying lengths
+                    if (regex_expr[i + 0] == '\\' && i + 3 < regex_expr.size() &&
                         regex_expr[i + 1] == 'p' &&
-                        regex_expr[i + 2] == '{' &&
-                        regex_expr[i + 4] == '}') {
-                        const std::string pat = regex_expr.substr(i, 5);
-                        if (k_ucat_enum.find(pat) != k_ucat_enum.end()) {
-                            if (!inside) {
-                                regex_expr_collapsed += '[';
+                        regex_expr[i + 2] == '{') {
+                        // Find the closing brace
+                        size_t closing_brace = regex_expr.find('}', i + 3);
+                        if (closing_brace != std::string::npos && closing_brace <= i + 10) { // reasonable limit
+                            const std::string pat = regex_expr.substr(i, closing_brace - i + 1);
+                            if (k_ucat_enum.find(pat) != k_ucat_enum.end()) {
+                                if (!inside) {
+                                    regex_expr_collapsed += '[';
+                                }
+                                regex_expr_collapsed += k_ucat_cpt.at(k_ucat_enum.at(pat));
+                                regex_expr_collapsed += k_ucat_map.at(k_ucat_enum.at(pat));
+                                if (!inside) {
+                                    regex_expr_collapsed += ']';
+                                }
+                                i = closing_brace;
+                                continue;
                             }
-                            regex_expr_collapsed += k_ucat_cpt.at(k_ucat_enum.at(pat));
-                            regex_expr_collapsed += k_ucat_map.at(k_ucat_enum.at(pat));
-                            if (!inside) {
-                                regex_expr_collapsed += ']';
-                            }
-                            i += 4;
-                            continue;
                         }
                     }
 
@@ -1024,7 +1097,7 @@ std::vector unicode_regex_split(const std::string & text, const std
                 bpe_offsets = unicode_regex_split_stl(text_collapsed, regex_expr_collapsed, bpe_offsets);
             } else {
                 // no unicode category used, we can use std::wregex directly
-                const std::wstring wregex_expr = unicode_wstring_from_utf8(regex_expr);
+                std::wstring wregex_expr(cpts_regex.begin(), cpts_regex.end());
 
                 // std::wregex \s does not mach non-ASCII whitespaces, using 0x0B as fallback
                 std::wstring wtext(cpts.begin(), cpts.end());
@@ -1057,5 +1130,9 @@ std::vector unicode_regex_split(const std::string & text, const std
         start += offset;
     }
 
-    return unicode_byte_encoding_process(bpe_words);
+    if (byte_encode) {
+        return unicode_byte_encoding_process(bpe_words);
+    }
+
+    return bpe_words;
 }
diff --git a/src/unicode.h b/src/unicode.h
index 0a5fa2a7..1d702d0a 100644
--- a/src/unicode.h
+++ b/src/unicode.h
@@ -15,6 +15,10 @@ struct unicode_cpt_flags {
         SYMBOL          = 0x0040,  // regex: \p{S}
         CONTROL         = 0x0080,  // regex: \p{C}
         MASK_CATEGORIES = 0x00FF,
+        WHITESPACE      = 0x0100,
+        LOWERCASE       = 0x0200,
+        UPPERCASE       = 0x0400,
+        NFD             = 0x0800,
     };
 
     // codepoint type
@@ -34,11 +38,49 @@ struct unicode_cpt_flags {
 
     // decode from uint16
     inline unicode_cpt_flags(const uint16_t flags = 0) {
+#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
         *reinterpret_cast(this) = flags;
+#elif __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+        is_undefined   = (flags & UNDEFINED)   ? 1 : 0;
+        is_number      = (flags & NUMBER)      ? 1 : 0;
+        is_letter      = (flags & LETTER)      ? 1 : 0;
+        is_separator   = (flags & SEPARATOR)   ? 1 : 0;
+        is_accent_mark = (flags & ACCENT_MARK) ? 1 : 0;
+        is_punctuation = (flags & PUNCTUATION) ? 1 : 0;
+        is_symbol      = (flags & SYMBOL)      ? 1 : 0;
+        is_control     = (flags & CONTROL)     ? 1 : 0;
+        is_whitespace  = (flags & WHITESPACE)  ? 1 : 0;
+        is_lowercase   = (flags & LOWERCASE)   ? 1 : 0;
+        is_uppercase   = (flags & UPPERCASE)   ? 1 : 0;
+        is_nfd         = (flags & NFD)         ? 1 : 0;
+#else
+#error Unexpected or undefined __BYTE_ORDER__
+#endif
     }
 
     inline uint16_t as_uint() const {
+#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
         return *reinterpret_cast(this);
+#elif __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+        uint16_t result =
+              is_undefined   * UNDEFINED
+            + is_number      * NUMBER
+            + is_letter      * LETTER
+            + is_separator   * SEPARATOR
+            + is_accent_mark * ACCENT_MARK
+            + is_punctuation * PUNCTUATION
+            + is_symbol      * SYMBOL
+            + is_control     * CONTROL
+            + is_whitespace  * WHITESPACE
+            + is_lowercase   * LOWERCASE
+            + is_uppercase   * UPPERCASE
+            + is_nfd         * NFD
+            ;
+
+        return result;
+#else
+#error Unexpected or undefined __BYTE_ORDER__
+#endif
     }
 
     inline uint16_t category_flag() const {
@@ -65,4 +107,4 @@ uint32_t unicode_tolower(uint32_t cpt);
 
 bool unicode_cpt_is_han(uint32_t cpt);
 
-std::vector unicode_regex_split(const std::string & text, const std::vector & regex_exprs);
+std::vector unicode_regex_split(const std::string & text, const std::vector & regex_exprs, bool byte_encode = true);