echodict/llama.cpp
version https://git-lfs.github.com/spec/v1 oid sha256:cfc44b7ba25614df70e6b65e3341cae0310163bd32fd31a6b928a542df433faf size 30786
0479
1#pragma once2 3#include "llama-arch.h"4#include "llama-batch.h"5#include "llama-hparams.h"6#include "llama-adapter.h"7 8#include <cstdint>9#include <vector>10#include <memory>11#include <set>12#include <functional>13#include <map>14 15struct ggml_cgraph;16struct ggml_context;17struct ggml_tensor;18 19struct llama_cparams;20struct llama_layer;21 22struct llama_memory_context_i;23 24class llama_kv_cache_context;25class llama_kv_cache_iswa_context;26class llama_memory_recurrent_context;27class llama_memory_hybrid_context;28class llama_memory_hybrid_iswa_context;29 30// certain models (typically multi-modal) can produce different types of graphs31enum llm_graph_type {32 LLM_GRAPH_TYPE_DEFAULT,33 LLM_GRAPH_TYPE_ENCODER,34 LLM_GRAPH_TYPE_DECODER,35};36 37enum llm_ffn_op_type {38 LLM_FFN_SILU,39 LLM_FFN_GELU,40 LLM_FFN_RELU,41 LLM_FFN_RELU_SQR,42 LLM_FFN_SWIGLU,43 LLM_FFN_GEGLU,44 LLM_FFN_REGLU,45 LLM_FFN_SWIGLU_OAI_MOE,46};47 48enum llm_ffn_gate_type {49 LLM_FFN_SEQ,50 LLM_FFN_PAR, // ffn_gate is parallel to ffn_up51};52 53enum llm_norm_type {54 LLM_NORM,55 LLM_NORM_RMS,56 LLM_NORM_GROUP,57};58 59// TODO: tmp - need something better to pass the data from the encoder to the decoder60struct llama_cross {61 // the output embeddings from the encoder as a ggml tensor62 // TODO: this needs more work to be correct, for now copy the embeddings data to host memory63 // ref: https://github.com/ggml-org/llama.cpp/pull/11213#discussion_r196989252464 //ggml_tensor * t_embd = nullptr;65 66 int64_t n_embd = 0;67 int64_t n_enc = 0;68 69 // embeddings data copied to host memory (tmp)70 std::vector<float> v_embd;71 72 // needed to construct the cross-attention mask in the decoder73 std::vector<std::set<llama_seq_id>> seq_ids_enc;74};75 76struct llm_graph_params;77 78//79// llm_graph_input80//81 82class llm_graph_input_i {83public:84 llm_graph_input_i() {85 const char * LLAMA_GRAPH_INPUT_DEBUG = getenv("LLAMA_GRAPH_INPUT_DEBUG");86 debug = LLAMA_GRAPH_INPUT_DEBUG ? atoi(LLAMA_GRAPH_INPUT_DEBUG) : 0;87 }88 89 virtual ~llm_graph_input_i() = default;90 91 virtual void set_input(const llama_ubatch * ubatch) = 0;92 93 // return true if the resulting input tensors using the provided graph parameters would be94 // the same as the previous input tensors that we have currently stored in the object95 virtual bool can_reuse(const llm_graph_params & params) {96 // returning false here by default will prevent from reusing the graph if the check97 // for the input type has not been implemented yet98 GGML_UNUSED(params);99 return false;100 }101protected:102 // env: LLAMA_GRAPH_INPUT_DEBUG103 int debug = 0;104};105 106using llm_graph_input_ptr = std::unique_ptr<llm_graph_input_i>;107 108class llm_graph_input_embd : public llm_graph_input_i {109public:110 llm_graph_input_embd(int64_t n_embd) : n_embd(n_embd) {}111 virtual ~llm_graph_input_embd() = default;112 113 void set_input(const llama_ubatch * ubatch) override;114 115 bool can_reuse(const llm_graph_params & params) override;116 117 ggml_tensor * tokens = nullptr; // I32 [n_batch]118 ggml_tensor * embd = nullptr; // F32 [n_embd, n_batch]119 120 const int64_t n_embd = 0;121};122 123class llm_graph_input_pos : public llm_graph_input_i {124public:125 llm_graph_input_pos(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {}126 virtual ~llm_graph_input_pos() = default;127 128 void set_input(const llama_ubatch * ubatch) override;129 130 bool can_reuse(const llm_graph_params & params) override;131 132 ggml_tensor * pos = nullptr; // I32 [n_batch]133 134 const uint32_t n_pos_per_embd = 1;135};136 137// temperature tuning, used by llama4138class llm_graph_input_attn_temp : public llm_graph_input_i {139public:140 llm_graph_input_attn_temp(uint32_t n_attn_temp_floor_scale, float f_attn_temp_scale, float f_attn_temp_offset)141 : n_attn_temp_floor_scale(n_attn_temp_floor_scale), f_attn_temp_scale(f_attn_temp_scale), f_attn_temp_offset(f_attn_temp_offset) {}142 virtual ~llm_graph_input_attn_temp() = default;143 144 void set_input(const llama_ubatch * ubatch) override;145 146 ggml_tensor * attn_scale = nullptr; // F32 [n_batch]147 148 const uint32_t n_attn_temp_floor_scale;149 const float f_attn_temp_scale;150 const float f_attn_temp_offset;151};152 153class llm_graph_input_pos_bucket : public llm_graph_input_i {154public:155 llm_graph_input_pos_bucket(const llama_hparams & hparams) : hparams(hparams) {}156 virtual ~llm_graph_input_pos_bucket() = default;157 158 void set_input(const llama_ubatch * ubatch) override;159 160 ggml_tensor * pos_bucket = nullptr; // I32 [n_batch, n_batch]161 162 const llama_hparams hparams;163};164 165class llm_graph_input_pos_bucket_kv : public llm_graph_input_i {166public:167 llm_graph_input_pos_bucket_kv(168 const llama_hparams & hparams,169 const llama_kv_cache_context * mctx) : hparams(hparams), mctx(mctx) {}170 virtual ~llm_graph_input_pos_bucket_kv() = default;171 172 void set_input(const llama_ubatch * ubatch) override;173 174 ggml_tensor * pos_bucket = nullptr; // I32 [n_kv, n_batch]175 176 const llama_hparams hparams;177 178 const llama_kv_cache_context * mctx;179};180 181class llm_graph_input_out_ids : public llm_graph_input_i {182public:183 llm_graph_input_out_ids(184 const llama_hparams & hparams,185 const llama_cparams & cparams,186 uint32_t n_outputs) : hparams(hparams), cparams(cparams), n_outputs(n_outputs) {}187 virtual ~llm_graph_input_out_ids() = default;188 189 void set_input(const llama_ubatch * ubatch) override;190 191 bool can_reuse(const llm_graph_params & params) override;192 193 ggml_tensor * out_ids; // I32 [n_outputs]194 195 const llama_hparams hparams;196 const llama_cparams cparams;197 198 const uint32_t n_outputs;199};200 201class llm_graph_input_mean : public llm_graph_input_i {202public:203 llm_graph_input_mean(const llama_cparams & cparams) : cparams(cparams) {}204 virtual ~llm_graph_input_mean() = default;205 206 void set_input(const llama_ubatch * ubatch) override;207 208 ggml_tensor * mean; // F32 [n_batch, n_batch]209 210 const llama_cparams cparams;211};212 213class llm_graph_input_cls : public llm_graph_input_i {214public:215 llm_graph_input_cls(const llama_cparams & cparams, const llm_arch arch) : cparams(cparams), arch(arch) {}216 virtual ~llm_graph_input_cls() = default;217 218 void set_input(const llama_ubatch * ubatch) override;219 220 ggml_tensor * cls; // I32 [n_batch]221 222 const llama_cparams cparams;223 const llm_arch arch;224};225 226class llm_graph_input_rs : public llm_graph_input_i {227public:228 llm_graph_input_rs(const llama_memory_recurrent_context * mctx) : mctx(mctx) {}229 virtual ~llm_graph_input_rs() = default;230 231 void set_input(const llama_ubatch * ubatch) override;232 233 bool can_reuse(const llm_graph_params & params) override;234 235 ggml_tensor * s_copy; // I32 [n_rs]236 237 // views of s_copy, computed once per graph238 // and shared across layers which use build_rs239 ggml_tensor * s_copy_main; // I32 [n_seqs]240 ggml_tensor * s_copy_extra; // I32 [n_rs - n_seqs]241 242 const llama_memory_recurrent_context * mctx;243 244 // used in view offsets, need to match for valid graph reuse245 uint32_t head;246 int32_t rs_z;247};248 249class llm_graph_input_cross_embd : public llm_graph_input_i {250public:251 llm_graph_input_cross_embd(252 const llama_cross * cross) : cross(cross) {}253 virtual ~llm_graph_input_cross_embd() = default;254 255 void set_input(const llama_ubatch * ubatch) override;256 257 ggml_tensor * cross_embd; // F32 [n_embd, n_outputs_enc]258 259 const llama_cross * cross;260};261 262class llm_graph_input_attn_no_cache : public llm_graph_input_i {263public:264 llm_graph_input_attn_no_cache(const llama_hparams & hparams, const llama_cparams & cparams) :265 hparams(hparams),266 cparams(cparams) {267 }268 ~llm_graph_input_attn_no_cache() = default;269 270 void set_input(const llama_ubatch * ubatch) override;271 272 ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }273 ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }274 275 // n_tokens == n_batch276 ggml_tensor * self_kq_mask = nullptr; // F32 [n_tokens, n_batch/n_stream, 1, n_stream]277 ggml_tensor * self_kq_mask_cnv = nullptr; // [n_tokens, n_batch/n_stream, 1, n_stream]278 ggml_tensor * self_kq_mask_swa = nullptr; // F32 [n_tokens, n_batch/n_stream, 1, n_stream]279 ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_tokens, n_batch/n_stream, 1, n_stream]280 281 const llama_hparams hparams;282 const llama_cparams cparams;283};284 285class llm_graph_input_attn_kv : public llm_graph_input_i {286public:287 llm_graph_input_attn_kv(288 const llama_hparams & hparams,289 const llama_cparams & cparams,290 const llama_kv_cache_context * mctx) :291 hparams(hparams),292 cparams(cparams),293 mctx(mctx) {294 }295 ~llm_graph_input_attn_kv() = default;296 297 void set_input(const llama_ubatch * ubatch) override;298 299 bool can_reuse(const llm_graph_params & params) override;300 301 ggml_tensor * get_k_idxs() const { return self_k_idxs; }302 ggml_tensor * get_v_idxs() const { return self_v_idxs; }303 304 ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }305 306 ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]307 ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]308 309 ggml_tensor * self_kq_mask = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]310 ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]311 312 // note: assumes v_rot^2 == I313 ggml_tensor * self_k_rot = nullptr;314 ggml_tensor * self_v_rot = nullptr;315 316 // note: these have to be copies because in order to be able to reuse a graph, its inputs317 // need to carry these parameters with them. otherwise, they can point to freed318 // llm_graph_params from a previous batch, causing stack-use-after-return319 const llama_hparams hparams;320 const llama_cparams cparams;321 322 const llama_kv_cache_context * mctx;323};324 325// V-less input for the KV cache326// ref: https://github.com/ggml-org/llama.cpp/pull/19067327class llm_graph_input_attn_k : public llm_graph_input_i {328public:329 llm_graph_input_attn_k(330 const llama_hparams & hparams,331 const llama_cparams & cparams,332 const llama_kv_cache_context * mctx) :333 hparams(hparams),334 cparams(cparams),335 mctx(mctx) {336 }337 ~llm_graph_input_attn_k() = default;338 339 void set_input(const llama_ubatch * ubatch) override;340 341 bool can_reuse(const llm_graph_params & params) override;342 343 ggml_tensor * get_k_idxs() const { return self_k_idxs; }344 345 ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }346 347 ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]348 349 ggml_tensor * self_kq_mask = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]350 ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]351 352 const llama_hparams hparams;353 const llama_cparams cparams;354 355 const llama_kv_cache_context * mctx;356};357 358class llm_graph_input_attn_kv_iswa : public llm_graph_input_i {359public:360 llm_graph_input_attn_kv_iswa(361 const llama_hparams & hparams,362 const llama_cparams & cparams,363 const llama_kv_cache_iswa_context * mctx) :364 hparams(hparams),365 cparams(cparams),366 mctx(mctx) {367 }368 ~llm_graph_input_attn_kv_iswa() = default;369 370 void set_input(const llama_ubatch * ubatch) override;371 372 bool can_reuse(const llm_graph_params & params) override;373 374 ggml_tensor * get_k_idxs() const { return self_k_idxs; }375 ggml_tensor * get_v_idxs() const { return self_v_idxs; }376 ggml_tensor * get_k_idxs_swa() const { return self_k_idxs_swa; }377 ggml_tensor * get_v_idxs_swa() const { return self_v_idxs_swa; }378 379 ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }380 ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }381 382 ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]383 ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]384 ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch]385 ggml_tensor * self_v_idxs_swa = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]386 387 ggml_tensor * self_kq_mask = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]388 ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]389 ggml_tensor * self_kq_mask_swa = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]390 ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]391 392 ggml_tensor * self_k_rot = nullptr;393 ggml_tensor * self_v_rot = nullptr;394 395 ggml_tensor * self_k_rot_swa = nullptr;396 ggml_tensor * self_v_rot_swa = nullptr;397 398 const llama_hparams hparams;399 const llama_cparams cparams;400 401 const llama_kv_cache_iswa_context * mctx;402};403 404class llm_graph_input_attn_cross : public llm_graph_input_i {405public:406 llm_graph_input_attn_cross(const llama_cross * cross) : cross(cross) {}407 ~llm_graph_input_attn_cross() = default;408 409 void set_input(const llama_ubatch * ubatch) override;410 411 ggml_tensor * get_kq_mask_cross() const { return cross_kq_mask_cnv; }412 413 ggml_tensor * cross_kq_mask = nullptr; // F32 [n_outputs_enc, n_batch, 1, 1]414 ggml_tensor * cross_kq_mask_cnv = nullptr; // F32 [n_outputs_enc, n_batch, 1, 1]415 416 const llama_cross * cross = nullptr;417};418 419class llm_graph_input_mem_hybrid : public llm_graph_input_i {420public:421 llm_graph_input_mem_hybrid(422 const llama_cparams & cparams,423 std::unique_ptr<llm_graph_input_attn_kv> inp_attn,424 std::unique_ptr<llm_graph_input_rs> inp_rs,425 const llama_memory_hybrid_context * mctx) :426 inp_attn(std::move(inp_attn)),427 inp_rs(std::move(inp_rs)),428 cparams(cparams),429 mctx(mctx) { }430 virtual ~llm_graph_input_mem_hybrid() = default;431 432 void set_input(const llama_ubatch * ubatch) override;433 434 bool can_reuse(const llm_graph_params & params) override;435 436 std::unique_ptr<llm_graph_input_attn_kv> inp_attn;437 std::unique_ptr<llm_graph_input_rs> inp_rs;438 439 llm_graph_input_attn_kv * get_attn() const { return inp_attn.get(); }440 llm_graph_input_rs * get_recr() const { return inp_rs.get(); }441 442 const llama_cparams cparams;443 444 const llama_memory_hybrid_context * mctx;445};446 447class llm_graph_input_mem_hybrid_k : public llm_graph_input_i {448public:449 llm_graph_input_mem_hybrid_k(450 const llama_cparams & cparams,451 std::unique_ptr<llm_graph_input_attn_k> inp_attn,452 std::unique_ptr<llm_graph_input_rs> inp_rs,453 const llama_memory_hybrid_context * mctx) :454 inp_attn(std::move(inp_attn)),455 inp_rs(std::move(inp_rs)),456 cparams(cparams),457 mctx(mctx) { }458 virtual ~llm_graph_input_mem_hybrid_k() = default;459 460 void set_input(const llama_ubatch * ubatch) override;461 462 bool can_reuse(const llm_graph_params & params) override;463 464 std::unique_ptr<llm_graph_input_attn_k> inp_attn;465 std::unique_ptr<llm_graph_input_rs> inp_rs;466 467 llm_graph_input_attn_k * get_attn() const { return inp_attn.get(); }468 llm_graph_input_rs * get_recr() const { return inp_rs.get(); }469 470 const llama_cparams cparams;471 472 const llama_memory_hybrid_context * mctx;473};474 475class llm_graph_input_mem_hybrid_iswa : public llm_graph_input_i {476public:477 llm_graph_input_mem_hybrid_iswa(478 const llama_cparams & cparams,479 std::unique_ptr<llm_graph_input_attn_kv_iswa> inp_attn,480 std::unique_ptr<llm_graph_input_rs> inp_rs,481 const llama_memory_hybrid_iswa_context * mctx) :482 inp_attn(std::move(inp_attn)),483 inp_rs(std::move(inp_rs)),484 cparams(cparams),485 mctx(mctx) { }486 virtual ~llm_graph_input_mem_hybrid_iswa() = default;487 488 void set_input(const llama_ubatch * ubatch) override;489 490 bool can_reuse(const llm_graph_params & params) override;491 492 std::unique_ptr<llm_graph_input_attn_kv_iswa> inp_attn;493 std::unique_ptr<llm_graph_input_rs> inp_rs;494 495 llm_graph_input_attn_kv_iswa * get_attn() const { return inp_attn.get(); }496 llm_graph_input_rs * get_recr() const { return inp_rs.get(); }497 498 const llama_cparams cparams;499 500 const llama_memory_hybrid_iswa_context * mctx;501};502 503class llm_graph_input_sampling : public llm_graph_input_i {504public:505 llm_graph_input_sampling(std::map<llama_seq_id, llama_sampler *> samplers) :506 samplers(std::move(samplers)) { }507 virtual ~llm_graph_input_sampling() = default;508 509 void set_input(const llama_ubatch * ubatch) override;510 bool can_reuse(const llm_graph_params & params) override;511 512 std::map<llama_seq_id, llama_sampler *> samplers;513};514 515//516// llm_graph_result517//518 519// these objects deliver the result from the graph build process back to the llama_context520// note that the input tensors created for the graph are referenced here - the goal is to be able to populate their521// specific data, by calling the set_inputs() method522// along with the input tensors, the object also provides commonly used outputs tensors, such as logits, embeddings, etc.523// these are used by the llama_context to extact the relevant data, based on the compute parameters524 525// callback that allows us to apply custom logic to each tensor (e.g. ggml-alloc, offloading, etc.)526using llm_graph_cb = std::function<void(const llama_ubatch & ubatch, ggml_tensor * cur, const char * name, int il)>;527 528class llm_graph_result;529 530struct llm_graph_params {531 llm_arch arch = LLM_ARCH_UNKNOWN;532 533 llama_hparams hparams;534 llama_cparams cparams;535 536 llama_ubatch ubatch; // note: intentionally make a copy537 538 llm_graph_type gtype;539 540 ggml_backend_sched_t sched;541 ggml_backend_t backend_cpu;542 543 const llama_adapter_cvec * cvec;544 const llama_adapter_loras * loras;545 const llama_memory_context_i * mctx;546 const llama_cross * cross;547 548 std::map<llama_seq_id, llama_sampler *> samplers;549 550 static bool samplers_equal(551 const std::map<llama_seq_id, llama_sampler *> & lhs,552 const std::map<llama_seq_id, llama_sampler *> & rhs) {553 if (lhs.size() != rhs.size()) {554 return false;555 }556 for (const auto & [seq_id, sampler] : lhs) {557 auto it = rhs.find(seq_id);558 if (it == rhs.end() || it->second != sampler) {559 return false;560 }561 }562 return true;563 }564 565 uint32_t n_outputs;566 567 llm_graph_cb cb;568 569 llm_graph_result * res;570 571 // return true if the "other" params would result in a graph with the same topology as with the current params572 // having the same topology allows us to reuse the graph in some cases573 bool allow_reuse(const llm_graph_params & other) const {574 // first check the ubatch575 bool can_reuse_ubatch =576 ubatch.equal_seqs() == other.ubatch.equal_seqs() &&577 ubatch.n_tokens == other.ubatch.n_tokens &&578 ubatch.n_seq_tokens == other.ubatch.n_seq_tokens &&579 ubatch.n_seqs == other.ubatch.n_seqs &&580 ubatch.n_seqs_unq == other.ubatch.n_seqs_unq &&581 (582 (!ubatch.token && !other.ubatch.token) ||583 (!ubatch.embd && !other.ubatch.embd)584 );585 586 // when we split the batch using "equal_seqs" we have to verify that the participating sequences are the same587 // the reason is because the set of attention streams would be different for different sequences588 if (can_reuse_ubatch && ubatch.equal_seqs()) {589 if (!ubatch.data) {590 // if the old ubatch does not own it's data, then we cannot guarantee that it is still alive, and591 // therefore we cannot perform the sequence id check. normally should never happen592 can_reuse_ubatch = false;593 } else {594 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {595 can_reuse_ubatch &= ubatch.seq_id_unq[s] == other.ubatch.seq_id_unq[s];596 }597 }598 }599 600 if (!can_reuse_ubatch) {601 return false;602 }603 604 if (n_outputs != other.n_outputs) {605 return false;606 }607 608 if (!samplers_equal(samplers, other.samplers)) {609 return false;610 }611 612 if (samplers.size() > 0) {613 if (!ubatch.data || !other.ubatch.data) {614 return false;615 }616 617 // check that the outputs are the same for all samplers618 for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {619 if (ubatch.output[i] != other.ubatch.output[i] ||620 ubatch.seq_id[i][0] != other.ubatch.seq_id[i][0]) {621 return false;622 }623 }624 }625 626 return627 cparams.embeddings == other.cparams.embeddings &&628 cparams.causal_attn == other.cparams.causal_attn &&629 arch == other.arch &&630 gtype == other.gtype &&631 cvec == other.cvec &&632 loras == other.loras &&633 cross == other.cross;634 }635};636 637class llm_graph_result {638public:639 llm_graph_result(int64_t max_nodes);640 641 virtual ~llm_graph_result() = default;642 643 ggml_tensor * get_inp_tokens() const { return t_inp_tokens; }644 ggml_tensor * get_logits() const { return t_logits; }645 ggml_tensor * get_embd() const { return t_embd; }646 ggml_tensor * get_embd_pooled() const { return t_embd_pooled; }647 648 ggml_cgraph * get_gf() const { return gf; }649 ggml_context * get_ctx() const { return ctx_compute.get(); }650 651 int64_t get_max_nodes() const;652 653 void reset();654 655 void set_inputs(const llama_ubatch * ubatch);656 void set_outputs();657 658 // try to update the existing graph result using the new graph parameters in order to reuse it659 // this can only be done if we determine that the resulting graph using the new graph parameters660 // would be identical to the existing graph. in that case, we simply have to update the memory661 // contexts of the input tensors of the graph and we can reuse it for another computation662 // return true if the graph was updated and can be reused663 bool can_reuse(const llm_graph_params & params);664 665 llm_graph_input_i * add_input(llm_graph_input_ptr input);666 667 void set_params(const llm_graph_params & params);668 669 // important graph nodes670 ggml_tensor * t_inp_tokens = nullptr;671 ggml_tensor * t_inp_embd = nullptr; // [n_embd_inp, n_tokens]672 ggml_tensor * t_logits = nullptr;673 ggml_tensor * t_embd = nullptr;674 ggml_tensor * t_embd_pooled = nullptr;675 676 std::map<llama_seq_id, ggml_tensor*> t_sampled_logits;677 std::map<llama_seq_id, ggml_tensor*> t_candidates;678 std::map<llama_seq_id, ggml_tensor*> t_sampled;679 std::map<llama_seq_id, ggml_tensor*> t_sampled_probs;680 681 std::vector<llm_graph_input_ptr> inputs;682 683 ggml_context_ptr ctx_compute;684 685 // memory buffers used to evaluate the model686 std::vector<uint8_t> buf_compute_meta;687 688 ggml_cgraph * gf;689 690 int64_t max_nodes;691 692private:693 // keep a copy of the previous graph parameters694 // we will use this to determine whether the graph can be reused by comparing them with the new parameters695 // note: these are updated after constructing the new graph696 llm_graph_params params;697 698 // env: LLAMA_GRAPH_RESULT_DEBUG699 int debug = 0;700};701 702using llm_graph_result_ptr = std::unique_ptr<llm_graph_result>;703 704//705// llm_graph_context706//707 708// used in build_rs to properly order writes and avoid unnecessary copies709using llm_graph_get_rows_fn = std::function<ggml_tensor * (ggml_context *, ggml_tensor * states, ggml_tensor * ids)>;710 711struct llm_graph_qkv {712 ggml_tensor * q; // [n_embd_head, n_head, n_tokens]713 ggml_tensor * k; // [n_embd_head, n_head_kv, n_tokens]714 ggml_tensor * v; // [n_embd_head, n_head_kv, n_tokens]715};716 717struct llm_graph_context {718 const llm_arch arch;719 720 const llama_hparams & hparams;721 const llama_cparams & cparams;722 const llama_ubatch & ubatch;723 724 const int64_t n_embd;725 const int64_t n_layer;726 const int64_t n_rot;727 const int64_t n_ctx; // user-specified context size (can be different from n_ctx_train)728 const int64_t n_head;729 const int64_t n_head_kv;730 const int64_t n_embd_head_k;731 const int64_t n_embd_k_gqa;732 const int64_t n_embd_head_v;733 const int64_t n_embd_v_gqa;734 const int64_t n_expert;735 const int64_t n_expert_used;736 737 const float freq_base;738 const float freq_scale;739 const float ext_factor;740 const float attn_factor;741 const float beta_fast;742 const float beta_slow;743 const float norm_eps;744 const float norm_rms_eps;745 746 const int64_t n_tokens;747 const int64_t n_outputs;748 const int32_t n_ctx_orig; // yarn749 750 const enum llama_pooling_type pooling_type;751 const enum llama_rope_type rope_type;752 753 ggml_backend_sched_t sched;754 755 ggml_backend_t backend_cpu; // TODO: needed by build_attn_mha, figure out a way to remove?756 757 const llama_adapter_cvec * cvec;758 const llama_adapter_loras * loras;759 const llama_memory_context_i * mctx;760 const llama_cross * cross;761 762 std::map<llama_seq_id, llama_sampler *> samplers;763 764 const llm_graph_cb & cb_func;765 766 llm_graph_result * res;767 768 ggml_context * ctx0 = nullptr;769 ggml_cgraph * gf = nullptr;770 771 llm_graph_context(const llm_graph_params & params);772 virtual ~llm_graph_context() = default;773 774 void cb(ggml_tensor * cur, const char * name, int il) const;775 776 //777 // common778 //779 780 ggml_tensor * build_cvec(781 ggml_tensor * cur,782 int il) const;783 784 // do mat_mul, while optionally apply lora and per-tensor scale785 ggml_tensor * build_lora_mm(786 ggml_tensor * w,787 ggml_tensor * cur,788 ggml_tensor * w_s = nullptr) const;789 790 // do mat_mul_id, while optionally apply lora791 ggml_tensor * build_lora_mm_id(792 ggml_tensor * w, // ggml_tensor * as793 ggml_tensor * cur, // ggml_tensor * b794 ggml_tensor * ids) const;795 796 ggml_tensor * build_norm(797 ggml_tensor * cur,798 ggml_tensor * mw,799 ggml_tensor * mb,800 llm_norm_type type,801 int il) const;802 803 804 // compute Q, K, V projections with optional bias and reshape805 // supports both fused wqkv and separate wq/wk/wv paths806 llm_graph_qkv build_qkv(807 const llama_layer & layer,808 ggml_tensor * cur,809 int64_t n_embd_head,810 int64_t n_head,811 int64_t n_head_kv,812 int il) const;813 814 ggml_tensor * build_ffn(815 ggml_tensor * cur,816 ggml_tensor * up,817 ggml_tensor * up_b,818 ggml_tensor * up_s,819 ggml_tensor * gate,820 ggml_tensor * gate_b,821 ggml_tensor * gate_s,822 ggml_tensor * down,823 ggml_tensor * down_b,824 ggml_tensor * down_s,825 ggml_tensor * act_scales,826 llm_ffn_op_type type_op,827 llm_ffn_gate_type type_gate,828 int il) const;829 830 // build MoE FFN without bias tensors831 ggml_tensor * build_moe_ffn(832 ggml_tensor * cur,833 ggml_tensor * gate_inp,834 ggml_tensor * up_exps,835 ggml_tensor * gate_exps,836 ggml_tensor * down_exps,837 ggml_tensor * exp_probs_b,838 int64_t n_expert,839 int64_t n_expert_used,840 llm_ffn_op_type type_op,841 bool norm_w,842 float w_scale,843 llama_expert_gating_func_type gating_op,844 int il,845 ggml_tensor * probs_in = nullptr,846 ggml_tensor * gate_up_exps = nullptr,847 ggml_tensor * up_exps_s = nullptr,848 ggml_tensor * gate_exps_s = nullptr,849 ggml_tensor * down_exps_s = nullptr) const;850 851 ggml_tensor * build_moe_ffn(852 ggml_tensor * cur,853 ggml_tensor * gate_inp,854 ggml_tensor * gate_inp_b,855 ggml_tensor * up_exps,856 ggml_tensor * up_exps_b,857 ggml_tensor * gate_exps,858 ggml_tensor * gate_exps_b,859 ggml_tensor * down_exps,860 ggml_tensor * down_exps_b,861 ggml_tensor * exp_probs_b,862 int64_t n_expert,863 int64_t n_expert_used,864 llm_ffn_op_type type_op,865 bool norm_w,866 float w_scale,867 llama_expert_gating_func_type gating_op,868 int il,869 ggml_tensor * probs_in = nullptr,870 ggml_tensor * gate_up_exps = nullptr,871 ggml_tensor * gate_up_exps_b = nullptr,872 ggml_tensor * up_exps_s = nullptr,873 ggml_tensor * gate_exps_s = nullptr,874 ggml_tensor * down_exps_s = nullptr) const;875 876 //877 // inputs878 //879 880 ggml_tensor * build_inp_embd(ggml_tensor * tok_embd) const;881 ggml_tensor * build_inp_pos() const;882 ggml_tensor * build_inp_attn_scale() const;883 ggml_tensor * build_inp_out_ids() const;884 ggml_tensor * build_inp_mean() const;885 ggml_tensor * build_inp_cls() const;886 887 ggml_tensor * build_inp_cross_embd() const;888 ggml_tensor * build_inp_pos_bucket_enc() const;889 ggml_tensor * build_inp_pos_bucket_dec() const;890 ggml_tensor * build_pos_bias(ggml_tensor * pos_bucket, ggml_tensor * attn_rel_b) const;891 892 //893 // attention894 //895 896 ggml_tensor * build_attn_mha(897 ggml_tensor * q, // [n_embd_head_q, n_head_q, n_tokens]898 ggml_tensor * k, // [n_embd_head_k, n_head_k, n_tokens]899 ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans == false)900 ggml_tensor * kq_b,901 ggml_tensor * kq_mask,902 ggml_tensor * sinks, // [n_head_q]903 ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]904 float kq_scale,905 int il) const;906 907 llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const;908 909 ggml_tensor * build_attn(910 llm_graph_input_attn_no_cache * inp,911 ggml_tensor * wo,912 ggml_tensor * wo_b,913 ggml_tensor * wo_s,914 ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens]915 ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens]916 ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens]917 ggml_tensor * kq_b,918 ggml_tensor * sinks, // [n_head_q]919 ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]920 float kq_scale,921 int il) const;922 923 llm_graph_input_attn_kv * build_attn_inp_kv() const;924 925 ggml_tensor * build_attn(926 llm_graph_input_attn_kv * inp,927 ggml_tensor * wo,928 ggml_tensor * wo_b,929 ggml_tensor * wo_s,930 ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens]931 ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens]932 ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens]933 ggml_tensor * kq_b,934 ggml_tensor * sinks, // [n_head_q]935 ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] // TODO: remove936 float kq_scale,937 int il) const;938 939 llm_graph_input_attn_k * build_attn_inp_k() const;940 941 ggml_tensor * build_attn(942 llm_graph_input_attn_k * inp,943 ggml_tensor * wo,944 ggml_tensor * wo_b,945 ggml_tensor * wo_s,946 ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens]947 ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens]948 ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens]949 ggml_tensor * kq_b,950 ggml_tensor * sinks, // [n_head_q]951 ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]952 float kq_scale,953 int il) const;954 955 llm_graph_input_attn_kv_iswa * build_attn_inp_kv_iswa() const;956 957 // note: if k_cur or v_cur are not provided, they will not be stored in the memory958 ggml_tensor * build_attn(959 llm_graph_input_attn_kv_iswa * inp,960 ggml_tensor * wo,961 ggml_tensor * wo_b,962 ggml_tensor * wo_s,963 ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens]964 ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens] optional965 ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] optional966 ggml_tensor * kq_b,967 ggml_tensor * sinks, // [n_head_q]968 ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]969 float kq_scale,970 int il) const;971 972 llm_graph_input_attn_cross * build_attn_inp_cross() const;973 974 ggml_tensor * build_attn(975 llm_graph_input_attn_cross * inp,976 ggml_tensor * wo,977 ggml_tensor * wo_b,978 ggml_tensor * wo_s,979 ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens]980 ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens]981 ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens]982 ggml_tensor * kq_b,983 ggml_tensor * sinks, // [n_head_q]984 ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]985 float kq_scale,986 int il) const;987 988 //989 // recurrent990 //991 992 // TODO: move this implementation to llama_memory_recurrent.993 // this is analogous to llama_kv_cache::cpy_k / cpy_v994 // when moving, avoid passing `ggml_cgraph` - only pass `ggml_context`. would likely need to split the995 // implementation in 2 separate methods. the goal is to avoid calling `ggml_build_forward_expand` in996 // `llama_memory_recurrent`997 ggml_tensor * build_rs(998 ggml_tensor * s,999 ggml_tensor * state_copy_main,1000 ggml_tensor * state_copy_extra,1001 int32_t state_size,1002 int32_t n_seqs,1003 uint32_t n_rs,1004 uint32_t rs_head,1005 uint32_t rs_size,1006 int32_t rs_zero,1007 const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows) const;1008 1009 llm_graph_input_rs * build_rs_inp() const;1010 1011 ggml_tensor * build_rs(1012 llm_graph_input_rs * inp,1013 ggml_tensor * s,1014 int32_t state_size,1015 int32_t n_seqs,1016 const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows) const;1017 1018 ggml_tensor * build_rwkv_token_shift_load(1019 llm_graph_input_rs * inp,1020 const llama_ubatch & ubatch,1021 int il) const;1022 1023 ggml_tensor * build_rwkv_token_shift_store(1024 ggml_tensor * token_shift,1025 const llama_ubatch & ubatch,1026 int il) const;1027 //1028 // hybrid1029 //1030 1031 llm_graph_input_mem_hybrid * build_inp_mem_hybrid() const;1032 llm_graph_input_mem_hybrid_k * build_inp_mem_hybrid_k() const;1033 1034 llm_graph_input_mem_hybrid_iswa * build_inp_mem_hybrid_iswa() const;1035 1036 //1037 // pooling1038 //1039 1040 void build_pooling(1041 ggml_tensor * cls,1042 ggml_tensor * cls_b,1043 ggml_tensor * cls_out,1044 ggml_tensor * cls_out_b,1045 ggml_tensor * cls_norm) const;1046 1047 //1048 // sampling (backend sampling)1049 //1050 1051 void build_sampling() const;1052 1053 //1054 // dense (out)1055 //1056 1057 void build_dense_out(1058 ggml_tensor * dense_2,1059 ggml_tensor * dense_2_b,1060 ggml_tensor * dense_3) const;1061};1062 1063// TODO: better name1064int32_t llama_relative_position_bucket(llama_pos x, llama_pos y, uint64_t n_buckets, bool bidirectional);1065