Brunobkr/llama.cpp_AlgMor24_github

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