context : initial abstraction

ggml-ci
This commit is contained in:
Georgi Gerganov
2025-02-11 11:25:18 +02:00
parent 2cd8a903c8
commit 02ef4be975
3 changed files with 1471 additions and 1128 deletions

File diff suppressed because it is too large Load Diff

View File

@ -16,38 +16,245 @@
using llama_loras = std::unordered_map<struct llama_adapter_lora *, float>;
struct llama_batch_manager_i;
// TODO: make implementation details private
// TODO: become abstract base class, split the current implementation into different child classes
struct llama_context {
// TODO: tmp until llama-model starts implementing the graph build function
typedef std::function<ggml_cgraph *(llama_context &, const llama_ubatch &, bool worst_case)> build_graph_callback;
llama_context(const llama_model & model);
virtual ~llama_context();
llama_context(
const llama_model & model,
const llama_context_params & params,
build_graph_callback && cb_build_graph);
virtual void synchronize();
virtual ~llama_context() = default;
virtual uint32_t n_ctx() const = 0;
virtual uint32_t n_batch() const = 0;
virtual uint32_t n_ubatch() const = 0;
virtual uint32_t n_seq_max() const = 0;
const struct llama_model & model;
virtual llama_kv_cache * get_kv_self() = 0;
virtual const llama_kv_cache * get_kv_self() const = 0;
virtual void kv_self_update() = 0;
virtual enum llama_pooling_type pooling_type() const = 0;
virtual float * get_logits() = 0;
virtual float * get_logits_ith(int32_t i) = 0;
virtual float * get_embeddings() = 0;
virtual float * get_embeddings_ith(int32_t i) = 0;
virtual float * get_embeddings_seq(llama_seq_id seq_id) = 0;
int64_t n_pos_per_token() const; // vision
virtual ggml_context_ptr init();
virtual int decode(llama_batch & inp_batch) = 0;
virtual int encode(llama_batch & inp_batch) = 0;
// graph build API (generic)
// do mat_mul, while optionally apply lora
virtual ggml_tensor * build_lora_mm(
ggml_context * ctx0,
ggml_tensor * w,
ggml_tensor * cur);
// do mat_mul_id, while optionally apply lora
virtual ggml_tensor * build_lora_mm_id(
ggml_context * ctx0,
ggml_tensor * w, // struct ggml_tensor * as
ggml_tensor * cur, // struct ggml_tensor * b
ggml_tensor * ids);
// graph build API (context-specific)
virtual ggml_tensor * build_inp_embd(
ggml_context * ctx0,
ggml_tensor * tok_embd,
const llama_ubatch & ubatch) = 0;
virtual ggml_tensor * build_inp_pos(
ggml_context * ctx0,
int32_t n_tokens) = 0;
virtual ggml_tensor * build_inp_out_ids(
ggml_context * ctx0,
int32_t n_tokens,
bool worst_case) = 0;
virtual ggml_tensor * build_inp_mean(
ggml_context * ctx0,
int32_t n_tokens) = 0;
virtual ggml_tensor * build_inp_cls(
ggml_context * ctx0,
int32_t n_tokens) = 0;
virtual void build_attn_inp(
ggml_context * ctx0,
int32_t n_tokens,
bool causal,
bool swa,
bool worst_case) = 0;
virtual void build_attn_kv_store(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * k_cur,
ggml_tensor * v_cur,
int32_t n_tokens,
int64_t il,
bool worst_case) = 0;
virtual ggml_tensor * build_attn_qkv(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * wo,
ggml_tensor * wo_b,
ggml_tensor * q_cur,
int32_t n_tokens,
float kq_scale,
int il,
bool worst_case) = 0;
virtual ggml_tensor * build_soft_max_ext(
ggml_context * ctx0,
ggml_tensor * kq,
float kq_scale) = 0;
virtual ggml_tensor * get_rope_factors(int il) = 0;
virtual void build_k_shift(
ggml_context * ctx0,
ggml_cgraph * graph) = 0;
// find holes from the beginning of the KV cache and fill them by moving data from the end of the cache
virtual void build_defrag(
ggml_context * ctx0,
ggml_cgraph * graph) = 0;
virtual ggml_tensor * build_inp_embd_enc(
ggml_context * ctx0,
int32_t n_tokens,
bool worst_case) = 0;
virtual ggml_tensor * build_inp_KQ_mask_cross(
ggml_context * ctx0,
int32_t n_tokens,
bool worst_case) = 0;
virtual ggml_tensor * build_inp_s_copy(
ggml_context * ctx0,
bool worst_case) = 0;
virtual ggml_tensor * build_inp_s_mask(
ggml_context * ctx0,
bool worst_case) = 0;
virtual ggml_tensor * build_copy_mask_state(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * s,
ggml_tensor * state_copy,
ggml_tensor * state_mask,
int32_t n_tokens,
int32_t n_state,
int32_t n_seqs,
bool worst_case) = 0;
virtual ggml_tensor * build_mamba_layer(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * cur,
ggml_tensor * state_copy,
ggml_tensor * state_mask,
const llama_ubatch & ubatch,
int il,
bool worst_case) = 0;
virtual ggml_tensor * build_rwkv_token_shift_load(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * state_copy,
ggml_tensor * state_mask,
const llama_ubatch & ubatch,
int il,
bool worst_case) = 0;
virtual ggml_tensor * build_rwkv_token_shift_store(
ggml_context * ctx0,
ggml_tensor * token_shift,
const llama_ubatch & ubatch,
int il,
bool worst_case) = 0;
virtual ggml_tensor * build_rwkv6_time_mix(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * cur,
ggml_tensor * x_prev,
ggml_tensor * state_copy,
ggml_tensor * state_mask,
const llama_ubatch & ubatch,
int il,
bool worst_case) = 0;
// state save/load
virtual size_t state_get_size() = 0;
virtual size_t state_get_data( uint8_t * dst, size_t size) = 0;
virtual size_t state_set_data(const uint8_t * src, size_t size) = 0;
virtual size_t state_seq_get_size(llama_seq_id seq_id) = 0;
virtual size_t state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, size_t size) = 0;
virtual size_t state_seq_set_data(llama_seq_id seq_id, const uint8_t * src, size_t size) = 0;
virtual bool state_load_file(
const char * filepath,
llama_token * tokens_out,
size_t n_token_capacity,
size_t * n_token_count_out) = 0;
virtual bool state_save_file(
const char * filepath,
const llama_token * tokens,
size_t n_token_count) = 0;
virtual size_t state_seq_load_file(
llama_seq_id seq_id,
const char * filepath,
llama_token * tokens_out,
size_t n_token_capacity,
size_t * n_token_count_out) = 0;
virtual size_t state_seq_save_file(
llama_seq_id seq_id,
const char * filepath,
const llama_token * tokens,
size_t n_token_count) = 0;
// members
const llama_model & model;
llama_cparams cparams;
llama_sbatch sbatch; // TODO: revisit if needed
llama_adapter_cvec cvec;
llama_loras loras;
build_graph_callback cb_build_graph;
ggml_threadpool_t threadpool = nullptr;
ggml_threadpool_t threadpool_batch = nullptr;
ggml_abort_callback abort_callback = nullptr;
void * abort_callback_data = nullptr;
std::vector<ggml_backend_ptr> backends;
std::vector<std::pair<ggml_backend_t, ggml_backend_set_n_threads_t>> set_n_threads_fns;
ggml_backend_t backend_cpu = nullptr;
ggml_threadpool_t threadpool = nullptr;
ggml_threadpool_t threadpool_batch = nullptr;
ggml_backend_sched_ptr sched;
// memory buffers used to evaluate the model
std::vector<uint8_t> buf_compute_meta;
// perf
bool has_evaluated_once = false;
mutable int64_t t_start_us;
@ -60,6 +267,49 @@ struct llama_context {
mutable int32_t n_p_eval = 0; // number of tokens in eval calls for the prompt (with batch size > 1)
mutable int32_t n_eval = 0; // number of eval calls
};
// TODO: make implementation details private
struct llama_context_unified : public llama_context {
struct batch_manager;
// TODO: tmp until llama-model starts implementing the graph build function
typedef std::function<ggml_cgraph *(llama_context &, const llama_ubatch &, bool worst_case)> build_graph_callback;
llama_context_unified(
const llama_model & model,
const llama_context_params & params,
build_graph_callback && cb_build_graph);
virtual ~llama_context_unified();
virtual uint32_t n_ctx() const override;
virtual uint32_t n_batch() const override;
virtual uint32_t n_ubatch() const override;
virtual uint32_t n_seq_max() const override;
virtual llama_kv_cache * get_kv_self() override;
virtual const llama_kv_cache * get_kv_self() const override;
virtual void kv_self_update() override;
virtual enum llama_pooling_type pooling_type() const override;
virtual float * get_logits() override;
virtual float * get_logits_ith(int32_t i) override;
virtual float * get_embeddings() override;
virtual float * get_embeddings_ith(int32_t i) override;
virtual float * get_embeddings_seq(llama_seq_id seq_id) override;
virtual ggml_context_ptr init() override;
virtual int decode(llama_batch & inp_batch) override;
virtual int encode(llama_batch & inp_batch) override;
llama_sbatch sbatch;
build_graph_callback cb_build_graph;
// host buffer for the model output (logits and embeddings)
ggml_backend_buffer_ptr buf_output;
@ -72,7 +322,7 @@ struct llama_context {
size_t output_size = 0; // capacity (of tokens positions) for the output buffers
int32_t n_outputs = 0; // number of actually-used outputs in the current ubatch or last logical batch
bool logits_all = false;
bool logits_all = false;
bool need_reserve = false;
// embeddings output (2-dimensional array: [n_outputs][n_embd])
@ -84,17 +334,7 @@ struct llama_context {
// populated only when pooling_type != LLAMA_POOLING_TYPE_NONE
std::map<llama_seq_id, std::vector<float>> embd_seq;
// memory buffers used to evaluate the model
std::vector<uint8_t> buf_compute_meta;
ggml_backend_sched_ptr sched;
ggml_abort_callback abort_callback = nullptr;
void * abort_callback_data = nullptr;
virtual std::unique_ptr<llama_batch_manager_i> prepare_batch(const llama_batch & batch);
virtual int decode(llama_batch & inp_batch);
virtual int encode(llama_batch & inp_batch);
virtual std::unique_ptr<batch_manager> prepare_batch(const llama_batch & batch);
// returns the result of ggml_backend_sched_graph_compute_async execution
enum ggml_status compute_graph(
@ -107,32 +347,19 @@ struct llama_context {
// certain implementations could require a padding for the context size
uint32_t get_ctx_padding(const llama_cparams & cparams) const;
void reset();
void prepare_k_shift();
void prepare_defrag();
void set_inputs(const llama_ubatch & ubatch);
// make the outputs have the same order they had in the user-provided batch
// TODO: maybe deprecate this
// TODO: maybe remove this
void reorder_outputs();
// Make sure enough space is available for outputs.
// Returns max number of outputs for which space was reserved.
size_t reserve_outputs(size_t n_outputs);
ggml_tensor * build_lora_mm(
ggml_context * ctx0,
ggml_tensor * w,
ggml_tensor * cur);
ggml_tensor * build_lora_mm_id(
ggml_context * ctx0,
ggml_tensor * w, // struct ggml_tensor * as
ggml_tensor * cur, // struct ggml_tensor * b
ggml_tensor * ids);
// input tensors
struct ggml_tensor * inp_tokens; // I32 [n_batch]
struct ggml_tensor * inp_embd; // F32 [n_embd, n_batch]
@ -141,6 +368,81 @@ struct llama_context {
struct ggml_tensor * inp_mean; // F32 [n_batch, n_batch]
struct ggml_tensor * inp_cls; // I32 [n_batch]
// === unified KV cache ===
llama_kv_cache kv_self;
struct ggml_tensor * inp_KQ_mask; // F32 [kv_size, n_batch]
struct ggml_tensor * inp_KQ_mask_cnv; // [kv_size, n_batch]
struct ggml_tensor * inp_KQ_mask_swa; // F32 [kv_size, n_batch]
struct ggml_tensor * inp_KQ_mask_swa_cnv; // [kv_size, n_batch]
struct ggml_tensor * inp_K_shift; // I32 [kv_size]
virtual ggml_tensor * build_inp_embd(
ggml_context * ctx0,
ggml_tensor * tok_embd,
const llama_ubatch & ubatch) override;
virtual ggml_tensor * build_inp_pos(
ggml_context * ctx0,
int32_t n_tokens) override;
virtual ggml_tensor * build_inp_out_ids(
ggml_context * ctx0,
int32_t n_tokens,
bool worst_case) override;
virtual ggml_tensor * build_inp_mean(
ggml_context * ctx0,
int32_t n_tokens) override;
virtual ggml_tensor * build_inp_cls(
ggml_context * ctx0,
int32_t n_tokens) override;
virtual void build_attn_inp(
ggml_context * ctx0,
int32_t n_tokens,
bool causal,
bool swa,
bool worst_case) override;
virtual void build_attn_kv_store(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * k_cur,
ggml_tensor * v_cur,
int32_t n_tokens,
int64_t il,
bool worst_case) override;
virtual ggml_tensor * build_attn_qkv(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * wo,
ggml_tensor * wo_b,
ggml_tensor * q_cur,
int32_t n_tokens,
float kq_scale,
int il,
bool worst_case) override;
virtual ggml_tensor * build_soft_max_ext(
ggml_context * ctx0,
ggml_tensor * kq,
float kq_scale) override;
virtual ggml_tensor * get_rope_factors(int il) override;
virtual void build_k_shift(
ggml_context * ctx0,
ggml_cgraph * graph) override;
// find holes from the beginning of the KV cache and fill them by moving data from the end of the cache
virtual void build_defrag(
ggml_context * ctx0,
ggml_cgraph * graph) override;
// === encoder-decoder ===
// whether we are computing encoder output or decoder output
@ -152,79 +454,36 @@ struct llama_context {
struct ggml_tensor * inp_embd_enc; // F32 [n_embd, n_outputs_enc]
struct ggml_tensor * inp_pos_bucket; // I32 [n_batch|n_kv, n_batch]
struct ggml_tensor * inp_KQ_mask_cross; // F32 [n_outputs_enc, n_batch]
// === unified KV cache ===
llama_kv_cache kv_self;
struct ggml_tensor * inp_KQ_mask; // F32 [kv_size, n_batch]
struct ggml_tensor * inp_KQ_mask_cnv; // [kv_size, n_batch]
struct ggml_tensor * inp_KQ_mask_swa; // F32 [kv_size, n_batch]
struct ggml_tensor * inp_KQ_mask_swa_cnv; // [kv_size, n_batch]
struct ggml_tensor * inp_KQ_mask_cross; // F32 [n_outputs_enc, n_batch]
struct ggml_tensor * inp_K_shift; // I32 [kv_size]
// return true if need to reserve new worst-case graph
void kv_self_update();
void build_attn_inp(
virtual ggml_tensor * build_inp_embd_enc(
ggml_context * ctx0,
int32_t n_tokens,
bool causal,
bool swa,
bool worst_case);
bool worst_case) override;
void build_attn_kv_store(
virtual ggml_tensor * build_inp_KQ_mask_cross(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * k_cur,
ggml_tensor * v_cur,
int32_t n_tokens,
int64_t il,
bool worst_case);
ggml_tensor * build_attn_qkv(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * wo,
ggml_tensor * wo_b,
ggml_tensor * q_cur,
int32_t n_tokens,
float kq_scale,
int il,
bool worst_case);
ggml_tensor * build_soft_max_ext(
ggml_context * ctx0,
ggml_tensor * kq,
float kq_scale);
ggml_tensor * get_rope_factors(int il);
void build_k_shift(
ggml_context * ctx0,
ggml_cgraph * graph);
// find holes from the beginning of the KV cache and fill them by moving data from the end of the cache
void build_defrag(
ggml_context * ctx0,
ggml_cgraph * graph);
bool worst_case) override;
// === recurrent ===
struct ggml_tensor * inp_s_copy; // I32 [kv_size]
struct ggml_tensor * inp_s_mask; // F32 [1, n_kv]
// TODO: add recurrent cache
// TODO: add mamba-specific llama_context
// TODO: change these to build_mamba_inp and hide `state_copy` and `state_mask` inside the llama_context impl
ggml_tensor * build_inp_s_copy(
virtual ggml_tensor * build_inp_s_copy(
ggml_context * ctx0,
bool worst_case);
bool worst_case) override;
ggml_tensor * build_inp_s_mask(
virtual ggml_tensor * build_inp_s_mask(
ggml_context * ctx0,
bool worst_case);
bool worst_case) override;
ggml_tensor * build_copy_mask_state(
virtual ggml_tensor * build_copy_mask_state(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * s,
@ -233,9 +492,9 @@ struct llama_context {
int32_t n_tokens,
int32_t n_state,
int32_t n_seqs,
bool worst_case);
bool worst_case) override;
ggml_tensor * build_mamba_layer(
virtual ggml_tensor * build_mamba_layer(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * cur,
@ -243,25 +502,25 @@ struct llama_context {
ggml_tensor * state_mask,
const llama_ubatch & ubatch,
int il,
bool worst_case);
bool worst_case) override;
ggml_tensor * build_rwkv_token_shift_load(
virtual ggml_tensor * build_rwkv_token_shift_load(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * state_copy,
ggml_tensor * state_mask,
const llama_ubatch & ubatch,
int il,
bool worst_case);
bool worst_case) override;
ggml_tensor * build_rwkv_token_shift_store(
virtual ggml_tensor * build_rwkv_token_shift_store(
ggml_context * ctx0,
ggml_tensor * token_shift,
const llama_ubatch & ubatch,
int il,
bool worst_case);
bool worst_case) override;
ggml_tensor * build_rwkv6_time_mix(
virtual ggml_tensor * build_rwkv6_time_mix(
ggml_context * ctx0,
ggml_cgraph * graph,
ggml_tensor * cur,
@ -270,17 +529,48 @@ struct llama_context {
ggml_tensor * state_mask,
const llama_ubatch & ubatch,
int il,
bool worst_case);
bool worst_case) override;
struct ggml_tensor * inp_s_copy; // I32 [kv_size]
struct ggml_tensor * inp_s_mask; // F32 [1, n_kv]
// state save/load
// === vision ===
virtual size_t state_get_size() override;
virtual size_t state_get_data( uint8_t * dst, size_t size) override;
virtual size_t state_set_data(const uint8_t * src, size_t size) override;
// TODO: find a better way to accommodate mutli-dimension position encoding methods
// number of position id each token get, 1 for each token in most cases.
// when using m-rope, it will be 3 position ids per token to representing 3 dimension coordinate.
int n_pos_per_token = 1;
virtual size_t state_seq_get_size(llama_seq_id seq_id) override;
virtual size_t state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, size_t size) override;
virtual size_t state_seq_set_data(llama_seq_id seq_id, const uint8_t * src, size_t size) override;
virtual bool state_load_file(
const char * filepath,
llama_token * tokens_out,
size_t n_token_capacity,
size_t * n_token_count_out) override;
virtual bool state_save_file(
const char * filepath,
const llama_token * tokens,
size_t n_token_count) override;
virtual size_t state_seq_load_file(
llama_seq_id seq_id,
const char * filepath,
llama_token * tokens_out,
size_t n_token_capacity,
size_t * n_token_count_out) override;
virtual size_t state_seq_save_file(
llama_seq_id seq_id,
const char * filepath,
const llama_token * tokens,
size_t n_token_count) override;
private:
size_t state_get_data(struct llama_data_write & data_ctx);
size_t state_set_data(struct llama_data_read & data_ctx);
size_t state_seq_get_data(struct llama_data_write & data_ctx, llama_seq_id seq_id);
size_t state_seq_set_data(struct llama_data_read & data_ctx, llama_seq_id seq_id);
};
// For internal test use

View File

@ -8,7 +8,6 @@
#include "llama-model.h"
#include "ggml.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml-cpp.h"
@ -86,8 +85,6 @@ struct llm_build_context {
const float norm_rms_eps;
const int32_t n_tokens;
const int32_t n_outputs;
const int32_t n_outputs_enc;
const int32_t n_ctx_orig;
const bool worst_case;
@ -98,9 +95,8 @@ struct llm_build_context {
const llm_build_cb & cb;
std::vector<uint8_t> & buf_compute_meta;
struct ggml_context * ctx0 = nullptr;
const ggml_context_ptr ctx = nullptr;
ggml_context * ctx0 = nullptr;
// TODO: consider making the entire interface noexcept
llm_build_context(
@ -136,132 +132,37 @@ struct llm_build_context {
norm_eps (hparams.f_norm_eps),
norm_rms_eps (hparams.f_norm_rms_eps),
n_tokens (ubatch.n_tokens),
n_outputs (worst_case ? n_tokens : lctx.n_outputs),
n_outputs_enc (worst_case ? n_tokens : lctx.embd_enc.size() / hparams.n_embd),
n_ctx_orig (cparams.n_ctx_orig_yarn),
worst_case (worst_case),
flash_attn (cparams.flash_attn),
pooling_type (cparams.pooling_type),
rope_type (hparams.rope_type),
cb (cb),
buf_compute_meta (lctx.buf_compute_meta) {
// all initializations should be done in init()
ctx (lctx.init()),
ctx0 (ctx.get()) {
}
void init() {
struct ggml_init_params params = {
/*.mem_size =*/ buf_compute_meta.size(),
/*.mem_buffer =*/ buf_compute_meta.data(),
/*.no_alloc =*/ true,
};
ctx0 = ggml_init(params);
lctx.reset();
}
void free() {
ggml_free(ctx0);
ctx0 = nullptr;
}
// TODO: tmp
struct ggml_tensor * build_inp_embd(struct ggml_tensor * tok_embd) {
struct ggml_tensor * inpL;
if (ubatch.token) {
lctx.inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens);
cb(lctx.inp_tokens, "inp_tokens", -1);
ggml_set_input(lctx.inp_tokens);
inpL = ggml_get_rows(ctx0, tok_embd, lctx.inp_tokens);
// apply lora for embedding tokens if needed
for (const auto & lora : loras) {
struct llama_adapter_lora_weight * lw = lora.first->get_weight(tok_embd);
if (lw == nullptr) {
continue;
}
const float adapter_scale = lora.second;
const float scale = lw->get_scale(lora.first->alpha, adapter_scale);
struct ggml_tensor * inpL_delta = ggml_scale(ctx0, ggml_mul_mat(
ctx0, lw->b, // non-transposed lora_b
ggml_get_rows(ctx0, lw->a, lctx.inp_tokens)
), scale);
inpL = ggml_add(ctx0, inpL, inpL_delta);
}
} else {
lctx.inp_embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, ubatch.n_tokens);
inpL = lctx.inp_embd;
ggml_set_input(lctx.inp_embd);
}
// For Granite architecture
if (hparams.f_embedding_scale != 0.0f) {
inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale);
}
struct ggml_tensor * inpL = lctx.build_inp_embd(ctx0, tok_embd, ubatch);
cb(inpL, "inp_embd", -1);
return inpL;
}
// do mat_mul, while optionally apply lora
// TODO: tmp
struct ggml_tensor * build_lora_mm(
struct ggml_tensor * w,
struct ggml_tensor * cur) {
struct ggml_tensor * res = ggml_mul_mat(ctx0, w, cur);
for (const auto & lora : loras) {
struct llama_adapter_lora_weight * lw = lora.first->get_weight(w);
if (lw == nullptr) {
continue;
}
const float adapter_scale = lora.second;
const float scale = lw->get_scale(lora.first->alpha, adapter_scale);
struct ggml_tensor * ab_cur = ggml_mul_mat(
ctx0, lw->b,
ggml_mul_mat(ctx0, lw->a, cur)
);
ab_cur = ggml_scale(ctx0, ab_cur, scale);
res = ggml_add(ctx0, res, ab_cur);
}
return res;
return lctx.build_lora_mm(ctx0, w, cur);
}
// do mat_mul_id, while optionally apply lora
// TODO: tmp
struct ggml_tensor * build_lora_mm_id(
struct ggml_tensor * w, // struct ggml_tensor * as
struct ggml_tensor * cur, // struct ggml_tensor * b
struct ggml_tensor * ids) {
struct ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids);
for (const auto & lora : loras) {
struct llama_adapter_lora_weight * lw = lora.first->get_weight(w);
if (lw == nullptr) {
continue;
}
const float alpha = lora.first->alpha;
const float rank = (float) lw->b->ne[0];
const float scale = alpha ? lora.second * alpha / rank : lora.second;
struct ggml_tensor * ab_cur = ggml_mul_mat_id(
ctx0, lw->b,
ggml_mul_mat_id(ctx0, lw->a, cur, ids),
ids
);
ab_cur = ggml_scale(ctx0, ab_cur, scale);
res = ggml_add(ctx0, res, ab_cur);
}
return res;
return lctx.build_lora_mm_id(ctx0, w, cur, ids);
}
struct ggml_tensor * build_norm(
@ -620,31 +521,31 @@ struct llm_build_context {
}
struct ggml_tensor * build_inp_pos() {
lctx.inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
cb(lctx.inp_pos, "inp_pos", -1);
ggml_set_input(lctx.inp_pos);
return lctx.inp_pos;
ggml_tensor * cur = lctx.build_inp_pos(ctx0, n_tokens);
cb(cur, "inp_pos", -1);
return cur;
}
struct ggml_tensor * build_inp_out_ids() {
lctx.inp_out_ids = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_outputs);
cb(lctx.inp_out_ids, "inp_out_ids", -1);
ggml_set_input(lctx.inp_out_ids);
return lctx.inp_out_ids;
ggml_tensor * cur = lctx.build_inp_out_ids(ctx0, n_tokens, worst_case);
cb(cur, "inp_out_ids", -1);
return cur;
}
struct ggml_tensor * build_inp_mean() {
lctx.inp_mean = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tokens, n_tokens);
cb(lctx.inp_mean, "inp_mean", -1);
ggml_set_input(lctx.inp_mean);
return lctx.inp_mean;
ggml_tensor * cur = lctx.build_inp_mean(ctx0, n_tokens);
cb(cur, "inp_mean", -1);
return cur;
}
struct ggml_tensor * build_inp_cls() {
lctx.inp_cls = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
cb(lctx.inp_cls, "inp_cls", -1);
ggml_set_input(lctx.inp_cls);
return lctx.inp_cls;
ggml_tensor * cur = lctx.build_inp_cls(ctx0, n_tokens);
cb(cur, "inp_cls", -1);
return cur;
}
struct ggml_cgraph * append_pooling(struct ggml_cgraph * gf) {
@ -745,26 +646,22 @@ struct llm_build_context {
//}
struct ggml_tensor * build_inp_embd_enc() {
const int64_t n_embd = hparams.n_embd;
lctx.inp_embd_enc = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_outputs_enc);
ggml_set_input(lctx.inp_embd_enc);
cb(lctx.inp_embd_enc, "embd_enc", -1);
return lctx.inp_embd_enc;
ggml_tensor * cur = lctx.build_inp_embd_enc(ctx0, n_tokens, worst_case);
cb(cur, "embd_enc", -1);
return cur;
}
struct ggml_tensor * build_inp_KQ_mask_cross() {
lctx.inp_KQ_mask_cross = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_outputs_enc, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD));
ggml_set_input(lctx.inp_KQ_mask_cross);
cb(lctx.inp_KQ_mask_cross, "KQ_mask_cross", -1);
return lctx.inp_KQ_mask_cross;
ggml_tensor * cur = lctx.build_inp_KQ_mask_cross(ctx0, n_tokens, worst_case);
cb(cur, "KQ_mask_cross", -1);
return cur;
}
struct ggml_cgraph * build_llama() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -838,7 +735,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -927,9 +823,6 @@ struct llm_build_context {
struct ggml_cgraph * build_deci() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -1014,7 +907,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -1422,9 +1314,6 @@ struct llm_build_context {
struct ggml_cgraph * build_grok() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -1498,7 +1387,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -1580,9 +1468,6 @@ struct llm_build_context {
struct ggml_cgraph * build_dbrx() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
const int64_t n_embd_gqa = hparams.n_embd_v_gqa();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
@ -1649,7 +1534,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -2716,10 +2600,7 @@ struct llm_build_context {
inpL = build_inp_embd(model.tok_embd);
// inp_pos - contains the positions
lctx.inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens * 4);
cb(lctx.inp_pos, "inp_pos", -1);
ggml_set_input(lctx.inp_pos);
struct ggml_tensor * inp_pos = lctx.inp_pos;
struct ggml_tensor * inp_pos = build_inp_pos();
lctx.build_attn_inp(ctx0, n_tokens, true, false, worst_case);
@ -2825,9 +2706,6 @@ struct llm_build_context {
struct ggml_cgraph * build_qwen2moe() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -2891,7 +2769,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -4685,9 +4562,6 @@ struct llm_build_context {
struct ggml_cgraph * build_olmo() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -4757,7 +4631,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -4808,9 +4681,6 @@ struct llm_build_context {
struct ggml_cgraph * build_olmo2() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -4880,7 +4750,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -4935,9 +4804,6 @@ struct llm_build_context {
struct ggml_cgraph * build_olmoe() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -5006,7 +4872,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -5325,9 +5190,6 @@ struct llm_build_context {
struct ggml_cgraph * build_arctic() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -5385,7 +5247,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -5458,9 +5319,6 @@ struct llm_build_context {
struct ggml_cgraph * build_deepseek() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -5535,7 +5393,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -5616,9 +5473,6 @@ struct llm_build_context {
struct ggml_cgraph * build_deepseek2() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
bool is_lite = (hparams.n_layer == 27);
// We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.
@ -5767,7 +5621,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -5996,9 +5849,6 @@ struct llm_build_context {
//struct ggml_cgraph * build_t5_enc() {
// struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// // mutable variable, needed during the last layer of the computation to skip unused tokens
// int32_t n_tokens = this->n_tokens;
// const int64_t n_embd_head = hparams.n_embd_head_v;
// const int64_t n_embd_gqa = hparams.n_embd_v_gqa();
// GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
@ -6072,7 +5922,6 @@ struct llm_build_context {
// if (il == n_layer - 1) {
// // skip computing output for unused tokens
// struct ggml_tensor * inp_out_ids = build_inp_out_ids();
// n_tokens = n_outputs;
// cur = ggml_get_rows(ctx0, cur, inp_out_ids);
// inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
// }
@ -6128,9 +5977,6 @@ struct llm_build_context {
//struct ggml_cgraph * build_t5_dec() {
// struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// // mutable variable, needed during the last layer of the computation to skip unused tokens
// int32_t n_tokens = this->n_tokens;
// const int64_t n_embd_head = hparams.n_embd_head_v;
// const int64_t n_embd_gqa = hparams.n_embd_v_gqa();
// GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
@ -6272,7 +6118,6 @@ struct llm_build_context {
// if (il == n_layer - 1) {
// // skip computing output for unused tokens
// struct ggml_tensor * inp_out_ids = build_inp_out_ids();
// n_tokens = n_outputs;
// cur = ggml_get_rows(ctx0, cur, inp_out_ids);
// inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
// inpCA = ggml_get_rows(ctx0, inpCA, inp_out_ids);
@ -6673,9 +6518,6 @@ struct llm_build_context {
struct ggml_cgraph * build_exaone() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -6748,7 +6590,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -6978,9 +6819,6 @@ struct llm_build_context {
struct ggml_cgraph * build_chameleon() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(), false);
// mutable variable, needed during the last layer of the computation to skip unused tokens
int32_t n_tokens = this->n_tokens;
const int64_t n_embd_head = hparams.n_embd_head_v;
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
GGML_ASSERT(n_embd_head == hparams.n_rot);
@ -7076,7 +6914,6 @@ struct llm_build_context {
if (il == n_layer - 1) {
// skip computing output for unused tokens
struct ggml_tensor * inp_out_ids = build_inp_out_ids();
n_tokens = n_outputs;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@ -7341,8 +7178,6 @@ static struct ggml_cgraph * llama_build_graph(
struct llm_build_context llm(lctx, ubatch, cb, worst_case);
llm.init();
switch (model.arch) {
case LLM_ARCH_LLAMA:
case LLM_ARCH_MINICPM:
@ -7403,7 +7238,6 @@ static struct ggml_cgraph * llama_build_graph(
} break;
case LLM_ARCH_QWEN2VL:
{
lctx.n_pos_per_token = 4;
result = llm.build_qwen2vl();
} break;
case LLM_ARCH_QWEN2MOE:
@ -7564,8 +7398,6 @@ static struct ggml_cgraph * llama_build_graph(
result = llm.append_pooling(result);
}
llm.free();
return result;
}
@ -7908,7 +7740,7 @@ struct llama_context * llama_init_from_model(
try {
// TODO: add logic which llama_context implementation to construct
ctx = new llama_context(*model, params,
ctx = new llama_context_unified(*model, params,
[](llama_context & lctx, const llama_ubatch & ubatch, bool worst_case) {
return llama_build_graph(lctx, ubatch, worst_case);
});