Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

model: add Kimi-K3 text model (#26185)

Browse Source 
* model: add Kimi-K3 text model

Hybrid KDA (linear) + MLA (full) attention as in Kimi-Linear-48B, plus five
things that architecture does not have:

  1. cross-layer residual attention  (attn_res_block_size)
  2. latent MoE                      (routed experts run at n_expert_latent)
  3. situ activation                 (replaces SwiGLU everywhere)
  4. MLA output gate                 (sigmoid gate before o_proj)
  5. full-rank KDA gate              (single ssm_g instead of ssm_g_a/ssm_g_b)

K3's text_config reports KimiLinearForCausalLM - the older 48B architecture -
so get_model_architecture routes on the top-level name instead.

The KDA decay gate has two forms, selected by linear_attn_config's
gate_lower_bound. It is not a clamp: when set it swaps the activation entirely
(fla/ops/kda/gate.py), from -exp(A_log)*softplus(x) to
lower_bound*sigmoid(exp(A_log)*x). K3 sets it to -5.0; kimi-linear leaves it
unset, so that path is unchanged.

Cross-layer residuals reuse ggml_dsv4_hc_pre for the weighted sum. That op is
CPU + CUDA only, so Metal/Vulkan will fall back per-node until those kernels
exist.

The routed experts ship as compressed-tensors "mxfp4-pack-quantized". That is
bit-compatible with ggml's MXFP4 - same E2M1 code assignment, same E8M0 scale
byte, only the nibble positions within a block differ - so they are repacked
rather than dequantized, losslessly and without a ~5.5 TB bf16 round-trip.
The repack is built lazily because gguf_writer holds every added tensor until
the final write. DeepSeek-V4 was already doing the identical bit-shuffling, so
it now shares the helper.

Verified against Moonshot's own code path (transformers + fla's Triton KDA
kernels) on a tiny model exercising every K3-specific feature. Final-position
logits vs the fp32 reference: 6.7e-05 rel / corr 1.00000000 for both the
chunked and the recurrent delta-net path. MXFP4 blocks dequantize to the source
weights with 0.0e+00 error.

Assisted-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* model: fix ty errors in the Kimi-K3 converter

- `_res_parts` buffers (kind, tensor) pairs, not bare tensors
- `get_tensors` must return an Iterator, matching ModelBase
- LazyBase's `func` takes one argument, so pass the expert loaders through
  `args` instead of the closure
- borrowing KimiLinearModel.set_vocab from an unrelated TextModel is
  deliberate and safe, but not expressible in the signature

No behaviour change: the MXFP4 repack still dequantizes to the source weights
with 0.0e+00 error and end-to-end logits are unchanged (8.386e-03 rel,
corr 0.99996630).

Assisted-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

* Update conversion/kimi_k3.py

Co-authored-by: Boris Dvorkin  <b_dvorkin@niuitmo.ru>

* Increase LLAMA_MAX_EXPERTS from 512 to 1024

* tests : support for Kimi K3 in archs test

* chat : add Kimi K3 chat format (reasoning, content, typed tool calls)

K3's assistant output is an XTML-ish tagged format built by the template's
open_tag/close_tag macros. Two properties break generic parsing:

1. The generation prompt ends with open_tag('think'), so the completion
   starts inside the think section with no opening marker in the output
   (thinking_forced_open).
2. Only <|open|>/<|close|>/<|sep|>/<|end_of_msg|> are special tokens; tag
   names ("think", "response", "message") are ordinary text tokens.

Adds common_chat_params_init_kimi_k3 (PEG_NATIVE) with detection on the
marker trio, reasoning extraction, response unwrapping, and tool-call
parsing of the tools/call/argument tag structure with argument types
taken from the tool schema. Includes the K3 chat template fixture and 9
test-chat cases derived from real generations of the full 2.8T model.

Verified end-to-end against Kimi-K3-Q2_K (GrEarl/Kimi-K3-GGUF) on 8x B200:
content, reasoning_content, streaming deltas, and tool_calls all correct;
finish_reason stop/tool_calls as appropriate.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* chat : add message_delimiters for Kimi K3

Per-role message-start markers for token-level span splitting. User and
assistant messages carry only the role attribute, so their full opener
(through <|sep|>) is used; system and tool messages continue with more
attributes (type=/tool=/index=), so those delimiters stop after the
role's closing quote. Verified against the K3 tiktoken vocabulary that
the closing quote is always a standalone token across all attribute
variants, so the token-level prefix match stays exact.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* fix: apply nits from @ngxson and text fixes from @danielhanchen

* tests : added missing hyperparameters and tensors for Kimi K3 in test-llama-archs

* chore : move overly verbose header file comments to Kimi K3 source file

* tests : re-enabled KIMI_K3 in test-llama-archs for WebGPU backend

* model-saver : emit kda_gate_lower_bound for Kimi K3

Quick fix. The Kimi K3 loader reads kda_gate_lower_bound and gates a graph branch on it (it scales the KDA gate when the bound is above -INFINITY), but the model
saver never wrote the key, so a save->load roundtrip silently dropped it back to the -INFINITY default and changed the model's output. The real K3 config sets gate_lower_bound = -5.0.

I propose to emit it from the saver, and set it to -5.0 in the test-llama-archs K3 case so the roundtrip check exercises it (the roundtrip fails without the saver line).

* Refactor conditional for model architecture check

* tests : re-enabled (again) KIMI_K3 and MINIMAX_M3 in test-llama-archs for WebGPU backend

* fix code comments

* add template on conversion

* move repack_mxfp4_blocks to model base

* nits

* add_value_length

* optimize res_stack construction

* nits

---------

Co-authored-by: Boris Dvorkin <b_dvorkin@niuitmo.ru>
Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com>
Co-authored-by: Deepankar Singh <singh.deepankar39@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Caleb DeLeeuw <caleb.deleeuw@gmail.com>
Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
This commit is contained in:
Piotr Wilkin (ilintar)
2026-08-15 17:11:05 +02:00
committed by GitHub
co-authored by Boris Dvorkin Stanisław Szymczyk Deepankar Singh Claude Fable 5 Caleb DeLeeuw Xuan Son Nguyen
parent 22b8e310b9
commit ad1de39e07
22 changed files with 1875 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-arch.cpp
+22
View File
@@ -144,6 +144,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_LLAMA_EMBED, "llama-embed" },
{ LLM_ARCH_MAINCODER, "maincoder" },
{ LLM_ARCH_KIMI_LINEAR, "kimi-linear" },
{ LLM_ARCH_KIMI_K3, "kimi-k3" },
{ LLM_ARCH_TALKIE, "talkie" },
{ LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_NANBEIGE, "nanbeige" },
@@ -187,6 +188,9 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_FEATURES_LENGTH, "%s.features_length" },
{ LLM_KV_BLOCK_COUNT, "%s.block_count" },
{ LLM_KV_LEADING_DENSE_BLOCK_COUNT, "%s.leading_dense_block_count" },
{ LLM_KV_ATTN_RES_BLOCK_SIZE, "%s.attn_res.block_size" },
{ LLM_KV_ACTIVATION_SITU_BETA, "%s.activation.situ_beta" },
{ LLM_KV_ACTIVATION_SITU_LINEAR_BETA, "%s.activation.situ_linear_beta" },
{ LLM_KV_FEED_FORWARD_LENGTH, "%s.feed_forward_length" },
{ LLM_KV_EXPERT_FEED_FORWARD_LENGTH, "%s.expert_feed_forward_length" },
{ LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, "%s.expert_shared_feed_forward_length" },
@@ -202,6 +206,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_EXPERT_GROUP_USED_COUNT, "%s.expert_group_used_count" },
{ LLM_KV_EXPERT_WEIGHTS_SCALE, "%s.expert_weights_scale" },
{ LLM_KV_EXPERT_WEIGHTS_NORM, "%s.expert_weights_norm" },
{ LLM_KV_EXPERT_LATENT_LENGTH, "%s.expert_latent_length" },
{ LLM_KV_EXPERT_GATING_FUNC, "%s.expert_gating_func" },
{ LLM_KV_EXPERT_GROUP_SCALE, "%s.expert_group_scale" },
{ LLM_KV_EXPERTS_PER_GROUP, "%s.experts_per_group" },
@@ -313,6 +318,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_SSM_DT_B_C_RMS, "%s.ssm.dt_b_c_rms" },
{ LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" },
{ LLM_KV_KDA_GATE_LOWER_BOUND, "%s.kda.gate_lower_bound" },
{ LLM_KV_WKV_HEAD_SIZE, "%s.wkv.head_size" },
@@ -463,6 +469,13 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_SSM_F_B, "blk.%d.ssm_f_b" },
{ LLM_TENSOR_SSM_BETA, "blk.%d.ssm_beta" },
{ LLM_TENSOR_SSM_G_A, "blk.%d.ssm_g_a" },
{ LLM_TENSOR_SSM_G, "blk.%d.ssm_g" },
{ LLM_TENSOR_ATTN_RES_SCORE, "blk.%d.attn_res_score" },
{ LLM_TENSOR_FFN_RES_SCORE, "blk.%d.ffn_res_score" },
{ LLM_TENSOR_OUTPUT_RES_SCORE, "output_res_score" },
{ LLM_TENSOR_FFN_ROUTED_DOWN, "blk.%d.ffn_routed_down" },
{ LLM_TENSOR_FFN_ROUTED_UP, "blk.%d.ffn_routed_up" },
{ LLM_TENSOR_FFN_ROUTED_NORM, "blk.%d.ffn_routed_norm" },
{ LLM_TENSOR_SSM_G_B, "blk.%d.ssm_g_b" },
{ LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" },
{ LLM_TENSOR_ATTN_Q_A_NORM, "blk.%d.attn_q_a_norm" },
@@ -756,6 +769,13 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_SSM_F_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_SSM_BETA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_SSM_G_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_SSM_G, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_ATTN_RES_SCORE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_FFN_RES_SCORE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_OUTPUT_RES_SCORE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}},
{LLM_TENSOR_FFN_ROUTED_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_FFN_ROUTED_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_FFN_ROUTED_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_SSM_G_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_TIME_MIX_LERP_X, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_TIME_MIX_LN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
@@ -976,6 +996,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) {
case LLM_ARCH_NEMOTRON_H_MOE:
case LLM_ARCH_QWEN3NEXT:
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_KIMI_K3:
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_DEEPSEEK4:
@@ -1040,6 +1061,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
case LLM_ARCH_MINIMAX_M3:
case LLM_ARCH_MISTRAL4:
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_KIMI_K3:
case LLM_ARCH_QWEN3TTS:
return false;
default:
src/llama-arch.h
+13
View File
@@ -145,6 +145,7 @@ enum llm_arch {
LLM_ARCH_LLAMA_EMBED,
LLM_ARCH_MAINCODER,
LLM_ARCH_KIMI_LINEAR,
LLM_ARCH_KIMI_K3,
LLM_ARCH_TALKIE,
LLM_ARCH_MELLUM,
LLM_ARCH_EAGLE3,
@@ -192,6 +193,9 @@ enum llm_kv {
LLM_KV_FEATURES_LENGTH,
LLM_KV_BLOCK_COUNT,
LLM_KV_LEADING_DENSE_BLOCK_COUNT,
LLM_KV_ATTN_RES_BLOCK_SIZE,
LLM_KV_ACTIVATION_SITU_BETA,
LLM_KV_ACTIVATION_SITU_LINEAR_BETA,
LLM_KV_FEED_FORWARD_LENGTH,
LLM_KV_EXPERT_FEED_FORWARD_LENGTH,
LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH,
@@ -207,6 +211,7 @@ enum llm_kv {
LLM_KV_EXPERT_GROUP_USED_COUNT,
LLM_KV_EXPERT_WEIGHTS_SCALE,
LLM_KV_EXPERT_WEIGHTS_NORM,
LLM_KV_EXPERT_LATENT_LENGTH,
LLM_KV_EXPERT_GATING_FUNC,
LLM_KV_EXPERT_GROUP_SCALE,
LLM_KV_EXPERTS_PER_GROUP,
@@ -318,6 +323,7 @@ enum llm_kv {
LLM_KV_SSM_DT_B_C_RMS,
LLM_KV_KDA_HEAD_DIM,
LLM_KV_KDA_GATE_LOWER_BOUND,
LLM_KV_WKV_HEAD_SIZE,
@@ -492,6 +498,13 @@ enum llm_tensor {
LLM_TENSOR_SSM_BETA, // kimi: beta mixing coefficient and qwen3.5
LLM_TENSOR_SSM_G_A, // kimi: output gate projection A
LLM_TENSOR_SSM_G_B, // kimi: output gate projection B
LLM_TENSOR_SSM_G, // kimi-k3: full-rank KDA gate
LLM_TENSOR_ATTN_RES_SCORE, // kimi-k3: fused res_norm*res_proj (pre-attn)
LLM_TENSOR_FFN_RES_SCORE, // kimi-k3: fused res_norm*res_proj (pre-ffn)
LLM_TENSOR_OUTPUT_RES_SCORE, // kimi-k3: fused res_norm*res_proj (final)
LLM_TENSOR_FFN_ROUTED_DOWN, // kimi-k3: latent MoE down
LLM_TENSOR_FFN_ROUTED_UP, // kimi-k3: latent MoE up
LLM_TENSOR_FFN_ROUTED_NORM, // kimi-k3: latent MoE norm
LLM_TENSOR_TIME_MIX_W0,
LLM_TENSOR_TIME_MIX_W1,
LLM_TENSOR_TIME_MIX_W2,
src/llama-context.cpp
+4 -1
View File
@@ -2293,7 +2293,10 @@ void llama_context::output_reorder() {
uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
uint32_t res;
if (model.arch == LLM_ARCH_QWEN3NEXT ||
if (model.arch == LLM_ARCH_KIMI_K3) {
// the n_tokens*40 budget below is exhausted at ubatch 3840
res = std::max<uint32_t>(n_tokens * 160, 64u * model.n_tensors());
} else if (model.arch == LLM_ARCH_QWEN3NEXT ||
model.arch == LLM_ARCH_KIMI_LINEAR ||
model.arch == LLM_ARCH_QWEN35 ||
model.arch == LLM_ARCH_QWEN35MOE ||
src/llama-graph.cpp
+17
View File
@@ -1835,6 +1835,8 @@ ggml_tensor * llm_graph_context::build_ffn(
cur = ggml_reglu(ctx0, cur);
cb(cur, "ffn_reglu", il);
} break;
case LLM_FFN_SITU:
GGML_ABORT("not yet supported");
default:
GGML_ABORT("fatal error");
}
@@ -2174,6 +2176,21 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
cur = ggml_silu(ctx0, cur);
cb(cur, "ffn_moe_silu", il);
} break;
case LLM_FFN_SITU:
{
// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * lb*tanh(up/lb)
GGML_ASSERT(has_gate);
const float beta = hparams.situ_beta;
const float lb = hparams.situ_linear_beta;
ggml_tensor * act = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, cur, 1.0f/beta)), beta);
act = ggml_mul(ctx0, act, ggml_sigmoid(ctx0, cur));
if (lb > 0.0f) {
up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/lb)), lb);
}
cur = ggml_mul(ctx0, act, up);
cb(cur, "ffn_moe_situ", il);
} break;
case LLM_FFN_GELU:
if (has_gate) {
cur = ggml_geglu_split(ctx0, cur, up);
src/llama-graph.h
+1
View File
@@ -59,6 +59,7 @@ enum llm_ffn_op_type : int {
LLM_FFN_GEGLU,
LLM_FFN_REGLU,
LLM_FFN_SWIGLU_OAI_MOE,
LLM_FFN_SITU, // kimi-k3
};
enum llm_ffn_gate_type {
src/llama-hparams.h
+9 -1
View File
@@ -4,10 +4,11 @@
#include <array>
#include <cassert>
#include <cmath>
// bump if necessary
#define LLAMA_MAX_LAYERS 512
#define LLAMA_MAX_EXPERTS 512 // Qwen3 Next
#define LLAMA_MAX_EXPERTS 1024 // Kimi K3
enum llama_expert_gating_func_type {
LLAMA_EXPERT_GATING_FUNC_TYPE_NONE = 0,
@@ -170,6 +171,13 @@ struct llama_hparams {
// for Kimi Linear KDA
uint32_t n_embd_head_kda = 0;
// kimi-k3
uint32_t n_expert_latent = 0; // routed_expert_hidden_size (0 = experts run at n_embd)
uint32_t attn_res_block_size = 0; // 0 = no cross-layer attention residuals
float kda_gate_lower_bound = -INFINITY;
float situ_beta = 1.0f;
float situ_linear_beta = 0.0f; // 0 = no linear-beta transform on the up branch
bool ssm_dt_b_c_rms = false;
float f_clamp_kqv = 0.0f;
src/llama-model-saver.cpp
+7
View File
@@ -213,6 +213,7 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead);
add_kv(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, true);
add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
add_kv(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent);
add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH, hparams.n_ff_chexp);
add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp);
@@ -319,6 +320,7 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_SSM_DT_B_C_RMS, hparams.ssm_dt_b_c_rms);
add_kv(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda);
add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound);
add_kv(LLM_KV_WKV_HEAD_SIZE, hparams.wkv_head_size);
@@ -376,6 +378,10 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_XIELU_BETA, hparams.xielu_beta);
add_kv(LLM_KV_XIELU_EPS, hparams.xielu_eps);
add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size);
add_kv(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta);
add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta);
// deprecated
// add_kv(LLM_KV_TOKENIZER_PREFIX_ID, ???);
// add_kv(LLM_KV_TOKENIZER_SUFFIX_ID, ???);
@@ -403,6 +409,7 @@ void llama_model_saver::add_tensors_from_model() {
add_tensor(model->output_norm_enc);
add_tensor(model->output_s);
add_tensor(model->output_in_s);
add_tensor(model->output_res_score);
add_tensor(model->cls);
add_tensor(model->cls_b);
add_tensor(model->cls_out);
src/llama-model.cpp
+4
View File
@@ -322,6 +322,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_mimo2(params);
case LLM_ARCH_KIMI_LINEAR:
return new llama_model_kimi_linear(params);
case LLM_ARCH_KIMI_K3:
return new llama_model_kimi_k3(params);
case LLM_ARCH_STEP35:
return new llama_model_step35(params);
default:
@@ -845,6 +847,7 @@ const char * llm_type_name(llm_type type) {
case LLM_TYPE_397B_A17B: return "397B.A17B";
case LLM_TYPE_685B_A37B: return "685B.A37B";
case LLM_TYPE_744B_A40B: return "744B.A40B";
case LLM_TYPE_2_8T_A50B: return "2.8T.A50B";
case LLM_TYPE_E2B: return "E2B";
case LLM_TYPE_E4B: return "E4B";
default: return "?B";
@@ -2602,6 +2605,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_NEMOTRON_H:
case LLM_ARCH_NEMOTRON_H_MOE:
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_KIMI_K3:
return LLAMA_ROPE_TYPE_NONE;
// use what we call a normal RoPE, operating on pairs of consecutive head values
src/llama-model.h
+10
View File
@@ -144,6 +144,7 @@ enum llm_type {
LLM_TYPE_397B_A17B, // Qwen3.5
LLM_TYPE_685B_A37B, // DeepSeek V3.2
LLM_TYPE_744B_A40B, // GLM-5
LLM_TYPE_2_8T_A50B, // Kimi-K3
LLM_TYPE_E2B,
LLM_TYPE_E4B,
};
@@ -530,6 +531,14 @@ struct llama_layer {
struct ggml_tensor * ssm_g_b = nullptr;
struct ggml_tensor * ssm_o_norm = nullptr;
// kimi-k3
struct ggml_tensor * ssm_g = nullptr; // full-rank KDA gate (replaces ssm_g_a/ssm_g_b)
struct ggml_tensor * attn_res_score = nullptr; // fused res_norm*res_proj, pre-attention
struct ggml_tensor * ffn_res_score = nullptr; // fused res_norm*res_proj, pre-FFN
struct ggml_tensor * ffn_routed_down = nullptr; // latent MoE: n_embd -> n_expert_latent
struct ggml_tensor * ffn_routed_up = nullptr; // latent MoE: n_expert_latent -> n_embd
struct ggml_tensor * ffn_routed_norm = nullptr;
// DSA (deepseek sparse attention)
struct ggml_tensor * indexer_k_norm = nullptr;
struct ggml_tensor * indexer_k_norm_b = nullptr;
@@ -589,6 +598,7 @@ struct llama_model {
struct ggml_tensor * tok_norm_b = nullptr;
struct ggml_tensor * output_norm = nullptr;
struct ggml_tensor * output_res_score = nullptr; // kimi-k3: final cross-layer residual mix
struct ggml_tensor * output_norm_b = nullptr;
struct ggml_tensor * output = nullptr;
struct ggml_tensor * output_b = nullptr;
src/models/kimi-k3.cpp
+614
View File
@@ -0,0 +1,614 @@
#include "models.h"
#include "llama-memory-recurrent.h"
//
// Kimi-K3 text model: hybrid KDA (linear) + MLA (full) attention, as in kimi-linear.
// Parts that kimi-linear does not have:
// 1. cross-layer residual attention (attn_res_block_size)
// 2. latent MoE (routed experts run at n_expert_latent)
// 3. situ activation (replaces SwiGLU everywhere)
// 4. MLA output gate (sigmoid gate before o_proj)
// 5. full-rank KDA gate (single ssm_g instead of ssm_g_a/ssm_g_b)
//
void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);
ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false);
ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);
ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);
ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda);
ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false);
// the MLA cache holds the compressed latent
// set it here too, as older GGUFs have no value_length key
hparams.n_embd_head_v_full = hparams.n_lora_kv;
// n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-linear
for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0;
}
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent, false);
ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size);
ml.get_key(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta);
ml.get_key(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta);
switch (hparams.n_layer()) {
case 93: type = LLM_TYPE_2_8T_A50B; break; // Kimi-K3
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_kimi_k3::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
if (hparams.attn_res_block_size > 0) {
output_res_score = create_tensor(tn(LLM_TENSOR_OUTPUT_RES_SCORE, "weight"), {n_embd}, 0);
}
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
if (hparams.attn_res_block_size > 0) {
layer.attn_res_score = create_tensor(tn(LLM_TENSOR_ATTN_RES_SCORE, "weight", i), {n_embd}, 0);
layer.ffn_res_score = create_tensor(tn(LLM_TENSOR_FFN_RES_SCORE, "weight", i), {n_embd}, 0);
}
const int64_t head_dim = hparams.n_embd_head_kda;
const int64_t d_conv = hparams.ssm_d_conv;
const int64_t d_inner = head_dim * n_head;
if (hparams.is_recr(i)) {
// conv1d may be stored 4D [d_conv, 1, d_inner, 1] or 3D (quantization drops the trailing 1)
auto conv = [&](llm_tensor tid) {
ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED);
return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0);
};
layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q);
layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K);
layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V);
create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0);
layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, 0);
layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0);
layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0);
// K3's A_log is a plain 1-D [n_head] tensor (kimi-linear's is padded)
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, 0);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0);
// K3 uses a single full-rank gate instead of kimi-linear's g_a/g_b pair
layer.ssm_g = create_tensor(tn(LLM_TENSOR_SSM_G, "weight", i), {n_embd, d_inner}, 0);
layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0);
} else {
const int64_t q_lora_rank = hparams.n_lora_q;
const int64_t kv_lora_rank = hparams.n_lora_kv;
const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();
const int64_t n_embd_head_v = hparams.n_embd_head_v_mla();
const int64_t qk_rope_head_dim = hparams.n_rot();
const int64_t qk_nope_head_dim = n_embd_head_k - qk_rope_head_dim;
layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, TENSOR_NOT_REQUIRED);
layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0);
if (layer.attn_q_a_norm) {
layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, 0);
} else {
layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k}, 0);
}
layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, 0);
layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i),
{kv_lora_rank, n_head * (qk_nope_head_dim + n_embd_head_v)},
TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);
if (!layer.wkv_b) {
layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {qk_nope_head_dim, kv_lora_rank, n_head}, 0);
layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, 0);
}
// K3: sigmoid output gate applied to the attention output before o_proj
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head * n_embd_head_v}, TENSOR_NOT_REQUIRED);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, 0);
}
if (i < (int) hparams.n_layer_dense_lead) {
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
} else {
const int64_t n_ff_exp = hparams.n_ff_exp;
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
// routed experts live in the latent space
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd_latent, n_expert}, 0);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0);
if (hparams.n_expert_latent > 0) {
layer.ffn_routed_down = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_DOWN, "weight", i), {n_embd, n_embd_latent}, 0);
layer.ffn_routed_up = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_UP, "weight", i), {n_embd_latent, n_embd}, 0);
layer.ffn_routed_norm = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_NORM, "weight", i), {n_embd_latent}, TENSOR_NOT_REQUIRED);
}
// shared experts stay at n_embd, width = moe_intermediate_size * n_expert_shared
const int64_t n_ff_shexp = n_ff_exp * (hparams.n_expert_shared > 0 ? hparams.n_expert_shared : 1);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED);
}
}
}
std::unique_ptr<llm_graph_context> llama_model_kimi_k3::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta)
// linear_beta <= 0 disables the transform on the up branch
static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_tensor * up,
float beta, float linear_beta) {
ggml_tensor * a = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, gate, 1.0f/beta)), beta);
a = ggml_mul(ctx0, a, ggml_sigmoid(ctx0, gate));
if (linear_beta > 0.0f) {
up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/linear_beta)), linear_beta);
}
return ggml_mul(ctx0, a, up);
}
//
// cross-layer residual attention
//
// layout is [n_embd, n_ckpt, n_tokens]: rms_norm reduces over ne0, dsv4_hc_pre over ne1
// append the new checkpoint, do not re-fold the whole chain
void llama_model_kimi_k3::graph::res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens) {
ggml_tensor * ckpt = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
resi_stack = resi_stack ? ggml_concat(ctx0, resi_stack, ckpt, 1) : ckpt;
}
ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor * score_w,
int64_t n_tokens, int il) {
if (!resi_stack) {
return cur; // layer 0: nothing banked yet
}
const int n_ckpt = (int) resi_stack->ne[1];
const float eps = hparams.f_norm_rms_eps;
ggml_tensor * src = resi_stack; // [n_embd, n_ckpt, n_tokens]
// one rms_norm scores all checkpoints at once
// note: the scores use the normalized values, but the sum below uses the raw ones
ggml_tensor * sc_src = ggml_rms_norm(ctx0, src, eps);
sc_src = ggml_mul(ctx0, sc_src, score_w);
sc_src = ggml_sum_rows(ctx0, sc_src); // [1, n_ckpt, n_tokens]
sc_src = ggml_reshape_2d(ctx0, sc_src, n_ckpt, n_tokens);
// the current residual stream is scored apart, so the stack stays append-only
ggml_tensor * sc_cur = ggml_rms_norm(ctx0, cur, eps);
sc_cur = ggml_mul(ctx0, sc_cur, score_w);
sc_cur = ggml_sum_rows(ctx0, sc_cur); // [1, n_tokens]
ggml_tensor * scores = ggml_concat(ctx0, sc_src, sc_cur, 0); // [n_ckpt+1, n_tokens]
ggml_tensor * probs = ggml_soft_max(ctx0, scores); // over ne0 = n_ckpt+1
cb(probs, "res_probs", il);
// split the sum: hc_pre handles the stack, a broadcast-multiply the current stream
ggml_tensor * p_src = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, n_ckpt, n_tokens, probs->nb[1], 0));
ggml_tensor * p_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, 1, n_tokens, probs->nb[1],
probs->nb[0] * n_ckpt));
ggml_tensor * out = ggml_dsv4_hc_pre(ctx0, src, p_src);
out = ggml_add(ctx0, out, ggml_mul(ctx0, cur, p_cur));
return out;
}
llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_params & params) :
llm_build_delta_net_base(params), model(model) {
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
cb(inpL, "inp_embd", -1);
// K3 MLA is nope-only, so there is no position input
auto * inp_kv = !hparams.is_mla() ? build_inp_mem_hybrid() : nullptr;
auto * inp_k = hparams.is_mla() ? build_inp_mem_hybrid_k() : nullptr;
auto * inp_rs = hparams.is_mla() ? inp_k->get_recr() : inp_kv->get_recr();
auto * inp_attn_kv = !hparams.is_mla() ? inp_kv->get_attn() : nullptr;
auto * inp_attn_k = hparams.is_mla() ? inp_k->get_attn() : nullptr;
ggml_tensor * inp_out_ids = build_inp_out_ids();
const int64_t n_head_kda = hparams.n_head();
const int64_t head_dim = hparams.n_embd_head_kda;
const int64_t d_conv = hparams.ssm_d_conv;
const int64_t d_inner = n_head_kda * head_dim;
const int64_t n_seqs = ubatch.n_seqs;
const int64_t n_seq_tokens = ubatch.n_seq_tokens;
GGML_ASSERT(n_seqs != 0);
GGML_ASSERT(ubatch.equal_seqs());
GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
const int64_t kv_lora_rank = hparams.n_lora_kv;
const int64_t n_embd_head_qk_rope = hparams.n_rot();
const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope;
const float kq_scale_mla = 1.0f / sqrtf((float) n_embd_head_k_mla);
const uint32_t res_bs = hparams.attn_res_block_size;
const bool use_attn_res = res_bs > 0;
const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd;
for (int il = 0; il < n_layer; ++il) {
const auto & layer = model.layers[il];
// the residual stream, banked on checkpoint layers and then restarted
// from the attention output alone
ggml_tensor * prefix_sum = inpL;
cur = use_attn_res ? res_mix(prefix_sum, layer.attn_res_score, n_tokens, il)
: prefix_sum;
bool banked = false;
if (use_attn_res && (uint32_t) il % res_bs == 0) {
res_push(prefix_sum, n_embd, n_tokens); // banks the RAW layer input, not `cur`
banked = true;
}
cur = build_norm(cur, layer.attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
ggml_build_forward_expand(gf, cur);
if (hparams.is_recr(il)) {
cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda,
d_inner, n_seq_tokens, n_seqs, il);
} else {
cur = build_mla_layer(cur, layer, inp_attn_k, inp_attn_kv,
n_embd_head_k_mla, n_embd_head_v_mla, kv_lora_rank,
n_embd_head_qk_rope, n_embd_head_qk_nope, kq_scale_mla, il);
}
prefix_sum = banked ? cur : ggml_add(ctx0, prefix_sum, cur);
cb(prefix_sum, "prefix_sum_attn", il);
cur = use_attn_res ? res_mix(prefix_sum, layer.ffn_res_score, n_tokens, il)
: prefix_sum;
cur = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
if ((uint32_t) il < hparams.n_layer_dense_lead) {
ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate, cur);
ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up, cur);
cur = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta);
cur = ggml_mul_mat(ctx0, layer.ffn_down, cur);
cb(cur, "ffn_out", il);
} else {
cur = build_latent_moe(cur, layer, n_embd_latent, il);
}
prefix_sum = ggml_add(ctx0, prefix_sum, cur);
prefix_sum = build_cvec(prefix_sum, il);
cb(prefix_sum, "l_out", il);
inpL = prefix_sum;
}
cur = inpL;
// final mix, then narrow to the output tokens
if (use_attn_res) {
cur = res_mix(cur, model.output_res_score, n_tokens, -1);
}
if (inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = ggml_mul_mat(ctx0, model.output, cur);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
//
// KDA layer
//
// causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use
static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
int64_t d_conv, int64_t head_dim, int64_t n_head,
int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head) {
const int64_t d_inner = head_dim * n_head;
const int64_t conv_state_size = (d_conv - 1) * d_inner;
const int64_t n_embd_r_total = 3 * conv_state_size;
ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs,
(d_conv - 1) * ggml_element_size(conv_state_all),
n_embd_r_total * ggml_element_size(conv_state_all),
qkv * conv_state_size * ggml_element_size(conv_state_all));
ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x);
ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0);
ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]);
ggml_build_forward_expand(gf,
ggml_cpy(ctx0, last_conv_x,
ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
(d_conv - 1) * ggml_element_size(conv_states_all),
n_embd_r_total * ggml_element_size(conv_states_all),
(kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);
Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens);
Xcur = ggml_silu(ctx0, Xcur);
return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs);
}
ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer(
ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs,
int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) {
const auto * mctx_cur = inp_rs->mctx;
const auto kv_head = mctx_cur->get_head();
ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
ggml_tensor * Qcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head);
ggml_tensor * Kcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head);
ggml_tensor * Vcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head);
cb(Qcur, "kda_q_conv", il);
cb(Kcur, "kda_k_conv", il);
cb(Vcur, "kda_v_conv", il);
// gate_lower_bound is not a clamp - when set, it swaps the decay gate activation:
// unset (kimi-linear): g = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias)
// set (K3, -5.0): g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias))
// ssm_a holds -exp(A_log) (folded at conversion time), so exp(A_log) == -ssm_a
ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a);
g1 = ggml_add(ctx0, g1, layer.ssm_dt_b);
ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1);
if (hparams.kda_gate_lower_bound > -INFINITY) {
g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens);
g1 = ggml_mul(ctx0, g1, A); // -exp(A_log) * (...)
g1 = ggml_sigmoid(ctx0, ggml_scale(ctx0, g1, -1.0f));
g1 = ggml_scale(ctx0, g1, hparams.kda_gate_lower_bound);
} else {
g1 = ggml_softplus(ctx0, g1);
g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens);
g1 = ggml_mul(ctx0, g1, A);
}
cb(g1, "kda_g1", il);
g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head_kda, n_seq_tokens, n_seqs);
ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur);
beta = ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs);
beta = ggml_sigmoid(ctx0, beta);
cb(beta, "kda_beta", il);
ggml_tensor * cur_3d = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);
ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);
ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs);
state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs);
const float eps = hparams.f_norm_rms_eps;
Qcur = ggml_l2_norm(ctx0, Qcur, eps);
Kcur = ggml_l2_norm(ctx0, Kcur, eps);
auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il);
ggml_tensor * output = ggml_cont(ctx0, attn_out.first);
cb(output, "kda_scan_out", il);
ggml_tensor * new_state = attn_out.second;
ggml_build_forward_expand(gf,
ggml_cpy(ctx0, new_state,
ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs,
kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all))));
// K3: single full-rank gate (kimi-linear factors this as g_b(g_a(x)))
ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur_3d, cur_3d->ne[0], n_seq_tokens * n_seqs);
ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g, cur_2d);
g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head_kda, n_seq_tokens * n_seqs);
ggml_tensor * o = ggml_reshape_3d(ctx0, output, head_dim, n_head_kda, n_seq_tokens * n_seqs);
ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il);
cb(g2, "kda_g2", il);
cb(normed, "kda_normed", il);
ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2));
gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens);
cur = ggml_mul_mat(ctx0, layer.wo, gated);
cb(cur, "kda_out", il);
return cur;
}
//
// MLA layer (nope-only, with K3's sigmoid output gate)
//
ggml_tensor * llama_model_kimi_k3::graph::build_mla_layer(
ggml_tensor * cur, const llama_layer & layer,
llm_graph_input_attn_k * inp_attn_k, llm_graph_input_attn_kv * inp_attn_kv,
int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla, int64_t kv_lora_rank,
int64_t n_embd_head_qk_rope, int64_t n_embd_head_qk_nope, float kq_scale, int il) {
ggml_tensor * inp_gate = cur; // the output gate reads the *normed* layer input
ggml_tensor * Qcur;
if (layer.wq_a) {
Qcur = ggml_mul_mat(ctx0, layer.wq_a, cur);
Qcur = build_norm(Qcur, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
Qcur = ggml_mul_mat(ctx0, layer.wq_b, Qcur);
} else {
Qcur = ggml_mul_mat(ctx0, layer.wq, cur);
}
ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);
ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
// no RoPE: mla_use_nope is asserted at conversion time
kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
ggml_tensor * out;
if (layer.wk_b && layer.wv_b) {
ggml_tensor * q_nope = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_nope, n_head, n_tokens,
ggml_row_size(Qcur->type, n_embd_head_k_mla),
ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, 0);
ggml_tensor * q_pe = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_rope, n_head, n_tokens,
ggml_row_size(Qcur->type, n_embd_head_k_mla),
ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head,
ggml_row_size(Qcur->type, n_embd_head_qk_nope));
q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
ggml_tensor * Q = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
ggml_tensor * kv_cmpr_3d = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
ggml_tensor * K = ggml_concat(ctx0, kv_cmpr_3d, k_pe, 0);
ggml_tensor * V = kv_cmpr_3d;
// wo == NULL: the output projection is applied after the gate below
out = build_attn(inp_attn_k, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, layer.wv_b, kq_scale, il);
} else {
ggml_tensor * Q = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens);
ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr);
const int64_t kv_per_head = n_embd_head_qk_nope + n_embd_head_v_mla;
ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens,
ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head), 0);
ggml_tensor * V = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens,
ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head),
ggml_row_size(kv->type, n_embd_head_qk_nope)));
ggml_tensor * k_pe_t = ggml_new_tensor_3d(ctx0, k_pe->type, n_embd_head_qk_rope, n_head, n_tokens);
ggml_tensor * K = ggml_concat(ctx0, ggml_repeat(ctx0, k_pe, k_pe_t), k_nope, 0);
out = build_attn(inp_attn_kv, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, nullptr, kq_scale, il);
}
// K3: attn_output *= sigmoid(g_proj(x)), then o_proj
if (layer.wqkv_gate) {
ggml_tensor * g = ggml_sigmoid(ctx0, ggml_mul_mat(ctx0, layer.wqkv_gate, inp_gate));
out = ggml_mul(ctx0, out, g);
cb(out, "mla_gated", il);
}
out = ggml_mul_mat(ctx0, layer.wo, out);
cb(out, "mla_out", il);
return out;
}
//
// latent MoE: down-project, run the routed experts in the latent space, norm, up-project;
// shared experts stay at n_embd and read the un-projected input.
//
ggml_tensor * llama_model_kimi_k3::graph::build_latent_moe(
ggml_tensor * cur, const llama_layer & layer, int64_t n_embd_latent, int il) {
ggml_tensor * identity = cur;
ggml_tensor * routed_in = layer.ffn_routed_down
? ggml_mul_mat(ctx0, layer.ffn_routed_down, cur)
: cur;
// the router scores the full-width input while the experts take the latent one,
// so the logits are computed here and passed to build_moe_ffn
ggml_tensor * logits = ggml_mul_mat(ctx0, layer.ffn_gate_inp, identity);
cb(logits, "ffn_moe_logits", il);
ggml_tensor * moe_out = build_moe_ffn(routed_in,
nullptr, // gate_inp unused: the logits above are passed instead
layer.ffn_up_exps,
layer.ffn_gate_exps,
layer.ffn_down_exps,
layer.ffn_exp_probs_b,
hparams.n_expert,
hparams.n_expert_used,
LLM_FFN_SITU, hparams.expert_weights_norm,
hparams.expert_weights_scale,
(llama_expert_gating_func_type) hparams.expert_gating_func,
il,
logits);
cb(moe_out, "ffn_moe_out", il);
if (layer.ffn_routed_norm) {
moe_out = build_norm(moe_out, layer.ffn_routed_norm, NULL, LLM_NORM_RMS, il);
}
if (layer.ffn_routed_up) {
moe_out = ggml_mul_mat(ctx0, layer.ffn_routed_up, moe_out);
}
GGML_UNUSED(n_embd_latent);
if (layer.ffn_gate_shexp) {
ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate_shexp, identity);
ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up_shexp, identity);
ggml_tensor * sh = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta);
sh = ggml_mul_mat(ctx0, layer.ffn_down_shexp, sh);
cb(sh, "ffn_shexp", il);
moe_out = ggml_add(ctx0, moe_out, sh);
}
cb(moe_out, "ffn_out", il);
return moe_out;
}
src/models/models.h
+36
View File
@@ -2285,6 +2285,42 @@ struct llama_model_mimo2 : public llama_model_base {
};
struct llama_model_kimi_k3 : public llama_model_base {
llama_model_kimi_k3(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_build_delta_net_base {
graph(const llama_model & model, const llm_graph_params & params);
const llama_model & model;
// Cross-layer residual attention (K3's `_apply_attn_res`).
ggml_tensor * resi_stack = nullptr;
void res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens);
ggml_tensor * res_mix(ggml_tensor * cur, ggml_tensor * score_w,
int64_t n_tokens, int il);
ggml_tensor * build_kda_layer(ggml_tensor * cur, const llama_layer & layer,
llm_graph_input_rs * inp_rs,
int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il);
ggml_tensor * build_mla_layer(ggml_tensor * cur, const llama_layer & layer,
llm_graph_input_attn_k * inp_attn_k,
llm_graph_input_attn_kv * inp_attn_kv,
int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla,
int64_t kv_lora_rank, int64_t n_embd_head_qk_rope,
int64_t n_embd_head_qk_nope, float kq_scale, int il);
ggml_tensor * build_latent_moe(ggml_tensor * cur, const llama_layer & layer,
int64_t n_embd_latent, int il);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_kimi_linear : public llama_model_base {
llama_model_kimi_linear(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;