Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3#include "llama-kv-cache.h"4#include "llama-kv-cache-dsa.h"5 6void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) {7 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);8 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);9 hparams.f_norm_eps = 1e-6; // eps for layer norm10 ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);11 12 // MoE parameters13 ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);14 ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);15 ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);16 ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);17 18 // deepseek MLA parameters19 ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);20 ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);21 ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl, false);22 ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl, false);23 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);24 ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);25 26 // DSA parameters27 ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);28 ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);29 ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);30 31 // Expert gating function32 ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);33 34 if (ml.get_key(LLM_KV_ROPE_SCALING_YARN_LOG_MUL, hparams.rope_yarn_log_mul, 0.0f)) {35 // [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]36 // cancel the factor from the convert script37 hparams.rope_yarn_log_mul /= 0.1f;38 }39 40 switch (hparams.n_layer()) {41 case 61: type = LLM_TYPE_685B_A37B; break;42 default: type = LLM_TYPE_UNKNOWN;43 }44}45 46void llama_model_deepseek32::load_arch_tensors(llama_model_loader & ml) {47 LLAMA_LOAD_LOCALS;48 49 const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);50 const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";51 const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);52 const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;53 int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;54 55 if (!ml.load_mtp) {56 mtp_flags |= TENSOR_SKIP;57 }58 59 const bool is_mla = hparams.is_mla();60 if (!is_mla) {61 throw std::runtime_error("DEEPSEEK32 architecture requires MLA");62 }63 64 // note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA65 const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();66 const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();67 68 const int64_t n_embd_head_qk_rope = hparams.n_rot();69 const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope;70 71 const int64_t q_lora_rank = hparams.n_lora_q;72 const int64_t kv_lora_rank = hparams.n_lora_kv;73 74 const int64_t n_ff_exp = hparams.n_ff_exp();75 const int64_t n_expert_shared = hparams.n_expert_shared;76 77 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);78 79 // output80 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);81 // try to load output.weight, if not found, use token_embd (tied embeddings)82 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);83 if (!output) {84 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);85 }86 87 for (int i = 0; i < n_layer_all; ++i) {88 const int flags = (i >= n_layer) ? mtp_flags : trunk_flags;89 90 auto & layer = layers[i];91 92 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);93 layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags);94 layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags);95 96 layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags);97 layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k_mla}, flags);98 99 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);100 101 // note: only old legacy GGUF files will have the unsplit wkv_b tensor in102 layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, flags);103 layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, flags);104 105 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, flags);106 107 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);108 109 // DSA indexer110 layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, flags);111 layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {hparams.indexer_head_size}, flags);112 layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, flags);113 layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, hparams.indexer_head_size}, flags);114 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);115 if (i < (int) hparams.n_layer_dense_lead) {116 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags);117 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, flags);118 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags);119 } else {120 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags);121 layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);122 123 if (n_expert == 0) {124 throw std::runtime_error("n_expert must be > 0");125 }126 if (n_expert_used == 0) {127 throw std::runtime_error("n_expert_used must be > 0");128 }129 130 // MoE branch131 layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags);132 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags);133 layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags);134 135 // Shared expert branch136 layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);137 layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, flags);138 layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);139 }140 141 // NextN/MTP tensors - conditionally load for last nextn_predict_layers142 if (i >= n_layer) {143 layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);144 layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);145 layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags);146 147 // Optional tensors148 layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED);149 layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab }, flags | TENSOR_NOT_REQUIRED);150 layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, flags | TENSOR_NOT_REQUIRED);151 }152 }153}154 155std::unique_ptr<llm_graph_context> llama_model_deepseek32::build_arch_graph(const llm_graph_params & params) const {156 if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {157 return std::make_unique<graph_mtp>(*this, params);158 }159 return std::make_unique<graph>(*this, params);160}161 162llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_params & params) :163 llm_graph_context(params) {164 const bool is_mla = hparams.is_mla();165 GGML_ASSERT(is_mla);166 167 // note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA168 const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();169 const int64_t n_embd_head_v = hparams.n_embd_head_v_mla();170 GGML_UNUSED(n_embd_head_v);171 172 const int64_t n_embd_head_qk_rope = hparams.n_rot();173 const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;174 175 const int64_t n_indexer_head = hparams.indexer_n_head;176 const int64_t n_embd_indexer_head = hparams.indexer_head_size;177 const uint32_t n_indexer_top_k = hparams.indexer_top_k;178 179 // the indexer head layous is [rope | nope]180 GGML_ASSERT(hparams.n_rot() <= n_embd_indexer_head);181 182 const uint32_t kv_lora_rank = hparams.n_lora_kv;183 184 // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.185 // See https://github.com/ggml-org/llama.cpp/discussions/7416 for detailed explanation.186 // And also: https://github.com/ggml-org/llama.cpp/pull/17945 [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]187 188 // first cancel the adjustment from llama_hparams::yarn_attn_factor_adjust to get the original attn_factor189 GGML_ASSERT(ext_factor >= 0.0f);190 const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale));191 192 // use the original attn_factor to pre-scale the kq_scale193 const float mscale = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));194 const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k));195 196 ggml_tensor * cur;197 ggml_tensor * inpL;198 199 // {n_embd, n_tokens}200 inpL = build_inp_embd(model.tok_embd);201 202 // inp_pos - contains the positions203 ggml_tensor * inp_pos = build_inp_pos();204 205 llm_graph_input_attn_k_dsa * inp_attn_dsa = build_attn_inp_k_dsa();206 207 ggml_tensor * inp_out_ids = build_inp_out_ids();208 209 for (int il = 0; il < n_layer; ++il) {210 ggml_tensor * inpSA = inpL;211 212 // norm213 cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);214 cb(cur, "attn_norm", il);215 216 // self_attention217 {218 ggml_tensor * qr = ggml_mul_mat(ctx0, model.layers[il].wq_a, cur);219 cb(qr, "qr", il);220 221 qr = build_norm(qr, model.layers[il].attn_q_a_norm, nullptr, LLM_NORM_RMS, il);222 cb(qr, "qr", il);223 224 ggml_tensor * top_k = nullptr;225 226 // lightning indexer227 {228 ggml_tensor * indexer_q = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_q_b, qr);229 cb(indexer_q, "indexer_q", il);230 231 // {n_embd_indexer_head, n_indexer_head, n_tokens}232 indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens);233 indexer_q = ggml_rope_ext(ctx0, indexer_q, inp_pos, nullptr, n_rot,234 LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale,235 ext_factor, attn_factor, beta_fast, beta_slow);236 cb(indexer_q, "indexer_q", il);237 238 ggml_tensor * indexer_k = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_k, cur);239 cb(indexer_k, "indexer_k", il);240 241 indexer_k = build_norm(indexer_k, model.layers[il].indexer_k_norm, model.layers[il].indexer_k_norm_b, LLM_NORM, il);242 cb(indexer_k, "indexer_k", il);243 244 // {n_embd_indexer_head, 1, n_tokens}245 indexer_k = ggml_reshape_3d(ctx0, indexer_k, n_embd_indexer_head, 1, n_tokens);246 indexer_k = ggml_rope_ext(ctx0, indexer_k, inp_pos, nullptr, n_rot,247 LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale,248 ext_factor, attn_factor, beta_fast, beta_slow);249 cb(indexer_k, "indexer_k", il);250 251 // perform Hadamard transform on indexer q and k252 indexer_q = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_q);253 cb(indexer_q, "indexer_q", il);254 indexer_k = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_k);255 cb(indexer_k, "indexer_k", il);256 257 // store indexer keys to KV cache258 const auto * mctx_lid = inp_attn_dsa->mctx->get_lid();259 const auto & k_idxs_lid = inp_attn_dsa->get_k_idxs_lid();260 ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, indexer_k, k_idxs_lid, il));261 262 // prepare indexer weights263 ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, model.layers[il].indexer_proj, cur);264 cb(indexer_weights, "indexer_weights", il);265 266 // get cached indexer keys267 indexer_k = mctx_lid->get_k(ctx0, il);268 269 // split the batch into streams if needed270 const auto n_stream = indexer_k->ne[3];271 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);272 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);273 274 // pre-scale weights to avoid scaling operations on huge indexer_score tensor275 indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head * n_indexer_head)));276 cb(indexer_weights, "indexer_weights", il);277 278 ggml_tensor * indexer_score = nullptr;279 if (cparams.fused_lid) {280 indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());281 cb(indexer_score, "indexer_score", il);282 res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});283 } else {284 // calculate indexer kq285 indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);286 cb(indexer_q, "indexer_q", il);287 indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);288 cb(indexer_k, "indexer_k", il);289 290 ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);291 cb(indexer_kq, "indexer_kq", il);292 293 // ReLU requires contiguous tensors294 indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));295 cb(indexer_kq, "indexer_kq", il);296 297 // apply ReLU298 indexer_score = ggml_relu(ctx0, indexer_kq);299 cb(indexer_score, "indexer_score", il);300 301 // multiply scores by indexer weights302 indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);303 cb(indexer_score, "indexer_score", il);304 305 // sum by q n_indexer_head dimension306 indexer_score = ggml_sum_rows(ctx0, indexer_score);307 cb(indexer_score, "indexer_score", il);308 309 // permute result to match KQ mask310 indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));311 cb(indexer_score, "indexer_score", il);312 313 // mask indexer scores314 ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();315 indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);316 cb(indexer_score, "indexer_score", il);317 }318 319 // get indices of top k indexer scores320 uint32_t n_top_k = indexer_score->ne[0] < n_indexer_top_k ? indexer_score->ne[0] : n_indexer_top_k;321 top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));322 cb(top_k, "top_k", il);323 }324 325 ggml_tensor * q = ggml_mul_mat(ctx0, model.layers[il].wq_b, qr);326 cb(q, "q", il);327 328 // split into {n_embd_head_qk_nope, n_head, n_tokens}329 ggml_tensor * q_nope =330 ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),331 ggml_row_size(q->type, n_embd_head_k) * n_head, 0);332 cb(q_nope, "q_nope", il);333 334 // and {n_embd_head_qk_rope, n_head, n_tokens}335 ggml_tensor * q_pe = ggml_view_3d(336 ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),337 ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope));338 cb(q_pe, "q_pe", il);339 340 ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur);341 cb(kv_cmpr_pe, "kv_cmpr_pe", il);342 343 // split into {kv_lora_rank, n_tokens}344 ggml_tensor * kv_cmpr =345 ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,346 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);347 cb(kv_cmpr, "kv_cmpr", il);348 349 // and {n_embd_head_qk_rope, 1, n_tokens}350 ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,351 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),352 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),353 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));354 cb(k_pe, "k_pe", il);355 356 q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,357 ext_factor, attn_factor, beta_fast, beta_slow);358 cb(q_pe, "q_pe", il);359 360 k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,361 ext_factor, attn_factor, beta_fast, beta_slow);362 cb(k_pe, "k_pe", il);363 364 kv_cmpr = build_norm(kv_cmpr, model.layers[il].attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);365 cb(kv_cmpr, "kv_cmpr", il);366 367 // MLA attention368 {369 // {n_embd_head_qk_nope, n_tokens, n_head}370 q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);371 cb(q_nope, "q_nope_perm", il);372 373 // {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}374 ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, model.layers[il].wk_b, q_nope);375 cb(q_nope_absorbed, "q_nope_absorbed", il);376 377 // {kv_lora_rank, n_head, n_tokens}378 q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);379 cb(q_nope_absorbed, "q_nope_absorbed_perm", il);380 381 // {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens}382 // note: rope must go first for in-place context shifting in build_rope_shift()383 ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);384 cb(Qcur, "Qcur", il);385 386 kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);387 cb(kv_cmpr, "kv_cmpr_reshape", il);388 389 // {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens}390 ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);391 cb(Kcur, "Kcur", il);392 393 // {kv_lora_rank, 1, n_tokens}394 ggml_tensor * Vcur = kv_cmpr;395 cb(Vcur, "Vcur", il);396 397 // note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)398 cur = build_attn(inp_attn_dsa,399 model.layers[il].wo, NULL, model.layers[il].wo_s,400 Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, il);401 }402 }403 // when unmasked nextn embeddings are requested, t_h_nextn must keep all rows,404 // so the early output masking has to be skipped (it is applied after the final norm instead)405 if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {406 cur = ggml_get_rows(ctx0, cur, inp_out_ids);407 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);408 }409 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);410 cb(ffn_inp, "ffn_inp", il);411 412 cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);413 cb(cur, "ffn_norm", il);414 415 if ((uint32_t) il < hparams.n_layer_dense_lead) {416 cur = build_ffn(cur,417 model.layers[il].ffn_up, NULL, model.layers[il].ffn_up_s,418 model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,419 model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,420 NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);421 cb(cur, "ffn_out", il);422 } else {423 // MoE branch424 ggml_tensor * moe_out = build_moe_ffn(cur,425 model.layers[il].ffn_gate_inp,426 model.layers[il].ffn_up_exps,427 model.layers[il].ffn_gate_exps,428 model.layers[il].ffn_down_exps,429 model.layers[il].ffn_exp_probs_b,430 n_expert, n_expert_used,431 LLM_FFN_SILU, hparams.expert_weights_norm,432 hparams.expert_weights_scale,433 (llama_expert_gating_func_type) hparams.expert_gating_func,434 il,435 nullptr,436 model.layers[il].ffn_gate_up_exps,437 model.layers[il].ffn_up_exps_s,438 model.layers[il].ffn_gate_exps_s,439 model.layers[il].ffn_down_exps_s);440 cb(moe_out, "ffn_moe_out", il);441 442 // FFN shared expert443 {444 ggml_tensor * ffn_shexp =445 build_ffn(cur,446 model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,447 model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,448 model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,449 NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);450 cb(ffn_shexp, "ffn_shexp", il);451 452 cur = ggml_add(ctx0, moe_out, ffn_shexp);453 cb(cur, "ffn_out", il);454 }455 }456 cur = ggml_add(ctx0, cur, ffn_inp);457 458 cur = build_cvec(cur, il);459 cb(cur, "l_out", il);460 461 // input for next layer462 inpL = cur;463 }464 cur = inpL;465 466 cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);467 468 // post-norm hidden state feeds the NextN/MTP draft head469 cb(cur, "h_nextn", -1);470 res->t_h_nextn = cur;471 472 if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {473 cur = ggml_get_rows(ctx0, cur, inp_out_ids);474 }475 476 cb(cur, "result_norm", -1);477 res->t_embd = cur;478 479 // lm_head480 cur = ggml_mul_mat(ctx0, model.output, cur);481 482 cb(cur, "result_output", -1);483 res->t_logits = cur;484 485 ggml_build_forward_expand(gf, cur);486}487 488// LLM_GRAPH_TYPE_DECODER_MTP draft head for DeepSeek V3.2 (DEEPSEEK32).489// Semantics mirror the deepseek-family NextN/MTP layer:490// enorm(embed) + hnorm(prev_hidden) -> concat(e, h) -> eh_proj ->491// full deepseek32 decoder block (dense MLA attention + sigmoid-gated MoE FFN492// with shared expert, exactly as the trunk deepseek2 graph builds it) ->493// shared_head_norm (fallback output_norm) -> shared LM head.494// The DSA indexer is not used at runtime.495llama_model_deepseek32::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)496 : llm_graph_context(params) {497 GGML_ASSERT(hparams.n_layer_nextn > 0 && "DEEPSEEK32 MTP requires n_layer_nextn > 0");498 GGML_ASSERT(hparams.n_layer_nextn == 1 && "DEEPSEEK32 MTP currently only supports a single MTP block");499 GGML_ASSERT(hparams.is_mla() && "DEEPSEEK32 MTP requires MLA");500 501 const int il = hparams.n_layer() + cparams.nextn_layer_offset;502 GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&503 cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&504 "nextn_layer_offset out of range [0, n_layer_nextn)");505 const auto & layer = model.layers[il];506 507 GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");508 GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");509 GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");510 GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");511 512 // note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA513 const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();514 515 const int64_t n_embd_head_qk_rope = hparams.n_rot();516 const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;517 518 const uint32_t kv_lora_rank = hparams.n_lora_kv;519 520 // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.521 // See the deepseek2 trunk graph for the detailed explanation - this must match it EXACTLY.522 GGML_ASSERT(ext_factor >= 0.0f);523 const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale));524 525 const float mscale = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));526 const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k));527 528 // TODO: extract in a common llm_graph_context::build_inp_embd_h()529 auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);530 531 inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);532 ggml_set_input(inp->tokens);533 534 inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);535 ggml_set_input(inp->embd);536 537 ggml_tensor * tok_embd;538 if (ubatch.token) {539 ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;540 541 tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);542 } else {543 tok_embd = inp->embd;544 }545 cb(tok_embd, "mtp_tok_embd", il);546 547 inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);548 ggml_set_input(inp->h);549 ggml_set_name(inp->h, "mtp_h_input");550 551 ggml_tensor * h_embd = inp->h;552 553 res->add_input(std::move(inp));554 555 ggml_tensor * inp_pos = build_inp_pos();556 ggml_tensor * inp_out_ids = build_inp_out_ids();557 558 // MLA with the absorption optimization uses a K-only cache (V is a view of K)559 auto * inp_attn = build_attn_inp_k();560 561 ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);562 cb(h_norm, "mtp_hnorm", il);563 564 ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);565 cb(e_norm, "mtp_enorm", il);566 567 ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);568 cb(concat, "mtp_concat", il);569 570 ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);571 cb(cur, "mtp_eh_proj", il);572 573 ggml_tensor * inpSA = cur;574 575 cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);576 cb(cur, "mtp_attn_norm", il);577 578 // self-attention: dense MLA, same construction as the deepseek2 trunk graph579 {580 ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur);581 cb(q, "mtp_q", il);582 583 q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);584 cb(q, "mtp_q", il);585 586 q = ggml_mul_mat(ctx0, layer.wq_b, q);587 cb(q, "mtp_q", il);588 589 // split into {n_embd_head_qk_nope, n_head, n_tokens}590 ggml_tensor * q_nope =591 ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),592 ggml_row_size(q->type, n_embd_head_k) * n_head, 0);593 cb(q_nope, "mtp_q_nope", il);594 595 // and {n_embd_head_qk_rope, n_head, n_tokens}596 ggml_tensor * q_pe = ggml_view_3d(597 ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),598 ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope));599 cb(q_pe, "mtp_q_pe", il);600 601 ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);602 cb(kv_cmpr_pe, "mtp_kv_cmpr_pe", il);603 604 // split into {kv_lora_rank, n_tokens}605 ggml_tensor * kv_cmpr =606 ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,607 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);608 cb(kv_cmpr, "mtp_kv_cmpr", il);609 610 // and {n_embd_head_qk_rope, 1, n_tokens}611 ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,612 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),613 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),614 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));615 cb(k_pe, "mtp_k_pe", il);616 617 q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,618 ext_factor, attn_factor, beta_fast, beta_slow);619 cb(q_pe, "mtp_q_pe", il);620 621 k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,622 ext_factor, attn_factor, beta_fast, beta_slow);623 cb(k_pe, "mtp_k_pe", il);624 625 kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);626 cb(kv_cmpr, "mtp_kv_cmpr", il);627 628 // {n_embd_head_qk_nope, n_tokens, n_head}629 q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);630 cb(q_nope, "mtp_q_nope_perm", il);631 632 // {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}633 ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);634 cb(q_nope_absorbed, "mtp_q_nope_absorbed", il);635 636 // {kv_lora_rank, n_head, n_tokens}637 q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);638 cb(q_nope_absorbed, "mtp_q_nope_absorbed_perm", il);639 640 // {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens}641 // note: rope must go first for in-place context shifting in build_rope_shift()642 ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);643 cb(Qcur, "mtp_Qcur", il);644 645 kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);646 cb(kv_cmpr, "mtp_kv_cmpr_reshape", il);647 648 // {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens}649 ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);650 cb(Kcur, "mtp_Kcur", il);651 652 // {kv_lora_rank, 1, n_tokens}653 ggml_tensor * Vcur = kv_cmpr;654 cb(Vcur, "mtp_Vcur", il);655 656 // note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)657 cur = build_attn(inp_attn,658 layer.wo, NULL, layer.wo_s,659 Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale, il);660 cb(cur, "mtp_attn_out", il);661 }662 663 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);664 cb(ffn_inp, "mtp_ffn_inp", il);665 666 cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);667 cb(cur, "mtp_ffn_norm", il);668 669 // MoE FFN with shared expert - same construction as the deepseek2 trunk graph670 ggml_tensor * moe_out = build_moe_ffn(cur,671 layer.ffn_gate_inp,672 layer.ffn_up_exps,673 layer.ffn_gate_exps,674 layer.ffn_down_exps,675 layer.ffn_exp_probs_b,676 n_expert, n_expert_used,677 LLM_FFN_SILU, hparams.expert_weights_norm,678 hparams.expert_weights_scale,679 (llama_expert_gating_func_type) hparams.expert_gating_func,680 il,681 nullptr,682 layer.ffn_gate_up_exps,683 layer.ffn_up_exps_s,684 layer.ffn_gate_exps_s,685 layer.ffn_down_exps_s);686 cb(moe_out, "mtp_ffn_moe_out", il);687 688 // FFN shared expert689 ggml_tensor * ffn_shexp =690 build_ffn(cur,691 layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,692 layer.ffn_gate_shexp, NULL, layer.ffn_gate_shexp_s,693 layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,694 NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);695 cb(ffn_shexp, "mtp_ffn_shexp", il);696 697 cur = ggml_add(ctx0, moe_out, ffn_shexp);698 cb(cur, "mtp_ffn_out", il);699 700 cur = ggml_add(ctx0, cur, ffn_inp);701 cb(cur, "mtp_post_ffn", il);702 703 // shared_head_norm applied after the decoder block, before the shared LM head.704 // The post-norm hidden state seeds the next MTP step.705 ggml_tensor * head_norm_w = layer.nextn.shared_head_norm706 ? layer.nextn.shared_head_norm707 : model.output_norm;708 GGML_ASSERT(head_norm_w && "DEEPSEEK32 MTP: missing both nextn.shared_head_norm and output_norm");709 cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);710 711 cb(cur, "h_nextn", -1);712 res->t_h_nextn = cur;713 714 cur = ggml_get_rows(ctx0, cur, inp_out_ids);715 cb(cur, "mtp_shared_head_norm", -1);716 717 ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;718 ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;719 GGML_ASSERT(head_w && "DEEPSEEK32 MTP: missing LM head (nextn.shared_head_head or model.output)");720 cur = build_lora_mm(head_w, cur, head_s);721 cb(cur, "result_output", -1);722 723 res->t_logits = cur;724 ggml_build_forward_expand(gf, cur);725}726 727