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