Brunobkr/llama.cpp_AlgMor24_github

1#include "models.h"2 3#include "llama-kv-cache-dsv4.h"4 5#include <algorithm>6#include <cmath>7#include <stdexcept>8#include <string>9 10static float dsv4_rope_attn_factor(float freq_scale, float ext_factor) {11    if (ext_factor == 0.0f) {12        return 1.0f;13    }14 15    return 1.0f / (1.0f + 0.1f*logf(1.0f/freq_scale));16}17 18void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {19    ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);20    if (hparams.n_layer_nextn > 0 && hparams.n_layer_nextn < hparams.n_layer_all) {21        const uint32_t n_layer_main = hparams.n_layer_all - hparams.n_layer_nextn;22        const std::string mtp_probe = "blk." + std::to_string(n_layer_main) + ".nextn.eh_proj.weight";23        if (ml.get_weight(mtp_probe.c_str()) == nullptr) {24            hparams.n_layer_nextn = 0;25        }26    }27    GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < block_count");28 29    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);30    ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK,       hparams.n_lora_q);31    ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW,    hparams.n_swa);32 33    ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,  hparams.n_ff_exp);34    ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,         hparams.n_expert_shared);35    ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,        hparams.expert_weights_scale);36    ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,         hparams.expert_weights_norm);37    ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP,     hparams.swiglu_clamp_exp,   hparams.n_layer_all);38    if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP,   hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) {39        hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp;40    }41 42    ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);43    ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);44    ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K,      hparams.indexer_top_k);45 46    ml.get_key(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,         hparams.dsv4_o_group_count);47    ml.get_key(LLM_KV_ATTENTION_OUTPUT_LORA_RANK,           hparams.dsv4_o_lora_rank);48    ml.get_key(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,    hparams.dsv4_compress_rope_base);49    ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT,               hparams.dsv4_hc_mult);50    ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);51    ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON,             hparams.dsv4_hc_eps);52    ml.get_key(LLM_KV_HASH_LAYER_COUNT,                     hparams.dsv4_hash_layer_count);53 54    hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd;55 56    uint32_t n_compress_ratios = 0;57    ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_compress_ratios);58    if (n_compress_ratios < hparams.n_layer_all) {59        throw std::runtime_error("DeepSeek-V4 compress_ratios is shorter than block_count");60    }61    ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios);62 63    ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);64    if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {65        throw std::runtime_error("DeepSeek-V4 loader currently expects sqrtsoftplus MoE scoring");66    }67    hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;68    hparams.set_swa_pattern(0);69    for (uint32_t il = hparams.n_layer(); il < hparams.n_layer_all; ++il) {70        hparams.is_swa_impl[il] = true;71    }72 73    switch (hparams.n_layer()) {74        case 43: type = LLM_TYPE_UNKNOWN; break;75        default: type = LLM_TYPE_UNKNOWN;76    }77}78 79void llama_model_deepseek4::load_arch_tensors(llama_model_loader & ml) {80    LLAMA_LOAD_LOCALS;81 82    const int64_t q_lora_rank     = hparams.n_lora_q;83    const int64_t n_ff_exp        = hparams.n_ff_exp;84    const int64_t n_expert_shared = hparams.n_expert_shared;85 86    const int64_t n_embd_head = hparams.n_embd_head_k();87    const int64_t o_groups    = hparams.dsv4_o_group_count;88    const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;89    const int64_t hc_mult     = hparams.dsv4_hc_mult;90    const int64_t hc_dim      = hc_mult * n_embd;91    const int64_t hc_mix_dim  = (2 + hc_mult) * hc_mult;92 93    const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);94    const int trunk_flags = mtp_only    ? TENSOR_NOT_REQUIRED : 0;95    const int mtp_flags   = ml.load_mtp ? 0 : TENSOR_SKIP;96 97    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);98 99    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);100    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, 0);101 102    hc_head_fn    = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"),    {hc_dim, hc_mult}, 0);103    hc_head_base  = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"),  {hc_mult}, 0);104    hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);105 106    for (int i = 0; i < n_layer_all; ++i) {107        auto & layer = layers[i];108        const int flags = i < n_layer ? trunk_flags : mtp_flags;109 110        layer.attn_norm     = create_tensor(tn(LLM_TENSOR_ATTN_NORM,     "weight", i), {n_embd}, flags);111        layer.attn_sinks    = create_tensor(tn(LLM_TENSOR_ATTN_SINKS,    "weight", i), {n_head}, flags);112        layer.wq_a          = create_tensor(tn(LLM_TENSOR_ATTN_Q_A,      "weight", i), {n_embd, q_lora_rank}, flags);113        layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags);114        layer.wq_b          = create_tensor(tn(LLM_TENSOR_ATTN_Q_B,      "weight", i), {q_lora_rank, n_head * n_embd_head}, flags);115        layer.wkv           = create_tensor(tn(LLM_TENSOR_ATTN_KV,       "weight", i), {n_embd, n_embd_head}, flags);116        layer.attn_kv_norm  = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM,  "weight", i), {n_embd_head}, flags);117        // for wo_a, the shape in the file is (n_head * n_embd_head / o_groups, o_lora_rank*o_groups)118        // so we reshape here, to avoid reshaping the tensor in the graph119        layer.wo_a          = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A,    "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank, o_groups}, flags | TENSOR_ALLOW_RESHAPE);120        layer.wo_b          = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B,    "weight", i), {o_groups * o_lora_rank, n_embd}, flags);121 122        layer.hc_attn_fn    = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN,    "weight", i), {hc_dim, hc_mix_dim}, flags);123        layer.hc_attn_base  = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE,  "weight", i), {hc_mix_dim}, flags);124        layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, flags);125        layer.hc_ffn_fn     = create_tensor(tn(LLM_TENSOR_HC_FFN_FN,     "weight", i), {hc_dim, hc_mix_dim}, flags);126        layer.hc_ffn_base   = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE,   "weight", i), {hc_mix_dim}, flags);127        layer.hc_ffn_scale  = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE,  "weight", i), {3}, flags);128 129        const int64_t ratio = hparams.dsv4_compress_ratios[i];130        if (ratio != 0) {131            const int64_t coff = ratio == 4 ? 2 : 1;132 133            layer.attn_comp_wkv   = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV,   "weight", i), {n_embd, coff * n_embd_head}, flags);134            layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, coff * n_embd_head}, flags);135            layer.attn_comp_ape   = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_APE,   "weight", i), {coff * n_embd_head, ratio}, flags);136            layer.attn_comp_norm  = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM,  "weight", i), {n_embd_head}, flags);137 138            if (ratio == 4) {139                const int64_t n_embd_indexer = hparams.indexer_head_size;140 141                layer.indexer_proj     = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ,     "weight", i), {n_embd, hparams.indexer_n_head}, flags);142                layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, flags);143 144                layer.indexer_comp_wkv   = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WKV,   "weight", i), {n_embd, 2 * n_embd_indexer}, flags);145                layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, 2 * n_embd_indexer}, flags);146                layer.indexer_comp_ape   = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE,   "weight", i), {2 * n_embd_indexer, ratio}, flags);147                layer.indexer_comp_norm  = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_NORM,  "weight", i), {n_embd_indexer}, flags);148            } else if (ratio != 128) {149                throw std::runtime_error("DeepSeek-V4 loader only supports compression ratios 0, 4, and 128");150            }151        }152 153        layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags);154        if ((uint32_t) i < hparams.dsv4_hash_layer_count) {155            layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, "weight", i), {n_expert_used, n_vocab}, flags);156        } else {157            layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags);158        }159        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);160 161        layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd,   n_ff_exp, n_expert}, flags);162        layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd,   n_expert}, flags);163        layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {n_embd,   n_ff_exp, n_expert}, flags);164 165        layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd,                     n_ff_exp * n_expert_shared}, flags);166        layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd                    }, flags);167        layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd,                     n_ff_exp * n_expert_shared}, flags);168 169        if (i >= n_layer) {170            layer.nextn.eh_proj          = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ,          "weight", i), {2 * n_embd, n_embd}, flags);171            layer.nextn.enorm            = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM,            "weight", i), {n_embd},             flags);172            layer.nextn.hnorm            = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM,            "weight", i), {n_embd},             flags);173            layer.nextn.embed_tokens     = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS,     "weight", i), {n_embd, n_vocab},    TENSOR_NOT_REQUIRED | flags);174            layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab},    TENSOR_NOT_REQUIRED | flags);175            layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd},             TENSOR_NOT_REQUIRED | flags);176        }177    }178}179 180std::unique_ptr<llm_graph_context> llama_model_deepseek4::build_arch_graph(const llm_graph_params & params) const {181    if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {182        return std::make_unique<graph_mtp>(*this, params);183    }184    return std::make_unique<graph>(*this, params);185}186 187static size_t dsv4_elem_offset(const ggml_tensor * t, int64_t i) {188    return ggml_row_size(t->type, i);189}190 191static ggml_tensor * dsv4_view_1d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t i0) {192    return ggml_view_1d(ctx, t, ne0, dsv4_elem_offset(t, i0));193}194 195static ggml_tensor * dsv4_view_2d(196        ggml_context * ctx,197        ggml_tensor  * t,198        int64_t        ne0,199        int64_t        ne1,200        int64_t        i0) {201    return ggml_view_2d(ctx, t, ne0, ne1, t->nb[1], dsv4_elem_offset(t, i0));202}203 204static ggml_tensor * dsv4_append_zero_row(ggml_context * ctx, ggml_tensor * t, bool neg_inf) {205    ggml_tensor * row = ggml_view_1d(ctx, t, t->ne[0], 0);206    row = neg_inf ? ggml_scale_bias(ctx, row, 0.0f, -INFINITY) : ggml_scale(ctx, row, 0.0f);207    row = ggml_reshape_2d(ctx, row, t->ne[0], 1);208 209    return ggml_concat(ctx, t, row, 1);210}211 212struct dsv4_state_tensors {213    ggml_tensor * kv;214    ggml_tensor * score;215};216 217static dsv4_state_tensors dsv4_build_state_restore(218        ggml_context * ctx,219        const llm_graph_input_dsv4::comp_input & inp,220        const llama_dsv4_comp_state * state,221        int32_t il) {222    dsv4_state_tensors restored = {223        state->get_kv_all(ctx, il),224        state->get_score_all(ctx, il),225    };226 227    if (inp.state_restore_src_idxs == nullptr || inp.state_restore_dst_idxs == nullptr) {228        return restored;229    }230 231    ggml_tensor * kv_rows = ggml_get_rows(ctx, restored.kv, inp.state_restore_src_idxs);232    restored.kv = state->cpy_kv(ctx, kv_rows, inp.state_restore_dst_idxs, il);233 234    ggml_tensor * score_rows = ggml_get_rows(ctx, restored.score, inp.state_restore_src_idxs);235    restored.score = state->cpy_score(ctx, score_rows, inp.state_restore_dst_idxs, il);236 237    return restored;238}239 240static dsv4_state_tensors dsv4_build_state_snapshot(241        ggml_context * ctx,242        const llm_graph_input_dsv4::comp_input & inp,243        const llama_dsv4_comp_state * state,244        ggml_tensor * source_kv,245        ggml_tensor * source_score,246        int32_t il) {247    if (inp.state_snapshot_src_idxs == nullptr || inp.state_snapshot_dst_idxs == nullptr ||248            source_kv == nullptr || source_score == nullptr) {249        return {};250    }251 252    ggml_tensor * kv_rows = ggml_get_rows(ctx, source_kv, inp.state_snapshot_src_idxs);253    ggml_tensor * kv = state->cpy_kv(ctx, kv_rows, inp.state_snapshot_dst_idxs, il);254 255    ggml_tensor * score_rows = ggml_get_rows(ctx, source_score, inp.state_snapshot_src_idxs);256    ggml_tensor * score = state->cpy_score(ctx, score_rows, inp.state_snapshot_dst_idxs, il);257 258    return { kv, score };259}260 261static constexpr int64_t DSV4_CSA_RATIO  = 4;262static constexpr int64_t DSV4_HCA_RATIO  = 128;263 264// mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens]265static ggml_tensor * dsv4_hc_mean(ggml_context * ctx, ggml_tensor * x) {266    const int64_t hc = x->ne[1];267 268    ggml_tensor * acc = ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], 0);269    for (int64_t s = 1; s < hc; ++s) {270        acc = ggml_add(ctx, acc, ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1]));271    }272    return ggml_scale(ctx, acc, 1.0f/hc);273}274 275static ggml_tensor * dsv4_hc_affine(276        ggml_context * ctx,277        ggml_tensor  * x,278        ggml_tensor  * scale,279        ggml_tensor  * base) {280    x = ggml_mul(ctx, x, scale);281    x = ggml_add(ctx, x, base);282    return x;283}284 285ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(286        ggml_tensor * x,287        ggml_tensor * weights,288        int           il) const {289    GGML_ASSERT(x->ne[0] == n_embd);290    GGML_ASSERT(x->ne[1] == hparams.dsv4_hc_mult);291 292    const int64_t hc = hparams.dsv4_hc_mult;293    const int64_t nt = x->ne[2];294 295    if (cparams.fused_dsv4_hc_pre && il >= 0) {296        ggml_tensor * result = ggml_dsv4_hc_pre(ctx0, x, weights);297        res->add_fused_node({LLM_FUSED_OP_DSV4_HC_PRE, result, il});298        return result;299    }300 301    ggml_tensor * result = nullptr;302    for (int64_t ih = 0; ih < hc; ++ih) {303        ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih*x->nb[1]);304        ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih*weights->nb[0]);305        ggml_tensor * cur = ggml_mul(ctx0, xh, wh);306        result = result ? ggml_add(ctx0, result, cur) : cur;307    }308 309    return result;310}311 312ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn(313        ggml_tensor * comb,314        int           il) const {315    GGML_UNUSED(il);316 317    // comb is [dst_hc, src_hc, n_tokens]. Sinkhorn follows the reference:318    // row softmax over dst, one column normalization, then repeated row/column normalization.319    comb = ggml_soft_max(ctx0, comb);320 321    ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);322    eps = ggml_fill(ctx0, eps, hparams.dsv4_hc_eps);323 324    comb = ggml_add(ctx0, comb, eps);325 326    auto norm_cols = [&]() {327        ggml_tensor * comb_src_dst = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3));328        ggml_tensor * col_sum = ggml_sum_rows(ctx0, comb_src_dst);329        col_sum = ggml_add(ctx0, col_sum, eps);330        col_sum = ggml_permute(ctx0, col_sum, 1, 0, 2, 3);331        comb = ggml_div(ctx0, comb, col_sum);332    };333 334    auto norm_rows = [&]() {335        ggml_tensor * row_sum = ggml_sum_rows(ctx0, comb);336        row_sum = ggml_add(ctx0, row_sum, eps);337        comb = ggml_div(ctx0, comb, row_sum);338    };339 340    norm_cols();341    for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) {342        norm_rows();343        norm_cols();344    }345 346    return comb;347}348 349ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(350        ggml_tensor * x,351        ggml_tensor * hc_fn,352        ggml_tensor * hc_scale,353        ggml_tensor * hc_base,354        ggml_tensor ** post,355        ggml_tensor ** comb,356        int il) const {357    const int64_t hc         = hparams.dsv4_hc_mult;358    const int64_t hc_dim     = hc*n_embd;359    const int64_t hc_mix_dim = (2 + hc)*hc;360    const int64_t nt         = x->ne[2];361 362    GGML_ASSERT(hc == 4);363    GGML_ASSERT(hc_fn->ne[1] == hc_mix_dim);364 365    ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt);366    ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps);367    ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm);368    cb(mixes, "hc_mixes", il);369 370    ggml_tensor * scale_pre  = dsv4_view_1d(ctx0, hc_scale, 1, 0);371    ggml_tensor * scale_post = dsv4_view_1d(ctx0, hc_scale, 1, 1);372 373    ggml_tensor * base_pre  = dsv4_view_1d(ctx0, hc_base, hc, 0);374    ggml_tensor * base_post = dsv4_view_1d(ctx0, hc_base, hc, hc);375 376    ggml_tensor * pre = dsv4_view_2d(ctx0, mixes, hc, nt, 0);377    pre = dsv4_hc_affine(ctx0, pre, scale_pre, base_pre);378    pre = ggml_sigmoid(ctx0, pre);379    pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps);380    cb(pre, "hc_pre", il);381 382    *post = dsv4_view_2d(ctx0, mixes, hc, nt, hc);383    *post = dsv4_hc_affine(ctx0, *post, scale_post, base_post);384    *post = ggml_sigmoid(ctx0, *post);385    *post = ggml_scale(ctx0, *post, 2.0f);386    cb(*post, "hc_post", il);387 388    if (cparams.fused_dsv4_hc_comb) {389        *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps,390                (int32_t) hparams.dsv4_hc_sinkhorn_iters);391        res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il});392    } else {393        ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2);394        ggml_tensor * base_comb  = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc);395 396        *comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);397        *comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb);398        *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt);399        *comb = build_hc_sinkhorn(*comb, il);400    }401    cb(*comb, "hc_comb", il);402 403    ggml_tensor * result = build_hc_pre(x, pre, il);404    return result;405}406 407ggml_tensor * llama_model_deepseek4::graph::build_hc_post(408        ggml_tensor * x,409        ggml_tensor * residual,410        ggml_tensor * post,411        ggml_tensor * comb,412        int il) const {413    GGML_ASSERT(x->ne[0] == n_embd);414    GGML_ASSERT(residual->ne[1] == hparams.dsv4_hc_mult);415 416    if (cparams.fused_dsv4_hc_post) {417        ggml_tensor * result = ggml_dsv4_hc_post(ctx0, x, residual, post, comb);418        res->add_fused_node({LLM_FUSED_OP_DSV4_HC_POST, result, il});419        return result;420    }421 422    const int64_t hc = hparams.dsv4_hc_mult;423    const int64_t nt = x->ne[1];424 425    ggml_tensor * out = nullptr;426    for (int64_t dst = 0; dst < hc; ++dst) {427        ggml_tensor * post_dst = ggml_view_2d(ctx0, post, 1, nt, post->nb[1], dst*post->nb[0]);428        ggml_tensor * cur = ggml_mul(ctx0, x, post_dst);429 430        for (int64_t src = 0; src < hc; ++src) {431            ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src*residual->nb[1]);432            ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2],433                    dst*comb->nb[0] + src*comb->nb[1]);434            cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst));435        }436 437        cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt);438        out = out ? ggml_concat(ctx0, out, cur, 1) : cur;439    }440 441    return out;442}443 444ggml_tensor * llama_model_deepseek4::graph::build_hc_head(445        ggml_tensor * x,446        ggml_tensor * hc_fn,447        ggml_tensor * hc_scale,448        ggml_tensor * hc_base) const {449    const int64_t hc     = hparams.dsv4_hc_mult;450    const int64_t hc_dim = hc*n_embd;451    const int64_t nt     = x->ne[2];452 453    ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt);454    ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps);455    ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm);456    cb(mixes, "hc_head_mixes", -1);457 458    ggml_tensor * pre = dsv4_hc_affine(ctx0, mixes, hc_scale, hc_base);459    pre = ggml_sigmoid(ctx0, pre);460    pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps);461    cb(pre, "hc_head_pre", -1);462 463    return build_hc_pre(x, pre, -1);464}465 466ggml_tensor * llama_model_deepseek4::graph::build_hca_compressed_kv_from_state(467        ggml_tensor * kv_state,468        ggml_tensor * score_state,469        ggml_tensor * state_read_idxs,470        ggml_tensor * comp_pos,471        ggml_tensor * norm,472        int64_t n_embd_head,473        const char * name,474        int il) const {475    const int64_t n_embd_head_rope = hparams.n_rot();476    const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;477    const int64_t n_blocks         = comp_pos ? comp_pos->ne[0] : 0;478 479    GGML_ASSERT(n_blocks > 0);480    GGML_ASSERT(state_read_idxs);481    GGML_ASSERT(state_read_idxs->ne[0] == DSV4_HCA_RATIO*n_blocks);482    GGML_ASSERT(n_embd_head >= n_embd_head_rope);483 484    ggml_tensor * kv = ggml_get_rows(ctx0, kv_state, state_read_idxs);485    kv = ggml_reshape_3d(ctx0, kv, n_embd_head, DSV4_HCA_RATIO, n_blocks);486    cb(kv, name, il);487 488    ggml_tensor * score = ggml_get_rows(ctx0, score_state, state_read_idxs);489    score = ggml_reshape_3d(ctx0, score, n_embd_head, DSV4_HCA_RATIO, n_blocks);490    cb(score, name, il);491 492    ggml_tensor * values = ggml_cont(ctx0, ggml_permute(ctx0, kv, 1, 0, 2, 3));493    ggml_tensor * scores = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3));494 495    ggml_tensor * weights = ggml_soft_max(ctx0, scores);496    ggml_tensor * comp = ggml_mul(ctx0, values, weights);497    comp = ggml_sum_rows(ctx0, comp);498    comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3));499    cb(comp, name, il);500 501    comp = build_norm(comp, norm, nullptr, LLM_NORM_RMS, il);502    cb(comp, name, il);503 504    ggml_tensor * comp_nope = ggml_view_3d(ctx0, comp, n_embd_head_nope, 1, n_blocks,505            ggml_row_size(comp->type, n_embd_head),506            ggml_row_size(comp->type, n_embd_head),507            0);508    ggml_tensor * comp_pe = ggml_view_3d(ctx0, comp, n_embd_head_rope, 1, n_blocks,509            ggml_row_size(comp->type, n_embd_head),510            ggml_row_size(comp->type, n_embd_head),511            ggml_row_size(comp->type, n_embd_head_nope));512 513    comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig,514            hparams.dsv4_compress_rope_base, freq_scale, ext_factor,515            dsv4_rope_attn_factor(freq_scale, ext_factor), beta_fast, beta_slow);516    cb(comp_pe, name, il);517 518    comp = ggml_concat(ctx0, comp_nope, comp_pe, 0);519    cb(comp, name, il);520 521    return comp;522}523 524ggml_tensor * llama_model_deepseek4::graph::build_overlap_compressed_kv_from_state(525        ggml_tensor * kv_state,526        ggml_tensor * score_state,527        ggml_tensor * state_read_idxs,528        ggml_tensor * comp_pos,529        ggml_tensor * norm,530        int64_t ratio,531        int64_t n_embd_head,532        const char * name,533        int il) const {534    const int64_t n_embd_head_rope = hparams.n_rot();535    const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;536    const int64_t n_blocks         = comp_pos ? comp_pos->ne[0] : 0;537 538    GGML_ASSERT(n_blocks > 0);539    GGML_ASSERT(state_read_idxs);540    GGML_ASSERT(state_read_idxs->ne[0] == 2*ratio*n_blocks);541    GGML_ASSERT(kv_state->ne[0] == 2*n_embd_head);542    GGML_ASSERT(score_state->ne[0] == 2*n_embd_head);543    GGML_ASSERT(n_embd_head >= n_embd_head_rope);544 545    kv_state    = dsv4_append_zero_row(ctx0, kv_state,    false);546    score_state = dsv4_append_zero_row(ctx0, score_state, true);547 548    const int64_t n_read = ratio*n_blocks;549 550    ggml_tensor * kv_rows = ggml_get_rows(ctx0, kv_state, state_read_idxs);551    ggml_tensor * score_rows = ggml_get_rows(ctx0, score_state, state_read_idxs);552 553    ggml_tensor * kv_prev = ggml_cont(ctx0,554            ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], 0));555    kv_prev = ggml_reshape_3d(ctx0, kv_prev, n_embd_head, ratio, n_blocks);556    cb(kv_prev, name, il);557 558    ggml_tensor * score_prev = ggml_cont(ctx0,559            ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], 0));560    score_prev = ggml_reshape_3d(ctx0, score_prev, n_embd_head, ratio, n_blocks);561    cb(score_prev, name, il);562 563    ggml_tensor * kv_cur = ggml_cont(ctx0,564            ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1],565                n_read*kv_rows->nb[1] + ggml_row_size(kv_rows->type, n_embd_head)));566    kv_cur = ggml_reshape_3d(ctx0, kv_cur, n_embd_head, ratio, n_blocks);567 568    ggml_tensor * score_cur = ggml_cont(ctx0,569            ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1],570                n_read*score_rows->nb[1] + ggml_row_size(score_rows->type, n_embd_head)));571    score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks);572 573    ggml_tensor * values = ggml_concat(ctx0, kv_prev, kv_cur, 1);574    ggml_tensor * scores = ggml_concat(ctx0, score_prev, score_cur, 1);575 576    values = ggml_cont(ctx0, ggml_permute(ctx0, values, 1, 0, 2, 3));577    scores = ggml_cont(ctx0, ggml_permute(ctx0, scores, 1, 0, 2, 3));578 579    ggml_tensor * weights = ggml_soft_max(ctx0, scores);580    ggml_tensor * comp = ggml_mul(ctx0, values, weights);581    comp = ggml_sum_rows(ctx0, comp);582    comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3));583    cb(comp, name, il);584 585    comp = build_norm(comp, norm, nullptr, LLM_NORM_RMS, il);586    cb(comp, name, il);587 588    ggml_tensor * comp_nope = ggml_view_3d(ctx0, comp, n_embd_head_nope, 1, n_blocks,589            ggml_row_size(comp->type, n_embd_head),590            ggml_row_size(comp->type, n_embd_head),591            0);592    ggml_tensor * comp_pe = ggml_view_3d(ctx0, comp, n_embd_head_rope, 1, n_blocks,593            ggml_row_size(comp->type, n_embd_head),594            ggml_row_size(comp->type, n_embd_head),595            ggml_row_size(comp->type, n_embd_head_nope));596 597    comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig,598            hparams.dsv4_compress_rope_base, freq_scale, ext_factor,599            dsv4_rope_attn_factor(freq_scale, ext_factor), beta_fast, beta_slow);600    cb(comp_pe, name, il);601 602    comp = ggml_concat(ctx0, comp_nope, comp_pe, 0);603    cb(comp, name, il);604 605    return comp;606}607 608ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k(609        const llama_model & model,610        llm_graph_input_dsv4 * inp_dsv4,611        ggml_tensor * qr,612        ggml_tensor * cur,613        ggml_tensor * inp_pos,614        int il) const {615    const auto & layer = model.layers[il];616    const auto & inp_lid = inp_dsv4->get_lid();617    const int64_t n_embd_indexer_head      = hparams.indexer_head_size;618    const int64_t n_embd_indexer_head_rope = hparams.n_rot();619    const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope;620    const int64_t n_indexer_head           = hparams.indexer_n_head;621    const int64_t nt                       = cur->ne[1];622 623    GGML_ASSERT(inp_lid.kq_mask);624    GGML_ASSERT(inp_lid.k_rot);625    GGML_ASSERT(n_embd_indexer_head >= n_embd_indexer_head_rope);626 627    ggml_tensor * indexer_q = build_lora_mm(layer.indexer_attn_q_b, qr);628    indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, nt);629    cb(indexer_q, "lid_q", il);630 631    ggml_tensor * indexer_q_nope = ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, nt,632            ggml_row_size(indexer_q->type, n_embd_indexer_head),633            ggml_row_size(indexer_q->type, n_embd_indexer_head)*n_indexer_head,634            0);635    ggml_tensor * indexer_q_pe = ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, nt,636            ggml_row_size(indexer_q->type, n_embd_indexer_head),637            ggml_row_size(indexer_q->type, n_embd_indexer_head)*n_indexer_head,638            ggml_row_size(indexer_q->type, n_embd_indexer_head_nope));639 640    indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_embd_indexer_head_rope,641            rope_type, n_ctx_orig, hparams.dsv4_compress_rope_base, freq_scale,642            ext_factor, dsv4_rope_attn_factor(freq_scale, ext_factor), beta_fast, beta_slow);643    cb(indexer_q_pe, "lid_q_pe", il);644 645    indexer_q = ggml_concat(ctx0, indexer_q_nope, indexer_q_pe, 0);646    indexer_q = llama_mul_mat_hadamard(ctx0, indexer_q, inp_lid.k_rot);647    cb(indexer_q, "lid_q_rot", il);648 649    ggml_tensor * indexer_weights = build_lora_mm(layer.indexer_proj, cur);650    indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f/sqrtf(float(n_embd_indexer_head*n_indexer_head)));651    cb(indexer_weights, "lid_weights", il);652 653    ggml_tensor * indexer_k = inp_dsv4->mctx->get_lid()->get_k(ctx0, il);654    const int64_t n_lid = inp_lid.kq_mask->ne[0];655    GGML_ASSERT(n_lid > 0);656    GGML_ASSERT(n_lid <= indexer_k->ne[2]);657 658    indexer_k = ggml_view_4d(ctx0, indexer_k,659            indexer_k->ne[0], indexer_k->ne[1], n_lid, indexer_k->ne[3],660            indexer_k->nb[1], indexer_k->nb[2], indexer_k->nb[3], 0);661    cb(indexer_k, "lid_k", il);662 663    const int64_t n_stream = indexer_k->ne[3];664    indexer_q = ggml_view_4d(ctx0, indexer_q,665            indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream,666            indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0);667    indexer_weights = ggml_view_4d(ctx0, indexer_weights,668            indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream,669            indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0);670 671    ggml_tensor * indexer_score = nullptr;672    if (cparams.fused_lid) {673        indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_lid.kq_mask);674        cb(indexer_score, "lid_score_masked", il);675        res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});676    } else {677        indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);678        cb(indexer_q, "lid_q", il);679        indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);680        cb(indexer_k, "lid_k", il);681 682        ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);683        cb(indexer_kq, "lid_kq", il);684 685        indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));686        cb(indexer_kq, "lid_kq", il);687 688        indexer_score = ggml_relu(ctx0, indexer_kq);689        indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);690        indexer_score = ggml_sum_rows(ctx0, indexer_score);691        indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));692        cb(indexer_score, "lid_score", il);693 694        indexer_score = ggml_add(ctx0, indexer_score, inp_lid.kq_mask);695        cb(indexer_score, "lid_score_masked", il);696    }697 698    const uint32_t n_top_k = indexer_score->ne[0] < hparams.indexer_top_k ? indexer_score->ne[0] : hparams.indexer_top_k;699    ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));700    cb(top_k, "lid_top_k", il);701 702    return top_k;703}704 705ggml_tensor * llama_model_deepseek4::graph::build_top_k_mask(706        ggml_tensor * kq_mask,707        ggml_tensor * top_k,708        const char * name,709        int il) const {710    GGML_ASSERT(kq_mask);711    GGML_ASSERT(top_k);712 713    ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY);714    kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3],715            kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0);716 717    ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1,718            top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0);719 720    ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]);721    zeros = ggml_fill(ctx0, zeros, 0.0f);722 723    ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d);724    kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k,725            kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3],726            kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0);727 728    kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask);729    cb(kq_mask_top_k, name, il);730 731    return kq_mask_top_k;732}733 734ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention(735        const llama_model & model,736        llm_graph_input_dsv4 * inp_dsv4,737        llm_graph_input_dsv4_raw * inp_attn,738        ggml_tensor * q,739        ggml_tensor * kv,740        ggml_tensor * qr,741        ggml_tensor * cur,742        ggml_tensor * inp_pos,743        ggml_tensor * sinks,744        float kq_scale,745        int il) const {746    const auto & inp_csa = inp_dsv4->get_csa();747    GGML_ASSERT(inp_csa.kq_mask);748 749    ggml_tensor * top_k = build_lid_top_k(model, inp_dsv4, qr, cur, inp_pos, il);750 751    ggml_tensor * k_rot = inp_attn->self_k_rot;752    if (k_rot) {753        q  = llama_mul_mat_hadamard(ctx0, q, k_rot);754        kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);755    }756 757    ggml_build_forward_expand(gf, q);758    ggml_build_forward_expand(gf, kv);759 760    const llama_kv_cache_dsv4_raw_context * mctx_raw = inp_attn->mctx;761 762    ggml_build_forward_expand(gf, mctx_raw->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il));763 764    ggml_tensor * raw_k = mctx_raw->get_k(ctx0, il);765    cb(raw_k, "csa_raw_k", il);766 767    ggml_tensor * csa_k = inp_dsv4->mctx->get_csa()->get_k(ctx0, il);768    const int64_t n_csa = inp_csa.kq_mask->ne[0];769    GGML_ASSERT(n_csa > 0);770    GGML_ASSERT(n_csa <= csa_k->ne[2]);771 772    csa_k = ggml_view_4d(ctx0, csa_k,773            csa_k->ne[0], csa_k->ne[1], n_csa, csa_k->ne[3],774            csa_k->nb[1], csa_k->nb[2], csa_k->nb[3], 0);775    cb(csa_k, "csa_comp_k", il);776 777    ggml_tensor * k_all = ggml_concat(ctx0, raw_k, csa_k, 2);778    cb(k_all, "csa_k_all", il);779 780    ggml_tensor * raw_mask = inp_attn->get_kq_mask();781    ggml_tensor * csa_mask = build_top_k_mask(inp_csa.kq_mask, top_k, "csa_top_k_mask", il);782 783    ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0);784    cb(kq_mask, "csa_lid_kq_mask", il);785 786    ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il);787    if (k_rot) {788        out = llama_mul_mat_hadamard(ctx0, out, k_rot);789    }790    cb(out, "attn_csa_lid", il);791 792    return out;793}794 795ggml_tensor * llama_model_deepseek4::graph::build_hca_attention(796        llm_graph_input_dsv4 * inp_dsv4,797        llm_graph_input_dsv4_raw * inp_attn,798        ggml_tensor * q,799        ggml_tensor * kv,800        ggml_tensor * sinks,801        float kq_scale,802        int il) const {803    const auto & inp_hca = inp_dsv4->get_hca();804    GGML_ASSERT(inp_hca.kq_mask);805 806    ggml_tensor * k_rot = inp_attn->self_k_rot;807    if (k_rot) {808        q  = llama_mul_mat_hadamard(ctx0, q, k_rot);809        kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);810    }811 812    ggml_build_forward_expand(gf, q);813    ggml_build_forward_expand(gf, kv);814 815    const llama_kv_cache_dsv4_raw_context * mctx_raw = inp_attn->mctx;816 817    ggml_build_forward_expand(gf, mctx_raw->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il));818 819    ggml_tensor * raw_k = mctx_raw->get_k(ctx0, il);820    cb(raw_k, "hca_raw_k", il);821 822    ggml_tensor * hca_k = inp_dsv4->mctx->get_hca()->get_k(ctx0, il);823    const int64_t n_hca = inp_hca.kq_mask->ne[0];824    GGML_ASSERT(n_hca > 0);825    GGML_ASSERT(n_hca <= hca_k->ne[2]);826 827    hca_k = ggml_view_4d(ctx0, hca_k,828            hca_k->ne[0], hca_k->ne[1], n_hca, hca_k->ne[3],829            hca_k->nb[1], hca_k->nb[2], hca_k->nb[3], 0);830    cb(hca_k, "hca_comp_k", il);831 832    ggml_tensor * k_all = ggml_concat(ctx0, raw_k, hca_k, 2);833    cb(k_all, "hca_k_all", il);834 835    ggml_tensor * raw_mask = inp_attn->get_kq_mask();836    ggml_tensor * hca_mask = inp_hca.kq_mask;837 838    ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0);839    cb(kq_mask, "hca_kq_mask", il);840 841    ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il);842    if (k_rot) {843        out = llama_mul_mat_hadamard(ctx0, out, k_rot);844    }845    cb(out, "attn_hca", il);846 847    return out;848}849 850ggml_tensor * llama_model_deepseek4::graph::build_raw_attention(851        llm_graph_input_dsv4_raw * inp_attn,852        ggml_tensor * q,853        ggml_tensor * kv,854        ggml_tensor * sinks,855        float kq_scale,856        int il) const {857    GGML_ASSERT(hparams.is_swa(il));858 859    ggml_tensor * k_rot = inp_attn->self_k_rot;860 861    if (k_rot) {862        q  = llama_mul_mat_hadamard(ctx0, q, k_rot);863        kv = llama_mul_mat_hadamard(ctx0, kv, k_rot);864    }865 866    ggml_build_forward_expand(gf, q);867    ggml_build_forward_expand(gf, kv);868 869    const llama_kv_cache_dsv4_raw_context * mctx_cur = inp_attn->mctx;870 871    ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il));872 873    ggml_tensor * kq_mask = inp_attn->get_kq_mask();874 875    ggml_tensor * k = mctx_cur->get_k(ctx0, il);876 877    ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, kq_scale, il);878    if (k_rot) {879        out = llama_mul_mat_hadamard(ctx0, out, k_rot);880    }881    cb(out, "attn_raw", il);882 883    return out;884}885 886ggml_tensor * llama_model_deepseek4::graph::build_attention(887        const llama_model & model,888        llm_graph_input_dsv4 * inp_dsv4,889        ggml_tensor * cur,890        ggml_tensor * inp_pos,891        int il) const {892    return build_attention_impl(model, inp_dsv4, nullptr, cur, inp_pos, il);893}894 895ggml_tensor * llama_model_deepseek4::graph::build_attention(896        const llama_model & model,897        llm_graph_input_attn_k_iswa * inp_mtp,898        ggml_tensor * cur,899        ggml_tensor * inp_pos,900        int il) const {901    return build_attention_impl(model, nullptr, inp_mtp, cur, inp_pos, il);902}903 904ggml_tensor * llama_model_deepseek4::graph::build_attention_impl(905        const llama_model & model,906        llm_graph_input_dsv4 * inp_dsv4,907        llm_graph_input_attn_k_iswa * inp_mtp,908        ggml_tensor * cur,909        ggml_tensor * inp_pos,910        int il) const {911    GGML_ASSERT((inp_dsv4 == nullptr) != (inp_mtp == nullptr));912 913    const auto & layer = model.layers[il];914    llm_graph_input_dsv4_raw * inp_attn = inp_dsv4 ? inp_dsv4->get_raw() : nullptr;915 916    const int64_t n_embd_head      = hparams.n_embd_head_k();917    const int64_t n_embd_head_rope = hparams.n_rot();918    const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;919    const int64_t n_groups         = hparams.dsv4_o_group_count;920    const int64_t n_heads_group    = n_head / n_groups;921    const int64_t o_lora_rank      = hparams.dsv4_o_lora_rank;922    const int64_t o_group_dim      = n_heads_group*n_embd_head;923    const int64_t nt               = cur->ne[1];924 925    GGML_ASSERT(n_embd_head == n_embd_head_v);926    GGML_ASSERT(n_head % n_groups == 0);927 928    const bool use_compress_rope = hparams.dsv4_compress_ratios[il] != 0;929    const float freq_base_l      = use_compress_rope ? hparams.dsv4_compress_rope_base : freq_base;930    const float freq_scale_l     = use_compress_rope ? freq_scale : 1.0f;931    const float ext_factor_l     = use_compress_rope ? ext_factor : 0.0f;932    const float attn_factor_l    = dsv4_rope_attn_factor(freq_scale_l, ext_factor_l);933    const float beta_fast_l      = use_compress_rope ? beta_fast : 0.0f;934    const float beta_slow_l      = use_compress_rope ? beta_slow : 0.0f;935    const int32_t n_ctx_orig_l   = use_compress_rope ? n_ctx_orig : 0;936 937    ggml_tensor * qr = build_lora_mm(layer.wq_a, cur);938    cb(qr, "qr", il);939 940    qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);941    cb(qr, "qr_norm", il);942 943    ggml_tensor * q = build_lora_mm(layer.wq_b, qr);944    q = ggml_reshape_3d(ctx0, q, n_embd_head, n_head, nt);945    q = ggml_rms_norm(ctx0, q, norm_rms_eps);946    cb(q, "q_norm", il);947 948    ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_nope, n_head, nt,949            ggml_row_size(q->type, n_embd_head),950            ggml_row_size(q->type, n_embd_head)*n_head,951            0);952    ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_embd_head_rope, n_head, nt,953            ggml_row_size(q->type, n_embd_head),954            ggml_row_size(q->type, n_embd_head)*n_head,955            ggml_row_size(q->type, n_embd_head_nope));956    q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l,957            freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);958    cb(q_pe, "q_pe", il);959    q = ggml_concat(ctx0, q_nope, q_pe, 0);960    cb(q, "q", il);961 962    ggml_tensor * kv = build_lora_mm(layer.wkv, cur);963    kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);964    kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, nt);965    cb(kv, "kv_norm", il);966 967    ggml_tensor * kv_nope = ggml_view_3d(ctx0, kv, n_embd_head_nope, 1, nt,968            ggml_row_size(kv->type, n_embd_head),969            ggml_row_size(kv->type, n_embd_head),970            0);971    ggml_tensor * kv_pe = ggml_view_3d(ctx0, kv, n_embd_head_rope, 1, nt,972            ggml_row_size(kv->type, n_embd_head),973            ggml_row_size(kv->type, n_embd_head),974            ggml_row_size(kv->type, n_embd_head_nope));975    kv_pe = ggml_rope_ext(ctx0, kv_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l,976            freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);977    cb(kv_pe, "kv_pe", il);978    kv = ggml_concat(ctx0, kv_nope, kv_pe, 0);979    cb(kv, "kv", il);980 981    const int64_t ratio = hparams.dsv4_compress_ratios[il];982    GGML_ASSERT(inp_dsv4 || ratio == 0);983 984    ggml_tensor * hca_state_kv    = nullptr;985    ggml_tensor * hca_state_score = nullptr;986    ggml_tensor * hca_source_kv   = nullptr;987    ggml_tensor * hca_source_score = nullptr;988    if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) {989        hca_state_kv = build_lora_mm(layer.attn_comp_wkv, cur);990        cb(hca_state_kv, "hca_state_kv", il);991 992        hca_state_score = build_lora_mm(layer.attn_comp_wgate, cur);993        cb(hca_state_score, "hca_state_score", il);994 995        ggml_tensor * ape = layer.attn_comp_ape;996 997        ggml_tensor * ape_rows = ggml_get_rows(ctx0, ape, inp_dsv4->get_hca().state_pos);998        hca_state_score = ggml_add(ctx0, hca_state_score, ape_rows);999        cb(hca_state_score, "hca_state_score_ape", il);1000 1001    }1002 1003    if (ratio == DSV4_CSA_RATIO && inp_dsv4->get_csa().state_pos) {1004        ggml_tensor * csa_state_kv = build_lora_mm(layer.attn_comp_wkv, cur);1005        cb(csa_state_kv, "csa_state_kv", il);1006 1007        ggml_tensor * csa_state_score = build_lora_mm(layer.attn_comp_wgate, cur);1008        cb(csa_state_score, "csa_state_score", il);1009 1010        ggml_tensor * csa_ape = layer.attn_comp_ape;1011 1012        ggml_tensor * csa_ape_rows = ggml_get_rows(ctx0, csa_ape, inp_dsv4->get_csa().state_pos);1013        csa_state_score = ggml_add(ctx0, csa_state_score, csa_ape_rows);1014        cb(csa_state_score, "csa_state_score_ape", il);1015 1016        GGML_ASSERT(inp_dsv4->get_csa().state_write_idxs);1017 1018        const auto * csa_state = inp_dsv4->mctx->get_csa_state();1019        const dsv4_state_tensors csa_restored = dsv4_build_state_restore(1020                ctx0, inp_dsv4->get_csa(), csa_state, il);1021        ggml_tensor * csa_base_kv = dsv4_view_2d(1022                ctx0, csa_restored.kv, csa_restored.kv->ne[0], csa_state->get_n_rows(), 0);1023        ggml_tensor * csa_base_score = dsv4_view_2d(1024                ctx0, csa_restored.score, csa_restored.score->ne[0], csa_state->get_n_rows(), 0);1025 1026        ggml_tensor * csa_source_kv = ggml_concat(ctx0, csa_base_kv, csa_state_kv, 1);1027        ggml_tensor * csa_source_score = ggml_concat(ctx0, csa_base_score, csa_state_score, 1);1028 1029        ggml_tensor * kv_comp_csa_state = build_overlap_compressed_kv_from_state(1030                csa_source_kv,1031                csa_source_score,1032                inp_dsv4->get_csa().state_read_idxs,1033                inp_dsv4->get_csa().state_write_pos,1034                layer.attn_comp_norm,1035                DSV4_CSA_RATIO,1036                n_embd_head,1037                "csa_state_compress",1038                il);1039 1040        if (inp_dsv4->get_csa().k_rot) {1041            kv_comp_csa_state = llama_mul_mat_hadamard(ctx0, kv_comp_csa_state, inp_dsv4->get_csa().k_rot);1042            cb(kv_comp_csa_state, "csa_state_compress_rot", il);1043        }1044 1045        ggml_build_forward_expand(gf, inp_dsv4->mctx->get_csa()->cpy_k(ctx0,1046                    kv_comp_csa_state, inp_dsv4->get_csa().state_write_idxs, il));1047 1048        ggml_tensor * csa_snapshot_source_kv = ggml_concat(ctx0,1049                csa_restored.kv, csa_state_kv, 1);1050        ggml_tensor * csa_snapshot_source_score = ggml_concat(ctx0,1051                csa_restored.score, csa_state_score, 1);1052 1053        const dsv4_state_tensors csa_snapshot = dsv4_build_state_snapshot(1054                ctx0, inp_dsv4->get_csa(), csa_state, csa_snapshot_source_kv, csa_snapshot_source_score, il);1055        if (csa_snapshot.kv != nullptr) {1056            ggml_build_forward_expand(gf, csa_snapshot.kv);1057        }1058        if (csa_snapshot.score != nullptr) {1059            ggml_build_forward_expand(gf, csa_snapshot.score);1060        }1061 1062        ggml_tensor * csa_persist_kv = ggml_get_rows(ctx0, csa_state_kv, inp_dsv4->get_csa().state_persist_src_idxs);1063        ggml_tensor * csa_persist_score = ggml_get_rows(ctx0, csa_state_score, inp_dsv4->get_csa().state_persist_src_idxs);1064 1065        csa_state_kv = inp_dsv4->mctx->get_csa_state()->cpy_kv(ctx0,1066                csa_persist_kv, inp_dsv4->get_csa().state_persist_dst_idxs, il);1067        csa_state_score = inp_dsv4->mctx->get_csa_state()->cpy_score(ctx0,1068                csa_persist_score, inp_dsv4->get_csa().state_persist_dst_idxs, il);1069 1070        ggml_build_forward_expand(gf, csa_state_kv);1071        ggml_build_forward_expand(gf, csa_state_score);1072 1073        ggml_tensor * lid_state_kv = build_lora_mm(layer.indexer_comp_wkv, cur);1074        cb(lid_state_kv, "lid_state_kv", il);1075 1076        ggml_tensor * lid_state_score = build_lora_mm(layer.indexer_comp_wgate, cur);1077        cb(lid_state_score, "lid_state_score", il);1078 1079        ggml_tensor * lid_ape = layer.indexer_comp_ape;1080 1081        ggml_tensor * lid_ape_rows = ggml_get_rows(ctx0, lid_ape, inp_dsv4->get_lid().state_pos);1082        lid_state_score = ggml_add(ctx0, lid_state_score, lid_ape_rows);1083        cb(lid_state_score, "lid_state_score_ape", il);1084 1085        GGML_ASSERT(inp_dsv4->get_lid().state_write_idxs);1086 1087        const auto * lid_state = inp_dsv4->mctx->get_lid_state();1088        const dsv4_state_tensors lid_restored = dsv4_build_state_restore(1089                ctx0, inp_dsv4->get_lid(), lid_state, il);1090        ggml_tensor * lid_base_kv = dsv4_view_2d(1091                ctx0, lid_restored.kv, lid_restored.kv->ne[0], lid_state->get_n_rows(), 0);1092        ggml_tensor * lid_base_score = dsv4_view_2d(1093                ctx0, lid_restored.score, lid_restored.score->ne[0], lid_state->get_n_rows(), 0);1094 1095        ggml_tensor * lid_source_kv = ggml_concat(ctx0, lid_base_kv, lid_state_kv, 1);1096        ggml_tensor * lid_source_score = ggml_concat(ctx0, lid_base_score, lid_state_score, 1);1097 1098        ggml_tensor * kv_comp_lid_state = build_overlap_compressed_kv_from_state(1099                lid_source_kv,1100                lid_source_score,1101                inp_dsv4->get_lid().state_read_idxs,1102                inp_dsv4->get_lid().state_write_pos,1103                layer.indexer_comp_norm,1104                DSV4_CSA_RATIO,1105                hparams.indexer_head_size,1106                "lid_state_compress",1107                il);1108 1109        if (inp_dsv4->get_lid().k_rot) {1110            kv_comp_lid_state = llama_mul_mat_hadamard(ctx0, kv_comp_lid_state, inp_dsv4->get_lid().k_rot);1111            cb(kv_comp_lid_state, "lid_state_compress_rot", il);1112        }1113 1114        ggml_build_forward_expand(gf, inp_dsv4->mctx->get_lid()->cpy_k(ctx0,1115                    kv_comp_lid_state, inp_dsv4->get_lid().state_write_idxs, il));1116 1117        ggml_tensor * lid_snapshot_source_kv = ggml_concat(ctx0,1118                lid_restored.kv, lid_state_kv, 1);1119        ggml_tensor * lid_snapshot_source_score = ggml_concat(ctx0,1120                lid_restored.score, lid_state_score, 1);1121 1122        const dsv4_state_tensors lid_snapshot = dsv4_build_state_snapshot(1123                ctx0, inp_dsv4->get_lid(), lid_state, lid_snapshot_source_kv, lid_snapshot_source_score, il);1124        if (lid_snapshot.kv != nullptr) {1125            ggml_build_forward_expand(gf, lid_snapshot.kv);1126        }1127        if (lid_snapshot.score != nullptr) {1128            ggml_build_forward_expand(gf, lid_snapshot.score);1129        }1130 1131        ggml_tensor * lid_persist_kv = ggml_get_rows(ctx0, lid_state_kv, inp_dsv4->get_lid().state_persist_src_idxs);1132        ggml_tensor * lid_persist_score = ggml_get_rows(ctx0, lid_state_score, inp_dsv4->get_lid().state_persist_src_idxs);1133 1134        lid_state_kv = inp_dsv4->mctx->get_lid_state()->cpy_kv(ctx0,1135                lid_persist_kv, inp_dsv4->get_lid().state_persist_dst_idxs, il);1136        lid_state_score = inp_dsv4->mctx->get_lid_state()->cpy_score(ctx0,1137                lid_persist_score, inp_dsv4->get_lid().state_persist_dst_idxs, il);1138 1139        ggml_build_forward_expand(gf, lid_state_kv);1140        ggml_build_forward_expand(gf, lid_state_score);1141    }1142 1143    const llama_dsv4_comp_state * hca_state = nullptr;1144    dsv4_state_tensors hca_restored = {};1145    if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_write_idxs) {1146        GGML_ASSERT(hca_state_kv);1147        GGML_ASSERT(hca_state_score);1148 1149        hca_state = inp_dsv4->mctx->get_hca_state();1150        hca_restored = dsv4_build_state_restore(ctx0, inp_dsv4->get_hca(), hca_state, il);1151        ggml_tensor * hca_base_kv = dsv4_view_2d(1152                ctx0, hca_restored.kv, hca_restored.kv->ne[0], hca_state->get_n_rows(), 0);1153        ggml_tensor * hca_base_score = dsv4_view_2d(1154                ctx0, hca_restored.score, hca_restored.score->ne[0], hca_state->get_n_rows(), 0);1155 1156        hca_source_kv = ggml_concat(ctx0, hca_base_kv, hca_state_kv, 1);1157        hca_source_score = ggml_concat(ctx0, hca_base_score, hca_state_score, 1);1158 1159        ggml_tensor * kv_comp_hca = build_hca_compressed_kv_from_state(1160                hca_source_kv,1161                hca_source_score,1162                inp_dsv4->get_hca().state_read_idxs,1163                inp_dsv4->get_hca().state_write_pos,1164                layer.attn_comp_norm,1165                n_embd_head,1166                "hca_state_compress",1167                il);1168 1169        if (inp_dsv4->get_hca().k_rot) {1170            kv_comp_hca = llama_mul_mat_hadamard(ctx0, kv_comp_hca, inp_dsv4->get_hca().k_rot);1171            cb(kv_comp_hca, "hca_state_compress_rot", il);1172        }1173 1174        ggml_build_forward_expand(gf, inp_dsv4->mctx->get_hca()->cpy_k(ctx0,1175                    kv_comp_hca, inp_dsv4->get_hca().state_write_idxs, il));1176    }1177 1178    if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) {1179        GGML_ASSERT(hca_state_kv);1180        GGML_ASSERT(hca_state_score);1181 1182        if (hca_state == nullptr) {1183            hca_state = inp_dsv4->mctx->get_hca_state();1184        }1185        if (hca_restored.kv == nullptr) {1186            hca_restored = dsv4_build_state_restore(ctx0, inp_dsv4->get_hca(), hca_state, il);1187        }1188        if (hca_source_kv == nullptr || hca_source_score == nullptr) {1189            ggml_tensor * hca_base_kv = dsv4_view_2d(1190                    ctx0, hca_restored.kv, hca_restored.kv->ne[0], hca_state->get_n_rows(), 0);1191            ggml_tensor * hca_base_score = dsv4_view_2d(1192                    ctx0, hca_restored.score, hca_restored.score->ne[0], hca_state->get_n_rows(), 0);1193 1194            hca_source_kv = ggml_concat(ctx0, hca_base_kv, hca_state_kv, 1);1195            hca_source_score = ggml_concat(ctx0, hca_base_score, hca_state_score, 1);1196        }1197 1198        ggml_tensor * hca_snapshot_source_kv = ggml_concat(ctx0,1199                hca_restored.kv, hca_state_kv, 1);1200        ggml_tensor * hca_snapshot_source_score = ggml_concat(ctx0,

Showing the first 1,200 of 1547 lines. Download the file for the rest.