ggml-org/llama.cpp · #31779
model: add Kimi-K3.1-Pro text model
conversion/__init__.py1 + / 0 −
@@ -158,6 +158,7 @@ "KORMoForCausalLM": "qwen", "KimiK25ForConditionalGeneration": "deepseek", "KimiK3ForConditionalGeneration": "kimi_k3",+ "KimiK31ForConditionalGeneration": "kimi_k3", "KimiLinearForCausalLM": "kimi_linear", "KimiLinearModel": "kimi_linear", "KimiVLForConditionalGeneration": "deepseek",conversion/base.py1 + / 1 −
@@ -3404,7 +3404,7 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals, # latent MoE, situ, ...). Route on the top-level architecture instead.- if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration", "Nemotron3_5TranscribeForConditionalGeneration"):+ if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration", "KimiK31ForConditionalGeneration", "Nemotron3_5TranscribeForConditionalGeneration"): return arch # if "architectures" is found in the sub-config, use that insteadconversion/kimi_k3.py30 + / 0 −
@@ -343,3 +343,33 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter return yield from super().modify_tensors(data_torch, name, bid)+++@ModelBase.register("KimiK31ForConditionalGeneration")+@ModelBase.example("moonshotai/Kimi-K3.1-Pro")+class KimiK31Model(KimiK3Model):+ """+ Kimi-K3.1: K3 with the full-attention MLA layers replaced by gated DSA.++ The tensor set is K3's plus a lightning indexer per full-attention layer,+ and `dsa_config` adds the attention.indexer.* metadata.+ """++ model_arch = gguf.MODEL_ARCH.KIMI_K31++ def set_gguf_parameters(self):+ super().set_gguf_parameters()++ # --- DSA indexer ---+ dsa = self.hparams["dsa_config"]+ dsa_layers = dsa["dsa_layers"]+ # the layer lists are 1-indexed, like linear_attn_config's+ assert sorted(dsa_layers) == sorted(self.hparams["linear_attn_config"]["full_attn_layers"]), \+ "every full-attention layer is expected to own a DSA indexer"+ self.gguf_writer.add_indexer_head_count(dsa["index_num_heads"])+ self.gguf_writer.add_indexer_key_length(dsa["index_head_dim"])+ self.gguf_writer.add_indexer_top_k(dsa["index_topk"])+ self.gguf_writer.add_indexer_types([(il + 1) in dsa_layers+ for il in range(self.hparams["num_hidden_layers"])])+ # eps of the indexer key LayerNorm; the config declares no separate value+ self.gguf_writer.add_layer_norm_eps(dsa.get("index_norm_eps", self.hparams["rms_norm_eps"]))gguf-py/gguf/constants.py56 + / 0 −
@@ -752,6 +752,7 @@ class MODEL_ARCH(IntEnum): MAION_CODER = auto() KIMI_LINEAR = auto() KIMI_K3 = auto()+ KIMI_K31 = auto() TALKIE = auto() MELLUM = auto() NANBEIGE = auto()@@ -1669,6 +1670,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.MAION_CODER: "maion-coder", MODEL_ARCH.KIMI_LINEAR: "kimi-linear", MODEL_ARCH.KIMI_K3: "kimi-k3",+ MODEL_ARCH.KIMI_K31: "kimi-k31", MODEL_ARCH.TALKIE: "talkie", MODEL_ARCH.MELLUM: "mellum", MODEL_ARCH.NANBEIGE: "nanbeige",@@ -6536,6 +6538,60 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_ROUTED_UP, MODEL_TENSOR.FFN_ROUTED_NORM, ],+ MODEL_ARCH.KIMI_K31: [+ MODEL_TENSOR.TOKEN_EMBD,+ MODEL_TENSOR.OUTPUT_NORM,+ MODEL_TENSOR.OUTPUT,+ MODEL_TENSOR.OUTPUT_RES_SCORE,+ MODEL_TENSOR.ATTN_NORM,+ MODEL_TENSOR.ATTN_RES_SCORE,+ MODEL_TENSOR.FFN_RES_SCORE,+ # MLA + DSA indexer+ MODEL_TENSOR.ATTN_Q,+ MODEL_TENSOR.ATTN_K,+ MODEL_TENSOR.ATTN_V,+ MODEL_TENSOR.ATTN_OUT,+ MODEL_TENSOR.ATTN_GATE,+ MODEL_TENSOR.ATTN_Q_A,+ MODEL_TENSOR.ATTN_Q_B,+ MODEL_TENSOR.ATTN_KV_A_MQA,+ MODEL_TENSOR.ATTN_KV_B,+ MODEL_TENSOR.ATTN_K_B,+ MODEL_TENSOR.ATTN_V_B,+ MODEL_TENSOR.ATTN_Q_A_NORM,+ MODEL_TENSOR.ATTN_KV_A_NORM,+ MODEL_TENSOR.INDEXER_ATTN_Q_B,+ MODEL_TENSOR.INDEXER_ATTN_K,+ MODEL_TENSOR.INDEXER_K_NORM,+ MODEL_TENSOR.INDEXER_PROJ,+ # KDA (linear-attention layers)+ MODEL_TENSOR.SSM_CONV1D_Q,+ MODEL_TENSOR.SSM_CONV1D_K,+ MODEL_TENSOR.SSM_CONV1D_V,+ MODEL_TENSOR.SSM_F_A,+ MODEL_TENSOR.SSM_F_B,+ MODEL_TENSOR.SSM_BETA,+ MODEL_TENSOR.SSM_A,+ MODEL_TENSOR.SSM_G,+ MODEL_TENSOR.SSM_DT,+ MODEL_TENSOR.SSM_NORM,+ # FFN+ MODEL_TENSOR.FFN_NORM,+ MODEL_TENSOR.FFN_GATE,+ MODEL_TENSOR.FFN_DOWN,+ MODEL_TENSOR.FFN_UP,+ MODEL_TENSOR.FFN_GATE_INP,+ MODEL_TENSOR.FFN_EXP_PROBS_B,+ MODEL_TENSOR.FFN_GATE_EXP,+ MODEL_TENSOR.FFN_DOWN_EXP,+ MODEL_TENSOR.FFN_UP_EXP,+ MODEL_TENSOR.FFN_GATE_SHEXP,+ MODEL_TENSOR.FFN_DOWN_SHEXP,+ MODEL_TENSOR.FFN_UP_SHEXP,+ MODEL_TENSOR.FFN_ROUTED_DOWN,+ MODEL_TENSOR.FFN_ROUTED_UP,+ MODEL_TENSOR.FFN_ROUTED_NORM,+ ], MODEL_ARCH.TALKIE: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT,src/llama-arch.cpp4 + / 0 −
@@ -168,6 +168,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = { { LLM_ARCH_MAION_CODER, "maion-coder" }, { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, { LLM_ARCH_KIMI_K3, "kimi-k3" },+ { LLM_ARCH_KIMI_K31, "kimi-k31" }, { LLM_ARCH_GLM5_NEXT, "glm5-next" }, { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_MELLUM, "mellum" },@@ -1364,6 +1365,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_KIMI_K3:+ case LLM_ARCH_KIMI_K31: case LLM_ARCH_GLM5_NEXT: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE:@@ -1393,6 +1395,7 @@ bool llm_arch_is_diffusion(const llm_arch & arch) { bool llm_arch_supports_rs_rollback(const llm_arch & arch) { switch (arch) { case LLM_ARCH_KIMI_K3:+ case LLM_ARCH_KIMI_K31: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_CLOVER_MOE:@@ -1465,6 +1468,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_KIMI_K3:+ case LLM_ARCH_KIMI_K31: case LLM_ARCH_QWEN3TTS: case LLM_ARCH_LIMITE: return false;src/llama-arch.h1 + / 0 −
@@ -166,6 +166,7 @@ enum llm_arch { LLM_ARCH_MAION_CODER, LLM_ARCH_KIMI_LINEAR, LLM_ARCH_KIMI_K3,+ LLM_ARCH_KIMI_K31, LLM_ARCH_GLM5_NEXT, LLM_ARCH_TALKIE, LLM_ARCH_MELLUM,src/llama-context.cpp1 + / 1 −
@@ -2724,7 +2724,7 @@ 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_KIMI_K3 || model.arch == LLM_ARCH_GLM5_NEXT) {+ if (model.arch == LLM_ARCH_KIMI_K3 || model.arch == LLM_ARCH_KIMI_K31 || model.arch == LLM_ARCH_GLM5_NEXT) { // the n_tokens*40 budget below is exhausted at ubatch 3840 res = std::max<uint32_t>(n_tokens * 160, 64u * model.n_tensors()); } else if (model.arch == LLM_ARCH_HRM_TEXT) {src/llama-graph.cpp36 + / 30 −
@@ -3456,6 +3456,41 @@ ggml_tensor * llm_graph_context::build_fused_top_k_mask( return cur; } +ggml_tensor * llm_graph_context::build_top_k_mask(+ ggml_tensor * kq_mask,+ ggml_tensor * top_k,+ int il) const {+ if (cparams.fused_top_k_mask) {+ return build_fused_top_k_mask(kq_mask, top_k, il);+ }++ // prepare new kq mask - starts filled with -INFINITY+ ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);++ // reshape KQ mask into tensor with rows of size 1:+ // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]+ kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);++ // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]+ ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);++ // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]+ // this will be our source of zero values for unmasking top k mask elements+ ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);+ zeros = ggml_fill(ctx0, zeros, 0.0f);++ // modify KQ mask by unmasking elements that are in top_k indices+ // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])+ ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);++ // reshape to restore the original shape of KQ mask:+ // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]+ kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);++ // combine with the original kq mask+ return ggml_add(ctx0, kq_mask_top_k, kq_mask);+}+ ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_k_dsa * inp, ggml_tensor * wo,@@ -3500,36 +3535,7 @@ ggml_tensor * llm_graph_context::build_attn( kq_mask = inp->self_attn_mask_lid.raw; } - ggml_tensor * kq_mask_top_k = nullptr;- if (cparams.fused_top_k_mask) {- kq_mask_top_k = build_fused_top_k_mask(kq_mask, top_k, il);- } else {- // prepare new kq mask - starts filled with -INFINITY- ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);-- // reshape KQ mask into tensor with rows of size 1:- // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream]- kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);-- // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1]- ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);-- // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream]- // this will be our source of zero values for unmasking top k mask elements- ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);- zeros = ggml_fill(ctx0, zeros, 0.0f);-- // modify KQ mask by unmasking elements that are in top_k indices- // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1])- kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);-- // reshape to restore the original shape of KQ mask:- // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream]- kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);-- // combine with the original kq mask- kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);- }+ ggml_tensor * kq_mask_top_k = build_top_k_mask(kq_mask, top_k, il); ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il);src/llama-graph.h6 + / 0 −
@@ -1357,6 +1357,12 @@ struct llm_graph_context { ggml_tensor * top_k, // [n_top, n_tokens] int il) const; + // kq_mask with only the cells named by top_k left visible, for sparse attention+ ggml_tensor * build_top_k_mask(+ ggml_tensor * kq_mask, // [n_kv, n_tokens/n_stream, 1, n_stream]+ ggml_tensor * top_k, // [n_top, n_tokens]+ int il) const;+ llm_graph_input_attn_kv_iswa * build_attn_inp_kv_iswa() const; llm_graph_input_dsv4 * build_inp_dsv4() const;src/llama-model.cpp17 + / 2 −
@@ -372,6 +372,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_kimi_linear(params); case LLM_ARCH_KIMI_K3: return new llama_model_kimi_k3(params);+ case LLM_ARCH_KIMI_K31:+ return new llama_model_kimi_k31(params); case LLM_ARCH_GLM5_NEXT: return new llama_model_glm5_next(params); case LLM_ARCH_STEP35:@@ -3282,8 +3284,20 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // only the sparse-attention architectures use llama_memory_hybrid_idx // a null filter_idx means the GGUF has no indexer tensors llama_memory_hybrid::layer_filter_cb filter_idx = nullptr;- const bool needs_mem_idx = (arch == LLM_ARCH_QWEN4EXP);- if (arch == LLM_ARCH_FALCON_H1) {+ const bool needs_mem_idx = (arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_KIMI_K31);+ if (arch == LLM_ARCH_KIMI_K31) {+ // DSA layers keep a K-only MLA cache plus a per-token indexer key cache,+ // the KDA layers stay recurrent+ filter_attn = [&](uint32_t il) {+ return il < hparams.n_layer() && !hparams.is_recr(il);+ };+ filter_recr = [&](uint32_t il) {+ return il < hparams.n_layer() && hparams.is_recr(il);+ };+ filter_idx = [&](uint32_t il) {+ return il < hparams.n_layer() && !hparams.is_recr(il) && hparams.is_indexer_full(il);+ };+ } else if (arch == LLM_ARCH_FALCON_H1) { filter_attn = [&](uint32_t) { return true; }; filter_recr = [&](uint32_t) { return true; }; } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {@@ -3660,6 +3674,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_METIS16: case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_KIMI_K3:+ case LLM_ARCH_KIMI_K31: case LLM_ARCH_GLM5_NEXT: case LLM_ARCH_STARCODER: return LLAMA_ROPE_TYPE_NONE;src/llama-quant.cpp7 + / 0 −
@@ -371,6 +371,13 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param quantize &= name.find("hc_") == std::string::npos; } + // kimi-k31: the indexer picks the top-k cells, so keep its projections accurate+ if (arch == LLM_ARCH_KIMI_K31) {+ quantize &= name.find("indexer.attn_q_b") == std::string::npos;+ quantize &= name.find("indexer.attn_k") == std::string::npos;+ quantize &= name.find("indexer.proj") == std::string::npos;+ }+ // glm5-next if (arch == LLM_ARCH_GLM5_NEXT) { quantize &= name.find("hc_") == std::string::npos;src/models/kimi-k31.cppadded701 + / 0 −
@@ -0,0 +1,701 @@+#include "models.h"++#include <algorithm>+#include "llama-memory-recurrent.h"+#include "llama-memory-hybrid-idx.h"++//+// Kimi-K3.1 text model: Kimi-K3 with DSA (indexer + top-k) on the full-attention layers.+// KDA layers, cross-layer residuals, latent MoE, situ and the MLA shapes are unchanged.+// 1. full-attention layers run a lightning indexer over the cached latent, keep the+// top-k cells and attend only to those (the "Gated DSA" of the model card)+// 2. the indexer keeps its own per-token key cache, one cell per attention cell+//++void llama_model_kimi_k31::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_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);+ 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);++ // DSA indexer, present on the dense-attention layers only+ hparams.f_norm_eps = 1e-5f; // eps of the indexer key LayerNorm+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps, false);+ ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);+ ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);+ ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);+ // every DSA layer scores its own selection unless the GGUF says otherwise+ std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1);+ ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false);++ if (hparams.indexer_top_k > 0) {+ GGML_ASSERT(hparams.indexer_n_head > 0 && hparams.indexer_head_size > 0);+ GGML_ASSERT(hparams.n_lora_q > 0 && "the indexer queries come from the MLA q latent");+ }++ switch (hparams.n_layer()) {+ case 93: type = LLM_TYPE_2_8T_A50B; break; // Kimi-K3.1-Pro+ default: type = LLM_TYPE_UNKNOWN;+ }+}++void llama_model_kimi_k31::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_NOSCAN, 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);++ // DSA lightning indexer, present on the layers the GGUF marks as indexer layers+ if (hparams.indexer_top_k > 0 && hparams.is_indexer_full(i)) {+ const int64_t n_idx_head = hparams.indexer_n_head;+ const int64_t n_idx_dim = hparams.indexer_head_size;++ layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, n_idx_head * n_idx_dim}, 0);+ layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, n_idx_dim}, 0);+ layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_idx_dim}, 0);+ layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {n_idx_dim}, 0);+ layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, n_idx_head}, 0);+ }++ 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<llm_graph_context> llama_model_kimi_k31::build_arch_graph(const llm_graph_params & params) const {+ return std::make_unique<graph>(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_k31_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_k31::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_k31::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_k31::graph::graph(const llm_graph_params & params) :+ llm_build_delta_net_base(params) {++ ggml_tensor * cur;+ ggml_tensor * inpL;++ inpL = build_inp_embd(model.tok_embd);+ cb(inpL, "inp_embd", -1);++ // K3.1 MLA is nope-only, so there is no position input++ auto * inp_k = build_inp_mem_hybrid_k();+ auto * inp_rs = inp_k->get_recr();+ auto * inp_attn = inp_k->get_attn();++ // the indexer cache sits next to the attention cache and shares its cells+ const auto * mctx_hyb = static_cast<const llama_memory_hybrid_idx_context *>(inp_k->mctx);+ const auto * mctx_idx = mctx_hyb->get_idx();+ if (mctx_idx) {+ GGML_ASSERT(mctx_idx->get_n_kv() == inp_attn->mctx->get_n_kv() &&+ "the indexer cache must track the attention cache cell for cell");+ }++ 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 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;++ ggml_tensor * prev_sel = nullptr;++ 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 {+ dsa_out dsa = build_dsa_layer(cur, layer, mctx_idx, inp_attn, prev_sel,+ n_embd_head_k_mla, kv_lora_rank,+ n_embd_head_qk_rope, n_embd_head_qk_nope, kq_scale_mla, il);+ cur = dsa.out;+ prev_sel = dsa.sel;+ }++ 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_k31_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;+}++//+// KDA layer+//++// causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use+ggml_tensor * llama_model_kimi_k31::graph::build_kda_conv1d(+ llm_graph_input_rs * inp_rs, 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_kda, int64_t n_seq_tokens, int64_t n_seqs) {+ const int64_t d_inner = head_dim * n_head_kda;+ 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);++ build_conv_state_store(inp_rs, conv_states_all, conv_x, d_conv - 1,+ ggml_row_size(conv_states_all->type, qkv * conv_state_size));++ 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_kda, n_seq_tokens, n_seqs);+}++ggml_tensor * llama_model_kimi_k31::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;++ 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 = build_kda_conv1d(inp_rs, 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);+ ggml_tensor * Kcur = build_kda_conv1d(inp_rs, 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);+ ggml_tensor * Vcur = build_kda_conv1d(inp_rs, 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);+ 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_norm = hparams.f_norm_rms_eps;+ Qcur = build_gdn_l2_norm(ctx0, Qcur, eps_norm);+ Kcur = build_gdn_l2_norm(ctx0, Kcur, eps_norm);++ ggml_tensor * output = build_recurrent_attn(inp_rs, ssm_states_all, Qcur, Kcur, Vcur, g1, beta, state, il);+ output = ggml_cont(ctx0, output);+ cb(output, "kda_scan_out", il);++ // 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;+}++//+// Gated DSA layer: nope MLA over the top-k cells picked by the lightning indexer+//++llama_model_kimi_k31::graph::dsa_out+llama_model_kimi_k31::graph::build_dsa_layer(+ ggml_tensor * cur, const llama_layer & layer,+ const llama_kv_cache_context * mctx_idx, llm_graph_input_attn_k * inp_attn,+ ggml_tensor * prev_sel,+ int64_t n_embd_head_k_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) {++ const auto * mctx_mla = inp_attn->mctx;++ const int64_t n_idx_head = hparams.indexer_n_head;+ const int64_t n_idx_dim = hparams.indexer_head_size;++ ggml_tensor * inp_gate = cur; // the output gate reads the *normed* layer input++ // --- MLA query; the indexer reads the same q latent ---+ GGML_ASSERT(layer.wq_a && layer.wq_b && "K3.1 MLA is served from the q latent");++ ggml_tensor * qr = ggml_mul_mat(ctx0, layer.wq_a, cur);+ qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);+ ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.wq_b, qr);++ 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);++ // the cache holds the compressed latent, so the MLA is served absorbed+ GGML_ASSERT(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);++ // --- DSA: keep only the top-k cells of this layer ---+ ggml_tensor * sel = prev_sel;+ if (hparams.is_indexer_full(il) && layer.indexer_attn_q_b) {+ GGML_ASSERT(mctx_idx != nullptr);++ ggml_tensor * iq = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr);+ iq = ggml_reshape_3d(ctx0, iq, n_idx_dim, n_idx_head, n_tokens);++ ggml_tensor * ik = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur);+ ik = build_norm(ik, layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, il);+ ik = ggml_reshape_3d(ctx0, ik, n_idx_dim, 1, n_tokens);+ cb(ik, "indexer_k", il);++ // the indexer keys go to the cells of the matching attention layer+ ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, ik, inp_attn->get_k_idxs(), il));++ ggml_tensor * iw = ggml_mul_mat(ctx0, layer.indexer_proj, cur);+ // pre-scale so the huge score tensor is not scaled afterwards+ iw = ggml_scale(ctx0, iw, 1.0f / sqrtf(float(n_idx_dim * n_idx_head)));++ ik = mctx_idx->get_k(ctx0, il);++ // split the ubatch into its streams, k carries one row per stream+ const int64_t n_stream = ik->ne[3];+ iq = ggml_view_4d(ctx0, iq, n_idx_dim, n_idx_head, n_tokens/n_stream, n_stream,+ iq->nb[1], iq->nb[2], iq->nb[3]/n_stream, 0);+ iw = ggml_view_4d(ctx0, iw, n_idx_head, n_tokens/n_stream, 1, n_stream,+ iw->nb[1], iw->nb[2]/n_stream, iw->nb[2]/n_stream, 0);++ ggml_tensor * kq_mask_lid = inp_attn->self_attn_mask.raw;+ if (cparams.fused_lid && kq_mask_lid->type != GGML_TYPE_F16) {+ kq_mask_lid = ggml_cast(ctx0, kq_mask_lid, GGML_TYPE_F16);+ }++ ggml_tensor * score;+ if (cparams.fused_lid) {+ score = ggml_lightning_indexer(ctx0, iq, ik, iw, kq_mask_lid);+ res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, score, il});+ } else {+ ggml_tensor * q_p = ggml_permute(ctx0, iq, 0, 2, 1, 3); // [n_idx_dim, n_tokens, n_idx_head]+ ggml_tensor * k_p = ggml_permute(ctx0, ik, 0, 2, 1, 3); // [n_idx_dim, n_kv, 1]++ ggml_tensor * kq = ggml_mul_mat(ctx0, k_p, q_p); // [n_kv, n_tokens, n_idx_head]+ kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); // [n_idx_head, n_tokens, n_kv]++ score = ggml_relu(ctx0, kq);+ score = ggml_mul(ctx0, score, iw);+ score = ggml_sum_rows(ctx0, score); // [1, n_tokens, n_kv]+ score = ggml_cont(ctx0, ggml_permute(ctx0, score, 2, 1, 0, 3)); // [n_kv, n_tokens, 1]++ score = ggml_add(ctx0, score, kq_mask_lid);+ }+ cb(score, "indexer_score", il);++ const int64_t n_top_k = std::min<int64_t>(score->ne[0], hparams.indexer_top_k);+ ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, score, n_top_k));+ cb(top_k, "indexer_top_k", il);++ sel = build_top_k_mask(inp_attn->self_attn_mask.raw, top_k, il);+ cb(sel, "kq_mask_dsa", il);+ } else {+ GGML_ASSERT(sel != nullptr && "shared indexer layer must follow a full indexer layer");+ }++ ggml_build_forward_expand(gf, Q);+ ggml_build_forward_expand(gf, K);+ ggml_build_forward_expand(gf, mctx_mla->cpy_k(ctx0, K, inp_attn->get_k_idxs(), il));++ ggml_tensor * k = mctx_mla->get_k(ctx0, il);+ ggml_tensor * v = ggml_view_4d(ctx0, k, kv_lora_rank, k->ne[1], k->ne[2], k->ne[3],+ k->nb[1], k->nb[2], k->nb[3], 0);++ // the selection is applied as an additive mask, so attention only gets cheaper when the backend+ // can skip masked columns: cuda flash attention gathers the visible cells when n_kv_max > 0 and+ // the cache holds at least 4096 cells. on the dense fallbacks the mask keeps the result correct+ // but the kq and kqv products still run over the whole cache, and the indexer is pure overhead+ const int64_t n_sel = std::min<int64_t>(mctx_mla->get_n_kv(), hparams.indexer_top_k);+ ggml_tensor * out = build_attn_mha(Q, k, v, nullptr, sel, nullptr, nullptr, layer.wv_b,+ n_sel, kq_scale, il);+ cb(out, "kqv_out", il);++ // K3.1: 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, sel };+}++//+// 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_k31::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_k31_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;+}src/models/models.h47 + / 0 −
@@ -2961,6 +2961,53 @@ struct llama_model_kimi_k3 : public llama_model_base { std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override; }; +// Kimi-K3.1: Kimi-K3 with a lightning indexer and top-k sparse attention on the full-attention layers+struct llama_model_kimi_k31 : public llama_model_base {+ llama_model_kimi_k31(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 llm_graph_params & params);++ // Cross-layer residual attention (kimi-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_conv1d(llm_graph_input_rs * inp_rs,+ 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_kda,+ int64_t n_seq_tokens, int64_t n_seqs);++ 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);++ struct dsa_out {+ ggml_tensor * out;+ ggml_tensor * sel;+ };++ dsa_out build_dsa_layer(ggml_tensor * cur, const llama_layer & layer,+ const llama_kv_cache_context * mctx_idx,+ llm_graph_input_attn_k * inp_attn,+ ggml_tensor * prev_sel,+ int64_t n_embd_head_k_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<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;+};+ struct llama_model_glm5_next : public llama_model_base { llama_model_glm5_next(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override;tests/fusion/MTL.csv6 + / 0 −
@@ -168,6 +168,12 @@ kimi-k3 ,0 ,any ,MUL+ADD , 2 kimi-k3 ,0 ,any ,RMS_NORM+MUL , 17 kimi-k3 ,0 ,any ,RMS_NORM+SCALE , 2 kimi-k3 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1+kimi-k31 ,0 ,any ,GATED_DELTA_NET+CPY , 1+kimi-k31 ,0 ,any ,MUL+ADD , 2+kimi-k31 ,0 ,any ,NORM+MUL+ADD , 1+kimi-k31 ,0 ,any ,RMS_NORM+MUL , 17+kimi-k31 ,0 ,any ,RMS_NORM+SCALE , 2+kimi-k31 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1 kimi-linear ,0 ,any ,ADD+ADD , 1 kimi-linear ,0 ,any ,GATED_DELTA_NET+CPY , 1 kimi-linear ,0 ,any ,MUL+ADD , 2tests/test-llama-archs.cpp5 + / 1 −
@@ -290,6 +290,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe, const || arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3+ || arch == LLM_ARCH_KIMI_K31 || arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_MISTRAL_LARGE4 || arch == LLM_ARCH_HY_V4@@ -365,7 +366,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe, const if (arch == LLM_ARCH_PLAMO2 || arch == LLM_ARCH_JAMBA || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE || arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_LFM3 ||- arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 || arch == LLM_ARCH_GLM5_NEXT+ arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 || arch == LLM_ARCH_KIMI_K31 ||+ arch == LLM_ARCH_GLM5_NEXT ) { GGML_ASSERT(n_layer >= 2); std::vector<uint32_t> n_head_per_layer;@@ -410,6 +412,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe, const || arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3+ || arch == LLM_ARCH_KIMI_K31 || arch == LLM_ARCH_GLM5_NEXT || arch == LLM_ARCH_HY_V4 || arch == LLM_ARCH_MISTRAL4@@ -1024,6 +1027,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_MIMO2: case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_KIMI_K3:+ case LLM_ARCH_KIMI_K31: case LLM_ARCH_GLM5_NEXT: case LLM_ARCH_STEP35: case LLM_ARCH_MISTRAL4: