Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3#include <algorithm>4#include "llama-memory-recurrent.h"5 6//7// Kimi-K3 text model: hybrid KDA (linear) + MLA (full) attention, as in kimi-linear.8// Parts that kimi-linear does not have:9// 1. cross-layer residual attention (attn_res_block_size)10// 2. latent MoE (routed experts run at n_expert_latent)11// 3. situ activation (replaces SwiGLU everywhere)12// 4. MLA output gate (sigmoid gate before o_proj)13// 5. full-rank KDA gate (single ssm_g instead of ssm_g_a/ssm_g_b)14//15 16void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) {17 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);18 ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);19 ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);20 ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false);21 ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);22 ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);23 ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda);24 ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false);25 26 // the MLA cache holds the compressed latent27 // set it here too, as older GGUFs have no value_length key28 hparams.n_embd_head_v_full = hparams.n_lora_kv;29 30 // n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-linear31 for (uint32_t i = 0; i < hparams.n_layer(); ++i) {32 hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0;33 }34 35 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);36 ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);37 ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);38 ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);39 ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);40 ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);41 ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent, false);42 43 ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size);44 ml.get_key(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta);45 ml.get_key(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta);46 47 switch (hparams.n_layer()) {48 case 93: type = LLM_TYPE_2_8T_A50B; break; // Kimi-K349 default: type = LLM_TYPE_UNKNOWN;50 }51}52 53void llama_model_kimi_k3::load_arch_tensors(llama_model_loader &) {54 LLAMA_LOAD_LOCALS;55 56 const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd;57 58 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);59 60 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);61 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);62 63 if (hparams.attn_res_block_size > 0) {64 output_res_score = create_tensor(tn(LLM_TENSOR_OUTPUT_RES_SCORE, "weight"), {n_embd}, 0);65 }66 67 for (int i = 0; i < n_layer; ++i) {68 auto & layer = layers[i];69 70 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);71 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);72 73 if (hparams.attn_res_block_size > 0) {74 layer.attn_res_score = create_tensor(tn(LLM_TENSOR_ATTN_RES_SCORE, "weight", i), {n_embd}, 0);75 layer.ffn_res_score = create_tensor(tn(LLM_TENSOR_FFN_RES_SCORE, "weight", i), {n_embd}, 0);76 }77 78 const int64_t head_dim = hparams.n_embd_head_kda;79 const int64_t d_conv = hparams.ssm_d_conv;80 const int64_t d_inner = head_dim * n_head;81 82 if (hparams.is_recr(i)) {83 // conv1d may be stored 4D [d_conv, 1, d_inner, 1] or 3D (quantization drops the trailing 1)84 auto conv = [&](llm_tensor tid) {85 ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED);86 return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0);87 };88 layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q);89 layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K);90 layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V);91 92 create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0);93 94 layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, 0);95 layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0);96 layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0);97 98 // K3's A_log is a plain 1-D [n_head] tensor (kimi-linear's is padded)99 layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, 0);100 layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0);101 102 // K3 uses a single full-rank gate instead of kimi-linear's g_a/g_b pair103 layer.ssm_g = create_tensor(tn(LLM_TENSOR_SSM_G, "weight", i), {n_embd, d_inner}, 0);104 layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0);105 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0);106 } else {107 const int64_t q_lora_rank = hparams.n_lora_q;108 const int64_t kv_lora_rank = hparams.n_lora_kv;109 const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();110 const int64_t n_embd_head_v = hparams.n_embd_head_v_mla();111 const int64_t qk_rope_head_dim = hparams.n_rot();112 const int64_t qk_nope_head_dim = n_embd_head_k - qk_rope_head_dim;113 114 layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, TENSOR_NOT_REQUIRED);115 layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0);116 117 if (layer.attn_q_a_norm) {118 layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);119 layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, 0);120 } else {121 layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k}, 0);122 }123 124 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);125 layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i),126 {kv_lora_rank, n_head * (qk_nope_head_dim + n_embd_head_v)},127 TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);128 if (!layer.wkv_b) {129 layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {qk_nope_head_dim, kv_lora_rank, n_head}, 0);130 layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, 0);131 }132 133 // K3: sigmoid output gate applied to the attention output before o_proj134 layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head * n_embd_head_v}, TENSOR_NOT_REQUIRED);135 136 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, 0);137 }138 139 if (i < (int) hparams.n_layer_dense_lead) {140 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);141 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);142 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);143 } else {144 const int64_t n_ff_exp = hparams.n_ff_exp();145 146 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);147 layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);148 149 // routed experts live in the latent space150 layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0);151 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd_latent, n_expert}, 0);152 layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0);153 154 if (hparams.n_expert_latent > 0) {155 layer.ffn_routed_down = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_DOWN, "weight", i), {n_embd, n_embd_latent}, 0);156 layer.ffn_routed_up = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_UP, "weight", i), {n_embd_latent, n_embd}, 0);157 layer.ffn_routed_norm = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_NORM, "weight", i), {n_embd_latent}, TENSOR_NOT_REQUIRED);158 }159 160 // shared experts stay at n_embd, width = moe_intermediate_size * n_expert_shared161 const int64_t n_ff_shexp = n_ff_exp * (hparams.n_expert_shared > 0 ? hparams.n_expert_shared : 1);162 layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED);163 layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED);164 layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED);165 }166 }167}168 169std::unique_ptr<llm_graph_context> llama_model_kimi_k3::build_arch_graph(const llm_graph_params & params) const {170 return std::make_unique<graph>(*this, params);171}172 173// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta)174// linear_beta <= 0 disables the transform on the up branch175static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_tensor * up,176 float beta, float linear_beta) {177 ggml_tensor * a = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, gate, 1.0f/beta)), beta);178 a = ggml_mul(ctx0, a, ggml_sigmoid(ctx0, gate));179 180 if (linear_beta > 0.0f) {181 up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/linear_beta)), linear_beta);182 }183 return ggml_mul(ctx0, a, up);184}185 186//187// cross-layer residual attention188//189 190// layout is [n_embd, n_ckpt, n_tokens]: rms_norm reduces over ne0, dsv4_hc_pre over ne1191// append the new checkpoint, do not re-fold the whole chain192void llama_model_kimi_k3::graph::res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens) {193 ggml_tensor * ckpt = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);194 195 resi_stack = resi_stack ? ggml_concat(ctx0, resi_stack, ckpt, 1) : ckpt;196}197 198ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor * score_w,199 int64_t n_tokens, int il) {200 if (!resi_stack) {201 return cur; // layer 0: nothing banked yet202 }203 204 const int n_ckpt = (int) resi_stack->ne[1];205 const float eps = hparams.f_norm_rms_eps;206 207 ggml_tensor * src = resi_stack; // [n_embd, n_ckpt, n_tokens]208 209 // one rms_norm scores all checkpoints at once210 // note: the scores use the normalized values, but the sum below uses the raw ones211 ggml_tensor * sc_src = ggml_rms_norm(ctx0, src, eps);212 sc_src = ggml_mul(ctx0, sc_src, score_w);213 sc_src = ggml_sum_rows(ctx0, sc_src); // [1, n_ckpt, n_tokens]214 sc_src = ggml_reshape_2d(ctx0, sc_src, n_ckpt, n_tokens);215 216 // the current residual stream is scored apart, so the stack stays append-only217 ggml_tensor * sc_cur = ggml_rms_norm(ctx0, cur, eps);218 sc_cur = ggml_mul(ctx0, sc_cur, score_w);219 sc_cur = ggml_sum_rows(ctx0, sc_cur); // [1, n_tokens]220 221 ggml_tensor * scores = ggml_concat(ctx0, sc_src, sc_cur, 0); // [n_ckpt+1, n_tokens]222 ggml_tensor * probs = ggml_soft_max(ctx0, scores); // over ne0 = n_ckpt+1223 cb(probs, "res_probs", il);224 225 // split the sum: hc_pre handles the stack, a broadcast-multiply the current stream226 ggml_tensor * p_src = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, n_ckpt, n_tokens, probs->nb[1], 0));227 ggml_tensor * p_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, 1, n_tokens, probs->nb[1],228 probs->nb[0] * n_ckpt));229 230 ggml_tensor * out = ggml_dsv4_hc_pre(ctx0, src, p_src);231 out = ggml_add(ctx0, out, ggml_mul(ctx0, cur, p_cur));232 233 return out;234}235 236llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_params & params) :237 llm_build_delta_net_base(params), model(model) {238 239 ggml_tensor * cur;240 ggml_tensor * inpL;241 242 inpL = build_inp_embd(model.tok_embd);243 cb(inpL, "inp_embd", -1);244 245 // K3 MLA is nope-only, so there is no position input246 247 auto * inp_kv = !hparams.is_mla() ? build_inp_mem_hybrid() : nullptr;248 auto * inp_k = hparams.is_mla() ? build_inp_mem_hybrid_k() : nullptr;249 auto * inp_rs = hparams.is_mla() ? inp_k->get_recr() : inp_kv->get_recr();250 auto * inp_attn_kv = !hparams.is_mla() ? inp_kv->get_attn() : nullptr;251 auto * inp_attn_k = hparams.is_mla() ? inp_k->get_attn() : nullptr;252 253 ggml_tensor * inp_out_ids = build_inp_out_ids();254 255 const int64_t n_head_kda = hparams.n_head();256 const int64_t head_dim = hparams.n_embd_head_kda;257 const int64_t d_conv = hparams.ssm_d_conv;258 const int64_t d_inner = n_head_kda * head_dim;259 const int64_t n_seqs = ubatch.n_seqs;260 const int64_t n_seq_tokens = ubatch.n_seq_tokens;261 262 GGML_ASSERT(n_seqs != 0);263 GGML_ASSERT(ubatch.equal_seqs());264 GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);265 266 const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();267 const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();268 const int64_t kv_lora_rank = hparams.n_lora_kv;269 const int64_t n_embd_head_qk_rope = hparams.n_rot();270 const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope;271 const float kq_scale_mla = 1.0f / sqrtf((float) n_embd_head_k_mla);272 273 const uint32_t res_bs = hparams.attn_res_block_size;274 const bool use_attn_res = res_bs > 0;275 const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd;276 277 for (int il = 0; il < n_layer; ++il) {278 const auto & layer = model.layers[il];279 280 // the residual stream, banked on checkpoint layers and then restarted281 // from the attention output alone282 ggml_tensor * prefix_sum = inpL;283 284 cur = use_attn_res ? res_mix(prefix_sum, layer.attn_res_score, n_tokens, il)285 : prefix_sum;286 287 bool banked = false;288 if (use_attn_res && (uint32_t) il % res_bs == 0) {289 res_push(prefix_sum, n_embd, n_tokens); // banks the RAW layer input, not `cur`290 banked = true;291 }292 293 cur = build_norm(cur, layer.attn_norm, NULL, LLM_NORM_RMS, il);294 cb(cur, "attn_norm", il);295 ggml_build_forward_expand(gf, cur);296 297 if (hparams.is_recr(il)) {298 cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda,299 d_inner, n_seq_tokens, n_seqs, il);300 } else {301 cur = build_mla_layer(cur, layer, inp_attn_k, inp_attn_kv,302 n_embd_head_k_mla, n_embd_head_v_mla, kv_lora_rank,303 n_embd_head_qk_rope, n_embd_head_qk_nope, kq_scale_mla, il);304 }305 306 prefix_sum = banked ? cur : ggml_add(ctx0, prefix_sum, cur);307 cb(prefix_sum, "prefix_sum_attn", il);308 309 cur = use_attn_res ? res_mix(prefix_sum, layer.ffn_res_score, n_tokens, il)310 : prefix_sum;311 312 cur = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il);313 cb(cur, "ffn_norm", il);314 315 if ((uint32_t) il < hparams.n_layer_dense_lead) {316 ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate, cur);317 ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up, cur);318 cur = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta);319 cur = ggml_mul_mat(ctx0, layer.ffn_down, cur);320 cb(cur, "ffn_out", il);321 } else {322 cur = build_latent_moe(cur, layer, n_embd_latent, il);323 }324 325 prefix_sum = ggml_add(ctx0, prefix_sum, cur);326 prefix_sum = build_cvec(prefix_sum, il);327 cb(prefix_sum, "l_out", il);328 329 inpL = prefix_sum;330 }331 332 cur = inpL;333 334 // final mix, then narrow to the output tokens335 if (use_attn_res) {336 cur = res_mix(cur, model.output_res_score, n_tokens, -1);337 }338 if (inp_out_ids) {339 cur = ggml_get_rows(ctx0, cur, inp_out_ids);340 }341 342 cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);343 cb(cur, "result_norm", -1);344 res->t_embd = cur;345 346 cur = ggml_mul_mat(ctx0, model.output, cur);347 cb(cur, "result_output", -1);348 res->t_logits = cur;349 350 ggml_build_forward_expand(gf, cur);351}352 353//354// KDA layer355//356 357// causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use358static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0,359 ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,360 int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,361 int64_t d_conv, int64_t head_dim, int64_t n_head,362 int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head,363 int64_t mem_size, int64_t K_rs) {364 const int64_t d_inner = head_dim * n_head;365 const int64_t conv_state_size = (d_conv - 1) * d_inner;366 const int64_t n_embd_r_total = 3 * conv_state_size;367 368 ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs,369 (d_conv - 1) * ggml_element_size(conv_state_all),370 n_embd_r_total * ggml_element_size(conv_state_all),371 qkv * conv_state_size * ggml_element_size(conv_state_all));372 373 ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x);374 ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);375 ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0);376 377 // group s holds the conv window s tokens back.378 // [TAG_RECURRENT_ROLLBACK_SPLITS]: the last K_rs tokens must share one ubatch.379 for (int64_t s = 0; s < K_rs; ++s) {380 const int64_t s_idx = std::max<int64_t>(0, n_seq_tokens - s);381 ggml_tensor * conv_x_s = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,382 conv_x->nb[1], conv_x->nb[2], s_idx * conv_x->nb[0]);383 ggml_build_forward_expand(gf,384 ggml_cpy(ctx0, conv_x_s,385 ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,386 (d_conv - 1) * ggml_element_size(conv_states_all),387 n_embd_r_total * ggml_element_size(conv_states_all),388 ((s * mem_size + kv_head) * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));389 }390 391 ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);392 ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);393 Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens);394 Xcur = ggml_silu(ctx0, Xcur);395 396 return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs);397}398 399ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer(400 ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs,401 int64_t d_conv, int64_t head_dim, int64_t n_head_kda,402 int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) {403 404 const auto * mctx_cur = inp_rs->mctx;405 const auto kv_head = mctx_cur->get_head();406 407 ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);408 ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);409 410 const int64_t mem_size = mctx_cur->get_size();411 const int64_t K_rs = (int64_t) cparams.n_rs_seq + 1;412 413 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, mem_size, K_rs);414 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, mem_size, K_rs);415 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, mem_size, K_rs);416 cb(Qcur, "kda_q_conv", il);417 cb(Kcur, "kda_k_conv", il);418 cb(Vcur, "kda_v_conv", il);419 420 // gate_lower_bound is not a clamp - when set, it swaps the decay gate activation:421 // unset (kimi-linear): g = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias)422 // set (K3, -5.0): g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias))423 // ssm_a holds -exp(A_log) (folded at conversion time), so exp(A_log) == -ssm_a424 ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);425 ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a);426 g1 = ggml_add(ctx0, g1, layer.ssm_dt_b);427 428 ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1);429 430 if (hparams.kda_gate_lower_bound > -INFINITY) {431 g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens);432 g1 = ggml_mul(ctx0, g1, A); // -exp(A_log) * (...)433 g1 = ggml_sigmoid(ctx0, ggml_scale(ctx0, g1, -1.0f));434 g1 = ggml_scale(ctx0, g1, hparams.kda_gate_lower_bound);435 } else {436 g1 = ggml_softplus(ctx0, g1);437 g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens);438 g1 = ggml_mul(ctx0, g1, A);439 }440 cb(g1, "kda_g1", il);441 442 g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head_kda, n_seq_tokens, n_seqs);443 444 ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur);445 beta = ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs);446 beta = ggml_sigmoid(ctx0, beta);447 cb(beta, "kda_beta", il);448 449 ggml_tensor * cur_3d = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs);450 451 ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);452 ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs);453 state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs);454 455 const float eps_norm = hparams.f_norm_rms_eps;456 Qcur = build_gdn_l2_norm(ctx0, Qcur, eps_norm);457 Kcur = build_gdn_l2_norm(ctx0, Kcur, eps_norm);458 459 ggml_tensor * output = build_recurrent_attn(inp_rs, ssm_states_all, Qcur, Kcur, Vcur, g1, beta, state, il);460 output = ggml_cont(ctx0, output);461 cb(output, "kda_scan_out", il);462 463 // K3: single full-rank gate (kimi-linear factors this as g_b(g_a(x)))464 ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur_3d, cur_3d->ne[0], n_seq_tokens * n_seqs);465 ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g, cur_2d);466 g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head_kda, n_seq_tokens * n_seqs);467 468 ggml_tensor * o = ggml_reshape_3d(ctx0, output, head_dim, n_head_kda, n_seq_tokens * n_seqs);469 ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il);470 cb(g2, "kda_g2", il);471 cb(normed, "kda_normed", il);472 ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2));473 474 gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens);475 cur = ggml_mul_mat(ctx0, layer.wo, gated);476 cb(cur, "kda_out", il);477 478 return cur;479}480 481//482// MLA layer (nope-only, with K3's sigmoid output gate)483//484 485ggml_tensor * llama_model_kimi_k3::graph::build_mla_layer(486 ggml_tensor * cur, const llama_layer & layer,487 llm_graph_input_attn_k * inp_attn_k, llm_graph_input_attn_kv * inp_attn_kv,488 int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla, int64_t kv_lora_rank,489 int64_t n_embd_head_qk_rope, int64_t n_embd_head_qk_nope, float kq_scale, int il) {490 491 ggml_tensor * inp_gate = cur; // the output gate reads the *normed* layer input492 493 ggml_tensor * Qcur;494 if (layer.wq_a) {495 Qcur = ggml_mul_mat(ctx0, layer.wq_a, cur);496 Qcur = build_norm(Qcur, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);497 Qcur = ggml_mul_mat(ctx0, layer.wq_b, Qcur);498 } else {499 Qcur = ggml_mul_mat(ctx0, layer.wq, cur);500 }501 502 ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);503 504 ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,505 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);506 ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,507 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),508 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),509 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));510 511 // no RoPE: mla_use_nope is asserted at conversion time512 kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);513 514 ggml_tensor * out;515 if (layer.wk_b && layer.wv_b) {516 ggml_tensor * q_nope = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_nope, n_head, n_tokens,517 ggml_row_size(Qcur->type, n_embd_head_k_mla),518 ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, 0);519 ggml_tensor * q_pe = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_rope, n_head, n_tokens,520 ggml_row_size(Qcur->type, n_embd_head_k_mla),521 ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head,522 ggml_row_size(Qcur->type, n_embd_head_qk_nope));523 524 q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);525 ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);526 q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);527 528 ggml_tensor * Q = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);529 ggml_tensor * kv_cmpr_3d = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);530 ggml_tensor * K = ggml_concat(ctx0, kv_cmpr_3d, k_pe, 0);531 ggml_tensor * V = kv_cmpr_3d;532 533 // wo == NULL: the output projection is applied after the gate below534 out = build_attn(inp_attn_k, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, layer.wv_b, kq_scale, il);535 } else {536 ggml_tensor * Q = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens);537 ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr);538 const int64_t kv_per_head = n_embd_head_qk_nope + n_embd_head_v_mla;539 540 ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens,541 ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head), 0);542 ggml_tensor * V = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens,543 ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head),544 ggml_row_size(kv->type, n_embd_head_qk_nope)));545 546 ggml_tensor * k_pe_t = ggml_new_tensor_3d(ctx0, k_pe->type, n_embd_head_qk_rope, n_head, n_tokens);547 ggml_tensor * K = ggml_concat(ctx0, ggml_repeat(ctx0, k_pe, k_pe_t), k_nope, 0);548 549 out = build_attn(inp_attn_kv, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, nullptr, kq_scale, il);550 }551 552 // K3: attn_output *= sigmoid(g_proj(x)), then o_proj553 if (layer.wqkv_gate) {554 ggml_tensor * g = ggml_sigmoid(ctx0, ggml_mul_mat(ctx0, layer.wqkv_gate, inp_gate));555 out = ggml_mul(ctx0, out, g);556 cb(out, "mla_gated", il);557 }558 559 out = ggml_mul_mat(ctx0, layer.wo, out);560 cb(out, "mla_out", il);561 562 return out;563}564 565//566// latent MoE: down-project, run the routed experts in the latent space, norm, up-project;567// shared experts stay at n_embd and read the un-projected input.568//569 570ggml_tensor * llama_model_kimi_k3::graph::build_latent_moe(571 ggml_tensor * cur, const llama_layer & layer, int64_t n_embd_latent, int il) {572 573 ggml_tensor * identity = cur;574 575 ggml_tensor * routed_in = layer.ffn_routed_down576 ? ggml_mul_mat(ctx0, layer.ffn_routed_down, cur)577 : cur;578 579 // the router scores the full-width input while the experts take the latent one,580 // so the logits are computed here and passed to build_moe_ffn581 ggml_tensor * logits = ggml_mul_mat(ctx0, layer.ffn_gate_inp, identity);582 cb(logits, "ffn_moe_logits", il);583 584 ggml_tensor * moe_out = build_moe_ffn(routed_in,585 nullptr, // gate_inp unused: the logits above are passed instead586 layer.ffn_up_exps,587 layer.ffn_gate_exps,588 layer.ffn_down_exps,589 layer.ffn_exp_probs_b,590 hparams.n_expert,591 hparams.n_expert_used(),592 LLM_FFN_SITU, hparams.expert_weights_norm,593 hparams.expert_weights_scale,594 (llama_expert_gating_func_type) hparams.expert_gating_func,595 il,596 logits);597 cb(moe_out, "ffn_moe_out", il);598 599 if (layer.ffn_routed_norm) {600 moe_out = build_norm(moe_out, layer.ffn_routed_norm, NULL, LLM_NORM_RMS, il);601 }602 if (layer.ffn_routed_up) {603 moe_out = ggml_mul_mat(ctx0, layer.ffn_routed_up, moe_out);604 }605 GGML_UNUSED(n_embd_latent);606 607 if (layer.ffn_gate_shexp) {608 ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate_shexp, identity);609 ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up_shexp, identity);610 ggml_tensor * sh = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta);611 sh = ggml_mul_mat(ctx0, layer.ffn_down_shexp, sh);612 cb(sh, "ffn_shexp", il);613 moe_out = ggml_add(ctx0, moe_out, sh);614 }615 616 cb(moe_out, "ffn_out", il);617 return moe_out;618}619