Commit bcb015287 for llama.cpp
commit bcb0152876e8f0b684ba70d3defcab6cdf35adb5
Author: usrlocalben <benjamin@rqdq.com>
Date: Fri Nov 6 15:42:46 2026 -0700
model: add Kimi-K3.1-Pro text model (#31779)
* model: add Kimi-K3.1-Pro text model
K3 with the full-attention MLA layers replaced by DSA: on this checkpoint 24
of the 93 layers carry a lightning indexer (32 heads x 128, top_k 2048, no
rope). The KDA (linear-attention) layers and the rest of the K3 skeleton -
cross-layer residuals, latent MoE, situ, MLA output gate, full-rank KDA gate -
are unchanged, so this reuses the K3 graph and converter and adds the indexer.
Indexer:
- the selection is an additive mask over the whole cache, not a gather of
the top-k rows, so it reuses the existing sparse flash attention path;
off that path the mask stays correct but still costs a kq over the cache
- the indexer key LayerNorm takes its eps from the model's rms_norm_eps,
because dsa_config declares no eps of its own
Graph and cache:
- new arch with four indexer tensors per DSA layer, and a hybrid memory
that keeps a K-only MLA cache for those layers, recurrent state for the
KDA layers, and the indexer key cache
- the graph node budget shares the K3 entry
Conversion:
- KimiK31Model(KimiK3Model), picked by KimiK31ForConditionalGeneration
- dsa_config becomes attention.indexer.{head_count,key_length,top_k,types}
plus the indexer LayerNorm eps; the layer list is 1-indexed like
linear_attn_config's, and the converter asserts that it covers exactly
the full-attention layers
- the routed experts stay MXFP4 and are repacked as in K3, not dequantized
- the vision tower and mm_projector are skipped: text only for now
- the arch key follows the checkpoint's own model_type ("kimi_k31") and the
loader registers KimiK31ForConditionalGeneration; the release name carries
a -Pro suffix that nothing in the config or the tensor names uses
Quantisation: the indexer's q_b, k and proj are kept out of the quant
overrides, since the top-k selection is sensitive to their precision.
Verified on an 8-layer slice against the HF checkpoint: 38 tensors checked
exactly, the indexer metadata against the config, and the selection against
full attention - identical PPL at top_k beyond the context length, divergent
once top_k is forced below it. Arch test fixture and the full arch suite pass.
* converter: avoid clutter in the k3 base
* narrow comment to match model
* update metal fusion expectations
* restore
* sync with PR 31733 refactoring
diff --git a/conversion/__init__.py b/conversion/__init__.py
index fcbfb6a3a..de568f58d 100644
--- a/conversion/__init__.py
+++ b/conversion/__init__.py
@@ -158,6 +158,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"KORMoForCausalLM": "qwen",
"KimiK25ForConditionalGeneration": "deepseek",
"KimiK3ForConditionalGeneration": "kimi_k3",
+ "KimiK31ForConditionalGeneration": "kimi_k3",
"KimiLinearForCausalLM": "kimi_linear",
"KimiLinearModel": "kimi_linear",
"KimiVLForConditionalGeneration": "deepseek",
diff --git a/conversion/base.py b/conversion/base.py
index 841efe729..ecce23b9c 100644
--- a/conversion/base.py
+++ b/conversion/base.py
@@ -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 instead
diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py
index 70aabb705..6ea151f72 100644
--- a/conversion/kimi_k3.py
+++ b/conversion/kimi_k3.py
@@ -343,3 +343,33 @@ class KimiK3Model(TextModel):
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"]))
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index 69af1b909..cdcccd4be 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -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 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
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 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
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,
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index 0702fc72d..6dcf34e77 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -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;
diff --git a/src/llama-arch.h b/src/llama-arch.h
index 225803cb3..3b8f45d26 100644
--- a/src/llama-arch.h
+++ b/src/llama-arch.h
@@ -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,
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index ea20062a6..f6e7e9ca2 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -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) {
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index b99370e2d..d1ee7c090 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -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);
diff --git a/src/llama-graph.h b/src/llama-graph.h
index 0f8ad8ac9..6d25dab8c 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -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;
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 0b063e845..7e845bc30 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -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;
diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp
index 688a6f0cc..4ae396101 100644
--- a/src/llama-quant.cpp
+++ b/src/llama-quant.cpp
@@ -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;
diff --git a/src/models/kimi-k31.cpp b/src/models/kimi-k31.cpp
new file mode 100644
index 000000000..0c0ca2af2
--- /dev/null
+++ b/src/models/kimi-k31.cpp
@@ -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;
+}
diff --git a/src/models/models.h b/src/models/models.h
index 3e338c26b..f2c43b800 100644
--- a/src/models/models.h
+++ b/src/models/models.h
@@ -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;
diff --git a/tests/fusion/MTL.csv b/tests/fusion/MTL.csv
index 6d7599e4d..89e884183 100644
--- a/tests/fusion/MTL.csv
+++ b/tests/fusion/MTL.csv
@@ -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 , 2
diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp
index 64637ed8d..ea9afd660 100644
--- a/tests/test-llama-archs.cpp
+++ b/tests/test-llama-archs.cpp
@@ -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: