Felipe97/llama-cpp-compiled
01.1k
1#pragma once2 3#include "llama.h"4#include "llama-ext.h"5#include "llama-cparams.h"6#include "llama-graph.h"7#include "llama-adapter.h"8#include "llama-impl.h"9#include "llama-memory.h"10 11#include "ggml-cpp.h"12#include "ggml-opt.h"13 14#include <array>15#include <map>16#include <vector>17 18struct llama_model;19class llama_batch_allocr;20 21class llama_io_read_i;22class llama_io_write_i;23 24// "memory" as in abstract memory for the context25struct llama_memory_i;26struct llama_memory_context_i;27 28// stores copy of the memory in device buffer. used for fast state save/load29struct llama_memory_buffer {30 int n_tensors = 0;31 size_t total_size = 0;32 33 ggml_backend_buffer_ptr buf;34 35 ggml_context_ptr ctx;36 37 std::vector<ggml_tensor *> org;38 std::vector<ggml_tensor *> cpy;39};40 41using llama_memory_buffers = std::map<ggml_backend_buffer_type_t, llama_memory_buffer>;42 43struct llama_context {44 // init scheduler and compute buffers, reserve worst-case graphs45 llama_context(46 const llama_model & model,47 llama_context_params params);48 49 ~llama_context();50 51 // reserve a new backend scheduler (if needed)52 // for example, when:53 // - changing loras54 // - changing samplers55 // - changing attention type56 // - etc.57 void sched_reserve();58 59 void synchronize();60 61 const llama_model & get_model() const;62 const llama_cparams & get_cparams() const;63 64 ggml_backend_sched_t get_sched() const;65 66 uint32_t n_ctx() const;67 uint32_t n_ctx_seq() const;68 uint32_t n_batch() const;69 uint32_t n_ubatch() const;70 uint32_t n_seq_max() const;71 72 uint32_t n_threads() const;73 uint32_t n_threads_batch() const;74 75 llama_memory_t get_memory() const;76 77 // return true if the memory was updated78 bool memory_update(bool optimize);79 80 enum llama_pooling_type pooling_type() const;81 82 float * get_logits();83 float * get_logits_ith(int32_t i);84 85 float * get_embeddings();86 float * get_embeddings_ith(int32_t i);87 float * get_embeddings_seq(llama_seq_id seq_id);88 89 float * get_embeddings_nextn();90 float * get_embeddings_nextn_ith(int32_t i);91 92 float * get_embeddings_layer_inp(uint32_t lid);93 94 llama_token * get_sampled_tokens() const;95 llama_token get_sampled_token_ith(int32_t idx);96 97 float * get_sampled_logits_ith(int32_t idx);98 size_t get_sampled_logits_count(int32_t idx);99 100 float * get_sampled_probs_ith(int32_t idx);101 size_t get_sampled_probs_count(int32_t idx);102 103 const llama_token * get_sampled_candidates_ith(int32_t idx);104 size_t get_sampled_candidates_count(int32_t idx);105 106 void attach_threadpool(107 ggml_threadpool_t threadpool,108 ggml_threadpool_t threadpool_batch);109 110 void detach_threadpool();111 112 void set_n_threads(int32_t n_threads, int32_t n_threads_batch);113 114 void set_abort_callback(bool (*abort_callback)(void * data), void * abort_callback_data);115 116 void set_embeddings (bool value);117 void set_embeddings_nextn(bool value, bool masked);118 void set_embeddings_layer_inp(uint32_t lid, bool enable);119 void set_nextn_layer_offset(int32_t offset);120 void set_causal_attn(bool value);121 void set_warmup(bool value);122 123 void set_adapters_lora(llama_adapter_lora ** adapters, size_t n_adapters, float * scales);124 125 bool adapters_lora_are_same(llama_adapter_lora ** adapters, size_t n_adapters, float * scales);126 127 bool set_adapter_cvec(128 const float * data,129 size_t len,130 int32_t n_embd,131 int32_t il_start,132 int32_t il_end);133 134 // process a single ubatch with a specific graph type135 // if memory_context is provided, it will be applied first to the context's memory136 // ret contains the status of the graph computation137 // returns nullptr only if ret != GGML_STATUS_SUCCESS138 llm_graph_result * process_ubatch(139 const llama_ubatch & ubatch,140 llm_graph_type gtype,141 llama_memory_context_i * mctx,142 ggml_status & ret);143 144 int encode(const llama_batch & batch_inp);145 int decode(const llama_batch & batch_inp);146 147 //148 // state save/load149 //150 151 size_t state_get_size();152 size_t state_get_data( uint8_t * dst, size_t size);153 size_t state_set_data(const uint8_t * src, size_t size);154 155 size_t state_seq_get_size(llama_seq_id seq_id, llama_state_seq_flags flags);156 157 size_t state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, size_t size, llama_state_seq_flags flags);158 size_t state_seq_set_data(llama_seq_id seq_id, const uint8_t * src, size_t size, llama_state_seq_flags flags);159 160 bool state_load_file(161 const char * filepath,162 llama_token * tokens_out,163 size_t n_token_capacity,164 size_t * n_token_count_out);165 166 bool state_save_file(167 const char * filepath,168 const llama_token * tokens,169 size_t n_token_count);170 171 size_t state_seq_load_file(172 llama_seq_id seq_id,173 const char * filepath,174 llama_token * tokens_out,175 size_t n_token_capacity,176 size_t * n_token_count_out);177 178 size_t state_seq_save_file(179 llama_seq_id seq_id,180 const char * filepath,181 const llama_token * tokens,182 size_t n_token_count);183 184 //185 // perf186 //187 188 llama_perf_context_data perf_get_data() const;189 void perf_reset();190 191 llama_memory_breakdown memory_breakdown() const;192 193 //194 // training195 //196 197 void opt_init(struct llama_model * model, struct llama_opt_params lopt_params);198 199 // TODO: more flexible combinations of logical/physical batch size and context size200 void opt_epoch(201 ggml_opt_dataset_t dataset,202 ggml_opt_result_t result_train,203 ggml_opt_result_t result_eval,204 int64_t idata_split,205 ggml_opt_epoch_callback callback_train,206 ggml_opt_epoch_callback callback_eval);207 208 void opt_epoch_iter(209 ggml_opt_dataset_t dataset,210 ggml_opt_result_t result,211 const std::vector<llama_token> & tokens,212 const std::vector<llama_token> & labels_sparse,213 llama_batch & batch,214 ggml_opt_epoch_callback callback,215 bool train,216 int64_t idata_in_loop,217 int64_t ndata_in_loop,218 int64_t t_loop_start);219 220private:221 //222 // output223 //224 225 // Make sure enough space is available for outputs.226 // Returns max number of outputs for which space was reserved.227 uint32_t output_reserve(int32_t n_outputs);228 229 void output_reorder();230 231 // map the output row index `i` to batch index232 int64_t output_resolve_row(int32_t i) const;233 234 // async-copy enabled layer-input tensors (per cparams.output_layer_inp)235 // from backend into host-side embd_layer_inp buffers236 void extract_layer_inputs(const llm_graph_result * res, size_t token_offset, size_t n_tokens);237 238 //239 // graph240 //241 242public:243 uint32_t graph_max_nodes(uint32_t n_tokens) const;244 245 // can reuse the llm_graph_result instance of the context (for example to update a memory module)246 llm_graph_result * get_gf_res_reserve() const;247 248 // returns the result of ggml_backend_sched_graph_compute_async execution249 ggml_status graph_compute(ggml_cgraph * gf, bool batched);250 251 // reserve a graph with a dummy ubatch of the specified size252 ggml_cgraph * graph_reserve(253 uint32_t n_tokens, uint32_t n_seqs, uint32_t n_outputs, const llama_memory_context_i * mctx, bool split_only = false, size_t * sizes = nullptr);254 255 bool set_sampler(llama_seq_id seq_id, llama_sampler * sampler);256 257private:258 llm_graph_result * get_gf_res_prev();259 260 llm_graph_params graph_params(261 llm_graph_result * res,262 const llama_ubatch & ubatch,263 const llama_memory_context_i * mctx,264 llm_graph_type gtype) const;265 266 llm_graph_cb graph_get_cb() const;267 268 // disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device269 // that differs from the layer it belongs to (usually due to missing backend support)270 void resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs);271 272 // TODO: read/write lora adapters and cvec273 size_t state_write_data(llama_io_write_i & io);274 size_t state_read_data (llama_io_read_i & io);275 276 size_t state_seq_write_data(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags);277 size_t state_seq_read_data (llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags);278 279 //280 // members281 //282 283 const llama_model & model;284 285 llama_cparams cparams;286 287 llama_adapter_cvec_ptr cvec;288 llama_adapter_loras_ptr loras;289 290 llama_cross cross; // TODO: tmp for handling cross-attention - need something better probably291 292 llama_memory_ptr memory;293 294 // decode output (2-dimensional array: [n_outputs][n_vocab])295 buffer_view<float> logits = {nullptr, 0};296 297 // embeddings output (2-dimensional array: [n_outputs][n_embd])298 // populated only when pooling_type == LLAMA_POOLING_TYPE_NONE299 buffer_view<float> embd = {nullptr, 0};300 301 // hidden state required by the nextn layers (2-dimensional array: [n_outputs][n_embd])302 // populated only when cparams.embeddings_nextn is enabled and the model graph303 // sets llm_graph_result::t_h_nextn304 buffer_view<float> embd_nextn = {nullptr, 0};305 306 // host buffers for output layer input embeddings, per layer307 // populated when cparams.output_layer_inp[il] is true308 std::vector<buffer_view<float>> embd_layer_inp;309 310 struct sampling_info {311 // !samplers.empty() to check if any samplers are active312 std::map<llama_seq_id, llama_sampler *> samplers;313 314 buffer_view<float> logits = {nullptr, 0};315 buffer_view<llama_token> sampled = {nullptr, 0};316 buffer_view<float> probs = {nullptr, 0};317 buffer_view<llama_token> candidates = {nullptr, 0};318 319 std::vector<uint32_t> logits_count;320 std::vector<uint32_t> probs_count;321 std::vector<uint32_t> candidates_count;322 323 // optimization324 std::vector<llama_token> token_ids_full_vocab;325 };326 327 sampling_info sampling;328 329 // sequence embeddings output (map of [n_embd] vectors)330 // populated only when pooling_type != LLAMA_POOLING_TYPE_NONE331 std::map<llama_seq_id, std::vector<float>> embd_seq;332 333 // reuse the batch_allocr to avoid unnecessary memory allocations334 std::unique_ptr<llama_batch_allocr> balloc;335 336 uint32_t n_outputs = 0; // number of actually-used outputs in the current ubatch or last logical batch337 338 std::vector<int32_t> output_ids; // map batch token positions to ids of the logits and embd buffers339 340 struct swap_info {341 uint32_t i0;342 uint32_t i1;343 };344 345 std::vector<swap_info> output_swaps;346 347 ggml_backend_sched_ptr sched;348 349 bool sched_need_reserve = true;350 351 ggml_backend_t backend_cpu = nullptr;352 std::vector<ggml_backend_ptr> backends;353 354 // training355 ggml_opt_context_t opt_ctx = nullptr;356 357 ggml_threadpool_t threadpool = nullptr;358 ggml_threadpool_t threadpool_batch = nullptr;359 360 ggml_abort_callback abort_callback = nullptr;361 void * abort_callback_data = nullptr;362 363 std::vector<std::pair<ggml_backend_t, ggml_backend_set_n_threads_t>> set_n_threads_fns;364 365 // pointers and buffer types used for the compute buffer of each backend366 std::vector<ggml_backend_t> backend_ptrs;367 std::vector<ggml_backend_buffer_type_t> backend_buft;368 std::vector<size_t> backend_buf_exp_size; // expected buffer sizes369 370 // Separate arenas give batches with and without outputs distinct CUDA graph cache keys.371 std::array<llm_graph_result_ptr, 2> gf_res_prev;372 llm_graph_result_ptr gf_res_reserve;373 374 llm_graph_result * gf_res_prev_active = nullptr;375 376 // host buffer for the model output (logits and embeddings)377 ggml_backend_buffer_ptr buf_output;378 379 // keep copies of the per-sequence memory on the device380 std::map<llama_seq_id, llama_memory_buffers> mem_storage;381 382 bool has_evaluated_once = false;383 384 // env: LLAMA_GRAPH_REUSE_DISABLE385 bool graph_reuse_disable = false;386 387 // perf388 mutable int64_t t_start_us = 0;389 mutable int64_t t_load_us = 0;390 mutable int64_t t_p_eval_us = 0;391 mutable int64_t t_eval_us = 0;392 393 mutable int64_t t_compute_start_us = 0;394 mutable int64_t n_queued_tokens = 0;395 396 mutable int32_t n_p_eval = 0; // number of tokens in eval calls for the prompt (with batch size > 1)397 mutable int32_t n_eval = 0; // number of eval calls398 399 mutable int32_t n_reused = 0; // number of times the previous graph was reused400};401