Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3#include "llama-kv-cache-dsa.h"4 5// https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json#L266const std::array<uint32_t, LLAMA_MAX_LAYERS> GLM_5_2_DEFAULT_INDEXER_TYPES = {7 1, 1,8 1, 0, 0, 0,9 1, 0, 0, 0,10 1, 0, 0, 0,11 1, 0, 0, 0,12 1, 0, 0, 0,13 1, 0, 0, 0,14 1, 0, 0, 0,15 1, 0, 0, 0,16 1, 0, 0, 0,17 1, 0, 0, 0,18 1, 0, 0, 0,19 1, 0, 0, 0,20 1, 0, 0, 0,21 1, 0, 0, 0,22 1, 0, 0, 0,23 1, 0, 0, 0,24 1, 0, 0, 0,25 1, 0, 0, 0,26 1, 0, 0, 0,27};28 29void llama_model_glm_dsa::load_arch_hparams(llama_model_loader & ml) {30 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);31 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);32 ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);33 34 // MoE parameters35 ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);36 ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);37 ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);38 ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);39 40 // deepseek MLA parameters41 ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);42 ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);43 ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl, false);44 ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl, false);45 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);46 ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);47 48 // DSA parameters49 ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);50 ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);51 ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);52 53 // Expert gating function (GLM-4.5 uses sigmoid)54 ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);55 if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {56 hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;57 }58 59 // BC for GLM 5, 5.1 (full indexers) without indexer_types metadata60 const bool is_pre_5_2 = hparams.n_ctx_train < 1048576;61 if (is_pre_5_2) {62 std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1);63 } else {64 hparams.is_indexer_full_impl = GLM_5_2_DEFAULT_INDEXER_TYPES;65 }66 ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false);67 68 switch (hparams.n_layer()) {69 case 78: type = LLM_TYPE_744B_A40B; break;70 default: type = LLM_TYPE_UNKNOWN;71 }72}73 74void llama_model_glm_dsa::load_arch_tensors(llama_model_loader & ml) {75 LLAMA_LOAD_LOCALS;76 const int64_t n_expert_shared = hparams.n_expert_shared;77 78 // MTP-only: the GGUF carries only the NextN/MTP block(s) (user split target/draft).79 const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);80 // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP81 // tensors live in a separate file (or were stripped at conversion). Mark82 // MTP tensors NOT_REQUIRED so the trunk loads cleanly.83 const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";84 const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);85 const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;86 int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;87 88 if (!ml.load_mtp) {89 mtp_flags |= TENSOR_SKIP;90 }91 92 const bool is_mla = hparams.is_mla();93 if (!is_mla) {94 throw std::runtime_error("GLM_DSA architecture requires MLA");95 }96 97 // note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA98 const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();99 const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();100 101 const int64_t n_embd_head_qk_rope = hparams.n_rot();102 const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope;103 104 const int64_t q_lora_rank = hparams.n_lora_q;105 const int64_t kv_lora_rank = hparams.n_lora_kv;106 107 const int64_t n_ff_exp = hparams.n_ff_exp();108 109 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);110 111 // output112 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);113 // try to load output.weight, if not found, use token_embd (tied embeddings)114 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);115 if (!output) {116 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);117 }118 119 for (int i = 0; i < n_layer_all; ++i) {120 // NextN/MTP layers (i >= n_layer) are full decoder blocks used by the121 // LLM_GRAPH_TYPE_DECODER_MTP draft head; load them like qwen35moe/step35/hy_v3.122 const int flags = (i >= n_layer) ? mtp_flags : trunk_flags;123 124 auto & layer = layers[i];125 126 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);127 layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags);128 layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags);129 130 layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags);131 layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k_mla}, flags);132 133 layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + n_embd_head_qk_rope}, flags);134 135 // note: only old legacy GGUF files will have the unsplit wkv_b tensor in136 layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, flags);137 layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, flags);138 139 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, flags);140 141 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);142 143 // DSA indexer144 layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, flags | TENSOR_NOT_REQUIRED);145 layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {hparams.indexer_head_size}, flags | TENSOR_NOT_REQUIRED);146 layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, flags | TENSOR_NOT_REQUIRED);147 layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, hparams.indexer_head_size}, flags | TENSOR_NOT_REQUIRED);148 layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * hparams.indexer_head_size}, flags | TENSOR_NOT_REQUIRED);149 if (i < (int) hparams.n_layer_dense_lead) {150 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags);151 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, flags);152 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags);153 } else {154 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags);155 layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);156 157 if (n_expert == 0) {158 throw std::runtime_error("n_expert must be > 0");159 }160 if (n_expert_used == 0) {161 throw std::runtime_error("n_expert_used must be > 0");162 }163 164 // MoE branch165 layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags);166 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags);167 layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags);168 169 // Shared expert branch170 layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);171 layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, flags);172 layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);173 }174 175 // NextN/MTP tensors - the NextN-specific wiring around the extra decoder block176 if (i >= n_layer) {177 layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);178 layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);179 layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags);180 181 // Optional tensors182 layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED);183 layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED);184 layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, flags | TENSOR_NOT_REQUIRED);185 }186 }187}188 189std::unique_ptr<llm_graph_context> llama_model_glm_dsa::build_arch_graph(const llm_graph_params & params) const {190 if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {191 return std::make_unique<graph_mtp>(*this, params);192 }193 return std::make_unique<graph>(*this, params);194}195 196llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_params & params) :197 llm_graph_context(params) {198 const bool is_mla = hparams.is_mla();199 GGML_ASSERT(is_mla);200 201 // note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA202 const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();203 const int64_t n_embd_head_v = hparams.n_embd_head_v_mla();204 GGML_UNUSED(n_embd_head_v);205 206 const int64_t n_embd_head_qk_rope = hparams.n_rot();207 const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;208 209 const int64_t n_indexer_head = hparams.indexer_n_head;210 const int64_t n_embd_indexer_head = hparams.indexer_head_size;211 const uint32_t n_indexer_top_k = hparams.indexer_top_k;212 213 // the indexer head layout is [rope | nope]214 GGML_ASSERT(hparams.n_rot() <= n_embd_indexer_head);215 216 const uint32_t kv_lora_rank = hparams.n_lora_kv;217 218 // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.219 // See https://github.com/ggml-org/llama.cpp/discussions/7416 for detailed explanation.220 // And also: https://github.com/ggml-org/llama.cpp/pull/17945 [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]221 222 // first cancel the adjustment from llama_hparams::yarn_attn_factor_adjust to get the original attn_factor223 GGML_ASSERT(ext_factor >= 0.0f);224 const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale));225 226 // use the original attn_factor to pre-scale the kq_scale227 const float mscale = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));228 const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k));229 230 ggml_tensor * cur;231 ggml_tensor * inpL;232 233 // {n_embd, n_tokens}234 inpL = build_inp_embd(model.tok_embd);235 236 // inp_pos - contains the positions237 ggml_tensor * inp_pos = build_inp_pos();238 239 llm_graph_input_attn_k_dsa * inp_attn_dsa = build_attn_inp_k_dsa();240 241 ggml_tensor * inp_out_ids = build_inp_out_ids();242 243 // Difference vs Deepseek 3.2: shared indexer layers reuse the top_k from the previous full indexer layers244 // See https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json#L30245 ggml_tensor * prev_top_k = nullptr;246 for (int il = 0; il < n_layer; ++il) {247 ggml_tensor * inpSA = inpL;248 249 // norm250 cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);251 cb(cur, "attn_norm", il);252 253 // self_attention254 {255 ggml_tensor * qr = ggml_mul_mat(ctx0, model.layers[il].wq_a, cur);256 cb(qr, "qr", il);257 258 qr = build_norm(qr, model.layers[il].attn_q_a_norm, nullptr, LLM_NORM_RMS, il);259 cb(qr, "qr", il);260 261 ggml_tensor * top_k = nullptr;262 263 // lightning indexer264 if (hparams.is_indexer_full(il)) {265 // "full" layer266 ggml_tensor * indexer_q = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_q_b, qr);267 cb(indexer_q, "indexer_q", il);268 269 // {n_embd_indexer_head, n_indexer_head, n_tokens}270 indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens);271 indexer_q = ggml_rope_ext(ctx0, indexer_q, inp_pos, nullptr, n_rot,272 LLAMA_ROPE_TYPE_NORM, n_ctx_orig, freq_base, freq_scale,273 ext_factor, attn_factor, beta_fast, beta_slow);274 cb(indexer_q, "indexer_q", il);275 276 ggml_tensor * indexer_k = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_k, cur);277 cb(indexer_k, "indexer_k", il);278 279 indexer_k = build_norm(indexer_k, model.layers[il].indexer_k_norm, model.layers[il].indexer_k_norm_b, LLM_NORM, il);280 cb(indexer_k, "indexer_k", il);281 282 // {n_embd_indexer_head, 1, n_tokens}283 indexer_k = ggml_reshape_3d(ctx0, indexer_k, n_embd_indexer_head, 1, n_tokens);284 indexer_k = ggml_rope_ext(ctx0, indexer_k, inp_pos, nullptr, n_rot,285 LLAMA_ROPE_TYPE_NORM, n_ctx_orig, freq_base, freq_scale,286 ext_factor, attn_factor, beta_fast, beta_slow);287 cb(indexer_k, "indexer_k", il);288 289 // perform Hadamard transform on indexer q and k290 indexer_q = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_q);291 cb(indexer_q, "indexer_q", il);292 indexer_k = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_k);293 cb(indexer_k, "indexer_k", il);294 295 // store indexer keys to KV cache296 const auto * mctx_lid = inp_attn_dsa->mctx->get_lid();297 const auto & k_idxs_lid = inp_attn_dsa->get_k_idxs_lid();298 ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, indexer_k, k_idxs_lid, il));299 300 // prepare indexer weights301 ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, model.layers[il].indexer_proj, cur);302 cb(indexer_weights, "indexer_weights", il);303 304 // get cached indexer keys305 indexer_k = mctx_lid->get_k(ctx0, il);306 307 // split the batch into streams if needed308 const auto n_stream = indexer_k->ne[3];309 indexer_q = ggml_view_4d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0);310 indexer_weights = ggml_view_4d(ctx0, indexer_weights, indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0);311 312 // pre-scale weights to avoid scaling operations on huge indexer_score tensor313 indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head * n_indexer_head)));314 cb(indexer_weights, "indexer_weights", il);315 316 ggml_tensor * indexer_score = nullptr;317 if (cparams.fused_lid) {318 indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());319 cb(indexer_score, "indexer_score", il);320 res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});321 } else {322 // calculate indexer kq323 indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);324 cb(indexer_q, "indexer_q", il);325 indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);326 cb(indexer_k, "indexer_k", il);327 328 ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);329 cb(indexer_kq, "indexer_kq", il);330 331 // ReLU requires contiguous tensors332 indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));333 cb(indexer_kq, "indexer_kq", il);334 335 // apply ReLU336 indexer_score = ggml_relu(ctx0, indexer_kq);337 cb(indexer_score, "indexer_score", il);338 339 // multiply scores by indexer weights340 indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);341 cb(indexer_score, "indexer_score", il);342 343 // sum by q n_indexer_head dimension344 indexer_score = ggml_sum_rows(ctx0, indexer_score);345 cb(indexer_score, "indexer_score", il);346 347 // permute result to match KQ mask348 indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));349 cb(indexer_score, "indexer_score", il);350 351 // mask indexer scores352 ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();353 indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);354 cb(indexer_score, "indexer_score", il);355 }356 357 // get indices of top k indexer scores358 uint32_t n_top_k = indexer_score->ne[0] < n_indexer_top_k ? indexer_score->ne[0] : n_indexer_top_k;359 top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));360 prev_top_k = top_k;361 cb(top_k, "top_k", il);362 } else {363 // "shared" indexer layer - reuse top-k from a previous full layer364 GGML_ASSERT(prev_top_k != nullptr && "shared indexer layer must follow a previous full indexer layer");365 top_k = prev_top_k;366 cb(top_k, "top_k", il);367 }368 369 ggml_tensor * q = ggml_mul_mat(ctx0, model.layers[il].wq_b, qr);370 cb(q, "q", il);371 372 // split into {n_embd_head_qk_nope, n_head, n_tokens}373 ggml_tensor * q_nope =374 ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),375 ggml_row_size(q->type, n_embd_head_k) * n_head, 0);376 cb(q_nope, "q_nope", il);377 378 // and {n_embd_head_qk_rope, n_head, n_tokens}379 ggml_tensor * q_pe = ggml_view_3d(380 ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),381 ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope));382 cb(q_pe, "q_pe", il);383 384 ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur);385 cb(kv_cmpr_pe, "kv_cmpr_pe", il);386 387 // split into {kv_lora_rank, n_tokens}388 ggml_tensor * kv_cmpr =389 ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,390 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);391 cb(kv_cmpr, "kv_cmpr", il);392 393 // and {n_embd_head_qk_rope, 1, n_tokens}394 ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,395 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),396 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),397 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));398 cb(k_pe, "k_pe", il);399 400 q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,401 ext_factor, attn_factor, beta_fast, beta_slow);402 cb(q_pe, "q_pe", il);403 404 k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,405 ext_factor, attn_factor, beta_fast, beta_slow);406 cb(k_pe, "k_pe", il);407 408 kv_cmpr = build_norm(kv_cmpr, model.layers[il].attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);409 cb(kv_cmpr, "kv_cmpr", il);410 411 // MLA attention412 {413 // {n_embd_head_qk_nope, n_tokens, n_head}414 q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);415 cb(q_nope, "q_nope_perm", il);416 417 // {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}418 ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, model.layers[il].wk_b, q_nope);419 cb(q_nope_absorbed, "q_nope_absorbed", il);420 421 // {kv_lora_rank, n_head, n_tokens}422 q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);423 cb(q_nope_absorbed, "q_nope_absorbed_perm", il);424 425 // {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens}426 // note: rope must go first for in-place context shifting in build_rope_shift()427 ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);428 cb(Qcur, "Qcur", il);429 430 kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);431 cb(kv_cmpr, "kv_cmpr_reshape", il);432 433 // {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens}434 ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);435 cb(Kcur, "Kcur", il);436 437 // {kv_lora_rank, 1, n_tokens}438 ggml_tensor * Vcur = kv_cmpr;439 cb(Vcur, "Vcur", il);440 441 // note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)442 cur = build_attn(inp_attn_dsa,443 model.layers[il].wo, NULL, model.layers[il].wo_s,444 Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, il);445 }446 }447 // when unmasked nextn embeddings are requested, t_h_nextn must keep all rows,448 // so the early output masking has to be skipped (it is applied after the final norm instead)449 if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {450 cur = ggml_get_rows(ctx0, cur, inp_out_ids);451 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);452 }453 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);454 cb(ffn_inp, "ffn_inp", il);455 456 cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);457 cb(cur, "ffn_norm", il);458 459 if ((uint32_t) il < hparams.n_layer_dense_lead) {460 cur = build_ffn(cur,461 model.layers[il].ffn_up, NULL, model.layers[il].ffn_up_s,462 model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,463 model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,464 NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);465 cb(cur, "ffn_out", il);466 } else {467 // MoE branch468 ggml_tensor * moe_out = build_moe_ffn(cur,469 model.layers[il].ffn_gate_inp,470 model.layers[il].ffn_up_exps,471 model.layers[il].ffn_gate_exps,472 model.layers[il].ffn_down_exps,473 model.layers[il].ffn_exp_probs_b,474 n_expert, n_expert_used,475 LLM_FFN_SILU, hparams.expert_weights_norm,476 hparams.expert_weights_scale,477 (llama_expert_gating_func_type) hparams.expert_gating_func,478 il,479 nullptr,480 model.layers[il].ffn_gate_up_exps,481 model.layers[il].ffn_up_exps_s,482 model.layers[il].ffn_gate_exps_s,483 model.layers[il].ffn_down_exps_s);484 cb(moe_out, "ffn_moe_out", il);485 486 // FFN shared expert487 {488 ggml_tensor * ffn_shexp =489 build_ffn(cur,490 model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,491 model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,492 model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,493 NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);494 cb(ffn_shexp, "ffn_shexp", il);495 496 cur = ggml_add(ctx0, moe_out, ffn_shexp);497 cb(cur, "ffn_out", il);498 }499 }500 cur = ggml_add(ctx0, cur, ffn_inp);501 502 cur = build_cvec(cur, il);503 cb(cur, "l_out", il);504 505 // input for next layer506 inpL = cur;507 }508 cur = inpL;509 510 cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);511 512 // post-norm hidden state feeds the NextN/MTP draft head513 cb(cur, "h_nextn", -1);514 res->t_h_nextn = cur;515 516 if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {517 cur = ggml_get_rows(ctx0, cur, inp_out_ids);518 }519 520 cb(cur, "result_norm", -1);521 res->t_embd = cur;522 523 // lm_head524 cur = ggml_mul_mat(ctx0, model.output, cur);525 526 cb(cur, "result_output", -1);527 res->t_logits = cur;528 529 ggml_build_forward_expand(gf, cur);530}531 532// LLM_GRAPH_TYPE_DECODER_MTP draft head for GLM-5.2 (GLM_DSA).533// Semantics mirror the deepseek-family NextN/MTP layer:534// enorm(embed) + hnorm(prev_hidden) -> concat(e, h) -> eh_proj ->535// full glm_dsa decoder block (dense MLA attention + sigmoid-gated MoE FFN536// with shared expert, exactly as the trunk deepseek2 graph builds it) ->537// shared_head_norm (fallback output_norm) -> shared LM head.538// The DSA indexer is not used at runtime (same as the trunk graph).539llama_model_glm_dsa::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)540 : llm_graph_context(params) {541 GGML_ASSERT(hparams.n_layer_nextn > 0 && "GLM_DSA MTP requires n_layer_nextn > 0");542 GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM_DSA MTP currently only supports a single MTP block");543 GGML_ASSERT(hparams.is_mla() && "GLM_DSA MTP requires MLA");544 545 const int il = hparams.n_layer() + cparams.nextn_layer_offset;546 GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&547 cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&548 "nextn_layer_offset out of range [0, n_layer_nextn)");549 const auto & layer = model.layers[il];550 551 GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");552 GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");553 GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");554 GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");555 556 // note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA557 const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();558 559 const int64_t n_embd_head_qk_rope = hparams.n_rot();560 const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;561 562 const uint32_t kv_lora_rank = hparams.n_lora_kv;563 564 // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.565 // See the deepseek2 trunk graph for the detailed explanation - this must match it EXACTLY.566 GGML_ASSERT(ext_factor >= 0.0f);567 const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale));568 569 const float mscale = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));570 const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k));571 572 // TODO: extract in a common llm_graph_context::build_inp_embd_h()573 auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);574 575 inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);576 ggml_set_input(inp->tokens);577 578 inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);579 ggml_set_input(inp->embd);580 581 ggml_tensor * tok_embd;582 if (ubatch.token) {583 ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;584 585 tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);586 } else {587 tok_embd = inp->embd;588 }589 cb(tok_embd, "mtp_tok_embd", il);590 591 inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);592 ggml_set_input(inp->h);593 ggml_set_name(inp->h, "mtp_h_input");594 595 ggml_tensor * h_embd = inp->h;596 597 res->add_input(std::move(inp));598 599 ggml_tensor * inp_pos = build_inp_pos();600 ggml_tensor * inp_out_ids = build_inp_out_ids();601 602 // MLA with the absorption optimization uses a K-only cache (V is a view of K)603 auto * inp_attn = build_attn_inp_k();604 605 ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);606 cb(h_norm, "mtp_hnorm", il);607 608 ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);609 cb(e_norm, "mtp_enorm", il);610 611 ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);612 cb(concat, "mtp_concat", il);613 614 ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);615 cb(cur, "mtp_eh_proj", il);616 617 ggml_tensor * inpSA = cur;618 619 cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);620 cb(cur, "mtp_attn_norm", il);621 622 // self-attention: dense MLA, same construction as the deepseek2 trunk graph623 {624 ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur);625 cb(q, "mtp_q", il);626 627 q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);628 cb(q, "mtp_q", il);629 630 q = ggml_mul_mat(ctx0, layer.wq_b, q);631 cb(q, "mtp_q", il);632 633 // split into {n_embd_head_qk_nope, n_head, n_tokens}634 ggml_tensor * q_nope =635 ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),636 ggml_row_size(q->type, n_embd_head_k) * n_head, 0);637 cb(q_nope, "mtp_q_nope", il);638 639 // and {n_embd_head_qk_rope, n_head, n_tokens}640 ggml_tensor * q_pe = ggml_view_3d(641 ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),642 ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope));643 cb(q_pe, "mtp_q_pe", il);644 645 ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);646 cb(kv_cmpr_pe, "mtp_kv_cmpr_pe", il);647 648 // split into {kv_lora_rank, n_tokens}649 ggml_tensor * kv_cmpr =650 ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,651 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);652 cb(kv_cmpr, "mtp_kv_cmpr", il);653 654 // and {n_embd_head_qk_rope, 1, n_tokens}655 ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,656 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),657 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),658 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));659 cb(k_pe, "mtp_k_pe", il);660 661 q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,662 ext_factor, attn_factor, beta_fast, beta_slow);663 cb(q_pe, "mtp_q_pe", il);664 665 k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,666 ext_factor, attn_factor, beta_fast, beta_slow);667 cb(k_pe, "mtp_k_pe", il);668 669 kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);670 cb(kv_cmpr, "mtp_kv_cmpr", il);671 672 // {n_embd_head_qk_nope, n_tokens, n_head}673 q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);674 cb(q_nope, "mtp_q_nope_perm", il);675 676 // {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}677 ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);678 cb(q_nope_absorbed, "mtp_q_nope_absorbed", il);679 680 // {kv_lora_rank, n_head, n_tokens}681 q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);682 cb(q_nope_absorbed, "mtp_q_nope_absorbed_perm", il);683 684 // {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens}685 // note: rope must go first for in-place context shifting in build_rope_shift()686 ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);687 cb(Qcur, "mtp_Qcur", il);688 689 kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);690 cb(kv_cmpr, "mtp_kv_cmpr_reshape", il);691 692 // {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens}693 ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);694 cb(Kcur, "mtp_Kcur", il);695 696 // {kv_lora_rank, 1, n_tokens}697 ggml_tensor * Vcur = kv_cmpr;698 cb(Vcur, "mtp_Vcur", il);699 700 // note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)701 cur = build_attn(inp_attn,702 layer.wo, NULL, layer.wo_s,703 Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale, il);704 cb(cur, "mtp_attn_out", il);705 }706 707 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);708 cb(ffn_inp, "mtp_ffn_inp", il);709 710 cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);711 cb(cur, "mtp_ffn_norm", il);712 713 // MoE FFN with shared expert - same construction as the deepseek2 trunk graph714 ggml_tensor * moe_out = build_moe_ffn(cur,715 layer.ffn_gate_inp,716 layer.ffn_up_exps,717 layer.ffn_gate_exps,718 layer.ffn_down_exps,719 layer.ffn_exp_probs_b,720 n_expert, n_expert_used,721 LLM_FFN_SILU, hparams.expert_weights_norm,722 hparams.expert_weights_scale,723 (llama_expert_gating_func_type) hparams.expert_gating_func,724 il,725 nullptr,726 layer.ffn_gate_up_exps,727 layer.ffn_up_exps_s,728 layer.ffn_gate_exps_s,729 layer.ffn_down_exps_s);730 cb(moe_out, "mtp_ffn_moe_out", il);731 732 // FFN shared expert733 ggml_tensor * ffn_shexp =734 build_ffn(cur,735 layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,736 layer.ffn_gate_shexp, NULL, layer.ffn_gate_shexp_s,737 layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,738 NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);739 cb(ffn_shexp, "mtp_ffn_shexp", il);740 741 cur = ggml_add(ctx0, moe_out, ffn_shexp);742 cb(cur, "mtp_ffn_out", il);743 744 cur = ggml_add(ctx0, cur, ffn_inp);745 cb(cur, "mtp_post_ffn", il);746 747 // shared_head_norm applied after the decoder block, before the shared LM head.748 // The post-norm hidden state seeds the next MTP step.749 ggml_tensor * head_norm_w = layer.nextn.shared_head_norm750 ? layer.nextn.shared_head_norm751 : model.output_norm;752 GGML_ASSERT(head_norm_w && "GLM_DSA MTP: missing both nextn.shared_head_norm and output_norm");753 cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);754 755 cb(cur, "h_nextn", -1);756 res->t_h_nextn = cur;757 758 cur = ggml_get_rows(ctx0, cur, inp_out_ids);759 cb(cur, "mtp_shared_head_norm", -1);760 761 ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;762 ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;763 GGML_ASSERT(head_w && "GLM_DSA MTP: missing LM head (nextn.shared_head_head or model.output)");764 cur = build_lora_mm(head_w, cur, head_s);765 cb(cur, "result_output", -1);766 767 res->t_logits = cur;768 ggml_build_forward_expand(gf, cur);769}770