From 4834a2327d008ace3ec5a9ed00f51454bcabbc1c Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 18 Aug 2026 11:32:12 +0300 Subject: [PATCH] talk-llama : sync llama.cpp --- examples/talk-llama/llama-adapter.cpp | 7 +- examples/talk-llama/llama-arch.cpp | 31 +- examples/talk-llama/llama-arch.h | 16 + examples/talk-llama/llama-context.cpp | 7 +- examples/talk-llama/llama-graph.cpp | 17 + examples/talk-llama/llama-graph.h | 1 + examples/talk-llama/llama-hparams.cpp | 7 + examples/talk-llama/llama-hparams.h | 14 +- examples/talk-llama/llama-model-loader.cpp | 26 +- examples/talk-llama/llama-model-saver.cpp | 15 +- examples/talk-llama/llama-model.cpp | 23 +- examples/talk-llama/llama-model.h | 14 + examples/talk-llama/llama-quant.cpp | 7 +- examples/talk-llama/llama-vocab.cpp | 49 +- examples/talk-llama/models/bailingmoe3.cpp | 532 ++++++++++++++++++ examples/talk-llama/models/deepseek32.cpp | 51 +- examples/talk-llama/models/dflash.cpp | 57 +- examples/talk-llama/models/glm-dsa.cpp | 51 +- examples/talk-llama/models/kimi-k3.cpp | 614 +++++++++++++++++++++ examples/talk-llama/models/minimax-01.cpp | 520 +++++++++++++++++ examples/talk-llama/models/minimax-m3.cpp | 2 + examples/talk-llama/models/models.h | 68 +++ 22 files changed, 2017 insertions(+), 112 deletions(-) create mode 100644 examples/talk-llama/models/bailingmoe3.cpp create mode 100644 examples/talk-llama/models/kimi-k3.cpp create mode 100644 examples/talk-llama/models/minimax-01.cpp diff --git a/examples/talk-llama/llama-adapter.cpp b/examples/talk-llama/llama-adapter.cpp index 3e0fe66af..e6678a66d 100644 --- a/examples/talk-llama/llama-adapter.cpp +++ b/examples/talk-llama/llama-adapter.cpp @@ -396,8 +396,11 @@ static void llama_adapter_lora_init_impl(llama_model & model, const char * path_ llama_file gguf_file(path_lora, "rb"); std::vector read_buf; auto set_tensor = [&](ggml_tensor * orig, ggml_tensor * dev) { - size_t offs = gguf_get_data_offset(ctx_gguf.get()) + gguf_get_tensor_offset(ctx_gguf.get(), gguf_find_tensor(ctx_gguf.get(), orig->name)); - size_t size = ggml_nbytes(orig); + const size_t offs = gguf_get_data_offset(ctx_gguf.get()) + gguf_get_tensor_offset(ctx_gguf.get(), gguf_find_tensor(ctx_gguf.get(), orig->name)); + const size_t size = ggml_nbytes(orig); + if (offs + size < offs || offs + size > gguf_file.size()) { + throw std::runtime_error(format("LoRA tensor '%s' data is not within the file bounds, file is corrupted or incomplete", orig->name)); + } read_buf.resize(size); gguf_file.seek(offs, SEEK_SET); gguf_file.read_raw(read_buf.data(), size); diff --git a/examples/talk-llama/llama-arch.cpp b/examples/talk-llama/llama-arch.cpp index 292ab2610..5b88bde14 100644 --- a/examples/talk-llama/llama-arch.cpp +++ b/examples/talk-llama/llama-arch.cpp @@ -107,6 +107,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_PLM, "plm" }, { LLM_ARCH_BAILINGMOE, "bailingmoe" }, { LLM_ARCH_BAILINGMOE2, "bailingmoe2" }, + { LLM_ARCH_BAILINGMOE3, "bailingmoe3" }, { LLM_ARCH_DOTS1, "dots1" }, { LLM_ARCH_ARCEE, "arcee" }, { LLM_ARCH_AFMOE, "afmoe" }, @@ -128,6 +129,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_SEED_OSS, "seed_oss" }, { LLM_ARCH_GROVEMOE, "grovemoe" }, { LLM_ARCH_APERTUS, "apertus" }, + { LLM_ARCH_MINIMAX_01, "minimax-01" }, { LLM_ARCH_MINIMAX_M2, "minimax-m2" }, { LLM_ARCH_MINIMAX_M3, "minimax-m3" }, { LLM_ARCH_COGVLM, "cogvlm" }, @@ -143,6 +145,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_LLAMA_EMBED, "llama-embed" }, { LLM_ARCH_MAINCODER, "maincoder" }, { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, + { LLM_ARCH_KIMI_K3, "kimi-k3" }, { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_NANBEIGE, "nanbeige" }, @@ -186,6 +189,9 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_FEATURES_LENGTH, "%s.features_length" }, { LLM_KV_BLOCK_COUNT, "%s.block_count" }, { LLM_KV_LEADING_DENSE_BLOCK_COUNT, "%s.leading_dense_block_count" }, + { LLM_KV_ATTN_RES_BLOCK_SIZE, "%s.attn_res.block_size" }, + { LLM_KV_ACTIVATION_SITU_BETA, "%s.activation.situ_beta" }, + { LLM_KV_ACTIVATION_SITU_LINEAR_BETA, "%s.activation.situ_linear_beta" }, { LLM_KV_FEED_FORWARD_LENGTH, "%s.feed_forward_length" }, { LLM_KV_EXPERT_FEED_FORWARD_LENGTH, "%s.expert_feed_forward_length" }, { LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, "%s.expert_shared_feed_forward_length" }, @@ -201,6 +207,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_EXPERT_GROUP_USED_COUNT, "%s.expert_group_used_count" }, { LLM_KV_EXPERT_WEIGHTS_SCALE, "%s.expert_weights_scale" }, { LLM_KV_EXPERT_WEIGHTS_NORM, "%s.expert_weights_norm" }, + { LLM_KV_EXPERT_LATENT_LENGTH, "%s.expert_latent_length" }, { LLM_KV_EXPERT_GATING_FUNC, "%s.expert_gating_func" }, { LLM_KV_EXPERT_GROUP_SCALE, "%s.expert_group_scale" }, { LLM_KV_EXPERTS_PER_GROUP, "%s.experts_per_group" }, @@ -311,7 +318,9 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_SSM_GROUP_COUNT, "%s.ssm.group_count" }, { LLM_KV_SSM_DT_B_C_RMS, "%s.ssm.dt_b_c_rms" }, - { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" }, + { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" }, + { LLM_KV_KDA_SAFE_GATE, "%s.kda.safe_gate" }, + { LLM_KV_KDA_GATE_LOWER_BOUND, "%s.kda.gate_lower_bound" }, { LLM_KV_WKV_HEAD_SIZE, "%s.wkv.head_size" }, @@ -462,6 +471,13 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_SSM_F_B, "blk.%d.ssm_f_b" }, { LLM_TENSOR_SSM_BETA, "blk.%d.ssm_beta" }, { LLM_TENSOR_SSM_G_A, "blk.%d.ssm_g_a" }, + { LLM_TENSOR_SSM_G, "blk.%d.ssm_g" }, + { LLM_TENSOR_ATTN_RES_SCORE, "blk.%d.attn_res_score" }, + { LLM_TENSOR_FFN_RES_SCORE, "blk.%d.ffn_res_score" }, + { LLM_TENSOR_OUTPUT_RES_SCORE, "output_res_score" }, + { LLM_TENSOR_FFN_ROUTED_DOWN, "blk.%d.ffn_routed_down" }, + { LLM_TENSOR_FFN_ROUTED_UP, "blk.%d.ffn_routed_up" }, + { LLM_TENSOR_FFN_ROUTED_NORM, "blk.%d.ffn_routed_norm" }, { LLM_TENSOR_SSM_G_B, "blk.%d.ssm_g_b" }, { LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" }, { LLM_TENSOR_ATTN_Q_A_NORM, "blk.%d.attn_q_a_norm" }, @@ -755,6 +771,13 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_SSM_F_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_BETA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_G_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_SSM_G, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_RES_SCORE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_RES_SCORE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_OUTPUT_RES_SCORE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_ROUTED_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_ROUTED_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_ROUTED_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_SSM_G_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_TIME_MIX_LERP_X, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_TIME_MIX_LN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, @@ -975,9 +998,12 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_BAILINGMOE3: + case LLM_ARCH_KIMI_K3: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_MINIMAX_01: return true; default: return false; @@ -1033,10 +1059,13 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_GRANITE_HYBRID: case LLM_ARCH_LFM2: case LLM_ARCH_LFM2MOE: + case LLM_ARCH_MINIMAX_01: case LLM_ARCH_MINIMAX_M2: case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_BAILINGMOE3: + case LLM_ARCH_KIMI_K3: case LLM_ARCH_QWEN3TTS: return false; default: diff --git a/examples/talk-llama/llama-arch.h b/examples/talk-llama/llama-arch.h index 18d9de186..8042120a2 100644 --- a/examples/talk-llama/llama-arch.h +++ b/examples/talk-llama/llama-arch.h @@ -112,6 +112,7 @@ enum llm_arch { LLM_ARCH_PLM, LLM_ARCH_BAILINGMOE, LLM_ARCH_BAILINGMOE2, + LLM_ARCH_BAILINGMOE3, LLM_ARCH_DOTS1, LLM_ARCH_ARCEE, LLM_ARCH_AFMOE, @@ -145,6 +146,7 @@ enum llm_arch { LLM_ARCH_LLAMA_EMBED, LLM_ARCH_MAINCODER, LLM_ARCH_KIMI_LINEAR, + LLM_ARCH_KIMI_K3, LLM_ARCH_TALKIE, LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, @@ -153,6 +155,7 @@ enum llm_arch { LLM_ARCH_NANBEIGE, LLM_ARCH_QWEN3TTS, LLM_ARCH_POCKETTTS, + LLM_ARCH_MINIMAX_01, LLM_ARCH_UNKNOWN, }; @@ -191,6 +194,9 @@ enum llm_kv { LLM_KV_FEATURES_LENGTH, LLM_KV_BLOCK_COUNT, LLM_KV_LEADING_DENSE_BLOCK_COUNT, + LLM_KV_ATTN_RES_BLOCK_SIZE, + LLM_KV_ACTIVATION_SITU_BETA, + LLM_KV_ACTIVATION_SITU_LINEAR_BETA, LLM_KV_FEED_FORWARD_LENGTH, LLM_KV_EXPERT_FEED_FORWARD_LENGTH, LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, @@ -206,6 +212,7 @@ enum llm_kv { LLM_KV_EXPERT_GROUP_USED_COUNT, LLM_KV_EXPERT_WEIGHTS_SCALE, LLM_KV_EXPERT_WEIGHTS_NORM, + LLM_KV_EXPERT_LATENT_LENGTH, LLM_KV_EXPERT_GATING_FUNC, LLM_KV_EXPERT_GROUP_SCALE, LLM_KV_EXPERTS_PER_GROUP, @@ -317,6 +324,8 @@ enum llm_kv { LLM_KV_SSM_DT_B_C_RMS, LLM_KV_KDA_HEAD_DIM, + LLM_KV_KDA_SAFE_GATE, + LLM_KV_KDA_GATE_LOWER_BOUND, LLM_KV_WKV_HEAD_SIZE, @@ -491,6 +500,13 @@ enum llm_tensor { LLM_TENSOR_SSM_BETA, // kimi: beta mixing coefficient and qwen3.5 LLM_TENSOR_SSM_G_A, // kimi: output gate projection A LLM_TENSOR_SSM_G_B, // kimi: output gate projection B + LLM_TENSOR_SSM_G, // kimi-k3: full-rank KDA gate + LLM_TENSOR_ATTN_RES_SCORE, // kimi-k3: fused res_norm*res_proj (pre-attn) + LLM_TENSOR_FFN_RES_SCORE, // kimi-k3: fused res_norm*res_proj (pre-ffn) + LLM_TENSOR_OUTPUT_RES_SCORE, // kimi-k3: fused res_norm*res_proj (final) + LLM_TENSOR_FFN_ROUTED_DOWN, // kimi-k3: latent MoE down + LLM_TENSOR_FFN_ROUTED_UP, // kimi-k3: latent MoE up + LLM_TENSOR_FFN_ROUTED_NORM, // kimi-k3: latent MoE norm LLM_TENSOR_TIME_MIX_W0, LLM_TENSOR_TIME_MIX_W1, LLM_TENSOR_TIME_MIX_W2, diff --git a/examples/talk-llama/llama-context.cpp b/examples/talk-llama/llama-context.cpp index cd013cdb1..52f8d5367 100644 --- a/examples/talk-llama/llama-context.cpp +++ b/examples/talk-llama/llama-context.cpp @@ -2293,13 +2293,18 @@ void llama_context::output_reorder() { uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { uint32_t res; - if (model.arch == LLM_ARCH_QWEN3NEXT || + if (model.arch == LLM_ARCH_KIMI_K3) { + // the n_tokens*40 budget below is exhausted at ubatch 3840 + res = std::max(n_tokens * 160, 64u * model.n_tensors()); + } else if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || + model.arch == LLM_ARCH_BAILINGMOE3 || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4 || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || + model.arch == LLM_ARCH_MINIMAX_01 || model.arch == LLM_ARCH_MINIMAX_M3) { res = std::max(n_tokens * 40, 32u * model.n_tensors()); } else { diff --git a/examples/talk-llama/llama-graph.cpp b/examples/talk-llama/llama-graph.cpp index 55d858024..1896758c5 100644 --- a/examples/talk-llama/llama-graph.cpp +++ b/examples/talk-llama/llama-graph.cpp @@ -1835,6 +1835,8 @@ ggml_tensor * llm_graph_context::build_ffn( cur = ggml_reglu(ctx0, cur); cb(cur, "ffn_reglu", il); } break; + case LLM_FFN_SITU: + GGML_ABORT("not yet supported"); default: GGML_ABORT("fatal error"); } @@ -2174,6 +2176,21 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cur = ggml_silu(ctx0, cur); cb(cur, "ffn_moe_silu", il); } break; + case LLM_FFN_SITU: + { + // situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * lb*tanh(up/lb) + GGML_ASSERT(has_gate); + const float beta = hparams.situ_beta; + const float lb = hparams.situ_linear_beta; + + ggml_tensor * act = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, cur, 1.0f/beta)), beta); + act = ggml_mul(ctx0, act, ggml_sigmoid(ctx0, cur)); + if (lb > 0.0f) { + up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/lb)), lb); + } + cur = ggml_mul(ctx0, act, up); + cb(cur, "ffn_moe_situ", il); + } break; case LLM_FFN_GELU: if (has_gate) { cur = ggml_geglu_split(ctx0, cur, up); diff --git a/examples/talk-llama/llama-graph.h b/examples/talk-llama/llama-graph.h index 75bc0fe80..94324c745 100644 --- a/examples/talk-llama/llama-graph.h +++ b/examples/talk-llama/llama-graph.h @@ -59,6 +59,7 @@ enum llm_ffn_op_type : int { LLM_FFN_GEGLU, LLM_FFN_REGLU, LLM_FFN_SWIGLU_OAI_MOE, + LLM_FFN_SITU, // kimi-k3 }; enum llm_ffn_gate_type { diff --git a/examples/talk-llama/llama-hparams.cpp b/examples/talk-llama/llama-hparams.cpp index 781277f3f..e3f0cf0ed 100644 --- a/examples/talk-llama/llama-hparams.cpp +++ b/examples/talk-llama/llama-hparams.cpp @@ -217,6 +217,13 @@ uint32_t llama_hparams::n_embd_s() const { return n_embd_head_kda * n_embd_head_kda * n_head(); // 128 * 128 * 32 = 524288 } + if (n_embd_head_la != 0) { + // for MiniMax-Text-01 linear attention layers + // Full recurrent state: head_dim * head_dim * n_head + // tensor shape for linear attention: [head_dim, head_dim, n_head] + return n_embd_head_la * n_embd_head_la * n_head(); // 128 * 128 * 64 = 1048576 + } + // corresponds to Mamba's ssm_states size return ssm_d_state * ssm_d_inner; } diff --git a/examples/talk-llama/llama-hparams.h b/examples/talk-llama/llama-hparams.h index 57de80824..e91ce1cc3 100644 --- a/examples/talk-llama/llama-hparams.h +++ b/examples/talk-llama/llama-hparams.h @@ -4,10 +4,11 @@ #include #include +#include // bump if necessary #define LLAMA_MAX_LAYERS 512 -#define LLAMA_MAX_EXPERTS 512 // Qwen3 Next +#define LLAMA_MAX_EXPERTS 1024 // Kimi K3 enum llama_expert_gating_func_type { LLAMA_EXPERT_GATING_FUNC_TYPE_NONE = 0, @@ -164,8 +165,19 @@ struct llama_hparams { uint32_t ssm_dt_rank = 0; uint32_t ssm_n_group = 0; + // for MiniMax-Text-01 linear attention + uint32_t n_embd_head_la = 0; + // for Kimi Linear KDA uint32_t n_embd_head_kda = 0; + bool kda_safe_gate = false; + + // kimi-k3 + uint32_t n_expert_latent = 0; // routed_expert_hidden_size (0 = experts run at n_embd) + uint32_t attn_res_block_size = 0; // 0 = no cross-layer attention residuals + float kda_gate_lower_bound = -INFINITY; + float situ_beta = 1.0f; + float situ_linear_beta = 0.0f; // 0 = no linear-beta transform on the up branch bool ssm_dt_b_c_rms = false; diff --git a/examples/talk-llama/llama-model-loader.cpp b/examples/talk-llama/llama-model-loader.cpp index 5c5e97fbc..1ca698704 100644 --- a/examples/talk-llama/llama-model-loader.cpp +++ b/examples/talk-llama/llama-model-loader.cpp @@ -316,15 +316,19 @@ namespace GGUFMeta { struct GGUFMeta::ArrayInfo arr_info = GGUFMeta::GKV::get_kv(ctx, kid); + bool type_ok = false; switch (arr_info.gt) { case GGUF_TYPE_UINT32: - case GGUF_TYPE_INT32: GGML_ASSERT((std::is_same::value) || - (std::is_same::value)); break; - case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same::value)); break; - case GGUF_TYPE_STRING: GGML_ASSERT((std::is_same::value)); break; + case GGUF_TYPE_INT32: type_ok = (std::is_same::value) || + (std::is_same::value); break; + case GGUF_TYPE_FLOAT32: type_ok = (std::is_same::value); break; + case GGUF_TYPE_STRING: type_ok = (std::is_same::value); break; default: throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str())); } + if (!type_ok) { + throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt))); + } if constexpr (std::is_same::value) { const size_t n_items = gguf_get_arr_n(ctx, kid); @@ -357,16 +361,20 @@ namespace GGUFMeta { struct GGUFMeta::ArrayInfo arr_info = GGUFMeta::GKV::get_kv(ctx, kid); + bool type_ok = false; switch (arr_info.gt) { case GGUF_TYPE_BOOL: case GGUF_TYPE_UINT32: - case GGUF_TYPE_INT32: GGML_ASSERT((std::is_same::value) || - (std::is_same::value)); break; - case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same::value)); break; - case GGUF_TYPE_STRING: GGML_ASSERT((std::is_same::value)); break; + case GGUF_TYPE_INT32: type_ok = (std::is_same::value) || + (std::is_same::value); break; + case GGUF_TYPE_FLOAT32: type_ok = (std::is_same::value); break; + case GGUF_TYPE_STRING: type_ok = (std::is_same::value); break; default: throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str())); } + if (!type_ok) { + throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt))); + } if (arr_info.length > N_MAX) { throw std::runtime_error(format("array length %u for key %s exceeds max %u", (uint32_t) arr_info.length, key.c_str(), (uint32_t) N_MAX)); @@ -1178,7 +1186,7 @@ struct ggml_tensor * llama_model_loader::create_tensor( if (use_mmap) { static std::once_flag once; std::call_once(once, [] { - LLAMA_LOG_WARN("llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --no-mmap for better performance\n"); + LLAMA_LOG_WARN("llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --load-mode none for better performance\n"); }); } } else { diff --git a/examples/talk-llama/llama-model-saver.cpp b/examples/talk-llama/llama-model-saver.cpp index abca773a9..be9524d40 100644 --- a/examples/talk-llama/llama-model-saver.cpp +++ b/examples/talk-llama/llama-model-saver.cpp @@ -121,6 +121,7 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c } // instantiate for external usage: template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); +template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); void llama_model_saver::add_kv(const enum llm_kv key, const std::vector & value) { std::vector tmp(value.size()); @@ -213,10 +214,13 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); add_kv(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, true); add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + add_kv(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent); add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp); add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH, hparams.n_ff_chexp); - add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp); - add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp); + add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector( + hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.begin() + hparams.n_layer_all)); + add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector( + hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.begin() + hparams.n_layer_all)); add_kv(LLM_KV_USE_PARALLEL_RESIDUAL, hparams.use_par_res); // add_kv(LLM_KV_TENSOR_DATA_LAYOUT, ???); add_kv(LLM_KV_EXPERT_COUNT, hparams.n_expert); @@ -319,6 +323,8 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_SSM_DT_B_C_RMS, hparams.ssm_dt_b_c_rms); add_kv(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + add_kv(LLM_KV_KDA_SAFE_GATE, hparams.kda_safe_gate); + add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); add_kv(LLM_KV_WKV_HEAD_SIZE, hparams.wkv_head_size); @@ -376,6 +382,10 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_XIELU_BETA, hparams.xielu_beta); add_kv(LLM_KV_XIELU_EPS, hparams.xielu_eps); + add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size); + add_kv(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta); + add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta); + // deprecated // add_kv(LLM_KV_TOKENIZER_PREFIX_ID, ???); // add_kv(LLM_KV_TOKENIZER_SUFFIX_ID, ???); @@ -403,6 +413,7 @@ void llama_model_saver::add_tensors_from_model() { add_tensor(model->output_norm_enc); add_tensor(model->output_s); add_tensor(model->output_in_s); + add_tensor(model->output_res_score); add_tensor(model->cls); add_tensor(model->cls_b); add_tensor(model->cls_out); diff --git a/examples/talk-llama/llama-model.cpp b/examples/talk-llama/llama-model.cpp index c81005505..0d74a2135 100644 --- a/examples/talk-llama/llama-model.cpp +++ b/examples/talk-llama/llama-model.cpp @@ -256,6 +256,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_bailingmoe(params); case LLM_ARCH_BAILINGMOE2: return new llama_model_bailingmoe2(params); + case LLM_ARCH_BAILINGMOE3: + return new llama_model_bailingmoe3(params); case LLM_ARCH_SEED_OSS: return new llama_model_seed_oss(params); case LLM_ARCH_DOTS1: @@ -296,6 +298,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_grovemoe(params); case LLM_ARCH_APERTUS: return new llama_model_apertus(params); + case LLM_ARCH_MINIMAX_01: + return new llama_model_minimax_01(params); case LLM_ARCH_MINIMAX_M2: return new llama_model_minimax_m2(params); case LLM_ARCH_MINIMAX_M3: @@ -320,6 +324,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_mimo2(params); case LLM_ARCH_KIMI_LINEAR: return new llama_model_kimi_linear(params); + case LLM_ARCH_KIMI_K3: + return new llama_model_kimi_k3(params); case LLM_ARCH_STEP35: return new llama_model_step35(params); default: @@ -798,6 +804,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_290B: return "290B"; case LLM_TYPE_314B: return "314B"; case LLM_TYPE_405B: return "405B"; + case LLM_TYPE_456B: return "456B"; case LLM_TYPE_671B: return "671B"; case LLM_TYPE_SMALL: return "0.1B"; case LLM_TYPE_MEDIUM: return "0.4B"; @@ -816,6 +823,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_A13B: return "A13B"; case LLM_TYPE_7B_A1B: return "7B.A1B"; case LLM_TYPE_8B_A1B: return "8B.A1B"; + case LLM_TYPE_7_9B_A1_3B: return "7.9B.A1.3B"; case LLM_TYPE_12B_A2_5B: return "12B.A2.5B"; case LLM_TYPE_16B_A1B: return "16B.A1B"; case LLM_TYPE_21B_A3B: return "21B.A3B"; @@ -832,6 +840,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_118B_A8B: return "118B.A8B"; case LLM_TYPE_120B_A12B: return "120B.A12B"; case LLM_TYPE_122B_A10B: return "122B.A10B"; + case LLM_TYPE_124B_A5_1B: return "124B.A5.1B"; case LLM_TYPE_196B_A11B: return "196B.A11B"; case LLM_TYPE_230B_A10B: return "230B.A10B"; case LLM_TYPE_428B_A23B: return "428B.A23B"; @@ -842,6 +851,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_397B_A17B: return "397B.A17B"; case LLM_TYPE_685B_A37B: return "685B.A37B"; case LLM_TYPE_744B_A40B: return "744B.A40B"; + case LLM_TYPE_2_8T_A50B: return "2.8T.A50B"; case LLM_TYPE_E2B: return "E2B"; case LLM_TYPE_E4B: return "E4B"; default: return "?B"; @@ -1954,7 +1964,7 @@ void llama_model::print_info() const { LLAMA_LOG_INFO("%s: expert_weights_norm = %d\n", __func__, hparams.expert_weights_norm); } - if (arch == LLM_ARCH_BAILINGMOE2) { + if (arch == LLM_ARCH_BAILINGMOE2 || arch == LLM_ARCH_BAILINGMOE3) { LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead); LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp); LLAMA_LOG_INFO("%s: n_ff_shexp = %d\n", __func__, hparams.n_ff_shexp); @@ -2249,11 +2259,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // checks default: { - // The MTP head is dense-attention only on hybrid Qwen3-Next/3.5/3.6, so use a plain - // attention KV cache for the MTP context instead of the hybrid wrapper. + // Dense MTP heads use a plain attention KV cache instead of the hybrid wrapper. const bool mtp_on_hybrid_qwen = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && - (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE); + (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || + arch == LLM_ARCH_BAILINGMOE3); const bool mtp_on_hybrid_nemotron = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE; @@ -2283,7 +2293,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_recr = [&](uint32_t il) { return hparams.is_recr(il) && hparams.n_ff(il) == 0; }; - } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) { + } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01) { filter_attn = [&](uint32_t il) { return il < hparams.n_layer() && !hparams.is_recr(il); }; @@ -2599,6 +2609,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_KIMI_K3: return LLAMA_ROPE_TYPE_NONE; // use what we call a normal RoPE, operating on pairs of consecutive head values @@ -2630,6 +2641,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GRANITE_SWITCH: case LLM_ARCH_CHAMELEON: case LLM_ARCH_BAILINGMOE: + case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_NEO_BERT: case LLM_ARCH_SMOLLM3: case LLM_ARCH_ARCEE: @@ -2704,6 +2716,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_SEED_OSS: case LLM_ARCH_GROVEMOE: case LLM_ARCH_APERTUS: + case LLM_ARCH_MINIMAX_01: case LLM_ARCH_MINIMAX_M2: case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_COGVLM: diff --git a/examples/talk-llama/llama-model.h b/examples/talk-llama/llama-model.h index 341cb66fb..4412ef08e 100644 --- a/examples/talk-llama/llama-model.h +++ b/examples/talk-llama/llama-model.h @@ -99,6 +99,7 @@ enum llm_type { LLM_TYPE_290B, LLM_TYPE_314B, LLM_TYPE_405B, + LLM_TYPE_456B, LLM_TYPE_671B, LLM_TYPE_SMALL, LLM_TYPE_MEDIUM, @@ -117,6 +118,7 @@ enum llm_type { LLM_TYPE_A13B, LLM_TYPE_7B_A1B, LLM_TYPE_8B_A1B, // lfm2moe + LLM_TYPE_7_9B_A1_3B, // Ling-3.0-tiny LLM_TYPE_12B_A2_5B, LLM_TYPE_16B_A1B, LLM_TYPE_21B_A3B, // Ernie MoE small @@ -133,6 +135,7 @@ enum llm_type { LLM_TYPE_118B_A8B, // Laguna-S-2 LLM_TYPE_120B_A12B, // Nemotron 3 Super LLM_TYPE_122B_A10B, // Qwen3.5 + LLM_TYPE_124B_A5_1B, // Ling-3.0-flash LLM_TYPE_196B_A11B, // Step3.5-Flash LLM_TYPE_230B_A10B, // Minimax M2 LLM_TYPE_428B_A23B, // Minimax M3 @@ -143,6 +146,7 @@ enum llm_type { LLM_TYPE_397B_A17B, // Qwen3.5 LLM_TYPE_685B_A37B, // DeepSeek V3.2 LLM_TYPE_744B_A40B, // GLM-5 + LLM_TYPE_2_8T_A50B, // Kimi-K3 LLM_TYPE_E2B, LLM_TYPE_E4B, }; @@ -271,6 +275,7 @@ struct llama_layer { struct ggml_tensor * wv = nullptr; struct ggml_tensor * wo = nullptr; struct ggml_tensor * wqkv = nullptr; + struct ggml_tensor * wg = nullptr; struct ggml_tensor * wq_a = nullptr; struct ggml_tensor * wq_b = nullptr; struct ggml_tensor * wkv_a_mqa = nullptr; @@ -528,6 +533,14 @@ struct llama_layer { struct ggml_tensor * ssm_g_b = nullptr; struct ggml_tensor * ssm_o_norm = nullptr; + // kimi-k3 + struct ggml_tensor * ssm_g = nullptr; // full-rank KDA gate (replaces ssm_g_a/ssm_g_b) + struct ggml_tensor * attn_res_score = nullptr; // fused res_norm*res_proj, pre-attention + struct ggml_tensor * ffn_res_score = nullptr; // fused res_norm*res_proj, pre-FFN + struct ggml_tensor * ffn_routed_down = nullptr; // latent MoE: n_embd -> n_expert_latent + struct ggml_tensor * ffn_routed_up = nullptr; // latent MoE: n_expert_latent -> n_embd + struct ggml_tensor * ffn_routed_norm = nullptr; + // DSA (deepseek sparse attention) struct ggml_tensor * indexer_k_norm = nullptr; struct ggml_tensor * indexer_k_norm_b = nullptr; @@ -587,6 +600,7 @@ struct llama_model { struct ggml_tensor * tok_norm_b = nullptr; struct ggml_tensor * output_norm = nullptr; + struct ggml_tensor * output_res_score = nullptr; // kimi-k3: final cross-layer residual mix struct ggml_tensor * output_norm_b = nullptr; struct ggml_tensor * output = nullptr; struct ggml_tensor * output_b = nullptr; diff --git a/examples/talk-llama/llama-quant.cpp b/examples/talk-llama/llama-quant.cpp index fd6e787bd..7f99e96bc 100644 --- a/examples/talk-llama/llama-quant.cpp +++ b/examples/talk-llama/llama-quant.cpp @@ -474,7 +474,12 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type } else if (ftype == LLAMA_FTYPE_MOSTLY_MXFP4_MOE) { // MoE tensors -> MXFP4 // other tensors -> Q8_0 - if (tensor->ne[2] > 1) { + // MLA projection tensors are also 3D, so match expert tensor roles explicitly. + const bool is_bailingmoe3_expert = arch == LLM_ARCH_BAILINGMOE3 && + (category == tensor_category::FFN_UP || + category == tensor_category::FFN_GATE || + category == tensor_category::FFN_DOWN); + if (tensor->ne[2] > 1 && (arch != LLM_ARCH_BAILINGMOE3 || is_bailingmoe3_expert)) { new_type = GGML_TYPE_MXFP4; } else { new_type = GGML_TYPE_Q8_0; diff --git a/examples/talk-llama/llama-vocab.cpp b/examples/talk-llama/llama-vocab.cpp index 4a01dfd4c..ff926ceec 100644 --- a/examples/talk-llama/llama-vocab.cpp +++ b/examples/talk-llama/llama-vocab.cpp @@ -1989,6 +1989,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { // Kimi-K2 doesn't need merges, skip LLAMA_LOG_INFO("%s: Kimi-K2 tokenizer detected, skipping BPE merges\n", __func__); } else { + if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str())); + } 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); @@ -2028,8 +2032,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { const int precompiled_charsmap_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str()); if (precompiled_charsmap_keyidx != -1) { + if (gguf_get_kv_type(ctx, precompiled_charsmap_keyidx) != GGUF_TYPE_ARRAY) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str())); + } const gguf_type pc_type = gguf_get_arr_type(ctx, precompiled_charsmap_keyidx); - GGML_ASSERT(pc_type == GGUF_TYPE_INT8 || pc_type == GGUF_TYPE_UINT8); + if (pc_type != GGUF_TYPE_INT8 && pc_type != GGUF_TYPE_UINT8) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str())); + } 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); @@ -2081,6 +2090,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { throw std::runtime_error("cannot find tokenizer merges in model file\n"); } { + if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str())); + } 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); @@ -2407,21 +2420,41 @@ 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"); } + if (gguf_get_kv_type(ctx, token_idx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, token_idx) != GGUF_TYPE_STRING) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_LIST).c_str())); + } + const uint32_t n_tokens = gguf_get_arr_n(ctx, token_idx); const float * scores = nullptr; + const int * iscores = nullptr; const int score_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SCORES).c_str()); if (score_idx != -1) { + const gguf_type kv_type = gguf_get_kv_type(ctx, score_idx); + const gguf_type arr_type = kv_type == GGUF_TYPE_ARRAY ? gguf_get_arr_type(ctx, score_idx) : GGUF_TYPE_COUNT; + if (arr_type != GGUF_TYPE_INT32 && + arr_type != GGUF_TYPE_FLOAT32) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SCORES).c_str())); + } 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); + if (arr_type == GGUF_TYPE_INT32) { + iscores = (const int *) gguf_get_arr_data(ctx, score_idx); + } else { + 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) { + if (gguf_get_kv_type(ctx, toktype_idx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, toktype_idx) != GGUF_TYPE_INT32) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str())); + } 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"); @@ -2443,7 +2476,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { auto & token_data = id_to_token[i]; token_data.text = std::move(word); - token_data.score = scores ? scores[i] : 0.0f; + if (scores) { + token_data.score = scores[i]; + } else if (iscores) { + token_data.score = static_cast(iscores[i]); + } else { + token_data.score = 0.0f; + } token_data.attr = LLAMA_TOKEN_ATTR_NORMAL; if (toktypes) { //TODO: remove, required until per token attributes are available from GGUF file @@ -2584,6 +2623,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { { const int suppress_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str()); if (suppress_idx != -1) { + if (gguf_get_kv_type(ctx, suppress_idx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, suppress_idx) != GGUF_TYPE_INT32) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str())); + } const int n = gguf_get_arr_n(ctx, suppress_idx); const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx); // drop out-of-range ids diff --git a/examples/talk-llama/models/bailingmoe3.cpp b/examples/talk-llama/models/bailingmoe3.cpp new file mode 100644 index 000000000..f5855696e --- /dev/null +++ b/examples/talk-llama/models/bailingmoe3.cpp @@ -0,0 +1,532 @@ +#include "models.h" +#include "llama-memory-recurrent.h" + +void llama_model_bailingmoe3::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false); + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + if (!ml.get_key(LLM_KV_KDA_SAFE_GATE, hparams.kda_safe_gate, false)) { + hparams.kda_safe_gate = true; + } + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false); + + if (hparams.n_ff_shexp == 0) { + hparams.n_ff_shexp = hparams.n_ff_exp * std::max(1u, hparams.n_expert_shared); + } + + GGML_ASSERT(hparams.kda_safe_gate); + GGML_ASSERT(hparams.kda_gate_lower_bound < 0.0f); + + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0; + } + + switch (hparams.n_layer()) { + case 24: type = hparams.n_embd == 1536 && hparams.n_expert == 128 ? LLM_TYPE_7_9B_A1_3B : LLM_TYPE_UNKNOWN; break; + case 42: type = hparams.n_embd == 2560 && hparams.n_expert == 512 ? LLM_TYPE_124B_A5_1B : LLM_TYPE_UNKNOWN; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_bailingmoe3::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + if (output == nullptr) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); + } + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = head_dim * n_head; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + + const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); + const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; + const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; + + if (!ml.load_mtp) { + mtp_flags |= TENSOR_SKIP; + } + + for (int il = 0; il < n_layer; ++il) { + auto & layer = layers[il]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, trunk_flags); + + if (hparams.is_recr(il)) { + layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags); + layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags); + layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags); + + create_tensor_qkv(layer, il, n_embd, d_inner, d_inner, d_inner, trunk_flags); + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", il), { n_embd, d_inner }, trunk_flags); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_head }, trunk_flags); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, il), { 1, n_head }, trunk_flags); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { d_inner }, trunk_flags); + layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", il), { n_embd, d_inner }, trunk_flags); + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_dim }, trunk_flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { d_inner, n_embd }, trunk_flags); + } else { + if (q_lora_rank > 0) { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, trunk_flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, trunk_flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head * qk_head_dim }, trunk_flags); + } else { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", il), { n_embd, n_head * qk_head_dim }, trunk_flags); + } + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il), { n_embd, kv_lora_rank + qk_rope_head_dim }, trunk_flags); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { kv_lora_rank }, trunk_flags); + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", il), { qk_head_dim - qk_rope_head_dim, kv_lora_rank, n_head }, trunk_flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), { kv_lora_rank, v_head_dim, n_head }, trunk_flags); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, n_head }, trunk_flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_head * v_head_dim, n_embd }, trunk_flags); + } + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, trunk_flags); + if ((uint32_t) il < hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), { n_embd, n_ff }, trunk_flags); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), { n_embd, n_ff }, trunk_flags); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd }, trunk_flags); + } else { + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, trunk_flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, trunk_flags); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, trunk_flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, trunk_flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { hparams.n_ff_exp, n_embd, n_expert }, trunk_flags); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, trunk_flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, trunk_flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { hparams.n_ff_shexp, n_embd }, trunk_flags); + } + } + + for (int il = n_layer; il < n_layer_all; ++il) { + auto & layer = layers[il]; + const int flags = mtp_flags; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags); + if (q_lora_rank > 0) { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head * qk_head_dim }, flags); + } else { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", il), { n_embd, n_head * qk_head_dim }, flags); + } + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il), { n_embd, kv_lora_rank + qk_rope_head_dim }, flags); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { kv_lora_rank }, flags); + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", il), { qk_head_dim - qk_rope_head_dim, kv_lora_rank, n_head }, flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), { kv_lora_rank, v_head_dim, n_head }, flags); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, n_head }, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_head * v_head_dim, n_embd }, flags); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, flags); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, flags); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { hparams.n_ff_exp, n_embd, n_expert }, flags); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { hparams.n_ff_shexp, n_embd }, flags); + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", il), { n_embd }, flags); + } +} + +std::unique_ptr llama_model_bailingmoe3::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } + return std::make_unique(*this, params); +} + +static ggml_tensor * bailingmoe3_causal_conv1d( + ggml_cgraph * gf, + ggml_context * ctx0, + ggml_tensor * conv_states_all, + ggml_tensor * conv_state_all, + int64_t qkv, + ggml_tensor * x, + ggml_tensor * proj_w, + ggml_tensor * conv_w, + int64_t d_conv, + int64_t head_dim, + int64_t n_head, + int64_t n_seq_tokens, + int64_t n_seqs, + int64_t n_tokens, + int64_t cache_head) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t total_state_size = 3 * conv_state_size; + + ggml_tensor * conv_state = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), + total_state_size * ggml_element_size(conv_state_all), + qkv * conv_state_size * ggml_element_size(conv_state_all)); + + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0); + + ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv_x, + ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + total_state_size * ggml_element_size(conv_states_all), + (cache_head * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight); + out = ggml_silu(ctx0, ggml_reshape_2d(ctx0, out, d_inner, n_tokens)); + return ggml_reshape_4d(ctx0, out, head_dim, n_head, n_seq_tokens, n_seqs); +} + +llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + ggml_tensor * inpL = build_inp_embd(model.tok_embd); + cb(inpL, "model.input_embed", -1); + + auto * inp = build_inp_mem_hybrid_k(); + auto * inp_rs = inp->get_recr(); + auto * inp_attn = inp->get_attn(); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t n_head = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = n_head * head_dim; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = qk_head_dim - qk_rope_head_dim; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const float kq_scale = 1.0f / sqrtf((float) qk_head_dim); + + GGML_ASSERT(n_seqs > 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + ggml_tensor * inpSA = inpL; + ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + if (hparams.is_recr(il)) { + const auto * mctx_cur = inp_rs->mctx; + const auto cache_head = mctx_cur->get_head(); + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + + ggml_tensor * q = bailingmoe3_causal_conv1d( + gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + ggml_tensor * k = bailingmoe3_causal_conv1d( + gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + ggml_tensor * v = bailingmoe3_causal_conv1d( + gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + + ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); + gate = ggml_add(ctx0, gate, layer.ssm_dt_b); + gate = ggml_reshape_3d(ctx0, gate, head_dim, n_head, n_tokens); + ggml_tensor * a = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1); + gate = ggml_scale(ctx0, ggml_sigmoid(ctx0, ggml_mul(ctx0, gate, a)), hparams.kda_gate_lower_bound); + gate = ggml_reshape_4d(ctx0, gate, head_dim, n_head, n_seq_tokens, n_seqs); + cb(gate, "kda_gate", il); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_sigmoid(ctx0, ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs)); + + q = ggml_l2_norm(ctx0, q, hparams.f_norm_rms_eps); + k = ggml_l2_norm(ctx0, k, hparams.f_norm_rms_eps); + + ggml_tensor * states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs); + + auto result = build_delta_net(q, k, v, gate, beta, state, il); + ggml_tensor * out = ggml_cont(ctx0, result.first); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, result.second, + ggml_view_1d(ctx0, states_all, hparams.n_embd_s() * n_seqs, + cache_head * hparams.n_embd_s() * ggml_element_size(states_all)))); + + ggml_tensor * out_gate = ggml_mul_mat(ctx0, layer.ssm_g_a, cur); + out_gate = ggml_reshape_3d(ctx0, out_gate, head_dim, n_head, n_tokens); + out = ggml_reshape_3d(ctx0, out, head_dim, n_head, n_tokens); + out = build_norm(out, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + out = ggml_mul(ctx0, out, ggml_sigmoid(ctx0, out_gate)); + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, out, d_inner, n_tokens)); + cb(cur, "kda_out", il); + } else { + ggml_tensor * attn_input = cur; + ggml_tensor * q_all; + if (layer.wq_a) { + q_all = ggml_mul_mat(ctx0, layer.wq_a, cur); + cb(q_all, "q_a", il); + q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(q_all, "q_a_norm", il); + q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all); + cb(q_all, "q_b", il); + } else { + q_all = ggml_mul_mat(ctx0, layer.wq, cur); + } + ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, q_all, qk_rope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, + ggml_row_size(q_all->type, qk_nope_head_dim)); + + ggml_tensor * kv_all = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + ggml_tensor * kv = ggml_view_2d(ctx0, kv_all, kv_lora_rank, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_all, qk_rope_head_dim, 1, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank)); + + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + q_nope = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + + ggml_tensor * q = ggml_concat(ctx0, q_nope, q_pe, 0); + kv = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens); + ggml_tensor * k = ggml_concat(ctx0, kv, k_pe, 0); + + cur = build_attn(inp_attn, nullptr, nullptr, nullptr, + q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il); + + ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input); + attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens)); + cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens); + cur = ggml_mul(ctx0, cur, attn_gate); + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens)); + cb(cur, "mla_out", il); + } + + if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + cur = build_ffn(cur, + layer.ffn_up, nullptr, nullptr, + layer.ffn_gate, nullptr, nullptr, + layer.ffn_down, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + } else { + ggml_tensor * moe = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, + hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + ggml_tensor * shared = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cur = ggml_add(ctx0, moe, shared); + } + + cur = ggml_add(ctx0, cur, ffn_inp); + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + inpL = cur; + } + + ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (!cparams.embeddings_nextn_masked && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} + +llama_model_bailingmoe3::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "BailingMoE3 MTP requires one NextN layer"); + + const int il = hparams.n_layer() + cparams.nextn_layer_offset; + GGML_ASSERT(cparams.nextn_layer_offset >= 0 && + cparams.nextn_layer_offset < (int) hparams.n_layer_nextn && + "nextn_layer_offset out of range"); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj"); + GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); + GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); + GGML_ASSERT(layer.nextn.shared_head_norm && "MTP block missing final norm"); + + const int64_t n_head = hparams.n_head(); + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = qk_head_dim - qk_rope_head_dim; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const float kq_scale = 1.0f / sqrtf((float) qk_head_dim); + + auto inp = std::make_unique(hparams.n_embd); + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->embd); + ggml_set_name(inp->embd, "mtp_h_input"); + + ggml_tensor * tok_embd = ggml_get_rows(ctx0, model.tok_embd, inp->tokens); + ggml_tensor * h_norm = build_norm(inp->embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * cur = ggml_mul_mat(ctx0, layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0)); + cb(cur, "mtp_eh_proj", il); + + res->add_input(std::move(inp)); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + auto * inp_attn = build_attn_inp_k(); + + ggml_tensor * inpSA = cur; + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * attn_input = cur; + + ggml_tensor * q_all; + if (layer.wq_a) { + q_all = ggml_mul_mat(ctx0, layer.wq_a, cur); + cb(q_all, "q_a", il); + q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(q_all, "q_a_norm", il); + q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all); + cb(q_all, "q_b", il); + } else { + q_all = ggml_mul_mat(ctx0, layer.wq, cur); + } + ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, q_all, qk_rope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, + ggml_row_size(q_all->type, qk_nope_head_dim)); + + ggml_tensor * kv_all = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + ggml_tensor * kv = ggml_view_2d(ctx0, kv_all, kv_lora_rank, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_all, qk_rope_head_dim, 1, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank)); + + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + q_nope = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + + ggml_tensor * q = ggml_concat(ctx0, q_nope, q_pe, 0); + kv = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens); + ggml_tensor * k = ggml_concat(ctx0, kv, k_pe, 0); + + cur = build_attn(inp_attn, nullptr, nullptr, nullptr, + q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il); + + ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input); + attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens)); + cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens); + cur = ggml_mul(ctx0, cur, attn_gate); + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens)); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + + ggml_tensor * moe = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, + hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + ggml_tensor * shared = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cur = ggml_add(ctx0, moe, shared); + cur = ggml_add(ctx0, cur, ffn_inp); + cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM_RMS, -1); + + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} diff --git a/examples/talk-llama/models/deepseek32.cpp b/examples/talk-llama/models/deepseek32.cpp index 8a07a0b71..08555a801 100644 --- a/examples/talk-llama/models/deepseek32.cpp +++ b/examples/talk-llama/models/deepseek32.cpp @@ -180,10 +180,11 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ const int64_t n_indexer_head = hparams.indexer_n_head; const int64_t n_embd_indexer_head = hparams.indexer_head_size; - const int64_t n_embd_indexer_head_rope = hparams.n_rot(); - const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope; const uint32_t n_indexer_top_k = hparams.indexer_top_k; + // the indexer head layous is [rope | nope] + GGML_ASSERT(hparams.n_rot() <= n_embd_indexer_head); + const uint32_t kv_lora_rank = hparams.n_lora_kv; // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly. @@ -233,28 +234,11 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * indexer_q = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_q_b, qr); cb(indexer_q, "indexer_q", il); - // split into {n_embd_indexer_head_rope, n_indexer_head, n_tokens} - ggml_tensor * indexer_q_pe = - ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, n_tokens, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, 0); - cb(indexer_q_pe, "indexer_q_pe", il); - - // and {n_embd_indexer_head_nope, n_indexer_head, n_tokens} - ggml_tensor * indexer_q_nope = - ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, n_tokens, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, - ggml_row_size(indexer_q->type, n_embd_indexer_head_nope)); - cb(indexer_q_nope, "indexer_q_nope", il); - - indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot, + // {n_embd_indexer_head, n_indexer_head, n_tokens} + indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens); + indexer_q = ggml_rope_ext(ctx0, indexer_q, inp_pos, nullptr, n_rot, LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - cb(indexer_q_pe, "indexer_q_pe", il); - - // {n_embd_indexer_head_rope + n_embd_indexer_head_nope, n_head, n_tokens} - indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0); cb(indexer_q, "indexer_q", il); ggml_tensor * indexer_k = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_k, cur); @@ -263,28 +247,11 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_ indexer_k = build_norm(indexer_k, model.layers[il].indexer_k_norm, model.layers[il].indexer_k_norm_b, LLM_NORM, il); cb(indexer_k, "indexer_k", il); - // split into {n_embd_indexer_head_rope, 1, n_tokens} - ggml_tensor * indexer_k_pe = - ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_rope, 1, n_tokens, - ggml_row_size(indexer_k->type, n_embd_indexer_head), - ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, 0); - cb(indexer_k_pe, "indexer_k_pe", il); - - // and {n_embd_indexer_head_nope, 1, n_tokens} - ggml_tensor * indexer_k_nope = - ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_nope, 1, n_tokens, - ggml_row_size(indexer_k->type, n_embd_indexer_head), - ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, - ggml_row_size(indexer_k->type, n_embd_indexer_head_nope)); - cb(indexer_k_nope, "indexer_k_nope", il); - - indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot, + // {n_embd_indexer_head, 1, n_tokens} + indexer_k = ggml_reshape_3d(ctx0, indexer_k, n_embd_indexer_head, 1, n_tokens); + indexer_k = ggml_rope_ext(ctx0, indexer_k, inp_pos, nullptr, n_rot, LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - cb(indexer_k_pe, "indexer_k_pe", il); - - // {n_embd_indexer_head_rope + n_embd_indexer_head_nope, 1, n_tokens} - indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0); cb(indexer_k, "indexer_k", il); // perform Hadamard transform on indexer q and k diff --git a/examples/talk-llama/models/dflash.cpp b/examples/talk-llama/models/dflash.cpp index daaa20826..5b70a5179 100644 --- a/examples/talk-llama/models/dflash.cpp +++ b/examples/talk-llama/models/dflash.cpp @@ -43,6 +43,8 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, false); + GGML_ASSERT(hparams.dsv4_o_group_count > 0); // avoid div by zero + if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) { throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring"); } @@ -83,6 +85,16 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t n_embd_inp = hparams.n_embd_inp_enc(); tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + + // reduced draft vocab (optional): d2t maps draft rows to target token ids + int64_t n_vocab_draft = n_vocab; + const struct ggml_tensor * d2t_meta = ml->get_tensor_meta("d2t"); + if (d2t_meta) { + n_vocab_draft = d2t_meta->ne[0]; + d2t = create_tensor(tn(LLM_TENSOR_D2T), { n_vocab_draft }, 0); + LLAMA_LOG_INFO("%s: DFlash using d2t mapping (draft_vocab_size = %lld)\n", __func__, (long long) n_vocab_draft); + } + // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head // // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) @@ -92,7 +104,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t dspark_markov_rank = markov_meta->ne[0]; dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0); - dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab }, 0); + dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0); dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0); dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); @@ -155,6 +167,9 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { return; } + // optional: reduced-vocab drafts ship their own, full-vocab drafts share the target's via ctx_other + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED); + for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; @@ -240,6 +255,11 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & const int64_t block_size = std::stoi(it->second); GGML_ASSERT(block_size > 0); + // bonus anchor (SpecForge exports): slot 0 is a bonus token, not a prediction slot + const auto it_anchor = model.gguf_kv.find("dflash.sample_from_anchor"); + const bool sample_from_anchor = it_anchor == model.gguf_kv.end() || it_anchor->second == "true"; + const int64_t i_draft_beg = sample_from_anchor ? 0 : 1; + const int64_t n_blocks = g.ubatch.n_seqs_unq; GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DSpark markov head requires equal-size blocks"); // runtime tokens per block in this ubatch (anchor + drafted positions), bounded by training block_size @@ -261,11 +281,26 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_tensor * cat = nullptr; ggml_tensor * cat_conf = nullptr; + if (!sample_from_anchor) { + // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column + cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0)); + cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0))); + } + // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final // token pick, not the Markov conditioning path - for (int64_t i = 0; i < block_drafts; ++i) { + for (int64_t i = i_draft_beg; i < block_drafts; ++i) { ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks] - ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab, n_blocks] + ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab_draft, n_blocks] + if (model.d2t) { + // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others) + const int64_t n_draft_vocab = bias->ne[0]; + ggml_tensor * full = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_blocks), 0.0f); + bias = ggml_set_rows(ctx0, full, + ggml_reshape_3d(ctx0, bias, 1, n_draft_vocab, n_blocks), + ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1)); + bias = ggml_reshape_2d(ctx0, bias, n_vocab, n_blocks); + } // position i of every block: strided view [n_vocab, n_blocks] ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, i*base->nb[1]); @@ -495,6 +530,22 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra } cur = build_lora_mm(output, cur, output_s); + + // reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t + if (model.d2t) { + const int64_t n_draft_vocab = cur->ne[0]; + const int64_t n_outputs = cur->ne[1]; + const int64_t n_vocab = (int64_t) model.vocab.n_tokens(); + + GGML_ASSERT(model.d2t->type == GGML_TYPE_I64); + GGML_ASSERT(model.d2t->ne[0] == n_draft_vocab); + + ggml_tensor * logits = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_outputs), -INFINITY); + cur = ggml_set_rows(ctx0, logits, + ggml_reshape_3d(ctx0, cur, 1, n_draft_vocab, n_outputs), + ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1)); + cur = ggml_reshape_2d(ctx0, cur, n_vocab, n_outputs); + } cb(cur, "result_output", -1); res->t_logits = cur; diff --git a/examples/talk-llama/models/glm-dsa.cpp b/examples/talk-llama/models/glm-dsa.cpp index 360c2ee77..803ef7674 100644 --- a/examples/talk-llama/models/glm-dsa.cpp +++ b/examples/talk-llama/models/glm-dsa.cpp @@ -216,10 +216,11 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par const int64_t n_indexer_head = hparams.indexer_n_head; const int64_t n_embd_indexer_head = hparams.indexer_head_size; - const int64_t n_embd_indexer_head_rope = hparams.n_rot(); - const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope; const uint32_t n_indexer_top_k = hparams.indexer_top_k; + // the indexer head layout is [rope | nope] + GGML_ASSERT(hparams.n_rot() <= n_embd_indexer_head); + const uint32_t kv_lora_rank = hparams.n_lora_kv; // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly. @@ -273,28 +274,11 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par ggml_tensor * indexer_q = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_q_b, qr); cb(indexer_q, "indexer_q", il); - // split into {n_embd_indexer_head_rope, n_indexer_head, n_tokens} - ggml_tensor * indexer_q_pe = - ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, n_tokens, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, 0); - cb(indexer_q_pe, "indexer_q_pe", il); - - // and {n_embd_indexer_head_nope, n_indexer_head, n_tokens} - ggml_tensor * indexer_q_nope = - ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, n_tokens, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, - ggml_row_size(indexer_q->type, n_embd_indexer_head_nope)); - cb(indexer_q_nope, "indexer_q_nope", il); - - indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot, + // {n_embd_indexer_head, n_indexer_head, n_tokens} + indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens); + indexer_q = ggml_rope_ext(ctx0, indexer_q, inp_pos, nullptr, n_rot, LLAMA_ROPE_TYPE_NORM, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - cb(indexer_q_pe, "indexer_q_pe", il); - - // {n_embd_indexer_head_rope + n_embd_indexer_head_nope, n_head, n_tokens} - indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0); cb(indexer_q, "indexer_q", il); ggml_tensor * indexer_k = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_k, cur); @@ -303,28 +287,11 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par indexer_k = build_norm(indexer_k, model.layers[il].indexer_k_norm, model.layers[il].indexer_k_norm_b, LLM_NORM, il); cb(indexer_k, "indexer_k", il); - // split into {n_embd_indexer_head_rope, 1, n_tokens} - ggml_tensor * indexer_k_pe = - ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_rope, 1, n_tokens, - ggml_row_size(indexer_k->type, n_embd_indexer_head), - ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, 0); - cb(indexer_k_pe, "indexer_k_pe", il); - - // and {n_embd_indexer_head_nope, 1, n_tokens} - ggml_tensor * indexer_k_nope = - ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_nope, 1, n_tokens, - ggml_row_size(indexer_k->type, n_embd_indexer_head), - ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, - ggml_row_size(indexer_k->type, n_embd_indexer_head_nope)); - cb(indexer_k_nope, "indexer_k_nope", il); - - indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot, + // {n_embd_indexer_head, 1, n_tokens} + indexer_k = ggml_reshape_3d(ctx0, indexer_k, n_embd_indexer_head, 1, n_tokens); + indexer_k = ggml_rope_ext(ctx0, indexer_k, inp_pos, nullptr, n_rot, LLAMA_ROPE_TYPE_NORM, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - cb(indexer_k_pe, "indexer_k_pe", il); - - // {n_embd_indexer_head_rope + n_embd_indexer_head_nope, 1, n_tokens} - indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0); cb(indexer_k, "indexer_k", il); // perform Hadamard transform on indexer q and k diff --git a/examples/talk-llama/models/kimi-k3.cpp b/examples/talk-llama/models/kimi-k3.cpp new file mode 100644 index 000000000..d952d72cd --- /dev/null +++ b/examples/talk-llama/models/kimi-k3.cpp @@ -0,0 +1,614 @@ +#include "models.h" +#include "llama-memory-recurrent.h" + +// +// Kimi-K3 text model: hybrid KDA (linear) + MLA (full) attention, as in kimi-linear. +// Parts that kimi-linear does not have: +// 1. cross-layer residual attention (attn_res_block_size) +// 2. latent MoE (routed experts run at n_expert_latent) +// 3. situ activation (replaces SwiGLU everywhere) +// 4. MLA output gate (sigmoid gate before o_proj) +// 5. full-rank KDA gate (single ssm_g instead of ssm_g_a/ssm_g_b) +// + +void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false); + + // the MLA cache holds the compressed latent + // set it here too, as older GGUFs have no value_length key + hparams.n_embd_head_v_full = hparams.n_lora_kv; + + // n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-linear + for (uint32_t i = 0; i < hparams.n_layer(); ++i) { + hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0; + } + + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent, false); + + ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size); + ml.get_key(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta); + ml.get_key(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta); + + switch (hparams.n_layer()) { + case 93: type = LLM_TYPE_2_8T_A50B; break; // Kimi-K3 + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_kimi_k3::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + if (hparams.attn_res_block_size > 0) { + output_res_score = create_tensor(tn(LLM_TENSOR_OUTPUT_RES_SCORE, "weight"), {n_embd}, 0); + } + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + if (hparams.attn_res_block_size > 0) { + layer.attn_res_score = create_tensor(tn(LLM_TENSOR_ATTN_RES_SCORE, "weight", i), {n_embd}, 0); + layer.ffn_res_score = create_tensor(tn(LLM_TENSOR_FFN_RES_SCORE, "weight", i), {n_embd}, 0); + } + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = head_dim * n_head; + + if (hparams.is_recr(i)) { + // conv1d may be stored 4D [d_conv, 1, d_inner, 1] or 3D (quantization drops the trailing 1) + auto conv = [&](llm_tensor tid) { + ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED); + return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0); + }; + layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q); + layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K); + layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V); + + create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0); + + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, 0); + layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); + + // K3's A_log is a plain 1-D [n_head] tensor (kimi-linear's is padded) + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, 0); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0); + + // K3 uses a single full-rank gate instead of kimi-linear's g_a/g_b pair + layer.ssm_g = create_tensor(tn(LLM_TENSOR_SSM_G, "weight", i), {n_embd, d_inner}, 0); + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0); + } else { + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = n_embd_head_k - qk_rope_head_dim; + + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, TENSOR_NOT_REQUIRED); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0); + + if (layer.attn_q_a_norm) { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, 0); + } else { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k}, 0); + } + + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, 0); + layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), + {kv_lora_rank, n_head * (qk_nope_head_dim + n_embd_head_v)}, + TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); + if (!layer.wkv_b) { + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {qk_nope_head_dim, kv_lora_rank, n_head}, 0); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, 0); + } + + // K3: sigmoid output gate applied to the attention output before o_proj + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head * n_embd_head_v}, TENSOR_NOT_REQUIRED); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, 0); + } + + if (i < (int) hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } else { + const int64_t n_ff_exp = hparams.n_ff_exp; + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + + // routed experts live in the latent space + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd_latent, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0); + + if (hparams.n_expert_latent > 0) { + layer.ffn_routed_down = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_DOWN, "weight", i), {n_embd, n_embd_latent}, 0); + layer.ffn_routed_up = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_UP, "weight", i), {n_embd_latent, n_embd}, 0); + layer.ffn_routed_norm = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_NORM, "weight", i), {n_embd_latent}, TENSOR_NOT_REQUIRED); + } + + // shared experts stay at n_embd, width = moe_intermediate_size * n_expert_shared + const int64_t n_ff_shexp = n_ff_exp * (hparams.n_expert_shared > 0 ? hparams.n_expert_shared : 1); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED); + } + } +} + +std::unique_ptr llama_model_kimi_k3::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) +// linear_beta <= 0 disables the transform on the up branch +static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_tensor * up, + float beta, float linear_beta) { + ggml_tensor * a = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, gate, 1.0f/beta)), beta); + a = ggml_mul(ctx0, a, ggml_sigmoid(ctx0, gate)); + + if (linear_beta > 0.0f) { + up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/linear_beta)), linear_beta); + } + return ggml_mul(ctx0, a, up); +} + +// +// cross-layer residual attention +// + +// layout is [n_embd, n_ckpt, n_tokens]: rms_norm reduces over ne0, dsv4_hc_pre over ne1 +// append the new checkpoint, do not re-fold the whole chain +void llama_model_kimi_k3::graph::res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens) { + ggml_tensor * ckpt = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens); + + resi_stack = resi_stack ? ggml_concat(ctx0, resi_stack, ckpt, 1) : ckpt; +} + +ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor * score_w, + int64_t n_tokens, int il) { + if (!resi_stack) { + return cur; // layer 0: nothing banked yet + } + + const int n_ckpt = (int) resi_stack->ne[1]; + const float eps = hparams.f_norm_rms_eps; + + ggml_tensor * src = resi_stack; // [n_embd, n_ckpt, n_tokens] + + // one rms_norm scores all checkpoints at once + // note: the scores use the normalized values, but the sum below uses the raw ones + ggml_tensor * sc_src = ggml_rms_norm(ctx0, src, eps); + sc_src = ggml_mul(ctx0, sc_src, score_w); + sc_src = ggml_sum_rows(ctx0, sc_src); // [1, n_ckpt, n_tokens] + sc_src = ggml_reshape_2d(ctx0, sc_src, n_ckpt, n_tokens); + + // the current residual stream is scored apart, so the stack stays append-only + ggml_tensor * sc_cur = ggml_rms_norm(ctx0, cur, eps); + sc_cur = ggml_mul(ctx0, sc_cur, score_w); + sc_cur = ggml_sum_rows(ctx0, sc_cur); // [1, n_tokens] + + ggml_tensor * scores = ggml_concat(ctx0, sc_src, sc_cur, 0); // [n_ckpt+1, n_tokens] + ggml_tensor * probs = ggml_soft_max(ctx0, scores); // over ne0 = n_ckpt+1 + cb(probs, "res_probs", il); + + // split the sum: hc_pre handles the stack, a broadcast-multiply the current stream + ggml_tensor * p_src = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, n_ckpt, n_tokens, probs->nb[1], 0)); + ggml_tensor * p_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, 1, n_tokens, probs->nb[1], + probs->nb[0] * n_ckpt)); + + ggml_tensor * out = ggml_dsv4_hc_pre(ctx0, src, p_src); + out = ggml_add(ctx0, out, ggml_mul(ctx0, cur, p_cur)); + + return out; +} + +llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + cb(inpL, "inp_embd", -1); + + // K3 MLA is nope-only, so there is no position input + + auto * inp_kv = !hparams.is_mla() ? build_inp_mem_hybrid() : nullptr; + auto * inp_k = hparams.is_mla() ? build_inp_mem_hybrid_k() : nullptr; + auto * inp_rs = hparams.is_mla() ? inp_k->get_recr() : inp_kv->get_recr(); + auto * inp_attn_kv = !hparams.is_mla() ? inp_kv->get_attn() : nullptr; + auto * inp_attn_k = hparams.is_mla() ? inp_k->get_attn() : nullptr; + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t n_head_kda = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = n_head_kda * head_dim; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla(); + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_qk_rope = hparams.n_rot(); + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; + const float kq_scale_mla = 1.0f / sqrtf((float) n_embd_head_k_mla); + + const uint32_t res_bs = hparams.attn_res_block_size; + const bool use_attn_res = res_bs > 0; + const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd; + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + // the residual stream, banked on checkpoint layers and then restarted + // from the attention output alone + ggml_tensor * prefix_sum = inpL; + + cur = use_attn_res ? res_mix(prefix_sum, layer.attn_res_score, n_tokens, il) + : prefix_sum; + + bool banked = false; + if (use_attn_res && (uint32_t) il % res_bs == 0) { + res_push(prefix_sum, n_embd, n_tokens); // banks the RAW layer input, not `cur` + banked = true; + } + + cur = build_norm(cur, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + ggml_build_forward_expand(gf, cur); + + if (hparams.is_recr(il)) { + cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda, + d_inner, n_seq_tokens, n_seqs, il); + } else { + cur = build_mla_layer(cur, layer, inp_attn_k, inp_attn_kv, + n_embd_head_k_mla, n_embd_head_v_mla, kv_lora_rank, + n_embd_head_qk_rope, n_embd_head_qk_nope, kq_scale_mla, il); + } + + prefix_sum = banked ? cur : ggml_add(ctx0, prefix_sum, cur); + cb(prefix_sum, "prefix_sum_attn", il); + + cur = use_attn_res ? res_mix(prefix_sum, layer.ffn_res_score, n_tokens, il) + : prefix_sum; + + cur = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate, cur); + ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up, cur); + cur = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta); + cur = ggml_mul_mat(ctx0, layer.ffn_down, cur); + cb(cur, "ffn_out", il); + } else { + cur = build_latent_moe(cur, layer, n_embd_latent, il); + } + + prefix_sum = ggml_add(ctx0, prefix_sum, cur); + prefix_sum = build_cvec(prefix_sum, il); + cb(prefix_sum, "l_out", il); + + inpL = prefix_sum; + } + + cur = inpL; + + // final mix, then narrow to the output tokens + if (use_attn_res) { + cur = res_mix(cur, model.output_res_score, n_tokens, -1); + } + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +// +// KDA layer +// + +// causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use +static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0, + ggml_tensor * conv_states_all, ggml_tensor * conv_state_all, + int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w, + int64_t d_conv, int64_t head_dim, int64_t n_head, + int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t n_embd_r_total = 3 * conv_state_size; + + ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), + n_embd_r_total * ggml_element_size(conv_state_all), + qkv * conv_state_size * ggml_element_size(conv_state_all)); + + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0); + + ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, last_conv_x, + ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + n_embd_r_total * ggml_element_size(conv_states_all), + (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight); + Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens); + Xcur = ggml_silu(ctx0, Xcur); + + return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs); +} + +ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer( + ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs, + int64_t d_conv, int64_t head_dim, int64_t n_head_kda, + int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) { + + const auto * mctx_cur = inp_rs->mctx; + const auto kv_head = mctx_cur->get_head(); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + + ggml_tensor * Qcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Kcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Vcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + cb(Qcur, "kda_q_conv", il); + cb(Kcur, "kda_k_conv", il); + cb(Vcur, "kda_v_conv", il); + + // gate_lower_bound is not a clamp - when set, it swaps the decay gate activation: + // unset (kimi-linear): g = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias) + // set (K3, -5.0): g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias)) + // ssm_a holds -exp(A_log) (folded at conversion time), so exp(A_log) == -ssm_a + ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); + ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a); + g1 = ggml_add(ctx0, g1, layer.ssm_dt_b); + + ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1); + + if (hparams.kda_gate_lower_bound > -INFINITY) { + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens); + g1 = ggml_mul(ctx0, g1, A); // -exp(A_log) * (...) + g1 = ggml_sigmoid(ctx0, ggml_scale(ctx0, g1, -1.0f)); + g1 = ggml_scale(ctx0, g1, hparams.kda_gate_lower_bound); + } else { + g1 = ggml_softplus(ctx0, g1); + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens); + g1 = ggml_mul(ctx0, g1, A); + } + cb(g1, "kda_g1", il); + + g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head_kda, n_seq_tokens, n_seqs); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs); + beta = ggml_sigmoid(ctx0, beta); + cb(beta, "kda_beta", il); + + ggml_tensor * cur_3d = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs); + + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs); + + const float eps = hparams.f_norm_rms_eps; + Qcur = ggml_l2_norm(ctx0, Qcur, eps); + Kcur = ggml_l2_norm(ctx0, Kcur, eps); + + auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il); + + ggml_tensor * output = ggml_cont(ctx0, attn_out.first); + cb(output, "kda_scan_out", il); + ggml_tensor * new_state = attn_out.second; + + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, new_state, + ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs, + kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)))); + + // K3: single full-rank gate (kimi-linear factors this as g_b(g_a(x))) + ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur_3d, cur_3d->ne[0], n_seq_tokens * n_seqs); + ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g, cur_2d); + g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head_kda, n_seq_tokens * n_seqs); + + ggml_tensor * o = ggml_reshape_3d(ctx0, output, head_dim, n_head_kda, n_seq_tokens * n_seqs); + ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + cb(g2, "kda_g2", il); + cb(normed, "kda_normed", il); + ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2)); + + gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens); + cur = ggml_mul_mat(ctx0, layer.wo, gated); + cb(cur, "kda_out", il); + + return cur; +} + +// +// MLA layer (nope-only, with K3's sigmoid output gate) +// + +ggml_tensor * llama_model_kimi_k3::graph::build_mla_layer( + ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_attn_k * inp_attn_k, llm_graph_input_attn_kv * inp_attn_kv, + int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla, int64_t kv_lora_rank, + int64_t n_embd_head_qk_rope, int64_t n_embd_head_qk_nope, float kq_scale, int il) { + + ggml_tensor * inp_gate = cur; // the output gate reads the *normed* layer input + + ggml_tensor * Qcur; + if (layer.wq_a) { + Qcur = ggml_mul_mat(ctx0, layer.wq_a, cur); + Qcur = build_norm(Qcur, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + Qcur = ggml_mul_mat(ctx0, layer.wq_b, Qcur); + } else { + Qcur = ggml_mul_mat(ctx0, layer.wq, cur); + } + + ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + + ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank)); + + // no RoPE: mla_use_nope is asserted at conversion time + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + + ggml_tensor * out; + if (layer.wk_b && layer.wv_b) { + ggml_tensor * q_nope = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_nope, n_head, n_tokens, + ggml_row_size(Qcur->type, n_embd_head_k_mla), + ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_rope, n_head, n_tokens, + ggml_row_size(Qcur->type, n_embd_head_k_mla), + ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, + ggml_row_size(Qcur->type, n_embd_head_qk_nope)); + + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3); + + ggml_tensor * Q = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0); + ggml_tensor * kv_cmpr_3d = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); + ggml_tensor * K = ggml_concat(ctx0, kv_cmpr_3d, k_pe, 0); + ggml_tensor * V = kv_cmpr_3d; + + // wo == NULL: the output projection is applied after the gate below + out = build_attn(inp_attn_k, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, layer.wv_b, kq_scale, il); + } else { + ggml_tensor * Q = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens); + ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr); + const int64_t kv_per_head = n_embd_head_qk_nope + n_embd_head_v_mla; + + ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens, + ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head), 0); + ggml_tensor * V = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens, + ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head), + ggml_row_size(kv->type, n_embd_head_qk_nope))); + + ggml_tensor * k_pe_t = ggml_new_tensor_3d(ctx0, k_pe->type, n_embd_head_qk_rope, n_head, n_tokens); + ggml_tensor * K = ggml_concat(ctx0, ggml_repeat(ctx0, k_pe, k_pe_t), k_nope, 0); + + out = build_attn(inp_attn_kv, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, nullptr, kq_scale, il); + } + + // K3: attn_output *= sigmoid(g_proj(x)), then o_proj + if (layer.wqkv_gate) { + ggml_tensor * g = ggml_sigmoid(ctx0, ggml_mul_mat(ctx0, layer.wqkv_gate, inp_gate)); + out = ggml_mul(ctx0, out, g); + cb(out, "mla_gated", il); + } + + out = ggml_mul_mat(ctx0, layer.wo, out); + cb(out, "mla_out", il); + + return out; +} + +// +// latent MoE: down-project, run the routed experts in the latent space, norm, up-project; +// shared experts stay at n_embd and read the un-projected input. +// + +ggml_tensor * llama_model_kimi_k3::graph::build_latent_moe( + ggml_tensor * cur, const llama_layer & layer, int64_t n_embd_latent, int il) { + + ggml_tensor * identity = cur; + + ggml_tensor * routed_in = layer.ffn_routed_down + ? ggml_mul_mat(ctx0, layer.ffn_routed_down, cur) + : cur; + + // the router scores the full-width input while the experts take the latent one, + // so the logits are computed here and passed to build_moe_ffn + ggml_tensor * logits = ggml_mul_mat(ctx0, layer.ffn_gate_inp, identity); + cb(logits, "ffn_moe_logits", il); + + ggml_tensor * moe_out = build_moe_ffn(routed_in, + nullptr, // gate_inp unused: the logits above are passed instead + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + hparams.n_expert, + hparams.n_expert_used, + LLM_FFN_SITU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + logits); + cb(moe_out, "ffn_moe_out", il); + + if (layer.ffn_routed_norm) { + moe_out = build_norm(moe_out, layer.ffn_routed_norm, NULL, LLM_NORM_RMS, il); + } + if (layer.ffn_routed_up) { + moe_out = ggml_mul_mat(ctx0, layer.ffn_routed_up, moe_out); + } + GGML_UNUSED(n_embd_latent); + + if (layer.ffn_gate_shexp) { + ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate_shexp, identity); + ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up_shexp, identity); + ggml_tensor * sh = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta); + sh = ggml_mul_mat(ctx0, layer.ffn_down_shexp, sh); + cb(sh, "ffn_shexp", il); + moe_out = ggml_add(ctx0, moe_out, sh); + } + + cb(moe_out, "ffn_out", il); + return moe_out; +} diff --git a/examples/talk-llama/models/minimax-01.cpp b/examples/talk-llama/models/minimax-01.cpp new file mode 100644 index 000000000..a6ccee191 --- /dev/null +++ b/examples/talk-llama/models/minimax-01.cpp @@ -0,0 +1,520 @@ +#include "models.h" +#include "llama-memory-recurrent.h" + +void llama_model_minimax_01::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale); + + // we use n_embd_head_la to set recurrent memory n_embd_s + hparams.n_embd_head_la = hparams.n_embd_head_k_full; + + // Mark recurrent layers (lightning attention layers). + if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) { + uint32_t full_attn_interval = 8; + ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0); + } + } + + switch (hparams.n_layer()) { + case 80: type = LLM_TYPE_456B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_minimax_01::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + // output + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + + // if output is NULL, init from the input tok embed + if (output == NULL) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + if (!hparams.is_recr(i)) { + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0); + } else { + layer.attn_norm_2 = create_tensor(tn(LLM_TENSOR_ATTN_NORM_2, "weight", i), {n_embd_head_k * n_head}, 0); + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, 3 * n_embd_head_k * n_head}, 0); + layer.wg = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0); + } + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff, n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff, n_embd, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff, n_expert}, 0); + } +} + +std::unique_ptr llama_model_minimax_01::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +class llm_graph_input_la : public llm_graph_input_i { +public: + llm_graph_input_la(const llama_hparams & hparams) : hparams(hparams) {} + + void set_input(const llama_ubatch * ubatch) override { + // this operates on assumption that we have an equal ubatch split + + const int64_t n_head = hparams.n_head(); + const int32_t n_seqs = ubatch->n_seqs; + const int32_t n_seqs_unq = ubatch->n_seqs_unq; + const int32_t n_tokens = ubatch->n_tokens; + const int32_t n_seq_tokens = ubatch->n_seq_tokens; + + std::vector p0(n_seqs_unq); + std::fill(p0.begin(), p0.end(), std::numeric_limits::max()); + + // get lowest token position in a ubatch for each stream + for (int i = 0; i < n_tokens; ++i) { + llama_seq_id seq_id = ubatch->seq_id[i][0]; + int32_t seq_idx = ubatch->seq_idx[seq_id]; + llama_pos pos = ubatch->pos[i]; + if (p0[seq_idx] > pos) { + p0[seq_idx] = pos; + } + } + + if (inp_slopes) { + GGML_ASSERT(ggml_backend_buffer_is_host(inp_slopes->buffer)); + + float * data = (float *) inp_slopes->data; + + float start = powf(2, -powf(2, -(log2f(n_head) - 3))); + float ratio = start; + + for (int h = 0; h < n_head; ++h) { + data[h] = start * powf(ratio, h); + } + } + + if (inp_q_decay) { + GGML_ASSERT(ggml_backend_buffer_is_host(inp_q_decay->buffer)); + + float * slopes = (float *) inp_slopes->data; + float * data = (float *) inp_q_decay->data; + + for (int s = 0; s < n_seqs; ++s) { + for (int i = 0; i < n_seq_tokens; ++i) { + llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + i][0]; + int32_t seq_idx = ubatch->seq_idx[seq_id]; + llama_pos pos = ubatch->pos[s * n_seq_tokens + i]; + int pos_rel = pos - p0[seq_idx]; + + for (int h = 0; h < n_head; ++h) { + data[seq_idx * n_head * n_seq_tokens + i * n_head + h] = -slopes[h] * (pos_rel + 1); + } + } + } + } + + if (inp_k_decay) { + GGML_ASSERT(ggml_backend_buffer_is_host(inp_k_decay->buffer)); + + float * slopes = (float *) inp_slopes->data; + float * data = (float *) inp_k_decay->data; + + for (int s = 0; s < n_seqs; ++s) { + for (int i = 0; i < n_seq_tokens; ++i) { + llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + i][0]; + int32_t seq_idx = ubatch->seq_idx[seq_id]; + llama_pos pos = ubatch->pos[s * n_seq_tokens + i]; + int pos_rel = pos - p0[seq_idx]; + + for (int h = 0; h < n_head; ++h) { + data[seq_idx * n_head * n_seq_tokens + i * n_head + h] = -slopes[h] * (n_seq_tokens - pos_rel - 1); + } + } + } + } + + if (inp_diag_decay) { + GGML_ASSERT(ggml_backend_buffer_is_host(inp_diag_decay->buffer)); + + float * slopes = (float *) inp_slopes->data; + float * data = (float *) inp_diag_decay->data; + + for (int s = 0; s < n_seqs; ++s) { + for (int h = 0; h < n_head; ++h) { + for (int j = 0; j < n_seq_tokens; ++j) { + llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + j][0]; + int32_t seq_idx = ubatch->seq_idx[seq_id]; + llama_pos pos_j = ubatch->pos[s * n_seq_tokens + j]; + int pos_rel_j = pos_j - p0[seq_idx]; + + for (int i = 0; i < n_seq_tokens; ++i) { + llama_pos pos_i = ubatch->pos[s * n_seq_tokens + i]; + int pos_rel_i = pos_i - p0[seq_idx]; + + int index = pos_rel_j - pos_rel_i; + float s_index = index >= 0 ? -slopes[h] * index : -INFINITY; + data[seq_idx * n_head * n_seq_tokens * n_seq_tokens + h * n_seq_tokens * n_seq_tokens + j * n_seq_tokens + i] = s_index; + } + } + } + } + } + } + + bool can_reuse(const llm_graph_params & params) override { + bool res = true; + + if (params.ubatch.n_seq_tokens > 1) { + res &= ( inp_q_decay && inp_q_decay->ne[2] == params.ubatch.n_seq_tokens); + res &= ( inp_k_decay && inp_k_decay->ne[2] == params.ubatch.n_seq_tokens); + res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens); + } + + return res; + } + + const llama_hparams & hparams; + + ggml_tensor * inp_slopes = nullptr; // F32 [n_head] + ggml_tensor * inp_q_decay = nullptr; // F32 [1, n_head, n_batch] + ggml_tensor * inp_k_decay = nullptr; // F32 [1, n_head, n_batch] + ggml_tensor * inp_diag_decay = nullptr; // F32 [n_batch, n_batch, n_head] +}; + +llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + // GGML_ASSERT(n_embd_head == n_rot); this is wrong in case of minimax, head_dim = 128, n_rot = 64 + + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + auto * inp_hybrid = build_inp_mem_hybrid(); + auto * inp_rs = inp_hybrid->get_recr(); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + llm_graph_input_la * la = nullptr; + + auto inp = std::make_unique(hparams); + + inp->inp_slopes = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_head); + ggml_set_input(inp->inp_slopes); + cb(inp->inp_slopes, "slopes", -1); + + if (n_seq_tokens != 1) { + inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); + ggml_set_input(inp->inp_q_decay); + cb(inp->inp_q_decay, "q_decay_exp", -1); + + inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); + ggml_set_input(inp->inp_k_decay); + cb(inp->inp_k_decay, "k_decay_exp", -1); + + inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs); + ggml_set_input(inp->inp_diag_decay); + cb(inp->inp_diag_decay, "diag_decay_exp", -1); + } + + la = (llm_graph_input_la *) res->add_input(std::move(inp)); + + ggml_tensor * slopes = la->inp_slopes; + + for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; + + ggml_tensor * inpSA = inpL; + + cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + ggml_tensor * residual = cur; + + // self_attention + if (!hparams.is_recr(il)) { + // softmax attention layer + + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, il); + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + cur = build_attn(inp_hybrid->get_attn(), + model.layers[il].wo, NULL, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il); + } else { + // lightning attention layer + + const auto * mctx_cur = inp_rs->mctx; + const auto kv_head = mctx_cur->get_head(); + + // TODO unneeded - any way to make conv states optional in recurrent memory? + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + ggml_build_forward_expand(gf, conv_state_all); + + float slope_scale = 1.0 - 1.0 * il / (n_layer - 1) + 1e-5; + ggml_tensor * slope_rate = ggml_scale(ctx0, slopes, slope_scale); + cb(slope_rate, "slope_rate", il); + + cur = ggml_reshape_4d(ctx0, cur, cur->ne[0], n_seq_tokens, 1, n_seqs); + + ggml_tensor * QKVcur = build_lora_mm(model.layers[il].wqkv, cur); + cb(QKVcur, "QKVcur", il); + + QKVcur = ggml_silu(ctx0, QKVcur); + cb(QKVcur, "QKVcur_silu", il); + + QKVcur = ggml_reshape_4d(ctx0, QKVcur, n_embd_head * 3, n_head, n_seq_tokens, n_seqs); + + ggml_tensor * Qcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 0*ggml_element_size(QKVcur)*n_embd_head); + ggml_tensor * Kcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 1*ggml_element_size(QKVcur)*n_embd_head); + ggml_tensor * Vcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 2*ggml_element_size(QKVcur)*n_embd_head); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + // get previous KV + ggml_tensor * la_states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, la_states_all, hparams.n_embd_s(), n_seqs); + + ggml_tensor * kv_old = ggml_reshape_4d(ctx0, state, n_embd_head, n_embd_head, n_head, n_seqs); + cb(kv_old, "kv_old", il); + + ggml_tensor * qkv = nullptr; + ggml_tensor * kv_new = nullptr; + + if (n_seq_tokens == 1) { + // lightning attention - optimized single token case for TG + + ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0); + cb(slopes_neg, "slopes_neg", il); + + ggml_tensor * ratio = ggml_exp(ctx0, slopes_neg); + cb(ratio, "ratio", il); + + ggml_tensor * ratio_3d = ggml_reshape_3d(ctx0, ratio, 1, 1, n_head); + cb(ratio_3d, "ratio3d", il); + + ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3)); + cb(v_trans, "v_trans", il); + + ggml_tensor * k_trans = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 1, 2, 0, 3)); + cb(k_trans, "k_trans", il); + + ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_trans, v_trans); + cb(kv_cur, "kv_cur", il); + + ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, ratio_3d); + cb(kv_old_s, "kv_old_s", il); + + kv_new = ggml_add(ctx0, kv_old_s, kv_cur); + cb(kv_new, "kv_new", il); + + ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + cb(q_trans, "q_trans", il); + + qkv = ggml_mul_mat(ctx0, kv_new, q_trans); + cb(qkv, "qkv", il); + } else if(n_seq_tokens > 1) { + // lightning attention - general multi token case for PP + + ggml_tensor * q_decay_exp = la->inp_q_decay; + ggml_tensor * k_decay_exp = la->inp_k_decay; + ggml_tensor * diag_decay_exp = la->inp_diag_decay; + + ggml_tensor * q_decay = ggml_exp(ctx0, ggml_scale(ctx0, q_decay_exp, slope_scale)); + cb(q_decay, "q_decay", il); + ggml_tensor * k_decay = ggml_exp(ctx0, ggml_scale(ctx0, k_decay_exp, slope_scale)); + cb(k_decay, "k_decay", il); + ggml_tensor * diag_decay = ggml_exp(ctx0, ggml_scale(ctx0, diag_decay_exp, slope_scale)); + cb(diag_decay, "diag_decay", il); + + ggml_tensor * q_s = ggml_mul(ctx0, Qcur, q_decay); + cb(q_s, "q_s", il); + + ggml_tensor * q_s_trans = ggml_permute(ctx0, q_s, 0, 2, 1, 3); + cb(q_s_trans, "q_s_trans", il); + + ggml_tensor * qkv_none_diag = ggml_mul_mat(ctx0, kv_old, q_s_trans); + cb(qkv_none_diag, "qkv_none_diag", il); + + ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + cb(q_trans, "q_trans", il); + + ggml_tensor * k_trans = ggml_permute(ctx0, Kcur, 0, 2, 1, 3); + cb(k_trans, "k_trans", il); + + ggml_tensor * qk = ggml_mul_mat(ctx0, k_trans, q_trans); + cb(qk, "qk", il); + + qk = ggml_mul(ctx0, qk, diag_decay); + cb(qk, "qk_s", il); + + ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3)); + cb(v_trans, "v_trans", il); + + ggml_tensor * qkv_diag = ggml_mul_mat(ctx0, v_trans, qk); + cb(qkv_diag, "qkv_diag", il); + + qkv = ggml_add(ctx0, qkv_none_diag, qkv_diag); + cb(qkv, "qkv", il); + + ggml_build_forward_expand(gf, qkv); + + ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0*n_seq_tokens); + cb(slopes_neg, "slopes_neg", il); + + ggml_tensor * block_decay = ggml_exp(ctx0, slopes_neg); + cb(block_decay, "block_decay", il); + + ggml_tensor * block_decay_3d = ggml_reshape_3d(ctx0, block_decay, 1, 1, n_head); + cb(block_decay_3d, "block_decay_3d", il); + + ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, block_decay_3d); + cb(kv_old_s, "kv_old_s", il); + + ggml_tensor * k_after_decay = ggml_mul(ctx0, Kcur, k_decay); + cb(k_after_decay, "k_after_decay", il); + + ggml_tensor * k_after_decay_trans = ggml_cont(ctx0, ggml_permute(ctx0, k_after_decay, 1, 2, 0, 3)); + cb(k_after_decay_trans, "k_after_decay_trans", il); + + ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_after_decay_trans, v_trans); + cb(kv_cur, "kv_cur", il); + + kv_new = ggml_add(ctx0, kv_old_s, kv_cur); + cb(kv_new, "kv_new", il); + } + + // store new KV + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, kv_new, + ggml_view_1d(ctx0, la_states_all, hparams.n_embd_s() * n_seqs, + kv_head * hparams.n_embd_s() * ggml_element_size(la_states_all)))); + + qkv = ggml_cont(ctx0, ggml_permute(ctx0, qkv, 0, 2, 1, 3)); + cb(qkv, "qkv_permuted", il); + + qkv = ggml_reshape_4d(ctx0, qkv, qkv->ne[0]*qkv->ne[1], qkv->ne[2], 1, qkv->ne[3]); + + // norm + ggml_tensor * qkv_norm = build_norm(qkv, + model.layers[il].attn_norm_2, NULL, + LLM_NORM_RMS, il); + cb(qkv_norm, "qkv_norm", il); + + ggml_tensor * g = build_lora_mm(model.layers[il].wg, cur); + cb(g, "g", il); + + g = ggml_sigmoid(ctx0, g); + cb(g, "g_sigm", il); + + cur = ggml_mul(ctx0, g, qkv_norm); + + cur = build_lora_mm(model.layers[il].wo, cur); + cb(cur, "attn_out", il); + + cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], n_seq_tokens*n_seqs); + cb(cur, "attn_out", il); + } + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + residual = ggml_get_rows(ctx0, residual, inp_out_ids); + } + + residual = ggml_scale(ctx0, residual, hparams.f_residual_scale); + cb(residual, "residual_scaled_attn", il); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, residual); + cb(ffn_inp, "ffn_inp", il); + + // MoE branch + cur = build_norm(ffn_inp, + model.layers[il].ffn_norm, NULL, + LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + residual = cur; + + cur = build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + model.layers[il].ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, true, + hparams.expert_weights_scale, + LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, + il); + cb(cur, "ffn_moe_out", il); + + residual = ggml_scale(ctx0, residual, hparams.f_residual_scale); + cb(residual, "residual_scaled_ffn", il); + + cur = ggml_add(ctx0, cur, residual); + cb(cur, "ffn_out", il); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + // input for next layer + inpL = cur; + } + + cur = inpL; + + cur = build_norm(cur, + model.output_norm, NULL, + LLM_NORM_RMS, -1); + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + // lm_head + cur = build_lora_mm(model.output, cur, model.output_s); + + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/examples/talk-llama/models/minimax-m3.cpp b/examples/talk-llama/models/minimax-m3.cpp index 854d5aed0..1ba699d01 100644 --- a/examples/talk-llama/models/minimax-m3.cpp +++ b/examples/talk-llama/models/minimax-m3.cpp @@ -25,6 +25,8 @@ void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks); msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks }; + GGML_ASSERT(hparams.indexer_block_size > 0); // avoid div by zero + switch (hparams.n_layer()) { case 60: type = LLM_TYPE_428B_A23B; break; default: type = LLM_TYPE_UNKNOWN; diff --git a/examples/talk-llama/models/models.h b/examples/talk-llama/models/models.h index ddb9ae2f1..180b30a46 100644 --- a/examples/talk-llama/models/models.h +++ b/examples/talk-llama/models/models.h @@ -1784,6 +1784,25 @@ struct llama_model_bailingmoe2 : public llama_model_base { }; +struct llama_model_bailingmoe3 : public llama_model_base { + llama_model_bailingmoe3(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_build_delta_net_base { + graph(const llama_model & model, const llm_graph_params & params); + + const llama_model & model; + }; + + struct graph_mtp : public llm_graph_context { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_seed_oss : public llama_model_base { llama_model_seed_oss(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; @@ -2043,6 +2062,19 @@ struct llama_model_apertus : public llama_model_base { }; +struct llama_model_minimax_01 : public llama_model_base { + llama_model_minimax_01(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_minimax_m2 : public llama_model_base { llama_model_minimax_m2(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; @@ -2272,6 +2304,42 @@ struct llama_model_mimo2 : public llama_model_base { }; +struct llama_model_kimi_k3 : public llama_model_base { + llama_model_kimi_k3(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_build_delta_net_base { + graph(const llama_model & model, const llm_graph_params & params); + + const llama_model & model; + + // Cross-layer residual attention (K3's `_apply_attn_res`). + ggml_tensor * resi_stack = nullptr; + + void res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens); + ggml_tensor * res_mix(ggml_tensor * cur, ggml_tensor * score_w, + int64_t n_tokens, int il); + + ggml_tensor * build_kda_layer(ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_rs * inp_rs, + int64_t d_conv, int64_t head_dim, int64_t n_head_kda, + int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il); + + ggml_tensor * build_mla_layer(ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_attn_k * inp_attn_k, + llm_graph_input_attn_kv * inp_attn_kv, + int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla, + int64_t kv_lora_rank, int64_t n_embd_head_qk_rope, + int64_t n_embd_head_qk_nope, float kq_scale, int il); + + ggml_tensor * build_latent_moe(ggml_tensor * cur, const llama_layer & layer, + int64_t n_embd_latent, int il); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_kimi_linear : public llama_model_base { llama_model_kimi_linear(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override;