KBaba7/llama.cpp
0
1#include "llama-impl.h"2 3#include "llama-chat.h"4#include "llama-mmap.h"5#include "llama-context.h"6#include "llama-vocab.h"7#include "llama-sampling.h"8#include "llama-kv-cache.h"9#include "llama-model-loader.h"10#include "llama-model.h"11 12#include "ggml.h"13#include "ggml-alloc.h"14#include "ggml-backend.h"15#include "ggml-cpp.h"16 17#include <algorithm>18#include <array>19#include <cassert>20#include <cfloat>21#include <cmath>22#include <cstddef>23#include <cstdint>24#include <cstdio>25#include <cstring>26#include <ctime>27#include <functional>28 29#if defined(_MSC_VER)30#pragma warning(disable: 4244 4267) // possible loss of data31#endif32 33// Returns 0 on success, -1 on error, and -2 on cancellation via llama_progress_callback34static int llama_model_load(const std::string & fname, std::vector<std::string> & splits, llama_model & model, llama_model_params & params) {35 // loading time will be recalculated after the first eval, so36 // we take page faults deferred by mmap() into consideration37 model.t_load_us = 0;38 time_meas tm(model.t_load_us);39 40 model.t_start_us = tm.t_start_us;41 42 try {43 llama_model_loader ml(fname, splits, params.use_mmap, params.check_tensors, params.kv_overrides);44 45 ml.print_info();46 47 model.hparams.vocab_only = params.vocab_only;48 49 try {50 model.load_arch(ml);51 } catch(const std::exception & e) {52 throw std::runtime_error("error loading model architecture: " + std::string(e.what()));53 }54 try {55 model.load_hparams(ml);56 } catch(const std::exception & e) {57 throw std::runtime_error("error loading model hyperparameters: " + std::string(e.what()));58 }59 try {60 model.load_vocab(ml);61 } catch(const std::exception & e) {62 throw std::runtime_error("error loading model vocabulary: " + std::string(e.what()));63 }64 65 model.load_stats(ml);66 model.print_info();67 68 if (params.vocab_only) {69 LLAMA_LOG_INFO("%s: vocab only - skipping tensors\n", __func__);70 return 0;71 }72 73 if (!model.load_tensors(ml)) {74 return -2;75 }76 } catch (const std::exception & err) {77 LLAMA_LOG_ERROR("%s: error loading model: %s\n", __func__, err.what());78 return -1;79 }80 81 return 0;82}83 84//85// llm_build86//87 88using llm_build_cb = std::function<void(struct ggml_tensor * cur, const char * name, int nl)>;89 90enum llm_ffn_op_type {91 LLM_FFN_SILU,92 LLM_FFN_GELU,93 LLM_FFN_RELU,94 LLM_FFN_RELU_SQR,95 LLM_FFN_SWIGLU,96};97 98enum llm_ffn_gate_type {99 LLM_FFN_SEQ,100 LLM_FFN_PAR, // ffn_gate is parallel to ffn_up101};102 103enum llm_norm_type {104 LLM_NORM,105 LLM_NORM_RMS,106 LLM_NORM_GROUP,107};108 109static struct ggml_tensor * llm_build_inp_embd(110 struct ggml_context * ctx,111 struct llama_context & lctx,112 const llama_hparams & hparams,113 const llama_ubatch & ubatch,114 struct ggml_tensor * tok_embd,115 const llm_build_cb & cb) {116 const int64_t n_embd = hparams.n_embd;117 118 struct ggml_tensor * inpL;119 120 if (ubatch.token) {121 lctx.inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ubatch.n_tokens);122 cb(lctx.inp_tokens, "inp_tokens", -1);123 ggml_set_input(lctx.inp_tokens);124 125 inpL = ggml_get_rows(ctx, tok_embd, lctx.inp_tokens);126 127 // apply lora for embedding tokens if needed128 for (auto & it : lctx.lora) {129 struct llama_adapter_lora_weight * lw = it.first->get_weight(tok_embd);130 if (lw == nullptr) {131 continue;132 }133 const float adapter_scale = it.second;134 const float scale = lw->get_scale(it.first->alpha, adapter_scale);135 struct ggml_tensor * inpL_delta = ggml_scale(ctx, ggml_mul_mat(136 ctx, lw->b, // non-transposed lora_b137 ggml_get_rows(ctx, lw->a, lctx.inp_tokens)138 ), scale);139 inpL = ggml_add(ctx, inpL, inpL_delta);140 }141 } else {142 lctx.inp_embd = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, ubatch.n_tokens);143 inpL = lctx.inp_embd;144 ggml_set_input(lctx.inp_embd);145 }146 147 // For Granite architecture148 if (hparams.f_embedding_scale != 0.0f) {149 inpL = ggml_scale(ctx, inpL, hparams.f_embedding_scale);150 }151 152 cb(inpL, "inp_embd", -1);153 154 return inpL;155}156 157static void llm_build_kv_store(158 struct ggml_context * ctx,159 const llama_hparams & hparams,160 const llama_cparams & cparams,161 const llama_kv_cache & kv,162 struct ggml_cgraph * graph,163 struct ggml_tensor * k_cur,164 struct ggml_tensor * v_cur,165 int32_t n_tokens,166 int32_t kv_head,167 const llm_build_cb & cb,168 int64_t il) {169 const int64_t n_ctx = cparams.n_ctx;170 171 const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);172 const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);173 174 GGML_ASSERT(kv.size == n_ctx);175 176 struct ggml_tensor * k_cache_view = ggml_view_1d(ctx, kv.k_l[il], n_tokens*n_embd_k_gqa, ggml_row_size(kv.k_l[il]->type, n_embd_k_gqa)*kv_head);177 cb(k_cache_view, "k_cache_view", il);178 179 // note: storing RoPE-ed version of K in the KV cache180 ggml_build_forward_expand(graph, ggml_cpy(ctx, k_cur, k_cache_view));181 182 assert(v_cur->ne[0] == n_embd_v_gqa && v_cur->ne[1] == n_tokens);183 184 struct ggml_tensor * v_cache_view = nullptr;185 186 if (cparams.flash_attn) {187 v_cache_view = ggml_view_1d(ctx, kv.v_l[il], n_tokens*n_embd_v_gqa, ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa)*kv_head);188 } else {189 // note: the V cache is transposed when not using flash attention190 v_cache_view = ggml_view_2d(ctx, kv.v_l[il], n_tokens, n_embd_v_gqa,191 ( n_ctx)*ggml_element_size(kv.v_l[il]),192 (kv_head)*ggml_element_size(kv.v_l[il]));193 194 v_cur = ggml_transpose(ctx, v_cur);195 }196 cb(v_cache_view, "v_cache_view", il);197 198 ggml_build_forward_expand(graph, ggml_cpy(ctx, v_cur, v_cache_view));199}200 201// do mat_mul, while optionally apply lora202static struct ggml_tensor * llm_build_lora_mm(203 struct llama_context & lctx,204 struct ggml_context * ctx0,205 struct ggml_tensor * w,206 struct ggml_tensor * cur) {207 struct ggml_tensor * res = ggml_mul_mat(ctx0, w, cur);208 for (auto & it : lctx.lora) {209 struct llama_adapter_lora_weight * lw = it.first->get_weight(w);210 if (lw == nullptr) {211 continue;212 }213 const float adapter_scale = it.second;214 const float scale = lw->get_scale(it.first->alpha, adapter_scale);215 struct ggml_tensor * ab_cur = ggml_mul_mat(216 ctx0, lw->b,217 ggml_mul_mat(ctx0, lw->a, cur)218 );219 ab_cur = ggml_scale(ctx0, ab_cur, scale);220 res = ggml_add(ctx0, res, ab_cur);221 }222 return res;223}224 225// do mat_mul_id, while optionally apply lora226static struct ggml_tensor * llm_build_lora_mm_id(227 struct llama_context & lctx,228 struct ggml_context * ctx0,229 struct ggml_tensor * w, // struct ggml_tensor * as230 struct ggml_tensor * cur, // struct ggml_tensor * b231 struct ggml_tensor * ids) {232 struct ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids);233 for (auto & it : lctx.lora) {234 struct llama_adapter_lora_weight * lw = it.first->get_weight(w);235 if (lw == nullptr) {236 continue;237 }238 const float alpha = it.first->alpha;239 const float rank = (float) lw->b->ne[0];240 const float scale = alpha ? it.second * alpha / rank : it.second;241 struct ggml_tensor * ab_cur = ggml_mul_mat_id(242 ctx0, lw->b,243 ggml_mul_mat_id(ctx0, lw->a, cur, ids),244 ids245 );246 ab_cur = ggml_scale(ctx0, ab_cur, scale);247 res = ggml_add(ctx0, res, ab_cur);248 }249 return res;250}251 252static struct ggml_tensor * llm_build_norm(253 struct ggml_context * ctx,254 struct ggml_tensor * cur,255 const llama_hparams & hparams,256 struct ggml_tensor * mw,257 struct ggml_tensor * mb,258 llm_norm_type type,259 const llm_build_cb & cb,260 int il) {261 switch (type) {262 case LLM_NORM: cur = ggml_norm (ctx, cur, hparams.f_norm_eps); break;263 case LLM_NORM_RMS: cur = ggml_rms_norm (ctx, cur, hparams.f_norm_rms_eps); break;264 case LLM_NORM_GROUP:265 {266 cur = ggml_reshape_3d(ctx, cur, cur->ne[0], 1, cur->ne[1]);267 cur = ggml_group_norm(ctx, cur, hparams.n_norm_groups, hparams.f_norm_group_eps);268 cur = ggml_reshape_2d(ctx, cur, cur->ne[0], cur->ne[2]);269 } break;270 }271 272 if (mw || mb) {273 cb(cur, "norm", il);274 }275 276 if (mw) {277 cur = ggml_mul(ctx, cur, mw);278 if (mb) {279 cb(cur, "norm_w", il);280 }281 }282 283 if (mb) {284 cur = ggml_add(ctx, cur, mb);285 }286 287 return cur;288}289 290static struct ggml_tensor * llm_build_ffn(291 struct ggml_context * ctx,292 struct llama_context & lctx,293 struct ggml_tensor * cur,294 struct ggml_tensor * up,295 struct ggml_tensor * up_b,296 struct ggml_tensor * up_s,297 struct ggml_tensor * gate,298 struct ggml_tensor * gate_b,299 struct ggml_tensor * gate_s,300 struct ggml_tensor * down,301 struct ggml_tensor * down_b,302 struct ggml_tensor * down_s,303 struct ggml_tensor * act_scales,304 llm_ffn_op_type type_op,305 llm_ffn_gate_type type_gate,306 const llm_build_cb & cb,307 int il) {308 struct ggml_tensor * tmp = up ? llm_build_lora_mm(lctx, ctx, up, cur) : cur;309 cb(tmp, "ffn_up", il);310 311 if (up_b) {312 tmp = ggml_add(ctx, tmp, up_b);313 cb(tmp, "ffn_up_b", il);314 }315 316 if (up_s) {317 tmp = ggml_mul(ctx, tmp, up_s);318 cb(tmp, "ffn_up_s", il);319 }320 321 if (gate) {322 switch (type_gate) {323 case LLM_FFN_SEQ:324 {325 cur = llm_build_lora_mm(lctx, ctx, gate, tmp);326 cb(cur, "ffn_gate", il);327 } break;328 case LLM_FFN_PAR:329 {330 cur = llm_build_lora_mm(lctx, ctx, gate, cur);331 cb(cur, "ffn_gate", il);332 } break;333 }334 335 if (gate_b) {336 cur = ggml_add(ctx, cur, gate_b);337 cb(cur, "ffn_gate_b", il);338 }339 340 if (gate_s) {341 cur = ggml_mul(ctx, cur, gate_s);342 cb(cur, "ffn_gate_s", il);343 }344 345 } else {346 cur = tmp;347 }348 349 switch (type_op) {350 case LLM_FFN_SILU:351 {352 cur = ggml_silu(ctx, cur);353 cb(cur, "ffn_silu", il);354 } break;355 case LLM_FFN_GELU:356 {357 cur = ggml_gelu(ctx, cur);358 cb(cur, "ffn_gelu", il);359 if (act_scales != NULL) {360 cur = ggml_div(ctx, cur, act_scales);361 cb(cur, "ffn_act", il);362 }363 } break;364 case LLM_FFN_RELU:365 {366 cur = ggml_relu(ctx, cur);367 cb(cur, "ffn_relu", il);368 } break;369 case LLM_FFN_RELU_SQR:370 {371 cur = ggml_relu(ctx, cur);372 cb(cur, "ffn_relu", il);373 374 cur = ggml_sqr(ctx, cur);375 cb(cur, "ffn_sqr(relu)", il);376 } break;377 case LLM_FFN_SWIGLU:378 {379 // Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf380 int64_t split_point = cur->ne[0] / 2;381 struct ggml_tensor * x0 = ggml_cont(ctx, ggml_view_2d(ctx, cur, split_point, cur->ne[1], cur->nb[1], 0));382 struct ggml_tensor * x1 = ggml_cont(ctx, ggml_view_2d(ctx, cur, split_point, cur->ne[1], cur->nb[1], split_point * ggml_element_size(cur)));383 384 x0 = ggml_silu(ctx, x0);385 cb(cur, "ffn_silu", il);386 387 cur = ggml_mul(ctx, x0, x1);388 cb(cur, "ffn_mul", il);389 } break;390 }391 392 if (type_gate == LLM_FFN_PAR) {393 cur = ggml_mul(ctx, cur, tmp);394 cb(cur, "ffn_gate_par", il);395 }396 397 if (down) {398 cur = llm_build_lora_mm(lctx, ctx, down, cur);399 }400 401 if (down_b) {402 cb(cur, "ffn_down", il);403 }404 405 if (down_b) {406 cur = ggml_add(ctx, cur, down_b);407 }408 409 if (down_s) {410 cur = ggml_mul(ctx, cur, down_s);411 cb(cur, "ffn_down_s", il);412 }413 414 return cur;415}416 417static struct ggml_tensor * llm_build_moe_ffn(418 struct ggml_context * ctx,419 struct llama_context & lctx,420 struct ggml_tensor * cur,421 struct ggml_tensor * gate_inp,422 struct ggml_tensor * up_exps,423 struct ggml_tensor * gate_exps,424 struct ggml_tensor * down_exps,425 struct ggml_tensor * exp_probs_b,426 int64_t n_expert,427 int64_t n_expert_used,428 llm_ffn_op_type type_op,429 bool norm_w,430 bool scale_w,431 float w_scale,432llama_expert_gating_func_type gating_op,433 const llm_build_cb & cb,434 int il) {435 int64_t n_embd = cur->ne[0];436 int64_t n_tokens = cur->ne[1];437 438 ggml_tensor * logits = llm_build_lora_mm(lctx, ctx, gate_inp, cur); // [n_expert, n_tokens]439 cb(logits, "ffn_moe_logits", il);440 441 ggml_tensor * probs = nullptr;442 switch (gating_op) {443 case LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX:444 {445 probs = ggml_soft_max(ctx, logits); // [n_expert, n_tokens]446 } break;447 case LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID:448 {449 probs = ggml_sigmoid(ctx, logits); // [n_expert, n_tokens]450 } break;451 default:452 GGML_ABORT("fatal error");453 }454 cb(probs, "ffn_moe_probs", il);455 456 // add experts selection bias - introduced in DeepSeek V3457 // leave probs unbiased as it's later used to get expert weights458 ggml_tensor * selection_probs = probs;459 if (exp_probs_b != nullptr) {460 selection_probs = ggml_add(ctx, probs, exp_probs_b);461 cb(selection_probs, "ffn_moe_probs_biased", il);462 }463 464 // select experts465 ggml_tensor * selected_experts = ggml_top_k(ctx, selection_probs, n_expert_used); // [n_expert_used, n_tokens]466 cb(selected_experts->src[0], "ffn_moe_argsort", il);467 cb(selected_experts, "ffn_moe_topk", il);468 469 ggml_tensor * weights = ggml_get_rows(ctx,470 ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens), selected_experts); // [1, n_expert_used, n_tokens]471 cb(weights, "ffn_moe_weights", il);472 473 if (norm_w) {474 weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens);475 476 ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); // [1, n_tokens]477 cb(weights_sum, "ffn_moe_weights_sum", il);478 479 weights = ggml_div(ctx, weights, weights_sum); // [n_expert_used, n_tokens]480 cb(weights, "ffn_moe_weights_norm", il);481 482 weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens);483 }484 if (scale_w) {485 weights = ggml_scale(ctx, weights, w_scale);486 cb(weights, "ffn_moe_weights_scaled", il);487 }488 489 cur = ggml_reshape_3d(ctx, cur, n_embd, 1, n_tokens);490 ggml_tensor * up = llm_build_lora_mm_id(lctx, ctx, up_exps, cur, selected_experts); // [n_ff, n_expert_used, n_tokens]491 cb(up, "ffn_moe_up", il);492 493 ggml_tensor * gate = llm_build_lora_mm_id(lctx, ctx, gate_exps, cur, selected_experts); // [n_ff, n_expert_used, n_tokens]494 cb(gate, "ffn_moe_gate", il);495 496 switch (type_op) {497 case LLM_FFN_SILU:498 {499 gate = ggml_silu(ctx, gate);500 cb(gate, "ffn_moe_silu", il);501 } break;502 case LLM_FFN_GELU:503 {504 gate = ggml_gelu(ctx, gate);505 cb(gate, "ffn_moe_gelu", il);506 } break;507 default:508 GGML_ABORT("fatal error");509 }510 511 ggml_tensor * par = ggml_mul(ctx, up, gate); // [n_ff, n_expert_used, n_tokens]512 cb(par, "ffn_moe_gate_par", il);513 514 ggml_tensor * experts = llm_build_lora_mm_id(lctx, ctx, down_exps, par, selected_experts); // [n_embd, n_expert_used, n_tokens]515 cb(experts, "ffn_moe_down", il);516 517 experts = ggml_mul(ctx, experts, weights);518 519 // aggregate experts520 ggml_tensor * moe_out = nullptr;521 for (int i = 0; i < n_expert_used; ++i) {522 ggml_tensor * cur_expert = ggml_view_2d(ctx, experts, n_embd, n_tokens,523 experts->nb[2], i*experts->nb[1]);524 525 if (i == 0) {526 moe_out = cur_expert;527 } else {528 moe_out = ggml_add(ctx, moe_out, cur_expert);529 }530 }531 532 if (n_expert_used == 1) {533 // avoid returning a non-contiguous tensor534 moe_out = ggml_cont(ctx, moe_out);535 }536 537 return moe_out;538}539 540static struct ggml_tensor * llm_build_kqv(541 struct ggml_context * ctx,542 struct llama_context & lctx,543 const llama_kv_cache & kv,544 struct ggml_cgraph * graph,545 struct ggml_tensor * wo,546 struct ggml_tensor * wo_b,547 struct ggml_tensor * q_cur,548 struct ggml_tensor * kq_mask,549 int32_t n_tokens,550 int32_t n_kv,551 float kq_scale,552 const llm_build_cb & cb,553 int il) {554 const llama_model & model = lctx.model;555 const llama_hparams & hparams = lctx.model.hparams;556 const llama_cparams & cparams = lctx.cparams;557 558 const int64_t n_ctx = cparams.n_ctx;559 const int64_t n_head = hparams.n_head(il);560 const int64_t n_head_kv = hparams.n_head_kv(il);561 const int64_t n_embd_head_k = hparams.n_embd_head_k;562 const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);563 const int64_t n_embd_head_v = hparams.n_embd_head_v;564 const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(il);565 566 struct ggml_tensor * q = ggml_permute(ctx, q_cur, 0, 2, 1, 3);567 cb(q, "q", il);568 569 struct ggml_tensor * k =570 ggml_view_3d(ctx, kv.k_l[il],571 n_embd_head_k, n_kv, n_head_kv,572 ggml_row_size(kv.k_l[il]->type, n_embd_k_gqa),573 ggml_row_size(kv.k_l[il]->type, n_embd_head_k),574 0);575 cb(k, "k", il);576 577 struct ggml_tensor * cur;578 579 if (cparams.flash_attn) {580 GGML_UNUSED(model);581 GGML_UNUSED(n_ctx);582 583 // split cached v into n_head heads (not transposed)584 struct ggml_tensor * v =585 ggml_view_3d(ctx, kv.v_l[il],586 n_embd_head_v, n_kv, n_head_kv,587 ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa),588 ggml_row_size(kv.v_l[il]->type, n_embd_head_v),589 0);590 cb(v, "v", il);591 592 cur = ggml_flash_attn_ext(ctx, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,593 hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);594 595 ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32);596 597 cur = ggml_reshape_2d(ctx, cur, n_embd_head_v*n_head, n_tokens);598 } else {599 struct ggml_tensor * kq = ggml_mul_mat(ctx, k, q);600 cb(kq, "kq", il);601 602 // note: this op tends to require high floating point range603 // while for some models F16 is enough, for others it is not, so we default to F32 here604 ggml_mul_mat_set_prec(kq, GGML_PREC_F32);605 606 if (model.arch == LLM_ARCH_GROK) {607 // need to do the following:608 // multiply by attn_output_multiplyer of 0.08838834764831845609 // and then :610 // kq = 30 * tanh(kq / 30)611 // before the softmax below612 613 kq = ggml_tanh(ctx, ggml_scale(ctx, kq, 0.08838834764831845f/30.0f));614 kq = ggml_scale(ctx, kq, 30);615 }616 617 if (hparams.attn_soft_cap) {618 kq = ggml_scale(ctx, kq, 1.0f / hparams.f_attn_logit_softcapping);619 kq = ggml_tanh(ctx, kq);620 kq = ggml_scale(ctx, kq, hparams.f_attn_logit_softcapping);621 }622 623 kq = ggml_soft_max_ext(ctx, kq, kq_mask, kq_scale, hparams.f_max_alibi_bias);624 cb(kq, "kq_soft_max_ext", il);625 626 GGML_ASSERT(kv.size == n_ctx);627 628 // split cached v into n_head heads629 struct ggml_tensor * v =630 ggml_view_3d(ctx, kv.v_l[il],631 n_kv, n_embd_head_v, n_head_kv,632 ggml_element_size(kv.v_l[il])*n_ctx,633 ggml_element_size(kv.v_l[il])*n_ctx*n_embd_head_v,634 0);635 cb(v, "v", il);636 637 struct ggml_tensor * kqv = ggml_mul_mat(ctx, v, kq);638 cb(kqv, "kqv", il);639 640 struct ggml_tensor * kqv_merged = ggml_permute(ctx, kqv, 0, 2, 1, 3);641 cb(kqv_merged, "kqv_merged", il);642 643 cur = ggml_cont_2d(ctx, kqv_merged, n_embd_head_v*n_head, n_tokens);644 cb(cur, "kqv_merged_cont", il);645 }646 647 ggml_build_forward_expand(graph, cur);648 649 if (wo) {650 cur = llm_build_lora_mm(lctx, ctx, wo, cur);651 }652 653 if (wo_b) {654 cb(cur, "kqv_wo", il);655 }656 657 if (wo_b) {658 cur = ggml_add(ctx, cur, wo_b);659 }660 661 return cur;662}663 664static struct ggml_tensor * llm_build_kv(665 struct ggml_context * ctx,666 struct llama_context & lctx,667 const llama_kv_cache & kv,668 struct ggml_cgraph * graph,669 struct ggml_tensor * wo,670 struct ggml_tensor * wo_b,671 struct ggml_tensor * k_cur,672 struct ggml_tensor * v_cur,673 struct ggml_tensor * q_cur,674 struct ggml_tensor * kq_mask,675 int32_t n_tokens,676 int32_t kv_head,677 int32_t n_kv,678 float kq_scale,679 const llm_build_cb & cb,680 int il) {681 const llama_hparams & hparams = lctx.model.hparams;682 const llama_cparams & cparams = lctx.cparams;683 684 // these nodes are added to the graph together so that they are not reordered685 // by doing so, the number of splits in the graph is reduced686 ggml_build_forward_expand(graph, q_cur);687 ggml_build_forward_expand(graph, k_cur);688 ggml_build_forward_expand(graph, v_cur);689 690 llm_build_kv_store(ctx, hparams, cparams, kv, graph, k_cur, v_cur, n_tokens, kv_head, cb, il);691 692 struct ggml_tensor * cur;693 694 cur = llm_build_kqv(ctx, lctx, kv, graph, wo, wo_b, q_cur, kq_mask, n_tokens, n_kv, kq_scale, cb, il);695 cb(cur, "kqv_out", il);696 697 return cur;698}699 700static struct ggml_tensor * llm_build_copy_mask_state(701 struct ggml_context * ctx,702 struct ggml_cgraph * graph,703 struct ggml_tensor * s,704 struct ggml_tensor * state_copy,705 struct ggml_tensor * state_mask,706 int32_t n_state,707 int32_t kv_size,708 int32_t kv_head,709 int32_t n_kv,710 int32_t n_seqs) {711 struct ggml_tensor * states = ggml_reshape_2d(ctx, s, n_state, kv_size);712 713 // copy states714 // NOTE: assuming the copy destinations are ALL contained between kv_head and kv_head + n_kv715 // this shrinks the tensors's ne[1] to n_kv716 states = ggml_get_rows(ctx, states, state_copy);717 718 // clear states of sequences which are starting at the beginning of this batch719 // FIXME: zero-out NANs?720 states = ggml_mul(ctx, states, state_mask);721 722 // copy states which won't be changed further (between n_seqs and n_kv)723 ggml_build_forward_expand(graph,724 ggml_cpy(ctx,725 ggml_view_1d(ctx, states, n_state*(n_kv - n_seqs), n_seqs*n_state*ggml_element_size(states)),726 ggml_view_1d(ctx, s, n_state*(n_kv - n_seqs), (kv_head + n_seqs)*n_state*ggml_element_size(s))));727 728 // the part of the states that will be used and modified729 return ggml_view_2d(ctx, states, n_state, n_seqs, states->nb[1], 0);730}731 732// TODO: split733static struct ggml_tensor * llm_build_mamba(734 struct ggml_context * ctx,735 struct llama_context & lctx,736 const llama_ubatch & ubatch,737 struct ggml_cgraph * graph,738 struct ggml_tensor * cur,739 struct ggml_tensor * state_copy,740 struct ggml_tensor * state_mask,741 int32_t kv_head,742 int32_t n_kv,743 const llm_build_cb & cb,744 int il) {745 const llama_model & model = lctx.model;746 const llama_hparams & hparams = model.hparams;747 const llama_kv_cache & kv = lctx.kv_self;748 const int64_t d_conv = hparams.ssm_d_conv;749 const int64_t d_inner = hparams.ssm_d_inner;750 const int64_t d_state = hparams.ssm_d_state;751 const int64_t dt_rank = hparams.ssm_dt_rank;752 const int64_t n_seqs = ubatch.n_seqs;753 // Some variants of Mamba arch (e.g. FalconMamba do apply layer norm on B and Dt layers)754 const bool ssm_dt_b_c_rms = hparams.ssm_dt_b_c_rms;755 // Use the same RMS norm as the final layer norm756 const float norm_rms_eps = hparams.f_norm_rms_eps;757 758 const int64_t n_seq_tokens = ubatch.n_seq_tokens;759 760 GGML_ASSERT(n_seqs != 0);761 GGML_ASSERT(ubatch.equal_seqs);762 GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);763 764 struct ggml_tensor * conv_states_all = kv.k_l[il];765 struct ggml_tensor * ssm_states_all = kv.v_l[il];766 767 // (ab)using the KV cache to store the states768 struct ggml_tensor * conv = llm_build_copy_mask_state(ctx,769 graph, conv_states_all, state_copy, state_mask,770 hparams.n_embd_k_s(), kv.size, kv_head, n_kv, n_seqs);771 conv = ggml_reshape_3d(ctx, conv, d_conv - 1, d_inner, n_seqs);772 struct ggml_tensor * ssm = llm_build_copy_mask_state(ctx,773 graph, ssm_states_all, state_copy, state_mask,774 hparams.n_embd_v_s(), kv.size, kv_head, n_kv, n_seqs);775 ssm = ggml_reshape_3d(ctx, ssm, d_state, d_inner, n_seqs);776 777 // {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs}778 cur = ggml_reshape_3d(ctx, cur, cur->ne[0], n_seq_tokens, n_seqs);779 780 // {n_embd, 2*d_inner} @ {n_embd, n_seq_tokens, n_seqs} => {2*d_inner, n_seq_tokens, n_seqs}781 struct ggml_tensor * xz = llm_build_lora_mm(lctx, ctx, model.layers[il].ssm_in, cur);782 // split the above in two783 // => {d_inner, n_seq_tokens, n_seqs}784 struct ggml_tensor * x = ggml_view_3d(ctx, xz, d_inner, xz->ne[1], xz->ne[2], xz->nb[1], xz->nb[2], 0);785 struct ggml_tensor * z = ggml_view_3d(ctx, xz, d_inner, xz->ne[1], xz->ne[2], xz->nb[1], xz->nb[2], d_inner*ggml_element_size(xz));786 787 // conv788 {789 // => {d_conv - 1 + n_seq_tokens, d_inner, n_seqs}790 struct ggml_tensor * conv_x = ggml_concat(ctx, conv, ggml_transpose(ctx, x), 0);791 792 // copy last (d_conv - 1) columns back into the state cache793 struct ggml_tensor * last_conv = ggml_view_3d(ctx, conv_x, d_conv - 1, d_inner, n_seqs, conv_x->nb[1], conv_x->nb[2], n_seq_tokens*(conv_x->nb[0]));794 795 ggml_build_forward_expand(graph,796 ggml_cpy(ctx, last_conv,797 ggml_view_1d(ctx, conv_states_all,798 (d_conv - 1)*(d_inner)*(n_seqs),799 kv_head*(d_conv - 1)*(d_inner)*ggml_element_size(conv_states_all))));800 801 // 1D convolution802 // The equivalent is to make a self-overlapping view of conv_x803 // over d_conv columns at each stride in the 3rd dimension,804 // then element-wise multiply that with the conv1d weight,805 // then sum the elements of each row,806 // (the last two steps are a dot product over rows (also doable with mul_mat))807 // then permute away the ne[0] dimension,808 // and then you're left with the resulting x tensor.809 // For simultaneous sequences, all sequences need to have the same length.810 x = ggml_ssm_conv(ctx, conv_x, model.layers[il].ssm_conv1d);811 812 // bias813 x = ggml_add(ctx, x, model.layers[il].ssm_conv1d_b);814 815 x = ggml_silu(ctx, x);816 }817 818 // ssm819 {820 // {d_inner, dt_rank + 2*d_state} @ {d_inner, n_seq_tokens, n_seqs} => {dt_rank + 2*d_state, n_seq_tokens, n_seqs}821 struct ggml_tensor * x_db = llm_build_lora_mm(lctx, ctx, model.layers[il].ssm_x, x);822 // split823 struct ggml_tensor * dt = ggml_view_3d(ctx, x_db, dt_rank, n_seq_tokens, n_seqs, x_db->nb[1], x_db->nb[2], 0);824 struct ggml_tensor * B = ggml_view_3d(ctx, x_db, d_state, n_seq_tokens, n_seqs, x_db->nb[1], x_db->nb[2], ggml_element_size(x_db)*dt_rank);825 struct ggml_tensor * C = ggml_view_3d(ctx, x_db, d_state, n_seq_tokens, n_seqs, x_db->nb[1], x_db->nb[2], ggml_element_size(x_db)*(dt_rank+d_state));826 827 // Some Mamba variants (e.g. FalconMamba) apply RMS norm in B, C & Dt layers828 if (ssm_dt_b_c_rms) {829 dt = ggml_rms_norm(ctx, dt, norm_rms_eps);830 B = ggml_rms_norm(ctx, B, norm_rms_eps);831 C = ggml_rms_norm(ctx, C, norm_rms_eps);832 }833 834 // {dt_rank, d_inner} @ {dt_rank, n_seq_tokens, n_seqs} => {d_inner, n_seq_tokens, n_seqs}835 dt = llm_build_lora_mm(lctx, ctx, model.layers[il].ssm_dt, dt);836 dt = ggml_add(ctx, dt, model.layers[il].ssm_dt_b);837 838 // Custom operator to optimize the parallel associative scan839 // as described in the Annex D of the Mamba paper.840 // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}841 struct ggml_tensor * y_ssm = ggml_ssm_scan(ctx, ssm, x, dt, model.layers[il].ssm_a, B, C);842 843 // store last states844 ggml_build_forward_expand(graph,845 ggml_cpy(ctx,846 ggml_view_1d(ctx, y_ssm, d_state*d_inner*n_seqs, x->nb[3]),847 ggml_view_1d(ctx, ssm_states_all, d_state*d_inner*n_seqs, kv_head*d_state*d_inner*ggml_element_size(ssm_states_all))));848 849 struct ggml_tensor * y = ggml_view_3d(ctx, y_ssm, d_inner, n_seq_tokens, n_seqs, x->nb[1], x->nb[2], 0);850 851 // TODO: skip computing output earlier for unused tokens852 853 // {d_inner, n_seq_tokens, n_seqs} * {d_inner} => {d_inner, n_seq_tokens, n_seqs}854 y = ggml_add(ctx, y, ggml_mul(ctx, x, model.layers[il].ssm_d));855 y = ggml_mul(ctx, y, ggml_silu(ctx, ggml_cont(ctx, z)));856 857 // {d_inner, n_embd} @ {d_inner, n_seq_tokens, n_seqs} => {n_embd, n_seq_tokens, n_seqs}858 cur = llm_build_lora_mm(lctx, ctx, model.layers[il].ssm_out, y);859 }860 861 // {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens}862 cur = ggml_reshape_2d(ctx, cur, cur->ne[0], n_seq_tokens * n_seqs);863 cb(cur, "mamba_out", il);864 865 return cur;866}867 868static struct ggml_tensor * llm_build_rwkv6_time_mix(869 struct llama_context & lctx,870 struct ggml_context * ctx,871 const struct llama_layer * layer,872 struct ggml_tensor * cur,873 struct ggml_tensor * x_prev,874 struct ggml_tensor ** wkv_state,875 size_t wkv_head_size,876 size_t head_count_kv) {877 size_t n_embd = cur->ne[0];878 size_t n_seq_tokens = cur->ne[1];879 size_t n_seqs = cur->ne[2];880 881 size_t head_size = wkv_head_size;882 size_t head_count = n_embd / head_size;883 884 size_t n_tokens = n_seqs * n_seq_tokens;885 886 bool is_qrwkv = layer->time_mix_first == nullptr;887 888 struct ggml_tensor * sx = ggml_sub(ctx, x_prev, cur);889 890 sx = ggml_reshape_2d(ctx, sx, n_embd, n_tokens);891 cur = ggml_reshape_2d(ctx, cur, n_embd, n_tokens);892 893 struct ggml_tensor * xxx = ggml_add(ctx, ggml_mul(ctx, sx, layer->time_mix_lerp_x), cur);894 895 xxx = ggml_reshape_4d(896 ctx,897 ggml_tanh(898 ctx,899 ggml_mul_mat(ctx, layer->time_mix_w1, xxx)900 ),901 layer->time_mix_w1->ne[1] / 5, 1, 5, n_tokens902 );903 904 xxx = ggml_cont(ctx, ggml_permute(ctx, xxx, 0, 1, 3, 2));905 906 xxx = ggml_mul_mat(907 ctx,908 ggml_reshape_4d(909 ctx,910 layer->time_mix_w2,911 layer->time_mix_w2->ne[0], layer->time_mix_w2->ne[1], 1, 5912 ),913 xxx914 );915 916 struct ggml_tensor *xw, *xk, *xv, *xr, *xg;917 if (layer->time_mix_lerp_fused) {918 // fusing these weights makes some performance improvement919 sx = ggml_reshape_3d(ctx, sx, n_embd, 1, n_tokens);920 cur = ggml_reshape_3d(ctx, cur, n_embd, 1, n_tokens);921 xxx = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xxx, layer->time_mix_lerp_fused), sx), cur);922 xw = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], 0);923 xk = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * sizeof(float));924 xv = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 2 * sizeof(float));925 xr = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 3 * sizeof(float));926 xg = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 4 * sizeof(float));927 } else {928 // for backward compatibility929 xw = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], 0);930 xk = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * sizeof(float));931 xv = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 2 * sizeof(float));932 xr = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 3 * sizeof(float));933 xg = ggml_view_2d(ctx, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 4 * sizeof(float));934 935 xw = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xw, layer->time_mix_lerp_w), sx), cur);936 xk = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xk, layer->time_mix_lerp_k), sx), cur);937 xv = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xv, layer->time_mix_lerp_v), sx), cur);938 xr = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xr, layer->time_mix_lerp_r), sx), cur);939 xg = ggml_add(ctx, ggml_mul(ctx, ggml_add(ctx, xg, layer->time_mix_lerp_g), sx), cur);940 }941 942 struct ggml_tensor * r = llm_build_lora_mm(lctx, ctx, layer->time_mix_receptance, xr);943 struct ggml_tensor * k = llm_build_lora_mm(lctx, ctx, layer->time_mix_key, xk);944 struct ggml_tensor * v = llm_build_lora_mm(lctx, ctx, layer->time_mix_value, xv);945 if (layer->time_mix_receptance_b) {946 r = ggml_add(ctx, r, layer->time_mix_receptance_b);947 }948 if (layer->time_mix_key_b) {949 k = ggml_add(ctx, k, layer->time_mix_key_b);950 }951 if (layer->time_mix_value_b) {952 v = ggml_add(ctx, v, layer->time_mix_value_b);953 }954 955 struct ggml_tensor * g = llm_build_lora_mm(lctx, ctx, layer->time_mix_gate, xg);956 if (is_qrwkv) {957 g = ggml_sigmoid(ctx, g);958 } else {959 g = ggml_silu(ctx, g);960 }961 962 if (head_count_kv != head_count) {963 GGML_ASSERT(head_count % head_count_kv == 0);964 k = ggml_reshape_4d(ctx, k, head_size, 1, head_count_kv, n_tokens);965 v = ggml_reshape_4d(ctx, v, head_size, 1, head_count_kv, n_tokens);966 struct ggml_tensor * tmp = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, head_size, head_count / head_count_kv, head_count_kv, n_tokens);967 k = ggml_repeat(ctx, k, tmp);968 v = ggml_repeat(ctx, v, tmp);969 }970 971 k = ggml_reshape_3d(ctx, k, head_size, head_count, n_tokens);972 v = ggml_reshape_3d(ctx, v, head_size, head_count, n_tokens);973 r = ggml_reshape_3d(ctx, r, head_size, head_count, n_tokens);974 975 struct ggml_tensor * w = ggml_mul_mat(976 ctx,977 layer->time_mix_decay_w2,978 ggml_tanh(979 ctx,980 ggml_mul_mat(ctx, layer->time_mix_decay_w1, xw)981 )982 );983 984 w = ggml_add(ctx, w, layer->time_mix_decay);985 w = ggml_exp(ctx, ggml_neg(ctx, ggml_exp(ctx, w)));986 w = ggml_reshape_3d(ctx, w, head_size, head_count, n_tokens);987 988 if (is_qrwkv) {989 // k = k * (1 - w)990 k = ggml_sub(ctx, k, ggml_mul(ctx, k, w));991 }992 993 struct ggml_tensor * wkv_output;994 if (!layer->time_mix_first) {995 wkv_output = ggml_gated_linear_attn(ctx, k, v, r, w, *wkv_state, pow(head_size, -0.5f));996 } else {997 wkv_output = ggml_rwkv_wkv6(ctx, k, v, r, layer->time_mix_first, w, *wkv_state);998 }999 cur = ggml_view_1d(ctx, wkv_output, n_embd * n_tokens, 0);1000 *wkv_state = ggml_view_1d(ctx, wkv_output, n_embd * head_size * n_seqs, n_embd * n_tokens * sizeof(float));1001 1002 if (!is_qrwkv) {1003 // group norm with head_count groups1004 cur = ggml_reshape_3d(ctx, cur, n_embd / head_count, head_count, n_tokens);1005 cur = ggml_norm(ctx, cur, 64e-5f);1006 1007 // Convert back to regular vectors.1008 cur = ggml_reshape_2d(ctx, cur, n_embd, n_tokens);1009 cur = ggml_add(ctx, ggml_mul(ctx, cur, layer->time_mix_ln), layer->time_mix_ln_b);1010 } else {1011 cur = ggml_reshape_2d(ctx, cur, n_embd, n_tokens);1012 }1013 1014 cur = ggml_mul(ctx, cur, g);1015 cur = llm_build_lora_mm(lctx, ctx, layer->time_mix_output, cur);1016 1017 return ggml_reshape_3d(ctx, cur, n_embd, n_seq_tokens, n_seqs);1018}1019 1020static struct ggml_tensor * llm_build_rwkv6_channel_mix(1021 struct llama_context & lctx,1022 struct ggml_context * ctx,1023 const struct llama_layer * layer,1024 struct ggml_tensor * cur,1025 struct ggml_tensor * x_prev) {1026 struct ggml_tensor * sx = ggml_sub(ctx, x_prev, cur);1027 struct ggml_tensor * xk = ggml_add(ctx, ggml_mul(ctx, sx, layer->channel_mix_lerp_k), cur);1028 struct ggml_tensor * xr = ggml_add(ctx, ggml_mul(ctx, sx, layer->channel_mix_lerp_r), cur);1029 1030 struct ggml_tensor * r = ggml_sigmoid(ctx, llm_build_lora_mm(lctx, ctx, layer->channel_mix_receptance, xr));1031 struct ggml_tensor * k = ggml_sqr(1032 ctx,1033 ggml_relu(1034 ctx,1035 llm_build_lora_mm(lctx, ctx, layer->channel_mix_key, xk)1036 )1037 );1038 1039 return ggml_mul(ctx, r, llm_build_lora_mm(lctx, ctx, layer->channel_mix_value, k));1040}1041 1042struct llm_build_context {1043 const llama_model & model;1044 llama_context & lctx;1045 const llama_hparams & hparams;1046 const llama_cparams & cparams;1047 const llama_ubatch & ubatch;1048 const llama_kv_cache & kv_self;1049 1050 const int64_t n_embd;1051 const int64_t n_layer;1052 const int64_t n_rot;1053 const int64_t n_ctx; // user-specified context size (can be different from n_ctx_train)1054 const int64_t n_head;1055 const int64_t n_head_kv;1056 const int64_t n_embd_head_k;1057 const int64_t n_embd_k_gqa;1058 const int64_t n_embd_head_v;1059 const int64_t n_embd_v_gqa;1060 const int64_t n_expert;1061 const int64_t n_expert_used;1062 1063 const float freq_base;1064 const float freq_scale;1065 const float ext_factor;1066 const float attn_factor;1067 const float beta_fast;1068 const float beta_slow;1069 const float norm_eps;1070 const float norm_rms_eps;1071 1072 const int32_t n_tokens;1073 const int32_t n_kv; // size of KV cache to consider (n_kv <= kv_self.size)1074 const int32_t n_outputs;1075 const int32_t n_outputs_enc;1076 const int32_t kv_head; // index of where we store new KV data in the cache1077 const int32_t n_ctx_orig;1078 1079 const bool flash_attn;1080 1081 const enum llama_pooling_type pooling_type;1082 const enum llama_rope_type rope_type;1083 1084 const llm_build_cb & cb;1085 1086 std::vector<uint8_t> & buf_compute_meta;1087 1088 struct ggml_context * ctx0 = nullptr;1089 1090 // TODO: consider making the entire interface noexcept1091 llm_build_context(1092 llama_context & lctx,1093 const llama_ubatch & ubatch,1094 const llm_build_cb & cb,1095 bool worst_case) :1096 model (lctx.model),1097 lctx (lctx),1098 hparams (model.hparams),1099 cparams (lctx.cparams),1100 ubatch (ubatch),1101 kv_self (lctx.kv_self),1102 n_embd (hparams.n_embd),1103 n_layer (hparams.n_layer),1104 n_rot (hparams.n_rot),1105 n_ctx (cparams.n_ctx),1106 n_head (hparams.n_head()),1107 n_head_kv (hparams.n_head_kv()),1108 n_embd_head_k (hparams.n_embd_head_k),1109 n_embd_k_gqa (hparams.n_embd_k_gqa()),1110 n_embd_head_v (hparams.n_embd_head_v),1111 n_embd_v_gqa (hparams.n_embd_v_gqa()),1112 n_expert (hparams.n_expert),1113 n_expert_used (hparams.n_expert_used),1114 freq_base (cparams.rope_freq_base),1115 freq_scale (cparams.rope_freq_scale),1116 ext_factor (cparams.yarn_ext_factor),1117 attn_factor (cparams.yarn_attn_factor),1118 beta_fast (cparams.yarn_beta_fast),1119 beta_slow (cparams.yarn_beta_slow),1120 norm_eps (hparams.f_norm_eps),1121 norm_rms_eps (hparams.f_norm_rms_eps),1122 n_tokens (ubatch.n_tokens),1123 n_kv (worst_case ? kv_self.size : kv_self.n),1124 n_outputs (worst_case ? n_tokens : lctx.n_outputs),1125 n_outputs_enc (worst_case ? n_tokens : lctx.embd_enc.size() / hparams.n_embd),1126 kv_head (worst_case ? (kv_self.recurrent ? 0 : kv_self.size - n_tokens) : kv_self.head),1127 n_ctx_orig (cparams.n_ctx_orig_yarn),1128 flash_attn (cparams.flash_attn),1129 pooling_type (cparams.pooling_type),1130 rope_type (hparams.rope_type),1131 cb (cb),1132 buf_compute_meta (lctx.buf_compute_meta) {1133 // all initializations should be done in init()1134 }1135 1136 void init() {1137 struct ggml_init_params params = {1138 /*.mem_size =*/ buf_compute_meta.size(),1139 /*.mem_buffer =*/ buf_compute_meta.data(),1140 /*.no_alloc =*/ true,1141 };1142 1143 ctx0 = ggml_init(params);1144 1145 lctx.inp_tokens = nullptr;1146 lctx.inp_embd = nullptr;1147 lctx.inp_pos = nullptr;1148 lctx.inp_out_ids = nullptr;1149 lctx.inp_KQ_mask = nullptr;1150 lctx.inp_KQ_mask_swa = nullptr;1151 lctx.inp_K_shift = nullptr;1152 lctx.inp_mean = nullptr;1153 lctx.inp_cls = nullptr;1154 lctx.inp_s_copy = nullptr;1155 lctx.inp_s_mask = nullptr;1156 lctx.inp_s_seq = nullptr;1157 lctx.inp_pos_bucket = nullptr;1158 lctx.inp_embd_enc = nullptr;1159 lctx.inp_KQ_mask_cross = nullptr;1160 }1161 1162 void free() {1163 ggml_free(ctx0);1164 ctx0 = nullptr;1165 }1166 1167 struct ggml_cgraph * build_k_shift() {1168 struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);1169 1170 GGML_ASSERT(kv_self.size == n_ctx);1171 1172 lctx.inp_K_shift = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_ctx);1173 cb(lctx.inp_K_shift, "K_shift", -1);1174 ggml_set_input(lctx.inp_K_shift);1175 1176 for (int il = 0; il < n_layer; ++il) {1177 const int64_t n_head_kv = hparams.n_head_kv(il);1178 const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);1179 struct ggml_tensor * rope_factors = build_rope_factors(il);1180 struct ggml_tensor * k =1181 ggml_view_3d(ctx0, kv_self.k_l[il],1182 n_embd_head_k, n_head_kv, n_ctx,1183 ggml_row_size(kv_self.k_l[il]->type, n_embd_head_k),1184 ggml_row_size(kv_self.k_l[il]->type, n_embd_k_gqa),1185 0);1186 1187 struct ggml_tensor * tmp;1188 if (ggml_is_quantized(k->type)) {1189 // dequantize to f32 -> RoPE -> quantize back1190 tmp = ggml_cast(ctx0, k, GGML_TYPE_F32);1191 cb(tmp, "K_f32", il);1192 for (auto & backend : lctx.backends) {1193 // Figure out which backend KV cache belongs to1194 if (ggml_backend_supports_buft(backend.get(), ggml_backend_buffer_get_type(kv_self.k_l[il]->buffer))) {1195 ggml_backend_sched_set_tensor_backend(lctx.sched.get(), tmp, backend.get());1196 break;1197 }1198 }1199 tmp = ggml_rope_ext_inplace(ctx0, tmp,1200 lctx.inp_K_shift, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,