Felipe97/llama-cpp-compiled
01.1k
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 <cstdlib>10#include <vector>11#include <memory>12#include <set>13#include <functional>14#include <map>15 16struct ggml_cgraph;17struct ggml_context;18struct ggml_tensor;19 20struct llama_cparams;21struct llama_layer;22 23struct llama_memory_context_i;24 25class llama_kv_cache_context;26class llama_kv_cache_dsa_context;27class llama_kv_cache_dsa_iswa_context;28class llama_kv_cache_msa_context;29class llama_kv_cache_dsv4_raw_context;30class llama_kv_cache_dsv4_context;31class llama_kv_cache_iswa_context;32class llama_memory_recurrent_context;33class llama_memory_hybrid_context;34class llama_memory_hybrid_iswa_context;35 36// certain models (typically multi-modal) can produce different types of graphs37enum llm_graph_type {38 LLM_GRAPH_TYPE_DEFAULT,39 LLM_GRAPH_TYPE_ENCODER,40 LLM_GRAPH_TYPE_DECODER,41 LLM_GRAPH_TYPE_DECODER_MTP,42};43 44enum llm_fused_op {45 LLM_FUSED_OP_FLASH_ATTN,46 LLM_FUSED_OP_GDN_AR,47 LLM_FUSED_OP_GDN_CH,48 LLM_FUSED_OP_LIGHTNING_INDEXER,49 LLM_FUSED_OP_DSV4_HC_PRE,50 LLM_FUSED_OP_DSV4_HC_COMB,51 LLM_FUSED_OP_DSV4_HC_POST,52};53 54enum llm_ffn_op_type : int {55 LLM_FFN_NONE = 0, // sentinel: unset; archs must assign before use56 LLM_FFN_SILU,57 LLM_FFN_GELU,58 LLM_FFN_RELU,59 LLM_FFN_RELU_SQR,60 LLM_FFN_SWIGLU,61 LLM_FFN_GEGLU,62 LLM_FFN_REGLU,63 LLM_FFN_SWIGLU_OAI_MOE,64 LLM_FFN_SITU, // kimi-k365};66 67enum llm_ffn_gate_type {68 LLM_FFN_SEQ,69 LLM_FFN_PAR, // ffn_gate is parallel to ffn_up70};71 72enum llm_norm_type {73 LLM_NORM,74 LLM_NORM_RMS,75 LLM_NORM_GROUP,76};77 78// TODO: tmp - need something better to pass the data from the encoder to the decoder79struct llama_cross {80 // the output embeddings from the encoder as a ggml tensor81 // TODO: this needs more work to be correct, for now copy the embeddings data to host memory82 // ref: https://github.com/ggml-org/llama.cpp/pull/11213#discussion_r196989252483 //ggml_tensor * t_embd = nullptr;84 85 int64_t n_embd = 0;86 int64_t n_enc = 0;87 88 // embeddings data copied to host memory (tmp)89 std::vector<float> v_embd;90 91 // needed to construct the cross-attention mask in the decoder92 std::vector<std::set<llama_seq_id>> seq_ids_enc;93};94 95struct llm_graph_params;96 97//98// llm_graph_input99//100 101class llm_graph_input_i {102public:103 llm_graph_input_i() {104 const char * LLAMA_GRAPH_INPUT_DEBUG = getenv("LLAMA_GRAPH_INPUT_DEBUG");105 debug = LLAMA_GRAPH_INPUT_DEBUG ? atoi(LLAMA_GRAPH_INPUT_DEBUG) : 0;106 }107 108 virtual ~llm_graph_input_i() = default;109 110 virtual void set_input(const llama_ubatch * ubatch) = 0;111 112 // return true if the resulting input tensors using the provided graph parameters would be113 // the same as the previous input tensors that we have currently stored in the object114 virtual bool can_reuse(const llm_graph_params & params) {115 // returning false here by default will prevent from reusing the graph if the check116 // for the input type has not been implemented yet117 GGML_UNUSED(params);118 return false;119 }120protected:121 // env: LLAMA_GRAPH_INPUT_DEBUG122 int debug = 0;123};124 125using llm_graph_input_ptr = std::unique_ptr<llm_graph_input_i>;126 127class llm_graph_input_embd : public llm_graph_input_i {128public:129 llm_graph_input_embd(int64_t n_embd) : n_embd(n_embd) {}130 virtual ~llm_graph_input_embd() = default;131 132 void set_input(const llama_ubatch * ubatch) override;133 134 bool can_reuse(const llm_graph_params & params) override;135 136 ggml_tensor * tokens = nullptr; // I32 [n_batch]137 ggml_tensor * embd = nullptr; // F32 [n_embd, n_batch]138 139 const int64_t n_embd = 0;140};141 142// similar to llm_graph_input_embd but with an additional hidden state input143class llm_graph_input_embd_h : public llm_graph_input_i {144public:145 llm_graph_input_embd_h(int64_t n_embd) : n_embd(n_embd) {}146 virtual ~llm_graph_input_embd_h() = default;147 148 void set_input(const llama_ubatch * ubatch) override;149 150 bool can_reuse(const llm_graph_params & params) override;151 152 ggml_tensor * tokens = nullptr; // I32 [n_batch]153 ggml_tensor * embd = nullptr; // F32 [n_embd, n_batch]154 ggml_tensor * h = nullptr; // F32 [n_embd, n_batch]155 156 const int64_t n_embd = 0;157};158 159class llm_graph_input_pos : public llm_graph_input_i {160public:161 llm_graph_input_pos(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {}162 virtual ~llm_graph_input_pos() = default;163 164 void set_input(const llama_ubatch * ubatch) override;165 166 bool can_reuse(const llm_graph_params & params) override;167 168 ggml_tensor * pos = nullptr; // I32 [n_batch]169 170 const uint32_t n_pos_per_embd = 1;171};172 173// temperature tuning, used by llama4174class llm_graph_input_attn_temp : public llm_graph_input_i {175public:176 llm_graph_input_attn_temp(uint32_t n_attn_temp_floor_scale, float f_attn_temp_scale, float f_attn_temp_offset)177 : 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) {}178 virtual ~llm_graph_input_attn_temp() = default;179 180 void set_input(const llama_ubatch * ubatch) override;181 182 ggml_tensor * attn_scale = nullptr; // F32 [n_batch]183 184 const uint32_t n_attn_temp_floor_scale;185 const float f_attn_temp_scale;186 const float f_attn_temp_offset;187};188 189class llm_graph_input_pos_bucket : public llm_graph_input_i {190public:191 llm_graph_input_pos_bucket(const llama_hparams & hparams) : hparams(hparams) {}192 virtual ~llm_graph_input_pos_bucket() = default;193 194 void set_input(const llama_ubatch * ubatch) override;195 196 ggml_tensor * pos_bucket = nullptr; // I32 [n_batch, n_batch]197 198 const llama_hparams hparams;199};200 201class llm_graph_input_pos_bucket_kv : public llm_graph_input_i {202public:203 llm_graph_input_pos_bucket_kv(204 const llama_hparams & hparams,205 const llama_kv_cache_context * mctx) : hparams(hparams), mctx(mctx) {}206 virtual ~llm_graph_input_pos_bucket_kv() = default;207 208 void set_input(const llama_ubatch * ubatch) override;209 210 ggml_tensor * pos_bucket = nullptr; // I32 [n_kv, n_batch]211 212 const llama_hparams hparams;213 214 const llama_kv_cache_context * mctx;215};216 217class llm_graph_input_out_ids : public llm_graph_input_i {218public:219 llm_graph_input_out_ids(220 const llama_hparams & hparams,221 const llama_cparams & cparams,222 uint32_t n_outputs) : hparams(hparams), cparams(cparams), n_outputs(n_outputs) {}223 virtual ~llm_graph_input_out_ids() = default;224 225 void set_input(const llama_ubatch * ubatch) override;226 227 bool can_reuse(const llm_graph_params & params) override;228 229 ggml_tensor * out_ids; // I32 [n_outputs]230 231 const llama_hparams hparams;232 const llama_cparams cparams;233 234 const uint32_t n_outputs;235};236 237class llm_graph_input_mean : public llm_graph_input_i {238public:239 llm_graph_input_mean(const llama_cparams & cparams) : cparams(cparams) {}240 virtual ~llm_graph_input_mean() = default;241 242 void set_input(const llama_ubatch * ubatch) override;243 244 ggml_tensor * mean; // F32 [n_batch, n_batch]245 246 const llama_cparams cparams;247};248 249class llm_graph_input_cls : public llm_graph_input_i {250public:251 llm_graph_input_cls(const llama_cparams & cparams, const llm_arch arch) : cparams(cparams), arch(arch) {}252 virtual ~llm_graph_input_cls() = default;253 254 void set_input(const llama_ubatch * ubatch) override;255 256 ggml_tensor * cls; // I32 [n_batch]257 258 const llama_cparams cparams;259 const llm_arch arch;260};261 262class llm_graph_input_rs : public llm_graph_input_i {263public:264 llm_graph_input_rs(const llama_memory_recurrent_context * mctx) : mctx(mctx) {}265 virtual ~llm_graph_input_rs() = default;266 267 void set_input(const llama_ubatch * ubatch) override;268 269 bool can_reuse(const llm_graph_params & params) override;270 271 ggml_tensor * s_copy; // I32 [n_rs]272 273 // views of s_copy, computed once per graph274 // and shared across layers which use build_rs275 ggml_tensor * s_copy_main; // I32 [n_seqs]276 ggml_tensor * s_copy_extra; // I32 [n_rs - n_seqs]277 278 const llama_memory_recurrent_context * mctx;279 280 // used in view offsets, need to match for valid graph reuse281 uint32_t head;282 int32_t rs_z;283};284 285class llm_graph_input_cross_embd : public llm_graph_input_i {286public:287 llm_graph_input_cross_embd(288 const llama_cross * cross) : cross(cross) {}289 virtual ~llm_graph_input_cross_embd() = default;290 291 void set_input(const llama_ubatch * ubatch) override;292 293 ggml_tensor * cross_embd; // F32 [n_embd, n_outputs_enc]294 295 const llama_cross * cross;296};297 298class llm_graph_input_attn_no_cache : public llm_graph_input_i {299public:300 llm_graph_input_attn_no_cache(const llama_hparams & hparams, const llama_cparams & cparams) :301 hparams(hparams),302 cparams(cparams) {303 }304 ~llm_graph_input_attn_no_cache() = default;305 306 void set_input(const llama_ubatch * ubatch) override;307 308 ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }309 ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }310 311 // n_tokens == n_batch312 ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_tokens, n_batch/n_stream, 1, n_stream]313 ggml_tensor * self_kq_mask_cnv = nullptr; // [n_tokens, n_batch/n_stream, 1, n_stream]314 ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_tokens, n_batch/n_stream, 1, n_stream]315 ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_tokens, n_batch/n_stream, 1, n_stream]316 317 const llama_hparams hparams;318 const llama_cparams cparams;319};320 321class llm_graph_input_attn_kv : public llm_graph_input_i {322public:323 llm_graph_input_attn_kv(324 const llama_hparams & hparams,325 const llama_cparams & cparams,326 const llama_kv_cache_context * mctx) :327 hparams(hparams),328 cparams(cparams),329 mctx(mctx) {330 }331 ~llm_graph_input_attn_kv() = default;332 333 void set_input(const llama_ubatch * ubatch) override;334 335 bool can_reuse(const llm_graph_params & params) override;336 337 ggml_tensor * get_k_idxs() const { return self_k_idxs; }338 ggml_tensor * get_v_idxs() const { return self_v_idxs; }339 340 ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }341 342 ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]343 ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]344 345 ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]346 ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]347 348 // note: assumes v_rot^2 == I349 ggml_tensor * self_k_rot = nullptr;350 ggml_tensor * self_v_rot = nullptr;351 352 // note: these have to be copies because in order to be able to reuse a graph, its inputs353 // need to carry these parameters with them. otherwise, they can point to freed354 // llm_graph_params from a previous batch, causing stack-use-after-return355 const llama_hparams hparams;356 const llama_cparams cparams;357 358 const llama_kv_cache_context * mctx;359};360 361// V-less input for the KV cache362// ref: https://github.com/ggml-org/llama.cpp/pull/19067363class llm_graph_input_attn_k : public llm_graph_input_i {364public:365 llm_graph_input_attn_k(366 const llama_hparams & hparams,367 const llama_cparams & cparams,368 const llama_kv_cache_context * mctx) :369 hparams(hparams),370 cparams(cparams),371 mctx(mctx) {372 }373 ~llm_graph_input_attn_k() = default;374 375 void set_input(const llama_ubatch * ubatch) override;376 377 bool can_reuse(const llm_graph_params & params) override;378 379 // like can_reuse, but does not re-bind mctx380 bool can_reuse_impl(const llm_graph_params & params);381 382 ggml_tensor * get_k_idxs() const { return self_k_idxs; }383 384 ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }385 386 ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]387 388 ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]389 ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]390 391 const llama_hparams hparams;392 const llama_cparams cparams;393 394 const llama_kv_cache_context * mctx;395};396 397class llm_graph_input_attn_k_dsa : public llm_graph_input_i {398public:399 llm_graph_input_attn_k_dsa(400 const llama_hparams & hparams,401 const llama_cparams & cparams,402 const llama_kv_cache_dsa_context * mctx) :403 hparams(hparams),404 cparams(cparams),405 mctx(mctx) {406 }407 ~llm_graph_input_attn_k_dsa() = default;408 409 void set_input(const llama_ubatch * ubatch) override;410 411 bool can_reuse(const llm_graph_params & params) override;412 413 // like can_reuse, but does not re-bind mctx414 bool can_reuse_impl(const llm_graph_params & params);415 416 ggml_tensor * get_k_idxs_mla() const { return self_k_idxs_mla; }417 ggml_tensor * get_k_idxs_lid() const { return self_k_idxs_lid; }418 419 ggml_tensor * get_kq_mask_mla() const { return self_kq_mask_mla_cnv; }420 ggml_tensor * get_kq_mask_lid() const { return self_kq_mask_lid; }421 422 ggml_tensor * self_k_idxs_mla = nullptr; // I64 [n_batch]423 ggml_tensor * self_k_idxs_lid = nullptr; // I64 [n_batch]424 425 ggml_tensor * self_kq_mask_mla = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]426 ggml_tensor * self_kq_mask_mla_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]427 ggml_tensor * self_kq_mask_lid = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]428 ggml_tensor * self_kq_mask_lid_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]429 430 ggml_tensor * self_k_rot_lid = nullptr;431 432 const llama_hparams hparams;433 const llama_cparams cparams;434 435 const llama_kv_cache_dsa_context * mctx;436};437 438// DSA input (full-attention layers + indexer) with K-only input for the SWA layers439class llm_graph_input_attn_k_dsa_iswa : public llm_graph_input_i {440public:441 llm_graph_input_attn_k_dsa_iswa(442 std::unique_ptr<llm_graph_input_attn_k_dsa> inp_dsa,443 std::unique_ptr<llm_graph_input_attn_k> inp_swa,444 const llama_kv_cache_dsa_iswa_context * mctx) :445 inp_dsa(std::move(inp_dsa)),446 inp_swa(std::move(inp_swa)),447 mctx(mctx) {448 }449 ~llm_graph_input_attn_k_dsa_iswa() = default;450 451 void set_input(const llama_ubatch * ubatch) override;452 453 bool can_reuse(const llm_graph_params & params) override;454 455 llm_graph_input_attn_k_dsa * get_dsa() const { return inp_dsa.get(); }456 llm_graph_input_attn_k * get_swa() const { return inp_swa.get(); }457 458 std::unique_ptr<llm_graph_input_attn_k_dsa> inp_dsa;459 std::unique_ptr<llm_graph_input_attn_k> inp_swa;460 461 const llama_kv_cache_dsa_iswa_context * mctx;462};463 464// standard K/V attention input against the base cache, plus destination indices for the indexer key cache465class llm_graph_input_attn_kv_msa : public llm_graph_input_attn_kv {466public:467 llm_graph_input_attn_kv_msa(468 const llama_hparams & hparams,469 const llama_cparams & cparams,470 const llama_kv_cache_msa_context * mctx);471 ~llm_graph_input_attn_kv_msa() = default;472 473 void set_input(const llama_ubatch * ubatch) override;474 475 bool can_reuse(const llm_graph_params & params) override;476 477 ggml_tensor * get_k_idxs_idx() const { return self_k_idxs_idx; }478 479 ggml_tensor * self_k_idxs_idx = nullptr; // I64 [n_batch]480 481 const llama_kv_cache_msa_context * mctx_msa;482};483 484class llm_graph_input_attn_kv_iswa : public llm_graph_input_i {485public:486 llm_graph_input_attn_kv_iswa(487 const llama_hparams & hparams,488 const llama_cparams & cparams,489 const llama_kv_cache_iswa_context * mctx) :490 hparams(hparams),491 cparams(cparams),492 mctx(mctx) {493 }494 ~llm_graph_input_attn_kv_iswa() = default;495 496 void set_input(const llama_ubatch * ubatch) override;497 498 bool can_reuse(const llm_graph_params & params) override;499 500 ggml_tensor * get_k_idxs() const { return self_k_idxs; }501 ggml_tensor * get_v_idxs() const { return self_v_idxs; }502 ggml_tensor * get_k_idxs_swa() const { return self_k_idxs_swa; }503 ggml_tensor * get_v_idxs_swa() const { return self_v_idxs_swa; }504 505 ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }506 ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }507 508 ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]509 ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]510 ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch]511 ggml_tensor * self_v_idxs_swa = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]512 513 ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]514 ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]515 ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]516 ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]517 518 ggml_tensor * self_k_rot = nullptr;519 ggml_tensor * self_v_rot = nullptr;520 521 ggml_tensor * self_k_rot_swa = nullptr;522 ggml_tensor * self_v_rot_swa = nullptr;523 524 const llama_hparams hparams;525 const llama_cparams cparams;526 527 const llama_kv_cache_iswa_context * mctx;528};529 530class llm_graph_input_attn_k_iswa : public llm_graph_input_i {531public:532 llm_graph_input_attn_k_iswa(533 const llama_hparams & hparams,534 const llama_cparams & cparams,535 const llama_kv_cache_iswa_context * mctx) :536 hparams(hparams),537 cparams(cparams),538 mctx(mctx) {539 }540 ~llm_graph_input_attn_k_iswa() = default;541 542 void set_input(const llama_ubatch * ubatch) override;543 544 bool can_reuse(const llm_graph_params & params) override;545 546 ggml_tensor * get_k_idxs() const { return self_k_idxs; }547 ggml_tensor * get_k_idxs_swa() const { return self_k_idxs_swa; }548 549 ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }550 ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }551 552 ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]553 ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch]554 555 ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]556 ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]557 ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]558 ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]559 560 ggml_tensor * self_k_rot = nullptr;561 ggml_tensor * self_k_rot_swa = nullptr;562 563 const llama_hparams hparams;564 const llama_cparams cparams;565 566 const llama_kv_cache_iswa_context * mctx;567};568 569// DSV4 raw graph inputs are SWA-only, but their mask may be stream-shaped570// so raw K can be concatenated with DSV4 compressed K in one attention op.571class llm_graph_input_dsv4_raw {572public:573 llm_graph_input_dsv4_raw(574 const llama_cparams & cparams,575 const llama_kv_cache_dsv4_raw_context * mctx) :576 cparams(cparams),577 mctx(mctx) {578 }579 580 void set_input(const llama_ubatch * ubatch);581 582 ggml_tensor * get_k_idxs() const { return self_k_idxs; }583 ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }584 585 ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]586 587 ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]588 ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]589 590 ggml_tensor * self_k_rot = nullptr;591 592 const llama_cparams cparams;593 594 const llama_kv_cache_dsv4_raw_context * mctx;595};596 597class llm_graph_input_dsv4 : public llm_graph_input_i {598public:599 struct comp_input {600 ggml_tensor * state_pos = nullptr; // I32 [n_state]601 ggml_tensor * state_persist_src_idxs = nullptr; // I32 [n_state_persist]602 ggml_tensor * state_persist_dst_idxs = nullptr; // I32 [n_state_persist]603 ggml_tensor * state_restore_src_idxs = nullptr; // I32 [n_state_restore]604 ggml_tensor * state_restore_dst_idxs = nullptr; // I32 [n_state_restore]605 ggml_tensor * state_snapshot_src_idxs = nullptr; // I32 [n_state_snapshot]606 ggml_tensor * state_snapshot_dst_idxs = nullptr; // I32 [n_state_snapshot]607 ggml_tensor * state_read_idxs = nullptr; // I32 [ratio*n_state_write]608 ggml_tensor * state_write_idxs = nullptr; // I64 [n_state_write]609 ggml_tensor * state_write_pos = nullptr; // I32 [n_state_write]610 611 ggml_tensor * kq_mask = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]612 613 ggml_tensor * k_rot = nullptr;614 };615 616 llm_graph_input_dsv4(617 const llama_cparams & cparams,618 std::unique_ptr<llm_graph_input_dsv4_raw> inp_raw,619 const llama_kv_cache_dsv4_context * mctx) :620 inp_raw(std::move(inp_raw)),621 cparams(cparams),622 mctx(mctx) {623 }624 ~llm_graph_input_dsv4() = default;625 626 void set_input(const llama_ubatch * ubatch) override;627 628 bool can_reuse(const llm_graph_params & params) override;629 630 llm_graph_input_dsv4_raw * get_raw() const { return inp_raw.get(); }631 const comp_input & get_csa() const { return inp_csa; }632 const comp_input & get_hca() const { return inp_hca; }633 const comp_input & get_lid() const { return inp_lid; }634 635 std::unique_ptr<llm_graph_input_dsv4_raw> inp_raw;636 637 comp_input inp_csa;638 comp_input inp_hca;639 comp_input inp_lid;640 641 const llama_cparams cparams;642 643 const llama_kv_cache_dsv4_context * mctx;644};645 646class llm_graph_input_attn_cross : public llm_graph_input_i {647public:648 llm_graph_input_attn_cross(const llama_cross * cross) : cross(cross) {}649 ~llm_graph_input_attn_cross() = default;650 651 void set_input(const llama_ubatch * ubatch) override;652 653 ggml_tensor * get_kq_mask_cross() const { return cross_kq_mask_cnv; }654 655 ggml_tensor * cross_kq_mask = nullptr; // F32/F16 [n_outputs_enc, n_batch, 1, 1]656 ggml_tensor * cross_kq_mask_cnv = nullptr; // F32/F16 [n_outputs_enc, n_batch, 1, 1]657 658 const llama_cross * cross = nullptr;659};660 661class llm_graph_input_mem_hybrid : public llm_graph_input_i {662public:663 llm_graph_input_mem_hybrid(664 const llama_cparams & cparams,665 std::unique_ptr<llm_graph_input_attn_kv> inp_attn,666 std::unique_ptr<llm_graph_input_rs> inp_rs,667 const llama_memory_hybrid_context * mctx) :668 inp_attn(std::move(inp_attn)),669 inp_rs(std::move(inp_rs)),670 cparams(cparams),671 mctx(mctx) { }672 virtual ~llm_graph_input_mem_hybrid() = default;673 674 void set_input(const llama_ubatch * ubatch) override;675 676 bool can_reuse(const llm_graph_params & params) override;677 678 std::unique_ptr<llm_graph_input_attn_kv> inp_attn;679 std::unique_ptr<llm_graph_input_rs> inp_rs;680 681 llm_graph_input_attn_kv * get_attn() const { return inp_attn.get(); }682 llm_graph_input_rs * get_recr() const { return inp_rs.get(); }683 684 const llama_cparams cparams;685 686 const llama_memory_hybrid_context * mctx;687};688 689class llm_graph_input_mem_hybrid_k : public llm_graph_input_i {690public:691 llm_graph_input_mem_hybrid_k(692 const llama_cparams & cparams,693 std::unique_ptr<llm_graph_input_attn_k> inp_attn,694 std::unique_ptr<llm_graph_input_rs> inp_rs,695 const llama_memory_hybrid_context * mctx) :696 inp_attn(std::move(inp_attn)),697 inp_rs(std::move(inp_rs)),698 cparams(cparams),699 mctx(mctx) { }700 virtual ~llm_graph_input_mem_hybrid_k() = default;701 702 void set_input(const llama_ubatch * ubatch) override;703 704 bool can_reuse(const llm_graph_params & params) override;705 706 std::unique_ptr<llm_graph_input_attn_k> inp_attn;707 std::unique_ptr<llm_graph_input_rs> inp_rs;708 709 llm_graph_input_attn_k * get_attn() const { return inp_attn.get(); }710 llm_graph_input_rs * get_recr() const { return inp_rs.get(); }711 712 const llama_cparams cparams;713 714 const llama_memory_hybrid_context * mctx;715};716 717class llm_graph_input_mem_hybrid_iswa : public llm_graph_input_i {718public:719 llm_graph_input_mem_hybrid_iswa(720 const llama_cparams & cparams,721 std::unique_ptr<llm_graph_input_attn_kv_iswa> inp_attn,722 std::unique_ptr<llm_graph_input_rs> inp_rs,723 const llama_memory_hybrid_iswa_context * mctx) :724 inp_attn(std::move(inp_attn)),725 inp_rs(std::move(inp_rs)),726 cparams(cparams),727 mctx(mctx) { }728 virtual ~llm_graph_input_mem_hybrid_iswa() = default;729 730 void set_input(const llama_ubatch * ubatch) override;731 732 bool can_reuse(const llm_graph_params & params) override;733 734 std::unique_ptr<llm_graph_input_attn_kv_iswa> inp_attn;735 std::unique_ptr<llm_graph_input_rs> inp_rs;736 737 llm_graph_input_attn_kv_iswa * get_attn() const { return inp_attn.get(); }738 llm_graph_input_rs * get_recr() const { return inp_rs.get(); }739 740 const llama_cparams cparams;741 742 const llama_memory_hybrid_iswa_context * mctx;743};744 745class llm_graph_input_sampling : public llm_graph_input_i {746public:747 llm_graph_input_sampling(std::map<llama_seq_id, llama_sampler *> samplers) :748 samplers(std::move(samplers)) { }749 virtual ~llm_graph_input_sampling() = default;750 751 void set_input(const llama_ubatch * ubatch) override;752 bool can_reuse(const llm_graph_params & params) override;753 754 std::map<llama_seq_id, llama_sampler *> samplers;755};756 757//758// llm_graph_result759//760 761// these objects deliver the result from the graph build process back to the llama_context762// note that the input tensors created for the graph are referenced here - the goal is to be able to populate their763// specific data, by calling the set_inputs() method764// along with the input tensors, the object also provides commonly used outputs tensors, such as logits, embeddings, etc.765// these are used by the llama_context to extact the relevant data, based on the compute parameters766 767// callback that allows us to apply custom logic to each tensor (e.g. ggml-alloc, offloading, etc.)768using llm_graph_cb = std::function<void(const llama_ubatch & ubatch, ggml_tensor * cur, const char * name, int il)>;769 770class llm_graph_result;771 772struct llm_graph_params {773 llm_arch arch = LLM_ARCH_UNKNOWN;774 775 llama_hparams hparams;776 llama_cparams cparams;777 778 llama_ubatch ubatch; // note: intentionally make a copy779 780 llm_graph_type gtype;781 782 ggml_backend_sched_t sched;783 ggml_backend_t backend_cpu;784 785 const llama_adapter_cvec * cvec;786 const llama_adapter_loras * loras;787 const llama_memory_context_i * mctx;788 const llama_cross * cross;789 790 std::map<llama_seq_id, llama_sampler *> samplers;791 792 static bool samplers_equal(793 const std::map<llama_seq_id, llama_sampler *> & lhs,794 const std::map<llama_seq_id, llama_sampler *> & rhs) {795 if (lhs.size() != rhs.size()) {796 return false;797 }798 for (const auto & [seq_id, sampler] : lhs) {799 auto it = rhs.find(seq_id);800 if (it == rhs.end() || it->second != sampler) {801 return false;802 }803 }804 return true;805 }806 807 uint32_t n_outputs;808 809 llm_graph_cb cb;810 811 llm_graph_result * res;812 813 // return true if the "other" params would result in a graph with the same topology as with the current params814 // having the same topology allows us to reuse the graph in some cases815 bool allow_reuse(const llm_graph_params & other) const {816 // first check the ubatch817 bool can_reuse_ubatch =818 ubatch.equal_seqs() == other.ubatch.equal_seqs() &&819 ubatch.n_tokens == other.ubatch.n_tokens &&820 ubatch.n_seq_tokens == other.ubatch.n_seq_tokens &&821 ubatch.n_seqs == other.ubatch.n_seqs &&822 ubatch.n_seqs_unq == other.ubatch.n_seqs_unq &&823 (824 (!ubatch.token && !other.ubatch.token) ||825 (!ubatch.embd && !other.ubatch.embd) ||826 (ubatch.token && other.ubatch.token && ubatch.embd && other.ubatch.embd)827 );828 829 // when we split the batch using "equal_seqs" we have to verify that the participating sequences are the same830 // the reason is because the set of attention streams would be different for different sequences831 if (can_reuse_ubatch && ubatch.equal_seqs()) {832 if (!ubatch.data) {833 // if the old ubatch does not own it's data, then we cannot guarantee that it is still alive, and834 // therefore we cannot perform the sequence id check. normally should never happen835 can_reuse_ubatch = false;836 } else {837 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {838 can_reuse_ubatch &= ubatch.seq_id_unq[s] == other.ubatch.seq_id_unq[s];839 }840 }841 }842 843 if (!can_reuse_ubatch) {844 return false;845 }846 847 if (n_outputs != other.n_outputs) {848 return false;849 }850 851 if (!samplers_equal(samplers, other.samplers)) {852 return false;853 }854 855 if (samplers.size() > 0) {856 if (!ubatch.data || !other.ubatch.data) {857 return false;858 }859 860 // check that the outputs are the same for all samplers861 for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {862 if (ubatch.output[i] != other.ubatch.output[i] ||863 ubatch.seq_id[i][0] != other.ubatch.seq_id[i][0]) {864 return false;865 }866 }867 }868 869 // TODO: https://github.com/ggml-org/llama.cpp/pull/24340#discussion_r3448035248870 if (cparams.nextn_layer_offset != other.cparams.nextn_layer_offset) {871 return false;872 }873 874 return875 cparams.embeddings == other.cparams.embeddings &&876 cparams.embeddings_nextn == other.cparams.embeddings_nextn &&877 cparams.embeddings_nextn_masked == other.cparams.embeddings_nextn_masked &&878 cparams.causal_attn == other.cparams.causal_attn &&879 arch == other.arch &&880 gtype == other.gtype &&881 cvec == other.cvec &&882 loras == other.loras &&883 cross == other.cross;884 }885};886 887struct llm_graph_fused_node {888 llm_fused_op op;889 ggml_tensor * tensor;890 int il;891};892 893class llm_graph_result {894public:895 llm_graph_result(int64_t max_nodes);896 897 virtual ~llm_graph_result() = default;898 899 ggml_tensor * get_inp_tokens() const { return t_inp_tokens; }900 ggml_tensor * get_logits() const { return t_logits; }901 ggml_tensor * get_embd() const { return t_embd; }902 ggml_tensor * get_embd_pooled() const { return t_embd_pooled; }903 ggml_tensor * get_h_nextn() const { return t_h_nextn; }904 905 ggml_tensor * get_layer_inp(int il) const { return t_layer_inp[il]; }906 907 ggml_cgraph * get_gf() const { return gf; }908 ggml_context * get_ctx() const { return ctx_compute.get(); }909 910 int64_t get_max_nodes() const;911 912 void reset();913 914 void set_inputs(const llama_ubatch * ubatch);915 void set_outputs(const llm_graph_params & params);916 917 // try to update the existing graph result using the new graph parameters in order to reuse it918 // this can only be done if we determine that the resulting graph using the new graph parameters919 // would be identical to the existing graph. in that case, we simply have to update the memory920 // contexts of the input tensors of the graph and we can reuse it for another computation921 // return true if the graph was updated and can be reused922 bool can_reuse(const llm_graph_params & params);923 924 llm_graph_input_i * add_input(llm_graph_input_ptr input);925 926 void add_fused_node(llm_graph_fused_node result);927 928 const std::vector<llm_graph_fused_node> & get_fused_nodes() const { return fused_nodes; }929 930 void set_params(const llm_graph_params & params);931 932 // important graph nodes933 ggml_tensor * t_inp_tokens = nullptr;934 ggml_tensor * t_inp_embd = nullptr; // [n_embd_inp, n_tokens]935 ggml_tensor * t_logits = nullptr;936 ggml_tensor * t_embd = nullptr;937 ggml_tensor * t_embd_pooled = nullptr;938 ggml_tensor * t_h_nextn = nullptr; // [n_embd, n_outputs] hidden state before final output norm939 940 std::vector<ggml_tensor *> t_layer_inp;941 942 std::vector<ggml_tensor *> t_sampled;943 std::vector<ggml_tensor *> t_sampled_probs;944 std::vector<ggml_tensor *> t_sampled_logits;945 std::vector<ggml_tensor *> t_candidates;946 947 std::vector<llm_graph_input_ptr> inputs;948 std::vector<llm_graph_fused_node> fused_nodes;949 950 ggml_context_ptr ctx_compute;951 952 // memory buffers used to evaluate the model953 std::vector<uint8_t> buf_compute_meta;954 955 ggml_cgraph * gf;956 957 int64_t max_nodes;958 959private:960 // keep a copy of the previous graph parameters961 // we will use this to determine whether the graph can be reused by comparing them with the new parameters962 // note: these are updated after constructing the new graph963 llm_graph_params params;964 965 // env: LLAMA_GRAPH_RESULT_DEBUG966 int debug = 0;967};968 969using llm_graph_result_ptr = std::unique_ptr<llm_graph_result>;970 971//972// llm_graph_context973//974 975// used in build_rs to properly order writes and avoid unnecessary copies976using llm_graph_get_rows_fn = std::function<ggml_tensor * (ggml_context *, ggml_tensor * states, ggml_tensor * ids)>;977 978struct llm_graph_qkv {979 ggml_tensor * q; // [n_embd_head, n_head, n_tokens]980 ggml_tensor * k; // [n_embd_head, n_head_kv, n_tokens]981 ggml_tensor * v; // [n_embd_head, n_head_kv, n_tokens]982};983 984struct llm_graph_context {985 const llm_arch arch;986 987 const llama_hparams & hparams;988 const llama_cparams & cparams;989 const llama_ubatch & ubatch;990 991 const int64_t n_embd;992 const int64_t n_layer;993 const int64_t n_layer_nextn;994 const int64_t n_rot;995 const int64_t n_ctx; // user-specified context size (can be different from n_ctx_train)996 const int64_t n_head;997 const int64_t n_head_kv;998 const int64_t n_embd_head_k;999 const int64_t n_embd_k_gqa;1000 const int64_t n_embd_head_v;1001 const int64_t n_embd_v_gqa;1002 const int64_t n_expert;1003 const int64_t n_expert_used;1004 1005 const float freq_base;1006 const float freq_scale;1007 const float ext_factor;1008 const float attn_factor;1009 const float beta_fast;1010 const float beta_slow;1011 const float norm_eps;1012 const float norm_rms_eps;1013 1014 const int64_t n_tokens;1015 const int64_t n_outputs;1016 const int32_t n_ctx_orig; // yarn1017 1018 const enum llama_pooling_type pooling_type;1019 const enum llama_rope_type rope_type;1020 1021 ggml_backend_sched_t sched;1022 1023 ggml_backend_t backend_cpu; // TODO: needed by build_attn_mha, figure out a way to remove?1024 1025 const llama_adapter_cvec * cvec;1026 const llama_adapter_loras * loras;1027 const llama_memory_context_i * mctx;1028 const llama_cross * cross;1029 1030 std::map<llama_seq_id, llama_sampler *> samplers;1031 1032 const llm_graph_cb & cb_func;1033 1034 llm_graph_result * res;1035 1036 ggml_context * ctx0 = nullptr;1037 ggml_cgraph * gf = nullptr;1038 1039 llm_graph_context(const llm_graph_params & params);1040 virtual ~llm_graph_context() = default;1041 1042 void cb(ggml_tensor * cur, const char * name, int il) const;1043 1044 //1045 // common1046 //1047 1048 ggml_tensor * build_cvec(1049 ggml_tensor * cur,1050 int il) const;1051 1052 // do mat_mul, while optionally apply lora and per-tensor scale1053 ggml_tensor * build_lora_mm(1054 ggml_tensor * w,1055 ggml_tensor * cur,1056 ggml_tensor * w_s = nullptr) const;1057 1058 // do mat_mul_id, while optionally apply lora and per-expert scale1059 ggml_tensor * build_lora_mm_id(1060 ggml_tensor * w, // ggml_tensor * as1061 ggml_tensor * cur, // ggml_tensor * b1062 ggml_tensor * ids,1063 ggml_tensor * w_s = nullptr) const;1064 1065 ggml_tensor * build_norm(1066 ggml_tensor * cur,1067 ggml_tensor * mw,1068 ggml_tensor * mb,1069 llm_norm_type type,1070 int il) const;1071 1072 1073 // compute Q, K, V projections with optional bias and reshape1074 // supports both fused wqkv and separate wq/wk/wv paths1075 llm_graph_qkv build_qkv(1076 const llama_layer & layer,1077 ggml_tensor * cur,1078 int64_t n_embd_head,1079 int64_t n_head,1080 int64_t n_head_kv,1081 int il) const;1082 1083 // Set reshape to false to return contiguous projections before clamp/reshape.1084 llm_graph_qkv build_qkv(1085 const llama_layer & layer,1086 ggml_tensor * cur,1087 int64_t n_embd_head_q,1088 int64_t n_head_q,1089 int64_t n_embd_head_k,1090 int64_t n_head_k,1091 int64_t n_embd_head_v,1092 int64_t n_head_v,1093 int il,1094 bool reshape = true) const;1095 1096 ggml_tensor * build_ffn(1097 ggml_tensor * cur,1098 ggml_tensor * up,1099 ggml_tensor * up_b,1100 ggml_tensor * up_s,1101 ggml_tensor * gate,1102 ggml_tensor * gate_b,1103 ggml_tensor * gate_s,1104 ggml_tensor * down,1105 ggml_tensor * down_b,1106 ggml_tensor * down_s,1107 ggml_tensor * act_scales,1108 llm_ffn_op_type type_op,1109 llm_ffn_gate_type type_gate,1110 int il) const;1111 1112 // build MoE FFN without bias tensors1113 ggml_tensor * build_moe_ffn(1114 ggml_tensor * cur,1115 ggml_tensor * gate_inp,1116 ggml_tensor * up_exps,1117 ggml_tensor * gate_exps,1118 ggml_tensor * down_exps,1119 ggml_tensor * exp_probs_b,1120 int64_t n_expert,1121 int64_t n_expert_used,1122 llm_ffn_op_type type_op,1123 bool norm_w,1124 float w_scale,1125 llama_expert_gating_func_type gating_op,1126 int il,1127 ggml_tensor * probs_in = nullptr,1128 ggml_tensor * gate_up_exps = nullptr,1129 ggml_tensor * up_exps_s = nullptr,1130 ggml_tensor * gate_exps_s = nullptr,1131 ggml_tensor * down_exps_s = nullptr,1132 ggml_tensor * selected_experts_in = nullptr) const;1133 1134 ggml_tensor * build_moe_ffn(1135 ggml_tensor * cur,1136 ggml_tensor * gate_inp,1137 ggml_tensor * gate_inp_b,1138 ggml_tensor * up_exps,1139 ggml_tensor * up_exps_b,1140 ggml_tensor * gate_exps,1141 ggml_tensor * gate_exps_b,1142 ggml_tensor * down_exps,1143 ggml_tensor * down_exps_b,1144 ggml_tensor * exp_probs_b,1145 int64_t n_expert,1146 int64_t n_expert_used,1147 llm_ffn_op_type type_op,1148 bool norm_w,1149 float w_scale,1150 llama_expert_gating_func_type gating_op,1151 int il,1152 ggml_tensor * probs_in = nullptr,1153 ggml_tensor * gate_up_exps = nullptr,1154 ggml_tensor * gate_up_exps_b = nullptr,1155 ggml_tensor * up_exps_s = nullptr,1156 ggml_tensor * gate_exps_s = nullptr,1157 ggml_tensor * down_exps_s = nullptr,1158 ggml_tensor * selected_experts_in = nullptr) const;1159 1160 //1161 // inputs1162 //1163 1164 ggml_tensor * build_inp_embd(ggml_tensor * tok_embd) const;1165 ggml_tensor * build_inp_pos() const;1166 ggml_tensor * build_inp_attn_scale() const;1167 ggml_tensor * build_inp_out_ids() const;1168 ggml_tensor * build_inp_mean() const;1169 ggml_tensor * build_inp_cls() const;1170 1171 ggml_tensor * build_inp_cross_embd() const;1172 ggml_tensor * build_inp_pos_bucket_enc() const;1173 ggml_tensor * build_inp_pos_bucket_dec() const;1174 ggml_tensor * build_pos_bias(ggml_tensor * pos_bucket, ggml_tensor * attn_rel_b) const;1175 1176 //1177 // attention1178 //1179 1180 ggml_tensor * build_attn_mha(1181 ggml_tensor * q, // [n_embd_head_q, n_head_q, n_tokens]1182 ggml_tensor * k, // [n_embd_head_k, n_head_k, n_tokens]1183 ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans = false)1184 ggml_tensor * kq_b,1185 ggml_tensor * kq_mask,1186 ggml_tensor * sinks, // [n_head_q]1187 ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]1188 int64_t n_kv_max,1189 float kq_scale,1190 int il) const;1191 1192 llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const;1193 1194 ggml_tensor * build_attn(1195 llm_graph_input_attn_no_cache * inp,1196 ggml_tensor * wo,1197 ggml_tensor * wo_b,1198 ggml_tensor * wo_s,1199 ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens]1200 ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens]